From a0a762a12f614498d08ca29144c8457c20a5c552 Mon Sep 17 00:00:00 2001 From: vahid-ahmadi Date: Tue, 11 Aug 2026 19:22:05 +0100 Subject: [PATCH 1/6] Add observed firm frame: real sponsor records calibrated to SUSB (#192) Builds the firm population from observed Form 5500 sponsor records instead of generating one. No synthetic firm row is created. Form 5500 alone populates all 97 SUSB sector x canonical-band cells with at least ten records each, so the frame uses a single consistent unit (plan sponsor) and OSHA ITA stays an independent measurement reference. Unioning the two would silently mix a sponsor unit with an establishment unit: they agree on the canonical band for only 66.3% of the 43,001 EINs they share. The calibration took three attempts and the first two are recorded because the failures are informative: 1. post_stratify (firm-count margin only) matches firms exactly and overshoots employment by 2.06x. A single weight per cell treats a 307,086-participant enterprise as representative of twenty ordinary 500+ firms. 2. Finer stratification on SUSB's 22 detail size classes makes it worse (2.52x), because SUSB's top class is also unbounded. 3. calibrate_dual_margin lets weights vary within a cell via a maximum-entropy tilt, matching firm and employment margins at once. Employment ratio 0.973 across 92 of 97 cells. The five remaining cells fail closed rather than being forced, and they are a SUSB data-quality artifact: noise infusion distorts thin published cells enough to make them arithmetically impossible for their own size class. NAICS 11's 2,000-2,499 class reports 5 firms and 292 employees (flag H); its 5,000+ class reports 28 firms and 6,778 (flag J). Where the implied cell mean falls outside the band it belongs to, no reweighting can reach it. Post-stratification is retained alongside the working method because its employment_coverage diagnostic is what falsified the naive design. Co-Authored-By: Claude Opus 5 (1M context) --- src/populace_dynamics/firms/frame.py | 401 +++++++++++++++++++++++++++ 1 file changed, 401 insertions(+) create mode 100644 src/populace_dynamics/firms/frame.py diff --git a/src/populace_dynamics/firms/frame.py b/src/populace_dynamics/firms/frame.py new file mode 100644 index 00000000..d39a1ddb --- /dev/null +++ b/src/populace_dynamics/firms/frame.py @@ -0,0 +1,401 @@ +"""Observed firm frame — real sponsor records reweighted to SUSB (#192). + +The employer-firm design originally assumed a **generated** firm +population, on the premise that no public firm microdata existed. It +does (``data/form5500.py``, ``data/osha_ita.py``), and it covers the +calibration grid densely enough to reweight instead of generate. +Measured against the committed SUSB 2022 extract, Form 5500 sponsors +alone populate **all 97** sector x canonical-band cells with at least +ten records each and 0.000% of SUSB employment in an under-covered +cell. So this module builds the firm population by **post-stratifying +observed sponsor records**, and no synthetic firm row is created. + +**One source defines the frame; the other validates it.** Form 5500 +sponsors and OSHA ITA establishments are *different units* — the two +sources agree on the canonical band for only 66.3% of the 43,001 EINs +they share, with disagreement running in both directions. Unioning +them would silently mix a plan-sponsor unit with an establishment +unit. The frame is therefore Form 5500 only, and OSHA ITA is retained +as an independent measurement reference (see +:func:`band_agreement_reference`). Mixing is a design error, not a +coverage improvement. + +**What post-stratification does and does not fix.** Cell weights are +``SUSB firms in cell / observed records in cell``, so the weighted +firm count matches SUSB **exactly, by construction**, in every cell. +It does not fix *within-cell* selection: a plan-sponsoring firm may +differ systematically from a non-sponsor of the same size and sector. +That is the same limitation RAND COMPARE carries when it reweights +Kaiser/HRET records, and it is documented rather than solved. + +**Employment is deliberately not forced to match — and measuring it +falsifies the naive design.** ``Active participants`` count +plan-covered workers, a lower bound on employment *per firm*, so the +weighted participant total was expected to fall short of SUSB +employment. Measured on the 2023 files it does the opposite: + +====================== ========== ========== ====== +band sponsor SUSB ratio + mean size mean size +====================== ========== ========== ====== +``LT10`` 4.62 2.59 1.78 +``B10_49`` 21.99 19.76 1.11 +``B50_99`` 69.13 64.60 1.07 +``B100_499`` 213.33 161.15 1.32 +``B500_PLUS`` 4,364.97 1,675.47 2.61 +====================== ========== ========== ====== + +Weighted participants total **279,300,780 against SUSB's 135,748,407 +— 2.06x**. Firm counts match exactly and employment is off by more +than a factor of two. + +The cause is **within-cell selection**, the limitation named above, +and it is large rather than marginal. A firm must be big enough to +sponsor an ERISA plan at all, so the sponsors sitting in a given band +are systematically the *larger* firms in that band, while SUSB's band +is dominated by the many non-sponsoring firms below them. Post- +stratifying on firm count then multiplies those over-large records by +the cell weight. ``B500_PLUS`` is worst because a second effect +compounds it: the sponsor EIN is frequently a parent enterprise +aggregating several operating subsidiaries (``data/form5500.py``). + +So :func:`post_stratify` as written produces a population that is +**exact on the firm margin and unusable on the employment margin**. +It is retained because the diagnostic is the point: the ratios are +reported per band in ``frame.attrs["employment_coverage"]`` rather +than silently raked away. Closing the gap needs a size-measure +correction or a sponsorship-propensity model, not a second raking +step, and that is a design question for the referee round — raking to +both margins would hide a definitional and selection problem inside a +weight. +""" + +from __future__ import annotations + +from pathlib import Path + +import pandas as pd + +from ..data.form5500 import read_sponsors +from .banding import CANONICAL_BANDS, susb_entrsize_to_canonical +from .targets import load_susb_sector_size + +__all__ = [ + "CELL_KEYS", + "susb_cell_targets", + "sponsor_frame", + "post_stratify", + "calibrate_dual_margin", + "band_agreement_reference", +] + +#: The post-stratification cell: NAICS sector x canonical firm-size band. +CELL_KEYS = ("naics_sector", "canonical_band") + +_BAND_ORDER = [band.name for band in CANONICAL_BANDS] + + +def susb_cell_targets(path: str | None = None) -> pd.DataFrame: + """SUSB firm counts and employment per sector x canonical band. + + Built from the *detail* enterprise-size classes only — the canonical + edges were chosen so every detail class nests exactly in one band, + so no straddling allocation is needed. Total and subtotal rows + (``ENTRSIZE`` 01/33/37) and the all-sector ``--`` margin are + dropped; including either would double-count. + """ + susb = load_susb_sector_size(path) + susb = susb[susb["naics_sector"] != "--"].copy() + spans = susb["entrsize_code"].map(susb_entrsize_to_canonical) + susb = susb[spans.notna()].copy() + susb["canonical_band"] = [span.band.name for span in spans[spans.notna()]] + # SUSB publishes combined sector ranges such as "31-33"; the first + # two digits are the sector key both sides join on. + susb["naics_sector"] = susb["naics_sector"].astype(str).str[:2] + out = ( + susb.groupby(list(CELL_KEYS))[["firms", "employment"]] + .sum() + .reset_index() + ) + return out[out["firms"] > 0].reset_index(drop=True) + + +def sponsor_frame(year: int, data_dir: Path | None = None) -> pd.DataFrame: + """Observed Form 5500 sponsor records, keyed to the SUSB cell. + + One row per sponsor EIN. Sponsors whose business code does not + yield a usable NAICS sector are dropped rather than pooled into an + "unknown" cell, because an unknown-sector row cannot be weighted to + any SUSB target and would silently dilute whichever cell absorbed + it. The count dropped is recorded on the returned frame. + """ + sponsors = read_sponsors(year, data_dir) + usable = sponsors["naics_sector"].notna() + frame = sponsors[usable].reset_index(drop=True).copy() + frame.attrs["dropped_unknown_sector"] = int((~usable).sum()) + return frame + + +def post_stratify( + frame: pd.DataFrame, + targets: pd.DataFrame | None = None, + *, + require_full_coverage: bool = True, +) -> pd.DataFrame: + """Attach cell weights so weighted firm counts match SUSB exactly. + + ``weight = SUSB firms in cell / observed records in cell``. Every + record in a cell carries the same weight; the weighted count then + reproduces the SUSB firm margin in that cell by construction. + + ``require_full_coverage`` (default True) makes an **empty target + cell a hard error**. No weight can populate a stratum with zero + observed records, so silently returning a frame that under-counts + a whole sector x band cell would produce a population that looks + calibrated and is not. Set it False only to inspect a deliberately + partial frame. + + Diagnostics are attached to ``result.attrs``: + + ``weighted_firms`` / ``target_firms`` + totals, equal by construction when coverage is full. + ``employment_coverage`` + per band, weighted active participants against SUSB + employment. Measured at **1.07-2.61, total 2.06x** on the 2023 + files — the naive design overshoots employment badly because + plan sponsors are the larger firms within any band (module + docstring). Reported, never corrected: a caller that needs a + usable employment margin must address the size measure, not + rake this away. + ``uncovered_cells`` + target cells with no observed record, empty when coverage is + full. + """ + if targets is None: + targets = susb_cell_targets() + for column in CELL_KEYS: + if column not in frame.columns: + raise ValueError(f"Frame lacks {column!r}; use sponsor_frame().") + + observed = ( + frame.groupby(list(CELL_KEYS)).size().rename("records").reset_index() + ) + merged = targets.merge(observed, on=list(CELL_KEYS), how="left") + merged["records"] = merged["records"].fillna(0).astype(int) + + uncovered = merged[merged["records"] == 0] + if require_full_coverage and not uncovered.empty: + raise ValueError( + f"{len(uncovered)} SUSB target cell(s) have no observed " + "record; no weight can populate an empty stratum. Cells: " + + ", ".join( + f"{row.naics_sector}/{row.canonical_band}" + for row in uncovered.itertuples() + ) + ) + + merged["weight"] = merged["firms"] / merged["records"].where( + merged["records"] > 0 + ) + out = frame.merge( + merged[[*CELL_KEYS, "weight", "firms", "employment"]], + on=list(CELL_KEYS), + how="left", + ) + # A sponsor in a cell SUSB does not publish has no target and is + # given zero weight rather than being dropped: keeping the row + # makes the exclusion visible in the frame instead of silently + # shrinking it. + out["weight"] = out["weight"].fillna(0.0) + + coverage = [] + for band in _BAND_ORDER: + rows = out[out["canonical_band"] == band] + target_emp = targets.loc[ + targets["canonical_band"] == band, "employment" + ].sum() + weighted_participants = float( + (rows["active_participants"] * rows["weight"]).sum() + ) + coverage.append( + { + "band": band, + "weighted_active_participants": weighted_participants, + "susb_employment": int(target_emp), + "ratio": ( + weighted_participants / target_emp if target_emp else None + ), + } + ) + + out.attrs["weighted_firms"] = float(out["weight"].sum()) + out.attrs["target_firms"] = int(targets["firms"].sum()) + out.attrs["employment_coverage"] = coverage + out.attrs["uncovered_cells"] = [ + f"{row.naics_sector}/{row.canonical_band}" + for row in uncovered.itertuples() + ] + return out + + +def _cell_weights(sizes, target_firms: float, target_employment: float): + """Maximum-entropy weights matching a cell's firm and size margins. + + Solves ``w_i = exp(l * (p_i - mean))`` scaled so ``sum(w) == + target_firms`` and ``sum(w * p) == target_employment``. The + exponential tilt keeps every weight strictly positive, which a + linear calibration does not guarantee. + + Returns ``(weights, None)`` on success, or ``(None, reason)`` when + the target mean lies outside the observed support — no reweighting + of the observed records can reach a mean they do not bracket. + """ + import numpy as np + + sizes = np.asarray(sizes, dtype=float) + if target_firms <= 0 or len(sizes) == 0: + return None, "empty cell or non-positive firm target" + target_mean = target_employment / target_firms + if not (sizes.min() <= target_mean <= sizes.max()): + return None, ( + f"target mean {target_mean:.1f} outside observed support " + f"[{sizes.min():.0f}, {sizes.max():.0f}]" + ) + tilt = 0.0 + for _ in range(200): + weights = np.exp(tilt * (sizes - target_mean)) + total = weights.sum() + mean = (weights * sizes).sum() / total + variance = (weights * sizes * sizes).sum() / total - mean * mean + if variance <= 0: + break + step = (target_mean - mean) / variance + tilt += step + if abs(step) < 1e-12: + break + weights = np.exp(tilt * (sizes - target_mean)) + return weights * (target_firms / weights.sum()), None + + +def calibrate_dual_margin( + frame: pd.DataFrame, + targets: pd.DataFrame | None = None, + *, + size_column: str = "active_participants", +) -> pd.DataFrame: + """Weight records to match SUSB firm **and** employment margins. + + :func:`post_stratify` matches firm counts only, which overshoots + employment by 2.06x (module docstring) because a single weight per + cell treats a 307,086-participant enterprise as representative of + twenty ordinary ``500+`` firms. Stratifying more finely does not + help — it makes the ratio worse (2.52), because SUSB's top class is + also unbounded. The fix is to let weights vary *within* a cell so + both margins can be met at once. + + On the 2023 files this lands the weighted employment ratio at + **0.973** with 92 of 97 cells calibrated. + + **The five failures are a SUSB data-quality artifact, not a + modelling one, and they fail closed.** SUSB infuses noise into + published employment, and in thin cells the distortion is extreme: + NAICS 11's ``2,000-2,499`` class reports 5 firms and 292 employees + (flag ``H``), and its ``5,000+`` class 28 firms and 6,778 (flag + ``J``). Both are arithmetically impossible for their own size + class. Where the implied cell mean falls outside the band it + belongs to, no reweighting can reach it, so the cell is reported in + ``attrs["infeasible_cells"]`` with its reason and given zero weight + rather than being forced. + """ + if targets is None: + targets = susb_cell_targets() + if size_column not in frame.columns: + raise ValueError(f"Frame lacks {size_column!r}; use sponsor_frame().") + + weights = pd.Series(0.0, index=frame.index) + infeasible: list[dict[str, object]] = [] + calibrated = 0 + indexed = targets.set_index(list(CELL_KEYS)) + for cell, group in frame.groupby(list(CELL_KEYS)): + if cell not in indexed.index: + continue + row = indexed.loc[cell] + cell_weights, reason = _cell_weights( + group[size_column].to_numpy(), + float(row["firms"]), + float(row["employment"]), + ) + if cell_weights is None: + infeasible.append( + { + "cell": "/".join(cell), + "reason": reason, + "susb_firms": int(row["firms"]), + "susb_employment": int(row["employment"]), + } + ) + continue + weights.loc[group.index] = cell_weights + calibrated += 1 + + out = frame.copy() + out["weight"] = weights + weighted_employment = float((out[size_column] * out["weight"]).sum()) + out.attrs["calibrated_cells"] = calibrated + out.attrs["infeasible_cells"] = infeasible + out.attrs["weighted_firms"] = float(out["weight"].sum()) + out.attrs["target_firms"] = int(targets["firms"].sum()) + out.attrs["weighted_employment"] = weighted_employment + out.attrs["target_employment"] = int(targets["employment"].sum()) + out.attrs["employment_ratio"] = ( + weighted_employment / targets["employment"].sum() + ) + return out + + +def band_agreement_reference( + sponsors: pd.DataFrame, establishments: pd.DataFrame +) -> dict[str, object]: + """Cross-source band agreement — a measurement floor, not a gate. + + Joins Form 5500 sponsors to OSHA ITA establishments on EIN (OSHA + establishments summed to the EIN) and reports how often the two + *administrative* firm-size measures land in the same canonical + band. On the 2023/2025 vintages that rate is 66.3% over 43,001 + shared EINs. + + The number bounds how tightly any firm-size gate can be set: a + threshold finer than the agreement rate between two administrative + instruments is finer than the instruments resolve. It is evidence + for the floor batteries and is deliberately **not** used to + reweight or correct either source, which measure different units. + """ + if "canonical_band" not in establishments.columns: + raise ValueError( + "Establishments lack 'canonical_band'; pass the output of " + "osha_ita.apply_quality_rule()." + ) + left = sponsors.set_index("sponsor_ein")["canonical_band"] + right = ( + establishments.dropna(subset=["ein"]) + .groupby("ein")["annual_average_employees"] + .sum() + ) + from .banding import band_of_count + + right_bands = right[right >= 1].map(lambda n: band_of_count(int(n)).name) + joined = pd.concat( + [left.rename("sponsor_band"), right_bands.rename("estab_band")], + axis=1, + join="inner", + ) + if joined.empty: + return {"matched_eins": 0, "agreement_rate": None} + agree = float((joined["sponsor_band"] == joined["estab_band"]).mean()) + return { + "matched_eins": int(len(joined)), + "agreement_rate": agree, + "confusion": pd.crosstab( + joined["sponsor_band"], joined["estab_band"] + ).to_dict(), + } From 41ef2a4e57a55f4e15d5212b3295dd4d8db92f6b Mon Sep 17 00:00:00 2001 From: vahid-ahmadi Date: Wed, 12 Aug 2026 09:20:28 +0100 Subject: [PATCH 2/6] Firm frame: bounded dual-margin calibration and scope exclusions (#192) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the frame's first commit, resolving the two defects it left open. Both choices are measured rather than assumed. ## Weight bound: 140, from the sensitivity curve The unbounded exponential tilt put a weight of 3,866 on one NAICS 99 record in a 35-record cell, beside another at 1.3e-08. Bounded calibration is the Deville-Sarndal (1992) remedy the project already cites. The bound is not a guess — p99 is 76.5, so bounds above ~140 never bind and give identical margins, while tighter ones degrade the firm margin sharply: bound firm ratio emp ratio 50 0.7023 0.9766 100 0.9558 1.0151 130 0.9960 1.0226 140 0.9988 1.0237 250 0.9988 1.0237 1000 0.9990 1.0237 140 is the tightest bound that costs nothing on either margin. It is a referee parameter, like the OSHA employment cap, not a default to inherit silently. Fixed along the way: the first bounded implementation clipped and then rescaled to hit the firm total, which pushes weights straight back over the bound — it left two records at 4,138 while the bound was nominally 250. The final clip no longer rescales, and the firm-margin residual that leaves is reported instead. ## NAICS 55 employment is excluded, firms are kept NAICS 55 (Management of Companies) is a holding-company sector: one Form 5500 covers the whole enterprise's workforce, while SUSB attributes those workers to the operating sectors and records only head-office staff under 55. Counting them again double-counts, and the effect is not marginal: 42,676,524 weighted employees against SUSB's 3,661,977 (11.7x), which is 24% of all weighted employment. It is systematic across every band (1.0x, 2.6x, 4.2x, 6.0x, 12.4x), not a thin-cell artifact. These sponsors stay in the firm margin — SUSB does count 25,413 of them — but carry employment_in_scope=False. The flag travels on the frame rather than only in attrs, so a downstream consumer cannot sum double-counted employment by accident. Also corrected: the five impossible cells were earlier attributed to SUSB noise infusion. That was wrong — 507 of 532 rows carry the low-noise G flag. They are a documented SUSB scope rule: NAICS 55 as above, and NAICS 11 because SUSB excludes crop and animal production. ## Result, and a held-out check firms 6,453,598 / 6,461,497 = 0.9988 employment (in-scope) 1.0237 weights min 0.100 median 2.99 max 140.0 Per band, in-scope: 0.9993 / 1.0000 / 1.0000 / 1.0024 / 1.0440. Validated against BDS 2022, which is never used in the calibration, with SUSB-vs-BDS as a control to separate our error from source disagreement. Firm counts: our/BDS 1.176, 1.270, 1.892 against SUSB/BDS 1.177, 1.270, 1.888 — our own error is 0.999, 1.000, 1.002. The whole gap is SUSB and BDS disagreeing, not the calibration. No random draw, sampling or generation anywhere in the firm path: every row traces to a real DOL filing. Co-Authored-By: Claude Opus 5 (1M context) --- src/populace_dynamics/firms/frame.py | 219 ++++++++++++++++++++++++--- 1 file changed, 197 insertions(+), 22 deletions(-) diff --git a/src/populace_dynamics/firms/frame.py b/src/populace_dynamics/firms/frame.py index d39a1ddb..65b12d0f 100644 --- a/src/populace_dynamics/firms/frame.py +++ b/src/populace_dynamics/firms/frame.py @@ -82,6 +82,8 @@ __all__ = [ "CELL_KEYS", + "DEFAULT_WEIGHT_BOUNDS", + "EMPLOYMENT_OUT_OF_SCOPE_SECTORS", "susb_cell_targets", "sponsor_frame", "post_stratify", @@ -92,6 +94,27 @@ #: The post-stratification cell: NAICS sector x canonical firm-size band. CELL_KEYS = ("naics_sector", "canonical_band") +#: Sectors whose sponsor employment is **already counted elsewhere** and +#: must be excluded from the employment margin. +#: +#: NAICS 55 (Management of Companies and Enterprises) is a holding-company +#: sector: one Form 5500 covers the whole enterprise's workforce, while +#: SUSB attributes those workers to the *operating* sectors and records +#: only head-office staff under 55. Including them double-counts. The +#: effect is not marginal — measured on the 2023 files NAICS 55 yields +#: 42,676,524 weighted employees against SUSB's 3,661,977 (**11.7x**), +#: which is **24% of all weighted employment in the frame**. It is +#: systematic across every band, not a thin-cell artifact:: +#: +#: 1-9 1.0x 100-499 6.0x +#: 10-49 2.6x 500+ 12.4x +#: 50-99 4.2x +#: +#: These sponsors are kept in the **firm** margin — SUSB does count +#: 25,413 NAICS 55 firms, and the frame reproduces them — but their +#: employment is flagged out of scope rather than summed. +EMPLOYMENT_OUT_OF_SCOPE_SECTORS = frozenset({"55"}) + _BAND_ORDER = [band.name for band in CANONICAL_BANDS] @@ -238,43 +261,130 @@ def post_stratify( return out -def _cell_weights(sizes, target_firms: float, target_employment: float): - """Maximum-entropy weights matching a cell's firm and size margins. +def _solve_tilt(sizes, target_mean: float, free) -> float: + """Exponential-tilt parameter matching ``target_mean`` on ``free``. + + Newton iteration on the tilt: the derivative of the weighted mean + with respect to the tilt is the weighted variance, so the step is + ``(target - mean) / variance``. + """ + import numpy as np + + free_sizes = sizes[free] + tilt = 0.0 + for _ in range(200): + weights = np.exp(np.clip(tilt * (free_sizes - target_mean), -700, 700)) + total = weights.sum() + if total <= 0 or not np.isfinite(total): + break + mean = (weights * free_sizes).sum() / total + variance = (weights * free_sizes**2).sum() / total - mean**2 + if variance <= 0 or not np.isfinite(variance): + break + step = (target_mean - mean) / variance + tilt += step + if abs(step) < 1e-12: + break + return tilt + + +#: Default weight bounds. An unbounded exponential tilt produced a +#: 3,866 weight next to a 1.3e-08 one in the 35-record NAICS 99 +#: (unclassified) cell — a single record standing in for 3,866 firms is +#: not a calibration, it is a leverage point. Bounded calibration is +#: the Deville-Sarndal (1992) remedy the project already cites. +#: +#: 140 is the knee of the measured sensitivity curve, not a guess. The +#: p99 weight is 76.5, so bounds above ~140 never bind and give +#: identical margins; below it the firm margin degrades sharply: +#: +#: ====== ============ =========== +#: bound firm ratio emp ratio +#: ====== ============ =========== +#: 50 0.7023 0.9766 +#: 100 0.9558 1.0151 +#: 120 0.9705 1.0173 +#: 130 0.9960 1.0226 +#: **140** 0.9988 1.0237 +#: 250 0.9988 1.0237 +#: 1000 0.9990 1.0237 +#: ====== ============ =========== +#: +#: So 140 is the tightest bound that costs nothing on either margin, +#: and it is a referee parameter like the OSHA employment cap, not an +#: implementation default to inherit silently. +DEFAULT_WEIGHT_BOUNDS = (0.1, 140.0) + + +def _cell_weights( + sizes, + target_firms: float, + target_employment: float, + bounds: tuple[float, float] | None = DEFAULT_WEIGHT_BOUNDS, +): + """Bounded maximum-entropy weights for one cell's two margins. Solves ``w_i = exp(l * (p_i - mean))`` scaled so ``sum(w) == target_firms`` and ``sum(w * p) == target_employment``. The exponential tilt keeps every weight strictly positive, which a linear calibration does not guarantee. - Returns ``(weights, None)`` on success, or ``(None, reason)`` when - the target mean lies outside the observed support — no reweighting - of the observed records can reach a mean they do not bracket. + ``bounds`` clips weights to ``[lo, hi]`` and re-solves the tilt on + the records still free, iterating until the clipped set is stable. + Because the tilt is monotone in ``p_i``, clipping only ever removes + the extremes, so the procedure terminates. Bounding trades exactness + on the employment margin for the absence of leverage points; the + residual is returned so the caller can report it rather than + discover it later. + + Returns ``(weights, None, residual)`` on success, or + ``(None, reason, None)`` when the target mean lies outside the + observed support — no reweighting of the observed records can reach + a mean they do not bracket. """ import numpy as np sizes = np.asarray(sizes, dtype=float) if target_firms <= 0 or len(sizes) == 0: - return None, "empty cell or non-positive firm target" + return None, "empty cell or non-positive firm target", None target_mean = target_employment / target_firms if not (sizes.min() <= target_mean <= sizes.max()): - return None, ( - f"target mean {target_mean:.1f} outside observed support " - f"[{sizes.min():.0f}, {sizes.max():.0f}]" + return ( + None, + ( + f"target mean {target_mean:.1f} outside observed support " + f"[{sizes.min():.0f}, {sizes.max():.0f}]" + ), + None, ) - tilt = 0.0 - for _ in range(200): - weights = np.exp(tilt * (sizes - target_mean)) - total = weights.sum() - mean = (weights * sizes).sum() / total - variance = (weights * sizes * sizes).sum() / total - mean * mean - if variance <= 0: + + free = np.ones(len(sizes), dtype=bool) + weights = np.ones(len(sizes), dtype=float) + for _ in range(20): + tilt = _solve_tilt(sizes, target_mean, free) + # Clip the exponent before exponentiating: a wide cell (500 to + # 2.5M participants) can overflow float64 mid-solve, and an inf + # weight silently poisons the rescale to NaN. + weights = np.exp(np.clip(tilt * (sizes - target_mean), -700, 700)) + weights *= target_firms / weights.sum() + if bounds is None: break - step = (target_mean - mean) / variance - tilt += step - if abs(step) < 1e-12: + lo, hi = bounds + clipped = (weights < lo) | (weights > hi) + if not clipped.any() or not (free & ~clipped).any(): break - weights = np.exp(tilt * (sizes - target_mean)) - return weights * (target_firms / weights.sum()), None + weights = np.clip(weights, lo, hi) + free = ~clipped + if bounds is not None: + # Final clip is NOT followed by a rescale. Rescaling to hit the + # firm total after clipping pushes weights straight back over + # the bound — that bug left two NAICS 99 records at 4,138 while + # the bound was nominally 250. The firm-margin residual this + # leaves is returned instead of being papered over. + weights = np.clip(weights, *bounds) + residual = float((weights * sizes).sum() - target_employment) + firm_residual = float(weights.sum() - target_firms) + return weights, None, (residual, firm_residual) def calibrate_dual_margin( @@ -282,6 +392,8 @@ def calibrate_dual_margin( targets: pd.DataFrame | None = None, *, size_column: str = "active_participants", + weight_bounds: tuple[float, float] | None = DEFAULT_WEIGHT_BOUNDS, + firm_only_fallback: bool = True, ) -> pd.DataFrame: """Weight records to match SUSB firm **and** employment margins. @@ -314,16 +426,19 @@ def calibrate_dual_margin( weights = pd.Series(0.0, index=frame.index) infeasible: list[dict[str, object]] = [] + residuals: list[dict[str, object]] = [] calibrated = 0 + firm_only = 0 indexed = targets.set_index(list(CELL_KEYS)) for cell, group in frame.groupby(list(CELL_KEYS)): if cell not in indexed.index: continue row = indexed.loc[cell] - cell_weights, reason = _cell_weights( + cell_weights, reason, residual = _cell_weights( group[size_column].to_numpy(), float(row["firms"]), float(row["employment"]), + weight_bounds, ) if cell_weights is None: infeasible.append( @@ -332,17 +447,65 @@ def calibrate_dual_margin( "reason": reason, "susb_firms": int(row["firms"]), "susb_employment": int(row["employment"]), + "firm_margin_held": bool(firm_only_fallback), } ) + if firm_only_fallback: + # The employment target is out of scope for this cell + # (module docstring), but the firm count is sound. Hold + # the firm margin and drop only the employment + # constraint, rather than losing the cell's firms. + # + # The same bound applies here. NAICS 55's 500+ cell has + # 7,373 SUSB firms behind very few sponsors, so an + # unbounded fallback weight reaches 4,138 — a worse + # leverage point than the one bounding was added to + # remove. Bounding it under-counts that cell's firms + # instead, which is recorded rather than hidden. + flat = float(row["firms"]) / len(group) + if weight_bounds is not None: + flat = min(max(flat, weight_bounds[0]), weight_bounds[1]) + weights.loc[group.index] = flat + infeasible[-1]["fallback_weight"] = flat + infeasible[-1]["firms_represented"] = flat * len(group) + firm_only += 1 continue weights.loc[group.index] = cell_weights calibrated += 1 + if residual and any(residual): + residuals.append( + { + "cell": "/".join(cell), + "employment_residual": residual[0], + "firm_residual": residual[1], + } + ) out = frame.copy() out["weight"] = weights weighted_employment = float((out[size_column] * out["weight"]).sum()) + scoped = targets.copy() + excluded = {entry["cell"] for entry in infeasible} + # Whole-sector exclusions (NAICS 55) join the cell-level ones, so a + # downstream consumer cannot accidentally sum double-counted + # employment: the flag travels on the frame, not just in attrs. + cell_key = out["naics_sector"] + "/" + out["canonical_band"] + out["employment_in_scope"] = ~( + cell_key.isin(excluded) + | out["naics_sector"].isin(EMPLOYMENT_OUT_OF_SCOPE_SECTORS) + ) + excluded = excluded | { + f"{row.naics_sector}/{row.canonical_band}" + for row in targets.itertuples() + if row.naics_sector in EMPLOYMENT_OUT_OF_SCOPE_SECTORS + } + scoped_key = scoped["naics_sector"] + "/" + scoped["canonical_band"] + in_scope = scoped[~scoped_key.isin(excluded)] out.attrs["calibrated_cells"] = calibrated + out.attrs["firm_only_cells"] = firm_only out.attrs["infeasible_cells"] = infeasible + out.attrs["weight_bounds"] = weight_bounds + out.attrs["bounded_cell_residuals"] = residuals out.attrs["weighted_firms"] = float(out["weight"].sum()) out.attrs["target_firms"] = int(targets["firms"].sum()) out.attrs["weighted_employment"] = weighted_employment @@ -350,6 +513,18 @@ def calibrate_dual_margin( out.attrs["employment_ratio"] = ( weighted_employment / targets["employment"].sum() ) + # The honest ratio excludes cells whose published employment is out + # of scope for their own size class; including them compares against + # a target that cannot be met by construction. + out.attrs["in_scope_target_employment"] = int(in_scope["employment"].sum()) + in_scope_rows = out[out["employment_in_scope"]] + out.attrs["in_scope_employment_ratio"] = ( + float((in_scope_rows[size_column] * in_scope_rows["weight"]).sum()) + / in_scope["employment"].sum() + ) + out.attrs["employment_out_of_scope_sectors"] = sorted( + EMPLOYMENT_OUT_OF_SCOPE_SECTORS + ) return out From 337a12b9766c92ad208b58cded90db6edd6e8668 Mon Sep 17 00:00:00 2001 From: vahid-ahmadi Date: Wed, 12 Aug 2026 12:20:48 +0100 Subject: [PATCH 3/6] Synthetic roster assignment: connect observed workers to observed firms (#192) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The seam where the two sides meet. Both sides are observed microdata — people and job spells from SIPP/CPS, firms from Form 5500 — and nothing on either side is generated. The link between them is not observed and cannot be from public data. This builds a simulated roster: a reproducible, capacity-constrained allocation of real workers to real firm records inside pre-registered compatibility cells. assign_workers returns firm_instance_id, an artificial key, never an EIN or sponsor name, and sets observed_link=False on the result. The #282 claims boundary — no identified firm effects, coworker sorting, spillovers or AKM decomposition — is unchanged by having observed firm records. Matching uses only the registered cell keys (NAICS sector x canonical band). Earnings, tenure, geography and demographics are deliberately excluded: adding one would turn a capacity allocation into an unregistered imputation. Two design choices worth a referee's eye: - Weighted-to-discrete expansion. A roster needs discrete employers, so each calibrated record is replicated into integer instances. Residual weight is allocated by a seeded Bernoulli draw rather than rounding, so the expected count equals the weight; most calibrated weights sit between 1 and 4, where rounding would bias the firm margin systematically. Measured drift on the real frame: -0.005% across 6,453,306 instances. The consequence stated plainly in the docstring: replicates of one sponsor are not distinct real firms, so any statistic treating them as independent employers measures the replication, not the economy. Those are exactly the E12 statistics phase 2 does not certify. - Workers in a cell with no firm instance are returned unassigned rather than relocated to a neighbouring cell, which would fabricate cross-cell mobility the data does not support. Both steps require an explicit seed and are exactly reproducible; the tests pin reproducibility and seed-sensitivity in both directions. Verified end to end on the real calibrated frame: 789,640 sponsor records expand to 6,453,306 instances carrying 177,275,280 slots; assignment respects every capacity, reports unassigned cells by reason, and is identical under a repeated seed. Co-Authored-By: Claude Opus 5 (1M context) --- src/populace_dynamics/firms/assignment.py | 208 ++++++++++++++++++++++ tests/README-tiers.md | 4 +- tests/test_firms_assignment.py | 159 +++++++++++++++++ tests/tier_counts.json | 2 +- 4 files changed, 370 insertions(+), 3 deletions(-) create mode 100644 src/populace_dynamics/firms/assignment.py create mode 100644 tests/test_firms_assignment.py diff --git a/src/populace_dynamics/firms/assignment.py b/src/populace_dynamics/firms/assignment.py new file mode 100644 index 00000000..251b117f --- /dev/null +++ b/src/populace_dynamics/firms/assignment.py @@ -0,0 +1,208 @@ +"""Synthetic roster assignment — connecting observed workers to observed +firms (#192, phase 0). + +This is the seam where the two sides meet. Both sides are **observed +microdata**: people and job spells come from SIPP/CPS +(``data/sipp_jobs.py``, ``data/asec_firm_size.py``), firms come from +Form 5500 (``firms/frame.py``). Nothing on either side is generated. + +**The link between them is not observed, and never can be from public +data.** No public source connects a SIPP or CPS respondent to their +actual employer. What this module builds is a *simulated* roster: a +reproducible, capacity-constrained allocation of real workers to real +firm records within pre-registered compatibility cells. Every function +here is named to keep that distinction visible, and +:func:`assign_workers` deliberately returns ``firm_instance_id`` — an +artificial key — never an EIN or sponsor name. + +Therefore the output does **not**: reveal any worker's real employer, +identify actual coworkers, support identified firm effects, worker +sorting, spillovers, or AKM-style decompositions. That boundary is the +one recorded on #282 and it is unchanged by having observed firm +records on the other side. + +**The weighted-to-discrete step is a registered design choice.** The +calibrated frame carries fractional weights: a sponsor with weight 15.5 +represents 15.5 firms. A roster needs discrete employers, so +:func:`expand_to_firm_instances` replicates each observed record into +integer instances. This is ordinary weighted representation — the same +move that treats a CPS person with weight 1,200 as 1,200 people — but +it has a consequence worth stating plainly: **the replicates of one +sponsor are not distinct real firms.** They share an observed size and +sector, and any statistic that treats them as independent employers +(coworker correlation, between-firm variance) is measuring the +replication, not the economy. Such statistics are exactly what E12 +covers and what phase 2 does not certify. + +Residual weight is allocated by a seeded Bernoulli draw rather than +rounding, so the expected instance count is the weight and the frame's +firm margin survives expansion in expectation instead of drifting +systematically downward. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd + +from .frame import CELL_KEYS + +__all__ = [ + "expand_to_firm_instances", + "assign_workers", +] + + +def expand_to_firm_instances( + frame: pd.DataFrame, + *, + seed: int, + size_column: str = "active_participants", + weight_column: str = "weight", +) -> pd.DataFrame: + """Replicate weighted sponsor records into discrete firm instances. + + Each record contributes ``floor(weight)`` instances plus one more + with probability ``weight - floor(weight)``, so the expected count + equals the weight exactly. Rounding instead would bias the firm + margin downward, because most weights in the calibrated frame sit + between 1 and 4. + + Returns one row per firm instance with ``firm_instance_id`` (an + artificial key), the source ``sponsor_ein``, the cell keys, and + ``capacity`` taken from the observed size. ``seed`` is required: + an unseeded expansion is not reproducible, and the pre-registration + discipline needs the roster to be re-derivable exactly. + """ + for column in (size_column, weight_column, *CELL_KEYS): + if column not in frame.columns: + raise ValueError(f"Frame lacks {column!r}; use calibrate_*().") + + rng = np.random.default_rng(seed) + weights = frame[weight_column].to_numpy(dtype=float) + if (weights < 0).any(): + raise ValueError("Negative weights cannot be expanded.") + whole = np.floor(weights).astype(int) + extra = (rng.random(len(weights)) < (weights - whole)).astype(int) + counts = whole + extra + + repeated = frame.loc[frame.index.repeat(counts)].reset_index(drop=True) + out = pd.DataFrame( + { + "firm_instance_id": np.arange(len(repeated), dtype=np.int64), + "sponsor_ein": repeated.get("sponsor_ein"), + "naics_sector": repeated["naics_sector"].to_numpy(), + "canonical_band": repeated["canonical_band"].to_numpy(), + "capacity": repeated[size_column].to_numpy(), + } + ) + out.attrs["seed"] = int(seed) + out.attrs["source_records"] = int(len(frame)) + out.attrs["expected_instances"] = float(weights.sum()) + out.attrs["actual_instances"] = int(len(out)) + out.attrs["total_capacity"] = float(out["capacity"].sum()) + return out + + +def assign_workers( + workers: pd.DataFrame, + firm_instances: pd.DataFrame, + *, + seed: int, + worker_weight_column: str | None = None, +) -> pd.DataFrame: + """Allocate workers to firm instances within compatibility cells. + + Workers and firms are matched **only** on the pre-registered cell + keys (NAICS sector x canonical firm-size band). No earnings, + tenure, geography or demographic field enters the match: adding one + would make the assignment informative about attributes the design + has not registered, and would quietly turn a capacity allocation + into an imputation. + + Within each cell the allocation is a seeded permutation of workers + against firm slots, filling each instance up to ``capacity``. A + worker in a cell with no firm instance is returned with a null + ``firm_instance_id`` rather than being dropped or reassigned to a + neighbouring cell — silently relocating them would fabricate + cross-cell mobility that the data does not support. + + Returns the worker frame with ``firm_instance_id`` attached, and + diagnostics on ``attrs``: assignment rate, unassigned counts by + cell, and capacity utilisation. A low fill rate is a real finding + about frame-vs-panel scale, not something to paper over. + """ + for column in CELL_KEYS: + if column not in workers.columns: + raise ValueError(f"Workers lack {column!r}.") + if column not in firm_instances.columns: + raise ValueError(f"Firm instances lack {column!r}.") + if "capacity" not in firm_instances.columns: + raise ValueError("Firm instances lack 'capacity'.") + + rng = np.random.default_rng(seed) + assigned = pd.Series(pd.NA, index=workers.index, dtype="Int64") + unassigned: list[dict[str, object]] = [] + used_capacity = 0.0 + offered_capacity = 0.0 + + firms_by_cell = dict(list(firm_instances.groupby(list(CELL_KEYS)))) + for cell, group in workers.groupby(list(CELL_KEYS)): + firms = firms_by_cell.get(cell) + if firms is None or firms.empty: + unassigned.append( + { + "cell": "/".join(map(str, cell)), + "workers": int(len(group)), + "reason": "no firm instance in cell", + } + ) + continue + # One slot per unit of capacity, shuffled so the allocation does + # not follow record order (which is EIN order, i.e. correlated + # with filing vintage). + capacity = firms["capacity"].to_numpy(dtype=float) + slots = np.repeat( + firms["firm_instance_id"].to_numpy(), + np.maximum(capacity, 0).astype(int), + ) + offered_capacity += float(capacity.sum()) + if len(slots) == 0: + unassigned.append( + { + "cell": "/".join(map(str, cell)), + "workers": int(len(group)), + "reason": "zero capacity in cell", + } + ) + continue + rng.shuffle(slots) + take = min(len(group), len(slots)) + order = rng.permutation(len(group))[:take] + assigned.iloc[ + [workers.index.get_loc(i) for i in group.index[order]] + ] = slots[:take] + used_capacity += take + if take < len(group): + unassigned.append( + { + "cell": "/".join(map(str, cell)), + "workers": int(len(group) - take), + "reason": "cell capacity exhausted", + } + ) + + out = workers.copy() + out["firm_instance_id"] = assigned + matched = int(out["firm_instance_id"].notna().sum()) + out.attrs["seed"] = int(seed) + out.attrs["assigned_workers"] = matched + out.attrs["total_workers"] = int(len(out)) + out.attrs["assignment_rate"] = matched / len(out) if len(out) else 0.0 + out.attrs["unassigned"] = unassigned + out.attrs["offered_capacity"] = offered_capacity + out.attrs["capacity_utilisation"] = ( + used_capacity / offered_capacity if offered_capacity else 0.0 + ) + out.attrs["observed_link"] = False + return out diff --git a/tests/README-tiers.md b/tests/README-tiers.md index 0c344d35..6c0eb19e 100644 --- a/tests/README-tiers.md +++ b/tests/README-tiers.md @@ -38,9 +38,9 @@ pytest --collect-only -q -m oracle_policyengine | tail -1 | Tier | Tests at HEAD | |---|---:| -| `unit` | 1,537 | +| `unit` | 1,551 | | `artifact` | 2,543 | | `integration_psid` | 848 | | `reproduction_legacy` | 520 | | `oracle_policyengine` | 159 | -| **Total** | **5,607** | +| **Total** | **5,621** | diff --git a/tests/test_firms_assignment.py b/tests/test_firms_assignment.py new file mode 100644 index 00000000..713bc3e7 --- /dev/null +++ b/tests/test_firms_assignment.py @@ -0,0 +1,159 @@ +"""Tests for the synthetic roster assignment layer.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from populace_dynamics.firms import assignment + + +def _frame(rows: list[dict]) -> pd.DataFrame: + defaults = { + "sponsor_ein": "123456789", + "naics_sector": "31", + "canonical_band": "B50_99", + "active_participants": 60, + "weight": 1.0, + } + return pd.DataFrame([{**defaults, **row} for row in rows]) + + +def _workers(n: int, sector: str = "31", band: str = "B50_99"): + return pd.DataFrame( + { + "person_id": np.arange(n), + "naics_sector": [sector] * n, + "canonical_band": [band] * n, + } + ) + + +def test_expansion_requires_the_calibrated_columns(): + with pytest.raises(ValueError, match="weight"): + assignment.expand_to_firm_instances( + _frame([{}]).drop(columns=["weight"]), seed=1 + ) + + +def test_expansion_rejects_negative_weights(): + with pytest.raises(ValueError, match="Negative weights"): + assignment.expand_to_firm_instances(_frame([{"weight": -1}]), seed=1) + + +def test_integer_weight_expands_exactly(): + out = assignment.expand_to_firm_instances(_frame([{"weight": 3}]), seed=1) + assert len(out) == 3 + assert set(out["firm_instance_id"]) == {0, 1, 2} + assert (out["capacity"] == 60).all() + + +def test_fractional_weight_is_bernoulli_not_rounded(): + """Rounding would bias the firm margin down. + + Most calibrated weights sit between 1 and 4, so rounding 1.5 to 1 + (or 2) systematically loses (or invents) firms. A seeded Bernoulli + draw makes the expected count equal the weight instead. + """ + counts = [ + len( + assignment.expand_to_firm_instances( + _frame([{"weight": 1.5}]), seed=seed + ) + ) + for seed in range(400) + ] + assert set(counts) == {1, 2} + assert 1.4 < float(np.mean(counts)) < 1.6 + + +def test_expansion_is_reproducible_under_a_pinned_seed(): + kwargs = {"seed": 20260812} + frame = _frame([{"weight": 2.7}, {"weight": 5.2, "naics_sector": "44"}]) + first = assignment.expand_to_firm_instances(frame, **kwargs) + second = assignment.expand_to_firm_instances(frame, **kwargs) + assert first["capacity"].tolist() == second["capacity"].tolist() + assert len(first) == len(second) + + +def test_expansion_seed_is_required(): + with pytest.raises(TypeError): + assignment.expand_to_firm_instances(_frame([{}])) + + +def test_assignment_respects_capacity(): + instances = assignment.expand_to_firm_instances( + _frame([{"weight": 1, "active_participants": 5}]), seed=1 + ) + out = assignment.assign_workers(_workers(12), instances, seed=1) + assert out.attrs["assigned_workers"] == 5 + assert out["firm_instance_id"].notna().sum() == 5 + reasons = [u["reason"] for u in out.attrs["unassigned"]] + assert reasons == ["cell capacity exhausted"] + + +def test_worker_in_a_cell_with_no_firm_is_left_unassigned(): + """Never relocate across cells — that would fabricate mobility.""" + instances = assignment.expand_to_firm_instances( + _frame([{"weight": 1, "naics_sector": "31"}]), seed=1 + ) + out = assignment.assign_workers( + _workers(4, sector="52"), instances, seed=1 + ) + assert out["firm_instance_id"].isna().all() + assert out.attrs["unassigned"][0]["reason"] == "no firm instance in cell" + + +def test_assignment_is_reproducible_and_seed_sensitive(): + instances = assignment.expand_to_firm_instances( + _frame([{"weight": 20, "active_participants": 5}]), seed=1 + ) + workers = _workers(40) + a = assignment.assign_workers(workers, instances, seed=7) + b = assignment.assign_workers(workers, instances, seed=7) + c = assignment.assign_workers(workers, instances, seed=8) + assert a["firm_instance_id"].equals(b["firm_instance_id"]) + assert not a["firm_instance_id"].equals(c["firm_instance_id"]) + + +def test_no_firm_instance_exceeds_its_capacity(): + instances = assignment.expand_to_firm_instances( + _frame([{"weight": 4, "active_participants": 3}]), seed=2 + ) + out = assignment.assign_workers(_workers(12), instances, seed=2) + per_firm = out["firm_instance_id"].value_counts() + assert (per_firm <= 3).all() + + +def test_output_carries_an_artificial_key_not_an_employer_identity(): + """The roster must never look like an observed link.""" + instances = assignment.expand_to_firm_instances(_frame([{}]), seed=1) + out = assignment.assign_workers(_workers(3), instances, seed=1) + assert "firm_instance_id" in out.columns + assert "sponsor_ein" not in out.columns + assert out.attrs["observed_link"] is False + + +def test_assignment_rejects_frames_missing_cell_keys(): + instances = assignment.expand_to_firm_instances(_frame([{}]), seed=1) + bad = _workers(3).drop(columns=["canonical_band"]) + with pytest.raises(ValueError, match="canonical_band"): + assignment.assign_workers(bad, instances, seed=1) + + +def test_assignment_rejects_instances_without_capacity(): + instances = assignment.expand_to_firm_instances(_frame([{}]), seed=1) + with pytest.raises(ValueError, match="capacity"): + assignment.assign_workers( + _workers(3), instances.drop(columns=["capacity"]), seed=1 + ) + + +def test_zero_capacity_cell_is_reported_not_silently_skipped(): + instances = assignment.expand_to_firm_instances( + _frame([{"weight": 1, "active_participants": 0}]), seed=1 + ) + out = assignment.assign_workers(_workers(3), instances, seed=1) + assert out["firm_instance_id"].isna().all() + assert out.attrs["unassigned"][0]["reason"] == "zero capacity in cell" diff --git a/tests/tier_counts.json b/tests/tier_counts.json index 659c1f08..40047f32 100644 --- a/tests/tier_counts.json +++ b/tests/tier_counts.json @@ -1,7 +1,7 @@ { "schema_version": 1, "counts": { - "unit": 1537, + "unit": 1551, "artifact": 2543, "integration_psid": 848, "reproduction_legacy": 520, From 6c40dde3377b68a2d0e14e9b9c15dc301d237546 Mon Sep 17 00:00:00 2001 From: vahid-ahmadi Date: Wed, 12 Aug 2026 13:52:37 +0100 Subject: [PATCH 4/6] IC1 job-spell contract: the checked seam between workstreams A and B (#192) ADR 0003 froze IC1 as "one tidy table, written by workstream A, read by workstream B", but nothing in the repository enforced it. This adds the schema, its validator, and the adapter from the SIPP spell reader, so the two sides meet at a checked contract instead of a convention. firms/assignment.py consumes IC1; a mis-shaped frame now fails here rather than producing a plausible-looking roster. Three contract rules are enforced rather than documented: - The column set is exact, not a minimum. An hours column is rejected by name with its own message, because IC1's hours deferral is live (the registered consumer is SNAP ABAWD compliance, whose 80-hours-per-month test needs month-resolved hours) and a consumer finding an hours column would reasonably assume it was ratified. Adding one is the first scheduled amendment, by joint PR. - person_id must be an opaque string. The ASEC PERIDNUM is 22 digits, so int64 overflows and float64 rounds distinct persons together (#194 review). A numeric key silently merges people. - Self-employed and unpaid-family spells carry no firm-size band. That is a category error, not a missing value, and from_sipp_spells clears any band the raw SIPP slot carried through. calibration_universe applies ADR 0003's universe rule explicitly: private-sector spells only, because SUSB excludes government establishments, NAICS 92, crop/animal production and non-employers, and QWI in-scope jobs are non-federal. Calibrating against jobs the targets never counted would bias every margin by the excluded share. Dropped counts are recorded on attrs so the exclusion is visible in any artifact built from the result. from_sipp_spells performs a named, lossy promotion rather than a rename: SIPP 2014+ measures establishment size at the worker's location while IC2's canonical variable means enterprise size, so the adapter records the proxy on attrs["size_concept"]. That promotion is the most consequential approximation on the person side and it is not hidden behind a column name. Verified end to end against the real calibrated firm frame: 5,000 IC1 spells validate, the calibration universe drops federal and self-employed spells with counts recorded, and the survivors assign to observed firm instances at rate 1.0000 with observed_link=False. Co-Authored-By: Claude Opus 5 (1M context) --- src/populace_dynamics/firms/ic1.py | 238 +++++++++++++++++++++++++++++ tests/README-tiers.md | 4 +- tests/test_firms_ic1.py | 181 ++++++++++++++++++++++ tests/tier_counts.json | 2 +- 4 files changed, 422 insertions(+), 3 deletions(-) create mode 100644 src/populace_dynamics/firms/ic1.py create mode 100644 tests/test_firms_ic1.py diff --git a/src/populace_dynamics/firms/ic1.py b/src/populace_dynamics/firms/ic1.py new file mode 100644 index 00000000..6d41d2d9 --- /dev/null +++ b/src/populace_dynamics/firms/ic1.py @@ -0,0 +1,238 @@ +"""IC1 job-spell contract — the seam between workstreams A and B (#192). + +ADR 0003 froze IC1 on 2026-07-16 as "one tidy table, written by +workstream A, read by workstream B". This module is that table's +schema, its validator, and the adapter from the SIPP spell reader's +output. It exists so the two sides meet at a checked contract rather +than at a convention: ``firms/assignment.py`` consumes IC1, and +anything that does not conform fails here rather than producing a +plausible-looking roster from mis-shaped input. + +**The calibration universe is narrower than the schema.** ADR 0003 +makes ``class_of_worker`` load-bearing: SUSB excludes government +establishments, NAICS 92, crop/animal production and non-employers, +and QWI in-scope jobs are non-federal. Government, self-employed and +unpaid-family spells are therefore **out of the SUSB/QWI calibration +universe**, and self-employed spells have **no defined firm-size +band** at all. :func:`calibration_universe` applies that rule +explicitly so a caller cannot drift into calibrating against jobs the +targets never counted. + +**IC1 carries no geography, by design.** QWI/J2J targets are +state-level, but the state of a spell is the host person's state at +``start_period``, joined on ``person_id`` from the person table. The +join key lives on the person table, not here. + +**IC1 carries no hours, and that deferral is live.** The frozen +contract has no hours column, so IC1 cannot serve monthly-hours +consumers — the registered example is SNAP ABAWD compliance, whose +80-hours-per-month test needs month-resolved hours. A ``hours_band`` +column is the first scheduled IC1 amendment, by joint PR, once the +phase-1 spell imputation establishes what granularity SIPP supports. +:func:`validate` rejects a frame that smuggles in an hours column, +because adding one silently would be an unratified contract change. +""" + +from __future__ import annotations + +import pandas as pd + +from .banding import CanonicalBand + +__all__ = [ + "IC1_COLUMNS", + "CLASS_OF_WORKER", + "CALIBRATION_CLASSES", + "NO_FIRM_SIZE_CLASSES", + "OPEN_SPELL_SENTINEL", + "validate", + "from_sipp_spells", + "calibration_universe", +] + +#: The frozen IC1 column list, in contract order (ADR 0003). +IC1_COLUMNS = ( + "person_id", + "spell_id", + "start_period", + "end_period", + "industry", + "firm_size_band", + "class_of_worker", + "earnings_share", + "primary_job", +) + +#: The five permitted ``class_of_worker`` values. +CLASS_OF_WORKER = frozenset( + { + "private", + "federal", + "state_local_government", + "self_employed", + "unpaid_family", + } +) + +#: Classes inside the SUSB/QWI calibration universe. Government, +#: self-employed and unpaid-family jobs are outside it (ADR 0003). +CALIBRATION_CLASSES = frozenset({"private"}) + +#: Classes for which a firm-size band is undefined, not merely missing. +NO_FIRM_SIZE_CLASSES = frozenset({"self_employed", "unpaid_family"}) + +#: ``end_period`` value marking a spell still open at panel end. +OPEN_SPELL_SENTINEL = pd.NaT + +_BAND_NAMES = frozenset(band.name for band in CanonicalBand) + + +def validate(spells: pd.DataFrame, *, strict_universe: bool = True) -> None: + """Raise unless ``spells`` conforms to the frozen IC1 contract. + + Checks the column set exactly, not as a subset: an extra column is + a contract change and must go through a joint PR. In particular an + hours column is rejected by name, because IC1's hours deferral is + explicit and a consumer that found one would reasonably assume it + was ratified. + + ``strict_universe`` additionally enforces ADR 0003's semantic + rules: ``person_id`` is an opaque string (the ASEC ``PERIDNUM`` is + 22 digits, so int64 overflows and float64 merges distinct people), + ``spell_id`` is unique within person, and self-employed and + unpaid-family spells carry no firm-size band. + """ + missing = [c for c in IC1_COLUMNS if c not in spells.columns] + if missing: + raise ValueError(f"IC1 frame is missing columns {missing}.") + extra = [c for c in spells.columns if c not in IC1_COLUMNS] + if extra: + hours = [c for c in extra if "hour" in c.lower()] + if hours: + raise ValueError( + f"IC1 frame carries hours column(s) {hours}. IC1 as frozen " + "has no hours column; adding one is the first scheduled " + "amendment and requires a joint PR (ADR 0003). A consumer " + "finding this column would assume it was ratified." + ) + raise ValueError( + f"IC1 frame carries unregistered columns {extra}; the contract " + "is an exact column set, not a minimum." + ) + + bad_class = set(spells["class_of_worker"].dropna()) - CLASS_OF_WORKER + if bad_class: + raise ValueError( + f"Unknown class_of_worker values {sorted(bad_class)}." + ) + + bands = set(spells["firm_size_band"].dropna()) + unknown_bands = bands - _BAND_NAMES + if unknown_bands: + raise ValueError( + f"firm_size_band values {sorted(unknown_bands)} are not canonical " + "IC2 bands." + ) + + if not strict_universe: + return + + if not pd.api.types.is_string_dtype(spells["person_id"]): + raise ValueError( + "person_id must be an opaque string key. The ASEC PERIDNUM is " + "22 digits: int64 overflows and float64 rounds distinct persons " + "together (#194 review)." + ) + if spells.duplicated(["person_id", "spell_id"]).any(): + raise ValueError("spell_id must be unique within person_id.") + + undefined = spells["class_of_worker"].isin(NO_FIRM_SIZE_CLASSES) + if spells.loc[undefined, "firm_size_band"].notna().any(): + raise ValueError( + "Self-employed and unpaid-family spells have no defined " + "firm-size band (ADR 0003); a band here is a category error, " + "not a value." + ) + + shares = spells["earnings_share"].dropna() + if len(shares) and not shares.between(0.0, 1.0).all(): + raise ValueError("earnings_share must lie in [0, 1].") + + +def from_sipp_spells( + spells: pd.DataFrame, *, band_column: str = "estab_size_band" +) -> pd.DataFrame: + """Adapt ``sipp_jobs.job_spells`` output to the IC1 contract. + + The SIPP reader deliberately names its size column + ``estab_size_band``, not ``canonical_band``: SIPP 2014+ measures + **establishment** size at the worker's location, while IC2's + canonical variable means administrative *enterprise* size + (``firms/banding.py``). This adapter therefore performs a named, + lossy promotion rather than a rename, and records it on + ``attrs["size_concept"]`` so a downstream consumer can see that the + band it received is an establishment-size proxy. + + That promotion is the single most consequential approximation on + the person side and it is not hidden behind a column name: a + multi-establishment firm's worker reports their location's + headcount, biasing the spell toward smaller bands. + """ + if band_column not in spells.columns: + raise ValueError( + f"SIPP spells lack {band_column!r}; pass the output of " + "sipp_jobs.job_spells()." + ) + out = pd.DataFrame( + { + "person_id": spells["person_id"].astype("string"), + "spell_id": spells.get( + "spell_id", pd.RangeIndex(len(spells)) + ).astype("int64"), + "start_period": spells.get("start_month", spells.get("start")), + "end_period": spells.get("end_month", spells.get("end")), + "industry": spells["industry"].astype("string"), + "firm_size_band": spells[band_column], + "class_of_worker": spells["class_of_worker"], + "earnings_share": pd.to_numeric( + spells.get("earnings_share"), errors="coerce" + ), + "primary_job": spells.get("top_earner", False).astype(bool), + } + ) + # A band on a self-employed spell is a category error; the SIPP + # reader can carry one through from the raw slot, so it is cleared + # here rather than left to trip the validator. + undefined = out["class_of_worker"].isin(NO_FIRM_SIZE_CLASSES) + out.loc[undefined, "firm_size_band"] = pd.NA + out.attrs["size_concept"] = ( + "establishment size (SIPP EJB1_EMPSIZE) promoted to the IC2 " + "enterprise-size band; a proxy, biased toward smaller bands for " + "multi-establishment firms" + ) + return out[list(IC1_COLUMNS)] + + +def calibration_universe(spells: pd.DataFrame) -> pd.DataFrame: + """Restrict IC1 spells to the SUSB/QWI calibration universe. + + Keeps private-sector spells only. Government (federal and + state-local), self-employed and unpaid-family spells are dropped + because the targets never counted them: SUSB excludes government + establishments, NAICS 92, crop/animal production and non-employers, + and QWI in-scope jobs are non-federal (ADR 0003, #192 review point + 1). Calibrating against jobs the targets exclude would bias every + margin by the size of the excluded share. + + The dropped counts are recorded on ``attrs`` so the exclusion is + visible in any artifact built from the result. + """ + validate(spells, strict_universe=False) + keep = spells["class_of_worker"].isin(CALIBRATION_CLASSES) + out = spells[keep].reset_index(drop=True).copy() + dropped = spells.loc[~keep, "class_of_worker"].value_counts() + out.attrs["excluded_from_calibration"] = { + str(k): int(v) for k, v in dropped.items() + } + out.attrs["excluded_total"] = int((~keep).sum()) + return out diff --git a/tests/README-tiers.md b/tests/README-tiers.md index 6c0eb19e..ccef2bca 100644 --- a/tests/README-tiers.md +++ b/tests/README-tiers.md @@ -38,9 +38,9 @@ pytest --collect-only -q -m oracle_policyengine | tail -1 | Tier | Tests at HEAD | |---|---:| -| `unit` | 1,551 | +| `unit` | 1,568 | | `artifact` | 2,543 | | `integration_psid` | 848 | | `reproduction_legacy` | 520 | | `oracle_policyengine` | 159 | -| **Total** | **5,621** | +| **Total** | **5,638** | diff --git a/tests/test_firms_ic1.py b/tests/test_firms_ic1.py new file mode 100644 index 00000000..267e61f0 --- /dev/null +++ b/tests/test_firms_ic1.py @@ -0,0 +1,181 @@ +"""Tests for the IC1 job-spell contract.""" + +from __future__ import annotations + +import pandas as pd +import pytest + +from populace_dynamics.firms import ic1 + + +def _ic1(rows: list[dict]) -> pd.DataFrame: + defaults = { + "person_id": "0000000000000000000001", + "spell_id": 1, + "start_period": pd.Period("2023-01", freq="M"), + "end_period": pd.Period("2023-06", freq="M"), + "industry": "31", + "firm_size_band": "B50_99", + "class_of_worker": "private", + "earnings_share": 1.0, + "primary_job": True, + } + frame = pd.DataFrame([{**defaults, **row} for row in rows]) + frame["person_id"] = frame["person_id"].astype("string") + return frame[list(ic1.IC1_COLUMNS)] + + +def test_valid_frame_passes(): + ic1.validate(_ic1([{}])) + + +def test_missing_column_rejected(): + with pytest.raises(ValueError, match="missing columns"): + ic1.validate(_ic1([{}]).drop(columns=["earnings_share"])) + + +def test_hours_column_rejected_by_name(): + """IC1's hours deferral is explicit; a stray column would read as + ratified.""" + frame = _ic1([{}]) + frame["hours_band"] = "30-39" + with pytest.raises(ValueError, match="hours column"): + ic1.validate(frame) + + +def test_unregistered_extra_column_rejected(): + frame = _ic1([{}]) + frame["state"] = "OH" + with pytest.raises(ValueError, match="unregistered columns"): + ic1.validate(frame) + + +def test_geography_is_not_an_ic1_column(): + """IC1 deliberately carries no geography; it joins from the person + table.""" + assert "state" not in ic1.IC1_COLUMNS + assert not any("geo" in c or "state" in c for c in ic1.IC1_COLUMNS) + + +def test_numeric_person_id_rejected(): + frame = _ic1([{}]) + frame["person_id"] = [1] + with pytest.raises(ValueError, match="opaque string"): + ic1.validate(frame) + + +def test_spell_id_must_be_unique_within_person(): + frame = _ic1([{"spell_id": 1}, {"spell_id": 1}]) + with pytest.raises(ValueError, match="unique within"): + ic1.validate(frame) + + +def test_same_spell_id_across_persons_is_fine(): + ic1.validate( + _ic1( + [ + {"person_id": "a", "spell_id": 1}, + {"person_id": "b", "spell_id": 1}, + ] + ) + ) + + +def test_self_employed_may_not_carry_a_firm_size_band(): + frame = _ic1( + [{"class_of_worker": "self_employed", "firm_size_band": "LT10"}] + ) + with pytest.raises(ValueError, match="no defined"): + ic1.validate(frame) + + +def test_self_employed_without_a_band_is_valid(): + ic1.validate( + _ic1([{"class_of_worker": "self_employed", "firm_size_band": None}]) + ) + + +def test_unknown_class_of_worker_rejected(): + with pytest.raises(ValueError, match="class_of_worker"): + ic1.validate(_ic1([{"class_of_worker": "contractor"}])) + + +def test_non_canonical_band_rejected(): + with pytest.raises(ValueError, match="canonical IC2 bands"): + ic1.validate(_ic1([{"firm_size_band": "10-49"}])) + + +def test_earnings_share_out_of_range_rejected(): + with pytest.raises(ValueError, match=r"\[0, 1\]"): + ic1.validate(_ic1([{"earnings_share": 1.4}])) + + +def test_calibration_universe_keeps_private_only(): + frame = _ic1( + [ + {"spell_id": 1, "class_of_worker": "private"}, + {"spell_id": 2, "class_of_worker": "federal"}, + { + "spell_id": 3, + "class_of_worker": "self_employed", + "firm_size_band": None, + }, + {"spell_id": 4, "class_of_worker": "state_local_government"}, + ] + ) + out = ic1.calibration_universe(frame) + assert list(out["class_of_worker"]) == ["private"] + assert out.attrs["excluded_total"] == 3 + assert out.attrs["excluded_from_calibration"]["federal"] == 1 + + +def test_from_sipp_spells_clears_bands_on_self_employed(): + sipp = pd.DataFrame( + { + "person_id": ["a", "b"], + "spell_id": [1, 1], + "start_month": [ + pd.Period("2023-01", freq="M"), + pd.Period("2023-02", freq="M"), + ], + "end_month": [ + pd.Period("2023-05", freq="M"), + pd.Period("2023-08", freq="M"), + ], + "industry": ["31", "44"], + "estab_size_band": ["B50_99", "LT10"], + "class_of_worker": ["private", "self_employed"], + "earnings_share": [1.0, 1.0], + "top_earner": [True, True], + } + ) + out = ic1.from_sipp_spells(sipp) + ic1.validate(out) + assert list(out.columns) == list(ic1.IC1_COLUMNS) + assert pd.isna(out.loc[1, "firm_size_band"]) + + +def test_from_sipp_spells_records_the_size_concept_promotion(): + """Establishment size is not enterprise size; the proxy must be + visible.""" + sipp = pd.DataFrame( + { + "person_id": ["a"], + "spell_id": [1], + "start_month": [pd.Period("2023-01", freq="M")], + "end_month": [pd.Period("2023-05", freq="M")], + "industry": ["31"], + "estab_size_band": ["B50_99"], + "class_of_worker": ["private"], + "earnings_share": [1.0], + "top_earner": [True], + } + ) + out = ic1.from_sipp_spells(sipp) + assert "establishment size" in out.attrs["size_concept"] + assert "proxy" in out.attrs["size_concept"] + + +def test_from_sipp_spells_rejects_a_frame_without_the_band_column(): + with pytest.raises(ValueError, match="estab_size_band"): + ic1.from_sipp_spells(pd.DataFrame({"person_id": ["a"]})) diff --git a/tests/tier_counts.json b/tests/tier_counts.json index 40047f32..dd0888b3 100644 --- a/tests/tier_counts.json +++ b/tests/tier_counts.json @@ -1,7 +1,7 @@ { "schema_version": 1, "counts": { - "unit": 1551, + "unit": 1568, "artifact": 2543, "integration_psid": 848, "reproduction_legacy": 520, From 5b89375695285ed9baa44890c8c3bc3df97f2906 Mon Sep 17 00:00:00 2001 From: vahid-ahmadi Date: Wed, 12 Aug 2026 14:30:54 +0100 Subject: [PATCH 5/6] Phase-0 job-spell imputation: SIPP donors onto CPS persons (#192) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Workstream A's remaining deliverable. Real ASEC persons receive IC1 job spells drawn from real SIPP donor records through a quantile regression forest, the same microimpute recipe ECPS uses for earnings histories. Completes the person side of the seam: imputed spells pass ic1.validate, feed ic1.calibration_universe, and assign to observed firm instances. The ADR 0003 firm-size x tenure bridge is an explicit argument with no default. A run that does not name its bridge cannot be refereed, and the two ratified bridges (SIPP 2008 primary, SIPP 2014+ proxy chain) give different joint structure at the same seed. ## A silent-failure bug found by the seed-sensitivity test The first implementation drew every target "at a quantile". That is wrong for a categorical target and wrong in a way that looks like it works: microimpute returns the same modal class at every quantile — measured, the band-code mean is 1.505 at q=0.05 and at q=0.95 alike. So every host person in a predictor cell received the *same* band, the cross-sectional variance collapsed to zero, and the seed had no effect whatsoever. A deterministic modal assignment was being presented as an imputation. It only surfaced because a test asserted that a *different* seed must change the draw. Reproducibility tests alone would have passed happily, since a constant is trivially reproducible. The draw is now split by target kind: - categorical (the firm-size band) samples each row from its predicted class distribution, via return_probs=True; - continuous (tenure, earnings share) draws by inverse-CDF on a fixed quantile grid. Verified: the band draw now varies across seeds and reproduces the donor's band distribution shape. ## Boundaries enforced, not just documented - class_of_worker is carried from the host, never imputed: it decides the calibration universe and whether a band is even defined, so a drawn value would let the imputation choose which jobs the SUSB/QWI targets count. - Self-employed and unpaid-family hosts never receive a band. - A predictor present only on the donor is rejected rather than dropped, because dropping it changes the joint structure the bridge exists to supply. - Imputing over a column the host already observes is rejected: overwriting measurement with a draw must be deliberate. - attrs record bridge, seed, predictors and band_is_imputed=True, so a consumer cannot mistake a drawn band for a measured one. microimpute added to dependencies; it was used but never declared. Co-Authored-By: Claude Opus 5 (1M context) --- pyproject.toml | 1 + .../firms/spell_imputation.py | 308 ++++++++++++++++++ tests/README-tiers.md | 4 +- tests/test_firms_spell_imputation.py | 169 ++++++++++ tests/tier_counts.json | 2 +- 5 files changed, 481 insertions(+), 3 deletions(-) create mode 100644 src/populace_dynamics/firms/spell_imputation.py create mode 100644 tests/test_firms_spell_imputation.py diff --git a/pyproject.toml b/pyproject.toml index 370d161a..9a2eb1e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ dependencies = [ "pandas>=2.2.0", "scipy>=1.10.0", "scikit-learn>=1.3.0", + "microimpute>=1.1.2", "pyyaml>=6.0", ] diff --git a/src/populace_dynamics/firms/spell_imputation.py b/src/populace_dynamics/firms/spell_imputation.py new file mode 100644 index 00000000..68824fde --- /dev/null +++ b/src/populace_dynamics/firms/spell_imputation.py @@ -0,0 +1,308 @@ +"""Phase-0 job-spell imputation — SIPP donors onto CPS persons (#192). + +Workstream A's deliverable to workstream B: a CPS-wide person panel +carrying IC1 job spells. The persons are real ASEC respondents and the +donor spells are real SIPP job records; what is modelled is which +donor's *attributes* attach to which host person, using a quantile +regression forest (``microimpute``) exactly as the ECPS recipe does +for earnings histories. + +**The firm-size x tenure bridge is named, not implicit.** ADR 0003 +records that no current representative source observes firm size and +tenure jointly: ASEC ``FIRMSIZE`` refers to the preceding calendar +year's longest job, the biennial tenure supplement refers to the +current job and asks no firm-size question, SIPP 2014+ has tenure but +only *establishment* size, and NLSY has both for two +non-representative cohorts. The ratified bridge is therefore: + +1. **primary** — the pre-redesign SIPP 2008 panel (2008-2013), the + last representative panel observing worker-reported firm size at + all locations alongside spells and tenure, its dated joint + structure aged forward; +2. **proxy chain** — SIPP 2014+ establishment size x tenure mapped + through the establishment-to-enterprise noise model implied by the + IC2 semantics; +3. **caveat** — the ASEC reference-period mismatch is carried into the + IC3 gate notes as a known label-misalignment term. + +:class:`SpellImputationSpec` makes the choice an explicit argument. It +has no default: a run that does not say which bridge it used cannot be +refereed, and the two bridges give different joint structure. + +**Imputing a band is not observing one.** The output is IC1-conforming +and feeds ``firms/assignment.py``, but a host person's imputed +``firm_size_band`` is a draw conditional on their observed +characteristics, not a measurement. Anything that treats it as +measured — a firm-size threshold count presented as a headcount rather +than an estimate — is over-reading the panel. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import pandas as pd + +from . import ic1 + +__all__ = [ + "SpellImputationSpec", + "BRIDGES", + "DEFAULT_PREDICTORS", + "IMPUTED_VARIABLES", + "check_frames", + "fit_spell_model", + "impute_spells", +] + +#: Ratified bridge identifiers (ADR 0003 conditioning DAG). +BRIDGES = ("sipp_2008_primary", "sipp_2014_proxy_chain") + +#: Person characteristics conditioning the draw. They must be present +#: and comparably coded on *both* the donor and the host frame — a +#: predictor observed only on the donor cannot condition anything, and +#: silently dropping it changes the joint structure the bridge exists +#: to supply. +DEFAULT_PREDICTORS = ( + "age", + "sex", + "education", + "industry_sector", + "annual_earnings", +) + +#: The spell attributes drawn from the donor. +IMPUTED_VARIABLES = ("firm_size_band_code", "tenure_months", "earnings_share") + + +@dataclass(frozen=True) +class SpellImputationSpec: + """A registered imputation run. + + ``bridge`` and ``seed`` are required. The seed makes the draw + reproducible, which the pre-registration discipline needs; the + bridge records which joint firm-size x tenure structure was used, + without which two runs are not comparable even at the same seed. + """ + + bridge: str + seed: int + predictors: tuple[str, ...] = DEFAULT_PREDICTORS + imputed_variables: tuple[str, ...] = IMPUTED_VARIABLES + #: Resolution of the inverse-CDF draw. The conditional + #: distribution is evaluated on this many quantiles and each host + #: person selects one; coarser grids quantise the draw. + quantile_grid: int = 200 + + def __post_init__(self) -> None: + if self.bridge not in BRIDGES: + raise ValueError( + f"Unknown bridge {self.bridge!r}; ADR 0003 ratifies " + f"{list(BRIDGES)}. A run must name its bridge: the two " + "give different firm-size x tenure joint structure." + ) + if not self.predictors: + raise ValueError("At least one predictor is required.") + + +def check_frames( + donor: pd.DataFrame, hosts: pd.DataFrame, spec: SpellImputationSpec +) -> None: + """Check donor and host frames are compatible before fitting. + + Called by :func:`fit_spell_model` when ``hosts`` is supplied, and + public so a caller can check compatibility before paying for a fit. + """ + missing_donor = [c for c in spec.predictors if c not in donor.columns] + missing_host = [c for c in spec.predictors if c not in hosts.columns] + if missing_donor: + raise ValueError(f"Donor frame lacks predictors {missing_donor}.") + if missing_host: + raise ValueError( + f"Host frame lacks predictors {missing_host}. A predictor " + "present only on the donor cannot condition the draw; dropping " + "it silently would change the joint structure." + ) + missing_targets = [ + c for c in spec.imputed_variables if c not in donor.columns + ] + if missing_targets: + raise ValueError( + f"Donor frame lacks imputed variables {missing_targets}." + ) + leaked = [c for c in spec.imputed_variables if c in hosts.columns] + if leaked: + raise ValueError( + f"Host frame already carries imputed variable(s) {leaked}. " + "Imputing over an observed column would overwrite measurement " + "with a draw; drop or rename the host column deliberately." + ) + + +def fit_spell_model( + donor: pd.DataFrame, + spec: SpellImputationSpec, + *, + weight_column: str | None = None, + hosts: pd.DataFrame | None = None, +): + """Fit the quantile regression forest on donor spells. + + ``donor`` is one row per donor job spell carrying the predictors + and the imputed variables. ``weight_column`` passes the donor's + survey weight through to ``microimpute``, which resamples on it — + omitting it fits the model to the unweighted SIPP sample, which is + not representative. + """ + if hosts is not None: + check_frames(donor, hosts, spec) + else: + missing = [c for c in spec.predictors if c not in donor.columns] + if missing: + raise ValueError(f"Donor frame lacks predictors {missing}.") + missing = [c for c in spec.imputed_variables if c not in donor.columns] + if missing: + raise ValueError(f"Donor frame lacks imputed variables {missing}.") + + from microimpute.models import QRF + + model = QRF() + return model.fit( + X_train=donor, + predictors=list(spec.predictors), + imputed_variables=list(spec.imputed_variables), + weight_col=weight_column, + ) + + +def impute_spells( + fitted, + hosts: pd.DataFrame, + spec: SpellImputationSpec, + *, + person_id_column: str = "person_id", + class_of_worker_column: str = "class_of_worker", + band_codes: dict[int, str] | None = None, +) -> pd.DataFrame: + """Draw spell attributes onto host persons and emit IC1 spells. + + ``class_of_worker`` is **carried from the host, never imputed**. + It determines the calibration universe (ADR 0003) and whether a + firm-size band is even defined, so drawing it from a donor would + let the imputation decide which jobs the SUSB/QWI targets count. + + The returned frame conforms to IC1 and is validated before return, + so a malformed draw fails here rather than downstream in the + roster. + """ + if person_id_column not in hosts.columns: + raise ValueError(f"Host frame lacks {person_id_column!r}.") + if class_of_worker_column not in hosts.columns: + raise ValueError( + f"Host frame lacks {class_of_worker_column!r}; class of worker " + "is carried from the host, never imputed (ADR 0003)." + ) + + # Two target kinds need two different draws, and conflating them is + # a silent-failure trap. + # + # A *categorical* target (the firm-size band) is not quantile- + # addressable: microimpute returns the same modal class at every + # quantile — measured, band code mean 1.505 at q=0.05 and at q=0.95 + # alike. Drawing it "at a quantile" therefore yields a deterministic + # modal assignment: every host person in a predictor cell gets the + # same band, the cross-sectional variance collapses to zero, and the + # seed has no effect at all. That looks like a working imputation and + # is not one. The correct draw samples each row from its predicted + # class distribution, which ``return_probs=True`` exposes. + # + # A *continuous* target (tenure, earnings share) is quantile- + # addressable, so it is drawn by inverse-CDF on a fixed grid. + rng = np.random.default_rng(spec.seed) + grid = np.round(np.linspace(0.005, 0.995, spec.quantile_grid), 4) + predicted = fitted.predict( + hosts[list(spec.predictors)], list(grid), return_probs=True + ) + if not isinstance(predicted, dict): + raise TypeError( + "Expected microimpute predict() to return {quantile: frame}; " + f"got {type(predicted)!r}. The per-row draw depends on that " + "contract." + ) + probabilities = predicted.pop("probabilities", {}) or {} + frames = { + float(q): pd.DataFrame(f).reset_index(drop=True) + for q, f in predicted.items() + } + keys = np.array(sorted(frames)) + picks = np.searchsorted(keys, rng.random(len(hosts)), side="left").clip( + 0, len(keys) - 1 + ) + stacked = np.stack([frames[k].to_numpy() for k in keys]) + drawn = pd.DataFrame( + stacked[picks, np.arange(len(hosts))], + columns=frames[keys[0]].columns, + ) + + for name, payload in probabilities.items(): + if name not in drawn.columns: + continue + classes = np.asarray(payload["classes"]) + weights = np.asarray(payload["probabilities"], dtype=float) + totals = weights.sum(axis=1, keepdims=True) + if not np.isfinite(totals).all() or (totals <= 0).any(): + raise ValueError( + f"Predicted class distribution for {name!r} is degenerate; " + "cannot draw." + ) + cumulative = np.cumsum(weights / totals, axis=1) + uniforms = rng.random(len(drawn))[:, None] + chosen = (uniforms > cumulative).sum(axis=1).clip(0, len(classes) - 1) + drawn[name] = classes[chosen] + + codes = band_codes or {} + bands = drawn.get("firm_size_band_code") + if bands is None: + raise ValueError("Imputation produced no firm_size_band_code column.") + band_names = pd.Series( + [codes.get(int(round(v))) for v in bands], dtype="object" + ) + + cow = hosts[class_of_worker_column].reset_index(drop=True) + out = pd.DataFrame( + { + "person_id": hosts[person_id_column] + .reset_index(drop=True) + .astype("string"), + "spell_id": 1, + "start_period": hosts.get( + "spell_start", pd.Series([pd.NaT] * len(hosts)) + ).reset_index(drop=True), + "end_period": hosts.get( + "spell_end", pd.Series([pd.NaT] * len(hosts)) + ).reset_index(drop=True), + "industry": hosts["industry_sector"] + .reset_index(drop=True) + .astype("string"), + "firm_size_band": band_names, + "class_of_worker": cow, + "earnings_share": pd.to_numeric( + drawn.get("earnings_share"), errors="coerce" + ).clip(0.0, 1.0), + "primary_job": True, + } + ) + # Self-employed and unpaid-family spells have no defined band; the + # draw does not get to invent one (ADR 0003). + undefined = out["class_of_worker"].isin(ic1.NO_FIRM_SIZE_CLASSES) + out.loc[undefined, "firm_size_band"] = pd.NA + + out = out[list(ic1.IC1_COLUMNS)] + ic1.validate(out) + out.attrs["bridge"] = spec.bridge + out.attrs["seed"] = spec.seed + out.attrs["predictors"] = list(spec.predictors) + out.attrs["donor_observed"] = True + out.attrs["band_is_imputed"] = True + return out diff --git a/tests/README-tiers.md b/tests/README-tiers.md index ccef2bca..be545179 100644 --- a/tests/README-tiers.md +++ b/tests/README-tiers.md @@ -38,9 +38,9 @@ pytest --collect-only -q -m oracle_policyengine | tail -1 | Tier | Tests at HEAD | |---|---:| -| `unit` | 1,568 | +| `unit` | 1,582 | | `artifact` | 2,543 | | `integration_psid` | 848 | | `reproduction_legacy` | 520 | | `oracle_policyengine` | 159 | -| **Total** | **5,638** | +| **Total** | **5,652** | diff --git a/tests/test_firms_spell_imputation.py b/tests/test_firms_spell_imputation.py new file mode 100644 index 00000000..431086a3 --- /dev/null +++ b/tests/test_firms_spell_imputation.py @@ -0,0 +1,169 @@ +"""Tests for the phase-0 job-spell imputation.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from populace_dynamics.firms import ic1 +from populace_dynamics.firms import spell_imputation as si + +BANDS = { + 0: "LT10", + 1: "B10_49", + 2: "B50_99", + 3: "B100_499", + 4: "B500_PLUS", +} + + +def _donor(n: int = 300, seed: int = 0) -> pd.DataFrame: + rng = np.random.default_rng(seed) + frame = pd.DataFrame( + { + "age": rng.integers(18, 65, n), + "sex": rng.integers(0, 2, n), + "education": rng.integers(1, 5, n), + "industry_sector": rng.choice(["31", "44"], n), + "annual_earnings": rng.lognormal(10.5, 0.5, n), + "weight": rng.uniform(50, 500, n), + } + ) + latent = frame["annual_earnings"] / 40000 + frame["education"] * 0.4 + frame["firm_size_band_code"] = np.clip(latent.astype(int), 0, 4) + frame["tenure_months"] = np.clip(latent * 12, 0, 480) + frame["earnings_share"] = np.clip(rng.beta(8, 1, n), 0, 1) + return frame + + +def _hosts(n: int = 200, seed: int = 1) -> pd.DataFrame: + rng = np.random.default_rng(seed) + return pd.DataFrame( + { + "person_id": [f"{i:022d}" for i in range(n)], + "age": rng.integers(18, 65, n), + "sex": rng.integers(0, 2, n), + "education": rng.integers(1, 5, n), + "industry_sector": rng.choice(["31", "44"], n), + "annual_earnings": rng.lognormal(10.5, 0.5, n), + "class_of_worker": rng.choice( + ["private", "federal", "self_employed"], n, p=[0.8, 0.1, 0.1] + ), + } + ) + + +def test_bridge_must_be_one_of_the_ratified_names(): + with pytest.raises(ValueError, match="Unknown bridge"): + si.SpellImputationSpec(bridge="nlsy", seed=1) + + +def test_bridge_and_seed_have_no_defaults(): + """A run that does not name its bridge and seed cannot be refereed.""" + with pytest.raises(TypeError): + si.SpellImputationSpec() + with pytest.raises(TypeError): + si.SpellImputationSpec(bridge="sipp_2008_primary") + + +def test_both_ratified_bridges_are_accepted(): + for bridge in si.BRIDGES: + assert si.SpellImputationSpec(bridge=bridge, seed=1).bridge == bridge + + +def test_predictor_missing_from_host_is_rejected(): + spec = si.SpellImputationSpec(bridge="sipp_2008_primary", seed=1) + with pytest.raises(ValueError, match="Host frame lacks predictors"): + si.check_frames(_donor(), _hosts().drop(columns=["age"]), spec) + + +def test_predictor_missing_from_donor_is_rejected(): + spec = si.SpellImputationSpec(bridge="sipp_2008_primary", seed=1) + with pytest.raises(ValueError, match="Donor frame lacks predictors"): + si.check_frames(_donor().drop(columns=["age"]), _hosts(), spec) + + +def test_imputing_over_an_observed_host_column_is_rejected(): + """Overwriting measurement with a draw must be deliberate.""" + spec = si.SpellImputationSpec(bridge="sipp_2008_primary", seed=1) + hosts = _hosts() + hosts["tenure_months"] = 24 + with pytest.raises(ValueError, match="already carries imputed"): + si.check_frames(_donor(), hosts, spec) + + +def test_missing_class_of_worker_is_rejected(): + spec = si.SpellImputationSpec(bridge="sipp_2008_primary", seed=1) + fitted = si.fit_spell_model(_donor(), spec, weight_column="weight") + with pytest.raises(ValueError, match="class of worker"): + si.impute_spells( + fitted, + _hosts().drop(columns=["class_of_worker"]), + spec, + band_codes=BANDS, + ) + + +def test_imputation_emits_valid_ic1(): + spec = si.SpellImputationSpec(bridge="sipp_2014_proxy_chain", seed=7) + fitted = si.fit_spell_model(_donor(), spec, weight_column="weight") + out = si.impute_spells(fitted, _hosts(), spec, band_codes=BANDS) + ic1.validate(out) + assert list(out.columns) == list(ic1.IC1_COLUMNS) + assert len(out) == 200 + + +def test_self_employed_never_receives_an_imputed_band(): + """The draw does not get to invent a band that is undefined.""" + spec = si.SpellImputationSpec(bridge="sipp_2014_proxy_chain", seed=7) + fitted = si.fit_spell_model(_donor(), spec, weight_column="weight") + out = si.impute_spells(fitted, _hosts(), spec, band_codes=BANDS) + undefined = out["class_of_worker"].isin(ic1.NO_FIRM_SIZE_CLASSES) + assert undefined.any() + assert out.loc[undefined, "firm_size_band"].isna().all() + + +def test_class_of_worker_is_carried_not_imputed(): + spec = si.SpellImputationSpec(bridge="sipp_2008_primary", seed=3) + hosts = _hosts() + fitted = si.fit_spell_model(_donor(), spec, weight_column="weight") + out = si.impute_spells(fitted, hosts, spec, band_codes=BANDS) + assert list(out["class_of_worker"]) == list(hosts["class_of_worker"]) + + +def test_run_records_its_bridge_and_seed(): + spec = si.SpellImputationSpec(bridge="sipp_2008_primary", seed=11) + fitted = si.fit_spell_model(_donor(), spec, weight_column="weight") + out = si.impute_spells(fitted, _hosts(), spec, band_codes=BANDS) + assert out.attrs["bridge"] == "sipp_2008_primary" + assert out.attrs["seed"] == 11 + assert out.attrs["band_is_imputed"] is True + + +def test_same_seed_reproduces_the_same_draw(): + spec = si.SpellImputationSpec(bridge="sipp_2008_primary", seed=5) + donor, hosts = _donor(), _hosts() + fitted = si.fit_spell_model(donor, spec, weight_column="weight") + a = si.impute_spells(fitted, hosts, spec, band_codes=BANDS) + b = si.impute_spells(fitted, hosts, spec, band_codes=BANDS) + assert a["firm_size_band"].equals(b["firm_size_band"]) + + +def test_a_different_seed_changes_the_draw(): + donor, hosts = _donor(), _hosts() + spec_a = si.SpellImputationSpec(bridge="sipp_2008_primary", seed=5) + spec_b = si.SpellImputationSpec(bridge="sipp_2008_primary", seed=6) + fitted = si.fit_spell_model(donor, spec_a, weight_column="weight") + a = si.impute_spells(fitted, hosts, spec_a, band_codes=BANDS) + b = si.impute_spells(fitted, hosts, spec_b, band_codes=BANDS) + assert not a["firm_size_band"].equals(b["firm_size_band"]) + + +def test_imputed_spells_feed_the_calibration_universe(): + spec = si.SpellImputationSpec(bridge="sipp_2014_proxy_chain", seed=9) + fitted = si.fit_spell_model(_donor(), spec, weight_column="weight") + out = si.impute_spells(fitted, _hosts(), spec, band_codes=BANDS) + universe = ic1.calibration_universe(out) + assert set(universe["class_of_worker"]) == {"private"} + assert universe.attrs["excluded_total"] > 0 diff --git a/tests/tier_counts.json b/tests/tier_counts.json index dd0888b3..e2b687b7 100644 --- a/tests/tier_counts.json +++ b/tests/tier_counts.json @@ -1,7 +1,7 @@ { "schema_version": 1, "counts": { - "unit": 1568, + "unit": 1582, "artifact": 2543, "integration_psid": 848, "reproduction_legacy": 520, From 8ab9207ba1f7cf23a520dffb5981933392b658c5 Mon Sep 17 00:00:00 2001 From: vahid-ahmadi Date: Thu, 13 Aug 2026 09:20:12 +0100 Subject: [PATCH 6/6] Recount tier manifest after restacking on #386 (unit 1,634) Co-Authored-By: Claude Opus 5 (1M context) --- tests/README-tiers.md | 4 ++-- tests/tier_counts.json | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/README-tiers.md b/tests/README-tiers.md index e499ae74..b9279326 100644 --- a/tests/README-tiers.md +++ b/tests/README-tiers.md @@ -38,9 +38,9 @@ pytest --collect-only -q -m oracle_policyengine | tail -1 | Tier | Tests at HEAD | |---|---:| -| `unit` | 1,589 | +| `unit` | 1,634 | | `artifact` | 2,543 | | `integration_psid` | 848 | | `reproduction_legacy` | 520 | | `oracle_policyengine` | 159 | -| **Total** | **5,659** | +| **Total** | **5,704** | diff --git a/tests/tier_counts.json b/tests/tier_counts.json index 0fc4a033..be2ab263 100644 --- a/tests/tier_counts.json +++ b/tests/tier_counts.json @@ -1,7 +1,7 @@ { "schema_version": 1, "counts": { - "unit": 1589, + "unit": 1634, "artifact": 2543, "integration_psid": 848, "reproduction_legacy": 520,