From 8370d5cdef555a65944b671457f418a6d49643a9 Mon Sep 17 00:00:00 2001
From: Jason DeBacker <jdebacker@users.noreply.github.com>
Date: Tue, 16 Jun 2026 15:02:52 -0400
Subject: [PATCH 1/9] add script to endogenously estimate parameters

---
 ogusa/estimate_lifecycle_params.py | 941 +++++++++++++++++++++++++++++
 1 file changed, 941 insertions(+)
 create mode 100644 ogusa/estimate_lifecycle_params.py

diff --git a/ogusa/estimate_lifecycle_params.py b/ogusa/estimate_lifecycle_params.py
new file mode 100644
index 00000000..fd13a460
--- /dev/null
+++ b/ogusa/estimate_lifecycle_params.py
@@ -0,0 +1,941 @@
+"""
+Joint SMM calibration of beta, chi_b, and chi_n parameters.
+
+This module estimates the preference parameters that govern savings,
+bequests, and labor supply in one steady-state SMM problem.  It is designed
+as a standalone calibration layer; callers can wire the returned parameter
+values into the standard OG-USA calibration flow after validating the fit.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field, replace
+from typing import Literal
+
+import numpy as np
+import pandas as pd
+import scipy.optimize as opt
+from ogcore import SS
+from ogcore.utils import Inequality
+
+from ogusa import compute_moments, wealth
+
+SS.VERBOSE = False
+
+WeightingMethod = Literal["identity", "diagonal", "optimal"]
+TailMethod = Literal["scaled_default", "flat"]
+WealthProfileMoment = Literal["level", "mean_normalized"]
+WEALTH_MOMENT_BIN_WEIGHTS = np.array(
+    [0.25, 0.25, 0.20, 0.10, 0.10, 0.09, 0.01]
+)
+
+
+@dataclass(frozen=True)
+class LifecycleCalibrationConfig:
+    """
+    Configuration for the joint lifecycle preference calibration.
+
+    The default age window is 20 through 79.  That gives 60 chi_n parameters
+    and 60 labor and wealth-profile moments, matching the intended 80
+    parameters and 130 moments when J=10.
+    """
+
+    min_age: int = 20
+    max_age: int = 79
+    estimate_chi_n_min_age: int = 20
+    estimate_chi_n_max_age: int = 79
+    chi_n_tail_method: TailMethod = "scaled_default"
+    wealth_profile_moment: WealthProfileMoment = "mean_normalized"
+    include_labor_profile: bool = True
+    include_wealth_profile: bool = True
+    include_income_gini: bool = True
+    include_wealth_distribution: bool = True
+    include_inheritance_moments: bool = False
+    cps_years: tuple[int, ...] = (2023, 2022)
+    scf_yrs_list: tuple[int, ...] = (2019,)
+    cps_directory: str | None = None
+    scf_directory: str | None = None
+    scf_web: bool = False
+    bootstrap_iterations: int = 1000
+    weighting_method: WeightingMethod = "identity"
+    weighting_ridge: float = 1e-8
+    optimizer_method: str = "L-BFGS-B"
+    optimizer_tol: float = 1e-10
+    optimizer_options: dict = field(
+        default_factory=lambda: {"maxiter": 100, "maxfun": 250}
+    )
+
+    @property
+    def moment_ages(self) -> np.ndarray:
+        """Return the model-age labels used in age-profile moments."""
+        return np.arange(self.min_age, self.max_age + 1)
+
+    @property
+    def estimated_chi_n_ages(self) -> np.ndarray:
+        """Return the age labels for directly estimated chi_n values."""
+        return np.arange(
+            self.estimate_chi_n_min_age,
+            self.estimate_chi_n_max_age + 1,
+        )
+
+    def validate(self, p) -> None:
+        """Validate age and dimension settings against an OG-Core spec."""
+        starting_age = int(getattr(p, "starting_age", 20))
+        ending_age = int(getattr(p, "ending_age", 100))
+        model_ages = np.arange(starting_age, ending_age)
+        min_model_age = int(model_ages[0])
+        max_model_age = int(model_ages[-1])
+        requested_ages = np.concatenate(
+            [self.moment_ages, self.estimated_chi_n_ages]
+        )
+        if requested_ages.min() < min_model_age:
+            raise ValueError("Requested ages start before model ages.")
+        if requested_ages.max() > max_model_age:
+            raise ValueError("Requested ages extend beyond model ages.")
+        if self.estimate_chi_n_max_age < self.estimate_chi_n_min_age:
+            raise ValueError(
+                "estimate_chi_n_max_age must be at least min age."
+            )
+        if self.max_age < self.min_age:
+            raise ValueError("max_age must be at least min_age.")
+
+
+@dataclass(frozen=True)
+class MomentSet:
+    """Named vector of moments."""
+
+    names: tuple[str, ...]
+    values: np.ndarray
+
+    def __post_init__(self) -> None:
+        values = np.asarray(self.values, dtype=float)
+        object.__setattr__(self, "values", values)
+        if len(self.names) != values.size:
+            raise ValueError("Moment names and values have different lengths.")
+
+
+@dataclass
+class LifecycleCalibrationResult:
+    """Container for joint lifecycle preference calibration results."""
+
+    beta_annual: np.ndarray
+    chi_b: np.ndarray
+    chi_n: np.ndarray
+    objective_value: float
+    optimizer_result: opt.OptimizeResult
+    data_moments: MomentSet
+    model_moments: MomentSet
+    weighting_matrix: np.ndarray
+
+    @property
+    def parameter_dict(self) -> dict[str, list[float]]:
+        """Return calibrated values in OG-Core update_specifications format."""
+        return {
+            "beta_annual": self.beta_annual.tolist(),
+            "chi_b": self.chi_b.tolist(),
+            "chi_n": self.chi_n.tolist(),
+        }
+
+
+def _as_vector(values) -> np.ndarray:
+    """Return values as a one-dimensional float array."""
+    return np.asarray(values, dtype=float).reshape(-1)
+
+
+def _lambdas(p) -> np.ndarray:
+    """Return lifetime-income weights as a one-dimensional array."""
+    return _as_vector(p.lambdas)
+
+
+def _ss_chi_n(p) -> np.ndarray:
+    """Return the steady-state chi_n age profile from a spec object."""
+    chi_n = np.asarray(p.chi_n, dtype=float)
+    if chi_n.ndim == 1:
+        return chi_n.copy()
+    return chi_n[-1, :].copy()
+
+
+def _age_to_index(age: int, p) -> int:
+    """Map an age label to its model age index."""
+    return int(age) - int(getattr(p, "starting_age", 20))
+
+
+def _age_indices(ages: np.ndarray, p) -> np.ndarray:
+    """Map age labels to model indices."""
+    return np.array([_age_to_index(age, p) for age in ages], dtype=int)
+
+
+def _weighted_mean(values, weights) -> float:
+    """Return a weighted mean after dropping nonfinite observations."""
+    data = pd.DataFrame({"value": values, "weight": weights})
+    data = data.replace([np.inf, -np.inf], np.nan).dropna()
+    data = data[data["weight"] > 0]
+    if data.empty:
+        return np.nan
+    return float((data["value"] * data["weight"]).sum() / data["weight"].sum())
+
+
+def _weighted_mean_by_age(
+    data: pd.DataFrame,
+    value_col: str,
+    weight_col: str | None,
+    ages: np.ndarray,
+    age_col: str = "age",
+) -> np.ndarray:
+    """Compute weighted means by single-year age."""
+    columns = [age_col, value_col]
+    if weight_col is not None:
+        columns.append(weight_col)
+    age_data = data[columns].copy()
+    age_data[age_col] = pd.to_numeric(age_data[age_col], errors="coerce")
+    age_data[value_col] = pd.to_numeric(age_data[value_col], errors="coerce")
+    age_data = age_data.replace([np.inf, -np.inf], np.nan).dropna()
+    age_data = age_data[age_data[age_col].isin(ages)].copy()
+    age_data[age_col] = age_data[age_col].astype(int)
+
+    if weight_col is None:
+        profile = age_data.groupby(age_col)[value_col].mean()
+    else:
+        age_data[weight_col] = pd.to_numeric(
+            age_data[weight_col], errors="coerce"
+        )
+        age_data = age_data[age_data[weight_col] > 0].copy()
+        age_data["weighted_value"] = age_data[value_col] * age_data[weight_col]
+        by_age = age_data.groupby(age_col)[
+            ["weighted_value", weight_col]
+        ].sum()
+        profile = by_age["weighted_value"] / by_age[weight_col]
+
+    return profile.reindex(ages).to_numpy(dtype=float)
+
+
+def _require_finite(values: np.ndarray, label: str) -> np.ndarray:
+    """Validate that all values are finite."""
+    values = np.asarray(values, dtype=float)
+    if not np.all(np.isfinite(values)):
+        raise ValueError(f"{label} contains missing or nonfinite values.")
+    return values
+
+
+def _normalize_wealth_profile(
+    profile: np.ndarray,
+    method: WealthProfileMoment,
+) -> np.ndarray:
+    """Scale a wealth age profile according to the requested convention."""
+    profile = np.asarray(profile, dtype=float)
+    if method == "level":
+        return profile
+    if method == "mean_normalized":
+        mean = np.nanmean(profile)
+        if not np.isfinite(mean) or np.isclose(mean, 0.0):
+            raise ValueError("Cannot normalize wealth profile with zero mean.")
+        return profile / mean
+    raise ValueError(f"Unsupported wealth profile moment: {method}")
+
+
+def _wealth_distribution_moment_names() -> tuple[str, ...]:
+    """Return names for the nine SCF/model wealth distribution moments."""
+    return (
+        "wealth_share_0_25",
+        "wealth_share_25_50",
+        "wealth_share_50_70",
+        "wealth_share_70_80",
+        "wealth_share_80_90",
+        "wealth_share_90_99",
+        "wealth_share_99_100",
+        "wealth_gini",
+        "wealth_var_log",
+    )
+
+
+def _model_wealth_distribution_moments(ss_output: dict, p) -> np.ndarray:
+    """Compute the model moments matching wealth.compute_wealth_moments."""
+    b_sp1 = np.asarray(ss_output["b_sp1"], dtype=float)
+    wealth_ineq = Inequality(b_sp1, p.omega_SS, _lambdas(p), p.S, p.J)
+    return np.array(
+        [
+            1 - wealth_ineq.top_share(0.75),
+            wealth_ineq.top_share(0.75) - wealth_ineq.top_share(0.50),
+            wealth_ineq.top_share(0.50) - wealth_ineq.top_share(0.30),
+            wealth_ineq.top_share(0.30) - wealth_ineq.top_share(0.20),
+            wealth_ineq.top_share(0.20) - wealth_ineq.top_share(0.10),
+            wealth_ineq.top_share(0.10) - wealth_ineq.top_share(0.01),
+            wealth_ineq.top_share(0.01),
+            wealth_ineq.gini(),
+            wealth_ineq.var_of_logs(),
+        ],
+        dtype=float,
+    )
+
+
+def load_cps_hours_data(
+    cps_years: tuple[int, ...] = (2023, 2022),
+    cps_directory: str | None = None,
+) -> pd.DataFrame:
+    """Load packaged CPS ASEC hours data used for labor moments."""
+    if cps_directory is None:
+        cps_directory = compute_moments.CPS_DATA_DIR
+    data = []
+    for year in cps_years:
+        path = f"{cps_directory}/cps_asec_hours_{year}.csv"
+        data.append(pd.read_csv(path))
+    return pd.concat(data, ignore_index=True)
+
+
+def load_scf_wealth_data(config: LifecycleCalibrationConfig) -> pd.DataFrame:
+    """Load SCF wealth data with ages for wealth moments."""
+    return wealth.get_wealth_data(
+        scf_yrs_list=list(config.scf_yrs_list),
+        web=config.scf_web,
+        directory=config.scf_directory,
+        include_age=True,
+    )
+
+
+def labor_profile_from_cps(
+    cps: pd.DataFrame,
+    config: LifecycleCalibrationConfig,
+) -> np.ndarray:
+    """Compute labor supply by age from CPS data."""
+    cps = cps.copy()
+    if "hours" not in cps:
+        if "hours_per_week" in cps:
+            cps["hours"] = cps["hours_per_week"]
+        else:
+            raise ValueError("CPS data must include hours or hours_per_week.")
+    weight_col = None
+    for possible_weight in ("weight", "wtsupp", "s006", "wgt"):
+        if possible_weight in cps:
+            weight_col = possible_weight
+            break
+    hours = _weighted_mean_by_age(
+        cps,
+        "hours",
+        weight_col,
+        config.moment_ages,
+    )
+    labor = hours / ((24 - 8) * 7)
+    return _require_finite(labor, "labor profile")
+
+
+def wealth_profile_from_scf(
+    scf: pd.DataFrame,
+    config: LifecycleCalibrationConfig,
+) -> np.ndarray:
+    """Compute net-wealth age profile moments from SCF data."""
+    profile = _weighted_mean_by_age(
+        scf,
+        "networth_infadj",
+        "wgt",
+        config.moment_ages,
+    )
+    profile = _normalize_wealth_profile(
+        profile,
+        config.wealth_profile_moment,
+    )
+    return _require_finite(profile, "wealth profile")
+
+
+def income_gini_data_moment(
+    income_year: int | None = None,
+) -> float:
+    """Compute the before-tax income Gini data moment."""
+    moments = compute_moments._taxcalc_cps_income_ginis(income_year)
+    return float(moments["Gini coefficient, income"])
+
+
+def compute_inheritance_moments_from_scf(
+    scf: pd.DataFrame,
+    amount_col: str,
+    received_col: str | None = None,
+    weight_col: str = "wgt",
+    networth_col: str = "networth_infadj",
+) -> MomentSet:
+    """
+    Compute optional inherited-transfer moments from full SCF extracts.
+
+    The trimmed SCF files packaged in OG-USA do not include the inheritance
+    variables needed here.  Pass a full SCF extract with the relevant Section X
+    variables and specify the amount and receipt indicator columns.
+    """
+    if amount_col not in scf:
+        raise ValueError(f"SCF data are missing {amount_col}.")
+    if weight_col not in scf:
+        raise ValueError(f"SCF data are missing {weight_col}.")
+
+    data = scf.copy()
+    data[amount_col] = pd.to_numeric(data[amount_col], errors="coerce")
+    data[weight_col] = pd.to_numeric(data[weight_col], errors="coerce")
+    data = data.replace([np.inf, -np.inf], np.nan).dropna(
+        subset=[amount_col, weight_col]
+    )
+    data = data[data[weight_col] > 0].copy()
+    if received_col is None:
+        received = data[amount_col] > 0
+    else:
+        if received_col not in data:
+            raise ValueError(f"SCF data are missing {received_col}.")
+        received = pd.to_numeric(data[received_col], errors="coerce") > 0
+
+    names = ["inheritance_received_rate"]
+    values = [_weighted_mean(received.astype(float), data[weight_col])]
+
+    recipients = data[received].copy()
+    if not recipients.empty:
+        names.append("inheritance_amount_conditional_mean")
+        values.append(
+            _weighted_mean(recipients[amount_col], recipients[weight_col])
+        )
+
+    if networth_col in data:
+        data[networth_col] = pd.to_numeric(data[networth_col], errors="coerce")
+        positive_networth = data[networth_col] > 0
+        ratio_data = data[positive_networth].copy()
+        if not ratio_data.empty:
+            names.append("inheritance_to_networth_mean")
+            values.append(
+                _weighted_mean(
+                    ratio_data[amount_col] / ratio_data[networth_col],
+                    ratio_data[weight_col],
+                )
+            )
+
+    return MomentSet(tuple(names), np.asarray(values, dtype=float))
+
+
+def compute_data_moments(
+    p,
+    config: LifecycleCalibrationConfig | None = None,
+    cps: pd.DataFrame | None = None,
+    scf: pd.DataFrame | None = None,
+    income_year: int | None = None,
+    inheritance_moments: MomentSet | None = None,
+) -> MomentSet:
+    """Compute the stacked data moment vector for the SMM objective."""
+    if config is None:
+        config = LifecycleCalibrationConfig()
+    config.validate(p)
+    names: list[str] = []
+    values: list[float] = []
+
+    if config.include_labor_profile:
+        if cps is None:
+            cps = load_cps_hours_data(config.cps_years, config.cps_directory)
+        labor = labor_profile_from_cps(cps, config)
+        names.extend(f"labor_supply_age_{age}" for age in config.moment_ages)
+        values.extend(labor)
+
+    if config.include_wealth_profile:
+        if scf is None:
+            scf = load_scf_wealth_data(config)
+        wealth_profile = wealth_profile_from_scf(scf, config)
+        names.extend(f"net_wealth_age_{age}" for age in config.moment_ages)
+        values.extend(wealth_profile)
+
+    if config.include_income_gini:
+        names.append("income_gini")
+        values.append(income_gini_data_moment(income_year=income_year))
+
+    if config.include_wealth_distribution:
+        if scf is None:
+            scf = load_scf_wealth_data(config)
+        wealth_dist = wealth.compute_wealth_moments(
+            scf.copy(),
+            WEALTH_MOMENT_BIN_WEIGHTS,
+        )
+        names.extend(_wealth_distribution_moment_names())
+        values.extend(wealth_dist)
+
+    if config.include_inheritance_moments:
+        if inheritance_moments is None:
+            raise ValueError(
+                "inheritance_moments must be supplied when "
+                "include_inheritance_moments is True."
+            )
+        names.extend(inheritance_moments.names)
+        values.extend(inheritance_moments.values)
+
+    return MomentSet(tuple(names), np.asarray(values, dtype=float))
+
+
+def compute_model_moments(
+    ss_output: dict,
+    p,
+    config: LifecycleCalibrationConfig | None = None,
+    inheritance_moments: MomentSet | None = None,
+) -> MomentSet:
+    """Compute model moments in the same order as compute_data_moments."""
+    if config is None:
+        config = LifecycleCalibrationConfig()
+    config.validate(p)
+    age_idx = _age_indices(config.moment_ages, p)
+    lambdas = _lambdas(p)
+    names: list[str] = []
+    values: list[float] = []
+
+    if config.include_labor_profile:
+        n = np.asarray(ss_output["n"], dtype=float)
+        labor = (n[age_idx, :] * lambdas.reshape(1, p.J)).sum(axis=1)
+        names.extend(f"labor_supply_age_{age}" for age in config.moment_ages)
+        values.extend(labor)
+
+    if config.include_wealth_profile:
+        b_sp1 = np.asarray(ss_output["b_sp1"], dtype=float)
+        factor = float(ss_output.get("factor", 1.0))
+        wealth_profile = (
+            b_sp1[age_idx, :] * factor * lambdas.reshape(1, p.J)
+        ).sum(axis=1)
+        wealth_profile = _normalize_wealth_profile(
+            wealth_profile,
+            config.wealth_profile_moment,
+        )
+        names.extend(f"net_wealth_age_{age}" for age in config.moment_ages)
+        values.extend(wealth_profile)
+
+    if config.include_income_gini:
+        income = np.asarray(ss_output["before_tax_income"], dtype=float)
+        income_ineq = Inequality(income, p.omega_SS, lambdas, p.S, p.J)
+        names.append("income_gini")
+        values.append(income_ineq.gini())
+
+    if config.include_wealth_distribution:
+        names.extend(_wealth_distribution_moment_names())
+        values.extend(_model_wealth_distribution_moments(ss_output, p))
+
+    if config.include_inheritance_moments:
+        if inheritance_moments is None:
+            raise ValueError(
+                "Model inheritance moments must be supplied when "
+                "include_inheritance_moments is True."
+            )
+        names.extend(inheritance_moments.names)
+        values.extend(inheritance_moments.values)
+
+    return MomentSet(tuple(names), np.asarray(values, dtype=float))
+
+
+def build_chi_n_profile(
+    estimated_chi_n: np.ndarray,
+    base_chi_n: np.ndarray,
+    p,
+    config: LifecycleCalibrationConfig | None = None,
+) -> np.ndarray:
+    """Build a full-S chi_n profile from directly estimated values."""
+    if config is None:
+        config = LifecycleCalibrationConfig()
+    config.validate(p)
+    full_chi_n = np.asarray(base_chi_n, dtype=float).reshape(-1).copy()
+    if full_chi_n.size != p.S:
+        raise ValueError("base_chi_n length must equal p.S.")
+
+    est_values = np.asarray(estimated_chi_n, dtype=float).reshape(-1)
+    est_ages = config.estimated_chi_n_ages
+    if est_values.size != est_ages.size:
+        raise ValueError("estimated_chi_n length does not match config ages.")
+
+    est_idx = _age_indices(est_ages, p)
+    full_chi_n[est_idx] = est_values
+    tail_start = int(est_idx[-1] + 1)
+    if tail_start >= p.S:
+        return full_chi_n
+
+    if config.chi_n_tail_method == "flat":
+        full_chi_n[tail_start:] = est_values[-1]
+    elif config.chi_n_tail_method == "scaled_default":
+        base_anchor = base_chi_n[tail_start - 1]
+        if np.isclose(base_anchor, 0.0):
+            scale = 1.0
+        else:
+            scale = est_values[-1] / base_anchor
+        full_chi_n[tail_start:] = base_chi_n[tail_start:] * scale
+    else:
+        raise ValueError(
+            f"Unsupported chi_n tail method: {config.chi_n_tail_method}"
+        )
+
+    return full_chi_n
+
+
+def pack_lifecycle_params(
+    beta_annual: np.ndarray,
+    chi_b: np.ndarray,
+    chi_n: np.ndarray,
+    p,
+    config: LifecycleCalibrationConfig | None = None,
+    transform: bool = True,
+) -> np.ndarray:
+    """Pack natural lifecycle parameters into an optimizer vector."""
+    if config is None:
+        config = LifecycleCalibrationConfig()
+    config.validate(p)
+    beta_annual = _as_vector(beta_annual)
+    chi_b = _as_vector(chi_b)
+    chi_n = _as_vector(chi_n)
+    est_idx = _age_indices(config.estimated_chi_n_ages, p)
+    chi_n_est = chi_n[est_idx]
+    params = np.concatenate([beta_annual, chi_b, chi_n_est])
+    if not transform:
+        return params
+    if np.any((beta_annual <= 0) | (beta_annual >= 1)):
+        raise ValueError(
+            "beta_annual values must be strictly between 0 and 1."
+        )
+    if np.any(params[beta_annual.size :] <= 0):
+        raise ValueError("chi_b and chi_n values must be positive.")
+    beta_trans = np.log(beta_annual / (1 - beta_annual))
+    chi_trans = np.log(params[beta_annual.size :])
+    return np.concatenate([beta_trans, chi_trans])
+
+
+def unpack_lifecycle_params(
+    theta: np.ndarray,
+    p,
+    config: LifecycleCalibrationConfig | None = None,
+    base_chi_n: np.ndarray | None = None,
+    transform: bool = True,
+) -> dict[str, np.ndarray]:
+    """Unpack an optimizer vector into natural lifecycle parameters."""
+    if config is None:
+        config = LifecycleCalibrationConfig()
+    config.validate(p)
+    if base_chi_n is None:
+        base_chi_n = _ss_chi_n(p)
+
+    theta = _as_vector(theta)
+    n_beta = p.J
+    n_chi_b = p.J
+    n_chi_n = config.estimated_chi_n_ages.size
+    expected = n_beta + n_chi_b + n_chi_n
+    if theta.size != expected:
+        raise ValueError(f"Expected {expected} parameters, got {theta.size}.")
+
+    beta_raw = theta[:n_beta]
+    chi_b_raw = theta[n_beta : n_beta + n_chi_b]
+    chi_n_raw = theta[n_beta + n_chi_b :]
+    if transform:
+        beta_annual = 1 / (1 + np.exp(-beta_raw))
+        chi_b = np.exp(chi_b_raw)
+        chi_n_est = np.exp(chi_n_raw)
+    else:
+        beta_annual = beta_raw
+        chi_b = chi_b_raw
+        chi_n_est = chi_n_raw
+
+    chi_n = build_chi_n_profile(chi_n_est, base_chi_n, p, config)
+    return {
+        "beta_annual": beta_annual,
+        "chi_b": chi_b,
+        "chi_n": chi_n,
+        "chi_n_estimated": chi_n_est,
+    }
+
+
+def initial_lifecycle_theta(
+    p,
+    config: LifecycleCalibrationConfig | None = None,
+    transform: bool = True,
+) -> np.ndarray:
+    """Build an optimizer vector from the current spec values."""
+    if config is None:
+        config = LifecycleCalibrationConfig()
+    return pack_lifecycle_params(
+        p.beta_annual,
+        p.chi_b,
+        _ss_chi_n(p),
+        p,
+        config,
+        transform=transform,
+    )
+
+
+def apply_lifecycle_params(
+    p,
+    beta_annual: np.ndarray,
+    chi_b: np.ndarray,
+    chi_n: np.ndarray,
+) -> None:
+    """Apply natural lifecycle parameters to a spec object in-place."""
+    p.update_specifications(
+        {
+            "beta_annual": _as_vector(beta_annual).tolist(),
+            "chi_b": _as_vector(chi_b).tolist(),
+            "chi_n": _as_vector(chi_n).tolist(),
+        }
+    )
+
+
+def weighting_matrix(
+    moment_count: int,
+    method: WeightingMethod = "identity",
+    bootstrap_moments: np.ndarray | None = None,
+    ridge: float = 1e-8,
+) -> np.ndarray:
+    """Construct a weighting matrix for the SMM objective."""
+    if method == "identity":
+        return np.eye(moment_count)
+    if bootstrap_moments is None:
+        raise ValueError("bootstrap_moments are required for this method.")
+
+    boot = np.asarray(bootstrap_moments, dtype=float)
+    if boot.ndim != 2 or boot.shape[1] != moment_count:
+        raise ValueError("bootstrap_moments must be n x moment_count.")
+    vcv = np.cov(boot.T)
+    vcv = np.atleast_2d(vcv)
+    if method == "diagonal":
+        diag = np.diag(vcv).copy()
+        diag[diag < ridge] = ridge
+        return np.diag(1 / diag)
+    if method == "optimal":
+        vcv = vcv + ridge * np.eye(moment_count)
+        return np.linalg.pinv(vcv)
+    raise ValueError(f"Unsupported weighting method: {method}")
+
+
+def bootstrap_data_moments(
+    p,
+    config: LifecycleCalibrationConfig | None = None,
+    cps: pd.DataFrame | None = None,
+    scf: pd.DataFrame | None = None,
+    seed: int | None = None,
+) -> np.ndarray:
+    """Bootstrap the data moments available from CPS and SCF microdata."""
+    if config is None:
+        config = LifecycleCalibrationConfig()
+    if cps is None and config.include_labor_profile:
+        cps = load_cps_hours_data(config.cps_years, config.cps_directory)
+    if scf is None and (
+        config.include_wealth_profile or config.include_wealth_distribution
+    ):
+        scf = load_scf_wealth_data(config)
+
+    point_moments = compute_data_moments(p, config, cps=cps, scf=scf)
+    resampled_config = replace(
+        config,
+        include_income_gini=False,
+        include_inheritance_moments=False,
+    )
+    rng = np.random.default_rng(seed)
+    boot = np.zeros((config.bootstrap_iterations, point_moments.values.size))
+    for i in range(config.bootstrap_iterations):
+        cps_boot = None
+        scf_boot = None
+        if cps is not None:
+            cps_boot = cps.iloc[
+                rng.integers(0, len(cps), size=len(cps))
+            ].reset_index(drop=True)
+        if scf is not None:
+            scf_boot = scf.iloc[
+                rng.integers(0, len(scf), size=len(scf))
+            ].reset_index(drop=True)
+        resampled_moments = compute_data_moments(
+            p,
+            resampled_config,
+            cps=cps_boot,
+            scf=scf_boot,
+        )
+        resampled_values = dict(
+            zip(resampled_moments.names, resampled_moments.values)
+        )
+        boot[i, :] = [
+            resampled_values.get(name, value)
+            for name, value in zip(point_moments.names, point_moments.values)
+        ]
+    return boot
+
+
+def smm_distance(
+    model_moments: MomentSet,
+    data_moments: MomentSet,
+    W: np.ndarray,
+) -> float:
+    """Compute the quadratic SMM distance."""
+    if model_moments.names != data_moments.names:
+        raise ValueError("Model and data moments are not aligned.")
+    diff = model_moments.values - data_moments.values
+    if not np.all(np.isfinite(diff)):
+        return np.inf
+    return float(diff.T @ W @ diff)
+
+
+def smm_objective(
+    theta: np.ndarray,
+    data_moments: MomentSet,
+    W: np.ndarray,
+    p,
+    config: LifecycleCalibrationConfig | None = None,
+    base_chi_n: np.ndarray | None = None,
+    client=None,
+    transform: bool = True,
+) -> float:
+    """Evaluate the joint lifecycle SMM objective."""
+    if config is None:
+        config = LifecycleCalibrationConfig()
+    if base_chi_n is None:
+        base_chi_n = _ss_chi_n(p)
+    try:
+        params = unpack_lifecycle_params(
+            theta,
+            p,
+            config,
+            base_chi_n=base_chi_n,
+            transform=transform,
+        )
+        apply_lifecycle_params(
+            p,
+            params["beta_annual"],
+            params["chi_b"],
+            params["chi_n"],
+        )
+        ss_output = SS.run_SS(p, client=client)
+        model_moments = compute_model_moments(ss_output, p, config)
+        distance = smm_distance(model_moments, data_moments, W)
+    except (
+        AssertionError,
+        FloatingPointError,
+        ValueError,
+        RuntimeError,
+        KeyError,
+    ):
+        distance = np.inf
+    if not np.isfinite(distance):
+        return 1e30
+    return distance
+
+
+def estimate_lifecycle_params(
+    p,
+    config: LifecycleCalibrationConfig | None = None,
+    theta0: np.ndarray | None = None,
+    data_moments: MomentSet | None = None,
+    W: np.ndarray | None = None,
+    bootstrap_moments: np.ndarray | None = None,
+    client=None,
+    transform: bool = True,
+) -> LifecycleCalibrationResult:
+    """Estimate beta_annual, chi_b, and chi_n jointly by SMM."""
+    if config is None:
+        config = LifecycleCalibrationConfig()
+    config.validate(p)
+    base_chi_n = _ss_chi_n(p)
+    if theta0 is None:
+        theta0 = initial_lifecycle_theta(p, config, transform=transform)
+    if data_moments is None:
+        data_moments = compute_data_moments(p, config)
+    if W is None:
+        W = weighting_matrix(
+            data_moments.values.size,
+            method=config.weighting_method,
+            bootstrap_moments=bootstrap_moments,
+            ridge=config.weighting_ridge,
+        )
+
+    est_output = opt.minimize(
+        smm_objective,
+        theta0,
+        args=(data_moments, W, p, config, base_chi_n, client, transform),
+        method=config.optimizer_method,
+        tol=config.optimizer_tol,
+        options=config.optimizer_options,
+    )
+    params = unpack_lifecycle_params(
+        est_output.x,
+        p,
+        config,
+        base_chi_n=base_chi_n,
+        transform=transform,
+    )
+    apply_lifecycle_params(
+        p,
+        params["beta_annual"],
+        params["chi_b"],
+        params["chi_n"],
+    )
+    ss_output = SS.run_SS(p, client=client)
+    model_moments = compute_model_moments(ss_output, p, config)
+
+    return LifecycleCalibrationResult(
+        beta_annual=params["beta_annual"],
+        chi_b=params["chi_b"],
+        chi_n=params["chi_n"],
+        objective_value=float(est_output.fun),
+        optimizer_result=est_output,
+        data_moments=data_moments,
+        model_moments=model_moments,
+        weighting_matrix=W,
+    )
+
+
+def compute_parameter_vcv(
+    theta_hat: np.ndarray,
+    W: np.ndarray,
+    p,
+    config: LifecycleCalibrationConfig | None = None,
+    base_chi_n: np.ndarray | None = None,
+    h: float = 1e-4,
+    client=None,
+    transform: bool = True,
+) -> np.ndarray:
+    """Compute a numerical GMM parameter VCV matrix."""
+    if config is None:
+        config = LifecycleCalibrationConfig()
+    if base_chi_n is None:
+        base_chi_n = _ss_chi_n(p)
+    theta_hat = _as_vector(theta_hat)
+    params = unpack_lifecycle_params(
+        theta_hat,
+        p,
+        config,
+        base_chi_n=base_chi_n,
+        transform=transform,
+    )
+    apply_lifecycle_params(
+        p,
+        params["beta_annual"],
+        params["chi_b"],
+        params["chi_n"],
+    )
+    ss_output = SS.run_SS(p, client=client)
+    base_moments = compute_model_moments(ss_output, p, config)
+    deriv = np.zeros((base_moments.values.size, theta_hat.size))
+
+    for i in range(theta_hat.size):
+        step = h * max(1.0, abs(theta_hat[i]))
+        high = theta_hat.copy()
+        low = theta_hat.copy()
+        high[i] += step
+        low[i] -= step
+
+        high_params = unpack_lifecycle_params(
+            high,
+            p,
+            config,
+            base_chi_n=base_chi_n,
+            transform=transform,
+        )
+        apply_lifecycle_params(
+            p,
+            high_params["beta_annual"],
+            high_params["chi_b"],
+            high_params["chi_n"],
+        )
+        high_output = SS.run_SS(p, client=client)
+        high_moments = compute_model_moments(high_output, p, config)
+
+        low_params = unpack_lifecycle_params(
+            low,
+            p,
+            config,
+            base_chi_n=base_chi_n,
+            transform=transform,
+        )
+        apply_lifecycle_params(
+            p,
+            low_params["beta_annual"],
+            low_params["chi_b"],
+            low_params["chi_n"],
+        )
+        low_output = SS.run_SS(p, client=client)
+        low_moments = compute_model_moments(low_output, p, config)
+        deriv[:, i] = (high_moments.values - low_moments.values) / (2 * step)
+
+    return np.linalg.pinv(deriv.T @ W @ deriv)

From c437ebcefbf7864b296ddfc81d417cee9fa6a9e7 Mon Sep 17 00:00:00 2001
From: Jason DeBacker <jdebacker@users.noreply.github.com>
Date: Wed, 17 Jun 2026 10:47:24 -0400
Subject: [PATCH 2/9] normalize wealth profile moments

---
 ogusa/estimate_lifecycle_params.py | 82 +++++++++++++++++++++++++-----
 1 file changed, 69 insertions(+), 13 deletions(-)

diff --git a/ogusa/estimate_lifecycle_params.py b/ogusa/estimate_lifecycle_params.py
index fd13a460..2c1a4f89 100644
--- a/ogusa/estimate_lifecycle_params.py
+++ b/ogusa/estimate_lifecycle_params.py
@@ -24,7 +24,7 @@
 
 WeightingMethod = Literal["identity", "diagonal", "optimal"]
 TailMethod = Literal["scaled_default", "flat"]
-WealthProfileMoment = Literal["level", "mean_normalized"]
+WealthProfileMoment = Literal["anchor_window", "level", "mean_normalized"]
 WEALTH_MOMENT_BIN_WEIGHTS = np.array(
     [0.25, 0.25, 0.20, 0.10, 0.10, 0.09, 0.01]
 )
@@ -35,17 +35,22 @@ class LifecycleCalibrationConfig:
     """
     Configuration for the joint lifecycle preference calibration.
 
-    The default age window is 20 through 79.  That gives 60 chi_n parameters
-    and 60 labor and wealth-profile moments, matching the intended 80
-    parameters and 130 moments when J=10.
+    The default labor age window is 20 through 79.  Wealth-profile moments
+    use ages 21 through 79 and are normalized by the mean wealth level from
+    ages 20 through 24.  That gives 80 parameters and 129 default moments
+    when J=10.
     """
 
     min_age: int = 20
     max_age: int = 79
+    wealth_profile_min_age: int = 21
+    wealth_profile_max_age: int = 79
+    wealth_anchor_min_age: int = 20
+    wealth_anchor_max_age: int = 24
     estimate_chi_n_min_age: int = 20
     estimate_chi_n_max_age: int = 79
     chi_n_tail_method: TailMethod = "scaled_default"
-    wealth_profile_moment: WealthProfileMoment = "mean_normalized"
+    wealth_profile_moment: WealthProfileMoment = "anchor_window"
     include_labor_profile: bool = True
     include_wealth_profile: bool = True
     include_income_gini: bool = True
@@ -67,9 +72,24 @@ class LifecycleCalibrationConfig:
 
     @property
     def moment_ages(self) -> np.ndarray:
-        """Return the model-age labels used in age-profile moments."""
+        """Return age labels used in labor-profile moments."""
         return np.arange(self.min_age, self.max_age + 1)
 
+    @property
+    def wealth_profile_ages(self) -> np.ndarray:
+        """Return age labels used in wealth-profile moments."""
+        return np.arange(
+            self.wealth_profile_min_age,
+            self.wealth_profile_max_age + 1,
+        )
+
+    @property
+    def wealth_anchor_ages(self) -> np.ndarray:
+        """Return age labels used to normalize wealth-profile moments."""
+        return np.arange(
+            self.wealth_anchor_min_age, self.wealth_anchor_max_age + 1
+        )
+
     @property
     def estimated_chi_n_ages(self) -> np.ndarray:
         """Return the age labels for directly estimated chi_n values."""
@@ -86,7 +106,12 @@ def validate(self, p) -> None:
         min_model_age = int(model_ages[0])
         max_model_age = int(model_ages[-1])
         requested_ages = np.concatenate(
-            [self.moment_ages, self.estimated_chi_n_ages]
+            [
+                self.moment_ages,
+                self.wealth_profile_ages,
+                self.wealth_anchor_ages,
+                self.estimated_chi_n_ages,
+            ]
         )
         if requested_ages.min() < min_model_age:
             raise ValueError("Requested ages start before model ages.")
@@ -98,6 +123,12 @@ def validate(self, p) -> None:
             )
         if self.max_age < self.min_age:
             raise ValueError("max_age must be at least min_age.")
+        if self.wealth_profile_max_age < self.wealth_profile_min_age:
+            raise ValueError(
+                "wealth_profile_max_age must be at least min age."
+            )
+        if self.wealth_anchor_max_age < self.wealth_anchor_min_age:
+            raise ValueError("wealth_anchor_max_age must be at least min age.")
 
 
 @dataclass(frozen=True)
@@ -219,6 +250,7 @@ def _require_finite(values: np.ndarray, label: str) -> np.ndarray:
 
 def _normalize_wealth_profile(
     profile: np.ndarray,
+    anchor_profile: np.ndarray,
     method: WealthProfileMoment,
 ) -> np.ndarray:
     """Scale a wealth age profile according to the requested convention."""
@@ -230,6 +262,13 @@ def _normalize_wealth_profile(
         if not np.isfinite(mean) or np.isclose(mean, 0.0):
             raise ValueError("Cannot normalize wealth profile with zero mean.")
         return profile / mean
+    if method == "anchor_window":
+        anchor_mean = np.nanmean(np.asarray(anchor_profile, dtype=float))
+        if not np.isfinite(anchor_mean) or np.isclose(anchor_mean, 0.0):
+            raise ValueError(
+                "Cannot normalize wealth profile with zero anchor mean."
+            )
+        return profile / anchor_mean
     raise ValueError(f"Unsupported wealth profile moment: {method}")
 
 
@@ -314,7 +353,7 @@ def labor_profile_from_cps(
         weight_col,
         config.moment_ages,
     )
-    labor = hours / ((24 - 8) * 7)
+    labor = hours / ((24 - 8) * 7)  # scale so fraction of time endowment
     return _require_finite(labor, "labor profile")
 
 
@@ -327,10 +366,17 @@ def wealth_profile_from_scf(
         scf,
         "networth_infadj",
         "wgt",
-        config.moment_ages,
+        config.wealth_profile_ages,
+    )
+    anchor_profile = _weighted_mean_by_age(
+        scf,
+        "networth_infadj",
+        "wgt",
+        config.wealth_anchor_ages,
     )
     profile = _normalize_wealth_profile(
         profile,
+        anchor_profile,
         config.wealth_profile_moment,
     )
     return _require_finite(profile, "wealth profile")
@@ -429,7 +475,9 @@ def compute_data_moments(
         if scf is None:
             scf = load_scf_wealth_data(config)
         wealth_profile = wealth_profile_from_scf(scf, config)
-        names.extend(f"net_wealth_age_{age}" for age in config.moment_ages)
+        names.extend(
+            f"net_wealth_age_{age}" for age in config.wealth_profile_ages
+        )
         values.extend(wealth_profile)
 
     if config.include_income_gini:
@@ -468,28 +516,36 @@ def compute_model_moments(
     if config is None:
         config = LifecycleCalibrationConfig()
     config.validate(p)
-    age_idx = _age_indices(config.moment_ages, p)
     lambdas = _lambdas(p)
     names: list[str] = []
     values: list[float] = []
 
     if config.include_labor_profile:
+        age_idx = _age_indices(config.moment_ages, p)
         n = np.asarray(ss_output["n"], dtype=float)
         labor = (n[age_idx, :] * lambdas.reshape(1, p.J)).sum(axis=1)
         names.extend(f"labor_supply_age_{age}" for age in config.moment_ages)
         values.extend(labor)
 
     if config.include_wealth_profile:
+        wealth_age_idx = _age_indices(config.wealth_profile_ages, p)
+        anchor_age_idx = _age_indices(config.wealth_anchor_ages, p)
         b_sp1 = np.asarray(ss_output["b_sp1"], dtype=float)
         factor = float(ss_output.get("factor", 1.0))
         wealth_profile = (
-            b_sp1[age_idx, :] * factor * lambdas.reshape(1, p.J)
+            b_sp1[wealth_age_idx, :] * factor * lambdas.reshape(1, p.J)
+        ).sum(axis=1)
+        anchor_profile = (
+            b_sp1[anchor_age_idx, :] * factor * lambdas.reshape(1, p.J)
         ).sum(axis=1)
         wealth_profile = _normalize_wealth_profile(
             wealth_profile,
+            anchor_profile,
             config.wealth_profile_moment,
         )
-        names.extend(f"net_wealth_age_{age}" for age in config.moment_ages)
+        names.extend(
+            f"net_wealth_age_{age}" for age in config.wealth_profile_ages
+        )
         values.extend(wealth_profile)
 
     if config.include_income_gini:

From 4362348dddae8353c742ae874a6f66be3beb3c78 Mon Sep 17 00:00:00 2001
From: Jason DeBacker <jdebacker@users.noreply.github.com>
Date: Wed, 17 Jun 2026 16:36:09 -0400
Subject: [PATCH 3/9] add ipykernel to depends

---
 pyproject.toml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index a177558e..8e683fe2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -119,3 +119,8 @@ markers = [
     "needs_tmd: marks tests that require the TMD data to run",
     "needs_fred: marks tests that require a FRED API key to run",
 ]
+
+[dependency-groups]
+dev = [
+    "ipykernel>=7.3.0",
+]

From 7af8b3bc75b5f051beab8f76c93dbce319e82e49 Mon Sep 17 00:00:00 2001
From: Jason DeBacker <jdebacker@users.noreply.github.com>
Date: Thu, 18 Jun 2026 00:04:21 -0400
Subject: [PATCH 4/9] use previous solutions

---
 ogusa/estimate_lifecycle_params.py | 94 ++++++++++++++++++++++++++++--
 1 file changed, 88 insertions(+), 6 deletions(-)

diff --git a/ogusa/estimate_lifecycle_params.py b/ogusa/estimate_lifecycle_params.py
index 2c1a4f89..df50f03b 100644
--- a/ogusa/estimate_lifecycle_params.py
+++ b/ogusa/estimate_lifecycle_params.py
@@ -69,6 +69,7 @@ class LifecycleCalibrationConfig:
     optimizer_options: dict = field(
         default_factory=lambda: {"maxiter": 100, "maxfun": 250}
     )
+    use_ss_solver_restart: bool = True
 
     @property
     def moment_ages(self) -> np.ndarray:
@@ -168,6 +169,18 @@ def parameter_dict(self) -> dict[str, list[float]]:
         }
 
 
+@dataclass
+class SSSolutionCache:
+    """Mutable cache for warm-starting repeated SS solves."""
+
+    previous_output: dict | None = None
+    use_ss_solver: bool = True
+
+    def reset(self) -> None:
+        """Clear the cached SS output."""
+        self.previous_output = None
+
+
 def _as_vector(values) -> np.ndarray:
     """Return values as a one-dimensional float array."""
     return np.asarray(values, dtype=float).reshape(-1)
@@ -720,6 +733,63 @@ def apply_lifecycle_params(
     )
 
 
+def solve_ss_with_cache(
+    p,
+    client=None,
+    ss_cache: SSSolutionCache | None = None,
+) -> dict:
+    """
+    Solve SS, optionally warm-starting from the previous SS output.
+
+    The direct SS_solver path keeps p.baseline unchanged, so baseline solves
+    still update the model scaling factor.  If the warm start fails, fall back
+    to SS.run_SS and refresh the cache with that solution.
+    """
+    use_cache = (
+        ss_cache is not None
+        and ss_cache.use_ss_solver
+        and ss_cache.previous_output is not None
+    )
+    if use_cache:
+        previous = ss_cache.previous_output
+        try:
+            ig_baseline = (
+                previous.get("I_g")
+                if getattr(p, "baseline_spending", False)
+                else None
+            )
+            ss_output = SS.SS_solver(
+                previous["b_sp1"],
+                previous["n"],
+                float(previous["r_p"]),
+                float(previous["r"]),
+                float(previous["w"]),
+                previous["p_m"],
+                float(previous["Y"]),
+                previous["BQ"],
+                float(previous["TR"]),
+                ig_baseline,
+                float(previous["factor"]),
+                p,
+                client,
+            )
+        except (
+            AssertionError,
+            FloatingPointError,
+            KeyError,
+            RuntimeError,
+            TypeError,
+            ValueError,
+        ):
+            ss_output = SS.run_SS(p, client=client)
+    else:
+        ss_output = SS.run_SS(p, client=client)
+
+    if ss_cache is not None:
+        ss_cache.previous_output = ss_output
+    return ss_output
+
+
 def weighting_matrix(
     moment_count: int,
     method: WeightingMethod = "identity",
@@ -822,6 +892,7 @@ def smm_objective(
     base_chi_n: np.ndarray | None = None,
     client=None,
     transform: bool = True,
+    ss_cache: SSSolutionCache | None = None,
 ) -> float:
     """Evaluate the joint lifecycle SMM objective."""
     if config is None:
@@ -842,7 +913,7 @@ def smm_objective(
             params["chi_b"],
             params["chi_n"],
         )
-        ss_output = SS.run_SS(p, client=client)
+        ss_output = solve_ss_with_cache(p, client=client, ss_cache=ss_cache)
         model_moments = compute_model_moments(ss_output, p, config)
         distance = smm_distance(model_moments, data_moments, W)
     except (
@@ -884,11 +955,21 @@ def estimate_lifecycle_params(
             bootstrap_moments=bootstrap_moments,
             ridge=config.weighting_ridge,
         )
+    ss_cache = SSSolutionCache(use_ss_solver=config.use_ss_solver_restart)
 
     est_output = opt.minimize(
         smm_objective,
         theta0,
-        args=(data_moments, W, p, config, base_chi_n, client, transform),
+        args=(
+            data_moments,
+            W,
+            p,
+            config,
+            base_chi_n,
+            client,
+            transform,
+            ss_cache,
+        ),
         method=config.optimizer_method,
         tol=config.optimizer_tol,
         options=config.optimizer_options,
@@ -906,7 +987,7 @@ def estimate_lifecycle_params(
         params["chi_b"],
         params["chi_n"],
     )
-    ss_output = SS.run_SS(p, client=client)
+    ss_output = solve_ss_with_cache(p, client=client, ss_cache=ss_cache)
     model_moments = compute_model_moments(ss_output, p, config)
 
     return LifecycleCalibrationResult(
@@ -936,6 +1017,7 @@ def compute_parameter_vcv(
         config = LifecycleCalibrationConfig()
     if base_chi_n is None:
         base_chi_n = _ss_chi_n(p)
+    ss_cache = SSSolutionCache(use_ss_solver=config.use_ss_solver_restart)
     theta_hat = _as_vector(theta_hat)
     params = unpack_lifecycle_params(
         theta_hat,
@@ -950,7 +1032,7 @@ def compute_parameter_vcv(
         params["chi_b"],
         params["chi_n"],
     )
-    ss_output = SS.run_SS(p, client=client)
+    ss_output = solve_ss_with_cache(p, client=client, ss_cache=ss_cache)
     base_moments = compute_model_moments(ss_output, p, config)
     deriv = np.zeros((base_moments.values.size, theta_hat.size))
 
@@ -974,7 +1056,7 @@ def compute_parameter_vcv(
             high_params["chi_b"],
             high_params["chi_n"],
         )
-        high_output = SS.run_SS(p, client=client)
+        high_output = solve_ss_with_cache(p, client=client, ss_cache=ss_cache)
         high_moments = compute_model_moments(high_output, p, config)
 
         low_params = unpack_lifecycle_params(
@@ -990,7 +1072,7 @@ def compute_parameter_vcv(
             low_params["chi_b"],
             low_params["chi_n"],
         )
-        low_output = SS.run_SS(p, client=client)
+        low_output = solve_ss_with_cache(p, client=client, ss_cache=ss_cache)
         low_moments = compute_model_moments(low_output, p, config)
         deriv[:, i] = (high_moments.values - low_moments.values) / (2 * step)
 

From 998fd4e7be35c3ff3baa53fd48aa6f4c0a5d5220 Mon Sep 17 00:00:00 2001
From: Jason DeBacker <jdebacker@users.noreply.github.com>
Date: Mon, 22 Jun 2026 13:25:28 -0400
Subject: [PATCH 5/9] use pct deviations for moments

---
 ogusa/compute_moments.py           |   4 +-
 ogusa/estimate_lifecycle_params.py | 267 ++++++++++++++++++++++++-----
 2 files changed, 230 insertions(+), 41 deletions(-)

diff --git a/ogusa/compute_moments.py b/ogusa/compute_moments.py
index 127f254a..1dbb0900 100644
--- a/ogusa/compute_moments.py
+++ b/ogusa/compute_moments.py
@@ -484,7 +484,7 @@ def get_macro_moments(year=2025):
         r"Investment rate $(I/K)$",
         r"Capital-Output ratio $(K/Y)$",
         r"Consumption-Output ratio $(C/Y)$",
-        r"Savings rate $(B/Y)$",
+        r"Gross savings rate $(S/Y)$",
         r"Interest rate $(r)$",
         r"Capital share of output",
         r"Labor share of output",
@@ -609,7 +609,7 @@ def get_macro_moments(year=2025):
         fred_data_q["Personal consumption expenditures"],
         fred_data_q["Nominal GDP"],
     )
-    macro_moments[r"Savings rate $(B/Y)$"] = _mean_ratio(
+    macro_moments[r"Gross savings rate $(S/Y)$"] = _mean_ratio(
         fred_data_q["Gross private savings"], fred_data_q["Nominal GDP"]
     )
     macro_moments[r"Interest rate $(r)$"] = _mean_real_rate(
diff --git a/ogusa/estimate_lifecycle_params.py b/ogusa/estimate_lifecycle_params.py
index df50f03b..7e4ddf7b 100644
--- a/ogusa/estimate_lifecycle_params.py
+++ b/ogusa/estimate_lifecycle_params.py
@@ -9,25 +9,30 @@
 
 from __future__ import annotations
 
+import logging
 from dataclasses import dataclass, field, replace
 from typing import Literal
 
 import numpy as np
 import pandas as pd
 import scipy.optimize as opt
+import ogcore
 from ogcore import SS
 from ogcore.utils import Inequality
 
 from ogusa import compute_moments, wealth
 
-SS.VERBOSE = False
+ogcore.config.VERBOSE = False
+logger = logging.getLogger(__name__)
 
 WeightingMethod = Literal["identity", "diagonal", "optimal"]
 TailMethod = Literal["scaled_default", "flat"]
 WealthProfileMoment = Literal["anchor_window", "level", "mean_normalized"]
+MomentDistanceMethod = Literal["absolute", "relative"]
 WEALTH_MOMENT_BIN_WEIGHTS = np.array(
     [0.25, 0.25, 0.20, 0.10, 0.10, 0.09, 0.01]
 )
+SAVINGS_RATE_DATA_LABEL = r"Gross savings rate $(S/Y)$"
 
 
 @dataclass(frozen=True)
@@ -37,8 +42,10 @@ class LifecycleCalibrationConfig:
 
     The default labor age window is 20 through 79.  Wealth-profile moments
     use ages 21 through 79 and are normalized by the mean wealth level from
-    ages 20 through 24.  That gives 80 parameters and 129 default moments
-    when J=10.
+    ages 20 through 24.  chi_n is parameterized as a cubic B-spline in log
+    space over all model ages, with chi_n_n_spline_knots basis functions.
+    That gives 2*J + chi_n_n_spline_knots parameters and ~130 default
+    moments when J=10 and chi_n_n_spline_knots=10 (30 parameters total).
     """
 
     min_age: int = 20
@@ -54,8 +61,10 @@ class LifecycleCalibrationConfig:
     include_labor_profile: bool = True
     include_wealth_profile: bool = True
     include_income_gini: bool = True
+    include_savings_rate: bool = True
     include_wealth_distribution: bool = True
     include_inheritance_moments: bool = False
+    macro_year: int = 2025
     cps_years: tuple[int, ...] = (2023, 2022)
     scf_yrs_list: tuple[int, ...] = (2019,)
     cps_directory: str | None = None
@@ -67,9 +76,14 @@ class LifecycleCalibrationConfig:
     optimizer_method: str = "L-BFGS-B"
     optimizer_tol: float = 1e-10
     optimizer_options: dict = field(
-        default_factory=lambda: {"maxiter": 100, "maxfun": 250}
+        default_factory=lambda: {"maxiter": 100, "maxfun": 500}
     )
+    chi_n_n_spline_knots: int = 10
+    chi_n_spline_degree: int = 3
+    moment_distance_method: MomentDistanceMethod = "relative"
+    moment_distance_floor: float = 1e-8
     use_ss_solver_restart: bool = True
+    log_optimizer_progress: bool = True
 
     @property
     def moment_ages(self) -> np.ndarray:
@@ -209,6 +223,58 @@ def _age_indices(ages: np.ndarray, p) -> np.ndarray:
     return np.array([_age_to_index(age, p) for age in ages], dtype=int)
 
 
+def _build_chi_n_spline_basis(
+    ages: np.ndarray,
+    n_basis: int,
+    degree: int = 3,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Build a B-spline design matrix for the chi_n age profile.
+
+    The spline is defined in log space: evaluating ``B @ gamma`` gives
+    ``log(chi_n)`` at each age, so ``chi_n = exp(B @ gamma)`` is always
+    positive regardless of the coefficient values.
+
+    Args:
+        ages: Age values at which to evaluate the basis (length N).
+        n_basis: Number of B-spline basis functions (= number of free
+            coefficients).  Must satisfy n_basis >= degree + 1.
+        degree: Polynomial degree of the spline (default 3 = cubic).
+
+    Returns:
+        B: Design matrix of shape (N, n_basis).
+        knots: Full knot vector used to construct the basis.
+    """
+    from scipy.interpolate import BSpline
+
+    ages_f = np.asarray(ages, dtype=float)
+    age_min, age_max = ages_f[0], ages_f[-1]
+    n_internal = n_basis - degree - 1
+    if n_internal < 0:
+        raise ValueError(
+            f"n_basis={n_basis} is too small for degree={degree}. "
+            f"Need n_basis >= degree + 1 = {degree + 1}."
+        )
+    internal = (
+        np.linspace(age_min, age_max, n_internal + 2)[1:-1]
+        if n_internal > 0
+        else np.array([], dtype=float)
+    )
+    knots = np.concatenate(
+        [
+            np.repeat(age_min, degree + 1),
+            internal,
+            np.repeat(age_max, degree + 1),
+        ]
+    )
+    n_cols = len(knots) - degree - 1
+    B = np.zeros((len(ages_f), n_cols))
+    for i in range(n_cols):
+        c = np.zeros(n_cols)
+        c[i] = 1.0
+        B[:, i] = BSpline(knots, c, degree)(ages_f)
+    return B, knots
+
+
 def _weighted_mean(values, weights) -> float:
     """Return a weighted mean after dropping nonfinite observations."""
     data = pd.DataFrame({"value": values, "weight": weights})
@@ -403,6 +469,24 @@ def income_gini_data_moment(
     return float(moments["Gini coefficient, income"])
 
 
+def savings_rate_data_moment(macro_year: int = 2025) -> float:
+    """Compute the aggregate savings-rate data moment."""
+    moments = compute_moments.get_macro_moments(year=macro_year)
+    return float(moments[SAVINGS_RATE_DATA_LABEL])
+
+
+def model_savings_rate_moment(ss_output: dict, p) -> float:
+    """Compute the model gross aggregate savings rate."""
+    output = float(ss_output["Y"])
+    if np.isclose(output, 0.0):
+        raise ValueError("Cannot compute savings rate with zero output.")
+    growth = (1 + p.g_n_ss) * np.exp(p.g_y)
+    gross_saving_flow = (growth - 1.0) * float(ss_output["B"]) + (
+        p.delta * float(ss_output["K_d"])
+    )
+    return gross_saving_flow / output
+
+
 def compute_inheritance_moments_from_scf(
     scf: pd.DataFrame,
     amount_col: str,
@@ -469,6 +553,7 @@ def compute_data_moments(
     scf: pd.DataFrame | None = None,
     income_year: int | None = None,
     inheritance_moments: MomentSet | None = None,
+    savings_rate: float | None = None,
 ) -> MomentSet:
     """Compute the stacked data moment vector for the SMM objective."""
     if config is None:
@@ -497,6 +582,12 @@ def compute_data_moments(
         names.append("income_gini")
         values.append(income_gini_data_moment(income_year=income_year))
 
+    if config.include_savings_rate:
+        if savings_rate is None:
+            savings_rate = savings_rate_data_moment(config.macro_year)
+        names.append("savings_rate")
+        values.append(float(savings_rate))
+
     if config.include_wealth_distribution:
         if scf is None:
             scf = load_scf_wealth_data(config)
@@ -567,6 +658,10 @@ def compute_model_moments(
         names.append("income_gini")
         values.append(income_ineq.gini())
 
+    if config.include_savings_rate:
+        names.append("savings_rate")
+        values.append(model_savings_rate_moment(ss_output, p))
+
     if config.include_wealth_distribution:
         names.extend(_wealth_distribution_moment_names())
         values.extend(_model_wealth_distribution_moments(ss_output, p))
@@ -633,27 +728,45 @@ def pack_lifecycle_params(
     config: LifecycleCalibrationConfig | None = None,
     transform: bool = True,
 ) -> np.ndarray:
-    """Pack natural lifecycle parameters into an optimizer vector."""
+    """Pack natural lifecycle parameters into an optimizer vector.
+
+    chi_n (length p.S) is projected onto the B-spline basis in log space.
+    The returned vector layout is [beta_trans, chi_b_trans, gamma] where
+    gamma holds the chi_n spline coefficients (unconstrained when transform
+    is True, since the log-space spline already enforces positivity via exp).
+    """
     if config is None:
         config = LifecycleCalibrationConfig()
     config.validate(p)
     beta_annual = _as_vector(beta_annual)
     chi_b = _as_vector(chi_b)
     chi_n = _as_vector(chi_n)
-    est_idx = _age_indices(config.estimated_chi_n_ages, p)
-    chi_n_est = chi_n[est_idx]
-    params = np.concatenate([beta_annual, chi_b, chi_n_est])
+
+    starting_age = int(getattr(p, "starting_age", 20))
+    all_ages = np.arange(starting_age, starting_age + p.S)
+    B, _ = _build_chi_n_spline_basis(
+        all_ages, config.chi_n_n_spline_knots, config.chi_n_spline_degree
+    )
+
     if not transform:
-        return params
+        # Project chi_n directly (no log); coefficients are in natural space.
+        gamma, _, _, _ = np.linalg.lstsq(B, chi_n, rcond=None)
+        return np.concatenate([beta_annual, chi_b, gamma])
+
     if np.any((beta_annual <= 0) | (beta_annual >= 1)):
         raise ValueError(
             "beta_annual values must be strictly between 0 and 1."
         )
-    if np.any(params[beta_annual.size :] <= 0):
-        raise ValueError("chi_b and chi_n values must be positive.")
+    if np.any(chi_b <= 0):
+        raise ValueError("chi_b values must be positive.")
+    if np.any(chi_n <= 0):
+        raise ValueError("chi_n values must be positive.")
+
     beta_trans = np.log(beta_annual / (1 - beta_annual))
-    chi_trans = np.log(params[beta_annual.size :])
-    return np.concatenate([beta_trans, chi_trans])
+    chi_b_trans = np.log(chi_b)
+    # Fit spline to log(chi_n) over all model ages; gamma is unconstrained.
+    gamma, _, _, _ = np.linalg.lstsq(B, np.log(chi_n), rcond=None)
+    return np.concatenate([beta_trans, chi_b_trans, gamma])
 
 
 def unpack_lifecycle_params(
@@ -663,39 +776,48 @@ def unpack_lifecycle_params(
     base_chi_n: np.ndarray | None = None,
     transform: bool = True,
 ) -> dict[str, np.ndarray]:
-    """Unpack an optimizer vector into natural lifecycle parameters."""
+    """Unpack an optimizer vector into natural lifecycle parameters.
+
+    The chi_n block of theta holds B-spline coefficients (gamma).  When
+    transform=True the spline operates in log space and chi_n = exp(B @ gamma).
+    The base_chi_n argument is accepted for backward compatibility but is no
+    longer used; the spline covers all model ages directly.
+    """
     if config is None:
         config = LifecycleCalibrationConfig()
     config.validate(p)
-    if base_chi_n is None:
-        base_chi_n = _ss_chi_n(p)
 
     theta = _as_vector(theta)
     n_beta = p.J
     n_chi_b = p.J
-    n_chi_n = config.estimated_chi_n_ages.size
-    expected = n_beta + n_chi_b + n_chi_n
+    n_gamma = config.chi_n_n_spline_knots
+    expected = n_beta + n_chi_b + n_gamma
     if theta.size != expected:
         raise ValueError(f"Expected {expected} parameters, got {theta.size}.")
 
     beta_raw = theta[:n_beta]
     chi_b_raw = theta[n_beta : n_beta + n_chi_b]
-    chi_n_raw = theta[n_beta + n_chi_b :]
+    gamma = theta[n_beta + n_chi_b :]
+
+    starting_age = int(getattr(p, "starting_age", 20))
+    all_ages = np.arange(starting_age, starting_age + p.S)
+    B, _ = _build_chi_n_spline_basis(
+        all_ages, config.chi_n_n_spline_knots, config.chi_n_spline_degree
+    )
+
     if transform:
         beta_annual = 1 / (1 + np.exp(-beta_raw))
         chi_b = np.exp(chi_b_raw)
-        chi_n_est = np.exp(chi_n_raw)
+        chi_n = np.exp(B @ gamma)
     else:
         beta_annual = beta_raw
         chi_b = chi_b_raw
-        chi_n_est = chi_n_raw
+        chi_n = B @ gamma
 
-    chi_n = build_chi_n_profile(chi_n_est, base_chi_n, p, config)
     return {
         "beta_annual": beta_annual,
         "chi_b": chi_b,
         "chi_n": chi_n,
-        "chi_n_estimated": chi_n_est,
     }
 
 
@@ -823,6 +945,7 @@ def bootstrap_data_moments(
     cps: pd.DataFrame | None = None,
     scf: pd.DataFrame | None = None,
     seed: int | None = None,
+    savings_rate: float | None = None,
 ) -> np.ndarray:
     """Bootstrap the data moments available from CPS and SCF microdata."""
     if config is None:
@@ -834,10 +957,17 @@ def bootstrap_data_moments(
     ):
         scf = load_scf_wealth_data(config)
 
-    point_moments = compute_data_moments(p, config, cps=cps, scf=scf)
+    point_moments = compute_data_moments(
+        p,
+        config,
+        cps=cps,
+        scf=scf,
+        savings_rate=savings_rate,
+    )
     resampled_config = replace(
         config,
         include_income_gini=False,
+        include_savings_rate=False,
         include_inheritance_moments=False,
     )
     rng = np.random.default_rng(seed)
@@ -873,11 +1003,31 @@ def smm_distance(
     model_moments: MomentSet,
     data_moments: MomentSet,
     W: np.ndarray,
+    method: MomentDistanceMethod = "relative",
+    floor: float = 1e-8,
 ) -> float:
-    """Compute the quadratic SMM distance."""
+    """Compute the quadratic SMM distance.
+
+    Args:
+        model_moments: Simulated model moments.
+        data_moments: Empirical data moments.
+        W: Positive semi-definite weighting matrix.
+        method: ``"relative"`` divides each residual by ``|data_moment|``
+            (floored at ``floor``) so all moments are on a comparable
+            percentage-deviation scale.  ``"absolute"`` uses raw levels.
+        floor: Minimum absolute value used as the denominator when
+            method="relative", preventing division by zero for near-zero
+            moments such as bottom wealth shares.
+    """
     if model_moments.names != data_moments.names:
         raise ValueError("Model and data moments are not aligned.")
-    diff = model_moments.values - data_moments.values
+    m = model_moments.values
+    d = data_moments.values
+    if method == "relative":
+        safe_denom = np.where(np.abs(d) > floor, np.abs(d), floor)
+        diff = (m - d) / safe_denom
+    else:
+        diff = m - d
     if not np.all(np.isfinite(diff)):
         return np.inf
     return float(diff.T @ W @ diff)
@@ -915,7 +1065,13 @@ def smm_objective(
         )
         ss_output = solve_ss_with_cache(p, client=client, ss_cache=ss_cache)
         model_moments = compute_model_moments(ss_output, p, config)
-        distance = smm_distance(model_moments, data_moments, W)
+        distance = smm_distance(
+            model_moments,
+            data_moments,
+            W,
+            method=config.moment_distance_method,
+            floor=config.moment_distance_floor,
+        )
     except (
         AssertionError,
         FloatingPointError,
@@ -938,6 +1094,7 @@ def estimate_lifecycle_params(
     bootstrap_moments: np.ndarray | None = None,
     client=None,
     transform: bool = True,
+    savings_rate: float | None = None,
 ) -> LifecycleCalibrationResult:
     """Estimate beta_annual, chi_b, and chi_n jointly by SMM."""
     if config is None:
@@ -947,7 +1104,11 @@ def estimate_lifecycle_params(
     if theta0 is None:
         theta0 = initial_lifecycle_theta(p, config, transform=transform)
     if data_moments is None:
-        data_moments = compute_data_moments(p, config)
+        data_moments = compute_data_moments(
+            p,
+            config,
+            savings_rate=savings_rate,
+        )
     if W is None:
         W = weighting_matrix(
             data_moments.values.size,
@@ -956,23 +1117,51 @@ def estimate_lifecycle_params(
             ridge=config.weighting_ridge,
         )
     ss_cache = SSSolutionCache(use_ss_solver=config.use_ss_solver_restart)
+    objective_args = (
+        data_moments,
+        W,
+        p,
+        config,
+        base_chi_n,
+        client,
+        transform,
+        ss_cache,
+    )
+    optimizer_state = {
+        "function_evals": 0,
+        "iteration": 0,
+        "last_objective": None,
+    }
+
+    def tracked_objective(theta):
+        value = smm_objective(theta, *objective_args)
+        optimizer_state["function_evals"] += 1
+        optimizer_state["last_objective"] = value
+        return value
+
+    def optimizer_callback(_theta):
+        optimizer_state["iteration"] += 1
+        max_iterations = config.optimizer_options.get("maxiter", "?")
+        objective = optimizer_state["last_objective"]
+        if objective is None:
+            objective_text = "unavailable"
+        else:
+            objective_text = f"{objective:.6e}"
+        logger.info(
+            "Lifecycle SMM iteration %s/%s: objective=%s, function_evals=%s",
+            optimizer_state["iteration"],
+            max_iterations,
+            objective_text,
+            optimizer_state["function_evals"],
+        )
 
     est_output = opt.minimize(
-        smm_objective,
+        tracked_objective,
         theta0,
-        args=(
-            data_moments,
-            W,
-            p,
-            config,
-            base_chi_n,
-            client,
-            transform,
-            ss_cache,
-        ),
         method=config.optimizer_method,
         tol=config.optimizer_tol,
         options=config.optimizer_options,
+        callback=optimizer_callback if config.log_optimizer_progress else None,
     )
     params = unpack_lifecycle_params(
         est_output.x,

From 2ee85264d3847736e20134db0f3ce708678b586c Mon Sep 17 00:00:00 2001
From: Jason DeBacker <jdebacker@users.noreply.github.com>
Date: Mon, 22 Jun 2026 13:30:03 -0400
Subject: [PATCH 6/9] add unit tests

---
 tests/test_estimate_lifecycle_params.py | 463 ++++++++++++++++++++++++
 1 file changed, 463 insertions(+)
 create mode 100644 tests/test_estimate_lifecycle_params.py

diff --git a/tests/test_estimate_lifecycle_params.py b/tests/test_estimate_lifecycle_params.py
new file mode 100644
index 00000000..0e7147aa
--- /dev/null
+++ b/tests/test_estimate_lifecycle_params.py
@@ -0,0 +1,463 @@
+"""
+Tests for the joint lifecycle preference calibration helpers.
+"""
+
+import logging
+
+import numpy as np
+import pandas as pd
+
+from ogusa import estimate_lifecycle_params as elp
+
+
+class MockParams:
+    """
+    Minimal parameter object for lifecycle calibration helper tests.
+    """
+
+    S = 80
+    J = 10
+    starting_age = 20
+    ending_age = 100
+    lambdas = np.array(
+        [0.25, 0.25, 0.2, 0.1, 0.1, 0.09, 0.005, 0.004, 0.0009, 0.0001]
+    ).reshape(10, 1)
+    omega_SS = np.ones(80) / 80
+    g_n_ss = 0.01
+    g_y = 0.02
+    delta = 0.05
+    beta_annual = np.linspace(0.91, 0.995, 10)
+    chi_b = np.ones(10) * 80
+    chi_n = np.linspace(20, 80, 80)
+
+
+def _mock_ss_output(p, scale=1.0):
+    """
+    Build the subset of SS output used to warm-start repeated solves.
+    """
+    return {
+        "b_sp1": np.ones((p.S, p.J)) * scale,
+        "n": np.ones((p.S, p.J)) * 0.4 * scale,
+        "r_p": 0.04 * scale,
+        "r": 0.04 * scale,
+        "w": 1.2 * scale,
+        "p_m": np.ones(1) * scale,
+        "B": 2.0 * scale,
+        "K_d": 3.0 * scale,
+        "Y": 10.0 * scale,
+        "BQ": np.ones(p.J) * scale,
+        "TR": 0.2 * scale,
+        "factor": scale,
+    }
+
+
+def test_default_config_dimensions():
+    """
+    Test that the default setup gives 30 parameters and 130 moments.
+    """
+    p = MockParams()
+    config = elp.LifecycleCalibrationConfig()
+    assert config.estimated_chi_n_ages.size == 60
+    assert config.wealth_profile_ages.size == 59
+    assert config.chi_n_n_spline_knots == 10
+    assert config.chi_n_spline_degree == 3
+    assert config.moment_distance_method == "relative"
+    assert p.J + p.J + config.chi_n_n_spline_knots == 30
+
+    ss_output = {
+        "n": np.ones((p.S, p.J)) * 0.35,
+        "b_sp1": np.arange(1, p.S * p.J + 1).reshape(p.S, p.J),
+        "before_tax_income": (
+            np.arange(1, p.S * p.J + 1).reshape(p.S, p.J) + 1000
+        ),
+        "B": 2.0,
+        "K_d": 3.0,
+        "Y": 10.0,
+        "factor": 2.0,
+    }
+
+    moments = elp.compute_model_moments(ss_output, p, config)
+
+    assert len(moments.names) == 130
+    assert moments.values.shape == (130,)
+    assert moments.names[0] == "labor_supply_age_20"
+    assert moments.names[59] == "labor_supply_age_79"
+    assert moments.names[60] == "net_wealth_age_21"
+    assert moments.names[118] == "net_wealth_age_79"
+    assert moments.names[119] == "income_gini"
+    assert moments.names[120] == "savings_rate"
+    growth = (1 + p.g_n_ss) * np.exp(p.g_y)
+    expected_savings_rate = ((growth - 1.0) * 2.0 + p.delta * 3.0) / 10.0
+    assert np.allclose(moments.values[120], expected_savings_rate)
+    assert moments.names[-1] == "wealth_var_log"
+
+
+def test_pack_unpack_lifecycle_params_roundtrip():
+    """
+    Test transformed packing and unpacking of lifecycle parameters.
+    """
+    p = MockParams()
+    config = elp.LifecycleCalibrationConfig()
+    beta = np.linspace(0.92, 0.99, p.J)
+    chi_b = np.linspace(50, 100, p.J)
+    ages = np.arange(p.starting_age, p.starting_age + p.S)
+    basis, _ = elp._build_chi_n_spline_basis(
+        ages,
+        config.chi_n_n_spline_knots,
+        config.chi_n_spline_degree,
+    )
+    gamma = np.linspace(np.log(10), np.log(70), config.chi_n_n_spline_knots)
+    chi_n = np.exp(basis @ gamma)
+
+    theta = elp.pack_lifecycle_params(beta, chi_b, chi_n, p, config)
+    unpacked = elp.unpack_lifecycle_params(
+        theta,
+        p,
+        config,
+        base_chi_n=p.chi_n,
+    )
+
+    assert theta.size == 30
+    assert np.allclose(unpacked["beta_annual"], beta)
+    assert np.allclose(unpacked["chi_b"], chi_b)
+    assert np.allclose(unpacked["chi_n"], chi_n)
+
+
+def test_build_chi_n_profile_scales_default_tail():
+    """
+    Test that the default chi_n tail shape is preserved and scaled.
+    """
+    p = MockParams()
+    config = elp.LifecycleCalibrationConfig()
+    base_chi_n = np.arange(1, p.S + 1, dtype=float)
+    estimated = np.ones(config.estimated_chi_n_ages.size) * 10
+
+    full = elp.build_chi_n_profile(estimated, base_chi_n, p, config)
+
+    assert np.allclose(full[:60], 10)
+    expected_scale = 10 / base_chi_n[59]
+    assert np.allclose(full[60:], base_chi_n[60:] * expected_scale)
+
+
+def test_weighting_matrix_handles_singular_bootstrap_vcv():
+    """
+    Test pseudo-inverse/ridge handling for singular bootstrap VCV matrices.
+    """
+    boot = np.ones((5, 3))
+
+    W = elp.weighting_matrix(
+        3,
+        method="optimal",
+        bootstrap_moments=boot,
+        ridge=1e-6,
+    )
+
+    assert W.shape == (3, 3)
+    assert np.all(np.isfinite(W))
+    assert np.allclose(W, W.T)
+
+
+def test_smm_distance_supports_relative_and_absolute_residuals():
+    """
+    Test percentage-deviation and level-deviation distance calculations.
+    """
+    model = elp.MomentSet(("a", "b", "c"), np.array([2.0, 6.0, 2e-9]))
+    data = elp.MomentSet(("a", "b", "c"), np.array([1.0, 3.0, 0.0]))
+    W = np.eye(3)
+
+    absolute = elp.smm_distance(model, data, W, method="absolute")
+    relative = elp.smm_distance(
+        model,
+        data,
+        W,
+        method="relative",
+        floor=1e-8,
+    )
+
+    assert np.allclose(absolute, 1.0**2 + 3.0**2 + (2e-9) ** 2)
+    assert np.allclose(relative, 1.0**2 + 1.0**2 + 0.2**2)
+
+
+def test_compute_inheritance_moments_from_scf():
+    """
+    Test optional inherited-transfer moments from full SCF-like data.
+    """
+    scf = pd.DataFrame(
+        {
+            "inheritance": [0.0, 100.0, 200.0],
+            "received": [0, 1, 1],
+            "networth_infadj": [50.0, 1000.0, 2000.0],
+            "wgt": [1.0, 1.0, 2.0],
+        }
+    )
+
+    moments = elp.compute_inheritance_moments_from_scf(
+        scf,
+        amount_col="inheritance",
+        received_col="received",
+    )
+
+    assert moments.names == (
+        "inheritance_received_rate",
+        "inheritance_amount_conditional_mean",
+        "inheritance_to_networth_mean",
+    )
+    assert np.allclose(moments.values[0], 0.75)
+    assert np.allclose(moments.values[1], (100 + 2 * 200) / 3)
+
+
+def test_compute_data_moments_with_synthetic_microdata():
+    """
+    Test data moment construction without Tax-Calculator or web access.
+    """
+    p = MockParams()
+    ages = np.arange(20, 80)
+    cps = pd.DataFrame(
+        {
+            "age": ages,
+            "hours_per_week": np.linspace(20, 40, ages.size),
+            "weight": np.ones(ages.size),
+        }
+    )
+    scf = pd.DataFrame(
+        {
+            "age": ages,
+            "networth_infadj": np.linspace(1000, 100000, ages.size),
+            "networth": np.linspace(1000, 100000, ages.size),
+            "wgt": np.ones(ages.size),
+        }
+    )
+    config = elp.LifecycleCalibrationConfig(include_income_gini=False)
+
+    moments = elp.compute_data_moments(
+        p,
+        config,
+        cps=cps,
+        scf=scf,
+        savings_rate=0.17,
+    )
+
+    assert len(moments.names) == 129
+    assert moments.values.shape == (129,)
+    anchor_mean = scf.loc[scf["age"].between(20, 24), "networth_infadj"].mean()
+    age_21_wealth = scf.loc[scf["age"] == 21, "networth_infadj"].iloc[0]
+    assert np.allclose(moments.values[60], age_21_wealth / anchor_mean)
+    assert moments.names[119] == "savings_rate"
+    assert np.allclose(moments.values[119], 0.17)
+    assert moments.names[-9:] == elp._wealth_distribution_moment_names()
+
+
+def test_savings_rate_data_moment_uses_macro_moment(monkeypatch):
+    """
+    Test that the savings-rate data moment comes from macro moments.
+    """
+
+    def fake_get_macro_moments(year):
+        assert year == 2030
+        return {elp.SAVINGS_RATE_DATA_LABEL: 0.18}
+
+    monkeypatch.setattr(
+        elp.compute_moments,
+        "get_macro_moments",
+        fake_get_macro_moments,
+    )
+
+    assert np.allclose(elp.savings_rate_data_moment(macro_year=2030), 0.18)
+
+
+def test_estimate_lifecycle_params_logs_optimizer_progress(
+    monkeypatch, caplog
+):
+    """
+    Test that optimizer progress logging uses cached objective values.
+    """
+    p = MockParams()
+    config = elp.LifecycleCalibrationConfig(
+        include_labor_profile=False,
+        include_wealth_profile=False,
+        include_income_gini=False,
+        include_savings_rate=False,
+        include_wealth_distribution=False,
+        optimizer_options={"maxiter": 2, "maxfun": 4},
+        log_optimizer_progress=True,
+    )
+    data_moments = elp.MomentSet(("moment",), np.array([0.0]))
+    W = np.eye(1)
+    theta0 = np.array([0.1])
+    objective_values = [4.0, 3.0]
+
+    def fake_smm_objective(theta, *args):
+        return objective_values.pop(0)
+
+    def fake_minimize(
+        fun,
+        x0,
+        method=None,
+        tol=None,
+        options=None,
+        callback=None,
+    ):
+        first = fun(x0)
+        assert np.isclose(first, 4.0)
+        callback(x0)
+        second = fun(x0)
+        assert np.isclose(second, 3.0)
+        callback(x0)
+        return elp.opt.OptimizeResult(
+            x=x0,
+            fun=second,
+            success=True,
+            message="ok",
+            nit=2,
+            nfev=2,
+        )
+
+    def fake_unpack_lifecycle_params(*args, **kwargs):
+        return {
+            "beta_annual": np.ones(p.J) * 0.95,
+            "chi_b": np.ones(p.J),
+            "chi_n": np.ones(p.S),
+        }
+
+    monkeypatch.setattr(elp, "smm_objective", fake_smm_objective)
+    monkeypatch.setattr(elp.opt, "minimize", fake_minimize)
+    monkeypatch.setattr(
+        elp, "unpack_lifecycle_params", fake_unpack_lifecycle_params
+    )
+    monkeypatch.setattr(elp, "apply_lifecycle_params", lambda *args: None)
+    monkeypatch.setattr(
+        elp,
+        "solve_ss_with_cache",
+        lambda *args, **kwargs: _mock_ss_output(p),
+    )
+    monkeypatch.setattr(
+        elp,
+        "compute_model_moments",
+        lambda *args, **kwargs: data_moments,
+    )
+    caplog.set_level(logging.INFO, logger=elp.logger.name)
+
+    result = elp.estimate_lifecycle_params(
+        p,
+        config=config,
+        theta0=theta0,
+        data_moments=data_moments,
+        W=W,
+    )
+
+    assert result.objective_value == 3.0
+    assert objective_values == []
+    assert "Lifecycle SMM iteration 1/2" in caplog.text
+    assert "objective=4.000000e+00" in caplog.text
+    assert "Lifecycle SMM iteration 2/2" in caplog.text
+    assert "objective=3.000000e+00" in caplog.text
+    assert "function_evals=2" in caplog.text
+
+
+def test_solve_ss_with_cache_uses_run_ss_then_ss_solver(monkeypatch):
+    """
+    Test that repeated SS solves use the previous solution as a warm start.
+    """
+    p = MockParams()
+    calls = []
+
+    def fake_run_ss(params, client=None):
+        calls.append(("run", client))
+        return _mock_ss_output(params, scale=1.0)
+
+    def fake_ss_solver(
+        b_sp1,
+        n,
+        r_p,
+        r,
+        w,
+        p_m,
+        y,
+        bq,
+        tr,
+        ig_baseline,
+        factor,
+        params,
+        client=None,
+    ):
+        calls.append(("solver", factor, client))
+        assert np.allclose(b_sp1, np.ones((params.S, params.J)))
+        assert np.allclose(n, np.ones((params.S, params.J)) * 0.4)
+        assert np.isclose(r_p, 0.04)
+        assert np.isclose(r, 0.04)
+        assert np.isclose(w, 1.2)
+        assert np.allclose(p_m, np.ones(1))
+        assert np.isclose(y, 10.0)
+        assert np.allclose(bq, np.ones(params.J))
+        assert np.isclose(tr, 0.2)
+        assert ig_baseline is None
+        return _mock_ss_output(params, scale=2.0)
+
+    monkeypatch.setattr(elp.SS, "run_SS", fake_run_ss)
+    monkeypatch.setattr(elp.SS, "SS_solver", fake_ss_solver)
+
+    cache = elp.SSSolutionCache(use_ss_solver=True)
+    first = elp.solve_ss_with_cache(p, client="client", ss_cache=cache)
+    second = elp.solve_ss_with_cache(p, client="client", ss_cache=cache)
+
+    assert [call[0] for call in calls] == ["run", "solver"]
+    assert first["factor"] == 1.0
+    assert second["factor"] == 2.0
+    assert cache.previous_output is second
+
+
+def test_solve_ss_with_cache_falls_back_to_run_ss(monkeypatch):
+    """
+    Test fallback to SS.run_SS if the direct solver restart fails.
+    """
+    p = MockParams()
+    calls = []
+
+    def fake_run_ss(params, client=None):
+        calls.append("run")
+        return _mock_ss_output(params, scale=3.0)
+
+    def fake_ss_solver(*args, **kwargs):
+        calls.append("solver")
+        raise RuntimeError("failed warm start")
+
+    monkeypatch.setattr(elp.SS, "run_SS", fake_run_ss)
+    monkeypatch.setattr(elp.SS, "SS_solver", fake_ss_solver)
+
+    cache = elp.SSSolutionCache(
+        previous_output=_mock_ss_output(p, scale=1.0),
+        use_ss_solver=True,
+    )
+    output = elp.solve_ss_with_cache(p, ss_cache=cache)
+
+    assert calls == ["solver", "run"]
+    assert output["factor"] == 3.0
+    assert cache.previous_output is output
+
+
+def test_solve_ss_with_cache_can_disable_ss_solver(monkeypatch):
+    """
+    Test that the restart path can be disabled for robustness checks.
+    """
+    p = MockParams()
+    calls = []
+
+    def fake_run_ss(params, client=None):
+        calls.append("run")
+        return _mock_ss_output(params, scale=4.0)
+
+    def fake_ss_solver(*args, **kwargs):
+        calls.append("solver")
+        return _mock_ss_output(p, scale=5.0)
+
+    monkeypatch.setattr(elp.SS, "run_SS", fake_run_ss)
+    monkeypatch.setattr(elp.SS, "SS_solver", fake_ss_solver)
+
+    cache = elp.SSSolutionCache(
+        previous_output=_mock_ss_output(p, scale=1.0),
+        use_ss_solver=False,
+    )
+    output = elp.solve_ss_with_cache(p, ss_cache=cache)
+
+    assert calls == ["run"]
+    assert output["factor"] == 4.0

From d73565578a92451cb7aa0125190a87f7b2234890 Mon Sep 17 00:00:00 2001
From: Jason DeBacker <jdebacker@users.noreply.github.com>
Date: Mon, 22 Jun 2026 22:24:58 -0400
Subject: [PATCH 7/9] use dfo-ls as optimizer

---
 ogusa/estimate_lifecycle_params.py | 283 +++++++++++++++++++++++------
 pyproject.toml                     |   1 +
 2 files changed, 231 insertions(+), 53 deletions(-)

diff --git a/ogusa/estimate_lifecycle_params.py b/ogusa/estimate_lifecycle_params.py
index 7e4ddf7b..76305a20 100644
--- a/ogusa/estimate_lifecycle_params.py
+++ b/ogusa/estimate_lifecycle_params.py
@@ -15,7 +15,6 @@
 
 import numpy as np
 import pandas as pd
-import scipy.optimize as opt
 import ogcore
 from ogcore import SS
 from ogcore.utils import Inequality
@@ -73,11 +72,11 @@ class LifecycleCalibrationConfig:
     bootstrap_iterations: int = 1000
     weighting_method: WeightingMethod = "identity"
     weighting_ridge: float = 1e-8
-    optimizer_method: str = "L-BFGS-B"
-    optimizer_tol: float = 1e-10
-    optimizer_options: dict = field(
-        default_factory=lambda: {"maxiter": 100, "maxfun": 500}
-    )
+    n_starts: int = 3
+    start_radius: float = 0.2
+    dfols_rhoend: float = 1e-6
+    dfols_maxfun: int | None = None
+    bound_epsilon: float = 1e-4
     chi_n_n_spline_knots: int = 10
     chi_n_spline_degree: int = 3
     moment_distance_method: MomentDistanceMethod = "relative"
@@ -168,10 +167,12 @@ class LifecycleCalibrationResult:
     chi_b: np.ndarray
     chi_n: np.ndarray
     objective_value: float
-    optimizer_result: opt.OptimizeResult
+    optimizer_result: object
     data_moments: MomentSet
     model_moments: MomentSet
     weighting_matrix: np.ndarray
+    best_start_index: int = 0
+    all_start_results: list = field(default_factory=list)
 
     @property
     def parameter_dict(self) -> dict[str, list[float]]:
@@ -1085,6 +1086,165 @@ def smm_objective(
     return distance
 
 
+def _apply_weight_sqrt(residuals: np.ndarray, W: np.ndarray) -> np.ndarray:
+    """Return L.T @ residuals where W = L @ L.T, so ||result||^2 = r^T W r.
+
+    For identity W this is a no-op.  For diagonal W it multiplies element-wise
+    by sqrt of the diagonal.  For a full positive definite W it uses the
+    Cholesky factor, falling back to the diagonal if W is not PD.
+    """
+    if np.allclose(W, np.eye(len(W))):
+        return residuals
+    try:
+        L = np.linalg.cholesky(W)
+        return L.T @ residuals
+    except np.linalg.LinAlgError:
+        return np.sqrt(np.maximum(np.diag(W), 0.0)) * residuals
+
+
+def smm_residual(
+    theta: np.ndarray,
+    data_moments: MomentSet,
+    W: np.ndarray,
+    p,
+    config: LifecycleCalibrationConfig,
+    base_chi_n: np.ndarray | None = None,
+    client=None,
+    transform: bool = True,
+    ss_cache: SSSolutionCache | None = None,
+) -> np.ndarray:
+    """Return the weighted moment residual vector for DFO-LS.
+
+    DFO-LS minimises ``||f(x)||^2``.  Pre-multiplying by ``W^{1/2}`` ensures
+    that ``||result||^2 == r^T W r``, matching the SMM objective.  On solver
+    failure returns a large constant vector so DFO-LS avoids that region.
+    """
+    n_moments = data_moments.values.size
+    try:
+        params = unpack_lifecycle_params(
+            theta, p, config, base_chi_n=base_chi_n, transform=transform
+        )
+        apply_lifecycle_params(
+            p, params["beta_annual"], params["chi_b"], params["chi_n"]
+        )
+        ss_output = solve_ss_with_cache(p, client=client, ss_cache=ss_cache)
+        model_moments = compute_model_moments(ss_output, p, config)
+        m = model_moments.values
+        d = data_moments.values
+        if config.moment_distance_method == "relative":
+            safe_denom = np.where(
+                np.abs(d) > config.moment_distance_floor,
+                np.abs(d),
+                config.moment_distance_floor,
+            )
+            residuals = (m - d) / safe_denom
+        else:
+            residuals = m - d
+        weighted = _apply_weight_sqrt(residuals, W)
+        if not np.all(np.isfinite(weighted)):
+            return np.full(n_moments, 1e15)
+        return weighted
+    except (
+        AssertionError,
+        FloatingPointError,
+        KeyError,
+        RuntimeError,
+        TypeError,
+        ValueError,
+    ):
+        return np.full(n_moments, 1e15)
+
+
+def _extract_dfols_bounds(
+    p,
+    config: LifecycleCalibrationConfig,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Build DFO-LS bound arrays from the ParamTools validators in p.
+
+    Bounds are read from ``p._data[param]['validators']['range']`` and then
+    mapped into the same transformed space used by the optimizer:
+
+    * ``beta_annual`` — logit transform; validator min/max become the
+      logit of the natural bounds (with ``bound_epsilon`` as a floor so
+      log(0) is avoided).
+    * ``chi_b`` — log transform; same epsilon floor on the natural minimum.
+    * ``chi_n`` spline ``gamma`` — the B-spline has the convex-hull
+      property: if every coefficient satisfies
+      ``log(chi_n_min) ≤ γᵢ ≤ log(chi_n_max)`` then the evaluated
+      ``chi_n = exp(B @ gamma)`` is guaranteed to stay within
+      ``[chi_n_min, chi_n_max]`` at every model age.
+
+    Returns:
+        lower: 1-D array of length n_beta + n_chi_b + n_gamma.
+        upper: 1-D array of the same length.
+    """
+    eps = config.bound_epsilon
+
+    def _range(param_name):
+        validators = p._data.get(param_name, {}).get("validators", {})
+        r = validators.get("range", {})
+        lo = float(r.get("min", -np.inf))
+        hi = float(r.get("max", np.inf))
+        return lo, hi
+
+    # --- beta_annual (logit space) ---
+    b_lo_nat, b_hi_nat = _range("beta_annual")
+    b_lo_nat = max(b_lo_nat, eps)
+    b_hi_nat = min(b_hi_nat, 1.0 - eps)
+    beta_lo = np.full(p.J, np.log(b_lo_nat / (1.0 - b_lo_nat)))
+    beta_hi = np.full(p.J, np.log(b_hi_nat / (1.0 - b_hi_nat)))
+
+    # --- chi_b (log space) ---
+    cb_lo_nat, cb_hi_nat = _range("chi_b")
+    cb_lo_nat = max(cb_lo_nat, eps)
+    chi_b_lo = np.full(p.J, np.log(cb_lo_nat))
+    chi_b_hi = np.full(p.J, np.log(cb_hi_nat))
+
+    # --- chi_n spline coefficients (log space, convex-hull argument) ---
+    cn_lo_nat, cn_hi_nat = _range("chi_n")
+    cn_lo_nat = max(cn_lo_nat, eps)
+    gamma_lo = np.full(config.chi_n_n_spline_knots, np.log(cn_lo_nat))
+    gamma_hi = np.full(config.chi_n_n_spline_knots, np.log(cn_hi_nat))
+
+    lower = np.concatenate([beta_lo, chi_b_lo, gamma_lo])
+    upper = np.concatenate([beta_hi, chi_b_hi, gamma_hi])
+    return lower, upper
+
+
+def _generate_starts(
+    theta0: np.ndarray,
+    n_starts: int,
+    radius: float,
+    lower: np.ndarray | None = None,
+    upper: np.ndarray | None = None,
+    seed: int | None = None,
+) -> list[np.ndarray]:
+    """Return n_starts starting vectors centred on theta0.
+
+    The first element is always theta0 itself (the warm start from prior
+    calibration values).  Additional points are drawn by adding additive
+    noise scaled by ``radius * max(|theta0_i|, 1.0)`` in each dimension,
+    so near-zero transformed parameters still receive meaningful perturbations.
+    All generated points are clipped to ``[lower, upper]`` when provided.
+    """
+    def _clip(x):
+        if lower is not None:
+            x = np.maximum(x, lower)
+        if upper is not None:
+            x = np.minimum(x, upper)
+        return x
+
+    starts = [_clip(theta0.copy())]
+    if n_starts <= 1:
+        return starts
+    rng = np.random.default_rng(seed)
+    scale = np.maximum(np.abs(theta0), 1.0)
+    for _ in range(n_starts - 1):
+        noise = rng.uniform(-radius, radius, size=theta0.shape) * scale
+        starts.append(_clip(theta0 + noise))
+    return starts
+
+
 def estimate_lifecycle_params(
     p,
     config: LifecycleCalibrationConfig | None = None,
@@ -1096,7 +1256,9 @@ def estimate_lifecycle_params(
     transform: bool = True,
     savings_rate: float | None = None,
 ) -> LifecycleCalibrationResult:
-    """Estimate beta_annual, chi_b, and chi_n jointly by SMM."""
+    """Estimate beta_annual, chi_b, and chi_n jointly by SMM using DFO-LS."""
+    import dfols
+
     if config is None:
         config = LifecycleCalibrationConfig()
     config.validate(p)
@@ -1116,55 +1278,67 @@ def estimate_lifecycle_params(
             bootstrap_moments=bootstrap_moments,
             ridge=config.weighting_ridge,
         )
-    ss_cache = SSSolutionCache(use_ss_solver=config.use_ss_solver_restart)
-    objective_args = (
-        data_moments,
-        W,
-        p,
-        config,
-        base_chi_n,
-        client,
-        transform,
-        ss_cache,
+
+    lower, upper = _extract_dfols_bounds(p, config)
+    starts = _generate_starts(
+        theta0,
+        config.n_starts,
+        config.start_radius,
+        lower=lower,
+        upper=upper,
     )
-    optimizer_state = {
-        "function_evals": 0,
-        "iteration": 0,
-        "last_objective": None,
-    }
+    best_result = None
+    best_obj = np.inf
+    best_start_index = 0
+    all_start_results = []
 
-    def tracked_objective(theta):
-        value = smm_objective(theta, *objective_args)
-        optimizer_state["function_evals"] += 1
-        optimizer_state["last_objective"] = value
-        return value
-
-    def optimizer_callback(_theta):
-        optimizer_state["iteration"] += 1
-        max_iterations = config.optimizer_options.get("maxiter", "?")
-        objective = optimizer_state["last_objective"]
-        if objective is None:
-            objective_text = "unavailable"
-        else:
-            objective_text = f"{objective:.6e}"
+    for i, theta_start in enumerate(starts):
         logger.info(
-            "Lifecycle SMM iteration %s/%s: objective=%s, function_evals=%s",
-            optimizer_state["iteration"],
-            max_iterations,
-            objective_text,
-            optimizer_state["function_evals"],
+            "Lifecycle SMM: DFO-LS start %d/%d", i + 1, config.n_starts
         )
+        ss_cache = SSSolutionCache(use_ss_solver=config.use_ss_solver_restart)
+
+        def residual_fn(theta, _cache=ss_cache):
+            return smm_residual(
+                theta,
+                data_moments,
+                W,
+                p,
+                config,
+                base_chi_n=base_chi_n,
+                client=client,
+                transform=transform,
+                ss_cache=_cache,
+            )
+
+        maxfun = config.dfols_maxfun
+        dfols_result = dfols.solve(
+            residual_fn,
+            theta_start,
+            bounds=(lower, upper),
+            rhoend=config.dfols_rhoend,
+            maxfun=maxfun,
+            do_logging=config.log_optimizer_progress,
+            print_progress=False,
+        )
+        all_start_results.append(dfols_result)
+        obj = float(dfols_result.cost)
+        logger.info(
+            "Lifecycle SMM: start %d/%d finished — objective=%.6e, "
+            "evals=%d, msg=%s",
+            i + 1,
+            config.n_starts,
+            obj,
+            dfols_result.nf,
+            dfols_result.msg,
+        )
+        if obj < best_obj:
+            best_obj = obj
+            best_result = dfols_result
+            best_start_index = i
 
-    est_output = opt.minimize(
-        tracked_objective,
-        theta0,
-        method=config.optimizer_method,
-        tol=config.optimizer_tol,
-        options=config.optimizer_options,
-        callback=optimizer_callback if config.log_optimizer_progress else None,
-    )
     params = unpack_lifecycle_params(
-        est_output.x,
+        best_result.x,
         p,
         config,
         base_chi_n=base_chi_n,
@@ -1176,6 +1350,7 @@ def optimizer_callback(_theta):
         params["chi_b"],
         params["chi_n"],
     )
+    ss_cache = SSSolutionCache(use_ss_solver=config.use_ss_solver_restart)
     ss_output = solve_ss_with_cache(p, client=client, ss_cache=ss_cache)
     model_moments = compute_model_moments(ss_output, p, config)
 
@@ -1183,11 +1358,13 @@ def optimizer_callback(_theta):
         beta_annual=params["beta_annual"],
         chi_b=params["chi_b"],
         chi_n=params["chi_n"],
-        objective_value=float(est_output.fun),
-        optimizer_result=est_output,
+        objective_value=best_obj,
+        optimizer_result=best_result,
         data_moments=data_moments,
         model_moments=model_moments,
         weighting_matrix=W,
+        best_start_index=best_start_index,
+        all_start_results=all_start_results,
     )
 
 
diff --git a/pyproject.toml b/pyproject.toml
index 8e683fe2..34ff51ee 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,6 +45,7 @@ dependencies = [
     "statsmodels",
     "linearmodels",
     "ogcore>=0.14.6",
+    "dfo-ls",
 ]
 
 [project.urls]

From c6bb2186756ba691a4431136dcf8be3964c07a97 Mon Sep 17 00:00:00 2001
From: Jason DeBacker <jdebacker@users.noreply.github.com>
Date: Sun, 19 Jul 2026 18:36:45 -0400
Subject: [PATCH 8/9] update est

---
 ogusa/estimate_lifecycle_params.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ogusa/estimate_lifecycle_params.py b/ogusa/estimate_lifecycle_params.py
index 76305a20..5304a798 100644
--- a/ogusa/estimate_lifecycle_params.py
+++ b/ogusa/estimate_lifecycle_params.py
@@ -1207,7 +1207,7 @@ def _range(param_name):
     gamma_hi = np.full(config.chi_n_n_spline_knots, np.log(cn_hi_nat))
 
     lower = np.concatenate([beta_lo, chi_b_lo, gamma_lo])
-    upper = np.concatenate([beta_hi, chi_b_hi, gamma_hi])
+    upper = np.concatenate([beta_hi - 0.00015, chi_b_hi / 2, gamma_hi / 2])  # TODO: check if want to divide in half, but want to keep away from parameter range extremes as it keeps hitting in the optimization routine even with these bounds
     return lower, upper
 
 
@@ -1322,7 +1322,7 @@ def residual_fn(theta, _cache=ss_cache):
             print_progress=False,
         )
         all_start_results.append(dfols_result)
-        obj = float(dfols_result.cost)
+        obj = float(dfols_result.obj)
         logger.info(
             "Lifecycle SMM: start %d/%d finished — objective=%.6e, "
             "evals=%d, msg=%s",

From 90149ce6a8d5367cf98a45eb7cd6167fef90f2b4 Mon Sep 17 00:00:00 2001
From: Jason DeBacker <jdebacker@users.noreply.github.com>
Date: Wed, 22 Jul 2026 18:13:44 -0400
Subject: [PATCH 9/9] fix messages

---
 ogusa/estimate_lifecycle_params.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ogusa/estimate_lifecycle_params.py b/ogusa/estimate_lifecycle_params.py
index 5304a798..67979863 100644
--- a/ogusa/estimate_lifecycle_params.py
+++ b/ogusa/estimate_lifecycle_params.py
@@ -1206,8 +1206,8 @@ def _range(param_name):
     gamma_lo = np.full(config.chi_n_n_spline_knots, np.log(cn_lo_nat))
     gamma_hi = np.full(config.chi_n_n_spline_knots, np.log(cn_hi_nat))
 
-    lower = np.concatenate([beta_lo, chi_b_lo, gamma_lo])
-    upper = np.concatenate([beta_hi - 0.00015, chi_b_hi / 2, gamma_hi / 2])  # TODO: check if want to divide in half, but want to keep away from parameter range extremes as it keeps hitting in the optimization routine even with these bounds
+    lower = np.concatenate([0.8, 0.1, gamma_lo])
+    upper = np.concatenate([0.999, 200, gamma_hi])  # TODO: check if want to divide in half, but want to keep away from parameter range extremes as it keeps hitting in the optimization routine even with these bounds
     return lower, upper