Split jax/cpp PEtab benchmark tests (#2794)

FFroehlich · web-flow · commit 4eebf1613352 · 2025-06-18T17:02:29.000Z
* Remove import guard from PEtab benchmark

* Apply suggestions from code review

* remove spurious inits

* move and update AGENTS.md

* Move benchmark fixture to conftest

* Update tests/benchmark-models/conftest.py

* Update tests/benchmark-models/test_petab_benchmark_jax.py
diff --git a/.github/workflows/test_benchmark_collection_models.yml b/.github/workflows/test_benchmark_collection_models.yml
@@ -14,8 +14,8 @@ on:
     - cron:  '48 4 * * *'
 
 jobs:
-  build:
-    name: Benchmark Collection
+  cpp:
+    name: Benchmark Collection CPP
 
     runs-on: ubuntu-24.04
 
@@ -66,11 +66,11 @@ jobs:
       env:
         AMICI_PARALLEL_COMPILE: ""
       run: |
-          cd tests/benchmark-models && pytest \
-              --durations=10 \
-              --cov=amici \
-              --cov-report=xml:"coverage_py.xml" \
-              --cov-append \
+        cd tests/benchmark-models && pytest test_petab_benchmark.py \
+          --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'
@@ -93,3 +93,67 @@ jobs:
         path: |
           tests/benchmark-models/computation_times.csv
           tests/benchmark-models/computation_times.png
+
+  jax:
+    name: Benchmark Collection JAX
+    runs-on: ubuntu-24.04
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [ "3.12" ]
+        extract_subexpressions: ["true", "false"]
+    env:
+      AMICI_EXTRACT_CSE: ${{ matrix.extract_subexpressions }}
+    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
+
+    - name: Install apt dependencies
+      uses: ./.github/actions/install-apt-dependencies
+
+    - run: echo "${HOME}/.local/bin/" >> $GITHUB_PATH
+
+    - name: Create AMICI sdist
+      run: pip3 install build && cd python/sdist && python3 -m build --sdist
+
+    - name: Install AMICI sdist
+      run: |
+        pip3 install --user petab[vis] && \
+        pip3 install -v --user \
+            $(ls -t python/sdist/dist/amici-*.tar.gz | head -1)[petab,test,vis,jax]
+
+    - name: Install test dependencies
+      run: |
+        python3 -m pip uninstall -y petab && python3 -m pip install git+https://github.com/petab-dev/libpetab-python.git@develop \
+        &&  python3 -m pip install -U sympy \
+        &&  python3 -m pip install git+https://github.com/ICB-DCM/fiddy.git
+
+    - name: Download benchmark collection
+      run: |
+        pip install git+https://github.com/Benchmarking-Initiative/Benchmark-Models-PEtab.git@master#subdirectory=src/python
+
+    - name: Run JAX tests
+      env:
+        AMICI_PARALLEL_COMPILE: ""
+      run: |
+        cd tests/benchmark-models && pytest test_petab_benchmark_jax.py \
+          --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@v5
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        files: coverage_py.xml
+        flags: python
+        fail_ci_if_error: true
+        verbose: true
diff --git a/doc/AGENTS.md b/doc/AGENTS.md
@@ -0,0 +1,63 @@
+# Project Agents.md Guide for OpenAI Codex
+
+This Agents.md file provides comprehensive guidance for OpenAI Codex and other AI agents working with this codebase.
+
+## Project Structure for OpenAI Codex Navigation
+
+AMICI is a python package that uses SWIG to generate python bindings to C++ code. There are also a deprecated matlab interface to the C++ code, which will be removed at some point in the future.
+
+- `/binder`: binder configuration
+- `/cmake`: various cmake utility functions
+- `/container`: docker configuration
+- `/doc`: high level documentation, all API documentation is automatically generated using doxygen/sphinx
+- `/include`: C++ header files
+- `/matlab`: matlab interface
+- `/models`: pre-generated c++ models for testing
+- `/python`: python source code
+  - `/benchmark`: helper scripts for benchmarking
+  - `/sdist`: python package
+    - `/amici`: python module
+      - `/_codegen` helper functions for C++ code generation
+      - `/debugging` helpfer functions for debugging simulation failures
+      - `/include` symlink to C++ headers
+      - `/jax` code for JAX backend, alternative to C++ backend
+      - `/petab` interface for simulating models/problems in the PEtab format
+      - `/testing` helpfer functions for tests
+    - `/tests`: self contained python package specific tests
+- `/scripts`: bash scripts for testing and installing the package
+- `/src`: C++ source files
+- `/swig`: definition of the SWIG interface
+- `/tests`: C++ tests and python tests that require third party resources, base directory contains tests for the SBML testsuite
+  - `/benchmark-models` regression tests for the PEtab benchmark collection
+  - `/cpp` C++ tests
+  - `/generateTestConfg` helper functions to generate python test configurations for pregenerated C++ models
+  - `/performace` performance tests for the PEtab benchmark collection
+  - `/petab_test_suite` regression tests for the PEtab testsuite
+
+## Contribution Guidelines
+
+Please see instructions in `doc/CONTRIBUTING.md`
+
+# Agent Instructions
+
+To ensure all tools and dependencies are available, activate the virtual environment before running any commands:
+
+```bash
+source ./venv/bin/activate
+```
+
+This project uses `pre-commit` for linting and `pytest` for tests. Run them on changed files whenever you make modifications.
+
+When running the tests locally, change into the test directory first:
+
+```bash
+cd tests/benchmark-models
+pytest test_petab_benchmark.py
+pytest test_petab_benchmark_jax.py
+```
+
+To quickly verify the benchmark tests, you can limit execution to a small model:
+
+```bash
+pytest -k Boehm_JProteomeRes2014 test_petab_benchmark.py
+```
diff --git a/doc/conf.py b/doc/conf.py
@@ -259,6 +259,7 @@ def install_doxygen():
     "MATLAB_.md",
     "CPP_.md",
     "gfx",
+    "AGENTS.md",
 ]
 
 # The name of the Pygments (syntax highlighting) style to use.
diff --git a/tests/benchmark-models/conftest.py b/tests/benchmark-models/conftest.py
@@ -0,0 +1,33 @@
+import copy
+from pathlib import Path
+
+import pytest
+import petab.v1 as petab
+from petab.v1.lint import measurement_table_has_timepoint_specific_mappings
+
+import benchmark_models_petab
+from amici.petab.petab_import import import_petab_problem
+
+from test_petab_benchmark import problems
+
+script_dir = Path(__file__).parent
+repo_root = script_dir.parent.parent
+benchmark_outdir = repo_root / "test_bmc"
+
+
+@pytest.fixture(scope="session", params=problems, ids=problems)
+def benchmark_problem(request):
+    """Fixture providing model and PEtab problem for a benchmark model."""
+    problem_id = request.param
+    petab_problem = benchmark_models_petab.get_problem(problem_id)
+    flat_petab_problem = copy.deepcopy(petab_problem)
+    if measurement_table_has_timepoint_specific_mappings(
+        petab_problem.measurement_df,
+    ):
+        petab.flatten_timepoint_specific_output_overrides(flat_petab_problem)
+
+    amici_model = import_petab_problem(
+        flat_petab_problem,
+        model_output_dir=benchmark_outdir / problem_id,
+    )
+    return problem_id, flat_petab_problem, petab_problem, amici_model
diff --git a/tests/benchmark-models/test_petab_benchmark.py b/tests/benchmark-models/test_petab_benchmark.py
@@ -5,40 +5,38 @@
 for a subset of the benchmark problems.
 """
 
-import copy
-from functools import partial
 from pathlib import Path
 
+import contextlib
+import logging
 import os
-import fiddy
-import amici
+from collections import defaultdict
+from dataclasses import dataclass, field
+
 import numpy as np
 import pandas as pd
 import petab.v1 as petab
 import pytest
-from amici.petab.petab_import import import_petab_problem
-import benchmark_models_petab
-from collections import defaultdict
-from dataclasses import dataclass, field
-from amici import SensitivityMethod
-from petab.v1.lint import measurement_table_has_timepoint_specific_mappings
-from fiddy import MethodId, get_derivative
-from fiddy.derivative_check import NumpyIsCloseDerivativeCheck
-from fiddy.extensions.amici import simulate_petab_to_cached_functions
-from fiddy.success import Consistency
-import contextlib
-import logging
 import yaml
+from petab.v1.lint import measurement_table_has_timepoint_specific_mappings
+from petab.v1.visualize import plot_problem
+
+import amici
+from amici import SensitivityMethod
 from amici.logging import get_logger
+from amici.petab.petab_import import import_petab_problem
 from amici.petab.simulations import (
     LLH,
-    SLLH,
     RDATAS,
     rdatas_to_measurement_df,
     simulate_petab,
 )
-
-from petab.v1.visualize import plot_problem
+import benchmark_models_petab
+import fiddy
+from fiddy import MethodId, get_derivative
+from fiddy.derivative_check import NumpyIsCloseDerivativeCheck
+from fiddy.extensions.amici import simulate_petab_to_cached_functions
+from fiddy.success import Consistency
 
 
 # Enable various debug output
@@ -241,133 +239,6 @@ class GradientCheckSettings:
 )
 
 
-@pytest.fixture(scope="session", params=problems, ids=problems)
-def benchmark_problem(request):
-    """Fixture providing model and PEtab problem for a problem from
-    the benchmark problem collection."""
-    problem_id = request.param
-    petab_problem = benchmark_models_petab.get_problem(problem_id)
-    flat_petab_problem = copy.deepcopy(petab_problem)
-    if measurement_table_has_timepoint_specific_mappings(
-        petab_problem.measurement_df,
-    ):
-        petab.flatten_timepoint_specific_output_overrides(flat_petab_problem)
-
-    # Setup AMICI objects.
-    amici_model = import_petab_problem(
-        flat_petab_problem,
-        model_output_dir=benchmark_outdir / problem_id,
-    )
-    return problem_id, flat_petab_problem, petab_problem, amici_model
-
-
-@pytest.mark.filterwarnings(
-    "ignore:The following problem parameters were not used *",
-    "ignore: The environment variable *",
-    "ignore:Adjoint sensitivity analysis for models with discontinuous ",
-)
-def test_jax_llh(benchmark_problem):
-    import jax
-    import equinox as eqx
-    import jax.numpy as jnp
-    from amici.jax.petab import run_simulations, JAXProblem
-
-    jax.config.update("jax_enable_x64", True)
-    from beartype import beartype
-
-    problem_id, flat_petab_problem, petab_problem, amici_model = (
-        benchmark_problem
-    )
-
-    amici_solver = amici_model.getSolver()
-    cur_settings = settings[problem_id]
-    amici_solver.setAbsoluteTolerance(1e-8)
-    amici_solver.setRelativeTolerance(1e-8)
-    amici_solver.setMaxSteps(10_000)
-
-    simulate_amici = partial(
-        simulate_petab,
-        petab_problem=flat_petab_problem,
-        amici_model=amici_model,
-        solver=amici_solver,
-        scaled_parameters=True,
-        scaled_gradients=True,
-        log_level=logging.DEBUG,
-    )
-
-    np.random.seed(cur_settings.rng_seed)
-
-    problem_parameters = None
-    if problem_id in problems_for_gradient_check:
-        point = flat_petab_problem.x_nominal_free_scaled
-        for _ in range(20):
-            amici_solver.setSensitivityMethod(amici.SensitivityMethod.adjoint)
-            amici_solver.setSensitivityOrder(amici.SensitivityOrder.first)
-            amici_model.setSteadyStateSensitivityMode(
-                cur_settings.ss_sensitivity_mode
-            )
-            point_noise = (
-                np.random.randn(len(point)) * cur_settings.noise_level
-            )
-            point += point_noise  # avoid small gradients at nominal value
-
-            problem_parameters = dict(
-                zip(flat_petab_problem.x_free_ids, point)
-            )
-
-            r_amici = simulate_amici(
-                problem_parameters=problem_parameters,
-            )
-            if np.isfinite(r_amici[LLH]):
-                break
-        else:
-            raise RuntimeError("Could not compute expected derivative.")
-    else:
-        r_amici = simulate_amici()
-    llh_amici = r_amici[LLH]
-
-    jax_model = import_petab_problem(
-        petab_problem,
-        model_output_dir=benchmark_outdir / (problem_id + "_jax"),
-        jax=True,
-    )
-    jax_problem = JAXProblem(jax_model, petab_problem)
-    if problem_parameters:
-        jax_problem = eqx.tree_at(
-            lambda x: x.parameters,
-            jax_problem,
-            jnp.array(
-                [problem_parameters[pid] for pid in jax_problem.parameter_ids]
-            ),
-        )
-
-    if problem_id in problems_for_gradient_check:
-        beartype(run_simulations)(jax_problem)
-        (llh_jax, _), sllh_jax = eqx.filter_value_and_grad(
-            run_simulations, has_aux=True
-        )(jax_problem)
-    else:
-        llh_jax, _ = beartype(run_simulations)(jax_problem)
-
-    np.testing.assert_allclose(
-        llh_jax,
-        llh_amici,
-        rtol=1e-3,
-        atol=1e-3,
-        err_msg=f"LLH mismatch for {problem_id}",
-    )
-
-    if problem_id in problems_for_gradient_check:
-        sllh_amici = r_amici[SLLH]
-        np.testing.assert_allclose(
-            sllh_jax.parameters,
-            np.array([sllh_amici[pid] for pid in jax_problem.parameter_ids]),
-            rtol=1e-2,
-            atol=1e-2,
-            err_msg=f"SLLH mismatch for {problem_id}, {dict(zip(jax_problem.parameter_ids, sllh_jax.parameters))}",
-        )
-
-
 @pytest.mark.filterwarnings(
     "ignore:divide by zero encountered in log",
     # https://github.com/AMICI-dev/AMICI/issues/18
diff --git a/tests/benchmark-models/test_petab_benchmark_jax.py b/tests/benchmark-models/test_petab_benchmark_jax.py

Original file line number	Diff line number	Diff line change
`@@ -259,6 +259,7 @@ def install_doxygen():`
`259`	`259`	`"MATLAB_.md",`
`260`	`260`	`"CPP_.md",`
`261`	`261`	`"gfx",`
	`262`	`+ "AGENTS.md",`
`262`	`263`	`]`
`263`	`264`
`264`	`265`	`# The name of the Pygments (syntax highlighting) style to use.`