OpenSourceEconomics · hmgaudecker · May 22, 2026 · May 22, 2026 · May 23, 2026 · May 23, 2026
diff --git a/.github/workflows/main.yml b/.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 feat/distributed-V-arrays / PR #364 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@feat/distributed-V-arrays"
       - name: Install aca-model with test deps
         run: pip install -e . pytest pdbp
       - name: Run pytest

diff --git a/.gitignore b/.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
diff --git a/.pre-commit-config.yaml b/.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:

diff --git a/pyproject.toml b/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]
+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 +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]

diff --git a/src/aca_model/__init__.py b/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),
+)
diff --git a/src/aca_model/_benchmark_data/benchmark_params.pkl b/src/aca_model/_benchmark_data/benchmark_params.pkl
diff --git a/src/aca_model/_version.py b/src/aca_model/_version.py
diff --git a/src/aca_model/aca/model.py b/src/aca_model/aca/model.py
@@ -25,6 +25,7 @@ def create_model(
     derived_categoricals: Mapping[str, DiscreteGrid],
     grid_config: GridConfig,
     pref_type_grid: DiscreteGrid,
+    subjects_batch_size: int = 0,
 ) -> Model:
     """Create an ACA policy variant model.
 
@@ -45,6 +46,10 @@ def create_model(
             `pref_type`.
         grid_config: Continuous-grid point counts.
         pref_type_grid: Pref-type `DiscreteGrid`.
+        subjects_batch_size: Per-device chunk size for the simulate-side
+            per-subject dispatch. `0` (default) keeps a single vmap over
+            all subjects; `>0` chunks each device's local shard via
+            `jax.lax.map`. Tune via `grid_config.get_subjects_batch_size(log_level)`.
 
     Returns:
         pylcm Model.
@@ -71,4 +76,5 @@ def create_model(
         fixed_params=fixed_params,
         derived_categoricals=derived_categoricals,
         n_subjects=n_subjects,
+        subjects_batch_size=subjects_batch_size,
     )
diff --git a/src/aca_model/agent/preferences.py b/src/aca_model/agent/preferences.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 ** (

diff --git a/src/aca_model/baseline/model.py b/src/aca_model/baseline/model.py
@@ -28,6 +28,7 @@ def create_model(
     derived_categoricals: Mapping[str, DiscreteGrid],
     grid_config: GridConfig,
     pref_type_grid: DiscreteGrid,
+    subjects_batch_size: int = 0,
 ) -> Model:
     """Create the baseline structural retirement model.
 
@@ -52,6 +53,10 @@ def create_model(
         pref_type_grid: Pref-type `DiscreteGrid`. Pass
             `DiscreteGrid(PrefType)` for the production 3-type layout,
             or a compact variant (e.g. `DiscreteGrid(BenchmarkPrefType)`).
+        subjects_batch_size: Per-device chunk size for the simulate-side
+            per-subject dispatch. `0` (default) keeps a single vmap over
+            all subjects; `>0` chunks each device's local shard via
+            `jax.lax.map`. Tune via `grid_config.get_subjects_batch_size(log_level)`.
 
     Returns:
         A pylcm Model with 19 regimes (18 non-terminal + dead) spanning
@@ -78,4 +83,5 @@ def create_model(
         fixed_params=fixed_params,
         derived_categoricals=derived_categoricals,
         n_subjects=n_subjects,
+        subjects_batch_size=subjects_batch_size,
     )