Add support in JAX backend for PEtab v2 (#3115)

* event assignments jax - sbml cases 348 - 404

* fix up sbml test cases - not implemented priority, update t_eps, fix heaviside init

* initialValue False not implemented

* try fix other test cases

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

* oops committed breakpoint

* looking for initialValue test cases

support explicit initial value event assignments

fix deltax tcl

* do not update h pre-solve

* handle_t0_event

* reinstate time skip (hack diffrax bug?)

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

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

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

This reverts commit 82caa94.

* rm clip controller

* handle t0 event near zero

* skip non-time dependent event assignment cases

* first pass petabv2 - updating JAXProblem init

* petabv2 test cases up to 15-ish

* petab v2 test cases up to 27-ish

* rework petabv2 jax test cases with ExperimentsToSbmlEvents and no v1 parameter_mapping

* fix some rebase issues

* add test skip for petab v1

* remaining petab v2 test cases

* update workflows - deactivate petab_sciml wf for now

* update tests to skip on petab v2 type error

* skip some more tests and reinstate more specific implicit triggers check

* fixup benchmark skipping check

* tidying - add docstrings - rm outputs in notebook

* skip implicit benchmark case - and prior distribution cases

* review feedback - rm simultaneous event check - implement sequential event assignments with assumed priority

* fix test_jax tests and sbml cases with no y0

* fix pysb test case for jax

* use h symbol and update example notebook

* rm v1 instance checks and improve preeq conditionals

* implement implicit triggers using fixed parameters check

* temp workaround for pysb build issue

* pysb workaround in petab workflow too

* fix indentation errors in gen JAX code - restore petabv1 conditional

* skip JAXProblems with v1 problems

* skip JAXProblems with v1 problems - again

* check implicit triggers in sep function

* use petabv2 constants - avoid df usage in loops

* restore pysb install to master

* tody petab file

---------

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

Author: BSnelling

.github/workflows/test_petab_sciml.yml

@@ -1,14 +1,14 @@
 name: PEtab SciML
-on:
-  push:
-    branches:
-      - main
-      - 'release*'
-  pull_request:
-    branches:
-      - main
-  merge_group:
-  workflow_dispatch:
+# on:
+#   push:
+#     branches:
+#       - main
+#       - 'release*'
+#   pull_request:
+#     branches:
+#       - main
+#   merge_group:
+#   workflow_dispatch:
 
 jobs:
   build:

.github/workflows/test_petab_test_suite.yml

@@ -172,7 +172,7 @@ jobs:
           git clone https://github.com/PEtab-dev/petab_test_suite \
             && source ./venv/bin/activate \
             && cd petab_test_suite \
-            && git checkout c12b9dc4e4c5585b1b83a1d6e89fd22447c46d03 \
+            && git checkout 9542847fb99bcbdffc236e2ef45ba90580a210fa \
             && pip3 install -e .
 
 # TODO: once there is a PEtab v2 benchmark collection
@@ -186,7 +186,7 @@ jobs:
         run: |
           source ./venv/bin/activate \
             && python3 -m pip uninstall -y petab \
-            && python3 -m pip install git+https://github.com/petab-dev/libpetab-python.git@8dc6c1c4b801fba5acc35fcd25308a659d01050e \
+            && python3 -m pip install git+https://github.com/petab-dev/libpetab-python.git@d57d9fed8d8d5f8592e76d0b15676e05397c3b4b \
             && python3 -m pip install git+https://github.com/pysb/pysb@master \
             && python3 -m pip install sympy>=1.12.1
 

doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb

@@ -32,7 +32,8 @@
    "outputs": [],
    "source": [
     "import petab.v1 as petab\n",
-    "from amici.importers.petab.v1 import import_petab_problem\n",
+    "from amici.importers.petab import *\n",
+    "from petab.v2 import Problem\n",
     "\n",
     "# Define the model name and YAML file location\n",
     "model_name = \"Boehm_JProteomeRes2014\"\n",
@@ -41,14 +42,20 @@
     "    f\"master/Benchmark-Models/{model_name}/{model_name}.yaml\"\n",
     ")\n",
     "\n",
-    "# Load the PEtab problem from the YAML file\n",
-    "petab_problem = petab.Problem.from_yaml(yaml_url)\n",
+    "# Load the PEtab problem from the YAML file as a PEtab v2 problem\n",
+    "# (the JAX backend only supports PEtab v2)\n",
+    "petab_problem = Problem.from_yaml(yaml_url)\n",
     "\n",
     "# Import the PEtab problem as a JAX-compatible AMICI problem\n",
-    "jax_problem = import_petab_problem(\n",
-    "    petab_problem,\n",
-    "    verbose=False,  # no text output\n",
-    "    jax=True,  # return jax problem\n",
+    "pi = PetabImporter(\n",
+    "    petab_problem=petab_problem,\n",
+    "    module_name=model_name,\n",
+    "    compile_=True,\n",
+    "    jax=True,\n",
+    ")\n",
+    "\n",
+    "jax_problem = pi.create_simulator(\n",
+    "    force_import=True,\n",
     ")"
    ]
   },
@@ -75,6 +82,16 @@
     "llh, results = run_simulations(jax_problem)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6c5b2980-13f0-42e9-b13e-0fce05793910",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "results"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "415962751301c64a",
@@ -90,11 +107,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Define the simulation condition\n",
-    "simulation_condition = (\"model1_data1\",)\n",
+    "# # Define the simulation condition\n",
+    "experiment_condition = \"_petab_experiment_condition___default__\"\n",
     "\n",
-    "# Access the results for the specified condition\n",
-    "ic = results[\"simulation_conditions\"].index(simulation_condition)\n",
+    "# # Access the results for the specified condition\n",
+    "ic = results[\"dynamic_conditions\"].index(experiment_condition)\n",
     "print(\"llh: \", results[\"llh\"][ic])\n",
     "print(\"state variables: \", results[\"x\"][ic, :])"
    ]
@@ -146,8 +163,8 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
-    "# Define the simulation condition\n",
-    "simulation_condition = (\"model1_data1\",)\n",
+    "# Define the experiment condition\n",
+    "experiment_condition = \"_petab_experiment_condition___default__\"\n",
     "\n",
     "\n",
     "def plot_simulation(results):\n",
@@ -158,7 +175,7 @@
     "        results (dict): Simulation results from run_simulations.\n",
     "    \"\"\"\n",
     "    # Extract the simulation results for the specific condition\n",
-    "    ic = results[\"simulation_conditions\"].index(simulation_condition)\n",
+    "    ic = results[\"dynamic_conditions\"].index(experiment_condition)\n",
     "\n",
     "    # Create a new figure for the state trajectories\n",
     "    plt.figure(figsize=(8, 6))\n",
@@ -172,7 +189,7 @@
     "    # Add labels, legend, and grid\n",
     "    plt.xlabel(\"Time\")\n",
     "    plt.ylabel(\"State Values\")\n",
-    "    plt.title(simulation_condition)\n",
+    "    plt.title(experiment_condition)\n",
     "    plt.legend()\n",
     "    plt.grid(True)\n",
     "    plt.show()\n",
@@ -187,18 +204,7 @@
    "id": "4fa97c33719c2277",
    "metadata": {},
    "source": [
-    "`run_simulations` enables users to specify the simulation conditions to be executed. For more complex models, this allows for restricting simulations to a subset of conditions. Since the Böhm model includes only a single condition, we demonstrate this functionality by simulating no condition at all."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "7950774a3e989042",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "llh, results = run_simulations(jax_problem, simulation_conditions=tuple())\n",
-    "results"
+    "`run_simulations` enables users to specify the simulation experiments to be executed. For more complex models, this allows for restricting simulations to a subset of experiments by passing a tuple of experiment ids under the keyword `simulation_experiments` to `run_simulations`."
    ]
   },
   {
@@ -384,8 +390,8 @@
     "from amici.jax import ReturnValue\n",
     "\n",
     "# Define the simulation condition\n",
-    "simulation_condition = (\"model1_data1\",)\n",
-    "ic = jax_problem.simulation_conditions.index(simulation_condition)\n",
+    "experiment_condition = \"_petab_experiment_condition___default__\"\n",
+    "ic = 0\n",
     "\n",
     "# Load condition-specific data\n",
     "ts_dyn = jax_problem._ts_dyn[ic, :]\n",
@@ -397,7 +403,7 @@
     "nps = jax_problem._np_numeric[ic, :]\n",
     "\n",
     "# Load parameters for the specified condition\n",
-    "p = jax_problem.load_model_parameters(simulation_condition[0])\n",
+    "p = jax_problem.load_model_parameters(jax_problem._petab_problem.experiments[0], is_preeq=False)\n",
     "\n",
     "\n",
     "# Define a function to compute the gradient with respect to dynamic timepoints\n",
@@ -431,13 +437,17 @@
    "cell_type": "markdown",
    "id": "19ca88c8900584ce",
    "metadata": {},
-   "source": "## Model training"
+   "source": [
+    "## Model training"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "7f99c046d7d4e225",
    "metadata": {},
-   "source": "This setup makes it pretty straightforward to train models using [equinox](https://docs.kidger.site/equinox/) and [optax](https://optax.readthedocs.io/en/latest/) frameworks. Below we provide barebones implementation that runs training for 5 steps using Adam."
+   "source": [
+    "This setup makes it pretty straightforward to train models using [equinox](https://docs.kidger.site/equinox/) and [optax](https://optax.readthedocs.io/en/latest/) frameworks. Below we provide barebones implementation that runs training for 5 steps using Adam."
+   ]
   },
   {
    "cell_type": "code",
@@ -569,16 +579,20 @@
     "from amici.sim.sundials.petab.v1 import simulate_petab\n",
     "\n",
     "# Import the PEtab problem as a standard AMICI model\n",
-    "amici_model = import_petab_problem(\n",
-    "    petab_problem,\n",
-    "    verbose=False,\n",
-    "    jax=False,  # load the amici model this time\n",
+    "pi = PetabImporter(\n",
+    "    petab_problem=petab_problem,\n",
+    "    module_name=model_name,\n",
+    "    compile_=True,\n",
+    "    jax=False,\n",
+    ")\n",
+    "\n",
+    "amici_model = pi.create_simulator(\n",
+    "    force_import=True,\n",
     ")\n",
     "\n",
     "# Configure the solver with appropriate tolerances\n",
-    "solver = amici_model.create_solver()\n",
-    "solver.set_absolute_tolerance(1e-8)\n",
-    "solver.set_relative_tolerance(1e-16)\n",
+    "amici_model.solver.set_absolute_tolerance(1e-8)\n",
+    "amici_model.solver.set_relative_tolerance(1e-16)\n",
     "\n",
     "# Prepare the parameters for the simulation\n",
     "problem_parameters = dict(\n",
@@ -594,86 +608,65 @@
    "outputs": [],
    "source": [
     "# Profile simulation only\n",
-    "solver.set_sensitivity_order(SensitivityOrder.none)"
+    "amici_model.solver.set_sensitivity_order(SensitivityOrder.none)"
    ]
   },
   {
-   "metadata": {},
    "cell_type": "code",
-   "outputs": [],
    "execution_count": null,
+   "id": "42cbc67bc09b67dc",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "%%timeit\n",
-    "simulate_petab(\n",
-    "    petab_problem,\n",
-    "    amici_model=amici_model,\n",
-    "    solver=solver,\n",
-    "    problem_parameters=problem_parameters,\n",
-    "    scaled_parameters=True,\n",
-    "    scaled_gradients=True,\n",
-    ")"
-   ],
-   "id": "42cbc67bc09b67dc"
+    "amici_model.simulate(petab_problem.get_x_nominal_dict())"
+   ]
   },
   {
-   "metadata": {},
    "cell_type": "code",
-   "outputs": [],
    "execution_count": null,
+   "id": "4f1c06c5893a9c07",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "# Profile gradient computation using forward sensitivity analysis\n",
-    "solver.set_sensitivity_order(SensitivityOrder.first)\n",
-    "solver.set_sensitivity_method(SensitivityMethod.forward)"
-   ],
-   "id": "4f1c06c5893a9c07"
+    "amici_model.solver.set_sensitivity_order(SensitivityOrder.first)\n",
+    "amici_model.solver.set_sensitivity_method(SensitivityMethod.forward)"
+   ]
   },
   {
-   "metadata": {},
    "cell_type": "code",
-   "outputs": [],
    "execution_count": null,
+   "id": "7367a19bcea98597",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "%%timeit\n",
-    "simulate_petab(\n",
-    "    petab_problem,\n",
-    "    amici_model=amici_model,\n",
-    "    solver=solver,\n",
-    "    problem_parameters=problem_parameters,\n",
-    "    scaled_parameters=True,\n",
-    "    scaled_gradients=True,\n",
-    ")"
-   ],
-   "id": "7367a19bcea98597"
+    "amici_model.simulate(petab_problem.get_x_nominal_dict())"
+   ]
   },
   {
-   "metadata": {},
    "cell_type": "code",
-   "outputs": [],
    "execution_count": null,
+   "id": "a31e8eda806c2d7",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "# Profile gradient computation using adjoint sensitivity analysis\n",
-    "solver.set_sensitivity_order(SensitivityOrder.first)\n",
-    "solver.set_sensitivity_method(SensitivityMethod.adjoint)"
-   ],
-   "id": "a31e8eda806c2d7"
+    "amici_model.solver.set_sensitivity_order(SensitivityOrder.first)\n",
+    "amici_model.solver.set_sensitivity_method(SensitivityMethod.adjoint)"
+   ]
   },
   {
-   "metadata": {},
    "cell_type": "code",
-   "outputs": [],
    "execution_count": null,
+   "id": "3f2ab1acb3ba818f",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "%%timeit\n",
-    "simulate_petab(\n",
-    "    petab_problem,\n",
-    "    amici_model=amici_model,\n",
-    "    solver=solver,\n",
-    "    problem_parameters=problem_parameters,\n",
-    "    scaled_parameters=True,\n",
-    "    scaled_gradients=True,\n",
-    ")"
-   ],
-   "id": "3f2ab1acb3ba818f"
+    "amici_model.simulate(petab_problem.get_x_nominal_dict())"
+   ]
   }
  ],
  "metadata": {
@@ -691,7 +684,8 @@
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3"
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,

python/sdist/amici/_symbolic/de_model.py

@@ -2679,11 +2679,16 @@ def has_priority_events(self) -> bool:
     def has_implicit_event_assignments(self) -> bool:
         """
         Checks whether the model has event assignments with implicit triggers
+        (i.e. triggers that are not time based).
 
         :return:
             boolean indicating if event assignments with implicit triggers are present
         """
-        return any(event.updates_state and not event.has_explicit_trigger_times({}) for event in self._events)
+        fixed_symbols = set([k._symbol for k in self._fixed_parameters])
+        allowed_symbols = fixed_symbols | {amici_time_symbol}
+        # TODO: update to use has_explicit_trigger_times once 
+        # https://github.com/AMICI-dev/AMICI/issues/3126 is resolved
+        return any(event.updates_state and event._has_implicit_triggers(allowed_symbols) for event in self._events)
 
     def toposort_expressions(
         self, reorder: bool = True