Merge pull request #194 from HopkinsIDD/gempyor-plugin

Gempyor plugins for `Seeding` and `InitialConditions`, fixes issue #101 🧩

.github/workflows/ci.yml

@@ -10,6 +10,7 @@ on:
     branches:
       - main
       - dev
+      - breaking-improvments
 
 jobs:
   unit-tests:

.gitignore

@@ -12,7 +12,7 @@ flepimop/gempyor_pkg/docs/full_graph*
 flepimop/gempyor_pkg/docs/readable_graph*
 Outcomes/test/model_output
 model_output/
-
+.idea/
 # History files
 .Rhistory
 .Rapp.history

flepimop/gempyor_pkg/docs/integration_benchmark.ipynb

@@ -449,8 +449,8 @@
     "    )\n",
     "\n",
     "with Timer(\"onerun_SEIR.seeding\"):\n",
-    "    initial_conditions = s.initial_conditions.draw(sim_id, setup=s)\n",
-    "    seeding_data, seeding_amounts = s.seeding.draw(sim_id, setup=s)\n",
+    "    initial_conditions = s.initial_conditions.get_from_config(sim_id, setup=s)\n",
+    "    seeding_data, seeding_amounts = s.seeding.get_from_config(sim_id, setup=s)\n",
     "\n",
     "mobility_subpop_indices = s.mobility.indices\n",
     "mobility_data_indices = s.mobility.indptr\n",

flepimop/gempyor_pkg/docs/integration_doc.ipynb

@@ -110,8 +110,8 @@
     "npi_seir = seir.build_npi_SEIR(s=gempyor_simulator.s, load_ID=load_ID, sim_id2load=sim_id2load, config=config)\n",
     "\n",
     "\n",
-    "initial_conditions = gempyor_simulator.s.initial_conditions.draw(sim_id2write, setup=gempyor_simulator.s)\n",
-    "seeding_data, seeding_amounts = gempyor_simulator.s.seeding.draw(sim_id2write, setup=gempyor_simulator.s)\n",
+    "initial_conditions = gempyor_simulator.s.initial_conditions.get_from_config(sim_id2write, setup=gempyor_simulator.s)\n",
+    "seeding_data, seeding_amounts = gempyor_simulator.s.seeding.get_from_config(sim_id2write, setup=gempyor_simulator.s)\n",
     "\n",
     "\n",
     "p_draw = gempyor_simulator.s.parameters.parameters_quick_draw(\n",

flepimop/gempyor_pkg/docs/interface.ipynb

@@ -257,17 +257,17 @@
     "### Run every time:\n",
     "with Timer(\"onerun_SEIR.seeding\"):\n",
     "    if load_ID:\n",
-    "        initial_conditions = gempyor_simulator.s.initial_conditions.load(\n",
+    "        initial_conditions = gempyor_simulator.s.initial_conditions.get_from_file(\n",
     "            sim_id2load, setup=gempyor_simulator.s\n",
     "        )\n",
-    "        seeding_data, seeding_amounts = gempyor_simulator.s.seeding.load(\n",
+    "        seeding_data, seeding_amounts = gempyor_simulator.s.seeding.get_from_file(\n",
     "            sim_id2load, setup=gempyor_simulator.s\n",
     "        )\n",
     "    else:\n",
-    "        initial_conditions = gempyor_simulator.s.initial_conditions.draw(\n",
+    "        initial_conditions = gempyor_simulator.s.initial_conditions.get_from_config(\n",
     "            sim_id2write, setup=gempyor_simulator.s\n",
     "        )\n",
-    "        seeding_data, seeding_amounts = gempyor_simulator.s.seeding.draw(\n",
+    "        seeding_data, seeding_amounts = gempyor_simulator.s.seeding.get_from_config(\n",
     "            sim_id2write, setup=gempyor_simulator.s\n",
     "        )\n",
     "\n",

flepimop/gempyor_pkg/src/gempyor/interface.py

@@ -222,13 +222,13 @@ def one_simulation(
 
             with Timer("onerun_SEIR.seeding"):
                 if load_ID:
-                    initial_conditions = self.modinf.initial_conditions.load(sim_id2load, setup=self.modinf)
-                    seeding_data, seeding_amounts = self.modinf.seeding.load(
+                    initial_conditions = self.modinf.initial_conditions.get_from_file(sim_id2load, setup=self.modinf)
+                    seeding_data, seeding_amounts = self.modinf.seeding.get_from_file(
                         sim_id2load, setup=self.modinf
                     )
                 else:
-                    initial_conditions = self.modinf.initial_conditions.draw(sim_id2write, setup=self.modinf)
-                    seeding_data, seeding_amounts = self.modinf.seeding.draw(
+                    initial_conditions = self.modinf.initial_conditions.get_from_config(sim_id2write, setup=self.modinf)
+                    seeding_data, seeding_amounts = self.modinf.seeding.get_from_config(
                         sim_id2write, setup=self.modinf
                     )
                 self.debug_seeding_data = seeding_data

flepimop/gempyor_pkg/src/gempyor/model_info.py

@@ -115,8 +115,8 @@ def __init__(
                 tf=self.tf,
                 subpop_names=self.subpop_struct.subpop_names,
             )
-            self.seeding = seeding_ic.Seeding(config = self.seeding_config)
-            self.initial_conditions = seeding_ic.InitialConditions(config = self.initial_conditions_config)
+            self.seeding = seeding_ic.SeedingFactory(config = self.seeding_config)
+            self.initial_conditions = seeding_ic.InitialConditionsFactory(config = self.initial_conditions_config)
             # really ugly references to the config globally here.
             if config["compartments"].exists() and self.seir_config is not None:
                 self.compartments = compartments.Compartments(

flepimop/gempyor_pkg/src/gempyor/outcomes.py

@@ -38,7 +38,7 @@ def run_parallel_outcomes(modinf, *, sim_id2write, nslots=1, n_jobs=1):
         tqdm.contrib.concurrent.process_map(
             onerun_delayframe_outcomes,
             sim_id2writes,
-            modinf,
+            itertools.repeat(modinf),
             max_workers=n_jobs,
         )
 
@@ -88,7 +88,6 @@ def build_outcome_modifiers(
 
 
 def onerun_delayframe_outcomes(
-    *,
     sim_id2write: int,
     modinf: model_info.ModelInfo,
     load_ID: bool = False,

flepimop/gempyor_pkg/src/gempyor/seeding_ic.py

@@ -9,7 +9,7 @@
 import confuse
 import logging
 from . import compartments
-from . import model_info
+from . import utils
 import numba as nb
 from .utils import read_df
 
@@ -343,10 +343,10 @@ class SimulationComponent:
     def __init__(self, config: confuse.ConfigView):
         raise NotImplementedError("This method should be overridden in subclasses.")
 
-    def load(self, sim_id: int, setup) -> np.ndarray:
+    def get_from_file(self, sim_id: int, setup) -> np.ndarray:
         raise NotImplementedError("This method should be overridden in subclasses.")
 
-    def draw(self, sim_id: int, setup) -> np.ndarray:
+    def get_from_config(self, sim_id: int, setup) -> np.ndarray:
         raise NotImplementedError("This method should be overridden in subclasses.")
 
     def write_to_file(self, sim_id: int, setup):
@@ -357,7 +357,7 @@ class Seeding(SimulationComponent):
     def __init__(self, config: confuse.ConfigView):
         self.seeding_config = config
 
-    def draw(self, sim_id: int, setup) -> nb.typed.Dict:
+    def get_from_config(self, sim_id: int, setup) -> nb.typed.Dict:
         method = "NoSeeding"
         if self.seeding_config is not None and "method" in self.seeding_config.keys():
             method = self.seeding_config["method"].as_str()
@@ -418,16 +418,15 @@ def draw(self, sim_id: int, setup) -> nb.typed.Dict:
 
         return _DataFrame2NumbaDict(df=seeding, amounts=amounts, setup=setup)
 
-    def load(self, sim_id: int, setup) -> nb.typed.Dict:
+    def get_from_file(self, sim_id: int, setup) -> nb.typed.Dict:
         """only difference with draw seeding is that the sim_id is now sim_id2load"""
-        return self.draw(sim_id=sim_id, setup=setup)
-
+        return self.get_from_config(sim_id=sim_id, setup=setup)
 
 class InitialConditions(SimulationComponent):
     def __init__(self, config: confuse.ConfigView):
         self.initial_conditions_config = config
 
-    def draw(self, sim_id: int, setup) -> np.ndarray:
+    def get_from_config(self, sim_id: int, setup) -> np.ndarray:
         method = "Default"
         if self.initial_conditions_config is not None and "method" in self.initial_conditions_config.keys():
             method = self.initial_conditions_config["method"].as_str()
@@ -483,7 +482,7 @@ def draw(self, sim_id: int, setup) -> np.ndarray:
                             else:
                                 raise ValueError(
                                     f"Initial Conditions: Could not set compartment {comp_name} (id: {comp_idx}) in subpop {pl} (id: {pl_idx}). The data from the init file is {states_pl}. \n \
-                                                 Use 'allow_missing_compartments' to default to 0 for compartments without initial conditions"
+                                                Use 'allow_missing_compartments' to default to 0 for compartments without initial conditions"
                                 )
                         if "rest" in str(ic_df_compartment_val).strip().lower():
                             rests.append([comp_idx, pl_idx])
@@ -607,5 +606,29 @@ def draw(self, sim_id: int, setup) -> np.ndarray:
             )
         return y0
 
-    def load(self, sim_id: int, setup) -> nb.typed.Dict:
-        return self.draw(sim_id=sim_id, setup=setup)
+    def get_from_file(self, sim_id: int, setup) -> np.ndarray:
+        return self.get_from_config(sim_id=sim_id, setup=setup)
+
+# TODO: rename config to initial_conditions_config as it shadows the global config
+
+def InitialConditionsFactory(config: confuse.ConfigView):
+    if config is not None and "method" in config.keys():
+            if config["method"].as_str() == "plugin":
+                klass = utils.search_and_import_plugins_class(
+                    plugin_file_path=config["plugin_file_path"].as_str(), 
+                    class_name="InitialConditions",
+                    config=config
+                    )
+                return klass
+    return InitialConditions(config)
+
+def SeedingFactory(config: confuse.ConfigView):
+    if config is not None and "method" in config.keys():
+            if config["method"].as_str() == "plugin":
+                klass = utils.search_and_import_plugins_class(
+                    plugin_file_path=config["plugin_file_path"].as_str(), 
+                    class_name="Seeding",
+                    config=config
+                    )
+                return klass
+    return Seeding(config)

flepimop/gempyor_pkg/src/gempyor/seir.py

@@ -238,11 +238,11 @@ def onerun_SEIR(
 
     with Timer("onerun_SEIR.seeding"):
         if load_ID:
-            initial_conditions = modinf.initial_conditions.load(sim_id2load, setup=modinf)
-            seeding_data, seeding_amounts = modinf.seeding.load(sim_id2load, setup=modinf)
+            initial_conditions = modinf.initial_conditions.get_from_file(sim_id2load, setup=modinf)
+            seeding_data, seeding_amounts = modinf.seeding.get_from_file(sim_id2load, setup=modinf)
         else:
-            initial_conditions = modinf.initial_conditions.draw(sim_id2write, setup=modinf)
-            seeding_data, seeding_amounts = modinf.seeding.draw(sim_id2write, setup=modinf)
+            initial_conditions = modinf.initial_conditions.get_from_config(sim_id2write, setup=modinf)
+            seeding_data, seeding_amounts = modinf.seeding.get_from_config(sim_id2write, setup=modinf)
 
     with Timer("onerun_SEIR.parameters"):
         # Draw or load parameters

flepimop/gempyor_pkg/src/gempyor/utils.py

@@ -59,6 +59,20 @@ def wrapper(*args, **kwargs):
     return decorator
 
 
+def search_and_import_plugins_class(plugin_file_path: str, class_name: str, **kwargs):
+    # Look for all possible plugins and import them
+    # https://stackoverflow.com/questions/67631/how-can-i-import-a-module-dynamically-given-the-full-path
+    # unfortunatelly very complicated, this is cpython only ??
+    import sys, os
+    sys.path.append(os.path.dirname(plugin_file_path))
+    # the following works, but these above lines seems necessary to pickle // runs
+
+    from pydoc import importfile
+    module = importfile(plugin_file_path)
+    klass = getattr(module, class_name)
+    return klass(**kwargs)
+
+
 ### Profile configuration
 import cProfile
 import pstats
@@ -223,6 +237,32 @@ def as_random_distribution(self):
             self.as_evaled_expression(),
         )
 
+def list_filenames(folder: str = '.', filters:list = []) -> list:
+    """
+    return the list of all filename and path in the provided folders. 
+    If filters [list] is provided, then only the files that contains each of the 
+    substrings in filter will be returned. Example to get all hosp file:
+    ```
+        gempyor.utils.list_filenames(folder="model_output/", filters=["hosp"])
+    ```
+        and be sure we only get parquet:
+    ```
+        gempyor.utils.list_filenames(folder="model_output/", filters=["hosp" , ".parquet"])
+    ```
+    """
+    from pathlib import Path
+    fn_list = []
+    for f in Path(str(folder)).rglob(f'*'):
+        if f.is_file(): # not a folder
+            f = str(f)
+            if not filters:
+                fn_list.append(f)
+            else:
+                if all(c in f for c in filters):
+                    fn_list.append(str(f))
+                else:
+                    pass
+    return fn_list
 
 def aws_disk_diagnosis():
     import os

flepimop/gempyor_pkg/tests/seir/interface.ipynb

@@ -253,13 +253,13 @@
     "### Run every time:\n",
     "with Timer(\"onerun_SEIR.seeding\"):\n",
     "    if load_ID:\n",
-    "        initial_conditions = gempyor_simulator.s.initial_conditions.load(sim_id2load, setup=gempyor_simulator.s)\n",
-    "        seeding_data, seeding_amounts = gempyor_simulator.s.seeding.load(\n",
+    "        initial_conditions = gempyor_simulator.s.initial_conditions.get_from_file(sim_id2load, setup=gempyor_simulator.s)\n",
+    "        seeding_data, seeding_amounts = gempyor_simulator.s.seeding.get_from_file(\n",
     "            sim_id2load, setup=gempyor_simulator.s\n",
     "        )\n",
     "    else:\n",
-    "        initial_conditions = gempyor_simulator.s.initial_conditions.draw(sim_id2write, setup=gempyor_simulator.s)\n",
-    "        seeding_data, seeding_amounts = gempyor_simulator.s.seeding.draw(\n",
+    "        initial_conditions = gempyor_simulator.s.initial_conditions.get_from_config(sim_id2write, setup=gempyor_simulator.s)\n",
+    "        seeding_data, seeding_amounts = gempyor_simulator.s.seeding.get_from_config(\n",
     "            sim_id2write, setup=gempyor_simulator.s\n",
     "        )\n",
     "\n",

flepimop/gempyor_pkg/tests/seir/test_ic.py

@@ -21,7 +21,7 @@ def test_IC_success(self):
             outcome_modifiers_scenario=None,
             write_csv=False,
         )
-        sic = seeding_ic.InitialConditions(config=s.initial_conditions_config)
+        sic = seeding_ic.InitialConditionsFactory(config=s.initial_conditions_config)
         assert sic.initial_conditions_config == s.initial_conditions_config
 
     def test_IC_allow_missing_node_compartments_success(self):
@@ -40,10 +40,10 @@ def test_IC_allow_missing_node_compartments_success(self):
 
         s.initial_conditions_config["allow_missing_nodes"] = True
         s.initial_conditions_config["allow_missing_compartments"] = True
-        sic = seeding_ic.InitialConditions(config=s.initial_conditions_config)
+        sic = seeding_ic.InitialConditionsFactory(config=s.initial_conditions_config)
         assert sic.initial_conditions_config == s.initial_conditions_config
 
-        initial_conditions = sic.draw(sim_id=100, setup=s)
+        initial_conditions = sic.get_from_config(sim_id=100, setup=s)
         print(initial_conditions)
 
     def test_IC_IC_notImplemented_fail(self):
@@ -61,6 +61,6 @@ def test_IC_IC_notImplemented_fail(self):
                 write_csv=False,
             )
             s.initial_conditions_config["method"] = "unknown"
-            sic = seeding_ic.InitialConditions(config=s.initial_conditions_config)
+            sic = seeding_ic.InitialConditionsFactory(config=s.initial_conditions_config)
 
-            sic.draw(sim_id=100, setup=s)
+            sic.get_from_config(sim_id=100, setup=s)

flepimop/gempyor_pkg/tests/seir/test_seeding.py

@@ -20,7 +20,7 @@ def test_Seeding_success(self):
             outcome_modifiers_scenario=None,
             write_csv=False,
         )
-        sic = seeding_ic.Seeding(config=s.seeding_config)
+        sic = seeding_ic.SeedingFactory(config=s.seeding_config)
         assert sic.seeding_config == s.seeding_config
 
     def test_Seeding_draw_success(self):
@@ -36,8 +36,8 @@ def test_Seeding_draw_success(self):
             outcome_modifiers_scenario=None,
             write_csv=False,
         )
-        sic = seeding_ic.Seeding(config=s.seeding_config)
+        sic = seeding_ic.SeedingFactory(config=s.seeding_config)
         s.seeding_config["method"] = "NoSeeding"
 
-        seeding = sic.draw(sim_id=100, setup=s)
+        seeding = sic.get_from_config(sim_id=100, setup=s)
         print(seeding)