Adapt aca-model to the pylcm #361 API restructure

Adopt pylcm's public `lcm/` / private `_lcm/` package split and the
accompanying API reorganisation: the `Regime` two-class split, the
grid renames (`Piece` → `PiecewiseGridSegment`, `*Process` classes),
the `FlatParams` rename, and the `regime` → `regime_id` / `regime_name`
distinction. Declare `distributed=True` on the assets grid for
multi-GPU sharding, activate the beartype claw on the package, and
apply the boilerplate update (hatch-vcs versioning, refreshed
pre-commit hooks, expanded `.gitignore`).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: hmgaudecker

.github/workflows/main.yml

@@ -1,5 +1,9 @@
 ---
 name: main
+# aca-model is a git submodule of the aca-dev workspace and has no pixi config
+# of its own — the pixi environments live in the parent workspace, whose
+# `tests-cpu` env has editable path-dependencies on private sibling repos that a
+# standalone CI runner cannot clone. CI therefore installs with pip directly.
 concurrency:
   group: ${{ github.head_ref || github.run_id }}
   cancel-in-progress: true
@@ -26,10 +30,10 @@ jobs:
       - uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
-      - name: Install pylcm (feature branch — revert to @main once pylcm#348/#350 merge)
+      - name: Install pylcm (pinned to the phase-2 branch until it merges to main)
         run: >-
           pip install "pylcm @
-          git+https://github.com/OpenSourceEconomics/pylcm.git@feat/runtime-grid-extra-params"
+          git+https://github.com/OpenSourceEconomics/pylcm.git@refactor/phase-2-api-reorganisation"
       - name: Install aca-model with test deps
         run: pip install -e . pytest pdbp
       - name: Run pytest

.gitignore

@@ -1,10 +1,59 @@
-__pycache__/
-*.py[cod]
+# Claude Code
+.claude/
+
+# Distribution / packaging
+*.egg
 *.egg-info/
-dist/
+*.manifest
+*.spec
+.eggs/
+.installed.cfg
 build/
-bld/
+dist/
+MANIFEST
+sdist/
+wheels/
+
+# IDE
+.idea/
+.vscode/
+
+# Jupyter / Jupyter Book
+.ipynb_checkpoints/
+_build
+
+# macOS
+.DS_Store
+
+# pixi
 .pixi/
+node_modules/
+
+# pytask
 .pytask/
+.pytask.sqlite3
+bld/
+out/
+pytask.lock
+pytask.lock.journal
+
+# Python
+__pycache__/
+*.py[cod]
+*.so
+*$py.class
+
+# Ruff
+.ruff_cache/
+
+# Testing
+.cache/
 .coverage
+.coverage.*
+.hypothesis/
+.pytest_cache/
+coverage.xml
 htmlcov/
+
+# Version file (generated by hatch-vcs)
+src/*/_version.py

.pre-commit-config.yaml

@@ -2,7 +2,9 @@
 repos:
   - repo: meta
     hooks:
-      - id: check-hooks-apply
+      # check-hooks-apply is omitted: aca-model ships no notebooks yet, so the
+      # boilerplate nbstripout hook matches nothing and that meta check would
+      # fail. Re-add it once the repo gains a notebook.
       - id: check-useless-excludes
   - repo: https://github.com/tox-dev/pyproject-fmt
     rev: v2.21.1
@@ -37,6 +39,7 @@ repos:
       - id: name-tests-test
         args:
           - --pytest-test-first
+        exclude: ^tests/helpers/
       - id: no-commit-to-branch
         args:
           - --branch
@@ -46,6 +49,10 @@ repos:
     rev: v1.38.0
     hooks:
       - id: yamllint
+  - repo: https://github.com/python-jsonschema/check-jsonschema
+    rev: 0.37.2
+    hooks:
+      - id: check-github-workflows
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.15.12
     hooks:
@@ -61,6 +68,13 @@ repos:
           - jupyter
           - pyi
           - python
+  - repo: https://github.com/kynan/nbstripout
+    rev: 0.9.1
+    hooks:
+      - id: nbstripout
+        args:
+          - --extra-keys
+          - metadata.kernelspec metadata.language_info.version metadata.vscode
   - repo: https://github.com/executablebooks/mdformat
     rev: 1.0.0
     hooks:

pyproject.toml

@@ -1,10 +1,9 @@
 [build-system]
 build-backend = "hatchling.build"
-requires = [ "hatchling" ]
+requires = [ "hatch-vcs", "hatchling" ]
 
 [project]
 name = "aca-model"
-version = "0.0.0"
 description = "Core lifecycle model for the ACA structural retirement project."
 readme = { file = "README.md", content-type = "text/markdown" }
 keywords = [
@@ -23,8 +22,10 @@ classifiers = [
   "Programming Language :: Python :: 3 :: Only",
   "Programming Language :: Python :: 3.14",
 ]
+dynamic = [ "version" ]
 dependencies = [
   "attrs",
+  "beartype",
   "cloudpickle",
   "dags",
   "estimagic",
@@ -43,13 +44,19 @@ email = "hmgaudecker@uni-bonn.de"
 [[project.maintainers]]
 name = "Hans-Martin von Gaudecker"
 email = "hmgaudecker@uni-bonn.de"
+[project.urls]
+Github = "https://github.com/OpenSourceEconomics/aca-model"
+Repository = "https://github.com/OpenSourceEconomics/aca-model"
+Tracker = "https://github.com/OpenSourceEconomics/aca-model/issues"
 
 [tool.hatch]
+build.hooks.vcs.version-file = "src/aca_model/_version.py"
 build.targets.sdist.exclude = [ "tests" ]
 build.targets.sdist.only-packages = true
 build.targets.wheel.only-include = [ "src" ]
 build.targets.wheel.sources = [ "src" ]
 metadata.allow-direct-references = true
+version.source = "vcs"
 
 [tool.ruff]
 fix = true
@@ -84,9 +91,21 @@ extend-ignore = [
   "RUF002",  # Ambiguous Unicode in docstrings (Greek letters in math)
   "RUF003",  # Ambiguous Unicode in comments (Greek letters in math)
 ]
-per-file-ignores."src/aca_model/models/*" = [ "E501" ]
-per-file-ignores."task_*.py" = [ "ANN", "ARG001" ]
-per-file-ignores."tests/*" = [ "D", "E501", "INP001", "PD011", "PLR2004", "S101" ]
+per-file-ignores."src/aca_model/models/*" = [
+  "E501", # Line too long (generated model files)
+]
+per-file-ignores."task_*.py" = [
+  "ANN",    # Type annotations (use ty instead)
+  "ARG001", # Unused function argument (pytask signatures)
+]
+per-file-ignores."tests/*" = [
+  "D",       # Docstrings
+  "E501",    # Line too long
+  "INP001",  # Implicit namespace package
+  "PD011",   # Use of .values (false positives on non-pandas objects)
+  "PLR2004", # Magic value used in comparison
+  "S101",    # Use of assert
+]
 pydocstyle.convention = "google"
 
 [tool.pyproject-fmt]

src/aca_model/__init__.py

@@ -1,3 +1,22 @@
 import jax
 
 jax.config.update("jax_enable_x64", True)
+
+# Import lcm before installing the claw so its `_jaxtyping_patch` (picklable
+# jaxtyping sentinel) and `MappingProxyType` pytree registration are in place.
+import lcm  # noqa: E402, F401
+
+# Install beartype's AST-rewriting claw on the whole `aca_model` package before
+# any submodule is imported. The claw transforms each module's AST at first
+# import to insert runtime type checks against its annotations; aca_model's
+# numerical DAG/transition/utility functions are otherwise unchecked, since
+# pylcm's own claw is scoped to `lcm.*`. Violations surface as beartype's
+# `BeartypeCallHintViolation` — aca_model is an application, not a library with
+# a documented exception contract.
+from beartype import BeartypeConf, BeartypeStrategy  # noqa: E402
+from beartype.claw import beartype_package  # noqa: E402
+
+beartype_package(
+    "aca_model",
+    conf=BeartypeConf(strategy=BeartypeStrategy.On, is_pep484_tower=True),
+)

src/aca_model/_benchmark_data/benchmark_params.pkl

@@ -140,9 +140,10 @@ def u_alive(
     coefficient_rra: FloatND,
     utility_scale_factor: FloatND,
 ) -> FloatND:
-    """Within-period utility for every non-dead regime: CES over consumption and leisure.
+    """Within-period utility for every non-dead regime.
 
-    `leisure` is a DAG input — supplied per-regime by `leisure_canwork_retiree_or_nongroup`,
+    CES over consumption and leisure. `leisure` is a DAG input — supplied
+    per-regime by `leisure_canwork_retiree_or_nongroup`,
     `leisure_canwork_tied`, or `leisure_forcedout`.
     """
     composite = consumption_equiv**consumption_weight * leisure ** (

src/aca_model/_version.py

@@ -10,20 +10,21 @@
 from typing import Any, Literal, TypedDict
 
 import jax.numpy as jnp
-import lcm.shocks.ar1
-import lcm.shocks.iid
 import numpy as np
+from _lcm.grids.continuous import ContinuousGrid
 from lcm import (
     DiscreteGrid,
     IrregSpacedGrid,
     LinSpacedGrid,
     MarkovTransition,
+    NormalIIDProcess,
+    PiecewiseGridSegment,
+    PiecewiseLinSpacedGrid,
     Regime,
+    RouwenhorstAR1Process,
     categorical,
 )
-from lcm.grids.continuous import ContinuousGrid
-from lcm.grids.piecewise import Piece, PiecewiseLinSpacedGrid
-from lcm.typing import BoolND, FloatND, RegimeName, ScalarInt, UserParams
+from lcm.typing import BoolND, FloatND, IntND, RegimeName, ScalarInt, UserParams
 
 from aca_model.agent import (
     assets_and_income,
@@ -34,7 +35,7 @@
 from aca_model.agent.health import Health, HealthWithDisability
 from aca_model.agent.labor_market import LaborSupply, LaggedLaborSupply, SpousalIncome
 from aca_model.baseline import health_insurance
-from aca_model.baseline.health_insurance import BuyPrivate, HealthInsuranceState
+from aca_model.baseline.health_insurance import BuyPrivate
 from aca_model.config import MODEL_CONFIG, GridConfig
 from aca_model.environment import social_security, taxes
 from aca_model.environment.social_security import ClaimedSS
@@ -237,14 +238,14 @@ def build_grids(
     # grid to have unconditional variance 1, the Rouwenhorst innovation
     # std must be √(1 − ρ²). Passing the σ_y itself (≈0.577 for hcc,
     # 0.5627 for wage) would mis-scale the grid.
-    wage_res = lcm.shocks.ar1.Rouwenhorst(
+    wage_res = RouwenhorstAR1Process(
         n_points=grid_config.n_wage_res_gridpoints,
         rho=_WAGE_RHO,
         sigma=(1.0 - _WAGE_RHO**2) ** 0.5,
         mu=0.0,
     )
     hcc_persistent = get_hcc_persistent_shock(grid_config=grid_config)
-    hcc_transitory = lcm.shocks.iid.Normal(
+    hcc_transitory = NormalIIDProcess(
         n_points=grid_config.n_hcc_transitory_gridpoints,
         gauss_hermite=True,
         mu=0.0,
@@ -261,11 +262,11 @@ def build_grids(
             stop=500_000.0,
             n_points=grid_config.n_assets_gridpoints,
             batch_size=grid_config.n_assets_batch_size,
+            distributed=True,
         ),
         aime=_build_aime_grid(grid_config=grid_config, fixed_params=fixed_params),
         consumption_dollars=IrregSpacedGrid(
             n_points=grid_config.n_consumption_dollars_gridpoints,
-            extra_param_names=("max_consumption_dollars",),
         ),
         wage_res=wage_res,
         hcc_persistent=hcc_persistent,
@@ -274,15 +275,15 @@ def build_grids(
     )
 
 
-def get_hcc_persistent_shock(*, grid_config: GridConfig) -> lcm.shocks.ar1.Rouwenhorst:
+def get_hcc_persistent_shock(*, grid_config: GridConfig) -> RouwenhorstAR1Process:
     """Return the persistent-HCC AR(1) shock grid for a given `grid_config`.
 
     Exposed so callers that need the shock's gridpoints / transition
     probs (e.g. `assemble_fixed_params`, the HCC insurer predictor)
     can derive them from `grid_config` alone without instantiating a
     full `Model`.
     """
-    return lcm.shocks.ar1.Rouwenhorst(
+    return RouwenhorstAR1Process(
         n_points=grid_config.n_hcc_persistent_gridpoints,
         rho=_HCC_RHO,
         sigma=(1.0 - _HCC_RHO**2) ** 0.5,
@@ -306,20 +307,26 @@ def _build_aime_grid(
     this path; the total is fixed by the PIA structure (32 points).
     """
     kinks = [float(k) for k in np.asarray(fixed_params["pia_aime_grid"])]
-    pieces = (
-        Piece(interval=f"[{kinks[0]}, {kinks[1]})", n_points=_AIME_PIECE_N_POINTS[0]),
-        Piece(interval=f"[{kinks[1]}, {kinks[2]})", n_points=_AIME_PIECE_N_POINTS[1]),
-        Piece(interval=f"[{kinks[2]}, {kinks[3]}]", n_points=_AIME_PIECE_N_POINTS[2]),
+    segments = (
+        PiecewiseGridSegment(
+            interval=f"[{kinks[0]}, {kinks[1]})", n_points=_AIME_PIECE_N_POINTS[0]
+        ),
+        PiecewiseGridSegment(
+            interval=f"[{kinks[1]}, {kinks[2]})", n_points=_AIME_PIECE_N_POINTS[1]
+        ),
+        PiecewiseGridSegment(
+            interval=f"[{kinks[2]}, {kinks[3]}]", n_points=_AIME_PIECE_N_POINTS[2]
+        ),
     )
     return PiecewiseLinSpacedGrid(
-        pieces=pieces, batch_size=grid_config.n_aime_batch_size
+        segments=segments, batch_size=grid_config.n_aime_batch_size
     )
 
 
 def _compute_max_annual_labor_income(
     *,
     wage_params: Mapping[str, Any],
-    wage_res_grid: lcm.shocks.ar1.Rouwenhorst,
+    wage_res_grid: RouwenhorstAR1Process,
 ) -> float:
     """Return the annual labor income at the top of the wage grid.
 
@@ -418,7 +425,7 @@ def build_actions(spec: RegimeSpec, grids: Grids) -> dict:
     return actions
 
 
-def build_regime_probs(target: FloatND, survival: FloatND) -> FloatND:
+def build_regime_probs(target: IntND, survival: FloatND) -> FloatND:
     """Build regime transition probability vector."""
     probs = jnp.zeros(19)
     probs = probs.at[RegimeId.dead].set(1.0 - survival)
@@ -603,10 +610,10 @@ def make_targets(name: str) -> tuple[dict[str, int], dict[str, int]]:
 
 
 def select_target_for_age(
-    next_age: int | FloatND,
+    next_age: int | IntND | FloatND,
     mc_next: bool | BoolND,
     tgts: dict[str, int],
-) -> FloatND:
+) -> IntND:
     """Select target regime ID based on next-period age bracket."""
     ss_choose = jnp.where(
         jnp.array(mc_next),