create_model: forward n_subjects through baseline + aca + benchmark

Lets callers opt in to pylcm's simulate-AOT path
(`Model(n_subjects=...)`) without bypassing the aca-model factories.

Author: hmgaudecker

src/aca_model/aca/model.py

@@ -24,6 +24,7 @@ def create_model(
     derived_categoricals: Mapping[str, DiscreteGrid | Mapping[str, DiscreteGrid]]
     | None = None,
     grid_config: GridConfig = GRID_CONFIG,
+    n_subjects: int | None = None,
 ) -> Model:
     """Create an ACA policy variant model.
 
@@ -67,4 +68,5 @@ def create_model(
         description=f"Structural retirement model ({policy.name})",
         fixed_params=fixed_params,
         derived_categoricals=derived_categoricals,
+        n_subjects=n_subjects,
     )

src/aca_model/baseline/model.py

@@ -32,6 +32,7 @@ def create_model(
     | None = None,
     grid_config: GridConfig = GRID_CONFIG,
     pref_type_grid: DiscreteGrid | None = None,
+    n_subjects: int | None = None,
 ) -> Model:
     """Create the baseline structural retirement model.
 
@@ -79,6 +80,7 @@ def create_model(
         description="Baseline structural retirement model (pre-ACA)",
         fixed_params=fixed_params,
         derived_categoricals=derived_categoricals,
+        n_subjects=n_subjects,
     )
 
 

src/aca_model/benchmark.py

@@ -71,7 +71,11 @@
 )
 
 
-def create_benchmark_model(*, pref_type_grid: DiscreteGrid | None = None) -> Model:
+def create_benchmark_model(
+    *,
+    pref_type_grid: DiscreteGrid | None = None,
+    n_subjects: int | None = None,
+) -> Model:
     """Create the aca baseline with `BENCHMARK_GRID_CONFIG` and frozen fixed_params.
 
     The benchmark uses a 2-type `BenchmarkPrefType`. No `batch_size != 0`
@@ -86,6 +90,9 @@ def create_benchmark_model(*, pref_type_grid: DiscreteGrid | None = None) -> Mod
             (or `PARTITION_VMAP`) to get the partition-lifted kernel — the
             recommended production setting for aca-model at scale, but only
             supported on pylcm versions that expose `DispatchStrategy`.
+        n_subjects: Forwarded to `lcm.Model(n_subjects=...)`. When set, the
+            first matching `simulate(...)` call AOT-compiles all simulate
+            functions for that batch shape.
     """
     if pref_type_grid is None:
         pref_type_grid = DiscreteGrid(BenchmarkPrefType)
@@ -95,6 +102,7 @@ def create_benchmark_model(*, pref_type_grid: DiscreteGrid | None = None) -> Mod
         fixed_params=fixed_params,
         derived_categoricals=_DERIVED_CATEGORICALS,
         pref_type_grid=pref_type_grid,
+        n_subjects=n_subjects,
     )