Merge branch 'main' into stochastic-transitions

.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
 

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

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.py

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

src/lcm/example_models_stochastic.py

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

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/next_state.py

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

src/lcm/sandbox/interpolation.ipynb

@@ -3,7 +3,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "da0bbbc9",
+   "id": "0",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -24,7 +24,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1c31af18",
+   "id": "1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -46,7 +46,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "adc65243",
+   "id": "2",
    "metadata": {},
    "source": [
     "## 1d Benchmarks"
@@ -55,7 +55,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "86f352f5",
+   "id": "3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -102,7 +102,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "46e61c87",
+   "id": "4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -114,7 +114,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8f9062cb",
+   "id": "5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -126,7 +126,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "55bc8605",
+   "id": "6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -138,7 +138,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "62bb0133",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -148,7 +148,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "8a071d57",
+   "id": "8",
    "metadata": {},
    "source": [
     "## 2d Benchmarks"
@@ -157,7 +157,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7db8a854",
+   "id": "9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -199,7 +199,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ff7b5ed0",
+   "id": "10",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -211,7 +211,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1a8a6787",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -223,7 +223,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5401c149",
+   "id": "12",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -235,7 +235,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "216b4e52",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -251,7 +251,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a76907f9",
+   "id": "14",
    "metadata": {},
    "outputs": [],
    "source": []