diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index ee3f39b..7afbac1 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -45,7 +45,7 @@ repos:
     hooks:
       - id: yamllint
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.10
+    rev: v0.14.11
     hooks:
       - id: ruff-check
         types_or:
@@ -54,7 +54,6 @@ repos:
           - jupyter
         args:
           - --fix
-          # - --unsafe-fixes
       - id: ruff-format
         types_or:
           - python
diff --git a/.yamllint.yml b/.yamllint.yml
index 72f64be..0bdfa07 100644
--- a/.yamllint.yml
+++ b/.yamllint.yml
@@ -3,6 +3,8 @@ yaml-files:
   - '*.yaml'
   - '*.yml'
   - .yamllint
+ignore:
+  - src/skillmodels/test_data/simplest_augmented_model.yaml
 rules:
   braces: enable
   brackets: enable
@@ -34,3 +36,4 @@ rules:
   trailing-spaces: enable
   truthy:
     level: warning
+    check-keys: false
diff --git a/CLAUDE.md b/CLAUDE.md
index 25ca239..ce3840b 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -86,10 +86,12 @@ The main package exports three functions:
 
 ## Code Style
 
-- Uses Ruff for linting (target: Python 3.13, line length: 88)
+- Require Python 3.14
+- Uses Ruff for linting (target: Python 3.14, line length: 88)
 - Google-style docstrings
 - Pre-commit hooks enforce formatting and linting
 - Type checking via `ty` with strict rules
+- Do not use `from __future__ import annotations`
 
 ## Testing
 
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 8c9b2d4..09c9be1 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -1,3 +1,5 @@
+"""Sphinx configuration file for skillmodels documentation."""
+
 #
 # Documentation build configuration file, created by sphinx-quickstart
 #
@@ -8,13 +10,13 @@
 #
 # All configuration values have a default; values that are commented out
 # serve to show the default.
-import os
 import sys
+from pathlib import Path
 
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
-# documentation root, use os.path.abspath to make it absolute, like shown here.
-sys.path.insert(0, os.path.abspath("../.."))
+# documentation root, use Path.resolve() to make it absolute, like shown here.
+sys.path.insert(0, str(Path("../..").resolve()))
 
 
 # -- General configuration ----------------------------------------------------
@@ -61,7 +63,7 @@
 
 # General information about the project.
 project = "skillmodels"
-copyright = "2016-2021, Janos Gabler"
+copyright = "2016-, Janos Gabler"  # noqa: A001
 
 # The version info for the project you"re documenting, acts as replacement for
 # |version| and |release|, also used in various other places throughout the
diff --git a/docs/source/getting_started/tutorial.ipynb b/docs/source/getting_started/tutorial.ipynb
index 15906e6..137c811 100644
--- a/docs/source/getting_started/tutorial.ipynb
+++ b/docs/source/getting_started/tutorial.ipynb
@@ -19,7 +19,7 @@
     "import pandas as pd\n",
     "import yaml\n",
     "\n",
-    "from skillmodels.config import TEST_DIR\n",
+    "from skillmodels.config import REGRESSION_VAULT, TEST_DATA_DIR\n",
     "from skillmodels.maximization_inputs import get_maximization_inputs"
    ]
   },
@@ -41,7 +41,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open(TEST_DIR / \"model2.yaml\") as y:\n",
+    "with (TEST_DATA_DIR / \"model2.yaml\").open() as y:\n",
     "    model_dict = yaml.load(y, Loader=yaml.SafeLoader)"
    ]
   },
@@ -51,7 +51,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "data = pd.read_stata(TEST_DIR / \"model2_simulated_data.dta\")\n",
+    "data = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\n",
     "data = data.set_index([\"caseid\", \"period\"])"
    ]
   },
@@ -110,7 +110,7 @@
    "outputs": [],
    "source": [
     "index_cols = [\"category\", \"period\", \"name1\", \"name2\"]\n",
-    "chs_path = TEST_DIR / \"regression_vault\" / \"chs_results.csv\"\n",
+    "chs_path = REGRESSION_VAULT / \"chs_results.csv\"\n",
     "chs_values = pd.read_csv(chs_path)\n",
     "chs_values = chs_values.set_index(index_cols)\n",
     "chs_values = chs_values[[\"chs_value\", \"good_start_value\", \"bad_start_value\"]]\n",
@@ -289,7 +289,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "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",
@@ -299,8 +298,6 @@
     "    + params.index.get_level_values(\"period\").astype(str)\n",
     ")\n",
     "params[\"group\"] = params[\"group\"].str.replace(\"_\", \"-\")\n",
-    "params[\"group\"] = params[\"group\"].astype(\"O\")\n",
-    "params.loc[~pp[\"_internal_free\"], \"group\"] = None\n",
     "params"
    ]
   },
diff --git a/docs/source/how_to_guides/how_to_simulate_dataset.ipynb b/docs/source/how_to_guides/how_to_simulate_dataset.ipynb
index bad7647..12905f6 100644
--- a/docs/source/how_to_guides/how_to_simulate_dataset.ipynb
+++ b/docs/source/how_to_guides/how_to_simulate_dataset.ipynb
@@ -9,7 +9,7 @@
     "import pandas as pd\n",
     "import yaml\n",
     "\n",
-    "from skillmodels.config import TEST_DIR\n",
+    "from skillmodels.config import REGRESSION_VAULT, TEST_DATA_DIR\n",
     "from skillmodels.simulate_data import simulate_dataset"
    ]
   },
@@ -34,13 +34,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open(TEST_DIR / \"model2.yaml\") as y:\n",
-    "    model_dict = yaml.load(y, Loader=yaml.FullLoader)\n",
+    "with (TEST_DATA_DIR / \"model2.yaml\").open() as y:\n",
+    "    model_dict = yaml.load(y, Loader=yaml.SafeLoader)\n",
     "\n",
-    "data = pd.read_stata(TEST_DIR / \"model2_simulated_data.dta\")\n",
+    "data = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\n",
     "data = data.set_index([\"caseid\", \"period\"])\n",
     "\n",
-    "params = pd.read_csv(TEST_DIR / \"regression_vault\" / \"one_stage_anchoring.csv\")\n",
+    "params = pd.read_csv(REGRESSION_VAULT / \"one_stage_anchoring.csv\")\n",
     "params = params.set_index([\"category\", \"period\", \"name1\", \"name2\"])"
    ]
   },
diff --git a/docs/source/how_to_guides/how_to_visualize_correlations.ipynb b/docs/source/how_to_guides/how_to_visualize_correlations.ipynb
index 299f320..45157de 100644
--- a/docs/source/how_to_guides/how_to_visualize_correlations.ipynb
+++ b/docs/source/how_to_guides/how_to_visualize_correlations.ipynb
@@ -16,7 +16,7 @@
     "import pandas as pd\n",
     "import yaml\n",
     "\n",
-    "from skillmodels.config import TEST_DIR\n",
+    "from skillmodels.config import REGRESSION_VAULT, TEST_DATA_DIR\n",
     "from skillmodels.correlation_heatmap import (\n",
     "    get_measurements_corr,\n",
     "    get_quasi_scores_corr,\n",
@@ -40,8 +40,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open(TEST_DIR / \"model2.yaml\") as y:\n",
-    "    model_dict = yaml.load(y, Loader=yaml.FullLoader)"
+    "with (TEST_DATA_DIR / \"model2.yaml\").open() as y:\n",
+    "    model_dict = yaml.load(y, Loader=yaml.SafeLoader)"
    ]
   },
   {
@@ -50,10 +50,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "params = pd.read_csv(TEST_DIR / \"regression_vault\" / \"one_stage_anchoring.csv\")\n",
+    "params = pd.read_csv(REGRESSION_VAULT / \"one_stage_anchoring.csv\")\n",
     "params = params.set_index([\"category\", \"period\", \"name1\", \"name2\"])\n",
     "\n",
-    "data = pd.read_stata(TEST_DIR / \"model2_simulated_data.dta\")\n",
+    "data = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\n",
     "data = data.set_index([\"caseid\", \"period\"])"
    ]
   },
@@ -170,7 +170,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from skillmodels.visualize_transition_equations import _get_pardict, _set_index_params"
+    "from skillmodels.visualize_transition_equations import (\n",
+    "    _get_parsed_params,\n",
+    "    _set_index_params,\n",
+    ")"
    ]
   },
   {
@@ -188,10 +191,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "_get_pardict(\n",
+    "_get_parsed_params(\n",
     "    params=_set_index_params(process_model(model_dict), params),\n",
     "    model=process_model(model_dict),\n",
-    ")[\"loadings\"]"
+    ").loadings"
    ]
   },
   {
diff --git a/docs/source/how_to_guides/how_to_visualize_pairwise_factor_distribution.ipynb b/docs/source/how_to_guides/how_to_visualize_pairwise_factor_distribution.ipynb
index db9ec04..05b7a52 100644
--- a/docs/source/how_to_guides/how_to_visualize_pairwise_factor_distribution.ipynb
+++ b/docs/source/how_to_guides/how_to_visualize_pairwise_factor_distribution.ipynb
@@ -21,7 +21,7 @@
     "import pandas as pd\n",
     "import yaml\n",
     "\n",
-    "from skillmodels.config import TEST_DIR\n",
+    "from skillmodels.config import REGRESSION_VAULT, TEST_DATA_DIR\n",
     "from skillmodels.maximization_inputs import get_maximization_inputs\n",
     "from skillmodels.simulate_data import simulate_dataset\n",
     "from skillmodels.visualize_factor_distributions import (\n",
@@ -57,12 +57,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open(TEST_DIR / \"model2.yaml\") as y:\n",
-    "    model_dict = yaml.load(y, Loader=yaml.FullLoader)\n",
-    "params = pd.read_csv(TEST_DIR / \"regression_vault\" / \"one_stage_anchoring.csv\")\n",
+    "with (TEST_DATA_DIR / \"model2.yaml\").open() as y:\n",
+    "    model_dict = yaml.load(y, Loader=yaml.SafeLoader)\n",
+    "params = pd.read_csv(REGRESSION_VAULT / \"one_stage_anchoring.csv\")\n",
     "params = params.set_index([\"category\", \"period\", \"name1\", \"name2\"])\n",
     "\n",
-    "data = pd.read_stata(TEST_DIR / \"model2_simulated_data.dta\")\n",
+    "data = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\n",
     "data = data.set_index([\"caseid\", \"period\"])"
    ]
   },
diff --git a/docs/source/how_to_guides/how_to_visualize_transition_equations.ipynb b/docs/source/how_to_guides/how_to_visualize_transition_equations.ipynb
index f67e28b..9f9f863 100644
--- a/docs/source/how_to_guides/how_to_visualize_transition_equations.ipynb
+++ b/docs/source/how_to_guides/how_to_visualize_transition_equations.ipynb
@@ -10,7 +10,7 @@
     "import pandas as pd\n",
     "import yaml\n",
     "\n",
-    "from skillmodels.config import TEST_DIR\n",
+    "from skillmodels.config import REGRESSION_VAULT, TEST_DATA_DIR\n",
     "from skillmodels.visualize_transition_equations import (\n",
     "    combine_transition_plots,\n",
     "    get_transition_plots,\n",
@@ -47,13 +47,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open(TEST_DIR / \"model2.yaml\") as y:\n",
-    "    model_dict = yaml.load(y, Loader=yaml.FullLoader)\n",
+    "with (TEST_DATA_DIR / \"model2.yaml\").open() as y:\n",
+    "    model_dict = yaml.load(y, Loader=yaml.SafeLoader)\n",
     "\n",
-    "params = pd.read_csv(TEST_DIR / \"regression_vault\" / \"one_stage_anchoring.csv\")\n",
+    "params = pd.read_csv(REGRESSION_VAULT / \"one_stage_anchoring.csv\")\n",
     "params = params.set_index([\"category\", \"period\", \"name1\", \"name2\"])\n",
     "\n",
-    "data = pd.read_stata(TEST_DIR / \"model2_simulated_data.dta\")\n",
+    "data = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\n",
     "data = data.set_index([\"caseid\", \"period\"])"
    ]
   },
diff --git a/pixi.lock b/pixi.lock
index ea662f2..39c7df6 100644
--- a/pixi.lock
+++ b/pixi.lock
@@ -14,13 +14,13 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/anyio-4.12.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/argon2-cffi-25.1.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/argon2-cffi-bindings-25.1.0-py313h07c4f96_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/argon2-cffi-bindings-25.1.0-py314h5bd0f2a_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/arrow-1.4.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/async-lru-2.0.5-pyh29332c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/babel-2.17.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/backports.zstd-1.3.0-py313h18e8e13_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.3.0-py314h680f03e_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.14.3-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/binutils_impl_linux-64-2.45-default_hfdba357_105.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/binutils_linux-64-2.45-default_h4852527_105.conda
@@ -28,20 +28,20 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-with-css-6.3.0-h5f6438b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.2.0-hed03a55_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.2.0-hb03c661_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py313hf159716_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py314h3de4e8d_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_8.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py313hf46b229_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py314h4a8dc5f_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/choreographer-1.2.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.3-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py313h7037e92_3.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.11-py313hd8ed1ab_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py314h9891dd4_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.14.2-py314hd8ed1ab_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cccl_linux-64-12.9.27-ha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-crt-dev_linux-64-12.9.86-ha770c72_2.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-crt-tools-12.9.86-ha770c72_2.conda
@@ -62,7 +62,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvvm-tools-12.9.86-h4bc722e_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-version-12.9-h4f385c5_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.18-py313h5d5ffb9_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.18-py314h42812f9_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/defusedxml-0.7.1-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/distlib-0.4.0-pyhd8ed1ab_0.conda
@@ -70,7 +70,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.20.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filterpy-1.4.5-pyhd8ed1ab_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/fonttools-4.61.1-py313h3dea7bd_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonttools-4.61.1-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/fqdn-1.5.1-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/freetype-2.14.1-ha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/gcc_impl_linux-64-14.3.0-he8b2097_16.conda
@@ -110,7 +110,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_server-2.28.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/kernel-headers_linux-64-4.18.0-he073ed8_9.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/keyutils-1.6.3-hb9d3cd8_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py313hc8edb43_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py314h97ea11e_2.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/krb5-1.21.3-h659f571_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/lark-1.3.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/lcms2-2.17-h717163a_0.conda
@@ -153,8 +153,8 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcb-1.17.0-h8a09558_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.3.1-hb9d3cd8_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/logistro-2.0.1-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/markupsafe-3.0.3-py313h3dea7bd_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.8-py313h683a580_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/markupsafe-3.0.3-pyh7db6752_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.8-py314h1194b4b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mistune-3.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
@@ -167,23 +167,23 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.6.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/nodeenv-1.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/notebook-shim-0.2.4-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py313hf6604e3_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py314h2b28147_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.4-h55fea9a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.0-h26f9b46_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.11.5-py313h541fbb8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.11.5-py314h3b757c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/overrides-7.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pandocfilters-1.5.0-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pexpect-4.9.0-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.1.0-py313h80991f8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.1.0-py314h8ec4b1a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/plotly-6.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.6.0-pyhf9edf01_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pre-commit-4.5.1-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prometheus_client-0.23.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.1-py313h54dd161_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.1-py314h0f05182_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ptyprocess-0.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
@@ -194,16 +194,16 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.0.2-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-timeout-2.4.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.11-hc97d973_100_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.14.2-h32b2ec7_100_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.2-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.13.11-h4df99d1_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.14.2-h4df99d1_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-json-logger-2.0.7-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-kaleido-1.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.14-8_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py313h3dea7bd_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pyyaml-6.0.3-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/pyzmq-27.1.0-py312hfb55c3c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/qhull-2020.2-h434a139_5.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda
@@ -212,11 +212,11 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3339-validator-0.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3986-validator-0.1.1-pyh9f0ad1d_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3987-syntax-1.1.0-pyhe01879c_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.30.0-py313h843e2db_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.16.3-py313h4b8bb8b_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.30.0-py314h2e6c369_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.16.3-py314hf07bd8e_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/send2trash-2.0.0-pyha191276_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.9.0-pyhff2d567_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/simplejson-3.20.2-py313h07c4f96_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/simplejson-3.20.2-py314h5bd0f2a_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/sniffio-1.3.1-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/soupsieve-2.8.1-pyhd8ed1ab_0.conda
@@ -226,16 +226,17 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/tinycss2-1.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_ha0e22de_103.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py313h07c4f96_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py314h5bd0f2a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_utils-0.1.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.0.1-py313h7037e92_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.0.1-py314h9891dd4_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/unicodedata2-17.0.0-py314h5bd0f2a_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/uri-template-1.3.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.6.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.35.4-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.36.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webcolors-25.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webencodings-0.5.1-pyhd8ed1ab_3.conda
@@ -251,13 +252,14 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/01/c1/e32fb13d9cb5afc5f25f0c84ca21d3ea2d5380f4c06417b814ecc9bf0f38/dags-0.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/fd/8e/424b8c6e78bd9837d14ff7df01a9829fc883ba2ab4ea787d4f848435f23f/greenlet-3.3.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/b8/14/bab308fc2c1b5228c3224ec2bf928ce2e4d21d8046c161e44a2012b5203e/greenlet-3.3.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/a8/f7/ae4ecf183d9693cd5fcce7ee063c5e54f173b66dc80a8a79951861e1b557/jax-0.8.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/53/f2/44ad0ce1d115f0f6be10f4af0ca05a18afb838b06e6ca6b01ba4b0137421/jax_cuda12_pjrt-0.8.2-py3-none-manylinux_2_27_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/1c/38/4ba2486f95fcf2120723932feacdded438e785258148b18a703cd1177e41/jax_cuda12_plugin-0.8.2-cp313-cp313-manylinux_2_27_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/6b/e0/91e5762a7ddb6351b07c742ca407cd28e26043d6945d6228b6c1b0881a45/jaxlib-0.8.2-cp313-cp313-manylinux_2_27_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/27/58/a5a27d4677d6890570f7e58cecd51891469cb620e6f64c8faed4935d93d0/jax_cuda12_plugin-0.8.2-cp314-cp314-manylinux_2_27_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/5e/27/2e6032727e41ce74914277478021140947af59127d68aa9e6f3776b428fd/jaxlib-0.8.2-cp314-cp314-manylinux_2_27_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/c6/bb/82c7dcf38070b46172a517e2334e665c5bf374a262f99a283ea454bece7c/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/77/3c/aa88abe01f3be3d1f8f787d1d33dc83e76fec05945f9a28fbb41cfb99cd5/nvidia_cublas_cu12-12.9.1.4-py3-none-manylinux_2_27_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/18/2a/d4cd8506d2044e082f8cd921be57392e6a9b5ccd3ffdf050362430a3d5d5/nvidia_cuda_cccl_cu12-12.9.27-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/c1/2e/b84e32197e33f39907b455b83395a017e697c07a449a2b15fd07fc1c9981/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_x86_64.whl
@@ -272,10 +274,10 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/46/0c/c75bbfb967457a0b7670b8ad267bfc4fffdf341c074e0a80db06c24ccfd4/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/64/b9/6ab941001c23cfb43499b5b0b7417b0bb4dfba3a29ffa2b06985422dad50/nvidia_nvshmem_cu12-3.5.19-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
-      - pypi: git+https://github.com/optimagic-dev/optimagic.git?branch=main#e02ea4743cac9f861a5813f3b4b1283fd2ade730
-      - pypi: https://files.pythonhosted.org/packages/15/07/284f757f63f8a8d69ed4472bfd85122bd086e637bf4ed09de572d575a693/pandas-2.3.3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: git+https://github.com/optimagic-dev/optimagic.git#522b8c9a21226569ffd25e950e44f0c5de308c9d
+      - pypi: https://files.pythonhosted.org/packages/15/b2/0e62f78c0c5ba7e3d2c5945a82456f4fac76c480940f805e0b97fcbc2f65/pandas-2.3.3-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/75/58/3af430d0de0b95d5adf7e576067e07d750ba76e28d142871982464fb40db/pdbp-1.8.1-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/0e/50/80a8d080ac7d3d321e5e5d420c9a522b0aa770ec7013ea91f9a8b7d36e4a/sqlalchemy-2.0.45-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/b3/27/caf606ee924282fe4747ee4fd454b335a72a6e018f97eab5ff7f28199e16/sqlalchemy-2.0.45-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/65/44/bb509c3d2c0b5a87e7a5af1d5917a402a32ff026f777a6d7cb6990746cbb/tabcompleter-1.4.0-py3-none-any.whl
       - pypi: ./
   default:
@@ -292,34 +294,34 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/anyio-4.12.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/argon2-cffi-25.1.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/argon2-cffi-bindings-25.1.0-py313h07c4f96_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/argon2-cffi-bindings-25.1.0-py314h5bd0f2a_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/arrow-1.4.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/async-lru-2.0.5-pyh29332c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/babel-2.17.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/backports.zstd-1.3.0-py313h18e8e13_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.3.0-py314h680f03e_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.14.3-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-6.3.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-with-css-6.3.0-h5f6438b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.2.0-hed03a55_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.2.0-hb03c661_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py313hf159716_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py314h3de4e8d_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_8.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py313hf46b229_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py314h4a8dc5f_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/choreographer-1.2.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.3-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py313h7037e92_3.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.11-py313hd8ed1ab_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py314h9891dd4_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.14.2-py314hd8ed1ab_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.18-py313h5d5ffb9_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.18-py314h42812f9_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/defusedxml-0.7.1-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/distlib-0.4.0-pyhd8ed1ab_0.conda
@@ -327,7 +329,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.20.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filterpy-1.4.5-pyhd8ed1ab_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/fonttools-4.61.1-py313h3dea7bd_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonttools-4.61.1-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/fqdn-1.5.1-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/freetype-2.14.1-ha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/h11-0.16.0-pyhcf101f3_1.conda
@@ -362,7 +364,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_pygments-0.3.0-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_server-2.28.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/keyutils-1.6.3-hb9d3cd8_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py313hc8edb43_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py314h97ea11e_2.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/krb5-1.21.3-h659f571_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/lark-1.3.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/lcms2-2.17-h717163a_0.conda
@@ -400,8 +402,8 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcb-1.17.0-h8a09558_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.3.1-hb9d3cd8_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/logistro-2.0.1-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/markupsafe-3.0.3-py313h3dea7bd_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.8-py313h683a580_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/markupsafe-3.0.3-pyh7db6752_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.8-py314h1194b4b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mistune-3.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
@@ -414,23 +416,23 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.6.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/nodeenv-1.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/notebook-shim-0.2.4-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py313hf6604e3_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py314h2b28147_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.4-h55fea9a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.0-h26f9b46_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.11.5-py313h541fbb8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.11.5-py314h3b757c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/overrides-7.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pandocfilters-1.5.0-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pexpect-4.9.0-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.1.0-py313h80991f8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.1.0-py314h8ec4b1a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/plotly-6.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.6.0-pyhf9edf01_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pre-commit-4.5.1-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prometheus_client-0.23.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.1-py313h54dd161_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.1-py314h0f05182_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ptyprocess-0.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
@@ -441,16 +443,16 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.0.2-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-timeout-2.4.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.11-hc97d973_100_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.14.2-h32b2ec7_100_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.2-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.13.11-h4df99d1_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.14.2-h4df99d1_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-json-logger-2.0.7-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-kaleido-1.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.14-8_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py313h3dea7bd_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pyyaml-6.0.3-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/pyzmq-27.1.0-py312hfb55c3c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/qhull-2020.2-h434a139_5.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda
@@ -459,11 +461,11 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3339-validator-0.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3986-validator-0.1.1-pyh9f0ad1d_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3987-syntax-1.1.0-pyhe01879c_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.30.0-py313h843e2db_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.16.3-py313h4b8bb8b_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.30.0-py314h2e6c369_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.16.3-py314hf07bd8e_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/send2trash-2.0.0-pyha191276_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.9.0-pyhff2d567_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/simplejson-3.20.2-py313h07c4f96_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/simplejson-3.20.2-py314h5bd0f2a_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/sniffio-1.3.1-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/soupsieve-2.8.1-pyhd8ed1ab_0.conda
@@ -472,16 +474,17 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/tinycss2-1.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_ha0e22de_103.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py313h07c4f96_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py314h5bd0f2a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_utils-0.1.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.0.1-py313h7037e92_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.0.1-py314h9891dd4_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/unicodedata2-17.0.0-py314h5bd0f2a_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/uri-template-1.3.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.6.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.35.4-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.36.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webcolors-25.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webencodings-0.5.1-pyhd8ed1ab_3.conda
@@ -497,16 +500,17 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/01/c1/e32fb13d9cb5afc5f25f0c84ca21d3ea2d5380f4c06417b814ecc9bf0f38/dags-0.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/fd/8e/424b8c6e78bd9837d14ff7df01a9829fc883ba2ab4ea787d4f848435f23f/greenlet-3.3.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/b8/14/bab308fc2c1b5228c3224ec2bf928ce2e4d21d8046c161e44a2012b5203e/greenlet-3.3.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/a8/f7/ae4ecf183d9693cd5fcce7ee063c5e54f173b66dc80a8a79951861e1b557/jax-0.8.2-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/6b/e0/91e5762a7ddb6351b07c742ca407cd28e26043d6945d6228b6c1b0881a45/jaxlib-0.8.2-cp313-cp313-manylinux_2_27_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/5e/27/2e6032727e41ce74914277478021140947af59127d68aa9e6f3776b428fd/jaxlib-0.8.2-cp314-cp314-manylinux_2_27_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/c6/bb/82c7dcf38070b46172a517e2334e665c5bf374a262f99a283ea454bece7c/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
-      - pypi: git+https://github.com/optimagic-dev/optimagic.git?branch=main#e02ea4743cac9f861a5813f3b4b1283fd2ade730
-      - pypi: https://files.pythonhosted.org/packages/15/07/284f757f63f8a8d69ed4472bfd85122bd086e637bf4ed09de572d575a693/pandas-2.3.3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: git+https://github.com/optimagic-dev/optimagic.git#522b8c9a21226569ffd25e950e44f0c5de308c9d
+      - pypi: https://files.pythonhosted.org/packages/15/b2/0e62f78c0c5ba7e3d2c5945a82456f4fac76c480940f805e0b97fcbc2f65/pandas-2.3.3-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/75/58/3af430d0de0b95d5adf7e576067e07d750ba76e28d142871982464fb40db/pdbp-1.8.1-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/0e/50/80a8d080ac7d3d321e5e5d420c9a522b0aa770ec7013ea91f9a8b7d36e4a/sqlalchemy-2.0.45-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/b3/27/caf606ee924282fe4747ee4fd454b335a72a6e018f97eab5ff7f28199e16/sqlalchemy-2.0.45-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/65/44/bb509c3d2c0b5a87e7a5af1d5917a402a32ff026f777a6d7cb6990746cbb/tabcompleter-1.4.0-py3-none-any.whl
       - pypi: ./
       osx-arm64:
@@ -515,34 +519,34 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/anyio-4.12.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/appnope-0.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/argon2-cffi-25.1.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/argon2-cffi-bindings-25.1.0-py313h6535dbc_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/argon2-cffi-bindings-25.1.0-py314h0612a62_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/arrow-1.4.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/async-lru-2.0.5-pyh29332c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/babel-2.17.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/backports.zstd-1.3.0-py313h48bb75e_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.3.0-py314h680f03e_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.14.3-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-6.3.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-with-css-6.3.0-h5f6438b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-1.2.0-h7d5ae5b_1.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-bin-1.2.0-hc919400_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-python-1.2.0-py313hde1f3bb_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-python-1.2.0-py314h3daef5d_1.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/bzip2-1.0.8-hd037594_8.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/cffi-2.0.0-py313h224173a_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/cffi-2.0.0-py314h44086f9_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/choreographer-1.2.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.3-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/contourpy-1.3.3-py313ha61f8ec_3.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.11-py313hd8ed1ab_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/contourpy-1.3.3-py314h784bc60_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.14.2-py314hd8ed1ab_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/debugpy-1.8.19-py313hc37fe24_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/debugpy-1.8.19-py314hf820bb6_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/defusedxml-0.7.1-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/distlib-0.4.0-pyhd8ed1ab_0.conda
@@ -550,7 +554,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.20.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filterpy-1.4.5-pyhd8ed1ab_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/fonttools-4.61.1-py313h7d74516_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonttools-4.61.1-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/fqdn-1.5.1-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/freetype-2.14.1-hce30654_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/h11-0.16.0-pyhcf101f3_1.conda
@@ -583,7 +587,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab-4.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_pygments-0.3.0-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_server-2.28.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/kiwisolver-1.4.9-py313h7add70c_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/kiwisolver-1.4.9-py314h42813c9_2.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/krb5-1.21.3-h237132a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/lark-1.3.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/lcms2-2.17-h7eeda09_0.conda
@@ -617,8 +621,8 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libzlib-1.3.1-h8359307_2.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/llvm-openmp-21.1.8-h4a912ad_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/logistro-2.0.1-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/markupsafe-3.0.3-py313h7d74516_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/matplotlib-base-3.10.8-py313h58042b9_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/markupsafe-3.0.3-pyh7db6752_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/matplotlib-base-3.10.8-py314hd63e3f0_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mistune-3.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
@@ -631,45 +635,45 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.6.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/nodeenv-1.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/notebook-shim-0.2.4-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/numpy-2.3.5-py313h16eae64_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/numpy-2.3.5-py314hae46ccb_1.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/openjpeg-2.5.4-hbfb3c88_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/openssl-3.6.0-h5503f6c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/orjson-3.11.5-py313hfea8034_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/orjson-3.11.5-py314hda6d10a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/overrides-7.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pandocfilters-1.5.0-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pexpect-4.9.0-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pillow-12.1.0-py313h45e5a15_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pillow-12.1.0-py314hab283cf_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/plotly-6.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.6.0-pyhf9edf01_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pre-commit-4.5.1-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prometheus_client-0.23.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/psutil-7.2.1-py313h6688731_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/psutil-7.2.1-py314ha14b1ff_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pthread-stubs-0.4-hd74edd7_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ptyprocess-0.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pybaum-0.1.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pycparser-2.22-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pygments-2.19.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-core-12.1-py313h40b429f_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-framework-cocoa-12.1-py313hcc5defa_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-core-12.1-py314h3a4d195_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-framework-cocoa-12.1-py314h36abed7_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.0.2-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-timeout-2.4.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.13.11-hfc2f54d_100_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.14.2-h40d2674_100_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.2-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.13.11-h4df99d1_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.14.2-h4df99d1_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-json-logger-2.0.7-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-kaleido-1.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.14-8_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyyaml-6.0.3-py313h7d74516_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pyyaml-6.0.3-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyzmq-27.1.0-py312hd65ceae_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/qhull-2020.2-h420ef59_5.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/readline-8.3-h46df422_0.conda
@@ -678,11 +682,11 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3339-validator-0.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3986-validator-0.1.1-pyh9f0ad1d_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3987-syntax-1.1.0-pyhe01879c_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/rpds-py-0.30.0-py313h2c089d5_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/scipy-1.16.3-py313h29d7d31_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/rpds-py-0.30.0-py314haad56a0_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/scipy-1.16.3-py314h725efaa_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/send2trash-2.0.0-pyh5552912_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.9.0-pyhff2d567_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/simplejson-3.20.2-py313h6535dbc_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/simplejson-3.20.2-py314h0612a62_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/sniffio-1.3.1-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/soupsieve-2.8.1-pyhd8ed1ab_0.conda
@@ -691,16 +695,17 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/tinycss2-1.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tk-8.6.13-h892fb3f_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tornado-6.5.4-py313h6535dbc_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tornado-6.5.4-py314h0612a62_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_utils-0.1.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ukkonen-1.0.1-py313hc50a443_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ukkonen-1.0.1-py314h6b18a25_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/unicodedata2-17.0.0-py314h0612a62_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/uri-template-1.3.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.6.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.35.4-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.36.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webcolors-25.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webencodings-0.5.1-pyhd8ed1ab_3.conda
@@ -716,13 +721,14 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/01/c1/e32fb13d9cb5afc5f25f0c84ca21d3ea2d5380f4c06417b814ecc9bf0f38/dags-0.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a8/f7/ae4ecf183d9693cd5fcce7ee063c5e54f173b66dc80a8a79951861e1b557/jax-0.8.2-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/c5/22/c0ec75e43a13b2457d78d509f49b49a57fa302ffced4f4a2778e428cb0a6/jaxlib-0.8.2-cp313-cp313-macosx_11_0_arm64.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/9d/dca93d916bf8664d7a2bb73ea3d219028dabbe382c31774348963287356a/jaxlib-0.8.2-cp314-cp314-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/d9/a1/4008f14bbc616cfb1ac5b39ea485f9c63031c4634ab3f4cf72e7541f816a/ml_dtypes-0.5.4-cp313-cp313-macosx_10_13_universal2.whl
+      - pypi: https://files.pythonhosted.org/packages/72/4e/1339dc6e2557a344f5ba5590872e80346f76f6cb2ac3dd16e4666e88818c/ml_dtypes-0.5.4-cp314-cp314-macosx_10_13_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
-      - pypi: git+https://github.com/optimagic-dev/optimagic.git?branch=main#e02ea4743cac9f861a5813f3b4b1283fd2ade730
-      - pypi: https://files.pythonhosted.org/packages/31/94/72fac03573102779920099bcac1c3b05975c2cb5f01eac609faf34bed1ca/pandas-2.3.3-cp313-cp313-macosx_11_0_arm64.whl
+      - pypi: git+https://github.com/optimagic-dev/optimagic.git#522b8c9a21226569ffd25e950e44f0c5de308c9d
+      - pypi: https://files.pythonhosted.org/packages/21/00/266d6b357ad5e6d3ad55093a7e8efc7dd245f5a842b584db9f30b0f0a287/pandas-2.3.3-cp314-cp314-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/75/58/3af430d0de0b95d5adf7e576067e07d750ba76e28d142871982464fb40db/pdbp-1.8.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/bf/e1/3ccb13c643399d22289c6a9786c1a91e3dcbb68bce4beb44926ac2c557bf/sqlalchemy-2.0.45-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/65/44/bb509c3d2c0b5a87e7a5af1d5917a402a32ff026f777a6d7cb6990746cbb/tabcompleter-1.4.0-py3-none-any.whl
@@ -732,34 +738,34 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/anyio-4.12.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/argon2-cffi-25.1.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/argon2-cffi-bindings-25.1.0-py313h5ea7bf4_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/argon2-cffi-bindings-25.1.0-py314h5a2d7ad_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/arrow-1.4.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/async-lru-2.0.5-pyh29332c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/babel-2.17.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/backports.zstd-1.3.0-py313h2a31948_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.3.0-py314h680f03e_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.14.3-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-6.3.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-with-css-6.3.0-h5f6438b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/brotli-1.2.0-h2d644bc_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/brotli-bin-1.2.0-hfd05255_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/brotli-python-1.2.0-py313h3ebfc14_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/brotli-python-1.2.0-py314he701e3d_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/bzip2-1.0.8-h0ad9c76_8.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-h4c7d964_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/cffi-2.0.0-py313h5ea7bf4_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/cffi-2.0.0-py314h5a2d7ad_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/choreographer-1.2.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.3-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/contourpy-1.3.3-py313hf069bd2_3.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.11-py313hd8ed1ab_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/contourpy-1.3.3-py314h909e829_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.14.2-py314hd8ed1ab_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/debugpy-1.8.19-py313h927ade5_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/debugpy-1.8.19-py314hb98de8c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/defusedxml-0.7.1-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/distlib-0.4.0-pyhd8ed1ab_0.conda
@@ -767,7 +773,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.20.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filterpy-1.4.5-pyhd8ed1ab_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/fonttools-4.61.1-py313hd650c13_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonttools-4.61.1-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/fqdn-1.5.1-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/freetype-2.14.1-h57928b3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/h11-0.16.0-pyhcf101f3_1.conda
@@ -801,7 +807,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab-4.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_pygments-0.3.0-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_server-2.28.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/kiwisolver-1.4.9-py313h1a38498_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/kiwisolver-1.4.9-py314hf309875_2.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/krb5-1.21.3-hdf4eb48_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/lark-1.3.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/lcms2-2.17-hbcf6048_0.conda
@@ -818,7 +824,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/win-64/libfreetype6-2.14.1-hdbac1cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/libgcc-15.2.0-h8ee18e1_16.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/libgomp-15.2.0-h8ee18e1_16.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/libhwloc-2.12.1-default_h4379cf1_1003.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/libhwloc-2.12.2-default_h4379cf1_1000.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/libiconv-1.18-hc1393d2_2.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/libjpeg-turbo-3.1.2-hfd05255_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/liblapack-3.11.0-5_hf9ab0e9_mkl.conda
@@ -836,8 +842,8 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/win-64/libzlib-1.3.1-h2466b09_2.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/llvm-openmp-21.1.8-h4fa8253_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/logistro-2.0.1-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/markupsafe-3.0.3-py313hd650c13_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/matplotlib-base-3.10.8-py313he1ded55_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/markupsafe-3.0.3-pyh7db6752_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/matplotlib-base-3.10.8-py314hfa45d96_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mistune-3.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/mkl-2025.3.0-hac47afa_455.conda
@@ -850,22 +856,22 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.6.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/nodeenv-1.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/notebook-shim-0.2.4-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/numpy-2.3.5-py313hce7ae62_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/numpy-2.3.5-py314h06c3c77_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/openjpeg-2.5.4-h24db6dd_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/openssl-3.6.0-h725018a_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/orjson-3.11.4-py313hfbe8231_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/orjson-3.11.5-py314h64f83cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/overrides-7.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pandocfilters-1.5.0-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/pillow-12.1.0-py313h38f99e1_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/pillow-12.1.0-py314h61b30b5_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/plotly-6.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.6.0-pyhf9edf01_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pre-commit-4.5.1-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prometheus_client-0.23.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/psutil-7.2.1-py313h5fd188c_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/psutil-7.2.1-py314hc5dbbe4_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/pthread-stubs-0.4-h0e40799_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pybaum-0.1.3-pyhd8ed1ab_1.conda
@@ -875,18 +881,18 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyh09c184e_7.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.0.2-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-timeout-2.4.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/python-3.13.11-h09917c8_100_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/python-3.14.2-h4b44e0e_100_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.2-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.13.11-h4df99d1_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.14.2-h4df99d1_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-json-logger-2.0.7-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-kaleido-1.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.14-8_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/pywin32-311-py313h40c08fc_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/pywinpty-2.0.15-py313h5813708_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/pyyaml-6.0.3-py313hd650c13_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/pywin32-311-py314h8f8f202_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/pywinpty-2.0.15-py314h51f0985_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pyyaml-6.0.3-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/pyzmq-27.1.0-py312hbb5da91_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/qhull-2020.2-hc790b64_5.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/referencing-0.37.0-pyhcf101f3_0.conda
@@ -894,34 +900,35 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3339-validator-0.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3986-validator-0.1.1-pyh9f0ad1d_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3987-syntax-1.1.0-pyhe01879c_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/rpds-py-0.30.0-py313hfbe8231_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/scipy-1.16.3-py313he51e9a2_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/rpds-py-0.30.0-py314h9f07db2_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/scipy-1.16.3-py314h221f224_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/send2trash-2.0.0-pyh6dadd2b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.9.0-pyhff2d567_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/simplejson-3.20.2-py313h5ea7bf4_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/simplejson-3.20.2-py314h5a2d7ad_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/sniffio-1.3.1-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/soupsieve-2.8.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/tbb-2022.3.0-hd094cb3_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/tbb-2022.3.0-h3155e25_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/terminado-0.18.1-pyh6dadd2b_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tinycss2-1.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/tk-8.6.13-h2c6b04d_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/tornado-6.5.4-py313h5ea7bf4_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/tornado-6.5.4-py314h5a2d7ad_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_utils-0.1.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/ucrt-10.0.26100.0-h57928b3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/ukkonen-1.0.1-py313hf069bd2_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/ukkonen-1.0.1-py314h909e829_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/unicodedata2-17.0.0-py314h5a2d7ad_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/uri-template-1.3.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.6.3-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/vc-14.3-h41ae7f8_34.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/vc14_runtime-14.44.35208-h818238b_34.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/vcomp14-14.44.35208-h818238b_34.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.35.4-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.36.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webcolors-25.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webencodings-0.5.1-pyhd8ed1ab_3.conda
@@ -939,17 +946,18 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/01/c1/e32fb13d9cb5afc5f25f0c84ca21d3ea2d5380f4c06417b814ecc9bf0f38/dags-0.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/7e/71/ba21c3fb8c5dce83b8c01f458a42e99ffdb1963aeec08fff5a18588d8fd7/greenlet-3.3.0-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/7c/9a/9030e6f9aa8fd7808e9c31ba4c38f87c4f8ec324ee67431d181fe396d705/greenlet-3.3.0-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/a8/f7/ae4ecf183d9693cd5fcce7ee063c5e54f173b66dc80a8a79951861e1b557/jax-0.8.2-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/85/68/25b38673b07a808616ce7b6efb3eed491f983f3373a09cbbd03f67178563/jaxlib-0.8.2-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/b3/8c/af5a00b07a446414edf6b84a7397eab02cf01ba44b6ae1fce7798ce4c127/jaxlib-0.8.2-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/e1/8b/200088c6859d8221454825959df35b5244fa9bdf263fd0249ac5fb75e281/ml_dtypes-0.5.4-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/e9/93/2bfed22d2498c468f6bcd0d9f56b033eaa19f33320389314c19ef6766413/ml_dtypes-0.5.4-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
-      - pypi: git+https://github.com/optimagic-dev/optimagic.git?branch=main#e02ea4743cac9f861a5813f3b4b1283fd2ade730
-      - pypi: https://files.pythonhosted.org/packages/4f/c7/e54682c96a895d0c808453269e0b5928a07a127a15704fedb643e9b0a4c8/pandas-2.3.3-cp313-cp313-win_amd64.whl
+      - pypi: git+https://github.com/optimagic-dev/optimagic.git#522b8c9a21226569ffd25e950e44f0c5de308c9d
+      - pypi: https://files.pythonhosted.org/packages/a6/3d/124ac75fcd0ecc09b8fdccb0246ef65e35b012030defb0e0eba2cbbbe948/pandas-2.3.3-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/75/58/3af430d0de0b95d5adf7e576067e07d750ba76e28d142871982464fb40db/pdbp-1.8.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/5a/dc/491b7661614ab97483abf2056be1deee4dc2490ecbf7bff9ab5cdbac86e1/pyreadline3-3.5.4-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/c0/c5/d17113020b2d43073412aeca09b60d2009442420372123b8d49cc253f8b8/sqlalchemy-2.0.45-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/89/a2/0e1590e9adb292b1d576dbcf67ff7df8cf55e56e78d2c927686d01080f4b/sqlalchemy-2.0.45-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/65/44/bb509c3d2c0b5a87e7a5af1d5917a402a32ff026f777a6d7cb6990746cbb/tabcompleter-1.4.0-py3-none-any.whl
       - pypi: ./
   test-cpu:
@@ -966,36 +974,36 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/anyio-4.12.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/argon2-cffi-25.1.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/argon2-cffi-bindings-25.1.0-py313h07c4f96_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/argon2-cffi-bindings-25.1.0-py314h5bd0f2a_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/arrow-1.4.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/async-lru-2.0.5-pyh29332c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/babel-2.17.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/backports.zstd-1.3.0-py313h18e8e13_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.3.0-py314h680f03e_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.14.3-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-6.3.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-with-css-6.3.0-h5f6438b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.2.0-hed03a55_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.2.0-hb03c661_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py313hf159716_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py314h3de4e8d_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_8.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/c-ares-1.34.6-hb03c661_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py313hf46b229_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py314h4a8dc5f_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/choreographer-1.2.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.3-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py313h7037e92_3.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/coverage-7.13.1-py313h3dea7bd_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.11-py313hd8ed1ab_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py314h9891dd4_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/coverage-7.13.1-py314h67df5f8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.14.2-py314hd8ed1ab_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.18-py313h5d5ffb9_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.18-py314h42812f9_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/defusedxml-0.7.1-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/distlib-0.4.0-pyhd8ed1ab_0.conda
@@ -1005,7 +1013,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.20.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filterpy-1.4.5-pyhd8ed1ab_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/fonttools-4.61.1-py313h3dea7bd_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonttools-4.61.1-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/fqdn-1.5.1-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/freetype-2.14.1-ha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/gmp-6.3.0-hac33072_2.conda
@@ -1042,7 +1050,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_pygments-0.3.0-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_server-2.28.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/keyutils-1.6.3-hb9d3cd8_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py313hc8edb43_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py314h97ea11e_2.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/krb5-1.21.3-h659f571_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/lark-1.3.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/lcms2-2.17-h717163a_0.conda
@@ -1096,12 +1104,12 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/lz4-c-1.10.0-h5888daf_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/lzo-2.10-h280c20c_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/markdown-it-py-4.0.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/markupsafe-3.0.3-py313h3dea7bd_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.8-py313h683a580_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/markupsafe-3.0.3-pyh7db6752_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.8-py314h1194b4b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mdit-py-plugins-0.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mdurl-0.1.2-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/memray-1.19.1-py313h422961c_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/memray-1.19.1-py314hef15ded_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mistune-3.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/narwhals-2.15.0-pyhcf101f3_0.conda
@@ -1114,24 +1122,24 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.6.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/nodeenv-1.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/notebook-shim-0.2.4-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py313hf6604e3_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py314h2b28147_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.4-h55fea9a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.0-h26f9b46_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.11.5-py313h541fbb8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.11.5-py314h3b757c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/overrides-7.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/p11-kit-0.25.10-h3435931_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pandocfilters-1.5.0-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pexpect-4.9.0-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.1.0-py313h80991f8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.1.0-py314h8ec4b1a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/plotly-6.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.6.0-pyhf9edf01_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pre-commit-4.5.1-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prometheus_client-0.23.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.1-py313h54dd161_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.1-py314h0f05182_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ptyprocess-0.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
@@ -1145,16 +1153,16 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-memray-1.8.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-timeout-2.4.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.11-hc97d973_100_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.14.2-h32b2ec7_100_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.2-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.13.11-h4df99d1_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.14.2-h4df99d1_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-json-logger-2.0.7-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-kaleido-1.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.14-8_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py313h3dea7bd_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pyyaml-6.0.3-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/pyzmq-27.1.0-py312hfb55c3c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/qhull-2020.2-h434a139_5.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda
@@ -1164,31 +1172,32 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3986-validator-0.1.1-pyh9f0ad1d_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3987-syntax-1.1.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/rich-14.2.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.30.0-py313h843e2db_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.16.3-py313h4b8bb8b_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.30.0-py314h2e6c369_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.16.3-py314hf07bd8e_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/send2trash-2.0.0-pyha191276_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.9.0-pyhff2d567_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/simplejson-3.20.2-py313h07c4f96_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/simplejson-3.20.2-py314h5bd0f2a_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/snakeviz-2.2.2-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/sniffio-1.3.1-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/soupsieve-2.8.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/terminado-0.18.1-pyhc90fa1f_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/textual-7.0.0-pyhcf101f3_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/textual-7.0.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tinycss2-1.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_ha0e22de_103.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py313h07c4f96_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py314h5bd0f2a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_utils-0.1.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.0.1-py313h7037e92_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.0.1-py314h9891dd4_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/unicodedata2-17.0.0-py314h5bd0f2a_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/uri-template-1.3.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.6.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.35.4-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.36.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webcolors-25.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webencodings-0.5.1-pyhd8ed1ab_3.conda
@@ -1204,16 +1213,17 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/01/c1/e32fb13d9cb5afc5f25f0c84ca21d3ea2d5380f4c06417b814ecc9bf0f38/dags-0.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/fd/8e/424b8c6e78bd9837d14ff7df01a9829fc883ba2ab4ea787d4f848435f23f/greenlet-3.3.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/b8/14/bab308fc2c1b5228c3224ec2bf928ce2e4d21d8046c161e44a2012b5203e/greenlet-3.3.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/a8/f7/ae4ecf183d9693cd5fcce7ee063c5e54f173b66dc80a8a79951861e1b557/jax-0.8.2-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/6b/e0/91e5762a7ddb6351b07c742ca407cd28e26043d6945d6228b6c1b0881a45/jaxlib-0.8.2-cp313-cp313-manylinux_2_27_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/5e/27/2e6032727e41ce74914277478021140947af59127d68aa9e6f3776b428fd/jaxlib-0.8.2-cp314-cp314-manylinux_2_27_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/c6/bb/82c7dcf38070b46172a517e2334e665c5bf374a262f99a283ea454bece7c/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
-      - pypi: git+https://github.com/optimagic-dev/optimagic.git?branch=main#e02ea4743cac9f861a5813f3b4b1283fd2ade730
-      - pypi: https://files.pythonhosted.org/packages/15/07/284f757f63f8a8d69ed4472bfd85122bd086e637bf4ed09de572d575a693/pandas-2.3.3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: git+https://github.com/optimagic-dev/optimagic.git#522b8c9a21226569ffd25e950e44f0c5de308c9d
+      - pypi: https://files.pythonhosted.org/packages/15/b2/0e62f78c0c5ba7e3d2c5945a82456f4fac76c480940f805e0b97fcbc2f65/pandas-2.3.3-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/75/58/3af430d0de0b95d5adf7e576067e07d750ba76e28d142871982464fb40db/pdbp-1.8.1-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/0e/50/80a8d080ac7d3d321e5e5d420c9a522b0aa770ec7013ea91f9a8b7d36e4a/sqlalchemy-2.0.45-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/b3/27/caf606ee924282fe4747ee4fd454b335a72a6e018f97eab5ff7f28199e16/sqlalchemy-2.0.45-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/65/44/bb509c3d2c0b5a87e7a5af1d5917a402a32ff026f777a6d7cb6990746cbb/tabcompleter-1.4.0-py3-none-any.whl
       - pypi: ./
       osx-arm64:
@@ -1222,35 +1232,35 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/anyio-4.12.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/appnope-0.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/argon2-cffi-25.1.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/argon2-cffi-bindings-25.1.0-py313h6535dbc_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/argon2-cffi-bindings-25.1.0-py314h0612a62_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/arrow-1.4.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/async-lru-2.0.5-pyh29332c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/babel-2.17.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/backports.zstd-1.3.0-py313h48bb75e_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.3.0-py314h680f03e_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.14.3-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-6.3.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-with-css-6.3.0-h5f6438b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-1.2.0-h7d5ae5b_1.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-bin-1.2.0-hc919400_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-python-1.2.0-py313hde1f3bb_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-python-1.2.0-py314h3daef5d_1.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/bzip2-1.0.8-hd037594_8.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/cffi-2.0.0-py313h224173a_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/cffi-2.0.0-py314h44086f9_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/choreographer-1.2.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.3-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/contourpy-1.3.3-py313ha61f8ec_3.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/coverage-7.13.1-py313h65a2061_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.11-py313hd8ed1ab_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/contourpy-1.3.3-py314h784bc60_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/coverage-7.13.1-py314h6e9b3f0_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.14.2-py314hd8ed1ab_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/debugpy-1.8.19-py313hc37fe24_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/debugpy-1.8.19-py314hf820bb6_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/defusedxml-0.7.1-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/distlib-0.4.0-pyhd8ed1ab_0.conda
@@ -1259,7 +1269,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.20.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filterpy-1.4.5-pyhd8ed1ab_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/fonttools-4.61.1-py313h7d74516_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonttools-4.61.1-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/fqdn-1.5.1-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/freetype-2.14.1-hce30654_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/h11-0.16.0-pyhcf101f3_1.conda
@@ -1292,7 +1302,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab-4.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_pygments-0.3.0-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_server-2.28.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/kiwisolver-1.4.9-py313h7add70c_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/kiwisolver-1.4.9-py314h42813c9_2.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/krb5-1.21.3-h237132a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/lark-1.3.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/lcms2-2.17-h7eeda09_0.conda
@@ -1328,12 +1338,12 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/logistro-2.0.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/lz4-c-1.10.0-h286801f_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/markdown-it-py-4.0.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/markupsafe-3.0.3-py313h7d74516_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/matplotlib-base-3.10.8-py313h58042b9_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/markupsafe-3.0.3-pyh7db6752_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/matplotlib-base-3.10.8-py314hd63e3f0_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mdit-py-plugins-0.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mdurl-0.1.2-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/memray-1.19.1-py313h78c9487_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/memray-1.19.1-py314habef2a7_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mistune-3.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/narwhals-2.15.0-pyhcf101f3_0.conda
@@ -1345,31 +1355,31 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.6.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/nodeenv-1.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/notebook-shim-0.2.4-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/numpy-2.3.5-py313h16eae64_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/numpy-2.3.5-py314hae46ccb_1.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/openjpeg-2.5.4-hbfb3c88_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/openssl-3.6.0-h5503f6c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/orjson-3.11.5-py313hfea8034_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/orjson-3.11.5-py314hda6d10a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/overrides-7.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pandocfilters-1.5.0-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pexpect-4.9.0-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pillow-12.1.0-py313h45e5a15_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pillow-12.1.0-py314hab283cf_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/plotly-6.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.6.0-pyhf9edf01_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pre-commit-4.5.1-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prometheus_client-0.23.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/psutil-7.2.1-py313h6688731_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/psutil-7.2.1-py314ha14b1ff_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pthread-stubs-0.4-hd74edd7_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ptyprocess-0.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pybaum-0.1.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pycparser-2.22-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pygments-2.19.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-core-12.1-py313h40b429f_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-framework-cocoa-12.1-py313hcc5defa_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-core-12.1-py314h3a4d195_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-framework-cocoa-12.1-py314h36abed7_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.0.2-pyhcf101f3_0.conda
@@ -1377,16 +1387,16 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-memray-1.8.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-timeout-2.4.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.13.11-hfc2f54d_100_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.14.2-h40d2674_100_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.2-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.13.11-h4df99d1_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.14.2-h4df99d1_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-json-logger-2.0.7-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-kaleido-1.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.14-8_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyyaml-6.0.3-py313h7d74516_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pyyaml-6.0.3-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyzmq-27.1.0-py312hd65ceae_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/qhull-2020.2-h420ef59_5.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/readline-8.3-h46df422_0.conda
@@ -1396,31 +1406,32 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3986-validator-0.1.1-pyh9f0ad1d_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3987-syntax-1.1.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/rich-14.2.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/rpds-py-0.30.0-py313h2c089d5_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/scipy-1.16.3-py313h29d7d31_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/rpds-py-0.30.0-py314haad56a0_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/scipy-1.16.3-py314h725efaa_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/send2trash-2.0.0-pyh5552912_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.9.0-pyhff2d567_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/simplejson-3.20.2-py313h6535dbc_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/simplejson-3.20.2-py314h0612a62_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/snakeviz-2.2.2-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/sniffio-1.3.1-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/soupsieve-2.8.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/terminado-0.18.1-pyhc90fa1f_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/textual-7.0.0-pyhcf101f3_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/textual-7.0.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tinycss2-1.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tk-8.6.13-h892fb3f_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tornado-6.5.4-py313h6535dbc_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tornado-6.5.4-py314h0612a62_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_utils-0.1.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ukkonen-1.0.1-py313hc50a443_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ukkonen-1.0.1-py314h6b18a25_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/unicodedata2-17.0.0-py314h0612a62_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/uri-template-1.3.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.6.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.35.4-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.36.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webcolors-25.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webencodings-0.5.1-pyhd8ed1ab_3.conda
@@ -1436,13 +1447,14 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/01/c1/e32fb13d9cb5afc5f25f0c84ca21d3ea2d5380f4c06417b814ecc9bf0f38/dags-0.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a8/f7/ae4ecf183d9693cd5fcce7ee063c5e54f173b66dc80a8a79951861e1b557/jax-0.8.2-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/c5/22/c0ec75e43a13b2457d78d509f49b49a57fa302ffced4f4a2778e428cb0a6/jaxlib-0.8.2-cp313-cp313-macosx_11_0_arm64.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/9d/dca93d916bf8664d7a2bb73ea3d219028dabbe382c31774348963287356a/jaxlib-0.8.2-cp314-cp314-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/d9/a1/4008f14bbc616cfb1ac5b39ea485f9c63031c4634ab3f4cf72e7541f816a/ml_dtypes-0.5.4-cp313-cp313-macosx_10_13_universal2.whl
+      - pypi: https://files.pythonhosted.org/packages/72/4e/1339dc6e2557a344f5ba5590872e80346f76f6cb2ac3dd16e4666e88818c/ml_dtypes-0.5.4-cp314-cp314-macosx_10_13_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
-      - pypi: git+https://github.com/optimagic-dev/optimagic.git?branch=main#e02ea4743cac9f861a5813f3b4b1283fd2ade730
-      - pypi: https://files.pythonhosted.org/packages/31/94/72fac03573102779920099bcac1c3b05975c2cb5f01eac609faf34bed1ca/pandas-2.3.3-cp313-cp313-macosx_11_0_arm64.whl
+      - pypi: git+https://github.com/optimagic-dev/optimagic.git#522b8c9a21226569ffd25e950e44f0c5de308c9d
+      - pypi: https://files.pythonhosted.org/packages/21/00/266d6b357ad5e6d3ad55093a7e8efc7dd245f5a842b584db9f30b0f0a287/pandas-2.3.3-cp314-cp314-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/75/58/3af430d0de0b95d5adf7e576067e07d750ba76e28d142871982464fb40db/pdbp-1.8.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/bf/e1/3ccb13c643399d22289c6a9786c1a91e3dcbb68bce4beb44926ac2c557bf/sqlalchemy-2.0.45-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/65/44/bb509c3d2c0b5a87e7a5af1d5917a402a32ff026f777a6d7cb6990746cbb/tabcompleter-1.4.0-py3-none-any.whl
@@ -1452,35 +1464,35 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/anyio-4.12.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/argon2-cffi-25.1.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/argon2-cffi-bindings-25.1.0-py313h5ea7bf4_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/argon2-cffi-bindings-25.1.0-py314h5a2d7ad_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/arrow-1.4.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/async-lru-2.0.5-pyh29332c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/babel-2.17.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/backports.zstd-1.3.0-py313h2a31948_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.3.0-py314h680f03e_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.14.3-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-6.3.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-with-css-6.3.0-h5f6438b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/brotli-1.2.0-h2d644bc_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/brotli-bin-1.2.0-hfd05255_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/brotli-python-1.2.0-py313h3ebfc14_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/brotli-python-1.2.0-py314he701e3d_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/bzip2-1.0.8-h0ad9c76_8.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-h4c7d964_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/cffi-2.0.0-py313h5ea7bf4_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/cffi-2.0.0-py314h5a2d7ad_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/choreographer-1.2.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.3-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/contourpy-1.3.3-py313hf069bd2_3.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/coverage-7.13.1-py313hd650c13_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.11-py313hd8ed1ab_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/contourpy-1.3.3-py314h909e829_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/coverage-7.13.1-py314h2359020_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.14.2-py314hd8ed1ab_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/debugpy-1.8.19-py313h927ade5_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/debugpy-1.8.19-py314hb98de8c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/defusedxml-0.7.1-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/distlib-0.4.0-pyhd8ed1ab_0.conda
@@ -1489,7 +1501,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.20.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filterpy-1.4.5-pyhd8ed1ab_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/fonttools-4.61.1-py313hd650c13_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonttools-4.61.1-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/fqdn-1.5.1-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/freetype-2.14.1-h57928b3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/h11-0.16.0-pyhcf101f3_1.conda
@@ -1523,7 +1535,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab-4.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_pygments-0.3.0-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_server-2.28.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/kiwisolver-1.4.9-py313h1a38498_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/kiwisolver-1.4.9-py314hf309875_2.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/krb5-1.21.3-hdf4eb48_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/lark-1.3.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/lcms2-2.17-hbcf6048_0.conda
@@ -1540,7 +1552,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/win-64/libfreetype6-2.14.1-hdbac1cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/libgcc-15.2.0-h8ee18e1_16.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/libgomp-15.2.0-h8ee18e1_16.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/libhwloc-2.12.1-default_h4379cf1_1003.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/libhwloc-2.12.2-default_h4379cf1_1000.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/libiconv-1.18-hc1393d2_2.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/libjpeg-turbo-3.1.2-hfd05255_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/liblapack-3.11.0-5_hf9ab0e9_mkl.conda
@@ -1558,8 +1570,8 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/win-64/libzlib-1.3.1-h2466b09_2.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/llvm-openmp-21.1.8-h4fa8253_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/logistro-2.0.1-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/markupsafe-3.0.3-py313hd650c13_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/matplotlib-base-3.10.8-py313he1ded55_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/markupsafe-3.0.3-pyh7db6752_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/matplotlib-base-3.10.8-py314hfa45d96_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mistune-3.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/mkl-2025.3.0-hac47afa_455.conda
@@ -1572,22 +1584,22 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.6.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/nodeenv-1.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/notebook-shim-0.2.4-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/numpy-2.3.5-py313hce7ae62_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/numpy-2.3.5-py314h06c3c77_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/openjpeg-2.5.4-h24db6dd_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/openssl-3.6.0-h725018a_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/orjson-3.11.4-py313hfbe8231_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/orjson-3.11.5-py314h64f83cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/overrides-7.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pandocfilters-1.5.0-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/pillow-12.1.0-py313h38f99e1_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/pillow-12.1.0-py314h61b30b5_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/plotly-6.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.6.0-pyhf9edf01_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pre-commit-4.5.1-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prometheus_client-0.23.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/psutil-7.2.1-py313h5fd188c_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/psutil-7.2.1-py314hc5dbbe4_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/pthread-stubs-0.4-h0e40799_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pybaum-0.1.3-pyhd8ed1ab_1.conda
@@ -1599,18 +1611,18 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-cov-7.0.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-timeout-2.4.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/python-3.13.11-h09917c8_100_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/python-3.14.2-h4b44e0e_100_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.2-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.13.11-h4df99d1_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.14.2-h4df99d1_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-json-logger-2.0.7-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-kaleido-1.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.14-8_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/pywin32-311-py313h40c08fc_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/pywinpty-2.0.15-py313h5813708_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/pyyaml-6.0.3-py313hd650c13_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/pywin32-311-py314h8f8f202_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/pywinpty-2.0.15-py314h51f0985_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pyyaml-6.0.3-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/pyzmq-27.1.0-py312hbb5da91_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/qhull-2020.2-hc790b64_5.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/referencing-0.37.0-pyhcf101f3_0.conda
@@ -1618,35 +1630,36 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3339-validator-0.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3986-validator-0.1.1-pyh9f0ad1d_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3987-syntax-1.1.0-pyhe01879c_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/rpds-py-0.30.0-py313hfbe8231_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/scipy-1.16.3-py313he51e9a2_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/rpds-py-0.30.0-py314h9f07db2_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/scipy-1.16.3-py314h221f224_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/send2trash-2.0.0-pyh6dadd2b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.9.0-pyhff2d567_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/simplejson-3.20.2-py313h5ea7bf4_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/simplejson-3.20.2-py314h5a2d7ad_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/snakeviz-2.2.2-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/sniffio-1.3.1-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/soupsieve-2.8.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/tbb-2022.3.0-hd094cb3_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/tbb-2022.3.0-h3155e25_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/terminado-0.18.1-pyh6dadd2b_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tinycss2-1.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/tk-8.6.13-h2c6b04d_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/tornado-6.5.4-py313h5ea7bf4_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/tornado-6.5.4-py314h5a2d7ad_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_utils-0.1.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/ucrt-10.0.26100.0-h57928b3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/ukkonen-1.0.1-py313hf069bd2_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/ukkonen-1.0.1-py314h909e829_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/unicodedata2-17.0.0-py314h5a2d7ad_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/uri-template-1.3.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.6.3-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/vc-14.3-h41ae7f8_34.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/vc14_runtime-14.44.35208-h818238b_34.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/vcomp14-14.44.35208-h818238b_34.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.35.4-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.36.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webcolors-25.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webencodings-0.5.1-pyhd8ed1ab_3.conda
@@ -1664,17 +1677,18 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/01/c1/e32fb13d9cb5afc5f25f0c84ca21d3ea2d5380f4c06417b814ecc9bf0f38/dags-0.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/7e/71/ba21c3fb8c5dce83b8c01f458a42e99ffdb1963aeec08fff5a18588d8fd7/greenlet-3.3.0-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/7c/9a/9030e6f9aa8fd7808e9c31ba4c38f87c4f8ec324ee67431d181fe396d705/greenlet-3.3.0-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/a8/f7/ae4ecf183d9693cd5fcce7ee063c5e54f173b66dc80a8a79951861e1b557/jax-0.8.2-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/85/68/25b38673b07a808616ce7b6efb3eed491f983f3373a09cbbd03f67178563/jaxlib-0.8.2-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/b3/8c/af5a00b07a446414edf6b84a7397eab02cf01ba44b6ae1fce7798ce4c127/jaxlib-0.8.2-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/e1/8b/200088c6859d8221454825959df35b5244fa9bdf263fd0249ac5fb75e281/ml_dtypes-0.5.4-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/e9/93/2bfed22d2498c468f6bcd0d9f56b033eaa19f33320389314c19ef6766413/ml_dtypes-0.5.4-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
-      - pypi: git+https://github.com/optimagic-dev/optimagic.git?branch=main#e02ea4743cac9f861a5813f3b4b1283fd2ade730
-      - pypi: https://files.pythonhosted.org/packages/4f/c7/e54682c96a895d0c808453269e0b5928a07a127a15704fedb643e9b0a4c8/pandas-2.3.3-cp313-cp313-win_amd64.whl
+      - pypi: git+https://github.com/optimagic-dev/optimagic.git#522b8c9a21226569ffd25e950e44f0c5de308c9d
+      - pypi: https://files.pythonhosted.org/packages/a6/3d/124ac75fcd0ecc09b8fdccb0246ef65e35b012030defb0e0eba2cbbbe948/pandas-2.3.3-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/75/58/3af430d0de0b95d5adf7e576067e07d750ba76e28d142871982464fb40db/pdbp-1.8.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/5a/dc/491b7661614ab97483abf2056be1deee4dc2490ecbf7bff9ab5cdbac86e1/pyreadline3-3.5.4-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/c0/c5/d17113020b2d43073412aeca09b60d2009442420372123b8d49cc253f8b8/sqlalchemy-2.0.45-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/89/a2/0e1590e9adb292b1d576dbcf67ff7df8cf55e56e78d2c927686d01080f4b/sqlalchemy-2.0.45-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/65/44/bb509c3d2c0b5a87e7a5af1d5917a402a32ff026f777a6d7cb6990746cbb/tabcompleter-1.4.0-py3-none-any.whl
       - pypi: ./
   test-gpu:
@@ -1691,13 +1705,13 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/anyio-4.12.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/argon2-cffi-25.1.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/argon2-cffi-bindings-25.1.0-py313h07c4f96_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/argon2-cffi-bindings-25.1.0-py314h5bd0f2a_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/arrow-1.4.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/async-lru-2.0.5-pyh29332c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/babel-2.17.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/backports.zstd-1.3.0-py313h18e8e13_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.3.0-py314h680f03e_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.14.3-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/binutils_impl_linux-64-2.45-default_hfdba357_105.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/binutils_linux-64-2.45-default_h4852527_105.conda
@@ -1705,22 +1719,22 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-with-css-6.3.0-h5f6438b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.2.0-hed03a55_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.2.0-hb03c661_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py313hf159716_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py314h3de4e8d_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_8.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/c-ares-1.34.6-hb03c661_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py313hf46b229_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py314h4a8dc5f_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/choreographer-1.2.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.3-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py313h7037e92_3.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/coverage-7.13.1-py313h3dea7bd_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.11-py313hd8ed1ab_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py314h9891dd4_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/coverage-7.13.1-py314h67df5f8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.14.2-py314hd8ed1ab_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cccl_linux-64-12.9.27-ha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-crt-dev_linux-64-12.9.86-ha770c72_2.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-crt-tools-12.9.86-ha770c72_2.conda
@@ -1741,7 +1755,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvvm-tools-12.9.86-h4bc722e_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-version-12.9-h4f385c5_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.18-py313h5d5ffb9_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.18-py314h42812f9_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/defusedxml-0.7.1-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/distlib-0.4.0-pyhd8ed1ab_0.conda
@@ -1751,7 +1765,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.20.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filterpy-1.4.5-pyhd8ed1ab_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/fonttools-4.61.1-py313h3dea7bd_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonttools-4.61.1-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/fqdn-1.5.1-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/freetype-2.14.1-ha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/gcc_impl_linux-64-14.3.0-he8b2097_16.conda
@@ -1793,7 +1807,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_server-2.28.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/kernel-headers_linux-64-4.18.0-he073ed8_9.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/keyutils-1.6.3-hb9d3cd8_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py313hc8edb43_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py314h97ea11e_2.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/krb5-1.21.3-h659f571_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/lark-1.3.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/lcms2-2.17-h717163a_0.conda
@@ -1852,12 +1866,12 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/lz4-c-1.10.0-h5888daf_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/lzo-2.10-h280c20c_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/markdown-it-py-4.0.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/markupsafe-3.0.3-py313h3dea7bd_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.8-py313h683a580_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/markupsafe-3.0.3-pyh7db6752_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.8-py314h1194b4b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mdit-py-plugins-0.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mdurl-0.1.2-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/memray-1.19.1-py313h422961c_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/memray-1.19.1-py314hef15ded_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mistune-3.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/narwhals-2.15.0-pyhcf101f3_0.conda
@@ -1870,24 +1884,24 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.6.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/nodeenv-1.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/notebook-shim-0.2.4-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py313hf6604e3_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py314h2b28147_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.4-h55fea9a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.0-h26f9b46_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.11.5-py313h541fbb8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.11.5-py314h3b757c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/overrides-7.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/p11-kit-0.25.10-h3435931_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pandocfilters-1.5.0-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pexpect-4.9.0-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.1.0-py313h80991f8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.1.0-py314h8ec4b1a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/plotly-6.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.6.0-pyhf9edf01_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pre-commit-4.5.1-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prometheus_client-0.23.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.1-py313h54dd161_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.1-py314h0f05182_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ptyprocess-0.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
@@ -1901,16 +1915,16 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-memray-1.8.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-timeout-2.4.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.11-hc97d973_100_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.14.2-h32b2ec7_100_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.2-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.13.11-h4df99d1_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.14.2-h4df99d1_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-json-logger-2.0.7-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-kaleido-1.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.14-8_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py313h3dea7bd_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pyyaml-6.0.3-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/pyzmq-27.1.0-py312hfb55c3c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/qhull-2020.2-h434a139_5.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda
@@ -1920,11 +1934,11 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3986-validator-0.1.1-pyh9f0ad1d_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3987-syntax-1.1.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/rich-14.2.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.30.0-py313h843e2db_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.16.3-py313h4b8bb8b_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.30.0-py314h2e6c369_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.16.3-py314hf07bd8e_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/send2trash-2.0.0-pyha191276_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.9.0-pyhff2d567_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/simplejson-3.20.2-py313h07c4f96_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/simplejson-3.20.2-py314h5bd0f2a_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/snakeviz-2.2.2-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/sniffio-1.3.1-pyhd8ed1ab_2.conda
@@ -1932,20 +1946,21 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/sysroot_linux-64-2.28-h4ee821c_9.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/terminado-0.18.1-pyhc90fa1f_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/textual-7.0.0-pyhcf101f3_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/textual-7.0.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tinycss2-1.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_ha0e22de_103.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py313h07c4f96_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py314h5bd0f2a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_utils-0.1.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.0.1-py313h7037e92_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.0.1-py314h9891dd4_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/unicodedata2-17.0.0-py314h5bd0f2a_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/uri-template-1.3.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.6.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.35.4-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.36.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webcolors-25.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webencodings-0.5.1-pyhd8ed1ab_3.conda
@@ -1961,13 +1976,14 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/01/c1/e32fb13d9cb5afc5f25f0c84ca21d3ea2d5380f4c06417b814ecc9bf0f38/dags-0.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/fd/8e/424b8c6e78bd9837d14ff7df01a9829fc883ba2ab4ea787d4f848435f23f/greenlet-3.3.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/b8/14/bab308fc2c1b5228c3224ec2bf928ce2e4d21d8046c161e44a2012b5203e/greenlet-3.3.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/a8/f7/ae4ecf183d9693cd5fcce7ee063c5e54f173b66dc80a8a79951861e1b557/jax-0.8.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/53/f2/44ad0ce1d115f0f6be10f4af0ca05a18afb838b06e6ca6b01ba4b0137421/jax_cuda12_pjrt-0.8.2-py3-none-manylinux_2_27_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/1c/38/4ba2486f95fcf2120723932feacdded438e785258148b18a703cd1177e41/jax_cuda12_plugin-0.8.2-cp313-cp313-manylinux_2_27_x86_64.whl
-      - pypi: https://files.pythonhosted.org/packages/6b/e0/91e5762a7ddb6351b07c742ca407cd28e26043d6945d6228b6c1b0881a45/jaxlib-0.8.2-cp313-cp313-manylinux_2_27_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/27/58/a5a27d4677d6890570f7e58cecd51891469cb620e6f64c8faed4935d93d0/jax_cuda12_plugin-0.8.2-cp314-cp314-manylinux_2_27_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/5e/27/2e6032727e41ce74914277478021140947af59127d68aa9e6f3776b428fd/jaxlib-0.8.2-cp314-cp314-manylinux_2_27_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/c6/bb/82c7dcf38070b46172a517e2334e665c5bf374a262f99a283ea454bece7c/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/77/3c/aa88abe01f3be3d1f8f787d1d33dc83e76fec05945f9a28fbb41cfb99cd5/nvidia_cublas_cu12-12.9.1.4-py3-none-manylinux_2_27_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/18/2a/d4cd8506d2044e082f8cd921be57392e6a9b5ccd3ffdf050362430a3d5d5/nvidia_cuda_cccl_cu12-12.9.27-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/c1/2e/b84e32197e33f39907b455b83395a017e697c07a449a2b15fd07fc1c9981/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_x86_64.whl
@@ -1982,10 +1998,10 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/46/0c/c75bbfb967457a0b7670b8ad267bfc4fffdf341c074e0a80db06c24ccfd4/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/64/b9/6ab941001c23cfb43499b5b0b7417b0bb4dfba3a29ffa2b06985422dad50/nvidia_nvshmem_cu12-3.5.19-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
-      - pypi: git+https://github.com/optimagic-dev/optimagic.git?branch=main#e02ea4743cac9f861a5813f3b4b1283fd2ade730
-      - pypi: https://files.pythonhosted.org/packages/15/07/284f757f63f8a8d69ed4472bfd85122bd086e637bf4ed09de572d575a693/pandas-2.3.3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: git+https://github.com/optimagic-dev/optimagic.git#522b8c9a21226569ffd25e950e44f0c5de308c9d
+      - pypi: https://files.pythonhosted.org/packages/15/b2/0e62f78c0c5ba7e3d2c5945a82456f4fac76c480940f805e0b97fcbc2f65/pandas-2.3.3-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/75/58/3af430d0de0b95d5adf7e576067e07d750ba76e28d142871982464fb40db/pdbp-1.8.1-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/0e/50/80a8d080ac7d3d321e5e5d420c9a522b0aa770ec7013ea91f9a8b7d36e4a/sqlalchemy-2.0.45-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/b3/27/caf606ee924282fe4747ee4fd454b335a72a6e018f97eab5ff7f28199e16/sqlalchemy-2.0.45-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/65/44/bb509c3d2c0b5a87e7a5af1d5917a402a32ff026f777a6d7cb6990746cbb/tabcompleter-1.4.0-py3-none-any.whl
       - pypi: ./
   ty:
@@ -2002,36 +2018,36 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/anyio-4.12.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/argon2-cffi-25.1.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/argon2-cffi-bindings-25.1.0-py313h07c4f96_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/argon2-cffi-bindings-25.1.0-py314h5bd0f2a_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/arrow-1.4.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/async-lru-2.0.5-pyh29332c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/babel-2.17.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/backports.zstd-1.3.0-py313h18e8e13_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.3.0-py314h680f03e_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.14.3-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-6.3.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-with-css-6.3.0-h5f6438b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.2.0-hed03a55_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.2.0-hb03c661_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py313hf159716_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py314h3de4e8d_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_8.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/c-ares-1.34.6-hb03c661_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py313hf46b229_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py314h4a8dc5f_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/choreographer-1.2.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.3-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py313h7037e92_3.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/coverage-7.13.1-py313h3dea7bd_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.11-py313hd8ed1ab_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py314h9891dd4_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/coverage-7.13.1-py314h67df5f8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.14.2-py314hd8ed1ab_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.18-py313h5d5ffb9_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.18-py314h42812f9_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/defusedxml-0.7.1-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/distlib-0.4.0-pyhd8ed1ab_0.conda
@@ -2041,7 +2057,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.20.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filterpy-1.4.5-pyhd8ed1ab_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/fonttools-4.61.1-py313h3dea7bd_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonttools-4.61.1-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/fqdn-1.5.1-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/freetype-2.14.1-ha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/gmp-6.3.0-hac33072_2.conda
@@ -2078,7 +2094,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_pygments-0.3.0-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_server-2.28.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/keyutils-1.6.3-hb9d3cd8_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py313hc8edb43_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py314h97ea11e_2.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/krb5-1.21.3-h659f571_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/lark-1.3.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/lcms2-2.17-h717163a_0.conda
@@ -2132,12 +2148,12 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/lz4-c-1.10.0-h5888daf_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/lzo-2.10-h280c20c_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/markdown-it-py-4.0.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/markupsafe-3.0.3-py313h3dea7bd_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.8-py313h683a580_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/markupsafe-3.0.3-pyh7db6752_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.8-py314h1194b4b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mdit-py-plugins-0.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mdurl-0.1.2-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/memray-1.19.1-py313h422961c_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/memray-1.19.1-py314hef15ded_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mistune-3.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/narwhals-2.15.0-pyhcf101f3_0.conda
@@ -2150,24 +2166,24 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.6.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/nodeenv-1.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/notebook-shim-0.2.4-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py313hf6604e3_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py314h2b28147_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.4-h55fea9a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.0-h26f9b46_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.11.5-py313h541fbb8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.11.5-py314h3b757c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/overrides-7.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/p11-kit-0.25.10-h3435931_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pandocfilters-1.5.0-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pexpect-4.9.0-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.1.0-py313h80991f8_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.1.0-py314h8ec4b1a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/plotly-6.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.6.0-pyhf9edf01_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pre-commit-4.5.1-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prometheus_client-0.23.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.1-py313h54dd161_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.1-py314h0f05182_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ptyprocess-0.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
@@ -2181,16 +2197,16 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-memray-1.8.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-timeout-2.4.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.11-hc97d973_100_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.14.2-h32b2ec7_100_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.2-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.13.11-h4df99d1_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.14.2-h4df99d1_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-json-logger-2.0.7-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-kaleido-1.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.14-8_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py313h3dea7bd_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pyyaml-6.0.3-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/pyzmq-27.1.0-py312hfb55c3c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/qhull-2020.2-h434a139_5.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda
@@ -2200,31 +2216,32 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3986-validator-0.1.1-pyh9f0ad1d_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3987-syntax-1.1.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/rich-14.2.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.30.0-py313h843e2db_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.16.3-py313h4b8bb8b_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.30.0-py314h2e6c369_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.16.3-py314hf07bd8e_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/send2trash-2.0.0-pyha191276_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.9.0-pyhff2d567_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/simplejson-3.20.2-py313h07c4f96_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/simplejson-3.20.2-py314h5bd0f2a_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/snakeviz-2.2.2-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/sniffio-1.3.1-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/soupsieve-2.8.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/terminado-0.18.1-pyhc90fa1f_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/textual-7.0.0-pyhcf101f3_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/textual-7.0.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tinycss2-1.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_ha0e22de_103.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py313h07c4f96_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py314h5bd0f2a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_utils-0.1.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.0.1-py313h7037e92_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.0.1-py314h9891dd4_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/linux-64/unicodedata2-17.0.0-py314h5bd0f2a_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/uri-template-1.3.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.6.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.35.4-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.36.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webcolors-25.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webencodings-0.5.1-pyhd8ed1ab_3.conda
@@ -2240,19 +2257,20 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/01/c1/e32fb13d9cb5afc5f25f0c84ca21d3ea2d5380f4c06417b814ecc9bf0f38/dags-0.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/fd/8e/424b8c6e78bd9837d14ff7df01a9829fc883ba2ab4ea787d4f848435f23f/greenlet-3.3.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/b8/14/bab308fc2c1b5228c3224ec2bf928ce2e4d21d8046c161e44a2012b5203e/greenlet-3.3.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/a8/f7/ae4ecf183d9693cd5fcce7ee063c5e54f173b66dc80a8a79951861e1b557/jax-0.8.2-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/6b/e0/91e5762a7ddb6351b07c742ca407cd28e26043d6945d6228b6c1b0881a45/jaxlib-0.8.2-cp313-cp313-manylinux_2_27_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/5e/27/2e6032727e41ce74914277478021140947af59127d68aa9e6f3776b428fd/jaxlib-0.8.2-cp314-cp314-manylinux_2_27_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/c6/bb/82c7dcf38070b46172a517e2334e665c5bf374a262f99a283ea454bece7c/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
-      - pypi: git+https://github.com/optimagic-dev/optimagic.git?branch=main#e02ea4743cac9f861a5813f3b4b1283fd2ade730
-      - pypi: https://files.pythonhosted.org/packages/15/07/284f757f63f8a8d69ed4472bfd85122bd086e637bf4ed09de572d575a693/pandas-2.3.3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: git+https://github.com/optimagic-dev/optimagic.git#522b8c9a21226569ffd25e950e44f0c5de308c9d
+      - pypi: https://files.pythonhosted.org/packages/15/b2/0e62f78c0c5ba7e3d2c5945a82456f4fac76c480940f805e0b97fcbc2f65/pandas-2.3.3-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/64/20/69f2a39792a653fd64d916cd563ed79ec6e5dcfa6408c4674021d810afcf/pandas_stubs-2.3.3.251219-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/75/58/3af430d0de0b95d5adf7e576067e07d750ba76e28d142871982464fb40db/pdbp-1.8.1-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/0e/50/80a8d080ac7d3d321e5e5d420c9a522b0aa770ec7013ea91f9a8b7d36e4a/sqlalchemy-2.0.45-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/b3/27/caf606ee924282fe4747ee4fd454b335a72a6e018f97eab5ff7f28199e16/sqlalchemy-2.0.45-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/65/44/bb509c3d2c0b5a87e7a5af1d5917a402a32ff026f777a6d7cb6990746cbb/tabcompleter-1.4.0-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/9e/4c/2f9ac5edbd0e67bf82f5cd04275c4e87cbbf69a78f43e5dcf90c1573d44e/ty-0.0.10-py3-none-manylinux_2_24_x86_64.whl
+      - pypi: https://files.pythonhosted.org/packages/74/18/8dd4fe6df1fd66f3e83b4798eddb1d8482d9d9b105f25099b76703402ebb/ty-0.0.11-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
       - pypi: https://files.pythonhosted.org/packages/e7/c1/56ef16bf5dcd255155cc736d276efa6ae0a5c26fd685e28f0412a4013c01/types_pytz-2025.2.0.20251108-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/bd/e0/1eed384f02555dde685fff1a1ac805c1c7dcb6dd019c916fe659b1c1f9ec/types_pyyaml-6.0.12.20250915-py3-none-any.whl
       - pypi: ./
@@ -2262,35 +2280,35 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/anyio-4.12.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/appnope-0.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/argon2-cffi-25.1.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/argon2-cffi-bindings-25.1.0-py313h6535dbc_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/argon2-cffi-bindings-25.1.0-py314h0612a62_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/arrow-1.4.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/async-lru-2.0.5-pyh29332c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/babel-2.17.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/backports.zstd-1.3.0-py313h48bb75e_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.3.0-py314h680f03e_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.14.3-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-6.3.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-with-css-6.3.0-h5f6438b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-1.2.0-h7d5ae5b_1.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-bin-1.2.0-hc919400_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-python-1.2.0-py313hde1f3bb_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-python-1.2.0-py314h3daef5d_1.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/bzip2-1.0.8-hd037594_8.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/cffi-2.0.0-py313h224173a_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/cffi-2.0.0-py314h44086f9_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/choreographer-1.2.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.3-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/contourpy-1.3.3-py313ha61f8ec_3.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/coverage-7.13.1-py313h65a2061_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.11-py313hd8ed1ab_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/contourpy-1.3.3-py314h784bc60_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/coverage-7.13.1-py314h6e9b3f0_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.14.2-py314hd8ed1ab_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/debugpy-1.8.19-py313hc37fe24_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/debugpy-1.8.19-py314hf820bb6_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/defusedxml-0.7.1-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/distlib-0.4.0-pyhd8ed1ab_0.conda
@@ -2299,7 +2317,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.20.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filterpy-1.4.5-pyhd8ed1ab_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/fonttools-4.61.1-py313h7d74516_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonttools-4.61.1-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/fqdn-1.5.1-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/freetype-2.14.1-hce30654_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/h11-0.16.0-pyhcf101f3_1.conda
@@ -2332,7 +2350,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab-4.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_pygments-0.3.0-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_server-2.28.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/kiwisolver-1.4.9-py313h7add70c_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/kiwisolver-1.4.9-py314h42813c9_2.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/krb5-1.21.3-h237132a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/lark-1.3.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/lcms2-2.17-h7eeda09_0.conda
@@ -2368,12 +2386,12 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/logistro-2.0.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/lz4-c-1.10.0-h286801f_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/markdown-it-py-4.0.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/markupsafe-3.0.3-py313h7d74516_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/matplotlib-base-3.10.8-py313h58042b9_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/markupsafe-3.0.3-pyh7db6752_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/matplotlib-base-3.10.8-py314hd63e3f0_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mdit-py-plugins-0.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mdurl-0.1.2-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/memray-1.19.1-py313h78c9487_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/memray-1.19.1-py314habef2a7_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mistune-3.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/narwhals-2.15.0-pyhcf101f3_0.conda
@@ -2385,31 +2403,31 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.6.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/nodeenv-1.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/notebook-shim-0.2.4-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/numpy-2.3.5-py313h16eae64_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/numpy-2.3.5-py314hae46ccb_1.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/openjpeg-2.5.4-hbfb3c88_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/openssl-3.6.0-h5503f6c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/orjson-3.11.5-py313hfea8034_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/orjson-3.11.5-py314hda6d10a_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/overrides-7.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pandocfilters-1.5.0-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pexpect-4.9.0-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pillow-12.1.0-py313h45e5a15_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pillow-12.1.0-py314hab283cf_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/plotly-6.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.6.0-pyhf9edf01_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pre-commit-4.5.1-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prometheus_client-0.23.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/psutil-7.2.1-py313h6688731_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/psutil-7.2.1-py314ha14b1ff_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pthread-stubs-0.4-hd74edd7_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ptyprocess-0.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pybaum-0.1.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pycparser-2.22-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pygments-2.19.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-core-12.1-py313h40b429f_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-framework-cocoa-12.1-py313hcc5defa_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-core-12.1-py314h3a4d195_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-framework-cocoa-12.1-py314h36abed7_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-9.0.2-pyhcf101f3_0.conda
@@ -2417,16 +2435,16 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-memray-1.8.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-timeout-2.4.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.13.11-hfc2f54d_100_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.14.2-h40d2674_100_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.2-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.13.11-h4df99d1_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.14.2-h4df99d1_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-json-logger-2.0.7-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-kaleido-1.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.14-8_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyyaml-6.0.3-py313h7d74516_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pyyaml-6.0.3-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyzmq-27.1.0-py312hd65ceae_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/qhull-2020.2-h420ef59_5.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/readline-8.3-h46df422_0.conda
@@ -2436,31 +2454,32 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3986-validator-0.1.1-pyh9f0ad1d_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3987-syntax-1.1.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/rich-14.2.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/rpds-py-0.30.0-py313h2c089d5_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/scipy-1.16.3-py313h29d7d31_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/rpds-py-0.30.0-py314haad56a0_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/scipy-1.16.3-py314h725efaa_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/send2trash-2.0.0-pyh5552912_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.9.0-pyhff2d567_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/simplejson-3.20.2-py313h6535dbc_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/simplejson-3.20.2-py314h0612a62_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/snakeviz-2.2.2-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/sniffio-1.3.1-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/soupsieve-2.8.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/terminado-0.18.1-pyhc90fa1f_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/textual-7.0.0-pyhcf101f3_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/textual-7.0.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tinycss2-1.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tk-8.6.13-h892fb3f_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tornado-6.5.4-py313h6535dbc_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tornado-6.5.4-py314h0612a62_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_utils-0.1.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ukkonen-1.0.1-py313hc50a443_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ukkonen-1.0.1-py314h6b18a25_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/osx-arm64/unicodedata2-17.0.0-py314h0612a62_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/uri-template-1.3.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.6.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.35.4-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.36.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webcolors-25.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webencodings-0.5.1-pyhd8ed1ab_3.conda
@@ -2476,18 +2495,19 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/01/c1/e32fb13d9cb5afc5f25f0c84ca21d3ea2d5380f4c06417b814ecc9bf0f38/dags-0.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a8/f7/ae4ecf183d9693cd5fcce7ee063c5e54f173b66dc80a8a79951861e1b557/jax-0.8.2-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/c5/22/c0ec75e43a13b2457d78d509f49b49a57fa302ffced4f4a2778e428cb0a6/jaxlib-0.8.2-cp313-cp313-macosx_11_0_arm64.whl
+      - pypi: https://files.pythonhosted.org/packages/d8/9d/dca93d916bf8664d7a2bb73ea3d219028dabbe382c31774348963287356a/jaxlib-0.8.2-cp314-cp314-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/d9/a1/4008f14bbc616cfb1ac5b39ea485f9c63031c4634ab3f4cf72e7541f816a/ml_dtypes-0.5.4-cp313-cp313-macosx_10_13_universal2.whl
+      - pypi: https://files.pythonhosted.org/packages/72/4e/1339dc6e2557a344f5ba5590872e80346f76f6cb2ac3dd16e4666e88818c/ml_dtypes-0.5.4-cp314-cp314-macosx_10_13_universal2.whl
       - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
-      - pypi: git+https://github.com/optimagic-dev/optimagic.git?branch=main#e02ea4743cac9f861a5813f3b4b1283fd2ade730
-      - pypi: https://files.pythonhosted.org/packages/31/94/72fac03573102779920099bcac1c3b05975c2cb5f01eac609faf34bed1ca/pandas-2.3.3-cp313-cp313-macosx_11_0_arm64.whl
+      - pypi: git+https://github.com/optimagic-dev/optimagic.git#522b8c9a21226569ffd25e950e44f0c5de308c9d
+      - pypi: https://files.pythonhosted.org/packages/21/00/266d6b357ad5e6d3ad55093a7e8efc7dd245f5a842b584db9f30b0f0a287/pandas-2.3.3-cp314-cp314-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/64/20/69f2a39792a653fd64d916cd563ed79ec6e5dcfa6408c4674021d810afcf/pandas_stubs-2.3.3.251219-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/75/58/3af430d0de0b95d5adf7e576067e07d750ba76e28d142871982464fb40db/pdbp-1.8.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/bf/e1/3ccb13c643399d22289c6a9786c1a91e3dcbb68bce4beb44926ac2c557bf/sqlalchemy-2.0.45-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/65/44/bb509c3d2c0b5a87e7a5af1d5917a402a32ff026f777a6d7cb6990746cbb/tabcompleter-1.4.0-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/e8/cd/9dd49e6d40e54d4b7d563f9e2a432c4ec002c0673a81266e269c4bc194ce/ty-0.0.10-py3-none-macosx_11_0_arm64.whl
+      - pypi: https://files.pythonhosted.org/packages/ad/01/3a563dba8b1255e474c35e1c3810b7589e81ae8c41df401b6a37c8e2cde9/ty-0.0.11-py3-none-macosx_11_0_arm64.whl
       - pypi: https://files.pythonhosted.org/packages/e7/c1/56ef16bf5dcd255155cc736d276efa6ae0a5c26fd685e28f0412a4013c01/types_pytz-2025.2.0.20251108-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/bd/e0/1eed384f02555dde685fff1a1ac805c1c7dcb6dd019c916fe659b1c1f9ec/types_pyyaml-6.0.12.20250915-py3-none-any.whl
       - pypi: ./
@@ -2496,35 +2516,35 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/_python_abi3_support-1.0-hd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/anyio-4.12.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/argon2-cffi-25.1.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/argon2-cffi-bindings-25.1.0-py313h5ea7bf4_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/argon2-cffi-bindings-25.1.0-py314h5a2d7ad_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/arrow-1.4.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/async-lru-2.0.5-pyh29332c3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/babel-2.17.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/backports.zstd-1.3.0-py313h2a31948_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.3.0-py314h680f03e_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.14.3-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-6.3.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/bleach-with-css-6.3.0-h5f6438b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/brotli-1.2.0-h2d644bc_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/brotli-bin-1.2.0-hfd05255_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/brotli-python-1.2.0-py313h3ebfc14_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/brotli-python-1.2.0-py314he701e3d_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/bzip2-1.0.8-h0ad9c76_8.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-h4c7d964_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/cffi-2.0.0-py313h5ea7bf4_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/cffi-2.0.0-py314h5a2d7ad_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/choreographer-1.2.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.3-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/contourpy-1.3.3-py313hf069bd2_3.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/coverage-7.13.1-py313hd650c13_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.11-py313hd8ed1ab_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/contourpy-1.3.3-py314h909e829_3.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/coverage-7.13.1-py314h2359020_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.14.2-py314hd8ed1ab_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/debugpy-1.8.19-py313h927ade5_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/debugpy-1.8.19-py314hb98de8c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/defusedxml-0.7.1-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/distlib-0.4.0-pyhd8ed1ab_0.conda
@@ -2533,7 +2553,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filelock-3.20.2-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/filterpy-1.4.5-pyhd8ed1ab_2.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/fonttools-4.61.1-py313hd650c13_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/fonttools-4.61.1-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/fqdn-1.5.1-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/freetype-2.14.1-h57928b3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/h11-0.16.0-pyhcf101f3_1.conda
@@ -2567,7 +2587,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab-4.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_pygments-0.3.0-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyterlab_server-2.28.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/kiwisolver-1.4.9-py313h1a38498_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/kiwisolver-1.4.9-py314hf309875_2.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/krb5-1.21.3-hdf4eb48_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/lark-1.3.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/lcms2-2.17-hbcf6048_0.conda
@@ -2584,7 +2604,7 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/win-64/libfreetype6-2.14.1-hdbac1cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/libgcc-15.2.0-h8ee18e1_16.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/libgomp-15.2.0-h8ee18e1_16.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/libhwloc-2.12.1-default_h4379cf1_1003.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/libhwloc-2.12.2-default_h4379cf1_1000.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/libiconv-1.18-hc1393d2_2.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/libjpeg-turbo-3.1.2-hfd05255_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/liblapack-3.11.0-5_hf9ab0e9_mkl.conda
@@ -2602,8 +2622,8 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/win-64/libzlib-1.3.1-h2466b09_2.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/llvm-openmp-21.1.8-h4fa8253_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/logistro-2.0.1-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/markupsafe-3.0.3-py313hd650c13_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/matplotlib-base-3.10.8-py313he1ded55_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/markupsafe-3.0.3-pyh7db6752_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/matplotlib-base-3.10.8-py314hfa45d96_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/mistune-3.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/mkl-2025.3.0-hac47afa_455.conda
@@ -2616,22 +2636,22 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.6.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/nodeenv-1.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/notebook-shim-0.2.4-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/numpy-2.3.5-py313hce7ae62_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/numpy-2.3.5-py314h06c3c77_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/openjpeg-2.5.4-h24db6dd_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/openssl-3.6.0-h725018a_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/orjson-3.11.4-py313hfbe8231_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/orjson-3.11.5-py314h64f83cb_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/overrides-7.7.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pandocfilters-1.5.0-pyhd8ed1ab_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/pillow-12.1.0-py313h38f99e1_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/pillow-12.1.0-py314h61b30b5_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/plotly-6.5.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pluggy-1.6.0-pyhf9edf01_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pre-commit-4.5.1-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prometheus_client-0.23.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/psutil-7.2.1-py313h5fd188c_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/psutil-7.2.1-py314hc5dbbe4_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/pthread-stubs-0.4-h0e40799_1002.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pybaum-0.1.3-pyhd8ed1ab_1.conda
@@ -2643,18 +2663,18 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-cov-7.0.0-pyhcf101f3_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-timeout-2.4.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytest-xdist-3.8.0-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/python-3.13.11-h09917c8_100_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/python-3.14.2-h4b44e0e_100_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-fastjsonschema-2.21.2-pyhe01879c_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.13.11-h4df99d1_100.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.14.2-h4df99d1_100.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-json-logger-2.0.7-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-kaleido-1.2.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.14-8_cp314.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/pywin32-311-py313h40c08fc_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/pywinpty-2.0.15-py313h5813708_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/pyyaml-6.0.3-py313hd650c13_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/pywin32-311-py314h8f8f202_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/pywinpty-2.0.15-py314h51f0985_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/pyyaml-6.0.3-pyh7db6752_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/pyzmq-27.1.0-py312hbb5da91_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/qhull-2020.2-hc790b64_5.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/referencing-0.37.0-pyhcf101f3_0.conda
@@ -2662,35 +2682,36 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3339-validator-0.1.4-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3986-validator-0.1.1-pyh9f0ad1d_0.tar.bz2
       - conda: https://conda.anaconda.org/conda-forge/noarch/rfc3987-syntax-1.1.0-pyhe01879c_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/rpds-py-0.30.0-py313hfbe8231_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/scipy-1.16.3-py313he51e9a2_2.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/rpds-py-0.30.0-py314h9f07db2_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/scipy-1.16.3-py314h221f224_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/send2trash-2.0.0-pyh6dadd2b_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.9.0-pyhff2d567_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/simplejson-3.20.2-py313h5ea7bf4_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/simplejson-3.20.2-py314h5a2d7ad_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/snakeviz-2.2.2-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/sniffio-1.3.1-pyhd8ed1ab_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/soupsieve-2.8.1-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/tbb-2022.3.0-hd094cb3_1.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/tbb-2022.3.0-h3155e25_2.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/terminado-0.18.1-pyh6dadd2b_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tinycss2-1.5.1-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/tk-8.6.13-h2c6b04d_3.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/tornado-6.5.4-py313h5ea7bf4_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/tornado-6.5.4-py314h5a2d7ad_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/typing_utils-0.1.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/ucrt-10.0.26100.0-h57928b3_0.conda
-      - conda: https://conda.anaconda.org/conda-forge/win-64/ukkonen-1.0.1-py313hf069bd2_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/ukkonen-1.0.1-py314h909e829_6.conda
+      - conda: https://conda.anaconda.org/conda-forge/win-64/unicodedata2-17.0.0-py314h5a2d7ad_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/uri-template-1.3.0-pyhd8ed1ab_1.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.6.3-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/vc-14.3-h41ae7f8_34.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/vc14_runtime-14.44.35208-h818238b_34.conda
       - conda: https://conda.anaconda.org/conda-forge/win-64/vcomp14-14.44.35208-h818238b_34.conda
-      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.35.4-pyhd8ed1ab_0.conda
+      - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.36.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webcolors-25.10.0-pyhd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/webencodings-0.5.1-pyhd8ed1ab_3.conda
@@ -2708,20 +2729,21 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/01/c1/e32fb13d9cb5afc5f25f0c84ca21d3ea2d5380f4c06417b814ecc9bf0f38/dags-0.4.2-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/43/f5/ee39c6e92acc742c052f137b47c210cd0a1b72dcd3f98495528bb4d27761/flatten_dict-0.4.2-py2.py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/7e/71/ba21c3fb8c5dce83b8c01f458a42e99ffdb1963aeec08fff5a18588d8fd7/greenlet-3.3.0-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/7c/9a/9030e6f9aa8fd7808e9c31ba4c38f87c4f8ec324ee67431d181fe396d705/greenlet-3.3.0-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/a8/f7/ae4ecf183d9693cd5fcce7ee063c5e54f173b66dc80a8a79951861e1b557/jax-0.8.2-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/85/68/25b38673b07a808616ce7b6efb3eed491f983f3373a09cbbd03f67178563/jaxlib-0.8.2-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/b3/8c/af5a00b07a446414edf6b84a7397eab02cf01ba44b6ae1fce7798ce4c127/jaxlib-0.8.2-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/e1/8b/200088c6859d8221454825959df35b5244fa9bdf263fd0249ac5fb75e281/ml_dtypes-0.5.4-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/e9/93/2bfed22d2498c468f6bcd0d9f56b033eaa19f33320389314c19ef6766413/ml_dtypes-0.5.4-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
-      - pypi: git+https://github.com/optimagic-dev/optimagic.git?branch=main#e02ea4743cac9f861a5813f3b4b1283fd2ade730
-      - pypi: https://files.pythonhosted.org/packages/4f/c7/e54682c96a895d0c808453269e0b5928a07a127a15704fedb643e9b0a4c8/pandas-2.3.3-cp313-cp313-win_amd64.whl
+      - pypi: git+https://github.com/optimagic-dev/optimagic.git#522b8c9a21226569ffd25e950e44f0c5de308c9d
+      - pypi: https://files.pythonhosted.org/packages/a6/3d/124ac75fcd0ecc09b8fdccb0246ef65e35b012030defb0e0eba2cbbbe948/pandas-2.3.3-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/64/20/69f2a39792a653fd64d916cd563ed79ec6e5dcfa6408c4674021d810afcf/pandas_stubs-2.3.3.251219-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/75/58/3af430d0de0b95d5adf7e576067e07d750ba76e28d142871982464fb40db/pdbp-1.8.1-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/5a/dc/491b7661614ab97483abf2056be1deee4dc2490ecbf7bff9ab5cdbac86e1/pyreadline3-3.5.4-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/c0/c5/d17113020b2d43073412aeca09b60d2009442420372123b8d49cc253f8b8/sqlalchemy-2.0.45-cp313-cp313-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/89/a2/0e1590e9adb292b1d576dbcf67ff7df8cf55e56e78d2c927686d01080f4b/sqlalchemy-2.0.45-cp314-cp314-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/65/44/bb509c3d2c0b5a87e7a5af1d5917a402a32ff026f777a6d7cb6990746cbb/tabcompleter-1.4.0-py3-none-any.whl
-      - pypi: https://files.pythonhosted.org/packages/42/36/82e66b9753a76964d26fd9bc3514ea0abce0a5ba5ad7d5f084070c6981da/ty-0.0.10-py3-none-win_amd64.whl
+      - pypi: https://files.pythonhosted.org/packages/df/04/5a5dfd0aec0ea99ead1e824ee6e347fb623c464da7886aa1e3660fb0f36c/ty-0.0.11-py3-none-win_amd64.whl
       - pypi: https://files.pythonhosted.org/packages/e7/c1/56ef16bf5dcd255155cc736d276efa6ae0a5c26fd685e28f0412a4013c01/types_pytz-2025.2.0.20251108-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/bd/e0/1eed384f02555dde685fff1a1ac805c1c7dcb6dd019c916fe659b1c1f9ec/types_pyyaml-6.0.12.20250915-py3-none-any.whl
       - pypi: ./
@@ -2835,43 +2857,43 @@ packages:
   - pkg:pypi/argon2-cffi?source=hash-mapping
   size: 18715
   timestamp: 1749017288144
-- conda: https://conda.anaconda.org/conda-forge/linux-64/argon2-cffi-bindings-25.1.0-py313h07c4f96_2.conda
-  sha256: ad188ccc06a06c633dc124b09e9e06fb9df4c32ffc38acc96ecc86e506062090
-  md5: 27bbec9f2f3a15d32b60ec5734f5b41c
+- conda: https://conda.anaconda.org/conda-forge/linux-64/argon2-cffi-bindings-25.1.0-py314h5bd0f2a_2.conda
+  sha256: 39234a99df3d2e3065383808ed8bfda36760de5ef590c54c3692bb53571ef02b
+  md5: 3cca1b74b2752917b5b65b81f61f0553
   depends:
   - __glibc >=2.17,<3.0.a0
-  - cffi >=1.0.1
+  - cffi >=2.0.0b1
   - libgcc >=14
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/argon2-cffi-bindings?source=hash-mapping
-  size: 35943
-  timestamp: 1762509452935
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/argon2-cffi-bindings-25.1.0-py313h6535dbc_2.conda
-  sha256: 05ea6fa7109235cfb4fc24526bae1fe82d88bbb5e697ab3945c313f5f041af5b
-  md5: e23e087109b2096db4cf9a3985bab329
+  size: 35598
+  timestamp: 1762509505285
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/argon2-cffi-bindings-25.1.0-py314h0612a62_2.conda
+  sha256: aab60bbaea5cc49dff37438d1ad469d64025cda2ce58103cf68da61701ed2075
+  md5: a240a79a49a95b388ef81ccda27a5e51
   depends:
   - __osx >=11.0
-  - cffi >=1.0.1
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
+  - cffi >=2.0.0b1
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/argon2-cffi-bindings?source=hash-mapping
-  size: 33947
-  timestamp: 1762510144907
-- conda: https://conda.anaconda.org/conda-forge/win-64/argon2-cffi-bindings-25.1.0-py313h5ea7bf4_2.conda
-  sha256: 3f8a1affdfeb2be5289d709e365fc6e386d734773895215cf8cbc5100fa6af9a
-  md5: eabb4b677b54874d7d6ab775fdaa3d27
-  depends:
-  - cffi >=1.0.1
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  size: 34218
+  timestamp: 1762509977830
+- conda: https://conda.anaconda.org/conda-forge/win-64/argon2-cffi-bindings-25.1.0-py314h5a2d7ad_2.conda
+  sha256: a742e7cd0d5534bfff3fd550a0c1e430411fad60a24f88930d261056ab08096f
+  md5: ffa247e46f47e157851dc547f4c513e4
+  depends:
+  - cffi >=2.0.0b1
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   - ucrt >=10.0.20348.0
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
@@ -2879,8 +2901,8 @@ packages:
   license_family: MIT
   purls:
   - pkg:pypi/argon2-cffi-bindings?source=hash-mapping
-  size: 38779
-  timestamp: 1762509796090
+  size: 38653
+  timestamp: 1762509771011
 - conda: https://conda.anaconda.org/conda-forge/noarch/arrow-1.4.0-pyhcf101f3_0.conda
   sha256: 792da8131b1b53ff667bd6fc617ea9087b570305ccb9913deb36b8e12b3b5141
   md5: 85c4f19f377424eafc4ed7911b291642
@@ -2945,49 +2967,16 @@ packages:
   - pkg:pypi/babel?source=hash-mapping
   size: 6938256
   timestamp: 1738490268466
-- conda: https://conda.anaconda.org/conda-forge/linux-64/backports.zstd-1.3.0-py313h18e8e13_0.conda
-  sha256: 9552afbec37c4d8d0e83a5c4c6b3c7f4b8785f935094ce3881e0a249045909ce
-  md5: d9e90792551a527200637e23a915dd79
-  depends:
-  - python
-  - libgcc >=14
-  - __glibc >=2.17,<3.0.a0
-  - python_abi 3.13.* *_cp313
-  - zstd >=1.5.7,<1.6.0a0
-  license: BSD-3-Clause AND MIT AND EPL-2.0
-  purls:
-  - pkg:pypi/backports-zstd?source=hash-mapping
-  size: 240943
-  timestamp: 1767044981366
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/backports.zstd-1.3.0-py313h48bb75e_0.conda
-  sha256: f3047ca3b41bb444b4b5a71a6eee182623192c77019746dd4685fd260becb249
-  md5: 54008c5cc8928e5cb5a0f9206b829451
-  depends:
-  - python
-  - python 3.13.* *_cp313
-  - __osx >=11.0
-  - zstd >=1.5.7,<1.6.0a0
-  - python_abi 3.13.* *_cp313
-  license: BSD-3-Clause AND MIT AND EPL-2.0
-  purls:
-  - pkg:pypi/backports-zstd?source=hash-mapping
-  size: 244371
-  timestamp: 1767045003420
-- conda: https://conda.anaconda.org/conda-forge/win-64/backports.zstd-1.3.0-py313h2a31948_0.conda
-  sha256: 1e76ed9bcf07ef1df9c964d73e9cda08a0380845d09c8da1678a1687dc087c34
-  md5: cdcdfe68c5bc9af9e908e35ebffc9fe1
+- conda: https://conda.anaconda.org/conda-forge/noarch/backports.zstd-1.3.0-py314h680f03e_0.conda
+  noarch: generic
+  sha256: c31ab719d256bc6f89926131e88ecd0f0c5d003fe8481852c6424f4ec6c7eb29
+  md5: a2ac7763a9ac75055b68f325d3255265
   depends:
-  - python
-  - vc >=14.3,<15
-  - vc14_runtime >=14.44.35208
-  - ucrt >=10.0.20348.0
-  - python_abi 3.13.* *_cp313
-  - zstd >=1.5.7,<1.6.0a0
+  - python >=3.14
   license: BSD-3-Clause AND MIT AND EPL-2.0
-  purls:
-  - pkg:pypi/backports-zstd?source=hash-mapping
-  size: 240406
-  timestamp: 1767045016907
+  purls: []
+  size: 7514
+  timestamp: 1767044983590
 - conda: https://conda.anaconda.org/conda-forge/noarch/beautifulsoup4-4.14.3-pyha770c72_0.conda
   sha256: bf1e71c3c0a5b024e44ff928225a0874fc3c3356ec1a0b6fe719108e6d1288f6
   md5: 5267bef8efea4127aacd1f4e1f149b6e
@@ -3128,46 +3117,46 @@ packages:
   purls: []
   size: 22714
   timestamp: 1764017952449
-- conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py313hf159716_1.conda
-  sha256: dadec2879492adede0a9af0191203f9b023f788c18efd45ecac676d424c458ae
-  md5: 6c4d3597cf43f3439a51b2b13e29a4ba
+- conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.2.0-py314h3de4e8d_1.conda
+  sha256: 3ad3500bff54a781c29f16ce1b288b36606e2189d0b0ef2f67036554f47f12b0
+  md5: 8910d2c46f7e7b519129f486e0fe927a
   depends:
   - __glibc >=2.17,<3.0.a0
   - libgcc >=14
   - libstdcxx >=14
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   constrains:
   - libbrotlicommon 1.2.0 hb03c661_1
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/brotli?source=hash-mapping
-  size: 367721
-  timestamp: 1764017371123
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-python-1.2.0-py313hde1f3bb_1.conda
-  sha256: 2e21dccccd68bedd483300f9ab87a425645f6776e6e578e10e0dd98c946e1be9
-  md5: b03732afa9f4f54634d94eb920dfb308
+  size: 367376
+  timestamp: 1764017265553
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/brotli-python-1.2.0-py314h3daef5d_1.conda
+  sha256: 5c2e471fd262fcc3c5a9d5ea4dae5917b885e0e9b02763dbd0f0d9635ed4cb99
+  md5: f9501812fe7c66b6548c7fcaa1c1f252
   depends:
   - __osx >=11.0
   - libcxx >=19
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
+  - python_abi 3.14.* *_cp314
   constrains:
   - libbrotlicommon 1.2.0 hc919400_1
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/brotli?source=hash-mapping
-  size: 359568
-  timestamp: 1764018359470
-- conda: https://conda.anaconda.org/conda-forge/win-64/brotli-python-1.2.0-py313h3ebfc14_1.conda
-  sha256: 3558006cd6e836de8dff53cbe5f0b9959f96ea6a6776b4e14f1c524916dd956c
-  md5: 916a39a0261621b8c33e9db2366dd427
-  depends:
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  size: 359854
+  timestamp: 1764018178608
+- conda: https://conda.anaconda.org/conda-forge/win-64/brotli-python-1.2.0-py314he701e3d_1.conda
+  sha256: 6854ee7675135c57c73a04849c29cbebc2fb6a3a3bfee1f308e64bf23074719b
+  md5: 1302b74b93c44791403cbeee6a0f62a3
+  depends:
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   - ucrt >=10.0.20348.0
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
@@ -3177,8 +3166,8 @@ packages:
   license_family: MIT
   purls:
   - pkg:pypi/brotli?source=hash-mapping
-  size: 335605
-  timestamp: 1764018132514
+  size: 335782
+  timestamp: 1764018443683
 - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_8.conda
   sha256: c30daba32ddebbb7ded490f0e371eae90f51e72db620554089103b4a6934b0d5
   md5: 51a19bba1b8ebfb60df25cde030b7ebc
@@ -3273,45 +3262,45 @@ packages:
   - pkg:pypi/certifi?source=compressed-mapping
   size: 150969
   timestamp: 1767500900768
-- conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py313hf46b229_1.conda
-  sha256: 2162a91819945c826c6ef5efe379e88b1df0fe9a387eeba23ddcf7ebeacd5bd6
-  md5: d0616e7935acab407d1543b28c446f6f
+- conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py314h4a8dc5f_1.conda
+  sha256: c6339858a0aaf5d939e00d345c98b99e4558f285942b27232ac098ad17ac7f8e
+  md5: cf45f4278afd6f4e6d03eda0f435d527
   depends:
   - __glibc >=2.17,<3.0.a0
   - libffi >=3.5.2,<3.6.0a0
   - libgcc >=14
   - pycparser
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/cffi?source=hash-mapping
-  size: 298357
-  timestamp: 1761202966461
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/cffi-2.0.0-py313h224173a_1.conda
-  sha256: 1fa69651f5e81c25d48ac42064db825ed1a3e53039629db69f86b952f5ce603c
-  md5: 050374657d1c7a4f2ea443c0d0cbd9a0
+  size: 300271
+  timestamp: 1761203085220
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/cffi-2.0.0-py314h44086f9_1.conda
+  sha256: 5b5ee5de01eb4e4fd2576add5ec9edfc654fbaf9293e7b7ad2f893a67780aa98
+  md5: 10dd19e4c797b8f8bdb1ec1fbb6821d7
   depends:
   - __osx >=11.0
   - libffi >=3.5.2,<3.6.0a0
   - pycparser
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/cffi?source=hash-mapping
-  size: 291376
-  timestamp: 1761203583358
-- conda: https://conda.anaconda.org/conda-forge/win-64/cffi-2.0.0-py313h5ea7bf4_1.conda
-  sha256: f867a11f42bb64a09b232e3decf10f8a8fe5194d7e3a216c6bac9f40483bd1c6
-  md5: 55b44664f66a2caf584d72196aa98af9
+  size: 292983
+  timestamp: 1761203354051
+- conda: https://conda.anaconda.org/conda-forge/win-64/cffi-2.0.0-py314h5a2d7ad_1.conda
+  sha256: 924f2f01fa7a62401145ef35ab6fc95f323b7418b2644a87fea0ea68048880ed
+  md5: c360170be1c9183654a240aadbedad94
   depends:
   - pycparser
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   - ucrt >=10.0.20348.0
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
@@ -3319,8 +3308,8 @@ packages:
   license_family: MIT
   purls:
   - pkg:pypi/cffi?source=hash-mapping
-  size: 292681
-  timestamp: 1761203203673
+  size: 294731
+  timestamp: 1761203441365
 - conda: https://conda.anaconda.org/conda-forge/noarch/cfgv-3.5.0-pyhd8ed1ab_0.conda
   sha256: aa589352e61bb221351a79e5946d56916e3c595783994884accdb3b97fe9d449
   md5: 381bd45fb7aa032691f3063aff47e3a1
@@ -3385,45 +3374,45 @@ packages:
   - pkg:pypi/comm?source=hash-mapping
   size: 14690
   timestamp: 1753453984907
-- conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py313h7037e92_3.conda
-  sha256: c545751fd48f119f2c28635514e6aa6ae784d9a1d4eb0e10be16c776e961f333
-  md5: 6186382cb34a9953bf2a18fc763dc346
+- conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.3-py314h9891dd4_3.conda
+  sha256: 54c79736927c787e535db184bb7f3bce13217cb7d755c50666cfc0da7c6c86f3
+  md5: 72d57382d0f63c20a16b1d514fcde6ff
   depends:
   - __glibc >=2.17,<3.0.a0
   - libgcc >=14
   - libstdcxx >=14
   - numpy >=1.25
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   license: BSD-3-Clause
   license_family: BSD
   purls:
   - pkg:pypi/contourpy?source=hash-mapping
-  size: 297459
-  timestamp: 1762525479137
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/contourpy-1.3.3-py313ha61f8ec_3.conda
-  sha256: a0e69aa3a039f0dab4af8c30933bcc6b718404263a002936c21c274b1f460958
-  md5: 5643cff3e9ab77999fba139465156e35
+  size: 299226
+  timestamp: 1762525516589
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/contourpy-1.3.3-py314h784bc60_3.conda
+  sha256: e5ca7f079f9bd49a9fce837dfe9014d96603600a29e5575cce19895d3639182c
+  md5: d75fae59fe0c8863de391e95959b2c65
   depends:
   - __osx >=11.0
   - libcxx >=19
   - numpy >=1.25
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
+  - python_abi 3.14.* *_cp314
   license: BSD-3-Clause
   license_family: BSD
   purls:
   - pkg:pypi/contourpy?source=hash-mapping
-  size: 259519
-  timestamp: 1762526242160
-- conda: https://conda.anaconda.org/conda-forge/win-64/contourpy-1.3.3-py313hf069bd2_3.conda
-  sha256: f5acc168a1f5eedd159bd1a89dc1dd4d901dc0502b769b4fca2bc5bdb4293fcf
-  md5: a1d5292683730418cd19b6e0cefcfc76
+  size: 262199
+  timestamp: 1762525837746
+- conda: https://conda.anaconda.org/conda-forge/win-64/contourpy-1.3.3-py314h909e829_3.conda
+  sha256: f014eb687eb8dd25cec124594f4e48cf85803ff1db85a2a1f95719f9ec6434d2
+  md5: 3647d90eea49efc6076729ef0ae81075
   depends:
   - numpy >=1.25
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   - ucrt >=10.0.20348.0
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
@@ -3431,44 +3420,44 @@ packages:
   license_family: BSD
   purls:
   - pkg:pypi/contourpy?source=hash-mapping
-  size: 225553
-  timestamp: 1762525633181
-- conda: https://conda.anaconda.org/conda-forge/linux-64/coverage-7.13.1-py313h3dea7bd_0.conda
-  sha256: 4275280f4fcef6cd0a0e5cd236120d7454a11390dd4c271378bf90bc563f6780
-  md5: 82315acb438e857f809f556e2dcdb822
+  size: 227536
+  timestamp: 1762525688384
+- conda: https://conda.anaconda.org/conda-forge/linux-64/coverage-7.13.1-py314h67df5f8_0.conda
+  sha256: 63b91c7308704819bc35747ed88097c391a75502921f7f3c9422d42e1ed07909
+  md5: a4525263f2fa741bffa4af1e40aec245
   depends:
   - __glibc >=2.17,<3.0.a0
   - libgcc >=14
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   - tomli
   license: Apache-2.0
   license_family: APACHE
   purls:
   - pkg:pypi/coverage?source=hash-mapping
-  size: 393234
-  timestamp: 1766951417242
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/coverage-7.13.1-py313h65a2061_0.conda
-  sha256: 46e4af43bd60580fda7955cc6c21b3a40465ef25a98c2a256419dc74caae56b0
-  md5: 3283d95f985c7f293cb13bb7e33500a5
+  size: 410205
+  timestamp: 1766951484026
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/coverage-7.13.1-py314h6e9b3f0_0.conda
+  sha256: 06311a6cb704c7c2db910ef4bda5f4d4f2c3a9e8bdffe4cc5c4481fc253a47d6
+  md5: 39869c1b0010c430849a7c2585c65f47
   depends:
   - __osx >=11.0
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
+  - python_abi 3.14.* *_cp314
   - tomli
   license: Apache-2.0
   license_family: APACHE
   purls:
   - pkg:pypi/coverage?source=hash-mapping
-  size: 393649
-  timestamp: 1766951606379
-- conda: https://conda.anaconda.org/conda-forge/win-64/coverage-7.13.1-py313hd650c13_0.conda
-  sha256: d41807f993eb1c097594f6481dc4a3ea1080ed57cfd1f0721216a3d7f7f3f949
-  md5: 6799738f6603dfddd97389ee3e65e891
-  depends:
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  size: 409230
+  timestamp: 1766951563419
+- conda: https://conda.anaconda.org/conda-forge/win-64/coverage-7.13.1-py314h2359020_0.conda
+  sha256: fd24db3e7d3407ae7a15cd636722c84ca26e4c274f639084cdd18afa6612fe5b
+  md5: c5cb6c314f63b0bd76c67775a515364d
+  depends:
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   - tomli
   - ucrt >=10.0.20348.0
   - vc >=14.3,<15
@@ -3477,19 +3466,19 @@ packages:
   license_family: APACHE
   purls:
   - pkg:pypi/coverage?source=hash-mapping
-  size: 418313
-  timestamp: 1766951491957
-- conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.11-py313hd8ed1ab_100.conda
+  size: 434074
+  timestamp: 1766951384017
+- conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.14.2-py314hd8ed1ab_100.conda
   noarch: generic
-  sha256: 63f677762304e6f8dc55e11dff6aafe71129cbbd0a77d176b99ba1f6a5053b77
-  md5: 5bf347916a543bcb290c780fa449bf73
+  sha256: 9e345f306446500956ffb1414b773f5476f497d7a2b5335a59edd2c335209dbb
+  md5: 30f999d06f347b0116f0434624b6e559
   depends:
-  - python >=3.13,<3.14.0a0
-  - python_abi * *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi * *_cp314
   license: Python-2.0
   purls: []
-  size: 48369
-  timestamp: 1765019689213
+  size: 49298
+  timestamp: 1765020324943
 - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cccl_linux-64-12.9.27-ha770c72_0.conda
   sha256: 2ee3b9564ca326226e5cda41d11b251482df8e7c757e333d28ec75213c75d126
   md5: 87ff6381e33b76e5b9b179a2cdd005ec
@@ -3732,51 +3721,51 @@ packages:
   - flatten-dict
   - networkx
   requires_python: '>=3.10'
-- conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.18-py313h5d5ffb9_0.conda
-  sha256: 29d10b4520846d3cbc511545552c11b726199013354e7517a53679272629c20d
-  md5: 80fd7ff9877570d12cabb5c5037dac89
+- conda: https://conda.anaconda.org/conda-forge/linux-64/debugpy-1.8.18-py314h42812f9_0.conda
+  sha256: 2803e9285da433a5d704a63ac9c64c87b5df9aaa1e2d48cc333e65d5a945912e
+  md5: 69635aa34b45d84c2599ff8b48094978
   depends:
   - python
+  - libgcc >=14
   - __glibc >=2.17,<3.0.a0
   - libstdcxx >=14
-  - libgcc >=14
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/debugpy?source=hash-mapping
-  size: 2870642
-  timestamp: 1765704059389
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/debugpy-1.8.19-py313hc37fe24_0.conda
-  sha256: 1eb7c9f5a994e273d714e945253fff40413fd63de9f6d5e01989d6d96199dad0
-  md5: 95287e5abbe8a588d2a8d234f3d591a7
+  size: 2888322
+  timestamp: 1765704065377
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/debugpy-1.8.19-py314hf820bb6_0.conda
+  sha256: 5c263dafa3660660087443ad37e32e0597067cf098b351230a76adf83e462e12
+  md5: 45961f5d077fca30eeff1a1973aca63d
   depends:
   - python
-  - python 3.13.* *_cp313
-  - libcxx >=19
   - __osx >=11.0
-  - python_abi 3.13.* *_cp313
+  - python 3.14.* *_cp314
+  - libcxx >=19
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
-  - pkg:pypi/debugpy?source=compressed-mapping
-  size: 2759061
-  timestamp: 1765840814720
-- conda: https://conda.anaconda.org/conda-forge/win-64/debugpy-1.8.19-py313h927ade5_0.conda
-  sha256: d6d62b00c9a81cf9f183b9f3929455f11e1906e37891a28b953237245df6a5f3
-  md5: a7e77991e54b031328253da027e2f3e1
+  - pkg:pypi/debugpy?source=hash-mapping
+  size: 2776268
+  timestamp: 1765840821598
+- conda: https://conda.anaconda.org/conda-forge/win-64/debugpy-1.8.19-py314hb98de8c_0.conda
+  sha256: 0ad7f50f664ede3aafcd23458ce4f669f63e32f7efb74c0938260bdb829679df
+  md5: 3361deac30d356844406fbe6def54d5b
   depends:
   - python
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
   - ucrt >=10.0.20348.0
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/debugpy?source=hash-mapping
-  size: 4002629
-  timestamp: 1765840845981
+  size: 4021751
+  timestamp: 1765840833937
 - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
   sha256: c17c6b9937c08ad63cb20a26f403a3234088e57d4455600974a0ce865cb14017
   md5: 9ce473d1d1be1cc3810856a48b3fab32
@@ -3896,55 +3885,22 @@ packages:
   - pathlib2>=2.3,<3.0 ; python_full_version < '3.4'
   - six>=1.12,<2.0
   requires_python: '>=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*'
-- conda: https://conda.anaconda.org/conda-forge/linux-64/fonttools-4.61.1-py313h3dea7bd_0.conda
-  sha256: 97f225199e6e5dfb93f551087c0951fee92db2d29a9dcb6a0346d66bff06fea4
-  md5: c0f36dfbb130da4f6ce2df31f6b25ea8
-  depends:
-  - __glibc >=2.17,<3.0.a0
-  - brotli
-  - libgcc >=14
-  - munkres
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
-  license: MIT
-  license_family: MIT
-  purls:
-  - pkg:pypi/fonttools?source=hash-mapping
-  size: 2988776
-  timestamp: 1765633043435
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/fonttools-4.61.1-py313h7d74516_0.conda
-  sha256: 52d4aacd7c154adff1f0e86609bf1b0e63b7049c947c4df1e78eedb9f2913091
-  md5: 894eb0c3e9a17643906a6da3209bf045
-  depends:
-  - __osx >=11.0
-  - brotli
-  - munkres
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
-  license: MIT
-  license_family: MIT
-  purls:
-  - pkg:pypi/fonttools?source=hash-mapping
-  size: 2897709
-  timestamp: 1765632961717
-- conda: https://conda.anaconda.org/conda-forge/win-64/fonttools-4.61.1-py313hd650c13_0.conda
-  sha256: da82b8e843103bf4aaab470e4b8025286357dc8c34cd47817350dcb14ad307fb
-  md5: c6fbf3a96192c26a75ed5755bd904fea
+- conda: https://conda.anaconda.org/conda-forge/noarch/fonttools-4.61.1-pyh7db6752_0.conda
+  sha256: bb74f1732065eb95c3ea4ae7f7ab29d6ddaafe6da32f009106bf9a335147cb77
+  md5: d5da976e963e70364b9e3ff270842b9f
   depends:
   - brotli
   - munkres
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
-  - ucrt >=10.0.20348.0
-  - vc >=14.3,<15
-  - vc14_runtime >=14.44.35208
+  - python >=3.10
+  - unicodedata2 >=15.1.0
+  track_features:
+  - fonttools_no_compile
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/fonttools?source=hash-mapping
-  size: 2523451
-  timestamp: 1765632913315
+  size: 834764
+  timestamp: 1765632669874
 - conda: https://conda.anaconda.org/conda-forge/noarch/fqdn-1.5.1-pyhd8ed1ab_1.conda
   sha256: 2509992ec2fd38ab27c7cdb42cf6cadc566a1cc0d1021a2673475d9fa87c6276
   md5: d3549fd50d450b6d9e7dddff25dd2110
@@ -3987,6 +3943,11 @@ packages:
   purls: []
   size: 184553
   timestamp: 1757946164012
+- pypi: https://files.pythonhosted.org/packages/38/74/f94141b38a51a553efef7f510fc213894161ae49b88bffd037f8d2a7cb2f/frozendict-2.4.7-py3-none-any.whl
+  name: frozendict
+  version: 2.4.7
+  sha256: 972af65924ea25cf5b4d9326d549e69a9a4918d8a76a9d3a7cd174d98b237550
+  requires_python: '>=3.6'
 - conda: https://conda.anaconda.org/conda-forge/linux-64/gcc_impl_linux-64-14.3.0-he8b2097_16.conda
   sha256: 4acf50b7d5673250d585a256a40aabdd922e0947ca12cdbad0cef960ee1a9509
   md5: d274bf1343507683e6eb2954d1871569
@@ -4043,10 +4004,10 @@ packages:
   purls: []
   size: 2009354
   timestamp: 1765814947748
-- pypi: https://files.pythonhosted.org/packages/7e/71/ba21c3fb8c5dce83b8c01f458a42e99ffdb1963aeec08fff5a18588d8fd7/greenlet-3.3.0-cp313-cp313-win_amd64.whl
+- pypi: https://files.pythonhosted.org/packages/7c/9a/9030e6f9aa8fd7808e9c31ba4c38f87c4f8ec324ee67431d181fe396d705/greenlet-3.3.0-cp314-cp314-win_amd64.whl
   name: greenlet
   version: 3.3.0
-  sha256: 9ee1942ea19550094033c35d25d20726e4f1c40d59545815e1128ac58d416d38
+  sha256: 73f51dd0e0bdb596fb0417e475fa3c5e32d4c83638296e560086b8d7da7c4170
   requires_dist:
   - sphinx ; extra == 'docs'
   - furo ; extra == 'docs'
@@ -4054,10 +4015,10 @@ packages:
   - psutil ; extra == 'test'
   - setuptools ; extra == 'test'
   requires_python: '>=3.10'
-- pypi: https://files.pythonhosted.org/packages/fd/8e/424b8c6e78bd9837d14ff7df01a9829fc883ba2ab4ea787d4f848435f23f/greenlet-3.3.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+- pypi: https://files.pythonhosted.org/packages/b8/14/bab308fc2c1b5228c3224ec2bf928ce2e4d21d8046c161e44a2012b5203e/greenlet-3.3.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
   name: greenlet
   version: 3.3.0
-  sha256: 087ea5e004437321508a8d6f20efc4cfec5e3c30118e1417ea96ed1d93950527
+  sha256: 5773edda4dc00e173820722711d043799d3adb4f01731f40619e07ea2750b955
   requires_dist:
   - sphinx ; extra == 'docs'
   - furo ; extra == 'docs'
@@ -4431,10 +4392,10 @@ packages:
   name: jax-cuda12-pjrt
   version: 0.8.2
   sha256: e3bab41ca7c48e4163db9e7efd271b3aa85f0fe45f5ed0708d6bbed93a59f977
-- pypi: https://files.pythonhosted.org/packages/1c/38/4ba2486f95fcf2120723932feacdded438e785258148b18a703cd1177e41/jax_cuda12_plugin-0.8.2-cp313-cp313-manylinux_2_27_x86_64.whl
+- pypi: https://files.pythonhosted.org/packages/27/58/a5a27d4677d6890570f7e58cecd51891469cb620e6f64c8faed4935d93d0/jax_cuda12_plugin-0.8.2-cp314-cp314-manylinux_2_27_x86_64.whl
   name: jax-cuda12-plugin
   version: 0.8.2
-  sha256: 82c6798be66bf8c773386918e4c8e5cd8119753f3bfb3ca4bbc46818283750c6
+  sha256: a5898bac1d8ab6020b54546440256409f2c66bcbbb3a1099ca473c84843addad
   requires_dist:
   - jax-cuda12-pjrt==0.8.2
   - nvidia-cublas-cu12>=12.1.3.1 ; sys_platform == 'linux' and extra == 'with-cuda'
@@ -4450,28 +4411,28 @@ packages:
   - nvidia-cuda-nvrtc-cu12>=12.1.55 ; sys_platform == 'linux' and extra == 'with-cuda'
   - nvidia-nvshmem-cu12>=3.2.5 ; sys_platform == 'linux' and extra == 'with-cuda'
   requires_python: '>=3.11'
-- pypi: https://files.pythonhosted.org/packages/6b/e0/91e5762a7ddb6351b07c742ca407cd28e26043d6945d6228b6c1b0881a45/jaxlib-0.8.2-cp313-cp313-manylinux_2_27_x86_64.whl
+- pypi: https://files.pythonhosted.org/packages/5e/27/2e6032727e41ce74914277478021140947af59127d68aa9e6f3776b428fd/jaxlib-0.8.2-cp314-cp314-manylinux_2_27_x86_64.whl
   name: jaxlib
   version: 0.8.2
-  sha256: 1bfbcf6c3de221784fa4cdb6765a09d71cb4298b15626b3d0409b3dfcd8a8667
+  sha256: e6a97dfb0232eed9a2bb6e3828e4f682dbac1a7fea840bfda574cae2dbf5faf9
   requires_dist:
   - scipy>=1.13
   - numpy>=2.0
   - ml-dtypes>=0.5.0
   requires_python: '>=3.11'
-- pypi: https://files.pythonhosted.org/packages/85/68/25b38673b07a808616ce7b6efb3eed491f983f3373a09cbbd03f67178563/jaxlib-0.8.2-cp313-cp313-win_amd64.whl
+- pypi: https://files.pythonhosted.org/packages/b3/8c/af5a00b07a446414edf6b84a7397eab02cf01ba44b6ae1fce7798ce4c127/jaxlib-0.8.2-cp314-cp314-win_amd64.whl
   name: jaxlib
   version: 0.8.2
-  sha256: f205e91c3a152a2a76c0bc59a6a2de03e87ec261b91e8812922777185e7b08f5
+  sha256: 05b958f497e49824c432e734bb059723b7dfe69e2ad696a9f9c8ad82fff7c3f8
   requires_dist:
   - scipy>=1.13
   - numpy>=2.0
   - ml-dtypes>=0.5.0
   requires_python: '>=3.11'
-- pypi: https://files.pythonhosted.org/packages/c5/22/c0ec75e43a13b2457d78d509f49b49a57fa302ffced4f4a2778e428cb0a6/jaxlib-0.8.2-cp313-cp313-macosx_11_0_arm64.whl
+- pypi: https://files.pythonhosted.org/packages/d8/9d/dca93d916bf8664d7a2bb73ea3d219028dabbe382c31774348963287356a/jaxlib-0.8.2-cp314-cp314-macosx_11_0_arm64.whl
   name: jaxlib
   version: 0.8.2
-  sha256: 4d006db96be020c8165212a1216372f8acac4ff4f8fb067743d694ef2b301ace
+  sha256: beffb004e7eeb5c9afb24439e2b2cf45a4ee3e3e8adf45e355edf2af62acf8b8
   requires_dist:
   - scipy>=1.13
   - numpy>=2.0
@@ -4540,6 +4501,7 @@ packages:
   - rpds-py >=0.25.0
   - python
   license: MIT
+  license_family: MIT
   purls:
   - pkg:pypi/jsonschema?source=compressed-mapping
   size: 82356
@@ -4572,6 +4534,7 @@ packages:
   - uri-template
   - webcolors >=24.6.0
   license: MIT
+  license_family: MIT
   purls: []
   size: 4740
   timestamp: 1767839954258
@@ -4782,39 +4745,39 @@ packages:
   purls: []
   size: 134088
   timestamp: 1754905959823
-- conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py313hc8edb43_2.conda
-  sha256: 60d7d525db89401f88f5c91bdbb79d3afbf005e7d7c1326318659fa097607e51
-  md5: 3e0e65595330e26515e31b7fc6d933c7
+- conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py314h97ea11e_2.conda
+  sha256: a707d08c095d02148201f2da9fba465054fb750e33117e215892a4fefcc1b54a
+  md5: 57f1ce4f7ba6bcd460be8f83c8f04c69
   depends:
   - python
-  - __glibc >=2.17,<3.0.a0
   - libstdcxx >=14
   - libgcc >=14
-  - python_abi 3.13.* *_cp313
+  - __glibc >=2.17,<3.0.a0
+  - python_abi 3.14.* *_cp314
   license: BSD-3-Clause
   license_family: BSD
   purls:
   - pkg:pypi/kiwisolver?source=hash-mapping
-  size: 77616
-  timestamp: 1762488778882
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/kiwisolver-1.4.9-py313h7add70c_2.conda
-  sha256: adc6b89070b6858b81fbe24dd034a73295e8fa9ccb68ed871bf04f1ed498f51c
-  md5: 9583687276aaa393e723f3b7970be69f
+  size: 78071
+  timestamp: 1762488742381
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/kiwisolver-1.4.9-py314h42813c9_2.conda
+  sha256: c4d7e6653d343e768110ec77ac1c6c89f313f77a19a1f2cd60b7c7b8b0758bdf
+  md5: 9aa431bf603c231e8c77a1b0842a85ed
   depends:
   - python
-  - libcxx >=19
-  - python 3.13.* *_cp313
+  - python 3.14.* *_cp314
   - __osx >=11.0
-  - python_abi 3.13.* *_cp313
+  - libcxx >=19
+  - python_abi 3.14.* *_cp314
   license: BSD-3-Clause
   license_family: BSD
   purls:
   - pkg:pypi/kiwisolver?source=hash-mapping
-  size: 68438
-  timestamp: 1762488945877
-- conda: https://conda.anaconda.org/conda-forge/win-64/kiwisolver-1.4.9-py313h1a38498_2.conda
-  sha256: 40eafae7e9cdbe97eeb56ab0882816d3f68a2af4080a822f7349f986de2adeb6
-  md5: f77249adfa3f0091e016610346affd09
+  size: 68534
+  timestamp: 1762489024029
+- conda: https://conda.anaconda.org/conda-forge/win-64/kiwisolver-1.4.9-py314hf309875_2.conda
+  sha256: ded907ab1ce24abcff20bc239e770ae7ef4cff6fdcfb8cc24ca59ebe736a1d3f
+  md5: e9d93271b021332f5492ff5478601614
   depends:
   - python
   - vc >=14.3,<15
@@ -4823,13 +4786,13 @@ packages:
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
   - ucrt >=10.0.20348.0
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   license: BSD-3-Clause
   license_family: BSD
   purls:
   - pkg:pypi/kiwisolver?source=hash-mapping
-  size: 73825
-  timestamp: 1762488792613
+  size: 73670
+  timestamp: 1762488752873
 - conda: https://conda.anaconda.org/conda-forge/linux-64/krb5-1.21.3-h659f571_0.conda
   sha256: 99df692f7a8a5c27cd14b5fb1374ee55e756631b9c3d659ed3ee60830249b238
   md5: 3f43953b7d3fb3aaa1d0d0723d91e368
@@ -5561,9 +5524,9 @@ packages:
   purls: []
   size: 663567
   timestamp: 1765260367147
-- conda: https://conda.anaconda.org/conda-forge/win-64/libhwloc-2.12.1-default_h4379cf1_1003.conda
-  sha256: 2d534c09f92966b885acb3f4a838f7055cea043165a03079a539b06c54e20a49
-  md5: d1699ce4fe195a9f61264a1c29b87035
+- conda: https://conda.anaconda.org/conda-forge/win-64/libhwloc-2.12.2-default_h4379cf1_1000.conda
+  sha256: 8cdf11333a81085468d9aa536ebb155abd74adc293576f6013fc0c85a7a90da3
+  md5: 3b576f6860f838f950c570f4433b086e
   depends:
   - libwinpthread >=12.0.0.r4.gg4f2fc60ca
   - libxml2
@@ -5574,8 +5537,8 @@ packages:
   license: BSD-3-Clause
   license_family: BSD
   purls: []
-  size: 2412642
-  timestamp: 1765090345611
+  size: 2411241
+  timestamp: 1765104337762
 - conda: https://conda.anaconda.org/conda-forge/linux-64/libiconv-1.18-h3b78370_2.conda
   sha256: c467851a7312765447155e071752d7bf9bf44d610a5687e32706f480aad2833f
   md5: 915f5995e94f60e9a4826e0b0920ee88
@@ -6370,58 +6333,24 @@ packages:
   - pkg:pypi/markdown-it-py?source=hash-mapping
   size: 64736
   timestamp: 1754951288511
-- conda: https://conda.anaconda.org/conda-forge/linux-64/markupsafe-3.0.3-py313h3dea7bd_0.conda
-  sha256: a530a411bdaaf0b1e4de8869dfaca46cb07407bc7dc0702a9e231b0e5ce7ca85
-  md5: c14389156310b8ed3520d84f854be1ee
-  depends:
-  - __glibc >=2.17,<3.0.a0
-  - libgcc >=14
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
-  constrains:
-  - jinja2 >=3.0.0
-  license: BSD-3-Clause
-  license_family: BSD
-  purls:
-  - pkg:pypi/markupsafe?source=hash-mapping
-  size: 25909
-  timestamp: 1759055357045
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/markupsafe-3.0.3-py313h7d74516_0.conda
-  sha256: e06902a1bf370fdd4ada0a8c81c504868fdb7e9971b72c6bd395aa4e5a497bd2
-  md5: 3df5979cc0b761dda0053ffdb0bca3ea
+- conda: https://conda.anaconda.org/conda-forge/noarch/markupsafe-3.0.3-pyh7db6752_0.conda
+  sha256: e0cbfea51a19b3055ca19428bd9233a25adca956c208abb9d00b21e7259c7e03
+  md5: fab1be106a50e20f10fe5228fd1d1651
   depends:
-  - __osx >=11.0
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
-  constrains:
-  - jinja2 >=3.0.0
-  license: BSD-3-Clause
-  license_family: BSD
-  purls:
-  - pkg:pypi/markupsafe?source=hash-mapping
-  size: 25778
-  timestamp: 1759055530601
-- conda: https://conda.anaconda.org/conda-forge/win-64/markupsafe-3.0.3-py313hd650c13_0.conda
-  sha256: 988d14095c1392e055fd75e24544da2db01ade73b0c2f99ddc8e2b8678ead4cc
-  md5: 47eaaa4405741beb171ea6edc6eaf874
-  depends:
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
-  - ucrt >=10.0.20348.0
-  - vc >=14.3,<15
-  - vc14_runtime >=14.44.35208
+  - python >=3.10
   constrains:
   - jinja2 >=3.0.0
+  track_features:
+  - markupsafe_no_compile
   license: BSD-3-Clause
   license_family: BSD
   purls:
   - pkg:pypi/markupsafe?source=hash-mapping
-  size: 28959
-  timestamp: 1759055685616
-- conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.8-py313h683a580_0.conda
-  sha256: b1117aa2c1d11ca70d1704054cdc8801cbcf2dfb846c565531edd417ddd82559
-  md5: ffe67570e1a9192d2f4c189b27f75f89
+  size: 15499
+  timestamp: 1759055275624
+- conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.8-py314h1194b4b_0.conda
+  sha256: ee773261fbd6c76fc8174b0e4e1ce272b0bbaa56610f130e9d3d1f575106f04f
+  md5: b8683e6068099b69c10dbfcf7204203f
   depends:
   - __glibc >=2.17,<3.0.a0
   - contourpy >=1.0.1
@@ -6438,20 +6367,20 @@ packages:
   - packaging >=20.0
   - pillow >=8
   - pyparsing >=2.3.1
-  - python >=3.13,<3.14.0a0
+  - python >=3.14,<3.15.0a0
   - python-dateutil >=2.7
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   - qhull >=2020.2,<2020.3.0a0
   - tk >=8.6.13,<8.7.0a0
   license: PSF-2.0
   license_family: PSF
   purls:
-  - pkg:pypi/matplotlib?source=compressed-mapping
-  size: 8405862
-  timestamp: 1763055358671
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/matplotlib-base-3.10.8-py313h58042b9_0.conda
-  sha256: 24767ca32ea9db74a4a5965d2df8c69c83c82583e8ba32b683123d406092e205
-  md5: 745c18472bc6d3dc9146c3dec18bb740
+  - pkg:pypi/matplotlib?source=hash-mapping
+  size: 8473358
+  timestamp: 1763055439346
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/matplotlib-base-3.10.8-py314hd63e3f0_0.conda
+  sha256: 198dcc0ed83e78bc7bf48e6ef8d4ecd220e9cf1f07db98508251b2bc0be067f9
+  md5: c84152e510d41378b8758826655b6ed7
   depends:
   - __osx >=11.0
   - contourpy >=1.0.1
@@ -6467,20 +6396,20 @@ packages:
   - packaging >=20.0
   - pillow >=8
   - pyparsing >=2.3.1
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
   - python-dateutil >=2.7
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   - qhull >=2020.2,<2020.3.0a0
   license: PSF-2.0
   license_family: PSF
   purls:
   - pkg:pypi/matplotlib?source=hash-mapping
-  size: 8197793
-  timestamp: 1763056104477
-- conda: https://conda.anaconda.org/conda-forge/win-64/matplotlib-base-3.10.8-py313he1ded55_0.conda
-  sha256: f63c4a5ded62cfb216c9d107a3c4527940036eef19cf481418080a0bd9bc11d8
-  md5: 05f96c429201a64ea752decf4b910a7c
+  size: 8286510
+  timestamp: 1763055937766
+- conda: https://conda.anaconda.org/conda-forge/win-64/matplotlib-base-3.10.8-py314hfa45d96_0.conda
+  sha256: 82a50284275e8a1818cd3323846f3032dc89bd23a3f80dcf44e34a62b016256b
+  md5: 9d491a60700e0e90e92607fcc4e2566c
   depends:
   - contourpy >=1.0.1
   - cycler >=0.10
@@ -6494,9 +6423,9 @@ packages:
   - packaging >=20.0
   - pillow >=8
   - pyparsing >=2.3.1
-  - python >=3.13,<3.14.0a0
+  - python >=3.14,<3.15.0a0
   - python-dateutil >=2.7
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   - qhull >=2020.2,<2020.3.0a0
   - ucrt >=10.0.20348.0
   - vc >=14.3,<15
@@ -6505,8 +6434,8 @@ packages:
   license_family: PSF
   purls:
   - pkg:pypi/matplotlib?source=hash-mapping
-  size: 8007333
-  timestamp: 1763055517579
+  size: 8185296
+  timestamp: 1763055983613
 - conda: https://conda.anaconda.org/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda
   sha256: 9d690334de0cd1d22c51bc28420663f4277cfa60d34fa5cad1ce284a13f1d603
   md5: 00e120ce3e40bad7bfc78861ce3c4a25
@@ -6542,44 +6471,44 @@ packages:
   - pkg:pypi/mdurl?source=hash-mapping
   size: 14465
   timestamp: 1733255681319
-- conda: https://conda.anaconda.org/conda-forge/linux-64/memray-1.19.1-py313h422961c_3.conda
-  sha256: 1a752d45a2c5da1289afac51ea5b89bde0a80f290708505b487f38d47b4e3267
-  md5: 6f9810aa09fbdab0c6b941d48a3b72bb
+- conda: https://conda.anaconda.org/conda-forge/linux-64/memray-1.19.1-py314hef15ded_3.conda
+  sha256: 43801200d3b8dcaa1f9ab47f527c9fe94028780b2760173a240e132e25be2194
+  md5: cc1bee6de727d07ce2dad51a5e8364b9
   depends:
   - python
   - rich >=11.2.0
   - jinja2
   - textual >=0.34.0
-  - libgcc >=14
   - libstdcxx >=14
+  - libgcc >=14
   - __glibc >=2.17,<3.0.a0
-  - python_abi 3.13.* *_cp313
   - lz4-c >=1.10.0,<1.11.0a0
   - elfutils >=0.194,<0.195.0a0
   - libunwind >=1.8.3,<1.9.0a0
+  - python_abi 3.14.* *_cp314
   license: Apache-2.0 AND BSD-3-Clause
   purls:
   - pkg:pypi/memray?source=hash-mapping
-  size: 1816303
-  timestamp: 1765821582847
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/memray-1.19.1-py313h78c9487_3.conda
-  sha256: eece155fd7c5f59226e24015ae08e5d8eb9a3e453f6c97bf16d04348e7f94c97
-  md5: f1dcaa6d7f501b2b8bd6294610c3982a
+  size: 1824670
+  timestamp: 1765821568349
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/memray-1.19.1-py314habef2a7_3.conda
+  sha256: 912a462c888a867a22e6ebf0607ad3a42d078260cafc7d4f25654989a17f41ac
+  md5: f8c08fd9eb42146b0489d42069fb270e
   depends:
   - python
   - rich >=11.2.0
   - jinja2
   - textual >=0.34.0
+  - python 3.14.* *_cp314
   - __osx >=11.0
-  - python 3.13.* *_cp313
   - libcxx >=19
   - lz4-c >=1.10.0,<1.11.0a0
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   license: Apache-2.0 AND BSD-3-Clause
   purls:
   - pkg:pypi/memray?source=hash-mapping
-  size: 1712578
-  timestamp: 1765821632543
+  size: 1721669
+  timestamp: 1765821674618
 - conda: https://conda.anaconda.org/conda-forge/noarch/mistune-3.2.0-pyhcf101f3_0.conda
   sha256: d3fb4beb5e0a52b6cc33852c558e077e1bfe44df1159eb98332d69a264b14bae
   md5: b11e360fc4de2b0035fc8aaa74f17fd6
@@ -6590,7 +6519,7 @@ packages:
   license: BSD-3-Clause
   license_family: BSD
   purls:
-  - pkg:pypi/mistune?source=compressed-mapping
+  - pkg:pypi/mistune?source=hash-mapping
   size: 74250
   timestamp: 1766504456031
 - conda: https://conda.anaconda.org/conda-forge/win-64/mkl-2025.3.0-hac47afa_455.conda
@@ -6607,10 +6536,10 @@ packages:
   purls: []
   size: 100224829
   timestamp: 1767634557029
-- pypi: https://files.pythonhosted.org/packages/d9/a1/4008f14bbc616cfb1ac5b39ea485f9c63031c4634ab3f4cf72e7541f816a/ml_dtypes-0.5.4-cp313-cp313-macosx_10_13_universal2.whl
+- pypi: https://files.pythonhosted.org/packages/72/4e/1339dc6e2557a344f5ba5590872e80346f76f6cb2ac3dd16e4666e88818c/ml_dtypes-0.5.4-cp314-cp314-macosx_10_13_universal2.whl
   name: ml-dtypes
   version: 0.5.4
-  sha256: 8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48
+  sha256: 2b857d3af6ac0d39db1de7c706e69c7f9791627209c3d6dedbfca8c7e5faec22
   requires_dist:
   - numpy>=1.21
   - numpy>=1.21.2 ; python_full_version >= '3.10'
@@ -6623,10 +6552,10 @@ packages:
   - pylint>=2.6.0 ; extra == 'dev'
   - pyink ; extra == 'dev'
   requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/e1/8b/200088c6859d8221454825959df35b5244fa9bdf263fd0249ac5fb75e281/ml_dtypes-0.5.4-cp313-cp313-win_amd64.whl
+- pypi: https://files.pythonhosted.org/packages/c6/bb/82c7dcf38070b46172a517e2334e665c5bf374a262f99a283ea454bece7c/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
   name: ml-dtypes
   version: 0.5.4
-  sha256: f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328
+  sha256: 14a4fd3228af936461db66faccef6e4f41c1d82fcc30e9f8d58a08916b1d811f
   requires_dist:
   - numpy>=1.21
   - numpy>=1.21.2 ; python_full_version >= '3.10'
@@ -6639,10 +6568,10 @@ packages:
   - pylint>=2.6.0 ; extra == 'dev'
   - pyink ; extra == 'dev'
   requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
+- pypi: https://files.pythonhosted.org/packages/e9/93/2bfed22d2498c468f6bcd0d9f56b033eaa19f33320389314c19ef6766413/ml_dtypes-0.5.4-cp314-cp314-win_amd64.whl
   name: ml-dtypes
   version: 0.5.4
-  sha256: 533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d
+  sha256: 8c6a2dcebd6f3903e05d51960a8058d6e131fe69f952a5397e5dbabc841b6d56
   requires_dist:
   - numpy>=1.21
   - numpy>=1.21.2 ; python_full_version >= '3.10'
@@ -6820,49 +6749,49 @@ packages:
   - pkg:pypi/notebook-shim?source=hash-mapping
   size: 16817
   timestamp: 1733408419340
-- conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py313hf6604e3_1.conda
-  sha256: 2f8aff2a17e4d43012e9863ef4392e6d5de3ae9da0c3e322831f8c5c3d86df71
-  md5: dce261869f78ba9b81b9091b084d328d
+- conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.3.5-py314h2b28147_1.conda
+  sha256: 81425306df4f0ddba159e80c8d91323a34df335079ca93a194201e57b337231c
+  md5: ab17cb5f388fa17c08937cb9cc24e7b6
   depends:
   - python
+  - __glibc >=2.17,<3.0.a0
   - libgcc >=14
   - libstdcxx >=14
-  - __glibc >=2.17,<3.0.a0
-  - python_abi 3.13.* *_cp313
-  - libcblas >=3.9.0,<4.0a0
   - liblapack >=3.9.0,<4.0a0
+  - libcblas >=3.9.0,<4.0a0
   - libblas >=3.9.0,<4.0a0
+  - python_abi 3.14.* *_cp314
   constrains:
   - numpy-base <0a0
   license: BSD-3-Clause
   license_family: BSD
   purls:
   - pkg:pypi/numpy?source=hash-mapping
-  size: 8919234
-  timestamp: 1766383469748
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/numpy-2.3.5-py313h16eae64_1.conda
-  sha256: d759e7fee853d8e18709a15b8fc8a6db90c96986cb9d316c4d5ccdf5a1d3f61f
-  md5: c72599556b49dc853839f4439c1eea32
+  size: 8983076
+  timestamp: 1766383421113
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/numpy-2.3.5-py314hae46ccb_1.conda
+  sha256: bc9dfe41ba4898365a82c485416fd4a572f86d94e606d89379766de70d34fc79
+  md5: d421394cf6758a6f27ead1530cfdfa6a
   depends:
   - python
   - libcxx >=19
   - __osx >=11.0
-  - python 3.13.* *_cp313
-  - liblapack >=3.9.0,<4.0a0
-  - python_abi 3.13.* *_cp313
+  - python 3.14.* *_cp314
   - libcblas >=3.9.0,<4.0a0
+  - python_abi 3.14.* *_cp314
   - libblas >=3.9.0,<4.0a0
+  - liblapack >=3.9.0,<4.0a0
   constrains:
   - numpy-base <0a0
   license: BSD-3-Clause
   license_family: BSD
   purls:
   - pkg:pypi/numpy?source=hash-mapping
-  size: 6792353
-  timestamp: 1766383288679
-- conda: https://conda.anaconda.org/conda-forge/win-64/numpy-2.3.5-py313hce7ae62_1.conda
-  sha256: c02d9587864174146bf0024051c76d368b2de18c94421e2f4e611fbb18576dd1
-  md5: 78749843445581c6dcc0cb80d146982d
+  size: 6861028
+  timestamp: 1766383292611
+- conda: https://conda.anaconda.org/conda-forge/win-64/numpy-2.3.5-py314h06c3c77_1.conda
+  sha256: 111a7af69521dce54ce6b4d89ef767ade9f3769576353a526174792de8702b5d
+  md5: 71dabea9914329c08b4864955c3793fc
   depends:
   - python
   - vc >=14.3,<15
@@ -6870,16 +6799,16 @@ packages:
   - ucrt >=10.0.20348.0
   - liblapack >=3.9.0,<4.0a0
   - libblas >=3.9.0,<4.0a0
+  - python_abi 3.14.* *_cp314
   - libcblas >=3.9.0,<4.0a0
-  - python_abi 3.13.* *_cp313
   constrains:
   - numpy-base <0a0
   license: BSD-3-Clause
   license_family: BSD
   purls:
-  - pkg:pypi/numpy?source=hash-mapping
-  size: 7524105
-  timestamp: 1766383318405
+  - pkg:pypi/numpy?source=compressed-mapping
+  size: 7584934
+  timestamp: 1766383321713
 - pypi: https://files.pythonhosted.org/packages/77/3c/aa88abe01f3be3d1f8f787d1d33dc83e76fec05945f9a28fbb41cfb99cd5/nvidia_cublas_cu12-12.9.1.4-py3-none-manylinux_2_27_x86_64.whl
   name: nvidia-cublas-cu12
   version: 12.9.1.4
@@ -7042,9 +6971,9 @@ packages:
   version: 3.4.0
   sha256: 69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd
   requires_python: '>=3.8'
-- pypi: git+https://github.com/optimagic-dev/optimagic.git?branch=main#e02ea4743cac9f861a5813f3b4b1283fd2ade730
+- pypi: git+https://github.com/optimagic-dev/optimagic.git#522b8c9a21226569ffd25e950e44f0c5de308c9d
   name: optimagic
-  version: 0.5.3.dev30+ge02ea4743
+  version: 0.5.3.dev31+g522b8c9a2
   requires_dist:
   - annotated-types
   - cloudpickle
@@ -7057,53 +6986,53 @@ packages:
   - sqlalchemy>=1.3
   - typing-extensions
   requires_python: '>=3.10'
-- conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.11.5-py313h541fbb8_0.conda
-  sha256: 6bb36f180ea4ba4f13f5e6ef8ec0b2fdd010d73430af53a05986ffc312091e8f
-  md5: 5dd1f02f38d71a29f3cfaf13c4cbf3dd
+- conda: https://conda.anaconda.org/conda-forge/linux-64/orjson-3.11.5-py314h3b757c3_0.conda
+  sha256: f8da6a925be44a867c172dd945049d7690ba6ae3a7905b61b1d5a4ba81fe0554
+  md5: 15ae5e4f52f2d9a98997e8859d35aa21
   depends:
   - python
-  - __glibc >=2.17,<3.0.a0
   - libgcc >=14
-  - python_abi 3.13.* *_cp313
+  - __glibc >=2.17,<3.0.a0
+  - python_abi 3.14.* *_cp314
   constrains:
   - __glibc >=2.17
   license: Apache-2.0
   license_family: APACHE
   purls:
   - pkg:pypi/orjson?source=hash-mapping
-  size: 317253
-  timestamp: 1765811463186
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/orjson-3.11.5-py313hfea8034_0.conda
-  sha256: 259cf50b358d2c1915123f0bf889db27d277efab7a3388c287f0dd4797764fe5
-  md5: d80421fc2b6f692925c82351f1c98407
+  size: 317280
+  timestamp: 1765811464445
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/orjson-3.11.5-py314hda6d10a_0.conda
+  sha256: 08f70edd4fc9f684083d18350c8a33c5a092e05aaaab7a97b09b381f8ca19eb7
+  md5: 21db7b1b5c5c04461bf40c33953f8cf7
   depends:
   - python
+  - python 3.14.* *_cp314
   - __osx >=11.0
-  - python 3.13.* *_cp313
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   constrains:
   - __osx >=11.0
   license: Apache-2.0
   license_family: APACHE
   purls:
   - pkg:pypi/orjson?source=hash-mapping
-  size: 288912
-  timestamp: 1765811468774
-- conda: https://conda.anaconda.org/conda-forge/win-64/orjson-3.11.4-py313hfbe8231_1.conda
-  sha256: cf55c2f55f7c0e8973da287217315c4b8652ca29dbcbcecfd0b3b8e48e784422
-  md5: db9e91caa5ee3f4891d340f8e323cc79
+  size: 288991
+  timestamp: 1765811524857
+- conda: https://conda.anaconda.org/conda-forge/win-64/orjson-3.11.5-py314h64f83cb_0.conda
+  sha256: 32014651690ee74eb65d4ef3f42f1ff679b216274f85c52b3be701cf16c6dff3
+  md5: 84b27320349d7fbf9fb6ad06141eec5b
   depends:
   - python
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
   - ucrt >=10.0.20348.0
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   license: Apache-2.0
   license_family: APACHE
   purls:
   - pkg:pypi/orjson?source=hash-mapping
-  size: 197832
-  timestamp: 1764441550892
+  size: 197828
+  timestamp: 1765811532648
 - conda: https://conda.anaconda.org/conda-forge/noarch/overrides-7.7.0-pyhd8ed1ab_1.conda
   sha256: 1840bd90d25d4930d60f57b4f38d4e0ae3f5b8db2819638709c36098c6ba770c
   md5: e51f1e4089cad105b6cac64bd8166587
@@ -7141,10 +7070,10 @@ packages:
   - pkg:pypi/packaging?source=hash-mapping
   size: 62477
   timestamp: 1745345660407
-- pypi: https://files.pythonhosted.org/packages/15/07/284f757f63f8a8d69ed4472bfd85122bd086e637bf4ed09de572d575a693/pandas-2.3.3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
+- pypi: https://files.pythonhosted.org/packages/15/b2/0e62f78c0c5ba7e3d2c5945a82456f4fac76c480940f805e0b97fcbc2f65/pandas-2.3.3-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
   name: pandas
   version: 2.3.3
-  sha256: 318d77e0e42a628c04dc56bcef4b40de67918f7041c2b061af1da41dcff670ac
+  sha256: ee67acbbf05014ea6c763beb097e03cd629961c8a632075eeb34247120abcb4b
   requires_dist:
   - numpy>=1.22.4 ; python_full_version < '3.11'
   - numpy>=1.23.2 ; python_full_version == '3.11.*'
@@ -7232,10 +7161,10 @@ packages:
   - xlsxwriter>=3.0.5 ; extra == 'all'
   - zstandard>=0.19.0 ; extra == 'all'
   requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/31/94/72fac03573102779920099bcac1c3b05975c2cb5f01eac609faf34bed1ca/pandas-2.3.3-cp313-cp313-macosx_11_0_arm64.whl
+- pypi: https://files.pythonhosted.org/packages/21/00/266d6b357ad5e6d3ad55093a7e8efc7dd245f5a842b584db9f30b0f0a287/pandas-2.3.3-cp314-cp314-macosx_11_0_arm64.whl
   name: pandas
   version: 2.3.3
-  sha256: bdcd9d1167f4885211e401b3036c0c8d9e274eee67ea8d0758a256d60704cfe8
+  sha256: 1611aedd912e1ff81ff41c745822980c49ce4a7907537be8692c8dbc31924593
   requires_dist:
   - numpy>=1.22.4 ; python_full_version < '3.11'
   - numpy>=1.23.2 ; python_full_version == '3.11.*'
@@ -7323,10 +7252,10 @@ packages:
   - xlsxwriter>=3.0.5 ; extra == 'all'
   - zstandard>=0.19.0 ; extra == 'all'
   requires_python: '>=3.9'
-- pypi: https://files.pythonhosted.org/packages/4f/c7/e54682c96a895d0c808453269e0b5928a07a127a15704fedb643e9b0a4c8/pandas-2.3.3-cp313-cp313-win_amd64.whl
+- pypi: https://files.pythonhosted.org/packages/a6/3d/124ac75fcd0ecc09b8fdccb0246ef65e35b012030defb0e0eba2cbbbe948/pandas-2.3.3-cp314-cp314-win_amd64.whl
   name: pandas
   version: 2.3.3
-  sha256: f8bfc0e12dc78f777f323f55c58649591b2cd0c43534e8355c51d3fede5f4dee
+  sha256: 1b07204a219b3b7350abaae088f451860223a52cfb8a6c53358e7948735158e5
   requires_dist:
   - numpy>=1.22.4 ; python_full_version < '3.11'
   - numpy>=1.23.2 ; python_full_version == '3.11.*'
@@ -7465,76 +7394,76 @@ packages:
   - pkg:pypi/pexpect?source=hash-mapping
   size: 53561
   timestamp: 1733302019362
-- conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.1.0-py313h80991f8_0.conda
-  sha256: bdad1e21cadd64154c45fa554247dd672288ad51982ca7d54b3fab63e40938df
-  md5: 183fe6b9e99e5c2b464c1573ec78eac8
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.1.0-py314h8ec4b1a_0.conda
+  sha256: 6d8e32dc44165cff96ec9c00383e998fd035983d971c5f35ebed6f5f51c4022a
+  md5: f9b6a8fbb8dcb840a0c1c052dc5092e4
   depends:
   - python
-  - libgcc >=14
   - __glibc >=2.17,<3.0.a0
-  - tk >=8.6.13,<8.7.0a0
-  - python_abi 3.13.* *_cp313
-  - libtiff >=4.7.1,<4.8.0a0
-  - libjpeg-turbo >=3.1.2,<4.0a0
+  - libgcc >=14
   - lcms2 >=2.17,<3.0a0
-  - libxcb >=1.17.0,<2.0a0
+  - libfreetype >=2.14.1
+  - libfreetype6 >=2.14.1
+  - libjpeg-turbo >=3.1.2,<4.0a0
   - zlib-ng >=2.3.2,<2.4.0a0
+  - libxcb >=1.17.0,<2.0a0
   - libwebp-base >=1.6.0,<2.0a0
   - openjpeg >=2.5.4,<3.0a0
-  - libfreetype >=2.14.1
-  - libfreetype6 >=2.14.1
+  - python_abi 3.14.* *_cp314
+  - libtiff >=4.7.1,<4.8.0a0
+  - tk >=8.6.13,<8.7.0a0
   license: HPND
   purls:
   - pkg:pypi/pillow?source=hash-mapping
-  size: 1043309
-  timestamp: 1767353193450
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/pillow-12.1.0-py313h45e5a15_0.conda
-  sha256: e5eaa7f00fca189848a0454303c56cc4edefd3e58a70bfd490d2cfe0d0aa525d
-  md5: 78a39731fd50dbd511de305934fe7e62
+  size: 1072995
+  timestamp: 1767353193452
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/pillow-12.1.0-py314hab283cf_0.conda
+  sha256: 3f88f2600862583c8bed3d37f4b95f0f96a459e9fdd36ca680472bc89a46e7bb
+  md5: 1f9dae6213643ac883e300c11df611eb
   depends:
   - python
   - __osx >=11.0
-  - python 3.13.* *_cp313
-  - libxcb >=1.17.0,<2.0a0
+  - python 3.14.* *_cp314
+  - libjpeg-turbo >=3.1.2,<4.0a0
   - openjpeg >=2.5.4,<3.0a0
-  - libtiff >=4.7.1,<4.8.0a0
+  - python_abi 3.14.* *_cp314
   - zlib-ng >=2.3.2,<2.4.0a0
-  - tk >=8.6.13,<8.7.0a0
-  - libjpeg-turbo >=3.1.2,<4.0a0
-  - python_abi 3.13.* *_cp313
+  - libxcb >=1.17.0,<2.0a0
   - lcms2 >=2.17,<3.0a0
   - libfreetype >=2.14.1
   - libfreetype6 >=2.14.1
+  - libtiff >=4.7.1,<4.8.0a0
+  - tk >=8.6.13,<8.7.0a0
   - libwebp-base >=1.6.0,<2.0a0
   license: HPND
   purls:
   - pkg:pypi/pillow?source=hash-mapping
-  size: 966296
-  timestamp: 1767353279679
-- conda: https://conda.anaconda.org/conda-forge/win-64/pillow-12.1.0-py313h38f99e1_0.conda
-  sha256: 181b4d169e7a671c387427ceb398d931802adace8808836b44295b07c3484abd
-  md5: 1927a42726a4ca0e94d5e8cb94c7a06d
+  size: 995543
+  timestamp: 1767353279681
+- conda: https://conda.anaconda.org/conda-forge/win-64/pillow-12.1.0-py314h61b30b5_0.conda
+  sha256: b30a83db337dab8579a46e3da7906851f53d6cf8c09695aef6d2a38b17636c1c
+  md5: 17dbdfedee39f31166b7e548f3ccc58a
   depends:
   - python
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
   - ucrt >=10.0.20348.0
-  - lcms2 >=2.17,<3.0a0
-  - libwebp-base >=1.6.0,<2.0a0
-  - python_abi 3.13.* *_cp313
   - libfreetype >=2.14.1
   - libfreetype6 >=2.14.1
-  - openjpeg >=2.5.4,<3.0a0
+  - tk >=8.6.13,<8.7.0a0
+  - libwebp-base >=1.6.0,<2.0a0
+  - lcms2 >=2.17,<3.0a0
+  - libtiff >=4.7.1,<4.8.0a0
   - zlib-ng >=2.3.2,<2.4.0a0
+  - openjpeg >=2.5.4,<3.0a0
   - libjpeg-turbo >=3.1.2,<4.0a0
-  - libtiff >=4.7.1,<4.8.0a0
   - libxcb >=1.17.0,<2.0a0
-  - tk >=8.6.13,<8.7.0a0
+  - python_abi 3.14.* *_cp314
   license: HPND
   purls:
   - pkg:pypi/pillow?source=hash-mapping
-  size: 946833
-  timestamp: 1767353195062
+  size: 973387
+  timestamp: 1767353195064
 - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
   sha256: 04c64fb78c520e5c396b6e07bc9082735a5cc28175dbe23138201d0a9441800b
   md5: 1bd2e65c8c7ef24f4639ae6e850dacc2
@@ -7557,6 +7486,7 @@ packages:
   constrains:
   - ipywidgets >=7.6
   license: MIT
+  license_family: MIT
   purls:
   - pkg:pypi/plotly?source=hash-mapping
   size: 4455861
@@ -7614,49 +7544,49 @@ packages:
   - pkg:pypi/prompt-toolkit?source=hash-mapping
   size: 273927
   timestamp: 1756321848365
-- conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.1-py313h54dd161_0.conda
-  sha256: 8a5f773e22ccd08fbda57c92f1d094533474db75f70db35311912cdcdb2f18ad
-  md5: d362949a1ed1ad4693b3928ad1d32c93
+- conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.1-py314h0f05182_0.conda
+  sha256: 324455a702ef721290de6e51d9af4f7ca057546d6398bbc6e88454db17cdaf6b
+  md5: 28af9719e28f0054e9aee68153899293
   depends:
   - python
-  - libgcc >=14
   - __glibc >=2.17,<3.0.a0
-  - python_abi 3.13.* *_cp313
+  - libgcc >=14
+  - python_abi 3.14.* *_cp314
   license: BSD-3-Clause
   license_family: BSD
   purls:
   - pkg:pypi/psutil?source=hash-mapping
-  size: 225429
-  timestamp: 1767012386804
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/psutil-7.2.1-py313h6688731_0.conda
-  sha256: 2abd12a0371836075a72e12fde44f63ea08b3781e5b6ec997233d50b9c9832d9
-  md5: c3a1b24571871fec4498a0226a3c22c1
+  size: 228170
+  timestamp: 1767012382363
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/psutil-7.2.1-py314ha14b1ff_0.conda
+  sha256: 686b643b97df8e7076b971820fb9b5d2ed0ea8a5a82922910da1600a6f462b79
+  md5: 6d799fc0d0178eb63202bf99ff7bc24f
   depends:
   - python
-  - python 3.13.* *_cp313
+  - python 3.14.* *_cp314
   - __osx >=11.0
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   license: BSD-3-Clause
   license_family: BSD
   purls:
   - pkg:pypi/psutil?source=hash-mapping
-  size: 238851
-  timestamp: 1767012473931
-- conda: https://conda.anaconda.org/conda-forge/win-64/psutil-7.2.1-py313h5fd188c_0.conda
-  sha256: 025574efd6e9d5b90d89ec1da8423132ab9c6131e21be7ec91b9fd7a14665a57
-  md5: 8732097a02c66f6b260dd15b705a014e
+  size: 241751
+  timestamp: 1767012600474
+- conda: https://conda.anaconda.org/conda-forge/win-64/psutil-7.2.1-py314hc5dbbe4_0.conda
+  sha256: d776855d47e14d8b1521a3949c1d1dc3848c690170253ecc439264e219859e22
+  md5: 65df3730bedf9c24f54414c8316f8e72
   depends:
   - python
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
   - ucrt >=10.0.20348.0
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   license: BSD-3-Clause
   license_family: BSD
   purls:
-  - pkg:pypi/psutil?source=hash-mapping
-  size: 243141
-  timestamp: 1767012395730
+  - pkg:pypi/psutil?source=compressed-mapping
+  size: 245991
+  timestamp: 1767012412984
 - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda
   sha256: 9c88f8c64590e9567c6c80823f0328e58d3b1efb0e1c539c0315ceca764e0973
   md5: b3c17d95b5a10c6e64a21fa17573e70e
@@ -7745,38 +7675,38 @@ packages:
   - pkg:pypi/pygments?source=hash-mapping
   size: 889287
   timestamp: 1750615908735
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-core-12.1-py313h40b429f_0.conda
-  sha256: 307ca29ebf2317bd2561639b1ee0290fd8c03c3450fa302b9f9437d8df6a5280
-  md5: 31a0a72f3466682d0ea2ebcbd7d319b8
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-core-12.1-py314h3a4d195_0.conda
+  sha256: df5af268c5a74b7160d772c263ece6f43257faff571783443e34b5f1d5a61cf2
+  md5: 75a84fc8337557347252cc4fd3ba2a93
   depends:
   - __osx >=11.0
   - libffi >=3.5.2,<3.6.0a0
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
+  - python_abi 3.14.* *_cp314
   - setuptools
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/pyobjc-core?source=hash-mapping
-  size: 481508
-  timestamp: 1763152124940
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-framework-cocoa-12.1-py313hcc5defa_0.conda
-  sha256: 194e188d8119befc952d04157079733e2041a7a502d50340ddde632658799fdc
-  md5: a6d28c8fc266a3d3c3dae183e25c4d31
+  size: 483374
+  timestamp: 1763151489724
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyobjc-framework-cocoa-12.1-py314h36abed7_0.conda
+  sha256: aa76ee4328d0514d7c1c455dcd2d3b547db1c59797e54ce0a3f27de5b970e508
+  md5: 4219bb3408016e22316cf8b443b5ef93
   depends:
   - __osx >=11.0
   - libffi >=3.5.2,<3.6.0a0
   - pyobjc-core 12.1.*
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/pyobjc-framework-cocoa?source=hash-mapping
-  size: 376136
-  timestamp: 1763160678792
+  size: 374792
+  timestamp: 1763160601898
 - conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.1-pyhcf101f3_0.conda
   sha256: 0c70bc577f5efa87501bdc841b88f594f4d3f3a992dfb851e2130fa5c817835b
   md5: d837065e4e0de4962c3462079c23f969
@@ -7901,10 +7831,10 @@ packages:
   - pkg:pypi/pytest-xdist?source=hash-mapping
   size: 39300
   timestamp: 1751452761594
-- conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.11-hc97d973_100_cp313.conda
+- conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.14.2-h32b2ec7_100_cp314.conda
   build_number: 100
-  sha256: 9cf014cf28e93ee242bacfbf664e8b45ae06e50b04291e640abeaeb0cba0364c
-  md5: 0cbb0010f1d8ecb64a428a8d4214609e
+  sha256: a120fb2da4e4d51dd32918c149b04a08815fd2bd52099dad1334647984bb07f1
+  md5: 1cef1236a05c3a98f68c33ae9425f656
   depends:
   - __glibc >=2.17,<3.0.a0
   - bzip2 >=1.0.8,<2.0a0
@@ -7919,19 +7849,20 @@ packages:
   - libzlib >=1.3.1,<2.0a0
   - ncurses >=6.5,<7.0a0
   - openssl >=3.5.4,<4.0a0
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   - readline >=8.2,<9.0a0
   - tk >=8.6.13,<8.7.0a0
   - tzdata
+  - zstd >=1.5.7,<1.6.0a0
   license: Python-2.0
   purls: []
-  size: 37226336
-  timestamp: 1765021889577
-  python_site_packages_path: lib/python3.13/site-packages
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.13.11-hfc2f54d_100_cp313.conda
+  size: 36790521
+  timestamp: 1765021515427
+  python_site_packages_path: lib/python3.14/site-packages
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.14.2-h40d2674_100_cp314.conda
   build_number: 100
-  sha256: c476f4e9b6d97c46b496b442878924868a54e5727251549ebfc82027aa52af68
-  md5: 18a8c69608151098a8fb75eea64cc266
+  sha256: 1a93782e90b53e04c2b1a50a0f8bf0887936649d19dba6a05b05c4b44dae96b7
+  md5: 14f15ab0d31a2ee5635aa56e77132594
   depends:
   - __osx >=11.0
   - bzip2 >=1.0.8,<2.0a0
@@ -7943,19 +7874,20 @@ packages:
   - libzlib >=1.3.1,<2.0a0
   - ncurses >=6.5,<7.0a0
   - openssl >=3.5.4,<4.0a0
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   - readline >=8.2,<9.0a0
   - tk >=8.6.13,<8.7.0a0
   - tzdata
+  - zstd >=1.5.7,<1.6.0a0
   license: Python-2.0
   purls: []
-  size: 12920650
-  timestamp: 1765020887340
-  python_site_packages_path: lib/python3.13/site-packages
-- conda: https://conda.anaconda.org/conda-forge/win-64/python-3.13.11-h09917c8_100_cp313.conda
+  size: 13575758
+  timestamp: 1765021280625
+  python_site_packages_path: lib/python3.14/site-packages
+- conda: https://conda.anaconda.org/conda-forge/win-64/python-3.14.2-h4b44e0e_100_cp314.conda
   build_number: 100
-  sha256: 0ee0402368783e1fad10025719530499c517a3dbbdfbe18351841d9b7aef1d6a
-  md5: 9e4c9a7ee9c4ab5b3778ab73e583283e
+  sha256: 6857d7c97cc71fe9ba298dcb1d3b66cc7df425132ab801babd655faa3df48f32
+  md5: c3c73414d5ae3f543c531c978d9cc8b8
   depends:
   - bzip2 >=1.0.8,<2.0a0
   - libexpat >=2.7.3,<3.0a0
@@ -7965,16 +7897,17 @@ packages:
   - libsqlite >=3.51.1,<4.0a0
   - libzlib >=1.3.1,<2.0a0
   - openssl >=3.5.4,<4.0a0
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   - tk >=8.6.13,<8.7.0a0
   - tzdata
   - ucrt >=10.0.20348.0
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
+  - zstd >=1.5.7,<1.6.0a0
   license: Python-2.0
   purls: []
-  size: 16617922
-  timestamp: 1765019627175
+  size: 16833248
+  timestamp: 1765020224759
   python_site_packages_path: Lib/site-packages
 - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
   sha256: d6a17ece93bbd5139e02d2bd7dbfa80bee1a4261dced63f65f679121686bf664
@@ -8001,16 +7934,16 @@ packages:
   - pkg:pypi/fastjsonschema?source=hash-mapping
   size: 244628
   timestamp: 1755304154927
-- conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.13.11-h4df99d1_100.conda
-  sha256: 4b08d4c2c4b956d306b4868d3faf724eebb5d6e6b170fad2eb0f2d4eb227f1af
-  md5: d1461b2e63b1909f4f5b41c823bd90ae
+- conda: https://conda.anaconda.org/conda-forge/noarch/python-gil-3.14.2-h4df99d1_100.conda
+  sha256: 8203dc90a5cb6687f5bfcf332eeaf494ec95d24ed13fca3c82ef840f0bb92a5d
+  md5: 0064ab66736c4814864e808169dc7497
   depends:
-  - cpython 3.13.11.*
-  - python_abi * *_cp313
+  - cpython 3.14.2.*
+  - python_abi * *_cp314
   license: Python-2.0
   purls: []
-  size: 48352
-  timestamp: 1765019767640
+  size: 49287
+  timestamp: 1765020424843
 - conda: https://conda.anaconda.org/conda-forge/noarch/python-json-logger-2.0.7-pyhd8ed1ab_0.conda
   sha256: 4790787fe1f4e8da616edca4acf6a4f8ed4e7c6967aa31b920208fc8f95efcca
   md5: a61bf9ec79426938ff785eb69dbb1960
@@ -8050,17 +7983,17 @@ packages:
   - pkg:pypi/tzdata?source=compressed-mapping
   size: 143542
   timestamp: 1765719982349
-- conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.13-8_cp313.conda
+- conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.14-8_cp314.conda
   build_number: 8
-  sha256: 210bffe7b121e651419cb196a2a63687b087497595c9be9d20ebe97dd06060a7
-  md5: 94305520c52a4aa3f6c2b1ff6008d9f8
+  sha256: ad6d2e9ac39751cc0529dd1566a26751a0bf2542adb0c232533d32e176e21db5
+  md5: 0539938c55b6b1a59b560e843ad864a4
   constrains:
-  - python 3.13.* *_cp313
+  - python 3.14.* *_cp314
   license: BSD-3-Clause
   license_family: BSD
   purls: []
-  size: 7002
-  timestamp: 1752805902938
+  size: 6989
+  timestamp: 1752805904792
 - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
   sha256: 8d2a8bf110cc1fc3df6904091dead158ba3e614d8402a83e51ed3a8aa93cdeb0
   md5: bc8e3267d44011051f2eb14d22fb0960
@@ -8072,9 +8005,9 @@ packages:
   - pkg:pypi/pytz?source=hash-mapping
   size: 189015
   timestamp: 1742920947249
-- conda: https://conda.anaconda.org/conda-forge/win-64/pywin32-311-py313h40c08fc_1.conda
-  sha256: 87eaeb79b5961e0f216aa840bc35d5f0b9b123acffaecc4fda4de48891901f20
-  md5: 1ce4f826332dca56c76a5b0cc89fb19e
+- conda: https://conda.anaconda.org/conda-forge/win-64/pywin32-311-py314h8f8f202_1.conda
+  sha256: 6918a8067f296f3c65d43e84558170c9e6c3f4dd735cfe041af41a7fdba7b171
+  md5: 2d7b7ba21e8a8ced0eca553d4d53f773
   depends:
   - python
   - vc >=14.3,<15
@@ -8083,19 +8016,19 @@ packages:
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
   - ucrt >=10.0.20348.0
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   license: PSF-2.0
   license_family: PSF
   purls:
   - pkg:pypi/pywin32?source=hash-mapping
-  size: 6695114
-  timestamp: 1756487139550
-- conda: https://conda.anaconda.org/conda-forge/win-64/pywinpty-2.0.15-py313h5813708_1.conda
-  sha256: d34a7cd0a4a7dc79662cb6005e01d630245d9a942e359eb4d94b2fb464ed2552
-  md5: 8f01ed27e2baa455e753301218e054fd
-  depends:
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  size: 6713155
+  timestamp: 1756487145487
+- conda: https://conda.anaconda.org/conda-forge/win-64/pywinpty-2.0.15-py314h51f0985_1.conda
+  sha256: 048e20641da680aedaab285640a2aca56b7b5baf7a18f8f164f2796e13628c1f
+  md5: dd84e8748bd3c85a5c751b0576488080
+  depends:
+  - python >=3.14.0rc3,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   - ucrt >=10.0.20348.0
   - vc >=14.2,<15
   - vc14_runtime >=14.29.30139
@@ -8104,54 +8037,22 @@ packages:
   license_family: MIT
   purls:
   - pkg:pypi/pywinpty?source=hash-mapping
-  size: 216075
-  timestamp: 1759556799508
-- conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py313h3dea7bd_0.conda
-  sha256: 40dcd6718dce5fbee8aabdd0519f23d456d8feb2e15ac352eaa88bbfd3a881af
-  md5: 4794ea0adaebd9f844414e594b142cb2
-  depends:
-  - __glibc >=2.17,<3.0.a0
-  - libgcc >=14
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
-  - yaml >=0.2.5,<0.3.0a0
-  license: MIT
-  license_family: MIT
-  purls:
-  - pkg:pypi/pyyaml?source=hash-mapping
-  size: 207109
-  timestamp: 1758892173548
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyyaml-6.0.3-py313h7d74516_0.conda
-  sha256: f5be0d84f72a567b7333b9efa74a65bfa44a25658cf107ffa3fc65d3ae6660d7
-  md5: 0e8e3235217b4483a7461b63dca5826b
-  depends:
-  - __osx >=11.0
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
-  - yaml >=0.2.5,<0.3.0a0
-  license: MIT
-  license_family: MIT
-  purls:
-  - pkg:pypi/pyyaml?source=hash-mapping
-  size: 191630
-  timestamp: 1758892258120
-- conda: https://conda.anaconda.org/conda-forge/win-64/pyyaml-6.0.3-py313hd650c13_0.conda
-  sha256: 5d9fd32d318b9da615524589a372b33a6f3d07db2708de16570d70360bf638c2
-  md5: c067122d76f8dcbe0848822942ba07be
-  depends:
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
-  - ucrt >=10.0.20348.0
-  - vc >=14.3,<15
-  - vc14_runtime >=14.44.35208
-  - yaml >=0.2.5,<0.3.0a0
+  size: 216325
+  timestamp: 1759557436167
+- conda: https://conda.anaconda.org/conda-forge/noarch/pyyaml-6.0.3-pyh7db6752_0.conda
+  sha256: 828af2fd7bb66afc9ab1c564c2046be391aaf66c0215f05afaf6d7a9a270fe2a
+  md5: b12f41c0d7fb5ab81709fcc86579688f
+  depends:
+  - python >=3.10.*
+  - yaml
+  track_features:
+  - pyyaml_no_compile
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/pyyaml?source=hash-mapping
-  size: 182043
-  timestamp: 1758892011955
+  size: 45223
+  timestamp: 1758891992558
 - conda: https://conda.anaconda.org/conda-forge/linux-64/pyzmq-27.1.0-py312hfb55c3c_0.conda
   noarch: python
   sha256: a00a41b66c12d9c60e66b391e9a4832b7e28743348cf4b48b410b91927cd7819
@@ -8347,56 +8248,56 @@ packages:
   - pkg:pypi/rich?source=hash-mapping
   size: 200840
   timestamp: 1760026188268
-- conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.30.0-py313h843e2db_0.conda
-  sha256: 076d26e51c62c8ecfca6eb19e3c1febdd7632df1990a7aa53da5df5e54482b1c
-  md5: 779e3307a0299518713765b83a36f4b1
+- conda: https://conda.anaconda.org/conda-forge/linux-64/rpds-py-0.30.0-py314h2e6c369_0.conda
+  sha256: e53b0cbf3b324eaa03ca1fe1a688fdf4ab42cea9c25270b0a7307d8aaaa4f446
+  md5: c1c368b5437b0d1a68f372ccf01cb133
   depends:
   - python
   - libgcc >=14
   - __glibc >=2.17,<3.0.a0
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   constrains:
   - __glibc >=2.17
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/rpds-py?source=hash-mapping
-  size: 383230
-  timestamp: 1764543223529
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/rpds-py-0.30.0-py313h2c089d5_0.conda
-  sha256: db63344f91e8bfe77703c6764aa9eeafb44d165e286053214722814eabda0264
-  md5: 190c2d0d4e98ec97df48cdb74caf44d8
+  size: 376121
+  timestamp: 1764543122774
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/rpds-py-0.30.0-py314haad56a0_0.conda
+  sha256: e161dd97403b8b8a083d047369a5cf854557dba1204d29e2f0250f5ac4403925
+  md5: 76a4f88d1b7748c477abf3c341edc64c
   depends:
   - python
   - __osx >=11.0
-  - python 3.13.* *_cp313
-  - python_abi 3.13.* *_cp313
+  - python 3.14.* *_cp314
+  - python_abi 3.14.* *_cp314
   constrains:
   - __osx >=11.0
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/rpds-py?source=hash-mapping
-  size: 358961
-  timestamp: 1764543165314
-- conda: https://conda.anaconda.org/conda-forge/win-64/rpds-py-0.30.0-py313hfbe8231_0.conda
-  sha256: 27bd383787c0df7a0a926b11014fd692d60d557398dcf1d50c55aa2378507114
-  md5: 58ae648b12cfa6df3923b5fd219931cb
+  size: 350976
+  timestamp: 1764543169524
+- conda: https://conda.anaconda.org/conda-forge/win-64/rpds-py-0.30.0-py314h9f07db2_0.conda
+  sha256: e4435368c5c25076dc0f5918ba531c5a92caee8e0e2f9912ef6810049cf00db2
+  md5: e86531e278ad304438e530953cd55d14
   depends:
   - python
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
   - ucrt >=10.0.20348.0
-  - python_abi 3.13.* *_cp313
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/rpds-py?source=hash-mapping
-  size: 243419
-  timestamp: 1764543047271
-- conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.16.3-py313h4b8bb8b_2.conda
-  sha256: a5ddc728be0589e770f59e45e3c6c670c56d96a801ddf76a304cc0af7bcef5c4
-  md5: 0be9bd58abfb3e8f97260bd0176d5331
+  size: 235780
+  timestamp: 1764543046065
+- conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.16.3-py314hf07bd8e_2.conda
+  sha256: 652f9a235051c1d39ccd2fe7e9326792b046a1d93de42171977fa1ba9668a0e8
+  md5: ee95e8bb52e35c3267a53d3ee1347cc4
   depends:
   - __glibc >=2.17,<3.0.a0
   - libblas >=3.9.0,<4.0a0
@@ -8409,17 +8310,17 @@ packages:
   - numpy <2.6
   - numpy >=1.23,<3
   - numpy >=1.25.2
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   license: BSD-3-Clause
   license_family: BSD
   purls:
   - pkg:pypi/scipy?source=compressed-mapping
-  size: 16785487
-  timestamp: 1766108773270
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/scipy-1.16.3-py313h29d7d31_2.conda
-  sha256: ee3cbddb7d598c78b592fafbfa3eaf8c89df353bbed56a1a9f32e9f7daa49bb4
-  md5: a3324bd937a39cbbf1cbe0940160e19e
+  size: 16982488
+  timestamp: 1766108668132
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/scipy-1.16.3-py314h725efaa_2.conda
+  sha256: 282f8b244f31d8c2e0ce401b0473e8090de4b59326018a360419693b629e6b87
+  md5: 6333b784ddfcccd3f5569f812f66c352
   depends:
   - __osx >=11.0
   - libblas >=3.9.0,<4.0a0
@@ -8431,18 +8332,18 @@ packages:
   - numpy <2.6
   - numpy >=1.23,<3
   - numpy >=1.25.2
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
+  - python_abi 3.14.* *_cp314
   license: BSD-3-Clause
   license_family: BSD
   purls:
-  - pkg:pypi/scipy?source=hash-mapping
-  size: 13929516
-  timestamp: 1766109298759
-- conda: https://conda.anaconda.org/conda-forge/win-64/scipy-1.16.3-py313he51e9a2_2.conda
-  sha256: 997a2202126425438a16de7ef1e5e924bd66feb43bda5b71326e281c7331489d
-  md5: a49556572438d5477f1eca06bb6d0770
+  - pkg:pypi/scipy?source=compressed-mapping
+  size: 13880523
+  timestamp: 1766109018710
+- conda: https://conda.anaconda.org/conda-forge/win-64/scipy-1.16.3-py314h221f224_2.conda
+  sha256: 99d6198dc05171610073083c9d218d2a9adfa756659b391183d21cca55f888f1
+  md5: b600c47282ee91e492b89f65708a5c9a
   depends:
   - libblas >=3.9.0,<4.0a0
   - libcblas >=3.9.0,<4.0a0
@@ -8450,8 +8351,8 @@ packages:
   - numpy <2.6
   - numpy >=1.23,<3
   - numpy >=1.25.2
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   - ucrt >=10.0.20348.0
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
@@ -8459,8 +8360,8 @@ packages:
   license_family: BSD
   purls:
   - pkg:pypi/scipy?source=hash-mapping
-  size: 15066293
-  timestamp: 1766109539389
+  size: 15082636
+  timestamp: 1766109482825
 - conda: https://conda.anaconda.org/conda-forge/noarch/send2trash-2.0.0-pyh5552912_0.conda
   sha256: 5893e203cb099c784bf5b08d29944b5402beebcc361d55e54b676e9b355c7844
   md5: dcff6f8ea9e86a0bda978b88f89f2310
@@ -8513,40 +8414,40 @@ packages:
   - pkg:pypi/setuptools?source=hash-mapping
   size: 748788
   timestamp: 1748804951958
-- conda: https://conda.anaconda.org/conda-forge/linux-64/simplejson-3.20.2-py313h07c4f96_1.conda
-  sha256: cf44d6bd3dc3be6b683fac251d6b53d508d041506a2101fd7cdb404468cf8be3
-  md5: 1cc1de04373b633177f4d367b8b75270
+- conda: https://conda.anaconda.org/conda-forge/linux-64/simplejson-3.20.2-py314h5bd0f2a_1.conda
+  sha256: fde24560898ecbb63edb6580fbf09fa07e10f55a89f8ae35f891f712f1d07872
+  md5: b2f9edf27e434edc6072e6f7c076015f
   depends:
   - __glibc >=2.17,<3.0.a0
   - libgcc >=14
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/simplejson?source=hash-mapping
-  size: 133692
-  timestamp: 1762506927030
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/simplejson-3.20.2-py313h6535dbc_1.conda
-  sha256: ef09659f0248066e8c06a0bd8bd1a360b8158cd2d73c65c969897e20344c6a2a
-  md5: 27a8bc65b5f0aecb87a01568e573e6ae
+  size: 135289
+  timestamp: 1762507017143
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/simplejson-3.20.2-py314h0612a62_1.conda
+  sha256: e6a6a2aab805c4c50464aecff3f752a78ce15bb1b9de006b1d929d0673f3a386
+  md5: 82c463d19f1d85e60d520d129c67b483
   depends:
   - __osx >=11.0
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/simplejson?source=hash-mapping
-  size: 133717
-  timestamp: 1762507593463
-- conda: https://conda.anaconda.org/conda-forge/win-64/simplejson-3.20.2-py313h5ea7bf4_1.conda
-  sha256: b1ea3625e7dcda6ea6121dc61461da9bc9be54a99aa20ed26a5ee5b43663b5c4
-  md5: bcdc4785e018f4325845f8217333a17e
-  depends:
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  size: 134426
+  timestamp: 1762507464057
+- conda: https://conda.anaconda.org/conda-forge/win-64/simplejson-3.20.2-py314h5a2d7ad_1.conda
+  sha256: 28f67233b03f8f1ebdcd5b35d1700d75101be0e9decf4975b8dc867609d4a507
+  md5: f5d14f3ecb62b185cb571b79034df477
+  depends:
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   - ucrt >=10.0.20348.0
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
@@ -8554,8 +8455,8 @@ packages:
   license_family: MIT
   purls:
   - pkg:pypi/simplejson?source=hash-mapping
-  size: 132684
-  timestamp: 1762507090611
+  size: 134026
+  timestamp: 1762507518751
 - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
   sha256: 458227f759d5e3fcec5d9b7acce54e10c9e1f4f4b7ec978f3bfd54ce4ee9853d
   md5: 3339e3b65d58accf4ca4fb8748ab16b3
@@ -8570,14 +8471,15 @@ packages:
   timestamp: 1753199211006
 - pypi: ./
   name: skillmodels
-  version: 0.0.53.dev37+gcc710a63e.d20260108
-  sha256: 167bd74677526ae18b099ed73f00727ec8b124bba6895ba9afe3862d04fbcd84
+  version: 0.0.24.dev248+gb15372a52.d20260111
+  sha256: 351c9cc8e9a879fe282da02d64b6d79d4326650bb3466ea8db23b29628fba895
   requires_dist:
   - dags
+  - frozendict
   - jax>=0.8
   - numpy
   - pandas
-  requires_python: '>=3.13,<3.14'
+  requires_python: '>=3.14,<3.15'
 - conda: https://conda.anaconda.org/conda-forge/noarch/snakeviz-2.2.2-pyhd8ed1ab_1.conda
   sha256: 833326122c18887b338262c13365cb146b6702c79d72da74a1c6b8af4c50e162
   md5: 421b7a950e384949ca1b0f04f0751ce0
@@ -8612,10 +8514,10 @@ packages:
   - pkg:pypi/soupsieve?source=compressed-mapping
   size: 37951
   timestamp: 1766075884412
-- pypi: https://files.pythonhosted.org/packages/0e/50/80a8d080ac7d3d321e5e5d420c9a522b0aa770ec7013ea91f9a8b7d36e4a/sqlalchemy-2.0.45-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
+- pypi: https://files.pythonhosted.org/packages/89/a2/0e1590e9adb292b1d576dbcf67ff7df8cf55e56e78d2c927686d01080f4b/sqlalchemy-2.0.45-cp314-cp314-win_amd64.whl
   name: sqlalchemy
   version: 2.0.45
-  sha256: 672c45cae53ba88e0dad74b9027dddd09ef6f441e927786b05bec75d949fbb2e
+  sha256: 4748601c8ea959e37e03d13dcda4a44837afcd1b21338e637f7c935b8da06177
   requires_dist:
   - importlib-metadata ; python_full_version < '3.8'
   - greenlet>=1 ; platform_machine == 'AMD64' or platform_machine == 'WIN32' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'ppc64le' or platform_machine == 'win32' or platform_machine == 'x86_64'
@@ -8650,10 +8552,10 @@ packages:
   - typing-extensions!=3.10.0.1 ; extra == 'aiosqlite'
   - sqlcipher3-binary ; extra == 'sqlcipher'
   requires_python: '>=3.7'
-- pypi: https://files.pythonhosted.org/packages/bf/e1/3ccb13c643399d22289c6a9786c1a91e3dcbb68bce4beb44926ac2c557bf/sqlalchemy-2.0.45-py3-none-any.whl
+- pypi: https://files.pythonhosted.org/packages/b3/27/caf606ee924282fe4747ee4fd454b335a72a6e018f97eab5ff7f28199e16/sqlalchemy-2.0.45-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
   name: sqlalchemy
   version: 2.0.45
-  sha256: 5225a288e4c8cc2308dbdd874edad6e7d0fd38eac1e9e5f23503425c8eee20d0
+  sha256: 883c600c345123c033c2f6caca18def08f1f7f4c3ebeb591a63b6fceffc95cce
   requires_dist:
   - importlib-metadata ; python_full_version < '3.8'
   - greenlet>=1 ; platform_machine == 'AMD64' or platform_machine == 'WIN32' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'ppc64le' or platform_machine == 'win32' or platform_machine == 'x86_64'
@@ -8688,10 +8590,10 @@ packages:
   - typing-extensions!=3.10.0.1 ; extra == 'aiosqlite'
   - sqlcipher3-binary ; extra == 'sqlcipher'
   requires_python: '>=3.7'
-- pypi: https://files.pythonhosted.org/packages/c0/c5/d17113020b2d43073412aeca09b60d2009442420372123b8d49cc253f8b8/sqlalchemy-2.0.45-cp313-cp313-win_amd64.whl
+- pypi: https://files.pythonhosted.org/packages/bf/e1/3ccb13c643399d22289c6a9786c1a91e3dcbb68bce4beb44926ac2c557bf/sqlalchemy-2.0.45-py3-none-any.whl
   name: sqlalchemy
   version: 2.0.45
-  sha256: afbf47dc4de31fa38fd491f3705cac5307d21d4bb828a4f020ee59af412744ee
+  sha256: 5225a288e4c8cc2308dbdd874edad6e7d0fd38eac1e9e5f23503425c8eee20d0
   requires_dist:
   - importlib-metadata ; python_full_version < '3.8'
   - greenlet>=1 ; platform_machine == 'AMD64' or platform_machine == 'WIN32' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'ppc64le' or platform_machine == 'win32' or platform_machine == 'x86_64'
@@ -8759,19 +8661,19 @@ packages:
   requires_dist:
   - pyreadline3 ; sys_platform == 'win32'
   requires_python: '>=3.8'
-- conda: https://conda.anaconda.org/conda-forge/win-64/tbb-2022.3.0-hd094cb3_1.conda
-  sha256: c31cac57913a699745d124cdc016a63e31c5749f16f60b3202414d071fc50573
-  md5: 17c38aaf14c640b85c4617ccb59c1146
+- conda: https://conda.anaconda.org/conda-forge/win-64/tbb-2022.3.0-h3155e25_2.conda
+  sha256: abd9a489f059fba85c8ffa1abdaa4d515d6de6a3325238b8e81203b913cf65a9
+  md5: 0f9817ffbe25f9e69ceba5ea70c52606
   depends:
-  - libhwloc >=2.12.1,<2.12.2.0a0
+  - libhwloc >=2.12.2,<2.12.3.0a0
   - ucrt >=10.0.20348.0
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
   license: Apache-2.0
   license_family: APACHE
   purls: []
-  size: 155714
-  timestamp: 1762510341121
+  size: 155869
+  timestamp: 1767886839029
 - conda: https://conda.anaconda.org/conda-forge/noarch/terminado-0.18.1-pyh6dadd2b_1.conda
   sha256: b375e8df0d5710717c31e7c8e93c025c37fa3504aea325c7a55509f64e5d4340
   md5: e43ca10d61e55d0a8ec5d8c62474ec9e
@@ -8802,9 +8704,9 @@ packages:
   - pkg:pypi/terminado?source=hash-mapping
   size: 24749
   timestamp: 1766513766867
-- conda: https://conda.anaconda.org/conda-forge/noarch/textual-7.0.0-pyhcf101f3_0.conda
-  sha256: 50ea42e243d349b8218168c06bfd408f4dcda68d4364de1f5866507e009e3cfd
-  md5: ca39d364b4f1b395bb6a70312d455c28
+- conda: https://conda.anaconda.org/conda-forge/noarch/textual-7.0.1-pyhcf101f3_0.conda
+  sha256: b601d7f7d200465547ed76fd6b95701d94b0bbf0ab1d9dae4beb2f7012947cdd
+  md5: 13e92b552eb58a0c243a967a7d9e4d78
   depends:
   - pygments >=2.19.2,<3.0.0
   - typing_extensions >=4.4.0,<5.0.0
@@ -8821,8 +8723,8 @@ packages:
   license_family: MIT
   purls:
   - pkg:pypi/textual?source=hash-mapping
-  size: 526014
-  timestamp: 1767448924135
+  size: 525875
+  timestamp: 1767859034631
 - conda: https://conda.anaconda.org/conda-forge/noarch/tinycss2-1.5.1-pyhcf101f3_0.conda
   sha256: 7c803480dbfb8b536b9bf6287fa2aa0a4f970f8c09075694174eb4550a4524cd
   md5: c0d0b883e97906f7524e2aac94be0e0d
@@ -8885,40 +8787,40 @@ packages:
   - pkg:pypi/tomli?source=compressed-mapping
   size: 20973
   timestamp: 1760014679845
-- conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py313h07c4f96_0.conda
-  sha256: 6006d4e5a6ff99be052c939e43adee844a38f2dc148f44a7c11aa0011fd3d811
-  md5: 82da2dcf1ea3e298f2557b50459809e0
+- conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py314h5bd0f2a_0.conda
+  sha256: b8f9f9ae508d79c9c697eb01b6a8d2ed4bc1899370f44aa6497c8abbd15988ea
+  md5: e35f08043f54d26a1be93fdbf90d30c3
   depends:
   - __glibc >=2.17,<3.0.a0
   - libgcc >=14
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   license: Apache-2.0
   license_family: Apache
   purls:
   - pkg:pypi/tornado?source=hash-mapping
-  size: 878109
-  timestamp: 1765458900582
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/tornado-6.5.4-py313h6535dbc_0.conda
-  sha256: a8130a361b7bc21190836ba8889276cc263fcb09f52bf22efcaed1de98179948
-  md5: 67a85c1b5c17124eaf9194206afd5159
+  size: 905436
+  timestamp: 1765458949518
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/tornado-6.5.4-py314h0612a62_0.conda
+  sha256: affbc6300e1baef5848f6e69569733a3e7a118aa642487c853f53d6f2bd23b89
+  md5: 83e1a2d7b0c1352870bbe9d9406135cf
   depends:
   - __osx >=11.0
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
+  - python_abi 3.14.* *_cp314
   license: Apache-2.0
   license_family: Apache
   purls:
   - pkg:pypi/tornado?source=hash-mapping
-  size: 877647
-  timestamp: 1765836696426
-- conda: https://conda.anaconda.org/conda-forge/win-64/tornado-6.5.4-py313h5ea7bf4_0.conda
-  sha256: 81b131db1bebed88f11a5f9891c0c0a7c6998dfd96cd96f54839f3a0cbebd5a0
-  md5: 1402782887fafaa117a8d76d2cfa4761
-  depends:
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  size: 909298
+  timestamp: 1765836779269
+- conda: https://conda.anaconda.org/conda-forge/win-64/tornado-6.5.4-py314h5a2d7ad_0.conda
+  sha256: 40fde32a4992ab0f875618f97d9aadf263d39c6c92ace7572c6b0a71c655abe1
+  md5: 00157f40fd3ea957a2616e9ffda6b84f
+  depends:
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   - ucrt >=10.0.20348.0
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
@@ -8926,8 +8828,8 @@ packages:
   license_family: Apache
   purls:
   - pkg:pypi/tornado?source=hash-mapping
-  size: 880049
-  timestamp: 1765836649731
+  size: 908399
+  timestamp: 1765836848636
 - conda: https://conda.anaconda.org/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda
   sha256: f39a5620c6e8e9e98357507262a7869de2ae8cc07da8b7f84e517c9fd6c2b959
   md5: 019a7385be9af33791c989871317e1ed
@@ -8939,20 +8841,20 @@ packages:
   - pkg:pypi/traitlets?source=hash-mapping
   size: 110051
   timestamp: 1733367480074
-- pypi: https://files.pythonhosted.org/packages/42/36/82e66b9753a76964d26fd9bc3514ea0abce0a5ba5ad7d5f084070c6981da/ty-0.0.10-py3-none-win_amd64.whl
+- pypi: https://files.pythonhosted.org/packages/74/18/8dd4fe6df1fd66f3e83b4798eddb1d8482d9d9b105f25099b76703402ebb/ty-0.0.11-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
   name: ty
-  version: 0.0.10
-  sha256: 16deb77a72cf93b89b4d29577829613eda535fbe030513dfd9fba70fe38bc9f5
+  version: 0.0.11
+  sha256: 25f88e8789072830348cb59b761d5ced70642ed5600673b4bf6a849af71eca8b
   requires_python: '>=3.8'
-- pypi: https://files.pythonhosted.org/packages/9e/4c/2f9ac5edbd0e67bf82f5cd04275c4e87cbbf69a78f43e5dcf90c1573d44e/ty-0.0.10-py3-none-manylinux_2_24_x86_64.whl
+- pypi: https://files.pythonhosted.org/packages/ad/01/3a563dba8b1255e474c35e1c3810b7589e81ae8c41df401b6a37c8e2cde9/ty-0.0.11-py3-none-macosx_11_0_arm64.whl
   name: ty
-  version: 0.0.10
-  sha256: e206a23bd887574302138b33383ae1edfcc39d33a06a12a5a00803b3f0287a45
+  version: 0.0.11
+  sha256: 121987c906e02264c3b511b95cb9f8a3cdd66f3283b8bbab678ca3525652e304
   requires_python: '>=3.8'
-- pypi: https://files.pythonhosted.org/packages/e8/cd/9dd49e6d40e54d4b7d563f9e2a432c4ec002c0673a81266e269c4bc194ce/ty-0.0.10-py3-none-macosx_11_0_arm64.whl
+- pypi: https://files.pythonhosted.org/packages/df/04/5a5dfd0aec0ea99ead1e824ee6e347fb623c464da7886aa1e3660fb0f36c/ty-0.0.11-py3-none-win_amd64.whl
   name: ty
-  version: 0.0.10
-  sha256: e4832f8879cb95fc725f7e7fcab4f22be0cf2550f3a50641d5f4409ee04176d4
+  version: 0.0.11
+  sha256: 1bb205db92715d4a13343bfd5b0c59ce8c0ca0daa34fb220ec9120fc66ccbda7
   requires_python: '>=3.8'
 - pypi: https://files.pythonhosted.org/packages/e7/c1/56ef16bf5dcd255155cc736d276efa6ae0a5c26fd685e28f0412a4013c01/types_pytz-2025.2.0.20251108-py3-none-any.whl
   name: types-pytz
@@ -9014,45 +8916,45 @@ packages:
   purls: []
   size: 694692
   timestamp: 1756385147981
-- conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.0.1-py313h7037e92_6.conda
-  sha256: bd1f3d159b204be5aeeb3dd165fad447d3a1c5df75fec64407a68f210a0cb722
-  md5: 1fa8d662361896873a165b051322073e
+- conda: https://conda.anaconda.org/conda-forge/linux-64/ukkonen-1.0.1-py314h9891dd4_6.conda
+  sha256: ef6753f6febaa74d35253e4e0dd09dc9497af8e370893bd97c479f59346daa57
+  md5: 28303a78c48916ab07b95ffdbffdfd6c
   depends:
   - __glibc >=2.17,<3.0.a0
   - cffi
   - libgcc >=14
   - libstdcxx >=14
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/ukkonen?source=hash-mapping
-  size: 14648
-  timestamp: 1761594865380
-- conda: https://conda.anaconda.org/conda-forge/osx-arm64/ukkonen-1.0.1-py313hc50a443_6.conda
-  sha256: 66596db68cd50d61af97b01de4fd6ba5b08c4f5c779c331888196253b4daf353
-  md5: 8e87b6fff522cabf8c02878c24d44312
+  size: 14762
+  timestamp: 1761594960135
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/ukkonen-1.0.1-py314h6b18a25_6.conda
+  sha256: 2ef342cc861c52ec3ac464e89b192a37fd7afd79740b2c0773d2588fd8acff26
+  md5: 452b75f09bc2a4c5eea4044b769bc659
   depends:
   - __osx >=11.0
   - cffi
   - libcxx >=19
-  - python >=3.13,<3.14.0a0
-  - python >=3.13,<3.14.0a0 *_cp313
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
+  - python_abi 3.14.* *_cp314
   license: MIT
   license_family: MIT
   purls:
   - pkg:pypi/ukkonen?source=hash-mapping
-  size: 14535
-  timestamp: 1761595088230
-- conda: https://conda.anaconda.org/conda-forge/win-64/ukkonen-1.0.1-py313hf069bd2_6.conda
-  sha256: f42cd55bd21746274d7074b93b53fb420b4ae0f8f1b6161cb2cc5004c20c7ec7
-  md5: 77444fe3f3004fe52c5ee70626d11d66
+  size: 14635
+  timestamp: 1761595172213
+- conda: https://conda.anaconda.org/conda-forge/win-64/ukkonen-1.0.1-py314h909e829_6.conda
+  sha256: f65b3bf31d22ae37300ed2521352107be830e7c5ba805a4c93e2ce0e0f739078
+  md5: 8528e182a2d9b5d14f0072734a24a6b9
   depends:
   - cffi
-  - python >=3.13,<3.14.0a0
-  - python_abi 3.13.* *_cp313
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
   - ucrt >=10.0.20348.0
   - vc >=14.3,<15
   - vc14_runtime >=14.44.35208
@@ -9060,8 +8962,51 @@ packages:
   license_family: MIT
   purls:
   - pkg:pypi/ukkonen?source=hash-mapping
-  size: 18266
-  timestamp: 1761595426854
+  size: 18357
+  timestamp: 1761595080794
+- conda: https://conda.anaconda.org/conda-forge/linux-64/unicodedata2-17.0.0-py314h5bd0f2a_1.conda
+  sha256: d1dafc15fc5d2b1dd5b0a525e8a815028de20dd53b2c775a1b56e8e4839fb736
+  md5: 58e2ee530005067c5db23f33c6ab43d2
+  depends:
+  - __glibc >=2.17,<3.0.a0
+  - libgcc >=14
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
+  license: Apache-2.0
+  license_family: Apache
+  purls:
+  - pkg:pypi/unicodedata2?source=hash-mapping
+  size: 409745
+  timestamp: 1763055060898
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/unicodedata2-17.0.0-py314h0612a62_1.conda
+  sha256: 48c51dd2ef696f7a1a3635716585a8e383a8c00e719305cfda2b480c36ee1283
+  md5: c673decfe1f120b0717d0aa193b10060
+  depends:
+  - __osx >=11.0
+  - python >=3.14,<3.15.0a0
+  - python >=3.14,<3.15.0a0 *_cp314
+  - python_abi 3.14.* *_cp314
+  license: Apache-2.0
+  license_family: Apache
+  purls:
+  - pkg:pypi/unicodedata2?source=hash-mapping
+  size: 416770
+  timestamp: 1763055099322
+- conda: https://conda.anaconda.org/conda-forge/win-64/unicodedata2-17.0.0-py314h5a2d7ad_1.conda
+  sha256: 47e061aec1487519c398e1c999ac3680f068f9e1d8574c8b365eac4787773250
+  md5: 1f90bb13fa5ced89ca4dcc0af3bbebf3
+  depends:
+  - python >=3.14,<3.15.0a0
+  - python_abi 3.14.* *_cp314
+  - ucrt >=10.0.20348.0
+  - vc >=14.3,<15
+  - vc14_runtime >=14.44.35208
+  license: Apache-2.0
+  license_family: Apache
+  purls:
+  - pkg:pypi/unicodedata2?source=hash-mapping
+  size: 405783
+  timestamp: 1763054877424
 - conda: https://conda.anaconda.org/conda-forge/noarch/uri-template-1.3.0-pyhd8ed1ab_1.conda
   sha256: e0eb6c8daf892b3056f08416a96d68b0a358b7c46b99c8a50481b22631a4dfc0
   md5: e7cb0f5745e4c5035a460248334af7eb
@@ -9125,12 +9070,12 @@ packages:
   purls: []
   size: 115235
   timestamp: 1767320173250
-- conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.35.4-pyhd8ed1ab_0.conda
-  sha256: 77193c99c6626c58446168d3700f9643d8c0dab1f6deb6b9dd039e6872781bfb
-  md5: cfccfd4e8d9de82ed75c8e2c91cab375
+- conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-20.36.0-pyhd8ed1ab_0.conda
+  sha256: cbb40ae88ccc72e95ce00911a73d9175eead4fb4e74925b0e9557bb60737317e
+  md5: c9a9b6e144b880308f5eedc905fe503d
   depends:
   - distlib >=0.3.7,<1
-  - filelock >=3.12.2,<4
+  - filelock >=3.20.1,<4
   - platformdirs >=3.9.1,<5
   - python >=3.10
   - typing_extensions >=4.13.2
@@ -9138,8 +9083,8 @@ packages:
   license_family: MIT
   purls:
   - pkg:pypi/virtualenv?source=hash-mapping
-  size: 4401341
-  timestamp: 1761726489722
+  size: 4403353
+  timestamp: 1767880093070
 - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda
   sha256: e311b64e46c6739e2a35ab8582c20fa30eb608da130625ed379f4467219d4813
   md5: 7e1e5ff31239f9cd5855714df8a3783d
diff --git a/pyproject.toml b/pyproject.toml
index 0b00efd..eac78f9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,9 +32,10 @@ classifiers = [
     "Programming Language :: Python :: 3 :: Only",
     "Topic :: Scientific/Engineering",
 ]
-requires-python = ">=3.13,<3.14"
+requires-python = ">=3.14,<3.15"
 dependencies = [
     "dags",
+    "frozendict",
     "jax>=0.8",
     "numpy",
     "pandas",
@@ -149,7 +150,7 @@ types-PyYAML = "*"
 types-pytz = "*"
 
 [tool.pixi.feature.ty.tasks]
-ty = "ty check"
+ty = "ty check src tests docs"
 
 # Environments
 # --------------------------------------------------------------------------------------
@@ -165,120 +166,51 @@ ty = {features = ["test", "ty"], solve-group = "default"}
 # ======================================================================================
 
 [tool.ruff]
-target-version = "py313"
+target-version = "py314"
 fix = true
 line-length = 88
-
+unsafe-fixes = false
 
 [tool.ruff.lint]
 select = ["ALL"]
 extend-ignore = [
-    # missing type annotation
-    "ANN001",
-
-    # missing type annotation for `*args`
-    "ANN002",
-
-    # missing type annotation for `**kwargs`
-    "ANN003",
-
-    # missing return type annotation for public function
-    "ANN201",
-
-    # missing return type annotation for private function
-    "ANN202",
-
-    # No explicit `stacklevel` keyword argument found
-    "B028",
-
-    # In conflict with formatter
-    "COM812",
-
-    # Missing docstring in public module
-    "D100",
-
-    # missing docstring in public function
-    "D103",
-
-    # missing docstring in public package
-    "D104",
-
-    # exception must not use a string literal
-    "EM101",
-
-    # exception must not use an f-string literal
-    "EM102",
-
-    # Boolean-typed positional arguments.
-    "FBT001",
-
-    # Boolean default positional argument in function definition
-    "FBT002",
-
-    # line contains a todo
-    "FIX002",
-
-    # In conflict with formatter
-    "ISC001",
-
-    # Leave Numpy's legacy RNG
-    "NPY002",
-
-    # array.at is perfectly valid Jax, but linter thinks it's Pandas...
-    "PD008",
-
-    # pd.merge is fine
-    "PD015",
-
-    # Many suggestions to use list comprehension are not helpful
-    "PERF401",
-
-    # Magic values are fine
-    "PLR2004",
-
-    # Too many arguments to function call
-    "PLR0913",
-
-    # Assignment before return statement is fine.
-    "RET504",
-
-    # use of `assert` detected
-    "S101",
-
-    # `pickle` module is unsafe
-    "S301",
-
-    # Private member accessed: `_stochastic_info`
-    "SLF001",
-
-    # long messages outside the exception class
-    "TRY003",
+    "ANN401", # Dynamically typed expressions (typing.Any) are disallowed - too strict
+    "COM812", # In conflict with formatter
+    "EM101", # exception must not use a string literal
+    "EM102", # exception must not use an f-string literal
+    "FBT002", # Boolean default positional argument in function definition
+    "ISC001", # In conflict with formatter
+    "NPY002", # Leave Numpy's legacy RNG
+    "PD015", # pd.merge is fine
+    "PERF401", # Many suggestions to use list comprehension are not helpful
+    "PLR0913", # Too many arguments to function call
+    "PLR2004", # Magic values are fine
+    "S301", # `pickle` module is unsafe
+    "TRY003", # long messages outside the exception class
 ]
 
 [tool.ruff.lint.per-file-ignores]
-"src/skillmodels/constraints.py" = ["D417"]
-"src/skillmodels/decorators.py" = ["D417"]
-"src/skillmodels/kalman_filters.py" = ["D417"]
-"src/skillmodels/likelihood_function.py" = ["D417"]
-"src/skillmodels/likelihood_function_debug.py" = ["D417"]
-"src/skillmodels/params_index.py" = ["D417"]
-"src/skillmodels/parse_params.py" = ["D417"]
-"src/skillmodels/process_data.py" = ["D417"]
-"src/skillmodels/simulate_data.py" = ["D417"]
-"src/skillmodels/visualize_*.py" = ["BLE001", "D417"]
-"src/skillmodels/*_heatmap*.py" = ["D417"]
-"**/*.ipynb" = ["B018", "T201", "E402", "PLR2004", "INP001", "PTH100"]
-"docs/**/*" = ["A001", "ERA001", "INP001", "PTH100", "PTH123", "S506"]
+"src/skillmodels/types.py" = ["TC"]  # Dataclasses need types at runtime
+"src/skillmodels/visualize_*.py" = ["BLE001"]
+"**/*.ipynb" = [
+    "B018", # Seemingly useless expression for printing.
+    "T201", # Printing is fine here.
+    "INP001", # No need for a namespace.
+]
+"docs/source/conf.py" = [
+    "ERA001", # Lots of erased code
+    "INP001", # No need for a namespace.
+]
 "tests/*" = [
-    "ARG001",
-    "E712",
-    "FBT003",
-    "INP001",
-    "PD002",
-    "PT011",
-    "NPY002",
-    "PTH123",
-    "S506"
+    "ANN",  # No type annotations needed for tests
+    "ARG001", # Unused arguments are common in fixture-heavy tests
+    "D100", # No module docstrings needed for tests
+    "D103", # No function docstrings needed for tests
+    "E712", # Comparison to True/False using == might be necessary for arrays.
+    "FBT003", # Boolean positional values are common in test setup
+    "INP001",  # No need for a namespace.
+    "PT011", # Broad pytest.raises() blocks are okay
+    "S101", # use of `assert` detected
 ]
 
 [tool.ruff.lint.pydocstyle]
diff --git a/src/skillmodels/__init__.py b/src/skillmodels/__init__.py
index e55229a..901e313 100644
--- a/src/skillmodels/__init__.py
+++ b/src/skillmodels/__init__.py
@@ -1,3 +1,5 @@
+"""Skillmodels: A Python package for estimating latent factor models."""
+
 import contextlib
 
 try:
diff --git a/src/skillmodels/check_model.py b/src/skillmodels/check_model.py
index 41dc1e4..28a618b 100644
--- a/src/skillmodels/check_model.py
+++ b/src/skillmodels/check_model.py
@@ -1,7 +1,21 @@
+"""Functions to validate model specifications."""
+
+from typing import TYPE_CHECKING
+
 import numpy as np
 
+if TYPE_CHECKING:
+    from skillmodels.types import Anchoring, Dimensions, Labels
 
-def check_model(model_dict, labels, dimensions, anchoring, has_endogenous_factors):
+
+def check_model(
+    model_dict: dict,
+    labels: Labels,
+    dimensions: Dimensions,
+    anchoring: Anchoring,
+    *,
+    has_endogenous_factors: bool,
+) -> None:
     """Check consistency and validity of the model specification.
 
     labels, dimensions and anchoring information are done before the model checking
@@ -12,27 +26,24 @@ def check_model(model_dict, labels, dimensions, anchoring, has_endogenous_factor
     that the assumptions we make during the processing are fulfilled.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`
-        dimensions (dict): Dimensional information like n_states, n_periods, n_controls,
-            n_mixtures. See :ref:`dimensions`.
-        labels (dict): Dict of lists with labels for the model quantities like
-            factors, periods, controls, stagemap and stages. See :ref:`labels`
-        anchoring (dict): Dictionary with information about anchoring.
-            See :ref:`anchoring`
-        has_endogenous_factors (bool): Whether the model has any endogenous factors
+        model_dict: The model specification. See: :ref:`model_specs`
+        dimensions: Dimensional information.
+        labels: Labels for model quantities.
+        anchoring: Information about anchoring.
+        has_endogenous_factors: Whether the model has any endogenous factors
 
     Raises:
         ValueError
 
     """
     report = check_stagemap(
-        stagemap=labels["aug_stagemap"],
-        stages=labels["aug_stages"],
-        n_periods=dimensions["n_aug_periods"],
+        stagemap=labels.aug_stagemap,
+        stages=labels.aug_stages,
+        n_periods=dimensions.n_aug_periods,
         is_augmented=has_endogenous_factors,
     )
     report += _check_anchoring(anchoring)
-    invalid_measurements = _check_measurements(model_dict, labels["latent_factors"])
+    invalid_measurements = _check_measurements(model_dict, labels.latent_factors)
     if invalid_measurements:
         report += invalid_measurements
     elif has_endogenous_factors:
@@ -40,22 +51,30 @@ def check_model(model_dict, labels, dimensions, anchoring, has_endogenous_factor
         report += _check_no_overlap_in_measurements_of_states_and_inv(
             model_dict, labels
         )
-    report += _check_normalizations(model_dict, labels["latent_factors"])
+    report += _check_normalizations(model_dict, labels.latent_factors)
 
     report = "\n".join(report)
     if report != "":
         raise ValueError(f"Invalid model specification: {report}")
 
 
-def check_stagemap(stagemap, stages, n_periods, is_augmented):
-    report = []
+def check_stagemap(
+    stagemap: tuple[int, ...],
+    stages: tuple[int, ...] | list[int],
+    n_periods: int,
+    *,
+    is_augmented: bool,
+) -> list[str]:
+    """Validate the stagemap configuration against model dimensions."""
+    report: list[str] = []
     step_size = 2 if is_augmented else 1
     if len(stagemap) != n_periods - step_size:
         report.append(
             f"The stagemap needs to be of length n_periods - {step_size}. "
             f" n_periods is {n_periods}, the stagemap has length {len(stagemap)}.",
         )
-    if stages != list(range(len(stages))):
+    # Convert to list for comparison (stages may be a tuple from dataclass)
+    if list(stages) != list(range(len(stages))):
         report.append("Stages need to be integers, start at zero and increase by 1.")
 
     # Hijacking the stagemap for endogenous factors leads to interleaved elements.
@@ -68,29 +87,32 @@ def check_stagemap(stagemap, stages, n_periods, is_augmented):
     return report
 
 
-def _check_anchoring(anchoring):
+def _check_anchoring(anchoring: Anchoring) -> list[str]:
     report = []
-    if not isinstance(anchoring["anchoring"], bool):
-        report.append("anchoring['anchoring'] must be a bool.")
-    if not isinstance(anchoring["outcomes"], dict):
-        report.append("anchoring['outcomes'] must be a dict")
+    if not isinstance(anchoring.anchoring, bool):
+        report.append("anchoring.anchoring must be a bool.")
+    if not isinstance(anchoring.outcomes, dict):
+        report.append("anchoring.outcomes must be a dict")
     else:
-        variables = list(anchoring["outcomes"].values())
+        variables = list(anchoring.outcomes.values())
         for var in variables:
             if not isinstance(var, str | int | tuple):
                 report.append("Outcomes variables have to be valid variable names.")
 
-    if not isinstance(anchoring["free_controls"], bool):
-        report.append("anchoring['use_controls'] must be a bool")
-    if not isinstance(anchoring["free_constant"], bool):
-        report.append("anchoring['use_constant'] must be a bool.")
-    if not isinstance(anchoring["free_loadings"], bool):
-        report.append("anchoring['free_loadings'] must be a bool.")
+    if not isinstance(anchoring.free_controls, bool):
+        report.append("anchoring.free_controls must be a bool")
+    if not isinstance(anchoring.free_constant, bool):
+        report.append("anchoring.free_constant must be a bool.")
+    if not isinstance(anchoring.free_loadings, bool):
+        report.append("anchoring.free_loadings must be a bool.")
     return report
 
 
-def _check_measurements(model_dict, factors):
-    report = []
+def _check_measurements(
+    model_dict: dict,
+    factors: tuple[str, ...],
+) -> list[str]:
+    report: list[str] = []
     for factor in factors:
         candidate = model_dict["factors"][factor]["measurements"]
         if not _is_list_of(candidate, list):
@@ -108,11 +130,13 @@ def _check_measurements(model_dict, factors):
     return report
 
 
-def _check_no_overlap_in_measurements_of_states_and_inv(model_dict, labels):
+def _check_no_overlap_in_measurements_of_states_and_inv(
+    model_dict: dict, labels: Labels
+) -> list[str]:
     report = []
-    for period in labels["periods"]:
+    for period in labels.periods:
         meas = {}
-        for factor in labels["latent_factors"]:
+        for factor in labels.latent_factors:
             props = model_dict["factors"][factor]
             if props.get("is_endogenous", False):
                 meas["endogenous_factors"] = set(props["measurements"][period])
@@ -126,8 +150,11 @@ def _check_no_overlap_in_measurements_of_states_and_inv(model_dict, labels):
     return report
 
 
-def _check_normalizations(model_dict, factors):
-    report = []
+def _check_normalizations(
+    model_dict: dict,
+    factors: tuple[str, ...],
+) -> list[str]:
+    report: list[str] = []
     for factor in factors:
         norminfo = model_dict["factors"][factor].get("normalizations", {})
         for norm_type in ["loadings", "intercepts"]:
@@ -152,8 +179,12 @@ def _check_normalizations(model_dict, factors):
     return report
 
 
-def _check_normalized_variables_are_present(list_of_normdicts, model_dict, factor):
-    report = []
+def _check_normalized_variables_are_present(
+    list_of_normdicts: list[dict],
+    model_dict: dict,
+    factor: str,
+) -> list[str]:
+    report: list[str] = []
     for period, norm_dict in enumerate(list_of_normdicts):
         for var in norm_dict:
             if var not in model_dict["factors"][factor]["measurements"][period]:
@@ -166,8 +197,11 @@ def _check_normalized_variables_are_present(list_of_normdicts, model_dict, facto
     return report
 
 
-def _check_loadings_are_not_normalized_to_zero(list_of_normdicts, factor):
-    report = []
+def _check_loadings_are_not_normalized_to_zero(
+    list_of_normdicts: list[dict],
+    factor: str,
+) -> list[str]:
+    report: list[str] = []
     for period, norm_dict in enumerate(list_of_normdicts):
         for var, val in norm_dict.items():
             if val == 0:
@@ -178,7 +212,7 @@ def _check_loadings_are_not_normalized_to_zero(list_of_normdicts, factor):
     return report
 
 
-def _is_list_of(candidate, type_):
+def _is_list_of(candidate: object, type_: type) -> bool:
     """Check if candidate is a list that only contains elements of type.
 
     Note that this is always falls if candidate is not a list and always true if
diff --git a/src/skillmodels/clipping.py b/src/skillmodels/clipping.py
index 6c32799..ce8dd00 100644
--- a/src/skillmodels/clipping.py
+++ b/src/skillmodels/clipping.py
@@ -1,8 +1,17 @@
+"""Soft clipping utilities for constraining values to bounded ranges."""
+
 import jax
 import jax.numpy as jnp
+from jax import Array
 
 
-def soft_clipping(arr, lower=None, upper=None, lower_hardness=1, upper_hardness=1):
+def soft_clipping(
+    arr: Array,
+    lower: float | None = None,
+    upper: float | None = None,
+    lower_hardness: float = 1,
+    upper_hardness: float = 1,
+) -> Array:
     """Clip values in an array elementwise using a soft maximum to avoid kinks.
 
     Clipping from below is taking a maximum between two values. Clipping
@@ -19,14 +28,13 @@ def soft_clipping(arr, lower=None, upper=None, lower_hardness=1, upper_hardness=
     ``scipy.special.logsumexp``. ``scipy.special.softmax`` is the gradient of
     ``scipy.special.logsumexp``.
 
-
     Args:
-        arr (jax.numpy.array): Array that is clipped elementwise.
-        lower (float): The value at which the array is clipped from below.
-        upper (float): The value at which the array is clipped from above.
-        lower_hardness (float): Scaling factor that is applied inside the soft maximum.
+        arr: Array that is clipped elementwise.
+        lower: The value at which the array is clipped from below.
+        upper: The value at which the array is clipped from above.
+        lower_hardness: Scaling factor that is applied inside the soft maximum.
             High values imply a closer approximation of the real maximum.
-        upper_hardness (float): Scaling factor that is applied inside the soft maximum.
+        upper_hardness: Scaling factor that is applied inside the soft maximum.
             High values imply a closer approximation of the real maximum.
 
     """
diff --git a/src/skillmodels/config.py b/src/skillmodels/config.py
index ddc66a4..cd7eb32 100644
--- a/src/skillmodels/config.py
+++ b/src/skillmodels/config.py
@@ -1,3 +1,8 @@
+"""Configuration constants and paths for skillmodels."""
+
 from pathlib import Path
 
-TEST_DIR = Path(__file__).resolve().parent / "tests"
+TEST_DATA_DIR = Path(__file__).resolve().parent / "test_data"
+REGRESSION_VAULT = (
+    Path(__file__).resolve().parent.parent.parent / "tests" / "regression_vault"
+)
diff --git a/src/skillmodels/constraints.py b/src/skillmodels/constraints.py
index 2fb1c70..308ed4d 100644
--- a/src/skillmodels/constraints.py
+++ b/src/skillmodels/constraints.py
@@ -3,38 +3,43 @@
 import functools
 import warnings
 from dataclasses import dataclass
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import optimagic as om
-import pandas as pd
 
 import skillmodels.transition_functions as t_f_module
+from skillmodels.types import MeasurementType
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+    from skillmodels.types import Anchoring, Dimensions, EndogenousFactorsInfo, Labels
 
 
 def get_constraints_dicts(
-    dimensions,
-    labels,
-    anchoring_info,
-    update_info,
-    normalizations,
-    endogenous_factors_info,
+    dimensions: Dimensions,
+    labels: Labels,
+    anchoring_info: Anchoring,
+    update_info: pd.DataFrame,
+    normalizations: dict[str, dict[str, list]],
+    endogenous_factors_info: EndogenousFactorsInfo,
 ) -> list[dict]:
     """Generate constraints implied by the model specification.
 
     The result can easily be converted to optimagic-style constraints.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`
-        dimensions (dict): Dimensional information like n_states, n_periods, n_controls,
+        dimensions: Dimensional information like n_states, n_periods, n_controls,
             n_mixtures. See :ref:`dimensions`.
-        labels (dict): Dict of lists with labels for the model quantities like
+        labels: Dict of lists with labels for the model quantities like
             factors, periods, controls, stagemap and stages. See :ref:`labels`
-        anchoring (dict): Information about anchoring. See :ref:`anchoring`
-        update_info (pandas.DataFrame): DataFrame with one row per Kalman update needed
+        anchoring_info: Information about anchoring. See :ref:`anchoring`
+        update_info: DataFrame with one row per Kalman update needed
             in the likelihood function. See :ref:`update_info`.
-        normalizations (dict): Nested dictionary with information on normalized factor
+        normalizations: Nested dictionary with information on normalized factor
             loadings and intercepts for each factor. See :ref:`normalizations`.
+        endogenous_factors_info: Information about endogenous factors in the model.
 
     Returns:
         A list of constraints dictionaries with entries:
@@ -49,26 +54,26 @@ def get_constraints_dicts(
     constraints_dicts = []
 
     constraints_dicts += _get_normalization_constraints(
-        normalizations, labels["latent_factors"]
+        normalizations, labels.latent_factors
     )
-    constraints_dicts += _get_mixture_weights_constraints(dimensions["n_mixtures"])
+    constraints_dicts += _get_mixture_weights_constraints(dimensions.n_mixtures)
     constraints_dicts += _get_stage_constraints(
-        stagemap=labels["aug_stagemap"],
-        stages=labels["aug_stages"],
+        stagemap=labels.aug_stagemap,
+        stages=labels.aug_stages,
     )
     constraints_dicts += _get_constant_factors_constraints(labels=labels)
     constraints_dicts += _get_initial_states_constraints(
-        n_mixtures=dimensions["n_mixtures"],
-        factors=labels["latent_factors"],
+        n_mixtures=dimensions.n_mixtures,
+        factors=labels.latent_factors,
     )
     constraints_dicts += _get_transition_constraints(labels=labels)
     constraints_dicts += _get_anchoring_constraints(
         update_info=update_info,
-        controls=labels["controls"],
+        controls=labels.controls,
         anchoring_info=anchoring_info,
-        periods=labels["aug_periods"],
+        periods=labels.aug_periods,
     )
-    if endogenous_factors_info["has_endogenous_factors"]:
+    if endogenous_factors_info.has_endogenous_factors:
         constraints_dicts += _get_constraints_for_augmented_periods(
             labels=labels,
             endogenous_factors_info=endogenous_factors_info,
@@ -116,7 +121,7 @@ def add_bounds(params: pd.DataFrame, bounds_distance: float) -> pd.DataFrame:
     return df
 
 
-def _is_diagonal_entry(ind_tup):
+def _is_diagonal_entry(ind_tup: tuple[str, ...]) -> bool:
     name2 = ind_tup[-1]
     middle_pos = int(len(name2) // 2)
     if (
@@ -130,12 +135,16 @@ def _is_diagonal_entry(ind_tup):
     return is_diag
 
 
-def _get_normalization_constraints(normalizations, factors) -> list[dict]:
+def _get_normalization_constraints(
+    normalizations: dict[str, dict[str, list]],
+    factors: tuple[str, ...],
+) -> list[dict]:
     """List of constraints to enforce normalizations.
 
     Args:
-        normalizations (dict): Nested dictionary with information on normalized factor
-        loadings and intercepts for each factor. See :ref:`normalizations`.
+        normalizations: Nested dictionary with information on normalized factor
+            loadings and intercepts for each factor. See :ref:`normalizations`.
+        factors: Tuple of factor names to process.
 
     Returns:
         constraints_dicts
@@ -171,7 +180,7 @@ def _get_normalization_constraints(normalizations, factors) -> list[dict]:
     return constraints_dicts
 
 
-def _get_mixture_weights_constraints(n_mixtures) -> list[dict]:
+def _get_mixture_weights_constraints(n_mixtures: int) -> list[dict]:
     """Constrain mixture weights to be between 0 and 1 and sum to 1."""
     if n_mixtures == 1:
         msg = "Set the mixture weight to 1 if there is only one mixture element."
@@ -191,12 +200,15 @@ def _get_mixture_weights_constraints(n_mixtures) -> list[dict]:
     return constraints_dicts
 
 
-def _get_stage_constraints(stagemap, stages) -> list[dict]:
+def _get_stage_constraints(
+    stagemap: tuple[int, ...],
+    stages: tuple[int, ...],
+) -> list[dict]:
     """Equality constraints for transition and shock parameters within stages.
 
     Args:
-        stagemap (list): map aug_periods to aug_stages
-        stages (list): aug_stages
+        stagemap: map aug_periods to aug_stages
+        stages: aug_stages
     Returns:
         constraints_dicts
 
@@ -232,11 +244,11 @@ def _get_stage_constraints(stagemap, stages) -> list[dict]:
     return constraints_dicts
 
 
-def _get_constant_factors_constraints(labels) -> list[dict]:
+def _get_constant_factors_constraints(labels: Labels) -> list[dict]:
     """Fix shock variances of constant factors to `bounds_distance`.
 
     Args:
-        labels (dict): Dict of lists with labels for the model quantities like
+        labels: Dict of lists with labels for the model quantities like
             factors, periods, controls, stagemap and stages. See :ref:`labels`
 
     Returns:
@@ -244,10 +256,10 @@ def _get_constant_factors_constraints(labels) -> list[dict]:
 
     """
     constraints_dicts = []
-    for f, factor in enumerate(labels["latent_factors"]):
-        if labels["transition_names"][f] == "constant":
+    for f, factor in enumerate(labels.latent_factors):
+        if labels.transition_names[f] == "constant":
             msg = f"This constraint was generated because {factor} is constant."
-            for aug_period in labels["aug_periods"][:-1]:
+            for aug_period in labels.aug_periods[:-1]:
                 constraints_dicts.append(
                     {
                         "loc": ("shock_sds", aug_period, factor, "-"),
@@ -259,14 +271,17 @@ def _get_constant_factors_constraints(labels) -> list[dict]:
     return constraints_dicts
 
 
-def _get_initial_states_constraints(n_mixtures, factors) -> list[dict]:
+def _get_initial_states_constraints(
+    n_mixtures: int,
+    factors: tuple[str, ...],
+) -> list[dict]:
     """Enforce that the x values of the first factor are increasing.
 
     Otherwise the model would only be identified up to the order of the start factors.
 
     Args:
-        n_mixtures (int): number of elements in the mixture of normal of the factors.
-        factors (list): the latent factors of the model
+        n_mixtures: number of elements in the mixture of normal of the factors.
+        factors: the latent factors of the model
 
     Returns:
         constraints_dicts
@@ -290,11 +305,11 @@ def _get_initial_states_constraints(n_mixtures, factors) -> list[dict]:
     return constraints_dicts
 
 
-def _get_transition_constraints(labels) -> list[dict]:
+def _get_transition_constraints(labels: Labels) -> list[dict]:
     """Collect possible constraints on transition parameters.
 
     Args:
-        labels (dict): Dict of lists with labels for the model quantities like
+        labels: Dict of lists with labels for the model quantities like
             factors, periods, controls, stagemap and stages. See :ref:`labels`
 
     Returns:
@@ -302,14 +317,14 @@ def _get_transition_constraints(labels) -> list[dict]:
 
     """
     constraints_dicts = []
-    for f, factor in enumerate(labels["latent_factors"]):
-        tname = labels["transition_names"][f]
+    for f, factor in enumerate(labels.latent_factors):
+        tname = labels.transition_names[f]
         msg = f"This constraint is inherent to the {tname} production function."
-        for aug_period in labels["aug_periods"][:-1]:
+        for aug_period in labels.aug_periods[:-1]:
             funcname = f"constraints_{tname}"
             if func := getattr(t_f_module, funcname, False):
                 c = func(  # ty: ignore[call-non-callable]
-                    factor=factor, factors=labels["all_factors"], aug_period=aug_period
+                    factor=factor, factors=labels.all_factors, aug_period=aug_period
                 )
                 if "description" not in c:
                     c["description"] = msg
@@ -318,16 +333,19 @@ def _get_transition_constraints(labels) -> list[dict]:
 
 
 def _get_anchoring_constraints(
-    update_info, controls, anchoring_info, periods
+    update_info: pd.DataFrame,
+    controls: tuple[str, ...],
+    anchoring_info: Anchoring,
+    periods: tuple[int, ...],
 ) -> list[dict]:
     """Constraints on anchoring parameters.
 
     Args:
-        update_info (pandas.DataFrame): DataFrame with one row per Kalman update needed
+        update_info: DataFrame with one row per Kalman update needed
             in the likelihood function. See :ref:`update_info`.
-        controls (list): List of control variables
-        anchoring_info (dict): Information about anchoring. See :ref:`anchoring`
-        periods (list): Period of the model
+        controls: List of control variables
+        anchoring_info: Information about anchoring. See :ref:`anchoring`
+        periods: Period of the model
 
     Returns:
         constraints_dicts
@@ -336,7 +354,7 @@ def _get_anchoring_constraints(
     anchoring_updates = update_info[update_info["purpose"] == "anchoring"].index
 
     constraints_dicts = []
-    if not anchoring_info["free_constant"]:
+    if not anchoring_info.free_constant:
         msg = (
             "This constraint was generated because free_constant in the anchoring "
             "section of the model specification is set to False."
@@ -348,7 +366,7 @@ def _get_anchoring_constraints(
             {"loc": locs, "type": "fixed", "value": 0, "description": msg},
         )
 
-    if not anchoring_info["free_controls"]:
+    if not anchoring_info.free_controls:
         msg = (
             "This constraint was generated because free_controls in the anchoring "
             "section of the model specification is set to False."
@@ -361,15 +379,15 @@ def _get_anchoring_constraints(
             {"loc": ind_tups, "type": "fixed", "value": 0, "description": msg},
         )
 
-    if not anchoring_info["free_loadings"]:
+    if not anchoring_info.free_loadings:
         msg = (
             "This constraint was generated because free_loadings in the anchoring "
             "section of the model specification is set to False."
         )
         ind_tups = []
         for period in periods:
-            for factor in anchoring_info["factors"]:
-                outcome = anchoring_info["outcomes"][factor]
+            for factor in anchoring_info.factors:
+                outcome = anchoring_info.outcomes[factor]  # ty: ignore[invalid-argument-type]
                 meas = f"{outcome}_{factor}"
                 ind_tups.append(("loadings", period, meas, factor))
 
@@ -377,13 +395,12 @@ def _get_anchoring_constraints(
             {"loc": ind_tups, "type": "fixed", "value": 1, "description": msg},
         )
 
-    constraints_dicts = [c for c in constraints_dicts if c["loc"] != []]
-
-    return constraints_dicts
+    return [c for c in constraints_dicts if c["loc"] != []]
 
 
 def _get_constraints_for_augmented_periods(
-    labels, endogenous_factors_info
+    labels: Labels,
+    endogenous_factors_info: EndogenousFactorsInfo,
 ) -> list[dict]:
     """Constraints for augmented periods.
 
@@ -394,30 +411,33 @@ def _get_constraints_for_augmented_periods(
     Both depend on the transition function.
 
     Args:
-        labels (dict): Dict of lists with labels for the model quantities like
+        labels: Dict of lists with labels for the model quantities like
             factors, periods, controls, stagemap and stages. See :ref:`labels`
+        endogenous_factors_info: Information about endogenous factors and their
+            relationship to augmented periods.
 
     Returns:
         constraints_dicts
 
     """
     constraints_dicts = []
-    for f, factor in enumerate(labels["latent_factors"]):
-        tname = labels["transition_names"][f]
+    for f, factor in enumerate(labels.latent_factors):
+        tname = labels.transition_names[f]
         if tname == "constant":
             continue
         # We are restricting transitions and shocks, not measurements. So this might
         # look counterintuitive...
         aug_period_meas_type_to_constrain = (
-            "states"
-            if endogenous_factors_info[factor]["is_state"]
-            else "endogenous_factors"
+            MeasurementType.STATES
+            if endogenous_factors_info.factor_info[factor].is_state  # ty: ignore[invalid-argument-type]
+            else MeasurementType.ENDOGENOUS_FACTORS
+        )
+        aug_period_meas_types = (
+            endogenous_factors_info.aug_periods_to_aug_period_meas_types
         )
         aug_periods_to_constrain = [
             k
-            for k, v in endogenous_factors_info[
-                "aug_periods_to_aug_period_meas_types"
-            ].items()
+            for k, v in aug_period_meas_types.items()
             if v == aug_period_meas_type_to_constrain
         ]
         for aug_period in aug_periods_to_constrain:
@@ -425,14 +445,14 @@ def _get_constraints_for_augmented_periods(
                 constraints_dicts += func(  # ty: ignore[call-non-callable]
                     factor=factor,
                     aug_period=aug_period,
-                    all_factors=labels["all_factors"],
+                    all_factors=labels.all_factors,
                 )
         for aug_period in aug_periods_to_constrain[:-1]:
             constraints_dicts.append(
                 {
                     "loc": ("shock_sds", aug_period, factor, "-"),
                     "type": "fixed",
-                    "value": endogenous_factors_info["bounds_distance"],
+                    "value": endogenous_factors_info.bounds_distance,
                     "description": "Identity constraint.",
                 }
             )
@@ -440,7 +460,7 @@ def _get_constraints_for_augmented_periods(
     return constraints_dicts
 
 
-def _sel(params, loc):
+def _sel(params: pd.DataFrame, loc: Any) -> pd.DataFrame:
     return params.loc[loc]
 
 
@@ -516,7 +536,7 @@ def constraints_dicts_to_om(
     """Convert constraints provided in dictionary form to optimagic constraints.
 
     Args:
-        constraints_dicts (list): see :ref:`get_constraints_dicts`.
+        constraints_dicts: see :ref:`get_constraints_dicts`.
 
     Returns:
         List of optimagic constraints.
@@ -555,8 +575,8 @@ def enforce_fixed_constraints(
     This means that any robust bounds will be overridden for fixed parameters.
 
     Args:
-        params_template (pd.DataFrame): see :ref:`params_df`.
-        constraints_dicts (list): see :ref:`get_constraints_dicts`.
+        params_template: see :ref:`params_df`.
+        constraints_dicts: see :ref:`get_constraints_dicts`.
 
     Returns:
         pd.DataFrame: modified copy of params_template
diff --git a/src/skillmodels/correlation_heatmap.py b/src/skillmodels/correlation_heatmap.py
index 9ea5d12..0286f1b 100644
--- a/src/skillmodels/correlation_heatmap.py
+++ b/src/skillmodels/correlation_heatmap.py
@@ -1,3 +1,7 @@
+"""Functions for creating correlation heatmap visualizations."""
+
+from typing import TYPE_CHECKING, Any
+
 import numpy as np
 import pandas as pd
 from plotly import graph_objects as go
@@ -5,56 +9,64 @@
 from skillmodels.process_data import pre_process_data
 from skillmodels.process_model import process_model
 
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+    from skillmodels.types import ProcessedModel
+
 
 def plot_correlation_heatmap(
-    corr,
-    heatmap_kwargs=None,
-    layout_kwargs=None,
-    rounding=2,
-    zmax=None,
-    zmin=None,
-    zmid=None,
-    colorscale="RdBu_r",
-    show_color_bar=True,
-    show_diagonal=True,
-    show_upper_triangle=True,
-    trim_heatmap=False,
-    annotate=True,
-    annotation_fontsize=13,
-    annotation_text_color="black",
-    annotation_text_angle=0,
-    axes_tick_fontsize=(12, 12),
-    axes_tick_label_angle=(90, 0),
-    axes_tick_label_color=("black", "black"),
-):
+    corr: pd.DataFrame,
+    heatmap_kwargs: dict[str, Any] | None = None,
+    layout_kwargs: dict[str, Any] | None = None,
+    rounding: int = 2,
+    zmax: float | None = None,
+    zmin: float | None = None,
+    zmid: float | None = None,
+    colorscale: str = "RdBu_r",
+    *,
+    show_color_bar: bool = True,
+    show_diagonal: bool = True,
+    show_upper_triangle: bool = True,
+    trim_heatmap: bool = False,
+    annotate: bool = True,
+    annotation_fontsize: int = 13,
+    annotation_text_color: str = "black",
+    annotation_text_angle: float = 0,
+    axes_tick_fontsize: tuple[int, int] = (12, 12),
+    axes_tick_label_angle: tuple[float, float] = (90, 0),
+    axes_tick_label_color: tuple[str, str] = ("black", "black"),
+) -> go.Figure:
     """Plot correlation heatmaps for factor measurements.
 
     Args:
-        corr (DataFrame): Data frame of measurement or factor score correlations.
-        heatmap_kwargs (dct): Dictionary of key word arguments to pass to go.Heatmap ().
+        corr: Data frame of measurement or factor score correlations.
+        heatmap_kwargs: Dictionary of key word arguments to pass to go.Heatmap ().
             If None, the default kwargs defined in the function will be used.
-        layout_kwargs (dct): Dictionary of key word arguments used to update layout of
+        layout_kwargs: Dictionary of key word arguments used to update layout of
             go.Figure object. If None, the default kwargs defined in the function will
             be used. Through layout_kwargs, you can edit figure properties such as
             - template
             - title
             - figsize
-        rounding (int): Number of digits after the decimal point to round the
+        rounding: Number of digits after the decimal point to round the
             correlation values to. Default 2.
-        zmax (float ot NoneType): Upper bound to set on correlation color map. If None,
+        zmax: Upper bound to set on correlation color map. If None,
             is set to maximum absolute correlation value.
-        zmin (float or NoneType): Lower bound to set on correlation color map. If None,
+        zmin: Lower bound to set on correlation color map. If None,
             is set to -zmax.
-        zmid (float or NoneType): Midpoint to set on correlation color map. If None,
+        zmid: Midpoint to set on correlation color map. If None,
             is set to 0.
-        colorscale (str): Name of the color palette to use in the heatmap.
+        colorscale: Name of the color palette to use in the heatmap.
             Default 'RdBu_r'.
-        show_color_bar (bool): A boolean variable for displaying heatmap colorbar.
+        show_color_bar: A boolean variable for displaying heatmap colorbar.
             Default True.
-        show_diagonal (bool): A boolean for displaying the correlations on the diagonal.
+        show_diagonal: A boolean for displaying the correlations on the diagonal.
             Default False.
-        show_upper_triangle (bool): A boolean for displaying upper triangular part
+        show_upper_triangle: A boolean for displaying upper triangular part
             of the correlation heatmap. Default False.
+        trim_heatmap: If True, trim empty rows/columns from the heatmap.
+            Default False.
 
         The following arguments are processed into dictionaries or special plotly
         objects and passed to layout_kwargs. Defining them as additional arguments
@@ -67,38 +79,38 @@ def plot_correlation_heatmap(
         defined in layout_kwargs will overwrite values passed via the individual
         arguments.
 
-        annotate (bool): If True, annotate the heatmap figure with correlation values.
+        annotate: If True, annotate the heatmap figure with correlation values.
             Default False.
-        annotation_font_size (int): Font size of the annotation text. Default 13.
-        annotation_font_color (str): Collor of the annotation text. Default 'black'.
-        annotation_text_angle (float): The angle at which to rotate annotation text.
+        annotation_fontsize: Font size of the annotation text. Default 13.
+        annotation_text_color: Color of the annotation text. Default 'black'.
+        annotation_text_angle: The angle at which to rotate annotation text.
             Default 0.
-        axes_tick_fontsize (list, tuple, other iterable or dict): Fontsize of axes
+        axes_tick_fontsize: Fontsize of axes
             ticks. Default (12,12)
-        axes_tick_label_angle (list, tuple, other iterable or dict): Rotation angles of
+        axes_tick_label_angle: Rotation angles of
             axes tick labels. Default (90,0).
-        axes_tick_label_color (list, tuple, other iterable or dict): Colors of the axes
+        axes_tick_label_color: Colors of the axes
             tick labels. Default ('black', 'black').
 
     Returns:
-        fig (plotly graph object): The figure with correlaiton heatmap.
+        fig: The figure with correlaiton heatmap.
 
     """
     corr = _process_corr_data_for_plotting(
         corr,
         rounding,
-        show_upper_triangle,
-        show_diagonal,
-        trim_heatmap,
+        show_upper_triangle=show_upper_triangle,
+        show_diagonal=show_diagonal,
+        trim_heatmap=trim_heatmap,
     )
     heatmap_kwargs = _get_heatmap_kwargs(
         corr,
         heatmap_kwargs,
         colorscale,
-        show_color_bar,
-        zmax,
-        zmin,
-        zmid,
+        show_color_bar=show_color_bar,
+        zmax=zmax,
+        zmin=zmin,
+        zmid=zmid,
     )
     layout_kwargs = _get_layout_kwargs(
         corr=corr,
@@ -122,25 +134,30 @@ def plot_correlation_heatmap(
     return fig
 
 
-def get_measurements_corr(data, model_dict, factors, periods):
+def get_measurements_corr(
+    data: pd.DataFrame,
+    model_dict: dict,
+    factors: list[str] | tuple[str, ...] | str | None,
+    periods: float | list[int] | None,
+) -> pd.DataFrame:
     """Get data frame with measurement correlations.
 
     Process data to retrieve measurements for each period and calculate correlations
     across period specific measurements.
 
     Args:
-        data (pd.DataFrame): DataFrame with observed measurements.
-        model_dict (dct): Dictionary of model attributes to be passed to process_model
+        data: DataFrame with observed measurements.
+        model_dict: Dictionary of model attributes to be passed to process_model
             and extract measurements for each period.
-        factors (list, str or NoneType): List of factors, to retrieve measurements for.
+        factors: List of factors, to retrieve measurements for.
             If None, then calculate correlations of measurements of all factors.
-        periods (int, float, list or NoneType): If int, the period within which to
+        periods: If int, the period within which to
             calculate measurement correlations. If a list, calculate correlations over
             periods. If None, calculate correlations across all periods. Note: Periods
             refer to originl periods, not the augmented periods.
 
     Returns:
-        corr (DataFrame): DataFrame with measurement correlations.
+        corr: DataFrame with measurement correlations.
 
     """
     data = data.copy(deep=True)
@@ -156,11 +173,15 @@ def get_measurements_corr(data, model_dict, factors, periods):
         latent_factors=latent_factors,
         observed_factors=observed_factors,
     )
-    corr = df.corr()
-    return corr
+    return df.corr()
 
 
-def get_quasi_scores_corr(data, model_dict, factors, periods):
+def get_quasi_scores_corr(
+    data: pd.DataFrame,
+    model_dict: dict,
+    factors: list[str] | tuple[str, ...] | str | None,
+    periods: float | list[int] | None,
+) -> pd.DataFrame:
     """Get data frame with correlations of factor scores.
 
     Process data to retrieve measurements for each period, standardize measurements
@@ -171,17 +192,17 @@ def get_quasi_scores_corr(data, model_dict, factors, periods):
     The calculated scores coincide with factor scores for linear models.
 
     Args:
-        data (pd.DataFrame): DataFrame with observed measurements.
-        model_dict (dct): Dictionary of model attributes to be passed to process_model
+        data: DataFrame with observed measurements.
+        model_dict: Dictionary of model attributes to be passed to process_model
             and extract measurements for each period.
-        factors (list, str or NoneType): List of factors, to retrieve measurements for.
+        factors: List of factors, to retrieve measurements for.
             If None, then calculate correlations of measurements of all factors.
-        periods (int,float, list or NoneType): If int, the period within which to
+        periods: If int, the period within which to
             calculate measurement correlations. If a list, calculate correlations over
             periods. If None, calculate correlations across all periods.
 
     Returns:
-        corr (DataFrame): DataFrame with score correlations.
+        corr: DataFrame with score correlations.
 
     """
     data = data.copy(deep=True)
@@ -197,11 +218,16 @@ def get_quasi_scores_corr(data, model_dict, factors, periods):
         latent_factors=latent_factors,
         observed_factors=observed_factors,
     )
-    corr = df.corr()
-    return corr
+    return df.corr()
 
 
-def get_scores_corr(data, params, model_dict, factors, periods):
+def get_scores_corr(
+    data: pd.DataFrame,
+    params: pd.DataFrame,
+    model_dict: dict,
+    factors: list[str] | tuple[str, ...] | str | None,
+    periods: float | list[int] | None,
+) -> pd.DataFrame:
     """Get data frame with correlations of factor scores.
 
     Process data to retrieve measurements for each period, standardize measurements
@@ -210,18 +236,18 @@ def get_scores_corr(data, params, model_dict, factors, periods):
     scores.
 
     Args:
-        data (pd.DataFrame): DataFrame with observed measurements.
-        params (pd.DataFrame): DataFrame with estimated model parameters
-        model_dict (dct): Dictionary of model attributes to be passed to process_model
+        data: DataFrame with observed measurements.
+        params: DataFrame with estimated model parameters
+        model_dict: Dictionary of model attributes to be passed to process_model
             and extract measurements for each period.
-        factors (list, str or NoneType): List of factors, to retrieve measurements for.
+        factors: List of factors, to retrieve measurements for.
             If None, then calculate correlations of measurements of all factors.
-        periods (int,float, list or NoneType): If int, the period within which to
+        periods: If int, the period within which to
             calculate measurement correlations. If a list, calculate correlations over
             periods. If None, calculate correlations across all periods.
 
     Returns:
-        corr (DataFrame): DataFrame with score correlations.
+        corr: DataFrame with score correlations.
 
     """
     data = data.copy(deep=True)
@@ -238,19 +264,21 @@ def get_scores_corr(data, params, model_dict, factors, periods):
         latent_factors=latent_factors,
         observed_factors=observed_factors,
     )
-    corr = df.corr()
-    return corr
+    return df.corr()
 
 
 def _process_corr_data_for_plotting(
-    corr,
-    rounding,
-    show_upper_triangle,
-    show_diagonal,
-    trim_heatmap,
-):
+    corr: pd.DataFrame,
+    rounding: int,
+    *,
+    show_upper_triangle: bool,
+    show_diagonal: bool,
+    trim_heatmap: bool,
+) -> pd.DataFrame:
     """Apply mask and rounding to correlation DataFrame."""
-    mask = _get_mask(corr, show_upper_triangle, show_diagonal)
+    mask = _get_mask(
+        corr, show_upper_triangle=show_upper_triangle, show_diagonal=show_diagonal
+    )
     corr = corr.where(mask).round(rounding)
     if trim_heatmap:
         keeprows = mask.any(axis=1) & corr.notna().any(axis="columns").to_numpy()
@@ -262,7 +290,12 @@ def _process_corr_data_for_plotting(
     return corr
 
 
-def _get_mask(corr, show_upper_triangle, show_diagonal):
+def _get_mask(
+    corr: pd.DataFrame,
+    *,
+    show_upper_triangle: bool,
+    show_diagonal: bool,
+) -> NDArray[np.bool_]:
     """Get array to mask the correlation DataFrame."""
     mask = np.zeros_like(corr, dtype=bool)
     mask[np.tril_indices_from(mask, k=-1)] = True
@@ -273,15 +306,15 @@ def _get_mask(corr, show_upper_triangle, show_diagonal):
     return mask
 
 
-def _get_update_info_for_periods(model):
+def _get_update_info_for_periods(model: ProcessedModel) -> pd.DataFrame:
     """Return update_info with user-provided periods instead of augmented periods."""
-    update_info = model["update_info"].copy()
+    update_info = model.update_info.copy()
 
     # Replace period level with user-provided period using set_codes
     period_values = update_info.index.get_level_values("aug_period").map(
-        model["labels"]["aug_periods_to_periods"]
+        model.labels.aug_periods_to_periods
     )
-    update_info.index = update_info.index.set_codes(period_values, level="aug_period")
+    update_info.index = update_info.index.set_codes(period_values, level="aug_period")  # ty: ignore[unresolved-attribute]
     update_info.index = update_info.index.set_names(["period", "variable"])
 
     # Group by period and variable, apply OR logic for boolean columns
@@ -293,26 +326,30 @@ def _get_update_info_for_periods(model):
 
 
 def _get_measurement_data(
-    data, update_info_by_period, periods, latent_factors, observed_factors
-):
+    data: pd.DataFrame,
+    update_info_by_period: pd.DataFrame,
+    periods: list[int],
+    latent_factors: list[str] | tuple[str, ...],
+    observed_factors: list[str] | tuple[str, ...],
+) -> pd.DataFrame:
     """Get data frame with factor measurements in each period, in wide format.
 
     For each factor, retrieve the data on measurements in each period and stack
     the data columns into a data frame.
 
     Args:
-        data (pd.DataFrame): Data with observable variables.
-        update_info (pd.DataFrame): DataFrame with information on measurements
+        data: Data with observable variables.
+        update_info_by_period: DataFrame with information on measurements
             for each factor in each model period.
-        periods (list): The list of periods that correlations are
+        periods: The list of periods that correlations are
             calculated for.
-        latent_factors (list): List of latent factors the measurements of which
+        latent_factors: List of latent factors the measurements of which
             correlations are calculated for.
-        observed_factors (list): List of observed factors the measurements of which
+        observed_factors: List of observed factors the measurements of which
             correlations are calculated for.
 
     Returns:
-        df (pd.DataFrame): Processed DataFrame to calculate correlations over.
+        df: Processed DataFrame to calculate correlations over.
 
     """
     if len(periods) == 1:
@@ -336,26 +373,26 @@ def _get_measurement_data(
 
 
 def _get_measurement_data_for_single_period(
-    data,
-    update_info_by_period,
-    period,
-    latent_factors,
-    observed_factors,
-):
+    data: pd.DataFrame,
+    update_info_by_period: pd.DataFrame,
+    period: int,
+    latent_factors: list[str] | tuple[str, ...],
+    observed_factors: list[str] | tuple[str, ...],
+) -> pd.DataFrame:
     """Extract measurements of factors for the given period.
 
     Args:
-        data (pd.DataFrame): Data with observable variables.
-        update_info (pd.DataFrame): DataFrame with information on measurements
+        data: Data with observable variables.
+        update_info_by_period: DataFrame with information on measurements
             for each factor in each model period.
-        periods (int or float): The period to extract measurements for.
-        latent_factors (list): List of latent factors the measurements of which
+        period: The period to extract measurements for.
+        latent_factors: List of latent factors the measurements of which
             correlations are calculated for.
-        observed_factors (list): List of observed factors the measurements of which
+        observed_factors: List of observed factors the measurements of which
             correlations are calculated for.
 
     Returns:
-        df (pd.DataFrame): DataFrame with measurements of factors for period 'period'.
+        df: DataFrame with measurements of factors for period 'period'.
 
     """
     period_info = update_info_by_period.loc[period].reset_index()
@@ -367,31 +404,30 @@ def _get_measurement_data_for_single_period(
         )["variable"].to_list()
     for fac in observed_factors:
         measurements.append(fac)
-    df = data.query(f"{update_info_by_period.index.names[0]}=={period}")[measurements]
-    return df
+    return data.query(f"{update_info_by_period.index.names[0]}=={period}")[measurements]
 
 
 def _get_measurement_data_for_multiple_periods(
-    data,
-    update_info_by_period,
-    periods,
-    latent_factors,
-    observed_factors,
-):
+    data: pd.DataFrame,
+    update_info_by_period: pd.DataFrame,
+    periods: list[int],
+    latent_factors: list[str] | tuple[str, ...],
+    observed_factors: list[str] | tuple[str, ...],
+) -> pd.DataFrame:
     """Extract measurements for factors for given periods.
 
     Args:
-        data (pd.DataFrame): Data with observable variables.
-        update_info_by_period (pd.DataFrame): DataFrame with information on measurements
+        data: Data with observable variables.
+        update_info_by_period: DataFrame with information on measurements
             for each factor in each user-provided period.
-        periods (list): The periods to extract measurements for.
-        latent_factors (list): List of latent factors the measurements of which
+        periods: The periods to extract measurements for.
+        latent_factors: List of latent factors the measurements of which
             correlations are calculated for.
-        observed_factors (list): List of observed factors the measurements of which
+        observed_factors: List of observed factors the measurements of which
             correlations are calculated for.
 
     Returns:
-        df (pd.DataFrame): DataFrame with measurements of factors in each period as
+        df: DataFrame with measurements of factors in each period as
             columns.
 
     """
@@ -408,17 +444,16 @@ def _get_measurement_data_for_multiple_periods(
             .add_suffix(f", {period}")
             .reset_index(drop=True),
         )
-    df = pd.concat(to_concat, axis=1)
-    return df
+    return pd.concat(to_concat, axis=1)
 
 
 def _get_quasi_factor_scores_data(
-    data,
-    update_info_by_period,
-    periods,
-    latent_factors,
-    observed_factors,
-):
+    data: pd.DataFrame,
+    update_info_by_period: pd.DataFrame,
+    periods: list[int],
+    latent_factors: list[str] | tuple[str, ...],
+    observed_factors: list[str] | tuple[str, ...],
+) -> pd.DataFrame:
     """Get data frame with summary information on factor measurements in each period.
 
     In each period, standardize factor measurements to zero mean and unit standard
@@ -427,18 +462,18 @@ def _get_quasi_factor_scores_data(
     models.
 
     Args:
-        data (pd.DataFrame): Data with observable variables.
-        update_info (pd.DataFrame): DataFrame with information on measurements
+        data: Data with observable variables.
+        update_info_by_period: DataFrame with information on measurements
             for each factor in each model period.
-        periods (list): The list of periods that correlations are
+        periods: The list of periods that correlations are
             calculated for.
-        latent_factors (list): List of latent factors the scores of which
+        latent_factors: List of latent factors the scores of which
             correlations are calculated for.
-        observed_factors (list): List of observed factors the scores of which
+        observed_factors: List of observed factors the scores of which
             correlations are calculated for.
 
     Returns:
-        df (pd.DataFrame): Processed DataFrame to calculate correlations over.
+        df: Processed DataFrame to calculate correlations over.
 
     """
     if len(periods) == 1:
@@ -463,27 +498,26 @@ def _get_quasi_factor_scores_data(
 
 
 def _get_quasi_factor_scores_data_for_single_period(
-    data,
-    update_info_by_period,
-    period,
-    latent_factors,
-    observed_factors,
-):
+    data: pd.DataFrame,
+    update_info_by_period: pd.DataFrame,
+    period: int,
+    latent_factors: list[str] | tuple[str, ...],
+    observed_factors: list[str] | tuple[str, ...],
+) -> pd.DataFrame:
     """Get frame with summary scores on factor measurements in a given period.
 
     Args:
-        data (pd.DataFrame): Data with observable variables.
-        update_info_by_period (pd.DataFrame): DataFrame with information on measurements
+        data: Data with observable variables.
+        update_info_by_period: DataFrame with information on measurements
             for each factor in each user-provided period.
-        periods (list): The list of periods that correlations are
-            calculated for.
-        latent_factors (list): List of latent factors the scores of which
+        period: The period that correlations are calculated for.
+        latent_factors: List of latent factors the scores of which
             correlations are calculated for.
-        observed_factors (list): List of observed factors the scores of which
+        observed_factors: List of observed factors the scores of which
             correlations are calculated for.
 
     Returns:
-        df (pd.DataFrame): Processed DataFrame to calculate correlations over.
+        df: Processed DataFrame to calculate correlations over.
 
     """
     period_info = update_info_by_period.loc[period].reset_index()
@@ -502,32 +536,31 @@ def _get_quasi_factor_scores_data_for_single_period(
     for factor in observed_factors:
         df = data.query(f"{update_info_by_period.index.names[0]}=={period}")[factor]
         to_concat.append(df)
-    df = pd.concat(to_concat, axis=1)
-    return df
+    return pd.concat(to_concat, axis=1)
 
 
 def _get_quasi_factor_scores_data_for_multiple_periods(
-    data,
-    update_info_by_period,
-    periods,
-    latent_factors,
-    observed_factors,
-):
+    data: pd.DataFrame,
+    update_info_by_period: pd.DataFrame,
+    periods: list[int],
+    latent_factors: list[str] | tuple[str, ...],
+    observed_factors: list[str] | tuple[str, ...],
+) -> pd.DataFrame:
     """Get frame with summary scores of factor measurements in a given period.
 
     Args:
-        data (pd.DataFrame): Data with observable variables.
-        update_info_by_period (pd.DataFrame): DataFrame with information on measurements
+        data: Data with observable variables.
+        update_info_by_period: DataFrame with information on measurements
             for each factor in each user-provided period.
-        periods (list): The list of periods that correlations are
+        periods: The list of periods that correlations are
             calculated for.
-        latent_factors (list): List of latent factors the scores of which
+        latent_factors: List of latent factors the scores of which
             correlations are calculated for.
-        observed_factors (list): List of observed factors the scores of which
+        observed_factors: List of observed factors the scores of which
             correlations are calculated for.
 
     Returns:
-        df (pd.DataFrame): Processed DataFrame to calculate correlations over.
+        df: Processed DataFrame to calculate correlations over.
 
     """
     to_concat = []
@@ -543,18 +576,17 @@ def _get_quasi_factor_scores_data_for_multiple_periods(
             .add_suffix(f", {period}")
             .reset_index(drop=True),
         )
-    df = pd.concat(to_concat, axis=1)
-    return df
+    return pd.concat(to_concat, axis=1)
 
 
 def _get_factor_scores_data(
-    data,
-    params,
-    model,
-    periods,
-    latent_factors,
-    observed_factors,
-):
+    data: pd.DataFrame,
+    params: pd.DataFrame,
+    model: ProcessedModel,
+    periods: list[int],
+    latent_factors: list[str] | tuple[str, ...],
+    observed_factors: list[str] | tuple[str, ...],
+) -> pd.DataFrame:
     """Get data frame with factor scores in each period.
 
     In each period, standardize factor measurements to with estimated intercepts and
@@ -562,19 +594,19 @@ def _get_factor_scores_data(
     a summary statistics.
 
     Args:
-        data (pd.DataFrame): Data with observable variables.
-        params (pd.DataFrame): Data frame with estimated measurement relevant
+        data: Data with observable variables.
+        params: Data frame with estimated measurement relevant
             model parameters.
-        model (dict): Processed model dict.
-        periods (list): The list of periods that correlations are
+        model: Processed model dict.
+        periods: The list of periods that correlations are
             calculated for.
-        latent_factors (list): List of latent factors the scores of which
+        latent_factors: List of latent factors the scores of which
             correlations are calculated for.
-        observed_factors (list): List of observed factors the scores of which
+        observed_factors: List of observed factors the scores of which
             correlations are calculated for.
 
     Returns:
-        df (pd.DataFrame): Processed DataFrame to calculate correlations over.
+        df: Processed DataFrame to calculate correlations over.
 
     """
     if len(periods) == 1:
@@ -601,13 +633,13 @@ def _get_factor_scores_data(
 
 
 def _get_factor_scores_data_for_single_period(
-    data,
-    params,
-    model,
-    period,
-    latent_factors,
-    observed_factors,
-):
+    data: pd.DataFrame,
+    params: pd.DataFrame,
+    model: ProcessedModel,
+    period: int,
+    latent_factors: list[str] | tuple[str, ...],
+    observed_factors: list[str] | tuple[str, ...],
+) -> pd.DataFrame:
     """Get frame with factor scores in a given period.
 
     Careful: When we have endogenous factors, *period* refers to the raw period, but the
@@ -615,27 +647,27 @@ def _get_factor_scores_data_for_single_period(
     augmented periods.
 
     Args:
-        data (pd.DataFrame): Data with observable variables.
-        params (pd.DataFrame): Data frame with estimated measurement relevant
+        data: Data with observable variables.
+        params: Data frame with estimated measurement relevant
             model parameters.
-        model (dict): Processed model dict.
-        period (int): The period that correlations are calculated for.
-        latent_factors (list): List of latent factors the scores of which
+        model: Processed model dict.
+        period: The period that correlations are calculated for.
+        latent_factors: List of latent factors the scores of which
             correlations are calculated for.
-        observed_factors (list): List of observed factors the scores of which
+        observed_factors: List of observed factors the scores of which
             correlations are calculated for.
 
     Returns:
-        df (pd.DataFrame): Processed DataFrame to calculate correlations over.
+        df: Processed DataFrame to calculate correlations over.
 
     """
-    aug_periods = model["endogenous_factors_info"]["aug_periods_from_period"](period)
+    aug_periods = model.endogenous_factors_info.aug_periods_from_period(period)
     df = pd.concat(
         [
             _get_factor_scores_data_for_single_model_period(
                 data=data,
                 params=params,
-                update_info=model["update_info"],
+                update_info=model.update_info,
                 aug_period=ap,
                 period=period,
                 latent_factors=latent_factors,
@@ -654,32 +686,32 @@ def _get_factor_scores_data_for_single_period(
 
 
 def _get_factor_scores_data_for_single_model_period(
-    data,
-    params,
-    update_info,
-    aug_period,
-    period,
-    latent_factors,
-    observed_factors,
-):
+    data: pd.DataFrame,
+    params: pd.DataFrame,
+    update_info: pd.DataFrame,
+    aug_period: int,
+    period: int,
+    latent_factors: list[str] | tuple[str, ...],
+    observed_factors: list[str] | tuple[str, ...],
+) -> pd.DataFrame:
     """Get frame with factor scores in a given model period.
 
     In this function, all calculations are at the augmented period level.
 
     Args:
-        data (pd.DataFrame): Data with observable variables.
-        params (pd.DataFrame): Data frame with estimated measurement relevant
-        update_info (pd.DataFrame): DataFrame with information on measurements
+        data: Data with observable variables.
+        params: Data frame with estimated measurement relevant
+        update_info: DataFrame with information on measurements
             for each factor in each model period.
-        aug_period (int): The (augmented) period that correlations are calculated for.
-        period (int): The (raw) period that correlations are calculated for.
-        latent_factors (list): List of latent factors the scores of which
+        aug_period: The (augmented) period that correlations are calculated for.
+        period: The (raw) period that correlations are calculated for.
+        latent_factors: List of latent factors the scores of which
             correlations are calculated for.
-        observed_factors (list): List of observed factors the scores of which
+        observed_factors: List of observed factors the scores of which
             correlations are calculated for.
 
     Returns:
-        df (pd.DataFrame): Processed DataFrame to calculate correlations over.
+        df: Processed DataFrame to calculate correlations over.
     """
     if aug_period not in update_info.index:
         return pd.DataFrame()
@@ -690,7 +722,7 @@ def _get_factor_scores_data_for_single_model_period(
         params.loc["controls"].query("name2 == 'constant'").droplevel("name2")["value"]
     )
     loadings_count = loadings.astype(bool).groupby("name1").sum()
-    leave_out_meas = loadings_count[loadings_count > 1].index.to_list()
+    leave_out_meas = loadings_count[loadings_count > 1].index.to_list()  # ty: ignore[unsupported-operator]
     to_concat = []
     for factor in latent_factors:
         period_factor_measurements = period_info.query(
@@ -712,28 +744,28 @@ def _get_factor_scores_data_for_single_model_period(
 
 
 def _get_factor_scores_data_for_multiple_periods(
-    data,
-    params,
-    model,
-    periods,
-    latent_factors,
-    observed_factors,
-):
+    data: pd.DataFrame,
+    params: pd.DataFrame,
+    model: ProcessedModel,
+    periods: list[int],
+    latent_factors: list[str] | tuple[str, ...],
+    observed_factors: list[str] | tuple[str, ...],
+) -> pd.DataFrame:
     """Get frame with factor scores in a given period.
 
     Args:
-        data (pd.DataFrame): Data with observable variables.
-        params (pd.DataFrame): Data frame with estimated model parameters.
-        model (dict): Processed model dict.
-        periods (list): The list of periods that correlations are
+        data: Data with observable variables.
+        params: Data frame with estimated model parameters.
+        model: Processed model dict.
+        periods: The list of periods that correlations are
             calculated for.
-        latent_factors (list): List of latent factors the scores of which
+        latent_factors: List of latent factors the scores of which
             correlations are calculated for.
-        observed_factors (list): List of observed factors the scores of which
+        observed_factors: List of observed factors the scores of which
             correlations are calculated for.
 
     Returns:
-        df (pd.DataFrame): Processed DataFrame to calculate correlations over.
+        df: Processed DataFrame to calculate correlations over.
 
     """
     to_concat = []
@@ -750,70 +782,75 @@ def _get_factor_scores_data_for_multiple_periods(
             .add_suffix(f", {period}")
             .reset_index(drop=True),
         )
-    df = pd.concat(to_concat, axis=1)
-    return df
+    return pd.concat(to_concat, axis=1)
 
 
-def _process_factors(model, factors):
-    """Process factors to get a tuple of lists."""
+def _process_factors(
+    model: ProcessedModel,
+    factors: list[str] | tuple[str, ...] | str | None,
+) -> tuple[tuple[str, ...], tuple[str, ...]]:
+    """Process factors to get a tuple of tuples."""
     if not factors:
-        latent_factors = model["labels"]["latent_factors"]
-        observed_factors = model["labels"]["observed_factors"]
+        latent_factors = model.labels.latent_factors
+        observed_factors = model.labels.observed_factors
     elif isinstance(factors, str):
-        if factors in model["labels"]["latent_factors"]:
-            latent_factors = [factors]
-            observed_factors = []
-        elif factors in model["labels"]["observed_factors"]:
-            observed_factors = [factors]
-            latent_factors = []
+        if factors in model.labels.latent_factors:
+            latent_factors = (factors,)
+            observed_factors = ()
+        elif factors in model.labels.observed_factors:
+            observed_factors = (factors,)
+            latent_factors = ()
     else:
-        observed_factors = []
-        latent_factors = []
-        for factor in factors:
-            if factor in model["labels"]["latent_factors"]:
-                latent_factors.append(factor)
-            elif factor in model["labels"]["observed_factors"]:
-                observed_factors.append(factor)
+        latent_factors = tuple(
+            fac for fac in factors if fac in model.labels.latent_factors
+        )
+        observed_factors = tuple(
+            fac for fac in factors if fac in model.labels.observed_factors
+        )
     return latent_factors, observed_factors  # ty: ignore[possibly-unresolved-reference]
 
 
-def _process_periods(periods, model):
+def _process_periods(
+    periods: float | list[int] | None,
+    model: ProcessedModel,
+) -> list[int]:
     """Process periods to get a list."""
     if periods is None:
-        periods = list(range(model["dimensions"]["n_periods"]))
-    elif isinstance(periods, int | float):
-        periods = [periods]
+        return list(range(model.dimensions.n_periods))
+    if isinstance(periods, int | float):
+        return [int(periods)]
     return periods
 
 
 def _get_layout_kwargs(
-    corr,
-    layout_kwargs,
-    annotate,
-    annotation_fontsize,
-    annotation_text_color,
-    annotation_text_angle,
-    axes_tick_fontsize,
-    axes_tick_label_angle,
-    axes_tick_label_color,
-):
+    corr: pd.DataFrame,
+    layout_kwargs: dict[str, Any] | None,
+    *,
+    annotate: bool,
+    annotation_fontsize: int,
+    annotation_text_color: str,
+    annotation_text_angle: float,
+    axes_tick_fontsize: tuple[int, int],
+    axes_tick_label_angle: tuple[float, float],
+    axes_tick_label_color: tuple[str, str],
+) -> dict[str, Any]:
     """Get kwargs to update figure layout.
 
     Args:
-        corr (DataFrame): The processed data frame with correlation coefficients.
-        layout_kwargs (dct): Dictionary of keyword arguments used to update layout of
+        corr: The processed data frame with correlation coefficients.
+        layout_kwargs: Dictionary of keyword arguments used to update layout of
             go.Figure object.
-        annotate (bool): Add annotations to the figure if True.
-        annotation_font_size (int): Fontsize of the annotation text.
-        annotation_font_color (str): Color of the annotation text.
-        annotation_text_angle (float): The angle at which to rotate annotation text.
+        annotate: Add annotations to the figure if True.
+        annotation_fontsize: Fontsize of the annotation text.
+        annotation_text_color: Color of the annotation text.
+        annotation_text_angle: The angle at which to rotate annotation text.
         axes_tick_fontsize(tuple,list or dict): Fontsizes of axes tick labels.
         axes_tick_label_angle(tuple,list or dict): The angle at which to rotate axes
             tick labels.
-        axes_tick_label_color(tuple,list or dict): Collor of axes labels.
+        axes_tick_label_color(tuple,list or dict): Color of axes labels.
 
     Returns:
-        default_layout_kwargs (dict): Dictionary to update figure layout.
+        default_layout_kwargs: Dictionary to update figure layout.
 
     """
     default_layout_kwargs = {
@@ -825,10 +862,10 @@ def _get_layout_kwargs(
     default_layout_kwargs.update(
         _get_annotations(
             corr,
-            annotate,
-            annotation_fontsize,
-            annotation_text_color,
-            annotation_text_angle,
+            annotate=annotate,
+            annotation_fontsize=annotation_fontsize,
+            annotation_text_color=annotation_text_color,
+            annotation_text_angle=annotation_text_angle,
         ),
     )
     default_layout_kwargs.update(
@@ -844,10 +881,10 @@ def _get_layout_kwargs(
 
 
 def _get_axes_ticks_kwargs(
-    axes_tick_fontsize,
-    axes_tick_label_angle,
-    axes_tick_label_color,
-):
+    axes_tick_fontsize: tuple[int, int] | dict[str, int],
+    axes_tick_label_angle: tuple[float, float] | dict[str, float],
+    axes_tick_label_color: tuple[str, str] | dict[str, str],
+) -> dict[str, Any]:
     """Get kwargs for axes ticks label formating."""
     axes_tick_fontsize = _process_axes_tick_args(axes_tick_fontsize)
     axes_tick_label_angle = _process_axes_tick_args(axes_tick_label_angle)
@@ -865,12 +902,13 @@ def _get_axes_ticks_kwargs(
 
 
 def _get_annotations(
-    df,
-    annotate,
-    annotation_fontsize,
-    annotation_text_color,
-    annotation_text_angle,
-):
+    df: pd.DataFrame,
+    *,
+    annotate: bool,
+    annotation_fontsize: int,
+    annotation_text_color: str,
+    annotation_text_angle: float,
+) -> dict[str, Any]:
     """Get annotations and formatting kwargs."""
     annotation_kwargs = {}
     if annotate:
@@ -897,31 +935,34 @@ def _get_annotations(
 
 
 def _get_heatmap_kwargs(
-    corr,
-    heatmap_kwargs,
-    colorscale,
-    show_color_bar,
-    zmax,
-    zmin,
-    zmid,
-):
+    corr: pd.DataFrame,
+    heatmap_kwargs: dict[str, Any] | None,
+    colorscale: str,
+    *,
+    show_color_bar: bool,
+    zmax: float | None,
+    zmin: float | None,
+    zmid: float | None,
+) -> dict[str, Any]:
     """Get kwargs to instantiate Heatmap object.
 
     Args:
-        heatmap_kwargs (dct): Dictionary of key word arguments to pass to go.Heatmap().
-        colorscale (str): Name of the color palette to use in the heatmap.
+        corr: Data frame with correlation coefficients.
+        heatmap_kwargs: Dictionary of key word arguments to pass to go.Heatmap().
+        colorscale: Name of the color palette to use in the heatmap.
             Default 'RdBu_r'.
-        show_color_bar (bool): A boolean variable for displayin heatmap colorbar.
-        zmax (float or None): Upper bound to set on correlation color map.
-        zmin (float or None): Lower bound to set on correlation color map.
-        zmid (float or None): Midpoint to set on correlation color map.
+        show_color_bar: A boolean variable for displaying heatmap colorbar.
+        zmax: Upper bound to set on correlation color map.
+        zmin: Lower bound to set on correlation color map.
+        zmid: Midpoint to set on correlation color map.
 
     Returns:
-        default_heatmap_kwargs (dict): Dictionary of kwargs to instantiate go.Heatmap.
+        default_heatmap_kwargs: Dictionary of kwargs to instantiate go.Heatmap.
 
     """
     if zmax is None:
-        zmax = np.abs(corr.to_numpy())[np.tril_indices_from(corr, k=-1)].max()
+        corr_arr = corr.to_numpy()
+        zmax = np.abs(corr_arr)[np.tril_indices_from(corr_arr, k=-1)].max()
     if zmin is None:
         zmin = -zmax
     if zmid is None:
@@ -938,7 +979,9 @@ def _get_heatmap_kwargs(
     return default_heatmap_kwargs
 
 
-def _process_axes_tick_args(args):
+def _process_axes_tick_args(
+    args: tuple[Any, Any] | list[Any] | dict[str, Any],
+) -> dict[str, Any]:
     if isinstance(args, tuple | list):
         args = {"x": args[0], "y": args[1]}
     return args
diff --git a/src/skillmodels/decorators.py b/src/skillmodels/decorators.py
index 6dcf779..e1c834a 100644
--- a/src/skillmodels/decorators.py
+++ b/src/skillmodels/decorators.py
@@ -1,27 +1,38 @@
+"""Decorators for parameter extraction and registration in transition functions."""
+
 import functools
+from collections.abc import Callable  # noqa: TC003
+from typing import Any
 
 import jax.numpy as jnp
+from jax import Array
 
 
-def extract_params(func=None, *, key=None, names=None):
+def extract_params(
+    func: Callable | None = None,
+    *,
+    key: str | None = None,
+    names: list[str] | None = None,
+) -> Callable:
     """Process params before passing them to func.
 
     Note: The resulting function is keyword only!
 
     Args:
-        key (str or None): If key is not None, we assume params is a dictionary of which
+        func: The function to be decorated, or None if using decorator with arguments.
+        key: If key is not None, we assume params is a dictionary of which
             only the params[key] should be passed into func.
-        names (list or None): If names is provided, we assume that params
+        names: If names is provided, we assume that params
             (or params[key]) should be converted to a dictionary with names as keys
             before passing them to func.
 
     """
 
-    def decorator_extract_params(func):
+    def decorator_extract_params(func: Callable) -> Callable:
         if key is not None and names is None:
 
             @functools.wraps(func)
-            def wrapper_extract_params(**kwargs):
+            def wrapper_extract_params(**kwargs: Any) -> Any:
                 internal_kwargs = kwargs.copy()
                 internal_kwargs["params"] = kwargs["params"][key]
                 return func(**internal_kwargs)
@@ -29,7 +40,7 @@ def wrapper_extract_params(**kwargs):
         elif key is None and names is not None:
 
             @functools.wraps(func)
-            def wrapper_extract_params(**kwargs):
+            def wrapper_extract_params(**kwargs: Any) -> Any:
                 internal_kwargs = kwargs.copy()
                 internal_kwargs["params"] = dict(
                     zip(names, kwargs["params"], strict=False)
@@ -39,7 +50,7 @@ def wrapper_extract_params(**kwargs):
         elif key is not None and names is not None:
 
             @functools.wraps(func)
-            def wrapper_extract_params(**kwargs):
+            def wrapper_extract_params(**kwargs: Any) -> Any:
                 internal_kwargs = kwargs.copy()
                 internal_kwargs["params"] = dict(
                     zip(names, kwargs["params"][key], strict=False)
@@ -56,21 +67,26 @@ def wrapper_extract_params(**kwargs):
     return decorator_extract_params
 
 
-def jax_array_output(func):
+def jax_array_output(func: Callable) -> Callable:
     """Convert tuple output to list output."""
 
     @functools.wraps(func)
-    def wrapper_jax_array_output(*args, **kwargs):
+    def wrapper_jax_array_output(*args: Any, **kwargs: Any) -> Array:
         raw = func(*args, **kwargs)
-        out = jnp.array(raw)
-        return out
+        return jnp.array(raw)
 
     return wrapper_jax_array_output
 
 
-def register_params(func=None, *, params=None):
-    def decorator_register_params(func):
-        func.__registered_params__ = params
+def register_params(
+    func: Callable | None = None,
+    *,
+    params: list[str] | None = None,
+) -> Callable:
+    """Register parameter names for a transition function."""
+
+    def decorator_register_params(func: Callable) -> Callable:
+        func.__registered_params__ = params  # ty: ignore[unresolved-attribute]
         return func
 
     if callable(func):
diff --git a/src/skillmodels/filtered_states.py b/src/skillmodels/filtered_states.py
index 2efb4c6..66e56f0 100644
--- a/src/skillmodels/filtered_states.py
+++ b/src/skillmodels/filtered_states.py
@@ -1,3 +1,7 @@
+"""Functions to compute and process filtered latent states."""
+
+from typing import TYPE_CHECKING, Any
+
 import jax.numpy as jnp
 import numpy as np
 
@@ -7,8 +11,16 @@
 from skillmodels.process_debug_data import create_state_ranges
 from skillmodels.process_model import process_model
 
+if TYPE_CHECKING:
+    import pandas as pd
+
 
-def get_filtered_states(model_dict, data, params):
+def get_filtered_states(
+    model_dict: dict,
+    data: pd.DataFrame,
+    params: pd.DataFrame,
+) -> dict[str, dict[str, Any]]:
+    """Compute filtered latent states given data and estimated parameters."""
     max_inputs = get_maximization_inputs(model_dict=model_dict, data=data)
     params = params.loc[max_inputs["params_template"].index]
     debug_loglike = max_inputs["debug_loglike"]
@@ -26,10 +38,10 @@ def get_filtered_states(model_dict, data, params):
 
     anchored_ranges = create_state_ranges(
         filtered_states=anchored_states_df,
-        factors=model["labels"]["latent_factors"],
+        factors=model.labels.latent_factors,
     )
 
-    out = {
+    return {
         "anchored_states": {
             "states": anchored_states_df,
             "state_ranges": anchored_ranges,
@@ -40,10 +52,14 @@ def get_filtered_states(model_dict, data, params):
         },
     }
 
-    return out
 
-
-def anchor_states_df(states_df, model_dict, params, use_aug_period):
+def anchor_states_df(
+    states_df: pd.DataFrame,
+    model_dict: dict,
+    params: pd.DataFrame,
+    *,
+    use_aug_period: bool,
+) -> pd.DataFrame:
     """Anchor states in a DataFrame.
 
     The DataFrame is expected to have a column called "period" as well as one column
@@ -58,40 +74,38 @@ def anchor_states_df(states_df, model_dict, params, use_aug_period):
     model = process_model(model_dict)
 
     p_index = get_params_index(
-        update_info=model["update_info"],
-        labels=model["labels"],
-        dimensions=model["dimensions"],
-        transition_info=model["transition_info"],
-        endogenous_factors_info=model["endogenous_factors_info"],
+        update_info=model.update_info,
+        labels=model.labels,
+        dimensions=model.dimensions,
+        transition_info=model.transition_info,
+        endogenous_factors_info=model.endogenous_factors_info,
     )
 
     params = params.loc[p_index]
 
     parsing_info = create_parsing_info(
         params_index=p_index,
-        update_info=model["update_info"],
-        labels=model["labels"],
-        anchoring=model["anchoring"],
-        has_endogenous_factors=model["endogenous_factors_info"][
-            "has_endogenous_factors"
-        ],
+        update_info=model.update_info,
+        labels=model.labels,
+        anchoring=model.anchoring,
+        has_endogenous_factors=model.endogenous_factors_info.has_endogenous_factors,
     )
 
-    *_, pardict = parse_params(
+    *_, parsed_params = parse_params(
         params=jnp.array(params["value"].to_numpy()),
         parsing_info=parsing_info,
-        dimensions=model["dimensions"],
-        labels=model["labels"],
+        dimensions=model.dimensions,
+        labels=model.labels,
         n_obs=1,
     )
 
-    n_latent = model["dimensions"]["n_latent_factors"]
+    n_latent = model.dimensions.n_latent_factors
 
-    _scaling_factors = np.array(pardict["anchoring_scaling_factors"][:, :n_latent])
-    _constants = np.array(pardict["anchoring_constants"][:, :n_latent])
+    _scaling_factors = np.array(parsed_params.anchoring_scaling_factors[:, :n_latent])
+    _constants = np.array(parsed_params.anchoring_constants[:, :n_latent])
     if use_aug_period:
         period_arr = states_df["aug_period"].to_numpy()
-        ap_to_p = model["labels"]["aug_periods_to_periods"]
+        ap_to_p = model.labels.aug_periods_to_periods
         scaling_factors = np.empty(shape=(len(ap_to_p), n_latent))
         constants = np.empty(shape=(len(ap_to_p), n_latent))
         for ap, p in ap_to_p.items():
@@ -106,9 +120,7 @@ def anchor_states_df(states_df, model_dict, params, use_aug_period):
     constants_arr = constants[period_arr]
 
     out = states_df.copy(deep=True)
-    for pos, factor in enumerate(model["labels"]["latent_factors"]):
+    for pos, factor in enumerate(model.labels.latent_factors):
         out[factor] = constants_arr[:, pos] + states_df[factor] * scaling_arr[:, pos]
 
-    out = out[states_df.columns]
-
-    return out
+    return out[states_df.columns]
diff --git a/src/skillmodels/kalman_filters.py b/src/skillmodels/kalman_filters.py
index f9cfae9..d1b2b97 100644
--- a/src/skillmodels/kalman_filters.py
+++ b/src/skillmodels/kalman_filters.py
@@ -1,5 +1,10 @@
+"""Kalman filter operations for state estimation using the square-root form."""
+
+from collections.abc import Callable  # noqa: TC003
+
 import jax
 import jax.numpy as jnp
+from jax import Array
 
 from skillmodels.qr import qr_gpu
 
@@ -9,44 +14,43 @@
     else jax.vmap(jax.vmap(jnp.linalg.qr))
 )
 
+
 # ======================================================================================
 # Update Step
 # ======================================================================================
-
-
 def kalman_update(
-    states,
-    upper_chols,
-    loadings,
-    control_params,
-    meas_sd,
-    measurements,
-    controls,
-    log_mixture_weights,
-):
+    states: Array,
+    upper_chols: Array,
+    loadings: Array,
+    control_params: Array,
+    meas_sd: Array,
+    measurements: Array,
+    controls: Array,
+    log_mixture_weights: Array,
+) -> tuple[Array, Array, Array, Array]:
     """Perform a Kalman update with likelihood evaluation.
 
     Args:
-        states (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states) with
+        states: Array of shape (n_obs, n_mixtures, n_states) with
             pre-update states estimates.
-        upper_chols (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states,
+        upper_chols: Array of shape (n_obs, n_mixtures, n_states,
             n_states) with the transpose of the lower triangular cholesky factor
             of the pre-update covariance matrix of the state estimates.
-        loadings (jax.numpy.array): 1d array of length n_states with factor loadings.
-        control_params (jax.numpy.array): 1d array of length n_controls.
-        meas_sd (float): Standard deviation of the measurement error.
-        measurements (jax.numpy.array): 1d array of length n_obs with measurements.
+        loadings: 1d array of length n_states with factor loadings.
+        control_params: 1d array of length n_controls.
+        meas_sd: Standard deviation of the measurement error.
+        measurements: 1d array of length n_obs with measurements.
             May contain NaNs if no measurement was observed.
-        controls (jax.numpy.array): Array of shape (n_obs, n_controls) with data on the
+        controls: Array of shape (n_obs, n_controls) with data on the
             control variables.
-        log_mixture_weights (jax.numpy.array): Array of shape (n_obs, n_mixtures) with
+        log_mixture_weights: Array of shape (n_obs, n_mixtures) with
             the natural logarithm of the weights of each element of the mixture of
             normals distribution.
 
     Returns:
-        states (jax.numpy.array): Same format as states.
-        new_states (jax.numpy.array): Same format as states.
-        new_upper_chols (jax.numpy.array): Same format as upper_chols
+        states: Same format as states.
+        new_states: Same format as states.
+        new_upper_chols: Same format as upper_chols
         new_log_mixture_weights: (jax.numpy.array): Same format as log_mixture_weights
         new_loglikes: (jax.numpy.array): 1d array of length n_obs
 
@@ -133,17 +137,18 @@ def kalman_update(
 # ======================================================================================
 # Predict Step
 # ======================================================================================
-
-
-def calculate_sigma_scaling_factor_and_weights(n_states, kappa=2):
+def calculate_sigma_scaling_factor_and_weights(
+    n_states: int,
+    kappa: float = 2,
+) -> tuple[Array, Array]:
     """Calculate the scaling factor and weights for sigma points according to Julier.
 
     There are other sigma point algorithms, but many of them possibly have negative
     weights which makes the unscented predict step more complicated.
 
     Args:
-        n_states (int): Number of states.
-        kappa (float): Spreading factor of the sigma points.
+        n_states: Number of states.
+        kappa: Spreading factor of the sigma points.
 
     Returns:
         float: Scaling factor
@@ -158,39 +163,41 @@ def calculate_sigma_scaling_factor_and_weights(n_states, kappa=2):
 
 
 def kalman_predict(
-    transition_func,
-    states,
-    upper_chols,
-    sigma_scaling_factor,
-    sigma_weights,
-    trans_coeffs,
-    shock_sds,
-    anchoring_scaling_factors,
-    anchoring_constants,
-    observed_factors,
-):
+    transition_func: Callable,
+    states: Array,
+    upper_chols: Array,
+    sigma_scaling_factor: float,
+    sigma_weights: Array,
+    trans_coeffs: dict[str, Array],
+    shock_sds: Array,
+    anchoring_scaling_factors: Array,
+    anchoring_constants: Array,
+    observed_factors: Array,
+) -> tuple[Array, Array]:
     """Make a unscented Kalman predict.
 
     Args:
-        transition_func (Callable): The transition function.
-        states (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states) with
+        transition_func: The transition function.
+        states: Array of shape (n_obs, n_mixtures, n_states) with
             pre-update states estimates.
-        upper_chols (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states,
+        upper_chols: Array of shape (n_obs, n_mixtures, n_states,
             n_states) with the transpose of the lower triangular cholesky factor
             of the pre-update covariance matrix of the state estimates.
-        sigma_scaling_factor (float): A scaling factor that controls the spread of the
+        sigma_scaling_factor: A scaling factor that controls the spread of the
             sigma points. Bigger means that sigma points are further apart. Depends on
             the sigma_point algorithm chosen.
-        sigma_weights (jax.numpy.array): 1d array of length n_sigma with non-negative
+        sigma_weights: 1d array of length n_sigma with non-negative
             sigma weights.
-        trans_coeffs (tuple): Tuple of 1d jax.numpy.arrays with transition parameters.
-        anchoring_scaling_factors (jax.numpy.array): Array of shape (2, n_fac) with
+        trans_coeffs: Tuple of 1d jax.numpy.arrays with transition parameters.
+        shock_sds: 1d array of length n_fac with shock standard
+            deviations.
+        anchoring_scaling_factors: Array of shape (2, n_fac) with
             the scaling factors for anchoring. The first row corresponds to the input
             period, the second to the output period (i.e. input period + 1).
-        anchoring_constants (jax.numpy.array): Array of shape (2, n_states) with the
+        anchoring_constants: Array of shape (2, n_states) with the
             constants for anchoring. The first row corresponds to the input
             period, the second to the output period (i.e. input period + 1).
-        observed_factors (jax.numpy.array): Array of shape (n_obs, n_observed_factors)
+        observed_factors: Array of shape (n_obs, n_observed_factors)
             with data on the observed factors in period t.
 
     Returns:
@@ -228,19 +235,24 @@ def kalman_predict(
     return predicted_states, predicted_covs
 
 
-def _calculate_sigma_points(states, upper_chols, scaling_factor, observed_factors):
+def _calculate_sigma_points(
+    states: Array,
+    upper_chols: Array,
+    scaling_factor: float,
+    observed_factors: Array,
+) -> Array:
     """Calculate the array of sigma_points for the unscented transform.
 
     Args:
-        states (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states) with
+        states: Array of shape (n_obs, n_mixtures, n_states) with
             pre-update states estimates.
-        upper_chols (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states,
+        upper_chols: Array of shape (n_obs, n_mixtures, n_states,
             n_states) with the transpose of the lower triangular cholesky factor
             of the pre-update covariance matrix of the state estimates.
-        scaling_factor (float): A scaling factor that controls the spread of the
+        scaling_factor: A scaling factor that controls the spread of the
             sigma points. Bigger means that sigma points are further apart. Depends on
             the sigma_point algorithm chosen.
-        observed_factors (jax.numpy.array): Array of shape (n_obs, n_observed_factors)
+        observed_factors: Array of shape (n_obs, n_observed_factors)
             with data on the observed factors in period t.
 
     Returns:
@@ -269,27 +281,26 @@ def _calculate_sigma_points(states, upper_chols, scaling_factor, observed_factor
         n_observed,
     )
 
-    sigma_points = jnp.concatenate([sigma_points, observed_part], axis=-1)
-    return sigma_points
+    return jnp.concatenate([sigma_points, observed_part], axis=-1)
 
 
 def transform_sigma_points(
-    sigma_points,
-    transition_func,
-    trans_coeffs,
-    anchoring_scaling_factors,
-    anchoring_constants,
-):
+    sigma_points: Array,
+    transition_func: Callable,
+    trans_coeffs: dict[str, Array],
+    anchoring_scaling_factors: Array,
+    anchoring_constants: Array,
+) -> Array:
     """Anchor sigma points, transform them and unanchor the transformed sigma points.
 
     Args:
-        sigma_points (jax.numpy.array) of shape n_obs, n_mixtures, n_sigma, n_fac.
-        transition_func (Callable): The transition function.
-        trans_coeffs (tuple): Tuple of 1d jax.numpy.arrays with transition parameters.
-        anchoring_scaling_factors (jax.numpy.array): Array of shape (2, n_states) with
+        sigma_points: Array of shape n_obs, n_mixtures, n_sigma, n_fac.
+        transition_func: The transition function.
+        trans_coeffs: Tuple of 1d jax.numpy.arrays with transition parameters.
+        anchoring_scaling_factors: Array of shape (2, n_states) with
             the scaling factors for anchoring. The first row corresponds to the input
             period, the second to the output period (i.e. input period + 1).
-        anchoring_constants (jax.numpy.array): Array of shape (2, n_states) with the
+        anchoring_constants: Array of shape (2, n_states) with the
             constants for anchoring. The first row corresponds to the input
             period, the second to the output period (i.e. input period + 1).
 
@@ -313,6 +324,4 @@ def transform_sigma_points(
     ) / anchoring_scaling_factors[1][:n_observed]
 
     out_shape = (n_obs, n_mixtures, n_sigma, -1)
-    out = transformed_unanchored.reshape(out_shape)
-
-    return out
+    return transformed_unanchored.reshape(out_shape)
diff --git a/src/skillmodels/kalman_filters_debug.py b/src/skillmodels/kalman_filters_debug.py
index 14b07e1..5d1ec44 100644
--- a/src/skillmodels/kalman_filters_debug.py
+++ b/src/skillmodels/kalman_filters_debug.py
@@ -1,5 +1,10 @@
+"""Debug versions of Kalman filter operations that return intermediate results."""
+
+from typing import Any
+
 import jax
 import jax.numpy as jnp
+from jax import Array
 
 array_qr_jax = jax.vmap(jax.vmap(jnp.linalg.qr))
 
@@ -7,44 +12,42 @@
 # ======================================================================================
 # Update Step
 # ======================================================================================
-
-
 def kalman_update(
-    states,
-    upper_chols,
-    loadings,
-    control_params,
-    meas_sd,
-    measurements,
-    controls,
-    log_mixture_weights,
-):
+    states: Array,
+    upper_chols: Array,
+    loadings: Array,
+    control_params: Array,
+    meas_sd: float,
+    measurements: Array,
+    controls: Array,
+    log_mixture_weights: Array,
+) -> tuple[Array, Array, Array, Array, dict[str, Any]]:
     """Perform a Kalman update with likelihood evaluation, returning debug info on top.
 
     Args:
-        states (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states) with
+        states: Array of shape (n_obs, n_mixtures, n_states) with
             pre-update states estimates.
-        upper_chols (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states,
+        upper_chols: Array of shape (n_obs, n_mixtures, n_states,
             n_states) with the transpose of the lower triangular cholesky factor
             of the pre-update covariance matrix of the state estimates.
-        loadings (jax.numpy.array): 1d array of length n_states with factor loadings.
-        control_params (jax.numpy.array): 1d array of length n_controls.
-        meas_sd (float): Standard deviation of the measurement error.
-        measurements (jax.numpy.array): 1d array of length n_obs with measurements.
+        loadings: 1d array of length n_states with factor loadings.
+        control_params: 1d array of length n_controls.
+        meas_sd: Standard deviation of the measurement error.
+        measurements: 1d array of length n_obs with measurements.
             May contain NaNs if no measurement was observed.
-        controls (jax.numpy.array): Array of shape (n_obs, n_controls) with data on the
+        controls: Array of shape (n_obs, n_controls) with data on the
             control variables.
-        log_mixture_weights (jax.numpy.array): Array of shape (n_obs, n_mixtures) with
+        log_mixture_weights: Array of shape (n_obs, n_mixtures) with
             the natural logarithm of the weights of each element of the mixture of
             normals distribution.
 
     Returns:
-        states (jax.numpy.array): Same format as states.
-        new_states (jax.numpy.array): Same format as states.
-        new_upper_chols (jax.numpy.array): Same format as upper_chols
+        states: Same format as states.
+        new_states: Same format as states.
+        new_upper_chols: Same format as upper_chols
         new_log_mixture_weights: (jax.numpy.array): Same format as log_mixture_weights
         new_loglikes: (jax.numpy.array): 1d array of length n_obs
-        debug_info (dict): Empty or containing residuals and residual_sds
+        debug_info: Empty or containing residuals and residual_sds
 
     """
     n_obs, n_mixtures, n_states = states.shape
diff --git a/src/skillmodels/likelihood_function.py b/src/skillmodels/likelihood_function.py
index 4e7eec2..d72055a 100644
--- a/src/skillmodels/likelihood_function.py
+++ b/src/skillmodels/likelihood_function.py
@@ -1,7 +1,12 @@
+"""Log-likelihood function for latent factor models."""
+
 import functools
+from collections.abc import Callable  # noqa: TC003
+from typing import Any
 
 import jax
 import jax.numpy as jnp
+from jax import Array
 
 from skillmodels.clipping import soft_clipping
 from skillmodels.kalman_filters import (
@@ -9,24 +14,63 @@
     kalman_update,
 )
 from skillmodels.parse_params import parse_params
+from skillmodels.types import (  # noqa: TC001
+    Dimensions,
+    EstimationOptions,
+    Labels,
+    ParsedParams,
+    ParsingInfo,
+)
 
 
 def log_likelihood(
-    params,
-    parsing_info,
-    measurements,
-    controls,
-    transition_func,
-    sigma_scaling_factor,
-    sigma_weights,
-    dimensions,
-    labels,
-    estimation_options,
-    is_measurement_iteration,
-    is_predict_iteration,
-    iteration_to_period,
-    observed_factors,
-):
+    params: Array,
+    parsing_info: ParsingInfo,
+    measurements: Array,
+    controls: Array,
+    transition_func: Callable,
+    sigma_scaling_factor: float,
+    sigma_weights: Array,
+    dimensions: Dimensions,
+    labels: Labels,
+    estimation_options: EstimationOptions,
+    is_measurement_iteration: Array,
+    is_predict_iteration: Array,
+    iteration_to_period: Array,
+    observed_factors: Array,
+) -> Array:
+    """Aggregated log likelihood of a skill formation model.
+
+    Wrapper around log_likelihood_obs that sums contributions across observations.
+
+    Args:
+        params: 1d array with model parameters.
+        parsing_info: Contains information how to parse parameter vector.
+        measurements: Array of shape (n_updates, n_obs) with data on
+            observed measurements. NaN if the measurement was not observed.
+        controls: Array of shape (n_periods, n_obs, n_controls)
+            with observed control variables for the measurement equations.
+        transition_func: The transition function.
+        sigma_scaling_factor: A scaling factor that controls the spread of the
+            sigma points.
+        sigma_weights: 1d array of length n_sigma with non-negative sigma weights.
+        dimensions: Dimensional information like n_states, n_periods, n_controls,
+            n_mixtures.
+        labels: Labels for the model quantities like factors, periods, controls,
+            stagemap and stages.
+        estimation_options: Options for estimation including clipping bounds.
+        is_measurement_iteration: Boolean array indicating which iterations are
+            measurement updates.
+        is_predict_iteration: Boolean array indicating which iterations are predict
+            steps.
+        iteration_to_period: Array mapping iteration index to period.
+        observed_factors: Array of shape (n_periods, n_obs, n_observed_factors) with
+            data on the observed factors.
+
+    Returns:
+        Scalar aggregated log likelihood.
+
+    """
     return log_likelihood_obs(
         params=params,
         parsing_info=parsing_info,
@@ -46,21 +90,21 @@ def log_likelihood(
 
 
 def log_likelihood_obs(
-    params,
-    parsing_info,
-    measurements,
-    controls,
-    transition_func,
-    sigma_scaling_factor,
-    sigma_weights,
-    dimensions,
-    labels,
-    estimation_options,
-    is_measurement_iteration,
-    is_predict_iteration,
-    iteration_to_period,
-    observed_factors,
-):
+    params: Array,
+    parsing_info: ParsingInfo,
+    measurements: Array,
+    controls: Array,
+    transition_func: Callable,
+    sigma_scaling_factor: float,
+    sigma_weights: Array,
+    dimensions: Dimensions,
+    labels: Labels,
+    estimation_options: EstimationOptions,
+    is_measurement_iteration: Array,
+    is_predict_iteration: Array,
+    iteration_to_period: Array,
+    observed_factors: Array,
+) -> Array:
     """Log likelihood of a skill formation model.
 
     This function is jax-differentiable and jax-jittable as long as all but the first
@@ -73,25 +117,29 @@ def log_likelihood_obs(
     with Jax.
 
     Args:
-        params (jax.numpy.array): 1d array with model parameters.
-        parsing_info (dict): Contains information how to parse parameter vector.
-        update_info (pandas.DataFrame): Contains information about number of updates in
-            each period and purpose of each update.
-        measurements (jax.numpy.array): Array of shape (n_updates, n_obs) with data on
+        params: 1d array with model parameters.
+        parsing_info: Contains information how to parse parameter vector.
+        measurements: Array of shape (n_updates, n_obs) with data on
             observed measurements. NaN if the measurement was not observed.
-        controls (jax.numpy.array): Array of shape (n_periods, n_obs, n_controls)
+        controls: Array of shape (n_periods, n_obs, n_controls)
             with observed control variables for the measurement equations.
-        transition_func (Callable): The transition function.
-        sigma_scaling_factor (float): A scaling factor that controls the spread of the
+        transition_func: The transition function.
+        sigma_scaling_factor: A scaling factor that controls the spread of the
             sigma points. Bigger means that sigma points are further apart. Depends on
             the sigma_point algorithm chosen.
-        sigma_weights (jax.numpy.array): 1d array of length n_sigma with non-negative
+        sigma_weights: 1d array of length n_sigma with non-negative
             sigma weights.
-        dimensions (dict): Dimensional information like n_states, n_periods, n_controls,
+        dimensions: Dimensional information like n_states, n_periods, n_controls,
             n_mixtures. See :ref:`dimensions`.
-        labels (dict): Dict of lists with labels for the model quantities like
+        labels: Dict of lists with labels for the model quantities like
             factors, periods, controls, stagemap and stages. See :ref:`labels`
-        observed_factors (jax.numpy.array): Array of shape (n_periods, n_obs,
+        estimation_options: Options for estimation including clipping bounds.
+        is_measurement_iteration: Boolean array indicating which
+            iterations are measurement updates.
+        is_predict_iteration: Boolean array indicating which
+            iterations are predict steps.
+        iteration_to_period: Array mapping iteration index to period.
+        observed_factors: Array of shape (n_periods, n_obs,
             n_observed_factors) with data on the observed factors.
 
     Returns:
@@ -99,7 +147,7 @@ def log_likelihood_obs(
 
     """
     n_obs = measurements.shape[1]
-    states, upper_chols, log_mixture_weights, pardict = parse_params(
+    states, upper_chols, log_mixture_weights, parsed_params = parse_params(
         params,
         parsing_info,
         dimensions,
@@ -115,9 +163,9 @@ def log_likelihood_obs(
 
     loop_args = {
         "period": iteration_to_period,
-        "loadings": pardict["loadings"],
-        "control_params": pardict["controls"],
-        "meas_sds": pardict["meas_sds"],
+        "loadings": parsed_params.loadings,
+        "control_params": parsed_params.controls,
+        "meas_sds": parsed_params.meas_sds,
         "measurements": measurements,
         "is_measurement_iteration": is_measurement_iteration,
         "is_predict_iteration": is_predict_iteration,
@@ -126,7 +174,7 @@ def log_likelihood_obs(
     _body = functools.partial(
         _scan_body,
         controls=controls,
-        pardict=pardict,
+        parsed_params=parsed_params,
         sigma_scaling_factor=sigma_scaling_factor,
         sigma_weights=sigma_weights,
         transition_func=transition_func,
@@ -139,23 +187,23 @@ def log_likelihood_obs(
     # possible.
     return soft_clipping(
         arr=static_out["loglikes"],
-        lower=estimation_options["clipping_lower_bound"],
-        upper=estimation_options["clipping_upper_bound"],
-        lower_hardness=estimation_options["clipping_lower_hardness"],
-        upper_hardness=estimation_options["clipping_upper_hardness"],
+        lower=estimation_options.clipping_lower_bound,
+        upper=estimation_options.clipping_upper_bound,
+        lower_hardness=estimation_options.clipping_lower_hardness,
+        upper_hardness=estimation_options.clipping_upper_hardness,
     ).sum(axis=0)
 
 
 def _scan_body(
-    carry,
-    loop_args,
-    controls,
-    pardict,
-    sigma_scaling_factor,
-    sigma_weights,
-    transition_func,
-    observed_factors,
-):
+    carry: dict[str, Array],
+    loop_args: dict[str, Array],
+    controls: Array,
+    parsed_params: ParsedParams,
+    sigma_scaling_factor: float,
+    sigma_weights: Array,
+    transition_func: Callable,
+    observed_factors: Array,
+) -> tuple[dict[str, Array], dict[str, Array]]:
     # ==================================================================================
     # create arguments needed for update
     # ==================================================================================
@@ -193,12 +241,12 @@ def _scan_body(
         "upper_chols": upper_chols,
         "sigma_scaling_factor": sigma_scaling_factor,
         "sigma_weights": sigma_weights,
-        "trans_coeffs": {k: arr[t] for k, arr in pardict["transition"].items()},
-        "shock_sds": pardict["shock_sds"][t],
-        "anchoring_scaling_factors": pardict["anchoring_scaling_factors"][
+        "trans_coeffs": {k: arr[t] for k, arr in parsed_params.transition.items()},
+        "shock_sds": parsed_params.shock_sds[t],
+        "anchoring_scaling_factors": parsed_params.anchoring_scaling_factors[
             jnp.array([t, t + 1])
         ],
-        "anchoring_constants": pardict["anchoring_constants"][jnp.array([t, t + 1])],
+        "anchoring_constants": parsed_params.anchoring_constants[jnp.array([t, t + 1])],
         "observed_factors": observed_factors[t],
     }
 
@@ -224,28 +272,36 @@ def _scan_body(
     return new_state, static_out
 
 
-def _one_arg_measurement_update(kwargs):
-    out = kalman_update(**kwargs)
-    return out
+def _one_arg_measurement_update(
+    kwargs: dict[str, Array],
+) -> tuple[Array, Array, Array, Array]:
+    return kalman_update(**kwargs)
 
 
-def _one_arg_anchoring_update(kwargs):
+def _one_arg_anchoring_update(
+    kwargs: dict[str, Array],
+) -> tuple[Array, Array, Array, Array]:
     _, _, new_log_mixture_weights, new_loglikes = kalman_update(**kwargs)
-    out = (
+    return (
         kwargs["states"],
         kwargs["upper_chols"],
         new_log_mixture_weights,
         new_loglikes,
     )
-    return out
 
 
-def _one_arg_no_predict(kwargs, transition_func):  # noqa: ARG001
+def _one_arg_no_predict(
+    kwargs: dict[str, Any],
+    transition_func: Callable,  # noqa: ARG001
+) -> tuple[Array, Array, Array]:
     """Just return the states cond chols without any changes."""
     return kwargs["states"], kwargs["upper_chols"], kwargs["states"]
 
 
-def _one_arg_predict(kwargs, transition_func):
+def _one_arg_predict(
+    kwargs: dict[str, Any],
+    transition_func: Callable,
+) -> tuple[Array, Array, Array]:
     """Do a predict step but also return the input states as filtered states."""
     new_states, new_upper_chols = kalman_predict(
         transition_func,
diff --git a/src/skillmodels/likelihood_function_debug.py b/src/skillmodels/likelihood_function_debug.py
index 6624050..800b5b6 100644
--- a/src/skillmodels/likelihood_function_debug.py
+++ b/src/skillmodels/likelihood_function_debug.py
@@ -1,72 +1,78 @@
+"""Debug version of log-likelihood function that returns intermediate results."""
+
 import functools
+from collections.abc import Callable  # noqa: TC003
+from typing import Any
 
 import jax
 import jax.numpy as jnp
+from jax import Array
 
 from skillmodels.clipping import soft_clipping
 from skillmodels.kalman_filters import kalman_predict
 from skillmodels.kalman_filters_debug import kalman_update
 from skillmodels.parse_params import parse_params
+from skillmodels.types import (  # noqa: TC001
+    Dimensions,
+    EstimationOptions,
+    Labels,
+    ParsedParams,
+    ParsingInfo,
+)
 
 
 def log_likelihood(
-    params,
-    parsing_info,
-    measurements,
-    controls,
-    transition_func,
-    sigma_scaling_factor,
-    sigma_weights,
-    dimensions,
-    labels,
-    estimation_options,
-    is_measurement_iteration,
-    is_predict_iteration,
-    iteration_to_period,
-    observed_factors,
-):
+    params: Array,
+    parsing_info: ParsingInfo,
+    measurements: Array,
+    controls: Array,
+    transition_func: Callable[..., Array],
+    sigma_scaling_factor: float,
+    sigma_weights: Array,
+    dimensions: Dimensions,
+    labels: Labels,
+    estimation_options: EstimationOptions,
+    is_measurement_iteration: Array,
+    is_predict_iteration: Array,
+    iteration_to_period: Array,
+    observed_factors: Array,
+) -> dict[str, Any]:
     """Log likelihood of a skill formation model, returning debug data on top.
 
     This function is jax-differentiable and jax-jittable as long as all but the first
     argument are marked as static.
 
-    The function returns both a tuple (float, dict). The first entry is the aggregated
-    log likelihood value. The second additional information like the log likelihood
-    contribution of each individual. Note that the dict also contains the aggregated
-    value. Returning that value separately is only needed to calculate a gradient with
-    Jax.
-
     Args:
-        params (jax.numpy.array): 1d array with model parameters. parsing_info (dict):
-        Contains information how to parse parameter vector. update_info
-        (pandas.DataFrame): Contains information about number of updates in
-            each period and purpose of each update.
-        measurements (jax.numpy.array): Array of shape (n_updates, n_obs) with data on
-            observed measurements. NaN if the measurement was not observed.
-        controls (jax.numpy.array): Array of shape (n_periods, n_obs, n_controls)
-            with observed control variables for the measurement equations.
-        transition_func (dict): Dict with the entries "func" (the actual transition
-            function) and "columns" (a dictionary mapping factors that are needed as
-            individual columns to positions in the factor array).
-        sigma_scaling_factor (float): A scaling factor that controls the spread of the
-            sigma points. Bigger means that sigma points are further apart. Depends on
-            the sigma_point algorithm chosen.
-        sigma_weights (jax.numpy.array): 1d array of length n_sigma with non-negative
-            sigma weights.
-        dimensions (dict): Dimensional information like n_states, n_periods, n_controls,
-            n_mixtures. See :ref:`dimensions`.
-        labels (dict): Dict of lists with labels for the model quantities like
-            factors, periods, controls, stagemap and stages. See :ref:`labels`
-        observed_factors (jax.numpy.array): Array of shape (n_periods, n_obs,
-            n_observed_factors) with data on the observed factors.
+        params: 1d array with model parameters.
+        parsing_info: Contains information how to parse parameter vector.
+        measurements: Array of shape (n_updates, n_obs) with data on observed
+            measurements. NaN if the measurement was not observed.
+        controls: Array of shape (n_periods, n_obs, n_controls) with observed
+            control variables for the measurement equations.
+        transition_func: The transition function.
+        sigma_scaling_factor: A scaling factor that controls the spread of the
+            sigma points. Bigger means that sigma points are further apart.
+        sigma_weights: 1d array of length n_sigma with non-negative sigma weights.
+        dimensions: Dimensional information like n_states, n_periods, n_controls,
+            n_mixtures.
+        labels: Labels for the model quantities like factors, periods, controls,
+            stagemap and stages.
+        estimation_options: Options for estimation including clipping bounds.
+        is_measurement_iteration: Boolean array indicating which iterations are
+            measurement updates.
+        is_predict_iteration: Boolean array indicating which iterations are predict
+            steps.
+        iteration_to_period: Array mapping iteration index to period.
+        observed_factors: Array of shape (n_periods, n_obs, n_observed_factors) with
+            data on the observed factors.
 
     Returns:
-        dict: All data relevant for debugging, e.g. the log likelihood contribution of
-            each Kalman update and additional information like the filtered states.
+        All data relevant for debugging, e.g. the log likelihood contribution of
+        each Kalman update and additional information like the filtered states.
 
     """
     n_obs = measurements.shape[1]
-    states, upper_chols, log_mixture_weights, pardict = parse_params(
+    states, upper_chols, log_mixture_weights, parsed_params = parse_params(
         params,
         parsing_info,
         dimensions,
@@ -82,9 +88,9 @@ def log_likelihood(
 
     loop_args = {
         "period": iteration_to_period,
-        "loadings": pardict["loadings"],
-        "control_params": pardict["controls"],
-        "meas_sds": pardict["meas_sds"],
+        "loadings": parsed_params.loadings,
+        "control_params": parsed_params.controls,
+        "meas_sds": parsed_params.meas_sds,
         "measurements": measurements,
         "is_measurement_iteration": is_measurement_iteration,
         "is_predict_iteration": is_predict_iteration,
@@ -93,7 +99,7 @@ def log_likelihood(
     _body = functools.partial(
         _scan_body,
         controls=controls,
-        pardict=pardict,
+        parsed_params=parsed_params,
         sigma_scaling_factor=sigma_scaling_factor,
         sigma_weights=sigma_weights,
         transition_func=transition_func,
@@ -106,10 +112,10 @@ def log_likelihood(
     # possible.
     clipped = soft_clipping(
         arr=static_out["loglikes"],
-        lower=estimation_options["clipping_lower_bound"],
-        upper=estimation_options["clipping_upper_bound"],
-        lower_hardness=estimation_options["clipping_lower_hardness"],
-        upper_hardness=estimation_options["clipping_upper_hardness"],
+        lower=estimation_options.clipping_lower_bound,
+        upper=estimation_options.clipping_upper_bound,
+        lower_hardness=estimation_options.clipping_lower_hardness,
+        upper_hardness=estimation_options.clipping_upper_hardness,
     )
 
     value = clipped.sum()
@@ -142,15 +148,15 @@ def log_likelihood(
 
 
 def _scan_body(
-    carry,
-    loop_args,
-    controls,
-    pardict,
-    sigma_scaling_factor,
-    sigma_weights,
-    transition_func,
-    observed_factors,
-):
+    carry: dict[str, Array],
+    loop_args: dict[str, Array],
+    controls: Array,
+    parsed_params: ParsedParams,
+    sigma_scaling_factor: float,
+    sigma_weights: Array,
+    transition_func: Callable[..., Array],
+    observed_factors: Array,
+) -> tuple[dict[str, Array], dict[str, Any]]:
     # ==================================================================================
     # create arguments needed for update
     # ==================================================================================
@@ -188,12 +194,12 @@ def _scan_body(
         "upper_chols": upper_chols,
         "sigma_scaling_factor": sigma_scaling_factor,
         "sigma_weights": sigma_weights,
-        "trans_coeffs": {k: arr[t] for k, arr in pardict["transition"].items()},
-        "shock_sds": pardict["shock_sds"][t],
-        "anchoring_scaling_factors": pardict["anchoring_scaling_factors"][
+        "trans_coeffs": {k: arr[t] for k, arr in parsed_params.transition.items()},
+        "shock_sds": parsed_params.shock_sds[t],
+        "anchoring_scaling_factors": parsed_params.anchoring_scaling_factors[
             jnp.array([t, t + 1])
         ],
-        "anchoring_constants": pardict["anchoring_constants"][jnp.array([t, t + 1])],
+        "anchoring_constants": parsed_params.anchoring_constants[jnp.array([t, t + 1])],
         "observed_factors": observed_factors[t],
     }
 
@@ -219,29 +225,37 @@ def _scan_body(
     return new_state, static_out
 
 
-def _one_arg_measurement_update(kwargs):
-    out = kalman_update(**kwargs)
-    return out
+def _one_arg_measurement_update(
+    kwargs: dict[str, Any],
+) -> tuple[Array, Array, Array, Array, dict[str, Any]]:
+    return kalman_update(**kwargs)
 
 
-def _one_arg_anchoring_update(kwargs):
+def _one_arg_anchoring_update(
+    kwargs: dict[str, Any],
+) -> tuple[Array, Array, Array, Array, dict[str, Any]]:
     _, _, new_log_mixture_weights, new_loglikes, debug_info = kalman_update(**kwargs)
-    out = (
+    return (
         kwargs["states"],
         kwargs["upper_chols"],
         new_log_mixture_weights,
         new_loglikes,
         debug_info,
     )
-    return out
 
 
-def _one_arg_no_predict(kwargs, transition_func):  # noqa: ARG001
+def _one_arg_no_predict(
+    kwargs: dict[str, Any],
+    transition_func: Callable[..., Array],  # noqa: ARG001
+) -> tuple[Array, Array, Array]:
     """Just return the states cond chols without any changes."""
     return kwargs["states"], kwargs["upper_chols"], kwargs["states"]
 
 
-def _one_arg_predict(kwargs, transition_func):
+def _one_arg_predict(
+    kwargs: dict[str, Any],
+    transition_func: Callable[..., Array],
+) -> tuple[Array, Array, Array]:
     """Do a predict step but also return the input states as filtered states."""
     new_states, new_upper_chols = kalman_predict(
         transition_func,
diff --git a/src/skillmodels/maximization_inputs.py b/src/skillmodels/maximization_inputs.py
index 333804b..028f2ba 100644
--- a/src/skillmodels/maximization_inputs.py
+++ b/src/skillmodels/maximization_inputs.py
@@ -1,9 +1,15 @@
+"""Functions to create inputs for optimization of the log-likelihood."""
+
 import functools
+from collections.abc import Callable  # noqa: TC003
+from typing import TYPE_CHECKING, Any
 
 import jax
 import jax.numpy as jnp
 import numpy as np
 import pandas as pd
+from jax import Array
+from numpy.typing import NDArray  # noqa: TC002
 
 import skillmodels.likelihood_function as lf
 import skillmodels.likelihood_function_debug as lfd
@@ -19,75 +25,77 @@
 from skillmodels.process_data import process_data
 from skillmodels.process_debug_data import process_debug_data
 from skillmodels.process_model import process_model
+from skillmodels.types import ParsingInfo  # noqa: TC001
+
+if TYPE_CHECKING:
+    from skillmodels.types import ProcessedModel
 
 jax.config.update("jax_enable_x64", True)  # noqa: FBT003
 
 
-def get_maximization_inputs(model_dict, data, split_dataset=1):
+def get_maximization_inputs(
+    model_dict: dict,
+    data: pd.DataFrame,
+    split_dataset: int = 1,
+) -> dict[str, Any]:
     """Create inputs for optimagic's maximize function.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`
-        data (DataFrame): dataset in long format.
+        model_dict: The model specification. See: :ref:`model_specs`
+        data: dataset in long format.
         split_dataset(Int): Controls into how many sclices to split the dataset
             during the gradient computation.
 
     Returns a dictionary with keys:
-        loglike (function): A jax jitted function that takes an optimagic-style
+        loglike: A jax jitted function that takes an optimagic-style
             params dataframe as only input and returns a dict with entries:
             - "value": The scalar log likelihood
             - "contributions": An array with the log likelihood per observation
-        debug_loglike (function): Similar to loglike, with the following differences:
+        debug_loglike: Similar to loglike, with the following differences:
             - It is not jitted and thus faster on the first call and debuggable
             - It will add intermediate results as additional entries in the returned
               dictionary. Those can be used for debugging and plotting.
-        gradient (function): The gradient of the scalar log likelihood
+        gradient: The gradient of the scalar log likelihood
             function with respect to the parameters.
-        loglike_and_gradient (function): Combination of loglike and
+        loglike_and_gradient: Combination of loglike and
             loglike_gradient that is faster than calling the two functions separately.
-        constraints (list): List of optimagic constraints that are implied by the
+        constraints: List of optimagic constraints that are implied by the
             model specification.
-        params_template (pd.DataFrame): Parameter DataFrame with correct index and
+        params_template: Parameter DataFrame with correct index and
             bounds. The value column is empty except for the fixed constraints, which
             are set including the bounds.
-        data_aug (pd.DataFrame): DataFrame with augmented data. If model contains
+        data_aug: DataFrame with augmented data. If model contains
             endogenous factors, we double up the number of periods in order to add
 
-
-
     """
     model = process_model(model_dict)
     p_index = get_params_index(
-        update_info=model["update_info"],
-        labels=model["labels"],
-        dimensions=model["dimensions"],
-        transition_info=model["transition_info"],
-        endogenous_factors_info=model["endogenous_factors_info"],
+        update_info=model.update_info,
+        labels=model.labels,
+        dimensions=model.dimensions,
+        transition_info=model.transition_info,
+        endogenous_factors_info=model.endogenous_factors_info,
     )
 
     parsing_info = create_parsing_info(
         params_index=p_index,
-        update_info=model["update_info"],
-        labels=model["labels"],
-        anchoring=model["anchoring"],
-        has_endogenous_factors=model["endogenous_factors_info"][
-            "has_endogenous_factors"
-        ],
+        update_info=model.update_info,
+        labels=model.labels,
+        anchoring=model.anchoring,
+        has_endogenous_factors=model.endogenous_factors_info.has_endogenous_factors,
     )
     processed_data = process_data(
         df=data,
-        has_endogenous_factors=model["endogenous_factors_info"][
-            "has_endogenous_factors"
-        ],
-        labels=model["labels"],
-        update_info=model["update_info"],
-        anchoring_info=model["anchoring"],
+        has_endogenous_factors=model.endogenous_factors_info.has_endogenous_factors,
+        labels=model.labels,
+        update_info=model.update_info,
+        anchoring_info=model.anchoring,
         purpose="estimation",
     )
 
     sigma_scaling_factor, sigma_weights = calculate_sigma_scaling_factor_and_weights(
-        model["dimensions"]["n_latent_factors"],
-        model["estimation_options"]["sigma_points_scale"],
+        model.dimensions.n_latent_factors,
+        model.estimation_options.sigma_points_scale,
     )
 
     partialed_get_jnp_params_vec = functools.partial(
@@ -116,15 +124,17 @@ def get_maximization_inputs(model_dict, data, split_dataset=1):
     _jitted_loglikeobs = jax.jit(partialed_loglikes["llo"])
     _gradient = jax.jit(jax.grad(partialed_loglikes["ll"]))
 
-    def loglike(params):
+    def loglike(params: pd.DataFrame) -> float:
         params_vec = partialed_get_jnp_params_vec(params)
         return float(_jitted_loglike(params_vec))
 
-    def loglikeobs(params):
+    def loglikeobs(params: pd.DataFrame) -> NDArray[np.floating]:
         params_vec = partialed_get_jnp_params_vec(params)
         return _to_numpy(_jitted_loglikeobs(params_vec))
 
-    def loglike_and_gradient(params):
+    def loglike_and_gradient(
+        params: pd.DataFrame,
+    ) -> tuple[float, NDArray[np.floating]]:
         params_vec = partialed_get_jnp_params_vec(params)
         crit = float(_jitted_loglike(params_vec))
         n_obs = processed_data["measurements"].shape[1]
@@ -150,7 +160,7 @@ def loglike_and_gradient(params):
         grad = _to_numpy(_grad)
         return crit, grad
 
-    def debug_loglike(params):
+    def debug_loglike(params: pd.DataFrame) -> dict[str, Any]:
         params_vec = partialed_get_jnp_params_vec(params)
         jax_output = partialed_loglikes["debug_ll"](params_vec)
         tmp = _to_numpy(jax_output)
@@ -158,12 +168,12 @@ def debug_loglike(params):
         return process_debug_data(debug_data=tmp, model=model)
 
     _constraints_dicts = get_constraints_dicts(
-        dimensions=model["dimensions"],
-        labels=model["labels"],
-        anchoring_info=model["anchoring"],
-        update_info=model["update_info"],
-        normalizations=model["normalizations"],
-        endogenous_factors_info=model["endogenous_factors_info"],
+        dimensions=model.dimensions,
+        labels=model.labels,
+        anchoring_info=model.anchoring,
+        update_info=model.update_info,
+        normalizations=model.normalizations,
+        endogenous_factors_info=model.endogenous_factors_info,
     )
 
     constraints = constraints_dicts_to_om(_constraints_dicts)
@@ -171,16 +181,15 @@ def debug_loglike(params):
     params_template = pd.DataFrame(columns=["value"], index=p_index)
     params_template = add_bounds(
         params=params_template,
-        bounds_distance=model["estimation_options"]["bounds_distance"],
+        bounds_distance=model.estimation_options.bounds_distance,
     )
     params_template = enforce_fixed_constraints(
         params_template=params_template,
         constraints_dicts=_constraints_dicts,
     )
-    assert params_template.index.equals(p_index), (
-        "params_template index is not equal to p_index"
-    )
-    out = {
+    if not params_template.index.equals(p_index):
+        raise ValueError("params_template index is not equal to p_index")
+    return {
         "loglike": loglike,
         "loglikeobs": loglikeobs,
         "debug_loglike": debug_loglike,
@@ -189,20 +198,18 @@ def debug_loglike(params):
         "params_template": params_template,
     }
 
-    return out
-
 
 def _partial_some_log_likelihood(
-    fun,
-    parsing_info,
-    measurements,
-    controls,
-    observed_factors,
-    model,
-    sigma_weights,
-    sigma_scaling_factor,
-):
-    update_info = model["update_info"]
+    fun: Callable,
+    parsing_info: ParsingInfo,
+    measurements: Array,
+    controls: Array,
+    observed_factors: Array,
+    model: ProcessedModel,
+    sigma_weights: Array,
+    sigma_scaling_factor: Array,
+) -> Callable:
+    update_info = model.update_info
     is_measurement_iteration = (update_info["purpose"] == "measurement").to_numpy()
     _aug_periods = pd.Series(
         update_info.index.get_level_values("aug_period").to_numpy()
@@ -215,24 +222,25 @@ def _partial_some_log_likelihood(
     # are endogenous factors, the last aug_period is found at index -2 (there should not
     # be measurements for endogenous factors in the "second half" of the last period).
     last_aug_period = (
-        model["labels"]["aug_periods"][-2]
-        if parsing_info["has_endogenous_factors"]
-        else model["labels"]["aug_periods"][-1]
+        model.labels.aug_periods[-2]
+        if parsing_info.has_endogenous_factors
+        else model.labels.aug_periods[-1]
     )
     iteration_to_period = _aug_periods.replace(last_aug_period, -1).to_numpy()
-    assert max(iteration_to_period) == last_aug_period - 1
+    if max(iteration_to_period) != last_aug_period - 1:
+        raise ValueError("Unexpected iteration_to_period configuration")
 
     return functools.partial(
         fun,
         parsing_info=parsing_info,
         measurements=measurements,
         controls=controls,
-        transition_func=model["transition_info"]["func"],
+        transition_func=model.transition_info.func,
         sigma_scaling_factor=sigma_scaling_factor,
         sigma_weights=sigma_weights,
-        dimensions=model["dimensions"],
-        labels=model["labels"],
-        estimation_options=model["estimation_options"],
+        dimensions=model.dimensions,
+        labels=model.labels,
+        estimation_options=model.estimation_options,
         is_measurement_iteration=is_measurement_iteration,
         is_predict_iteration=is_predict_iteration,
         iteration_to_period=iteration_to_period,
@@ -240,7 +248,7 @@ def _partial_some_log_likelihood(
     )
 
 
-def _to_numpy(obj):
+def _to_numpy(obj: Any) -> Any:
     if isinstance(obj, dict):
         res = {}
         for key, value in obj.items():
@@ -257,7 +265,7 @@ def _to_numpy(obj):
     return res
 
 
-def _get_jnp_params_vec(params, target_index):
+def _get_jnp_params_vec(params: pd.DataFrame, target_index: pd.MultiIndex) -> Array:
     if set(params.index) != set(target_index):
         additional_entries = params.index.difference(target_index).tolist()
         missing_entries = target_index.difference(params.index).tolist()
@@ -268,5 +276,4 @@ def _get_jnp_params_vec(params, target_index):
             msg += f"Your params have missing entries: {missing_entries}. "
         raise ValueError(msg)
 
-    vec = jnp.array(params.reindex(target_index)["value"].to_numpy())
-    return vec
+    return jnp.array(params.reindex(target_index)["value"].to_numpy())
diff --git a/src/skillmodels/params_index.py b/src/skillmodels/params_index.py
index 586744f..e1d12b8 100644
--- a/src/skillmodels/params_index.py
+++ b/src/skillmodels/params_index.py
@@ -1,9 +1,25 @@
+"""Functions to construct the parameter index for model estimation."""
+
+from typing import TYPE_CHECKING
+
 import pandas as pd
 
+if TYPE_CHECKING:
+    from skillmodels.types import (
+        Dimensions,
+        EndogenousFactorsInfo,
+        Labels,
+        TransitionInfo,
+    )
+
 
 def get_params_index(
-    update_info, labels, dimensions, transition_info, endogenous_factors_info
-):
+    update_info: pd.DataFrame,
+    labels: Labels,
+    dimensions: Dimensions,
+    transition_info: TransitionInfo,
+    endogenous_factors_info: EndogenousFactorsInfo,
+) -> pd.MultiIndex:
     """Generate index for the params_df for optimagic.
 
     The index has four levels. The first is the parameter category. The second is the
@@ -12,60 +28,59 @@ def get_params_index(
     it contains an empty string.
 
     Args:
-        update_info (pandas.DataFrame): DataFrame with one row per Kalman update needed
+        update_info: DataFrame with one row per Kalman update needed
             in the likelihood function. See :ref:`update_info`.
-        labels (dict): Dict of lists with labels for the model quantities like
-            factors, periods, controls, stagemap and stages. See :ref:`labels`
-        options (dict): Tuning parameters for the estimation.
-            See :ref:`estimation_options`.
-        transition_info (dict): Information about the transition equations.
-        endogenous_factors_info (dict): Information about endogenous factors, if any.
+        labels: Labels for model quantities.
+        dimensions: Dimensional information.
+        transition_info: Information about the transition equations.
+        endogenous_factors_info: Information about endogenous factors, if any.
 
     Returns:
         params_index (pd.MultiIndex)
 
     """
     ind_tups = get_control_params_index_tuples(
-        controls=labels["controls"], update_info=update_info
+        controls=labels.controls, update_info=update_info
     )
     ind_tups += get_loadings_index_tuples(
-        factors=labels["latent_factors"], update_info=update_info
+        factors=labels.latent_factors, update_info=update_info
     )
     ind_tups += get_meas_sds_index_tuples(update_info=update_info)
     ind_tups += get_shock_sds_index_tuples(
-        aug_periods=labels["aug_periods"],
-        factors=labels["latent_factors"],
-        has_endogenous_factors=endogenous_factors_info["has_endogenous_factors"],
+        aug_periods=labels.aug_periods,
+        factors=labels.latent_factors,
+        has_endogenous_factors=endogenous_factors_info.has_endogenous_factors,
     )
     ind_tups += initial_mean_index_tuples(
-        n_mixtures=dimensions["n_mixtures"],
-        factors=labels["latent_factors"],
+        n_mixtures=dimensions.n_mixtures,
+        factors=labels.latent_factors,
     )
-    ind_tups += get_mixture_weights_index_tuples(n_mixtures=dimensions["n_mixtures"])
+    ind_tups += get_mixture_weights_index_tuples(n_mixtures=dimensions.n_mixtures)
     ind_tups += get_initial_cholcovs_index_tuples(
-        n_mixtures=dimensions["n_mixtures"],
-        factors=labels["latent_factors"],
+        n_mixtures=dimensions.n_mixtures,
+        factors=labels.latent_factors,
     )
     ind_tups += get_transition_index_tuples(
         transition_info=transition_info,
-        aug_periods=labels["aug_periods"],
-        has_endogenous_factors=endogenous_factors_info["has_endogenous_factors"],
+        aug_periods=labels.aug_periods,
+        has_endogenous_factors=endogenous_factors_info.has_endogenous_factors,
     )
 
-    index = pd.MultiIndex.from_tuples(
+    return pd.MultiIndex.from_tuples(
         ind_tups,
         names=["category", "aug_period", "name1", "name2"],
     )
-    return index
 
 
-def get_control_params_index_tuples(controls, update_info):
+def get_control_params_index_tuples(
+    controls: tuple[str, ...],
+    update_info: pd.DataFrame,
+) -> list[tuple[str, int, str, str]]:
     """Index tuples for control coeffs.
 
     Args:
-        controls (list): List of lists. There is one sublist per period which contains
-            the names of the control variables in that period. Constant not included.
-        update_info (pandas.DataFrame): DataFrame with one row per Kalman update needed
+        controls: Names of the control variables. Constant not included.
+        update_info: DataFrame with one row per Kalman update needed
             in the likelihood function. See :ref:`update_info`.
 
     """
@@ -76,19 +91,19 @@ def get_control_params_index_tuples(controls, update_info):
     return ind_tups
 
 
-def get_loadings_index_tuples(factors, update_info):
+def get_loadings_index_tuples(
+    factors: tuple[str, ...],
+    update_info: pd.DataFrame,
+) -> list[tuple[str, int, str, str]]:
     """Index tuples for loading.
 
     Args:
-        factors (list): The latent factors of the model
-        update_info (pandas.DataFrame): DataFrame with one row per Kalman update needed
+        factors: The latent factors of the model.
+        update_info: DataFrame with one row per Kalman update needed
             in the likelihood function. See :ref:`update_info`.
 
-    Returns:
-        ind_tups (list)
-
     """
-    mask = update_info[factors].to_numpy()
+    mask = update_info[list(factors)].to_numpy()
     ind_tups = []
     for i, (aug_period, meas) in enumerate(update_info.index):
         for f, factor in enumerate(factors):
@@ -97,16 +112,15 @@ def get_loadings_index_tuples(factors, update_info):
     return ind_tups
 
 
-def get_meas_sds_index_tuples(update_info):
+def get_meas_sds_index_tuples(
+    update_info: pd.DataFrame,
+) -> list[tuple[str, int, str, str]]:
     """Index tuples for meas_sd.
 
     Args:
-        update_info (pandas.DataFrame): DataFrame with one row per Kalman update needed
+        update_info: DataFrame with one row per Kalman update needed
             in the likelihood function. See :ref:`update_info`.
 
-    Returns:
-        ind_tups (list)
-
     """
     ind_tups = []
     for aug_period, meas in update_info.index:
@@ -114,15 +128,18 @@ def get_meas_sds_index_tuples(update_info):
     return ind_tups
 
 
-def get_shock_sds_index_tuples(aug_periods, factors, has_endogenous_factors):
+def get_shock_sds_index_tuples(
+    aug_periods: tuple[int, ...],
+    factors: tuple[str, ...],
+    *,
+    has_endogenous_factors: bool,
+) -> list[tuple[str, int, str, str]]:
     """Index tuples for shock_sd.
 
     Args:
-        aug_periods (list): The augmented periods of the model.
-        factors (list): The latent factors of the model.
-
-    Returns:
-        ind_tups (list)
+        aug_periods: The augmented periods of the model.
+        factors: The latent factors of the model.
+        has_endogenous_factors: Whether the model has endogenous factors.
 
     """
     end = -2 if has_endogenous_factors else -1
@@ -133,15 +150,15 @@ def get_shock_sds_index_tuples(aug_periods, factors, has_endogenous_factors):
     return ind_tups
 
 
-def initial_mean_index_tuples(n_mixtures, factors):
+def initial_mean_index_tuples(
+    n_mixtures: int,
+    factors: tuple[str, ...],
+) -> list[tuple[str, int, str, str]]:
     """Index tuples for initial_mean.
 
     Args:
-        n_mixtures (int): Number of elements in the mixture distribution of the factors.
-        factors (list): The latent factors of the model
-
-    Returns:
-        ind_tups (list)
+        n_mixtures: Number of elements in the mixture distribution of the factors.
+        factors: The latent factors of the model.
 
     """
     ind_tups = []
@@ -151,14 +168,13 @@ def initial_mean_index_tuples(n_mixtures, factors):
     return ind_tups
 
 
-def get_mixture_weights_index_tuples(n_mixtures):
+def get_mixture_weights_index_tuples(
+    n_mixtures: int,
+) -> list[tuple[str, int, str, str]]:
     """Index tuples for mixture_weight.
 
     Args:
-        n_mixtures (int): Number of elements in the mixture distribution of the factors.
-
-    Returns:
-        ind_tups (list)
+        n_mixtures: Number of elements in the mixture distribution of the factors.
 
     """
     ind_tups = []
@@ -167,15 +183,15 @@ def get_mixture_weights_index_tuples(n_mixtures):
     return ind_tups
 
 
-def get_initial_cholcovs_index_tuples(n_mixtures, factors):
+def get_initial_cholcovs_index_tuples(
+    n_mixtures: int,
+    factors: tuple[str, ...],
+) -> list[tuple[str, int, str, str]]:
     """Index tuples for initial_cov.
 
     Args:
-        n_mixtures (int): Number of elements in the mixture distribution of the factors.
-        factors (list): The latent factors of the model
-
-    Returns:
-        ind_tups (list)
+        n_mixtures: Number of elements in the mixture distribution of the factors.
+        factors: The latent factors of the model.
 
     """
     ind_tups = []
@@ -194,23 +210,23 @@ def get_initial_cholcovs_index_tuples(n_mixtures, factors):
     return ind_tups
 
 
-def get_transition_index_tuples(transition_info, aug_periods, has_endogenous_factors):
+def get_transition_index_tuples(
+    transition_info: TransitionInfo,
+    aug_periods: tuple[int, ...],
+    *,
+    has_endogenous_factors: bool,
+) -> list[tuple[str, int, str, str]]:
     """Index tuples for transition equation coefficients.
 
     Args:
-        latent_factors (list): The latent factors of the model
-        all_factors (list): The latent and observed factors of the model.
-        aug_periods (list): The augmented periods of the model
-        transition_names (list): name of the transition equation of each factor
-        has_endogenous_factors (bool): Whether the model has endogenous factors.
-
-    Returns:
-        ind_tups (list)
+        transition_info: Information about transition equations.
+        aug_periods: The augmented periods of the model.
+        has_endogenous_factors: Whether the model has endogenous factors.
 
     """
     end = -2 if has_endogenous_factors else -1
     ind_tups = []
-    for factor, names in transition_info["param_names"].items():
+    for factor, names in transition_info.param_names.items():
         for aug_period in aug_periods[:end]:
             for name in names:
                 ind_tups.append(("transition", aug_period, factor, name))
diff --git a/src/skillmodels/parse_params.py b/src/skillmodels/parse_params.py
index 2515be2..e176234 100644
--- a/src/skillmodels/parse_params.py
+++ b/src/skillmodels/parse_params.py
@@ -1,89 +1,105 @@
+"""Functions to parse parameter vectors into structured dictionaries."""
+
 import warnings
+from typing import TYPE_CHECKING
 
 import jax.numpy as jnp
 import numpy as np
 import pandas as pd
+from jax import Array
+
+from skillmodels.types import LoadingsParsingInfo, ParsedParams, ParsingInfo
+
+if TYPE_CHECKING:
+    from skillmodels.types import Anchoring, Dimensions, Labels
 
 
 def create_parsing_info(
-    params_index, update_info, labels, anchoring, has_endogenous_factors
-):
-    """Create a dictionary with information how the parameter vector has to be parsed.
+    params_index: pd.MultiIndex,
+    update_info: pd.DataFrame,
+    labels: Labels,
+    anchoring: Anchoring,
+    *,
+    has_endogenous_factors: bool,
+) -> ParsingInfo:
+    """Create a dataclass with information how the parameter vector has to be parsed.
 
     Args:
-        params_index (pandas.MultiIndex): It has the levels ["category", "aug_period",
+        params_index: It has the levels ["category", "aug_period",
             "name1", "name2"]
-        update_info (pandas.DataFrame): DataFrame with one row per Kalman update needed
+        update_info: DataFrame with one row per Kalman update needed
             in the likelihood function. See :ref:`update_info`.
-        labels (dict): Dict of lists with labels for the model quantities like
+        labels: Labels dataclass with labels for the model quantities like
             factors, periods, controls, stagemap and stages. See :ref:`labels`
-        anchoring (dict): Dictionary with anchoring settings.
-        has_endogenous_factors (bool): Whether the model includes endogenous factors.
+        anchoring: Anchoring dataclass with anchoring settings.
+        has_endogenous_factors: Whether the model includes endogenous factors.
 
     Returns:
-        dict: dictionary that maps model quantities to positions or slices of the
+        ParsingInfo dataclass that maps model quantities to positions or slices of the
             parameter vector.
 
     """
     range_sr = pd.Series(data=np.arange(len(params_index)), index=params_index)
 
-    parsing_info = {}
-
-    simple_ones = [
-        "initial_states",
-        "initial_cholcovs",
-        "mixture_weights",
-        "controls",
-        "meas_sds",
-        "shock_sds",
-    ]
-
-    for quantity in simple_ones:
-        parsing_info[quantity] = _get_positional_selector_from_loc(range_sr, quantity)
+    # Simple quantities
+    initial_states = _get_positional_selector_from_loc(range_sr, "initial_states")
+    initial_cholcovs = _get_positional_selector_from_loc(range_sr, "initial_cholcovs")
+    mixture_weights = _get_positional_selector_from_loc(range_sr, "mixture_weights")
+    controls = _get_positional_selector_from_loc(range_sr, "controls")
+    meas_sds = _get_positional_selector_from_loc(range_sr, "meas_sds")
+    shock_sds = _get_positional_selector_from_loc(range_sr, "shock_sds")
 
     # loadings:
-    mask = update_info[labels["latent_factors"]].to_numpy()
+    mask = update_info[list(labels.latent_factors)].to_numpy()
     helper = np.arange(mask.size).reshape(mask.shape)
     flat_indices = helper[mask]
 
-    parsing_info["loadings"] = {
-        "slice": _get_positional_selector_from_loc(range_sr, "loadings"),
-        "flat_indices": jnp.array(flat_indices),
-        "shape": mask.shape,
-        "size": mask.size,
-    }
-
-    # "trans_coeffs"
-    pos_dict = {}
-    for factor in labels["latent_factors"]:
-        helper = pd.DataFrame(index=params_index)
-        loc = helper.query(f"category == 'transition' & name1 == '{factor}'").index
-        pos_dict[factor] = _get_positional_selector_from_loc(range_sr, loc)
+    loadings = LoadingsParsingInfo(
+        slice=_get_positional_selector_from_loc(range_sr, "loadings"),
+        flat_indices=jnp.array(flat_indices),
+        shape=mask.shape,
+        size=mask.size,
+    )
 
-    parsing_info["transition"] = pos_dict
+    # transition coefficients
+    transition: dict[str, Array | slice] = {}
+    for factor in list(labels.latent_factors):
+        helper_df = pd.DataFrame(index=params_index)
+        loc = helper_df.query(f"category == 'transition' & name1 == '{factor}'").index
+        transition[factor] = _get_positional_selector_from_loc(range_sr, loc)
 
     # anchoring_scaling_factors
-    is_free_loading = update_info[labels["latent_factors"]].to_numpy()
+    is_free_loading = update_info[list(labels.latent_factors)].to_numpy()
     is_anchoring = (update_info["purpose"] == "anchoring").to_numpy().reshape(-1, 1)
     is_anchoring_loading = jnp.array(is_free_loading & is_anchoring)
-    parsing_info["is_anchoring_loading"] = is_anchoring_loading
-    parsing_info["is_anchored_factor"] = jnp.array(
-        update_info.query("purpose == 'anchoring'")[labels["latent_factors"]].any(
+    is_anchored_factor = jnp.array(
+        update_info.query("purpose == 'anchoring'")[list(labels.latent_factors)].any(
             axis=0,
         ),
     )
-    parsing_info["is_anchoring_update"] = is_anchoring.flatten()
-    parsing_info["ignore_constant_when_anchoring"] = anchoring[
-        "ignore_constant_when_anchoring"
-    ]
-
-    # Add has_endogenous_factors to parsing_info
-    parsing_info["has_endogenous_factors"] = has_endogenous_factors
-
-    return parsing_info
+    is_anchoring_update = jnp.array(is_anchoring.flatten())
+
+    return ParsingInfo(
+        initial_states=initial_states,
+        initial_cholcovs=initial_cholcovs,
+        mixture_weights=mixture_weights,
+        controls=controls,
+        meas_sds=meas_sds,
+        shock_sds=shock_sds,
+        loadings=loadings,
+        transition=transition,
+        is_anchoring_loading=is_anchoring_loading,
+        is_anchored_factor=is_anchored_factor,
+        is_anchoring_update=is_anchoring_update,
+        ignore_constant_when_anchoring=anchoring.ignore_constant_when_anchoring,
+        has_endogenous_factors=has_endogenous_factors,
+    )
 
 
-def _get_positional_selector_from_loc(range_sr, loc):
+def _get_positional_selector_from_loc(
+    range_sr: pd.Series,
+    loc: str | pd.MultiIndex | pd.Index,
+) -> Array | slice:
     with warnings.catch_warnings():
         warnings.filterwarnings(
             "ignore",
@@ -98,176 +114,219 @@ def _get_positional_selector_from_loc(range_sr, loc):
     return ilocs
 
 
-def parse_params(params, parsing_info, dimensions, labels, n_obs):
+def parse_params(
+    params: Array,
+    parsing_info: ParsingInfo,
+    dimensions: Dimensions,
+    labels: Labels,
+    n_obs: int,
+) -> tuple[Array, Array, Array, ParsedParams]:
     """Parse params into the quantities that depend on it.
 
     Args:
-        params (jax.numpy.array): 1d array with model parameters.
-        parsing_info (dict): Dictionary with information on how the parameters
+        params: 1d array with model parameters.
+        parsing_info: ParsingInfo dataclass with information on how the parameters
             have to be parsed.
-        dimensions (dict): Dimensional information like n_states, n_periods, n_controls,
+        dimensions: Dimensional information like n_states, n_periods, n_controls,
             n_mixtures. See :ref:`dimensions`.
-        n_obs (int): Number of observations.
+        labels: Labels dataclass with labels for the model quantities like
+            factors, periods, controls, stagemap and stages. See :ref:`labels`
+        n_obs: Number of observations.
 
     Returns:
-        jax.numpy.array: Array of shape (n_obs, n_mixtures, n_states) with initial
-            state estimates.
-        jax.numpy.array: Array of shape (n_obs, n_mixtures, n_states, n_states) with the
-            transpose of the lower triangular cholesky factors of the initial covariance
-            matrices.
-        jax.numpy.array: Array of shape (n_obs, n_mixtures) with the log of the initial
-            weight for each element in the finite mixture of normals.
-        dict: Dictionary with other parameters. It has the following key-value pairs:
-            - "control_params":
-            - "loadings":
-            - "meas_sds":
-            - "shock_sds":
-            - "trans_params":
-            - "anchoring_scaling_factors":
-            - "anchoring_constants":
+        Tuple of:
+            - Array of shape (n_obs, n_mixtures, n_states) with initial state estimates.
+            - Array of shape (n_obs, n_mixtures, n_states, n_states) with the transpose
+              of the lower triangular cholesky factors of the initial covariance
+              matrices.
+            - Array of shape (n_obs, n_mixtures) with the log of the initial weight for
+              each element in the finite mixture of normals.
+            - ParsedParams dataclass with other model parameters.
 
     """
     states = _get_initial_states(params, parsing_info, dimensions, n_obs)
     upper_chols = _get_initial_upper_chols(params, parsing_info, dimensions, n_obs)
     log_weights = _get_initial_log_mixture_weights(params, parsing_info, n_obs)
-    pardict = {
-        "controls": _get_control_params(params, parsing_info, dimensions),
-        "loadings": _get_loadings(params, parsing_info),
-        "meas_sds": _get_meas_sds(params, parsing_info),
-        "shock_sds": _get_shock_sds(params, parsing_info, dimensions),
-        "transition": _get_transition_params(params, parsing_info, labels),
-    }
-
-    pardict["anchoring_scaling_factors"] = _get_anchoring_scaling_factors(
-        pardict["loadings"],
+
+    controls = _get_control_params(params, parsing_info, dimensions)
+    loadings = _get_loadings(params, parsing_info)
+    meas_sds = _get_meas_sds(params, parsing_info)
+    shock_sds = _get_shock_sds(params, parsing_info, dimensions)
+    transition = _get_transition_params(params, parsing_info, labels)
+
+    anchoring_scaling_factors = _get_anchoring_scaling_factors(
+        loadings,
         parsing_info,
         dimensions,
     )
 
-    pardict["anchoring_constants"] = _get_anchoring_constants(
-        pardict["controls"],
+    anchoring_constants = _get_anchoring_constants(
+        controls,
         parsing_info,
         dimensions,
     )
 
-    return states, upper_chols, log_weights, pardict
+    parsed = ParsedParams(
+        controls=controls,
+        loadings=loadings,
+        meas_sds=meas_sds,
+        shock_sds=shock_sds,
+        transition=transition,
+        anchoring_scaling_factors=anchoring_scaling_factors,
+        anchoring_constants=anchoring_constants,
+    )
+
+    return states, upper_chols, log_weights, parsed
 
 
-def _get_initial_states(params, info, dimensions, n_obs):
+def _get_initial_states(
+    params: Array,
+    info: ParsingInfo,
+    dimensions: Dimensions,
+    n_obs: int,
+) -> Array:
     """Create the array of initial states."""
-    state = params[info["initial_states"]].reshape(
+    state = params[info.initial_states].reshape(
         1,
-        dimensions["n_mixtures"],
-        dimensions["n_latent_factors"],
+        dimensions.n_mixtures,
+        dimensions.n_latent_factors,
     )
     return jnp.repeat(state, n_obs, axis=0)
 
 
-def _get_initial_upper_chols(params, info, dimensions, n_obs):
+def _get_initial_upper_chols(
+    params: Array,
+    info: ParsingInfo,
+    dimensions: Dimensions,
+    n_obs: int,
+) -> Array:
     """Create the array with cholesky factors of the initial states covariance matrix.
 
     Note: The matrices contain the transpose of the lower triangular cholesky factors.
 
     """
-    n_states, n_mixtures = dimensions["n_latent_factors"], dimensions["n_mixtures"]
-    chol_params = params[info["initial_cholcovs"]].reshape(n_mixtures, -1)
+    n_states, n_mixtures = dimensions.n_latent_factors, dimensions.n_mixtures
+    chol_params = params[info.initial_cholcovs].reshape(n_mixtures, -1)
     upper_chols = jnp.zeros((n_obs, n_mixtures, n_states, n_states))
     for i in range(n_mixtures):
         filler = jnp.zeros((n_states, n_states))
-        filler = filler.at[jnp.tril_indices(n_states)].set(chol_params[i])
-        upper_chols = upper_chols.at[:, i].set(filler.T)
+        filler = filler.at[jnp.tril_indices(n_states)].set(chol_params[i])  # noqa: PD008
+        upper_chols = upper_chols.at[:, i].set(filler.T)  # noqa: PD008
     return upper_chols
 
 
-def _get_initial_log_mixture_weights(params, info, n_obs):
+def _get_initial_log_mixture_weights(
+    params: Array,
+    info: ParsingInfo,
+    n_obs: int,
+) -> Array:
     """Create the array with the log of initial mixture weights."""
-    log_weights = jnp.log(params[info["mixture_weights"]]).reshape(1, -1)
+    log_weights = jnp.log(params[info.mixture_weights]).reshape(1, -1)
     return jnp.repeat(log_weights, n_obs, axis=0)
 
 
-def _get_control_params(params, info, dimensions):
+def _get_control_params(
+    params: Array,
+    info: ParsingInfo,
+    dimensions: Dimensions,
+) -> Array:
     """Create the parameters for control variables in measurement equations."""
-    return params[info["controls"]].reshape(-1, dimensions["n_controls"])
+    return params[info.controls].reshape(-1, dimensions.n_controls)
 
 
-def _get_loadings(params, info):
+def _get_loadings(
+    params: Array,
+    info: ParsingInfo,
+) -> Array:
     """Create the array of factor loadings."""
-    info = info["loadings"]
-    free = params[info["slice"]]
-    extended = jnp.zeros(info["size"]).at[info["flat_indices"]].set(free)
-    out = extended.reshape(info["shape"])
-    return out
+    loadings_info = info.loadings
+    free = params[loadings_info.slice]
+    extended = jnp.zeros(loadings_info.size).at[loadings_info.flat_indices].set(free)  # noqa: PD008
+    return extended.reshape(loadings_info.shape)
 
 
-def _get_meas_sds(params, info):
+def _get_meas_sds(
+    params: Array,
+    info: ParsingInfo,
+) -> Array:
     """Create the array of standard deviations of the measurement errors."""
-    return params[info["meas_sds"]]
+    return params[info.meas_sds]
 
 
-def _get_shock_sds(params, info, dimensions):
+def _get_shock_sds(
+    params: Array,
+    info: ParsingInfo,
+    dimensions: Dimensions,
+) -> Array:
     """Create the array of standard deviations of the shocks in transition functions."""
-    return params[info["shock_sds"]].reshape(-1, dimensions["n_latent_factors"])
+    return params[info.shock_sds].reshape(-1, dimensions.n_latent_factors)
 
 
-def _get_transition_params(params, info, labels):
+def _get_transition_params(
+    params: Array,
+    info: ParsingInfo,
+    labels: Labels,
+) -> dict[str, Array]:
     """Create a list of arrays with transition equation parameters."""
     trans_params = {}
-    t_info = info["transition"]
-    n_aug_periods = len(labels["aug_periods"])
+    n_aug_periods = len(labels.aug_periods)
 
-    # Use has_endogenous_factors from parsing_info instead of undefined global
-    len_reduction = 2 if info["has_endogenous_factors"] else 1
+    len_reduction = 2 if info.has_endogenous_factors else 1
 
-    for factor in labels["latent_factors"]:
-        ilocs = t_info[factor]
+    for factor in list(labels.latent_factors):
+        ilocs = info.transition[factor]
         trans_params[factor] = params[ilocs].reshape(n_aug_periods - len_reduction, -1)
     return trans_params
 
 
-def _get_anchoring_scaling_factors(loadings, info, dimensions):
+def _get_anchoring_scaling_factors(
+    loadings: Array,
+    info: ParsingInfo,
+    dimensions: Dimensions,
+) -> Array:
     """Create an array of anchoring scaling factors.
 
     Note: Parameters are not taken from the parameter vector but from the loadings.
 
     """
     scaling_factors = jnp.ones(
-        (dimensions["n_aug_periods"], dimensions["n_latent_factors"]),
+        (dimensions.n_aug_periods, dimensions.n_latent_factors),
     )
-    free_anchoring_loadings = loadings[info["is_anchoring_loading"]].reshape(
-        dimensions["n_aug_periods"],
+    free_anchoring_loadings = loadings[info.is_anchoring_loading].reshape(
+        dimensions.n_aug_periods,
         -1,
     )
-    scaling_factors = scaling_factors.at[:, info["is_anchored_factor"]].set(
+    scaling_factors = scaling_factors.at[:, info.is_anchored_factor].set(  # noqa: PD008
         free_anchoring_loadings,
     )
 
     scaling_for_observed = jnp.ones(
-        (dimensions["n_aug_periods"], dimensions["n_observed_factors"]),
+        (dimensions.n_aug_periods, dimensions.n_observed_factors),
     )
 
-    scaling_factors = jnp.hstack([scaling_factors, scaling_for_observed])
+    return jnp.hstack([scaling_factors, scaling_for_observed])
 
-    return scaling_factors
 
-
-def _get_anchoring_constants(controls, info, dimensions):
+def _get_anchoring_constants(
+    controls: Array,
+    info: ParsingInfo,
+    dimensions: Dimensions,
+) -> Array:
     """Create an array of anchoring constants.
 
     Note: Parameters are not taken from the parameter vector but from the controls.
 
     """
-    constants = jnp.zeros((dimensions["n_aug_periods"], dimensions["n_latent_factors"]))
-    if not info["ignore_constant_when_anchoring"]:
-        values = controls[:, 0][info["is_anchoring_update"]].reshape(
-            dimensions["n_aug_periods"],
+    constants = jnp.zeros((dimensions.n_aug_periods, dimensions.n_latent_factors))
+    if not info.ignore_constant_when_anchoring:
+        values = controls[:, 0][info.is_anchoring_update].reshape(
+            dimensions.n_aug_periods,
             -1,
         )
-        constants = constants.at[:, info["is_anchored_factor"]].set(values)
+        constants = constants.at[:, info.is_anchored_factor].set(values)  # noqa: PD008
 
     constants_for_observed = jnp.zeros(
-        (dimensions["n_aug_periods"], dimensions["n_observed_factors"]),
+        (dimensions.n_aug_periods, dimensions.n_observed_factors),
     )
 
-    constants = jnp.hstack([constants, constants_for_observed])
-
-    return constants
+    return jnp.hstack([constants, constants_for_observed])
diff --git a/src/skillmodels/process_data.py b/src/skillmodels/process_data.py
index 41dc0b3..289b435 100644
--- a/src/skillmodels/process_data.py
+++ b/src/skillmodels/process_data.py
@@ -1,45 +1,52 @@
+"""Functions to process and prepare data for model estimation."""
+
 import warnings
-from typing import Any
+from typing import TYPE_CHECKING, Any, Literal
 
 import jax.numpy as jnp
 import numpy as np
 import pandas as pd
+from jax import Array
+
+if TYPE_CHECKING:
+    from skillmodels.types import Anchoring, Labels
 
 
 def process_data(
-    df,
-    has_endogenous_factors,
-    labels,
-    update_info,
-    anchoring_info,
-    purpose="estimation",
-):
+    df: pd.DataFrame,
+    *,
+    has_endogenous_factors: bool,
+    labels: Labels,
+    update_info: pd.DataFrame,
+    anchoring_info: Anchoring,
+    purpose: Literal["estimation", "anything", "simulation"] = "estimation",
+) -> dict[str, Any]:
     """Process the data for estimation.
 
     Args:
-        df (DataFrame): panel dataset in long format. It has a MultiIndex
+        df: panel dataset in long format. It has a MultiIndex
             where the first level indicates the period and the second the individual.
-        has_endogenous_factors (bool):
-        labels (dict): Dict of lists with labels for the model quantities like
+        has_endogenous_factors: Whether the model includes endogenous factors.
+        labels: Dict of lists with labels for the model quantities like
             factors, periods, controls, stagemap and stages. See :ref:`labels`
-        update_info (pandas.DataFrame): DataFrame with one row per Kalman update needed
+        update_info: DataFrame with one row per Kalman update needed
             in the likelihood function. See :ref:`update_info`.
-        anchoring_qinfo (dict): Information about anchoring. See :ref:`anchoring`
-        purpose (Literal["estimation", "anything"]): Whether the data is used for
+        anchoring_info: Information about anchoring. See :ref:`anchoring`
+        purpose: Whether the data is used for
             estimation (default, includes measurement data) or not.
 
     Returns a dictionary with keys:
-        measurements (jax.numpy.array): Array of shape (n_updates, n_obs) with data on
+        measurements: Array of shape (n_updates, n_obs) with data on
             observed measurements. NaN if the measurement was not observed. Only
             returned if estimation==True
-        controls (jax.numpy.array): Array of shape (n_periods, n_obs, n_controls) with
+        controls: Array of shape (n_periods, n_obs, n_controls) with
             observed control variables for the measurement equations.
-        observed_factors (jax.numpy.array): Array of shape
+        observed_factors: Array of shape
             (n_periods, n_obs, n_observed_factors) with data on the observed factors.
             Only returned if estimation==True
 
     """
-    df = pre_process_data(df, labels["periods"])
+    df = pre_process_data(df, labels.periods)
     df["constant"] = 1
     out = {}
 
@@ -50,8 +57,8 @@ def process_data(
         df.index = df.index.set_names(["id", "aug_period"])
 
     _check_data(df, update_info, labels, purpose=purpose)
-    n_obs = int(len(df) / len(labels["aug_periods"]))
-    df = _handle_controls_with_missings(df, labels["controls"], update_info)
+    n_obs = int(len(df) / len(labels.aug_periods))
+    df = _handle_controls_with_missings(df, labels.controls, update_info)
     out["controls"] = _generate_controls_array(df, labels, n_obs)
     out["observed_factors"] = _generate_observed_factor_array(df, labels, n_obs)
 
@@ -60,16 +67,20 @@ def process_data(
     return out
 
 
-def pre_process_data(df, periods):
+def pre_process_data(
+    df: pd.DataFrame,
+    periods: tuple[int, ...] | list[int],
+) -> pd.DataFrame:
     """Balance panel data in long format, drop unnecessary periods and set index.
 
     Args:
-        df (DataFrame): panel dataset in long format. It has a MultiIndex
+        df: panel dataset in long format. It has a MultiIndex
             where the first level indicates the period and the second
             the individual.
+        periods: The periods to keep in the balanced panel.
 
     Returns:
-        balanced (DataFrame): balanced panel. It has a MultiIndex. The first
+        balanced: balanced panel. It has a MultiIndex. The first
             enumerates individuals. The second level counts periods, starting at 0.
 
     """
@@ -80,28 +91,27 @@ def pre_process_data(df, periods):
     # replace existing codes for periods and
     df.index.names = ["id", "period"]
     for level in [0, 1]:
-        df.index = df.index.set_levels(range(len(df.index.levels[level])), level=level)
+        # df.index is a MultiIndex but typed as Index
+        df.index = df.index.set_levels(range(len(df.index.levels[level])), level=level)  # ty: ignore[unresolved-attribute]
 
     # create new index
     ids = sorted(df.index.get_level_values("id").unique())
     new_index = pd.MultiIndex.from_product([ids, periods], names=["id", "period"])
 
     # set new index
-    df = df.reindex(new_index)
-
-    return df
+    return df.reindex(new_index)
 
 
 def _get_period_data_for_endogenous_factors(
     aug_period: int,
     period: int,
     df: pd.DataFrame,
-    labels: dict[str, Any],
+    labels: Labels,
     update_info: pd.DataFrame,
 ) -> pd.DataFrame:
     meas = _get_period_measurements(update_info, aug_period)
-    controls = labels["controls"]
-    observed = labels["observed_factors"]
+    controls = labels.controls
+    observed = labels.observed_factors
 
     out = df.query(f"period == {period}")[
         [
@@ -120,9 +130,9 @@ def _get_period_data_for_endogenous_factors(
 
 def _augment_data_for_endogenous_factors(
     df: pd.DataFrame,
-    labels: dict[str, Any],
+    labels: Labels,
     update_info: pd.DataFrame,
-):
+) -> pd.DataFrame:
     """Make room for endogenous factors by doubling up the periods.
 
     Endogeneity means that current states influence the factor. Typically, this comes
@@ -134,8 +144,10 @@ def _augment_data_for_endogenous_factors(
     # Make sure datset is balanced
     n_ids = df["id"].nunique()
     n_periods = df["period"].nunique()
-    assert n_ids * n_periods == df.shape[0]
-    assert set(df["period"]) == set(labels["aug_periods_to_periods"].values())
+    if n_ids * n_periods != df.shape[0]:
+        raise ValueError("Dataset is not balanced: n_ids * n_periods != n_rows")
+    if set(df["period"]) != set(labels.aug_periods_to_periods.values()):
+        raise ValueError("Periods in data don't match expected periods")
 
     out = pd.concat(
         [
@@ -146,25 +158,33 @@ def _augment_data_for_endogenous_factors(
                 update_info=update_info,
                 labels=labels,
             )
-            for aug_period, period in labels["aug_periods_to_periods"].items()
+            for aug_period, period in labels.aug_periods_to_periods.items()
         ]
     )
     return out.set_index(["id", "aug_period"]).sort_index()
 
 
-def _add_copies_of_anchoring_outcome(df, anchoring_info):
+def _add_copies_of_anchoring_outcome(
+    df: pd.DataFrame,
+    anchoring_info: Anchoring,
+) -> pd.DataFrame:
     df = df.copy()
-    for factor in anchoring_info["factors"]:
-        outcome = anchoring_info["outcomes"][factor]
+    for factor in anchoring_info.factors:
+        outcome = anchoring_info.outcomes[factor]  # ty: ignore[invalid-argument-type]
         df[f"{outcome}_{factor}"] = df[outcome]
     return df
 
 
-def _check_data(df, update_info, labels, purpose):  # noqa: C901
+def _check_data(  # noqa: C901
+    df: pd.DataFrame,
+    update_info: pd.DataFrame,
+    labels: Labels,
+    purpose: Literal["estimation", "anything", "simulation"],
+) -> None:
     var_report = pd.DataFrame(index=update_info.index[:0], columns=["problem"])
-    for aug_period in labels["aug_periods"]:
+    for aug_period in labels.aug_periods:
         period_data = df.query(f"aug_period == {aug_period}")
-        for cont in labels["controls"]:
+        for cont in labels.controls:
             if cont not in period_data.columns or period_data[cont].isna().all():
                 var_report.loc[(aug_period, cont), "problem"] = "Variable is missing"
 
@@ -179,7 +199,7 @@ def _check_data(df, update_info, labels, purpose):  # noqa: C901
                         "Variable has no variance"
                     )
 
-        for factor in labels["observed_factors"]:
+        for factor in labels.observed_factors:
             if factor not in period_data.columns:
                 var_report.loc[(aug_period, factor), "problem"] = "Variable is missing"
             elif period_data[factor].isna().any():
@@ -193,12 +213,16 @@ def _check_data(df, update_info, labels, purpose):  # noqa: C901
         raise ValueError(var_report)
 
 
-def _handle_controls_with_missings(df, controls, update_info):
+def _handle_controls_with_missings(
+    df: pd.DataFrame,
+    controls: tuple[str, ...],
+    update_info: pd.DataFrame,
+) -> pd.DataFrame:
     aug_periods = update_info.index.get_level_values(0).unique().tolist()
     problematic_index = df.index[:0]
     for aug_period in aug_periods:
         period_data = df.query(f"aug_period == {aug_period}")
-        control_data = period_data[controls]
+        control_data = period_data[list(controls)]
         meas_data = period_data[_get_period_measurements(update_info, aug_period)]
         problem = control_data.isna().any(axis=1) & meas_data.notna().any(axis=1)
         problematic_index = problematic_index.union(period_data[problem].index)
@@ -207,12 +231,15 @@ def _handle_controls_with_missings(df, controls, update_info):
         old_names = df.loc[problematic_index][["__old_id__", "__old_period__"]]
         msg = "Set measurements to NaN because there are NaNs in the controls for:\n{}"
         msg = msg.format(list(map(tuple, old_names.to_numpy().tolist())))
-        warnings.warn(msg)
+        warnings.warn(msg, stacklevel=2)
         df.loc[problematic_index] = np.nan
     return df
 
 
-def _get_period_measurements(update_info, aug_period):
+def _get_period_measurements(
+    update_info: pd.DataFrame,
+    aug_period: int,
+) -> list[str]:
     if aug_period in update_info.index:
         measurements = list(update_info.loc[aug_period].index)
     else:
@@ -220,26 +247,38 @@ def _get_period_measurements(update_info, aug_period):
     return measurements
 
 
-def _generate_measurements_array(df, update_info, n_obs):
+def _generate_measurements_array(
+    df: pd.DataFrame,
+    update_info: pd.DataFrame,
+    n_obs: int,
+) -> Array:
     arr = np.zeros((len(update_info), n_obs))
     for k, (aug_period, var) in enumerate(update_info.index):
         arr[k] = df.query(f"aug_period == {aug_period}")[var].to_numpy()
     return jnp.array(arr, dtype="float32")
 
 
-def _generate_controls_array(df, labels, n_obs):
-    arr = np.zeros((len(labels["aug_periods"]), n_obs, len(labels["controls"])))
-    for aug_period in labels["aug_periods"]:
+def _generate_controls_array(
+    df: pd.DataFrame,
+    labels: Labels,
+    n_obs: int,
+) -> Array:
+    arr = np.zeros((len(labels.aug_periods), n_obs, len(labels.controls)))
+    for aug_period in labels.aug_periods:
         arr[aug_period] = df.query(f"aug_period == {aug_period}")[
-            labels["controls"]
+            list(labels.controls)
         ].to_numpy()
     return jnp.array(arr, dtype="float32")
 
 
-def _generate_observed_factor_array(df, labels, n_obs):
-    arr = np.zeros((len(labels["aug_periods"]), n_obs, len(labels["observed_factors"])))
-    for aug_period in labels["aug_periods"]:
+def _generate_observed_factor_array(
+    df: pd.DataFrame,
+    labels: Labels,
+    n_obs: int,
+) -> Array:
+    arr = np.zeros((len(labels.aug_periods), n_obs, len(labels.observed_factors)))
+    for aug_period in labels.aug_periods:
         arr[aug_period] = df.query(f"aug_period == {aug_period}")[
-            labels["observed_factors"]
+            list(labels.observed_factors)
         ].to_numpy()
     return jnp.array(arr, dtype="float32")
diff --git a/src/skillmodels/process_debug_data.py b/src/skillmodels/process_debug_data.py
index 1d5e7dd..6a57215 100644
--- a/src/skillmodels/process_debug_data.py
+++ b/src/skillmodels/process_debug_data.py
@@ -1,31 +1,44 @@
+"""Functions to process debug output from likelihood function into DataFrames."""
+
+from typing import TYPE_CHECKING, Any
+
 import numpy as np
 import pandas as pd
 
+if TYPE_CHECKING:
+    from jax import Array
+    from numpy.typing import NDArray
+
+    from skillmodels.types import ProcessedModel
+
 
-def process_debug_data(debug_data, model):
+def process_debug_data(
+    debug_data: dict[str, Any],
+    model: ProcessedModel,
+) -> dict[str, Any]:
     """Process the raw debug data into pandas objects that make visualization easy.
 
     Args:
-        debug_data (dict): Dictionary containing the following entries (
+        debug_data: Dictionary containing the following entries (
         and potentially others which are not modified):
-        - filtered_states (jax.numpy.array): Array of shape (n_updates, n_obs,
+        - filtered_states: Array of shape (n_updates, n_obs,
             n_mixtures, n_states) containing the filtered states after each Kalman
             update.
-        - initial_states (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states)
+        - initial_states: Array of shape (n_obs, n_mixtures, n_states)
             with the state estimates before the first Kalman update.
-        - residuals (jax.numpy.array): Array of shape (n_updates, n_obs, n_mixtures)
+        - residuals: Array of shape (n_updates, n_obs, n_mixtures)
             containing the residuals of a Kalman update.
-        - residual_sds (jax.numpy.ndarray): Array of shape (n_updates, n_obs,
+        - residual_sds: Array of shape (n_updates, n_obs,
             n_mixtures) containing the theoretical standard deviation of the residuals.
-        - all_contributions (jax.numpy.array): Array of shape (n_updates, n_obs) with
+        - all_contributions: Array of shape (n_updates, n_obs) with
             the likelihood contributions per update and individual.
-        - log_mixture_weights (jax.numpy.array): Array of shape (n_updates, n_obs,
+        - log_mixture_weights: Array of shape (n_updates, n_obs,
             n_mixtures) containing the log mixture weights after each update.
-        - initial_log_mixture_weights (jax.numpy.array): Array of shape (n_obs,
+        - initial_log_mixture_weights: Array of shape (n_obs,
             n_mixtures) containing the log mixture weights before the first
             kalman update.
 
-        model (dict): Processed model dictionary.
+        model: Processed model dictionary.
 
     Returns:
         dict: Dictionary with processed debug data. It has the following entries:
@@ -36,23 +49,23 @@ def process_debug_data(debug_data, model):
             after the last update of each period. The columns are the factor names,
             "period" and "id". The filtered states are already aggregated over
             mixture distributions.
-        - state_ranges (dict): The keys are the names of the latent factors.
+        - state_ranges: The keys are the names of the latent factors.
             The values are DataFrames with the columns "period", "minimum", "maximum".
             Note that this aggregates over mixture distributions.
-        - residuals (pd.DataFrame): Tidy DataFrame with residuals of each Kalman update.
+        - residuals: Tidy DataFrame with residuals of each Kalman update.
             Columns are "residual", "mixture", "period", "measurement" and "id".
             "period" and "measurement" identify the Kalman update to which the residual
             belongs.
-        - residual_sds (pd.DataFrame): As residuals but containing the theoretical
+        - residual_sds: As residuals but containing the theoretical
             standard deviation of the corresponding residual.
-        - all_contributions (pd.DataFrame): Tidy DataFrame with log likelihood
+        - all_contributions: Tidy DataFrame with log likelihood
             contribution per individual and Kalman Update. The columns are
             "contribution", "period", "measurement" and "id". "period" and "measurement"
             identify the Kalman Update to which the likelihood contribution corresponds.
 
     """
-    update_info = model["update_info"]
-    factors = model["labels"]["latent_factors"]
+    update_info = model.update_info
+    factors = model.labels.latent_factors
 
     post_update_states = _create_post_update_states(
         debug_data["filtered_states"],
@@ -93,7 +106,11 @@ def process_debug_data(debug_data, model):
     return res
 
 
-def _create_post_update_states(filtered_states, factors, update_info):
+def _create_post_update_states(
+    filtered_states: Array,
+    factors: tuple[str, ...],
+    update_info: pd.DataFrame,
+) -> pd.DataFrame:
     to_concat = []
     for (aug_period, meas), data in zip(
         update_info.index, filtered_states, strict=False
@@ -104,30 +121,36 @@ def _create_post_update_states(filtered_states, factors, update_info):
         df["measurement"] = meas
         to_concat.append(df)
 
-    post_states = pd.concat(to_concat)
-
-    return post_states
+    return pd.concat(to_concat)
 
 
-def _convert_state_array_to_df(arr, factor_names):
+def _convert_state_array_to_df(
+    arr: NDArray[np.floating[Any]],
+    factor_names: tuple[str, ...],
+) -> pd.DataFrame:
     """Convert a 3d state array into a 2d DataFrame.
 
     Args:
-        arr (np.ndarray): Array of shape (n_obs, n_mixtures, n_states)
-        factor_names (list): Names of the latent factors.
+        arr: Array of shape (n_obs, n_mixtures, n_states)
+        factor_names: Names of the latent factors.
     """
     n_obs, n_mixtures, n_states = arr.shape
-    df = pd.DataFrame(data=arr.reshape(-1, n_states), columns=factor_names)
+    df = pd.DataFrame(data=arr.reshape(-1, n_states), columns=list(factor_names))
     df["mixture"] = np.full((n_obs, n_mixtures), np.arange(n_mixtures)).flatten()
     return df
 
 
-def _create_filtered_states(filtered_states, log_mixture_weights, update_info, factors):
-    filtered_states = np.array(filtered_states)
-    log_mixture_weights = np.array(log_mixture_weights)
-    weights = np.exp(log_mixture_weights)
+def _create_filtered_states(
+    filtered_states: Array,
+    log_mixture_weights: Array,
+    update_info: pd.DataFrame,
+    factors: tuple[str, ...],
+) -> pd.DataFrame:
+    filtered_states_np = np.array(filtered_states)
+    log_mixture_weights_np = np.array(log_mixture_weights)
+    weights = np.exp(log_mixture_weights_np)
 
-    agg_states = (filtered_states * weights.reshape(*weights.shape, 1)).sum(axis=-2)
+    agg_states = (filtered_states_np * weights.reshape(*weights.shape, 1)).sum(axis=-2)
 
     keep = []
     for i, (aug_period, measurement) in enumerate(update_info.index):
@@ -145,25 +168,30 @@ def _create_filtered_states(filtered_states, log_mixture_weights, update_info, f
         df["id"] = np.arange(len(df))
         to_concat.append(df)
 
-    filtered_states = pd.concat(to_concat)
-
-    return filtered_states
+    return pd.concat(to_concat)
 
 
-def create_state_ranges(filtered_states, factors):
-    ranges = {}
+def create_state_ranges(
+    filtered_states: pd.DataFrame,
+    factors: tuple[str, ...] | list[str],
+) -> dict[str, pd.DataFrame]:
+    """Compute minimum and maximum state values for each factor by period."""
+    ranges: dict[str, pd.DataFrame] = {}
     # Group by whichever period column is present
     period_col = "aug_period" if "aug_period" in filtered_states.columns else "period"
     minima = filtered_states.groupby(period_col).min()
     maxima = filtered_states.groupby(period_col).max()
     for factor in factors:
         df = pd.concat([minima[factor], maxima[factor]], axis=1)
-        df.columns = ["minimum", "maximum"]
+        df.columns = pd.Index(["minimum", "maximum"])
         ranges[factor] = df
     return ranges
 
 
-def _process_residuals(residuals, update_info):
+def _process_residuals(
+    residuals: Array,
+    update_info: pd.DataFrame,
+) -> pd.DataFrame:
     to_concat = []
     n_obs, n_mixtures = residuals[0].shape
     for (aug_period, meas), data in zip(update_info.index, residuals, strict=False):
@@ -176,11 +204,17 @@ def _process_residuals(residuals, update_info):
     return pd.concat(to_concat)
 
 
-def _process_residual_sds(residual_sds, update_info):
+def _process_residual_sds(
+    residual_sds: Array,
+    update_info: pd.DataFrame,
+) -> pd.DataFrame:
     return _process_residuals(residual_sds, update_info)
 
 
-def _process_all_contributions(all_contributions, update_info):
+def _process_all_contributions(
+    all_contributions: Array,
+    update_info: pd.DataFrame,
+) -> pd.DataFrame:
     to_concat = []
     for (period, meas), contribs in zip(
         update_info.index, all_contributions, strict=False
diff --git a/src/skillmodels/process_model.py b/src/skillmodels/process_model.py
index dd1771a..9ceb17e 100644
--- a/src/skillmodels/process_model.py
+++ b/src/skillmodels/process_model.py
@@ -1,22 +1,41 @@
+"""Functions to process model specifications from user-friendly to internal form."""
+
 from copy import deepcopy
+from dataclasses import replace
 from functools import partial
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import pandas as pd
 from dags import concatenate_functions
 from dags.signature import rename_arguments
-from jax import vmap
+from frozendict import frozendict
+from jax import Array, vmap
 from pandas import DataFrame
 
 import skillmodels.transition_functions as t_f_module
 from skillmodels.check_model import check_model, check_stagemap
 from skillmodels.decorators import extract_params, jax_array_output
+from skillmodels.types import (
+    Anchoring,
+    Dimensions,
+    EndogenousFactorsInfo,
+    EstimationOptions,
+    FactorInfo,
+    FactorType,
+    Labels,
+    MeasurementType,
+    ProcessedModel,
+    TransitionInfo,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import KeysView, Mapping
 
 pd.set_option("future.no_silent_downcasting", True)  # noqa:  FBT003
 
 
-def process_model(model_dict):
+def process_model(model_dict: dict) -> ProcessedModel:
     """Check, clean, extend and transform the model specs.
 
     Check the completeness, consistency and validity of the model specifications.
@@ -24,19 +43,19 @@ def process_model(model_dict):
     Set default values and extend the model specification where necessary.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`
+        model_dict: The model specification. See: :ref:`model_specs`
 
     Returns:
         dict: nested dictionary of model specs. It has the following entries:
-        - dimensions (dict): Dimensional information like n_states, n_periods,
+        - dimensions: Dimensional information like n_states, n_periods,
           n_controls, n_mixtures. See :ref:`dimensions`.
-        - labels (dict): Dict of lists with labels for the model quantities like
+        - labels: Dict of lists with labels for the model quantities like
           factors, periods, controls, stagemap and stages. See :ref:`labels`
-        - anchoring (dict): Information about anchoring. See :ref:`anchoring`
-        - transition_info (dict): Everything related to transition functions.
-        - update_info (pandas.DataFrame): DataFrame with one row per Kalman update
+        - anchoring: Information about anchoring. See :ref:`anchoring`
+        - transition_info: Everything related to transition functions.
+        - update_info: DataFrame with one row per Kalman update
           needed in the likelihood function. See :ref:`update_info`.
-        - normalizations (dict): Nested dictionary with information on normalized factor
+        - normalizations: Nested dictionary with information on normalized factor
           loadings and intercepts for each factor. See :ref:`normalizations`.
 
     """
@@ -64,7 +83,28 @@ def process_model(model_dict):
         )
     else:
         _model_dict_aug = model_dict
-        endogenous_factors_info = {"has_endogenous_factors": has_endogenous_factors}
+        endogenous_factors_info = EndogenousFactorsInfo(
+            has_endogenous_factors=has_endogenous_factors,
+            aug_periods_to_aug_period_meas_types=frozendict(
+                _get_aug_periods_to_aug_period_meas_types(
+                    aug_periods=labels.aug_periods_to_periods.keys(),
+                    has_endogenous_factors=has_endogenous_factors,
+                )
+            ),
+            bounds_distance=model_dict["estimation_options"].get(
+                "bounds_distance", 1e-3
+            ),
+            aug_periods_from_period=partial(
+                _aug_periods_from_period,
+                aug_periods_to_periods=labels.aug_periods_to_periods,
+            ),
+            factor_info=frozendict(
+                {
+                    fac: FactorInfo(factor_type=FactorType.STATE)
+                    for fac in labels.latent_factors
+                }
+            ),
+        )
     check_model(
         model_dict=_model_dict_aug,
         labels=labels,
@@ -73,19 +113,20 @@ def process_model(model_dict):
         has_endogenous_factors=has_endogenous_factors,
     )
     transition_info = _get_transition_info(_model_dict_aug, labels)
-    labels["transition_names"] = list(transition_info["function_names"].values())
-
-    processed = {
-        "dimensions": dims,
-        "labels": labels,
-        "anchoring": anchoring,
-        "estimation_options": _process_estimation_options(_model_dict_aug),
-        "transition_info": transition_info,
-        "update_info": _get_update_info(_model_dict_aug, dims, labels, anchoring),
-        "normalizations": _process_normalizations(_model_dict_aug, dims, labels),
-        "endogenous_factors_info": endogenous_factors_info,
-    }
-    return processed
+    labels = replace(
+        labels, transition_names=tuple(transition_info.function_names.values())
+    )
+
+    return ProcessedModel(
+        dimensions=dims,
+        labels=labels,
+        anchoring=anchoring,
+        estimation_options=_process_estimation_options(_model_dict_aug),
+        transition_info=transition_info,
+        update_info=_get_update_info(_model_dict_aug, dims, labels, anchoring),
+        normalizations=_process_normalizations(_model_dict_aug, dims, labels),
+        endogenous_factors_info=endogenous_factors_info,
+    )
 
 
 def get_has_endogenous_factors(factors: dict[str, Any]) -> bool:
@@ -115,36 +156,33 @@ def get_has_endogenous_factors(factors: dict[str, Any]) -> bool:
     return endogenous_factors["is_endogenous"].any()  # ty: ignore[invalid-return-type]
 
 
-def get_dimensions(model_dict, has_endogenous_factors):
+def get_dimensions(model_dict: dict, *, has_endogenous_factors: bool) -> Dimensions:
     """Extract the dimensions of the model.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`
-        has_endogenous_factors (bool): Whether endogenous factors are present.
+        model_dict: The model specification. See: :ref:`model_specs`
+        has_endogenous_factors: Whether endogenous factors are present.
 
     Returns:
-        dict: Dimensional information like n_states, n_periods, n_controls,
-            n_mixtures. See :ref:`dimensions`.
+        Dimensions dataclass with all dimensional information.
 
     """
     all_n_periods = [len(d["measurements"]) for d in model_dict["factors"].values()]
     n_periods = max(all_n_periods)
     n_aug_periods = 2 * n_periods if has_endogenous_factors else n_periods
 
-    dims = {
-        "n_latent_factors": len(model_dict["factors"]),
-        "n_observed_factors": len(model_dict.get("observed_factors", [])),
-        "n_controls": len(model_dict.get("controls", [])) + 1,  # plus 1: constant
-        "n_mixtures": model_dict["estimation_options"].get("n_mixtures", 1),
-        "n_aug_periods": n_aug_periods,
-        "n_periods": n_periods,
-    }
-    dims["n_all_factors"] = dims["n_latent_factors"] + dims["n_observed_factors"]
-    return dims
+    return Dimensions(
+        n_latent_factors=len(model_dict["factors"]),
+        n_observed_factors=len(model_dict.get("observed_factors", [])),
+        n_controls=len(model_dict.get("controls", [])) + 1,  # plus 1: constant
+        n_mixtures=model_dict["estimation_options"].get("n_mixtures", 1),
+        n_aug_periods=n_aug_periods,
+        n_periods=n_periods,
+    )
 
 
 def _get_aug_periods_to_periods(
-    n_aug_periods: int, has_endogenous_factors: bool
+    n_aug_periods: int, *, has_endogenous_factors: bool
 ) -> dict[int, int]:
     """Return mapper of (potentially) augmented periods to user-provided periods."""
     aug_periods = list(range(n_aug_periods))
@@ -162,39 +200,39 @@ def _aug_periods_from_period(
     return [ap for ap, p in aug_periods_to_periods.items() if p == period]
 
 
-def _get_labels(model_dict, has_endogenous_factors, dimensions):
+def _get_labels(
+    model_dict: dict, *, has_endogenous_factors: bool, dimensions: Dimensions
+) -> Labels:
     """Extract labels of the model quantities.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`
-        has_endogenous_factors (bool): Whether endogenous factors are present.
-        dimensions (dict): Dimensional information like n_states, n_periods, n_controls,
-            n_mixtures. See :ref:`dimensions`.
+        model_dict: The model specification. See: :ref:`model_specs`
+        has_endogenous_factors: Whether endogenous factors are present.
+        dimensions: Dimensional information.
 
     Returns:
-        dict: Dict of lists with labels for the model quantities like
-        factors, periods, controls, stagemap and stages. See :ref:`labels`
+        Labels dataclass with all label information.
 
     """
     aug_periods_to_periods = _get_aug_periods_to_periods(
-        n_aug_periods=dimensions["n_aug_periods"],
+        n_aug_periods=dimensions.n_aug_periods,
         has_endogenous_factors=has_endogenous_factors,
     )
 
-    stagemap = model_dict.get("stagemap", list(range(dimensions["n_periods"] - 1)))
+    stagemap = model_dict.get("stagemap", list(range(dimensions.n_periods - 1)))
     stages = sorted(int(v) for v in np.unique(stagemap))
 
     report = check_stagemap(
         stagemap=stagemap,
         stages=stages,
-        n_periods=dimensions["n_periods"],
+        n_periods=dimensions.n_periods,
         is_augmented=False,
     )
     if report:
         raise ValueError(f"Invalid stage map: {report}")
     if has_endogenous_factors:
-        aug_stagemap = []
-        aug_stages_to_stages = {}
+        aug_stagemap: list[int] = []
+        aug_stages_to_stages: dict[int, int] = {}
         relevant_aug_periods = sorted(aug_periods_to_periods.keys())[:-2]
         for aug_p in relevant_aug_periods:
             p = aug_periods_to_periods[aug_p]
@@ -203,84 +241,89 @@ def _get_labels(model_dict, has_endogenous_factors, dimensions):
             aug_stagemap.append(aug_s)
             aug_stages_to_stages[aug_s] = s
     else:
-        aug_stagemap = stagemap
+        aug_stagemap = list(stagemap)
         aug_stages_to_stages = {s: s for s in stages}
 
-    labels = {
-        "latent_factors": list(model_dict["factors"]),
-        "observed_factors": list(model_dict.get("observed_factors", [])),
-        "controls": ["constant", *list(model_dict.get("controls", []))],
-        "periods": sorted(set(aug_periods_to_periods.values())),
-        "stagemap": stagemap,
-        "stages": stages,
-        "aug_periods": list(aug_periods_to_periods.keys()),
-        "aug_periods_to_periods": aug_periods_to_periods,
-        "aug_stagemap": aug_stagemap,
-        "aug_stages": sorted(int(v) for v in np.unique(aug_stagemap)),
-        "aug_stages_to_stages": aug_stages_to_stages,
-    }
-
-    labels["all_factors"] = labels["latent_factors"] + labels["observed_factors"]  # ty: ignore[unsupported-operator]
-
-    return labels
+    return Labels(
+        latent_factors=tuple(model_dict["factors"]),
+        observed_factors=tuple(model_dict.get("observed_factors", [])),
+        controls=("constant", *model_dict.get("controls", [])),
+        periods=tuple(sorted(set(aug_periods_to_periods.values()))),
+        stagemap=tuple(stagemap),
+        stages=tuple(stages),
+        aug_periods=tuple(aug_periods_to_periods.keys()),
+        aug_periods_to_periods=frozendict(aug_periods_to_periods),
+        aug_stagemap=tuple(aug_stagemap),
+        aug_stages=tuple(sorted(int(v) for v in np.unique(aug_stagemap))),
+        aug_stages_to_stages=frozendict(aug_stages_to_stages),
+    )
 
 
-def _process_estimation_options(model_dict):
+def _process_estimation_options(model_dict: dict) -> EstimationOptions:
     """Process options.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`
+        model_dict: The model specification. See: :ref:`model_specs`
 
     Returns:
-        dict: Tuning parameters for the estimation. See :ref:`options`.
+        EstimationOptions dataclass with tuning parameters for the estimation.
 
     """
-    default_options = {
-        "sigma_points_scale": 2,
-        "robust_bounds": True,
-        "bounds_distance": 1e-3,
-        "clipping_lower_bound": -1e30,
-        "clipping_upper_bound": None,
-        "clipping_lower_hardness": 1,
-        "clipping_upper_hardness": 1,
-    }
-    default_options.update(model_dict.get("estimation_options", {}))
-
-    if not default_options["robust_bounds"]:
-        default_options["bounds_distance"] = 0
-
-    return default_options
+    user_opts = model_dict.get("estimation_options", {})
+
+    sigma_points_scale = user_opts.get("sigma_points_scale", 2)
+    robust_bounds = user_opts.get("robust_bounds", True)
+    bounds_distance = user_opts.get("bounds_distance", 1e-3)
+    clipping_lower_bound = user_opts.get("clipping_lower_bound", -1e30)
+    clipping_upper_bound = user_opts.get("clipping_upper_bound", None)
+    clipping_lower_hardness = user_opts.get("clipping_lower_hardness", 1)
+    clipping_upper_hardness = user_opts.get("clipping_upper_hardness", 1)
+
+    if not robust_bounds:
+        bounds_distance = 0
+
+    return EstimationOptions(
+        sigma_points_scale=sigma_points_scale,
+        robust_bounds=robust_bounds,
+        bounds_distance=bounds_distance,
+        clipping_lower_bound=clipping_lower_bound,
+        clipping_upper_bound=clipping_upper_bound,
+        clipping_lower_hardness=clipping_lower_hardness,
+        clipping_upper_hardness=clipping_upper_hardness,
+    )
 
 
-def _process_anchoring(model_dict):
+def _process_anchoring(model_dict: dict) -> Anchoring:
     """Process the specification that governs how latent factors are anchored.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`
+        model_dict: The model specification. See: :ref:`model_specs`
 
     Returns:
-        dict: Dictionary with information about anchoring. See :ref:`anchoring`
+        Anchoring dataclass with information about anchoring.
 
     """
-    anchinfo = {
-        "anchoring": False,
-        "outcomes": {},
-        "factors": [],
-        "free_controls": False,
-        "free_constant": False,
-        "free_loadings": False,
-        "ignore_constant_when_anchoring": False,
-    }
-
     if "anchoring" in model_dict:
-        anchinfo.update(model_dict["anchoring"])
-        anchinfo["anchoring"] = True
-        anchinfo["factors"] = list(anchinfo["outcomes"])  # ty: ignore[invalid-argument-type]
+        anch = model_dict["anchoring"]
+        return Anchoring.from_config(
+            outcomes=anch.get("outcomes", {}),
+            free_controls=anch.get("free_controls", False),
+            free_constant=anch.get("free_constant", False),
+            free_loadings=anch.get("free_loadings", False),
+            ignore_constant_when_anchoring=anch.get(
+                "ignore_constant_when_anchoring", False
+            ),
+        )
 
-    return anchinfo
+    return Anchoring.disabled()
 
 
-def _insert_empty_elements_into_list(old, insert_at_modulo, to_insert, aug_p_to_p):
+def _insert_empty_elements_into_list(
+    old: list,
+    insert_at_modulo: int,
+    to_insert: Any,
+    aug_p_to_p: Mapping[int, int],
+) -> list:
     return [
         to_insert if aug_p % 2 == insert_at_modulo else old[p]
         for aug_p, p in aug_p_to_p.items()
@@ -288,16 +331,14 @@ def _insert_empty_elements_into_list(old, insert_at_modulo, to_insert, aug_p_to_
 
 
 def _augment_periods_for_endogenous_factors(
-    model_dict: dict[str, Any], dimensions: dict[str, Any], labels: dict[str, Any]
+    model_dict: dict[str, Any], dimensions: Dimensions, labels: Labels
 ) -> dict[str, Any]:
     """Augment periods if endogenous factors are present.
 
     Args:
         model_dict: The model specification. See: :ref:`model_specs`
-        dimensions (dict): Dimensional information like n_states, n_periods, n_controls,
-            n_mixtures. See :ref:`dimensions`.
-        labels (dict): Dict of lists with labels for the model quantities like
-            factors, periods, controls, stagemap and stages. See :ref:`labels`
+        dimensions: Dimensional information.
+        labels: Labels for model quantities.
 
     Returns:
         Model dictionary with twice the amount of periods
@@ -308,7 +349,7 @@ def _augment_periods_for_endogenous_factors(
         insert_at_modulo = 0 if v.get("is_endogenous", False) else 1
 
         # Insert empty elements into measurements when we do not have those.
-        if len(v["measurements"]) != dimensions["n_periods"]:
+        if len(v["measurements"]) != dimensions.n_periods:
             raise ValueError(
                 "Measurements must be of length `n_periods`, "
                 f"got {v['measurements']} for {fac}"
@@ -317,12 +358,12 @@ def _augment_periods_for_endogenous_factors(
             old=v["measurements"],
             insert_at_modulo=insert_at_modulo,
             to_insert=[],
-            aug_p_to_p=labels["aug_periods_to_periods"],
+            aug_p_to_p=labels.aug_periods_to_periods,
         )
 
         # Insert empty elements into normalizations when we do not have those.
         for norm_type, normalizations in v.get("normalizations", {}).items():
-            if not len(normalizations) == dimensions["n_periods"]:
+            if not len(normalizations) == dimensions.n_periods:
                 raise ValueError(
                     "Normalizations must be lists of length `n_periods`, "
                     f"got {normalizations} for {fac}['normalizations']['{norm_type}']"
@@ -332,17 +373,17 @@ def _augment_periods_for_endogenous_factors(
                     old=normalizations,
                     insert_at_modulo=insert_at_modulo,
                     to_insert={},
-                    aug_p_to_p=labels["aug_periods_to_periods"],
+                    aug_p_to_p=labels.aug_periods_to_periods,
                 )
             )
     return aug
 
 
-def _get_transition_info(model_dict, labels):
+def _get_transition_info(model_dict: dict, labels: Labels) -> TransitionInfo:
     """Collect information about transition functions."""
     func_list, param_names = [], []
-    latent_factors = labels["latent_factors"]
-    all_factors = labels["all_factors"]
+    latent_factors = labels.latent_factors
+    all_factors = labels.all_factors
 
     for factor in latent_factors:
         spec = model_dict["factors"][factor]["transition_function"]
@@ -376,10 +417,10 @@ def _get_transition_info(model_dict, labels):
 
     # add functions to produce the individual factors out of the 1d states vector.
     # The dag will automatically sort out what we don't need.
-    def _extract_factor(states, pos):
+    def _extract_factor(states: Array, pos: int) -> Array:
         return states[pos]
 
-    for i, factor in enumerate(labels["all_factors"]):
+    for i, factor in enumerate(labels.all_factors):
         functions[factor] = partial(_extract_factor, pos=i)
 
     transition_function = concatenate_functions(
@@ -397,93 +438,99 @@ def _extract_factor(states, pos):
         func = vmap(func, in_axes=(None, 0))
         individual_functions[factor] = func
 
-    out = {
-        "func": transition_function,
-        "param_names": dict(zip(latent_factors, param_names, strict=False)),
-        "individual_functions": individual_functions,
-        "function_names": dict(zip(latent_factors, function_names, strict=False)),
-    }
-    return out
+    return TransitionInfo(
+        func=transition_function,
+        param_names=frozendict(zip(latent_factors, param_names, strict=False)),
+        individual_functions=frozendict(individual_functions),
+        function_names=frozendict(zip(latent_factors, function_names, strict=False)),
+    )
 
 
 def _get_endogenous_factors_info(
+    *,
     has_endogenous_factors: bool,
     model_dict: dict[str, Any],
-    labels: dict[str, Any],
+    labels: Labels,
     bounds_distance: float,
-) -> dict[str, Any]:
+) -> EndogenousFactorsInfo:
     """Collect information about endogenous factors."""
-    endogenous_factors_info = {
-        "has_endogenous_factors": has_endogenous_factors,
-        "aug_periods_to_aug_period_meas_types": _get_aug_periods_to_aug_period_meas_types(  # noqa: E501
-            aug_periods=labels["aug_periods_to_periods"].keys(),
-            has_endogenous_factors=has_endogenous_factors,
+    factor_info = {}
+    for fac, v in model_dict["factors"].items():
+        factor_info[fac] = FactorInfo.from_flags(
+            is_endogenous=v.get("is_endogenous", False),
+            is_correction=v.get("is_correction", False),
+        )
+
+    return EndogenousFactorsInfo(
+        has_endogenous_factors=has_endogenous_factors,
+        aug_periods_to_aug_period_meas_types=frozendict(
+            _get_aug_periods_to_aug_period_meas_types(
+                aug_periods=labels.aug_periods_to_periods.keys(),
+                has_endogenous_factors=has_endogenous_factors,
+            )
         ),
-        "bounds_distance": bounds_distance,
-        "aug_periods_from_period": partial(
+        bounds_distance=bounds_distance,
+        aug_periods_from_period=partial(
             _aug_periods_from_period,
-            aug_periods_to_periods=labels["aug_periods_to_periods"],
+            aug_periods_to_periods=labels.aug_periods_to_periods,
         ),
-    }
-    for fac, v in model_dict["factors"].items():
-        endogenous_factors_info[fac] = {
-            "is_state": (
-                not v.get("is_endogenous", False) and not v.get("is_correction", False)
-            ),
-            "is_endogenous": v.get("is_endogenous", False),
-            "is_correction": v.get("is_correction", False),
-        }
-    return endogenous_factors_info
+        factor_info=frozendict(factor_info),
+    )
 
 
 def _get_aug_periods_to_aug_period_meas_types(
-    aug_periods: list[int], has_endogenous_factors: bool
-) -> dict[int, Literal["states", "endogenous_factors"]]:
+    aug_periods: tuple[int, ...] | KeysView[int],
+    *,
+    has_endogenous_factors: bool,
+) -> dict[int, MeasurementType]:
     if has_endogenous_factors:
         return {
-            aug_p: ("states" if aug_p % 2 == 0 else "endogenous_factors")
+            aug_p: (
+                MeasurementType.STATES
+                if aug_p % 2 == 0
+                else MeasurementType.ENDOGENOUS_FACTORS
+            )
             for aug_p in aug_periods
         }
-    return dict.fromkeys(aug_periods, "states")
+    return dict.fromkeys(aug_periods, MeasurementType.STATES)
 
 
-def _get_update_info(model_dict, dimensions, labels, anchoring_info):
+def _get_update_info(
+    model_dict: dict, dimensions: Dimensions, labels: Labels, anchoring_info: Anchoring
+) -> DataFrame:
     """Construct a DataFrame with information on each Kalman update.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`
-        dimensions (dict): Dimensional information like n_states, n_periods, n_controls,
-            n_mixtures. See :ref:`dimensions`.
-        labels (dict): Dict of lists with labels for the model quantities like
-            factors, periods, controls, stagemap and stages. See :ref:`labels`
-        anchoring_info (dict): Information about anchoring. See :ref:`anchoring`
+        model_dict: The model specification. See: :ref:`model_specs`
+        dimensions: Dimensional information.
+        labels: Labels for model quantities.
+        anchoring_info: Information about anchoring. See :ref:`anchoring`
 
     Returns:
-        pandas.DataFrame: DataFrame with one row per Kalman update needed in
-            the likelihood function. See :ref:`update_info`.
+        DataFrame with one row per Kalman update needed in the likelihood function.
 
     """
     index = pd.MultiIndex(
         levels=[[], []], codes=[[], []], names=["aug_period", "variable"]
     )
-    uinfo = DataFrame(index=index, columns=labels["latent_factors"] + ["purpose"])
+    uinfo = DataFrame(index=index, columns=[*labels.latent_factors, "purpose"])
 
     measurements = {}
-    for factor in labels["latent_factors"]:
+    for factor in labels.latent_factors:
         measurements[factor] = model_dict["factors"][factor]["measurements"]
-        if len(measurements[factor]) != dimensions["n_aug_periods"]:
+        if len(measurements[factor]) != dimensions.n_aug_periods:
             raise ValueError(
                 "Measurements must be of length `n_aug_periods`, "
                 f"got {measurements[factor]} for {factor}"
             )
 
-    for aug_period in labels["aug_periods"]:
-        for factor in labels["latent_factors"]:
+    for aug_period in labels.aug_periods:
+        for factor in labels.latent_factors:
             for meas in measurements[factor][aug_period]:
                 uinfo.loc[(aug_period, meas), factor] = True
                 uinfo.loc[(aug_period, meas), "purpose"] = "measurement"
-        for factor in anchoring_info["factors"]:
-            outcome = anchoring_info["outcomes"][factor]
+        for factor in anchoring_info.factors:
+            outcome = anchoring_info.outcomes[factor]  # ty: ignore[invalid-argument-type]
             name = f"{outcome}_{factor}"
             uinfo.loc[(aug_period, name), factor] = True
             uinfo.loc[(aug_period, name), "purpose"] = "anchoring"
@@ -493,30 +540,30 @@ def _get_update_info(model_dict, dimensions, labels, anchoring_info):
     return uinfo
 
 
-def _process_normalizations(model_dict, dimensions, labels):
+def _process_normalizations(
+    model_dict: dict, dimensions: Dimensions, labels: Labels
+) -> dict[str, dict[str, list]]:
     """Process the normalizations of intercepts and factor loadings.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`
-        dimensions (dict): Dimensional information like n_states, n_periods, n_controls,
-            n_mixtures. See :ref:`dimensions`.
-        labels (dict): Dict of lists with labels for the model quantities like
-            factors, periods, controls, stagemap and stages. See :ref:`labels`
+        model_dict: The model specification. See: :ref:`model_specs`
+        dimensions: Dimensional information.
+        labels: Labels for model quantities.
 
     Returns:
-        normalizations (dict): Nested dictionary with information on normalized factor
-            loadings and intercepts for each factor. See :ref:`normalizations`.
+        Nested dictionary with information on normalized factor loadings and
+        intercepts for each factor.
 
     """
     normalizations = {}
-    for factor in labels["latent_factors"]:
+    for factor in labels.latent_factors:
         normalizations[factor] = {}
         norminfo = model_dict["factors"][factor].get("normalizations", {})
         for norm_type in ["loadings", "intercepts"]:
             candidate = norminfo.get(
-                norm_type, [{} for _ in range(dimensions["n_aug_periods"])]
+                norm_type, [{} for _ in range(dimensions.n_aug_periods)]
             )
-            if not len(candidate) == dimensions["n_aug_periods"]:
+            if not len(candidate) == dimensions.n_aug_periods:
                 raise ValueError(
                     "Normalizations must be of length `n_aug_periods`, "
                     f"got {norminfo} for {factor}['{norm_type}']"
diff --git a/src/skillmodels/qr.py b/src/skillmodels/qr.py
index c690eac..834ff53 100644
--- a/src/skillmodels/qr.py
+++ b/src/skillmodels/qr.py
@@ -1,9 +1,12 @@
+"""Custom QR decomposition implementation optimized for GPU."""
+
 import jax
 import jax.numpy as jnp
+from jax import Array
 
 
 @jax.custom_jvp
-def qr_gpu(a: jax.Array):
+def qr_gpu(a: Array) -> tuple[Array, Array]:
     """Custom implementation of the QR Decomposition."""
     r, tau = jnp.linalg.qr(a, mode="raw")
 
@@ -11,7 +14,7 @@ def qr_gpu(a: jax.Array):
     return q, jnp.triu(r.mT[: tau.shape[0]])
 
 
-def _householder(r: jax.Array, tau: jax.Array):
+def _householder(r: Array, tau: Array) -> Array:
     """Custom implementation of the Householder Product.
 
     Uses the outputs of jnp.linalg.qr with mode = "raw" to calculate Q. This is needed
@@ -33,17 +36,17 @@ def _householder(r: jax.Array, tau: jax.Array):
     return h[:, :n]
 
 
-def _t(x: jax.Array) -> jax.Array:
+def _t(x: Array) -> Array:
     """Transpose batched Matrix."""
     return jax.lax.transpose(x, (*range(x.ndim - 2), x.ndim - 1, x.ndim - 2))
 
 
-def _h(x: jax.Array) -> jax.Array:
+def _h(x: Array) -> Array:
     """Hermitian Transpose of a Matrix."""
     return _t(x).conj()
 
 
-def _tril(m: jax.Array, k: int = 0) -> jax.Array:
+def _tril(m: Array, k: int = 0) -> Array:
     """Select lower Triangle of a Matrix."""
     *_, dim_n, dim_m = m.shape
     mask = jnp.tri(dim_n, dim_m, k, bool)
@@ -51,7 +54,10 @@ def _tril(m: jax.Array, k: int = 0) -> jax.Array:
 
 
 @qr_gpu.defjvp
-def qr_jvp_rule(primals, tangents):
+def qr_jvp_rule(
+    primals: tuple[Array],
+    tangents: tuple[Array],
+) -> tuple[tuple[Array, Array], tuple[Array, Array]]:
     """Calculates the derivative of the custom QR composition."""
     # See j-towns.github.io/papers/qr-derivative.pdf for a terse derivation.
     (x,) = primals
diff --git a/src/skillmodels/simulate_data.py b/src/skillmodels/simulate_data.py
index 3ae3879..f294cd0 100644
--- a/src/skillmodels/simulate_data.py
+++ b/src/skillmodels/simulate_data.py
@@ -1,39 +1,60 @@
 """Functions to simulate a dataset generated by a latent factor model."""
 
 import warnings
+from typing import TYPE_CHECKING
 
 import jax.numpy as jnp
 import numpy as np
 import pandas as pd
+from jax import Array
 from numpy.random import choice, multivariate_normal
 
 from skillmodels.filtered_states import anchor_states_df
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from numpy.typing import NDArray
+
+    from skillmodels.types import (
+        Dimensions,
+        EndogenousFactorsInfo,
+        Labels,
+        TransitionInfo,
+    )
 from skillmodels.kalman_filters import transform_sigma_points
 from skillmodels.params_index import get_params_index
 from skillmodels.parse_params import create_parsing_info, parse_params
 from skillmodels.process_data import process_data
 from skillmodels.process_debug_data import create_state_ranges
 from skillmodels.process_model import process_model
+from skillmodels.types import MeasurementType, ParsedParams
 
 
-def simulate_dataset(model_dict, params, n_obs=None, data=None, policies=None):
+def simulate_dataset(
+    model_dict: dict,
+    params: pd.DataFrame,
+    n_obs: int | None = None,
+    data: pd.DataFrame | None = None,
+    policies: list[dict] | None = None,
+) -> dict:
     """Simulate datasets generated by a latent factor model.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`
-        params (pandas.DataFrame): DataFrame with model parameters.
-        n_obs (int): Number of simulated individuals
-        data (pd.DataFrame): Dataset in the same format as for estimation, containing
+        model_dict: The model specification. See: :ref:`model_specs`
+        params: DataFrame with model parameters.
+        n_obs: Number of simulated individuals
+        data: Dataset in the same format as for estimation, containing
             information about observed factors and control variables.
-        policies (list): list of dictionaries. Each dictionary specifies a
+        policies: list of dictionaries. Each dictionary specifies a
             a stochastic shock to a latent factor AT THE END of "period" for "factor"
             with mean "effect_size" and "standard deviation"
 
     Returns:
-        observed_data (pd.DataFrame): Dataset with measurements and control variables
+        observed_data: Dataset with measurements and control variables
             in long format
 
-        latent_data (pd.DataFrame): Dataset with latent factors in long format
+        latent_data: Dataset with latent factors in long format
 
     """
     if data is None and n_obs is None:
@@ -41,23 +62,21 @@ def simulate_dataset(model_dict, params, n_obs=None, data=None, policies=None):
 
     model = process_model(model_dict)
 
-    if model["labels"]["observed_factors"] and data is None:
+    if model.labels.observed_factors and data is None:
         raise ValueError(
             "To simulate a model with observed factors, data cannot be None.",
         )
 
-    if model["labels"]["controls"] != ["constant"] and data is None:
+    if model.labels.controls != ["constant"] and data is None:
         raise ValueError("To simulate a model with controls, data cannot be None.")
 
     if data is not None:
         processed_data = process_data(
             df=data,
-            has_endogenous_factors=model["endogenous_factors_info"][
-                "has_endogenous_factors"
-            ],
-            labels=model["labels"],
-            update_info=model["update_info"],
-            anchoring_info=model["anchoring"],
+            has_endogenous_factors=model.endogenous_factors_info.has_endogenous_factors,
+            labels=model.labels,
+            update_info=model.update_info,
+            anchoring_info=model.anchoring,
             purpose="simulation",
         )
         control_data = processed_data["controls"]
@@ -68,39 +87,40 @@ def simulate_dataset(model_dict, params, n_obs=None, data=None, policies=None):
             warnings.warn(
                 f"The number of observations inferred from data ({data_n_obs}) and "
                 f"n_obs ({n_obs}) are different. n_obs is ignored.",
+                stacklevel=2,
             )
         n_obs = data_n_obs
 
     else:
         control_data = jnp.ones((n_obs, 1))
-        n_periods = model["dimensions"]["n_periods"]
+        n_periods = model.dimensions.n_periods
         observed_factors = jnp.zeros((n_periods, n_obs, 0))
 
     params_index = get_params_index(
-        update_info=model["update_info"],
-        labels=model["labels"],
-        dimensions=model["dimensions"],
-        transition_info=model["transition_info"],
-        endogenous_factors_info=model["endogenous_factors_info"],
+        update_info=model.update_info,
+        labels=model.labels,
+        dimensions=model.dimensions,
+        transition_info=model.transition_info,
+        endogenous_factors_info=model.endogenous_factors_info,
     )
 
     params = params.reindex(params_index)
 
     parsing_info = create_parsing_info(
-        params_index=params.index,
-        update_info=model["update_info"],
-        labels=model["labels"],
-        anchoring=model["anchoring"],
-        has_endogenous_factors=model["endogenous_factors_info"][
-            "has_endogenous_factors"
-        ],
+        params_index=params.index,  # ty: ignore[invalid-argument-type]
+        update_info=model.update_info,
+        labels=model.labels,
+        anchoring=model.anchoring,
+        has_endogenous_factors=model.endogenous_factors_info.has_endogenous_factors,
     )
 
-    states, covs, log_weights, pardict = parse_params(
+    if n_obs is None:
+        raise ValueError("n_obs must be set by either data or argument")
+    states, covs, log_weights, parsed_params = parse_params(
         params=jnp.array(params["value"].to_numpy()),
         parsing_info=parsing_info,
-        dimensions=model["dimensions"],
-        labels=model["labels"],
+        dimensions=model.dimensions,
+        labels=model.labels,
         n_obs=n_obs,
     )
 
@@ -108,26 +128,24 @@ def simulate_dataset(model_dict, params, n_obs=None, data=None, policies=None):
         latent_states=states,
         covs=covs,
         log_weights=log_weights,
-        pardict=pardict,
-        labels=model["labels"],
-        dimensions=model["dimensions"],
+        parsed_params=parsed_params,
+        labels=model.labels,
+        dimensions=model.dimensions,
         n_obs=n_obs,
-        has_endogenous_factors=model["endogenous_factors_info"][
-            "has_endogenous_factors"
-        ],
-        update_info=model["update_info"],
+        has_endogenous_factors=model.endogenous_factors_info.has_endogenous_factors,
+        update_info=model.update_info,
         control_data=control_data,
         observed_factors=observed_factors,
-        policies=policies,
-        transition_info=model["transition_info"],
+        policies=policies,  # ty: ignore[invalid-argument-type]
+        transition_info=model.transition_info,
     )
 
     # Create collapsed versions with user-facing periods
     latent_data = _collapse_aug_periods_to_periods(
         df=aug_latent_data,
-        factors=model["labels"]["latent_factors"],
-        aug_periods_to_periods=model["labels"]["aug_periods_to_periods"],
-        endogenous_factors_info=model["endogenous_factors_info"],
+        factors=model.labels.latent_factors,
+        aug_periods_to_periods=model.labels.aug_periods_to_periods,
+        endogenous_factors_info=model.endogenous_factors_info,
     )
 
     # Anchor the collapsed version (anchoring only works with period, not aug_period)
@@ -138,102 +156,116 @@ def simulate_dataset(model_dict, params, n_obs=None, data=None, policies=None):
         use_aug_period=False,
     )
 
-    out = {
+    return {
         "unanchored_states": {
             "states": latent_data,
             "state_ranges": create_state_ranges(
                 latent_data,
-                model["labels"]["latent_factors"],
+                model.labels.latent_factors,
             ),
         },
         "anchored_states": {
             "states": anchored_latent_data,
             "state_ranges": create_state_ranges(
                 anchored_latent_data,
-                model["labels"]["latent_factors"],
+                model.labels.latent_factors,
             ),
         },
         "aug_unanchored_states": {
             "states": aug_latent_data,
             "state_ranges": create_state_ranges(
                 aug_latent_data,
-                model["labels"]["latent_factors"],
+                model.labels.latent_factors,
             ),
         },
         "aug_measurements": aug_measurements,
     }
 
-    return out
-
 
 def _simulate_dataset(
-    latent_states,
-    covs,
-    log_weights,
-    pardict,
-    labels,
-    dimensions,
-    n_obs,
-    has_endogenous_factors,
-    update_info,
-    control_data,
-    observed_factors,
-    policies,
-    transition_info,
-):
+    latent_states: Array,
+    covs: Array,
+    log_weights: Array,
+    parsed_params: ParsedParams,
+    labels: Labels,
+    dimensions: Dimensions,
+    n_obs: int,
+    *,
+    has_endogenous_factors: bool,
+    update_info: pd.DataFrame,
+    control_data: Array,
+    observed_factors: Array,
+    policies: list[dict],
+    transition_info: TransitionInfo,
+) -> tuple[pd.DataFrame, pd.DataFrame]:
     """Simulate datasets generated by a latent factor model.
 
     Args:
-        See simulate_data
+        latent_states: Array of shape (n_obs, n_mixtures, n_states) with initial
+            state estimates.
+        covs: Array of shape (n_obs, n_mixtures, n_states, n_states) with initial
+            covariance matrices.
+        log_weights: Array of shape (n_obs, n_mixtures) with log mixture weights.
+        parsed_params: ParsedParams dataclass with parsed parameters.
+        labels: Labels for the model quantities like factors, periods, controls.
+        dimensions: Dimensional information like n_states, n_periods, n_controls.
+        n_obs: Number of observations.
+        has_endogenous_factors: Whether the model includes endogenous factors.
+        update_info: DataFrame with information on measurements for each period.
+        control_data: Array of shape (n_periods, n_obs, n_controls) with controls.
+        observed_factors: Array of shape (n_periods, n_obs, n_observed_factors).
+        policies: List of policy dictionaries specifying stochastic shocks.
+        transition_info: Information about transition functions.
 
     Returns:
-        See simulate_data
+        observed_data: DataFrame with simulated measurements.
+        latent_data: DataFrame with simulated latent factors.
 
     """
     policies = policies if policies is not None else []
 
-    n_states = dimensions["n_latent_factors"]
+    n_states = dimensions.n_latent_factors
     if has_endogenous_factors:
-        n_aug_periods = dimensions["n_aug_periods"] - 1
+        n_aug_periods = dimensions.n_aug_periods - 1
     else:
-        n_aug_periods = dimensions["n_aug_periods"]
+        n_aug_periods = dimensions.n_aug_periods
 
     weights = np.exp(log_weights)[0]
     loadings_df = pd.DataFrame(
-        data=pardict["loadings"],
+        data=parsed_params.loadings,
         index=update_info.index,
-        columns=labels["latent_factors"],
+        columns=labels.latent_factors,
     )
 
     control_params_df = pd.DataFrame(
-        data=pardict["controls"],
+        data=parsed_params.controls,
         index=update_info.index,
-        columns=labels["controls"],
+        columns=labels.controls,
     )
     meas_sds = pd.DataFrame(
-        data=pardict["meas_sds"].reshape(-1, 1),
+        data=parsed_params.meas_sds.reshape(-1, 1),
         index=update_info.index,
     )
 
-    transition_params = pardict["transition"]
-    shock_sds = pardict["shock_sds"]
+    transition_params = parsed_params.transition
+    shock_sds = parsed_params.shock_sds
 
     dist_args = []
-    for mixture in range(dimensions["n_mixtures"]):
+    for mixture in range(dimensions.n_mixtures):
         args = {
             "mean": latent_states[0][mixture],
             "cov": covs[0][mixture].T @ covs[0][mixture],
         }
         dist_args.append(args)
 
-    latent_states = np.zeros((n_aug_periods, n_obs, n_states))
+    latent_states = np.zeros((n_aug_periods, n_obs, n_states))  # ty: ignore[invalid-assignment]
     latent_states[0] = generate_start_states(n_obs, dimensions, dist_args, weights)
 
     for t in range(n_aug_periods - 1):
         # if there is a shock in period t, add it here
         policies_t = [p for p in policies if p["aug_period"] == t]
         for policy in policies_t:
-            position = labels["latent_factors"].index(policy["factor"])
+            position = labels.latent_factors.index(policy["factor"])
             latent_states[t, :, position] += _get_shock(
                 mean=policy["effect_size"],
                 sd=policy["standard_deviation"],
@@ -250,17 +282,17 @@ def _simulate_dataset(
         trans_coeffs = {k: arr[t] for k, arr in transition_params.items()}
 
         # get anchoring_scaling_factors for the period
-        anchoring_scaling_factors = pardict["anchoring_scaling_factors"][
+        anchoring_scaling_factors = parsed_params.anchoring_scaling_factors[
             jnp.array([t, t + 1])
         ]
         # get anchoring constants for the period
-        anchoring_constants = pardict["anchoring_constants"][jnp.array([t, t + 1])]
+        anchoring_constants = parsed_params.anchoring_constants[jnp.array([t, t + 1])]
 
         # call transform_sigma_points and convert result to numpy
         next_states = np.array(
             transform_sigma_points(
                 sigma_points=states,
-                transition_func=transition_info["func"],
+                transition_func=transition_info.func,
                 trans_coeffs=trans_coeffs,
                 anchoring_scaling_factors=anchoring_scaling_factors,
                 anchoring_constants=anchoring_constants,
@@ -281,8 +313,8 @@ def _simulate_dataset(
     for t in range(n_aug_periods):
         meas = pd.DataFrame(
             data=measurements_from_states(
-                latent_states[t],
-                control_data[t],
+                latent_states[t],  # ty: ignore[invalid-argument-type]
+                control_data[t],  # ty: ignore[invalid-argument-type]
                 loadings_df.loc[t].to_numpy(),
                 control_params_df.loc[t].to_numpy(),
                 meas_sds.loc[t].to_numpy().flatten(),
@@ -298,7 +330,7 @@ def _simulate_dataset(
 
     latent_data_by_period = []
     for t in range(n_aug_periods):
-        lat = pd.DataFrame(data=latent_states[t], columns=labels["latent_factors"])
+        lat = pd.DataFrame(data=latent_states[t], columns=labels.latent_factors)
         lat["aug_period"] = t
         latent_data_by_period.append(lat)
 
@@ -309,56 +341,66 @@ def _simulate_dataset(
 
 
 def _collapse_aug_periods_to_periods(
-    df, factors, aug_periods_to_periods, endogenous_factors_info
-):
+    df: pd.DataFrame,
+    factors: tuple[str, ...],
+    aug_periods_to_periods: Mapping[int, int],
+    endogenous_factors_info: EndogenousFactorsInfo,
+) -> pd.DataFrame:
     """Collapse dataframe with aug_period index to user-facing period index.
 
     For each factor, extracts from the appropriate aug_period based on is_endogenous.
 
     Args:
-        df (pd.DataFrame): DataFrame with columns "aug_period" and "id"
-        latent_factors (list): List of latent factors
-        aug_periods_to_periods (dict): Mapping from aug_period to period
-        endogenous_factors_info (dict): Information about which factors are endogenous
+        df: DataFrame with columns "aug_period" and "id"
+        factors: Tuple of latent factors
+        aug_periods_to_periods: Mapping from aug_period to period
+        endogenous_factors_info: Information about which factors are endogenous
 
     Returns:
         pd.DataFrame: DataFrame with "period" column instead of "aug_period"
     """
     df = df.copy()
-    if not endogenous_factors_info["has_endogenous_factors"]:
+    if not endogenous_factors_info.has_endogenous_factors:
         return df.rename(columns={"aug_period": "period"})
 
     df["period"] = df["aug_period"].map(aug_periods_to_periods)
     df["_aug_period_meas_type"] = df["aug_period"].map(
-        endogenous_factors_info["aug_periods_to_aug_period_meas_types"]
+        endogenous_factors_info.aug_periods_to_aug_period_meas_types
     )
 
     endogenous_cols = [
-        fac for fac in factors if endogenous_factors_info[fac]["is_endogenous"]
+        fac
+        for fac in factors
+        if endogenous_factors_info.factor_info[fac].is_endogenous  # ty: ignore[invalid-argument-type]
     ]
     state_cols = [fac for fac in factors if fac not in endogenous_cols]
 
-    out = df.query('_aug_period_meas_type == "endogenous_factors"')[
-        ["id", "period", *endogenous_cols]
-    ]
+    is_endogenous = df["_aug_period_meas_type"] == MeasurementType.ENDOGENOUS_FACTORS
+    is_states = df["_aug_period_meas_type"] == MeasurementType.STATES
+
+    out = df.loc[is_endogenous, ["id", "period", *endogenous_cols]]
     return pd.merge(
         out,
-        df.query('_aug_period_meas_type == "states"')[["id", "period", *state_cols]],
+        df.loc[is_states, ["id", "period", *state_cols]],
         on=["id", "period"],
         how="outer",
     )
 
 
-def _get_shock(mean, sd, size):
+def _get_shock(
+    mean: float,
+    sd: float,
+    size: int,
+) -> NDArray[np.floating]:
     """Add stochastic effect to a  factor of length n_obs.
 
     Args:
-        mean (float): mean of the stochastic effect
-        sd (float): standard deviation of the effect
-        size (int): length of resulting array
+        mean: mean of the stochastic effect
+        sd: standard deviation of the effect
+        size: length of resulting array
 
     Returns:
-        shock (np.array): 1d array of length n_obs with the stochastic shock
+        shock: 1d array of length n_obs with the stochastic shock
 
     """
     if sd == 0:
@@ -370,22 +412,28 @@ def _get_shock(mean, sd, size):
     return shock
 
 
-def generate_start_states(n_obs, dimensions, dist_args, weights):
+def generate_start_states(
+    n_obs: int,
+    dimensions: Dimensions,
+    dist_args: list[dict],
+    weights: NDArray[np.floating],
+) -> NDArray[np.floating]:
     """Draw initial states and control variables from a (mixture of) normals.
 
     Args:
-        n_obs (int): number of observations
-        dimensions (dict): Dimensional information like n_states, n_periods, n_controls,
+        n_obs: number of observations
+        dimensions: Dimensional information like n_states, n_periods, n_controls,
             n_mixtures. See :ref:`dimensions`.
-        dist_args (list): list of dicts of length nmixtures of dictionaries with the
+        dist_args: list of dicts of length nmixtures of dictionaries with the
             entries "mean" and "cov" for each mixture distribution.
+        weights: Array of mixture weights.
 
     Returns:
-        start_states (np.ndarray): shape (n_obs, n_states),
-        controls (np.ndarray): shape (n_obs, n_controls),
+        start_states: shape (n_obs, n_states),
+        controls: shape (n_obs, n_controls),
 
     """
-    n_states = dimensions["n_latent_factors"]
+    n_states = dimensions.n_latent_factors
     if np.size(weights) == 1:
         out = multivariate_normal(size=n_obs, **dist_args[0])
     else:
@@ -397,24 +445,30 @@ def generate_start_states(n_obs, dimensions, dist_args, weights):
     return out
 
 
-def measurements_from_states(states, controls, loadings, control_params, sds):
+def measurements_from_states(
+    states: NDArray[np.floating],
+    controls: NDArray[np.floating],
+    loadings: NDArray[np.floating],
+    control_params: NDArray[np.floating],
+    sds: NDArray[np.floating],
+) -> NDArray[np.floating]:
     """Generate the variables that would be observed in practice.
 
     This generates the data for only one period. Let n_meas be the number
     of measurements in that period.
 
     Args:
-        states (pd.DataFrame or np.ndarray): DataFrame of shape (n_obs, n_states)
-        controls (pd.DataFrame or np.ndarray): DataFrame of shape
-            (n_obs, n_controlsrols)
-        loadings (np.ndarray): numpy array of size (n_meas, n_states)
-        control_coeffs (np.ndarray): numpy array of size (n_meas, n_states)
-        sds (np.ndarray): numpy array of size (n_meas) with the standard deviations
+        states: DataFrame of shape (n_obs, n_states)
+        controls: DataFrame of shape
+            (n_obs, n_controls)
+        loadings: numpy array of size (n_meas, n_states)
+        control_params: numpy array of size (n_meas, n_controls)
+        sds: numpy array of size (n_meas) with the standard deviations
             of the measurements. Measurement error is assumed to be independent
             across measurements.
 
     Returns:
-        measurements (np.ndarray): array of shape (n_obs, n_meas) with measurements.
+        measurements: array of shape (n_obs, n_meas) with measurements.
 
     """
     n_meas = loadings.shape[0]
@@ -422,5 +476,4 @@ def measurements_from_states(states, controls, loadings, control_params, sds):
     epsilon = multivariate_normal([0] * n_meas, np.diag(sds**2), n_obs)
     states_part = np.dot(states, loadings.T)
     control_part = np.dot(controls, control_params.T)
-    meas = states_part + control_part + epsilon
-    return meas
+    return states_part + control_part + epsilon
diff --git a/tests/model2.yaml b/src/skillmodels/test_data/model2.yaml
similarity index 100%
rename from tests/model2.yaml
rename to src/skillmodels/test_data/model2.yaml
diff --git a/tests/model2_correct_params_index.csv b/src/skillmodels/test_data/model2_correct_params_index.csv
similarity index 100%
rename from tests/model2_correct_params_index.csv
rename to src/skillmodels/test_data/model2_correct_params_index.csv
diff --git a/tests/model2_correct_update_info.csv b/src/skillmodels/test_data/model2_correct_update_info.csv
similarity index 100%
rename from tests/model2_correct_update_info.csv
rename to src/skillmodels/test_data/model2_correct_update_info.csv
diff --git a/tests/model2_simulated_data.dta b/src/skillmodels/test_data/model2_simulated_data.dta
similarity index 100%
rename from tests/model2_simulated_data.dta
rename to src/skillmodels/test_data/model2_simulated_data.dta
diff --git a/tests/model2_with_endog_correct_update_info.csv b/src/skillmodels/test_data/model2_with_endog_correct_update_info.csv
similarity index 100%
rename from tests/model2_with_endog_correct_update_info.csv
rename to src/skillmodels/test_data/model2_with_endog_correct_update_info.csv
diff --git a/tests/simplest_augmented_data_expected.csv b/src/skillmodels/test_data/simplest_augmented_data_expected.csv
similarity index 100%
rename from tests/simplest_augmented_data_expected.csv
rename to src/skillmodels/test_data/simplest_augmented_data_expected.csv
diff --git a/tests/simplest_augmented_model.yaml b/src/skillmodels/test_data/simplest_augmented_model.yaml
similarity index 100%
rename from tests/simplest_augmented_model.yaml
rename to src/skillmodels/test_data/simplest_augmented_model.yaml
diff --git a/src/skillmodels/transition_functions.py b/src/skillmodels/transition_functions.py
index 1642707..e8d08fc 100644
--- a/src/skillmodels/transition_functions.py
+++ b/src/skillmodels/transition_functions.py
@@ -14,8 +14,6 @@
 
 Returns:
         * float
-
-
 **names_example_func(** *factors* **)**:
 
     Generate a list of names for the params of the transition function.
@@ -33,21 +31,26 @@
 
 import jax
 import jax.numpy as jnp
+from jax import Array
 
 
-def linear(states, params):
+def linear(states: Array, params: Array) -> Array:
     """Linear production function where the constant is the last parameter."""
     constant = params[-1]
     betas = params[:-1]
     return jnp.dot(states, betas) + constant
 
 
-def params_linear(factors):
+def params_linear(factors: tuple[str, ...]) -> list[str]:
     """Index tuples for linear transition function."""
     return [*factors, "constant"]
 
 
-def identity_constraints_linear(factor, aug_period, all_factors) -> list[dict]:
+def identity_constraints_linear(
+    factor: str,
+    aug_period: int,
+    all_factors: tuple[str, ...],
+) -> list[dict]:
     """Identity constraints for linear transition function."""
     constraints_dicts = []
     for regressor in params_linear(all_factors):
@@ -63,7 +66,7 @@ def identity_constraints_linear(factor, aug_period, all_factors) -> list[dict]:
     return constraints_dicts
 
 
-def translog(states, params):
+def translog(states: Array, params: Array) -> Array:
     """Translog transition function.
 
     The name is a convention in the skill formation literature even though the function
@@ -85,18 +88,21 @@ def translog(states, params):
     return res
 
 
-def params_translog(factors):
+def params_translog(factors: tuple[str, ...]) -> list[str]:
     """Index tuples for the translog production function."""
-    names = (
-        factors
+    return (
+        list(factors)
         + [f"{factor} ** 2" for factor in factors]
         + [f"{a} * {b}" for a, b in combinations(factors, 2)]
         + ["constant"]
     )
-    return names
 
 
-def identity_constraints_translog(factor, aug_period, all_factors) -> list[dict]:
+def identity_constraints_translog(
+    factor: str,
+    aug_period: int,
+    all_factors: tuple[str, ...],
+) -> list[dict]:
     """Identity constraints for translog transition function."""
     constraints_dicts = []
     for regressor in params_translog(all_factors):
@@ -112,7 +118,7 @@ def identity_constraints_translog(factor, aug_period, all_factors) -> list[dict]
     return constraints_dicts
 
 
-def log_ces(states, params):
+def log_ces(states: Array, params: Array) -> Array:
     """Log CES production function (KLS version)."""
     phi = params[-1]
     gammas = params[:-1]
@@ -124,38 +130,45 @@ def log_ces(states, params):
     # the log step for gammas underflows for gamma = 0, but this is handled correctly
     # by logsumexp and does not raise a warning.
     unscaled = jax.scipy.special.logsumexp(jnp.log(gammas) + states * phi)
-    result = unscaled * scaling_factor
-    return result
+    return unscaled * scaling_factor
 
 
-def params_log_ces(factors):
+def params_log_ces(factors: tuple[str, ...]) -> list[str]:
     """Index tuples for the log_ces production function."""
     return [*factors, "phi"]
 
 
-def constraints_log_ces(factor, factors, aug_period):
+def constraints_log_ces(
+    factor: str,
+    factors: tuple[str, ...],
+    aug_period: int,
+) -> dict:
     """Constraints for log_ces production function."""
     names = params_log_ces(factors)
     loc = [("transition", aug_period, factor, name) for name in names[:-1]]
     return {"loc": loc, "type": "probability"}
 
 
-def identity_constraints_log_ces(factors, aug_period, all_factors):
+def identity_constraints_log_ces(
+    factors: tuple[str, ...],
+    aug_period: int,
+    all_factors: tuple[str, ...],
+) -> list[dict]:
     """Identity constraints for log_ces."""
     raise NotImplementedError
 
 
-def constant(state, params):  # noqa: ARG001
+def constant(state: Array, params: Array) -> Array:  # noqa: ARG001
     """Constant production function."""
     return state
 
 
-def params_constant(factors):  # noqa: ARG001
+def params_constant(factors: tuple[str, ...]) -> list[str]:  # noqa: ARG001
     """Index tuples for the constant production function."""
     return []
 
 
-def robust_translog(states, params):
+def robust_translog(states: Array, params: Array) -> Array:
     """Numerically robust version of the translog transition function.
 
     This function does a clipping of the state vector at +- 1e12 before calling
@@ -171,16 +184,21 @@ def robust_translog(states, params):
     return translog(clipped_states, params)
 
 
-def params_robust_translog(factors):
+def params_robust_translog(factors: tuple[str, ...]) -> list[str]:
+    """Return parameter names for robust translog transition function."""
     return params_translog(factors)
 
 
-def identity_constraints_robust_translog(factor, aug_period, all_factors) -> list[dict]:
+def identity_constraints_robust_translog(
+    factor: str,
+    aug_period: int,
+    all_factors: tuple[str, ...],
+) -> list[dict]:
     """Identity constraints for robust_translog."""
     return identity_constraints_translog(factor, aug_period, all_factors)
 
 
-def linear_and_squares(states, params):
+def linear_and_squares(states: Array, params: Array) -> Array:
     """linear_and_squares transition function."""
     nfac = len(states)
     constant = params[-1]
@@ -193,14 +211,15 @@ def linear_and_squares(states, params):
     return res
 
 
-def params_linear_and_squares(factors):
+def params_linear_and_squares(factors: tuple[str, ...]) -> list[str]:
     """Index tuples for the linear_and_squares production function."""
-    names = factors + [f"{factor} ** 2" for factor in factors] + ["constant"]
-    return names
+    return list(factors) + [f"{factor} ** 2" for factor in factors] + ["constant"]
 
 
 def identity_constraints_linear_and_squares(
-    factor, aug_period, all_factors
+    factor: str,
+    aug_period: int,
+    all_factors: tuple[str, ...],
 ) -> list[dict]:
     """Identity constraints for linear_and_squares transition function."""
     constraints_dicts = []
@@ -217,7 +236,7 @@ def identity_constraints_linear_and_squares(
     return constraints_dicts
 
 
-def log_ces_general(states, params):
+def log_ces_general(states: Array, params: Array) -> Array:
     """Generalized log_ces production function without known location and scale."""
     n = states.shape[-1]
     tfp = params[-1]
@@ -230,15 +249,18 @@ def log_ces_general(states, params):
     # the log step for gammas underflows for gamma = 0, but this is handled correctly
     # by logsumexp and does not raise a warning.
     unscaled = jax.scipy.special.logsumexp(jnp.log(gammas) + states * sigmas)
-    result = unscaled * tfp
-    return result
+    return unscaled * tfp
 
 
-def params_log_ces_general(factors):
+def params_log_ces_general(factors: tuple[str, ...]) -> list[str]:
     """Index tuples for the generalized log_ces production function."""
-    return factors + [f"sigma_{fac}" for fac in factors] + ["tfp"]
+    return list(factors) + [f"sigma_{fac}" for fac in factors] + ["tfp"]
 
 
-def identity_constraints_log_ces_general(factors, aug_period, all_factors):
+def identity_constraints_log_ces_general(
+    factors: tuple[str, ...],
+    aug_period: int,
+    all_factors: tuple[str, ...],
+) -> list[dict]:
     """Identity constraints for log_ces_general."""
     raise NotImplementedError
diff --git a/src/skillmodels/types.py b/src/skillmodels/types.py
new file mode 100644
index 0000000..c1e01e3
--- /dev/null
+++ b/src/skillmodels/types.py
@@ -0,0 +1,309 @@
+"""Dataclass definitions for skillmodels internal data structures."""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import NewType
+
+import pandas as pd
+from frozendict import frozendict
+from jax import Array
+
+# NewType definitions for domain safety
+# These prevent accidentally mixing up semantically different int values
+Period = NewType("Period", int)
+AugPeriod = NewType("AugPeriod", int)
+Stage = NewType("Stage", int)
+AugStage = NewType("AugStage", int)
+
+
+class FactorType(Enum):
+    """Type of a latent factor in the model."""
+
+    STATE = auto()  # Regular state factor
+    ENDOGENOUS = auto()  # Endogenous factor (not a correction)
+    CORRECTION = auto()  # Correction factor (is_endogenous=True, is_correction=True)
+
+
+class MeasurementType(Enum):
+    """Type of measurement in an augmented period."""
+
+    STATES = auto()
+    ENDOGENOUS_FACTORS = auto()
+
+
+@dataclass(frozen=True)
+class Dimensions:
+    """Dimensional information for a skill formation model.
+
+    All fields represent counts of model components.
+    """
+
+    n_latent_factors: int
+    n_observed_factors: int
+    n_controls: int
+    n_mixtures: int
+    n_aug_periods: int
+    n_periods: int
+
+    @property
+    def n_all_factors(self) -> int:
+        """Total number of factors (latent + observed)."""
+        return self.n_latent_factors + self.n_observed_factors
+
+
+@dataclass(frozen=True)
+class Labels:
+    """Labels for model quantities.
+
+    Contains string identifiers for factors, periods, controls, and stages.
+    """
+
+    latent_factors: tuple[str, ...]
+    observed_factors: tuple[str, ...]
+    controls: tuple[str, ...]
+    periods: tuple[int, ...]
+    stagemap: tuple[int, ...]
+    stages: tuple[int, ...]
+    aug_periods: tuple[int, ...]
+    aug_periods_to_periods: frozendict[int, int]
+    aug_stagemap: tuple[int, ...]
+    aug_stages: tuple[int, ...]
+    aug_stages_to_stages: frozendict[int, int]
+    transition_names: tuple[str, ...] = ()
+
+    @property
+    def all_factors(self) -> tuple[str, ...]:
+        """All factor names (latent + observed)."""
+        return self.latent_factors + self.observed_factors
+
+
+@dataclass(frozen=True)
+class Anchoring:
+    """Information about how latent factors are anchored to observed outcomes."""
+
+    anchoring: bool
+    outcomes: frozendict[str, str]
+    factors: tuple[str, ...]
+    free_controls: bool
+    free_constant: bool
+    free_loadings: bool
+    ignore_constant_when_anchoring: bool
+
+    @classmethod
+    def disabled(cls) -> Anchoring:
+        """Create an Anchoring config with anchoring disabled."""
+        return cls(
+            anchoring=False,
+            outcomes=frozendict({}),
+            factors=(),
+            free_controls=False,
+            free_constant=False,
+            free_loadings=False,
+            ignore_constant_when_anchoring=False,
+        )
+
+    @classmethod
+    def from_config(
+        cls,
+        outcomes: dict[str, str],
+        *,
+        free_controls: bool = False,
+        free_constant: bool = False,
+        free_loadings: bool = False,
+        ignore_constant_when_anchoring: bool = False,
+    ) -> Anchoring:
+        """Create an Anchoring config from a configuration dictionary.
+
+        Args:
+            outcomes: Mapping from factor names to outcome variable names.
+            free_controls: Whether control parameters are free in anchoring equations.
+            free_constant: Whether constant is free in anchoring equations.
+            free_loadings: Whether loadings are free in anchoring equations.
+            ignore_constant_when_anchoring: Whether to ignore constant when anchoring.
+
+        Returns:
+            Configured Anchoring instance with anchoring enabled.
+
+        """
+        return cls(
+            anchoring=True,
+            outcomes=frozendict(outcomes),
+            factors=tuple(outcomes.keys()),
+            free_controls=free_controls,
+            free_constant=free_constant,
+            free_loadings=free_loadings,
+            ignore_constant_when_anchoring=ignore_constant_when_anchoring,
+        )
+
+
+@dataclass(frozen=True)
+class EstimationOptions:
+    """Tuning parameters for the estimation."""
+
+    sigma_points_scale: float
+    robust_bounds: bool
+    bounds_distance: float
+    clipping_lower_bound: float | None
+    clipping_upper_bound: float | None
+    clipping_lower_hardness: float
+    clipping_upper_hardness: float
+
+
+@dataclass(frozen=True)
+class TransitionInfo:
+    """Information about transition functions."""
+
+    func: Callable
+    param_names: frozendict[str, list[str]]
+    individual_functions: frozendict[str, Callable]
+    function_names: frozendict[str, str]
+
+
+@dataclass(frozen=True)
+class FactorInfo:
+    """Information for a single factor."""
+
+    factor_type: FactorType
+
+    @property
+    def is_state(self) -> bool:
+        """Whether the factor is a regular state factor."""
+        return self.factor_type == FactorType.STATE
+
+    @property
+    def is_endogenous(self) -> bool:
+        """Whether the factor is endogenous (ENDOGENOUS or CORRECTION)."""
+        return self.factor_type in (FactorType.ENDOGENOUS, FactorType.CORRECTION)
+
+    @property
+    def is_correction(self) -> bool:
+        """Whether the factor is a correction factor."""
+        return self.factor_type == FactorType.CORRECTION
+
+    @classmethod
+    def from_flags(
+        cls, *, is_endogenous: bool = False, is_correction: bool = False
+    ) -> FactorInfo:
+        """Create FactorInfo from boolean flags.
+
+        Args:
+            is_endogenous: Whether the factor is endogenous.
+            is_correction: Whether the factor is a correction (must be endogenous).
+
+        Returns:
+            FactorInfo with the appropriate FactorType.
+
+        Raises:
+            ValueError: If is_correction is True but is_endogenous is False.
+
+        """
+        if is_correction and not is_endogenous:
+            msg = "A correction factor must also be endogenous"
+            raise ValueError(msg)
+        if is_correction:
+            return cls(factor_type=FactorType.CORRECTION)
+        if is_endogenous:
+            return cls(factor_type=FactorType.ENDOGENOUS)
+        return cls(factor_type=FactorType.STATE)
+
+
+@dataclass(frozen=True)
+class EndogenousFactorsInfo:
+    """Information about endogenous factors in the model."""
+
+    has_endogenous_factors: bool
+    aug_periods_to_aug_period_meas_types: frozendict[int, MeasurementType]
+    bounds_distance: float
+    aug_periods_from_period: Callable[[int], list[int]]
+    factor_info: frozendict[str, FactorInfo]
+
+
+@dataclass(frozen=True)
+class ProcessedModel:
+    """Complete processed model specification.
+
+    This is the main output of process_model() containing all information
+    needed for estimation.
+    """
+
+    dimensions: Dimensions
+    labels: Labels
+    anchoring: Anchoring
+    estimation_options: EstimationOptions
+    transition_info: TransitionInfo
+    update_info: pd.DataFrame
+    normalizations: dict[str, dict[str, list]]
+    endogenous_factors_info: EndogenousFactorsInfo
+
+
+@dataclass(frozen=True)
+class LoadingsParsingInfo:
+    """Information for parsing factor loadings from parameter vector."""
+
+    slice: Array | slice
+    flat_indices: Array
+    shape: tuple[int, ...]
+    size: int
+
+
+@dataclass(frozen=True)
+class ParsingInfo:
+    """Information for parsing the parameter vector.
+
+    Maps model quantities to positions or slices of the parameter vector.
+    """
+
+    initial_states: Array | slice
+    initial_cholcovs: Array | slice
+    mixture_weights: Array | slice
+    controls: Array | slice
+    meas_sds: Array | slice
+    shock_sds: Array | slice
+    loadings: LoadingsParsingInfo
+    transition: dict[str, Array | slice]
+    is_anchoring_loading: Array
+    is_anchored_factor: Array
+    is_anchoring_update: Array
+    ignore_constant_when_anchoring: bool
+    has_endogenous_factors: bool
+
+
+@dataclass(frozen=True)
+class ParsedParams:
+    """Parsed parameters from the flat parameter vector.
+
+    Contains all model parameters in structured arrays.
+    """
+
+    controls: Array
+    loadings: Array
+    meas_sds: Array
+    shock_sds: Array
+    transition: dict[str, Array]
+    anchoring_scaling_factors: Array
+    anchoring_constants: Array
+
+
+@dataclass(frozen=True)
+class ProcessedData:
+    """Processed data arrays for estimation.
+
+    All arrays are JAX arrays ready for use in the likelihood function.
+    """
+
+    measurements: Array
+    controls: Array
+    observed_factors: Array
+
+
+@dataclass(frozen=True)
+class KalmanState:
+    """State carried through Kalman filter iterations.
+
+    Used as the carry state in jax.lax.scan.
+    """
+
+    states: Array
+    upper_chols: Array
+    log_mixture_weights: Array
diff --git a/src/skillmodels/utilities.py b/src/skillmodels/utilities.py
index dd4f385..411b1c1 100644
--- a/src/skillmodels/utilities.py
+++ b/src/skillmodels/utilities.py
@@ -1,5 +1,8 @@
+"""Utility functions for manipulating model specifications and parameters."""
+
 import warnings
 from copy import deepcopy
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -12,15 +15,19 @@
 )
 
 
-def extract_factors(factors, model_dict, params=None):
+def extract_factors(
+    factors: str | list[str],
+    model_dict: dict[str, Any],
+    params: pd.DataFrame | None = None,
+) -> dict[str, Any] | tuple[dict[str, Any], pd.DataFrame]:
     """Reduce a specification to a model with fewer latent factors.
 
     If provided, a params DataFrame is also reduced correspondingly.
 
     Args:
-        factors (str or list): Name(s) of the factor(s) to extract.
-        model_dict (dict): The model specification. See: :ref:`model_specs`.
-        params (pandas.DataFrame or None): The params DataFrame for the full model.
+        factors: Name(s) of the factor(s) to extract.
+        model_dict: The model specification. See: :ref:`model_specs`.
+        params: The params DataFrame for the full model.
 
     Returns:
         dict: The reduced model dictionary
@@ -30,17 +37,19 @@ def extract_factors(factors, model_dict, params=None):
     if isinstance(factors, str):
         factors = [factors]
 
-    to_remove = set(model_dict["factors"]).difference(factors)
-    out = remove_factors(to_remove, model_dict, params)
-    return out
+    to_remove = list(set(model_dict["factors"]).difference(factors))
+    return remove_factors(to_remove, model_dict, params)
 
 
-def update_parameter_values(params, others):
+def update_parameter_values(
+    params: pd.DataFrame,
+    others: pd.DataFrame | list[pd.DataFrame],
+) -> pd.DataFrame:
     """Update the "value" column of params with values from other.
 
     Args:
-        params (pandas.DataFrame or None): The params DataFrame for the full model.
-        others (pandas.DataFrame or list): Another DataFrame with parameters or list
+        params: The params DataFrame for the full model.
+        others: Another DataFrame with parameters or list
             of thereof. The values from other are used to update the value column
             of ``params``. If other is a list, the updates will be in order, i.e.
             later elements overwrite earlier ones.
@@ -67,7 +76,11 @@ def update_parameter_values(params, others):
     return out
 
 
-def remove_factors(factors, model_dict, params=None):
+def remove_factors(
+    factors: str | list[str],
+    model_dict: dict[str, Any],
+    params: pd.DataFrame | None = None,
+) -> dict[str, Any] | tuple[dict[str, Any], pd.DataFrame]:
     """Remove factors from a model specification.
 
     If provided, a params DataFrame is also reduced correspondingly.
@@ -76,9 +89,9 @@ def remove_factors(factors, model_dict, params=None):
     This happens if the remaining factors do not have measurements in later periods.
 
     Args:
-        factors (str or list): Name(s) of the factor(s) to remove.
-        model_dict (dict): The model specification. See: :ref:`model_specs`.
-        params (pandas.DataFrame or None): The params DataFrame for the full model.
+        factors: Name(s) of the factor(s) to remove.
+        model_dict: The model specification. See: :ref:`model_specs`.
+        params: The params DataFrame for the full model.
 
     Returns:
         dict: The reduced model dictionary
@@ -104,25 +117,35 @@ def remove_factors(factors, model_dict, params=None):
     # Remove periods if necessary, but only if no endogenous factors are present.
     # (else we would mess up the mapping between raw periods model periods)
     if not has_endogenous_factors:
-        new_n_periods = get_dimensions(out, has_endogenous_factors)["n_periods"]
+        new_n_periods = get_dimensions(
+            out, has_endogenous_factors=has_endogenous_factors
+        ).n_periods
         out = reduce_n_periods(out, new_n_periods)
 
     if params is not None:
-        out_params = _reduce_params(params, out, has_endogenous_factors)
+        out_params = _reduce_params(
+            params,
+            out,  # ty: ignore[invalid-argument-type]
+            has_endogenous_factors=has_endogenous_factors,
+        )
         out = (out, out_params)
 
-    return out
+    return out  # ty: ignore[invalid-return-type]
 
 
-def remove_measurements(measurements, model_dict, params=None):
+def remove_measurements(
+    measurements: str | list[str],
+    model_dict: dict[str, Any],
+    params: pd.DataFrame | None = None,
+) -> dict[str, Any] | tuple[dict[str, Any], pd.DataFrame]:
     """Remove measurements from a model specification.
 
     If provided, a params DataFrame is also reduced correspondingly.
 
     Args:
-        measurements (str or list): Name(s) of the measurement(s) to remove.
-        model_dict (dict): The model specification. See: :ref:`model_specs`.
-        params (pandas.DataFrame or None): The params DataFrame for the full model.
+        measurements: Name(s) of the measurement(s) to remove.
+        model_dict: The model specification. See: :ref:`model_specs`.
+        params: The params DataFrame for the full model.
 
     Returns:
         dict: The reduced model dictionary
@@ -161,15 +184,19 @@ def remove_measurements(measurements, model_dict, params=None):
     return out
 
 
-def remove_controls(controls, model_dict, params=None):
+def remove_controls(
+    controls: str | list[str],
+    model_dict: dict[str, Any],
+    params: pd.DataFrame | None = None,
+) -> dict[str, Any] | tuple[dict[str, Any], pd.DataFrame]:
     """Remove control variables from a model specification.
 
     If provided, a params DataFrame is also reduced correspondingly.
 
     Args:
-        controls (str or list): Name(s) of the contral variable(s) to remove.
-        model_dict (dict): The model specification. See: :ref:`model_specs`.
-        params (pandas.DataFrame or None): The params DataFrame for the full model.
+        controls: Name(s) of the contral variable(s) to remove.
+        model_dict: The model specification. See: :ref:`model_specs`.
+        params: The params DataFrame for the full model.
 
     Returns:
         dict: The reduced model dictionary
@@ -189,14 +216,17 @@ def remove_controls(controls, model_dict, params=None):
     return out
 
 
-def switch_translog_to_linear(model_dict, params=None):
+def switch_translog_to_linear(
+    model_dict: dict[str, Any],
+    params: pd.DataFrame | None = None,
+) -> dict[str, Any] | tuple[dict[str, Any], pd.DataFrame]:
     """Switch all translog production functions to linear.
 
     If provided, a params DataFrame is also reduced correspondingly.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`.
-        params (pandas.DataFrame or None): The params DataFrame for the full model.
+        model_dict: The model specification. See: :ref:`model_specs`.
+        params: The params DataFrame for the full model.
 
     Returns:
         dict: The reduced model dictionary
@@ -216,7 +246,10 @@ def switch_translog_to_linear(model_dict, params=None):
     return out
 
 
-def switch_linear_to_translog(model_dict, params=None):
+def switch_linear_to_translog(
+    model_dict: dict[str, Any],
+    params: pd.DataFrame | None = None,
+) -> dict[str, Any] | tuple[dict[str, Any], pd.DataFrame]:
     """Switch all linear production functions to translog.
 
     If provided, a params DataFrame is also extended correspondingly. The fill value
@@ -225,8 +258,8 @@ def switch_linear_to_translog(model_dict, params=None):
     the additional parameters are not initialized at zero.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`.
-        params (pandas.DataFrame or None): The params DataFrame for the full model.
+        model_dict: The model specification. See: :ref:`model_specs`.
+        params: The params DataFrame for the full model.
 
     Returns:
         dict: The reduced model dictionary
@@ -244,13 +277,17 @@ def switch_linear_to_translog(model_dict, params=None):
     return out
 
 
-def reduce_n_periods(model_dict, new_n_periods, params=None):
+def reduce_n_periods(
+    model_dict: dict[str, Any],
+    new_n_periods: int,
+    params: pd.DataFrame | None = None,
+) -> dict[str, Any] | tuple[dict[str, Any], pd.DataFrame]:
     """Remove all periods after n_periods.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`.
-        new_n_periods (int): The new number of periods.
-        params (pandas.DataFrame or None): The params DataFrame for the full model.
+        model_dict: The model specification. See: :ref:`model_specs`.
+        new_n_periods: The new number of periods.
+        params: The params DataFrame for the full model.
 
     Returns:
         dict: The reduced model dictionary
@@ -285,29 +322,40 @@ def reduce_n_periods(model_dict, new_n_periods, params=None):
     return out
 
 
-def _remove_from_list(list_, to_remove):
+def _remove_from_list(
+    list_: list[Any],
+    to_remove: str | list[str],
+) -> list[Any]:
     if isinstance(to_remove, str):
         to_remove = [to_remove]
     return [element for element in list_ if element not in to_remove]
 
 
-def _remove_from_dict(dict_, to_remove):
+def _remove_from_dict(
+    dict_: dict[str, Any],
+    to_remove: str | list[str],
+) -> dict[str, Any]:
     if isinstance(to_remove, str):
         to_remove = [to_remove]
 
     return {key: val for key, val in dict_.items() if key not in to_remove}
 
 
-def _reduce_params(params, model_dict, has_endogenous_factors):
+def _reduce_params(
+    params: pd.DataFrame,
+    model_dict: dict[str, Any],
+    *,
+    has_endogenous_factors: bool,
+) -> pd.DataFrame:
     """Reduce a parameter DataFrame from a larger model to a reduced model.
 
     The reduced model must be nested in the original model for which the params
     DataFrame was constructed.
 
     Args:
-        params (pandas.DataFrame or None): The params DataFrame for the full model.
-        model_dict (dict): The model specification. See: :ref:`model_specs`.
-        has_endogenous_factors (bool): Whether the model has endogenous factors.
+        params: The params DataFrame for the full model.
+        model_dict: The model specification. See: :ref:`model_specs`.
+        has_endogenous_factors: Whether the model has endogenous factors.
 
     Returns:
         pandas.DataFrame: The reduced parameters DataFrame.
@@ -328,7 +376,11 @@ def _reduce_params(params, model_dict, has_endogenous_factors):
     return params.loc[index]
 
 
-def _extend_params(params, model_dict, fill_value):
+def _extend_params(
+    params: pd.DataFrame,
+    model_dict: dict[str, Any],
+    fill_value: float,
+) -> pd.DataFrame:
     index = _get_params_index_from_model_dict(model_dict)
     out = params.reindex(index)
     out["value"] = out["value"].fillna(fill_value)
@@ -341,29 +393,37 @@ def _extend_params(params, model_dict, fill_value):
     return out
 
 
-def _get_params_index_from_model_dict(model_dict):
+def _get_params_index_from_model_dict(
+    model_dict: dict[str, Any],
+) -> pd.MultiIndex:
     mod = process_model(model_dict)
-    index = get_params_index(
-        update_info=mod["update_info"],
-        labels=mod["labels"],
-        dimensions=mod["dimensions"],
-        transition_info=mod["transition_info"],
-        endogenous_factors_info=mod["endogenous_factors_info"],
+    return get_params_index(
+        update_info=mod.update_info,
+        labels=mod.labels,
+        dimensions=mod.dimensions,
+        transition_info=mod.transition_info,
+        endogenous_factors_info=mod.endogenous_factors_info,
     )
-    return index
 
 
-def _remove_measurements_from_normalizations(measurements, normalizations):
+def _remove_measurements_from_normalizations(
+    measurements: str | list[str],
+    normalizations: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
     reduced = [_remove_from_dict(norm, measurements) for norm in normalizations]
     if reduced != normalizations:
         warnings.warn(
             "Your removed a normalized measurement from a model. Make sure there are "
             "enough normalizations left to ensure identification.",
+            stacklevel=2,
         )
     return reduced
 
 
-def _shorten_if_necessary(list_, length):
+def _shorten_if_necessary(
+    list_: list[Any],
+    length: int,
+) -> list[Any]:
     if len(list_) > length:
         list_ = list_[:length]
     return list_
diff --git a/src/skillmodels/utils_plotting.py b/src/skillmodels/utils_plotting.py
index 491ac9e..b66b89b 100644
--- a/src/skillmodels/utils_plotting.py
+++ b/src/skillmodels/utils_plotting.py
@@ -1,14 +1,19 @@
+"""Utility functions for configuring plot layouts and subplots."""
+
+from typing import Any
+
 import numpy as np
 
 
 def get_layout_kwargs(
-    layout_kwargs=None,
-    legend_kwargs=None,
-    title_kwargs=None,
-    showlegend=False,
-    columns=None,
-    rows=None,
-):
+    layout_kwargs: dict[str, Any] | None = None,
+    legend_kwargs: dict[str, Any] | None = None,
+    title_kwargs: dict[str, Any] | None = None,
+    *,
+    showlegend: bool = False,
+    columns: list[str] | tuple[str, ...] | None = None,
+    rows: list[str] | tuple[str, ...] | None = None,
+) -> dict[str, Any]:
     """Define and update default kwargs for update_layout.
 
     Defines some default keyword arguments to update figure layout, such as
@@ -37,13 +42,14 @@ def get_layout_kwargs(
 
 
 def get_make_subplot_kwargs(
-    sharex,
-    sharey,
-    column_order,
-    row_order,
-    make_subplot_kwargs,
-    add_scenes=False,
-):
+    *,
+    sharex: bool,
+    sharey: bool,
+    column_order: list[str] | tuple[str, ...],
+    row_order: list[str] | tuple[str, ...],
+    make_subplot_kwargs: dict[str, Any] | None,
+    add_scenes: bool = False,
+) -> dict[str, Any]:
     """Define and update keywargs for instantiating figure with subplots."""
     nrows = len(row_order)
     ncols = len(column_order)
diff --git a/src/skillmodels/visualize_factor_distributions.py b/src/skillmodels/visualize_factor_distributions.py
index 8d9dea6..a12a39f 100644
--- a/src/skillmodels/visualize_factor_distributions.py
+++ b/src/skillmodels/visualize_factor_distributions.py
@@ -1,5 +1,8 @@
+"""Functions to visualize distributions of latent factors."""
+
 import warnings
 from copy import deepcopy
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import pandas as pd
@@ -13,62 +16,73 @@
 from skillmodels.process_model import process_model
 from skillmodels.utils_plotting import get_layout_kwargs, get_make_subplot_kwargs
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from numpy.typing import NDArray
+
+    from skillmodels.types import ProcessedModel
+
 
 def combine_distribution_plots(
-    kde_plots,
-    contour_plots,
-    surface_plots=None,
-    factor_order=None,
-    factor_mapping=None,
-    make_subplot_kwargs=None,
-    sharex=False,
-    sharey=False,
-    line_width=1.5,
-    showlegend=False,
-    layout_kwargs=None,
-    legend_kwargs=None,
-    title_kwargs=None,
-    eye_x=2.2,
-    eye_y=2.2,
-    eye_z=1,
-):
+    kde_plots: dict[str, go.Figure],
+    contour_plots: dict[tuple[str, str], go.Figure],
+    surface_plots: dict[tuple[str, str], go.Figure] | None = None,
+    factor_order: list[str] | tuple[str, ...] | None = None,
+    factor_mapping: dict[str, str] | None = None,
+    make_subplot_kwargs: dict[str, Any] | None = None,
+    *,
+    sharex: bool = False,
+    sharey: bool = False,
+    line_width: float = 1.5,
+    showlegend: bool = False,
+    layout_kwargs: dict[str, Any] | None = None,
+    legend_kwargs: dict[str, Any] | None = None,
+    title_kwargs: dict[str, Any] | None = None,
+    eye_x: float = 2.2,
+    eye_y: float = 2.2,
+    eye_z: float = 1,
+) -> go.Figure:
     """Combine individual plots into figure with subplots.
 
     Uses dictionary with plotly images as values to build plotly Figure with subplots.
 
     Args:
-        kde_plots (dict): Dictionary with plots of indivudal factor kde plots.
-        contour_plots (dict): Dictionary with plots of pairwise factor density
+        kde_plots: Dictionary with plots of indivudal factor kde plots.
+        contour_plots: Dictionary with plots of pairwise factor density
             contours.
-        surface_plots (dict): Dictionary with plots of pairwise factor density
+        surface_plots: Dictionary with plots of pairwise factor density
             3d plots.
-        make_subplot_kwargs (dict or NoneType): Dictionary of keyword arguments used
+        factor_order: List of factor names to define the order of
+            subplots. If None, uses the order from kde_plots keys.
+        make_subplot_kwargs: Dictionary of keyword arguments used
             to instantiate plotly Figure with multiple subplots. Is used to define
             properties such as, for example, the spacing between subplots. If None,
             default arguments defined in the function are used.
-        factor_mapping (dct): Dictionary to change displayed factor names.
-        sharex (bool): Whether to share the properties of x-axis across subplots.
+        factor_mapping: Dictionary to change displayed factor names.
+        sharex: Whether to share the properties of x-axis across subplots.
             Default False.
-        sharey (bool): Whether to share the properties ofy-axis across subplots.
+        sharey: Whether to share the properties ofy-axis across subplots.
             Default True.
-        line_width (float): A float used to set same line width across subplots.
-        showlegend (bool): Display legend if True.
-        layout_kwargs (dict or NoneType): Dictionary of key word arguments used to
+        line_width: A float used to set same line width across subplots.
+        showlegend: Display legend if True.
+        layout_kwargs: Dictionary of key word arguments used to
             update layout of plotly Figure object. If None, the default kwargs defined
             in the function will be used.
-        legend_kwargs (dict or NoneType): Dictionary of key word arguments used to
+        legend_kwargs: Dictionary of key word arguments used to
             update position, orientation and title of figure legend. If None, default
             position and orientation will be used with no title.
-        title_kwargs (dict or NoneType): Dictionary of key word arguments used to
+        title_kwargs: Dictionary of key word arguments used to
             update properties of the figure title. Use {'text': '<desired title>'}
             to set figure title. If None, infers title based on the value of
             `quntiles_of_other_factors`.
-        eye_x, eye_y and eye_z (float): Control camera (view point) of the 3d plots.
-            Together they form the a norm, and the larger the norm, the more zoomed out
-            is the view. Setting eye_z to a lower value lowers the view point.
+        eye_x: Control camera x position for the 3d plots. Default 2.2.
+        eye_y: Control camera y position for the 3d plots. Default 2.2.
+        eye_z: Control camera z position for the 3d plots. Default 1.
+            Setting eye_z to a lower value lowers the view point.
 
     Returns:
-        fig (plotly.Figure): Plotly figure with subplots that combines pairwise
+        fig: Plotly figure with subplots that combines pairwise
             distrubtion plots.
 
     """
@@ -149,61 +163,62 @@ def combine_distribution_plots(
 
 
 def univariate_densities(
-    data,
-    model_dict,
-    params,
-    period,
-    factors=None,
-    observed_factors=False,
-    states=None,
-    show_curve=True,
-    show_hist=False,
-    show_rug=False,
-    curve_type="kde",
-    colorscale="D3",
-    bin_size=1,
-    distplot_kwargs=None,
-    layout_kwargs=None,
-):
+    data: pd.DataFrame,
+    model_dict: dict[str, Any],
+    params: pd.DataFrame,
+    period: int,
+    factors: list[str] | tuple[str, ...] | None = None,
+    *,
+    observed_factors: bool = False,
+    states: pd.DataFrame | dict[str, pd.DataFrame] | list[pd.DataFrame] | None = None,
+    show_curve: bool = True,
+    show_hist: bool = False,
+    show_rug: bool = False,
+    curve_type: str = "kde",
+    colorscale: str = "D3",
+    bin_size: float = 1,
+    distplot_kwargs: dict[str, Any] | None = None,
+    layout_kwargs: dict[str, Any] | None = None,
+) -> dict[str, go.Figure]:
     """Get dictionary with kernel density estimate plots for each factor.
 
     Plots kernel densities for latent factors and collects them in a dictionary
     with factor names as keys.
 
     Args:
-        data (DataFrame): Model estimation input data.
-        model_dict (dict): Dictionary with model specifications.
-        params (DataFrame): DataFrame with estimated parameter values.
-        period (int or float): Model period for which to plot the distributions for.
-        factors (list or NoneType): List of factors for which to plot the densities.
+        data: Model estimation input data.
+        model_dict: Dictionary with model specifications.
+        params: DataFrame with estimated parameter values.
+        period: Model period for which to plot the distributions for.
+        factors: List of factors for which to plot the densities.
             If None, plot pairwise distributions for all latent factors.
-        observed_factors (bool): If True, plot densities of observed factors too.
-        states (dict, list, pd.DataFrame or NoneType): List or dictionary with tidy
+        observed_factors: If True, plot densities of observed factors too.
+        states: List or dictionary with tidy
             DataFrames with filtered or simulated states or only one DataFrame with
             filtered or simulated states. If None, retrieve data frame with filtered
             states using model_dict and data. States are used to estimate the state
             ranges in each period (if state_ranges are not given explicitly) and to
             estimate the distribution of the latent factors.
-        show_hist (bool): Add histogram to the distplot.
-        show_curve (bool): Add density curve to the displot.
-        show_rug (bool): Add rug to the distplot.
-        curve_type (str): Curve type, 'normal' or 'kde', to add to the distplot.
-        colorscale (str): The color palette used when plotting multiple data. Must be
+        show_hist: Add histogram to the distplot.
+        show_curve: Add density curve to the displot.
+        show_rug: Add rug to the distplot.
+        curve_type: Curve type, 'normal' or 'kde', to add to the distplot.
+        colorscale: The color palette used when plotting multiple data. Must be
             a valid attribute of px.colors.qualitative.
-        bin_size (float): Size of the histogram bins.
-        distplot_kwargs (NoneType or dict): Dictionary with additional keyword
+        bin_size: Size of the histogram bins.
+        distplot_kwargs: Dictionary with additional keyword
             arguments passed to ff.create_distplot() to initiate
             the distplot.
-        layout_kwargs (NoneType or dict): Dictionary of keyword arguments to update
+        layout_kwargs: Dictionary of keyword arguments to update
             layout of the plot figures. Some essential layout kwargs are:
-            - xaxis_title (str): label label
-            - yaxis_title (str): label of y axis
-            - xaxis_showgrid (bool): display axis grid
-            - yaxis_showgrid (bool): display axis grid
-            - template (str): figure background theme
-            - showlegend (bool): add legend
+            - xaxis_title: label label
+            - yaxis_title: label of y axis
+            - xaxis_showgrid: display axis grid
+            - yaxis_showgrid: display axis grid
+            - template: figure background theme
+            - showlegend: add legend
     Returns:
-        plots_dict (dict): Dictionary with density plots.
+        plots_dict: Dictionary with density plots.
 
     """
     if states is None:
@@ -221,20 +236,20 @@ def univariate_densities(
         states=states,
         period=period,
         factors=factors,
-        aug_periods_to_periods=model["labels"]["aug_periods_to_periods"],
+        aug_periods_to_periods=model.labels.aug_periods_to_periods,
         observed_states=observed_states,
     )
     scenarios = df["scenario"].unique()
     plots_dict = {}
     distplot_kwargs = _process_distplot_kwargs(
-        show_curve,
-        show_hist,
-        show_rug,
-        curve_type,
-        bin_size,
-        scenarios,
-        colorscale,
-        distplot_kwargs,
+        show_curve=show_curve,
+        show_hist=show_hist,
+        show_rug=show_rug,
+        curve_type=curve_type,
+        bin_size=bin_size,
+        scenarios=scenarios,
+        colorscale=colorscale,
+        distplot_kwargs=distplot_kwargs,
     )
     plots_dict = {}
     layout_kwargs = get_layout_kwargs(layout_kwargs)
@@ -246,6 +261,7 @@ def univariate_densities(
             warnings.warn(
                 f"""Plotting univariate density failed for {fac} in
                 period {period} with error:\n\n{e}""",
+                stacklevel=2,
             )
             fig = go.Figure()
         fig.update_layout(showlegend=False)
@@ -257,46 +273,47 @@ def univariate_densities(
 
 
 def bivariate_density_contours(
-    data,
-    model_dict,
-    params,
-    period,
-    factors=None,
-    observed_factors=False,
-    states=None,
-    n_points=50,
-    contour_kwargs=None,
-    layout_kwargs=None,
-    contours_showlabels=False,
-    contours_coloring="none",
-    contours_colorscale="RdBu_r",
-    lines_colorscale="D3",
-    showcolorbar=False,
-):
+    data: pd.DataFrame,
+    model_dict: dict[str, Any],
+    params: pd.DataFrame,
+    period: int,
+    factors: list[str] | tuple[str, ...] | None = None,
+    *,
+    observed_factors: bool = False,
+    states: pd.DataFrame | dict[str, pd.DataFrame] | list[pd.DataFrame] | None = None,
+    n_points: int = 50,
+    contour_kwargs: dict[str, Any] | None = None,
+    layout_kwargs: dict[str, Any] | None = None,
+    contours_showlabels: bool = False,
+    contours_coloring: str = "none",
+    contours_colorscale: str = "RdBu_r",
+    lines_colorscale: str = "D3",
+    showcolorbar: bool = False,
+) -> dict[tuple[str, str], go.Figure]:
     """Get dictionary with pariwise density contour plots.
 
     Plots pairwise bivariate density contours for latent factors
     and collects them in a dictionary with factor combinations as keys.
 
     Args:
-        data (DataFrame): Model estimation input data.
-        model_dict (dict): Dictionary with model specifications.
-        params (DataFrame): DataFrame with estimated parameter values.
-        period (int or float): Model period for which to plot the distributions for.
-        factors (list or NoneType): List of factors for which to plot the densities.
+        data: Model estimation input data.
+        model_dict: Dictionary with model specifications.
+        params: DataFrame with estimated parameter values.
+        period: Model period for which to plot the distributions for.
+        factors: List of factors for which to plot the densities.
             If None, plot pairwise distributions for all latent factors.
-        observed_factors (bool): If True, plot densities of observed factors too.
-        states (dict, list, pd.DataFrame or NoneType): List or dictionary with tidy
+        observed_factors: If True, plot densities of observed factors too.
+        states: List or dictionary with tidy
             DataFrames with filtered or simulated states or only one DataFrame with
             filtered or simulated states. If None, retrieve data frame with filtered
             states using model_dict and data. States are used to estimate the state
             ranges in each period (if state_ranges are not given explicitly) and to
             estimate the distribution of the latent factors.
-        n_points (int): Number of grid points used to create the mesh for calculation
+        n_points: Number of grid points used to create the mesh for calculation
             of kernel densities.
-        contour_kwargs (dict or NoneType): Dictionary with keyword arguments to set
+        contour_kwargs: Dictionary with keyword arguments to set
             contour line properties (such as annotation, colorscale).
-        layout_kwargs (dict or NoneType): Dictionary with keyword arguments to set
+        layout_kwargs: Dictionary with keyword arguments to set
             figure layout properties.
 
         The following are various essential keyword arguments defining various features
@@ -304,16 +321,19 @@ def bivariate_density_contours(
         'update_traces'. Some default figure layout properties (such as background
         theme) are defined if layout_kwargs is None.
 
-        contours_showlabels (bool): If True, annotate density contours.
-        contours_coloring (str): Defines how to apply color scale to density contours.
+        contours_showlabels: If True, annotate density contours.
+        contours_coloring: Defines how to apply color scale to density contours.
             Possible values are in ['lines', 'fill', 'heatmap', 'none']. Default is
             'none' which implies no colorscale.
-        contours_colorscale (str): The color scale to use for line legends. Must be
+        contours_colorscale: The color scale to use for line legends. Must be
             a valid plotly.express.colors.sequential attribute. Default 'RdBu_r'.
-        showcolorbar (bool): A boolean variable for displaying color bar.
+        lines_colorscale: The color palette used for contour lines when plotting
+            multiple scenarios. Must be a valid px.colors.qualitative attribute.
+            Default 'D3'.
+        showcolorbar: A boolean variable for displaying color bar.
 
     Returns:
-        plots_dict (dict): Dictionary with factor combinations as keys and respective
+        plots_dict: Dictionary with factor combinations as keys and respective
             pariwise plots of density contours as values.
 
     """
@@ -332,16 +352,16 @@ def bivariate_density_contours(
         states=states,
         period=period,
         factors=factors,
-        aug_periods_to_periods=model["labels"]["aug_periods_to_periods"],
+        aug_periods_to_periods=model.labels.aug_periods_to_periods,
         observed_states=observed_states,
     )
     plots_dict = {}
     contour_kwargs = _process_contour_kwargs(
         contour_kwargs,
-        contours_showlabels,
-        contours_coloring,
-        contours_colorscale,
-        showcolorbar,
+        contours_showlabels=contours_showlabels,
+        contours_coloring=contours_coloring,
+        contours_colorscale=contours_colorscale,
+        contours_showscale=showcolorbar,
     )
     layout_kwargs = _process_layout_kwargs(layout_kwargs)
     pairs = []
@@ -373,6 +393,7 @@ def bivariate_density_contours(
                     Contour plot failed for {pair} in period {period}
                     with error:\n\n{e}
                     """,
+                    stacklevel=2,
                 )
         fig.update_xaxes(title={"text": pair[0]})
         fig.update_yaxes(title={"text": pair[1]})
@@ -383,60 +404,63 @@ def bivariate_density_contours(
 
 
 def bivariate_density_surfaces(
-    data,
-    model_dict,
-    params,
-    period,
-    factors=None,
-    observed_factors=False,
-    states=None,
-    n_points=50,
-    layout_kwargs=None,
-    colorscale="RdBu_r",
-    opacity=0.9,
-    showcolorbar=False,
-    showgrids=True,
-    showaxlines=True,
-    showlabels=True,
-):
+    data: pd.DataFrame,
+    model_dict: dict[str, Any],
+    params: pd.DataFrame,
+    period: int,
+    factors: list[str] | tuple[str, ...] | None = None,
+    *,
+    observed_factors: bool = False,
+    states: pd.DataFrame | None = None,
+    n_points: int = 50,
+    layout_kwargs: dict[str, Any] | None = None,
+    colorscale: str = "RdBu_r",
+    opacity: float = 0.9,
+    showcolorbar: bool = False,
+    showgrids: bool = True,
+    showaxlines: bool = True,
+    showlabels: bool = True,
+) -> dict[tuple[str, str], go.Figure]:
     """Get dictionary with pariwise 3d density surface plots.
 
     Plots pairwise 3d density surfaces for latent factors
     and collects them in a dictionary with factor name combinations keys.
 
     Args:
-        data (DataFrame): Model estimation input data.
-        model_dict (dict): Dictionary with model specifications.
-        params (DataFrame): DataFrame with estimated parameter values.
-        period (int or float): Model period for which to plot the distributions for.
-        factors (list or NoneType): List of factors for which to plot the densities.
+        data: Model estimation input data.
+        model_dict: Dictionary with model specifications.
+        params: DataFrame with estimated parameter values.
+        period: Model period for which to plot the distributions for.
+        factors: List of factors for which to plot the densities.
             If None, plot pairwise distributions for all latent factors.
-        observed_factors (bool): If True, plot densities of observed factors too.
-        states (dict, list, pd.DataFrame or NoneType): List or dictionary with tidy
+        observed_factors: If True, plot densities of observed factors too.
+        states: List or dictionary with tidy
             DataFrames with filtered or simulated states or only one DataFrame with
             filtered or simulated states. If None, retrieve data frame with filtered
             states using model_dict and data. States are used to estimate the state
             ranges in each period (if state_ranges are not given explicitly) and to
             estimate the distribution of the latent factors.
-        n_points (int): Number of grid points used to create the mesh for calculation
+        n_points: Number of grid points used to create the mesh for calculation
             of kernel densities.
+
         The following are various essential keyword arguments defining various features
         of plots. All features can also be changed ex-post via 'update_layout' or
         'update_traces'. Some default figure layout properties (such as background
         theme) are defined if layout_kwargs is None.
 
-        layout_kwargs (dict or NoneType): Dictionary with keyword arguments to set
+        layout_kwargs: Dictionary with keyword arguments to set
             figure layout properties.
-        colorscale (str): The color scale to use for line legends. Must be a valid
+        colorscale: The color scale to use for line legends. Must be a valid
             plotly.express.colors.sequential attribute. Default 'RdBu_r'.
-        showcolorbar (bool): A boolean variable for displaying the colorbar associated
+        opacity: Opacity of the surface. Default 0.9.
+        showcolorbar: A boolean variable for displaying the colorbar associated
             with the surface color scale.
-        showgrids (bool): A boolean variable for showing axes grids.
-        showaxlines (bool): A boolean variable for showing axes lines.
-        showlabels (bool): A boolean variable for displaying axes labels.
+        showgrids: A boolean variable for showing axes grids.
+        showaxlines: A boolean variable for showing axes lines.
+        showlabels: A boolean variable for displaying axes labels.
 
     Returns:
-        plots_dict (dict): Dictionary with factor combinations as keys and respective
+        plots_dict: Dictionary with factor combinations as keys and respective
             pariwise plots of 3d density plots as values.
 
     """
@@ -457,15 +481,15 @@ def bivariate_density_surfaces(
         states=states,
         period=period,
         factors=factors,
-        aug_periods_to_periods=model["labels"]["aug_periods_to_periods"],
+        aug_periods_to_periods=model.labels.aug_periods_to_periods,
         observed_states=observed_states,
     )
     plots_dict = {}
     layout_kwargs = _process_layout_kwargs_3d(
         layout_kwargs,
-        showgrids,
-        showaxlines,
-        showlabels,
+        showgrids=showgrids,
+        showaxlines=showaxlines,
+        showlabels=showlabels,
     )
     pairs = []
     for fac1 in factors:
@@ -490,6 +514,7 @@ def bivariate_density_surfaces(
             warnings.warn(
                 f"""Plotting bivariate density surfaces for {pair} in
                 period {period} with error:\n\n{e}""",
+                stacklevel=2,
             )
             fig = go.Figure()
         fig.update_layout(
@@ -505,8 +530,12 @@ def bivariate_density_surfaces(
 
 
 def _process_data(
-    states, period, factors, aug_periods_to_periods, observed_states=None
-):
+    states: pd.DataFrame | dict[str, pd.DataFrame] | list[pd.DataFrame],
+    period: int,
+    factors: tuple[str, ...],
+    aug_periods_to_periods: Mapping[int, int],
+    observed_states: pd.DataFrame | None = None,
+) -> pd.DataFrame:
     ap_to_p = pd.Series(aug_periods_to_periods, name="period")
     ap_to_p.index.name = "aug_period"
     if isinstance(states, pd.DataFrame):
@@ -548,15 +577,16 @@ def _process_data(
 
 
 def _process_distplot_kwargs(
-    show_curve,
-    show_hist,
-    show_rug,
-    curve_type,
-    bin_size,
-    scenarios,
-    colorscale,
-    distplot_kwargs,
-):
+    *,
+    show_curve: bool,
+    show_hist: bool,
+    show_rug: bool,
+    curve_type: str,
+    bin_size: float,
+    scenarios: NDArray[Any],
+    colorscale: str,
+    distplot_kwargs: dict[str, Any] | None,
+) -> dict[str, Any]:
     """Define and update default distplot kwargs."""
     default_kwargs = {
         "show_hist": show_hist,
@@ -572,7 +602,13 @@ def _process_distplot_kwargs(
     return default_kwargs
 
 
-def _calculate_kde_for_3d(data, factors, n_points):
+def _calculate_kde_for_3d(
+    data: pd.DataFrame,
+    factors: tuple[str, str],
+    n_points: int,
+) -> tuple[
+    NDArray[np.floating[Any]], NDArray[np.floating[Any]], NDArray[np.floating[Any]]
+]:
     """Create grid mesh and calculate Gaussian kernel over the grid."""
     x = data[factors[0]]
     y = data[factors[1]]
@@ -588,12 +624,13 @@ def _calculate_kde_for_3d(data, factors, n_points):
 
 
 def _process_contour_kwargs(
-    contour_kwargs,
-    contours_showlabels,
-    contours_coloring,
-    contours_colorscale,
-    contours_showscale,
-):
+    contour_kwargs: dict[str, Any] | None,
+    *,
+    contours_showlabels: bool,
+    contours_coloring: str | None,
+    contours_colorscale: str,
+    contours_showscale: bool,
+) -> dict[str, Any]:
     """Define and update default density contour kwargs."""
     if contours_coloring is None:
         contours_coloring = "none"
@@ -609,9 +646,11 @@ def _process_contour_kwargs(
     return default_kwargs
 
 
-def _process_layout_kwargs(layout_kwargs):
+def _process_layout_kwargs(
+    layout_kwargs: dict[str, Any] | None,
+) -> dict[str, Any]:
     """Define and update default figure layout kwargs."""
-    default_kwargs = {
+    default_kwargs: dict[str, Any] = {
         "template": "simple_white",
         "xaxis_showgrid": False,
         "yaxis_showgrid": False,
@@ -621,12 +660,18 @@ def _process_layout_kwargs(layout_kwargs):
     return default_kwargs
 
 
-def _process_layout_kwargs_3d(layout_kwargs, showgrids, showaxlines, showlabels):
+def _process_layout_kwargs_3d(
+    layout_kwargs: dict[str, Any] | None,
+    *,
+    showgrids: bool,
+    showaxlines: bool,
+    showlabels: bool,
+) -> dict[str, Any]:
     """Define and update default figure layout kwargs for 3d plots."""
-    default_kwargs = {
+    default_kwargs: dict[str, Any] = {
         "template": "none",
     }
-    scene = {}
+    scene: dict[str, Any] = {}
     for ax in list("xyz"):
         scene[f"{ax}axis"] = {
             "showgrid": showgrids,
@@ -640,7 +685,10 @@ def _process_layout_kwargs_3d(layout_kwargs, showgrids, showaxlines, showlabels)
     return default_kwargs
 
 
-def _process_factor_mapping_dist(mapper, factors):
+def _process_factor_mapping_dist(
+    mapper: dict[str, str] | None,
+    factors: list[str] | tuple[str, ...],
+) -> dict[str, str]:
     """Process mapper to return dictionary with old and new factor names."""
     if mapper is None:
         mapper = {fac: fac for fac in factors}
@@ -651,28 +699,38 @@ def _process_factor_mapping_dist(mapper, factors):
     return mapper
 
 
-def _get_ordered_factors(factor_order, factors):
-    """Process factor orders to return list of strings."""
+def _get_ordered_factors(
+    factor_order: list[str] | tuple[str, ...] | str | None,
+    factors: list[str] | tuple[str, ...],
+) -> tuple[str, ...]:
+    """Process factor orders to return tuple of strings."""
     if factor_order is None:
-        ordered_factors = factors
+        ordered_factors = tuple(factors)
     elif isinstance(factor_order, str):
-        ordered_factors = [factor_order]
+        ordered_factors = (factor_order,)
     else:
-        ordered_factors = factor_order
+        ordered_factors = tuple(factor_order)
     return ordered_factors
 
 
-def _get_factors(factors, observed_factors, model):
-    """Proccess factor names to return list of strings."""
+def _get_factors(
+    factors: list[str] | tuple[str, ...] | None,
+    *,
+    observed_factors: bool,
+    model: ProcessedModel,
+) -> tuple[str, ...]:
+    """Proccess factor names to return tuple of strings."""
     if factors is None:
         if observed_factors:
-            factors = model["labels"]["all_factors"]
-        else:
-            factors = model["labels"]["latent_factors"]
-    return factors
+            return model.labels.all_factors
+        return model.labels.latent_factors
+    return tuple(factors)
 
 
-def _get_data_observed_factors(data, factors):
+def _get_data_observed_factors(
+    data: pd.DataFrame,
+    factors: tuple[str, ...],
+) -> pd.DataFrame | None:
     """Get data with observed factors if any."""
     to_concat = []
     for fac in factors:
diff --git a/src/skillmodels/visualize_transition_equations.py b/src/skillmodels/visualize_transition_equations.py
index 4dc5304..4203bec 100644
--- a/src/skillmodels/visualize_transition_equations.py
+++ b/src/skillmodels/visualize_transition_equations.py
@@ -1,9 +1,14 @@
+"""Functions to visualize transition equations and production functions."""
+
 import itertools
+from collections.abc import Callable  # noqa: TC003
 from copy import deepcopy
+from typing import TYPE_CHECKING, Any
 
 import jax.numpy as jnp
 import numpy as np
 import pandas as pd
+from jax import Array
 from plotly import express as px
 from plotly import graph_objects as go
 from plotly.subplots import make_subplots
@@ -14,59 +19,64 @@
 from skillmodels.process_data import process_data
 from skillmodels.process_debug_data import create_state_ranges
 from skillmodels.process_model import process_model
+from skillmodels.types import ParsedParams  # noqa: TC001
 from skillmodels.utils_plotting import get_layout_kwargs, get_make_subplot_kwargs
 
+if TYPE_CHECKING:
+    from skillmodels.types import ProcessedModel
+
 
 def combine_transition_plots(
-    plots_dict,
-    column_order=None,
-    row_order=None,
-    factor_mapping=None,
-    make_subplot_kwargs=None,
-    sharex=False,
-    sharey=True,
-    showlegend=True,
-    layout_kwargs=None,
-    legend_kwargs=None,
-    title_kwargs=None,
-):
+    plots_dict: dict[tuple[str, str], go.Figure],
+    column_order: list[str] | tuple[str, ...] | str | None = None,
+    row_order: list[str] | tuple[str, ...] | str | None = None,
+    factor_mapping: dict[str, str] | None = None,
+    make_subplot_kwargs: dict[str, Any] | None = None,
+    *,
+    sharex: bool = False,
+    sharey: bool = True,
+    showlegend: bool = True,
+    layout_kwargs: dict[str, Any] | None = None,
+    legend_kwargs: dict[str, Any] | None = None,
+    title_kwargs: dict[str, Any] | None = None,
+) -> go.Figure:
     """Combine individual plots into figure with subplots.
 
     Use dictionary with plotly images as values to build plotly figure with subplots.
 
     Args:
-        plots_dict (dict): Dictionary with plots of transition functions for each
+        plots_dict: Dictionary with plots of transition functions for each
             factor.
-        column_order (list, str or NoneType): List of (output) factor names according
+        column_order: List of (output) factor names according
             to which transition plots should be ordered horizontally. If None, infer
             from the keys of of plots_dict
-        row_order (list, str or NoneType): List of (input) factor names according
+        row_order: List of (input) factor names according
             to which transition plots should be ordered vertically. If None, infer
             from the keys of of plots_dict
-        factor_mapping (dict or NoneType): A dictionary with custom factor names to
+        factor_mapping: A dictionary with custom factor names to
             display as axes labels.
-        make_subplot_kwargs (dict or NoneType): Dictionary of keyword arguments used
+        make_subplot_kwargs: Dictionary of keyword arguments used
             to instantiate plotly Figure with multiple subplots. Is used to define
             properties such as, for example, the spacing between subplots. If None,
             default arguments defined in the function are used.
-        sharex (bool): Whether to share the properties of x-axis across subplots.
+        sharex: Whether to share the properties of x-axis across subplots.
             Default False.
-        sharey (bool): Whether to share the properties ofy-axis across subplots.
+        sharey: Whether to share the properties ofy-axis across subplots.
             Default True.
-        showlegend (bool): Display legend if True.
-        layout_kwargs (dict or NoneType): Dictionary of key word arguments used to
+        showlegend: Display legend if True.
+        layout_kwargs: Dictionary of key word arguments used to
             update layout of plotly Figure object. If None, the default kwargs defined
             in the function will be used.
-        legend_kwargs (dict or NoneType): Dictionary of key word arguments used to
+        legend_kwargs: Dictionary of key word arguments used to
             update position, orientation and title of figure legend. If None, default
             position and orientation will be used with no title.
-        title_kwargs (dict or NoneType): Dictionary of key word arguments used to
+        title_kwargs: Dictionary of key word arguments used to
             update properties of the figure title. Use {'text': '<desired title>'}
             to set figure title. If None, infers title based on the value of
             `quntiles_of_other_factors`.
 
     Returns:
-        fig (plotly.Figure): Plotly figure with subplots that combines individual
+        fig: Plotly figure with subplots that combines individual
             transition functions.
 
     """
@@ -74,11 +84,11 @@ def combine_transition_plots(
 
     column_order, row_order = _process_orders(column_order, row_order, plots_dict)
     make_subplot_kwargs = get_make_subplot_kwargs(
-        sharex,
-        sharey,
-        column_order,
-        row_order,
-        make_subplot_kwargs,
+        sharex=sharex,
+        sharey=sharey,
+        column_order=column_order,
+        row_order=row_order,
+        make_subplot_kwargs=make_subplot_kwargs,
     )
     factor_mapping = _process_factor_mapping_trans(
         factor_mapping,
@@ -119,56 +129,61 @@ def combine_transition_plots(
             )
 
     layout_kwargs = get_layout_kwargs(
-        layout_kwargs,
-        legend_kwargs,
-        title_kwargs,
-        showlegend,
-        column_order,
-        row_order,
+        layout_kwargs=layout_kwargs,
+        legend_kwargs=legend_kwargs,
+        title_kwargs=title_kwargs,
+        showlegend=showlegend,
+        columns=column_order,
+        rows=row_order,
     )
     fig.update_layout(**layout_kwargs)
     return fig
 
 
 def get_transition_plots(
-    model_dict,
-    params,
-    data,
-    period,
-    state_ranges=None,
-    quantiles_of_other_factors=(0.25, 0.5, 0.75),
-    n_points=50,
-    n_draws=50,
-    colorscale="Magenta_r",
-    layout_kwargs=None,
-    include_correction_factors=False,
-):
+    model_dict: dict[str, Any],
+    params: pd.DataFrame,
+    data: pd.DataFrame,
+    period: int,
+    state_ranges: dict[str, pd.DataFrame] | None = None,
+    quantiles_of_other_factors: tuple[float, ...] | list[float] | float | None = (
+        0.25,
+        0.5,
+        0.75,
+    ),
+    n_points: int = 50,
+    n_draws: int = 50,
+    colorscale: str = "Magenta_r",
+    layout_kwargs: dict[str, Any] | None = None,
+    *,
+    include_correction_factors: bool = False,
+) -> dict[tuple[str, str], go.Figure]:
     """Get dictionary with individual plots of transition equations for each factor.
 
     Args:
-        model_dict (dict): The model specification. See: :ref:`model_specs`
-        params (pandas.DataFrame): DataFrame with model parameters.
-        data (pd.DataFrame): Empirical dataset that is used to estimate the model.
-        period (int): The start period of the transition equations that are plotted.
-        state_ranges (dict or NoneType): The keys are the names of the latent factors.
+        model_dict: The model specification. See: :ref:`model_specs`
+        params: DataFrame with model parameters.
+        data: Empirical dataset that is used to estimate the model.
+        period: The start period of the transition equations that are plotted.
+        state_ranges: The keys are the names of the latent factors.
             The values are DataFrames with the columns "period", "minimum", "maximum".
             The state_ranges are used to define the axis limits of the plots.
-        quantiles_of_other_factors (float, list or None): Quantiles at which the factors
+        quantiles_of_other_factors: Quantiles at which the factors
             that are not varied in a given plot are fixed. If None, those factors are
             not fixed but integrated out.
-        n_points (int): Number of grid points per input. Default 50.
-        n_draws (int): Number of randomly drawn values of the factors that are averaged
+        n_points: Number of grid points per input. Default 50.
+        n_draws: Number of randomly drawn values of the factors that are averaged
             out. Only relevant if quantiles_of_other_factors is *None*. Default 50.
-        colorscale (str): The color scale to use for line legends. Must be a valid
+        colorscale: The color scale to use for line legends. Must be a valid
             plotly.express.colors.sequential attribute. Default 'Magenta_r'.
-        layout_kwargs (dict or NoneType): Dictionary of key word arguments used to
+        layout_kwargs: Dictionary of key word arguments used to
             update layout of plotly image object. If None, the default kwargs
             defined in the function will be used.
-        include_correction_factors (bool): Whether to include correction factors in the
+        include_correction_factors: Whether to include correction factors in the
             plots. Default False.
 
     Returns:
-        plots_dict (dict): Dictionary with individual plots of transition equations
+        plots_dict: Dictionary with individual plots of transition equations
             for each combination of input and output factors.
 
     """
@@ -178,33 +193,33 @@ def get_transition_plots(
 
     model = process_model(model_dict)
 
-    if period >= model["labels"]["periods"][-1]:
+    if period >= model.labels.periods[-1]:
         raise ValueError(
             "*period* must be the penultimate period of the model or earlier.",
         )
 
     if (
         include_correction_factors
-        or not model["endogenous_factors_info"]["has_endogenous_factors"]
+        or not model.endogenous_factors_info.has_endogenous_factors
     ):
-        latent_factors = model["labels"]["latent_factors"]
+        latent_factors = model.labels.latent_factors
     else:
         latent_factors = [
             lf
-            for lf in model["labels"]["latent_factors"]
-            if not model["endogenous_factors_info"][lf]["is_correction"]
+            for lf in model.labels.latent_factors
+            if not model.endogenous_factors_info.factor_info[lf].is_correction  # ty: ignore[invalid-argument-type]
         ]
-    all_factors = model["labels"]["all_factors"]
+    all_factors = model.labels.all_factors
     states = get_filtered_states(model_dict=model_dict, data=data, params=params)[
         "anchored_states"
     ]["states"]
-    plots_dict = _get_dictionary_with_plots(
+    return _get_dictionary_with_plots(
         model=model,
         data=data,
         params=params,
         states=states,
         state_ranges=state_ranges,
-        latent_factors=latent_factors,
+        latent_factors=latent_factors,  # ty: ignore[invalid-argument-type]
         all_factors=all_factors,
         quantiles_of_other_factors=quantiles_of_other_factors,
         period=period,
@@ -213,67 +228,69 @@ def get_transition_plots(
         colorscale=colorscale,
         layout_kwargs=layout_kwargs,
     )
-    return plots_dict
 
 
 def _get_dictionary_with_plots(
-    model,
-    data,
-    params,
-    states,
-    state_ranges,
-    latent_factors,
-    all_factors,
-    quantiles_of_other_factors,
-    period,
-    n_points,
-    n_draws,
-    colorscale,
-    layout_kwargs,
-    showlegend=True,
-):
+    model: ProcessedModel,
+    data: pd.DataFrame,
+    params: pd.DataFrame,
+    states: pd.DataFrame,
+    state_ranges: dict[str, pd.DataFrame] | None,
+    latent_factors: tuple[str, ...],
+    all_factors: tuple[str, ...],
+    quantiles_of_other_factors: list[float] | None,
+    period: int,
+    n_points: int,
+    n_draws: int,
+    colorscale: str,
+    layout_kwargs: dict[str, Any] | None,
+    *,
+    showlegend: bool = True,
+) -> dict[tuple[str, str], go.Figure]:
     """Get plots of transition functions for each input and output combination.
 
     Return a dictionary with individual plots of transition functions for each input
     and output factors.
 
     Args:
-        model (dict): The model specification. See: :ref:`model_specs`
-        params (pandas.DataFrame): DataFrame with model parameters.
-        states (pandas.DataFrame): Tidy DataFrame with filtered or simulated states.
+        model: The model specification. See: :ref:`model_specs`
+        data: Panel dataset in long format for getting observed factors.
+        params: DataFrame with model parameters.
+        states: Tidy DataFrame with filtered or simulated states.
             They are used to estimate the state ranges in each period (if state_ranges
             are not given explicitly) and to estimate the distribution of the factors
             that are not visualized.
-        state_ranges (dict): The keys are the names of the latent factors.
+        state_ranges: The keys are the names of the latent factors.
             The values are DataFrames with the columns "period", "minimum", "maximum".
             The state_ranges are used to define the axis limits of the plots.
 
-        latent_factors (list): Latent factors of the model that are outputs of
+        latent_factors: Latent factors of the model that are outputs of
             transition factors.
-        all_factors (list): All factors of the model that are the inputs of transition
+        all_factors: All factors of the model that are the inputs of transition
             functions.
-        quantiles_of_other_factors (float, list or None): Quantiles at which the factors
+        quantiles_of_other_factors: Quantiles at which the factors
             that are not varied in a given plot are fixed. If None, those factors are
             not fixed but integrated out.
-        period (int): The start period of the transition equations that are plotted.
-        n_points (int): Number of grid points per input. Default 50.
-        n_draws (int): Number of randomly drawn values of the factors that are averaged
+        period: The start period of the transition equations that are plotted.
+        n_points: Number of grid points per input. Default 50.
+        n_draws: Number of randomly drawn values of the factors that are averaged
             out. Only relevant if quantiles_of_other_factors is *None*. Default 50.
-        colorscale (str): The color scale to use for line legends. Must be a valid
+        colorscale: The color scale to use for line legends. Must be a valid
             plotly.express.colors.sequential attribute. Default 'Magenta_r'.
-        subfig_kwargs (dict or NoneType): Dictionary of key word arguments used to
+        layout_kwargs: Dictionary of key word arguments used to
             update layout of plotly image object. If None, the default kwargs defined
             in the function will be used.
+        showlegend: Display legend if True. Default True.
 
     Returns:
-        plots_dict (dict): Dictionary with individual plots of transition functions
+        plots_dict: Dictionary with individual plots of transition functions
             for each input and output factors.
 
     """
-    observed_factors = model["labels"]["observed_factors"]
+    observed_factors = model.labels.observed_factors
     states_data = _get_states_data(model, period, data, states, observed_factors)
     params = _set_index_params(model, params)
-    pardict = _get_pardict(model, params)
+    parsed_params = _get_parsed_params(model, params)
     state_ranges = _get_state_ranges(state_ranges, states_data, all_factors)
     layout_kwargs = get_layout_kwargs(
         layout_kwargs=layout_kwargs,
@@ -281,27 +298,23 @@ def _get_dictionary_with_plots(
         title_kwargs=None,
         showlegend=showlegend,
     )
-    has_endogenous_factors = model["endogenous_factors_info"]["has_endogenous_factors"]
+    has_endogenous_factors = model.endogenous_factors_info.has_endogenous_factors
     if has_endogenous_factors:
-        _aug_periods = model["endogenous_factors_info"]["aug_periods_from_period"](
-            period
-        )
+        _aug_periods = model.endogenous_factors_info.aug_periods_from_period(period)
     else:
         _aug_periods = [period]
     plots_dict = {}
     for output_factor, input_factor in itertools.product(latent_factors, all_factors):
-        transition_function = model["transition_info"]["individual_functions"][
-            output_factor
-        ]
+        transition_function = model.transition_info.individual_functions[output_factor]  # ty: ignore[invalid-argument-type]
         if (
             has_endogenous_factors
-            and model["endogenous_factors_info"][output_factor]["is_endogenous"]
+            and model.endogenous_factors_info.factor_info[output_factor].is_endogenous  # ty: ignore[invalid-argument-type]
         ):
             aug_period = min(_aug_periods)
         else:
             aug_period = max(_aug_periods)
         transition_params = {
-            output_factor: pardict["transition"][output_factor][aug_period]
+            output_factor: parsed_params.transition[output_factor][aug_period]
         }
 
         if quantiles_of_other_factors is not None:
@@ -354,50 +367,63 @@ def _get_dictionary_with_plots(
     return plots_dict
 
 
-def _get_state_ranges(state_ranges, states_data, all_factors):
+def _get_state_ranges(
+    state_ranges: dict[str, pd.DataFrame] | None,
+    states_data: pd.DataFrame,
+    all_factors: tuple[str, ...],
+) -> dict[str, pd.DataFrame]:
     """Create state ranges if none is given."""
     if state_ranges is None:
-        state_ranges = create_state_ranges(states_data, all_factors)
+        state_ranges = create_state_ranges(states_data, list(all_factors))
     return state_ranges
 
 
-def _get_pardict(model, params):
-    """Get parsed params dictionary."""
+def _get_parsed_params(
+    model: ProcessedModel,
+    params: pd.DataFrame,
+) -> ParsedParams:
+    """Get parsed params dataclass."""
     parsing_info = create_parsing_info(
-        params_index=params.index,
-        update_info=model["update_info"],
-        labels=model["labels"],
-        anchoring=model["anchoring"],
-        has_endogenous_factors=model["endogenous_factors_info"][
-            "has_endogenous_factors"
-        ],
+        params_index=params.index,  # ty: ignore[invalid-argument-type]
+        update_info=model.update_info,
+        labels=model.labels,
+        anchoring=model.anchoring,
+        has_endogenous_factors=model.endogenous_factors_info.has_endogenous_factors,
     )
 
-    _, _, _, pardict = parse_params(
+    _, _, _, parsed_params = parse_params(
         params=jnp.array(params["value"].to_numpy()),
         parsing_info=parsing_info,
-        dimensions=model["dimensions"],
-        labels=model["labels"],
+        dimensions=model.dimensions,
+        labels=model.labels,
         n_obs=1,
     )
-    return pardict
+    return parsed_params
 
 
-def _set_index_params(model, params):
+def _set_index_params(
+    model: ProcessedModel,
+    params: pd.DataFrame,
+) -> pd.DataFrame:
     """Reset index of params data frame to model implied values."""
     params_index = get_params_index(
-        update_info=model["update_info"],
-        labels=model["labels"],
-        dimensions=model["dimensions"],
-        transition_info=model["transition_info"],
-        endogenous_factors_info=model["endogenous_factors_info"],
+        update_info=model.update_info,
+        labels=model.labels,
+        dimensions=model.dimensions,
+        transition_info=model.transition_info,
+        endogenous_factors_info=model.endogenous_factors_info,
     )
 
-    params = params.reindex(params_index)
-    return params
+    return params.reindex(params_index)
 
 
-def _get_states_data(model, period, data, states, observed_factors):
+def _get_states_data(
+    model: ProcessedModel,
+    period: int,
+    data: pd.DataFrame,
+    states: pd.DataFrame,
+    observed_factors: tuple[str, ...],
+) -> pd.DataFrame:
     if observed_factors and data is None:
         raise ValueError(
             "The model has observed factors. You must pass the empirical data to "
@@ -407,19 +433,17 @@ def _get_states_data(model, period, data, states, observed_factors):
     if observed_factors:
         _observed_arr = process_data(
             df=data,
-            has_endogenous_factors=model["endogenous_factors_info"][
-                "has_endogenous_factors"
-            ],
-            labels=model["labels"],
-            update_info=model["update_info"],
-            anchoring_info=model["anchoring"],
+            has_endogenous_factors=model.endogenous_factors_info.has_endogenous_factors,
+            labels=model.labels,
+            update_info=model.update_info,
+            anchoring_info=model.anchoring,
         )["observed_factors"]
         # convert from jax to numpy
         _observed_arr = np.array(_observed_arr)
-        if model["endogenous_factors_info"]["has_endogenous_factors"]:
+        if model.endogenous_factors_info.has_endogenous_factors:
             both_aug_periods = [
                 aug_p
-                for aug_p, p in model["labels"]["aug_periods_to_periods"].items()
+                for aug_p, p in model.labels.aug_periods_to_periods.items()
                 if p == period
             ]
             to_concat = []
@@ -453,18 +477,18 @@ def _get_states_data(model, period, data, states, observed_factors):
 
 
 def _prepare_data_for_one_plot_fixed_quantile_2d(
-    states_data,
-    state_ranges,
-    aug_period,
-    input_factor,
-    output_factor,
-    n_points,
-    quantiles_of_other_factors,
-    transition_function,
-    transition_params,
-    all_factors,
-):
-    period_data = states_data.query(f"aug_period == {aug_period}")[all_factors]
+    states_data: pd.DataFrame,
+    state_ranges: dict[str, pd.DataFrame],
+    aug_period: int,
+    input_factor: str,
+    output_factor: str,
+    n_points: int,
+    quantiles_of_other_factors: list[float],
+    transition_function: Callable[..., Array],
+    transition_params: dict[str, Any],
+    all_factors: tuple[str, ...],
+) -> pd.DataFrame:
+    period_data = states_data.query(f"aug_period == {aug_period}")[list(all_factors)]
     input_min = state_ranges[input_factor].loc[aug_period]["minimum"]
     input_max = state_ranges[input_factor].loc[aug_period]["maximum"]
     to_concat = []
@@ -474,7 +498,7 @@ def _prepare_data_for_one_plot_fixed_quantile_2d(
         fixed_quantiles = period_data.drop(columns=input_factor).quantile(quantile)
         for col, val in fixed_quantiles.items():
             input_data[col] = val
-        input_arr = jnp.array(input_data[all_factors].to_numpy())
+        input_arr = jnp.array(input_data[list(all_factors)].to_numpy())
         # convert from jax to numpy array
         output_arr = np.array(transition_function(transition_params, input_arr))
         quantile_data = pd.DataFrame()
@@ -483,11 +507,12 @@ def _prepare_data_for_one_plot_fixed_quantile_2d(
         quantile_data["quantile"] = quantile
         to_concat.append(quantile_data)
 
-    out = pd.concat(to_concat).reset_index()
-    return out
+    return pd.concat(to_concat).reset_index()
 
 
-def _process_quantiles_of_other_factors(quantiles_of_other_factors):
+def _process_quantiles_of_other_factors(
+    quantiles_of_other_factors: tuple[float, ...] | list[float] | float | None,
+) -> list[float] | None:
     """Process quantiles of other factors to always have list as type."""
     if isinstance(quantiles_of_other_factors, float | int):
         quantiles_of_other_factors = [quantiles_of_other_factors]
@@ -497,17 +522,17 @@ def _process_quantiles_of_other_factors(quantiles_of_other_factors):
 
 
 def _prepare_data_for_one_plot_average_2d(
-    states_data,
-    state_ranges,
-    aug_period,
-    input_factor,
-    output_factor,
-    n_points,
-    n_draws,
-    transition_function,
-    transition_params,
-    all_factors,
-):
+    states_data: pd.DataFrame,
+    state_ranges: dict[str, pd.DataFrame],
+    aug_period: int,
+    input_factor: str,
+    output_factor: str,
+    n_points: int,
+    n_draws: int,
+    transition_function: Callable[..., Array],
+    transition_params: dict[str, Any],
+    all_factors: tuple[str, ...],
+) -> pd.DataFrame:
     period_data = states_data.query(f"aug_period == {aug_period}")
 
     sampled_factors = [factor for factor in all_factors if factor != input_factor]
@@ -521,7 +546,7 @@ def _prepare_data_for_one_plot_average_2d(
         input_data[input_factor] = np.linspace(input_min, input_max, n_points)
         for col, val in draw.items():
             input_data[col] = val
-        input_arr = jnp.array(input_data[all_factors].to_numpy())
+        input_arr = jnp.array(input_data[list(all_factors)].to_numpy())
         # convert from jax to numpy array
         output_arr = np.array(transition_function(transition_params, input_arr))
         draw_data = pd.DataFrame()
@@ -529,11 +554,14 @@ def _prepare_data_for_one_plot_average_2d(
         draw_data[f"output_{output_factor}"] = np.array(output_arr)
         to_concat.append(draw_data)
 
-    out = pd.concat(to_concat).groupby(f"input_{input_factor}").mean().reset_index()
-    return out
+    return pd.concat(to_concat).groupby(f"input_{input_factor}").mean().reset_index()
 
 
-def _process_factor_mapping_trans(factor_mapper, output_factors, input_factors):
+def _process_factor_mapping_trans(
+    factor_mapper: dict[str, str] | None,
+    output_factors: tuple[str, ...],
+    input_factors: tuple[str, ...],
+) -> dict[str, str]:
     """Process mapper to return dictionary with old and new factor names."""
     all_factors = input_factors + output_factors
     if factor_mapper is None:
@@ -545,20 +573,32 @@ def _process_factor_mapping_trans(factor_mapper, output_factors, input_factors):
     return factor_mapper
 
 
-def _process_orders(columns, rows, plots_dict):
-    """Process axes orders to return list of strings."""
+def _process_orders(
+    columns: list[str] | tuple[str, ...] | str | None,
+    rows: list[str] | tuple[str, ...] | str | None,
+    plots_dict: dict[tuple[str, str], go.Figure],
+) -> tuple[tuple[str, ...], tuple[str, ...]]:
+    """Process axes orders to return tuples of strings."""
+    out_columns: tuple[str, ...]
+    out_rows: tuple[str, ...]
     if columns is None:
-        columns = []
+        seen: list[str] = []
         for f in plots_dict:
-            if f[0] not in columns:
-                columns.append(f[0])
+            if f[0] not in seen:
+                seen.append(f[0])
+        out_columns = tuple(seen)
     elif isinstance(columns, str):
-        columns = [columns]
+        out_columns = (columns,)
+    else:
+        out_columns = tuple(columns)
     if rows is None:
-        rows = []
+        seen = []
         for f in plots_dict:
-            if f[1] not in rows:
-                rows.append(f[1])
+            if f[1] not in seen:
+                seen.append(f[1])
+        out_rows = tuple(seen)
     elif isinstance(rows, str):
-        rows = [rows]
-    return columns, rows
+        out_rows = (rows,)
+    else:
+        out_rows = tuple(rows)
+    return out_columns, out_rows
diff --git a/tests/test_clipping.py b/tests/test_clipping.py
index b6bae71..1afd17e 100644
--- a/tests/test_clipping.py
+++ b/tests/test_clipping.py
@@ -1,10 +1,13 @@
+"""Tests for soft clipping functions."""
+
 import jax.numpy as jnp
 import numpy as np
 
 from skillmodels.clipping import soft_clipping
 
 
-def test_one_sided_soft_maximum():
+def test_one_sided_soft_maximum() -> None:
+    """Test soft maximum clipping with lower bound."""
     arr = jnp.array([-10.0, -5, -1, 1, 5, 10])
     lower_bound = -8
     lower_hardness = 3
@@ -21,7 +24,8 @@ def test_one_sided_soft_maximum():
     np.testing.assert_allclose(res[1:], arr[1:], rtol=1e-05)
 
 
-def test_one_sided_soft_minimum():
+def test_one_sided_soft_minimum() -> None:
+    """Test soft minimum clipping with upper bound."""
     arr = jnp.array([-10.0, -5, -1, 1, 5, 10])
     upper_bound = 8
     upper_hardness = 3
diff --git a/tests/test_constraints.py b/tests/test_constraints.py
index ca3bb10..40ed681 100644
--- a/tests/test_constraints.py
+++ b/tests/test_constraints.py
@@ -1,11 +1,11 @@
-from pathlib import Path
-
 import numpy as np
 import pandas as pd
 import pytest
 import yaml
+from frozendict import frozendict
 from pandas.testing import assert_frame_equal
 
+from skillmodels.config import TEST_DATA_DIR
 from skillmodels.constraints import (
     _get_anchoring_constraints,
     _get_constant_factors_constraints,
@@ -18,12 +18,10 @@
     add_bounds,
 )
 from skillmodels.process_model import process_model
-
-# importing the TEST_DIR from config does not work for test run in conda build
-TEST_DIR = Path(__file__).parent.resolve()
+from skillmodels.types import Anchoring, Labels
 
 
-def test_add_bounds():
+def test_add_bounds() -> None:
     ind_tups = [("shock_sds", i) for i in range(5)] + [
         ("meas_sds", 4),
         ("bla", "blubb"),
@@ -47,7 +45,7 @@ def test_add_bounds():
 # ======================================================================================
 
 
-def test_normalization_constraints():
+def test_normalization_constraints() -> None:
     norm = {
         "fac1": {
             "loadings": [{"m1": 2, "m2": 1.5}, {"m1": 3}],
@@ -84,7 +82,7 @@ def test_normalization_constraints():
         },
     ]
 
-    calculated = _get_normalization_constraints(norm, factors=["fac1", "fac2"])
+    calculated = _get_normalization_constraints(norm, factors=("fac1", "fac2"))
     for c in calculated:
         del c["description"]
 
@@ -96,7 +94,7 @@ def test_normalization_constraints():
 # ======================================================================================
 
 
-def test_mixture_weight_constraints_mixture():
+def test_mixture_weight_constraints_mixture() -> None:
     calculated = _get_mixture_weights_constraints(n_mixtures=2)
     for c in calculated:
         del c["description"]
@@ -104,7 +102,7 @@ def test_mixture_weight_constraints_mixture():
     assert_list_equal_except_for_order(calculated, expected)
 
 
-def test_mixture_weight_constraints_normal():
+def test_mixture_weight_constraints_normal() -> None:
     calculated = _get_mixture_weights_constraints(n_mixtures=1)
     for c in calculated:
         del c["description"]
@@ -117,9 +115,9 @@ def test_mixture_weight_constraints_normal():
 # ======================================================================================
 
 
-def test_stage_constraints():
-    stages = [0]
-    stagemap = [0] * 3
+def test_stage_constraints() -> None:
+    stages = (0,)
+    stagemap = (0, 0, 0)
 
     expected = [
         {
@@ -138,9 +136,9 @@ def test_stage_constraints():
     assert_list_equal_except_for_order(calculated, expected)
 
 
-def test_stage_constraints_with_endogenous_factors():
-    stages = [0, 1, 2, 3]
-    stagemap = [0, 1, 0, 1, 2, 3]
+def test_stage_constraints_with_endogenous_factors() -> None:
+    stages = (0, 1, 2, 3)
+    stagemap = (0, 1, 0, 1, 2, 3)
     expected = [
         {
             "loc": [("transition", 0), ("transition", 2)],
@@ -171,12 +169,21 @@ def test_stage_constraints_with_endogenous_factors():
 # ======================================================================================
 
 
-def test_constant_factor_constraints():
-    labels = {
-        "latent_factors": ["fac1", "fac2"],
-        "aug_periods": [0, 1, 2],
-        "transition_names": ["bla", "constant"],
-    }
+def test_constant_factor_constraints() -> None:
+    labels = Labels(
+        latent_factors=("fac1", "fac2"),
+        observed_factors=(),
+        controls=("constant",),
+        periods=(0, 1, 2),
+        stagemap=(0, 0, 0),
+        stages=(0,),
+        aug_periods=(0, 1, 2),
+        aug_periods_to_periods=frozendict({0: 0, 1: 1, 2: 2}),
+        aug_stagemap=(0, 0, 0),
+        aug_stages=(0,),
+        aug_stages_to_stages=frozendict({0: 0}),
+        transition_names=("bla", "constant"),
+    )
 
     expected = [
         {"loc": ("shock_sds", 0, "fac2", "-"), "type": "fixed", "value": 0.0},
@@ -194,9 +201,9 @@ def test_constant_factor_constraints():
 # ======================================================================================
 
 
-def test_initial_mean_constraints():
+def test_initial_mean_constraints() -> None:
     nmixtures = 3
-    factors = ["fac1", "fac2", "fac3"]
+    factors = ("fac1", "fac2", "fac3")
     ind_tups = [
         ("initial_states", 0, "mixture_0", "fac1"),
         ("initial_states", 0, "mixture_1", "fac1"),
@@ -216,13 +223,21 @@ def test_initial_mean_constraints():
 # ======================================================================================
 
 
-def test_trans_coeff_constraints():
-    labels = {
-        "latent_factors": ["fac1", "fac2", "fac3"],
-        "transition_names": ["log_ces", "bla", "blubb"],
-        "aug_periods": [0, 1, 2],
-    }
-    labels["all_factors"] = labels["latent_factors"]
+def test_trans_coeff_constraints() -> None:
+    labels = Labels(
+        latent_factors=("fac1", "fac2", "fac3"),
+        observed_factors=(),
+        controls=("constant",),
+        periods=(0, 1, 2),
+        stagemap=(0, 0, 0),
+        stages=(0,),
+        aug_periods=(0, 1, 2),
+        aug_periods_to_periods=frozendict({0: 0, 1: 1, 2: 2}),
+        aug_stagemap=(0, 0, 0),
+        aug_stages=(0,),
+        aug_stages_to_stages=frozendict({0: 0}),
+        transition_names=("log_ces", "bla", "blubb"),
+    )
 
     expected = [
         {
@@ -271,24 +286,35 @@ def anch_uinfo():
 
 @pytest.fixture
 def base_anchoring_info():
-    anch_info = {
-        "factors": ["f1", "f2"],
-        "outcomes": {"f1": "outcome", "f2": "outcome"},
-        "free_controls": True,
-        "free_constant": True,
-        "free_loadings": True,
-    }
-    return anch_info
+    return Anchoring(
+        anchoring=True,
+        factors=("f1", "f2"),
+        outcomes=frozendict({"f1": "outcome", "f2": "outcome"}),
+        free_controls=True,
+        free_constant=True,
+        free_loadings=True,
+        ignore_constant_when_anchoring=False,
+    )
 
 
-def test_anchoring_constraints_no_constraint_needed(anch_uinfo, base_anchoring_info):
-    calculated = _get_anchoring_constraints(anch_uinfo, [], base_anchoring_info, (0, 1))
+def test_anchoring_constraints_no_constraint_needed(
+    anch_uinfo, base_anchoring_info
+) -> None:
+    calculated = _get_anchoring_constraints(anch_uinfo, (), base_anchoring_info, (0, 1))
     assert calculated == []
 
 
-def test_anchoring_constraints_for_constants(anch_uinfo, base_anchoring_info):
-    base_anchoring_info["free_constant"] = False
-    calculated = _get_anchoring_constraints(anch_uinfo, [], base_anchoring_info, (0, 1))
+def test_anchoring_constraints_for_constants(anch_uinfo) -> None:
+    anchoring_info = Anchoring(
+        anchoring=True,
+        factors=("f1", "f2"),
+        outcomes=frozendict({"f1": "outcome", "f2": "outcome"}),
+        free_controls=True,
+        free_constant=False,
+        free_loadings=True,
+        ignore_constant_when_anchoring=False,
+    )
+    calculated = _get_anchoring_constraints(anch_uinfo, (), anchoring_info, (0, 1))
 
     del calculated[0]["description"]
     expected = [
@@ -307,12 +333,20 @@ def test_anchoring_constraints_for_constants(anch_uinfo, base_anchoring_info):
     assert calculated == expected
 
 
-def test_anchoring_constraints_for_controls(anch_uinfo, base_anchoring_info):
-    base_anchoring_info["free_controls"] = False
+def test_anchoring_constraints_for_controls(anch_uinfo) -> None:
+    anchoring_info = Anchoring(
+        anchoring=True,
+        factors=("f1", "f2"),
+        outcomes=frozendict({"f1": "outcome", "f2": "outcome"}),
+        free_controls=False,
+        free_constant=True,
+        free_loadings=True,
+        ignore_constant_when_anchoring=False,
+    )
     calculated = _get_anchoring_constraints(
         anch_uinfo,
-        ["c1", "c2"],
-        base_anchoring_info,
+        ("c1", "c2"),
+        anchoring_info,
         (0, 1),
     )
 
@@ -339,9 +373,17 @@ def test_anchoring_constraints_for_controls(anch_uinfo, base_anchoring_info):
     assert calculated == expected
 
 
-def test_anchoring_constraints_for_loadings(anch_uinfo, base_anchoring_info):
-    base_anchoring_info["free_loadings"] = False
-    calculated = _get_anchoring_constraints(anch_uinfo, [], base_anchoring_info, (0, 1))
+def test_anchoring_constraints_for_loadings(anch_uinfo) -> None:
+    anchoring_info = Anchoring(
+        anchoring=True,
+        factors=("f1", "f2"),
+        outcomes=frozendict({"f1": "outcome", "f2": "outcome"}),
+        free_controls=True,
+        free_constant=True,
+        free_loadings=False,
+        ignore_constant_when_anchoring=False,
+    )
+    calculated = _get_anchoring_constraints(anch_uinfo, (), anchoring_info, (0, 1))
 
     expected = [
         {
@@ -362,7 +404,7 @@ def test_anchoring_constraints_for_loadings(anch_uinfo, base_anchoring_info):
     assert calculated == expected
 
 
-def assert_list_equal_except_for_order(list1, list2):
+def assert_list_equal_except_for_order(list1, list2) -> None:
     for item in list1:
         assert item in list2, f"{item} is in list1 but not in list2"
     for item in list2:
@@ -371,15 +413,15 @@ def assert_list_equal_except_for_order(list1, list2):
 
 @pytest.fixture
 def simplest_augmented_model():
-    with open(TEST_DIR / "simplest_augmented_model.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
+    with (TEST_DATA_DIR / "simplest_augmented_model.yaml").open() as y:
+        model_dict = yaml.load(y, Loader=yaml.SafeLoader)
     return process_model(model_dict)
 
 
-def test_get_constraints_for_augmented_periods(simplest_augmented_model):
+def test_get_constraints_for_augmented_periods(simplest_augmented_model) -> None:
     calculated = _get_constraints_for_augmented_periods(
-        labels=simplest_augmented_model["labels"],
-        endogenous_factors_info=simplest_augmented_model["endogenous_factors_info"],
+        labels=simplest_augmented_model.labels,
+        endogenous_factors_info=simplest_augmented_model.endogenous_factors_info,
     )
     for c in calculated:
         del c["description"]
diff --git a/tests/test_correlation_heatmap.py b/tests/test_correlation_heatmap.py
index 6f7e62a..1784c22 100644
--- a/tests/test_correlation_heatmap.py
+++ b/tests/test_correlation_heatmap.py
@@ -1,5 +1,8 @@
+from types import SimpleNamespace
+
 import numpy as np
 import pandas as pd
+from frozendict import frozendict
 from pandas.testing import assert_frame_equal as afe
 
 from skillmodels.correlation_heatmap import (
@@ -10,9 +13,10 @@
     _get_quasi_factor_scores_data_for_single_period,
     _process_factors,
 )
+from skillmodels.types import Labels
 
 
-def test_get_measurement_data_with_single_period():
+def test_get_measurement_data_with_single_period() -> None:
     period = 1
     factors = ["f3", "f1"]
     update_info = pd.DataFrame(
@@ -60,7 +64,7 @@ def test_get_measurement_data_with_single_period():
     afe(result, expected)
 
 
-def test_get_factor_scores_data_with_single_period():
+def test_get_factor_scores_data_with_single_period() -> None:
     period = 1
     factors = ["f1", "f2"]
     update_info = pd.DataFrame(
@@ -113,7 +117,7 @@ def test_get_factor_scores_data_with_single_period():
     afe(expected, result, check_dtype=False)
 
 
-def test_get_measurement_data_with_multiple_periods():
+def test_get_measurement_data_with_multiple_periods() -> None:
     period = [1, 2]
     factors = ["f3", "f1"]
     update_info = pd.DataFrame(
@@ -172,7 +176,7 @@ def test_get_measurement_data_with_multiple_periods():
     afe(result, expected)
 
 
-def test_get_factor_scores_data_with_multiple_period():
+def test_get_factor_scores_data_with_multiple_period() -> None:
     periods = [0, 1]
     factors = ["f1", "f2"]
     update_info = pd.DataFrame(
@@ -241,24 +245,36 @@ def test_get_factor_scores_data_with_multiple_period():
     afe(expected, result)
 
 
-def test_process_factors():
-    model = {
-        "labels": {"latent_factors": list("abcd"), "observed_factors": list("efg")},
-    }
+def test_process_factors() -> None:
+    model = SimpleNamespace(
+        labels=Labels(
+            latent_factors=tuple("abcd"),
+            observed_factors=tuple("efg"),
+            controls=("constant",),
+            periods=(0,),
+            stagemap=(0,),
+            stages=(0,),
+            aug_periods=(0,),
+            aug_periods_to_periods=frozendict({0: 0}),
+            aug_stagemap=(0,),
+            aug_stages=(0,),
+            aug_stages_to_stages=frozendict({0: 0}),
+        ),
+    )
     latent_factor = "c"
     observed_factor = "g"
     factors = ["b", "d", "g"]
     all_factors = None
-    assert list("abcd") == _process_factors(model, all_factors)[0]
-    assert list("efg") == _process_factors(model, all_factors)[1]
-    assert [latent_factor] == _process_factors(model, latent_factor)[0]
-    assert [observed_factor] == _process_factors(model, observed_factor)[1]
-    assert factors[:-1] == _process_factors(model, factors)[0]
-    assert [factors[-1] == _process_factors(model, factors)[1]]
+    assert tuple("abcd") == _process_factors(model, all_factors)[0]  # ty: ignore[invalid-argument-type]
+    assert tuple("efg") == _process_factors(model, all_factors)[1]  # ty: ignore[invalid-argument-type]
+    assert (latent_factor,) == _process_factors(model, latent_factor)[0]  # ty: ignore[invalid-argument-type]
+    assert (observed_factor,) == _process_factors(model, observed_factor)[1]  # ty: ignore[invalid-argument-type]
+    assert tuple(factors[:-1]) == _process_factors(model, factors)[0]  # ty: ignore[invalid-argument-type]
+    assert (factors[-1],) == _process_factors(model, factors)[1]  # ty: ignore[invalid-argument-type]
 
 
-def test_get_mask_lower_triangle_only():
-    corr = np.ones((4, 4))
+def test_get_mask_lower_triangle_only() -> None:
+    corr = pd.DataFrame(np.ones((4, 4)))
     show_upper = False
     show_diag = False
     expected = np.array(
@@ -269,12 +285,12 @@ def test_get_mask_lower_triangle_only():
             [True] * 3 + [False],
         ],
     )
-    result = _get_mask(corr, show_upper, show_diag)
+    result = _get_mask(corr, show_upper_triangle=show_upper, show_diagonal=show_diag)
     np.testing.assert_array_equal(result, expected)
 
 
-def test_get_mask_lower_triangle_and_diag():
-    corr = np.ones((4, 4))
+def test_get_mask_lower_triangle_and_diag() -> None:
+    corr = pd.DataFrame(np.ones((4, 4)))
     show_upper = False
     show_diag = True
     expected = np.array(
@@ -285,12 +301,12 @@ def test_get_mask_lower_triangle_and_diag():
             [True] * 4,
         ],
     )
-    result = _get_mask(corr, show_upper, show_diag)
+    result = _get_mask(corr, show_upper_triangle=show_upper, show_diagonal=show_diag)
     np.testing.assert_array_equal(result, expected)
 
 
-def test_get_mask_lower_and_upper_triangle_no_diag():
-    corr = np.ones((4, 4))
+def test_get_mask_lower_and_upper_triangle_no_diag() -> None:
+    corr = pd.DataFrame(np.ones((4, 4)))
     show_upper = True
     show_diag = False
     expected = np.array(
@@ -301,14 +317,14 @@ def test_get_mask_lower_and_upper_triangle_no_diag():
             [True] * 3 + [False],
         ],
     )
-    result = _get_mask(corr, show_upper, show_diag)
+    result = _get_mask(corr, show_upper_triangle=show_upper, show_diagonal=show_diag)
     np.testing.assert_array_equal(result, expected)
 
 
-def test_get_mask_full_square_matrix():
-    corr = np.ones((4, 4))
+def test_get_mask_full_square_matrix() -> None:
+    corr = pd.DataFrame(np.ones((4, 4)))
     show_upper = True
     show_diag = True
-    expected = corr.astype(bool)
-    result = _get_mask(corr, show_upper, show_diag)
+    expected = corr.to_numpy().astype(bool)
+    result = _get_mask(corr, show_upper_triangle=show_upper, show_diagonal=show_diag)
     np.testing.assert_array_equal(result, expected)
diff --git a/tests/test_decorators.py b/tests/test_decorators.py
index 85f9ede..fcfcc76 100644
--- a/tests/test_decorators.py
+++ b/tests/test_decorators.py
@@ -3,7 +3,7 @@
 from skillmodels.decorators import extract_params, jax_array_output, register_params
 
 
-def test_extract_params_decorator_only_key():
+def test_extract_params_decorator_only_key() -> None:
     @extract_params(key="a")
     def f(x, params):
         return x * params
@@ -11,7 +11,7 @@ def f(x, params):
     assert f(x=3, params={"a": 4, "b": 5}) == 12
 
 
-def test_extract_params_direct_call_only_key():
+def test_extract_params_direct_call_only_key() -> None:
     def f(x, params):
         return x * params
 
@@ -20,7 +20,7 @@ def f(x, params):
     assert g(x=3, params={"a": 4, "b": 5}) == 12
 
 
-def test_extract_params_decorator_only_names():
+def test_extract_params_decorator_only_names() -> None:
     @extract_params(names=["c", "d"])
     def f(x, params):
         return x * params["c"]
@@ -28,7 +28,7 @@ def f(x, params):
     assert f(x=3, params=[4, 5]) == 12
 
 
-def test_extract_params_direct_call_only_names():
+def test_extract_params_direct_call_only_names() -> None:
     def f(x, params):
         return x * params["c"]
 
@@ -36,7 +36,7 @@ def f(x, params):
     assert g(x=3, params=[4, 5]) == 12
 
 
-def test_extract_params_decorator_key_and_names():
+def test_extract_params_decorator_key_and_names() -> None:
     @extract_params(key="a", names=["c", "d"])
     def f(x, params):
         return x * params["c"]
@@ -44,7 +44,7 @@ def f(x, params):
     assert f(x=3, params={"a": [4, 5], "b": [5, 6]}) == 12
 
 
-def test_extract_params_direct_call_key_and_names():
+def test_extract_params_direct_call_key_and_names() -> None:
     def f(x, params):
         return x * params["c"]
 
@@ -52,7 +52,7 @@ def f(x, params):
     assert g(x=3, params={"a": [4, 5], "b": [5, 6]}) == 12
 
 
-def test_jax_array_output_decorator():
+def test_jax_array_output_decorator() -> None:
     @jax_array_output
     def f():
         return (1, 2, 3)
@@ -60,7 +60,7 @@ def f():
     assert isinstance(f(), jnp.ndarray)
 
 
-def test_jax_array_output_direct_call():
+def test_jax_array_output_direct_call() -> None:
     def f():
         return (1, 2, 3)
 
@@ -69,19 +69,19 @@ def f():
     assert isinstance(g(), jnp.ndarray)
 
 
-def test_register_params_decorator():
+def test_register_params_decorator() -> None:
     @register_params(params=["a", "b", "c"])
-    def f():
+    def f() -> str:
         return "bla"
 
     assert f.__registered_params__ == ["a", "b", "c"]
     assert f() == "bla"
 
 
-def test_register_params_direct_call():
-    def f():
+def test_register_params_direct_call() -> None:
+    def f() -> str:
         return "bla"
 
     g = register_params(f, params=["a", "b", "c"])
-    assert g.__registered_params__ == ["a", "b", "c"]
+    assert g.__registered_params__ == ["a", "b", "c"]  # ty: ignore[unresolved-attribute]
     assert g() == "bla"
diff --git a/tests/test_filtered_states.py b/tests/test_filtered_states.py
index b98f15b..bcfc0fc 100644
--- a/tests/test_filtered_states.py
+++ b/tests/test_filtered_states.py
@@ -5,29 +5,27 @@
 import pytest
 import yaml
 
+from skillmodels.config import TEST_DATA_DIR
 from skillmodels.filtered_states import get_filtered_states
 from skillmodels.maximization_inputs import get_maximization_inputs
 
-# importing the TEST_DIR from config does not work for test run in conda build
-TEST_DIR = Path(__file__).parent.resolve()
+REGRESSION_VAULT = Path(__file__).parent / "regression_vault"
 
 
 @pytest.fixture
 def model2():
-    with open(TEST_DIR / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
-    return model_dict
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        return yaml.load(y, Loader=yaml.SafeLoader)
 
 
 @pytest.fixture
 def model2_data():
-    data = pd.read_stata(TEST_DIR / "model2_simulated_data.dta")
-    data = data.set_index(["caseid", "period"])
-    return data
+    data = pd.read_stata(TEST_DATA_DIR / "model2_simulated_data.dta")
+    return data.set_index(["caseid", "period"])
 
 
-def test_get_filtered_states(model2, model2_data):
-    params = pd.read_csv(TEST_DIR / "regression_vault" / "one_stage_anchoring.csv")
+def test_get_filtered_states(model2, model2_data) -> None:
+    params = pd.read_csv(REGRESSION_VAULT / "one_stage_anchoring.csv")
     params = params.set_index(["category", "period", "name1", "name2"])
 
     max_inputs = get_maximization_inputs(model2, model2_data)
diff --git a/tests/test_kalman_filters.py b/tests/test_kalman_filters.py
index 4c71386..b9c45d3 100644
--- a/tests/test_kalman_filters.py
+++ b/tests/test_kalman_filters.py
@@ -28,7 +28,7 @@
 
 
 @pytest.mark.parametrize(("seed", "update_func"), product(SEEDS, UPDATE_FUNCS))
-def test_kalman_update(seed, update_func):
+def test_kalman_update(seed, update_func) -> None:
     np.random.seed(seed)
     dim = np.random.randint(low=1, high=10)
     n_obs = 5
@@ -86,7 +86,7 @@ def test_kalman_update(seed, update_func):
 
 
 @pytest.mark.parametrize("update_func", UPDATE_FUNCS)
-def test_kalman_update_with_missing(update_func):
+def test_kalman_update_with_missing(update_func) -> None:
     """State, cov and weights should not change, log likelihood should be zero."""
     n_mixtures = 2
     n_obs = 3
@@ -133,10 +133,10 @@ def test_kalman_update_with_missing(update_func):
 
 
 @pytest.mark.parametrize("seed", SEEDS)
-def test_sigma_points(seed):
+def test_sigma_points(seed: int) -> None:
     np.random.seed(seed)
     state, cov = _random_state_and_covariance()
-    observed_factors = np.arange(2).reshape(1, 2)
+    observed_factors = jnp.arange(2).reshape(1, 2)
     expected = JulierSigmaPoints(n=len(state), kappa=2).sigma_points(state, cov)
     observed_part = np.tile(observed_factors, len(expected)).reshape(-1, 2)
     expected = np.hstack([expected, observed_part])
@@ -157,7 +157,7 @@ def test_sigma_points(seed):
 
 
 @pytest.mark.parametrize("seed", SEEDS)
-def test_sigma_scaling_factor_and_weights(seed):
+def test_sigma_scaling_factor_and_weights(seed) -> None:
     np.random.seed(seed)
     dim = np.random.randint(low=1, high=15)
     kappa = np.random.uniform(low=0.5, high=5)
@@ -176,20 +176,19 @@ def test_sigma_scaling_factor_and_weights(seed):
 # ======================================================================================
 
 
-def test_transformation_of_sigma_points():
+def test_transformation_of_sigma_points() -> None:
     sp = jnp.arange(10).reshape(1, 1, 5, 2) + 1
 
     def f(params, states):
-        out = jnp.column_stack(
+        return jnp.column_stack(
             [(states * params["fac1"][0]).sum(axis=1), states[..., 1]],
         )
-        return out
 
     trans_coeffs = {"fac1": jnp.array([2]), "fac2": jnp.array([])}
 
     anch_scaling = jnp.array([[1, 1], [2, 1]])
 
-    anch_constants = np.array([[0, 0], [0, 0]])
+    anch_constants = jnp.array([[0, 0], [0, 0]])
 
     expected = jnp.array([[[[3, 2], [7, 4], [11, 6], [15, 8], [19, 10]]]])
 
@@ -213,7 +212,7 @@ def f(params, states):
 
 
 @pytest.mark.parametrize("seed", SEEDS)
-def test_predict_against_linear_filterpy(seed):
+def test_predict_against_linear_filterpy(seed) -> None:
     np.random.seed(seed)
     state, cov = _random_state_and_covariance()
     dim = len(state)
@@ -235,8 +234,7 @@ def linear(params, states):
         return jnp.dot(states, params)
 
     def transition_function(params, states):
-        out = jnp.column_stack([linear(params[f"fac{i}"], states) for i in range(dim)])
-        return out
+        return jnp.column_stack([linear(params[f"fac{i}"], states) for i in range(dim)])
 
     sm_state, sm_chol = _convert_predict_inputs_from_filterpy_to_skillmodels(state, cov)
     scaling_factor, weights = calculate_sigma_scaling_factor_and_weights(dim, 2)
@@ -249,13 +247,13 @@ def transition_function(params, states):
         transition_function,
         sm_state,
         sm_chol,
-        scaling_factor,
+        float(scaling_factor),
         weights,
         trans_coeffs,
         jnp.array(shock_sds),
         anch_scaling,
         anch_constants,
-        observed_factors,
+        jnp.asarray(observed_factors),
     )
 
     aaae(calc_states.flatten(), expected_state.flatten())
diff --git a/tests/test_likelihood_regression.py b/tests/test_likelihood_regression.py
index cd4cbc0..962c4a4 100644
--- a/tests/test_likelihood_regression.py
+++ b/tests/test_likelihood_regression.py
@@ -9,6 +9,7 @@
 import yaml
 from numpy.testing import assert_array_almost_equal as aaae
 
+from skillmodels.config import TEST_DATA_DIR
 from skillmodels.decorators import register_params
 from skillmodels.maximization_inputs import get_maximization_inputs
 from skillmodels.utilities import reduce_n_periods
@@ -23,22 +24,19 @@
     "one_stage_anchoring_custom_functions",
 ]
 
-# importing the TEST_DIR from config does not work for test run in conda build
-TEST_DIR = Path(__file__).parent.resolve()
+REGRESSION_VAULT = Path(__file__).parent / "regression_vault"
 
 
 @pytest.fixture
 def model2():
-    with open(TEST_DIR / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
-    return model_dict
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        return yaml.load(y, Loader=yaml.SafeLoader)
 
 
 @pytest.fixture
 def model2_data():
-    data = pd.read_stata(TEST_DIR / "model2_simulated_data.dta")
-    data = data.set_index(["caseid", "period"])
-    return data
+    data = pd.read_stata(TEST_DATA_DIR / "model2_simulated_data.dta")
+    return data.set_index(["caseid", "period"])
 
 
 def _convert_model(base_model, model_name):
@@ -60,8 +58,9 @@ def constant(fac3, params):
         @register_params(params=["fac1", "fac2", "fac3", "constant"])
         def linear(fac1, fac2, fac3, params):
             p = params
-            out = p["constant"] + fac1 * p["fac1"] + fac2 * p["fac2"] + fac3 * p["fac3"]
-            return out
+            return (
+                p["constant"] + fac1 * p["fac1"] + fac2 * p["fac2"] + fac3 * p["fac3"]
+            )
 
         model["factors"]["fac2"]["transition_function"] = linear
         model["factors"]["fac3"]["transition_function"] = constant
@@ -73,8 +72,10 @@ def linear(fac1, fac2, fac3, params):
 @pytest.mark.parametrize(
     ("model_name", "fun_key"), product(MODEL_NAMES, ["loglike", "debug_loglike"])
 )
-def test_likelihood_values_have_not_changed(model2, model2_data, model_name, fun_key):
-    regvault = TEST_DIR / "regression_vault"
+def test_likelihood_values_have_not_changed(
+    model2, model2_data, model_name, fun_key
+) -> None:
+    regvault = REGRESSION_VAULT
     model = _convert_model(model2, model_name)
     params = pd.read_csv(regvault / f"{model_name}.csv").set_index(
         ["category", "period", "name1", "name2"],
@@ -87,12 +88,12 @@ def test_likelihood_values_have_not_changed(model2, model2_data, model_name, fun
     fun = inputs[fun_key]
     new_loglike = fun(params)["value"] if "debug" in fun_key else fun(params)
 
-    with open(regvault / f"{model_name}_result.json") as j:
+    with (regvault / f"{model_name}_result.json").open() as j:
         old_loglike = np.array(json.load(j)).sum()
     aaae(new_loglike, old_loglike)
 
 
-def test_splitting_does_not_change_gradient(model2, model2_data):
+def test_splitting_does_not_change_gradient(model2, model2_data) -> None:
     inputs = get_maximization_inputs(model2, model2_data)
     inputs_split = get_maximization_inputs(model2, model2_data, 13)
 
@@ -110,8 +111,8 @@ def test_splitting_does_not_change_gradient(model2, model2_data):
 )
 def test_likelihood_contributions_have_not_changed(
     model2, model2_data, model_name, fun_key
-):
-    regvault = TEST_DIR / "regression_vault"
+) -> None:
+    regvault = REGRESSION_VAULT
     model = _convert_model(model2, model_name)
     params = pd.read_csv(regvault / f"{model_name}.csv").set_index(
         ["category", "period", "name1", "name2"],
@@ -124,7 +125,7 @@ def test_likelihood_contributions_have_not_changed(
     fun = inputs[fun_key]
     new_loglikes = fun(params)["contributions"] if "debug" in fun_key else fun(params)
 
-    with open(regvault / f"{model_name}_result.json") as j:
+    with (regvault / f"{model_name}_result.json").open() as j:
         old_loglikes = np.array(json.load(j))
     aaae(new_loglikes, old_loglikes)
 
@@ -133,8 +134,10 @@ def test_likelihood_contributions_have_not_changed(
     ("model_type", "fun_key"),
     product(["no_stages_anchoring", "with_missings"], ["loglike_and_gradient"]),
 )
-def test_likelihood_contributions_large_nobs(model2, model2_data, model_type, fun_key):
-    regvault = TEST_DIR / "regression_vault"
+def test_likelihood_contributions_large_nobs(
+    model2, model2_data, model_type, fun_key
+) -> None:
+    regvault = REGRESSION_VAULT
     model = _convert_model(model2, "no_stages_anchoring")
     params = pd.read_csv(regvault / "no_stages_anchoring.csv").set_index(
         ["category", "period", "name1", "name2"],
@@ -191,7 +194,7 @@ def test_likelihood_contributions_large_nobs(model2, model2_data, model_type, fu
     assert np.isfinite(loglike[1]).all()
 
 
-def test_likelihood_runs_with_empty_periods(model2, model2_data):
+def test_likelihood_runs_with_empty_periods(model2, model2_data) -> None:
     del model2["anchoring"]
     for factor in ["fac1", "fac2"]:
         model2["factors"][factor]["measurements"][-1] = []
@@ -206,9 +209,9 @@ def test_likelihood_runs_with_empty_periods(model2, model2_data):
     debug_loglike(params)
 
 
-def test_likelihood_runs_with_too_long_data(model2, model2_data):
+def test_likelihood_runs_with_too_long_data(model2, model2_data) -> None:
     model = reduce_n_periods(model2, 2)
-    func_dict = get_maximization_inputs(model, model2_data)
+    func_dict = get_maximization_inputs(model, model2_data)  # ty: ignore[invalid-argument-type]
 
     params = func_dict["params_template"]
     params["value"] = 0.1
@@ -217,7 +220,7 @@ def test_likelihood_runs_with_too_long_data(model2, model2_data):
     debug_loglike(params)
 
 
-def test_likelihood_runs_with_observed_factors(model2, model2_data):
+def test_likelihood_runs_with_observed_factors(model2, model2_data) -> None:
     model2["observed_factors"] = ["ob1", "ob2"]
     model2_data["ob1"] = np.arange(len(model2_data))
     model2_data["ob2"] = np.ones(len(model2_data))
diff --git a/tests/test_maximization_inputs.py b/tests/test_maximization_inputs.py
index 0f901f2..1f2dd6f 100644
--- a/tests/test_maximization_inputs.py
+++ b/tests/test_maximization_inputs.py
@@ -1,21 +1,26 @@
+"""Tests for maximization input functions."""
+
 import jax.numpy as jnp
 import numpy as np
 
 from skillmodels.maximization_inputs import _to_numpy
 
 
-def test_to_numpy_with_dict():
+def test_to_numpy_with_dict() -> None:
+    """Test _to_numpy with dictionary input."""
     dict_ = {"a": jnp.ones(3), "b": 4.5}
     calculated = _to_numpy(dict_)
     assert isinstance(calculated["a"], np.ndarray)
     assert isinstance(calculated["b"], float)
 
 
-def test_to_numpy_one_array():
+def test_to_numpy_one_array() -> None:
+    """Test _to_numpy with single array input."""
     calculated = _to_numpy(jnp.ones(3))
     assert isinstance(calculated, np.ndarray)
 
 
-def test_to_numpy_one_float():
+def test_to_numpy_one_float() -> None:
+    """Test _to_numpy with single float input."""
     calculated = _to_numpy(3.5)
     assert isinstance(calculated, float)
diff --git a/tests/test_params_index.py b/tests/test_params_index.py
index 9e10d61..b34b2b4 100644
--- a/tests/test_params_index.py
+++ b/tests/test_params_index.py
@@ -1,9 +1,9 @@
-from pathlib import Path
-
 import pandas as pd
 import pytest
 import yaml
+from frozendict import frozendict
 
+from skillmodels.config import TEST_DATA_DIR
 from skillmodels.params_index import (
     get_control_params_index_tuples,
     get_initial_cholcovs_index_tuples,
@@ -16,40 +16,38 @@
     initial_mean_index_tuples,
 )
 from skillmodels.process_model import process_model
+from skillmodels.types import TransitionInfo
 
 
 @pytest.fixture
 def model2_inputs():
-    test_dir = Path(__file__).parent.resolve()
-    with open(test_dir / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        model_dict = yaml.load(y, Loader=yaml.SafeLoader)
     processed = process_model(model_dict)
 
-    out = {
-        "update_info": processed["update_info"],
-        "labels": processed["labels"],
-        "dimensions": processed["dimensions"],
-        "transition_info": processed["transition_info"],
-        "endogenous_factors_info": processed["endogenous_factors_info"],
+    return {
+        "update_info": processed.update_info,
+        "labels": processed.labels,
+        "dimensions": processed.dimensions,
+        "transition_info": processed.transition_info,
+        "endogenous_factors_info": processed.endogenous_factors_info,
     }
-    return out
 
 
-def test_params_index_with_model2(model2_inputs):
-    test_dir = Path(__file__).parent.resolve()
+def test_params_index_with_model2(model2_inputs) -> None:
     calculated = get_params_index(**model2_inputs)
     expected = pd.read_csv(
-        test_dir / "model2_correct_params_index.csv",
+        TEST_DATA_DIR / "model2_correct_params_index.csv",
         index_col=["category", "period", "name1", "name2"],
     ).index
 
     assert calculated.equals(expected)
 
 
-def test_control_coeffs_index_tuples():
+def test_control_coeffs_index_tuples() -> None:
     uinfo_tups = [(0, "m1"), (0, "m2"), (0, "bla"), (1, "m1"), (1, "m2")]
     uinfo = pd.DataFrame(index=pd.MultiIndex.from_tuples(uinfo_tups))
-    controls = ["constant", "c1"]
+    controls = ("constant", "c1")
 
     expected = [
         ("controls", 0, "m1", "constant"),
@@ -68,14 +66,14 @@ def test_control_coeffs_index_tuples():
     assert calculated == expected
 
 
-def test_loading_index_tuples():
+def test_loading_index_tuples() -> None:
     uinfo_tups = [(0, "m1"), (0, "m2"), (0, "bla"), (1, "m1"), (1, "m2")]
     uinfo = pd.DataFrame(
         True,
         index=pd.MultiIndex.from_tuples(uinfo_tups),
         columns=["fac1", "fac2"],
     )
-    factors = ["fac1", "fac2"]
+    factors = ("fac1", "fac2")
     expected = [
         ("loadings", 0, "m1", "fac1"),
         ("loadings", 0, "m1", "fac2"),
@@ -93,7 +91,7 @@ def test_loading_index_tuples():
     assert calculated == expected
 
 
-def test_meas_sd_index_tuples():
+def test_meas_sd_index_tuples() -> None:
     uinfo_tups = [(0, "m1"), (0, "m2"), (0, "bla"), (1, "m1"), (1, "m2")]
     uinfo = pd.DataFrame(index=pd.MultiIndex.from_tuples(uinfo_tups))
 
@@ -109,9 +107,9 @@ def test_meas_sd_index_tuples():
     assert calculated == expected
 
 
-def test_shock_sd_index_tuples():
-    periods = [0, 1, 2]
-    factors = ["fac1", "fac2"]
+def test_shock_sd_index_tuples() -> None:
+    periods = (0, 1, 2)
+    factors = ("fac1", "fac2")
 
     expected = [
         ("shock_sds", 0, "fac1", "-"),
@@ -120,13 +118,15 @@ def test_shock_sd_index_tuples():
         ("shock_sds", 1, "fac2", "-"),
     ]
 
-    calculated = get_shock_sds_index_tuples(periods, factors, False)
+    calculated = get_shock_sds_index_tuples(
+        periods, factors, has_endogenous_factors=False
+    )
     assert calculated == expected
 
 
-def test_initial_mean_index_tuples():
+def test_initial_mean_index_tuples() -> None:
     nmixtures = 3
-    factors = ["fac1", "fac2"]
+    factors = ("fac1", "fac2")
 
     expected = [
         ("initial_states", 0, "mixture_0", "fac1"),
@@ -141,7 +141,7 @@ def test_initial_mean_index_tuples():
     assert calculated == expected
 
 
-def test_mixture_weight_index_tuples():
+def test_mixture_weight_index_tuples() -> None:
     nmixtures = 3
     expected = [
         ("mixture_weights", 0, "mixture_0", "-"),
@@ -152,9 +152,9 @@ def test_mixture_weight_index_tuples():
     assert calculated == expected
 
 
-def test_initial_cov_index_tuples():
+def test_initial_cov_index_tuples() -> None:
     nmixtures = 2
-    factors = ["fac1", "fac2", "fac3"]
+    factors = ("fac1", "fac2", "fac3")
     expected = [
         ("initial_cholcovs", 0, "mixture_0", "fac1-fac1"),
         ("initial_cholcovs", 0, "mixture_0", "fac2-fac1"),
@@ -174,15 +174,20 @@ def test_initial_cov_index_tuples():
     assert calculated == expected
 
 
-def test_trans_coeffs_index_tuples_no_endogenous_factors():
-    periods = [0, 1, 2]
+def test_trans_coeffs_index_tuples_no_endogenous_factors() -> None:
+    periods = (0, 1, 2)
 
     param_names = {
         "fac1": ["fac1", "fac2", "fac3", "constant"],
         "fac2": [],
         "fac3": ["fac1", "fac2", "fac3", "phi"],
     }
-    trans_info = {"param_names": param_names}
+    trans_info = TransitionInfo(
+        func=lambda x: x,  # dummy function
+        param_names=frozendict(param_names),
+        individual_functions=frozendict({}),
+        function_names=frozendict({}),
+    )
 
     expected = [
         ("transition", 0, "fac1", "fac1"),
@@ -212,15 +217,20 @@ def test_trans_coeffs_index_tuples_no_endogenous_factors():
     assert calculated == expected
 
 
-def test_trans_coeffs_index_tuples_has_endogenous_factors():
-    periods = [0, 1, 2, 3, 4, 5]
+def test_trans_coeffs_index_tuples_has_endogenous_factors() -> None:
+    periods = (0, 1, 2, 3, 4, 5)
 
     param_names = {
         "fac1": ["fac1", "fac2", "fac3", "constant"],
         "fac2": [],
         "fac3": ["fac1", "fac2", "fac3", "phi"],
     }
-    trans_info = {"param_names": param_names}
+    trans_info = TransitionInfo(
+        func=lambda x: x,  # dummy function
+        param_names=frozendict(param_names),
+        individual_functions=frozendict({}),
+        function_names=frozendict({}),
+    )
 
     expected = [
         ("transition", 0, "fac1", "fac1"),
diff --git a/tests/test_parse_params.py b/tests/test_parse_params.py
index 33ffbe8..35e8d2e 100644
--- a/tests/test_parse_params.py
+++ b/tests/test_parse_params.py
@@ -5,41 +5,49 @@
 
 """
 
-from pathlib import Path
-
 import jax.numpy as jnp
 import numpy as np
 import pandas as pd
 import pytest
 import yaml
+from frozendict import frozendict
 from numpy.testing import assert_array_equal as aae
 
+from skillmodels.config import TEST_DATA_DIR
 from skillmodels.parse_params import create_parsing_info, parse_params
 from skillmodels.process_model import process_model
+from skillmodels.types import Anchoring
 
 
 @pytest.fixture
 def parsed_parameters():
-    test_dir = Path(__file__).parent.resolve()
     p_index = pd.read_csv(
-        test_dir / "model2_correct_params_index.csv",
+        TEST_DATA_DIR / "model2_correct_params_index.csv",
         index_col=["category", "period", "name1", "name2"],
     ).index
 
-    with open(test_dir / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        model_dict = yaml.load(y, Loader=yaml.SafeLoader)
 
     processed = process_model(model_dict)
 
-    update_info = processed["update_info"]
-    labels = processed["labels"]
-    dimensions = processed["dimensions"]
+    update_info = processed.update_info
+    labels = processed.labels
+    dimensions = processed.dimensions
     # this overwrites the anchoring setting from the model specification to get a
     # more meaningful test
-    anchoring = {"ignore_constant_when_anchoring": False}
+    anchoring = Anchoring(
+        anchoring=False,
+        outcomes=frozendict({}),
+        factors=(),
+        free_controls=True,
+        free_constant=True,
+        free_loadings=True,
+        ignore_constant_when_anchoring=False,
+    )
 
     parsing_info = create_parsing_info(
-        params_index=p_index,
+        params_index=p_index,  # ty: ignore[invalid-argument-type]
         update_info=update_info,
         labels=labels,
         anchoring=anchoring,
@@ -49,41 +57,46 @@ def parsed_parameters():
     params_vec = jnp.arange(len(p_index))
     n_obs = 5
 
-    parsed = parse_params(params_vec, parsing_info, dimensions, labels, n_obs)
-
-    return dict(
-        zip(["states", "upper_chols", "log_weights", "pardict"], parsed, strict=False)
+    states, upper_chols, log_weights, parsed_params = parse_params(
+        params_vec, parsing_info, dimensions, labels, n_obs
     )
 
+    return {
+        "states": states,
+        "upper_chols": upper_chols,
+        "log_weights": log_weights,
+        "parsed_params": parsed_params,
+    }
+
 
-def test_controls(parsed_parameters):
+def test_controls(parsed_parameters) -> None:
     expected = jnp.arange(118).reshape(59, 2)
-    aae(parsed_parameters["pardict"]["controls"], expected)
+    aae(parsed_parameters["parsed_params"].controls, expected)
 
 
-def test_loadings(parsed_parameters):
+def test_loadings(parsed_parameters) -> None:
     expected_values = jnp.arange(118, 177)
-    calculated = parsed_parameters["pardict"]["loadings"]
+    calculated = parsed_parameters["parsed_params"].loadings
     calculated_values = calculated[calculated != 0]
     aae(expected_values, calculated_values)
 
 
-def test_meas_sds(parsed_parameters):
+def test_meas_sds(parsed_parameters) -> None:
     expected = jnp.arange(177, 236)
-    aae(parsed_parameters["pardict"]["meas_sds"], expected)
+    aae(parsed_parameters["parsed_params"].meas_sds, expected)
 
 
-def test_shock_sds(parsed_parameters):
+def test_shock_sds(parsed_parameters) -> None:
     expected = jnp.arange(236, 257).reshape(7, 3)
-    aae(parsed_parameters["pardict"]["shock_sds"], expected)
+    aae(parsed_parameters["parsed_params"].shock_sds, expected)
 
 
-def test_initial_states(parsed_parameters):
+def test_initial_states(parsed_parameters) -> None:
     expected = jnp.arange(257, 260).reshape(1, 3).repeat(5, axis=0).reshape(5, 1, 3)
     aae(parsed_parameters["states"], expected)
 
 
-def test_initial_upper_chols(parsed_parameters):
+def test_initial_upper_chols(parsed_parameters) -> None:
     expected = (
         jnp.array([[[261, 262, 264], [0, 263, 265], [0, 0, 266]]])
         .repeat(5, axis=0)
@@ -92,8 +105,8 @@ def test_initial_upper_chols(parsed_parameters):
     aae(parsed_parameters["upper_chols"], expected)
 
 
-def test_transition_parameters(parsed_parameters):
-    calculated = parsed_parameters["pardict"]["transition"]
+def test_transition_parameters(parsed_parameters) -> None:
+    calculated = parsed_parameters["parsed_params"].transition
 
     aae(calculated["fac1"], jnp.arange(385, 413).reshape(7, 4) - 118)
     aae(calculated["fac2"], jnp.arange(413, 441).reshape(7, 4) - 118)
@@ -102,15 +115,15 @@ def test_transition_parameters(parsed_parameters):
     assert isinstance(calculated, dict)
 
 
-def test_anchoring_scaling_factors(parsed_parameters):
-    calculated = parsed_parameters["pardict"]["anchoring_scaling_factors"]
+def test_anchoring_scaling_factors(parsed_parameters) -> None:
+    calculated = parsed_parameters["parsed_params"].anchoring_scaling_factors
     expected = np.ones((8, 3))
     expected[:, 0] = jnp.array([127 + 7 * i for i in range(8)])
     aae(calculated, expected)
 
 
-def test_anchoring_constants(parsed_parameters):
-    calculated = parsed_parameters["pardict"]["anchoring_constants"]
+def test_anchoring_constants(parsed_parameters) -> None:
+    calculated = parsed_parameters["parsed_params"].anchoring_constants
     expected = np.zeros((8, 3))
     expected[:, 0] = jnp.array([18 + i * 14 for i in range(8)])
     aae(calculated, expected)
diff --git a/tests/test_process_data.py b/tests/test_process_data.py
index 4ec797e..8910857 100644
--- a/tests/test_process_data.py
+++ b/tests/test_process_data.py
@@ -1,14 +1,15 @@
 import io
 import textwrap
-from pathlib import Path
 
 import jax.numpy as jnp
 import numpy as np
 import pandas as pd
 import pytest
 import yaml
+from frozendict import frozendict
 from numpy.testing import assert_array_equal as aae
 
+from skillmodels.config import TEST_DATA_DIR
 from skillmodels.process_data import (
     _augment_data_for_endogenous_factors,
     _generate_controls_array,
@@ -18,16 +19,14 @@
     pre_process_data,
 )
 from skillmodels.process_model import process_model
+from skillmodels.types import Labels
 
-# importing the TEST_DIR from config does not work for test run in conda build
-TEST_DIR = Path(__file__).parent.resolve()
 
-
-def test_pre_process_data():
+def test_pre_process_data() -> None:
     df = pd.DataFrame(data=np.arange(20).reshape(2, 10).T, columns=["var", "inv"])
     df["period"] = [1, 2, 3, 2, 3, 4, 2, 4, 3, 1]
     df["id"] = [1, 1, 1, 3, 3, 3, 4, 4, 5, 5]
-    df.set_index(["id", "period"], inplace=True)
+    df = df.set_index(["id", "period"])
 
     exp = pd.DataFrame()
     period = [0, 1, 2, 3] * 4
@@ -39,7 +38,7 @@ def test_pre_process_data():
     }
     data = np.column_stack([period, id_, data["var"], data["inv"]])
     exp = pd.DataFrame(data=data, columns=["__period__", "__id__", "var", "inv"])
-    exp.set_index(["__id__", "__period__"], inplace=True)
+    exp = exp.set_index(["__id__", "__period__"])
 
     res = pre_process_data(df, [0, 1, 2, 3])
     assert res[["var", "inv"]].equals(exp[["var", "inv"]])
@@ -48,36 +47,36 @@ def test_pre_process_data():
 @pytest.fixture
 def simplest_augmented():
     out = {}
-    with open(TEST_DIR / "simplest_augmented_model.yaml") as y:
-        out["model_dict"] = yaml.load(y, Loader=yaml.FullLoader)
+    with (TEST_DATA_DIR / "simplest_augmented_model.yaml").open() as y:
+        out["model_dict"] = yaml.load(y, Loader=yaml.SafeLoader)
     _df = pd.DataFrame(data=np.arange(15).reshape(3, 5).T, columns=["var", "inv", "of"])
     _df["period"] = [1, 1, 2, 1, 2]
     _df["id"] = [1, 3, 3, 5, 5]
     out["data_input"] = _df.set_index(["id", "period"])
     out["data_exp"] = pd.read_csv(
-        TEST_DIR / "simplest_augmented_data_expected.csv",
+        TEST_DATA_DIR / "simplest_augmented_data_expected.csv",
         index_col=["id", "aug_period"],
     )
     return out
 
 
-def test_augment_data_for_endogenous_factors(simplest_augmented):
+def test_augment_data_for_endogenous_factors(simplest_augmented) -> None:
     model = process_model(simplest_augmented["model_dict"])
     pre_processed_data = pre_process_data(
-        simplest_augmented["data_input"], model["labels"]["periods"]
+        simplest_augmented["data_input"], model.labels.periods
     )
     pre_processed_data["constant"] = 1
     res = _augment_data_for_endogenous_factors(
         df=pre_processed_data,
-        labels=model["labels"],
-        update_info=model["update_info"],
+        labels=model.labels,
+        update_info=model.update_info,
     )
     cols = ["var", "inv", "constant", "of"]
     pd.testing.assert_frame_equal(res[cols], simplest_augmented["data_exp"][cols])
 
 
-def test_handle_controls_with_missings():
-    controls = ["c1"]
+def test_handle_controls_with_missings() -> None:
+    controls = ("c1",)
     uinfo_ind_tups = [(0, "m1"), (0, "m2")]
     update_info = pd.DataFrame(index=pd.MultiIndex.from_tuples(uinfo_ind_tups))
     data = [[1, 1, 1], [np.nan, 1, 1], [np.nan, 1, np.nan], [np.nan, np.nan, np.nan]]
@@ -86,14 +85,14 @@ def test_handle_controls_with_missings():
     df["id"] = np.arange(4)
     df["__old_id__"] = df["id"]
     df["__old_period__"] = df["aug_period"] + 1
-    df.set_index(["id", "aug_period"], inplace=True)
+    df = df.set_index(["id", "aug_period"])
 
     with pytest.warns(UserWarning):  # noqa: PT030
         calculated = _handle_controls_with_missings(df, controls, update_info)
-    assert calculated.loc[(2, 0)].isna().all()
+    assert calculated.loc[(2, 0)].isna().all()  # ty: ignore[unresolved-attribute]
 
 
-def test_generate_measurements_array():
+def test_generate_measurements_array() -> None:
     uinfo_ind_tups = [(0, "m1"), (0, "m2"), (1, "m1"), (1, "m3")]
     update_info = pd.DataFrame(index=pd.MultiIndex.from_tuples(uinfo_ind_tups))
 
@@ -112,7 +111,7 @@ def test_generate_measurements_array():
     aae(calculated, expected)
 
 
-def test_generate_controls_array():
+def test_generate_controls_array() -> None:
     csv = """
     id,aug_period,c1,c2
     0, 0, 1, 2
@@ -122,14 +121,26 @@ def test_generate_controls_array():
     """
     data = _read_csv_string(csv, ["id", "aug_period"])
 
-    labels = {"controls": ["c1", "c2"], "aug_periods": [0, 1]}
+    labels = Labels(
+        latent_factors=(),
+        observed_factors=(),
+        controls=("c1", "c2"),
+        periods=(0, 1),
+        stagemap=(0, 0),
+        stages=(0,),
+        aug_periods=(0, 1),
+        aug_periods_to_periods=frozendict({0: 0, 1: 1}),
+        aug_stagemap=(0, 0),
+        aug_stages=(0,),
+        aug_stages_to_stages=frozendict({0: 0}),
+    )
 
     calculated = _generate_controls_array(data, labels, 2)
     expected = jnp.array([[[1, 2], [5, 8]], [[3, 4], [7, 8]]])
     aae(calculated, expected)
 
 
-def test_generate_observed_factor_array():
+def test_generate_observed_factor_array() -> None:
     csv = """
     id,aug_period,v1,v2
     0, 0, 1, 2
@@ -139,7 +150,19 @@ def test_generate_observed_factor_array():
     """
     data = _read_csv_string(csv, ["id", "aug_period"])
 
-    labels = {"observed_factors": ["v1", "v2"], "aug_periods": [0, 1]}
+    labels = Labels(
+        latent_factors=(),
+        observed_factors=("v1", "v2"),
+        controls=("constant",),
+        periods=(0, 1),
+        stagemap=(0, 0),
+        stages=(0,),
+        aug_periods=(0, 1),
+        aug_periods_to_periods=frozendict({0: 0, 1: 1}),
+        aug_stagemap=(0, 0),
+        aug_stages=(0,),
+        aug_stages_to_stages=frozendict({0: 0}),
+    )
 
     calculated = _generate_observed_factor_array(data, labels, 2)
     expected = jnp.array([[[1, 2], [5, 8]], [[3, 4], [7, 8]]])
diff --git a/tests/test_process_model.py b/tests/test_process_model.py
index ab80881..87aa0a5 100644
--- a/tests/test_process_model.py
+++ b/tests/test_process_model.py
@@ -1,92 +1,85 @@
 import inspect
-from pathlib import Path
 
 import pandas as pd
 import pytest
 import yaml
 from pandas.testing import assert_frame_equal
 
+from skillmodels.config import TEST_DATA_DIR
 from skillmodels.process_model import get_has_endogenous_factors, process_model
+from skillmodels.types import TransitionInfo
 
 # ======================================================================================
 # Integration test with model2 from the replication files of CHS2010
 # ======================================================================================
 
-# importing the TEST_DIR from config does not work for test run in conda build
-TEST_DIR = Path(__file__).parent.resolve()
-
 
 @pytest.fixture
 def model2():
-    with open(TEST_DIR / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
-    return model_dict
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        return yaml.load(y, Loader=yaml.SafeLoader)
 
 
-def test_has_endogenous_factors(model2):
-    assert (
-        process_model(model2)["endogenous_factors_info"]["has_endogenous_factors"]
-        == False
-    )
+def test_has_endogenous_factors(model2) -> None:
+    assert process_model(model2).endogenous_factors_info.has_endogenous_factors == False
 
 
-def test_dimensions(model2):
-    res = process_model(model2)["dimensions"]
-    assert res["n_latent_factors"] == 3
-    assert res["n_observed_factors"] == 0
-    assert res["n_all_factors"] == 3
-    assert res["n_periods"] == 8
-    assert res["n_controls"] == 2
-    assert res["n_mixtures"] == 1
+def test_dimensions(model2) -> None:
+    res = process_model(model2).dimensions
+    assert res.n_latent_factors == 3
+    assert res.n_observed_factors == 0
+    assert res.n_all_factors == 3
+    assert res.n_periods == 8
+    assert res.n_controls == 2
+    assert res.n_mixtures == 1
 
 
-def test_labels(model2):
-    res = process_model(model2)["labels"]
-    assert res["latent_factors"] == ["fac1", "fac2", "fac3"]
-    assert res["observed_factors"] == []
-    assert res["all_factors"] == ["fac1", "fac2", "fac3"]
-    assert res["controls"] == ["constant", "x1"]
-    assert res["periods"] == [0, 1, 2, 3, 4, 5, 6, 7]
-    assert res["stagemap"] == [0, 0, 0, 0, 0, 0, 0]
-    assert res["stages"] == [0]
+def test_labels(model2) -> None:
+    res = process_model(model2).labels
+    assert res.latent_factors == ("fac1", "fac2", "fac3")
+    assert res.observed_factors == ()
+    assert res.all_factors == ("fac1", "fac2", "fac3")
+    assert res.controls == ("constant", "x1")
+    assert res.periods == (0, 1, 2, 3, 4, 5, 6, 7)
+    assert res.stagemap == (0, 0, 0, 0, 0, 0, 0)
+    assert res.stages == (0,)
 
 
-def test_estimation_options(model2):
-    res = process_model(model2)["estimation_options"]
-    assert res["sigma_points_scale"] == 2
-    assert res["robust_bounds"]
-    assert res["bounds_distance"] == 0.001
+def test_estimation_options(model2) -> None:
+    res = process_model(model2).estimation_options
+    assert res.sigma_points_scale == 2
+    assert res.robust_bounds
+    assert res.bounds_distance == 0.001
 
 
-def test_anchoring(model2):
-    res = process_model(model2)["anchoring"]
-    assert res["outcomes"] == {"fac1": "Q1"}
-    assert res["factors"] == ["fac1"]
-    assert res["free_controls"]
-    assert res["free_constant"]
-    assert res["free_loadings"]
+def test_anchoring(model2) -> None:
+    res = process_model(model2).anchoring
+    assert res.outcomes == {"fac1": "Q1"}
+    assert res.factors == ("fac1",)
+    assert res.free_controls
+    assert res.free_constant
+    assert res.free_loadings
 
 
-def test_transition_info(model2):
-    res = process_model(model2)["transition_info"]
+def test_transition_info(model2) -> None:
+    res = process_model(model2).transition_info
 
-    assert isinstance(res, dict)
-    assert callable(res["func"])
+    assert isinstance(res, TransitionInfo)
+    assert callable(res.func)
 
-    assert list(inspect.signature(res["func"]).parameters) == ["params", "states"]
+    assert list(inspect.signature(res.func).parameters) == ["params", "states"]
 
 
-def test_update_info(model2):
-    res = process_model(model2)["update_info"]
-    test_dir = Path(__file__).parent.resolve()
+def test_update_info(model2) -> None:
+    res = process_model(model2).update_info
     expected = pd.read_csv(
-        test_dir / "model2_correct_update_info.csv",
+        TEST_DATA_DIR / "model2_correct_update_info.csv",
         index_col=["aug_period", "variable"],
     )
     assert_frame_equal(res, expected)
 
 
-def test_normalizations(model2):
+def test_normalizations(model2) -> None:
     expected = {
         "fac1": {
             "loadings": [
@@ -119,7 +112,7 @@ def test_normalizations(model2):
             "intercepts": [{}, {}, {}, {}, {}, {}, {}, {}],
         },
     }
-    res = process_model(model2)["normalizations"]
+    res = process_model(model2).normalizations
 
     assert res == expected
 
@@ -129,34 +122,32 @@ def test_normalizations(model2):
 # ======================================================================================
 
 
-def test_anchoring_and_endogenous_factors_work_together():
-    with open(TEST_DIR / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
+def test_anchoring_and_endogenous_factors_work_together() -> None:
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        model_dict = yaml.load(y, Loader=yaml.SafeLoader)
     # Set fac3 to be endogenous
     model_dict["factors"]["fac3"]["is_endogenous"] = True
     del model_dict["stagemap"]
     # Should not raise - anchoring and endogenous factors now work together
     result = process_model(model_dict)
     # Verify anchoring is enabled
-    assert result["anchoring"]["anchoring"]
-    assert result["anchoring"]["factors"] == ["fac1"]
+    assert result.anchoring.anchoring
+    assert result.anchoring.factors == ("fac1",)
     # Verify endogenous factors are enabled
-    assert result["endogenous_factors_info"]["has_endogenous_factors"]
+    assert result.endogenous_factors_info.has_endogenous_factors
     # Verify dimensions
-    assert result["dimensions"]["n_periods"] == 8
-    assert result["dimensions"]["n_aug_periods"] == 16
+    assert result.dimensions.n_periods == 8
+    assert result.dimensions.n_aug_periods == 16
     # Verify update_info has anchoring entries for all aug_periods
-    anchoring_updates = result["update_info"][
-        result["update_info"]["purpose"] == "anchoring"
-    ]
+    anchoring_updates = result.update_info[result.update_info["purpose"] == "anchoring"]
     assert (
         len(anchoring_updates) == 16
     )  # One per aug_period for the one anchored factor
 
 
-def test_stagemap_with_endogenous_factors_wrong_labels():
-    with open(TEST_DIR / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
+def test_stagemap_with_endogenous_factors_wrong_labels() -> None:
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        model_dict = yaml.load(y, Loader=yaml.SafeLoader)
     # Set fac3 to be endogenous
     model_dict["factors"]["fac3"]["is_endogenous"] = True
     model_dict["stagemap"] = [0, 0, 1, 1, 2, 2, 4]
@@ -165,23 +156,23 @@ def test_stagemap_with_endogenous_factors_wrong_labels():
         process_model(model_dict)
 
 
-def test_stagemap_with_endogenous_factors():
-    with open(TEST_DIR / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
+def test_stagemap_with_endogenous_factors() -> None:
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        model_dict = yaml.load(y, Loader=yaml.SafeLoader)
     # Set fac3 to be endogenous
     model_dict["factors"]["fac3"]["is_endogenous"] = True
     model_dict["stagemap"] = [0, 0, 1, 1, 2, 2, 3]
     del model_dict["anchoring"]
     model = process_model(model_dict)
-    assert model["labels"]["stagemap"] == model_dict["stagemap"]
-    assert model["labels"]["stages"] == [0, 1, 2, 3]
-    assert model["labels"]["aug_stagemap"] == [0, 1, 0, 1, 2, 3, 2, 3, 4, 5, 4, 5, 6, 7]
+    assert model.labels.stagemap == tuple(model_dict["stagemap"])
+    assert model.labels.stages == (0, 1, 2, 3)
+    assert model.labels.aug_stagemap == (0, 1, 0, 1, 2, 3, 2, 3, 4, 5, 4, 5, 6, 7)
 
 
 @pytest.fixture
 def model2_inv():
-    with open(TEST_DIR / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        model_dict = yaml.load(y, Loader=yaml.SafeLoader)
     # Set fac3 to be endogenous
     model_dict["factors"]["fac3"]["is_endogenous"] = True
     del model_dict["stagemap"]
@@ -189,73 +180,71 @@ def model2_inv():
     return model_dict
 
 
-def test_with_endog_has_endogenous_factors(model2_inv):
+def test_with_endog_has_endogenous_factors(model2_inv) -> None:
     assert (
-        process_model(model2_inv)["endogenous_factors_info"]["has_endogenous_factors"]
-        == True
+        process_model(model2_inv).endogenous_factors_info.has_endogenous_factors == True
     )
 
 
-def test_with_endog_dimensions(model2_inv):
-    res = process_model(model2_inv)["dimensions"]
-    assert res["n_latent_factors"] == 3
-    assert res["n_observed_factors"] == 0
-    assert res["n_all_factors"] == 3
-    assert res["n_aug_periods"] == 16
-    assert res["n_periods"] == 8
-    assert res["n_controls"] == 2
-    assert res["n_mixtures"] == 1
+def test_with_endog_dimensions(model2_inv) -> None:
+    res = process_model(model2_inv).dimensions
+    assert res.n_latent_factors == 3
+    assert res.n_observed_factors == 0
+    assert res.n_all_factors == 3
+    assert res.n_aug_periods == 16
+    assert res.n_periods == 8
+    assert res.n_controls == 2
+    assert res.n_mixtures == 1
 
 
-def test_with_endog_labels(model2_inv):
-    res = process_model(model2_inv)["labels"]
+def test_with_endog_labels(model2_inv) -> None:
+    res = process_model(model2_inv).labels
     n_aug_periods = 16
-    assert res["latent_factors"] == ["fac1", "fac2", "fac3"]
-    assert res["observed_factors"] == []
-    assert res["all_factors"] == ["fac1", "fac2", "fac3"]
-    assert res["controls"] == ["constant", "x1"]
-    assert res["aug_periods"] == list(range(n_aug_periods))
-    assert res["periods"] == [0, 1, 2, 3, 4, 5, 6, 7]
-    assert res["aug_stagemap"] == list(range(n_aug_periods - 2))
-    assert res["aug_stages"] == list(range(n_aug_periods - 2))
+    assert res.latent_factors == ("fac1", "fac2", "fac3")
+    assert res.observed_factors == ()
+    assert res.all_factors == ("fac1", "fac2", "fac3")
+    assert res.controls == ("constant", "x1")
+    assert res.aug_periods == tuple(range(n_aug_periods))
+    assert res.periods == (0, 1, 2, 3, 4, 5, 6, 7)
+    assert res.aug_stagemap == tuple(range(n_aug_periods - 2))
+    assert res.aug_stages == tuple(range(n_aug_periods - 2))
 
 
-def test_with_endog_estimation_options(model2_inv):
-    res = process_model(model2_inv)["estimation_options"]
-    assert res["sigma_points_scale"] == 2
-    assert res["robust_bounds"]
-    assert res["bounds_distance"] == 0.001
+def test_with_endog_estimation_options(model2_inv) -> None:
+    res = process_model(model2_inv).estimation_options
+    assert res.sigma_points_scale == 2
+    assert res.robust_bounds
+    assert res.bounds_distance == 0.001
 
 
-def test_with_endog_anchoring_is_empty(model2_inv):
-    res = process_model(model2_inv)["anchoring"]
-    assert res["outcomes"] == {}
-    assert res["factors"] == []
-    assert res["free_controls"] is False
-    assert res["free_constant"] is False
-    assert res["free_loadings"] is False
+def test_with_endog_anchoring_is_empty(model2_inv) -> None:
+    res = process_model(model2_inv).anchoring
+    assert res.outcomes == {}
+    assert res.factors == ()
+    assert res.free_controls is False
+    assert res.free_constant is False
+    assert res.free_loadings is False
 
 
-def test_with_endog_transition_info(model2_inv):
-    res = process_model(model2_inv)["transition_info"]
+def test_with_endog_transition_info(model2_inv) -> None:
+    res = process_model(model2_inv).transition_info
 
-    assert isinstance(res, dict)
-    assert callable(res["func"])
+    assert isinstance(res, TransitionInfo)
+    assert callable(res.func)
 
-    assert list(inspect.signature(res["func"]).parameters) == ["params", "states"]
+    assert list(inspect.signature(res.func).parameters) == ["params", "states"]
 
 
-def test_with_endog_update_info(model2_inv):
-    res = process_model(model2_inv)["update_info"]
-    test_dir = Path(__file__).parent.resolve()
+def test_with_endog_update_info(model2_inv) -> None:
+    res = process_model(model2_inv).update_info
     expected = pd.read_csv(
-        test_dir / "model2_with_endog_correct_update_info.csv",
+        TEST_DATA_DIR / "model2_with_endog_correct_update_info.csv",
         index_col=["aug_period", "variable"],
     )
     assert_frame_equal(res, expected)
 
 
-def test_with_endog_normalizations(model2_inv):
+def test_with_endog_normalizations(model2_inv) -> None:
     expected = {
         "fac1": {
             "loadings": [
@@ -372,7 +361,7 @@ def test_with_endog_normalizations(model2_inv):
             ],
         },
     }
-    res = process_model(model2_inv)["normalizations"]
+    res = process_model(model2_inv).normalizations
 
     assert res == expected
 
@@ -382,24 +371,24 @@ def test_with_endog_normalizations(model2_inv):
 # ======================================================================================
 
 
-def test_model_has_endogenous_factors_not_specified():
+def test_model_has_endogenous_factors_not_specified() -> None:
     factors = {"a": {}}
     assert get_has_endogenous_factors(factors) == False
 
 
-def test_get_has_endogenous_factors_wrong_type():
+def test_get_has_endogenous_factors_wrong_type() -> None:
     factors = {"a": {"is_endogenous": 3}}
     with pytest.raises(ValueError):
         get_has_endogenous_factors(factors)
 
 
-def test_get_has_endogenous_factors_wrong_constellation():
+def test_get_has_endogenous_factors_wrong_constellation() -> None:
     factors = {"a": {"is_endogenous": False, "is_correction": True}}
     with pytest.raises(ValueError):
         get_has_endogenous_factors(factors)
 
 
-def test_get_has_endogenous_factors_indeed():
+def test_get_has_endogenous_factors_indeed() -> None:
     factors = {
         "a": {"is_endogenous": True, "is_correction": False},
         "b": {"is_endogenous": False, "is_correction": False},
@@ -407,7 +396,7 @@ def test_get_has_endogenous_factors_indeed():
     assert get_has_endogenous_factors(factors) == True
 
 
-def test_get_has_endogenous_factors_and_correction():
+def test_get_has_endogenous_factors_and_correction() -> None:
     factors = {
         "a": {"is_endogenous": True, "is_correction": False},
         "b": {"is_endogenous": False, "is_correction": False},
diff --git a/tests/test_qr.py b/tests/test_qr.py
index 5a35475..ff231ac 100644
--- a/tests/test_qr.py
+++ b/tests/test_qr.py
@@ -1,3 +1,7 @@
+"""Tests for custom QR decomposition."""
+
+from typing import TYPE_CHECKING
+
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -6,30 +10,37 @@
 
 from skillmodels.qr import qr_gpu
 
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
 SEED = 20
 
 
 @pytest.fixture
-def cov_matrix():
+def cov_matrix() -> NDArray[np.floating]:
+    """Create a covariance matrix for testing."""
     fixedrng = np.random.default_rng(SEED)
     factorized = fixedrng.uniform(low=-1, high=3, size=(7, 7))
-    cov = factorized @ factorized.T * 0.5 + np.eye(7)
-    return cov
+    return factorized @ factorized.T * 0.5 + np.eye(7)
 
 
-def test_q(cov_matrix):
+def test_q(cov_matrix: NDArray[np.floating]) -> None:
+    """Test Q matrix from QR decomposition matches JAX implementation."""
     q_gpu, _ = qr_gpu(cov_matrix)
     q_jax, _ = jnp.linalg.qr(cov_matrix)
     aaae(q_gpu, q_jax)
 
 
-def test_r(cov_matrix):
+def test_r(cov_matrix: NDArray[np.floating]) -> None:
+    """Test R matrix from QR decomposition matches JAX implementation."""
     _, r_gpu = qr_gpu(cov_matrix)
     _, r_jax = jnp.linalg.qr(cov_matrix)
     aaae(r_gpu, r_jax)
 
 
-def test_grad_qr(cov_matrix):
+def test_grad_qr(cov_matrix: NDArray[np.floating]) -> None:
+    """Test gradient of QR decomposition matches JAX implementation."""
+
     def f_jax(a):
         q, r = jnp.linalg.qr(a)
         return jnp.sum(r) + jnp.sum(q)
diff --git a/tests/test_simulate_data.py b/tests/test_simulate_data.py
index f04b658..be7c6b4 100644
--- a/tests/test_simulate_data.py
+++ b/tests/test_simulate_data.py
@@ -8,29 +8,32 @@
 import yaml
 from numpy.testing import assert_array_almost_equal as aaae
 
-from skillmodels.simulate_data import measurements_from_states, simulate_dataset
+from skillmodels.config import TEST_DATA_DIR
+from skillmodels.process_model import process_model
+from skillmodels.simulate_data import (
+    _collapse_aug_periods_to_periods,
+    measurements_from_states,
+    simulate_dataset,
+)
 
-# importing the TEST_DIR from config does not work for test run in conda build
-TEST_DIR = Path(__file__).parent.resolve()
+REGRESSION_VAULT = Path(__file__).parent / "regression_vault"
 
 
 @pytest.fixture
 def model2():
-    with open(TEST_DIR / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
-    return model_dict
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        return yaml.load(y, Loader=yaml.SafeLoader)
 
 
 @pytest.fixture
 def model2_data():
-    data = pd.read_stata(TEST_DIR / "model2_simulated_data.dta")
-    data = data.set_index(["caseid", "period"])
-    return data
+    data = pd.read_stata(TEST_DATA_DIR / "model2_simulated_data.dta")
+    return data.set_index(["caseid", "period"])
 
 
-def test_simulate_dataset(model2, model2_data):
+def test_simulate_dataset(model2, model2_data) -> None:
     model_dict = model2
-    params = pd.read_csv(TEST_DIR / "regression_vault" / "one_stage_anchoring.csv")
+    params = pd.read_csv(REGRESSION_VAULT / "one_stage_anchoring.csv")
     params = params.set_index(["category", "period", "name1", "name2"])
 
     calculated = simulate_dataset(
@@ -48,7 +51,7 @@ def test_simulate_dataset(model2, model2_data):
         assert np.allclose(ratio, expected_ratio)
 
 
-def test_measurements_from_factors():
+def test_measurements_from_factors() -> None:
     inputs = {
         "states": np.array([[0, 0, 0], [1, 1, 1]]),
         "controls": np.array([[1, 1], [1, 1]]),
@@ -58,3 +61,63 @@ def test_measurements_from_factors():
     }
     expected = np.array([[1, 1, 1], [1.9, 1.9, 1.9]])
     aaae(measurements_from_states(**inputs), expected)
+
+
+@pytest.fixture
+def model2_with_endogenous():
+    """Model2 with fac3 set as endogenous factor."""
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        model_dict = yaml.load(y, Loader=yaml.SafeLoader)
+    model_dict["factors"]["fac3"]["is_endogenous"] = True
+    del model_dict["stagemap"]
+    del model_dict["anchoring"]
+    return model_dict
+
+
+def test_collapse_aug_periods_to_periods_with_endogenous_factors(
+    model2_with_endogenous,
+) -> None:
+    """Test that _collapse_aug_periods_to_periods works with endogenous factors.
+
+    This is a regression test for a bug where MeasurementType enum values were
+    compared against strings in pandas queries, causing empty results.
+    """
+    model = process_model(model2_with_endogenous)
+    factors = model.labels.latent_factors
+
+    # Create a mock aug_latent_data DataFrame with aug_period column
+    n_obs = 5
+    n_aug_periods = model.dimensions.n_aug_periods - 1  # Exclude last half-period
+    records = []
+    for aug_p in range(n_aug_periods):
+        for obs_id in range(n_obs):
+            record = {"id": obs_id, "aug_period": aug_p}
+            for fac in factors:
+                record[fac] = np.random.randn()
+            records.append(record)
+    aug_latent_data = pd.DataFrame(records)
+
+    result = _collapse_aug_periods_to_periods(
+        df=aug_latent_data,
+        factors=factors,
+        aug_periods_to_periods=model.labels.aug_periods_to_periods,
+        endogenous_factors_info=model.endogenous_factors_info,
+    )
+
+    # The result should not be empty
+    assert len(result) > 0, "Collapsed DataFrame should not be empty"
+
+    # Should have 'period' column, not 'aug_period'
+    assert "period" in result.columns
+    assert "aug_period" not in result.columns
+
+    # Should have all factor columns
+    for fac in factors:
+        assert fac in result.columns
+
+    # Should have correct number of unique periods (half of aug_periods)
+    expected_n_periods = model.dimensions.n_periods
+    assert result["period"].nunique() == expected_n_periods
+
+    # Should have all observations for each period
+    assert len(result) == n_obs * expected_n_periods
diff --git a/tests/test_transition_functions.py b/tests/test_transition_functions.py
index 1edc83c..0a84e7e 100644
--- a/tests/test_transition_functions.py
+++ b/tests/test_transition_functions.py
@@ -1,6 +1,5 @@
 import jax
 import jax.numpy as jnp
-import numpy as np
 from numpy.testing import assert_array_almost_equal as aaae
 
 from skillmodels.transition_functions import (
@@ -16,15 +15,15 @@
 jax.config.update("jax_enable_x64", True)
 
 
-def test_linear():
-    states = np.arange(3)
-    params = np.array([0.1, 0.2, 0.3, 0.4])
+def test_linear() -> None:
+    states = jnp.arange(3)
+    params = jnp.array([0.1, 0.2, 0.3, 0.4])
     expected = 1.2
     aaae(linear(states, params), expected)
 
 
-def test_translog():
-    all_states = np.array(
+def test_translog() -> None:
+    all_states = jnp.array(
         [
             [2, 0, 0],
             [0, 3, 0],
@@ -38,7 +37,7 @@ def test_translog():
         ],
     )
 
-    params = np.array(
+    params = jnp.array(
         [
             # linear terms
             0.2,
@@ -60,19 +59,19 @@ def test_translog():
     expected_translog = [0.76, 0.7, 1.32, 0.04, 0.77, 0.1, -0.07, 0.573, 76.72]
 
     for states, expected in zip(all_states, expected_translog, strict=False):
-        calculated = translog(states, params)
+        calculated = translog(jnp.asarray(states), params)
         aaae(calculated, expected)
 
 
-def test_log_ces():
-    states = np.array([3, 7.5])
+def test_log_ces() -> None:
+    states = jnp.array([3, 7.5])
     params = jnp.array([0.4, 0.6, 2])
     expected = 7.244628323025
     calculated = log_ces(states, params)
     aaae(calculated, expected)
 
 
-def test_where_all_but_one_gammas_are_zero():
+def test_where_all_but_one_gammas_are_zero() -> None:
     """This has to be tested, becaus it leads to an underflow in the log step."""
     states = jnp.ones(3)
     params = jnp.array([0, 0, 1, -0.5])
@@ -81,12 +80,12 @@ def test_where_all_but_one_gammas_are_zero():
     aaae(calculated, expected)
 
 
-def test_constant():
-    assert constant("bla", "blubb") == "bla"
+def test_constant() -> None:
+    assert constant("bla", "blubb") == "bla"  # ty: ignore[invalid-argument-type]
 
 
-def test_robust_translog():
-    all_states = np.array(
+def test_robust_translog() -> None:
+    all_states = jnp.array(
         [
             [2, 0, 0],
             [0, 3, 0],
@@ -100,7 +99,7 @@ def test_robust_translog():
         ],
     )
 
-    params = np.array(
+    params = jnp.array(
         [
             # linear terms
             0.2,
@@ -122,19 +121,19 @@ def test_robust_translog():
     expected_translog = [0.76, 0.7, 1.32, 0.04, 0.77, 0.1, -0.07, 0.573, 76.72]
 
     for states, expected in zip(all_states, expected_translog, strict=False):
-        calculated = robust_translog(states, params)
+        calculated = robust_translog(jnp.asarray(states), params)
         aaae(calculated, expected)
 
 
-def test_log_ces_general():
-    states = np.array([3, 7.5])
+def test_log_ces_general() -> None:
+    states = jnp.array([3, 7.5])
     params = jnp.array([0.4, 0.6, 2, 2, 0.5])
     expected = 7.244628323025
     calculated = log_ces_general(states, params)
     aaae(calculated, expected)
 
 
-def test_log_ces_general_where_all_but_one_gammas_are_zero():
+def test_log_ces_general_where_all_but_one_gammas_are_zero() -> None:
     """This has to be tested, becaus it leads to an underflow in the log step."""
     states = jnp.ones(3)
     params = jnp.array([0, 0, 1, -0.5, -0.5, -0.5, -2])
@@ -143,8 +142,8 @@ def test_log_ces_general_where_all_but_one_gammas_are_zero():
     aaae(calculated, expected)
 
 
-def test_param_names_log_ces_general():
-    factors = ["a", "b"]
+def test_param_names_log_ces_general() -> None:
+    factors = ("a", "b")
     expected = ["a", "b", "sigma_a", "sigma_b", "tfp"]
     calculated = params_log_ces_general(factors)
     assert calculated == expected
diff --git a/tests/test_utilities.py b/tests/test_utilities.py
index 8041fc9..35e3ca8 100644
--- a/tests/test_utilities.py
+++ b/tests/test_utilities.py
@@ -5,14 +5,13 @@
 
 """
 
-from pathlib import Path
-
 import numpy as np
 import pandas as pd
 import pytest
 import yaml
 from pandas.testing import assert_frame_equal, assert_index_equal
 
+from skillmodels.config import TEST_DATA_DIR
 from skillmodels.process_model import process_model
 from skillmodels.utilities import (
     _get_params_index_from_model_dict,
@@ -29,28 +28,24 @@
     update_parameter_values,
 )
 
-# importing the TEST_DIR from config does not work for test run in conda build
-TEST_DIR = Path(__file__).parent.resolve()
-
 
 @pytest.fixture
 def model2():
-    with open(TEST_DIR / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
-    return model_dict
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        return yaml.load(y, Loader=yaml.SafeLoader)
 
 
 @pytest.mark.parametrize("factors", ["fac2", ["fac2"]])
-def test_extract_factors_single(model2, factors):
+def test_extract_factors_single(model2, factors) -> None:
     reduced = extract_factors(factors, model2)
-    assert list(reduced["factors"]) == ["fac2"]
+    assert list(reduced["factors"]) == ["fac2"]  # ty: ignore[invalid-argument-type]
     assert list(model2["factors"]) == ["fac1", "fac2", "fac3"]
     assert "anchoring" not in reduced
     assert model2["anchoring"]["outcomes"] == {"fac1": "Q1"}
-    process_model(reduced)
+    process_model(reduced)  # ty: ignore[invalid-argument-type]
 
 
-def test_update_parameter_values():
+def test_update_parameter_values() -> None:
     params = pd.DataFrame()
     params["value"] = np.arange(5, dtype=np.int64)
 
@@ -67,50 +62,50 @@ def test_update_parameter_values():
 
 
 @pytest.mark.parametrize("factors", ["fac2", ["fac2"]])
-def test_remove_factors(model2, factors):
+def test_remove_factors(model2, factors) -> None:
     reduced = remove_factors(factors, model2)
-    assert list(reduced["factors"]) == ["fac1", "fac3"]
+    assert list(reduced["factors"]) == ["fac1", "fac3"]  # ty: ignore[invalid-argument-type]
     assert list(model2["factors"]) == ["fac1", "fac2", "fac3"]
     assert "anchoring" in reduced
-    process_model(reduced)
+    process_model(reduced)  # ty: ignore[invalid-argument-type]
 
 
 @pytest.mark.parametrize("measurements", ["y5", ["y5"]])
-def test_remove_measurements(model2, measurements):
+def test_remove_measurements(model2, measurements) -> None:
     reduced = remove_measurements(measurements, model2)
-    assert reduced["factors"]["fac2"]["measurements"] == [["y4", "y6"]] * 8
+    assert reduced["factors"]["fac2"]["measurements"] == [["y4", "y6"]] * 8  # ty: ignore[invalid-argument-type]
     assert "y5" in model2["factors"]["fac2"]["measurements"][0]
-    process_model(reduced)
+    process_model(reduced)  # ty: ignore[invalid-argument-type]
 
 
 @pytest.mark.parametrize("controls", ["x1", ["x1"]])
-def test_remove_controls(model2, controls):
+def test_remove_controls(model2, controls) -> None:
     reduced = remove_controls(controls, model2)
     assert "controls" not in reduced
     assert "controls" in model2
-    process_model(reduced)
+    process_model(reduced)  # ty: ignore[invalid-argument-type]
 
 
-def test_reduce_n_periods(model2):
+def test_reduce_n_periods(model2) -> None:
     reduced = reduce_n_periods(model2, 1)
-    assert reduced["factors"]["fac1"]["measurements"] == [["y1", "y2", "y3"]]
-    assert reduced["factors"]["fac2"]["normalizations"]["loadings"] == [{"y4": 1}]
-    process_model(reduced)
+    assert reduced["factors"]["fac1"]["measurements"] == [["y1", "y2", "y3"]]  # ty: ignore[invalid-argument-type]
+    assert reduced["factors"]["fac2"]["normalizations"]["loadings"] == [{"y4": 1}]  # ty: ignore[invalid-argument-type]
+    process_model(reduced)  # ty: ignore[invalid-argument-type]
 
 
-def test_switch_linear_to_translog(model2):
+def test_switch_linear_to_translog(model2) -> None:
     switched = switch_linear_to_translog(model2)
-    assert switched["factors"]["fac2"]["transition_function"] == "translog"
+    assert switched["factors"]["fac2"]["transition_function"] == "translog"  # ty: ignore[invalid-argument-type]
 
 
-def test_switch_linear_and_translog_back_and_forth(model2):
+def test_switch_linear_and_translog_back_and_forth(model2) -> None:
     with_translog = switch_linear_to_translog(model2)
-    with_linear = switch_translog_to_linear(with_translog)
+    with_linear = switch_translog_to_linear(with_translog)  # ty: ignore[invalid-argument-type]
     assert model2 == with_linear
 
 
 @pytest.mark.parametrize("to_remove", ["a", ["a"]])
-def test_remove_from_list(to_remove):
+def test_remove_from_list(to_remove) -> None:
     list_ = ["a", "b", "c"]
     calculated = _remove_from_list(list_, to_remove)
     assert calculated == ["b", "c"]
@@ -118,20 +113,20 @@ def test_remove_from_list(to_remove):
 
 
 @pytest.mark.parametrize("to_remove", ["a", ["a"]])
-def test_remove_from_dict(to_remove):
+def test_remove_from_dict(to_remove) -> None:
     dict_ = {"a": 1, "b": 2, "c": 3}
     calculated = _remove_from_dict(dict_, to_remove)
     assert calculated == {"b": 2, "c": 3}
     assert dict_ == {"a": 1, "b": 2, "c": 3}
 
 
-def test_reduce_params_via_extract_factors(model2):
+def test_reduce_params_via_extract_factors(model2) -> None:
     model_dict = reduce_n_periods(model2, 2)
 
-    full_index = _get_params_index_from_model_dict(model_dict)
+    full_index = _get_params_index_from_model_dict(model_dict)  # ty: ignore[invalid-argument-type]
     params = pd.DataFrame(columns=["value"], index=full_index)
 
-    _, reduced_params = extract_factors("fac3", model_dict, params)
+    _, reduced_params = extract_factors("fac3", model_dict, params)  # ty: ignore[invalid-argument-type]
 
     expected_index = pd.MultiIndex.from_tuples(
         [
@@ -155,17 +150,17 @@ def test_reduce_params_via_extract_factors(model2):
         names=["category", "aug_period", "name1", "name2"],
     )
 
-    assert_index_equal(reduced_params.index, expected_index)
+    assert_index_equal(reduced_params.index, expected_index)  # ty: ignore[invalid-argument-type]
 
 
-def test_extend_params_via_switch_to_translog(model2):
+def test_extend_params_via_switch_to_translog(model2) -> None:
     model_dict = reduce_n_periods(model2, 2)
-    normal_index = _get_params_index_from_model_dict(model_dict)
+    normal_index = _get_params_index_from_model_dict(model_dict)  # ty: ignore[invalid-argument-type]
     params = pd.DataFrame(columns=["value"], index=normal_index)
 
-    _, extended_params = switch_linear_to_translog(model_dict, params)
+    _, extended_params = switch_linear_to_translog(model_dict, params)  # ty: ignore[invalid-argument-type]
 
-    added_index = extended_params.index.difference(normal_index)
+    added_index = extended_params.index.difference(normal_index)  # ty: ignore[possibly-missing-attribute]
 
     expected_added_index = pd.MultiIndex.from_tuples(
         [
@@ -181,10 +176,10 @@ def test_extend_params_via_switch_to_translog(model2):
 
     assert_index_equal(added_index, expected_added_index)
 
-    assert extended_params.loc[added_index, "value"].unique()[0] == 0.05
+    assert extended_params.loc[added_index, "value"].unique()[0] == 0.05  # ty: ignore[possibly-missing-attribute]
 
 
-def test_shorten_if_necessary():
+def test_shorten_if_necessary() -> None:
     list_ = list(range(3))
     not_necessary = _shorten_if_necessary(list_, 5)
     assert not_necessary == list_
diff --git a/tests/test_visualize_factor_distributions.py b/tests/test_visualize_factor_distributions.py
index 1895301..6bb08ef 100644
--- a/tests/test_visualize_factor_distributions.py
+++ b/tests/test_visualize_factor_distributions.py
@@ -3,6 +3,7 @@
 import pandas as pd
 import yaml
 
+from skillmodels.config import TEST_DATA_DIR
 from skillmodels.maximization_inputs import get_maximization_inputs
 from skillmodels.simulate_data import simulate_dataset
 from skillmodels.visualize_factor_distributions import (
@@ -12,19 +13,18 @@
     univariate_densities,
 )
 
-# importing the TEST_DIR from config does not work for test run in conda build
-TEST_DIR = Path(__file__).parent.resolve()
+REGRESSION_VAULT = Path(__file__).parent / "regression_vault"
 
 
-def test_visualize_factor_distributions_runs_with_filtered_states():
-    with open(TEST_DIR / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
+def test_visualize_factor_distributions_runs_with_filtered_states() -> None:
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        model_dict = yaml.load(y, Loader=yaml.SafeLoader)
 
-    params = pd.read_csv(TEST_DIR / "regression_vault" / "one_stage_anchoring.csv")
+    params = pd.read_csv(REGRESSION_VAULT / "one_stage_anchoring.csv")
     params = params.set_index(["category", "period", "name1", "name2"])
 
-    data = pd.read_stata(TEST_DIR / "model2_simulated_data.dta")
-    data.set_index(["caseid", "period"], inplace=True)
+    data = pd.read_stata(TEST_DATA_DIR / "model2_simulated_data.dta")
+    data = data.set_index(["caseid", "period"])
 
     max_inputs = get_maximization_inputs(model_dict, data)
     params = params.loc[max_inputs["params_template"].index]
@@ -53,14 +53,14 @@ def test_visualize_factor_distributions_runs_with_filtered_states():
     )
 
 
-def test_visualize_factor_distributions_runs_with_simulated_states():
-    with open(TEST_DIR / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
+def test_visualize_factor_distributions_runs_with_simulated_states() -> None:
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        model_dict = yaml.load(y, Loader=yaml.SafeLoader)
 
-    data = pd.read_stata(TEST_DIR / "model2_simulated_data.dta")
-    data.set_index(["caseid", "period"], inplace=True)
+    data = pd.read_stata(TEST_DATA_DIR / "model2_simulated_data.dta")
+    data = data.set_index(["caseid", "period"])
 
-    params = pd.read_csv(TEST_DIR / "regression_vault" / "one_stage_anchoring.csv")
+    params = pd.read_csv(REGRESSION_VAULT / "one_stage_anchoring.csv")
     params = params.set_index(["category", "period", "name1", "name2"])
 
     max_inputs = get_maximization_inputs(model_dict, data)
diff --git a/tests/test_visualize_transition_equations.py b/tests/test_visualize_transition_equations.py
index cb81163..65f3df3 100644
--- a/tests/test_visualize_transition_equations.py
+++ b/tests/test_visualize_transition_equations.py
@@ -3,26 +3,27 @@
 import pandas as pd
 import yaml
 
+from skillmodels.config import TEST_DATA_DIR
 from skillmodels.maximization_inputs import get_maximization_inputs
 from skillmodels.visualize_transition_equations import (
     combine_transition_plots,
     get_transition_plots,
 )
 
-TEST_DIR = Path(__file__).parent.resolve()
+REGRESSION_VAULT = Path(__file__).parent / "regression_vault"
 
 
-def test_visualize_transition_equations_runs():
-    with open(TEST_DIR / "model2.yaml") as y:
-        model_dict = yaml.load(y, Loader=yaml.FullLoader)
+def test_visualize_transition_equations_runs() -> None:
+    with (TEST_DATA_DIR / "model2.yaml").open() as y:
+        model_dict = yaml.load(y, Loader=yaml.SafeLoader)
 
     model_dict["observed_factors"] = ["ob1"]
 
-    params = pd.read_csv(TEST_DIR / "regression_vault" / "one_stage_anchoring.csv")
+    params = pd.read_csv(REGRESSION_VAULT / "one_stage_anchoring.csv")
     params = params.set_index(["category", "period", "name1", "name2"])
 
-    data = pd.read_stata(TEST_DIR / "model2_simulated_data.dta")
-    data.set_index(["caseid", "period"], inplace=True)
+    data = pd.read_stata(TEST_DATA_DIR / "model2_simulated_data.dta")
+    data = data.set_index(["caseid", "period"])
     data["ob1"] = 0
 
     max_inputs = get_maximization_inputs(model_dict, data)