Bump optimagic, boilerplate

.github/workflows/main.yml

@@ -28,7 +28,7 @@ jobs:
       - uses: actions/checkout@v7
       - uses: prefix-dev/setup-pixi@v0.10.0
         with:
-          pixi-version: v0.71.2
+          pixi-version: v0.74.0
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: tests-cpu
@@ -54,7 +54,7 @@ jobs:
       - uses: actions/checkout@v7
       - uses: prefix-dev/setup-pixi@v0.10.0
         with:
-          pixi-version: v0.71.2
+          pixi-version: v0.74.0
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           frozen: true

.github/workflows/publish-to-pypi.yml

@@ -16,7 +16,7 @@ jobs:
           # (including tags) must be fetched.
           fetch-depth: 0
       - name: Set up Python 3.12
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@v7
         with:
           python-version: '3.12'
       - name: Install pypa/build

.pre-commit-config.yaml

@@ -5,7 +5,7 @@ repos:
       - id: check-hooks-apply
       - id: check-useless-excludes
   - repo: https://github.com/tox-dev/pyproject-fmt
-    rev: v2.25.1
+    rev: v2.26.0
     hooks:
       - id: pyproject-fmt
   - repo: https://github.com/lyz-code/yamlfix
@@ -46,8 +46,12 @@ repos:
     rev: v1.38.0
     hooks:
       - id: yamllint
+  - repo: https://github.com/python-jsonschema/check-jsonschema
+    rev: 0.37.4
+    hooks:
+      - id: check-github-workflows
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.15.20
+    rev: v0.16.0
     hooks:
       - id: ruff-check
         args:
@@ -62,7 +66,7 @@ repos:
           - pyi
           - python
   - repo: https://github.com/astral-sh/ty-pre-commit
-    rev: v0.0.55
+    rev: v0.0.64
     hooks:
       - id: ty
         # `--no-project` stops uv from creating a `.venv`/`uv.lock` in this
@@ -79,7 +83,9 @@ repos:
         exclude: ^docs/getting_started/tutorial\.ipynb$
         args:
           - --extra-keys
-          - metadata.kernelspec metadata.language_info.version metadata.vscode
+          # metadata.kernelspec is kept: this project uses Jupyter Book 2
+          # (mystmd), which needs kernelspec to select the execution kernel.
+          - metadata.language_info.version metadata.vscode
   - repo: https://github.com/executablebooks/mdformat
     rev: 1.0.0
     hooks:
@@ -92,6 +98,14 @@ repos:
           - --wrap
           - '88'
         files: (AGENTS\.md|CLAUDE\.md|README\.md|modules/.*\.md|profiles/.*\.md)
+      - id: mdformat
+        additional_dependencies:
+          - mdformat-myst
+          - mdformat-ruff
+        args:
+          - --wrap
+          - '88'
+        files: (docs/.*)
 ci:
   autoupdate_schedule: monthly
   # pre-commit.ci has no pixi environments and blocks network at hook runtime;

AGENTS.md

@@ -0,0 +1,269 @@
+@.ai-instructions/profiles/tier-a.md @.ai-instructions/modules/jax.md
+@.ai-instructions/modules/optimagic.md @.ai-instructions/modules/pandas.md
+@.ai-instructions/modules/plotting.md
+
+# skillmodels
+
+This file provides guidance to AI coding agents when working with code in this
+repository.
+
+## Project Overview
+
+skillmodels is a Python implementation of estimators for nonlinear dynamic latent factor
+models, primarily used for skill formation research in economics. It implements Kalman
+filter-based maximum likelihood estimation following Cunha, Heckman, Schennach (2010).
+
+Used as the core estimation engine by sibling application projects (`skane-struct-bw`,
+`health-cognition`) in the parent workspace.
+
+## Development Commands
+
+```bash
+# Run tests
+pixi run -e tests-cpu tests
+
+# Run tests with coverage
+pixi run -e tests-cpu tests-with-cov
+
+# Run a single test file
+pixi run -e tests-cpu pytest tests/test_kalman_filters.py
+
+# Run a single test
+pixi run -e tests-cpu pytest tests/test_kalman_filters.py::test_function_name
+
+# Type checking (ty runs as a pre-commit hook, not a pixi task)
+prek run ty --all-files
+
+# Quality checks (linting, formatting, type checking)
+prek run --all-files
+
+# Build documentation (mystmd, from docs/ directory)
+myst build
+```
+
+## Command Rules
+
+Always use these command mappings:
+
+- **Python**: Use `pixi run python` instead of `python` or `python3`
+- **Type checker**: ty runs as a pre-commit hook (`prek run ty --all-files`), not as a
+  pixi task — it resolves imports from the env in `[tool.ty] environment.python`
+- **Tests**: Use `pixi run -e tests-cpu tests` instead of `pytest` directly
+- **Linting/formatting**: Use `prek run --all-files` instead of `ruff` directly
+- **All quality checks**: Use `prek run --all-files`
+
+Before finishing any task that modifies code, always run:
+
+1. `pixi run -e tests-cpu tests` (tests)
+1. `prek run --all-files` (quality checks, including the ty pre-commit hook)
+
+## Architecture
+
+### Core Pipeline Flow
+
+```
+ModelSpec + Data
+       ↓
+process_model() → Validates/extends model specification → ProcessedModel
+       ↓
+get_maximization_inputs() → Creates optimization problem (likelihood, gradients,
+                            constraints, params_template)
+       ↓
+[optimagic maximize / estimagic estimate_ml with fides algorithm]
+       ↓
+get_filtered_states() → Extract estimated latent factors
+simulate_dataset() → Simulate states (with optional policy effects)
+```
+
+### Key Modules
+
+- **model_spec.py**: User-facing frozen dataclasses (`ModelSpec`, `FactorSpec`,
+  `AnchoringSpec`). Re-exports `EstimationOptions` and `Normalizations` from `types.py`.
+  `ModelSpec` supports construction via `__init__` or `ModelSpec.from_dict()`, and
+  fluent builder methods: `with_transition_functions()`, `with_added_factor()`,
+  `with_added_observed_factors()`, `with_estimation_options()`, `with_anchoring()`,
+  `with_controls()`, `with_stagemap()`.
+- **types.py**: Internal frozen dataclasses (`ProcessedModel`, `Labels`, `Dimensions`,
+  `Anchoring`, `ParsingInfo`, `ParsedParams`, `EndogenousFactorsInfo`, etc.),
+  `EstimationOptions`, `Normalizations`, and immutability utilities.
+- **process_model.py**: Model specification validation and preprocessing. Converts
+  `ModelSpec` into `ProcessedModel`.
+- **kalman_filters.py**: Core Kalman filter implementation (predict/update steps). Uses
+  square-root form for numerical stability.
+- **likelihood_function.py** / **likelihood_function_debug.py**: Log-likelihood
+  computation using Kalman filtering. The debug variant is not jitted and returns
+  intermediate results (residuals, contributions, filtered states).
+- **constraints.py**: Generates parameter constraints (bounds, equalities from stagemap,
+  fixed values) for optimization. Exports `get_constraints()`,
+  `enforce_fixed_constraints()`, `add_bounds()`, `FixedConstraintWithValue`.
+- **parse_params.py**: Converts flat parameter vectors to structured model parameters.
+  Exports `create_parsing_info()` and `parse_params()`.
+- **params_index.py**: Builds the `pd.MultiIndex` for the params DataFrame via
+  `get_params_index()`.
+- **transition_functions.py**: Pre-built transition equations: `linear`, `translog`,
+  `robust_translog`, `linear_and_squares`, `log_ces`, `log_ces_general`, `constant`.
+- **decorators.py**: `register_params` decorator for custom transition functions. Tags a
+  callable with `__registered_params__` so skillmodels knows its parameter names.
+- **process_data.py**: `process_data()` for internal estimation format,
+  `pre_process_data()` for reshaping data to long format with period indexing.
+- **simulate_data.py**: `simulate_dataset()` and `simulate_policy_effect()`.
+- **diagnostic_plots.py**: `plot_likelihood_contributions()` and
+  `plot_residual_boxplots()` (Plotly-based).
+- **variance_decomposition.py**: `decompose_measurement_variance()` and
+  `summarize_measurement_reliability()`.
+- **process_debug_data.py**: `create_state_ranges()` and `process_debug_data()` for
+  converting raw debug output into DataFrames.
+- **utilities.py**: Model manipulation helpers (`extract_factors`, `remove_factors`,
+  `update_parameter_values`, `switch_translog_to_linear`, etc.).
+- **Visualization modules** (not in `__all__`, imported by module path):
+  `correlation_heatmap.py`, `visualize_factor_distributions.py`,
+  `visualize_transition_equations.py`, `utils_plotting.py`.
+
+### `get_maximization_inputs()` Return Dict
+
+Returns a dict with 6 keys:
+
+- `"loglike"`: `(params: pd.DataFrame) -> float` — jitted scalar log-likelihood
+- `"loglikeobs"`: `(params: pd.DataFrame) -> NDArray` — jitted per-observation
+  log-likelihood
+- `"debug_loglike"`: `(params: pd.DataFrame) -> dict` — non-jitted, returns dict with
+  keys `value`, `contributions`, `residuals`, `residual_sds`, `filtered_states`,
+  `state_ranges`, etc.
+- `"loglike_and_gradient"`: `(params: pd.DataFrame) -> tuple[float, NDArray]`
+- `"constraints"`: list of optimagic constraint objects
+- `"params_template"`: `pd.DataFrame` with correct MultiIndex and bounds; fixed
+  constraints pre-applied
+
+### ProcessedModel Key Attributes
+
+Applications frequently access these after calling `process_model(model_spec)`:
+
+- `processed.labels` — `.latent_factors`, `.observed_factors`, `.all_factors`,
+  `.controls`, `.stagemap`, `.stages`, `.aug_periods_to_periods`
+- `processed.dimensions` — `.n_periods`, `.n_latent_factors`
+- `processed.update_info` — DataFrame indexed by `(aug_period, variable)` with factor
+  columns and a `purpose` column
+- `processed.endogenous_factors_info` — `.has_endogenous_factors`,
+  `.aug_periods_from_period(period)`, `.factor_info`
+- `processed.normalizations` — dict of factor name to `Normalizations`
+- `processed.transition_info` — `.func` (vectorized), `.individual_functions[factor]`
+
+### JAX Usage
+
+All computation-heavy code uses JAX for automatic differentiation and JIT compilation.
+The codebase uses:
+
+- `jax.vmap` for vectorization across observations
+- `jax.jit` for compilation
+- JAX arrays throughout the estimation pipeline
+- Optional GPU support via CUDA or Metal
+
+### Public API (`__init__.py`)
+
+**Model specification classes:**
+
+- `ModelSpec`, `FactorSpec`, `AnchoringSpec`, `EstimationOptions`, `Normalizations`
+
+**Core estimation:**
+
+- `get_maximization_inputs(model_spec, data, split_dataset=1)` — prepare optimization
+  problem
+- `get_filtered_states(model_spec, data, params)` — returns nested dict with
+  `"anchored_states"` and `"unanchored_states"`, each containing `"states"` (DataFrame)
+  and `"state_ranges"`
+
+**Simulation:**
+
+- `simulate_dataset(model_spec, params, n_obs=None, data=None, policies=None, seed=None)`
+  — returns dict with `"unanchored_states"`, `"anchored_states"`
+- `simulate_policy_effect(model_spec, params, data, policies, seed=None)` — returns
+  DataFrame of factor mean differences between policy and baseline
+
+**Diagnostics and visualization:**
+
+- `plot_likelihood_contributions(model_spec, data, params, period=None)`
+- `plot_residual_boxplots(model_spec, data, params, period=None)`
+- `decompose_measurement_variance(model_spec, params, data)` — returns DataFrame indexed
+  by `(period, measurement, factor)` with signal/noise columns
+- `summarize_measurement_reliability(variance_decomposition)`
+- `create_state_ranges(filtered_states, factors, quantile_cutoff=None)`
+
+### Frequently Used Internal APIs
+
+These are not in `__all__` but are imported directly by application projects:
+
+- `skillmodels.process_model.process_model` — central to all application code
+- `skillmodels.types.ProcessedModel`, `EndogenousFactorsInfo`
+- `skillmodels.decorators.register_params` — essential for custom transition functions
+- `skillmodels.constraints.get_constraints`, `enforce_fixed_constraints`,
+  `FixedConstraintWithValue`, `select_by_loc`
+- `skillmodels.utilities.extract_factors`, `update_parameter_values`
+- `skillmodels.process_data.pre_process_data`
+- `skillmodels.correlation_heatmap.get_measurements_corr`, `get_quasi_scores_corr`,
+  `get_scores_corr`, `plot_correlation_heatmap`
+- `skillmodels.visualize_factor_distributions.univariate_densities`,
+  `bivariate_density_contours`, `combine_distribution_plots`
+- `skillmodels.visualize_transition_equations.get_transition_plots`,
+  `combine_transition_plots`
+- `skillmodels.parse_params.create_parsing_info`, `parse_params`
+- `skillmodels.params_index.get_params_index`
+
+## Code Style
+
+- Require Python 3.14
+- Uses Ruff for linting (target: Python 3.14, line length: 88)
+- Google-style docstrings with imperative mood ("Return" not "Returns")
+- Use MyST syntax in docstrings (single backticks `like this`), not reStructuredText (no
+  double backticks, no `:ref:`, `:func:`, etc.)
+- Dataclass attributes use inline docstrings (docstring on the line after the field):
+  ```python
+  name: str
+  """Description of name."""
+  ```
+- Pre-commit hooks enforce formatting and linting
+- Type checking via `ty` with strict rules
+- Do not use `from __future__ import annotations`
+- Use modern numpy random API: `rng = np.random.default_rng(seed)` instead of
+  `np.random.seed()` or legacy functions like `np.random.randn()`
+
+### Immutability Conventions
+
+- All model configuration and internal data structures use frozen dataclasses
+- Dict fields on internal dataclasses use `MappingProxyType` (not `Mapping`); wrap at
+  the call site with `MappingProxyType(...)`
+- Dict fields on user-facing dataclasses (`AnchoringSpec`, `Normalizations`) use
+  `Mapping` with `__post_init__` conversion via `ensure_containers_are_immutable()`
+- List fields use `tuple`, set fields use `frozenset`
+- `ensure_containers_are_immutable()` recursively converts dict→MappingProxyType,
+  list→tuple, set→frozenset
+
+### Period vs Aug_period
+
+Models with endogenous factors split each calendar period into multiple **augmented
+periods** (`aug_period`). The public API uses `period` (user-facing); `aug_period` is
+strictly internal. All public functions now return `period`:
+
+- `ModelSpec` — clean, no `aug_period` exposure.
+- `get_transition_plots()` — clean, accepts `period`/`periods`.
+- `get_filtered_states()` — clean, returns `period` column.
+- `simulate_dataset()` — clean, returns `period` in states DataFrames.
+- `plot_residual_boxplots()` / `plot_likelihood_contributions()` — clean, accept and
+  return `period`.
+- `decompose_measurement_variance()` — clean, indexed by
+  `(period, measurement, factor)`.
+- `simulate_policy_effect()` / `simulate_dataset()` policies — accept `"period"` key.
+- `ProcessedModel.labels` — exposes `aug_periods_to_periods` mapping (acceptable for
+  internal/advanced use).
+
+When writing new public-facing code, always accept and return `period`. Convert to
+`aug_period` internally using `ProcessedModel.labels.aug_periods_to_periods`.
+
+## Testing
+
+- pytest with markers: `wip`, `unit`, `integration`, `end_to_end`, `long_running`
+- Test files mirror source structure in `tests/`
+- Memory profiling available via pytest-memray (Unix only)
+- MATLAB AF CES / translog reproduction tests live in the parent workspace at
+  `../matlab_ces_repro/` (alongside `sim_repro/`), not in this library. They depend on
+  reference data at `/home/hmg/sciebo/Skill estimation/` and the CNLSY xls bundled
+  beside them. Run from the workspace root.