Demo vignettes and a few related changes. (#87)

Vignettes originally created for Joe's Italy conference in late 2023.
Changes the SIRH IPM param `hospitalization_rate` to `hospitalization_prob`.
Adds some utility cache/plots functions.
Adds star exports to `epymorph.geo`
Fixes a minor potential bug in Census ADRIO.
Adds a few more numpy functions to the epymorph function namespace.

epymorph/data/ipm/sirh.py

@@ -22,12 +22,12 @@ def load() -> CompartmentModel:
             param('beta', shape=Shapes.TxN),
             param('gamma', shape=Shapes.TxN),
             param('xi', shape=Shapes.TxN),
-            param('hospitalization_rate', shape=Shapes.TxN),
+            param('hospitalization_prob', shape=Shapes.TxN),
             param('hospitalization_duration', shape=Shapes.TxN)
         ])
 
     [S, I, R, H] = symbols.compartment_symbols
-    [β, γ, ξ, h_rate, h_dur] = symbols.attribute_symbols
+    [β, γ, ξ, h_prob, h_dur] = symbols.attribute_symbols
 
     # formulate N so as to avoid dividing by zero;
     # this is safe in this instance because if the denominator is zero,
@@ -39,8 +39,8 @@ def load() -> CompartmentModel:
         transitions=[
             edge(S, I, rate=β * S * I / N),
             fork(
-                edge(I, H, rate=γ * I * h_rate),
-                edge(I, R, rate=γ * I * (1 - h_rate)),
+                edge(I, H, rate=γ * I * h_prob),
+                edge(I, R, rate=γ * I * (1 - h_prob)),
             ),
             edge(H, R, rate=H / h_dur),
             edge(R, S, rate=ξ * R),

epymorph/geo/__init__.py

@@ -0,0 +1,25 @@
+"""epymorph's geo package and exports"""
+from epymorph.geo.spec import (AttribDef, CentroidDType, DateAndDuration,
+                               DateRange, DynamicGeoSpec, Geography, GeoSpec,
+                               NonspecificDuration, SpecificTimePeriod,
+                               StaticGeoSpec, TimePeriod, Year)
+
+__all__ = [
+    'AttribDef',
+    'CentroidDType',
+    'DateAndDuration',
+    'DateRange',
+    'DynamicGeoSpec',
+    'Geography',
+    'GeoSpec',
+    'NonspecificDuration',
+    'SpecificTimePeriod',
+    'StaticGeoSpec',
+    'TimePeriod',
+    'Year',
+    # TODO: these imports cause circular deps
+    # 'DynamicGeo',
+    # 'StaticGeo',
+    # 'save_to_cache',
+    # 'load_from_cache',
+]

epymorph/geo/adrio/census/adrio_census.py

@@ -7,7 +7,7 @@
 from census import Census
 from geopandas import GeoDataFrame
 from numpy.typing import NDArray
-from pandas import DataFrame, read_excel
+from pandas import DataFrame, concat, read_excel
 from pygris import block_groups, counties, states, tracts
 
 from epymorph.data_shape import Shapes
@@ -228,7 +228,10 @@ def fetch_sf(self, granularity: int, nodes: dict[str, list[str]], year: int) ->
             sort_param = ['state']
 
         elif granularity == Granularity.COUNTY.value:
-            data_df = counties(state=state_fips, year=year)
+            if state_fips is None or '*' in state_fips:
+                data_df = counties(year=year)
+            else:
+                data_df = counties(state=state_fips, year=year)
             data_df = data_df.rename(columns={'STATEFP': 'state', 'COUNTYFP': 'county'})
 
             if county_fips is not None and county_fips[0] != '*':

epymorph/geo/cache.py

@@ -4,7 +4,9 @@
 from platformdirs import user_cache_path
 
 from epymorph.data import geo_library_dynamic
+from epymorph.geo.dynamic import DynamicGeo
 from epymorph.geo.geo import Geo
+from epymorph.geo.static import StaticGeo
 from epymorph.geo.static import StaticGeoFileOps as F
 from epymorph.geo.util import convert_to_static_geo
 
@@ -20,12 +22,12 @@ def fetch(geo_name: str) -> None:
     Caches all attribute data for a dynamic geo from the library.
     Raises GeoCacheException if spec not found.
     """
-    filepath = CACHE_PATH / F.to_archive_filename(geo_name)
+    file_path = CACHE_PATH / F.to_archive_filename(geo_name)
     geo_load = geo_library_dynamic.get(geo_name)
     if geo_load is not None:
         geo = geo_load()
         static_geo = convert_to_static_geo(geo)
-        static_geo.save(filepath)
+        F.save_as_archive(static_geo, file_path)
     else:
         raise GeoCacheException(f'spec file for {geo_name} not found.')
 
@@ -35,9 +37,9 @@ def remove(geo_name: str) -> None:
     Removes a geo's data from the cache.
     Raises GeoCacheException if geo not found in cache.
     """
-    filepath = CACHE_PATH / F.to_archive_filename(geo_name)
-    if os.path.exists(filepath):
-        os.remove(filepath)
+    file_path = CACHE_PATH / F.to_archive_filename(geo_name)
+    if os.path.exists(file_path):
+        os.remove(file_path)
     else:
         msg = f'{geo_name} not found in cache, check your spelling or use the list subcommand to view all currently cached geos'
         raise GeoCacheException(msg)
@@ -55,6 +57,19 @@ def clear():
         os.remove(CACHE_PATH / file[0])
 
 
+def save_to_cache(geo: Geo, geo_name: str) -> None:
+    """Save a Geo to the cache (if you happen to already have it as a Geo object)."""
+    match geo:
+        case DynamicGeo():
+            static_geo = convert_to_static_geo(geo)
+        case StaticGeo():
+            static_geo = geo
+        case _:
+            raise GeoCacheException('Unable to cache given geo.')
+    file_path = CACHE_PATH / F.to_archive_filename(geo_name)
+    F.save_as_archive(static_geo, file_path)
+
+
 def load_from_cache(geo_name: str) -> Geo | None:
     """Checks whether a dynamic geo has already been cached and returns it if so."""
     file_path = CACHE_PATH / F.to_archive_filename(geo_name)

epymorph/plots.py

@@ -1,8 +1,14 @@
 """Our built-in plotting functions and helpful utilities."""
+from typing import Any, Literal
+
 import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
+import pygris
+from numpy.typing import NDArray
 
 from epymorph.engine.standard_sim import Output
+from epymorph.geo.geo import Geo
 
 
 def plot_event(out: Output, event_idx: int) -> None:
@@ -43,3 +49,87 @@ def plot_pop(out: Output, pop_idx: int, log_scale: bool = False) -> None:
 
     fig.tight_layout()
     plt.show()
+
+
+def map_data_by_county(
+    geo: Geo,
+    data: NDArray,
+    *,
+    df_borders: pd.DataFrame | None = None,
+    title: str,
+    cmap: Any | None = None,
+    vmin: float | None = None,
+    vmax: float | None = None,
+    outline: Literal['states', 'counties'] = 'counties',
+) -> None:
+    """
+    Draw a county-level choropleth map using the given `data`. This must be a numpy array whose
+    ordering is the same as the nodes in the geo.
+    Assumes that the geo contains an attribute (`geoid`) containing the geoids of its nodes.
+    (This information is needed to fetch the map shapes.)
+    """
+    state_fips = {s[:2] for s in geo['geoid']}
+    df_states = pygris.states(cb=True, resolution='5m', cache=True, year=2020)
+    df_states = df_states.loc[df_states['GEOID'].isin(state_fips)]
+    df_counties = pd.concat([
+        pygris.counties(state=s, cb=True, resolution='5m', cache=True, year=2020)
+        for s in state_fips
+    ])
+    if outline == 'counties':
+        df_borders = df_counties
+    else:
+        df_borders = df_states
+    return _map_data_by_geo(geo, data, df_counties, df_borders=df_borders, title=title, cmap=cmap, vmin=vmin, vmax=vmax)
+
+
+def map_data_by_state(
+    geo: Geo,
+    data: NDArray,
+    *,
+    title: str,
+    cmap: Any | None = None,
+    vmin: float | None = None,
+    vmax: float | None = None,
+) -> None:
+    """
+    Draw a state-level choropleth map using the given `data`. This must be a numpy array whose
+    ordering is the same as the nodes in the geo.
+    Assumes that the geo contains an attribute (`geoid`) containing the geoids of its nodes.
+    (This information is needed to fetch the map shapes.)
+    """
+    state_fips = {s[:2] for s in geo['geoid']}
+    df_states = pygris.states(cb=True, resolution='5m', cache=True, year=2020)
+    df_states = df_states.loc[df_states['GEOID'].isin(state_fips)]
+    return _map_data_by_geo(geo, data, df_states, title=title, cmap=cmap, vmin=vmin, vmax=vmax)
+
+
+def _map_data_by_geo(
+    geo: Geo,
+    data: NDArray,
+    df_nodes: pd.DataFrame,
+    *,
+    df_borders: pd.DataFrame | None = None,
+    title: str,
+    cmap: Any | None = None,
+    vmin: float | None = None,
+    vmax: float | None = None,
+) -> None:
+    """
+    Draw a choropleth map for the geo information given.
+    This is an internal function to abstract commonalities -- see `map_data_by_county()` and `map_data_by_state()`.
+    """
+    df_merged = pd.merge(
+        on="GEOID",
+        left=df_nodes,
+        right=pd.DataFrame({'GEOID': geo['geoid'], 'data': data}),
+    )
+
+    fig, ax = plt.subplots(figsize=(8, 6))
+    ax.axis('off')
+    ax.set_title(title)
+    df_merged.plot(ax=ax, column='data', legend=True, cmap=cmap, vmin=vmin, vmax=vmax)
+    if df_borders is None:
+        df_borders = df_nodes
+    df_borders.plot(ax=ax, linewidth=1, edgecolor='black', color='none', alpha=0.8)
+    fig.tight_layout()
+    plt.show()

epymorph/simulation.py

@@ -225,7 +225,11 @@ def epymorph_namespace() -> dict[str, Any]:
             'array': partial(np.array, dtype=SimDType),
             'zeros': partial(np.zeros, dtype=SimDType),
             'zeros_like': partial(np.zeros_like, dtype=SimDType),
+            'ones': partial(np.ones, dtype=SimDType),
+            'ones_like': partial(np.ones_like, dtype=SimDType),
             'full': partial(np.full, dtype=SimDType),
+            'arange': partial(np.arange, dtype=SimDType),
+            'concatenate': partial(np.concatenate, dtype=SimDType),
             'sum': partial(np.sum, dtype=SimDType),
             'newaxis': np.newaxis,
             # numpy math functions