From 4ec7ce0425cf364920630f24a1c3caff3fc7736e Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 25 Sep 2026 11:42:17 +0200 Subject: [PATCH 1/3] TST Pin float32 felicity at the leisure floor against a 60-digit reference Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- tests/test_leisure_floor.py | 190 ++++++++++++++++++++++++++++++++++++ 1 file changed, 190 insertions(+) create mode 100644 tests/test_leisure_floor.py diff --git a/tests/test_leisure_floor.py b/tests/test_leisure_floor.py new file mode 100644 index 0000000..2c1e470 --- /dev/null +++ b/tests/test_leisure_floor.py @@ -0,0 +1,190 @@ +"""Float32 felicity at the leisure floor, against a 60-digit reference. + +The class covers the canwork cells at production preference values: every pref +type, health, lagged labor supply, hours choice and canwork age, at consumption +points spanning the NB-EGM numeric-inverse bracket (`1e-8` to the top of the +action range). Marginal utility is `jax.grad` of the consumption-dollar +felicity, the route pylcm's NB-EGM uses. +""" + +import itertools +from decimal import Decimal, getcontext + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from aca_model.agent import preferences + +getcontext().prec = 60 + +CONSUMPTION_WEIGHTS = (0.6776541629520845, 0.8805772328591686, 0.0718086283445225) +COEFFICIENTS_RRA = (3.841252231680976, 0.9990771146810682, 3.8328505891095137) +TIME_ENDOWMENT = 3926.9478390365557 +BAD_HEALTH_COST = 408.9190313043897 +REENTRY_COST = 119.61095526191421 +FIXED_COST_INTERCEPT = 337.5223349543413 +FIXED_COST_AGE_TREND = 84.09242960636178 +REFERENCE_AGE = 50 +REFERENCE_HOURS = 1000.0 +AVERAGE_CONSUMPTION = 20000.0 +HOURS = (0.0, 1000.0, 1500.0, 2000.0, 2500.0) +CANWORK_AGES = tuple(range(51, 72)) +CONSUMPTION = (1e-8, 1597.0921419521899, 20000.0, 300_000.0, 721271853.846792) +FLOAT32_REL_TOL = 1e-4 +# Leisure-map widths as fractions of the endowment: smoothing width and floor. +SMOOTHING_FRACTION = 0.01 +FLOOR_FRACTION = 1e-5 + + +def _f32(x: float) -> jnp.ndarray: + return jnp.asarray(x, dtype=jnp.float32) + + +def _felicity_of_dollars(consumption, leisure, weight, rra, scale): + return preferences.u_alive( + consumption_equiv=preferences.consumption_equiv( + consumption_dollars=consumption, + equivalence_scale=jnp.ones_like(consumption), + ), + leisure=leisure, + consumption_weight=weight, + coefficient_rra=rra, + utility_scale_factor=scale, + ) + + +def _leisure_available(good: int, lagged: int, hours: float, age: int) -> float: + fixed_cost = FIXED_COST_INTERCEPT + FIXED_COST_AGE_TREND * (age - REFERENCE_AGE) + reentry = REENTRY_COST if lagged == 0 else 0.0 + work = hours + fixed_cost + reentry if hours > 0 else 0.0 + return TIME_ENDOWMENT - (0.0 if good else BAD_HEALTH_COST) - work + + +def _cells() -> list[tuple[float, float]]: + """All `(leisure_available, consumption)` pairs of the class.""" + return [ + (_leisure_available(good, lagged, hours, age), c) + for good, lagged, hours, age, c in itertools.product( + (0, 1), (0, 1), HOURS, CANWORK_AGES, CONSUMPTION + ) + ] + + +def _float32_u_and_marginal( + pref_type: int, cells: list[tuple[float, float]] +) -> tuple[np.ndarray, np.ndarray]: + available = jnp.asarray([a for a, _ in cells], dtype=jnp.float32) + consumption = jnp.asarray([c for _, c in cells], dtype=jnp.float32) + weight = _f32(CONSUMPTION_WEIGHTS[pref_type]) + rra = _f32(COEFFICIENTS_RRA[pref_type]) + leisure = _floored_leisure(available, _f32(TIME_ENDOWMENT)) + scale = preferences.utility_scale_factor( + average_consumption_equiv=_f32(AVERAGE_CONSUMPTION), + consumption_weight=weight, + coefficient_rra=rra, + time_endowment=_f32(TIME_ENDOWMENT), + fixed_cost_of_work_intercept=_f32(FIXED_COST_INTERCEPT), + reference_hours=_f32(REFERENCE_HOURS), + ) + args = (weight, rra, scale) + u = jax.vmap(_felicity_of_dollars, in_axes=(0, 0, None, None, None))( + consumption, leisure, *args + ) + marginal = jax.vmap( + jax.grad(_felicity_of_dollars), in_axes=(0, 0, None, None, None) + )(consumption, leisure, *args) + assert u.dtype == marginal.dtype == jnp.float32 + return np.asarray(u, dtype=np.float64), np.asarray(marginal, dtype=np.float64) + + +def _floored_leisure(available, endowment): + """Production leisure at a given `leisure_available`, via the tied-regime map. + + Good health and no fixed cost make the work loss the only deduction, so hours of + `endowment - available` leave exactly `available`. + """ + return preferences.leisure_canwork_tied( + working_hours_value=endowment - available, + good_health=jnp.ones(available.shape, dtype=jnp.int32), + time_endowment=endowment, + leisure_cost_of_bad_health=jnp.zeros_like(endowment), + fixed_cost_of_work=jnp.zeros_like(endowment), + ) + + +def _d(x: float) -> Decimal: + return Decimal(repr(float(x))) + + +def _reference_u_and_marginal( + pref_type: int, cells: list[tuple[float, float]] +) -> tuple[np.ndarray, np.ndarray]: + """Closed forms in 60-digit log space, independent of the JAX code path.""" + weight, rra = _d(CONSUMPTION_WEIGHTS[pref_type]), _d(COEFFICIENTS_RRA[pref_type]) + endowment = _d(TIME_ENDOWMENT) + smoothing = _d(SMOOTHING_FRACTION) * endowment + floor = _d(FLOOR_FRACTION) * endowment + average_leisure = endowment - _d(REFERENCE_HOURS) - _d(FIXED_COST_INTERCEPT) + log_average = weight * _d(AVERAGE_CONSUMPTION).ln() + (1 - weight) * ( + average_leisure.ln() + ) + scale = abs((1 - rra) / ((1 - rra) * log_average).exp()) + u, marginal = [], [] + for available, c in cells: + leisure = ( + smoothing + * ((floor / smoothing).exp() + (_d(available) / smoothing).exp()).ln() + ) + log_composite = weight * _d(c).ln() + (1 - weight) * leisure.ln() + power = ((1 - rra) * log_composite).exp() + u.append(float(scale * power / (1 - rra))) + marginal.append(float(scale * weight / _d(c) * power)) + return np.array(u), np.array(marginal) + + +OUTPUTS = {"u": 0, "marginal": 1} +WITNESS = [(_leisure_available(0, 0, 2500.0, 71), c) for c in CONSUMPTION] + + +@pytest.mark.parametrize("output", ["u", "marginal"]) +def test_float32_felicity_at_witness_cell_matches_reference(output: str) -> None: + """Pref type 2, bad health, did not work, 2500 h, age 71: finite and accurate.""" + got = _float32_u_and_marginal(2, WITNESS)[OUTPUTS[output]] + ref = _reference_u_and_marginal(2, WITNESS)[OUTPUTS[output]] + np.testing.assert_allclose(got, ref, rtol=FLOAT32_REL_TOL) + + +@pytest.mark.parametrize("output", ["u", "marginal"]) +@pytest.mark.parametrize("pref_type", [0, 1, 2]) +def test_float32_felicity_over_canwork_class_matches_reference( + pref_type: int, output: str +) -> None: + """Every canwork cell and consumption point: float32 equals the reference.""" + cells = _cells() + got = _float32_u_and_marginal(pref_type, cells)[OUTPUTS[output]] + ref = _reference_u_and_marginal(pref_type, cells)[OUTPUTS[output]] + np.testing.assert_allclose(got, ref, rtol=FLOAT32_REL_TOL) + + +@pytest.mark.parametrize("pref_type", [0, 1, 2]) +def test_leisure_floor_leaves_utility_unchanged_away_from_the_floor( + pref_type: int, +) -> None: + """With leisure_available above 10% of the endowment, u moves by at most 1e-7. + + Compared against the unfloored softplus `s * log(1 + e^(x/s))`. + """ + available = np.array([a for a, _ in _cells() if a > 0.1 * TIME_ENDOWMENT]) + assert available.size > 0 + smoothing = SMOOTHING_FRACTION * TIME_ENDOWMENT + unfloored = smoothing * np.logaddexp(0.0, available / smoothing) + floored = np.asarray( + _floored_leisure(jnp.asarray(available), jnp.asarray(TIME_ENDOWMENT)) + ) + exponent = (1.0 - CONSUMPTION_WEIGHTS[pref_type]) * ( + 1.0 - COEFFICIENTS_RRA[pref_type] + ) + relative_change = np.abs((floored / unfloored) ** exponent - 1.0) + np.testing.assert_array_less(relative_change, 1e-7) From 4f37ec2d0945592a641f4bb62129dae49608edc6 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 25 Sep 2026 11:42:17 +0200 Subject: [PATCH 2/3] ENH Bend leisure onto a positive floor with a smooth max Leisure is s*logaddexp(F/s, available/s) with F = 1e-5 * time_endowment, so float32 felicity and its consumption derivative stay finite when work costs exceed the endowment, while cells away from the floor are unchanged to ~1e-8. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- src/aca_model/agent/preferences.py | 33 ++++++++++++++++++++++-------- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/src/aca_model/agent/preferences.py b/src/aca_model/agent/preferences.py index 438b2c4..3bb4b80 100644 --- a/src/aca_model/agent/preferences.py +++ b/src/aca_model/agent/preferences.py @@ -23,6 +23,13 @@ # endowment; it only bends the map near and beyond the endowment. _LEISURE_SMOOTHING_FRACTION = 0.01 +# Positive floor that leisure approaches as work costs exceed the endowment, as a +# fraction of the time endowment. It keeps felicity and its consumption derivative +# inside float32's range when leisure_available is far below zero, down to the +# NB-EGM inverse bracket's lowest consumption; away from the floor its effect decays +# exponentially. +_LEISURE_FLOOR_FRACTION = 1e-5 + @categorical(ordered=False) class PrefType: @@ -71,17 +78,25 @@ def fixed_cost_of_work( def _smooth_leisure_floor( leisure_available: FloatND, time_endowment: ScalarFloat ) -> FloatND: - """Bend leisure to a strictly positive floor as work costs approach the endowment. - - `softplus(x) = log(1 + e^x)` via `jnp.logaddexp(0, x)`, scaled by a small fraction - of the endowment. Where `leisure_available` is large relative to the smoothing width - the map reduces to `leisure_available` (bulk unchanged); as it falls to zero leisure - bends to `0⁺` — never negative, never a kinked clamp — so the CRRA aggregator never - receives a non-positive base. The smoothing width scales with the endowment, so the - map is scale-invariant. + """Bend leisure smoothly onto a positive floor as work costs approach the endowment. + + Leisure is a smooth maximum of `leisure_available` and the floor + `F = _LEISURE_FLOOR_FRACTION * time_endowment`: + `s * logaddexp(F / s, leisure_available / s)`, with smoothing width + `s = _LEISURE_SMOOTHING_FRACTION * time_endowment`. + + - `leisure_available` large relative to `s`: leisure equals `leisure_available` + up to a term that decays like `e^(-leisure_available / s)`. + - `leisure_available` at or below zero: leisure bends to `F⁺` — never below the + floor, never a kinked clamp — so every work choice stays feasible and the CRRA + aggregator's base stays large enough for float32 felicity and marginal + felicity. + + Both widths scale with the endowment, so the map is scale-invariant. """ smoothing = _LEISURE_SMOOTHING_FRACTION * time_endowment - return smoothing * jnp.logaddexp(0.0, leisure_available / smoothing) + floor = _LEISURE_FLOOR_FRACTION * time_endowment + return smoothing * jnp.logaddexp(floor / smoothing, leisure_available / smoothing) def leisure_canwork_retiree_or_nongroup( From 05c4a4331d787234a0ba0d32b801f5e4cd065991 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 25 Sep 2026 13:47:51 +0200 Subject: [PATCH 3/3] CI Install pylcm at the commit aca-dev pins Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- .github/workflows/main.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 632668d..897c37f 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -30,12 +30,12 @@ jobs: - uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - - name: Install pylcm from feat/nb-egm - # This slice depends on the NB-EGM solver, which lives on pylcm's - # feat/nb-egm branch; revert to @main once that branch merges. + - name: Install pylcm at the revision aca-dev pins + # Pinned by commit so a deleted or rewritten branch cannot break the + # install; move it together with the aca-dev pylcm submodule pointer. run: >- pip install "pylcm @ - git+https://github.com/OpenSourceEconomics/pylcm.git@feat/nb-egm" + git+https://github.com/OpenSourceEconomics/pylcm.git@7e8e4c223a63a07725015ebdd5af84cd9092cee2" - name: Install aca-model with test deps run: pip install -e . pytest pdbp - name: Run pytest