Update model processing

.github/workflows/main.yml

@@ -28,7 +28,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: prefix-dev/[email protected]
         with:
-          pixi-version: v0.23.0
+          pixi-version: v0.29.0
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: test-cpu
@@ -48,7 +48,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: prefix-dev/[email protected]
         with:
-          pixi-version: v0.23.0
+          pixi-version: v0.29.0
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: mypy
@@ -62,7 +62,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: prefix-dev/[email protected]
         with:
-          pixi-version: v0.23.0
+          pixi-version: v0.29.0
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: test-cpu

.pre-commit-config.yaml

@@ -6,19 +6,9 @@ repos:
       - id: check-useless-excludes
         # - id: identity  # Prints all files passed to pre-commits. Debugging.
   - repo: https://github.com/lyz-code/yamlfix
-    rev: 1.16.0
+    rev: 1.17.0
     hooks:
       - id: yamlfix
-  - repo: local
-    hooks:
-      - id: check-nbqa-version-mismatch
-        name: Check for version mismatch between black, ruff, and nbQA
-        entry: python scripts/check_nbqa_version_mismatch.py
-        language: python
-        always_run: true
-        require_serial: true
-        additional_dependencies:
-          - pyyaml
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.6.0
     hooks:
@@ -55,41 +45,55 @@ repos:
     rev: v1.35.1
     hooks:
       - id: yamllint
-  - repo: https://github.com/psf/black
-    rev: 24.4.2
-    hooks:
-      - id: black
-        language_version: python3.12
-  - repo: https://github.com/asottile/blacken-docs
-    rev: 1.18.0
-    hooks:
-      - id: blacken-docs
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.5.0
+    rev: v0.6.3
     hooks:
+      # Run the linter.
       - id: ruff
-      #   args:
-      #     - --verbose
-  - repo: https://github.com/nbQA-dev/nbQA
-    rev: 1.8.5
-    hooks:
-      - id: nbqa-black
-        additional_dependencies:
-          - black==24.4.2
-      - id: nbqa-ruff
-        additional_dependencies:
-          - ruff==v0.5.0
+        types_or:
+          - python
+          - pyi
+          - jupyter
+        args:
+          - --fix
+      # Run the formatter.
+      - id: ruff-format
+        types_or:
+          - python
+          - pyi
+          - jupyter
   - repo: https://github.com/executablebooks/mdformat
     rev: 0.7.17
     hooks:
       - id: mdformat
         additional_dependencies:
           - mdformat-gfm
           - mdformat-gfm-alerts
-          - mdformat-black
+          - mdformat-ruff
         args:
           - --wrap
           - '88'
         files: (README\.md)
+  - repo: https://github.com/kynan/nbstripout
+    rev: 0.7.1
+    hooks:
+      - id: nbstripout
+        args:
+          - --drop-empty-cells
+          - --keep-output
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.11.2
+    hooks:
+      - id: mypy
+        files: src|tests
+        additional_dependencies:
+          - jax>=0.4.20
+          - numpy
+          - packaging
+          - pandas-stubs
+        args:
+          - --config=pyproject.toml
 ci:
   autoupdate_schedule: monthly
+  skip:
+    - mypy  # installing jax is not possible on pre-commit.ci due to size limits.

README.md

@@ -1,6 +1,6 @@
 # Life Cycle Models
 
-[![image](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 [![pre-commit.ci status](https://results.pre-commit.ci/badge/github/OpenSourceEconomics/lcm/main.svg)](https://results.pre-commit.ci/latest/github/OpenSourceEconomics/lcm/main)
 [![image](https://codecov.io/gh/OpenSourceEconomics/lcm/branch/main/graph/badge.svg)](https://codecov.io/gh/OpenSourceEconomics/lcm)
 

codecov.yml

@@ -14,5 +14,4 @@ coverage:
       default:
         target: 90%
 ignore:
-  - setup.py
   - tests/**/*

examples/long_running.py

@@ -2,17 +2,7 @@
 
 import jax.numpy as jnp
 
-# ======================================================================================
-# Numerical parameters and constants
-# ======================================================================================
-N_GRID_POINTS = {
-    "wealth": 100,
-    "health": 100,
-    "consumption": 100,
-    "exercise": 200,
-}
-
-RETIREMENT_AGE = 65
+from lcm import DiscreteGrid, LinspaceGrid, Model
 
 # ======================================================================================
 # Model functions
@@ -62,9 +52,12 @@ def consumption_constraint(consumption, wealth, labor_income):
 # ======================================================================================
 # Model specification and parameters
 # ======================================================================================
+RETIREMENT_AGE = 65
 
-MODEL_CONFIG = {
-    "functions": {
+
+MODEL_CONFIG = Model(
+    n_periods=RETIREMENT_AGE - 18,
+    functions={
         "utility": utility,
         "next_wealth": next_wealth,
         "next_health": next_health,
@@ -73,37 +66,32 @@ def consumption_constraint(consumption, wealth, labor_income):
         "wage": wage,
         "age": age,
     },
-    "choices": {
-        "working": {"options": [0, 1]},
-        "consumption": {
-            "grid_type": "linspace",
-            "start": 1,
-            "stop": 100,
-            "n_points": N_GRID_POINTS["consumption"],
-        },
-        "exercise": {
-            "grid_type": "linspace",
-            "start": 0,
-            "stop": 1,
-            "n_points": N_GRID_POINTS["exercise"],
-        },
+    choices={
+        "working": DiscreteGrid([0, 1]),
+        "consumption": LinspaceGrid(
+            start=1,
+            stop=100,
+            n_points=100,
+        ),
+        "exercise": LinspaceGrid(
+            start=0,
+            stop=1,
+            n_points=200,
+        ),
     },
-    "states": {
-        "wealth": {
-            "grid_type": "linspace",
-            "start": 1,
-            "stop": 100,
-            "n_points": N_GRID_POINTS["wealth"],
-        },
-        "health": {
-            "grid_type": "linspace",
-            "start": 0,
-            "stop": 1,
-            "n_points": N_GRID_POINTS["health"],
-        },
+    states={
+        "wealth": LinspaceGrid(
+            start=1,
+            stop=100,
+            n_points=100,
+        ),
+        "health": LinspaceGrid(
+            start=0,
+            stop=1,
+            n_points=100,
+        ),
     },
-    "n_periods": RETIREMENT_AGE - 18,
-}
+)
 
 PARAMS = {
     "beta": 0.95,