M1: config-driven BQSEGM stochastic-node splay

The BQSEGM continuation integrates the child's stochastic next-states over their
joint node mesh; reading it in one pass scales the peak intermediate with the
full ride-along x node x child-grid product, which on the M1 slice (aime=38 fixed
PIA-bend grid x many ride-along axes) is enormous. Expose
n_bqsegm_stochastic_node_batch_size on GridConfig and pass it into the solver, so
a memory-rich device reads the whole mesh in one pass (0, default) while a tighter
budget loops it in blocks. Resolves the M1 solve OOM together with the upstream
pylcm splay threading.

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

Author: hmgaudecker

src/aca_model/baseline/regimes/_bqsegm.py

@@ -40,6 +40,10 @@ def build_bqsegm_solver(grids: Grids) -> BQSEGM:
         continuous_state="assets",
         budget_target="resources",
         post_decision_function="savings",
+        # Splay the child stochastic-node expectation per the grid config: `0` (the
+        # default) reads the whole node mesh in one pass on a memory-rich device; a
+        # positive value loops it in blocks to fit a tighter budget (a CPU run).
+        stochastic_node_batch_size=grids.grid_config.n_bqsegm_stochastic_node_batch_size,
     )
 
 

src/aca_model/config.py

@@ -84,6 +84,14 @@ class GridConfig:
     # gathered grid (value function unchanged). `0` keeps the whole grid in one
     # kernel. Only consulted under `solver="dcegm"`.
     n_savings_batch_size: int = 0
+    # Block size for splaying the BQSEGM continuation's child stochastic-node
+    # expectation (health, health-cost shocks, the wage residual). `0` reads the
+    # whole joint node mesh in one pass — fast, but its peak intermediate scales
+    # with the full ride-along × node × child-grid product. A positive value loops
+    # the mesh in blocks of that size, trading runtime for a much smaller peak; `1`
+    # (one node at a time) is the memory-minimal setting for a CPU validation grid.
+    # Only consulted under `solver="bqsegm"`.
+    n_bqsegm_stochastic_node_batch_size: int = 0
 
 
 MODEL_CONFIG = ModelConfig()