Multi-strata simulation refactor.

Includes class-based refactor for initializer and parameter functions.
Add spec to StaticGeo. Regenerate geos.
Update vignettes for API changes.
New initializers vignette.

Author: Tyler Coles

.vscode/settings.json

@@ -26,7 +26,7 @@
 
   "python.formatting.provider": "none",
   "python.analysis.autoImportCompletions": true,
-  "python.analysis.typeCheckingMode": "basic",
+  "python.analysis.typeCheckingMode": "standard",
   "python.analysis.diagnosticMode": "workspace",
 
   "python.testing.pytestEnabled": false,

README.md

@@ -1,14 +1,18 @@
 # epymorph
 
-Prototype EpiMoRPH system written in Python. It is usable as a CLI program and also as a code library; for instance you may want to use it from within a Jupyter Notebook (see: `USAGE.ipynb`).
+Prototype EpiMoRPH system written in Python. It is usable as a code library, for instance, from within a Jupyter Notebook (see: `USAGE.ipynb`).
 
-The `epymorph/data` directory is the model library, containing named implementations of IPMs, MMs, and GEOs. Ultimately our goal is to allow users to bring-their-own models by loading specification files, but for now they need to be registered in the model library.
+The `epymorph/data` directory is the model library, containing named implementations of IPMs, MMs, and GEOs. Ultimately our goal is to allow users to bring their own models by loading specification files, but for now they need to be registered in the model library.
 
 The `doc/devlog` directory contains Jupyter Notebooks demonstrating features of epymorph and general development progress.
 
 Beware: much of this code structure is experimental and subject to change!
 
-## Project setup
+## Basic usage
+
+See the `USAGE.ipynb` Notebook for a simple usage example.
+
+## Development setup
 
 For starters, you should have Python 3.11 installed and we'll assume it's accessible via the command `python3.11`.
 
@@ -18,7 +22,7 @@ You may need to install additional system packages for virtual environments and
 sudo apt install python3.11-venv python3.11-tk
 ```
 
-Using VS Code, install the project's recommended IDE extensions. Then use the "Python - Create Environment" command (`Ctrl+Shift+P`) to create a Venv environment and install all dependencies (including `dev`).
+If you are using VS Code, install the project's recommended IDE extensions. Then use the "Python - Create Environment" command (`Ctrl+Shift+P`) to create a Venv environment and install all dependencies (including `dev`).
 
 Or you can set up from the command line:
 
@@ -29,64 +33,14 @@ cd $PROJECT_DIRECTORY
 python3.11 -m venv .venv
 
 # activate it (after which `python` should be bound to the venv python)
+# NOTE: activating venv on Windows is different; see documentation
 source .venv/bin/activate
 
 # then install the project in editable mode
 python -m pip install --editable ".[dev]"
 ```
 
-(The quotes in the last command are necessary on some platforms, but if the quotes give you issues you can also try without them.)
-
-## Running from the command line
-
-The most basic task epymorph can perform is to run a spatial, compartmental disease simulation and output the time-series data of compartment populations (prevalence) as well as new events (incidence).
-
-A commonly-cited model was proposed by [Sen Pei, et al. in 2018](https://www.pnas.org/doi/10.1073/pnas.1708856115), modeling influenza in six 
-southern US states. epymorph has an intra-population model (IPM), movement model (MM), and geographic model (GEO) that closely mimics Pei's experiment.
-
-To run a simulation, first we need an input file that describes the simulation and all its parameters. Thankfully epymorph has a subcommand to help us create such a file. This example will store files in a `scratch` folder within the project (this is just for convenience, you can opt to put the input files anywhere you like).
-
-```bash
-cd $PROJECT_DIRECTORY
-
-# Activate the venv (if it's not already):
-source .venv/bin/activate
-
-# scratch is a convenient place to put all sorts of temp files because our .gitignore excludes it
-mkdir scratch
-
-# Prepare the simulation input file:
-python -m epymorph prepare --ipm pei --mm pei --geo pei ./scratch/my-experiment.toml
-
-# Now we need to edit the input file to specify the parameters needed by our combo of IPM and MM:
-# (I'll use `cat` for this but you can use any text editor of course.)
-cat << EOF >> ./scratch/my-experiment.toml
-theta = 0.1
-move_control = 0.9
-infection_duration = 4.0
-immunity_duration = 90.0
-EOF
-
-# Now we can run the simulation:
-python -m epymorph run ./scratch/my-experiment.toml --out ./scratch/output.csv
-
-# Now if I open that csv file I see:
-# - for each time-step (t) and population (p)
-#   - prevalence data by compartment (c0, c1, c2)
-#   - incidence data by event (e0, e1, e2)
-
-# You can also run to display a chart:
-python -m epymorph run ./scratch/my-experiment.toml --chart p0
-```
-
-To learn more about these and all other commands, you can always consult the CLI help:
-
-```bash
-python -m epymorph --help
-
-# or for a specific subcommand
-python -m epymorph run --help
-```
+Make sure you have correctly configured auto-formatting in your development environment. We're currently using autopep8 and isort. These formatting tools should run every time you save a file.
 
 ### Other command-line tasks
 
@@ -95,17 +49,3 @@ Run all unit tests:
 ```bash
 python -m unittest discover -v -s ./epymorph -p '*_test.py'
 ```
-
-Run a simulation with the pdb debugger:
-
-```bash
-python -m pdb -m epymorph run ./scratch/my-experiment.toml
-```
-
-Profile the simulation and show the results in `snakeviz`:
-
-```bash
-TMP=$(mktemp /tmp/py-XXXXXXXX.prof)
-python -m cProfile -o $TMP -m epymorph run ./scratch/my-experiment.toml --profile
-snakeviz $TMP
-```

USAGE.ipynb

@@ -26,17 +26,22 @@
     "\n",
     "from epymorph.data_shape import Shapes\n",
     "from epymorph.data_type import CentroidDType\n",
-    "from epymorph.geo.spec import StaticGeoSpec, Year, attrib\n",
+    "from epymorph.geo.spec import StaticGeoSpec, Year\n",
+    "from epymorph.geography.us_census import StateScope\n",
+    "from epymorph.simulation import AttributeDef\n",
     "\n",
     "spec = StaticGeoSpec(\n",
     "    attributes=[\n",
-    "        attrib('label', str, Shapes.N),\n",
-    "        attrib('geoid', str, Shapes.N),\n",
-    "        attrib('centroid', CentroidDType, Shapes.N),\n",
-    "        attrib('population', int, Shapes.N),\n",
-    "        attrib('commuters', int, Shapes.NxN),\n",
-    "        attrib('humidity', float, Shapes.TxN),\n",
+    "        AttributeDef('label', str, Shapes.N),\n",
+    "        AttributeDef('geoid', str, Shapes.N),\n",
+    "        AttributeDef('centroid', CentroidDType, Shapes.N),\n",
+    "        AttributeDef('population', int, Shapes.N),\n",
+    "        AttributeDef('commuters', int, Shapes.NxN),\n",
+    "        AttributeDef('humidity', float, Shapes.TxN),\n",
     "    ],\n",
+    "    # critically: these states are listed here in GEOID order,\n",
+    "    # and we maintain that order when entering data below\n",
+    "    scope=StateScope.in_states_by_code(['FL', 'GA', 'MD', 'NC', 'SC', 'VA'], year=2015),\n",
     "    time_period=Year(2015),\n",
     ")"
    ]

doc/demo/01-SIRH-IPM.ipynb

@@ -29,23 +29,33 @@
     "from epymorph.data_shape import Shapes\n",
     "from epymorph.data_type import CentroidDType\n",
     "from epymorph.error import GeoValidationException\n",
-    "from epymorph.geo.spec import LABEL, StaticGeoSpec, Year, attrib\n",
+    "from epymorph.geo.spec import LABEL, StaticGeoSpec, Year\n",
     "from epymorph.geo.static import StaticGeo\n",
     "from epymorph.geo.static import StaticGeoFileOps as F\n",
+    "from epymorph.geography.us_census import StateScope\n",
+    "from epymorph.simulation import AttributeDef\n",
     "\n",
     "YEAR = 2015\n",
     "NUM_STATES = 52\n",
     "NUM_COUNTIES = 3220\n",
     "\n",
-    "spec = StaticGeoSpec(\n",
+    "state_scope = StateScope.all(year=YEAR)\n",
+    "county_scope = state_scope.lower_granularity()\n",
+    "\n",
+    "# Both state and county geo will have the same attributes, just different scope.\n",
+    "state_spec = StaticGeoSpec(\n",
     "    attributes=[\n",
     "        LABEL,\n",
-    "        attrib('geoid', str, Shapes.N),\n",
-    "        attrib('centroid', CentroidDType, Shapes.N),\n",
-    "        attrib('population', int, Shapes.N),\n",
-    "        attrib('commuters', int, Shapes.NxN)\n",
+    "        AttributeDef('geoid', str, Shapes.N),\n",
+    "        AttributeDef('centroid', CentroidDType, Shapes.N),\n",
+    "        AttributeDef('population', int, Shapes.N),\n",
+    "        AttributeDef('commuters', int, Shapes.NxN),\n",
     "    ],\n",
-    "    time_period=Year(YEAR))\n",
+    "    scope=state_scope,\n",
+    "    time_period=Year(YEAR),\n",
+    ")\n",
+    "\n",
+    "county_spec = dataclasses.replace(state_spec, scope=county_scope)\n",
     "\n",
     "# Initialize Census API\n",
     "census = Census(os.environ['CENSUS_API_KEY'])"
@@ -91,7 +101,7 @@
     "        \"wrk_state\",\n",
     "        \"wrk_county\",\n",
     "        \"workers\",\n",
-    "        \"moe\"\n",
+    "        \"moe\",\n",
     "    ],\n",
     "    dtype=str\n",
     ")\n",
@@ -142,7 +152,7 @@
     "d = pd.DataFrame.from_records(state_data).astype({\n",
     "    'NAME': np.str_,\n",
     "    'B01003_001E': np.int64,\n",
-    "    'state': np.str_\n",
+    "    'state': np.str_,\n",
     "})\n",
     "d.rename(columns={\n",
     "    'NAME': 'label',\n",
@@ -170,7 +180,7 @@
     "    'geoid': d['geoid'].to_numpy(dtype=np.str_),\n",
     "    'centroid': d['centroid'].to_numpy(dtype=CentroidDType),\n",
     "    'population': d['population'].to_numpy(dtype=np.int64),\n",
-    "    'commuters': c\n",
+    "    'commuters': c,\n",
     "}\n",
     "\n",
     "num_states = len(states_values['label'])\n",
@@ -186,7 +196,7 @@
    "source": [
     "geofile = Path('epymorph/data/geo') / F.to_archive_filename('us_states_2015')\n",
     "try:\n",
-    "    states_geo = StaticGeo(dataclasses.replace(spec), states_values)\n",
+    "    states_geo = StaticGeo(state_spec, states_values)\n",
     "    states_geo.validate()\n",
     "    states_geo.save(geofile)\n",
     "except GeoValidationException as e:\n",
@@ -201,7 +211,7 @@
     {
      "data": {
       "text/plain": [
-       "<epymorph.geo.static.StaticGeo at 0x7ff85c5ebc50>"
+       "<epymorph.geo.static.StaticGeo at 0x7f0c8a6dc390>"
       ]
      },
      "execution_count": 6,
@@ -286,7 +296,7 @@
    "source": [
     "geofile = Path('epymorph/data/geo') / F.to_archive_filename('us_counties_2015')\n",
     "try:\n",
-    "    counties_geo = StaticGeo(dataclasses.replace(spec), counties_values)\n",
+    "    counties_geo = StaticGeo(county_spec, counties_values)\n",
     "    counties_geo.validate()\n",
     "    counties_geo.save(geofile)\n",
     "except GeoValidationException as e:\n",
@@ -301,7 +311,7 @@
     {
      "data": {
       "text/plain": [
-       "<epymorph.geo.static.StaticGeo at 0x7ff833500290>"
+       "<epymorph.geo.static.StaticGeo at 0x7f0c8f1db850>"
       ]
      },
      "execution_count": 10,

doc/demo/02-states-GEO.ipynb

@@ -26,29 +26,31 @@
     "from epymorph.data_shape import Shapes\n",
     "from epymorph.data_type import CentroidDType\n",
     "from epymorph.error import GeoValidationException\n",
-    "from epymorph.geo.spec import LABEL, StaticGeoSpec, Year, attrib\n",
+    "from epymorph.geo.spec import LABEL, StaticGeoSpec, Year\n",
     "from epymorph.geo.static import StaticGeo\n",
     "from epymorph.geo.static import StaticGeoFileOps as F\n",
+    "from epymorph.geography.scope import ScopeFilter\n",
+    "from epymorph.geography.us_census import BlockGroupScope\n",
+    "from epymorph.simulation import AttributeDef\n",
     "\n",
     "YEAR = 2019\n",
     "NUM_COUNTIES = 2494\n",
     "\n",
-    "spec = StaticGeoSpec(\n",
-    "    attributes=[\n",
-    "        LABEL,\n",
-    "        attrib('geoid', str, Shapes.N),\n",
-    "        attrib('centroid', CentroidDType, Shapes.N),\n",
-    "        attrib('population', int, Shapes.N),\n",
-    "        attrib('population_by_age', int, Shapes.NxA(3)),\n",
-    "        attrib('population_by_age_x6', int, Shapes.NxA(6)),\n",
-    "        attrib('median_age', float, Shapes.N),\n",
-    "        attrib('median_income', int, Shapes.N),\n",
-    "        attrib('average_household_size', float, Shapes.N),\n",
-    "        attrib('pop_density_km2', float, Shapes.N),\n",
-    "        attrib('tract_gini_index', float, Shapes.N),\n",
-    "        attrib('tract_median_income', int, Shapes.N),\n",
-    "    ],\n",
-    "    time_period=Year(YEAR))\n",
+    "attributes: list[AttributeDef] = [\n",
+    "    LABEL,\n",
+    "    AttributeDef('geoid', str, Shapes.N),\n",
+    "    AttributeDef('centroid', CentroidDType, Shapes.N),\n",
+    "    AttributeDef('population', int, Shapes.N),\n",
+    "    # AttributeDef('population_by_age', int, Shapes.NxA(3)),\n",
+    "    # AttributeDef('population_by_age_x6', int, Shapes.NxA(6)),\n",
+    "    AttributeDef('median_age', float, Shapes.N),\n",
+    "    AttributeDef('median_income', int, Shapes.N),\n",
+    "    AttributeDef('average_household_size', float, Shapes.N),\n",
+    "    AttributeDef('pop_density_km2', float, Shapes.N),\n",
+    "    AttributeDef('tract_gini_index', float, Shapes.N),\n",
+    "    AttributeDef('tract_median_income', int, Shapes.N),\n",
+    "]\n",
+    "\n",
     "\n",
     "AGE_VARS = [\n",
     "    \"B01001_003E\",  # Population (Male) 0-4 years\n",
@@ -234,6 +236,26 @@
    "execution_count": 4,
    "metadata": {},
    "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "68      040139411001\n",
+       "870     040139805001\n",
+       "1223    040139804001\n",
+       "1268    040131167331\n",
+       "1444    040131138021\n",
+       "1523    040131134001\n",
+       "1661    040139807001\n",
+       "1808    040137233061\n",
+       "2054    040137233031\n",
+       "2056    040139801001\n",
+       "2273    040130610171\n",
+       "Name: geoid, dtype: object"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
     {
      "data": {
       "text/html": [
@@ -383,10 +405,14 @@
     "cbgs.drop(columns=['tract_geoid'], inplace=True)\n",
     "\n",
     "# Filter CBGs\n",
-    "cbgs.drop(cbgs[\n",
+    "dropped_cbgs = cbgs[\n",
     "    (cbgs['median_income'] == 0) &\n",
     "    (cbgs['tract_median_income'] == 0)\n",
-    "].index, inplace=True)\n",
+    "]\n",
+    "\n",
+    "display(dropped_cbgs['geoid'])\n",
+    "\n",
+    "cbgs.drop(dropped_cbgs.index, inplace=True)\n",
     "\n",
     "cbgs.sort_values(by='geoid', inplace=True)\n",
     "cbgs.reset_index(drop=True, inplace=True)\n",
@@ -547,8 +573,8 @@
     "    'geoid': cbgs['geoid'].to_numpy(dtype=np.str_),\n",
     "    'centroid': cbgs['centroid'].to_numpy(dtype=CentroidDType),\n",
     "    'population': cbgs['population'].to_numpy(dtype=np.int64),\n",
-    "    'population_by_age': cbgs_age_1.to_numpy(dtype=np.int64),\n",
-    "    'population_by_age_x6': cbgs_age_2.to_numpy(dtype=np.int64),\n",
+    "    # 'population_by_age': cbgs_age_1.to_numpy(dtype=np.int64),\n",
+    "    # 'population_by_age_x6': cbgs_age_2.to_numpy(dtype=np.int64),\n",
     "    'median_age': cbgs['median_age'].to_numpy(dtype=np.float64),\n",
     "    'median_income': cbgs['median_income'].to_numpy(dtype=np.int64),\n",
     "    'average_household_size': cbgs['average_household_size'].to_numpy(dtype=np.float64),\n",
@@ -568,6 +594,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "spec = StaticGeoSpec(\n",
+    "    attributes=attributes,\n",
+    "    # Maricopa County, AZ is GEOID 04013\n",
+    "    scope=ScopeFilter(\n",
+    "        parent=BlockGroupScope.in_counties(['04013'], year=YEAR),\n",
+    "        remove=dropped_cbgs['geoid'].to_numpy(np.str_),\n",
+    "    ),\n",
+    "    time_period=Year(YEAR),\n",
+    ")\n",
+    "\n",
     "try:\n",
     "    geo = StaticGeo(spec, values)\n",
     "    geo.validate()\n",
@@ -601,7 +637,6 @@
       "label: No diffs!\n",
       "population: No diffs!\n",
       "median_age: No diffs!\n",
-      "population_by_age_x6: No diffs!\n",
       "median_income: No diffs!\n",
       "average_household_size: No diffs!\n",
       "tract_gini_index: No diffs!\n",
@@ -650,8 +685,9 @@
     "diff('population', d1['population'], d2['population'], np.int_.__eq__)\n",
     "diff('median_age', d1['median_age'], d2['median_age'], np.isclose)\n",
     "# there's no equivalent in the old geo to our new 'population_by_age'\n",
-    "diff('population_by_age_x6', d1['population_by_age_x6'],\n",
-    "     d2['pop_by_age'], np.array_equal)\n",
+    "# NOTE: we can't check population_by_age_x6 anymore...\n",
+    "# diff('population_by_age_x6', d1['population_by_age_x6'],\n",
+    "#      d2['pop_by_age'], np.array_equal)\n",
     "diff('median_income', d1['median_income'], d2['median_income'], np.int_.__eq__)\n",
     "diff('average_household_size', d1['average_household_size'],\n",
     "     d2['average_household_size'], np.isclose)\n",