Add support in JAX backend for PEtab v2 by BSnelling · Pull Request #3115 · AMICI-dev/AMICI

BSnelling · 2026-01-09T14:02:32Z

No description provided.

codecov · 2026-01-09T14:23:36Z

Codecov Report

❌ Patch coverage is 83.33333% with 40 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.32%. Comparing base (126e936) to head (a1aefbc).

Files with missing lines	Patch %	Lines
python/sdist/amici/jax/petab.py	73.57%	37 Missing ⚠️
python/sdist/amici/sim/jax/__init__.py	96.07%	2 Missing ⚠️
python/sdist/amici/jax/model.py	90.90%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3115      +/-   ##
==========================================
- Coverage   78.24%   77.32%   -0.92%     
==========================================
  Files         315      316       +1     
  Lines       20482    20648     +166     
  Branches     1484     1483       -1     
==========================================
- Hits        16027    15967      -60     
- Misses       4447     4673     +226     
  Partials        8        8

Flag	Coverage Δ
cpp	`70.39% <24.05%> (-1.32%)`	⬇️
cpp_python	`37.85% <7.59%> (-0.01%)`	⬇️
petab	`47.14% <83.33%> (+0.58%)`	⬆️
petab_sciml	`?`
python	`68.85% <24.05%> (-1.52%)`	⬇️
sbmlsuite-jax	`32.76% <25.00%> (-0.65%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
python/sdist/amici/_symbolic/de_model.py	`91.68% <100.00%> (-0.24%)`	⬇️
...ython/sdist/amici/_symbolic/de_model_components.py	`93.25% <100.00%> (+0.15%)`	⬆️
...hon/sdist/amici/importers/petab/_petab_importer.py	`86.74% <ø> (-0.04%)`	⬇️
...dist/amici/importers/petab/v1/parameter_mapping.py	`77.83% <100.00%> (-7.55%)`	⬇️
python/sdist/amici/importers/pysb/__init__.py	`94.57% <ø> (-0.07%)`	⬇️
python/sdist/amici/importers/sbml/__init__.py	`94.26% <ø> (-0.10%)`	⬇️
python/sdist/amici/jax/_simulation.py	`94.44% <100.00%> (-2.62%)`	⬇️
python/sdist/amici/jax/ode_export.py	`88.05% <100.00%> (-1.18%)`	⬇️
python/sdist/amici/jax/model.py	`84.54% <90.90%> (-2.96%)`	⬇️
python/sdist/amici/sim/jax/__init__.py	`96.07% <96.07%> (ø)`
... and 1 more

... and 5 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

.github/workflows/test_petab_sciml.yml

.github/workflows/test_petab_test_suite.yml

doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb

python/sdist/amici/_symbolic/de_model.py

python/sdist/amici/jax/petab.py

tests/petab_test_suite/test_petab_v2_suite.py

python/sdist/amici/_symbolic/de_model.py

python/sdist/amici/importers/petab/v1/sbml_import.py

python/sdist/amici/jax/_simulation.py

sonarqubecloud · 2026-01-29T15:16:50Z

Quality Gate passed

Issues
2 New issues
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

python/sdist/amici/jax/petab.py

dweindl · 2026-02-03T21:41:07Z

Sorry for the delay, I'll try to review tomorrow or the day after.

dweindl

Cool, thanks. Here some initial comments. Didn't manage to go over everything yet.

doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb

python/sdist/amici/_symbolic/de_model.py

dweindl · 2026-02-04T18:29:50Z

python/sdist/amici/_symbolic/de_model_components.py

+            if (s.startswith("_petab_") and "indicator" in s) or s == "t":
+                continue
+            implicit_symbols.append(s)
+        return len(implicit_symbols) > 0


This sub-package should stay PEtab-independent. I think a better way to handle that would be passing all indicator variables as fixed_parameters to SbmlImporter._build_ode_model via SbmlImporter.sbml2jax and then use the same logic as for the non-jax event handling.

python/sdist/amici/_symbolic/de_model.py

python/sdist/amici/jax/ode_export.py

dweindl · 2026-02-04T18:53:17Z

python/sdist/amici/jax/petab.py

            self._initialize_model_with_nominal_values(model)
        )
-        self._parameter_mappings = self._get_parameter_mappings(scs)
+        if isinstance(petab_problem, petabv1.Problem):


Merge with the if isinstance above?

dweindl · 2026-02-04T18:58:20Z

python/sdist/amici/jax/petab.py

        """
-        scs = petab_problem.get_simulation_conditions_from_measurement_df()
-        self.simulation_conditions = tuple(tuple(sc) for sc in scs.values)
+        if isinstance(petab_problem, petabv2.Problem):


Is the plan to support both, petab v1 and v2, or is this just until petab-sciml is updated?

python/sdist/amici/jax/petab.py

…heaviside init

support explicit initial value event assignments fix deltax tcl

Co-authored-by: Fabian Fröhlich <fabian.frohlich@crick.ac.uk>

This reverts commit 82caa94.

…parameter_mapping

…event assignments with assumed priority

BSnelling

Thank you for the review @dweindl! I have implemented most points, but I have a couple of outstanding questions, if you could advise on those.

BSnelling · 2026-02-13T12:29:56Z

doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb

-    ")"
-   ],
-   "id": "3f2ab1acb3ba818f"
+    "amici_model.simulate(petab_problem.get_x_nominal_dict())"


When the notebook runs there is an error from this cell that the FIM was not computed: https://github.com/AMICI-dev/AMICI/actions/runs/21985309084/job/63517900881?pr=3115

@dweindl can you advise on that?

My fault. I hope #3125 fixes that.

BSnelling · 2026-02-13T13:29:32Z

python/sdist/amici/_symbolic/de_model_components.py

+        """Check whether the event has implicit triggers.
+        """
+        t = self.get_val()
+        return not t.free_symbols.issubset(allowed_symbols)


I removed the PEtab related references from this function and am passing the indicator variables as fixed_parameters. I wanted to include this logic in the has_explicit_trigger_times function but in all the JAX examples, self._t_root is not set. Should _t_root be set during the import of petab events?

Oh, okay. I think this might be a general issue, not just with JAX. My guess is that sympy isn't able to solve those petab event trigger piecewise expressions for $t$ (here

AMICI/python/sdist/amici/_symbolic/de_model_components.py

Lines 768 to 772 in 20b07e9

try:

self._t_root = sp.solve(self.get_val(), amici_time_symbol)

except NotImplementedError:

# the trigger can't be solved for `t`

pass

). That should be feasible, but I'm not sure how easy it is to implement. I won't be able to look into that for at least another week. Okay to keep your previous approach for now then. Please open an issue to address that later.

Indeed, sympy is failing to solve these expressions. Opened #3126

BSnelling · 2026-02-13T13:31:44Z

python/sdist/amici/jax/petab.py

-        self.simulation_conditions = tuple(tuple(sc) for sc in scs.values)
+        if isinstance(petab_problem, petabv1.Problem):
+            raise TypeError(
+               "JAXProblem does not support PEtab v1 problems. Upgrade the problem to PEtab v2."


The plan is just to support v2, so I've added this error if a v1 problem is encountered.

Great. Can you please replace all the petab.v1 constants by those in petab.v2 to make it clear that we don't want to deal with v1 here?
It seems there is still quite a bit of v1 code left here. Probably all occurrences of simulationConditionId/SIMULATION_CONDITION_ID. Is that still needed for petab-sciml? If not please remove.

dweindl

Thanks, @BSnelling. Here a few more comments. Okay for me to merge once all tests pass and to address the remaining issues together with the petab-sciml update if that's more convenient.

dweindl · 2026-02-13T15:00:24Z

python/sdist/amici/sim/jax/__init__.py

+                petabv1.OBSERVABLE_ID: obs,
+                petabv2.C.EXPERIMENT_ID: [experiment_id] * len(t),
+                petabv1.TIME: t[problem._ts_masks[ic, :]],
+                petabv1.SIMULATION: y[ic, problem._ts_masks[ic, :]],
+            },
+            index=problem._petab_measurement_indices[ic, :],
+        )
+        if (
+            petabv1.OBSERVABLE_PARAMETERS
+            in problem._petab_problem.measurement_df
+        ):                
+            df_sc[petabv2.C.OBSERVABLE_PARAMETERS] = (
+                problem._petab_problem.measurement_df.query(
+                    f"{petabv2.C.EXPERIMENT_ID} == '{experiment_id}'"
+                )[petabv2.C.OBSERVABLE_PARAMETERS]
+            )
+        if petabv1.NOISE_PARAMETERS in problem._petab_problem.measurement_df:


Suggested change

petabv1.OBSERVABLE_ID: obs,

petabv2.C.EXPERIMENT_ID: [experiment_id] * len(t),

petabv1.TIME: t[problem._ts_masks[ic, :]],

petabv1.SIMULATION: y[ic, problem._ts_masks[ic, :]],

},

index=problem._petab_measurement_indices[ic, :],

)

if (

petabv1.OBSERVABLE_PARAMETERS

in problem._petab_problem.measurement_df

):

df_sc[petabv2.C.OBSERVABLE_PARAMETERS] = (

problem._petab_problem.measurement_df.query(

f"{petabv2.C.EXPERIMENT_ID} == '{experiment_id}'"

)[petabv2.C.OBSERVABLE_PARAMETERS]

)

if petabv1.NOISE_PARAMETERS in problem._petab_problem.measurement_df:

petabv2.OBSERVABLE_ID: obs,

petabv2.C.EXPERIMENT_ID: [experiment_id] * len(t),

petabv2.TIME: t[problem._ts_masks[ic, :]],

petabv2.SIMULATION: y[ic, problem._ts_masks[ic, :]],

},

index=problem._petab_measurement_indices[ic, :],

)

if (

petabv2.OBSERVABLE_PARAMETERS

in problem._petab_problem.measurement_df

):

df_sc[petabv2.C.OBSERVABLE_PARAMETERS] = (

problem._petab_problem.measurement_df.query(

f"{petabv2.C.EXPERIMENT_ID} == '{experiment_id}'"

)[petabv2.C.OBSERVABLE_PARAMETERS]

)

if petabv2.NOISE_PARAMETERS in problem._petab_problem.measurement_df:

No effect, but since the function is specifically for petab v2, let's use the constants from there.

dweindl · 2026-02-13T15:07:40Z

python/sdist/amici/jax/petab.py

-        self.simulation_conditions = tuple(tuple(sc) for sc in scs.values)
+        if isinstance(petab_problem, petabv1.Problem):
+            raise TypeError(
+               "JAXProblem does not support PEtab v1 problems. Upgrade the problem to PEtab v2."


Great. Can you please replace all the petab.v1 constants by those in petab.v2 to make it clear that we don't want to deal with v1 here?
It seems there is still quite a bit of v1 code left here. Probably all occurrences of simulationConditionId/SIMULATION_CONDITION_ID. Is that still needed for petab-sciml? If not please remove.

dweindl · 2026-02-13T15:19:23Z

python/sdist/amici/jax/petab.py

                            self._petab_problem.observable_df.loc[
-                                oid, petab.OBSERVABLE_TRANSFORMATION
+                                oid, petabv1.OBSERVABLE_TRANSFORMATION
                            ]


As a general comment: With petab.v2.Problem, all .*_df properties construct a new DataFrame on each call. Try to avoid accessing those properties inside loops.

In many cases it might be more convenient to directly use the underlying objects instead of the DataFrames (e.g., Problem.observables instead of Problem.observable_df).

BSnelling force-pushed the bes/petabv2_jax branch from b2fdbec to 9d107b3 Compare January 22, 2026 17:06

BSnelling commented Jan 26, 2026

View reviewed changes

BSnelling marked this pull request as ready for review January 26, 2026 12:31

BSnelling requested a review from a team as a code owner January 26, 2026 12:31

BSnelling requested a review from FFroehlich January 26, 2026 12:31

FFroehlich requested changes Jan 26, 2026

View reviewed changes

BSnelling requested review from FFroehlich and dweindl January 29, 2026 15:41

BSnelling commented Jan 29, 2026

View reviewed changes

python/sdist/amici/jax/petab.py Show resolved Hide resolved

FFroehlich approved these changes Feb 3, 2026

View reviewed changes

dweindl reviewed Feb 4, 2026

View reviewed changes

BSnelling and others added 16 commits February 9, 2026 15:14

event assignments jax - sbml cases 348 - 404

1885911

fix up sbml test cases - not implemented priority, update t_eps, fix …

1e5bfeb

…heaviside init

initialValue False not implemented

ff35c15

try fix other test cases

973aaa1

params only in explicit triggers - and matrix only in JAX again

36f16e6

oops committed breakpoint

865c340

looking for initialValue test cases

2279cca

support explicit initial value event assignments fix deltax tcl

do not update h pre-solve

599b8b5

handle_t0_event

52fea34

reinstate time skip (hack diffrax bug?)

dbf43c1

Update python/sdist/amici/jax/_simulation.py

fab9f6f

Co-authored-by: Fabian Fröhlich <fabian.frohlich@crick.ac.uk>

Revert "Update python/sdist/amici/jax/_simulation.py"

ebd0c0c

This reverts commit 82caa94.

rm clip controller

c19b183

handle t0 event near zero

0961646

skip non-time dependent event assignment cases

b8c1a8c

first pass petabv2 - updating JAXProblem init

60e7cf5

BSnelling added 18 commits February 9, 2026 15:14

petabv2 test cases up to 15-ish

b4aed26

petab v2 test cases up to 27-ish

e63b748

rework petabv2 jax test cases with ExperimentsToSbmlEvents and no v1 …

6efc38e

…parameter_mapping

fix some rebase issues

b0d7ad1

add test skip for petab v1

c33f01d

remaining petab v2 test cases

5c1b30e

update workflows - deactivate petab_sciml wf for now

b4aecfd

update tests to skip on petab v2 type error

b382c1f

skip some more tests and reinstate more specific implicit triggers check

8223a61

fixup benchmark skipping check

a1d88c3

tidying - add docstrings - rm outputs in notebook

2746d5b

skip implicit benchmark case - and prior distribution cases

693e29b

review feedback - rm simultaneous event check - implement sequential …

607781b

…event assignments with assumed priority

fix test_jax tests and sbml cases with no y0

828016c

fix pysb test case for jax

e3126f5

use h symbol and update example notebook

df88e31

rm v1 instance checks and improve preeq conditionals

720980c

implement implicit triggers using fixed parameters check

8286033

BSnelling force-pushed the bes/petabv2_jax branch from cadaa43 to 8286033 Compare February 9, 2026 15:19

BSnelling added 7 commits February 11, 2026 11:45

temp workaround for pysb build issue

590fa97

pysb workaround in petab workflow too

781957c

fix indentation errors in gen JAX code - restore petabv1 conditional

5bbd612

skip JAXProblems with v1 problems

48b13d6

skip JAXProblems with v1 problems - again

bfba95c

pin optax version

89b52a1

check implicit triggers in sep function

a1aefbc

BSnelling commented Feb 13, 2026

View reviewed changes

dweindl reviewed Feb 13, 2026

View reviewed changes

	try:
	self._t_root = sp.solve(self.get_val(), amici_time_symbol)
	except NotImplementedError:
	# the trigger can't be solved for `t`
	pass

Conversation

BSnelling commented Jan 9, 2026

Uh oh!

codecov bot commented Jan 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sonarqubecloud bot commented Jan 29, 2026

Quality Gate passed

Uh oh!

Uh oh!

dweindl commented Feb 3, 2026

Uh oh!

dweindl left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

BSnelling left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

BSnelling Feb 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dweindl left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov bot commented Jan 9, 2026 •

edited

Loading

BSnelling Feb 13, 2026 •

edited

Loading