@categorical: require ScalarInt field annotations + runtime values

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
 

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
     )

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

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

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):

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

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"
    ]
   },
   {

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

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

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

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