Support for PEtab SciML in jax

.github/workflows/test_petab_sciml.yml

@@ -0,0 +1,92 @@
+name: PEtab
+on:
+  push:
+    branches:
+      - develop
+      - main
+  pull_request:
+    branches:
+      - main
+      - develop
+      - jax_sciml
+  merge_group:
+  workflow_dispatch:
+
+jobs:
+  build:
+    name: PEtab SciML Testsuite
+
+    runs-on: ubuntu-latest
+
+    env:
+      ENABLE_GCOV_COVERAGE: TRUE
+
+    strategy:
+      matrix:
+        python-version: ["3.12"]
+
+    steps:
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 20
+
+      # todo, update after https://github.com/sebapersson/petab_sciml_testsuite/issues/14 is merged
+      - name: Download PEtab SciML test suite
+        run: |
+          git clone --depth 1 --branch main \
+            https://github.com/FFroehlich/petab_sciml_testsuite \
+            tests/sciml/testsuite
+
+      - name: Install apt dependencies
+        uses: ./.github/actions/install-apt-dependencies
+
+      # install dependencies
+      - name: apt
+        run: |
+          sudo apt-get update \
+            && sudo apt-get install -y python3-venv
+
+      - run: |
+          echo "${HOME}/.local/bin/" >> $GITHUB_PATH
+
+      # install AMICI
+      - name: Install python package
+        run: scripts/installAmiciSource.sh
+
+      - name: Install petab
+        run: |
+          source ./venv/bin/activate \
+            && pip3 install wheel pytest shyaml pytest-cov
+
+      # retrieve test models
+      - name: Download and install PEtab SciML
+        run: |
+          source ./venv/bin/activate \
+            && python -m pip install git+https://github.com/petab-dev/petab_sciml.git@main#subdirectory=src/python \
+
+
+      - name: Install petab
+        run: |
+          source ./venv/bin/activate \
+            && python3 -m pip uninstall -y petab \
+            && python3 -m pip install git+https://github.com/petab-dev/libpetab-python.git@sciml \
+
+      - name: Run PEtab SciML testsuite
+        run: |
+          source ./venv/bin/activate \
+            && pytest --cov-report=xml:coverage_petab_sciml.xml \
+              --cov=amici tests/sciml/test_sciml.py
+
+      - name: Codecov
+        if: github.event_name == 'pull_request' || github.repository_owner == 'AMICI-dev'
+        uses: codecov/codecov-action@v5
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          file: coverage_petab_sciml.xml
+          flags: petab_sciml
+          fail_ci_if_error: true

.gitignore

@@ -38,6 +38,9 @@ models/model_calvetti/build/*
 
 amici_models/
 
+# PEtab SciML test suite (downloaded dynamically)
+tests/sciml/testsuite/
+
 simulate_model_*_hdf.m
 simulate_model_*.m
 
@@ -196,3 +199,4 @@ debug/*
 tests/benchmark_models/cache_fiddy/*
 venv/*
 .coverage
+tests/sciml/models/*

CHANGELOG.md

@@ -74,6 +74,8 @@ See also our [versioning policy](https://amici.readthedocs.io/en/latest/versioni
     This only works on shared file systems, as the solver state is stored in a
     temporary HDF5 file.
   * `amici.ExpData` is now picklable.
+* Implemented support for the [PEtab SciML](https://github.com/PEtab-dev/petab_sciml)
+  extension for the JAX interface.
 * The import function `sbml2amici`, `pysb2amici`, and `antimony2amici` now
   return an instance of the generated model class if called with `compile=True`
   (default).

doc/conf.py

@@ -31,6 +31,7 @@
     import exhale_multiproject_monkeypatch  # noqa: F401
 
 # need to import before setting typing.TYPE_CHECKING=True, fails otherwise
+
 import amici
 import pandas as pd  # noqa: F401
 import sympy as sp  # noqa: F401
@@ -365,6 +366,11 @@ def install_doxygen():
     "ExpDataPtrVector": ":class:`amici.amici.ExpData`",
 }
 
+# TODO: alias for forward type definition, remove after release of petab_sciml
+autodoc_type_aliases = {
+    "NNModel": "petab_sciml.NNModel",
+}
+
 
 def process_docstring(app, what, name, obj, options, lines):
     # only apply in the amici.amici module

doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb

@@ -93,8 +93,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Access the results\n",
-    "results"
+    "# Define the simulation condition\n",
+    "simulation_condition = (\"model1_data1\",)\n",
+    "\n",
+    "# Access the results for the specified condition\n",
+    "ic = results[\"simulation_conditions\"].index(simulation_condition)\n",
+    "print(\"llh: \", results[\"llh\"][ic])\n",
+    "print(\"state variables: \", results[\"x\"][ic, :])"
    ]
   },
   {
@@ -356,7 +361,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "grad._my"
+    "grad._my[ic, :]"
    ]
   },
   {
@@ -393,7 +398,7 @@
     "nps = jax_problem._np_numeric[ic, :]\n",
     "\n",
     "# Load parameters for the specified condition\n",
-    "p = jax_problem.load_parameters(simulation_condition[0])\n",
+    "p = jax_problem.load_model_parameters(simulation_condition[0])\n",
     "\n",
     "\n",
     "# Define a function to compute the gradient with respect to dynamic timepoints\n",
@@ -612,16 +617,16 @@
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b8382b0b2b68f49e",
    "metadata": {},
-   "outputs": [],
+   "cell_type": "code",
    "source": [
     "# Profile gradient computation using forward sensitivity analysis\n",
     "solver.set_sensitivity_order(amici.SensitivityOrder.first)\n",
     "solver.set_sensitivity_method(amici.SensitivityMethod.forward)"
-   ]
+   ],
+   "id": "81fe95a6e7f613f1",
+   "outputs": [],
+   "execution_count": null
   },
   {
    "cell_type": "code",
@@ -687,8 +692,7 @@
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.13.0"
+   "pygments_lexer": "ipython3"
   }
  },
  "nbformat": 4,

doc/rtd_requirements.txt

@@ -7,6 +7,8 @@ setuptools>=67.7.2
 #  https://github.com/pysb/pysb/pull/599
 #  for building the documentation, we don't care whether this fully works
 git+https://github.com/pysb/pysb@0afeaab385e9a1d813ecf6fdaf0153f4b91358af
+# For forward type definition in generate_equinox
+git+https://github.com/PEtab-dev/petab_sciml.git@727d177fd3f85509d0bdcc278b672e9eeafd2384#subdirectory=src/python
 matplotlib>=3.7.1
 optax
 nbsphinx
@@ -16,6 +18,7 @@ sphinx_rtd_theme>=1.2.0
 petab[vis]>=0.2.0
 sphinx-autodoc-typehints
 ipython>=8.13.2
+h5py>=3.14.0
 breathe>=4.35.0
 exhale>=0.3.7
 -e git+https://github.com/mithro/sphinx-contrib-mithro#egg=sphinx-contrib-exhale-multiproject&subdirectory=sphinx-contrib-exhale-multiproject

python/sdist/amici/de_export.py

@@ -165,6 +165,7 @@ def __init__(
         allow_reinit_fixpar_initcond: bool | None = True,
         generate_sensitivity_code: bool | None = True,
         model_name: str | None = "model",
+        hybridization: dict | None = None,
     ):
         """
         Generate AMICI C++ files for the DE provided to the constructor.
@@ -196,6 +197,10 @@ def __init__(
 
         :param model_name:
             name of the model to be used during code generation
+
+        :param hybridization:
+            dict representation of the hybridization information in the PEtab YAML file, see
+            https://petab-sciml.readthedocs.io/latest/format.html#problem-yaml-file
         """
         set_log_level(logger, verbose)
 
@@ -237,6 +242,7 @@ def __init__(
         self.allow_reinit_fixpar_initcond: bool = allow_reinit_fixpar_initcond
         self._build_hints = set()
         self.generate_sensitivity_code: bool = generate_sensitivity_code
+        self.hybridisation = hybridization
 
     @log_execution_time("generating cpp code", logger)
     def generate_model_code(self) -> None: