Restructure package (#301)

Co-authored-by: timmens <mensingertim@gmail.com>

Author: hmgaudecker

.pre-commit-config.yaml

@@ -95,5 +95,7 @@ repos:
           - --wrap
           - '88'
         files: (docs/.)
+        # mdformat-myst mangles sphinx-design grid-item-card directives
+        exclude: docs/index\.md
 ci:
   autoupdate_schedule: monthly

AGENTS.md

@@ -42,9 +42,12 @@ automation. Python 3.14+ is required.
 
 - `InternalRegime`: Internal representation after processing user regime
 - `StateActionSpace`: Manages state-action combinations for solution/simulation
-- `StateSpaceInfo`: Metadata for working with function outputs on state spaces
 - `PeriodRegimeSimulationData`: Raw simulation results for one period in one regime
 
+**Value Function Representation (`src/lcm/regime_building/V.py`)**
+
+- `VInterpolationInfo`: Metadata for working with function outputs on state spaces
+
 **Solution (`src/lcm/solution/`)**
 
 - `solve_brute.py`: Brute force dynamic programming solver using backward induction
@@ -109,7 +112,7 @@ to transition functions for target-dependent transitions.
 - `tests/test_models/`: Shared test models (deterministic, stochastic variants)
 - `tests/solution/`: Tests for solution algorithms
 - `tests/simulation/`: Tests for simulation functionality
-- `tests/input_processing/`: Tests for model processing pipeline
+- `tests/regime_building/`: Tests for regime compilation pipeline
 - `tests/data/`: Analytical solutions and regression test data
 
 ## Model and Regime Interface

docs/development/benchmarking.md

@@ -4,9 +4,9 @@ title: Benchmarking
 
 # Benchmarking
 
-pylcm uses [ASV (Airspeed Velocity)](https://asv.readthedocs.io/) to track
-performance across commits. Benchmarks run locally on GPU hardware and results are
-published to a [dashboard](https://open-econ.org/pylcm-benchmarks/).
+pylcm uses [ASV (Airspeed Velocity)](https://asv.readthedocs.io/) to track performance
+across commits. Benchmarks run locally on GPU hardware and results are published to a
+[dashboard](https://open-econ.org/pylcm-benchmarks/).
 
 ## Machine Setup
 
@@ -61,16 +61,16 @@ pixi run asv-publish
 ```
 
 The `asv-run` and `asv-quick` tasks set `XLA_PYTHON_CLIENT_PREALLOCATE=false` and
-`XLA_PYTHON_CLIENT_MEM_FRACTION=0.3` automatically to prevent JAX from grabbing all
-GPU memory.
+`XLA_PYTHON_CLIENT_MEM_FRACTION=0.3` automatically to prevent JAX from grabbing all GPU
+memory.
 
 ## Benchmark Scenarios
 
-| File | What it benchmarks |
-|---|---|
+| File                             | What it benchmarks                                                                                    |
+| -------------------------------- | ----------------------------------------------------------------------------------------------------- |
 | `bench_precautionary_savings.py` | Solve (varying grid sizes), simulate (varying subjects), solve+simulate, lin vs irreg grid comparison |
-| `bench_mortality.py` | Mortality model — solve + simulate |
-| `bench_mahler_yum.py` | Mahler & Yum (2024) replication (GPU only) |
+| `bench_mortality.py`             | Mortality model — solve + simulate                                                                    |
+| `bench_mahler_yum.py`            | Mahler & Yum (2024) replication (GPU only)                                                            |
 
 Each benchmark tracks three metrics:
 

docs/development/setup.md

@@ -6,8 +6,7 @@ title: Setup & Workflow
 
 ## Setup
 
-pylcm uses [pixi](https://pixi.sh/) for dependency management. Python 3.14+ is
-required.
+pylcm uses [pixi](https://pixi.sh/) for dependency management. Python 3.14+ is required.
 
 ```bash
 # Clone the repository

docs/examples/index.md

@@ -4,25 +4,25 @@ title: Examples
 
 # Examples
 
-Complete example models built with pylcm, ordered from simplest to most complex.
-All models live in the `lcm_examples` package and can be imported directly.
+Complete example models built with pylcm, ordered from simplest to most complex. All
+models live in the `lcm_examples` package and can be imported directly.
 
-1. **[Tiny Consumption-Savings](tiny.md)** — `lcm_examples.tiny`
-   Minimal 2-regime model. 3 periods, discrete labor, log-spaced grids,
-   tax-and-transfer system.
+1. **[Tiny Consumption-Savings](tiny.md)** — `lcm_examples.tiny` Minimal 2-regime model.
+   3 periods, discrete labor, log-spaced grids, tax-and-transfer system.
 
-2. **[Mortality](mortality.md)** — `lcm_examples.mortality`
-   3-regime model with death. Discrete labor, borrowing constraint.
+1. **[Mortality](mortality.md)** — `lcm_examples.mortality` 3-regime model with death.
+   Discrete labor, borrowing constraint.
 
-3. **[Precautionary Savings](precautionary_savings.md)** — `lcm_examples.precautionary_savings`
-   2-regime model with income shocks. IID, Rouwenhorst, and Tauchen shock types.
+1. **[Precautionary Savings](precautionary_savings.md)** —
+   `lcm_examples.precautionary_savings` 2-regime model with income shocks. IID,
+   Rouwenhorst, and Tauchen shock types.
 
-4. **[Precautionary Savings with Health](precautionary_savings_health.md)** — `lcm_examples.precautionary_savings_health`
-   2-regime model with health & exercise. Multiple continuous states and actions,
-   auxiliary functions, constraints.
+1. **[Precautionary Savings with Health](precautionary_savings_health.md)** —
+   `lcm_examples.precautionary_savings_health` 2-regime model with health & exercise.
+   Multiple continuous states and actions, auxiliary functions, constraints.
 
-5. **[Mahler & Yum (2024)](mahler_yum_2024.md)** — `lcm_examples.mahler_yum_2024`
-   Full Econometrica replication. 8 states, stochastic transitions, data files,
+1. **[Mahler & Yum (2024)](mahler_yum_2024.md)** — `lcm_examples.mahler_yum_2024` Full
+   Econometrica replication. 8 states, stochastic transitions, data files,
    discount-factor heterogeneity. **Requires GPU.**
 
 ## Customizing models

docs/examples/mahler_yum_2024.md

@@ -6,16 +6,14 @@ title: Mahler & Yum (2024)
 
 Full replication of the lifecycle model from @mahler2024.
 
-Two regimes (alive/dead), 8 states including health, education, productivity type,
-and health type. Three actions: labor supply, saving, and health effort. Features
-stochastic health and regime transitions, AR(1) productivity shocks, and
-discount-factor heterogeneity. Ships with calibrated data files for survival
-probabilities and initial distributions.
+Two regimes (alive/dead), 8 states including health, education, productivity type, and
+health type. Three actions: labor supply, saving, and health effort. Features stochastic
+health and regime transitions, AR(1) productivity shocks, and discount-factor
+heterogeneity. Ships with calibrated data files for survival probabilities and initial
+distributions.
 
-::::{important}
-This model is computationally intensive and requires GPU acceleration. Run it in a CUDA
-environment (e.g., `pixi run -e cuda13 python your_script.py`).
-::::
+::::\{important} This model is computationally intensive and requires GPU acceleration.
+Run it in a CUDA environment (e.g., `pixi run -e cuda13 python your_script.py`). ::::
 
 [View source on GitHub](https://github.com/OpenSourceEconomics/pylcm/blob/main/src/lcm_examples/mahler_yum_2024/_model.py)
 
@@ -42,10 +40,9 @@ beta_std = START_PARAMS["beta"]["std"]
 
 # Select initial states with high discount factor type
 selected_ids_high = jnp.flatnonzero(discount_factor_types)
-initial_states_high =   {
-                        state: values[selected_ids_high] for state, values
-                        in initial_states.items()
-                        }
+initial_states_high = {
+    state: values[selected_ids_high] for state, values in initial_states.items()
+}
 
 # Solve and simulate for high discount factor type
 result = MAHLER_YUM_MODEL.simulate(

docs/examples/mortality.md

@@ -5,9 +5,9 @@ title: Mortality
 # Mortality
 
 A three-regime consumption-savings model with stochastic mortality: working life,
-retirement, and death. The agent chooses labor supply and consumption. Log utility
-with work disutility and a borrowing constraint. Mortality is age-dependent: a vector
-of per-period survival probabilities (last entry = 0.0 for certain death) governs the
+retirement, and death. The agent chooses labor supply and consumption. Log utility with
+work disutility and a borrowing constraint. Mortality is age-dependent: a vector of
+per-period survival probabilities (last entry = 0.0 for certain death) governs the
 transition to the dead regime.
 
 [View source on GitHub](https://github.com/OpenSourceEconomics/pylcm/blob/main/src/lcm_examples/mortality.py)

docs/examples/precautionary_savings.md

@@ -4,12 +4,11 @@ title: Precautionary Savings
 
 # Precautionary Savings
 
-A two-regime consumption-savings model with income shocks (alive + dead). Supports
-IID shocks (Normal Gauss-Hermite) and persistent AR(1) shocks (Rouwenhorst, Tauchen).
+A two-regime consumption-savings model with income shocks (alive + dead). Supports IID
+shocks (Normal Gauss-Hermite) and persistent AR(1) shocks (Rouwenhorst, Tauchen).
 Configurable interest rate and wealth grid type.
 
-With FGP-calibrated parameters, this replicates the simplified benchmark of
-@fella2019.
+With FGP-calibrated parameters, this replicates the simplified benchmark of @fella2019.
 
 [View source on GitHub](https://github.com/OpenSourceEconomics/pylcm/blob/main/src/lcm_examples/precautionary_savings.py)
 

docs/examples/precautionary_savings_health.md

@@ -12,10 +12,8 @@ where utility depends only on remaining wealth and health.
 This model demonstrates multiple continuous states, multiple continuous actions,
 auxiliary functions (wage, labor income), constraints, and regime transitions.
 
-:::{note}
-The parameterization is chosen to showcase pylcm's features, not to match any empirical
-calibration.
-:::
+:::\{note} The parameterization is chosen to showcase pylcm's features, not to match any
+empirical calibration. :::
 
 [View source on GitHub](https://github.com/OpenSourceEconomics/pylcm/blob/main/src/lcm_examples/precautionary_savings_health.py)
 

docs/examples/tiny.md

@@ -4,8 +4,8 @@ title: Tiny Consumption-Savings
 
 # Tiny Consumption-Savings
 
-A minimal two-regime consumption-savings model with working life and retirement.
-Three periods, discrete labor supply, log-spaced consumption grid, and a simple
+A minimal two-regime consumption-savings model with working life and retirement. Three
+periods, discrete labor supply, log-spaced consumption grid, and a simple
 tax-and-transfer system that guarantees a consumption floor.
 
 This is the simplest complete model in `lcm_examples` and a good starting point for
@@ -23,11 +23,13 @@ from lcm_examples.tiny import get_model, get_params
 model = get_model()
 params = get_params()
 
-initial_df = pd.DataFrame({
-    "regime": "working_life",
-    "age": model.ages.values[0],
-    "wealth": np.linspace(1, 20, 100),
-})
+initial_df = pd.DataFrame(
+    {
+        "regime": "working_life",
+        "age": model.ages.values[0],
+        "wealth": np.linspace(1, 20, 100),
+    }
+)
 
 result = model.simulate(
     params=params,