From 186e430ef399fcd0db174fd4d3cb118856e29d39 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 11 May 2026 14:32:32 +0200 Subject: [PATCH 1/9] @categorical: require ScalarInt field annotations + ScalarInt runtime values (#349) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The decorator used to read `__annotations__` only for field names and never validate them — every consumer wrote `field: int` even though codes flowed through JAX as `jnp.int32` everywhere downstream. The mismatch was a consistent lie ty couldn't catch, and downstream `jnp.int32(...)` wraps papered over the dtype gap at runtime. This change closes the gap at all three layers in lockstep: * **Annotation gate.** `@categorical` requires every field to be annotated `ScalarInt` (from `lcm.typing`); other annotations raise the new `CategoricalDefinitionError` at decoration time, naming the offending fields and pointing at the import. * **Runtime values.** Class- and instance-level attribute access return 0-d `jnp.int32` scalars. The decorator assigns `field(default=i, init=False)` (Python int placeholders for `dataclass(frozen=True)`'s mutable-default check) then overrides the class attributes with `jnp.int32(i)` via `type.__setattr__` post-decoration; `init=False` keeps instance `__dict__` empty so attribute lookup falls through to the class scalar. * **Validator.** `validate_category_class` now checks `isinstance(value, jax.Array) and value.shape == () and jnp.issubdtype(value.dtype, jnp.integer)`; the consecutiveness check coerces via `int(v)`. Hashability fix: JAX 0-d arrays aren't hashable, so the two pylcm dict-inversion sites (`simulation/simulate.py`, `initial_conditions.py`) and `DiscreteGrid.codes` now coerce to Python `int` at the boundary where hashable keys are needed. New helper `lcm.invert_regime_ids` lives in `lcm.utils.containers` and centralises the inversion pattern; downstream consumers can call it instead of hand-rolling `{int(v): k for k, v in ...}`. `RegimeNamesToIds` tightens to `MappingProxyType[RegimeName, ScalarInt]`. Sweep: every `@categorical`-decorated class in `src/lcm_examples`, `tests/`, `tests/test_models/`, and the user-guide notebooks/docs moves from `field: int` to `field: ScalarInt`. The `Effort` example in `mahler_yum_2024/_model.py` switches to the `make_dataclass + type.__setattr__` shim so its 40 `ScalarInt` defaults bypass `dataclass`'s mutable-default check. `tests/conftest.py` and `tests/test_grids.py` gain a `_make_dc` helper for hand-crafted test dataclasses with `ScalarInt` values. Co-Authored-By: Claude Opus 4.7 (1M context) --- AGENTS.md | 12 +- docs/explanations/beta_delta.ipynb | 4 +- .../function_representation.ipynb | 69 +-------- docs/user_guide/defining_models.md | 21 +-- docs/user_guide/grids.md | 5 +- docs/user_guide/tiny_example.ipynb | 8 +- docs/user_guide/transitions.ipynb | 8 +- src/lcm/__init__.py | 2 + src/lcm/exceptions.py | 10 ++ src/lcm/grids/categorical.py | 112 ++++++++++----- src/lcm/grids/discrete.py | 6 +- src/lcm/pandas_utils.py | 2 +- src/lcm/regime_building/processing.py | 5 +- src/lcm/simulation/initial_conditions.py | 21 +-- src/lcm/simulation/result.py | 3 +- src/lcm/simulation/simulate.py | 9 +- src/lcm/simulation/transitions.py | 2 +- src/lcm/typing.py | 2 +- src/lcm/utils/containers.py | 13 ++ src/lcm_examples/mahler_yum_2024/_model.py | 46 +++--- src/lcm_examples/mortality.py | 11 +- src/lcm_examples/precautionary_savings.py | 4 +- .../precautionary_savings_health.py | 8 +- src/lcm_examples/tiny.py | 8 +- tests/conftest.py | 10 +- .../regime_building/test_regime_processing.py | 11 +- tests/simulation/test_initial_conditions.py | 32 ++--- tests/simulation/test_simulate.py | 2 +- tests/solution/test_beta_delta.py | 4 +- tests/solution/test_custom_aggregator.py | 14 +- tests/solution/test_diagnostics.py | 6 +- tests/test_Q_and_F.py | 12 +- tests/test_chained_state_transitions.py | 8 +- tests/test_error_handling_invalid_vf.py | 4 +- tests/test_function_output_grid_indexing.py | 30 ++-- tests/test_function_representation.py | 8 +- tests/test_grids.py | 132 ++++++++++++++---- tests/test_model.py | 68 ++++----- tests/test_models/basic_discrete.py | 13 +- tests/test_models/deterministic/base.py | 6 +- tests/test_models/deterministic/discrete.py | 14 +- tests/test_models/deterministic/regression.py | 4 +- tests/test_models/regime_markov.py | 10 +- tests/test_models/shock_grids.py | 14 +- tests/test_models/stochastic.py | 9 +- tests/test_nan_diagnostics.py | 4 +- tests/test_next_state.py | 13 +- tests/test_pandas_utils.py | 99 ++++++------- tests/test_persistence.py | 4 +- tests/test_regime.py | 16 +-- tests/test_regime_state_mismatch.py | 98 ++++++------- tests/test_runtime_params.py | 6 +- tests/test_runtime_shock_params.py | 6 +- tests/test_single_feasible_action.py | 24 ++-- ...est_solution_on_toy_model_deterministic.py | 12 +- .../test_solution_on_toy_model_stochastic.py | 6 +- tests/test_state_action_space.py | 16 +-- tests/test_static_params.py | 26 ++-- tests/test_stochastic.py | 8 +- tests/test_validate_param_types.py | 11 +- .../test_validate_regime_transition_probs.py | 10 +- tests/test_validation_scalar_actions.py | 6 +- 62 files changed, 653 insertions(+), 514 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 1929fced0..b311ee198 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -70,7 +70,9 @@ automation. Python 3.14+ is required. - `PiecewiseLogSpacedGrid`: Piecewise logarithmically spaced grid with breakpoints. - `AgeGrid`: Lifecycle age grid (start, stop, step or exact_values) - `@categorical(ordered=...)`: Decorator factory for creating categorical classes with - auto-assigned integer codes. Requires explicit `ordered=True` or `ordered=False`. + auto-assigned `ScalarInt` (0-d `jnp.int32`) codes. Requires explicit `ordered=True` or + `ordered=False`. Every field must be annotated as `ScalarInt` (from `lcm.typing`) — + other annotations raise `CategoricalDefinitionError` at decoration time. - **ShockGrids** (in `src/lcm/shocks/`): `Rouwenhorst`, `Tauchen`, `Normal`, `Uniform`. These have intrinsic transitions — do NOT accept entries in `state_transitions`. Import as modules (`import lcm.shocks.iid`) and use qualified access @@ -189,8 +191,8 @@ from lcm import AgeGrid, categorical @categorical(ordered=False) class RegimeId: - working: int - retired: int + working: ScalarInt + retired: ScalarInt Model( @@ -210,8 +212,8 @@ Model( - Must have at least one terminal regime and one non-terminal regime - `regime_id_class` must be a dataclass with fields matching regime names (use `@categorical`) -- Field values must be consecutive integers starting from 0 (auto-assigned by - `@categorical`) +- Field values are consecutive `ScalarInt` (0-d `jnp.int32`) scalars starting from 0, + auto-assigned by `@categorical` ### Core Methods diff --git a/docs/explanations/beta_delta.ipynb b/docs/explanations/beta_delta.ipynb index a6f05c8ef..57a6a51e4 100644 --- a/docs/explanations/beta_delta.ipynb +++ b/docs/explanations/beta_delta.ipynb @@ -114,8 +114,8 @@ "source": [ "@categorical(ordered=False)\n", "class RegimeId:\n", - " working: int\n", - " dead: int\n", + " working: ScalarInt\n", + " dead: ScalarInt\n", "\n", "\n", "def utility(consumption: ContinuousAction) -> FloatND:\n", diff --git a/docs/explanations/function_representation.ipynb b/docs/explanations/function_representation.ipynb index e8fe6ebfe..13aff1817 100644 --- a/docs/explanations/function_representation.ipynb +++ b/docs/explanations/function_representation.ipynb @@ -55,74 +55,7 @@ "id": "3", "metadata": {}, "outputs": [], - "source": [ - "import jax.numpy as jnp\n", - "import plotly.graph_objects as go\n", - "\n", - "from lcm import AgeGrid, LinSpacedGrid, Model, Regime, categorical\n", - "from lcm.typing import ContinuousAction, ContinuousState, FloatND\n", - "\n", - "blue, orange, green = \"#4C78A8\", \"#F58518\", \"#54A24B\"\n", - "\n", - "\n", - "def utility(consumption: ContinuousAction, risk_aversion: float) -> FloatND:\n", - " return consumption ** (1 - risk_aversion) / (1 - risk_aversion)\n", - "\n", - "\n", - "def next_wealth(\n", - " wealth: ContinuousState,\n", - " consumption: ContinuousAction,\n", - " interest_rate: float,\n", - ") -> ContinuousState:\n", - " return (1 + interest_rate) * (wealth - consumption)\n", - "\n", - "\n", - "def borrowing_constraint(\n", - " consumption: ContinuousAction, wealth: ContinuousState\n", - ") -> FloatND:\n", - " return consumption <= wealth\n", - "\n", - "\n", - "@categorical(ordered=False)\n", - "class RegimeId:\n", - " working_life: int\n", - " retirement: int\n", - "\n", - "\n", - "retirement_regime = Regime(\n", - " transition=None,\n", - " functions={\"utility\": utility},\n", - " constraints={\"borrowing_constraint\": borrowing_constraint},\n", - " actions={\"consumption\": LinSpacedGrid(start=1, stop=400, n_points=50)},\n", - " states={\"wealth\": LinSpacedGrid(start=1, stop=400, n_points=10)},\n", - ")\n", - "\n", - "working_life_regime = Regime(\n", - " transition=lambda: RegimeId.retirement,\n", - " functions={\"utility\": utility},\n", - " constraints={\"borrowing_constraint\": borrowing_constraint},\n", - " actions={\"consumption\": LinSpacedGrid(start=1, stop=400, n_points=50)},\n", - " states={\n", - " \"wealth\": LinSpacedGrid(start=1, stop=400, n_points=10),\n", - " },\n", - " state_transitions={\n", - " \"wealth\": next_wealth,\n", - " },\n", - ")\n", - "\n", - "model = Model(\n", - " description=\"Minimal consumption-savings model\",\n", - " ages=AgeGrid(start=25, stop=65, step=\"20Y\"),\n", - " regimes={\"working_life\": working_life_regime, \"retirement\": retirement_regime},\n", - " regime_id_class=RegimeId,\n", - ")\n", - "\n", - "params = {\n", - " \"discount_factor\": 0.95,\n", - " \"risk_aversion\": 1.5,\n", - " \"interest_rate\": 0.04,\n", - "}" - ] + "source": "import jax.numpy as jnp\nimport plotly.graph_objects as go\n\nfrom lcm import AgeGrid, LinSpacedGrid, Model, Regime, categorical\nfrom lcm.typing import ContinuousAction, ContinuousState, FloatND, ScalarInt\n\nblue, orange, green = \"#4C78A8\", \"#F58518\", \"#54A24B\"\n\n\ndef utility(consumption: ContinuousAction, risk_aversion: float) -> FloatND:\n return consumption ** (1 - risk_aversion) / (1 - risk_aversion)\n\n\ndef next_wealth(\n wealth: ContinuousState,\n consumption: ContinuousAction,\n interest_rate: float,\n) -> ContinuousState:\n return (1 + interest_rate) * (wealth - consumption)\n\n\ndef borrowing_constraint(\n consumption: ContinuousAction, wealth: ContinuousState\n) -> FloatND:\n return consumption <= wealth\n\n\n@categorical(ordered=False)\nclass RegimeId:\n working_life: ScalarInt\n retirement: ScalarInt\n\n\nretirement_regime = Regime(\n transition=None,\n functions={\"utility\": utility},\n constraints={\"borrowing_constraint\": borrowing_constraint},\n actions={\"consumption\": LinSpacedGrid(start=1, stop=400, n_points=50)},\n states={\"wealth\": LinSpacedGrid(start=1, stop=400, n_points=10)},\n)\n\nworking_life_regime = Regime(\n transition=lambda: RegimeId.retirement,\n functions={\"utility\": utility},\n constraints={\"borrowing_constraint\": borrowing_constraint},\n actions={\"consumption\": LinSpacedGrid(start=1, stop=400, n_points=50)},\n states={\n \"wealth\": LinSpacedGrid(start=1, stop=400, n_points=10),\n },\n state_transitions={\n \"wealth\": next_wealth,\n },\n)\n\nmodel = Model(\n description=\"Minimal consumption-savings model\",\n ages=AgeGrid(start=25, stop=65, step=\"20Y\"),\n regimes={\"working_life\": working_life_regime, \"retirement\": retirement_regime},\n regime_id_class=RegimeId,\n)\n\nparams = {\n \"discount_factor\": 0.95,\n \"risk_aversion\": 1.5,\n \"interest_rate\": 0.04,\n}" }, { "cell_type": "markdown", diff --git a/docs/user_guide/defining_models.md b/docs/user_guide/defining_models.md index 7af938d48..84f8769fc 100644 --- a/docs/user_guide/defining_models.md +++ b/docs/user_guide/defining_models.md @@ -15,7 +15,7 @@ from lcm import Model model = Model( regimes=regimes, # dict mapping names to Regime instances ages=ages, # AgeGrid defining the lifecycle timeline - regime_id_class=RegimeId, # @categorical dataclass mapping names to int indices + regime_id_class=RegimeId, # @categorical dataclass mapping names to ScalarInt indices enable_jit=True, # controls JAX compilation (default: True) fixed_params={}, # optional params baked in at init time description="", # optional description string @@ -32,18 +32,22 @@ decorator to create it: ```python from lcm import categorical +from lcm.typing import ScalarInt @categorical(ordered=False) class RegimeId: - retired: int - working: int + retired: ScalarInt + working: ScalarInt ``` Rules: +- Fields must be annotated as `ScalarInt` — the 0-d `jnp.int32` scalar pylcm produces + for category codes. Other annotations raise `CategoricalDefinitionError` at decoration + time. - Fields must match the keys of the `regimes` dict exactly (sorted alphabetically). -- Values are auto-assigned as consecutive integers starting from 0. +- Values are auto-assigned as consecutive `jnp.int32` scalars starting from 0. - Use `RegimeId.working` (class attribute access) to reference regime IDs in transition functions. @@ -116,18 +120,19 @@ Use `model.get_params_template()` to get a mutable copy of the parameter templat ```python import jax.numpy as jnp from lcm import AgeGrid, DiscreteGrid, LinSpacedGrid, Model, Regime, categorical +from lcm.typing import ScalarInt @categorical(ordered=False) class RegimeId: - retired: int - working: int + retired: ScalarInt + working: ScalarInt @categorical(ordered=True) class LaborSupply: - do_not_work: int - work: int + do_not_work: ScalarInt + work: ScalarInt def next_wealth(wealth, consumption, interest_rate): diff --git a/docs/user_guide/grids.md b/docs/user_guide/grids.md index 635585c4b..a15ae8c83 100644 --- a/docs/user_guide/grids.md +++ b/docs/user_guide/grids.md @@ -43,12 +43,13 @@ categories: ```python from lcm import DiscreteGrid, categorical +from lcm.typing import ScalarInt @categorical(ordered=True) class LaborSupply: - do_not_work: int - work: int + do_not_work: ScalarInt + work: ScalarInt actions = {"labor_supply": DiscreteGrid(LaborSupply)} diff --git a/docs/user_guide/tiny_example.ipynb b/docs/user_guide/tiny_example.ipynb index 527eaee32..7e2bafb52 100644 --- a/docs/user_guide/tiny_example.ipynb +++ b/docs/user_guide/tiny_example.ipynb @@ -117,14 +117,14 @@ "source": [ "@categorical(ordered=True)\n", "class LaborSupply:\n", - " do_not_work: int\n", - " work: int\n", + " do_not_work: ScalarInt\n", + " work: ScalarInt\n", "\n", "\n", "@categorical(ordered=False)\n", "class RegimeId:\n", - " working_life: int\n", - " retirement: int" + " working_life: ScalarInt\n", + " retirement: ScalarInt" ] }, { diff --git a/docs/user_guide/transitions.ipynb b/docs/user_guide/transitions.ipynb index cac68b722..11d82dccf 100644 --- a/docs/user_guide/transitions.ipynb +++ b/docs/user_guide/transitions.ipynb @@ -52,8 +52,8 @@ "source": [ "@categorical(ordered=False)\n", "class RegimeId:\n", - " working_life: int\n", - " retirement: int\n", + " working_life: ScalarInt\n", + " retirement: ScalarInt\n", "\n", "\n", "def next_regime(age: float, retirement_age: float) -> ScalarInt:\n", @@ -84,8 +84,8 @@ "source": [ "@categorical(ordered=False)\n", "class RegimeIdMortality:\n", - " alive: int\n", - " dead: int\n", + " alive: ScalarInt\n", + " dead: ScalarInt\n", "\n", "\n", "def survival_transition(survival_prob: float) -> FloatND:\n", diff --git a/src/lcm/__init__.py b/src/lcm/__init__.py index 09e37a1fa..d6d267568 100644 --- a/src/lcm/__init__.py +++ b/src/lcm/__init__.py @@ -52,6 +52,7 @@ ) from lcm.regime import MarkovTransition, Regime from lcm.simulation.result import SimulationResult +from lcm.utils.containers import invert_regime_ids from lcm.utils.error_handling import validate_transition_probs # Register MappingProxyType as a JAX pytree so it can be used in JIT-traced functions. @@ -80,6 +81,7 @@ "SolveSnapshot", "__version__", "categorical", + "invert_regime_ids", "load_snapshot", "load_solution", "save_solution", diff --git a/src/lcm/exceptions.py b/src/lcm/exceptions.py index 30fd1bd10..c00e55aff 100644 --- a/src/lcm/exceptions.py +++ b/src/lcm/exceptions.py @@ -50,6 +50,16 @@ class GridInitializationError(PyLCMError): """Raised when there is an error in the grid initialization.""" +class CategoricalDefinitionError(PyLCMError): + """Raised when an `@categorical`-decorated class fails the contract. + + `@categorical` requires every field to be annotated as `ScalarInt` + (the 0-d `jnp.int32` scalar pylcm produces for category codes). + Violations are caught at decoration time, before any grid, regime, + or derived-categorical mapping is built. + """ + + class FunctionDispatchError(PyLCMError): """Raised when there is an error during the function dispatch.""" diff --git a/src/lcm/grids/categorical.py b/src/lcm/grids/categorical.py index 34fc5f977..fef009754 100644 --- a/src/lcm/grids/categorical.py +++ b/src/lcm/grids/categorical.py @@ -1,50 +1,77 @@ import dataclasses from collections.abc import Callable -from dataclasses import dataclass, is_dataclass +from dataclasses import dataclass, field, is_dataclass +import jax +import jax.numpy as jnp import pandas as pd -from lcm.exceptions import GridInitializationError, format_messages +from lcm.exceptions import ( + CategoricalDefinitionError, + GridInitializationError, + format_messages, +) +from lcm.typing import ScalarInt from lcm.utils.containers import find_duplicates, get_field_names_and_values def categorical[T](*, ordered: bool) -> Callable[[type[T]], type[T]]: - """Create a categorical class with auto-assigned integer values. + """Create a categorical class with auto-assigned `ScalarInt` values. - Transforms a class with int annotations into a frozen dataclass where each - field is assigned a consecutive integer value starting from 0. + Transforms a class with `ScalarInt`-annotated fields into a frozen + dataclass where each field is assigned a consecutive 0-d `jnp.int32` + scalar starting from 0. Decoration fails with + `CategoricalDefinitionError` if any field is annotated differently. Example: - @categorical(ordered=False) - class LaborSupply: - work: int - retire: int + from lcm.typing import ScalarInt - # Equivalent to: - @dataclass(frozen=True) + @categorical(ordered=False) class LaborSupply: - work: int = 0 - retire: int = 1 + work: ScalarInt + retire: ScalarInt - # Usage: - LaborSupply.work # 0 - LaborSupply.retire # 1 + LaborSupply.work # Array(0, dtype=int32) + LaborSupply.retire # Array(1, dtype=int32) + LaborSupply().work # Array(0, dtype=int32) Args: ordered: Whether the categories have a meaningful ordering. Must be explicitly specified. Returns: - A decorator that creates a frozen dataclass with auto-assigned integer values. + A decorator that creates a frozen dataclass with auto-assigned + `ScalarInt` values. """ def decorator(cls: type[T]) -> type[T]: annotations = getattr(cls, "__annotations__", {}) - # Assign sequential integers as defaults + bad_fields = { + name: annot for name, annot in annotations.items() if annot is not ScalarInt + } + if bad_fields: + details = ", ".join( + f"`{name}: {getattr(annot, '__name__', annot)}`" + for name, annot in bad_fields.items() + ) + raise CategoricalDefinitionError( + f"@categorical-decorated class {cls.__qualname__!r} must annotate " + f"every field as `ScalarInt` (the 0-d int32 scalar pylcm produces " + f"for category codes at runtime). The following fields are " + f"annotated otherwise: {details}. Import via " + f"`from lcm.typing import ScalarInt`." + ) + + # Mark fields as `init=False` with sequential int defaults so the + # generated `__init__` doesn't write per-instance values. The class + # attributes are then swapped to `jnp.int32` scalars post-decoration + # (frozen dataclasses reject `jax.Array` defaults directly under + # Python 3.14's mutable-default check, but class-level overrides via + # `type.__setattr__` are allowed). for i, name in enumerate(annotations): - setattr(cls, name, i) + setattr(cls, name, field(default=i, init=False)) cls._ordered = ordered # ty: ignore[unresolved-attribute] @@ -58,8 +85,15 @@ def _to_categorical_dtype(cls: type) -> pd.CategoricalDtype: cls.to_categorical_dtype = _to_categorical_dtype # ty: ignore[unresolved-attribute] - # Apply dataclass decorator - return dataclass(frozen=True)(cls) + new_cls = dataclass(frozen=True)(cls) + + # Class-level access (`X.foo`) and instance-level fall-through + # (`X().foo`, with `init=False` leaving instance __dict__ empty) + # both resolve to these `ScalarInt` scalars. + for i, name in enumerate(annotations): + type.__setattr__(new_cls, name, jnp.int32(i)) + + return new_cls return decorator @@ -70,13 +104,13 @@ def validate_category_class(category_class: type) -> list[str]: This validates that: - The class is a dataclass - It has at least one field - - All field values are int + - All field values are `ScalarInt` (0-d `jnp.int32` scalars) - All field values are unique - Field values are consecutive integers starting from 0 Args: category_class: The category class to validate. Must be a dataclass with fields - that have unique int values. + whose values are unique `ScalarInt` scalars. Returns: A list of error messages. Empty list if validation passes. @@ -86,7 +120,7 @@ def validate_category_class(category_class: type) -> list[str]: if not is_dataclass(category_class): error_messages.append( - "category_class must be a dataclass with int fields, " + "category_class must be a dataclass with `ScalarInt` fields, " f"but is {category_class}." ) return error_messages @@ -96,26 +130,29 @@ def validate_category_class(category_class: type) -> list[str]: if not names_and_values: error_messages.append("category_class must have at least one field.") - names_with_non_int_values = [ - name for name, value in names_and_values.items() if not isinstance(value, int) + names_with_bad_values = [ + name for name, value in names_and_values.items() if not _is_scalar_int(value) ] - if names_with_non_int_values: + if names_with_bad_values: error_messages.append( - "Field values of the category_class can only be int values. " - f"The values to the following fields are not: " - f"{names_with_non_int_values}" + "Field values of the category_class must be `ScalarInt` " + "(0-d int32 jax scalars). The values to the following fields are " + f"not: {names_with_bad_values}" ) + # The remaining checks coerce via `int(...)`; bail out if any value + # cannot be coerced cleanly. + return error_messages - values = list(names_and_values.values()) + values_as_py = [int(v) for v in names_and_values.values()] - duplicated_values = find_duplicates(values) + duplicated_values = find_duplicates(values_as_py) if duplicated_values: error_messages.append( "Field values of the category_class must be unique. " f"The following values are duplicated: {duplicated_values}" ) - if values != list(range(len(values))): + if values_as_py != list(range(len(values_as_py))): error_messages.append( "Field values of the category_class must be consecutive integers " "starting from 0 (e.g., 0, 1, 2, ...)." @@ -124,6 +161,15 @@ def validate_category_class(category_class: type) -> list[str]: return error_messages +def _is_scalar_int(value: object) -> bool: + """Return True iff `value` is a 0-d integer jax array (`ScalarInt`).""" + return ( + isinstance(value, jax.Array) + and value.shape == () + and jnp.issubdtype(value.dtype, jnp.integer) + ) + + def _validate_discrete_grid(category_class: type) -> None: """Validate the field names and values of the category_class passed to DiscreteGrid. diff --git a/src/lcm/grids/discrete.py b/src/lcm/grids/discrete.py index 72ded5eae..573c9d9b4 100644 --- a/src/lcm/grids/discrete.py +++ b/src/lcm/grids/discrete.py @@ -23,7 +23,11 @@ def __init__(self, category_class: type, batch_size: int = 0) -> None: _validate_discrete_grid(category_class) names_and_values = get_field_names_and_values(category_class) self.__categories = tuple(names_and_values.keys()) - self.__codes = tuple(names_and_values.values()) + # Coerce `ScalarInt` field values to Python `int` for the `codes` + # property. `codes` is the Python-side API (the tuple flows into + # dict/set operations that need hashable members); the JAX-side + # representation comes from `to_jax()`. + self.__codes = tuple(int(v) for v in names_and_values.values()) self.__ordered: bool = getattr(category_class, "_ordered", False) self.__batch_size: int = batch_size diff --git a/src/lcm/pandas_utils.py b/src/lcm/pandas_utils.py index adc04eac6..68f8ad25e 100644 --- a/src/lcm/pandas_utils.py +++ b/src/lcm/pandas_utils.py @@ -658,7 +658,7 @@ def _build_outcome_mapping( return _LevelMapping( name="next_regime", size=len(regime_names_to_ids), - get_code_from_label=regime_names_to_ids.__getitem__, + get_code_from_label=lambda label: int(regime_names_to_ids[label]), valid_labels=tuple(regime_names_to_ids), ) diff --git a/src/lcm/regime_building/processing.py b/src/lcm/regime_building/processing.py index a1488da44..e3017032d 100644 --- a/src/lcm/regime_building/processing.py +++ b/src/lcm/regime_building/processing.py @@ -1261,9 +1261,10 @@ def _wrap_regime_transition_probs( A wrapped function that returns MappingProxyType[str, float|Array]. """ - # Get regime names in index order from regime_names_to_ids + # Get regime names in index order from regime_names_to_ids. Coerce + # `ScalarInt` ids to Python `int` so `sorted` has a comparable key. regime_names_by_id: list[tuple[int, str]] = sorted( - [(idx, name) for name, idx in regime_names_to_ids.items()], + [(int(idx), name) for name, idx in regime_names_to_ids.items()], key=lambda x: x[0], ) regime_names = [name for _, name in regime_names_by_id] diff --git a/src/lcm/simulation/initial_conditions.py b/src/lcm/simulation/initial_conditions.py index 3db3063ee..db713eb76 100644 --- a/src/lcm/simulation/initial_conditions.py +++ b/src/lcm/simulation/initial_conditions.py @@ -31,6 +31,7 @@ RegimeName, RegimeNamesToIds, ) +from lcm.utils.containers import invert_regime_ids from lcm.utils.functools import get_union_of_args # Sentinel for categorical states not in initial conditions. Using int32 min @@ -119,10 +120,10 @@ def validate_initial_conditions( InvalidInitialConditionsError: If any validation check fails. """ - # Build reverse lookup from regime IDs to names - ids_to_regime_names: dict[int, RegimeName] = { - v: k for k, v in regime_names_to_ids.items() - } + # Build reverse lookup from regime IDs to names. `regime_names_to_ids` + # values are `ScalarInt` (jax 0-d arrays), which can't serve as dict + # keys directly; `invert_regime_ids` coerces them to Python `int`. + ids_to_regime_names = invert_regime_ids(regime_names_to_ids) # Extract regime array regime_arr = initial_conditions.get("regime") @@ -227,7 +228,7 @@ def _collect_state_name_errors( *, initial_states: Mapping[str, Array], regime_id_arr: Array, - ids_to_regime_names: dict[int, RegimeName], + ids_to_regime_names: MappingProxyType[int, RegimeName], internal_regimes: MappingProxyType[RegimeName, InternalRegime], valid_regime_names: set[RegimeName], ) -> list[str]: @@ -285,7 +286,7 @@ def _collect_structural_errors( *, initial_states: Mapping[str, Array], regime_id_arr: Array, - ids_to_regime_names: dict[int, RegimeName], + ids_to_regime_names: MappingProxyType[int, RegimeName], regime_names_to_ids: RegimeNamesToIds, internal_regimes: MappingProxyType[RegimeName, InternalRegime], ages: AgeGrid, @@ -437,7 +438,7 @@ def _validate_discrete_state_values( initial_states: Mapping[str, Array], internal_regimes: MappingProxyType[RegimeName, InternalRegime], regime_id_arr: Array, - regime_names_to_ids: Mapping[RegimeName, int], + regime_names_to_ids: RegimeNamesToIds, ) -> None: """Validate that discrete state values are valid codes. @@ -454,10 +455,12 @@ def _validate_discrete_state_values( InvalidInitialConditionsError: If any discrete state contains invalid codes. """ - # Build per-state: valid codes + regime IDs that have this state + # Build per-state: valid codes + regime IDs that have this state. + # `regime_id` is `ScalarInt` (a 0-d jax array); coerce to Python `int` + # before set insertion. discrete_info: dict[str, tuple[set[int], set[int]]] = {} for regime_name, internal_regime in internal_regimes.items(): - regime_id = regime_names_to_ids[regime_name] + regime_id = int(regime_names_to_ids[regime_name]) for state_name in internal_regime.variable_info.query( "is_state and is_discrete" ).index: diff --git a/src/lcm/simulation/result.py b/src/lcm/simulation/result.py index ed501739e..dfbbde96e 100644 --- a/src/lcm/simulation/result.py +++ b/src/lcm/simulation/result.py @@ -28,6 +28,7 @@ FloatND, InternalParams, RegimeName, + RegimeNamesToIds, StateName, UserFunction, ) @@ -226,7 +227,7 @@ def __repr__(self) -> str: def get_simulation_output_dtypes( regimes: Mapping[RegimeName, Regime], - regime_names_to_ids: Mapping[RegimeName, int], + regime_names_to_ids: RegimeNamesToIds, ) -> MappingProxyType[str, pd.CategoricalDtype]: """Compute pandas CategoricalDtype for all discrete output columns. diff --git a/src/lcm/simulation/simulate.py b/src/lcm/simulation/simulate.py index 120406234..395a87b18 100644 --- a/src/lcm/simulation/simulate.py +++ b/src/lcm/simulation/simulate.py @@ -34,6 +34,7 @@ ScalarFloat, ScalarInt, ) +from lcm.utils.containers import invert_regime_ids from lcm.utils.error_handling import validate_V from lcm.utils.logging import ( format_duration, @@ -109,8 +110,10 @@ def simulate( simulation_results: dict[RegimeName, dict[int, PeriodRegimeSimulationData]] = { regime_name: {} for regime_name in internal_regimes } - # Build reverse lookup for regime transition logging - ids_to_names: dict[int, RegimeName] = {v: k for k, v in regime_names_to_ids.items()} + # Build reverse lookup for regime transition logging. `regime_names_to_ids` + # values are `ScalarInt` (jax 0-d arrays), which can't serve as dict keys + # directly; `invert_regime_ids` coerces them to Python `int`. + ids_to_names = invert_regime_ids(regime_names_to_ids) for period, age in enumerate(ages.values): period_start = time.monotonic() @@ -209,7 +212,7 @@ def _simulate_regime_in_period( int, MappingProxyType[RegimeName, FloatND] ], internal_params: InternalParams, - regime_names_to_ids: MappingProxyType[RegimeName, int], + regime_names_to_ids: RegimeNamesToIds, active_regimes_next_period: tuple[RegimeName, ...], key: Array, ) -> tuple[PeriodRegimeSimulationData, MappingProxyType[str, Array], Int1D, Array]: diff --git a/src/lcm/simulation/transitions.py b/src/lcm/simulation/transitions.py index 82c5f8ca5..3742cb8b0 100644 --- a/src/lcm/simulation/transitions.py +++ b/src/lcm/simulation/transitions.py @@ -152,7 +152,7 @@ def calculate_next_regime_membership( period: int, age: ScalarInt | ScalarFloat, regime_params: FlatRegimeParams, - regime_names_to_ids: MappingProxyType[RegimeName, int], + regime_names_to_ids: RegimeNamesToIds, new_subject_regime_ids: Int1D, active_regimes_next_period: tuple[RegimeName, ...], key: Array, diff --git a/src/lcm/typing.py b/src/lcm/typing.py index 9553ebea9..0aa045d41 100644 --- a/src/lcm/typing.py +++ b/src/lcm/typing.py @@ -36,7 +36,7 @@ type ShockName = str type FunctionName = str type TransitionFunctionName = str -type RegimeNamesToIds = MappingProxyType[RegimeName, int] +type RegimeNamesToIds = MappingProxyType[RegimeName, ScalarInt] type FunctionsMapping = MappingProxyType[FunctionName, InternalUserFunction] diff --git a/src/lcm/utils/containers.py b/src/lcm/utils/containers.py index 0df253064..0c1682d14 100644 --- a/src/lcm/utils/containers.py +++ b/src/lcm/utils/containers.py @@ -84,6 +84,19 @@ def get_field_names_and_values(dc: type) -> MappingProxyType[str, Any]: ) +def invert_regime_ids[K](mapping: Mapping[K, Any]) -> MappingProxyType[int, K]: + """Return the inverse of a regime-name → id mapping, with Python-`int` keys. + + `@categorical` assigns `jnp.int32` scalars to class attributes, so + `regime_names_to_ids` values are 0-d jax arrays — which aren't + hashable and therefore can't serve as `dict` keys. This helper + coerces each value to a Python `int` so the inverted lookup + (`id → name`) is usable wherever JAX promotion isn't already + happening. + """ + return MappingProxyType({int(v): k for k, v in mapping.items()}) + + def first_non_none(*args: T | None) -> T: """Return the first non-None argument. diff --git a/src/lcm_examples/mahler_yum_2024/_model.py b/src/lcm_examples/mahler_yum_2024/_model.py index 67fc0b0d8..a237877d9 100644 --- a/src/lcm_examples/mahler_yum_2024/_model.py +++ b/src/lcm_examples/mahler_yum_2024/_model.py @@ -35,6 +35,7 @@ FloatND, Int1D, Period, + ScalarInt, ) from lcm.utils.dispatchers import productmap @@ -67,59 +68,62 @@ def calc_savingsgrid(x: Float1D) -> Float1D: @categorical(ordered=True) class LaborSupply: - do_not_work: int - part_time: int - full_time: int + do_not_work: ScalarInt + part_time: ScalarInt + full_time: ScalarInt @categorical(ordered=True) class Education: - low: int - high: int + low: ScalarInt + high: ScalarInt Effort = make_dataclass( - "HealthEffort", [("class" + str(i), int, int(i)) for i in range(40)] + "HealthEffort", + [(f"class{i}", ScalarInt) for i in range(40)], ) +for i in range(40): + type.__setattr__(Effort, f"class{i}", jnp.int32(i)) @categorical(ordered=True) class Health: - bad: int - good: int + bad: ScalarInt + good: ScalarInt @categorical(ordered=True) class ProductivityType: - low: int - high: int + low: ScalarInt + high: ScalarInt @categorical(ordered=True) class HealthType: - low: int - high: int + low: ScalarInt + high: ScalarInt @categorical(ordered=True) class ProductivityShock: - val0: int - val1: int - val2: int - val3: int - val4: int + val0: ScalarInt + val1: ScalarInt + val2: ScalarInt + val3: ScalarInt + val4: ScalarInt @categorical(ordered=True) class DiscountType: - small: int - large: int + small: ScalarInt + large: ScalarInt @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def discount_factor( diff --git a/src/lcm_examples/mortality.py b/src/lcm_examples/mortality.py index e58a9540b..4681b8d5d 100644 --- a/src/lcm_examples/mortality.py +++ b/src/lcm_examples/mortality.py @@ -25,20 +25,21 @@ DiscreteAction, FloatND, Period, + ScalarInt, ) @categorical(ordered=True) class LaborSupply: - work: int - retire: int + work: ScalarInt + retire: ScalarInt @categorical(ordered=False) class RegimeId: - working_life: int - retirement: int - dead: int + working_life: ScalarInt + retirement: ScalarInt + dead: ScalarInt def utility_working( diff --git a/src/lcm_examples/precautionary_savings.py b/src/lcm_examples/precautionary_savings.py index 44a94a959..57b41dbaf 100644 --- a/src/lcm_examples/precautionary_savings.py +++ b/src/lcm_examples/precautionary_savings.py @@ -74,8 +74,8 @@ def utility(consumption: ContinuousAction) -> FloatND: @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt @functools.cache diff --git a/src/lcm_examples/precautionary_savings_health.py b/src/lcm_examples/precautionary_savings_health.py index 2a0c3dcc9..08e7a2aef 100644 --- a/src/lcm_examples/precautionary_savings_health.py +++ b/src/lcm_examples/precautionary_savings_health.py @@ -23,14 +23,14 @@ @categorical(ordered=True) class LaborSupply: - do_not_work: int - work: int + do_not_work: ScalarInt + work: ScalarInt @categorical(ordered=False) class RegimeId: - working_life: int - retirement: int + working_life: ScalarInt + retirement: ScalarInt def utility( diff --git a/src/lcm_examples/tiny.py b/src/lcm_examples/tiny.py index a0982c235..0b310f6a6 100644 --- a/src/lcm_examples/tiny.py +++ b/src/lcm_examples/tiny.py @@ -30,14 +30,14 @@ @categorical(ordered=True) class LaborSupply: - do_not_work: int - work: int + do_not_work: ScalarInt + work: ScalarInt @categorical(ordered=False) class RegimeId: - working_life: int - retirement: int + working_life: ScalarInt + retirement: ScalarInt def utility( diff --git a/tests/conftest.py b/tests/conftest.py index b889345f2..2bc510b69 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,9 +1,12 @@ from collections.abc import Iterator from dataclasses import make_dataclass +import jax.numpy as jnp import pytest from jax import config as jax_config +from lcm.typing import ScalarInt + # Module-level precision settings (updated by pytest_configure based on --precision) X64_ENABLED: bool = True # 12 decimals (not 14): CI showed that 14 exceeded reproducible machine precision @@ -35,7 +38,12 @@ def pytest_configure(config): @pytest.fixture(scope="session") def binary_category_class(): - return make_dataclass("BinaryCategoryClass", [("cat0", int, 0), ("cat1", int, 1)]) + cls = make_dataclass( + "BinaryCategoryClass", [("cat0", ScalarInt), ("cat1", ScalarInt)] + ) + type.__setattr__(cls, "cat0", jnp.int32(0)) + type.__setattr__(cls, "cat1", jnp.int32(1)) + return cls @pytest.fixture(name="x64_disabled") diff --git a/tests/regime_building/test_regime_processing.py b/tests/regime_building/test_regime_processing.py index b2def9424..6801c9e10 100644 --- a/tests/regime_building/test_regime_processing.py +++ b/tests/regime_building/test_regime_processing.py @@ -17,6 +17,7 @@ process_regimes, ) from lcm.regime_building.variable_info import get_grids, get_variable_info +from lcm.typing import ScalarInt from tests.regime_mock import RegimeMock from tests.test_models.deterministic.base import dead, working_life @@ -112,7 +113,7 @@ def test_process_regimes(): ages = AgeGrid(start=0, stop=4, step="Y") regimes = {"working_life": working_life, "dead": dead} regime_names_to_ids = MappingProxyType( - {name: idx for idx, name in enumerate(regimes.keys())} + {name: jnp.int32(idx) for idx, name in enumerate(regimes.keys())} ) internal_regimes = process_regimes( regimes=regimes, @@ -192,8 +193,8 @@ def regime_transition(age, final_age): @categorical(ordered=False) class TwoRegimeId: - early: int - late: int + early: ScalarInt + late: ScalarInt early = Regime( transition=regime_transition, @@ -212,7 +213,9 @@ class TwoRegimeId: return process_regimes( regimes={"early": early, "late": late}, ages=AgeGrid(start=0, stop=2, step="Y"), - regime_names_to_ids=MappingProxyType({"early": 0, "late": 1}), + regime_names_to_ids=MappingProxyType( + {"early": jnp.int32(0), "late": jnp.int32(1)} + ), enable_jit=True, ) diff --git a/tests/simulation/test_initial_conditions.py b/tests/simulation/test_initial_conditions.py index 882fba505..bd8691bd6 100644 --- a/tests/simulation/test_initial_conditions.py +++ b/tests/simulation/test_initial_conditions.py @@ -1,7 +1,5 @@ """Tests for initial conditions validation utilities.""" -from dataclasses import dataclass - import jax.numpy as jnp import pytest @@ -31,15 +29,15 @@ def model() -> Model: """Minimal model with two states (wealth, health) for initial states tests.""" - @dataclass + @categorical(ordered=False) class Health: - healthy: int = 0 - sick: int = 1 + healthy: ScalarInt + sick: ScalarInt - @dataclass + @categorical(ordered=False) class RegimeId: - active: int = 0 - terminal: int = 1 + active: ScalarInt + terminal: ScalarInt def utility(wealth: ContinuousState, health: DiscreteState) -> FloatND: # noqa: ARG001 return jnp.array(0.0) @@ -259,8 +257,8 @@ def _make_constraint_model(wealth_grid) -> Model: @categorical(ordered=False) class RegimeId: - working_life: int - dead: int + working_life: ScalarInt + dead: ScalarInt def utility(consumption: ContinuousAction) -> FloatND: return jnp.log(consumption) @@ -443,8 +441,8 @@ def _make_constrained_asymmetric_model() -> Model: @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def utility(consumption: ContinuousAction) -> FloatND: return jnp.log(consumption) @@ -500,15 +498,15 @@ def _make_asymmetric_state_model() -> Model: Ages 0, 1, 2. alive is active for age < 2, dead is active for age >= 2. """ - @dataclass + @categorical(ordered=False) class Health: - healthy: int = 0 - sick: int = 1 + healthy: ScalarInt + sick: ScalarInt @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def utility( wealth: ContinuousState, diff --git a/tests/simulation/test_simulate.py b/tests/simulation/test_simulate.py index cb434a72a..66d62f225 100644 --- a/tests/simulation/test_simulate.py +++ b/tests/simulation/test_simulate.py @@ -39,7 +39,7 @@ def simulate_inputs(): ) regimes = {"working_life": updated_working_life, "dead": dead} regime_names_to_ids = MappingProxyType( - {name: idx for idx, name in enumerate(regimes.keys())} + {name: jnp.int32(idx) for idx, name in enumerate(regimes.keys())} ) internal_regimes = process_regimes( regimes=regimes, diff --git a/tests/solution/test_beta_delta.py b/tests/solution/test_beta_delta.py index 2d30fb6c9..d36b3d898 100644 --- a/tests/solution/test_beta_delta.py +++ b/tests/solution/test_beta_delta.py @@ -34,8 +34,8 @@ @categorical(ordered=False) class RegimeId: - working: int - dead: int + working: ScalarInt + dead: ScalarInt def utility(consumption: ContinuousAction) -> FloatND: diff --git a/tests/solution/test_custom_aggregator.py b/tests/solution/test_custom_aggregator.py index 6dcd227c0..1d5dc4166 100644 --- a/tests/solution/test_custom_aggregator.py +++ b/tests/solution/test_custom_aggregator.py @@ -19,14 +19,14 @@ @categorical(ordered=False) class LaborSupply: - work: int - retire: int + work: ScalarInt + retire: ScalarInt @categorical(ordered=False) class RegimeId: - working_life: int - dead: int + working_life: ScalarInt + dead: ScalarInt def utility( @@ -82,9 +82,9 @@ def ces_H( @categorical(ordered=False) class PrefType: - type_0: int - type_1: int - type_2: int + type_0: ScalarInt + type_1: ScalarInt + type_2: ScalarInt def discount_factor_from_type( diff --git a/tests/solution/test_diagnostics.py b/tests/solution/test_diagnostics.py index 262ae0e12..2310e9ee8 100644 --- a/tests/solution/test_diagnostics.py +++ b/tests/solution/test_diagnostics.py @@ -16,13 +16,13 @@ from lcm import AgeGrid, LinSpacedGrid, Model, Regime, categorical from lcm.exceptions import InvalidValueFunctionError -from lcm.typing import ContinuousAction, ContinuousState, FloatND +from lcm.typing import ContinuousAction, ContinuousState, FloatND, ScalarInt @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def _utility(consumption: ContinuousAction, wealth: ContinuousState) -> FloatND: diff --git a/tests/test_Q_and_F.py b/tests/test_Q_and_F.py index 2908f4e60..7d6a25fdf 100644 --- a/tests/test_Q_and_F.py +++ b/tests/test_Q_and_F.py @@ -45,7 +45,7 @@ def test_get_Q_and_F_function(): ages = AgeGrid(start=0, stop=4, step="Y") regimes = {"working_life": working_life, "dead": dead} regime_names_to_ids = MappingProxyType( - {name: idx for idx, name in enumerate(regimes.keys())} + {name: jnp.int32(idx) for idx, name in enumerate(regimes.keys())} ) internal_regimes = process_regimes( regimes=regimes, @@ -267,15 +267,15 @@ def _health_probs(health: DiscreteState, probs_array: FloatND) -> FloatND: @categorical(ordered=True) class _IncompleteTargetHealth: - bad: int = 0 - good: int = 1 + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class _IncompleteTargetRegimeId: - work: int - retire: int - dead: int + work: ScalarInt + retire: ScalarInt + dead: ScalarInt def _build_incomplete_target_model( diff --git a/tests/test_chained_state_transitions.py b/tests/test_chained_state_transitions.py index df1954aca..8bc4dc645 100644 --- a/tests/test_chained_state_transitions.py +++ b/tests/test_chained_state_transitions.py @@ -15,14 +15,14 @@ @categorical(ordered=False) class _LaborSupply: - work: int - rest: int + work: ScalarInt + rest: ScalarInt @categorical(ordered=False) class _RegimeId: - active: int - dead: int + active: ScalarInt + dead: ScalarInt def _next_aime(aime: float, labor_supply: DiscreteAction) -> FloatND: diff --git a/tests/test_error_handling_invalid_vf.py b/tests/test_error_handling_invalid_vf.py index e59dfb0fb..9b9fe1d39 100644 --- a/tests/test_error_handling_invalid_vf.py +++ b/tests/test_error_handling_invalid_vf.py @@ -22,8 +22,8 @@ def n_periods() -> int: @categorical(ordered=False) class RegimeId: - non_terminal: int - terminal: int + non_terminal: ScalarInt + terminal: ScalarInt @pytest.fixture diff --git a/tests/test_function_output_grid_indexing.py b/tests/test_function_output_grid_indexing.py index adc378161..86c9e3890 100644 --- a/tests/test_function_output_grid_indexing.py +++ b/tests/test_function_output_grid_indexing.py @@ -16,19 +16,25 @@ from lcm import AgeGrid, DiscreteGrid, LinSpacedGrid, Model, Regime, categorical from lcm.exceptions import RegimeInitializationError -from lcm.typing import ContinuousAction, DiscreteAction, DiscreteState, FloatND +from lcm.typing import ( + ContinuousAction, + DiscreteAction, + DiscreteState, + FloatND, + ScalarInt, +) @categorical(ordered=False) class PrefType: - type_0: int - type_1: int + type_0: ScalarInt + type_1: ScalarInt @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def _per_type_scale_takes_pref_type( @@ -163,8 +169,8 @@ def test_function_output_indexed_by_derived_categorical_raises(): @categorical(ordered=False) class IsMarried: - single: int - married: int + single: ScalarInt + married: ScalarInt def _is_married(spousal_income: DiscreteState) -> DiscreteState: return jnp.int32(spousal_income > 0) @@ -181,9 +187,9 @@ def _utility_clash( @categorical(ordered=True) class SpousalIncome: - single: int - married_no_inc: int - married_has_inc: int + single: ScalarInt + married_no_inc: ScalarInt + married_has_inc: ScalarInt with pytest.raises( RegimeInitializationError, @@ -209,8 +215,8 @@ def test_function_output_indexed_by_discrete_action_raises(): @categorical(ordered=False) class WorkChoice: - no_work: int - work: int + no_work: ScalarInt + work: ScalarInt def _per_choice_scale(labor_supply: DiscreteAction, some_param: FloatND) -> FloatND: return jnp.abs(1.0 / (1.0 - some_param[labor_supply])) diff --git a/tests/test_function_representation.py b/tests/test_function_representation.py index 12f55d8f5..4c6827a89 100644 --- a/tests/test_function_representation.py +++ b/tests/test_function_representation.py @@ -15,14 +15,16 @@ _get_lookup_function, get_V_interpolator, ) +from lcm.typing import ScalarInt from lcm.utils.dispatchers import productmap @pytest.fixture def binary_discrete_grid(): - return DiscreteGrid( - make_dataclass("BinaryCategory", [("a", bool, False), ("b", bool, True)]) - ) + cls = make_dataclass("BinaryCategory", [("a", ScalarInt), ("b", ScalarInt)]) + type.__setattr__(cls, "a", jnp.int32(0)) + type.__setattr__(cls, "b", jnp.int32(1)) + return DiscreteGrid(cls) @pytest.fixture diff --git a/tests/test_grids.py b/tests/test_grids.py index cd9cbfb17..ec652503c 100644 --- a/tests/test_grids.py +++ b/tests/test_grids.py @@ -6,7 +6,7 @@ import pytest from numpy.testing import assert_array_almost_equal as aaae -from lcm.exceptions import GridInitializationError +from lcm.exceptions import CategoricalDefinitionError, GridInitializationError from lcm.grids import ( DiscreteGrid, IrregSpacedGrid, @@ -20,10 +20,26 @@ ) from lcm.grids.categorical import _validate_discrete_grid from lcm.grids.continuous import _validate_continuous_grid +from lcm.typing import ScalarInt from lcm.utils.containers import get_field_names_and_values from tests.conftest import DECIMAL_PRECISION, X64_ENABLED +def _make_dc(name: str, *fields: tuple[str, object]) -> type: + """Build a dataclass with `ScalarInt` class attrs. + + Sidesteps `@dataclass(frozen=True)`'s rejection of `jax.Array` + defaults: we make the dataclass with no defaults, then write the + `ScalarInt` scalars onto the class via `type.__setattr__`. + `validate_category_class` reads field values via `getattr(cls, name)`, + so the class attrs are what get checked. + """ + cls = make_dataclass(name, [(fname, ScalarInt) for fname, _ in fields]) + for fname, fval in fields: + type.__setattr__(cls, fname, fval) + return cls + + def test_get_fields_with_defaults(): category_class = make_dataclass("Category", [("a", int, 1), ("b", int, 2)]) assert get_field_names_and_values(category_class) == {"a": 1, "b": 2} @@ -52,41 +68,43 @@ def test_validate_discrete_grid_empty(): def test_validate_discrete_grid_non_scalar_input(): - category_class = make_dataclass("Category", [("a", int, 1), ("b", str, "s")]) + category_class = _make_dc("Category", ("a", jnp.int32(1)), ("b", "s")) error_msg = ( - "Field values of the category_class can only be int " - r"values. The values to the following fields are not: \['b'\]" + "Field values of the category_class must be `ScalarInt` " + r"\(0-d int32 jax scalars\). The values to the following " + r"fields are not: \['b'\]" ) with pytest.raises(GridInitializationError, match=error_msg): _validate_discrete_grid(category_class) def test_validate_discrete_grid_none_input(): - category_class = make_dataclass("Category", [("a", int), ("b", int, 1)]) + category_class = _make_dc("Category", ("a", None), ("b", jnp.int32(1))) error_msg = ( - "Field values of the category_class can only be int " - r"values. The values to the following fields are not: \['a'\]" + "Field values of the category_class must be `ScalarInt` " + r"\(0-d int32 jax scalars\). The values to the following " + r"fields are not: \['a'\]" ) with pytest.raises(GridInitializationError, match=error_msg): _validate_discrete_grid(category_class) def test_validate_discrete_grid_non_unique(): - category_class = make_dataclass("Category", [("a", int, 1), ("b", int, 1)]) + category_class = _make_dc("Category", ("a", jnp.int32(1)), ("b", jnp.int32(1))) error_msg = "Field values of the category_class must be unique." with pytest.raises(GridInitializationError, match=error_msg): _validate_discrete_grid(category_class) def test_validate_discrete_grid_non_consecutive_unordered(): - category_class = make_dataclass("Category", [("a", int, 1), ("b", int, 0)]) + category_class = _make_dc("Category", ("a", jnp.int32(1)), ("b", jnp.int32(0))) error_msg = "Field values of the category_class must be consecutive integers" with pytest.raises(GridInitializationError, match=error_msg): _validate_discrete_grid(category_class) def test_validate_discrete_grid_non_consecutive_jumps(): - category_class = make_dataclass("Category", [("a", int, 0), ("b", int, 2)]) + category_class = _make_dc("Category", ("a", jnp.int32(0)), ("b", jnp.int32(2))) error_msg = "Field values of the category_class must be consecutive integers" with pytest.raises(GridInitializationError, match=error_msg): _validate_discrete_grid(category_class) @@ -94,7 +112,7 @@ def test_validate_discrete_grid_non_consecutive_jumps(): def test_validate_category_class_valid(): """Valid category class should return empty error list.""" - category_class = make_dataclass("Category", [("a", int, 0), ("b", int, 1)]) + category_class = _make_dc("Category", ("a", jnp.int32(0)), ("b", jnp.int32(1))) errors = validate_category_class(category_class) assert errors == [] @@ -113,7 +131,7 @@ class NotDataclass: def test_validate_category_class_non_consecutive(): """Non-consecutive values should return error.""" - category_class = make_dataclass("Category", [("a", int, 0), ("b", int, 2)]) + category_class = _make_dc("Category", ("a", jnp.int32(0)), ("b", jnp.int32(2))) errors = validate_category_class(category_class) assert len(errors) == 1 assert "consecutive integers" in errors[0] @@ -121,27 +139,28 @@ def test_validate_category_class_non_consecutive(): def test_validate_category_class_not_starting_at_zero(): """Values not starting at 0 should return error.""" - category_class = make_dataclass("Category", [("a", int, 1), ("b", int, 2)]) + category_class = _make_dc("Category", ("a", jnp.int32(1)), ("b", jnp.int32(2))) errors = validate_category_class(category_class) assert len(errors) == 1 assert "consecutive integers starting from 0" in errors[0] def test_discrete_grid_creation(): - category_class = make_dataclass( - "Category", [("a", int, 0), ("b", int, 1), ("c", int, 2)] + category_class = _make_dc( + "Category", + ("a", jnp.int32(0)), + ("b", jnp.int32(1)), + ("c", jnp.int32(2)), ) grid = DiscreteGrid(category_class) assert np.allclose(grid.to_jax(), np.arange(3)) def test_discrete_grid_invalid_category_class(): - category_class = make_dataclass( - "Category", [("a", int, 0), ("b", str, "wrong_type")] - ) + category_class = _make_dc("Category", ("a", jnp.int32(0)), ("b", "wrong_type")) with pytest.raises( GridInitializationError, - match="Field values of the category_class can only be int", + match="Field values of the category_class must be `ScalarInt`", ): DiscreteGrid(category_class) @@ -149,9 +168,9 @@ def test_discrete_grid_invalid_category_class(): def test_discrete_grid_ordered_true(): @categorical(ordered=True) class OrderedCat: - low: int - medium: int - high: int + low: ScalarInt + medium: ScalarInt + high: ScalarInt grid = DiscreteGrid(OrderedCat) assert grid.ordered is True @@ -160,13 +179,78 @@ class OrderedCat: def test_discrete_grid_ordered_false(): @categorical(ordered=False) class UnorderedCat: - a: int - b: int + a: ScalarInt + b: ScalarInt grid = DiscreteGrid(UnorderedCat) assert grid.ordered is False +# --- @categorical: ScalarInt annotation contract --- + + +def test_categorical_rejects_int_annotation(): + """Plain `int` annotations are rejected at decoration time.""" + with pytest.raises( + CategoricalDefinitionError, + match=r"must annotate every field as `ScalarInt`", + ): + # `a: int` MUST stay — this is the rejected case under test. + @categorical(ordered=False) + class _Bad: + a: int + + +def test_categorical_rejects_str_annotation(): + """Non-`ScalarInt` annotations of any kind are rejected.""" + with pytest.raises( + CategoricalDefinitionError, + match=r"must annotate every field as `ScalarInt`", + ): + + @categorical(ordered=False) + class _Bad: + a: str + + +def test_categorical_error_lists_all_offending_fields(): + """The error message names every field with a wrong annotation.""" + with pytest.raises(CategoricalDefinitionError, match=r"`x: int`.*`y: str`"): + # `x: int` / `y: str` MUST stay — both are rejected cases under test. + @categorical(ordered=False) + class _Bad: + x: int + y: str + + +def test_categorical_class_attr_is_scalar_int(): + """Class-level access returns a 0-d int32 jax scalar.""" + + @categorical(ordered=False) + class Cat: + first: ScalarInt + second: ScalarInt + + assert Cat.first.shape == () + assert Cat.first.dtype == jnp.int32 + assert int(Cat.first) == 0 + assert int(Cat.second) == 1 + + +def test_categorical_instance_attr_is_scalar_int(): + """Instance-level access also returns a 0-d int32 jax scalar.""" + + @categorical(ordered=False) + class Cat: + first: ScalarInt + second: ScalarInt + + instance = Cat() + assert instance.first.shape == () + assert instance.first.dtype == jnp.int32 + assert int(instance.second) == 1 + + def test_lin_spaced_grid_rejects_non_numeric_start(): """Non-numeric `start` is rejected at the boundary cast.""" with pytest.raises((TypeError, ValueError)): diff --git a/tests/test_model.py b/tests/test_model.py index cd672710b..8897fe431 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -169,7 +169,7 @@ def test_model_requires_terminal_regime(binary_category_class): @categorical(ordered=False) class RegimeId: - test: int + test: ScalarInt regime = Regime( states={ @@ -194,7 +194,7 @@ def test_model_requires_non_terminal_regime(binary_category_class): @categorical(ordered=False) class RegimeId: - dead: int + dead: ScalarInt dead = Regime( transition=None, @@ -224,9 +224,9 @@ def test_model_accepts_multiple_terminal_regimes(binary_category_class): @categorical(ordered=False) class RegimeId: - alive: int - dead1: int - dead2: int + alive: ScalarInt + dead1: ScalarInt + dead2: ScalarInt alive = Regime( states={ @@ -267,8 +267,8 @@ def test_model_regime_id_mapping_created_from_dict_keys(binary_category_class): @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt alive = Regime( states={ @@ -302,8 +302,8 @@ def test_model_regime_name_validation(binary_category_class): @categorical(ordered=False) class RegimeId: - alive__bad: int - dead: int + alive__bad: ScalarInt + dead: ScalarInt alive = Regime( states={ @@ -336,14 +336,14 @@ def test_unused_state_raises_error(): @categorical(ordered=False) class RegimeId: - working_life: int - retirement: int + working_life: ScalarInt + retirement: ScalarInt @categorical(ordered=False) class UnusedState: - low: int - medium: int - high: int + low: ScalarInt + medium: ScalarInt + high: ScalarInt # Define a regime where 'unused_state' is not used in any function working_life = Regime( @@ -393,13 +393,13 @@ def test_unused_action_raises_error(): @categorical(ordered=False) class RegimeId: - working_life: int - retirement: int + working_life: ScalarInt + retirement: ScalarInt @categorical(ordered=False) class UnusedAction: - option_a: int - option_b: int + option_a: ScalarInt + option_b: ScalarInt working_life = Regime( functions={ @@ -452,18 +452,18 @@ def test_constraint_depending_on_transition_output(): @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt @categorical(ordered=False) class EmploymentLastPeriod: - unemployed: int - employed: int + unemployed: ScalarInt + employed: ScalarInt @categorical(ordered=False) class EmploymentStatus: - not_employed: int - employed: int + not_employed: ScalarInt + employed: ScalarInt def next_regime(age: float, model_end_age: int) -> ScalarInt: return jnp.where(age == model_end_age, RegimeId.dead, RegimeId.alive) @@ -533,18 +533,18 @@ def test_state_only_used_in_transitions(): @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt @categorical(ordered=False) class EmploymentLastPeriod: - unemployed: int - employed: int + unemployed: ScalarInt + employed: ScalarInt @categorical(ordered=False) class EmploymentStatus: - not_employed: int - employed: int + not_employed: ScalarInt + employed: ScalarInt def next_regime(age: float, model_end_age: int) -> ScalarInt: return jnp.where(age == model_end_age, RegimeId.dead, RegimeId.alive) @@ -617,13 +617,13 @@ def test_state_only_in_transitions_with_terminal_regime(): @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt @categorical(ordered=False) class TypeVar: - low: int - high: int + low: ScalarInt + high: ScalarInt def utility(consumption, wealth): return jnp.log(consumption) + 0.01 * wealth diff --git a/tests/test_models/basic_discrete.py b/tests/test_models/basic_discrete.py index 4f7cfab37..e5f9f4a1c 100644 --- a/tests/test_models/basic_discrete.py +++ b/tests/test_models/basic_discrete.py @@ -1,22 +1,23 @@ """Basic model with discrete + continuous states, no stochastic transitions.""" from lcm import AgeGrid, DiscreteGrid, LinSpacedGrid, Model, Regime, categorical +from lcm.typing import ScalarInt @categorical(ordered=True) class Health: - bad: int - good: int + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class RegimeId: - working_life: int - retirement: int - dead: int + working_life: ScalarInt + retirement: ScalarInt + dead: ScalarInt -def _next_regime() -> int: +def _next_regime() -> ScalarInt: return RegimeId.dead diff --git a/tests/test_models/deterministic/base.py b/tests/test_models/deterministic/base.py index 1670ff7b4..e4f3c6aa7 100644 --- a/tests/test_models/deterministic/base.py +++ b/tests/test_models/deterministic/base.py @@ -27,9 +27,9 @@ @categorical(ordered=False) class RegimeId: - working_life: int - retirement: int - dead: int + working_life: ScalarInt + retirement: ScalarInt + dead: ScalarInt def next_regime_from_working( diff --git a/tests/test_models/deterministic/discrete.py b/tests/test_models/deterministic/discrete.py index 8f6537d21..52247d33c 100644 --- a/tests/test_models/deterministic/discrete.py +++ b/tests/test_models/deterministic/discrete.py @@ -32,21 +32,21 @@ @categorical(ordered=True) class DiscreteConsumption: - low: int - high: int + low: ScalarInt + high: ScalarInt @categorical(ordered=True) class DiscreteWealth: - low: int - medium: int - high: int + low: ScalarInt + medium: ScalarInt + high: ScalarInt @categorical(ordered=False) class RegimeId: - working_life: int - dead: int + working_life: ScalarInt + dead: ScalarInt def utility_discrete( diff --git a/tests/test_models/deterministic/regression.py b/tests/test_models/deterministic/regression.py index b845aa374..3b213f668 100644 --- a/tests/test_models/deterministic/regression.py +++ b/tests/test_models/deterministic/regression.py @@ -37,8 +37,8 @@ @categorical(ordered=False) class RegimeId: - working_life: int - dead: int + working_life: ScalarInt + dead: ScalarInt def wage(age: float) -> float | FloatND: diff --git a/tests/test_models/regime_markov.py b/tests/test_models/regime_markov.py index 3e78815d0..6f755d75e 100644 --- a/tests/test_models/regime_markov.py +++ b/tests/test_models/regime_markov.py @@ -9,19 +9,19 @@ Regime, categorical, ) -from lcm.typing import DiscreteState, FloatND, Period +from lcm.typing import DiscreteState, FloatND, Period, ScalarInt @categorical(ordered=True) class Health: - bad: int - good: int + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def _next_regime_probs( diff --git a/tests/test_models/shock_grids.py b/tests/test_models/shock_grids.py index 3150ff4be..32f005ce4 100644 --- a/tests/test_models/shock_grids.py +++ b/tests/test_models/shock_grids.py @@ -67,14 +67,14 @@ def utility( @categorical(ordered=True) class Health: - bad: int = 0 - good: int = 1 + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt @functools.cache @@ -120,9 +120,9 @@ def get_model( @categorical(ordered=False) class MultiRegimeId: - work: int - retire: int - dead: int + work: ScalarInt + retire: ScalarInt + dead: ScalarInt def _next_regime_multi( diff --git a/tests/test_models/stochastic.py b/tests/test_models/stochastic.py index c1c94fae1..618dace67 100644 --- a/tests/test_models/stochastic.py +++ b/tests/test_models/stochastic.py @@ -28,6 +28,7 @@ DiscreteState, FloatND, Period, + ScalarInt, ) from lcm_examples.mortality import ( LaborSupply, @@ -41,14 +42,14 @@ @categorical(ordered=True) class Health: - bad: int - good: int + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class PartnerStatus: - single: int - partnered: int + single: ScalarInt + partnered: ScalarInt def utility_working( diff --git a/tests/test_nan_diagnostics.py b/tests/test_nan_diagnostics.py index 47c42f32a..dd6e0b5f9 100644 --- a/tests/test_nan_diagnostics.py +++ b/tests/test_nan_diagnostics.py @@ -96,8 +96,8 @@ def _build_nan_model() -> tuple[Model, dict]: @categorical(ordered=False) class _Rid: - non_terminal: int - terminal: int + non_terminal: ScalarInt + terminal: ScalarInt def utility( consumption: ContinuousAction, diff --git a/tests/test_next_state.py b/tests/test_next_state.py index ccb79b396..1f3067ae9 100644 --- a/tests/test_next_state.py +++ b/tests/test_next_state.py @@ -1,18 +1,17 @@ -from dataclasses import dataclass from types import MappingProxyType import jax.numpy as jnp import pandas as pd from lcm.ages import AgeGrid -from lcm.grids import DiscreteGrid +from lcm.grids import DiscreteGrid, categorical from lcm.regime_building import process_regimes from lcm.regime_building.next_state import ( _create_discrete_stochastic_next_func, get_next_state_function_for_simulation, get_next_state_function_for_solution, ) -from lcm.typing import ContinuousState +from lcm.typing import ContinuousState, ScalarInt from tests.test_models.deterministic.regression import dead, working_life @@ -20,7 +19,7 @@ def test_get_next_state_function_with_solve_target(): ages = AgeGrid(start=0, stop=4, step="Y") regimes = {"working_life": working_life, "dead": dead} regime_names_to_ids = MappingProxyType( - {name: idx for idx, name in enumerate(regimes.keys())} + {name: jnp.int32(idx) for idx, name in enumerate(regimes.keys())} ) internal_regimes = process_regimes( regimes=regimes, @@ -69,10 +68,10 @@ def f_a(state: ContinuousState) -> ContinuousState: def f_b(state: ContinuousState) -> ContinuousState: return state + 1.0 - @dataclass + @categorical(ordered=False) class MockCategory: - cat_0: int = 0 - cat_1: int = 1 + cat_0: ScalarInt + cat_1: ScalarInt all_grids = MappingProxyType( {"mock": MappingProxyType({"b": DiscreteGrid(MockCategory)})} diff --git a/tests/test_pandas_utils.py b/tests/test_pandas_utils.py index a8dba7b6f..17bd50a23 100644 --- a/tests/test_pandas_utils.py +++ b/tests/test_pandas_utils.py @@ -23,6 +23,7 @@ initial_conditions_from_dataframe, ) from lcm.params.processing import broadcast_to_template +from lcm.typing import ScalarInt from lcm.utils.error_handling import validate_transition_probs from tests.test_models.basic_discrete import ( Health, @@ -40,8 +41,8 @@ @categorical(ordered=False) class Occupation: - blue_collar: int - white_collar: int + blue_collar: ScalarInt + white_collar: ScalarInt def test_to_categorical_dtype_returns_correct_type(): @@ -67,9 +68,9 @@ def test_to_categorical_dtype_is_not_ordered(): def test_to_categorical_dtype_ordered(): @categorical(ordered=True) class Severity: - mild: int - moderate: int - severe: int + mild: ScalarInt + moderate: ScalarInt + severe: ScalarInt result = Severity.to_categorical_dtype() # ty: ignore[unresolved-attribute] assert result.ordered is True @@ -108,8 +109,8 @@ def test_build_discrete_grid_lookup_ignores_continuous(): def test_build_discrete_grid_lookup_inconsistent_raises(): @categorical(ordered=False) class HealthAlt: - sick: int - healthy: int + sick: ScalarInt + healthy: ScalarInt regimes = { "a": Regime( @@ -405,19 +406,19 @@ def test_round_trip_with_discrete_model(): @categorical(ordered=True) class HealthWithDisability: - disabled: int - bad: int - good: int + disabled: ScalarInt + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class _HetRegimeId: - pre65: int - post65: int - dead: int + pre65: ScalarInt + post65: ScalarInt + dead: ScalarInt -def _het_next_regime() -> int: +def _het_next_regime() -> ScalarInt: return _HetRegimeId.dead @@ -504,16 +505,16 @@ def test_initial_conditions_heterogeneous_state_sets() -> None: @categorical(ordered=False) class _Rid: - with_status: int - without_status: int - dead: int + with_status: ScalarInt + without_status: ScalarInt + dead: ScalarInt @categorical(ordered=False) class _Status: - low: int - high: int + low: ScalarInt + high: ScalarInt - def _next_regime() -> int: + def _next_regime() -> ScalarInt: return _Rid.dead def _utility_with_status(wealth: float, status: int) -> float: @@ -577,11 +578,11 @@ def test_initial_conditions_shock_grid_heterogeneous_state_sets() -> None: @categorical(ordered=False) class _Rid: - earner: int - retiree: int - dead: int + earner: ScalarInt + retiree: ScalarInt + dead: ScalarInt - def _next_regime() -> int: + def _next_regime() -> ScalarInt: return _Rid.dead def _earner_utility(wealth: float, income: float) -> float: @@ -656,10 +657,10 @@ def test_convert_series_next_function_no_outcome_axis() -> None: @categorical(ordered=False) class _RId: - a: int - dead: int + a: ScalarInt + dead: ScalarInt - def _next_regime() -> int: + def _next_regime() -> ScalarInt: return _RId.dead def _next_wealth(wealth: float, rate: float, period: int) -> float: # noqa: ARG001 @@ -1665,8 +1666,8 @@ def test_convert_series_per_target_transition() -> None: @categorical(ordered=False) class _RId: - working: int - retired: int + working: ScalarInt + retired: ScalarInt def _health_probs( period: Period, health: DiscreteState, probs_array: FloatND @@ -1763,19 +1764,19 @@ def test_convert_series_structured_derived_categoricals() -> None: @categorical(ordered=False) class _RId: - regime_a: int - regime_b: int + regime_a: ScalarInt + regime_b: ScalarInt @categorical(ordered=False) class _ChoiceA: - x: int - y: int + x: ScalarInt + y: ScalarInt @categorical(ordered=False) class _ChoiceB: - x: int - y: int - z: int + x: ScalarInt + y: ScalarInt + z: ScalarInt # "derived" is a DAG function output used as an index in both regimes, # but with different cardinalities (2 vs 3 categories). Since it's a @@ -1843,8 +1844,8 @@ def test_convert_series_runtime_grid_param() -> None: @categorical(ordered=False) class _RId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt alive = Regime( transition=lambda age: jnp.where(age >= 1, _RId.dead, _RId.alive), @@ -1909,8 +1910,8 @@ def test_resolve_categoricals_conflict_raises() -> None: @categorical(ordered=False) class WrongPartner: - x: int - y: int + x: ScalarInt + y: ScalarInt conflicting_regime = dataclasses.replace( model.regimes["working_life"], @@ -1935,19 +1936,19 @@ def test_convert_series_cross_grid_transition() -> None: @categorical(ordered=True) class _HealthPre: - disabled: int - bad: int - good: int + disabled: ScalarInt + bad: ScalarInt + good: ScalarInt @categorical(ordered=True) class _HealthPost: - bad: int - good: int + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class _RId: - pre65: int - post65: int + pre65: ScalarInt + post65: ScalarInt def _health_probs_same( period: Period, health: DiscreteState, health_trans_probs: FloatND @@ -2032,8 +2033,8 @@ def test_resolve_categoricals_includes_derived_when_no_regime_name() -> None: @categorical(ordered=False) class ExtraVar: - a: int - b: int + a: ScalarInt + b: ScalarInt updated_regimes = { name: dataclasses.replace( diff --git a/tests/test_persistence.py b/tests/test_persistence.py index 561fc1a5e..d99ac6d5a 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -21,8 +21,8 @@ @categorical(ordered=False) class _RegimeId: - working: int - retired: int + working: ScalarInt + retired: ScalarInt def _retired_utility(wealth: ContinuousState) -> FloatND: diff --git a/tests/test_regime.py b/tests/test_regime.py index cf6b86bf4..94b9e768b 100644 --- a/tests/test_regime.py +++ b/tests/test_regime.py @@ -40,8 +40,8 @@ def test_regime_name_does_not_contain_separator(): @categorical(ordered=False) class RegimeId: - work__test: int # Contains separator - but RegimeId class has matching field - dead: int + work__test: ScalarInt # Contains separator — validated at Model level + dead: ScalarInt working = Regime( functions={"utility": utility}, @@ -235,8 +235,8 @@ def test_get_all_functions_includes_identity_for_fixed_discrete_state(): @categorical(ordered=False) class Edu: - low: int - high: int + low: ScalarInt + high: ScalarInt regime = Regime( transition=lambda: 0, @@ -342,8 +342,8 @@ def test_discrete_state_grid_without_explicit_transition_raises(): @categorical(ordered=False) class Status: - low: int - high: int + low: ScalarInt + high: ScalarInt with pytest.raises( RegimeInitializationError, match="must have an entry in state_transitions" @@ -373,8 +373,8 @@ def test_regime_with_fixed_states_only(): @categorical(ordered=False) class FixedRegimeId: - working_life: int - dead: int + working_life: ScalarInt + dead: ScalarInt def fixed_utility( consumption: ContinuousAction, wealth: ContinuousState diff --git a/tests/test_regime_state_mismatch.py b/tests/test_regime_state_mismatch.py index 11abb8e9e..20e22fed2 100644 --- a/tests/test_regime_state_mismatch.py +++ b/tests/test_regime_state_mismatch.py @@ -28,22 +28,22 @@ @categorical(ordered=False) class HealthWorkingLife: - disabled: int - bad: int - good: int + disabled: ScalarInt + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class HealthRetirement: - bad: int - good: int + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class RegimeId: - working_life: int - retirement: int - dead: int + working_life: ScalarInt + retirement: ScalarInt + dead: ScalarInt def hm_utility_working(consumption: ContinuousAction, health: DiscreteState) -> FloatND: @@ -140,13 +140,13 @@ def test_deterministic_target_only_state() -> None: @categorical(ordered=False) class HeirPresent: - no: int - yes: int + no: ScalarInt + yes: ScalarInt @categorical(ordered=False) class _RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def next_regime(age: float) -> ScalarInt: return jnp.where(age >= 1, _RegimeId.dead, _RegimeId.alive) @@ -218,13 +218,13 @@ def test_stochastic_target_only_state() -> None: @categorical(ordered=False) class HeirPresent: - no: int - yes: int + no: ScalarInt + yes: ScalarInt @categorical(ordered=False) class _RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def next_regime(age: float) -> ScalarInt: return jnp.where(age >= 1, _RegimeId.dead, _RegimeId.alive) @@ -389,19 +389,19 @@ def test_discrete_state_same_count_different_names(): @categorical(ordered=False) class StatusA: - employed: int - unemployed: int + employed: ScalarInt + unemployed: ScalarInt @categorical(ordered=False) class StatusB: - married: int - single: int + married: ScalarInt + single: ScalarInt @categorical(ordered=False) class _RegimeId: - work: int - retire: int - dead: int + work: ScalarInt + retire: ScalarInt + dead: ScalarInt def next_regime(age: float) -> ScalarInt: return jnp.where( @@ -447,21 +447,21 @@ def test_mixed_ordered_flags_raises(): @categorical(ordered=True) class HealthOrdered: - bad: int - good: int + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class HealthUnordered: - bad: int - good: int + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class _RegimeId: - a: int - b: int - dead: int + a: ScalarInt + b: ScalarInt + dead: ScalarInt - def next_regime() -> int: + def next_regime() -> ScalarInt: return _RegimeId.dead a = Regime( @@ -491,21 +491,21 @@ def test_both_ordered_same_categories_passes(): @categorical(ordered=True) class HealthA: - bad: int - good: int + bad: ScalarInt + good: ScalarInt @categorical(ordered=True) class HealthB: - bad: int - good: int + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class _RegimeId: - a: int - b: int - dead: int + a: ScalarInt + b: ScalarInt + dead: ScalarInt - def next_regime() -> int: + def next_regime() -> ScalarInt: return _RegimeId.dead a = Regime( @@ -619,9 +619,9 @@ def test_incomplete_per_target_reachable_target(): @categorical(ordered=False) class _RegimeId: - regime_a: int - regime_b: int - dead: int + regime_a: ScalarInt + regime_b: ScalarInt + dead: ScalarInt def next_regime_a(age: float) -> ScalarInt: """A → B at age 1. B IS reachable.""" @@ -696,9 +696,9 @@ def test_complete_per_target_stochastic_cross_grid() -> None: @categorical(ordered=False) class _RegimeId: - regime_a: int - regime_b: int - dead: int + regime_a: ScalarInt + regime_b: ScalarInt + dead: ScalarInt def next_regime_a(age: float) -> ScalarInt: return jnp.where( @@ -769,10 +769,10 @@ def test_incomplete_per_target_unreachable_target() -> None: @categorical(ordered=False) class _RegimeId: - regime_a: int - regime_b: int - regime_c: int - dead: int + regime_a: ScalarInt + regime_b: ScalarInt + regime_c: ScalarInt + dead: ScalarInt def next_regime_a(age: float) -> ScalarInt: """A → B at age 1, A otherwise. Never produces C.""" diff --git a/tests/test_runtime_params.py b/tests/test_runtime_params.py index 61b2414a4..58fdfc27c 100644 --- a/tests/test_runtime_params.py +++ b/tests/test_runtime_params.py @@ -6,13 +6,13 @@ from lcm import AgeGrid, LinSpacedGrid, Model, Regime, categorical from lcm.grids import IrregSpacedGrid -from lcm.typing import ContinuousAction, ContinuousState, FloatND +from lcm.typing import ContinuousAction, ContinuousState, FloatND, ScalarInt @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def _utility(consumption: ContinuousAction, wealth: ContinuousState) -> FloatND: diff --git a/tests/test_runtime_shock_params.py b/tests/test_runtime_shock_params.py index 1fdf5ffd9..1e130e866 100644 --- a/tests/test_runtime_shock_params.py +++ b/tests/test_runtime_shock_params.py @@ -11,13 +11,13 @@ Regime, categorical, ) -from lcm.typing import ContinuousAction, ContinuousState, FloatND +from lcm.typing import ContinuousAction, ContinuousState, FloatND, ScalarInt @categorical(ordered=False) class RegimeIdShock: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def _utility( diff --git a/tests/test_single_feasible_action.py b/tests/test_single_feasible_action.py index 53a185941..08b61ef7a 100644 --- a/tests/test_single_feasible_action.py +++ b/tests/test_single_feasible_action.py @@ -28,13 +28,19 @@ from lcm.grids import IrregSpacedGrid from lcm.grids.coordinates import get_irreg_coordinate from lcm.regime_building.ndimage import map_coordinates -from lcm.typing import ContinuousAction, ContinuousState, DiscreteState, FloatND +from lcm.typing import ( + ContinuousAction, + ContinuousState, + DiscreteState, + FloatND, + ScalarInt, +) @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def _utility(consumption: ContinuousAction) -> FloatND: @@ -245,14 +251,14 @@ def test_simulate_with_constrained_action_grid(wealth_lo, consumption_lo, label) @categorical(ordered=False) class PrefType: - type_0: int - type_1: int + type_0: ScalarInt + type_1: ScalarInt @categorical(ordered=False) class AliveDeadRegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def _crra_bequest( @@ -442,8 +448,8 @@ def _runtime_state_grid_model() -> tuple[Model, dict, dict]: @categorical(ordered=False) class RuntimeRegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def utility(consumption, wealth): return jnp.log(consumption) + 0.0 * wealth diff --git a/tests/test_solution_on_toy_model_deterministic.py b/tests/test_solution_on_toy_model_deterministic.py index c58c2e35a..98996a52b 100644 --- a/tests/test_solution_on_toy_model_deterministic.py +++ b/tests/test_solution_on_toy_model_deterministic.py @@ -30,20 +30,20 @@ @categorical(ordered=False) class DiscreteConsumption: - low: int - high: int + low: ScalarInt + high: ScalarInt @categorical(ordered=False) class LaborSupply: - do_not_work: int - work: int + do_not_work: ScalarInt + work: ScalarInt @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def utility( diff --git a/tests/test_solution_on_toy_model_stochastic.py b/tests/test_solution_on_toy_model_stochastic.py index a6c83d9a1..fdb175b1e 100644 --- a/tests/test_solution_on_toy_model_stochastic.py +++ b/tests/test_solution_on_toy_model_stochastic.py @@ -17,7 +17,7 @@ Model, categorical, ) -from lcm.typing import DiscreteState, FloatND +from lcm.typing import DiscreteState, FloatND, ScalarInt from tests.conftest import DECIMAL_PRECISION from tests.test_solution_on_toy_model_deterministic import ( RegimeId, @@ -33,8 +33,8 @@ @categorical(ordered=False) class Health: - bad: int - good: int + bad: ScalarInt + good: ScalarInt def next_health(health: DiscreteState, probs_array: FloatND) -> FloatND: diff --git a/tests/test_state_action_space.py b/tests/test_state_action_space.py index 5957e7019..e70d7e319 100644 --- a/tests/test_state_action_space.py +++ b/tests/test_state_action_space.py @@ -1,4 +1,3 @@ -from dataclasses import dataclass from types import MappingProxyType import jax.numpy as jnp @@ -10,6 +9,7 @@ from lcm.regime import Regime from lcm.regime_building.V import VInterpolationInfo, create_v_interpolation_info from lcm.state_action_space import create_state_action_space +from lcm.typing import ScalarInt def _create_variable_info( @@ -32,8 +32,8 @@ def _create_variable_info( def test_create_state_action_space_solution_discrete_action_continuous_state(): @categorical(ordered=False) class WorkChoice: - no_work: int - work: int + no_work: ScalarInt + work: ScalarInt variable_info = _create_variable_info( continuous_states=["wealth"], @@ -91,8 +91,8 @@ def test_create_state_action_space_solution_continuous_action(): def test_state_action_space_replace_method(): @categorical(ordered=False) class WorkChoice: - no_work: int - work: int + no_work: ScalarInt + work: ScalarInt variable_info = _create_variable_info( continuous_states=["wealth"], @@ -121,10 +121,10 @@ class WorkChoice: def test_create_v_interpolation_info(): - @dataclass + @categorical(ordered=False) class Health: - good: int = 0 - bad: int = 1 + good: ScalarInt + bad: ScalarInt regime = Regime( transition=lambda: 0, # non-terminal diff --git a/tests/test_static_params.py b/tests/test_static_params.py index d61f3d3b0..4669145ca 100644 --- a/tests/test_static_params.py +++ b/tests/test_static_params.py @@ -17,8 +17,8 @@ @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def _utility(consumption: ContinuousAction, wealth: ContinuousState) -> FloatND: @@ -314,8 +314,8 @@ def test_mixed_series_and_scalar_fixed_params(): @categorical(ordered=False) class _WealthGroup: - low: int - high: int + low: ScalarInt + high: ScalarInt def _wealth_group(wealth: ContinuousState) -> ScalarInt: @@ -433,9 +433,9 @@ def test_model_broadcast_conflict_raises(): @categorical(ordered=False) class _OtherGroup: - a: int - b: int - c: int + a: ScalarInt + b: ScalarInt + c: ScalarInt alive = Regime( functions={"utility": _utility_with_group, "wealth_group": _wealth_group}, @@ -466,18 +466,18 @@ def test_different_regime_derived_categoricals_with_model_broadcast(): @categorical(ordered=False) class _GroupA: - x: int - y: int + x: ScalarInt + y: ScalarInt @categorical(ordered=False) class _GroupB: - p: int - q: int + p: ScalarInt + q: ScalarInt @categorical(ordered=False) class _Shared: - lo: int - hi: int + lo: ScalarInt + hi: ScalarInt alive = Regime( functions={"utility": lambda: 0.0}, diff --git a/tests/test_stochastic.py b/tests/test_stochastic.py index 7dde47a99..2ec8ec107 100644 --- a/tests/test_stochastic.py +++ b/tests/test_stochastic.py @@ -226,13 +226,13 @@ def _make_minimal_stochastic_model(next_draw: Callable[..., FloatND]) -> Model: @categorical(ordered=False) class ShockStatus: - bad: int - good: int + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class ShockRegimeId: - working_life: int - dead: int + working_life: ScalarInt + dead: ScalarInt def utility(consumption: ContinuousAction, draw: DiscreteState) -> FloatND: bonus = jnp.where(draw == ShockStatus.good, 1.0, 0.0) diff --git a/tests/test_validate_param_types.py b/tests/test_validate_param_types.py index 174286065..06a95d597 100644 --- a/tests/test_validate_param_types.py +++ b/tests/test_validate_param_types.py @@ -12,21 +12,22 @@ from lcm import AgeGrid, DiscreteGrid, LinSpacedGrid, Model, Regime, categorical from lcm.dtypes import canonical_float_dtype +from lcm.typing import ScalarInt @categorical(ordered=True) class Health: - bad: int - good: int + bad: ScalarInt + good: ScalarInt @categorical(ordered=False) class RegimeId: - working: int - dead: int + working: ScalarInt + dead: ScalarInt -def _next_regime() -> int: +def _next_regime() -> ScalarInt: return RegimeId.dead diff --git a/tests/test_validate_regime_transition_probs.py b/tests/test_validate_regime_transition_probs.py index cacf20f20..9d97f1e87 100644 --- a/tests/test_validate_regime_transition_probs.py +++ b/tests/test_validate_regime_transition_probs.py @@ -13,7 +13,7 @@ categorical, ) from lcm.exceptions import InvalidRegimeTransitionProbabilitiesError -from lcm.typing import DiscreteAction, FloatND +from lcm.typing import DiscreteAction, FloatND, ScalarInt from lcm.utils.error_handling import ( _format_sum_violation, validate_regime_transition_probs, @@ -181,14 +181,14 @@ def test_format_sum_violation_with_scalar_input_and_state_action_values(): @categorical(ordered=False) class _Action: - stay: int - leave: int + stay: ScalarInt + leave: ScalarInt @categorical(ordered=False) class _RegimeId: - active: int - terminal: int + active: ScalarInt + terminal: ScalarInt def _next_regime_only_fails_for_leave(action: DiscreteAction) -> FloatND: diff --git a/tests/test_validation_scalar_actions.py b/tests/test_validation_scalar_actions.py index 89289f5b0..e53e521b3 100644 --- a/tests/test_validation_scalar_actions.py +++ b/tests/test_validation_scalar_actions.py @@ -16,13 +16,13 @@ from lcm import AgeGrid, LinSpacedGrid, Model, Regime, categorical from lcm.params import MappingLeaf, as_leaf -from lcm.typing import ContinuousAction, ContinuousState, FloatND +from lcm.typing import ContinuousAction, ContinuousState, FloatND, ScalarInt @categorical(ordered=False) class RegimeId: - alive: int - dead: int + alive: ScalarInt + dead: ScalarInt def _aggregate_with_ids( From 6e187a7f8cb2113699607a1a923e1fe8a7f95088 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 11 May 2026 15:12:10 +0200 Subject: [PATCH 2/9] bench: bump aca-model pin to feature/categorical-scalarint HEAD PR #350 tightens `@categorical` to require `ScalarInt` field annotations. The benchmarks-cuda12 env still pulls aca-model from the pre-#350 sha (`9ac2043`), which annotates fields as `int` and fails the new decoration-time gate during `setup_cache`. Point the pin at the aca-model #350 cascade branch (`feature/categorical-scalarint`, head `b807b28`) until both PRs land. After aca-model PR#10 and pylcm#350 merge, revert to a `main`-tracking sha. Co-Authored-By: Claude Opus 4.7 (1M context) --- pixi.lock | 8 ++++---- pyproject.toml | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pixi.lock b/pixi.lock index cc25be60b..a73e96fba 100644 --- a/pixi.lock +++ b/pixi.lock @@ -270,7 +270,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/noarch/zipp-3.23.1-pyhcf101f3_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - - pypi: git+https://github.com/OpenSourceEconomics/aca-model.git?rev=9ac20430f499a8b1cdb056af85bc2a26e850bad2#9ac20430f499a8b1cdb056af85bc2a26e850bad2 + - pypi: git+https://github.com/OpenSourceEconomics/aca-model.git?rev=b807b286175e00f8a721f039f4afd657bbb10e2f#b807b286175e00f8a721f039f4afd657bbb10e2f - pypi: https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/ae/44/c1221527f6a71a01ec6fbad7fa78f1d50dfa02217385cf0fa3eec7087d59/click-8.3.3-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/2c/1a/aff8bb287a4b1400f69e09a53bd65de96aa5cee5691925b38731c67fc695/click_default_group-1.2.4-py2.py3-none-any.whl @@ -5328,7 +5328,7 @@ packages: purls: [] size: 8191 timestamp: 1744137672556 -- pypi: git+https://github.com/OpenSourceEconomics/aca-model.git?rev=9ac20430f499a8b1cdb056af85bc2a26e850bad2#9ac20430f499a8b1cdb056af85bc2a26e850bad2 +- pypi: git+https://github.com/OpenSourceEconomics/aca-model.git?rev=b807b286175e00f8a721f039f4afd657bbb10e2f#b807b286175e00f8a721f039f4afd657bbb10e2f name: aca-model version: 0.0.0 requires_dist: @@ -13962,8 +13962,8 @@ packages: timestamp: 1774796815820 - pypi: ./ name: pylcm - version: 0.0.2.dev263+g6417124d0.d20260511 - sha256: bf60af4d734b824cea1c2a6bd1baf0311aa23945778c0ad0ce087b2696df5f5a + version: 0.0.2.dev131+g4a05a8283.d20260511 + sha256: 4f354caf09e81582a5adb942f82bd7f91aafbac5a9eeaacb931f28bd1d0499dc requires_dist: - cloudpickle>=3.1.2 - dags>=0.5.1 diff --git a/pyproject.toml b/pyproject.toml index 352f2e766..89846aec9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -98,7 +98,7 @@ tests-cuda13 = { features = [ "tests", "cuda13" ], solve-group = "cuda13" } tests-metal = { features = [ "tests", "metal" ], solve-group = "metal" } type-checking = { features = [ "type-checking", "tests" ], solve-group = "default" } [tool.pixi.feature.benchmarks.pypi-dependencies] -aca-model = { git = "https://github.com/OpenSourceEconomics/aca-model.git", rev = "9ac20430f499a8b1cdb056af85bc2a26e850bad2" } +aca-model = { git = "https://github.com/OpenSourceEconomics/aca-model.git", rev = "b807b286175e00f8a721f039f4afd657bbb10e2f" } [tool.pixi.feature.cuda12] platforms = [ "linux-64" ] system-requirements = { cuda = "12" } From 3f1d1e41f16a927fb8e3846d05e9ce5862bdb035 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 11 May 2026 15:20:47 +0200 Subject: [PATCH 3/9] bench: unpack 3-tuple from get_benchmark_params aca-model's `get_benchmark_params` returns three values (likely a moments tuple alongside the two pre-existing entries); the benchmark file was stuck on the older 2-tuple shape. aca-model's own `test_benchmark.py` runs against the up-to-date in-tree helper, so the drift only surfaces here once pylcm's benchmark CI pulls a fresh aca-model. Co-Authored-By: Claude Opus 4.7 (1M context) --- benchmarks/bench_aca_baseline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/bench_aca_baseline.py b/benchmarks/bench_aca_baseline.py index 8b15efabd..57e1ad02c 100644 --- a/benchmarks/bench_aca_baseline.py +++ b/benchmarks/bench_aca_baseline.py @@ -71,7 +71,7 @@ def _build() -> tuple[object, object, object]: n_subjects=_N_SUBJECTS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) - _, model_params = get_benchmark_params(model=model) + _, _, model_params = get_benchmark_params(model=model) initial_conditions = get_benchmark_initial_conditions( model=model, n_subjects=_N_SUBJECTS, seed=0 ) From e639f18589221d9da97fc350d75584f7123ca1ba Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 11 May 2026 15:17:23 +0200 Subject: [PATCH 4/9] IrregSpacedGrid: extra_param_names for user-side runtime scalars (#348) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `broadcast_to_template` rejects any `fixed_params` key that no DAG function lists in its signature. That assumption breaks for runtime-supplied grids whose points are computed by user-side code from per-iteration scalars (e.g., a grid upper bound) — pylcm itself never reads those scalars, but they still need to flow through the params machinery instead of being funneled through `simulate(params=...)` where they compete with estimation parameters. This adds an `extra_param_names: tuple[str, ...]` keyword to `IrregSpacedGrid`. When `pass_points_at_runtime=True`, each name in the tuple surfaces as a `ScalarFloat` slot in the action/state's pseudo-function template entry alongside `points`. The user's injection code reads the resolved values from `model.fixed_params` (or per-iteration params) before computing the points; pylcm just threads them through `broadcast_to_template` unchanged. Rejected on fixed-points grids — a baked-in grid has no user-side runtime computation to feed. Closes #348. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/lcm/grids/continuous.py | 29 ++++++++++++++++ src/lcm/params/regime_template.py | 19 +++++++++-- tests/test_runtime_params.py | 55 +++++++++++++++++++++++++++++++ 3 files changed, 101 insertions(+), 2 deletions(-) diff --git a/src/lcm/grids/continuous.py b/src/lcm/grids/continuous.py index 839d4cefb..25503a00e 100644 --- a/src/lcm/grids/continuous.py +++ b/src/lcm/grids/continuous.py @@ -227,10 +227,21 @@ class IrregSpacedGrid(ContinuousGrid): omitted and only `n_points` is given, the points must be supplied at runtime via the params. + `extra_param_names` declares additional scalar parameters consumed by + user-side code that *constructs* the runtime points (e.g., a grid + upper bound that changes across optimizer iterations). These names + enter the params template alongside `points` and pass through + `broadcast_to_template` even though no DAG function references them + — the user's injection code reads them from the resolved params / + `model.fixed_params` to derive the points before calling + `solve` / `simulate`. + Example: -------- Fixed grid: `IrregSpacedGrid(points=[-1.73, -0.58, 0.58, 1.73])` Grid that is only completed at runtime via params: `IrregSpacedGrid(n_points=4)` + Grid that pairs runtime `points` with an extra scalar bound: + `IrregSpacedGrid(n_points=4, extra_param_names=("max_consumption",))` """ @@ -240,12 +251,23 @@ class IrregSpacedGrid(ContinuousGrid): n_points: int """Number of points. Derived from `len(points)` when points are given.""" + extra_param_names: tuple[str, ...] + """Names of additional scalar params surfaced in the template. + + Only meaningful when points are supplied at runtime + (`pass_points_at_runtime=True`); pylcm itself never reads these + values — they're carried through the params machinery so user-side + injection code can pick them up without fighting the `Unknown keys` + validator. + """ + def __init__( self, *, points: Sequence[float] | Float1D | None = None, n_points: int | None = None, batch_size: int = 0, + extra_param_names: tuple[str, ...] = (), ) -> None: if points is not None: _validate_irreg_spaced_grid(points) @@ -269,9 +291,16 @@ def __init__( ) else: stored_points = None + if extra_param_names and stored_points is not None: + raise GridInitializationError( + "`extra_param_names` is only valid when points are supplied at " + "runtime (i.e. `points=None`); a fixed-points grid has no " + "user-side params to thread through." + ) object.__setattr__(self, "points", stored_points) object.__setattr__(self, "n_points", n_points) object.__setattr__(self, "batch_size", batch_size) + object.__setattr__(self, "extra_param_names", tuple(extra_param_names)) @property def pass_points_at_runtime(self) -> bool: diff --git a/src/lcm/params/regime_template.py b/src/lcm/params/regime_template.py index 67ddfe75b..97c6c74e4 100644 --- a/src/lcm/params/regime_template.py +++ b/src/lcm/params/regime_template.py @@ -92,7 +92,7 @@ def _add_runtime_grid_params( _fail_if_runtime_grid_shadows_function( function_params=function_params, name=state_name, kind="state" ) - function_params[state_name] = {"points": "Float1D"} + function_params[state_name] = _irreg_grid_template_entry(grid) elif isinstance(grid, _ShockGrid) and grid.params_to_pass_at_runtime: _fail_if_runtime_grid_shadows_function( function_params=function_params, @@ -108,7 +108,22 @@ def _add_runtime_grid_params( _fail_if_runtime_grid_shadows_function( function_params=function_params, name=action_name, kind="action" ) - function_params[action_name] = {"points": "Float1D"} + function_params[action_name] = _irreg_grid_template_entry(grid) + + +def _irreg_grid_template_entry(grid: IrregSpacedGrid) -> dict[str, str]: + """Template slots for a runtime-points `IrregSpacedGrid`. + + Always exposes `points: Float1D`. Each entry in + `grid.extra_param_names` adds a `ScalarFloat` slot so user-side + injection code can thread its scalar bounds through `fixed_params` + or per-iteration params without tripping `broadcast_to_template`'s + unknown-keys check. + """ + entry: dict[str, str] = {"points": "Float1D"} + for name in grid.extra_param_names: + entry[name] = "ScalarFloat" + return entry def _fail_if_runtime_grid_shadows_function( diff --git a/tests/test_runtime_params.py b/tests/test_runtime_params.py index 58fdfc27c..5064ab0e4 100644 --- a/tests/test_runtime_params.py +++ b/tests/test_runtime_params.py @@ -313,3 +313,58 @@ def test_simulate_with_runtime_action_grid_no_nan() -> None: ) df = result.to_dataframe() assert not df["value"].isna().any() + + +# --- extra_param_names: scalar params consumed by user-side injection code --- + + +def test_extra_param_names_added_to_template(): + """`extra_param_names` show up in the action's template alongside `points`.""" + model = _make_action_grid_model( + consumption_grid=IrregSpacedGrid( + n_points=5, + extra_param_names=("max_consumption",), + ), + ) + alive_template = model._params_template["alive"] + assert alive_template["consumption"] == { + "points": "Float1D", + "max_consumption": "ScalarFloat", + } + + +def test_extra_param_names_accepted_via_fixed_params(): + """Model-level `fixed_params` with extra-grid-param keys broadcast cleanly.""" + grid = IrregSpacedGrid(n_points=5, extra_param_names=("max_consumption",)) + model_fixed = _make_action_grid_model(consumption_grid=grid) + # Build via fresh Model to inject `fixed_params`. + alive = model_fixed.regimes["alive"] + dead = model_fixed.regimes["dead"] + model = Model( + regimes={"alive": alive, "dead": dead}, + ages=AgeGrid(start=0, stop=2, step="Y"), + regime_id_class=RegimeId, + fixed_params={"max_consumption": 5.0}, + ) + params = { + "discount_factor": 0.95, + "interest_rate": 0.05, + "alive": {"consumption": {"points": jnp.linspace(0.1, 5.0, 5)}}, + } + period_to_regime_to_V_arr = model.solve(params=params, log_level="off") + assert len(period_to_regime_to_V_arr) > 0 + + +def test_extra_param_names_rejected_on_fixed_points_grid(): + """`extra_param_names` is meaningless when points are baked in at construction.""" + with pytest.raises(Exception, match="only valid when points are supplied"): + IrregSpacedGrid(points=[1.0, 2.0, 3.0], extra_param_names=("foo",)) + + +def test_extra_param_names_empty_by_default(): + """No `extra_param_names` keeps the template's grid entry to just `points`.""" + model = _make_action_grid_model( + consumption_grid=IrregSpacedGrid(n_points=5), + ) + alive_template = model._params_template["alive"] + assert alive_template["consumption"] == {"points": "Float1D"} From a45d3c21a098693dddfadcc392cf09f2538ef9e8 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 11 May 2026 19:52:11 +0200 Subject: [PATCH 5/9] docs(categorical): drop redundant instance-attr example Class-attr access already documents the `Array(0, dtype=int32)` shape; the instance fall-through line repeats the same value without adding information about the decorator contract. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/lcm/grids/categorical.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/lcm/grids/categorical.py b/src/lcm/grids/categorical.py index fef009754..9515e39b5 100644 --- a/src/lcm/grids/categorical.py +++ b/src/lcm/grids/categorical.py @@ -33,7 +33,6 @@ class LaborSupply: LaborSupply.work # Array(0, dtype=int32) LaborSupply.retire # Array(1, dtype=int32) - LaborSupply().work # Array(0, dtype=int32) Args: ordered: Whether the categories have a meaningful ordering. Must be From 04c9533084cbfaa0df2d9483feba41e6eecd38a5 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 11 May 2026 19:53:34 +0200 Subject: [PATCH 6/9] =?UTF-8?q?refactor(initial=5Fconditions):=20rename=20?= =?UTF-8?q?ids=5Fto=5Fregime=5Fnames=20=E2=86=92=20regime=5Fids=5Fto=5Fnam?= =?UTF-8?q?es?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirrors the forward direction `regime_names_to_ids` so the inverse is read symmetrically. No behavior change. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/lcm/simulation/initial_conditions.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/lcm/simulation/initial_conditions.py b/src/lcm/simulation/initial_conditions.py index db713eb76..1f5486fe4 100644 --- a/src/lcm/simulation/initial_conditions.py +++ b/src/lcm/simulation/initial_conditions.py @@ -123,7 +123,7 @@ def validate_initial_conditions( # Build reverse lookup from regime IDs to names. `regime_names_to_ids` # values are `ScalarInt` (jax 0-d arrays), which can't serve as dict # keys directly; `invert_regime_ids` coerces them to Python `int`. - ids_to_regime_names = invert_regime_ids(regime_names_to_ids) + regime_ids_to_names = invert_regime_ids(regime_names_to_ids) # Extract regime array regime_arr = initial_conditions.get("regime") @@ -133,7 +133,7 @@ def validate_initial_conditions( ) # Vectorized regime ID validity check - valid_ids_arr = jnp.array(sorted(ids_to_regime_names.keys())) + valid_ids_arr = jnp.array(sorted(regime_ids_to_names.keys())) invalid_mask = ~jnp.isin(regime_arr, valid_ids_arr) if jnp.any(invalid_mask): invalid_ids = sorted({int(i) for i in jnp.unique(regime_arr[invalid_mask])}) @@ -141,7 +141,7 @@ def validate_initial_conditions( format_messages( [ f"Invalid regime IDs {invalid_ids}. " - f"Valid IDs: {sorted(ids_to_regime_names.keys())}" + f"Valid IDs: {sorted(regime_ids_to_names.keys())}" ] ) ) @@ -153,7 +153,7 @@ def validate_initial_conditions( structural_errors = _collect_structural_errors( initial_states=initial_states, regime_id_arr=regime_arr, - ids_to_regime_names=ids_to_regime_names, + regime_ids_to_names=regime_ids_to_names, regime_names_to_ids=regime_names_to_ids, internal_regimes=internal_regimes, ages=ages, @@ -228,7 +228,7 @@ def _collect_state_name_errors( *, initial_states: Mapping[str, Array], regime_id_arr: Array, - ids_to_regime_names: MappingProxyType[int, RegimeName], + regime_ids_to_names: MappingProxyType[int, RegimeName], internal_regimes: MappingProxyType[RegimeName, InternalRegime], valid_regime_names: set[RegimeName], ) -> list[str]: @@ -241,7 +241,7 @@ def _collect_state_name_errors( Args: initial_states: Mapping of state names to arrays. regime_id_arr: Array of integer regime IDs. - ids_to_regime_names: Mapping from integer IDs to regime names. + regime_ids_to_names: Mapping from integer IDs to regime names. internal_regimes: Immutable mapping of regime names to internal regime instances. valid_regime_names: Set of valid regime names. @@ -261,7 +261,7 @@ def _collect_state_name_errors( required_states: set[str] = set(PSEUDO_STATE_NAMES) used_ids = jnp.unique(regime_id_arr) used_regime_names = { - ids_to_regime_names[int(i)] for i in used_ids if int(i) in ids_to_regime_names + regime_ids_to_names[int(i)] for i in used_ids if int(i) in regime_ids_to_names } & valid_regime_names for regime_name in used_regime_names: required_states.update(_get_regime_state_names(internal_regimes[regime_name])) @@ -286,7 +286,7 @@ def _collect_structural_errors( *, initial_states: Mapping[str, Array], regime_id_arr: Array, - ids_to_regime_names: MappingProxyType[int, RegimeName], + regime_ids_to_names: MappingProxyType[int, RegimeName], regime_names_to_ids: RegimeNamesToIds, internal_regimes: MappingProxyType[RegimeName, InternalRegime], ages: AgeGrid, @@ -296,7 +296,7 @@ def _collect_structural_errors( Args: initial_states: Mapping of state names to arrays. regime_id_arr: Array of integer regime IDs. - ids_to_regime_names: Mapping from integer IDs to regime names. + regime_ids_to_names: Mapping from integer IDs to regime names. regime_names_to_ids: Immutable mapping of regime names to integer IDs. internal_regimes: Immutable mapping of regime names to internal regime instances. @@ -317,7 +317,7 @@ def _collect_structural_errors( _collect_state_name_errors( initial_states=initial_states, regime_id_arr=regime_id_arr, - ids_to_regime_names=ids_to_regime_names, + regime_ids_to_names=regime_ids_to_names, internal_regimes=internal_regimes, valid_regime_names=valid_regime_names, ) @@ -367,7 +367,7 @@ def _collect_structural_errors( if not jnp.all(active_mask): invalid_indices = jnp.where(~active_mask)[0].astype(jnp.int32) invalid_combos = { - (ids_to_regime_names[int(regime_id_arr[i])], float(age_values[i])) + (regime_ids_to_names[int(regime_id_arr[i])], float(age_values[i])) for i in invalid_indices } details = "\n".join( From 67434db024639be86d45ad5dbf19befbbb6e8645 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 11 May 2026 19:54:34 +0200 Subject: [PATCH 7/9] =?UTF-8?q?refactor(simulate,=20logging):=20rename=20i?= =?UTF-8?q?ds=5Fto=5Fnames=20=E2=86=92=20regime=5Fids=5Fto=5Fnames?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Matches the inverse direction's symmetry with `regime_names_to_ids` and the recent rename in `simulation.initial_conditions`. No behavior change. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/lcm/simulation/simulate.py | 4 ++-- src/lcm/utils/logging.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/lcm/simulation/simulate.py b/src/lcm/simulation/simulate.py index 395a87b18..22799c202 100644 --- a/src/lcm/simulation/simulate.py +++ b/src/lcm/simulation/simulate.py @@ -113,7 +113,7 @@ def simulate( # Build reverse lookup for regime transition logging. `regime_names_to_ids` # values are `ScalarInt` (jax 0-d arrays), which can't serve as dict keys # directly; `invert_regime_ids` coerces them to Python `int`. - ids_to_names = invert_regime_ids(regime_names_to_ids) + regime_ids_to_names = invert_regime_ids(regime_names_to_ids) for period, age in enumerate(ages.values): period_start = time.monotonic() @@ -172,7 +172,7 @@ def simulate( logger=logger, prev_regime_ids=prev_regime_ids, new_regime_ids=subject_regime_ids, - ids_to_names=ids_to_names, + regime_ids_to_names=regime_ids_to_names, ) elapsed = time.monotonic() - period_start diff --git a/src/lcm/utils/logging.py b/src/lcm/utils/logging.py index c924efe89..d17815b05 100644 --- a/src/lcm/utils/logging.py +++ b/src/lcm/utils/logging.py @@ -113,7 +113,7 @@ def log_regime_transitions( logger: logging.Logger, prev_regime_ids: Int1D, new_regime_ids: Int1D, - ids_to_names: Mapping[int, str], + regime_ids_to_names: Mapping[int, str], ) -> None: """Log regime transition counts at debug level. @@ -121,18 +121,18 @@ def log_regime_transitions( logger: Logger instance. prev_regime_ids: Regime IDs before the transition. new_regime_ids: Regime IDs after the transition. - ids_to_names: Mapping from regime integer IDs to regime names. + regime_ids_to_names: Mapping from regime integer IDs to regime names. """ if not logger.isEnabledFor(logging.DEBUG): return parts: list[str] = [] - for from_id, from_name in sorted(ids_to_names.items()): + for from_id, from_name in sorted(regime_ids_to_names.items()): mask = prev_regime_ids == from_id if not jnp.any(mask): continue - for to_id, to_name in sorted(ids_to_names.items()): + for to_id, to_name in sorted(regime_ids_to_names.items()): count = int(jnp.sum(mask & (new_regime_ids == to_id))) if count > 0: parts.append(f" - {from_name} \u2192 {to_name} = {count}") From 87432d1a43631f2f0316409c4666d4fe2f2ab1e7 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 11 May 2026 19:58:45 +0200 Subject: [PATCH 8/9] Apply review comments from PR #350 - bench: simplify get_benchmark_params unpacking via index access. - categorical: tighten validate_category_class docstring to refer to `ScalarInt`s directly without restating the type. Co-Authored-By: Claude Opus 4.7 (1M context) --- benchmarks/bench_aca_baseline.py | 2 +- src/lcm/grids/categorical.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/benchmarks/bench_aca_baseline.py b/benchmarks/bench_aca_baseline.py index 57e1ad02c..91ea8ecdc 100644 --- a/benchmarks/bench_aca_baseline.py +++ b/benchmarks/bench_aca_baseline.py @@ -71,7 +71,7 @@ def _build() -> tuple[object, object, object]: n_subjects=_N_SUBJECTS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) - _, _, model_params = get_benchmark_params(model=model) + model_params = get_benchmark_params(model=model)[2] initial_conditions = get_benchmark_initial_conditions( model=model, n_subjects=_N_SUBJECTS, seed=0 ) diff --git a/src/lcm/grids/categorical.py b/src/lcm/grids/categorical.py index 9515e39b5..b45834424 100644 --- a/src/lcm/grids/categorical.py +++ b/src/lcm/grids/categorical.py @@ -103,13 +103,13 @@ def validate_category_class(category_class: type) -> list[str]: This validates that: - The class is a dataclass - It has at least one field - - All field values are `ScalarInt` (0-d `jnp.int32` scalars) + - All field values are `ScalarInt`s - All field values are unique - Field values are consecutive integers starting from 0 Args: category_class: The category class to validate. Must be a dataclass with fields - whose values are unique `ScalarInt` scalars. + whose values are unique `ScalarInt`s. Returns: A list of error messages. Empty list if validation passes. From fd647da464dd29905f67178d92b2b4f14681205a Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 11 May 2026 20:03:06 +0200 Subject: [PATCH 9/9] docs(grids): document IrregSpacedGrid.extra_param_names workflow Adds a runtime-points subsection under the IrregSpacedGrid docs in `docs/user_guide/grids.md`: explains what `extra_param_names` declares, why it's needed (carries scalars through the params template without `Unknown keys` errors), and shows a minimal inject-points helper plus the call site against `solve()`. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/user_guide/grids.md | 50 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/docs/user_guide/grids.md b/docs/user_guide/grids.md index a15ae8c83..6d84e0844 100644 --- a/docs/user_guide/grids.md +++ b/docs/user_guide/grids.md @@ -98,6 +98,56 @@ are then supplied via the params dict: IrregSpacedGrid(n_points=4) ``` +#### Runtime points driven by `extra_param_names` + +When the gridpoints depend on a value that varies across solver iterations — a borrowing +limit, a per-iteration upper bound, a calibrated cap — declare the dependency explicitly +with `extra_param_names`: + +```python +IrregSpacedGrid(n_points=64, extra_param_names=("max_consumption",)) +``` + +`extra_param_names` is only valid together with runtime points (`points=None`). Each +name is added to the params template alongside the grid's `points` slot, and +`broadcast_to_template` carries it through even though no DAG function references it. +The names exist so user-side code that *constructs* the gridpoints can read them without +tripping the template validator's `Unknown keys` check. + +The typical workflow looks like this. Declare the grid with `extra_param_names`, +populate the extras (e.g. as `model.fixed_params["max_consumption"]`), then inject the +points before solving or simulating: + +```python +import jax.numpy as jnp +from lcm import IrregSpacedGrid, Model + + +def inject_consumption_points(*, params: dict, model: Model) -> dict: + """Fill the runtime `consumption` gridpoints on every non-terminal regime.""" + max_consumption = jnp.asarray(model.fixed_params["max_consumption"]) + out: dict = dict(params) + for regime_name, regime in model.regimes.items(): + if regime.terminal: + continue + grid = regime.actions["consumption"] + assert isinstance(grid, IrregSpacedGrid) and grid.pass_points_at_runtime + points = jnp.geomspace(1e-3, max_consumption, num=grid.n_points) + regime_entry = dict(out.get(regime_name, {})) + regime_entry["consumption"] = {"points": points} + out[regime_name] = regime_entry + return out + + +params = inject_consumption_points(params=model.get_params_template(), model=model) +model.solve(params=params) +``` + +`extra_param_names` carries any number of scalars (`("max_consumption", "floor")`, +etc.). The construction logic in the user-side helper is free to combine them however +the model needs — pylcm only validates that each declared name is present in the +resolved params. + ### PiecewiseLinSpacedGrid Multiple linearly spaced segments joined at breakpoints. Dense where you need precision,