Endogenize investments (#80)

* Revert to using 64-bit precision.

* Add function to check whether the model has investments.

* [WIP] Refactor: Include stuff in process_model / process_data.

* Backup commit before removing fill_list() feature.

* Fix tests, model processing seems to work with additional periods now. Missing data and constraints.

* Remove fill_list().

* [WIP] re-write process_data s.t. it returns a dictionary. Not working yet with investments.

* [WIP] Test for augmenting data with investments.

* Stack data for investments.

* Fix interfaces.

* Readability.

* Update environment.

* Jax 0.5 installed via pypi, only way that works right now.

* [WIP] implement fixed constraints right away.

* Refactor, enforce bounds on fixed parameters.

* Refactor test for augmented data so we can re-use the fixture.

* Implement restrictions for augmented periods. Allow stages + investments.

* Update pixi on GHA.

* Fix requirements for MacOS.

* Correct misnomer.

* Update & run hooks.

* In case investments are present, leave out entire 'raw' period (=2 skillmodels periods) for transitions.

* Fast QR Factorization on GPU (#81)

* Implement QR Factorization

* Move QR and add tests

* Fix Seed, Split tests

* Jax via conda.

* Fix (?) periods for shocks.

* Remove the last two internal periods for transitions if investments are present.

* Do not accidentally modify the params_template when enforcing constraints.

* Use local optimagic.

* Clip at -1e30 in case we end up running in 32bit mode. Use bounds_distance instead of 0, stick to dicts for the moment though code for custom constraints is there.

* Reduce memory usage (#82)

* Move Checkpoints

* Add splitting for gradient

* Speed up QR

* Update hooks and GHA pipeline.

* Use optimagic from GitHub.

* Fix heatmaps, update hooks etc.

* Fix heatmaps also for factor correlations.

* Got started updating transition equations.

* Be pedantic about using 'periods' for interface and 'aug_periods' for internal usage.

* Implement Janos' suggestion to only ever talk of endogenous factors.

* Get rid of deprecation warning.

* Use 'aug_period' as the index.

* More changes of periods -> aug_periods.

* Use 'aug_period' also in DataFrame.

* Use 'aug_period' as index also in update_info.

* More fixes to aug_period / period.

* Do not include correction factors as dependend variables in transition plots by default.

* Fix tests, basically allowing non-augmented models again.

* Update actions/pyproject.

* Use codecov token again.

* Fix docstring.

* Update GHA, pre-commit hooks.

* Add CLAUDE.md

* Move from mypy -> ty. Allow anchoring and endogenous factors at the same time.

---------

Co-authored-by: Max Jahn <max.jahn45@gmail.com>

Author: hmgaudecker

.github/workflows/main.yml

@@ -23,51 +23,27 @@ jobs:
           - macos-latest
           - windows-latest
         python-version:
-          - '3.12'
+          - '3.13'
     steps:
-      - uses: actions/checkout@v4
-      - uses: prefix-dev/setup-pixi@v0.8.0
+      - uses: actions/checkout@v6
+      - uses: prefix-dev/setup-pixi@v0.9.3
         with:
-          pixi-version: v0.29.0
+          pixi-version: v0.62.2
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: test-cpu
           activate-environment: true
+          frozen: true
       - name: Run pytest
+        if: runner.os != 'Linux' || matrix.python-version != '3.13'
+        shell: bash -l {0}
+        run: pixi run -e test-cpu tests
+      - name: Run pytest with coverage
+        if: runner.os == 'Linux' && matrix.python-version == '3.13'
         shell: bash -l {0}
         run: pixi run -e test-cpu tests-with-cov
       - name: Upload coverage report
-        if: runner.os == 'Linux' && matrix.python-version == '3.12'
-        uses: codecov/codecov-action@v4
+        if: runner.os == 'Linux' && matrix.python-version == '3.13'
+        uses: codecov/codecov-action@v5
         env:
           CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
-  # run-mypy:
-  #   name: Run mypy on Python 3.12
-  #   runs-on: ubuntu-latest
-  #   strategy:
-  #     fail-fast: false
-  #   steps:
-  #     - uses: actions/checkout@v4
-  #     - uses: prefix-dev/setup-pixi@v0.8.0
-  #       with:
-  #         pixi-version: v0.28.2
-  #         cache: true
-  #         cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
-  #         environments: mypy
-  #     - name: Run mypy
-  #       shell: bash -l {0}
-  #       run: pixi run mypy
-  # run-explanation-notebooks:
-  #   name: Run explanation notebooks on Python 3.12
-  #   runs-on: ubuntu-latest
-  #   steps:
-  #     - uses: actions/checkout@v4
-  #     - uses: prefix-dev/setup-pixi@v0.8.0
-  #       with:
-  #         pixi-version: v0.28.2
-  #         cache: true
-  #         cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
-  #         environments: test
-  #     - name: Run explanation notebooks
-  #       shell: bash -l {0}
-  #       run: pixi run -e test explanation-notebooks

.pre-commit-config.yaml

@@ -6,11 +6,11 @@ repos:
       - id: check-useless-excludes
       # - id: identity  # Prints all files passed to pre-commits. Debugging.
   - repo: https://github.com/lyz-code/yamlfix
-    rev: 1.17.0
+    rev: 1.19.1
     hooks:
       - id: yamlfix
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v6.0.0
     hooks:
       - id: check-added-large-files
         args:
@@ -30,43 +30,45 @@ repos:
         args:
           - --fix=lf
         description: Forces to replace line ending by the UNIX 'lf' character.
+      - id: no-commit-to-branch
+        args:
+          - --branch
+          - main
       - id: name-tests-test
         args:
           - --pytest-test-first
       - id: trailing-whitespace
       - id: check-ast
       - id: check-docstring-first
   - repo: https://github.com/adrienverge/yamllint.git
-    rev: v1.35.1
+    rev: v1.37.1
     hooks:
       - id: yamllint
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.4
+    rev: v0.14.10
     hooks:
-      # Run the linter.
-      - id: ruff
+      - id: ruff-check
         types_or:
           - python
           - pyi
           - jupyter
         args:
           - --fix
           # - --unsafe-fixes
-      # Run the formatter.
       - id: ruff-format
         types_or:
           - python
           - pyi
           - jupyter
   - repo: https://github.com/kynan/nbstripout
-    rev: 0.7.1
+    rev: 0.8.2
     hooks:
       - id: nbstripout
         args:
           - --extra-keys
           - metadata.kernelspec metadata.language_info.version metadata.vscode
   - repo: https://github.com/executablebooks/mdformat
-    rev: 0.7.17
+    rev: 1.0.0
     hooks:
       - id: mdformat
         additional_dependencies:

CLAUDE.md

@@ -0,0 +1,98 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) 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).
+
+## Development Commands
+
+```bash
+# Run tests
+pixi run tests
+
+# Run tests with coverage
+pixi run tests-with-cov
+
+# Run a single test file
+pixi run -e test-cpu pytest tests/test_kalman_filters.py
+
+# Run a single test
+pixi run -e test-cpu pytest tests/test_kalman_filters.py::test_function_name
+
+# Type checking
+pixi run ty
+
+# Install pre-commit hooks (required before committing)
+pre-commit install
+
+# Build documentation (from docs/ directory)
+make html
+```
+
+## Architecture
+
+### Core Pipeline Flow
+
+```
+Model Dict + Data
+       ↓
+process_model() → Validates/extends model specification
+       ↓
+process_data() → Transforms data to estimation format
+       ↓
+get_maximization_inputs() → Creates optimization problem (likelihood, gradients, constraints)
+       ↓
+[optimagic optimization]
+       ↓
+get_filtered_states() → Extract estimated latent factors
+```
+
+### Key Modules
+
+- **process_model.py**: Model specification validation and preprocessing. Handles
+  dimensions, labels, stagemap, anchoring, and endogenous factors.
+- **kalman_filters.py**: Core Kalman filter implementation (predict/update steps). Uses
+  square-root form for numerical stability.
+- **likelihood_function.py**: Log-likelihood computation using Kalman filtering.
+  Includes soft clipping for numerical stability.
+- **constraints.py**: Generates parameter constraints (bounds, equalities from stagemap,
+  fixed values) for optimization.
+- **parse_params.py**: Converts flat parameter vectors to structured model parameters.
+- **transition_functions.py**: Pre-built transition equations (`linear`, `log_ces`,
+  `constant`). Custom functions can be added.
+
+### 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
+
+### Public API
+
+The main package exports three functions:
+
+- `get_maximization_inputs()`: Prepare optimization problem for parameter estimation
+- `get_filtered_states()`: Extract filtered latent factor estimates
+- `simulate_dataset()`: Generate synthetic data from model specification
+
+## Code Style
+
+- Uses Ruff for linting (target: Python 3.13, line length: 88)
+- Google-style docstrings
+- Pre-commit hooks enforce formatting and linting
+- Type checking via `ty` with strict rules
+
+## Testing
+
+- pytest with markers: `wip`, `unit`, `integration`, `end_to_end`
+- Test files mirror source structure in `tests/`
+- Memory profiling available via pytest-memray (Unix only)

docs/source/getting_started/tutorial.ipynb

@@ -289,9 +289,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from om.process_constraints import process_constraints\n",
-    "\n",
-    "pc, pp = process_constraints(constraints, params)\n",
+    "pc, pp = om.process_constraints(constraints, params)\n",
     "params[\"group\"] = params.index.get_level_values(\"category\")\n",
     "params.loc[\"controls\", \"group\"] = params.loc[\"controls\"].index.get_level_values(\"name2\")\n",
     "\n",

docs/source/how_to_guides/model_specs.rst

@@ -96,14 +96,8 @@ Any python string is possible as factor name. The values are dictionaries with t
 entries:
 
 - measurements: A nested list that is as long as the number of periods of the model.
-  Each sublist contains the names of the measurements in that period. If A factor has
-  no measurements in a period, it has to be an empty list. If a factor only has
-  measurements up to a certain period you can leave out the empty lists at the end. In
-  the example this is done for factor 3. Note that even in that case, the measurements
-  have to be specified as nested list. If a factor only starts having measurements in
-  some period, you still have to specify the empty lists for all periods before that
-  period.
-
+  Each sublist contains the names of the measurements in that period. If a factor has
+  no measurements in a period, it has to be an empty list.
 
 - normalizations: This entry is optional. It is a dictionary that can have the keys
   ``"loadings"`` and ``"intercepts"``. The values are lists of dictionaries. The list
@@ -114,7 +108,7 @@ entries:
 - transition_equation: A string with the name of a pre-implemented transition equation
   or a custom transition equation. Pre-implemented transition equations are
   linear, log_ces (in the known location and scale version), constant and translog.
-  The example model dictionary only uses pre-implement transition functions.
+  The example model dictionary only uses pre-implemented transition functions.
 
   To see how to use custom transition functions, assume that the yaml file shown above
   has been loaded into a python dictionary called ``model`` and look at the following