Format long-run calibration files

Author: MaxGhenis

policyengine_us_data/datasets/cps/long_term/assess_calibration_frontier.py

@@ -86,13 +86,17 @@ def reorder_to_households(values, order, n_households: int) -> np.ndarray:
     return ordered
 
 
-def build_constraint_inputs(year: int, hh_id_to_idx: dict, n_households: int, profile) -> dict:
+def build_constraint_inputs(
+    year: int, hh_id_to_idx: dict, n_households: int, profile
+) -> dict:
     sim = Microsimulation(dataset=BASE_DATASET_PATH)
     if profile.use_h6_reform:
         raise NotImplementedError(
             "Frontier assessment for H6-enabled profiles is not yet implemented."
         )
-    household_ids = sim.calculate("household_id", period=year, map_to="household").values
+    household_ids = sim.calculate(
+        "household_id", period=year, map_to="household"
+    ).values
     if len(household_ids) != n_households:
         raise ValueError(
             f"Household count mismatch for {year}: {len(household_ids)} vs {n_households}"
@@ -226,10 +230,7 @@ def main() -> int:
         }
         rows.append(row)
         best_case_display = "n/a" if best_case is None else f"{best_case:.3f}%"
-        print(
-            f"{year}: best-case max error {best_case_display} -> "
-            f"{row['quality']}"
-        )
+        print(f"{year}: best-case max error {best_case_display} -> {row['quality']}")
 
     if args.output:
         output_path = Path(args.output)

policyengine_us_data/datasets/cps/long_term/assess_publishable_horizon.py

@@ -119,9 +119,7 @@ def benchmark_tob_values(
             else (oasdi_achieved - oasdi_target) / oasdi_target * 100
         ),
         "hi_tob_benchmark_pct_error": (
-            0.0
-            if hi_target == 0
-            else (hi_achieved - hi_target) / hi_target * 100
+            0.0 if hi_target == 0 else (hi_achieved - hi_target) / hi_target * 100
         ),
     }
 
@@ -287,9 +285,7 @@ def assess_years(
                 "validation_issues": str(error),
                 "runtime_error": str(error),
             }
-            best_case_match = re.search(
-                r"([0-9.]+)%\s*>\s*([0-9.]+)%", str(error)
-            )
+            best_case_match = re.search(r"([0-9.]+)%\s*>\s*([0-9.]+)%", str(error))
             if best_case_match:
                 row["reported_best_case_constraint_error_pct"] = float(
                     best_case_match.group(1)
@@ -315,7 +311,9 @@ def assess_years(
             h6_income_values=None,
             h6_revenue_target=None,
             oasdi_tob_values=oasdi_tob_values if profile.use_tob else None,
-            oasdi_tob_target=load_oasdi_tob_projections(year) if profile.use_tob else None,
+            oasdi_tob_target=load_oasdi_tob_projections(year)
+            if profile.use_tob
+            else None,
             hi_tob_values=hi_tob_values if profile.use_tob else None,
             hi_tob_target=load_hi_tob_projections(year) if profile.use_tob else None,
         )

policyengine_us_data/datasets/cps/long_term/benchmark_trustees_bracket_indexing.py

@@ -45,9 +45,7 @@ def modify_parameters(parameters):
             bracket_node = thresholds.get_child(bracket)
             for filing_status in FILING_STATUSES:
                 parameter = bracket_node.get_child(filing_status)
-                interval = float(
-                    parameter.metadata["uprating"]["rounding"]["interval"]
-                )
+                interval = float(parameter.metadata["uprating"]["rounding"]["interval"])
 
                 for year in range(start_year, end_year + 1):
                     previous_value = float(parameter(f"{year - 1}-01-01"))
@@ -66,6 +64,8 @@ def apply(self):
             self.modify_parameters(modify_parameters)
 
     return reform
+
+
 def _coerce_h5_path(raw: str) -> Path:
     path = Path(raw).expanduser()
     if path.is_dir():

policyengine_us_data/datasets/cps/long_term/build_long_term_target_sources.py

@@ -45,9 +45,7 @@ def build_oact_source(trustees: pd.DataFrame) -> pd.DataFrame:
         missing_years = merged.loc[
             merged["oasdi_nominal_delta_billions"].isna(), "year"
         ].tolist()
-        raise ValueError(
-            f"Missing OACT OASDI deltas for years: {missing_years}"
-        )
+        raise ValueError(f"Missing OACT OASDI deltas for years: {missing_years}")
 
     merged["oasdi_tob_billions_nominal_usd"] = (
         merged["oasdi_tob_billions_nominal_usd"]

policyengine_us_data/datasets/cps/long_term/calibration.py

@@ -279,9 +279,7 @@ def calibrate_entropy(
     baseline_weights = np.asarray(baseline_weights, dtype=float)
     gram = A_scaled.T @ (baseline_weights[:, None] * A_scaled)
     gram += np.eye(gram.shape[0]) * 1e-12
-    beta0 = np.linalg.solve(
-        gram, targets_scaled - (A_scaled.T @ baseline_weights)
-    )
+    beta0 = np.linalg.solve(gram, targets_scaled - (A_scaled.T @ baseline_weights))
 
     def objective_gradient_hessian(beta):
         eta = np.clip(A_scaled @ beta, -700, 700)
@@ -313,11 +311,7 @@ def _cached_ogh(z):
             return None
         _, gradient, _ = objective_gradient_hessian(result.x)
         max_error = float(
-            np.max(
-                100
-                * np.abs(gradient)
-                / np.maximum(np.abs(targets_scaled), 1e-12)
-            )
+            np.max(100 * np.abs(gradient) / np.maximum(np.abs(targets_scaled), 1e-12))
         )
         if max_error > tol * 100:
             return None
@@ -343,11 +337,7 @@ def infeasibility_error(prefix):
     for iterations in range(1, max_iters + 1):
         objective, gradient, hessian = objective_gradient_hessian(beta)
         final_max_error = float(
-            np.max(
-                100
-                * np.abs(gradient)
-                / np.maximum(np.abs(targets_scaled), 1e-12)
-            )
+            np.max(100 * np.abs(gradient) / np.maximum(np.abs(targets_scaled), 1e-12))
         )
         if final_max_error <= tol * 100:
             break
@@ -361,8 +351,8 @@ def infeasibility_error(prefix):
         step = 1.0
         while step >= 1e-8:
             candidate = beta - step * delta
-            candidate_objective, candidate_gradient, _ = (
-                objective_gradient_hessian(candidate)
+            candidate_objective, candidate_gradient, _ = objective_gradient_hessian(
+                candidate
             )
             candidate_max_error = float(
                 np.max(
@@ -477,9 +467,7 @@ def objective_and_gradient(z):
         exp_eta = np.exp(eta)
         weights = baseline_weights * exp_eta
         achieved = A_scaled.T @ weights
-        objective = float(
-            np.sum(weights) + upper_bounds @ alpha - lower_bounds @ gamma
-        )
+        objective = float(np.sum(weights) + upper_bounds @ alpha - lower_bounds @ gamma)
         gradient = np.concatenate(
             [
                 upper_bounds - achieved,
@@ -561,9 +549,7 @@ def objective_with_gradient(z):
             )
 
     if best_result is None or best_weights is None:
-        raise RuntimeError(
-            "Approximate bounded entropy calibration did not run."
-        )
+        raise RuntimeError("Approximate bounded entropy calibration did not run.")
 
     raise RuntimeError(
         "Approximate bounded entropy calibration failed: "
@@ -610,8 +596,7 @@ def densify_lp_solution(
         lam = (lo + hi) / 2.0
         candidate_weights = (1.0 - lam) * lp_weights + lam * baseline_weights
         candidate_error_pct = float(
-            np.max(np.abs(A_scaled.T @ candidate_weights - targets_scaled))
-            * 100
+            np.max(np.abs(A_scaled.T @ candidate_weights - targets_scaled)) * 100
         )
         if candidate_error_pct <= max_constraint_error_pct + 1e-6:
             best_lambda = lam
@@ -670,14 +655,11 @@ def calibrate_lp_minimax(
 
     A = aux_df.to_numpy(dtype=float)
     targets = np.array(list(controls.values()), dtype=float)
-    feasibility = assess_nonnegative_feasibility(
-        A, targets, return_weights=True
-    )
+    feasibility = assess_nonnegative_feasibility(A, targets, return_weights=True)
     weights = feasibility.get("weights")
     if not feasibility["success"] or weights is None:
         raise RuntimeError(
-            "Approximate nonnegative calibration failed: "
-            f"{feasibility['message']}"
+            f"Approximate nonnegative calibration failed: {feasibility['message']}"
         )
 
     return np.asarray(weights, dtype=float), 1, feasibility
@@ -706,9 +688,7 @@ def assess_nonnegative_feasibility(A, targets, *, return_weights=False):
     b_rel = targets / scales
 
     constraint_matrix = sparse.csr_matrix(A_rel)
-    epsilon_column = sparse.csc_matrix(
-        np.ones((constraint_matrix.shape[0], 1))
-    )
+    epsilon_column = sparse.csc_matrix(np.ones((constraint_matrix.shape[0], 1)))
     A_ub = sparse.vstack(
         [
             sparse.hstack([constraint_matrix, -epsilon_column]),
@@ -886,9 +866,7 @@ def calibrate_weights(
                 hi_tob_target=hi_tob_target,
                 n_ages=n_ages,
             )
-            approximate_error_pct = float(
-                feasibility["best_case_max_pct_error"]
-            )
+            approximate_error_pct = float(feasibility["best_case_max_pct_error"])
             if approximate_error_pct <= max(tol * 100, 1e-6):
                 audit["lp_fallback_used"] = True
                 audit["approximation_method"] = "lp_minimax_exact"
@@ -1019,9 +997,7 @@ def build_calibration_audit(
 ):
     achieved_ages = X.T @ weights
     age_errors = (
-        np.abs(achieved_ages - y_target)
-        / np.maximum(np.abs(y_target), 1e-10)
-        * 100
+        np.abs(achieved_ages - y_target) / np.maximum(np.abs(y_target), 1e-10) * 100
     )
 
     neg_mask = weights < 0
@@ -1031,12 +1007,8 @@ def build_calibration_audit(
     abs_weight_sum = float(np.sum(np.abs(weights)))
     if weight_sum > 0:
         sorted_weights = np.sort(weights)
-        top_10_weight_share_pct = float(
-            sorted_weights[-10:].sum() / weight_sum * 100
-        )
-        top_100_weight_share_pct = float(
-            sorted_weights[-100:].sum() / weight_sum * 100
-        )
+        top_10_weight_share_pct = float(sorted_weights[-10:].sum() / weight_sum * 100)
+        top_100_weight_share_pct = float(sorted_weights[-100:].sum() / weight_sum * 100)
     else:
         top_10_weight_share_pct = 0.0
         top_100_weight_share_pct = 0.0
@@ -1051,9 +1023,7 @@ def build_calibration_audit(
         {
             "age_max_pct_error": float(age_errors.max()),
             "negative_weight_count": int(neg_mask.sum()),
-            "negative_weight_household_pct": float(
-                100 * neg_mask.sum() / len(weights)
-            ),
+            "negative_weight_household_pct": float(100 * neg_mask.sum() / len(weights)),
             "negative_weight_pct": (
                 float(100 * negative_values.sum() / abs_weight_sum)
                 if abs_weight_sum > 0
@@ -1063,9 +1033,7 @@ def build_calibration_audit(
                 float(negative_values.max()) if negative_values.size else 0.0
             ),
             "positive_weight_count": int(positive_mask.sum()),
-            "positive_weight_pct": float(
-                100 * positive_mask.sum() / len(weights)
-            ),
+            "positive_weight_pct": float(100 * positive_mask.sum() / len(weights)),
             "effective_sample_size": effective_sample_size,
             "top_10_weight_share_pct": top_10_weight_share_pct,
             "top_100_weight_share_pct": top_100_weight_share_pct,
@@ -1097,10 +1065,7 @@ def build_calibration_audit(
 
     if audit["constraints"]:
         audit["max_constraint_pct_error"] = float(
-            max(
-                abs(stats["pct_error"])
-                for stats in audit["constraints"].values()
-            )
+            max(abs(stats["pct_error"]) for stats in audit["constraints"].values())
         )
 
     return audit

policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py

@@ -39,10 +39,7 @@ def normalize_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
     if "max_constraint_pct_error" not in audit:
         audit["max_constraint_pct_error"] = float(
             max(
-                (
-                    abs(stats.get("pct_error", 0.0))
-                    for stats in constraints.values()
-                ),
+                (abs(stats.get("pct_error", 0.0)) for stats in constraints.values()),
                 default=0.0,
             )
         )
@@ -186,11 +183,9 @@ def update_dataset_manifest(
         )
     manifest_profile = json.loads(json.dumps(manifest["profile"]))
     if manifest_profile != profile:
-        if manifest_profile.get("name") == profile.get(
-            "name"
-        ) and manifest_profile.get("calibration_method") == profile.get(
+        if manifest_profile.get("name") == profile.get("name") and manifest_profile.get(
             "calibration_method"
-        ):
+        ) == profile.get("calibration_method"):
             manifest["profile"] = profile
         else:
             raise ValueError(
@@ -230,17 +225,13 @@ def update_dataset_manifest(
         "method_used": calibration_audit.get("method_used"),
         "fell_back_to_ipf": calibration_audit.get("fell_back_to_ipf"),
         "age_max_pct_error": calibration_audit.get("age_max_pct_error"),
-        "max_constraint_pct_error": calibration_audit.get(
-            "max_constraint_pct_error"
-        ),
+        "max_constraint_pct_error": calibration_audit.get("max_constraint_pct_error"),
         "negative_weight_pct": calibration_audit.get("negative_weight_pct"),
         "negative_weight_household_pct": calibration_audit.get(
             "negative_weight_household_pct"
         ),
         "validation_passed": calibration_audit.get("validation_passed"),
-        "validation_issue_count": len(
-            calibration_audit.get("validation_issues", [])
-        ),
+        "validation_issue_count": len(calibration_audit.get("validation_issues", [])),
     }
 
     year_set = {int(value) for value in manifest.get("years", [])}

policyengine_us_data/datasets/cps/long_term/calibration_profiles.py

@@ -203,7 +203,9 @@ def get_profile(name: str) -> CalibrationProfile:
         return NAMED_PROFILES[name]
     except KeyError as error:
         valid = ", ".join(sorted(NAMED_PROFILES))
-        raise ValueError(f"Unknown calibration profile '{name}'. Valid profiles: {valid}") from error
+        raise ValueError(
+            f"Unknown calibration profile '{name}'. Valid profiles: {valid}"
+        ) from error
 
 
 def approximate_window_for_year(
@@ -459,8 +461,7 @@ def _collect_threshold_issues(
         and ess < min_effective_sample_size
     ):
         issues.append(
-            f"Effective sample size {ess:.3f} is below "
-            f"{min_effective_sample_size:.3f}"
+            f"Effective sample size {ess:.3f} is below {min_effective_sample_size:.3f}"
         )
 
     top_10_share = audit.get("top_10_weight_share_pct")

policyengine_us_data/datasets/cps/long_term/compare_tob_shares.py

@@ -83,9 +83,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "paths",
         nargs="+",
-        help=(
-            "Metadata files or directories containing *.metadata.json sidecars."
-        ),
+        help=("Metadata files or directories containing *.metadata.json sidecars."),
     )
     parser.add_argument(
         "--format",

policyengine_us_data/datasets/cps/long_term/diagnose_support_augmentation_translation.py

@@ -36,7 +36,9 @@ def main() -> int:
             "report_file named in metadata, or <output_dir>/support_augmentation_report.json."
         ),
     )
-    parser.add_argument("--year", type=int, required=True, help="Output year to inspect.")
+    parser.add_argument(
+        "--year", type=int, required=True, help="Output year to inspect."
+    )
     parser.add_argument(
         "--age-bucket-size",
         type=int,

policyengine_us_data/datasets/cps/long_term/evaluate_support_augmentation.py

@@ -49,9 +49,7 @@ def _evaluate_dataset(
 
     approximate_window = approximate_window_for_year(profile, year)
     age_bucket_size = (
-        approximate_window.age_bucket_size
-        if approximate_window is not None
-        else None
+        approximate_window.age_bucket_size if approximate_window is not None else None
     )
     if age_bucket_size and age_bucket_size > 1:
         age_bins = build_age_bins(n_ages=n_ages, bucket_size=age_bucket_size)
@@ -68,7 +66,9 @@ def _evaluate_dataset(
     ss_values = None
     ss_target = None
     if profile.use_ss:
-        ss_values = sim.calculate("social_security", period=year, map_to="household").values
+        ss_values = sim.calculate(
+            "social_security", period=year, map_to="household"
+        ).values
         ss_target = load_ssa_benefit_projections(year)
 
     payroll_values = None