Allow intermediate geographies between meta and seed levels

docs/application_configuration.rst

@@ -345,6 +345,15 @@ These settings control the functionality of the PopulationSim algorithm. The set
 | seed_geography | PUMA                | Seed geography name from the list of geographies                                |
 +----------------+---------------------+---------------------------------------------------------------------------------+
 
+The first geography is the meta geography. One or more intermediate geographies
+may appear between it and the seed geography, for example
+``geographies: [REGION, DISTRICT, PUMA, TRACT, TAZ]`` with
+``seed_geography: PUMA``. Include each level in the geographic crosswalk.
+Controls assigned to intermediate levels remain available in their control tables
+and aggregate summaries, but are not distributed to seed zones or enforced during
+seed balancing or integerization. Meta-level controls are still factored into the
+seed controls as usual.
+
 
 **Tracing**:
 

examples/example_test/configs_intermediate/controls.csv

@@ -0,0 +1,11 @@
+target,geography,seed_table,importance,control_field,expression
+num_hh,TAZ,households,1000000000,HHBASE,(households.WGTP > 0) & (households.WGTP < np.inf)
+hh_size_1,TAZ,households,1000,HHSIZE1,households.NP == 1
+hh_size_2,TAZ,households,1000,HHSIZE2,households.NP == 2
+hh_size_3,TAZ,households,1000,HHSIZE3,households.NP == 3
+hh_size_4_plus,TAZ,households,1000,HHSIZE4,households.NP >= 4
+students_by_housing_type,TAZ,persons,1000,OSUFAM,persons.OSUTAG == 1
+hh_by_type,TRACT,households,100,SF,households.HTYPE == 1
+persons_occ_1,DISTRICT,persons,1000,OCCP1,persons.OCCP == 1
+persons_occ_2,REGION,persons,100,OCCP2,persons.OCCP == 2
+persons_occ_3,REGION,persons,100,OCCP3,persons.OCCP == 3

examples/example_test/configs_intermediate/settings.yaml

@@ -0,0 +1,29 @@
+# Reuse the base example; only the geography hierarchy and input list differ.
+inherit_settings: True
+
+geographies: [REGION, DISTRICT, PUMA, TRACT, TAZ]
+seed_geography: PUMA
+
+input_table_list:
+  - tablename: households
+    filename: seed_households.csv
+    index_col: hh_id
+    rename_columns:
+      SERIALNO: hh_id
+  - tablename: persons
+    filename: seed_persons.csv
+    rename_columns:
+      SERIALNO: hh_id
+      SPORDER: per_num
+  - tablename: geo_cross_walk
+    filename: geo_cross_walk.csv
+    rename_columns:
+      TRACTCE: TRACT
+  - tablename: TAZ_control_data
+    filename: taz_controls.csv
+  - tablename: TRACT_control_data
+    filename: tract_controls.csv
+  - tablename: DISTRICT_control_data
+    filename: district_controls.csv
+  - tablename: REGION_control_data
+    filename: region_controls.csv

examples/example_test/data_intermediate/district_controls.csv

@@ -0,0 +1,4 @@
+DISTRICT,OCCP1,HHBASE,HHSIZE1,HHSIZE2,HHSIZE3,HHSIZE4,OSUFAM
+1,720,450,90,144,144,72,318
+2,720,450,90,144,144,72,318
+3,960,600,120,192,192,96,424

examples/example_test/data_intermediate/geo_cross_walk.csv

@@ -0,0 +1,39 @@
+TAZ,TRACTCE,PUMA,DISTRICT,REGION
+100,1,600,1,1
+101,1,600,1,1
+102,1,600,1,1
+103,2,600,1,1
+104,2,600,1,1
+105,2,600,1,1
+200,21,601,1,1
+201,21,601,1,1
+202,21,601,1,1
+203,22,601,1,1
+204,22,601,1,1
+205,22,601,1,1
+300,31,602,2,1
+301,31,602,2,1
+302,31,602,2,1
+303,32,602,2,1
+304,32,602,2,1
+305,32,602,2,1
+400,41,603,2,1
+401,41,603,2,1
+402,41,603,2,1
+403,42,603,2,1
+404,42,603,2,1
+405,42,603,2,1
+500,51,604,3,1
+501,51,604,3,1
+502,51,604,3,1
+503,52,604,3,1
+504,52,604,3,1
+505,52,604,3,1
+600,61,605,3,1
+601,61,605,3,1
+602,61,605,3,1
+603,62,605,3,1
+604,62,605,3,1
+605,62,605,3,1
+#,,,,
+700,72,606,1,1

examples/example_test/data_intermediate/region_controls.csv

@@ -0,0 +1,2 @@
+REGION,OCCP2,OCCP3
+1,300,1560

populationsim/steps/final_seed_balancing.py

@@ -12,7 +12,6 @@
 )
 from populationsim.balancing import do_balancing
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -58,8 +57,10 @@ def final_seed_balancing(settings, crosswalk, control_spec, incidence_table):
             )
         return
 
-    # we use all control_spec rows, so no need to filter on geography as for initial_seed_balancing
+    # Seed controls include factored meta targets, but no intermediate targets.
     seed_controls_df = get_control_table(seed_geography)
+    intermediate_geographies = geographies[1 : geographies.index(seed_geography)]
+    control_spec = control_spec[~control_spec.geography.isin(intermediate_geographies)]
     assert (seed_controls_df.columns == control_spec.target).all()
 
     # determine master_control_index if specified in settings

populationsim/steps/integerize_final_seed_weights.py

@@ -49,7 +49,10 @@ def integerize_final_seed_weights(settings, crosswalk, control_spec, incidence_t
 
     seed_weights_df = get_weight_table(seed_geography)
 
-    # FIXME - I assume we want to integerize using meta controls too?
+    # Match the seed controls, including meta targets but excluding intermediate ones.
+    geographies = settings["geographies"]
+    intermediate_geographies = geographies[1 : geographies.index(seed_geography)]
+    control_spec = control_spec[~control_spec.geography.isin(intermediate_geographies)]
     control_cols = control_spec.target
     assert (seed_controls_df.columns == control_cols).all()
 

populationsim/steps/meta_control_factoring.py

@@ -122,6 +122,10 @@ def meta_control_factoring(settings, control_spec, incidence_table):
 
     seed_controls_df = get_control_table(seed_geography)
 
+    # Remove controls for geographies between the seed and the meta geographies
+    intermediate_geographies = geographies[1 : geographies.index(seed_geography)]
+    control_spec = control_spec[~control_spec.geography.isin(intermediate_geographies)]
+
     assert len(seed_controls_df.index) == len(seed_level_meta_controls.index)
     seed_controls_df = pd.concat([seed_controls_df, seed_level_meta_controls], axis=1)
 

tests/test_intermediate_geography.py

@@ -0,0 +1,185 @@
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from populationsim.core import config, inject, pipeline, tracing
+
+# examples/example_test/configs_intermediate declares
+#     geographies: [REGION, DISTRICT, PUMA, TRACT, TAZ]
+#     seed_geography: PUMA
+# so DISTRICT sits between the meta geography (REGION) and the seed geography
+# (PUMA). Its controls.csv puts persons_occ_1 at DISTRICT and persons_occ_2 and
+# persons_occ_3 at REGION, which is what lets this example tell the two apart.
+SEED_GEOGRAPHY = "PUMA"
+INTERMEDIATE_GEOGRAPHY = "DISTRICT"
+INTERMEDIATE_CONTROL = "persons_occ_1"
+META_CONTROLS = ("persons_occ_2", "persons_occ_3")
+
+_MODELS = [
+    "input_pre_processor",
+    "setup_data_structures",
+    "initial_seed_balancing",
+    "meta_control_factoring",
+    "final_seed_balancing",
+    "integerize_final_seed_weights",
+    "sub_balancing.geography=TRACT",
+    "sub_balancing.geography=TAZ",
+    "expand_households",
+    "summarize",
+    "write_tables",
+    "write_synthetic_population",
+]
+
+
+def setup_function():
+    # Other test modules pin `settings` as a plain injectable via
+    # config.override_setting; restore the decorated one so this example's own
+    # settings -- in particular its five-level `geographies` list -- are read.
+    inject.reinject_decorated_tables()
+
+
+def teardown_function(func):
+    if pipeline.is_open():
+        pipeline.close_pipeline()
+    inject.clear_cache()
+    inject.reinject_decorated_tables()
+
+
+@pytest.mark.parametrize("integerize", [False, True], ids=["fractional", "integerized"])
+@pytest.mark.parametrize("intermediate_levels", [1, 2], ids=["one-level", "two-levels"])
+def test_intermediate_geography(tmp_path, integerize, intermediate_levels):
+
+    example_dir = Path(__file__).parent.parent / "examples" / "example_test"
+
+    # configs_dir and data_dir both cascade: the first directory holding a
+    # given file wins. The *_intermediate directories therefore carry only what
+    # this example actually changes -- its settings, controls, and the crosswalk
+    # with the extra geography levels -- and everything shared (seed households
+    # and persons, TAZ and TRACT controls, logging) resolves from the base
+    # example. Copying those would leave this example silently testing stale
+    # inputs whenever the originals changed.
+    inject.add_injectable(
+        "data_dir", [example_dir / "data_intermediate", example_dir / "data"]
+    )
+    inject.add_injectable(
+        "configs_dir", [example_dir / "configs_intermediate", example_dir / "configs"]
+    )
+    inject.add_injectable("output_dir", tmp_path)
+
+    inject.clear_cache()
+
+    config.override_setting("NO_INTEGERIZATION_EVER", not integerize)
+    intermediate_targets = {INTERMEDIATE_GEOGRAPHY: INTERMEDIATE_CONTROL}
+    meta_controls = META_CONTROLS
+    geographies = ["REGION", INTERMEDIATE_GEOGRAPHY, SEED_GEOGRAPHY, "TRACT", "TAZ"]
+
+    if intermediate_levels == 2:
+        # Insert AREA above DISTRICT, with its own target. Both intermediate
+        # targets must be excluded; filtering only the level next to seed fails.
+        inputs_dir = tmp_path / "inputs"
+        inputs_dir.mkdir()
+        crosswalk = pd.read_csv(
+            example_dir / "data_intermediate" / "geo_cross_walk.csv", comment="#"
+        )
+        crosswalk["AREA"] = crosswalk["DISTRICT"].map({1: 1, 2: 1, 3: 2})
+        crosswalk.to_csv(inputs_dir / "geo_cross_walk.csv", index=False)
+        pd.DataFrame({"AREA": [1, 2], "OCCP2": [180, 120]}).to_csv(
+            inputs_dir / "area_controls.csv", index=False
+        )
+        controls = pd.read_csv(example_dir / "configs_intermediate" / "controls.csv")
+        controls.loc[controls.target == "persons_occ_2", "geography"] = "AREA"
+        controls.to_csv(inputs_dir / "controls.csv", index=False)
+        geographies.insert(1, "AREA")
+        config.override_setting("geographies", geographies)
+        settings = inject.get_injectable("settings")
+        config.override_setting(
+            "input_table_list",
+            settings["input_table_list"]
+            + [{"tablename": "AREA_control_data", "filename": "area_controls.csv"}],
+        )
+        for directory in ("data_dir", "configs_dir"):
+            inject.add_injectable(
+                directory, [inputs_dir] + inject.get_injectable(directory)
+            )
+        intermediate_targets["AREA"] = "persons_occ_2"
+        meta_controls = ("persons_occ_3",)
+
+    tracing.config_logger()
+
+    # Completing at all is the primary regression. Without the control_spec
+    # filtering in meta_control_factoring / final_seed_balancing /
+    # integerize_final_seed_weights, the seed-level controls are indexed with
+    # the intermediate geography's target and this raises
+    # KeyError: "['persons_occ_1'] not in index".
+    pipeline.run(models=_MODELS, resume_after=None)
+
+    crosswalk = pipeline.get_table("crosswalk")
+    assert list(crosswalk.columns) == geographies
+
+    # The behaviour under test: controls belonging to a geography between the
+    # meta and seed geographies are kept out of the seed-level control table,
+    # while the meta geography's controls are still factored into it.
+    seed_controls = pipeline.get_table(f"{SEED_GEOGRAPHY}_controls")
+    for target in intermediate_targets.values():
+        assert target not in seed_controls.columns
+    for target in meta_controls:
+        assert target in seed_controls.columns
+
+    # Intermediate targets remain available for reporting. This feature does
+    # not distribute them to seed zones or enforce them during balancing.
+    for geography, target in intermediate_targets.items():
+        intermediate_controls = pipeline.get_table(f"{geography}_controls")
+        assert target in intermediate_controls.columns
+        for sub_geography in ("TRACT", "TAZ"):
+            summary = pipeline.get_table(f"summary_{sub_geography}_{geography}")
+            assert set(summary["id"]) == set(crosswalk[geography])
+            np.testing.assert_allclose(
+                summary[f"{target}_control"],
+                intermediate_controls.loc[summary["id"], target],
+            )
+            weights = pipeline.get_table(f"{sub_geography}_weights")
+            incidence = pipeline.get_table("incidence_table")
+            weight_col = "integer_weight" if integerize else "balanced_weight"
+            expected_results = (
+                (weights[weight_col] * weights["hh_id"].map(incidence[target]))
+                .groupby(weights[geography])
+                .sum()
+            )
+            np.testing.assert_allclose(
+                summary[f"{target}_result"], expected_results.loc[summary["id"]]
+            )
+
+    # Sub-geography balancing ran and preserved the seed-level household total.
+    seed_total = pipeline.get_table(f"{SEED_GEOGRAPHY}_weights")[
+        "balanced_weight"
+    ].sum()
+    assert seed_total > 0
+    for geography in ("TRACT", "TAZ"):
+        weights = pipeline.get_table(f"{geography}_weights")["balanced_weight"]
+        assert (weights >= 0).all()
+        assert weights.sum() == pytest.approx(seed_total, rel=1e-6)
+
+    expanded = pipeline.get_table("expanded_household_ids")
+    if integerize:
+        seed_weights = pipeline.get_table(f"{SEED_GEOGRAPHY}_weights")
+        assert (seed_weights.integer_weight >= 0).all()
+        assert (seed_weights.integer_weight % 1 == 0).all()
+        seed_totals = seed_weights.groupby(SEED_GEOGRAPHY).integer_weight.sum()
+        pd.testing.assert_series_equal(
+            seed_totals, seed_controls.num_hh, check_dtype=False, check_names=False
+        )
+        taz_weights = pipeline.get_table("TAZ_weights")
+        expected_counts = taz_weights.groupby("TAZ").integer_weight.sum()
+        pd.testing.assert_series_equal(
+            expanded.groupby("TAZ").size(),
+            expected_counts,
+            check_dtype=False,
+            check_names=False,
+        )
+        assert len(expanded) == expected_counts.sum()
+        assert (tmp_path / "synthetic_households.csv").exists()
+        assert (tmp_path / "synthetic_persons.csv").exists()
+    else:
+        assert expanded.empty