diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 9d029ce6..a974ec92 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -13,17 +13,53 @@ on: - '*' jobs: run-tests: - name: Run tests for ${{ matrix.os }} on ${{ matrix.python-version }} + name: Run tests for ${{ matrix.os }} (${{ matrix.shard }}) runs-on: ${{ matrix.os }} strategy: fail-fast: false + # The three estimation tests in tests/test_chs_estimate.py take about four + # minutes each; the rest of the suite takes about the same in total. They + # run on their own Linux runners; the other legs deselect them. + # yamllint disable rule:line-length matrix: - os: - - ubuntu-latest - - macos-latest - - windows-latest - python-version: - - '3.14' + include: + - os: ubuntu-latest + shard: rest + pytest_args: >- + tests --cov-report=xml --cov=./ + --deselect tests/test_chs_estimate.py::test_estimate_chs_returns_conforming_result + --deselect tests/test_chs_estimate.py::test_estimate_chs_provides_ml_inference + --deselect tests/test_chs_estimate.py::test_estimate_chs_enforces_user_fixed_constraint_value + - os: ubuntu-latest + shard: long-1 + pytest_args: >- + tests/test_chs_estimate.py::test_estimate_chs_returns_conforming_result + --cov-report=xml --cov=./ + - os: ubuntu-latest + shard: long-2 + pytest_args: >- + tests/test_chs_estimate.py::test_estimate_chs_provides_ml_inference + --cov-report=xml --cov=./ + - os: ubuntu-latest + shard: long-3 + pytest_args: >- + tests/test_chs_estimate.py::test_estimate_chs_enforces_user_fixed_constraint_value + --cov-report=xml --cov=./ + - os: macos-latest + shard: rest + pytest_args: >- + tests -m 'not long_running' + --deselect tests/test_chs_estimate.py::test_estimate_chs_returns_conforming_result + --deselect tests/test_chs_estimate.py::test_estimate_chs_provides_ml_inference + --deselect tests/test_chs_estimate.py::test_estimate_chs_enforces_user_fixed_constraint_value + - os: windows-latest + shard: rest + pytest_args: >- + tests -m 'not long_running' + --deselect tests/test_chs_estimate.py::test_estimate_chs_returns_conforming_result + --deselect tests/test_chs_estimate.py::test_estimate_chs_provides_ml_inference + --deselect tests/test_chs_estimate.py::test_estimate_chs_enforces_user_fixed_constraint_value + # yamllint enable rule:line-length steps: - uses: actions/checkout@v7 - uses: prefix-dev/setup-pixi@v0.10.0 @@ -35,15 +71,10 @@ jobs: activate-environment: true frozen: true - name: Run pytest - if: runner.os != 'Linux' || matrix.python-version != '3.14' shell: bash -l {0} - run: pixi run -e tests-cpu tests - - name: Run pytest with coverage - if: runner.os == 'Linux' && matrix.python-version == '3.14' - shell: bash -l {0} - run: pixi run -e tests-cpu tests-with-cov + run: pixi run -e tests-cpu pytest ${{ matrix.pytest_args }} - name: Upload coverage report - if: runner.os == 'Linux' && matrix.python-version == '3.14' + if: runner.os == 'Linux' uses: codecov/codecov-action@v7 env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} diff --git a/docs/reference_guides/transition_functions.md b/docs/reference_guides/transition_functions.md index d2f3d8ee..e5434461 100644 --- a/docs/reference_guides/transition_functions.md +++ b/docs/reference_guides/transition_functions.md @@ -70,6 +70,48 @@ See [When to use the AF variants](#when-to-use-the-af-variants) below. Same as `translog` but clips state values at ±10^12 before computation. Use this when states might grow very large and cause numerical overflow. +### asymmetric_translog + +Like `translog`, but each factor's own square is split at a knot $\kappa_j$ into a part +above and a part below it, each with its own coefficient: + +$$ +f_{t+1} = \sum_j \beta_j s_j + \sum_j \gamma^+_j \max(s_j - \kappa_j, 0)^2 ++ \sum_j \gamma^-_j \min(s_j - \kappa_j, 0)^2 + \sum_{a knot")` like any other parameter, e.g. via +`fixed_params`. A knot is not identified when its two curvature coefficients are equal +(zero included): the knot is then absorbed by the linear term and the constant. + +There is no clipping. A convex tail can still produce an unstable fixed point far from +the data, so check the fixed points of the estimated equations before using them as a +long-horizon law of motion. + +**Parameters** (in this order): linear terms (`fac`), squares above the knot +(`fac ** 2 (above)`), squares below the knot (`fac ** 2 (below)`), knots (`fac knot`), +interaction terms (`fac_a * fac_b`), and `constant`. As for `translog`, they are +enumerated over all factors, latent and observed. In carry-forward (identity) periods +the own linear coefficient is fixed to 1 and every other parameter, the knots included, +to 0. + ### linear_and_squares Like `translog` but without interaction terms: @@ -138,10 +180,10 @@ $$ ## When to use the AF variants The general `translog` and `log_ces` (and the other built-in production functions: -`linear`, `robust_translog`, `linear_and_squares`, `log_ces_with_constant`, -`log_ces_general`) enumerate parameters over **all** factors, including observed -factors. For the CHS estimator this is the intended behaviour: observed factors are -allowed to enter the production function with free coefficients. +`linear`, `robust_translog`, `asymmetric_translog`, `linear_and_squares`, +`log_ces_with_constant`, `log_ces_general`) enumerate parameters over **all** factors, +including observed factors. For the CHS estimator this is the intended behaviour: +observed factors are allowed to enter the production function with free coefficients. For an AF production function this is usually wrong. The AF model assumes that observed factors (e.g. income) affect skills **only** through the investment equation, not @@ -152,13 +194,13 @@ production factors (skill + investment) and match AF equations (6) and (7) respe To make the leakage visible, `validate_af_model` emits a loud `UserWarning` when a built-in production transition (`linear`, `linear_and_squares`, `translog`, -`robust_translog`, `log_ces`, `log_ces_with_constant`, `log_ces_general`) is used on a -non-endogenous production factor while observed factors are present. The warning is not -an error — intentionally-leaky models still run — but it flags the wrong-estimand risk. -Switch to `translog_af` / `log_ces_af`, or pin every observed-factor transition -coefficient to `0.0` via `fixed_params`, to remove the leakage. (The endogenous -investment equation legitimately uses observed factors, so endogenous factors do not -trigger the warning.) +`robust_translog`, `asymmetric_translog`, `log_ces`, `log_ces_with_constant`, +`log_ces_general`) is used on a non-endogenous production factor while observed factors +are present. The warning is not an error — intentionally-leaky models still run — but it +flags the wrong-estimand risk. Switch to `translog_af` / `log_ces_af`, or pin every +observed-factor transition coefficient to `0.0` via `fixed_params`, to remove the +leakage. (The endogenous investment equation legitimately uses observed factors, so +endogenous factors do not trigger the warning.) ## Custom Transition Functions diff --git a/docs/superpowers/specs/2026-10-16-asymmetric-translog-design.md b/docs/superpowers/specs/2026-10-16-asymmetric-translog-design.md new file mode 100644 index 00000000..1772de6b --- /dev/null +++ b/docs/superpowers/specs/2026-10-16-asymmetric-translog-design.md @@ -0,0 +1,111 @@ +# `asymmetric_translog`: translog with sign-split squares around an estimated knot + +Date: 2026-10-16. Status: design, to be implemented on branch `asymmetric-translog` from +`stable-translog` at 9cd73db. + +## 1. Why + +In the health-cognition application the estimated translog own-square of physical +capacity is positive (0.07 to 0.10 at ages 74+ with own linear coefficient 0.97). Binned +next-period means of the underlying measures show why: mobility and large-muscle counts +revert strongly at the unhealthy floor (slope 0.6) and are near-linear at the healthy +end (slope 0.83). A single quadratic term has to buy the floor reversion with a slope +above one at the healthy end (own slope ρ + 2γx exceeds 1 for x > 0.15, inside the +data), and extrapolates to an unstable fixed point at +1.5. The mirror pattern (ceiling +at zero depressive symptoms, negative own-square, slope above one at the unhealthy end) +appears for psychological capacity. A symmetric clip (`stable_translog`) cannot separate +the two ends. The remedy is one curvature coefficient per side of a knot. + +## 2. Function + +For states x (all factors, latent and observed, in the model's factor order) and a +parameter vector: + +``` +f(x) = c + Σ_j β_j x_j + + Σ_j γ⁺_j max(x_j − κ_j, 0)² + + Σ_j γ⁻_j min(x_j − κ_j, 0)² + + Σ_{a 0 the own derivative (jax.grad) above the + knot equals β exactly and below the knot equals β + 2γ⁻(x − κ). +1. C¹ at the knot: jax.grad with respect to x and with respect to κ is continuous across + the knot (values at κ ± 1e-7 within 1e-6), and the function value is continuous. +1. Parameter companions: `params_asymmetric_translog` names and length + `4*nfac + nfac*(nfac-1)/2 + 1`; identity constraints as in section 3. +1. JIT, vmap and a 50-step `lax.scan` give the same values as eager evaluation. +1. Fixed-point witness for the application pattern: one-factor own map with β = 0.97, γ⁻ + = 0.1, γ⁺ = 0, κ = −0.5, c = −0.03 iterated 200 steps from x ∈ {−3, −1, 0, 0.7, 2} + stays finite and converges to the same point from every start (slope ≤ β < 1 above + the knot, and below it β + 2γ⁻(x − κ) stays in (0, 1) on this range). + +Estimation, on the fixture of `tests/test_missing_observed_factors.py` as in +`tests/test_stable_translog_estimation.py`: a registered same-law reference callable +(explicit `(fac1, fac2, ob1, ob2, params)` signature) reproduces `loglikeobs` to 1e-10 +and the gradient to rtol 1e-8; with γ⁺ = γ⁻ and knots 0 the contributions equal those of +`"translog"`; with γ⁺ ≠ γ⁻ they differ by more than 1e-8; the endogenous carry-forward +fixture accepts the function with its identity constraints written explicitly for the +reference callable (copy the `_endogenous_inputs` helper). Dispatch: the public model +and the AMN resolver return the asymmetric law (copy +`tests/test_stable_translog_dispatch.py` with one value that distinguishes the two sides +of the knot). + +## 6. Documentation + +A section in `docs/reference_guides/transition_functions.md` after `stable_translog`: +formula, parameter names, nesting, the knot as a parameter and how to fix it, the +one-sided-tail property, and the recommendation to run a fixed-point check on estimated +laws. `pixi run build-docs` must exit 0. diff --git a/src/skillmodels/af/validate.py b/src/skillmodels/af/validate.py index c328d366..cac476b4 100644 --- a/src/skillmodels/af/validate.py +++ b/src/skillmodels/af/validate.py @@ -18,6 +18,7 @@ "translog", "translog_af", "robust_translog", + "asymmetric_translog", "log_ces", "log_ces_af", "log_ces_with_constant", @@ -40,6 +41,7 @@ "linear_and_squares", "translog", "robust_translog", + "asymmetric_translog", "log_ces", "log_ces_with_constant", "log_ces_general", diff --git a/src/skillmodels/amn/simulate_and_regress.py b/src/skillmodels/amn/simulate_and_regress.py index 3e7191ab..305b38f5 100644 --- a/src/skillmodels/amn/simulate_and_regress.py +++ b/src/skillmodels/amn/simulate_and_regress.py @@ -7,8 +7,8 @@ Specialised fitters: closed-form OLS for `linear`; softmax-constrained Levenberg-Marquardt for `log_ces` and `log_ces_with_constant` (keeps gammas on the simplex). Everything else (translog, robust_translog, -linear_and_squares, log_ces_general, and any user -`@register_params`-decorated transition) goes through a generic NLS +asymmetric_translog, linear_and_squares, log_ces_general, +and any user `@register_params`-decorated transition) goes through a generic NLS path that calls the transition function directly via `jax.vmap`. This mirrors the per-factor NLS in `Monte Carlo Simulations/master_approx_simulationces2periodrho_5.R` @@ -255,6 +255,7 @@ def _resolve_transition_callable( "linear", "translog", "robust_translog", + "asymmetric_translog", "linear_and_squares", "log_ces", "log_ces_with_constant", @@ -356,8 +357,8 @@ def _fit_generic_nls( """Generic Levenberg-Marquardt NLS via `jax.vmap` over the panel. Works for any `(states, params) -> scalar` callable, including - translog, robust_translog, linear_and_squares, log_ces_general, and - user-registered transitions. + translog, robust_translog, asymmetric_translog, + linear_and_squares, log_ces_general, and user-registered transitions. Args: transition_func: callable taking a 1D state vector and a 1D @@ -470,8 +471,8 @@ def _fit_transition( `linear` and `log_ces`-family functions get specialised fitters for speed / simplex constraints; everything else (translog, - robust_translog, linear_and_squares, log_ces_general, user) falls - through to a generic `jax.vmap`-based NLS. + robust_translog, asymmetric_translog, linear_and_squares, + log_ces_general, user) falls through to a generic `jax.vmap`-based NLS. When ``cf`` is given (the investment control-function residual), an extra additive linear term ``kappa * cf`` is fitted jointly with the diff --git a/src/skillmodels/common/transition_functions.py b/src/skillmodels/common/transition_functions.py index ef54b9b4..bd16d6ac 100644 --- a/src/skillmodels/common/transition_functions.py +++ b/src/skillmodels/common/transition_functions.py @@ -381,6 +381,65 @@ def identity_constraints_robust_translog( ) +def asymmetric_translog(states: Array, params: Array) -> Array: + """Translog with separate own-square curvature above and below a knot per factor. + + f = c + Σ_j β_j x_j + Σ_j γ⁺_j max(x_j − κ_j, 0)² + Σ_j γ⁻_j min(x_j − κ_j, 0)² + + Σ_{a list[str]: + """Parameter names for asymmetric_translog.""" + return ( + list(factors) + + [f"{factor} ** 2 (above)" for factor in factors] + + [f"{factor} ** 2 (below)" for factor in factors] + + [f"{factor} knot" for factor in factors] + + [f"{a} * {b}" for a, b in combinations(factors, 2)] + + ["constant"] + ) + + +def identity_constraints_asymmetric_translog( + factor: str, + aug_period: int, + all_factors: tuple[str, ...], +) -> list[FixedConstraintWithValue]: + """Identity constraints for asymmetric_translog: own linear 1, all else 0.""" + constraints: list[FixedConstraintWithValue] = [] + for regressor in params_asymmetric_translog(all_factors): + val = 1.0 if factor == regressor else 0.0 + loc = ("transition", aug_period, factor, regressor) + constraints.append(FixedConstraintWithValue(loc=loc, value=val)) + return constraints + + def linear_and_squares(states: Array, params: Array) -> Array: """linear_and_squares transition function.""" nfac = len(states) diff --git a/tests/test_asymmetric_translog.py b/tests/test_asymmetric_translog.py new file mode 100644 index 00000000..eb9bc272 --- /dev/null +++ b/tests/test_asymmetric_translog.py @@ -0,0 +1,190 @@ +"""Kernel tests for asymmetric_translog.""" + +import itertools + +import jax +import jax.numpy as jnp +import numpy as np +from numpy.testing import assert_allclose + +from skillmodels.common.transition_functions import ( + asymmetric_translog, + identity_constraints_asymmetric_translog, + params_asymmetric_translog, + translog, +) + +jax.config.update("jax_enable_x64", True) + +# Toy parameters of tests/test_transition_functions.py::test_translog. +LIN = jnp.array([0.2, 0.1, 0.12]) +SQ = jnp.array([0.08, 0.04, 0.05]) +INTER = jnp.array([0.05, 0.03, 0.06]) +CONST = 0.04 +TRANSLOG_PARAMS = jnp.concatenate([LIN, SQ, INTER, jnp.array([CONST])]) + +STATES = jnp.array( + [ + [2, 0, 0], + [0, 3, 0], + [0, 0, 4], + [0, 0, 0], + [1, 1, 1], + [0, -3, 0], + [-1, -1, -1], + [1.5, -2, 1.8], + [12, -34, 48], + *itertools.product([-2.5, -0.3, 0.4, 2.2], repeat=3), + ], + dtype=float, +) + + +def _asym_params(lin, above, below, knots, inter, const): + return jnp.concatenate( + [lin, above, below, knots, inter, jnp.atleast_1d(jnp.asarray(const))] + ) + + +def _eval(params): + return jax.vmap(asymmetric_translog, in_axes=(0, None))(STATES, params) + + +def test_equal_curvature_zero_knots_is_translog() -> None: + params = _asym_params(LIN, SQ, SQ, jnp.zeros(3), INTER, CONST) + expected = jax.vmap(translog, in_axes=(0, None))(STATES, TRANSLOG_PARAMS) + assert_allclose(_eval(params), expected, rtol=1e-12, atol=1e-12) + + +def test_equal_curvature_nonzero_knots_is_shifted_translog() -> None: + knots = jnp.array([0.7, -1.3, 0.25]) + params = _asym_params(LIN, SQ, SQ, knots, INTER, CONST) + shifted = jnp.concatenate( + [LIN - 2 * SQ * knots, SQ, INTER, jnp.array([CONST + jnp.sum(SQ * knots**2)])] + ) + expected = jax.vmap(translog, in_axes=(0, None))(STATES, shifted) + assert_allclose(_eval(params), expected, rtol=1e-12, atol=1e-12) + + +def test_one_sided_curvature_derivative() -> None: + below = jnp.array([0.3, 0.2, 0.15]) + knots = jnp.array([0.5, -0.4, 1.0]) + params = _asym_params(LIN, jnp.zeros(3), below, knots, jnp.zeros(3), CONST) + grad = jax.grad(asymmetric_translog) + for j in range(3): + for offset in (0.01, 0.8, 3.0): + x_above = jnp.zeros(3).at[j].set(knots[j] + offset) + assert_allclose(grad(x_above, params)[j], LIN[j], rtol=0, atol=1e-15) + x_below = jnp.zeros(3).at[j].set(knots[j] - offset) + expected = LIN[j] + 2 * below[j] * (x_below[j] - knots[j]) + assert_allclose(grad(x_below, params)[j], expected, rtol=1e-14) + + +def test_c1_at_the_knot() -> None: + above = jnp.array([-0.07, 0.11, 0.02]) + below = jnp.array([0.2, -0.05, 0.09]) + knots = jnp.array([0.5, -0.4, 1.0]) + params = _asym_params(LIN, above, below, knots, INTER, CONST) + grad_x = jax.grad(asymmetric_translog, argnums=0) + grad_p = jax.grad(asymmetric_translog, argnums=1) + base = jnp.array([0.3, -0.2, 0.6]) + for j in range(3): + lo = base.at[j].set(knots[j] - 1e-7) + hi = base.at[j].set(knots[j] + 1e-7) + assert_allclose( + asymmetric_translog(lo, params), + asymmetric_translog(hi, params), + atol=1e-6, + ) + assert_allclose(grad_x(lo, params), grad_x(hi, params), atol=1e-6) + knot_slice = slice(9, 12) + assert_allclose( + grad_p(lo, params)[knot_slice], grad_p(hi, params)[knot_slice], atol=1e-6 + ) + + +def test_param_names_and_identity_constraints() -> None: + factors = ("a", "b", "c") + names = params_asymmetric_translog(factors) + nfac = len(factors) + assert len(names) == 4 * nfac + nfac * (nfac - 1) // 2 + 1 + assert names == [ + "a", + "b", + "c", + "a ** 2 (above)", + "b ** 2 (above)", + "c ** 2 (above)", + "a ** 2 (below)", + "b ** 2 (below)", + "c ** 2 (below)", + "a knot", + "b knot", + "c knot", + "a * b", + "a * c", + "b * c", + "constant", + ] + constraints = identity_constraints_asymmetric_translog("b", 3, factors) + assert [c.loc for c in constraints] == [("transition", 3, "b", n) for n in names] + assert [c.value for c in constraints] == [float(n == "b") for n in names] + + +def _law_params(): + rng = np.random.default_rng(0) + out = [] + for j in range(3): + lin = jnp.zeros(3).at[j].set(0.8) + p = _asym_params( + lin, + jnp.array(rng.uniform(-0.05, 0.05, 3)), + jnp.array(rng.uniform(-0.05, 0.05, 3)), + jnp.array(rng.uniform(-0.5, 0.5, 3)), + jnp.array(rng.uniform(-0.02, 0.02, 3)), + 0.01, + ) + out.append(p) + return jnp.stack(out) + + +def _step(x, law): + return jax.vmap(asymmetric_translog, in_axes=(None, 0))(x, law) + + +def test_jit_vmap_and_scan_match_eager() -> None: + params = _asym_params(LIN, SQ, -SQ, jnp.array([0.7, -1.3, 0.25]), INTER, CONST) + eager = np.array([float(asymmetric_translog(s, params)) for s in STATES]) + assert_allclose(_eval(params), eager, rtol=1e-14) + jitted = jax.jit(asymmetric_translog) + assert_allclose([jitted(s, params) for s in STATES], eager, rtol=1e-14) + + law = _law_params() + x0 = jnp.array([0.4, -0.6, 1.1]) + x = x0 + path = [] + for _ in range(50): + x = _step(x, law) + path.append(x) + + def body(carry, _): + new = _step(carry, law) + return new, new + + _, scanned = jax.lax.scan(body, x0, None, length=50) + assert np.all(np.isfinite(scanned)) + assert_allclose(scanned, jnp.stack(path), rtol=1e-12, atol=1e-14) + + +def test_fixed_point_witness_for_application_pattern() -> None: + params = jnp.array([0.97, 0.0, 0.1, -0.5, -0.03]) + # Below the knot: 0.1 x^2 + 0.07 x - 0.005 = 0; the root below the knot. + fixed_point = (-0.07 - np.sqrt(0.07**2 + 4 * 0.1 * 0.005)) / 0.2 + + def body(x, _): + return asymmetric_translog(jnp.atleast_1d(x), params), None + + for start in (-3.0, -1.0, 0.0, 0.7, 2.0): + end, _ = jax.lax.scan(body, jnp.asarray(start), None, length=200) + assert np.isfinite(end) + assert_allclose(end, fixed_point, atol=1e-3) diff --git a/tests/test_asymmetric_translog_dispatch.py b/tests/test_asymmetric_translog_dispatch.py new file mode 100644 index 00000000..2e10cfec --- /dev/null +++ b/tests/test_asymmetric_translog_dispatch.py @@ -0,0 +1,68 @@ +"""Public-model, AMN-resolver and AF-validation dispatch of asymmetric_translog.""" + +import jax +import jax.numpy as jnp +import pytest +from numpy.testing import assert_allclose + +from skillmodels.af.validate import ( + _AF_COMPATIBLE_TRANSITIONS, + _LEAKY_BUILTIN_PRODUCTION, + validate_af_model, +) +from skillmodels.amn.simulate_and_regress import _resolve_transition_callable +from skillmodels.common import transition_functions as tf +from skillmodels.common.model_spec import FactorSpec, ModelSpec +from skillmodels.common.process_model import process_model +from skillmodels.common.types import Normalizations + +jax.config.update("jax_enable_x64", True) + +NAMES = ("skill", "skill ** 2 (above)", "skill ** 2 (below)", "skill knot", "constant") +# β = 1, curvature 0.5 below the knot at 0 only: f(2) = 2, f(-2) = -2 + 0.5 * 4 = 0. +PARAMS = jnp.array([1.0, 0.0, 0.5, 0.0, 0.0]) + + +@pytest.fixture +def model_and_processed(): + normalizations = Normalizations( + loadings=({"m0": 1.0}, {"m0": 1.0}), + intercepts=({"m0": 0.0}, {"m0": 0.0}), + ) + factor = FactorSpec( + measurements=(("m0", "m1", "m2"), ("m0", "m1", "m2")), + normalizations=normalizations, + transition_function="asymmetric_translog", + ) + model = ModelSpec(factors={"skill": factor}) + return model, process_model(model) + + +def test_public_model_uses_asymmetric_law(model_and_processed) -> None: + _, processed = model_and_processed + info = processed.transition_info + assert info is not None + states = jnp.array([[2.0], [-2.0]]) + params = {"skill": PARAMS} + expected = jnp.array([[2.0], [0.0]]) + assert_allclose(info.func(params, states), expected, atol=1e-14) + assert_allclose(jax.jit(info.func)(params, states), expected, atol=1e-14) + assert tuple(info.param_names["skill"]) == NAMES + + +def test_amn_resolver_returns_the_asymmetric_law(model_and_processed) -> None: + model, processed = model_and_processed + func, names = _resolve_transition_callable( + "asymmetric_translog", "skill", processed, model + ) + assert func is tf.asymmetric_translog + assert names == NAMES + assert_allclose(func(jnp.array([2.0]), PARAMS), 2.0, atol=1e-14) + assert_allclose(func(jnp.array([-2.0]), PARAMS), 0.0, atol=1e-14) + + +def test_af_allow_lists_and_validation(model_and_processed) -> None: + assert "asymmetric_translog" in _AF_COMPATIBLE_TRANSITIONS + assert "asymmetric_translog" in _LEAKY_BUILTIN_PRODUCTION + model, _ = model_and_processed + validate_af_model(model) diff --git a/tests/test_asymmetric_translog_estimation.py b/tests/test_asymmetric_translog_estimation.py new file mode 100644 index 00000000..ff695433 --- /dev/null +++ b/tests/test_asymmetric_translog_estimation.py @@ -0,0 +1,174 @@ +"""asymmetric_translog through the CHS likelihood, against a same-law oracle. + +Fixture: the two-latent, two-observed, three-period panel of +tests/test_missing_observed_factors.py and its endogenous two-period model with +carry-forward substeps. +""" + +from dataclasses import replace + +import jax +import jax.numpy as jnp +import numpy as np +from numpy.testing import assert_allclose +from test_missing_observed_factors import ( + MOMENTS, + OPTIONS, + _data, + _endogenous_data, + _endogenous_model, + _model, + _params, +) + +from skillmodels.chs.maximization_inputs import get_maximization_inputs +from skillmodels.common.decorators import register_params +from skillmodels.common.transition_functions import params_asymmetric_translog + +jax.config.update("jax_enable_x64", True) + +ALL_FACTORS = ("fac1", "fac2", "ob1", "ob2") + + +def _reference_formula(x, names, params): + """Translog with own squares split at the knot, written term by term.""" + out = params["constant"] + for f, v in zip(names, x, strict=True): + d = v - params[f"{f} knot"] + out = out + params[f] * v + out = out + params[f"{f} ** 2 (above)"] * jnp.where(d > 0, d, 0.0) ** 2 + out = out + params[f"{f} ** 2 (below)"] * jnp.where(d < 0, d, 0.0) ** 2 + for i in range(len(x)): + for j in range(i + 1, len(x)): + out = out + params[f"{names[i]} * {names[j]}"] * x[i] * x[j] + return out + + +@register_params(params=params_asymmetric_translog(ALL_FACTORS)) +def REFERENCE(fac1, fac2, ob1, ob2, params): # noqa: N802 + """Registered custom callable; dags passes factors by argument name.""" + return _reference_formula([fac1, fac2, ob1, ob2], ALL_FACTORS, params) + + +def _fill(template, *, above: float, below: float, knot: float): + """`_params` of the fixture for everything but the asymmetric-only entries.""" + name2 = template.index.get_level_values(3).astype(str) + is_knot = np.asarray(name2.str.endswith(" knot")) + params = template.copy() + params.loc[~is_knot, "value"] = _params(template[~is_knot], nonlinear=0.05)["value"] + params.loc[is_knot, "value"] = knot + params.loc[np.asarray(name2.str.endswith("(above)")), "value"] = above + params.loc[np.asarray(name2.str.endswith("(below)")), "value"] = below + return params + + +def _loglikeobs_and_grad(transition, data, **values): + inputs = get_maximization_inputs( + _model(2, transition, MOMENTS), data, chs_options=OPTIONS + ) + template = inputs["params_template"] + params = _fill(template, **values) if values else _params(template, 0.05) + _, grad = inputs["loglike_and_gradient"](params) + return np.asarray(inputs["loglikeobs"](params)), np.asarray(grad) + + +def test_contributions_and_gradient_match_same_law_reference() -> None: + data = _data() + values = {"above": 0.08, "below": -0.04, "knot": 0.3} + asym, g_asym = _loglikeobs_and_grad("asymmetric_translog", data, **values) + reference, g_reference = _loglikeobs_and_grad(REFERENCE, data, **values) + assert np.all(np.isfinite(asym)) + assert np.all(np.isfinite(g_asym)) + assert_allclose(asym, reference, rtol=1e-10, atol=1e-10) + assert_allclose(g_asym, g_reference, rtol=1e-8, atol=1e-10) + # Sigma points straddle the knot: swapping the two sides must change the result. + swapped, _ = _loglikeobs_and_grad( + "asymmetric_translog", data, above=-0.04, below=0.08, knot=0.3 + ) + assert np.max(np.abs(asym - swapped)) > 1e-6 + + +def test_equal_curvature_zero_knots_reproduce_translog() -> None: + data = _data() + asym, _ = _loglikeobs_and_grad( + "asymmetric_translog", data, above=0.05, below=0.05, knot=0.0 + ) + translog, _ = _loglikeobs_and_grad("translog", data) + assert_allclose(asym, translog, rtol=1e-10, atol=1e-10) + + +def test_unequal_curvature_differs_from_translog() -> None: + data = _data() + asym, _ = _loglikeobs_and_grad( + "asymmetric_translog", data, above=0.05, below=0.02, knot=0.0 + ) + translog, _ = _loglikeobs_and_grad("translog", data) + assert np.max(np.abs(asym - translog)) > 1e-8 + + +ENDOGENOUS_FACTORS = ("s", "inv", "x", "z") + + +@register_params(params=params_asymmetric_translog(ENDOGENOUS_FACTORS)) +def REFERENCE_ENDOGENOUS(s, inv, x, z, params): # noqa: N802 + return _reference_formula([s, inv, x, z], ENDOGENOUS_FACTORS, params) + + +def _endogenous_inputs(transition, nonlinear: float): + model = _endogenous_model(None) + # with_transition_functions requires every factor; replace only `s`. + factors = dict(model.factors) | { + "s": replace(model.factors["s"], transition_function=transition) + } + model = model._replace(factors=factors) + inputs = get_maximization_inputs(model, _endogenous_data(), chs_options=OPTIONS) + params = inputs["params_template"].copy() + cat = params.index.get_level_values("category") + name2 = params.index.get_level_values(3).astype(str) + params["value"] = np.where( + cat.isin(["meas_sds", "shock_sds", "loadings"]), 0.5, 0.1 + ) + params.loc[(cat == "transition") & name2.str.contains(r"\*"), "value"] = nonlinear + # Sides of the knot differ (knots stay at 0.1). + params.loc[(cat == "transition") & name2.str.endswith("(below)"), "value"] = ( + nonlinear / 2 + ) + # Carry-forward substep of `s` (aug_period 0): identity, zero on everything else, + # knots included. Built-ins get this from identity_constraints_*; the custom + # reference callable has no such hook, so the values are written explicitly. + carry = (cat == "transition") & (params.index.get_level_values("aug_period") == 0) + carry &= params.index.get_level_values("name1") == "s" + params.loc[carry, "value"] = np.where(name2[carry] == "s", 1.0, 0.0) + for idx in params.query("category == 'initial_cholcovs'").index: + row, col = idx[3].split("-") + params.loc[idx, "value"] = float(row == col) + return inputs, params + + +def _endogenous_loglike(transition, nonlinear: float) -> float: + inputs, params = _endogenous_inputs(transition, nonlinear) + return float(inputs["loglike"](params)) + + +def _endogenous_loglikeobs_and_grad(transition, nonlinear: float): + inputs, params = _endogenous_inputs(transition, nonlinear) + _, grad = inputs["loglike_and_gradient"](params) + return np.asarray(inputs["loglikeobs"](params)), np.asarray(grad) + + +def test_endogenous_carry_forward_substeps_accept_asymmetric_translog() -> None: + linear = _endogenous_loglike("linear", nonlinear=0.0) + assert_allclose( + _endogenous_loglike("asymmetric_translog", nonlinear=0.0), linear, rtol=1e-10 + ) + asym = _endogenous_loglike("asymmetric_translog", nonlinear=0.1) + assert np.isfinite(asym) + assert abs(asym - _endogenous_loglike("translog", nonlinear=0.1)) > 1e-8 + + +def test_endogenous_contributions_and_gradient_match_same_law_reference() -> None: + asym, g_asym = _endogenous_loglikeobs_and_grad("asymmetric_translog", 0.1) + reference, g_reference = _endogenous_loglikeobs_and_grad(REFERENCE_ENDOGENOUS, 0.1) + assert np.all(np.isfinite(g_asym)) + assert_allclose(asym, reference, rtol=1e-10, atol=1e-10) + assert_allclose(g_asym, g_reference, rtol=1e-8, atol=1e-10)