Make downstream aggregate weighting explicit; seed regime-aware imputer

downstream.py
- Replace reliance on MicroSeries ``.sum()`` semantics with an
  explicit ``compute_downstream_weighted_aggregate`` helper that pulls
  the correct entity weight variable (tax_unit_weight /
  spm_unit_weight / person_weight / ...) from PE's variable metadata
  and takes the numpy dot product. Same numerics as ``.sum()`` on the
  v11 artifact, but test-covered and robust to simulator changes.
- ``ENTITY_WEIGHT_VARIABLES`` table maps PE entity keys to weight
  variable names.

RegimeAwareDonorImputer
- Add ``seed`` constructor arg and deterministic
  ``_reset_prediction_rngs`` during ``generate`` so repeated calls
  with the same seed produce byte-identical output.

scripts/run_b2_batched.py
- Classify each h5 variable by PE's variable metadata first, then
  fall back to length matching; raises on ambiguous length matches
  rather than silently picking one. Added structural-variable
  overrides for IDs / weights / link columns.
- Wire batched runner's per-chunk aggregate through
  ``compute_downstream_weighted_aggregate``.

scripts/run_b2_validation.py / run_b2_validation_single_var.py
- Use ``compute_downstream_weighted_aggregate`` for consistency with
  the other callers and explicit weighting.

Tests: 3 new entity-resolution tests in test_run_b2_batched.py; 3 new
weighted-aggregate tests in test_downstream.py; 2 new
seed-determinism tests in test_regime_aware_donor_imputer.py. 21
tests pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: MaxGhenis

AGENTS.md

@@ -84,7 +84,7 @@ To avoid rebuilding long prompts in chat:
 <!-- gitnexus:start -->
 # GitNexus — Code Intelligence
 
-This project is indexed by GitNexus as **microplex-us** (4732 symbols, 12777 relationships, 300 execution flows). Use the GitNexus MCP tools to understand code, assess impact, and navigate safely.
+This project is indexed by GitNexus as **microplex-us** (4778 symbols, 12879 relationships, 300 execution flows). Use the GitNexus MCP tools to understand code, assess impact, and navigate safely.
 
 > If any GitNexus tool warns the index is stale, run `npx gitnexus analyze` in terminal first.
 

CLAUDE.md

@@ -1,7 +1,7 @@
 <!-- gitnexus:start -->
 # GitNexus — Code Intelligence
 
-This project is indexed by GitNexus as **microplex-us** (4732 symbols, 12777 relationships, 300 execution flows). Use the GitNexus MCP tools to understand code, assess impact, and navigate safely.
+This project is indexed by GitNexus as **microplex-us** (4778 symbols, 12879 relationships, 300 execution flows). Use the GitNexus MCP tools to understand code, assess impact, and navigate safely.
 
 > If any GitNexus tool warns the index is stale, run `npx gitnexus analyze` in terminal first.
 

scripts/augment_targets_db_for_b2.py

@@ -0,0 +1,77 @@
+"""Copy the calibration targets DB and add direct targets on SSI / CTC / ACA PTC.
+
+The v11 downstream validation showed those three aggregates drifting
++64% / +32% / -76% from their benchmark totals. They weren't in the
+original calibration target set (which focuses on AGI / income
+marginals, not downstream-disbursed amounts). Adding them as direct
+national targets should drive their calibrated aggregates toward the
+benchmark values.
+
+Stratum 1 is "United States" (from the existing DB). Period 2024 and
+reform_id=0 (baseline) match the rest of the 2024 target set.
+"""
+
+from __future__ import annotations
+
+import argparse
+import shutil
+import sqlite3
+from pathlib import Path
+
+from microplex_us.validation.downstream import DOWNSTREAM_BENCHMARKS_2024
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source", required=True, type=Path)
+    parser.add_argument("--output", required=True, type=Path)
+    parser.add_argument(
+        "--variables",
+        nargs="+",
+        default=["ssi", "ctc", "aca_ptc"],
+    )
+    parser.add_argument("--period", default=2024, type=int)
+    args = parser.parse_args()
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(args.source, args.output)
+
+    benchmarks_by_name = {spec.name: spec for spec in DOWNSTREAM_BENCHMARKS_2024}
+
+    con = sqlite3.connect(args.output)
+    cur = con.cursor()
+    for variable in args.variables:
+        spec = benchmarks_by_name.get(variable)
+        if spec is None:
+            raise KeyError(f"No 2024 benchmark spec for {variable}")
+        cur.execute(
+            "SELECT COUNT(*) FROM targets WHERE variable=? AND period=? "
+            "AND stratum_id=1 AND reform_id=0",
+            (variable, args.period),
+        )
+        if cur.fetchone()[0] > 0:
+            print(f"[skip] {variable} already has a national 2024 target")
+            continue
+        cur.execute(
+            "INSERT INTO targets "
+            "(variable, period, stratum_id, reform_id, value, active, source, notes) "
+            "VALUES (?, ?, 1, 0, ?, 1, ?, ?)",
+            (
+                variable,
+                args.period,
+                float(spec.benchmark),
+                spec.source,
+                f"B2 follow-up direct target for {variable}",
+            ),
+        )
+        print(
+            f"[add ] {variable} @ 2024 national: ${spec.benchmark/1e9:.1f}B ({spec.source})"
+        )
+    con.commit()
+    con.close()
+    print(f"\nWrote augmented DB to {args.output}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/run_b2_batched.py

@@ -26,7 +26,6 @@
 import h5py
 import numpy as np
 
-
 HOUSEHOLD_ID = "household_id"
 
 ENTITY_ID_COLUMNS = {
@@ -44,6 +43,25 @@
     "family": "person_family_id",
     "marital_unit": "person_marital_unit_id",
 }
+STRUCTURAL_VARIABLE_ENTITIES = {
+    "household_id": "household",
+    "household_weight": "household",
+    "person_id": "person",
+    "person_household_id": "person",
+    "person_weight": "person",
+    "tax_unit_id": "tax_unit",
+    "person_tax_unit_id": "person",
+    "tax_unit_weight": "tax_unit",
+    "spm_unit_id": "spm_unit",
+    "person_spm_unit_id": "person",
+    "spm_unit_weight": "spm_unit",
+    "family_id": "family",
+    "person_family_id": "person",
+    "family_weight": "family",
+    "marital_unit_id": "marital_unit",
+    "person_marital_unit_id": "person",
+    "marital_unit_weight": "marital_unit",
+}
 
 
 def _load_all_arrays(h5_path: Path, period_key: str) -> dict[str, np.ndarray]:
@@ -55,17 +73,51 @@ def _load_all_arrays(h5_path: Path, period_key: str) -> dict[str, np.ndarray]:
         return out
 
 
-def _entity_of(variable: str, arrays: dict[str, np.ndarray]) -> str:
-    """Classify a variable by matching its array length to an entity's id column."""
+def _load_policyengine_variable_entities() -> dict[str, str]:
+    try:
+        from policyengine_us import (
+            system as policyengine_system_module,  # noqa: PLC0415
+        )
+    except ImportError:
+        return {}
+
+    tax_benefit_system = getattr(policyengine_system_module, "system", None)
+    if tax_benefit_system is None:
+        return {}
+    variables = getattr(tax_benefit_system, "variables", {})
+    entity_map: dict[str, str] = {}
+    for name, metadata in variables.items():
+        entity_key = getattr(getattr(metadata, "entity", None), "key", None)
+        if entity_key is not None:
+            entity_map[str(name)] = str(entity_key)
+    return entity_map
+
+
+def _entity_of(
+    variable: str,
+    arrays: dict[str, np.ndarray],
+    *,
+    variable_entities: dict[str, str] | None = None,
+) -> str:
+    """Classify a variable, preferring PE metadata over fragile length matching."""
+    explicit_entity = STRUCTURAL_VARIABLE_ENTITIES.get(variable)
+    if explicit_entity is not None:
+        return explicit_entity
+    if variable_entities is not None and variable in variable_entities:
+        return variable_entities[variable]
     n = len(arrays[variable])
     entity_lengths = {
         entity: len(arrays[id_col])
         for entity, id_col in ENTITY_ID_COLUMNS.items()
         if id_col in arrays
     }
-    for entity, length in entity_lengths.items():
-        if length == n:
-            return entity
+    matches = [entity for entity, length in entity_lengths.items() if length == n]
+    if len(matches) == 1:
+        return matches[0]
+    if len(matches) > 1:
+        raise ValueError(
+            f"Ambiguous entity for variable {variable!r}: matched {matches} by length"
+        )
     return "unknown"
 
 
@@ -74,7 +126,6 @@ def _build_entity_masks(
 ) -> dict[str, np.ndarray]:
     """Produce boolean masks into each entity array for the households in ``chunk_hh_ids``."""
     hh_id = arrays["household_id"]
-    chunk_set = set(chunk_hh_ids.tolist())
     masks: dict[str, np.ndarray] = {}
     masks["household"] = np.isin(hh_id, chunk_hh_ids)
     person_hh = arrays["person_household_id"]
@@ -94,11 +145,17 @@ def _write_chunk_h5(
     entity_masks: dict[str, np.ndarray],
     period_key: str,
     tmp_path: Path,
+    *,
+    variable_entities: dict[str, str] | None = None,
 ) -> None:
     """Write a subset h5 keeping only rows matching each variable's entity mask."""
     with h5py.File(tmp_path, "w") as f:
         for variable, values in arrays.items():
-            entity = _entity_of(variable, arrays)
+            entity = _entity_of(
+                variable,
+                arrays,
+                variable_entities=variable_entities,
+            )
             mask = entity_masks.get(entity)
             if mask is None or len(values) != len(mask):
                 continue
@@ -118,6 +175,7 @@ def main() -> int:
     period_key = str(args.period)
     print(f"[{time.strftime('%H:%M:%S')}] loading all arrays from {args.dataset}", flush=True)
     arrays = _load_all_arrays(args.dataset, period_key)
+    variable_entities = _load_policyengine_variable_entities()
     print(
         f"[{time.strftime('%H:%M:%S')}] loaded {len(arrays)} variables",
         flush=True,
@@ -132,6 +190,10 @@ def main() -> int:
 
     from policyengine_us import Microsimulation  # noqa: PLC0415
 
+    from microplex_us.validation.downstream import (  # noqa: PLC0415
+        compute_downstream_weighted_aggregate,
+    )
+
     for batch_idx in range(n_batches):
         start = batch_idx * args.batch_size
         end = min(start + args.batch_size, n_hh)
@@ -141,12 +203,21 @@ def main() -> int:
 
         with tempfile.TemporaryDirectory() as tmp:
             tmp_path = Path(tmp) / "chunk.h5"
-            _write_chunk_h5(arrays, entity_masks, period_key, tmp_path)
+            _write_chunk_h5(
+                arrays,
+                entity_masks,
+                period_key,
+                tmp_path,
+                variable_entities=variable_entities,
+            )
 
             t0 = time.time()
             sim = Microsimulation(dataset=str(tmp_path))
-            values = sim.calculate(args.variable, args.period)
-            chunk_sum = float(values.sum())
+            chunk_sum = compute_downstream_weighted_aggregate(
+                sim,
+                args.variable,
+                args.period,
+            )
             total += chunk_sum
             elapsed = time.time() - t0
 

scripts/run_b2_validation.py

@@ -16,6 +16,7 @@
 from microplex_us.validation.downstream import (
     DOWNSTREAM_BENCHMARKS_2024,
     compute_downstream_comparison,
+    compute_downstream_weighted_aggregate,
 )
 
 
@@ -42,7 +43,7 @@ def main() -> int:
         t0 = time.time()
         print(f"[{time.strftime('%H:%M:%S')}] computing {variable} ...", flush=True)
         try:
-            total = float(sim.calculate(variable, args.period).sum())
+            total = compute_downstream_weighted_aggregate(sim, variable, args.period)
         except Exception as exc:
             print(f"  {variable}: FAILED ({exc})", flush=True)
             aggregates[variable] = float("nan")

scripts/run_b2_validation_single_var.py

@@ -16,6 +16,7 @@
 from microplex_us.validation.downstream import (
     DOWNSTREAM_BENCHMARKS_2024,
     compute_downstream_comparison,
+    compute_downstream_weighted_aggregate,
 )
 
 
@@ -33,7 +34,7 @@ def main() -> int:
     sim = Microsimulation(dataset=str(args.dataset))
     print(f"[{time.strftime('%H:%M:%S')}] loaded — computing {args.variable}", flush=True)
     t0 = time.time()
-    total = float(sim.calculate(args.variable, args.period).sum())
+    total = compute_downstream_weighted_aggregate(sim, args.variable, args.period)
     elapsed = time.time() - t0
     print(
         f"[{time.strftime('%H:%M:%S')}] {args.variable} = ${total/1e9:.2f}B "
@@ -42,11 +43,6 @@ def main() -> int:
     )
 
     args.output.parent.mkdir(parents=True, exist_ok=True)
-    if args.output.exists():
-        existing = json.loads(args.output.read_text())
-    else:
-        existing = {}
-
     # Re-read intermediate file if present (accumulates across runs).
     raw_agg_path = args.output.with_suffix(".raw.json")
     raw_aggs = (