From 8e0d9d0b1ea9609ddeedaa0352e51f25b1b25025 Mon Sep 17 00:00:00 2001
From: Nicolas Mounet <nicolas.mounet@cern.ch>
Date: Wed, 28 Apr 2021 18:05:14 +0200
Subject: [PATCH] adding a test on exposure model, to check scalar vs vector

---
 cara/tests/models/test_exposure_model.py | 30 ++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/cara/tests/models/test_exposure_model.py b/cara/tests/models/test_exposure_model.py
index 1cfab871..d7223f0c 100644
--- a/cara/tests/models/test_exposure_model.py
+++ b/cara/tests/models/test_exposure_model.py
@@ -71,3 +71,33 @@ def test_exposure_model_vectorisation(override_params):
     assert isinstance(expected_new_cases, np.ndarray)
     assert expected_new_cases.shape == (2, )
 
+
+@pytest.mark.parametrize(
+    "vector_param", [
+        'volume', 'air_change', 'virus_halflife',
+        'viral_load_in_sputum', 'coefficient_of_infectivity', 'η_exhale',
+        'η_leaks', 'η_inhale',
+    ]
+)
+def test_exposure_model_compare_scalar_vector(vector_param):
+    defaults = {
+        'volume': 75,
+        'air_change': 100,
+        'virus_halflife': 1.1,
+        'viral_load_in_sputum': 1e9,
+        'coefficient_of_infectivity': 0.02,
+        'η_exhale': 0.95,
+        'η_leaks': 0.15,
+        'η_inhale': 0.3,
+    }
+    e_model_scalar = exposure_model_from_params(defaults)
+    expected_new_cases_scalar = e_model_scalar.expected_new_cases()
+    assert isinstance(expected_new_cases_scalar, float)
+
+    defaults[vector_param] = np.ones(3)*defaults[vector_param]
+    e_model_vector = exposure_model_from_params(defaults)
+    expected_new_cases_vector = e_model_vector.expected_new_cases()
+    assert isinstance(expected_new_cases_vector, np.ndarray)
+    assert expected_new_cases_vector.shape == (3, )
+    assert np.all(expected_new_cases_vector==expected_new_cases_scalar)
+