Update PEtab SciML support to PEtab v2 (#3165)

* bank changes - values ok - gradients broken

* sciml tests passing (one skipped)

* fix petab test suite

* refactor sweep - for readability mostly

* delete unused code

* fix updating parameters in notebook for unscaled

* update test_sciml tests

* Apply suggestions from code review

Co-authored-by: Daniel Weindl <dweindl@users.noreply.github.com>

* implement feedback from code review

* revert out_val mapping

* try sympify petab again

---------

Co-authored-by: Daniel Weindl <dweindl@users.noreply.github.com>

Author: BSnelling

…hub/workflows/test_petab_sciml.yml.FIXME .github/workflows/test_petab_sciml.yml.github/workflows/test_petab_sciml.yml.FIXME renamed to .github/workflows/test_petab_sciml.yml .github/workflows/test_petab_sciml.yml.FIXME renamed to .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:
@@ -33,11 +33,10 @@ jobs:
         with:
           fetch-depth: 20
 
-      # todo, update after https://github.com/sebapersson/petab_sciml_testsuite/issues/14 is merged
       - name: Download PEtab SciML test suite
         run: |
           git clone --depth 1 --branch main \
-            https://github.com/FFroehlich/petab_sciml_testsuite \
+            https://github.com/PEtab-dev/petab_sciml_testsuite.git \
             tests/sciml/testsuite
 
       - name: Install apt dependencies
@@ -74,7 +73,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@sciml
+            && python3 -m pip install git+https://github.com/petab-dev/libpetab-python.git@main
 
       - name: Run PEtab SciML testsuite
         run: |

doc/examples/example_jax_petab/ExampleJaxPEtab.ipynb

@@ -228,7 +228,7 @@
     "import jax\n",
     "\n",
     "# Generate random noise to update the parameters\n",
-    "noise = (\n",
+    "noise = jax.numpy.exp(\n",
     "    jax.random.normal(\n",
     "        key=jax.random.PRNGKey(0), shape=jax_problem.parameters.shape\n",
     "    )\n",
@@ -260,7 +260,7 @@
    "outputs": [],
    "source": [
     "# Update the parameters and create a new JAXProblem instance\n",
-    "jax_problem = jax_problem.update_parameters(jax_problem.parameters + noise)\n",
+    "jax_problem = jax_problem.update_parameters(jax_problem.parameters * noise)\n",
     "\n",
     "# Run simulations with the updated parameters\n",
     "llh, results = run_simulations(jax_problem)\n",

python/sdist/amici/_symbolic/de_model.py

@@ -27,6 +27,7 @@
     ObservableTransformation,
     _default_simplify,
     amici_time_symbol,
+    symbol_with_assumptions,
     toposort_symbols,
     unique_preserve_order,
 )
@@ -2499,26 +2500,67 @@ def _process_hybridization(self, hybridization: dict) -> None:
         added_expressions = False
         orig_obs = tuple([s.get_sym() for s in self._observables])
         for net_id, net in hybridization.items():
-            if net["static"]:
+            if net["pre_initialization"]:
                 # do not integrate into ODEs, handle in amici.sim.jax.petab
                 continue
-            inputs = [
-                comp
-                for comp in self._components
-                if str(comp.get_sym()) in net["input_vars"]
-            ]
-            # sort inputs by order in input_vars
-            inputs = sorted(
-                inputs,
-                key=lambda comp: net["input_vars"].index(str(comp.get_sym())),
-            )
+            comp_by_sym = {comp.get_id(): comp for comp in self._components}
+            sym_locals = {s: comp.get_sym() for s, comp in comp_by_sym.items()}
+
+            inputs = []
+            unresolved_vars = []
+            for input_var in net["input_vars"]:
+                if input_var in comp_by_sym:
+                    inputs.append(comp_by_sym[input_var])
+                else:
+                    try:
+                        expr = sp.sympify(input_var, locals=sym_locals)
+                    except (sp.SympifyError, Exception):
+                        unresolved_vars.append(input_var)
+                        continue
+
+                    if {str(s) for s in expr.free_symbols} - set(comp_by_sym):
+                        unresolved_vars.append(input_var)
+                        continue
+
+                    expr_sym = symbol_with_assumptions(
+                        f"_nn_{net_id}_input_{len(inputs)}"
+                    )
+                    new_expr_comp = Expression(
+                        symbol=expr_sym,
+                        name=f"{net_id}_input_{len(inputs)}",
+                        value=expr,
+                    )
+                    self.add_component(new_expr_comp)
+                    added_expressions = True
+                    inputs.append(new_expr_comp)
+
             if len(inputs) != len(net["input_vars"]):
-                found_vars = {str(comp.get_sym()) for comp in inputs}
-                missing_vars = set(net["input_vars"]) - found_vars
-                raise ValueError(
-                    f"Could not find all input variables for neural network {net_id}. "
-                    f"Missing variables: {sorted(missing_vars)}"
-                )
+                missing_vars = set(unresolved_vars)
+                if missing_vars == {"array"}:
+                    array_inputs = net.get("array_inputs", {})
+                    petab_ids = list(array_inputs.keys())
+                    for i, input_var in enumerate(net["input_vars"]):
+                        if input_var == "array":
+                            if not petab_ids:
+                                raise ValueError(
+                                    f"Array input specified for {net_id} but no "
+                                    f"array_inputs info provided in hybridization."
+                                )
+                            petab_id = petab_ids.pop(0)
+                            array_sym = symbol_with_assumptions(
+                                f"_nn_array_{petab_id}"
+                            )
+                            array_comp = Expression(
+                                symbol=array_sym,
+                                name=f"{net_id}_array_{petab_id}",
+                                value=sp.Integer(0),
+                            )
+                            inputs.insert(i, array_comp)
+                else:
+                    raise ValueError(
+                        f"Could not find all input variables for neural network {net_id}. "
+                        f"Missing variables: {sorted(missing_vars)}"
+                    )
             for inp in inputs:
                 if isinstance(
                     inp,
@@ -2547,12 +2589,13 @@ def _process_hybridization(self, hybridization: dict) -> None:
                     f"Could not find all output variables for neural network {net_id}. "
                     f"Missing variables: {sorted(missing_vars)}"
                 )
-
             for out_var, parts in outputs.items():
                 comp = parts["comp"]
                 # remove output from model components
                 if isinstance(comp, FreeParameter):
                     self._free_parameters.remove(comp)
+                elif isinstance(comp, FixedParameter):
+                    self._fixed_parameters.remove(comp)
                 elif isinstance(comp, Expression):
                     self._expressions.remove(comp)
                 elif isinstance(comp, DifferentialState):
@@ -2586,7 +2629,7 @@ def _process_hybridization(self, hybridization: dict) -> None:
                     added_expressions = True
 
             observables = {
-                ob_var: {"comp": comp, "ind": net["observable_vars"][ob_var]}
+                ob_var: {"comp": comp, **net["observable_vars"][ob_var]}
                 for comp in self._components
                 if (ob_var := str(comp.get_sym())) in net["observable_vars"]
                 # # TODO: SYNTAX NEEDS to CHANGE
@@ -2609,9 +2652,19 @@ def _process_hybridization(self, hybridization: dict) -> None:
                     raise ValueError(
                         f"{comp.get_name()} ({type(comp)}) is not an observable."
                     )
-                out_val = sp.Function(net_id)(
-                    *[input.get_sym() for input in inputs], parts["ind"]
+                nn_call = sp.Function(net_id)(
+                    *[input.get_sym() for input in inputs], parts["index"]
                 )
+                formula = parts["formula"]
+                petab_id = parts["petab_id"]
+                if formula == petab_id:
+                    out_val = nn_call
+                else:
+                    from petab.math.sympify import sympify_petab
+
+                    out_val = sympify_petab(formula).subs(
+                        symbol_with_assumptions(petab_id), nn_call
+                    )
                 # add to the model
                 self.add_component(
                     Observable(

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

@@ -21,6 +21,8 @@ def __init__(self):
         self.jax_py_file = Path(__file__).resolve()
         self.nns = {TPL_NETS}
         self.parameters = TPL_ALL_P_VALUES
+        TPL_ARRAY_INPUTS_INIT
+        self._array_input_index = jnp.int32(0)
         super().__init__()
 
     def _xdot(self, t, x, args):

python/sdist/amici/exporters/jax/jaxcodeprinter.py

@@ -44,11 +44,44 @@ def _print_Mul(self, expr: sp.Expr) -> str:
         return f"safe_div({self.doprint(numer)}, {self.doprint(denom)})"
 
     def _print_Function(self, expr: sp.Expr) -> str:
-        if isinstance(expr.func, UndefinedFunction):
-            return f"self.nns['{expr.func.__name__}'].forward(jnp.array([{', '.join(self.doprint(a) for a in expr.args[:-1])}]))[{expr.args[-1]}]"
-        else:
+        if not isinstance(expr.func, UndefinedFunction):
             return super()._print_Function(expr)
 
+        nn_name = expr.func.__name__
+        input_args = expr.args[:-1]
+        output_idx = expr.args[-1]
+        forward_arg = self._build_nn_forward_arg(input_args)
+        return f"self.nns['{nn_name}'].forward({forward_arg})[{output_idx}]"
+
+    def _build_nn_forward_arg(self, input_args) -> str:
+        array_prefix = "_nn_array_"
+        if not any(
+            hasattr(a, "name") and a.name.startswith(array_prefix)
+            for a in input_args
+        ):
+            return f"jnp.array([{', '.join(self.doprint(a) for a in input_args)}])"
+
+        groups = self._group_nn_inputs(input_args, array_prefix)
+        return groups[0] if len(groups) == 1 else f"[{', '.join(groups)}]"
+
+    def _group_nn_inputs(self, input_args, array_prefix: str) -> list[str]:
+        groups = []
+        scalar_group = []
+        for a in input_args:
+            if hasattr(a, "name") and a.name.startswith(array_prefix):
+                if scalar_group:
+                    groups.append(f"jnp.array([{', '.join(scalar_group)}])")
+                    scalar_group = []
+                petab_id = a.name[len(array_prefix) :]
+                groups.append(
+                    f"self._array_inputs['{petab_id}'][self._array_input_index]"
+                )
+            else:
+                scalar_group.append(self.doprint(a))
+        if scalar_group:
+            groups.append(f"jnp.array([{', '.join(scalar_group)}])")
+        return groups
+
     def _print_Max(self, expr: sp.Expr) -> str:
         """
         Print the max function, replacing it with jnp.max.

python/sdist/amici/exporters/jax/nn.py

@@ -136,7 +136,7 @@ def generate_equinox(
     filename.parent.mkdir(parents=True, exist_ok=True)
 
     apply_template(
-        Path(amiciModulePath) / "jax" / "nn.template.py",
+        Path(amiciModulePath) / "exporters" / "jax" / "nn.template.py",
         filename,
         tpl_data,
     )
@@ -184,7 +184,7 @@ def _generate_layer(layer: "Layer", indent: int, ilayer: int) -> str:  # noqa: F
     layer_map = {
         "Dropout1d": "eqx.nn.Dropout",
         "Dropout2d": "eqx.nn.Dropout",
-        "Flatten": "amici.export.jax.Flatten",
+        "Flatten": "amici.exporters.jax.Flatten",
     }
 
     # mapping of keyword argument names in sciml yaml format to equinox/custom amici implementations
@@ -320,9 +320,9 @@ def _process_activation_call(node: "Node") -> str:  # noqa: F821
         "hardtanh": "jax.nn.hard_tanh",
         "hardsigmoid": "jax.nn.hard_sigmoid",
         "hardswish": "jax.nn.hard_swish",
-        "tanhshrink": "amici.export.jax.tanhshrink",
+        "tanhshrink": "amici.exporters.jax.tanhshrink",
         "softsign": "jax.nn.soft_sign",
-        "cat": "amici.export.jax.cat",
+        "cat": "amici.exporters.jax.cat",
     }
 
     # Validate hardtanh parameters

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

@@ -1,7 +1,10 @@
 # ruff: noqa: F401, F821, F841
 import equinox as eqx
+import jax
 import jax.random as jr
 
+import amici
+
 
 class TPL_MODEL_ID(eqx.Module):
     layers: dict

python/sdist/amici/exporters/jax/ode_export.py

@@ -16,6 +16,7 @@
 import os
 from pathlib import Path
 
+import numpy as np
 import sympy as sp
 
 from amici._symbolic.de_model import DEModel
@@ -247,6 +248,8 @@ def _generate_jax_code(self) -> None:
         )
         subs = subs_heaviside | subs_observables
 
+        array_inputs_init = self._generate_array_inputs_init()
+
         tpl_data = {
             # assign named variable using corresponding algebraic formula (function body)
             **_jax_variable_equations(
@@ -298,6 +301,7 @@ def _generate_jax_code(self) -> None:
                 f'"{net}": {net}.net(jr.PRNGKey(0))'
                 for net in self.hybridization.keys()
             ),
+            "ARRAY_INPUTS_INIT": array_inputs_init,
         }
 
         apply_template(
@@ -306,6 +310,42 @@ def _generate_jax_code(self) -> None:
             tpl_data,
         )
 
+    def _generate_array_inputs_init(self) -> str:
+        """Generate code to initialize array inputs from HDF5 data.
+
+        Reads array data from HDF5 files at code-generation time and
+        embeds them as jnp.array constants in the generated model code.
+
+        :return:
+            Python code string for initializing self._array_inputs
+        """
+        array_entries = []
+        for _, net in self.hybridization.items():
+            if net.get("pre_initialization", False):
+                continue
+            for petab_id, hdf5_path in net.get("array_inputs", {}).items():
+                import h5py
+
+                with h5py.File(hdf5_path, "r") as f:
+                    group = f["inputs"][petab_id]
+                    keys = sorted(group.keys())
+                    # Stack with experiment dimension first
+                    data = np.stack([group[k][:] for k in keys])
+                    rows = []
+                    for row in data:
+                        rows.append(
+                            "[" + ", ".join(str(float(v)) for v in row) + "]"
+                        )
+                    data_str = ", ".join(rows)
+
+                array_entries.append(f"'{petab_id}': jnp.array([{data_str}])")
+
+        if not array_entries:
+            return "self._array_inputs = {}"
+
+        entries_str = ", ".join(array_entries)
+        return f"self._array_inputs = {{{entries_str}}}"
+
     def _get_all_p_syms(self) -> list[sp.Symbol]:
         return list(self.model.sym("p")) + list(self.model.sym("k"))