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/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/src/populace_dynamics/firms/frame.py b/src/populace_dynamics/firms/frame.py new file mode 100644 index 00000000..65b12d0f --- /dev/null +++ b/src/populace_dynamics/firms/frame.py @@ -0,0 +1,576 @@ +"""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", + "DEFAULT_WEIGHT_BOUNDS", + "EMPLOYMENT_OUT_OF_SCOPE_SECTORS", + "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") + +#: 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] + + +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 _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. + + ``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", 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}]" + ), + None, + ) + + 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 + lo, hi = bounds + clipped = (weights < lo) | (weights > hi) + if not clipped.any() or not (free & ~clipped).any(): + break + 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( + frame: pd.DataFrame, + 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. + + :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]] = [] + 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, residual = _cell_weights( + group[size_column].to_numpy(), + float(row["firms"]), + float(row["employment"]), + weight_bounds, + ) + if cell_weights is None: + infeasible.append( + { + "cell": "/".join(cell), + "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 + out.attrs["target_employment"] = int(targets["employment"].sum()) + 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 + + +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(), + } 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/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 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/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/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/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 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,