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/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"

python/sdist/amici/de_export.py

@@ -21,6 +21,7 @@
     TYPE_CHECKING,
     Literal,
 )
+
 import sympy as sp
 
 from . import (
@@ -54,6 +55,7 @@
     AmiciCxxCodePrinter,
     get_switch_statement,
 )
+from .jaxcodeprinter import AmiciJaxCodePrinter
 from .de_model import DEModel
 from .de_model_components import *
 from .import_utils import (
@@ -143,7 +145,10 @@ class DEExporter:
         If the given model uses special functions, this set contains hints for
         model building.
 
-    :ivar _code_printer:
+    :ivar _code_printer_jax:
+        Code printer to generate JAX code
+
+    :ivar _code_printer_cpp:
         Code printer to generate C++ code
 
     :ivar generate_sensitivity_code:
@@ -212,14 +217,15 @@ def __init__(
         self.set_name(model_name)
         self.set_paths(outdir)
 
-        self._code_printer = AmiciCxxCodePrinter()
+        self._code_printer_cpp = AmiciCxxCodePrinter()
+        self._code_printer_jax = AmiciJaxCodePrinter()
         for fun in CUSTOM_FUNCTIONS:
-            self._code_printer.known_functions[fun["sympy"]] = fun["c++"]
+            self._code_printer_cpp.known_functions[fun["sympy"]] = fun["c++"]
 
         # Signatures and properties of generated model functions (see
         # include/amici/model.h for details)
         self.model: DEModel = de_model
-        self._code_printer.known_functions.update(
+        self._code_printer_cpp.known_functions.update(
             splines.spline_user_functions(
                 self.model._splines, self._get_index("p")
             )
@@ -242,6 +248,7 @@ def generate_model_code(self) -> None:
             sp.Pow, "_eval_derivative", _custom_pow_eval_derivative
         ):
             self._prepare_model_folder()
+            self._generate_jax_code()
             self._generate_c_code()
             self._generate_m_code()
 
@@ -269,6 +276,105 @@ def _prepare_model_folder(self) -> None:
             if os.path.isfile(file_path):
                 os.remove(file_path)
 
+    @log_execution_time("generating jax code", logger)
+    def _generate_jax_code(self) -> None:
+        from amici.jax.model import JAXModel
+
+        eq_names = (
+            "xdot",
+            "w",
+            "x0",
+            "y",
+            "sigmay",
+            "Jy",
+            "x_solver",
+            "x_rdata",
+            "total_cl",
+        )
+        sym_names = ("x", "tcl", "w", "my", "y", "sigmay", "x_rdata")
+
+        indent = 8
+
+        def jnp_array_str(array) -> str:
+            elems = ", ".join(str(s) for s in array)
+
+            return f"jnp.array([{elems}])"
+
+        tpl_data = {
+            **{
+                f"{eq_name.upper()}_EQ": "\n".join(
+                    self._code_printer_jax._get_sym_lines(
+                        (str(strip_pysb(s)) for s in self.model.sym(eq_name)),
+                        self.model.eq(eq_name).subs(
+                            dict(
+                                zip(
+                                    list(self.model.sym("h"))
+                                    + list(self.model.sym("my")),
+                                    [
+                                        sp.Heaviside(x)
+                                        for x in self.model.eq("root")
+                                    ]
+                                    + [sp.Symbol("my")]
+                                    * len(self.model.sym("my")),
+                                )
+                            )
+                        ),
+                        indent,
+                    )
+                )[indent:]
+                for eq_name in eq_names
+            },
+            **{
+                f"{eq_name.upper()}_RET": jnp_array_str(
+                    strip_pysb(s) for s in self.model.sym(eq_name)
+                )
+                if self.model.sym(eq_name)
+                else "jnp.array([])"
+                for eq_name in eq_names
+            },
+            **{
+                f"{sym_name.upper()}_SYMS": "".join(
+                    str(strip_pysb(s)) + ", " for s in self.model.sym(sym_name)
+                )
+                if self.model.sym(sym_name)
+                else "_"
+                for sym_name in sym_names
+            },
+            **{
+                f"{sym_name.upper()}_IDS": "".join(
+                    f'"{strip_pysb(s)}", ' for s in self.model.sym(sym_name)
+                )
+                if self.model.sym(sym_name)
+                else "tuple()"
+                for sym_name in ("p", "k", "y", "x")
+            },
+            **{
+                "PK_SYMS": "".join(
+                    str(strip_pysb(s)) + ", "
+                    for s in list(self.model.sym("p"))
+                    + list(self.model.sym("k"))
+                ),
+                "PK_IDS": "".join(
+                    f'"{strip_pysb(s)}", '
+                    for s in list(self.model.sym("p"))
+                    + list(self.model.sym("k"))
+                ),
+            },
+            **{
+                "MODEL_NAME": self.model_name,
+                "MODEL_API_VERSION": f"'{JAXModel.MODEL_API_VERSION}'",
+            },
+        }
+        os.makedirs(
+            os.path.join(self.model_path, self.model_name), exist_ok=True
+        )
+
+        apply_template(
+            os.path.join(amiciModulePath, "jax.template.py"),
+            os.path.join(self.model_path, self.model_name, "jax.py"),
+            tpl_data,
+        )
+
     def _generate_c_code(self) -> None:
         """
         Create C++ code files for the model based on
@@ -729,7 +835,7 @@ def _get_function_body(
                                 f"reinitialization_state_idxs.cend(), {index}) != "
                                 "reinitialization_state_idxs.cend())",
                                 f"    {function}[{index}] = "
-                                f"{self._code_printer.doprint(formula)};",
+                                f"{self._code_printer_cpp.doprint(formula)};",
                             ]
                         )
                 cases[ipar] = expressions
@@ -744,12 +850,12 @@ def _get_function_body(
                     f"reinitialization_state_idxs.cend(), {index}) != "
                     "reinitialization_state_idxs.cend())\n        "
                     f"{function}[{index}] = "
-                    f"{self._code_printer.doprint(formula)};"
+                    f"{self._code_printer_cpp.doprint(formula)};"
                 )
 
         elif function in event_functions:
             cases = {
-                ie: self._code_printer._get_sym_lines_array(
+                ie: self._code_printer_cpp._get_sym_lines_array(
                     equations[ie], function, 0
                 )
                 for ie in range(self.model.num_events())
@@ -762,7 +868,7 @@ def _get_function_body(
             for ie, inner_equations in enumerate(equations):
                 inner_lines = []
                 inner_cases = {
-                    ipar: self._code_printer._get_sym_lines_array(
+                    ipar: self._code_printer_cpp._get_sym_lines_array(
                         inner_equations[:, ipar], function, 0
                     )
                     for ipar in range(self.model.num_par())
@@ -777,7 +883,7 @@ def _get_function_body(
             and equations.shape[1] == self.model.num_par()
         ):
             cases = {
-                ipar: self._code_printer._get_sym_lines_array(
+                ipar: self._code_printer_cpp._get_sym_lines_array(
                     equations[:, ipar], function, 0
                 )
                 for ipar in range(self.model.num_par())
@@ -787,15 +893,15 @@ def _get_function_body(
         elif function in multiobs_functions:
             if function == "dJydy":
                 cases = {
-                    iobs: self._code_printer._get_sym_lines_array(
+                    iobs: self._code_printer_cpp._get_sym_lines_array(
                         equations[iobs], function, 0
                     )
                     for iobs in range(self.model.num_obs())
                     if not smart_is_zero_matrix(equations[iobs])
                 }
             else:
                 cases = {
-                    iobs: self._code_printer._get_sym_lines_array(
+                    iobs: self._code_printer_cpp._get_sym_lines_array(
                         equations[:, iobs], function, 0
                     )
                     for iobs in range(equations.shape[1])
@@ -825,7 +931,7 @@ def _get_function_body(
                     tmp_equations = sp.Matrix(
                         [equations[i] for i in static_idxs]
                     )
-                    tmp_lines = self._code_printer._get_sym_lines_symbols(
+                    tmp_lines = self._code_printer_cpp._get_sym_lines_symbols(
                         tmp_symbols,
                         tmp_equations,
                         function,
@@ -851,7 +957,7 @@ def _get_function_body(
                         [equations[i] for i in dynamic_idxs]
                     )
 
-                    tmp_lines = self._code_printer._get_sym_lines_symbols(
+                    tmp_lines = self._code_printer_cpp._get_sym_lines_symbols(
                         tmp_symbols,
                         tmp_equations,
                         function,
@@ -863,12 +969,12 @@ def _get_function_body(
                         lines.extend(tmp_lines)
 
             else:
-                lines += self._code_printer._get_sym_lines_symbols(
+                lines += self._code_printer_cpp._get_sym_lines_symbols(
                     symbols, equations, function, 4
                 )
 
         else:
-            lines += self._code_printer._get_sym_lines_array(
+            lines += self._code_printer_cpp._get_sym_lines_array(
                 equations, function, 4
             )
 
@@ -1024,10 +1130,10 @@ def _write_model_header_cpp(self) -> None:
             "NK": self.model.num_const(),
             "O2MODE": "amici::SecondOrderMode::none",
             # using code printer ensures proper handling of nan/inf
-            "PARAMETERS": self._code_printer.doprint(self.model.val("p"))[
+            "PARAMETERS": self._code_printer_cpp.doprint(self.model.val("p"))[
                 1:-1
             ],
-            "FIXED_PARAMETERS": self._code_printer.doprint(
+            "FIXED_PARAMETERS": self._code_printer_cpp.doprint(
                 self.model.val("k")
             )[1:-1],
             "PARAMETER_NAMES_INITIALIZER_LIST": self._get_symbol_name_initializer_list(
@@ -1221,7 +1327,7 @@ def _get_symbol_id_initializer_list(self, name: str) -> str:
             Template initializer list of ids
         """
         return "\n".join(
-            f'"{self._code_printer.doprint(symbol)}", // {name}[{idx}]'
+            f'"{self._code_printer_cpp.doprint(symbol)}", // {name}[{idx}]'
             for idx, symbol in enumerate(self.model.sym(name))
         )