Jax export

.github/workflows/test_benchmark_collection_models.yml

@@ -50,7 +50,7 @@ jobs:
       run: |
         pip3 install --user petab[vis] && \
         AMICI_PARALLEL_COMPILE="" pip3 install -v --user \
-            $(ls -t python/sdist/dist/amici-*.tar.gz | head -1)[petab,test,vis]
+            $(ls -t python/sdist/dist/amici-*.tar.gz | head -1)[petab,test,vis,jax]
 
     - name: Install test dependencies
       run: |
@@ -60,14 +60,27 @@ jobs:
 
     - name: Download benchmark collection
       run: |
-        git clone --depth 1 https://github.com/benchmarking-initiative/Benchmark-Models-PEtab.git \
-          && python3 -m pip install -e Benchmark-Models-PEtab/src/python
+        pip install git+https://github.com/Benchmarking-Initiative/Benchmark-Models-PEtab.git@master#subdirectory=src/python
 
     - name: Run tests
       env:
         AMICI_PARALLEL_COMPILE: ""
       run: |
-          cd tests/benchmark-models && pytest --durations=10
+          cd tests/benchmark-models && pytest \
+              --durations=10
+              --cov=amici \
+              --cov-report=xml:"coverage_py.xml" \
+              --cov-append \
+
+    - name: Codecov Python
+      if: github.event_name == 'pull_request' || github.repository_owner == 'AMICI-dev'
+      uses: codecov/codecov-action@v4
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        file: coverage_py.xml
+        flags: python
+        fail_ci_if_error: true
+        verbose: true
 
     # collect & upload results
     - name: Aggregate results

.github/workflows/test_python_cplusplus.yml

@@ -231,11 +231,6 @@ jobs:
     - name: Install python package
       run: scripts/installAmiciSource.sh
 
-    - name: Install notebook dependencies
-      run: |
-        source venv/bin/activate \
-        && pip install jax[cpu]
-
     - name: example notebooks
       run: scripts/runNotebook.sh python/examples/example_*/
 

.gitignore

@@ -137,6 +137,8 @@ tests/test/*
 */tests/explicit_amici/*
 */tests/fixed_initial_amici/*
 */tests/localfunc_amici/*
+*/tests/conversion/*
+*/tests/dimerization/*
 tests/cpp/writeResults.h5
 tests/cpp/writeResults.h5.bak
 tests/sbml-test-suite/*

documentation/ExampleJaxPEtab.ipynb

@@ -0,0 +1 @@
+../python/examples/example_jax_petab/ExampleJaxPEtab.ipynb

documentation/conf.py

@@ -206,6 +206,7 @@ def install_doxygen():
     "numpy": ("https://numpy.org/devdocs/", None),
     "sympy": ("https://docs.sympy.org/latest/", None),
     "python": ("https://docs.python.org/3", None),
+    "jax": ["https://jax.readthedocs.io/en/latest/", None],
 }
 
 # Add notebooks prolog with binder links

documentation/python_examples.rst

@@ -17,5 +17,6 @@ Various example notebooks.
    example_errors.ipynb
    example_large_models/example_performance_optimization.ipynb
    ExampleJax.ipynb
+   ExampleJaxPEtab.ipynb
    ExampleSplines.ipynb
    ExampleSplinesSwameye2003.ipynb

documentation/python_modules.rst

@@ -25,6 +25,7 @@ AMICI Python API
    amici.petab_objective
    amici.petab_simulate
    amici.import_utils
+   amici.jax
    amici.de_export
    amici.de_model
    amici.de_model_components

documentation/rtd_requirements.txt

@@ -3,6 +3,8 @@ sphinx<8
 mock>=5.0.2
 setuptools>=67.7.2
 pysb>=1.11.0
+jax>=0.4.26
+diffrax>=0.5.0
 matplotlib==3.7.1
 nbsphinx==0.9.1
 nbformat==5.8.0

python/sdist/amici/__init__.py

@@ -121,6 +121,11 @@ def _imported_from_setup() -> bool:
         assignmentRules2observables,
     )
 
+    try:
+        from .jax import JAXModel
+    except (ImportError, ModuleNotFoundError):
+        JAXModel = object
+
     @runtime_checkable
     class ModelModule(Protocol):  # noqa: F811
         """Type of AMICI-generated model modules.
@@ -135,6 +140,11 @@ def get_model(self) -> amici.Model:
             """Create a model instance."""
             ...
 
+        def get_jax_model(self) -> JAXModel:
+            ...
+
+    AmiciModel = Union[amici.Model, amici.ModelPtr]
+
 
 class add_path:
     """Context manager for temporarily changing PYTHONPATH"""

python/sdist/amici/__init__.template.py

@@ -1,9 +1,12 @@
 """AMICI-generated module for model TPL_MODELNAME"""
 
 from pathlib import Path
-
+from typing import TYPE_CHECKING
 import amici
 
+if TYPE_CHECKING:
+    from amici.jax import JAXModel
+
 # Ensure we are binary-compatible, see #556
 if "TPL_AMICI_VERSION" != amici.__version__:
     raise amici.AmiciVersionError(
@@ -18,4 +21,11 @@
 from .TPL_MODELNAME import *  # noqa: F403, F401
 from .TPL_MODELNAME import getModel as get_model  # noqa: F401
 
+
+def get_jax_model() -> "JAXModel":
+    from .jax import JAXModel_TPL_MODELNAME
+
+    return JAXModel_TPL_MODELNAME()
+
+
 __version__ = "TPL_PACKAGE_VERSION"