WIP on NPIs

HopkinsIDD · Nov 13, 2024 · 7ba97ce · 7ba97ce
1 parent aa66077
commit 7ba97ce
Show file tree

Hide file tree

Showing 9 changed files with 54 additions and 75 deletions.
diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/MultiPeriodModifier.py
@@ -1,17 +1,18 @@
+import datetime
+
 import numpy as np
 import pandas as pd
 
 from . import helpers
 from .base import NPIBase
-from ..model_info import ModelInfo
-
 
 class MultiPeriodModifier(NPIBase):
     def __init__(
         self,
         *,
         npi_config,
-        modinf: ModelInfo,
+        modinf_ti: datetime.date,
+        modinf_tf: datetime.date,
         modifiers_library,
         subpops,
         loaded_df=None,
@@ -29,8 +30,8 @@ def __init__(
         )
 
         self.sanitize = sanitize
-        self.start_date = modinf.ti
-        self.end_date = modinf.tf
+        self.start_date = modinf_ti
+        self.end_date = modinf_tf
 
         self.subpops = subpops
 

diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/SinglePeriodModifier.py
@@ -1,17 +1,18 @@
+import datetime
+
 import numpy as np
 import pandas as pd
 
 from . import helpers
 from .base import NPIBase
-from ..model_info import ModelInfo
-
 
 class SinglePeriodModifier(NPIBase):
     def __init__(
         self,
         *,
         npi_config,
-        modinf: ModelInfo,
+        modinf_ti: datetime.date,
+        modinf_tf: datetime.date,
         modifiers_library,
         subpops,
         loaded_df=None,
@@ -27,8 +28,8 @@ def __init__(
             )
         )
 
-        self.start_date = modinf.ti
-        self.end_date = modinf.tf
+        self.start_date = modinf_ti
+        self.end_date = modinf_tf
 
         self.pnames_overlap_operation_sum = pnames_overlap_operation_sum
         self.pnames_overlap_operation_reductionprod = pnames_overlap_operation_reductionprod

diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/StackedModifier.py b/flepimop/gempyor_pkg/src/gempyor/NPI/StackedModifier.py
@@ -1,13 +1,12 @@
 import collections
 import os
 import warnings
+import datetime
 
 import confuse
 import pandas as pd
 
 from .base import NPIBase
-from ..model_info import ModelInfo
-
 
 debug_print = False
 
@@ -21,7 +20,8 @@ def __init__(
         self,
         *,
         npi_config,
-        modinf: ModelInfo,
+        modinf_ti: datetime.date,
+        modinf_tf: datetime.date,
         modifiers_library,
         subpops,
         loaded_df=None,
@@ -30,8 +30,8 @@ def __init__(
     ):
         super().__init__(name=npi_config.name)
 
-        self.start_date = modinf.ti
-        self.end_date = modinf.tf
+        self.start_date = modinf_ti
+        self.end_date = modinf_tf
 
         self.pnames_overlap_operation_sum = pnames_overlap_operation_sum
         self.pnames_overlap_operation_reductionprod = pnames_overlap_operation_reductionprod
@@ -63,7 +63,7 @@ def __init__(
 
             sub_npi = NPIBase.execute(
                 npi_config=scenario_npi_config,
-                modinf=modinf,
+                modinf_ti=modinf_ti, modinf_tf=modinf_tf,
                 modifiers_library=modifiers_library,
                 subpops=subpops,
                 loaded_df=loaded_df,

diff --git a/flepimop/gempyor_pkg/src/gempyor/NPI/base.py b/flepimop/gempyor_pkg/src/gempyor/NPI/base.py
@@ -1,7 +1,5 @@
 import abc
-
-from ..model_info import ModelInfo
-
+import datetime
 
 class NPIBase(abc.ABC):
     __plugins__ = {}
@@ -28,7 +26,8 @@ def getReductionDF(self):
     def execute(
         *,
         npi_config,
-        modinf: ModelInfo,
+        modinf_ti: datetime.date,
+        modinf_tf: datetime.date,
         modifiers_library,
         subpops,
         loaded_df=None,
@@ -46,7 +45,8 @@ def execute(
         npi_class = NPIBase.__plugins__[method]
         return npi_class(
             npi_config=npi_config,
-            modinf=modinf,
+            modinf_ti=modinf_ti,
+            modinf_tf=modinf_tf,
             modifiers_library=modifiers_library,
             subpops=subpops,
             loaded_df=loaded_df,

diff --git a/flepimop/gempyor_pkg/src/gempyor/dev/dev_seir.py b/flepimop/gempyor_pkg/src/gempyor/dev/dev_seir.py
@@ -44,7 +44,7 @@
 
 npi = NPI.NPIBase.execute(
     npi_config=modinf.npi_config_seir,
-    modinf=modinf,
+    modinf_ti=modinf.ti, modinf_tf=modinf.tf,
     modifiers_library=modinf.seir_modifiers_library,
     subpops=modinf.subpop_struct.subpop_names,
     pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"],

diff --git a/flepimop/gempyor_pkg/src/gempyor/outcomes.py b/flepimop/gempyor_pkg/src/gempyor/outcomes.py
@@ -73,23 +73,14 @@ def build_outcome_modifiers(
         elif load_ID == True:
             loaded_df = modinf.read_simID(ftype="hnpi", sim_id=sim_id2load)
 
-        if loaded_df is not None:
-            npi = NPI.NPIBase.execute(
-                npi_config=modinf.npi_config_outcomes,
-                modinf=modinf,
-                modifiers_library=modinf.outcome_modifiers_library,
-                subpops=modinf.subpop_struct.subpop_names,
-                loaded_df=loaded_df,
-                # TODO: support other operation than product
-            )
-        else:
-            npi = NPI.NPIBase.execute(
-                npi_config=modinf.npi_config_outcomes,
-                modinf=modinf,
-                modifiers_library=modinf.outcome_modifiers_library,
-                subpops=modinf.subpop_struct.subpop_names,
-                # TODO: support other operation than product
-            )
+        npi = NPI.NPIBase.execute(
+            npi_config=modinf.npi_config_outcomes,
+            modinf_ti=modinf.ti, modinf_tf=modinf.tf,
+            modifiers_library=modinf.outcome_modifiers_library,
+            subpops=modinf.subpop_struct.subpop_names,
+            loaded_df=loaded_df,
+            # TODO: support other operation than product
+        )
     return npi
 
 

diff --git a/flepimop/gempyor_pkg/src/gempyor/seeding.py b/flepimop/gempyor_pkg/src/gempyor/seeding.py
@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, TYPE_CHECKING
 
 import numpy as np
 import pandas as pd

diff --git a/flepimop/gempyor_pkg/src/gempyor/seir.py b/flepimop/gempyor_pkg/src/gempyor/seir.py
@@ -235,33 +235,19 @@ def build_npi_SEIR(
         elif load_ID == True:
             loaded_df = modinf.read_simID(ftype="snpi", sim_id=sim_id2load)
 
-        if loaded_df is not None:
-            npi = NPI.NPIBase.execute(
-                npi_config=modinf.npi_config_seir,
-                modinf=modinf,
-                modifiers_library=modinf.seir_modifiers_library,
-                subpops=modinf.subpop_struct.subpop_names,
-                loaded_df=loaded_df,
-                pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method[
-                    "sum"
-                ],
-                pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[
-                    "reduction_product"
-                ],
-            )
-        else:
-            npi = NPI.NPIBase.execute(
-                npi_config=modinf.npi_config_seir,
-                modinf=modinf,
-                modifiers_library=modinf.seir_modifiers_library,
-                subpops=modinf.subpop_struct.subpop_names,
-                pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method[
-                    "sum"
-                ],
-                pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[
-                    "reduction_product"
-                ],
-            )
+        npi = NPI.NPIBase.execute(
+            npi_config=modinf.npi_config_seir,
+            modinf_ti=modinf.ti, modinf_tf=modinf.tf,
+            modifiers_library=modinf.seir_modifiers_library,
+            subpops=modinf.subpop_struct.subpop_names,
+            loaded_df=loaded_df,
+            pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method[
+                "sum"
+            ],
+            pnames_overlap_operation_reductionprod=modinf.parameters.stacked_modifier_method[
+                "reduction_product"
+            ],
+        )
     return npi
 
 

diff --git a/flepimop/gempyor_pkg/tests/seir/test_seir.py b/flepimop/gempyor_pkg/tests/seir/test_seir.py
@@ -80,7 +80,7 @@ def test_constant_population_legacy_integration():
 
     npi = NPI.NPIBase.execute(
         npi_config=modinf.npi_config_seir,
-        modinf=modinf,
+        modinf_ti=modinf.ti, modinf_tf=modinf.tf,
         modifiers_library=modinf.seir_modifiers_library,
         subpops=modinf.subpop_struct.subpop_names,
         pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"],
@@ -157,7 +157,7 @@ def test_constant_population_rk4jit_integration_fail():
 
         npi = NPI.NPIBase.execute(
             npi_config=modinf.npi_config_seir,
-            modinf=modinf,
+            modinf_ti=modinf.ti, modinf_tf=modinf.tf,
             modifiers_library=modinf.seir_modifiers_library,
             subpops=modinf.subpop_struct.subpop_names,
             pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"],
@@ -233,7 +233,7 @@ def test_constant_population_rk4jit_integration():
 
     npi = NPI.NPIBase.execute(
         npi_config=modinf.npi_config_seir,
-        modinf=modinf,
+        modinf_ti=modinf.ti, modinf_tf=modinf.tf,
         modifiers_library=modinf.seir_modifiers_library,
         subpops=modinf.subpop_struct.subpop_names,
         pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"],
@@ -307,7 +307,7 @@ def test_steps_SEIR_nb_simple_spread_with_txt_matrices():
 
     npi = NPI.NPIBase.execute(
         npi_config=modinf.npi_config_seir,
-        modinf=modinf,
+        modinf_ti=modinf.ti, modinf_tf=modinf.tf,
         modifiers_library=modinf.seir_modifiers_library,
         subpops=modinf.subpop_struct.subpop_names,
         pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"],
@@ -415,7 +415,7 @@ def test_steps_SEIR_nb_simple_spread_with_csv_matrices():
 
     npi = NPI.NPIBase.execute(
         npi_config=modinf.npi_config_seir,
-        modinf=modinf,
+        modinf_ti=modinf.ti, modinf_tf=modinf.tf,
         modifiers_library=modinf.seir_modifiers_library,
         subpops=modinf.subpop_struct.subpop_names,
         pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"],
@@ -493,7 +493,7 @@ def test_steps_SEIR_no_spread():
 
     npi = NPI.NPIBase.execute(
         npi_config=modinf.npi_config_seir,
-        modinf=modinf,
+        modinf_ti=modinf.ti, modinf_tf=modinf.tf,
         modifiers_library=modinf.seir_modifiers_library,
         subpops=modinf.subpop_struct.subpop_names,
         pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"],
@@ -767,7 +767,7 @@ def test_parallel_compartments_with_vacc():
 
     npi = NPI.NPIBase.execute(
         npi_config=modinf.npi_config_seir,
-        modinf=modinf,
+        modinf_ti=modinf.ti, modinf_tf=modinf.tf,
         modifiers_library=modinf.seir_modifiers_library,
         subpops=modinf.subpop_struct.subpop_names,
         pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"],
@@ -860,7 +860,7 @@ def test_parallel_compartments_no_vacc():
 
     npi = NPI.NPIBase.execute(
         npi_config=modinf.npi_config_seir,
-        modinf=modinf,
+        modinf_ti=modinf.ti, modinf_tf=modinf.tf,
         modifiers_library=modinf.seir_modifiers_library,
         subpops=modinf.subpop_struct.subpop_names,
         pnames_overlap_operation_sum=modinf.parameters.stacked_modifier_method["sum"],