Structural Modeling

Structural Modeling

Reference for implementing structural econometric models: from economic model to moment conditions to estimated parameters. Covers the full workflow of taking a theoretical model, deriving its empirical content, and recovering structural parameters from data.

When to Use This Skill

Use when the user is:

• Specifying a structural model and deriving moment conditions
• Implementing NFXP or MPEC estimation routines
• Working with BLP-style demand systems (random coefficients logit)
• Building dynamic discrete choice models (Rust, Hotz-Miller CCP)
• Estimating auction models (first-price, ascending, common value)
• Debugging convergence failures in structural estimation
• Choosing between estimation approaches for a given model

Skip when:

• The task is reduced-form causal inference (use causal-inference skill)
• The task is pure simulation design (use numerical-auditor agent)
• The user just needs standard regression (statsmodels/linearmodels suffice)

Quick Reference: Structural Methods

The Structural Estimation Workflow

Every structural estimation follows the same logical arc:

Economic Model → Equilibrium/Decision Rule → Observable Implications
    → Moment Conditions → Estimator → Optimization → Inference

Step 1: Model Specification

Define primitives clearly before writing any code:

# model_spec.py — Document structural primitives
"""
Model: Single-agent optimal stopping (Rust 1987 bus engine replacement)

State:    x_t ∈ {0, 1, ..., X_max}  (mileage bin)
Action:   a_t ∈ {0, 1}  (0 = maintain, 1 = replace)
Flow payoff:
    u(x, 0; θ) = -θ_1 * x - θ_2 * x²     (maintenance cost)
    u(x, 1; θ) = -RC                        (replacement cost)
Discount:  β = 0.9999 (fixed)
Shocks:    ε ~ Type 1 Extreme Value (logit errors)
"""

Document these before writing estimation code: agents, information, timing, payoff functional form, equilibrium concept.

Step 2: Derive Moment Conditions

Key question: Just-identified → method of moments; over-identified → GMM with optimal weighting matrix; under-identified → revisit assumptions.

NFXP vs MPEC

Two dominant paradigms for models with latent quantities (unobserved heterogeneity, future expectations, equilibrium objects):

NFXP (Nested Fixed-Point): Solve the model in an inner loop for each parameter guess, evaluate likelihood/moments in an outer loop. Conceptually simple; inner loop must fully converge at every iteration — requires tight tolerance (1e-12, not 1e-6; see Su & Judd 2012).

MPEC (Mathematical Programming with Equilibrium Constraints): Reformulate as a single constrained optimization. No inner loop — solver handles everything; can be faster for large state spaces; requires IPOPT or KNITRO.

For full NFXP and MPEC code (Rust 1987 bus engine model), see references/estimation-methods.md.

BLP Demand Estimation

BLP (Berry, Levinsohn, Pakes 1995) is the workhorse for differentiated products demand. Use PyBLP whenever possible — it handles the difficult numerical details correctly.

import pyblp

# Define the problem
problem = pyblp.Problem(
    product_formulations=(
        pyblp.Formulation('1 + prices + x1 + x2'),       # linear (β)
        pyblp.Formulation('1 + prices + x1'),              # random coefficients (Σ)
    ),
    product_data=product_data,
    agent_data=agent_data
)

# Solve — always use multiple starting values; BLP objective is non-convex
results = problem.solve(
    sigma=sigma_init,
    optimization=pyblp.Optimization('l-bfgs-b', {'gtol': 1e-8}),
    iteration=pyblp.Iteration('squarem', {'atol': 1e-14}),
    method='2s'
)

For the full multi-start loop, two-step GMM, elasticity checks, instrument selection, and marginal cost computation, see references/estimation-methods.md.

BLP Diagnostics Checklist:

• [ ] First-stage F > 10 for price instruments
• [ ] Run 10+ random starts (objective is non-convex)
• [ ] Own-price elasticities all negative: results.compute_elasticities('prices')
• [ ] All markets converged: results.fp_converged.all()
• [ ] Marginal costs positive: results.compute_costs()
• [ ] Inner loop atol <= 1e-14 (tighter is safer)

Dynamic Discrete Choice

Rust (1987) NFXP: Full solution — solve the Bellman equation by value function iteration at every outer iteration. Use for models where counterfactuals require the full model.

Hotz-Miller CCP: Two-step semiparametric approach. Step 1: estimate conditional choice probabilities nonparametrically. Step 2: form pseudo-value functions for a linear regression. Faster; less efficient; sufficient when counterfactuals are not needed.

For full CCP implementation code, see references/estimation-methods.md.

Auction Models

First-price sealed-bid (GPV): Guerre, Perrigne, Vuong (2000) — invert the bidding equilibrium condition v(b) = b + G(b)/((n-1)g(b)) to recover latent values nonparametrically from observed bids.

Ascending (English): In IPV setting, transaction price = second-highest value. Use MLE on order statistics to recover the value distribution.

Common value: Requires accounting for winner's curse. Li-Perrigne-Vuong (2002) approach; typically requires parametric assumptions.

For full GPV estimator code, ascending auction MLE, and validation diagnostics, see references/estimation-methods.md.

Method Selection

When to use structural vs. reduced-form:

• Structural: Need to evaluate counterfactual policies, recover preference parameters, or model strategic interactions
• Reduced-form: Need a credible causal estimate of a specific treatment effect with minimal assumptions

Within structural: 1. Static discrete choice with heterogeneity? → BLP / mixed logit 2. Dynamic single-agent optimal stopping? → NFXP (small state) or MPEC (large state) 3. Counterfactuals not needed, data rich? → CCP estimator 4. Auction data? → GPV (first-price) or order statistics MLE (ascending) 5. Market-level entry/exit? → Bresnahan-Reiss or Ciliberto-Tamer (see game-theory skill)

Common Anti-Patterns

JAX Acceleration

For GPU-accelerated structural estimation (JIT compilation, autodiff, vmap for simulated moments, differentiable fixed-point iteration with lax.while_loop), see references/jax-guide.md.

Integration with compound-science

• numerical-auditor — Systematic convergence review: gradient norms, conditioning, tolerance sensitivity
• numerical-auditor — DGP formalization, Monte Carlo studies, convergence review
• identification-critic — Verify equilibrium existence, uniqueness, stability, comparative statics
• econometric-reviewer — Reviews moment-matching strategy, parameter identification, sensitivity to targets
• /estimate — Full estimation pipeline with quality gates

Additional References

• references/estimation-methods.md — Full code: BLP multi-start, NFXP Bellman solver, MPEC cyipopt formulation, Hotz-Miller CCP, GPV auction estimator
• references/diagnostics-and-se.md — Convergence failure diagnosis, numerical safeguards (logsumexp, conditioning), GMM sandwich SEs, parametric bootstrap
• references/jax-guide.md — JAX JIT/autodiff for structural objectives, vmap for simulation, lax.while_loop for differentiable contraction mappings

Convergence Diagnostics and Standard Errors for Structural Models

Convergence Diagnostics

Convergence failures are the most common problem in structural estimation.

Starting Values

# Strategy 1: Grid search over coarse parameter space
from itertools import product

param_grid = {
    'RC': [2.0, 5.0, 10.0, 20.0],
    'theta1': [0.001, 0.01, 0.05, 0.1]
}

best_obj = np.inf
best_start = None
for RC, theta1 in product(param_grid['RC'], param_grid['theta1']):
    try:
        obj = nfxp_objective([RC, theta1], data, beta, trans_mat, n_states)
        if obj < best_obj:
            best_obj = obj
            best_start = [RC, theta1]
    except (ValueError, np.linalg.LinAlgError):
        continue

# Strategy 2: Estimate simplified model first
# e.g., static version, or version without random coefficients

# Strategy 3: Use estimates from related data/specification

Diagnosing Convergence Failures

Numerical Safeguards

# Always work in log space for likelihoods
def safe_log_likelihood(log_prob):
    """Numerically stable log-likelihood computation."""
    return np.sum(log_prob)  # already in log space

# Use logsumexp for softmax/logit choice probabilities
from scipy.special import logsumexp

def logit_choice_probs(utilities):
    """Numerically stable logit probabilities."""
    # utilities: (n_states, n_actions)
    log_denom = logsumexp(utilities, axis=1, keepdims=True)
    log_probs = utilities - log_denom
    return np.exp(log_probs)

# Check condition number of key matrices
def check_conditioning(matrix, name="matrix"):
    cond = np.linalg.cond(matrix)
    if cond > 1e10:
        print(f"WARNING: {name} condition number = {cond:.2e} — near singular")
    return cond

Standard Errors for Structural Models

GMM Standard Errors

def gmm_standard_errors(theta_hat, moment_fn, data, W, epsilon=1e-5):
    """
    Sandwich standard errors for GMM.

    V(θ) = (G'WG)^{-1} G'W S W G (G'WG)^{-1} / N

    G = Jacobian of moment conditions (∂m/∂θ)
    S = Variance of moment conditions
    W = Weighting matrix
    """
    n_params = len(theta_hat)
    moments = moment_fn(theta_hat, data)  # (N, n_moments)
    N = moments.shape[0]

    # Numerical Jacobian
    G = np.zeros((moments.shape[1], n_params))
    for j in range(n_params):
        theta_plus = theta_hat.copy()
        theta_minus = theta_hat.copy()
        theta_plus[j] += epsilon
        theta_minus[j] -= epsilon
        G[:, j] = (moment_fn(theta_plus, data).mean(axis=0)
                    - moment_fn(theta_minus, data).mean(axis=0)) / (2 * epsilon)

    # Long-run variance of moments
    S = moments.T @ moments / N

    # Sandwich formula
    GWG_inv = np.linalg.inv(G.T @ W @ G)
    V = GWG_inv @ (G.T @ W @ S @ W @ G) @ GWG_inv / N

    se = np.sqrt(np.diag(V))
    return se

Bootstrap for Complex Models

When analytic standard errors are difficult (e.g., multi-step estimators, simulation-based estimators):

def parametric_bootstrap(estimate_fn, data, n_bootstrap=200, seed=42):
    """
    Parametric bootstrap: resample from estimated model.
    For structural models, often better than nonparametric bootstrap
    because it preserves the data structure (markets, panels).
    """
    rng = np.random.default_rng(seed)
    theta_hat = estimate_fn(data)

    boot_estimates = []
    for b in range(n_bootstrap):
        # Resample: cluster at appropriate level (market, individual, etc.)
        idx = rng.choice(len(data), size=len(data), replace=True)
        data_b = data.iloc[idx].reset_index(drop=True)

        try:
            theta_b = estimate_fn(data_b)
            boot_estimates.append(theta_b)
        except Exception:
            continue  # skip failed replications but log the count

    boot_estimates = np.array(boot_estimates)
    se = boot_estimates.std(axis=0)

    # Report: how many bootstrap replications converged
    convergence_rate = len(boot_estimates) / n_bootstrap
    if convergence_rate < 0.8:
        print(f"WARNING: Only {convergence_rate:.0%} of bootstrap samples converged")

    return se, boot_estimates

Structural Modeling: Full Estimation Method Code

Full implementation code for BLP demand estimation, NFXP/MPEC dynamic discrete choice, and auction models. Referenced from SKILL.md.

---

BLP Demand Estimation (Full)

Problem Setup

import pyblp

# Define the problem
problem = pyblp.Problem(
    product_formulations=(
        pyblp.Formulation('1 + prices + x1 + x2'),           # linear (β)
        pyblp.Formulation('1 + prices + x1'),                  # random coefficients (Σ)
        pyblp.Formulation('0 + demand_instruments0 + demand_instruments1')  # supply
    ),
    product_data=product_data,  # DataFrame with market_ids, shares, prices, etc.
    agent_formulation=pyblp.Formulation('0 + income'),         # demographics
    agent_data=agent_data
)

Estimation with Multiple Starting Values

import numpy as np

# Starting values matter — use multiple starting points
results_best = None
for _ in range(10):
    sigma_init = np.random.uniform(0.1, 2.0, size=(3, 3))
    sigma_init = np.tril(sigma_init)  # lower triangular for Cholesky

    results = problem.solve(
        sigma=sigma_init,
        optimization=pyblp.Optimization('l-bfgs-b', {'gtol': 1e-8}),
        iteration=pyblp.Iteration('squarem', {'atol': 1e-14}),  # accelerated contraction
        method='1s'  # start with 1-step GMM, then switch to 2-step
    )

    if results_best is None or results.objective < results_best.objective:
        results_best = results

# Two-step GMM with optimal weighting matrix
results_2s = problem.solve(
    sigma=results_best.sigma,
    optimization=pyblp.Optimization('l-bfgs-b', {'gtol': 1e-8}),
    iteration=pyblp.Iteration('squarem', {'atol': 1e-14}),
    method='2s',
    W=results_best.updated_W
)

BLP Post-Estimation Checks

# Own-price elasticities: must be negative
elasticities = results_2s.compute_elasticities('prices')
print("Own-price elasticities (diagonal):", np.diag(elasticities).describe())
assert (np.diag(elasticities) < 0).all(), "Some own-price elasticities are positive!"

# Cross-price elasticities: should be positive for substitutes
diversion = results_2s.compute_diversion_ratios()

# Marginal costs: must be positive
costs = results_2s.compute_costs()
assert (costs > 0).all(), "Negative marginal costs — check instruments or supply-side spec"

# Check inner loop convergence
assert results_2s.fp_converged.all(), "Not all markets converged in contraction mapping"
print(f"Contraction evaluations: {results_2s.contraction_evaluations.sum()}")

# Optimal instruments (improves efficiency)
updated_instruments = results_2s.compute_optimal_instruments(method='approximate')

BLP Instruments

Standard BLP instruments and when to use them:

---

NFXP Implementation (Full)

Inner Loop: Bellman Contraction

import numpy as np
from scipy.optimize import minimize

def solve_inner(theta, beta, trans_mat, n_states):
    """Solve Bellman equation by value function iteration (contraction mapping)."""
    RC, theta1 = theta
    flow_maintain = -theta1 * np.arange(n_states)
    EV = np.zeros(n_states)

    for _ in range(2000):  # generous iteration limit
        # Choice-specific value functions (logit shocks)
        cv_maintain = flow_maintain + beta * trans_mat @ EV
        cv_replace = -RC + beta * trans_mat[0, :] @ EV

        # Log-sum formula for expected value with Type 1 EV errors
        EV_new = np.log(np.exp(cv_maintain) + np.exp(cv_replace))

        if np.max(np.abs(EV_new - EV)) < 1e-12:  # tight tolerance (Su & Judd 2012)
            break
        EV = EV_new

    return EV

def nfxp_objective(theta, data, beta, trans_mat, n_states):
    """Negative log-likelihood for NFXP."""
    EV = solve_inner(theta, beta, trans_mat, n_states)
    RC, theta1 = theta

    flow_maintain = -theta1 * np.arange(n_states)
    cv_maintain = flow_maintain + beta * trans_mat @ EV
    cv_replace = -RC + beta * trans_mat[0, :] @ EV

    # Choice probabilities (logit)
    prob_replace = 1 / (1 + np.exp(cv_maintain - cv_replace))

    # Log-likelihood
    ll = np.sum(
        data['replace'] * np.log(prob_replace[data['state']] + 1e-15)
        + (1 - data['replace']) * np.log(1 - prob_replace[data['state']] + 1e-15)
    )
    return -ll

result = minimize(nfxp_objective, x0=[5.0, 0.01],
                  args=(data, beta, trans_mat, n_states),
                  method='Nelder-Mead',
                  options={'xatol': 1e-8, 'fatol': 1e-10})

---

MPEC Implementation (Full)

MPEC Formulation with cyipopt

import cyipopt
import numpy as np

class RustMPEC:
    """MPEC formulation of Rust (1987) bus engine model."""

    def __init__(self, data, beta, trans_mat, n_states):
        self.data = data
        self.beta = beta
        self.trans = trans_mat
        self.n_states = n_states
        # Decision variables: [RC, theta1, EV_0, ..., EV_{n-1}]
        self.n_vars = 2 + n_states

    def objective(self, x):
        """Negative log-likelihood."""
        RC, theta1 = x[0], x[1]
        EV = x[2:]

        flow_maintain = -theta1 * np.arange(self.n_states)
        cv_m = flow_maintain + self.beta * self.trans @ EV
        cv_r = -RC + self.beta * self.trans[0, :] @ EV

        prob_r = 1 / (1 + np.exp(cv_m - cv_r))
        ll = np.sum(
            self.data['replace'] * np.log(prob_r[self.data['state']] + 1e-15)
            + (1 - self.data['replace']) * np.log(1 - prob_r[self.data['state']] + 1e-15)
        )
        return -ll

    def gradient(self, x):
        """Gradient of objective (use autodiff in practice)."""
        # In practice, compute via JAX: jax.grad(self.objective)(x)
        raise NotImplementedError("Use JAX autodiff for gradient")

    def constraints(self, x):
        """Bellman equation constraints: EV = log-sum-exp(CV)."""
        RC, theta1 = x[0], x[1]
        EV = x[2:]

        flow_maintain = -theta1 * np.arange(self.n_states)
        cv_m = flow_maintain + self.beta * self.trans @ EV
        cv_r = -RC + self.beta * self.trans[0, :] @ EV

        EV_implied = np.log(np.exp(cv_m) + np.exp(cv_r))
        return EV - EV_implied  # should equal zero at solution

    def jacobianstructure(self):
        """Sparsity structure of the constraint Jacobian."""
        # Constraints: n_states equations
        # Variables: 2 structural params + n_states EV values
        rows = np.repeat(np.arange(self.n_states), self.n_states + 2)
        cols = np.tile(np.arange(self.n_vars), self.n_states)
        return rows, cols


# Run MPEC via cyipopt
mpec_problem = RustMPEC(data, beta, trans_mat, n_states)
x0 = np.concatenate([[5.0, 0.01], np.zeros(n_states)])  # initial point
bounds_lower = np.concatenate([[0.0, 0.0], -np.inf * np.ones(n_states)])
bounds_upper = np.concatenate([[np.inf, np.inf], np.inf * np.ones(n_states)])
constraint_lower = np.zeros(n_states)  # equality constraints
constraint_upper = np.zeros(n_states)

nlp = cyipopt.Problem(
    n=mpec_problem.n_vars,
    m=n_states,
    problem_obj=mpec_problem,
    lb=bounds_lower, ub=bounds_upper,
    cl=constraint_lower, cu=constraint_upper
)
nlp.add_option('tol', 1e-10)
nlp.add_option('max_iter', 1000)
x_opt, info = nlp.solve(x0)

---

Hotz-Miller CCP Estimation (Full)

import numpy as np

def hotz_miller_ccp(data, n_states, n_actions, beta, trans_mat):
    """
    Hotz-Miller (1993) CCP estimator.
    Step 1: Estimate CCPs nonparametrically.
    Step 2: Use CCPs to form pseudo-value functions, then run simple regression.
    """
    # Step 1: Estimate CCPs from frequency of actions in each state
    ccps = np.zeros((n_states, n_actions))
    for s in range(n_states):
        mask = data['state'] == s
        if mask.sum() > 0:
            for a in range(n_actions):
                ccps[s, a] = (data['action'][mask] == a).mean()

    # Smooth to avoid log(0) — add small probability mass
    ccps = np.clip(ccps, 0.001, 0.999)
    ccps = ccps / ccps.sum(axis=1, keepdims=True)

    # Step 2: Construct pseudo-value functions
    # With logit errors: E[ε | a chosen] = euler_constant - log(P(a))
    euler = 0.5772156649

    # Forward simulation of CCPs to get expected future utilities
    e_eps = euler - np.log(ccps[:, 0])  # expected shock conditional on maintain

    # Mapping matrix: expected transitions under estimated policy
    F = np.diag(ccps[:, 0]) @ trans_mat + np.diag(ccps[:, 1]) @ trans_mat[[0], :]

    # Pseudo-value: (I - beta * F)^{-1} * (flow_payoff + correction)
    # This gives a linear-in-parameters system for the structural parameters

    return ccps, F

---

Auction Models (Full)

First-Price Sealed-Bid: GPV Estimator

from scipy.interpolate import UnivariateSpline
from scipy.stats import gaussian_kde
import numpy as np

def gpv_estimate(bids, n_bidders):
    """
    GPV (2000) nonparametric estimation for symmetric IPV first-price auctions.

    Key insight: In equilibrium, bidder with value v bids:
        b(v) = v - G(b)/(n-1)*g(b)
    where G is the bid distribution and g its density.

    Inversion: v(b) = b + G(b)/((n-1)*g(b))
    """
    n = n_bidders

    # Step 1: Estimate bid distribution and density
    kde = gaussian_kde(bids, bw_method='silverman')

    # Evaluate on a grid
    b_grid = np.linspace(bids.min(), bids.max(), 200)
    g_hat = kde(b_grid)                                    # density
    G_hat = np.array([kde.integrate_box_1d(-np.inf, b) for b in b_grid])  # CDF

    # Step 2: Invert to recover pseudo-values
    v_hat = b_grid + G_hat / ((n - 1) * g_hat)

    # Step 3: Estimate value distribution from pseudo-values
    # (can use kernel density on v_hat, or fit parametric family)

    return b_grid, v_hat, g_hat, G_hat


def validate_gpv(bids, v_hat):
    """Diagnostics for GPV estimator."""
    # Pseudo-values must exceed bids (bidders shade down in equilibrium)
    assert (v_hat >= bids[:len(v_hat)]).all(), "Some pseudo-values below bids"

    # Check for negative values in cost auctions
    if v_hat.min() < 0:
        print("Warning: negative pseudo-values — check support assumption")

    # Boundary bias check
    print(f"Fraction of grid at boundaries: {(v_hat == v_hat[0]).mean():.3f}")

Ascending (English) Auction Estimation

from scipy.stats import rv_continuous
from scipy.optimize import minimize
import numpy as np

# Transaction prices = second-order statistics of value distribution
# Use order statistics theory to recover the parent distribution

# With N bidders, transaction price ~ F_{(N-1:N)} distribution
# f_{(k:n)}(x) = n!/(k-1)!(n-k)! * F(x)^{k-1} * (1-F(x))^{n-k} * f(x)

def order_stat_density(x, params, n_bidders, dist='lognormal'):
    """Density of the (n-1)-th order statistic for lognormal values."""
    from scipy.stats import lognorm
    k = n_bidders - 1  # second-highest
    n = n_bidders
    mu, sigma = params

    # Lognormal: F and f
    F = lognorm.cdf(x, s=sigma, scale=np.exp(mu))
    f = lognorm.pdf(x, s=sigma, scale=np.exp(mu))

    # Order statistic density
    coeff = np.math.factorial(n) / (np.math.factorial(k-1) * np.math.factorial(n-k))
    return coeff * F**(k-1) * (1-F)**(n-k) * f

def ascending_mle(prices, n_bidders):
    """MLE for value distribution from ascending auction prices."""
    def neg_ll(params):
        densities = np.array([order_stat_density(p, params, n_bidders) for p in prices])
        densities = np.clip(densities, 1e-15, None)
        return -np.sum(np.log(densities))

    result = minimize(neg_ll, x0=[np.log(prices.mean()), 0.5],
                      method='Nelder-Mead', options={'xatol': 1e-8})
    return result

Auction Model Selection

JAX for Structural Estimation

JAX provides automatic differentiation and JIT compilation — valuable for structural models where analytic gradients are tedious.

import jax
import jax.numpy as jnp
from jax import grad, jit

@jit
def bellman_operator(EV, theta, beta, trans_mat):
    """JAX-compatible Bellman operator with autodiff support."""
    RC, theta1 = theta[0], theta[1]
    n_states = EV.shape[0]

    flow_maintain = -theta1 * jnp.arange(n_states, dtype=float)
    cv_maintain = flow_maintain + beta * trans_mat @ EV
    cv_replace = -RC + beta * trans_mat[0, :] @ EV

    # logsumexp for numerical stability
    EV_new = jnp.logaddexp(cv_maintain, cv_replace)
    return EV_new

# Automatic gradient of the objective w.r.t. parameters
# — no hand-derived gradients needed
grad_objective = jit(grad(nfxp_objective_jax, argnums=0))

When to use JAX:

• Models with many parameters (gradient computation is expensive)
• Need second derivatives (Hessian) for standard errors or Newton steps
• Inner loop can be expressed as a differentiable fixed-point iteration
• Want GPU acceleration for large state spaces

When NOT to use JAX:

• Simple models where scipy.optimize works fine
• Models with non-differentiable components (discrete jumps, if-else logic)
• PyBLP already handles the specific model class

vmap for Simulated Moments

JAX's vmap vectorizes computation over simulation draws — useful for MSM and simulated MLE:

from jax import vmap

# Simulate moments for each draw in parallel (no Python loop)
simulate_one = lambda draw: compute_moments(theta, draw)
all_moments = vmap(simulate_one)(simulation_draws)  # (R, n_moments)
simulated_moments = all_moments.mean(axis=0)

Fixed-Point Iteration with JAX

For differentiating through contraction mappings (e.g., BLP inversion, Bellman iteration):

from jax.lax import while_loop

def contraction_jax(ev0, theta, beta, trans_mat, tol=1e-12):
    """Differentiable fixed-point via lax.while_loop."""
    def cond_fn(state):
        ev, ev_prev, _ = state
        return jnp.max(jnp.abs(ev - ev_prev)) > tol

    def body_fn(state):
        ev, _, i = state
        ev_new = bellman_operator(ev, theta, beta, trans_mat)
        return ev_new, ev, i + 1

    ev_init = jnp.zeros(trans_mat.shape[0])
    ev_final, _, n_iter = while_loop(cond_fn, body_fn, (ev_init, ev_init + 1.0, 0))
    return ev_final