Class-based syntax revamp: IPM/MM/ADRIOs

doc/devlog/2024-06-03.ipynb

@@ -8,9 +8,7 @@
     "\n",
     "_author: Trevor Johnson_\n",
     "\n",
-    "ADRIOMakerCensus has been refactored to utilize the recently added GeoScope class heirarchy. This notebook tests the correct functionality of Census-based ADRIOs post-refactor by creating a DynamicGeo for each granularity and populating them with every attribute that is valid for their granularity.\n",
-    "\n",
-    "Additional cases can be tested by changing the type of the scope objects, changing the year in time period or scope, or changing the includes attribute of the scope object."
+    "Tests ACS5 ADRIOs at a variety of granularities."
    ]
   },
   {
@@ -19,185 +17,233 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from epymorph.data_shape import Shapes\n",
-    "from epymorph.data_type import CentroidType\n",
-    "from epymorph.geo.adrio import adrio_maker_library\n",
-    "from epymorph.geo.dynamic import DynamicGeo\n",
-    "from epymorph.geo.spec import DynamicGeoSpec, Year\n",
-    "from epymorph.geography.us_census import DEFAULT_YEAR, StateScopeAll\n",
-    "from epymorph.simulation import AttributeDef\n",
+    "from epymorph import *\n",
+    "from epymorph.adrio import acs5, adrio, commuting_flows, us_tiger\n",
+    "from epymorph.data_shape import DataShapeMatcher\n",
+    "from epymorph.geography.us_census import (BlockGroupScope, CountyScope,\n",
+    "                                          StateScope, TractScope)\n",
+    "from epymorph.params import ParamValue\n",
+    "from epymorph.simulator.data import evaluate_param\n",
+    "from epymorph.util import NumpyTypeError, check_ndarray, match\n",
+    "\n",
+    "# This is the expected type and shape for every attribute we're going to test.\n",
+    "expected: list[AttributeDef] = [\n",
+    "    AttributeDef(\"label\", str, Shapes.N),\n",
+    "    AttributeDef(\"population\", int, Shapes.N),\n",
+    "    AttributeDef(\"population_by_age_table\", int, Shapes.NxA),\n",
+    "    AttributeDef(\"population_by_age\", int, Shapes.N),\n",
+    "    AttributeDef(\"average_household_size\", float, Shapes.N),\n",
+    "    AttributeDef(\"dissimilarity_index\", float, Shapes.N),\n",
+    "    AttributeDef(\"commuters\", int, Shapes.NxN),\n",
+    "    AttributeDef(\"gini_index\", float, Shapes.N),\n",
+    "    AttributeDef(\"median_age\", float, Shapes.N),\n",
+    "    AttributeDef(\"median_income\", float, Shapes.N),\n",
+    "    AttributeDef(\"pop_density_km2\", float, Shapes.N),\n",
+    "]\n",
+    "\n",
+    "# And here are the ADRIOs for each of those attributes.\n",
+    "params: dict[str, ParamValue] = {\n",
+    "    \"label\": us_tiger.Name(),\n",
+    "    \"population\": acs5.Population(),\n",
+    "    \"population_by_age_table\": acs5.PopulationByAgeTable(),\n",
+    "    \"population_by_age\": acs5.PopulationByAge(18, 24),\n",
+    "    \"average_household_size\": acs5.AverageHouseholdSize(),\n",
+    "    \"dissimilarity_index\": acs5.DissimilarityIndex(\"White\", \"Black\"),\n",
+    "    \"commuters\": commuting_flows.Commuters(),\n",
+    "    \"gini_index\": acs5.GiniIndex(),\n",
+    "    \"median_age\": acs5.MedianAge(),\n",
+    "    \"median_income\": acs5.MedianIncome(),\n",
+    "    \"land_area_km2\": adrio.Scale(us_tiger.LandAreaM2(), 1e-6),\n",
+    "    \"pop_density_km2\": adrio.PopulationPerKm2(),\n",
+    "}\n",
     "\n",
-    "spec = DynamicGeoSpec(\n",
-    "    attributes=[\n",
-    "        AttributeDef('label', str, 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('centroid', CentroidType, Shapes.N),\n",
-    "        AttributeDef('geoid', str, Shapes.N),\n",
-    "        AttributeDef('average_household_size', int, Shapes.N),\n",
-    "        AttributeDef('dissimilarity_index', float, Shapes.N),\n",
-    "        AttributeDef('commuters', int, Shapes.NxN),\n",
-    "        AttributeDef('gini_index', float, Shapes.N),\n",
-    "        AttributeDef('median_age', int, Shapes.N),\n",
-    "        AttributeDef('median_income', int, Shapes.N),\n",
-    "        AttributeDef('pop_density_km2', float, Shapes.N)\n",
-    "    ],\n",
-    "    time_period=Year(2020),\n",
-    "    scope=StateScopeAll(DEFAULT_YEAR),\n",
-    "    source={\n",
-    "        'label': 'Census:name',\n",
-    "        'population': 'Census',\n",
-    "        # 'population_by_age': 'Census',\n",
-    "        # 'population_by_age_x6': 'Census',\n",
-    "        'centroid': 'Census',\n",
-    "        'geoid': 'Census',\n",
-    "        'average_household_size': 'Census',\n",
-    "        'dissimilarity_index': 'Census',\n",
-    "        'commuters': 'Census',\n",
-    "        'gini_index': 'Census',\n",
-    "        'median_age': 'Census',\n",
-    "        'median_income': 'Census',\n",
-    "        'pop_density_km2': 'Census',\n",
-    "        'tract_median_income': 'Census'\n",
-    "    }\n",
-    ")\n",
     "\n",
-    "geo = DynamicGeo.from_library(spec, adrio_maker_library)"
+    "def run_test(rume: Rume, skip: tuple[str, ...] = ()):\n",
+    "    for attr in (a for a in expected if a.name not in skip):\n",
+    "        actual = evaluate_param(rume, attr.name)\n",
+    "        try:\n",
+    "            check_ndarray(\n",
+    "                actual,\n",
+    "                dtype=match.dtype(attr.dtype),\n",
+    "                shape=DataShapeMatcher(attr.shape, rume.dim, True),\n",
+    "            )\n",
+    "            print(f\"{attr.name}: good\")\n",
+    "        except NumpyTypeError as e:\n",
+    "            print(f\"{attr.name}: FAILED\")\n",
+    "            print(e)\n",
+    "\n",
+    "\n",
+    "def placeholder_rume(scope, time_frame):\n",
+    "    return SingleStrataRume.build(\n",
+    "        ipm=ipm_library['no'](),\n",
+    "        mm=mm_library['no'](),\n",
+    "        init=init.NoInfection(),\n",
+    "        scope=scope,\n",
+    "        time_frame=time_frame,\n",
+    "        params=params\n",
+    "    )"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 2,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "label: good\n",
+      "population: good\n",
+      "population_by_age_table: good\n",
+      "population_by_age: good\n",
+      "average_household_size: good\n",
+      "dissimilarity_index: good\n",
+      "commuters: good\n",
+      "gini_index: good\n",
+      "median_age: good\n",
+      "median_income: good\n",
+      "pop_density_km2: good\n"
+     ]
+    }
+   ],
    "source": [
-    "geo.validate()"
+    "rume = placeholder_rume(\n",
+    "    scope=StateScope.all(year=2020),\n",
+    "    time_frame=TimeFrame.year(2020),\n",
+    ")\n",
+    "\n",
+    "run_test(rume)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 3,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "label: good\n",
+      "population: good\n",
+      "population_by_age_table: good\n",
+      "population_by_age: good\n",
+      "average_household_size: good\n",
+      "dissimilarity_index: good\n",
+      "commuters: good\n",
+      "gini_index: good\n",
+      "median_age: good\n",
+      "median_income: good\n",
+      "pop_density_km2: good\n"
+     ]
+    }
+   ],
    "source": [
-    "from dataclasses import replace\n",
-    "\n",
-    "from epymorph.geography.us_census import (BlockGroupScope, CountyScope,\n",
-    "                                          StateScope, TractScope)\n",
+    "rume = placeholder_rume(\n",
+    "    scope=StateScope.in_states(['04', '08'], year=2020),\n",
+    "    time_frame=TimeFrame.year(2020),\n",
+    ")\n",
     "\n",
-    "spec = replace(spec, scope=StateScope.in_states(['04', '08']))\n",
-    "geo = DynamicGeo.from_library(spec, adrio_maker_library)"
+    "run_test(rume)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 4,
    "metadata": {},
-   "outputs": [],
-   "source": [
-    "geo.validate()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "spec = replace(spec, scope=CountyScope.in_counties(['35001', '04013', '04017']))\n",
-    "geo = DynamicGeo.from_library(spec, adrio_maker_library)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "geo.validate()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "label: good\n",
+      "population: good\n",
+      "population_by_age_table: good\n",
+      "population_by_age: good\n",
+      "average_household_size: good\n",
+      "dissimilarity_index: good\n",
+      "commuters: good\n",
+      "gini_index: good\n",
+      "median_age: good\n",
+      "median_income: good\n",
+      "pop_density_km2: good\n"
+     ]
+    }
+   ],
    "source": [
-    "spec = replace(spec, scope=TractScope.in_tracts(['35001000720', '35001000904', '35001000906',\n",
-    "                                                 '04027011405', '04027011407']), attributes=[\n",
-    "    AttributeDef('label', str, 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('centroid', CentroidType, Shapes.N),\n",
-    "    AttributeDef('geoid', str, Shapes.N),\n",
-    "    AttributeDef('average_household_size', int, Shapes.N),\n",
-    "    AttributeDef('dissimilarity_index', float, Shapes.N),\n",
-    "    AttributeDef('gini_index', float, Shapes.N),\n",
-    "    AttributeDef('median_age', int, Shapes.N),\n",
-    "    AttributeDef('median_income', int, Shapes.N),\n",
-    "    AttributeDef('pop_density_km2', float, Shapes.N)\n",
-    "])\n",
+    "rume = placeholder_rume(\n",
+    "    scope=CountyScope.in_counties(['35001', '04013', '04017'], year=2020),\n",
+    "    time_frame=TimeFrame.year(2020),\n",
+    ")\n",
     "\n",
-    "geo = DynamicGeo.from_library(spec, adrio_maker_library)"
+    "run_test(rume)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# tract and block group geos fetch shape file attributes prior to validating so that the kernel\n",
-    "# is not overwhelmed by several large shapefile requests in parallel\n",
-    "geo['centroid']\n",
-    "geo['pop_density_km2']\n",
-    "geo.validate()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 5,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "label: good\n",
+      "population: good\n",
+      "population_by_age_table: good\n",
+      "population_by_age: good\n",
+      "average_household_size: good\n",
+      "dissimilarity_index: good\n",
+      "gini_index: good\n",
+      "median_age: good\n",
+      "median_income: good\n",
+      "pop_density_km2: good\n"
+     ]
+    }
+   ],
    "source": [
-    "spec = replace(spec, scope=BlockGroupScope.in_block_groups(['350010007201', '350010009041', '350010009061',\n",
-    "                                                            '040270114053', '040270114072']), attributes=[\n",
-    "    AttributeDef('label', str, 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('centroid', CentroidType, Shapes.N),\n",
-    "    AttributeDef('geoid', str, Shapes.N),\n",
-    "    AttributeDef('average_household_size', int, Shapes.N),\n",
-    "    AttributeDef('gini_index', float, Shapes.N),\n",
-    "    AttributeDef('median_age', int, Shapes.N),\n",
-    "    AttributeDef('median_income', int, Shapes.N),\n",
-    "    AttributeDef('pop_density_km2', float, Shapes.N),\n",
-    "    AttributeDef('tract_median_income', int, Shapes.N)\n",
-    "])\n",
+    "rume = placeholder_rume(\n",
+    "    scope=TractScope.in_tracts([\n",
+    "        '35001000720', '35001000904', '35001000906', '04027011405', '04027011407'\n",
+    "    ], year=2020),\n",
+    "    time_frame=TimeFrame.year(2020),\n",
+    ")\n",
     "\n",
-    "geo = DynamicGeo.from_library(spec, adrio_maker_library)"
+    "run_test(rume, skip=(\"commuters\",))"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Gini Index cannot be retrieved for block group level, fetching tract level data instead.\n"
+      "label: good\n",
+      "population: good\n",
+      "population_by_age_table: good\n",
+      "population_by_age: good\n",
+      "average_household_size: good\n",
+      "Gini Index cannot be retrieved for block group level, fetching tract level data instead.\n",
+      "gini_index: good\n",
+      "median_age: good\n",
+      "median_income: good\n",
+      "pop_density_km2: good\n"
      ]
     }
    ],
    "source": [
-    "geo['centroid']\n",
-    "geo['pop_density_km2']\n",
-    "geo.validate()"
+    "rume = placeholder_rume(\n",
+    "    scope=BlockGroupScope.in_block_groups([\n",
+    "        '350010007201', '350010009041', '350010009061', '040270114053', '040270114072'\n",
+    "    ], year=2020),\n",
+    "    time_frame=TimeFrame.year(2020),\n",
+    ")\n",
+    "\n",
+    "run_test(rume, skip=(\"commuters\", \"dissimilarity_index\"))"
    ]
   }
  ],