Geo messaging (#88)

A new dynamic_geo_messaging() context was added to provide verbose geo console messaging, for instance when doing a geo.fetch_all().
The sim_messaging() context now has an option to enable verbose geo console messaging (if applicable).
Also fixes a minor bug fetching county-level commuter data from ACS Commuter Flows.

---------

Co-authored-by: Tyler Coles <tylercoles@gmail.com>

doc/devlog/2024-02-06-adrio-demo.ipynb

@@ -0,0 +1,213 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from epymorph.geo.spec import DynamicGeoSpec, AttribDef, CentroidDType, Year\n",
+    "from epymorph.geo.adrio.census.adrio_census import CensusGeography, Granularity\n",
+    "from epymorph.data_shape import Shapes\n",
+    "from pathlib import Path\n",
+    "import numpy as np\n",
+    "\n",
+    "maricopa = DynamicGeoSpec(\n",
+    "        attributes=[\n",
+    "            AttribDef('label', np.str_, Shapes.N),\n",
+    "            AttribDef('geoid', np.str_, Shapes.N),\n",
+    "            AttribDef('centroid', CentroidDType, Shapes.N),\n",
+    "            AttribDef('median_income', np.int64, Shapes.N),\n",
+    "            AttribDef('tract_median_income', np.int64, Shapes.N),\n",
+    "            AttribDef('population', np.int64, Shapes.N),\n",
+    "            AttribDef('population_by_age', np.int64, Shapes.NxA(3)),\n",
+    "            AttribDef('population_by_age_x6', np.int64, Shapes.NxA(6)),\n",
+    "            AttribDef('pop_density_km2', np.float64, Shapes.N),\n",
+    "            AttribDef('median_age', np.int64, Shapes.N),\n",
+    "            AttribDef('average_household_size', np.int64, Shapes.N),\n",
+    "            AttribDef('gini_index', np.float64, Shapes.N)\n",
+    "        ],\n",
+    "        geography=CensusGeography(granularity=Granularity.CBG, filter={\n",
+    "            'state': ['04'],\n",
+    "            'county': ['013'],\n",
+    "            'tract': ['*'],\n",
+    "            'block group': ['*']\n",
+    "        }),\n",
+    "        time_period=Year(2019),\n",
+    "        source={\n",
+    "            'label': 'Census:name',\n",
+    "            'population': 'Census',\n",
+    "            'geoid': 'Census',\n",
+    "            'centroid': 'Census',\n",
+    "            'median_income': 'Census',\n",
+    "            'tract_median_income': 'Census',\n",
+    "            'population_by_age': 'Census', \n",
+    "            'population_by_age_x6': 'Census',\n",
+    "            'pop_density_km2': 'Census',\n",
+    "            'median_age': 'Census', \n",
+    "            'average_household_size': 'Census',\n",
+    "            'gini_index': 'Census'\n",
+    "        }\n",
+    "    )\n",
+    "\n",
+    "json = maricopa.serialize()\n",
+    "with open(Path('./scratch/geo/maricopa_cbg_2019.geo'), mode='w', encoding='utf-8') as f:\n",
+    "    f.write(json)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from epymorph.geo.dynamic import DynamicGeoFileOps\n",
+    "from epymorph.geo.adrio import adrio_maker_library\n",
+    "\n",
+    "geo = DynamicGeoFileOps.load_from_spec(Path('./scratch/geo/maricopa_cbg_2019.geo'), adrio_maker_library)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Fetching dynamic geo data\n",
+      "• 12 attributes\n",
+      "Fetching name...[1/12]\n",
+      "Fetching geoid...[2/12]\n",
+      "Fetching centroid...[3/12]\n",
+      "Fetching median_income...[4/12]\n",
+      "Fetching tract_median_income...[5/12]\n",
+      "Fetching population...[6/12]\n",
+      "Fetching population_by_age...[7/12]\n",
+      "Fetching population_by_age_x6...[8/12]\n",
+      "Fetching pop_density_km2...[9/12]\n",
+      "Fetching median_age...[10/12]\n",
+      "Fetching average_household_size...[11/12]\n",
+      "Fetching gini_index...[12/12]\n",
+      "Gini Index cannot be retrieved for block group level, fetching tract level data instead.\n",
+      "Complete.\n",
+      "Total fetch time: 27.819s\n"
+     ]
+    }
+   ],
+   "source": [
+    "from epymorph.logging.messaging import dynamic_geo_messaging\n",
+    "\n",
+    "with dynamic_geo_messaging(geo):\n",
+    "    geo.fetch_all()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Loading requirements:\n",
+      "[✓] IPM (sirs)\n",
+      "[✓] MM (centroids)\n",
+      "[✓] GEO (us_sw_counties_2015)\n",
+      "Running simulation (StandardSimulation):\n",
+      "• 2010-01-01 to 2010-05-31 (150 days)\n",
+      "• 158 geo nodes\n",
+      "|####################| 100% \n",
+      "Runtime: 12.335s\n",
+      "Done\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "0"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from epymorph.cli.run import run\n",
+    "\n",
+    "run(input_path='./scratch/adrio_test.toml',\n",
+    "    out_path=None,\n",
+    "    chart=None,\n",
+    "    profiling=False,\n",
+    "    ignore_cache=False,\n",
+    "    geo_messaging=True\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Fetching dynamic geo data\n",
+      "• 8 attributes\n",
+      "Fetching name...[1/8]\n",
+      "Fetching population...[2/8]\n",
+      "Fetching population_by_age...[3/8]\n",
+      "Fetching centroid...[4/8]\n",
+      "Fetching geoid...[5/8]\n",
+      "Fetching dissimilarity_index...[6/8]\n",
+      "Fetching median_income...[7/8]\n",
+      "Fetching pop_density_km2...[8/8]\n",
+      "Complete.\n",
+      "Total fetch time: 22.568s\n",
+      "geo sucessfully cached.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "0"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from epymorph.cli.cache import fetch\n",
+    "\n",
+    "fetch('us_sw_counties_2015', True)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

doc/devlog/README.md

@@ -41,6 +41,7 @@ This folder is a handy place to put Jupyter notebooks or other documents which h
 | 2023-11-22-ipm-probs.ipynb | Tyler | | Analyzing statistical correctness of our IPM processing algorithms. |
 | 2023-12-05.ipynb | Tyler | | A brief tour of changes to epymorph due to the refactor effort. |
 | 2024-01-08.ipynb | Tyler | | Another functional parameters demonstration, revisiting the Bonus Example from 2023-10-10. |
+| 2024-02-06-adrio-demo.ipynb | Trevor | | Demonstrates the ADRIO system using code updated for latest changes. |
 | 2024-02-06.ipynb | Tyler | | Revisiting age-class IPMs, and thinking about modularity of approach. |
 | 2024-02-12.ipynb | Tyler | | Continued age-class IPM work, this time in more than one geo node. |
 

epymorph/__init__.py

@@ -4,9 +4,10 @@
 from epymorph.data import geo_library, ipm_library, mm_library
 from epymorph.data_shape import Shapes
 from epymorph.engine.standard_sim import StandardSimulation
+from epymorph.logging.messaging import sim_messaging
 from epymorph.plots import plot_event, plot_pop
 from epymorph.proxy import dim, geo
-from epymorph.simulation import SimDType, TimeFrame, default_rng, sim_messaging
+from epymorph.simulation import SimDType, TimeFrame, default_rng
 
 __all__ = [
     'IPM',

epymorph/cli/cache.py

@@ -31,6 +31,10 @@ def define_argparser(command_parser: _SubParsersAction):
         'geo',
         type=str,
         help='the name of the geo to fetch; must include a geo path if not already in the library')
+    fetch_command.add_argument(
+        '-p', '--path',
+        help='(optional) the path to a geo spec file not in the library'
+    )
     fetch_command.add_argument(
         '-f', '--force',
         action='store_true',

epymorph/cli/run.py

@@ -24,9 +24,9 @@
 from epymorph.geo.geo import Geo
 from epymorph.geo.static import StaticGeoFileOps
 from epymorph.initializer import initializer_library, normalize_init_params
+from epymorph.logging.messaging import sim_messaging
 from epymorph.movement.parser import MovementSpec, parse_movement_spec
-from epymorph.simulation import (TimeFrame, default_rng, enable_logging,
-                                 sim_messaging)
+from epymorph.simulation import TimeFrame, default_rng, enable_logging
 
 
 def define_argparser(command_parser: _SubParsersAction):
@@ -52,14 +52,19 @@ def define_argparser(command_parser: _SubParsersAction):
     p.add_argument(
         '-i', '--ignore_cache',
         action='store_true',
-        help='(optional) include this flag to run the simulation without utilizing the Geo cache.'
-    )
+        help='(optional) include this flag to run the simulation without utilizing the Geo cache.')
+    p.add_argument(
+        '-m', '--mute_geo',
+        action='store_false',
+        help='(optional) include this flag to silence geo data retreival messaging.')
+
     p.set_defaults(handler=lambda args: run(
         input_path=args.input,
         out_path=args.out,
         chart=args.chart,
         profiling=args.profile,
-        ignore_cache=args.ignore_cache
+        ignore_cache=args.ignore_cache,
+        geo_messaging=args.mute_geo
     ))
 
 
@@ -79,7 +84,8 @@ def run(input_path: str,
         out_path: str | None,
         chart: str | None,
         profiling: bool,
-        ignore_cache: bool) -> int:
+        ignore_cache: bool,
+        geo_messaging: bool) -> int:
     """CLI command handler: run a simulation."""
 
     # Exit codes:
@@ -150,7 +156,9 @@ def run(input_path: str,
     if not profiling:
         enable_logging()
 
-    with sim_messaging(sim):
+    # Run simulation with appropriate messaging contexts
+
+    with sim_messaging(sim, geo_messaging):
         out = sim.run()
 
     # Draw charts (if specified).

epymorph/engine/standard_sim.py

@@ -93,6 +93,7 @@ class StandardSimulation(SimulationEvents):
 
     _config: RumeConfig
     _params: ContextParams | None = None
+    geo: Geo
     on_tick: Event[SimTick]  # this class supports on_tick; so narrow the type def
 
     def __init__(self,
@@ -103,6 +104,7 @@ def __init__(self,
                  time_frame: TimeFrame,
                  initializer: Initializer | None = None,
                  rng: Callable[[], np.random.Generator] | None = None):
+        self.geo = geo
         if initializer is None:
             initializer = DEFAULT_INITIALIZER
         if rng is None:

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, concat, read_excel
+from pandas import DataFrame, read_excel
 from pygris import block_groups, counties, states, tracts
 
 from epymorph.data_shape import Shapes
@@ -370,10 +370,6 @@ def fetch_commuters(self, granularity: int, nodes: dict[str, list[str]], year: i
                 data = data.loc[data['wrk_state_code'] < '057']
                 data = data.loc[data['wrk_state_code'] != '011']
 
-            # filter out non-county locations
-            data = data.loc[data['res_county_code'] < '508']
-            data = data.loc[data['wrk_county_code'] < '508']
-
         if granularity == Granularity.COUNTY.value:
             if counties is not None and counties[0] != '*':
                 data = data.loc[data['res_county_code'].isin(counties)]

epymorph/geo/cache.py

@@ -11,6 +11,7 @@
 from epymorph.geo.static import StaticGeo
 from epymorph.geo.static import StaticGeoFileOps as F
 from epymorph.geo.util import convert_to_static_geo
+from epymorph.logging.messaging import dynamic_geo_messaging
 
 CACHE_PATH = user_cache_path(appname='epymorph', ensure_exists=True)
 
@@ -31,7 +32,8 @@ def fetch(geo_name_or_path: str) -> None:
         geo_load = geo_library_dynamic.get(geo_name_or_path)
         if geo_load is not None:
             geo = geo_load()
-            static_geo = convert_to_static_geo(geo)
+            with dynamic_geo_messaging(geo):
+                static_geo = convert_to_static_geo(geo)
             static_geo.save(file_path)
 
     # checks for geo spec at given path (path passed)
@@ -41,7 +43,8 @@ def fetch(geo_name_or_path: str) -> None:
             geo_name = geo_path.stem
             file_path = CACHE_PATH / F.to_archive_filename(geo_name)
             geo = DF.load_from_spec(geo_path, adrio_maker_library)
-            static_geo = convert_to_static_geo(geo)
+            with dynamic_geo_messaging(geo):
+                static_geo = convert_to_static_geo(geo)
             static_geo.save(file_path)
         else:
             raise GeoCacheException(f'spec file at {geo_name_or_path} not found.')