Add JAX SBML test suite (#2795)

* Add JAX SBML test suite

* Refine JAX SBML tests

* Move SBML tests to subfolder

* Remove unused test package markers

* Refactor SBML tests directory

* Move SBML tests back to tests/sbml

* Consolidate SBML test scripts with --jax option

* smaller fixes, add error for jax models with event assignemnts

* enable x64 precision, add error for algebraic systems

* fix t in x0 and expressionIds, specific lower tolerances

* Handle default parameters for JAX models (#2800)

* Apply suggestions from code review

* Update test_sbml_semantic_test_suite_jax.yml

* Update testSBMLSuiteJax.py

* fix inf/nan

* small fixes

* reintroduce test matrix

* update test readme

* move sbml testsuite

* Update conftest.py

* fix petab benchmark?

* Update test_petab_benchmark_jax.py

* Update .github/workflows/test_sbml_semantic_test_suite_jax.yml

* enable 01395

Author: FFroehlich

.github/workflows/test_sbml_semantic_test_suite.yml

@@ -14,8 +14,8 @@ on:
       - python/sdist/amici/sbml_import.py
       - python/sdist/amici/import_utils.py
       - scripts/run-SBMLTestsuite.sh
-      - tests/testSBMLSuite.py
-      - tests/conftest.py
+      - tests/sbml/testSBMLSuite.py
+      - tests/sbml/conftest.py
   check_suite:
    types: [requested]
   workflow_dispatch:
@@ -52,7 +52,7 @@ jobs:
       uses: actions/upload-artifact@v4
       with:
         name: amici-semantic-results-${{ matrix.cases }}
-        path: tests/amici-semantic-results
+        path: tests/sbml/amici-semantic-results
 
     - name: Codecov SBMLSuite
       if: github.event_name == 'pull_request' || github.repository_owner == 'AMICI-dev'

.github/workflows/test_sbml_semantic_test_suite_jax.yml

@@ -0,0 +1,60 @@
+name: SBML JAX
+on:
+  push:
+    branches:
+      - develop
+      - master
+      - release**
+  pull_request:
+    paths:
+      - .github/workflows/test_sbml_semantic_test_suite_jax.yml
+      - python/sdist/amici/jax/**
+      - scripts/run-SBMLTestsuite.sh
+      - tests/sbml/testSBMLSuiteJax.py
+      - tests/sbml/conftest.py
+  check_suite:
+   types: [requested]
+  workflow_dispatch:
+
+jobs:
+  build:
+    name: SBML Semantic Test Suite JAX
+    runs-on: ubuntu-24.04
+
+    strategy:
+      fail-fast: false
+      matrix:
+        cases: ["1-600", "601-1200", "1200-"]
+        python-version: [ "3.13" ]
+
+    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: 1
+
+    - name: Install apt dependencies
+      uses: ./.github/actions/install-apt-dependencies
+
+    - run: AMICI_PARALLEL_COMPILE="" ./scripts/installAmiciSource.sh
+    - run: ./scripts/run-SBMLTestsuite.sh --jax ${{ matrix.cases }}
+
+    - name: "Upload artifact: SBML semantic test suite results"
+      uses: actions/upload-artifact@v4
+      with:
+        name: amici-semantic-results-jax-${{ matrix.cases }}
+        path: tests/sbml/amici-semantic-results-jax
+
+    - name: Codecov SBMLSuiteJax
+      if: github.event_name == 'pull_request' || github.repository_owner == 'AMICI-dev'
+      uses: codecov/codecov-action@v5
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        files: coverage_SBMLSuite_jax.xml
+        flags: sbmlsuite-jax
+        fail_ci_if_error: true

.gitignore

@@ -133,11 +133,10 @@ tests/test/*
 */tests/dimerization/*
 tests/cpp/writeResults.h5
 tests/cpp/writeResults.h5.bak
-tests/sbml-test-suite/*
-tests/sbml-test-suite/
-tests/sedml-test-suite/
+tests/sbml/sbml-test-suite/*
+tests/sbml/sbml-test-suite/
 */sbml-semantic-test-cases/*
-tests/SBMLTestModels/
+tests/sbml/SBMLTestModels/
 tests/benchmark-models/test_bmc
 tests/petab_test_suite
 petab_test_suite

doc/CI.md

@@ -43,10 +43,11 @@ To run the SBML Test Suite test cases, the easiest way is:
 2. Running `scripts/run-SBMLTestsuite.sh`. This will download the test cases
    if necessary and run them all. A subset of test cases can be selected with
    an optional argument (e.g. `scripts/run-SBMLTestsuite.sh 1,3-6,8`, to run
-   cases 1, 3, 4, 5, 6 and 8).
+   cases 1, 3, 4, 5, 6 and 8). Use `--jax` as the first argument to run the
+   SBML tests using the JAX backend.
 
    Once the test cases are available locally, for debugging it might be easier
-   to directly use `pytest` with `tests/testSBMLSuite.py`.
+   to directly use `pytest` with `tests/sbml/testSBMLSuite.py`.
 
 
 ## Matlab tests (not included in CI pipeline)

python/sdist/amici/jax/jax.template.py

@@ -2,6 +2,8 @@
 import jax.numpy as jnp
 from interpax import interp1d
 from pathlib import Path
+from jax.numpy import inf as oo
+from jax.numpy import nan as nan
 
 from amici.jax.model import JAXModel, safe_log, safe_div
 
@@ -11,6 +13,7 @@ class JAXModel_TPL_MODEL_NAME(JAXModel):
 
     def __init__(self):
         self.jax_py_file = Path(__file__).resolve()
+        self.parameters = TPL_P_VALUES
         super().__init__()
 
     def _xdot(self, t, x, args):
@@ -34,7 +37,7 @@ def _w(self, t, x, p, tcl):
 
         return TPL_W_RET
 
-    def _x0(self, p):
+    def _x0(self, t, p):
         TPL_P_SYMS = p
 
         TPL_X0_EQ
@@ -86,6 +89,9 @@ def _sigmay(self, y, p, np):
 
     def _nllh(self, t, x, p, tcl, my, iy, op, np):
         y = self._y(t, x, p, tcl, op)
+        if not y.size:
+            return jnp.array(0.0)
+
         TPL_Y_SYMS = y
         TPL_SIGMAY_SYMS = self._sigmay(y, p, np)
 
@@ -105,5 +111,9 @@ def state_ids(self):
     def parameter_ids(self):
         return TPL_P_IDS
 
+    @property
+    def expression_ids(self):
+        return TPL_W_IDS
+
 
 Model = JAXModel_TPL_MODEL_NAME

python/sdist/amici/jax/jaxcodeprinter.py

@@ -8,6 +8,12 @@
 from sympy.printing.numpy import NumPyPrinter
 
 
+def _jnp_array_str(array) -> str:
+    elems = ", ".join(str(s) for s in array)
+
+    return f"jnp.array([{elems}])"
+
+
 class AmiciJaxCodePrinter(NumPyPrinter):
     """JAX code printer"""
 
@@ -36,6 +42,18 @@ def _print_Mul(self, expr: sp.Expr) -> str:
             return super()._print_Mul(expr)
         return f"safe_div({self.doprint(numer)}, {self.doprint(denom)})"
 
+    def _print_Max(self, expr: sp.Expr) -> str:
+        """
+        Print the max function, replacing it with jnp.max.
+        """
+        return f"jnp.max({_jnp_array_str(expr.args)})"
+
+    def _print_Min(self, expr: sp.Expr) -> str:
+        """
+        Print the min function, replacing it with jnp.min.
+        """
+        return f"jnp.min({_jnp_array_str(expr.args)})"
+
     def _get_sym_lines(
         self,
         symbols: sp.Matrix | Iterable[str],

python/sdist/amici/jax/model.py

@@ -5,6 +5,7 @@
 from abc import abstractmethod
 from pathlib import Path
 import enum
+from dataclasses import field
 
 import diffrax
 import equinox as eqx
@@ -43,9 +44,10 @@ class JAXModel(eqx.Module):
         Path to the JAX model file.
     """
 
-    MODEL_API_VERSION = "0.0.3"
+    MODEL_API_VERSION = "0.0.4"
     api_version: str
     jax_py_file: Path
+    parameters: jnp.ndarray = field(default_factory=lambda: jnp.array([]))
 
     def __init__(self):
         if self.api_version != self.MODEL_API_VERSION:
@@ -93,11 +95,16 @@ def _w(
         ...
 
     @abstractmethod
-    def _x0(self, p: jt.Float[jt.Array, "np"]) -> jt.Float[jt.Array, "nx"]:
+    def _x0(
+        self, t: jnp.float_, p: jt.Float[jt.Array, "np"]
+    ) -> jt.Float[jt.Array, "nx"]:
         """
         Compute the initial state vector.
 
+        :param t: initial time point
         :param p: parameters
+        :return:
+            Initial state vector.
         """
         ...
 
@@ -264,6 +271,17 @@ def parameter_ids(self) -> list[str]:
         """
         ...
 
+    @property
+    @abstractmethod
+    def expression_ids(self) -> list[str]:
+        """
+        Get the expression ids of the model.
+
+        :return:
+            Expression ids
+        """
+        ...
+
     def _eq(
         self,
         p: jt.Float[jt.Array, "np"],
@@ -496,7 +514,7 @@ def _sigmays(
     @eqx.filter_jit
     def simulate_condition(
         self,
-        p: jt.Float[jt.Array, "np"],
+        p: jt.Float[jt.Array, "np"] | None,
         ts_dyn: jt.Float[jt.Array, "nt_dyn"],
         ts_posteq: jt.Float[jt.Array, "nt_posteq"],
         my: jt.Float[jt.Array, "nt"],
@@ -521,7 +539,8 @@ def simulate_condition(
         Simulate a condition.
 
         :param p:
-            parameters for simulation ordered according to ids in :ivar parameter_ids:
+            parameters for simulation ordered according to ids in :ivar parameter_ids:. If ``None``,
+            the values stored in :attr:`parameters` are used.
         :param ts_dyn:
             time points for dynamic simulation. Sorted in monotonically increasing order but duplicate time points are
             allowed to facilitate the evaluation of multiple observables at specific time points.
@@ -564,10 +583,13 @@ def simulate_condition(
         :return:
             output according to `ret` and general results/statistics
         """
+        if p is None:
+            p = self.parameters
+
         if x_preeq.shape[0]:
             x = x_preeq
         else:
-            x = self._x0(p)
+            x = self._x0(0.0, p)
 
         if not ts_mask.shape[0]:
             ts_mask = jnp.ones_like(my, dtype=jnp.bool_)
@@ -675,7 +697,7 @@ def simulate_condition(
     @eqx.filter_jit
     def preequilibrate_condition(
         self,
-        p: jt.Float[jt.Array, "np"],
+        p: jt.Float[jt.Array, "np"] | None,
         x_reinit: jt.Float[jt.Array, "*nx"],
         mask_reinit: jt.Bool[jt.Array, "*nx"],
         solver: diffrax.AbstractSolver,
@@ -689,7 +711,8 @@ def preequilibrate_condition(
         Simulate a condition.
 
         :param p:
-            parameters for simulation ordered according to ids in :ivar parameter_ids:
+            parameters for simulation ordered according to ids in :ivar parameter_ids:. If ``None``,
+            the values stored in :attr:`parameters` are used.
         :param x_reinit:
             re-initialized state vector. If not provided, the state vector is not re-initialized.
         :param mask_reinit:
@@ -704,7 +727,10 @@ def preequilibrate_condition(
             pre-equilibrated state variables and statistics
         """
         # Pre-equilibration
-        x0 = self._x0(p)
+        if p is None:
+            p = self.parameters
+
+        x0 = self._x0(0.0, p)
         if x_reinit.shape[0]:
             x0 = jnp.where(mask_reinit, x_reinit, x0)
         tcl = self._tcl(x0, p)

python/sdist/amici/jax/ode_export.py

@@ -22,7 +22,7 @@
 )
 
 from amici._codegen.template import apply_template
-from amici.jax.jaxcodeprinter import AmiciJaxCodePrinter
+from amici.jax.jaxcodeprinter import AmiciJaxCodePrinter, _jnp_array_str
 from amici.jax.model import JAXModel
 from amici.de_model import DEModel
 from amici.de_export import is_valid_identifier
@@ -96,12 +96,6 @@ def _jax_variable_ids(model: DEModel, sym_names: tuple[str, ...]) -> dict:
     }
 
 
-def _jnp_array_str(array) -> str:
-    elems = ", ".join(str(s) for s in array)
-
-    return f"jnp.array([{elems}])"
-
-
 class ODEExporter:
     """
     The ODEExporter class generates AMICI jax files for a model as
@@ -148,6 +142,16 @@ def __init__(
         """
         set_log_level(logger, verbose)
 
+        if any(event.updates_state for event in ode_model._events):
+            raise NotImplementedError(
+                "The JAX backend does not support models with event assignments."
+            )
+
+        if ode_model._algebraic_equations:
+            raise NotImplementedError(
+                "The JAX backend does not support models with algebraic equations."
+            )
+
         self.verbose: bool = logger.getEffectiveLevel() <= logging.DEBUG
 
         self.model_path: Path = Path()
@@ -237,7 +241,8 @@ def _generate_jax_code(self) -> None:
             # assign named variables from a jax array
             **_jax_variable_assignments(self.model, sym_names),
             # tuple of variable names (ids as they are unique)
-            **_jax_variable_ids(self.model, ("p", "k", "y", "x_rdata")),
+            **_jax_variable_ids(self.model, ("p", "k", "y", "w", "x_rdata")),
+            "P_VALUES": _jnp_array_str(self.model.val("p")),
             **{
                 "MODEL_NAME": self.model_name,
                 # keep track of the API version that the model was generated with so we

python/sdist/amici/sbml_import.py

@@ -1029,7 +1029,8 @@ def _gather_base_locals(
             if not x_ref.isSetId():
                 continue
             if (
-                x_ref.isSetStoichiometry()
+                hasattr(x_ref, "isSetStoichiometry")
+                and x_ref.isSetStoichiometry()
                 and not self.is_assignment_rule_target(x_ref)
             ):
                 value = sp.Float(x_ref.getStoichiometry())
@@ -2986,7 +2987,7 @@ def _get_element_stoichiometry(self, ele: libsbml.SBase) -> sp.Expr:
             if self.is_assignment_rule_target(ele):
                 return _get_identifier_symbol(ele)
 
-        if ele.isSetStoichiometry():
+        if hasattr(ele, "isSetStoichiometry") and ele.isSetStoichiometry():
             stoichiometry: float = ele.getStoichiometry()
             return (
                 sp.Integer(stoichiometry)