Abstract single state runner

docker_template_configs/config_global.json

@@ -11,7 +11,7 @@
         65
     ],
     "NUM_STRAINS": 3,
-    "MAX_VAX_COUNT": 2,
+    "MAX_VACCINATION_COUNT": 2,
     "NUM_WANING_COMPARTMENTS": 5,
     "WANING_TIMES": [
         131,

docker_template_configs/config_inferer_covid.json

@@ -2,9 +2,9 @@
     "SCENARIO_NAME": "test covid run for testing suite",
     "CONTACT_MATRIX_PATH": "/input/data/demographic-data/contact_matrices",
     "SAVE_PATH": "/output/",
-    "HOSP_PATH": "/input/data/hospital_220213_220108.csv",
-    "VAX_MODEL_DATA": "/input/data/spline_fits.csv",
-    "VAX_MODEL_NUM_KNOTS": 18,
+    "HOSPITALIZATION_DATA_PATH": "/input/data/hospital_220213_220108.csv",
+    "VACCINATION_MODEL_DATA": "/input/data/spline_fits.csv",
+    "SEASONAL_VACCINATION": false,
     "STRAIN_R0s": [
         1.2,
         1.8,
@@ -52,14 +52,20 @@
             }
         }
     ],
-    "INTRODUCTION_PERCENTAGE": 0.01,
-    "INTRODUCTION_SCALE": 10,
+    "CONSTANT_STEP_SIZE": 0,
+    "INTRODUCTION_PCTS": [
+        0.01
+    ],
+    "INTRODUCTION_SCALES": [
+        10
+    ],
     "INTRODUCTION_AGE_MASK": [
         false,
         true,
         false,
         false
     ],
+    "MIN_HOMOLOGOUS_IMMUNITY": 0.1,
     "WANING_PROTECTIONS": [
         1.0,
         0.942,
@@ -84,7 +90,7 @@
             1.0
         ]
     ],
-    "VAX_EFF_MATRIX": [
+    "VACCINE_EFF_MATRIX": [
         [
             0,
             0.34,
@@ -111,6 +117,9 @@
         1.0,
         1.0
     ],
+    "SEASONALITY_AMPLITUDE": 0.15,
+    "SEASONALITY_SECOND_WAVE": 0.5,
+    "SEASONALITY_SHIFT": 0,
     "INFERENCE_PRNGKEY": 8675309,
     "INFERENCE_NUM_WARMUP": 100,
     "INFERENCE_NUM_SAMPLES": 500,

docker_template_configs/config_initializer_covid.json

@@ -1,7 +1,50 @@
 {
     "INITALIZER_NAME": "base initializer covid",
-    "SEROLOGICAL_DATA_PATH": "/input/data/serological-data/",
-    "SIM_DATA_PATH": "/input/data/abm-data/abm_population.csv",
+    "SEROLOGICAL_DATA_PATH": "/input/data/serological-data/sero-initializer-csvs",
+    "CONTACT_MATRIX_PATH": "/input/data/demographic-data/contact_matrices",
     "POP_SIZE": 1000000,
-    "INITIAL_INFECTIONS": 32000
+    "INITIAL_INFECTIONS": 32000,
+    "WANING_PROTECTIONS": [
+        1.0,
+        0.942,
+        0.942,
+        0.942,
+        0.0
+    ],
+    "INFECTIOUS_PERIOD": 7.0,
+    "EXPOSED_TO_INFECTIOUS": 3.6,
+    "STRAIN_INTERACTIONS": [
+        [
+            1.0,
+            0.7,
+            0.49
+        ],
+        [
+            0.7,
+            1.0,
+            0.7
+        ],
+        [
+            0.49,
+            0.7,
+            1.0
+        ]
+    ],
+    "VACCINE_EFF_MATRIX": [
+        [
+            0,
+            0.34,
+            0.68
+        ],
+        [
+            0,
+            0.24,
+            0.48
+        ],
+        [
+            0,
+            0.14,
+            0.28
+        ]
+    ]
 }

docker_template_configs/config_runner_covid.json

@@ -3,8 +3,8 @@
     "CONTACT_MATRIX_PATH": "/input/data/demographic-data/contact_matrices",
     "SAVE_PATH": "/output/",
     "HOSP_PATH": "/input/data/hospital_220213_220108.csv",
-    "VAX_MODEL_DATA": "/input/data/spline_fits.csv",
-    "VAX_MODEL_NUM_KNOTS": 18,
+    "VACCINATION_MODEL_DATA": "/input/data/spline_fits.csv",
+    "SEASONAL_VACCINATION": false,
     "STRAIN_R0s": [
         1.2,
         1.8,
@@ -16,14 +16,19 @@
     "INTRODUCTION_TIMES": [
         60
     ],
-    "INTRODUCTION_PERCENTAGE": 0.01,
-    "INTRODUCTION_SCALE": 10,
+    "INTRODUCTION_PCTS": [
+        0.01
+    ],
+    "INTRODUCTION_SCALES": [
+        10
+    ],
     "INTRODUCTION_AGE_MASK": [
         false,
         true,
         false,
         false
     ],
+    "MIN_HOMOLOGOUS_IMMUNITY": 0.1,
     "WANING_PROTECTIONS": [
         1.0,
         0.942,
@@ -48,7 +53,7 @@
             1.0
         ]
     ],
-    "VAX_EFF_MATRIX": [
+    "VACCINE_EFF_MATRIX": [
         [
             0,
             0.34,
@@ -65,6 +70,7 @@
             0.28
         ]
     ],
+    "CONSTANT_STEP_SIZE": 0,
     "BETA_TIMES": [
         0.0,
         120.0,
@@ -75,6 +81,9 @@
         1.0,
         1.0
     ],
+    "SEASONALITY_AMPLITUDE": 0.15,
+    "SEASONALITY_SECOND_WAVE": 0.0,
+    "SEASONALITY_SHIFT": 20,
     "MCMC_PRNGKEY": 8675309,
     "MCMC_NUM_WARMUP": 100,
     "MCMC_NUM_SAMPLES": 1000,

exp/example_azure_experiment/example_single_state_runner.py

@@ -0,0 +1,85 @@
+"""
+An example script similar to example_end_to_end_run.py
+but adapted to show the differences between Azure runs and local ones.
+"""
+
+import argparse
+
+from mechanistic_model.abstract_azure_runner import AbstractAzureRunner
+from mechanistic_model.covid_sero_initializer import CovidSeroInitializer
+from mechanistic_model.mechanistic_runner import MechanisticRunner
+from mechanistic_model.static_value_parameters import StaticValueParameters
+from model_odes.seip_model import seip_ode
+
+
+class ExampleRunner(AbstractAzureRunner):
+    # __init__ already implemented by the abstract case
+    def __init__(self, azure_output_dir):
+        super().__init__(azure_output_dir)
+
+    # override the process state command since this is what
+    # each individual runner must figure out for itself
+    def process_state(self, state):
+        """
+        A similar version to example_end_to_end_run.py but modified to run on Azure batch
+
+        Parameters
+        ---------------
+        state: str
+            the USPS postal code of the state, e.g. CA, OR, NY, TX
+        """
+        # step 1: define your paths NOTE: These are all within the docker container!
+        # /input is a MOUNTED drive that we upload these files into right before the job launched
+        config_path = "/app/exp/example_azure_experiment/states/%s/" % state
+        # global_config include definitions such as age bin bounds and strain definitions
+        # Any value or data structure that needs context to be interpretted is here.
+        GLOBAL_CONFIG_PATH = config_path + "config_global.json"
+        # defines the init conditions of the scenario: pop size, initial infections etc.
+        INITIALIZER_CONFIG_PATH = config_path + "config_initializer_covid.json"
+        # defines the running variables, strain R0s, external strain introductions etc.
+        RUNNER_CONFIG_PATH = config_path + "config_runner_covid.json"
+        # sets up the initial conditions, initializer.get_initial_state() passed to runner
+        initializer = CovidSeroInitializer(
+            INITIALIZER_CONFIG_PATH, GLOBAL_CONFIG_PATH
+        )
+        # reads and interprets values from config, sets up downstream parameters
+        # like beta = STRAIN_R0s / INFECTIOUS_PERIOD
+        static_params = StaticValueParameters(
+            initializer.get_initial_state(),
+            RUNNER_CONFIG_PATH,
+            GLOBAL_CONFIG_PATH,
+        )
+        # A runner that does ODE solving of a single run.
+        runner = MechanisticRunner(seip_ode)
+        # run for 200 days, using init state and parameters from StaticValueParameters
+        solution = runner.run(
+            initializer.get_initial_state(),
+            tf=200,
+            args=static_params.get_parameters(),
+        )
+        self.save_static_run_timelines(static_params, solution)
+        return solution
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "-s",
+    "--state",
+    type=str,
+    help="directory for the state to run, resembles USPS code of the state",
+    required=True,
+)
+
+parser.add_argument(
+    "-j",
+    "--jobid",
+    type=str,
+    help="job-id of the state being run on Azure",
+    required=True,
+)
+args = parser.parse_args()
+state = args.state
+jobid = args.jobid
+save_path = "/output/example_output/%s/%s/" % (jobid, state)
+runner = ExampleRunner(save_path)
+runner.process_state(state)

experiment_setup.py

@@ -208,13 +208,20 @@ def get_all_codes():
     state_names = pd.read_csv("data/fips_to_name.csv")
     pops = pd.read_csv("data/demographic-data/CenPop2020_Mean_ST.csv")
     # adding a USA row with the sum of all state pops
-    pops.loc[-1] = [
-        "US",
-        "United States",
-        sum(pops["POPULATION"]),
-        None,
-        None,
-    ]
+    usa_pop_row = pd.DataFrame(
+        [
+            [
+                "US",
+                "United States",
+                sum(pops["POPULATION"]),
+                None,
+                None,
+            ]
+        ],
+        columns=pops.columns,
+    )
+    pops = pd.concat([pops, usa_pop_row], ignore_index=True)
+    print(pops)
     args = parser.parse_args()
     exp = args.exp
     states = args.states

mechanistic_model/__init__.py

@@ -1 +1,9 @@
 # needs to exist to define a module
+import jax
+
+SEIC_Compartments = tuple[
+    jax.Array,
+    jax.Array,
+    jax.Array,
+    jax.Array,
+]