aca-model updates: sharding config, regimes, ACA overlay, tooling + ty

Squash-merge of feat/pytask-slurm-gpu-flags — GridConfig sharding cleanup
(discrete-only, ll×css flags removed), regime/ACA-overlay refinements, plus the
tooling bump (CI→pylcm main), ty fixes, and stale-test cleanup.

Co-Authored-By: Claude Opus 4.8 <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 from main
         run: >-
           pip install "pylcm @
-          git+https://github.com/OpenSourceEconomics/pylcm.git@feat/runtime-grid-extra-params"
+          git+https://github.com/OpenSourceEconomics/pylcm.git@main"
       - 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,10 +2,12 @@
 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
+    rev: v2.23.0
     hooks:
       - id: pyproject-fmt
   - repo: https://github.com/lyz-code/yamlfix
@@ -37,6 +39,7 @@ repos:
       - id: name-tests-test
         args:
           - --pytest-test-first
+        exclude: ^tests/helpers/
       - id: no-commit-to-branch
         args:
           - --branch
@@ -46,8 +49,12 @@ 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
+    rev: v0.15.16
     hooks:
       - id: ruff-check
         args:
@@ -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,11 +22,12 @@ classifiers = [
   "Programming Language :: Python :: 3 :: Only",
   "Programming Language :: Python :: 3.14",
 ]
+dynamic = [ "version" ]
 dependencies = [
   "attrs",
+  "beartype",
   "cloudpickle",
   "dags",
-  "estimagic",
   "jax>=0.9",
   "jaxtyping",
   "numpy>=2.2",
@@ -43,13 +43,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]
+version.source = "vcs"
+metadata.allow-direct-references = true
+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
 
 [tool.ruff]
 fix = true
@@ -84,9 +90,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]
@@ -111,13 +129,13 @@ rules.unused-ignore-comment = "error"
 
 [tool.pytest]
 ini_options.addopts = '--pdbcls=pdbp:Pdb --ignore-glob=*skip_* -m "not long_running"'
-ini_options.filterwarnings = []
 ini_options.markers = [
   """\
   long_running: Tests that take a lot of time to run (run them via pytest -m \
   long_running).\
   """,
 ]
+ini_options.filterwarnings = []
 
 [tool.yamlfix]
 line_length = 88

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

@@ -6,7 +6,9 @@
 swapping in the regime DAG.
 """
 
+from collections.abc import Mapping
 from enum import Enum, auto
+from typing import Any, cast
 
 import jax.numpy as jnp
 from lcm.params import MappingLeaf
@@ -43,7 +45,7 @@ def mandate_penalty(
     Penalty = clip(income * income_fraction, min, max) if uninsured and
     income above exemption threshold; 0 otherwise.
     """
-    sched = mandate_schedule.data
+    sched = cast("Mapping[str, Any]", mandate_schedule.data)
     is_uninsured = buy_private == BuyPrivate.no
     exempt = gross_income < sched["exempt_income"][spousal_income]
     raw = jnp.clip(
@@ -67,7 +69,7 @@ def premium_subsidy(
     FPL kink points, subsidy = max(0, premium - income * applicable_rate).
     Return 0 when buy_private==no or income outside 100-400% FPL range.
     """
-    sched = premium_credit_schedule.data
+    sched = cast("Mapping[str, Any]", premium_credit_schedule.data)
     kinks = sched["kinks"]  # [n_kinks, 3]
     frac_income = sched["frac_income"]  # [n_kinks]
 
@@ -92,7 +94,7 @@ def cost_sharing(
     Applied to deductible, coinsurance, and OOP max.
     Return 1.0 when buy_private==no (no reduction for uninsured).
     """
-    sched = cost_sharing_schedule.data
+    sched = cast("Mapping[str, Any]", cost_sharing_schedule.data)
     kinks = sched["kinks"]  # [n_kinks, 3]
     factors = sched["factors"]  # [n_kinks + 1]
     bracket = jnp.searchsorted(kinks[:, spousal_income], gross_income, side="right")
@@ -111,7 +113,8 @@ def is_medicaid_eligible(
     Unlike baseline SSI-based Medicaid, ACA expansion uses only income
     (no assets test, no Medicare requirement).
     """
-    threshold = medicaid_schedule.data["income_threshold"]
+    sched = cast("Mapping[str, Any]", medicaid_schedule.data)
+    threshold = sched["income_threshold"]
     return countable_income < threshold[spousal_income]
 
 

src/aca_model/_version.py

@@ -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 ** (