Adding check_parameter_positivity() function to seir.py

flepimop/gempyor_pkg/src/gempyor/seir.py

@@ -16,6 +16,51 @@
 logger = logging.getLogger(__name__)
 
 
+def check_parameter_positivity(
+    parsed_parameters: np.ndarray,
+    parameter_names: list[str],
+    dates: pd.DatetimeIndex,
+    subpop_names: list[str],
+) -> None:
+    """
+    Identifies and reports earliest negative values for parameters after modifiers have been applied.
+
+    Args:
+        parsed_parameters: An array of parameter values.
+        parameter_names: A list of the names of parameters.
+        dates: A pandas DatetimeIndex containing the dates.
+        subpop_names: A list of the names of subpopulations.
+
+    Raises:
+        ValueError: Negative parameter values were detected.
+
+    Returns:
+        None
+    """
+    if len(negative_index_parameters := np.argwhere(parsed_parameters < 0)) > 0:
+        unique_param_sp_combinations = []
+        row_index = -1
+        redundant_rows = []
+        for row in negative_index_parameters:
+            row_index += 1
+            if (row[0], row[2]) in unique_param_sp_combinations:
+                redundant_rows.append(row_index)
+            if (row[0], row[2]) not in unique_param_sp_combinations:
+                unique_param_sp_combinations.append((row[0], row[2]))
+        non_redundant_negative_parameters = np.delete(
+            negative_index_parameters, (redundant_rows), axis=0
+        )
+
+        neg_subpops = []
+        neg_params = []
+        first_neg_date = dates[0].date()
+        for param_idx, _, sp_idx in non_redundant_negative_parameters:
+            neg_subpops.append(subpop_names[sp_idx])
+            neg_params.append(parameter_names[param_idx])
+        error_message = f"There are negative parameter errors in subpops {neg_subpops}, starting from date {first_neg_date} in parameters {neg_params}."
+        raise ValueError(f"{error_message}")
+
+
 def build_step_source_arg(
     modinf: ModelInfo,
     parsed_parameters,
@@ -119,6 +164,11 @@ def build_step_source_arg(
         "population": modinf.subpop_pop,
         "stochastic_p": modinf.stoch_traj_flag,
     }
+
+    check_parameter_positivity(
+        fnct_args["parameters"], modinf.parameters.pnames, modinf.dates, modinf.subpop_pop
+    )
+
     return fnct_args
 
 

flepimop/gempyor_pkg/tests/seir/test_seir.py

@@ -3,6 +3,9 @@
 import pytest
 import warnings
 import shutil
+from random import randint
+import pandas as pd
+import re
 
 import pathlib
 import pyarrow as pa
@@ -16,8 +19,97 @@
 os.chdir(os.path.dirname(__file__))
 
 
+def test_check_parameter_positivity():
+
+    parameter_names = [
+        "alpha*1*1*1",
+        "sigma_OMICRON*1*1*1",
+        "3*gamma*1*1*1",
+        "epsilon+omegaph4*1*1*1",
+        "1*zeta*1*1",
+        "r0*gamma*theta10*1*chi_OMICRON*1",
+        "r0*gamma*theta9*1*chi_OMICRON*1",
+        "eta_X0toX3_highIE*1*1*nuage0to17",
+        "eta_X0toX3_highIE*1*1*nuage18to64LR",
+        "eta_X0toX3_highIE*1*1*nuage18to64HR",
+        "eta_X0toX3_highIE*1*1*nuage65to100",
+    ]
+    dates = pd.date_range("2023-03-19", "2025-04-30", freq="D")
+    subpop_names = [
+        "56000",
+        "50000",
+        "11000",
+        "02000",
+        "38000",
+        "46000",
+        "10000",
+        "30000",
+        "44000",
+        "23000",
+    ]
+
+    # No negative params
+    test_array1 = np.zeros(
+        (len(parameter_names) - 1, len(dates) - 1, len(subpop_names) - 1)
+    )
+    assert (
+        seir.check_parameter_positivity(test_array1, parameter_names, dates, subpop_names)
+        is None
+    )
+    # No Error
+
+    # Randomized negative params
+    test_array2 = np.zeros(
+        (len(parameter_names) - 1, len(dates) - 1, len(subpop_names) - 1)
+    )
+    for _ in range(5):
+        test_array2[randint(0, len(parameter_names) - 1)][randint(0, len(dates) - 1)][
+            randint(0, len(subpop_names) - 1)
+        ] = -1
+    test_2_negative_index_parameters = np.argwhere(test_array2 < 0)
+    test_2_neg_params = []
+    test_2_neg_subpops = []
+    test_2_first_neg_date = dates[0].date()
+    for param_idx, _, sp_idx in test_2_negative_index_parameters:
+        test_2_neg_subpops.append(subpop_names[sp_idx])
+        test_2_neg_params.append(parameter_names[param_idx])
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            rf"There are negative parameter errors in subpops {test_2_neg_subpops}, starting from date {test_2_first_neg_date} in parameters {test_2_neg_params}."
+        ),
+    ):
+        seir.check_parameter_positivity(
+            test_array2, parameter_names, dates, subpop_names
+        )  # ValueError
+
+    # Manually set negative params with intentional redundancy
+    test_array3 = np.zeros((len(parameter_names), len(dates), len(subpop_names)))
+    test_array3[0, 0, 0] = -1
+    test_array3[1, 1, 1] = -1
+    test_array3[2, 2, 2] = -1
+    test_array3[3, 3, 3] = -1
+    test_3_negative_index_parameters = np.argwhere(test_array3 < 0)
+    test_3_neg_params = []
+    test_3_neg_subpops = []
+    test_3_first_neg_date = dates[0].date()
+    for param_idx, _, sp_idx in test_3_negative_index_parameters:
+        test_3_neg_subpops.append(subpop_names[sp_idx])
+        test_3_neg_params.append(parameter_names[param_idx])
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            rf"There are negative parameter errors in subpops {test_3_neg_subpops}, starting from date {test_3_first_neg_date} in parameters {test_3_neg_params}."
+        ),
+    ):
+        seir.check_parameter_positivity(
+            test_array3, parameter_names, dates, subpop_names
+        )  # ValueError
+
+
 def test_check_values():
-    os.chdir(os.path.dirname(__file__))
     config.set_file(f"{DATA_DIR}/config.yml")
 
     modinf = model_info.ModelInfo(