@categorical fields: int → ScalarInt (pylcm #350 cascade)

.github/workflows/main.yml

@@ -26,10 +26,10 @@ jobs:
       - uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
-      - name: Install pylcm
+      - name: Install pylcm (feature branch — revert to @main once pylcm#350 merges)
         run: >-
           pip install "pylcm @
-          git+https://github.com/OpenSourceEconomics/pylcm.git@main"
+          git+https://github.com/OpenSourceEconomics/pylcm.git@feat/categorical-scalarint"
       - name: Install aca-model with test deps
         run: pip install -e . pytest pdbp
       - name: Run pytest

src/aca_model/aca/health_insurance.py

@@ -10,7 +10,14 @@
 
 import jax.numpy as jnp
 from lcm.params import MappingLeaf
-from lcm.typing import BoolND, ContinuousState, DiscreteAction, DiscreteState, FloatND
+from lcm.typing import (
+    BoolND,
+    ContinuousState,
+    DiscreteAction,
+    DiscreteState,
+    FloatND,
+    ScalarFloat,
+)
 
 from aca_model.baseline.health_insurance import BuyPrivate, oop_costs
 
@@ -136,9 +143,9 @@ def primary_oop(
     total_health_costs: FloatND,
     cost_sharing_scale: FloatND,
     buy_private: DiscreteAction,
-    deductible: float,
-    coinsurance_rate: float,
-    oop_max: float,
+    deductible: ScalarFloat,
+    coinsurance_rate: ScalarFloat,
+    oop_max: ScalarFloat,
 ) -> FloatND:
     """Compute primary OOP costs with ACA cost-sharing reductions.
 

src/aca_model/agent/assets_and_income.py

@@ -9,12 +9,13 @@
     ContinuousAction,
     ContinuousState,
     FloatND,
+    ScalarFloat,
 )
 
 
 def capital_income(
     assets: ContinuousState,
-    rate_of_return: float,
+    rate_of_return: ScalarFloat,
 ) -> FloatND:
     """Compute capital income from assets."""
     return assets * rate_of_return
@@ -36,7 +37,7 @@ def cash_on_hand(
 
 
 def consumption_dollars_floor(
-    consumption_equiv_floor: float,
+    consumption_equiv_floor: ScalarFloat,
     equivalence_scale: FloatND,
 ) -> FloatND:
     """Per-household $-floor on consumption."""

src/aca_model/agent/health.py

@@ -6,28 +6,28 @@
 
 import jax.numpy as jnp
 from lcm import categorical
-from lcm.typing import DiscreteState, FloatND, IntND, Period
+from lcm.typing import DiscreteState, FloatND, IntND, Period, ScalarInt
 
 
 @categorical(ordered=True)
 class HealthWithDisability:
-    disabled: int
-    bad: int
-    good: int
+    disabled: ScalarInt
+    bad: ScalarInt
+    good: ScalarInt
 
 
 @categorical(ordered=True)
 class Health:
-    bad: int
-    good: int
+    bad: ScalarInt
+    good: ScalarInt
 
 
 @categorical(ordered=True)
 class GoodHealth:
     """Derived categorical for good_health DAG output (0=no, 1=yes)."""
 
-    no: int
-    yes: int
+    no: ScalarInt
+    yes: ScalarInt
 
 
 def is_good_health_3(health: DiscreteState) -> IntND:

src/aca_model/agent/labor_market.py

@@ -12,29 +12,31 @@
     FloatND,
     IntND,
     Period,
+    ScalarFloat,
+    ScalarInt,
 )
 
 
 @categorical(ordered=True)
 class LaborSupply:
-    do_not_work: int
-    h1000: int
-    h1500: int
-    h2000: int
-    h2500: int
+    do_not_work: ScalarInt
+    h1000: ScalarInt
+    h1500: ScalarInt
+    h2000: ScalarInt
+    h2500: ScalarInt
 
 
 @categorical(ordered=False)
 class LaggedLaborSupply:
-    did_not_work: int
-    worked: int
+    did_not_work: ScalarInt
+    worked: ScalarInt
 
 
 @categorical(ordered=False)
 class SpousalIncome:
-    single: int
-    married_no_inc: int
-    married_has_inc: int
+    single: ScalarInt
+    married_no_inc: ScalarInt
+    married_has_inc: ScalarInt
 
 
 HOURS_VALUES = jnp.array([0.0, 1000.0, 1500.0, 2000.0, 2500.0])
@@ -52,8 +54,8 @@ def income(
     good_health: IntND,
     log_ft_wage_mean: FloatND,
     log_ft_wage_std: FloatND,
-    adj_wage_hours_exp: float,
-    adj_wage_hours_int: float,
+    adj_wage_hours_exp: ScalarFloat,
+    adj_wage_hours_int: ScalarFloat,
 ) -> FloatND:
     """Labor income with wage-hours interaction (French & Jones 2011).
 
@@ -88,8 +90,8 @@ def next_lagged_supply(labor_supply: DiscreteAction) -> DiscreteState:
 class IsMarried:
     """Derived categorical for is_married DAG output (0=no, 1=yes)."""
 
-    no: int
-    yes: int
+    no: ScalarInt
+    yes: ScalarInt
 
 
 def is_married(spousal_income: DiscreteState) -> IntND: