Add asymmetric_translog: translog with sign-split own squares around an estimated knot

.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 }}

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<b} \delta_{ab} s_a s_b + c
+$$
+
+Use it when the data revert strongly at one end of a factor's range (a floor or a
+ceiling) and are close to linear at the other. A single square coefficient has to buy
+the reversion at one end with a slope above one at the other end; here the curvature on
+each side of the knot is free. With $\gamma^+_j = 0$, for instance, the own slope above
+the knot is exactly $\beta_j$ (plus any interaction terms), and the convex or concave
+tail is confined to values below it.
+
+Both one-sided squares are continuously differentiable in $s_j$ and in $\kappa_j$, so
+the unscented filter, the gradient and the optimiser see a smooth function.
+
+**Nesting**: with $\gamma^+_j = \gamma^-_j = \gamma_j$ the function is a `translog`
+whose linear coefficients are shifted by $-2\gamma_j\kappa_j$ and whose constant is
+shifted by $\sum_j \gamma_j \kappa_j^2$; with all knots at zero as well it equals
+`translog` exactly.
+
+**Knots** are parameters, one per factor, estimated by default. To fix a knot (e.g. to
+profile the likelihood over a grid of knots), fix the parameter
+`("transition", period, factor, "<regressor> 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
 

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<b} δ_ab x_a x_b
+```
+
+- `max(·,0)²` and `min(·,0)²` are C¹ in x and in κ, so the unscented filter, the
+  gradient and the optimiser see a smooth function. Implement as
+  `jnp.maximum(d, 0.0) ** 2` and `jnp.minimum(d, 0.0) ** 2` with `d = x − κ`; no
+  `where`.
+- Nesting: γ⁺ = γ⁻ = γ gives `c + Σβx + Σγ(x−κ)² + Σδ x_a x_b`, a translog with the
+  linear and constant terms shifted by `−2γκ` and `γκ²`. With κ = 0 it is exactly
+  `translog`. Tests pin both.
+- The knot κ_j is a parameter per factor, estimated by default. Users fix it through the
+  ordinary constraints machinery (profiling over a grid of knots is the intended first
+  use).
+- No clipping. If a fixed-point check of an estimated law shows a spurious attractor far
+  in the convex tail, composing with `_soft_clip` is a later extension, not part of this
+  spec.
+
+## 3. Parameter layout and names
+
+`params_asymmetric_translog(factors)` returns, in this order:
+
+```
+[f for f in factors]                      # β, linear
++ [f"{f} ** 2 (above)" for f in factors]  # γ⁺
++ [f"{f} ** 2 (below)" for f in factors]  # γ⁻
++ [f"{f} knot" for f in factors]          # κ
++ [f"{a} * {b}" for a, b in combinations(factors, 2)]  # δ
++ ["constant"]
+```
+
+so with `nfac = len(states)`: `lin = params[:nfac]`, `above = params[nfac:2*nfac]`,
+`below = params[2*nfac:3*nfac]`, `knots = params[3*nfac:4*nfac]`,
+`inter = params[4*nfac:-1]`, `constant = params[-1]`.
+
+`identity_constraints_asymmetric_translog(factor, aug_period, all_factors)`: own linear
+coefficient 1, every other parameter including the knots 0 (same pattern as
+`identity_constraints_translog`).
+
+## 4. Registration
+
+Add `"asymmetric_translog"` wherever `"stable_translog"` is listed:
+`_AF_COMPATIBLE_TRANSITIONS` and `_LEAKY_BUILTIN_PRODUCTION` in
+`src/skillmodels/af/validate.py`, `builtin_names` in
+`src/skillmodels/amn/simulate_and_regress.py`, and the lists in
+`docs/reference_guides/transition_functions.md`. The function has no keyword-only
+configuration, so `_hide_keyword_only_config` is a no-op for it.
+
+## 5. Tests (`tests/test_asymmetric_translog.py`, `tests/test_asymmetric_translog_estimation.py`)
+
+Kernel, on the 3-factor toy states and parameters of
+`tests/test_transition_functions.py`:
+
+1. γ⁺ = γ⁻, κ = 0 equals `translog` to 1e-12 on a grid of states including both signs.
+1. γ⁺ = γ⁻ = γ, κ ≠ 0 equals `translog` with linear coefficients β − 2γκ and constant c
+   \+ Σγκ² to 1e-12.
+1. One-sided curvature: with γ⁺ = 0 and γ⁻ > 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.

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",

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

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<b} δ_ab x_a x_b.
+
+    Both one-sided squares are C¹ in x and κ. With γ⁺ = γ⁻ = γ it is a translog whose
+    linear and constant terms are shifted by −2γκ and γκ²; with κ = 0 as well it is
+    exactly `translog`. The knots κ are parameters (estimated unless fixed). No
+    clipping: check the fixed points of estimated laws before iterating them.
+
+    Parameter layout: linear, squares above, squares below, knots, interactions,
+    constant (see `params_asymmetric_translog`).
+
+    """  # noqa: RUF002
+    nfac = len(states)
+    constant = params[-1]
+    lin_beta = params[:nfac]
+    above_beta = params[nfac : 2 * nfac]
+    below_beta = params[2 * nfac : 3 * nfac]
+    knots = params[3 * nfac : 4 * nfac]
+    inter_beta = params[4 * nfac : -1]
+
+    d = states - knots
+    res = jnp.dot(states, lin_beta)
+    res += jnp.dot(jnp.maximum(d, 0.0) ** 2, above_beta)
+    res += jnp.dot(jnp.minimum(d, 0.0) ** 2, below_beta)
+    for p, (a, b) in zip(inter_beta, combinations(range(nfac), 2), strict=False):
+        res += p * states[a] * states[b]
+    res += constant
+    return res
+
+
+def params_asymmetric_translog(factors: tuple[str, ...]) -> 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)