Merge branch 'main' into long-running-test

.envs/testenv.yml

@@ -8,13 +8,6 @@ dependencies:
   - setuptools_scm
   - toml
 
-  # Package dependencies
-  - dags
-  - jax>=0.4.10
-  - jaxlib>=0.4.10
-  - numpy
-  - pandas
-
   # Testing dependencies
   - scipy
   - pybaum
@@ -24,6 +17,6 @@ dependencies:
   - pytest-cov
   - pytest-xdist
 
-  # Install lcm locally
+  # Install lcm and its dependencies locally
   - pip:
       - -e ../

.github/workflows/main.yml

@@ -31,16 +31,15 @@ jobs:
           environment-file: ./.envs/testenv.yml
           environment-name: lcm
           cache-environment: true
-          create-args: >
-          extra-specs: |
+          create-args: >-
             python=${{ matrix.python-version }}
       - name: run pytest
         shell: bash -l {0}
         run: |
           micromamba activate lcm
           pytest --cov-report=xml --cov=./
       - name: Upload coverage report.
-        if: runner.os == 'Linux' && matrix.python-version == '3.10'
+        if: runner.os == 'Linux' && matrix.python-version == '3.11'
         uses: codecov/codecov-action@v3
         with:
           token: ${{ secrets.CODECOV_TOKEN }}

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

@@ -7,10 +7,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
-      - name: Set up Python 3.10
+      - name: Set up Python 3.11
         uses: actions/setup-python@v4
         with:
-          python-version: '3.10'
+          python-version: '3.11'
       - name: Install pypa/build
         run: >-
           python -m

.pre-commit-config.yaml

@@ -45,7 +45,7 @@ repos:
     hooks:
       - id: yamllint
   - repo: https://github.com/psf/black
-    rev: 23.11.0
+    rev: 24.1.1
     hooks:
       - id: black
         language_version: python3.11
@@ -54,13 +54,13 @@ repos:
     hooks:
       - id: blacken-docs
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.1.6
+    rev: v0.2.0
     hooks:
       - id: ruff
       #   args:
       #     - --verbose
   - repo: https://github.com/kynan/nbstripout
-    rev: 0.6.1
+    rev: 0.7.1
     hooks:
       - id: nbstripout
         args:

docs/source/conf.py

@@ -9,6 +9,7 @@
 documentation root, use os.path.abspath to make it absolute, like shown here.
 
 """
+
 import pathlib
 import sys
 

pyproject.toml

@@ -54,11 +54,11 @@ extend-ignore = [
     "TRY003",
 ]
 
-[tool.ruff.per-file-ignores]
+[tool.ruff.lint.per-file-ignores]
 "docs/source/conf.py" = ["E501", "ERA001", "DTZ005"]
 "tests/test_*.py" = ["PLR2004"]
 
-[tool.ruff.pydocstyle]
+[tool.ruff.lint.pydocstyle]
 convention = "google"
 
 

setup.cfg

@@ -29,11 +29,11 @@ keywords =
 packages = find:
 install_requires =
     dags
-    jax>=0.3.0
-    jaxlib>=0.3.0
+    jax>=0.4.10
+    jaxlib>=0.4.10
     numpy
     pandas
-python_requires = >=3.10
+python_requires = >=3.11
 include_package_data = True
 package_dir =
     =src

setup.py

@@ -1,4 +1,5 @@
 """Setup script for the package."""
+
 from setuptools import setup
 
 if __name__ == "__main__":

src/lcm/create_params.py

@@ -1,4 +1,5 @@
 """Create a parameters for a model specification."""
+
 import inspect
 
 import numpy as np
@@ -102,12 +103,12 @@ def _create_shock_params(model, variable_info, grids):
             inspect.signature(model["functions"][f"next_{var}"]).parameters,
         )
 
-        _check_variables_are_all_discrete_states(
+        _check_variables_are_all_discrete_or_period(
             variables=dependencies,
             variable_info=variable_info,
             msg_suffix=(
-                f"The function next_{var} can only depend on discrete state variables "
-                f"or '_period'."
+                f"The function next_{var} can only depend on discrete variables or"
+                f"'_period'."
             ),
         )
 
@@ -128,10 +129,10 @@ def _create_standard_params():
     return {"beta": np.nan}
 
 
-def _check_variables_are_all_discrete_states(variables, variable_info, msg_suffix):
-    discrete_state_vars = variable_info.query("is_state and is_discrete").index.tolist()
+def _check_variables_are_all_discrete_or_period(variables, variable_info, msg_suffix):
+    discrete_vars = variable_info.query("is_discrete").index.tolist()
     for var in variables:
-        if var not in discrete_state_vars and var != "_period":
+        if var not in discrete_vars and var != "_period":
             raise ValueError(
-                f"Variable {var} is not a discrete state variable. {msg_suffix}",
+                f"Variable {var} is not a discrete variable. {msg_suffix}",
             )

src/lcm/discrete_emax.py

@@ -20,6 +20,7 @@
 https://github.com/google/jax/issues/6265
 
 """
+
 from functools import partial
 
 import jax

src/lcm/distributions.py

@@ -4,6 +4,7 @@
 in JAX.
 
 """
+
 import numpy as np
 
 

src/lcm/example_models/example_models.py

@@ -1,4 +1,5 @@
 """Define example model specifications."""
+
 import jax.numpy as jnp
 
 RETIREMENT_AGE = 65

src/lcm/example_models/example_models_stochastic.py

@@ -1,4 +1,5 @@
 """Define example model specifications."""
+
 import jax.numpy as jnp
 
 import lcm
@@ -21,7 +22,7 @@ def next_health(health, partner):  # noqa: ARG001
 
 
 @lcm.mark.stochastic
-def next_partner(_period):
+def next_partner(_period, working, partner):  # noqa: ARG001
     pass
 
 
@@ -41,7 +42,7 @@ def consumption_constraint(consumption, wealth):
         "working": {"options": [0, 1]},
         "consumption": {
             "grid_type": "linspace",
-            "start": 0,
+            "start": 1,
             "stop": 100,
             "n_points": N_CHOICE_GRID_POINTS,
         },
@@ -51,7 +52,7 @@ def consumption_constraint(consumption, wealth):
         "partner": {"options": [0, 1]},
         "wealth": {
             "grid_type": "linspace",
-            "start": 0,
+            "start": 1,
             "stop": 100,
             "n_points": N_STATE_GRID_POINTS,
         },
@@ -61,13 +62,67 @@ def consumption_constraint(consumption, wealth):
 
 
 PARAMS = {
-    "beta": 0.25,
-    "utility": {"delta": 0.25, "gamma": 0.25},
-    "next_wealth": {"interest_rate": 0.25, "wage": 0.25},
+    "beta": 0.95,
+    "utility": {"delta": 0.5, "gamma": 0.25},
+    "next_wealth": {"interest_rate": 0.05, "wage": 10.0},
     "next_health": {},
     "consumption_constraint": {},
     "shocks": {
-        "health": jnp.array([[[0.5, 0.5], [0.5, 0.5]], [[0.5, 0.5], [0.5, 0.5]]]),
-        "partner": jnp.array([[1.0, 0], [0.0, 1]]),
+        # Health shock:
+        # ------------------------------------------------------------------------------
+        # 1st dimension: Current health state
+        # 2nd dimension: Current Partner state
+        # 3rd dimension: Probability distribution over next period's health state
+        "health": jnp.array(
+            [
+                # Current health state 0
+                [
+                    # Current Partner state 0
+                    [0.9, 0.1],
+                    # Current Partner state 1
+                    [0.5, 0.5],
+                ],
+                # Current health state 1
+                [
+                    # Current Partner state 0
+                    [0.5, 0.5],
+                    # Current Partner state 1
+                    [0.1, 0.9],
+                ],
+            ],
+        ),
+        # Partner shock:
+        # ------------------------------------------------------------------------------
+        # 1st dimension: The period
+        # 2nd dimension: Current working decision
+        # 3rd dimension: Current partner state
+        # 4th dimension: Probability distribution over next period's partner state
+        "partner": jnp.array(
+            [
+                # Transition from period 0 to period 1
+                [
+                    # Current working decision 0
+                    [
+                        # Current partner state 0
+                        [0, 1.0],
+                        # Current partner state 1
+                        [1.0, 0],
+                    ],
+                    # Current working decision 1
+                    [
+                        # Current partner state 0
+                        [0, 1.0],
+                        # Current partner state 1
+                        [0.0, 1.0],
+                    ],
+                ],
+                # Transition from period 1 to period 2
+                [
+                    # Description is the same as above
+                    [[0, 1.0], [1.0, 0]],
+                    [[0, 1.0], [0.0, 1.0]],
+                ],
+            ],
+        ),
     },
 }

src/lcm/functools.py

@@ -67,9 +67,11 @@ def allow_kwargs(func):
 
     # Create new signature without positional-only arguments
     new_parameters = [
-        p.replace(kind=inspect.Parameter.POSITIONAL_OR_KEYWORD)
-        if p.kind == inspect.Parameter.POSITIONAL_ONLY
-        else p
+        (
+            p.replace(kind=inspect.Parameter.POSITIONAL_OR_KEYWORD)
+            if p.kind == inspect.Parameter.POSITIONAL_ONLY
+            else p
+        )
         for p in parameters.values()
     ]
     new_signature = signature.replace(parameters=new_parameters)
@@ -117,9 +119,11 @@ def allow_args(func):
 
     # Create new signature without keyword-only arguments
     new_parameters = [
-        p.replace(kind=inspect.Parameter.POSITIONAL_OR_KEYWORD)
-        if p.kind == inspect.Parameter.KEYWORD_ONLY
-        else p
+        (
+            p.replace(kind=inspect.Parameter.POSITIONAL_OR_KEYWORD)
+            if p.kind == inspect.Parameter.KEYWORD_ONLY
+            else p
+        )
         for p in parameters.values()
     ]
     new_signature = signature.replace(parameters=new_parameters)

src/lcm/get_model.py

@@ -1,4 +1,5 @@
 """Get a user model and parameters."""
+
 from typing import NamedTuple
 
 from pybaum import tree_update

src/lcm/grids.py

@@ -17,6 +17,7 @@
 it easy to call functions interchangeably.
 
 """
+
 import jax.numpy as jnp
 
 

src/lcm/mark.py

@@ -1,4 +1,5 @@
 """Collection of LCM marking decorators."""
+
 import functools
 from typing import NamedTuple
 

src/lcm/model_functions.py

@@ -107,7 +107,12 @@ def u_and_f(*args, **kwargs):
                 _period=period,
                 params=kwargs["params"],
             )
-            weights = next_weights(**states, _period=period, params=kwargs["params"])
+            weights = next_weights(
+                **states,
+                **choices,
+                _period=period,
+                params=kwargs["params"],
+            )
 
             value_function = productmap(
                 scalar_value_function,

src/lcm/next_state.py

@@ -6,6 +6,7 @@
 deteministic next states.
 
 """
+
 from dags import concatenate_functions
 from dags.signature import with_signature
 

src/lcm/process_model.py

@@ -393,12 +393,12 @@ def _get_stochastic_weight_function(raw_func, name, variable_info, grids):
     """
     function_parameters = list(inspect.signature(raw_func).parameters)
 
-    # Assert that stochastic next function only depends on state variables
+    # Assert that stochastic next function only depends on discrete variables or period
     for arg in function_parameters:
-        if arg != "_period" and not variable_info.loc[arg, "is_state"]:
+        if arg != "_period" and not variable_info.loc[arg, "is_discrete"]:
             raise ValueError(
-                f"Stochastic variables can only depend on state variables and '_period'"
-                f" but {name} depends on {arg}.",
+                f"Stochastic variables can only depend on discrete variables and "
+                f"'_period', but {name} depends on {arg}.",
             )
 
     label_translators = {