AMICI-dev · BSnelling · Jun 16, 2026 · Jun 8, 2026 · Jun 8, 2026 · Jun 8, 2026
@@ -54,7 +54,7 @@ jobs:
     - name: Install sdist
       working-directory: python/sdist
       shell: bash
-      run: pip install -v $(ls -t dist/amici-*.tar.gz | head -1)[petab,test]
+      run: pip install -v $(ls -t dist/amici-*.tar.gz | head -1)[petab,test,jax]
 
     - run: python -m amici
 

@@ -1,4 +1,5 @@
 from pathlib import Path
+from shutil import copyfile
 
 import equinox as eqx
 import jax
@@ -374,15 +375,16 @@ def cat(tensors, axis: int = 0):
 
 
 def generate_equinox(
-    nn_model: "NNModel",  # noqa: F821
+    nn_model: "NNModel | Path | str",  # noqa: F821
     filename: Path | str,
     frozen_layers: dict[str, bool] | None = None,
 ) -> None:
     """
-    Generate Equinox model file from petab_sciml neural network object.
+    Generate Equinox model file from petab_sciml neural network object or copy from
+    existing file if a path is provided.
 
     :param nn_model:
-        Neural network model in petab_sciml format
+        Neural network model in petab_sciml format or path to existing Equinox model file
     :param filename:
         output filename for generated Equinox model
     :param frozen_layers:
@@ -391,6 +393,15 @@ def generate_equinox(
     # TODO: move to top level import and replace forward type definitions
     from petab_sciml import Layer
 
+    filename.parent.mkdir(parents=True, exist_ok=True)
+
+    if isinstance(nn_model, str):
+        nn_model = Path(nn_model)
+
+    if isinstance(nn_model, Path):
+        copyfile(nn_model, filename)
+        return
+
     if frozen_layers is None:
         frozen_layers = {}
 
@@ -453,8 +464,6 @@ def generate_equinox(
         "N_LAYERS": len(nn_model.layers),
     }
 
-    filename.parent.mkdir(parents=True, exist_ok=True)
-
     apply_template(
         Path(amiciModulePath) / "exporters" / "jax" / "nn.template.py",
         filename,

@@ -685,10 +685,14 @@ def _build_hybridization(self) -> dict[str, dict]:
                 if mid.split(".")[1].startswith("output")
             ]
 
+            model = (
+                Path(net_config["location"]).resolve()
+                if net_config["format"] == "equinox"
+                else NNModelStandard.load_data(Path(net_config["location"]))
+            )
+
             hybridization[net_id] = {
-                "model": NNModelStandard.load_data(
-                    Path(net_config["location"])
-                ),
+                "model": model,
                 "input_vars": [
                     input_hybridization[petab_id]
                     for petab_id, _ in input_mappings

@@ -0,0 +1,2 @@
+experimentId	time	conditionId
+e1	0.0
@@ -0,0 +1,4 @@
+targetId	targetValue
+net1_input1	prey
+net1_input2	predator
+gamma	net1_output1
@@ -0,0 +1,118 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<sbml xmlns="http://www.sbml.org/sbml/level2/version4" level="2" version="4">
+  <model id="lv_ude" name="lv_ude">
+    <notes>
+      <body xmlns="http://www.w3.org/1999/xhtml">
+        <pre>PEtab implementation of the simple model</pre>
+      </body>
+    </notes>
+    <annotation>
+      <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#" xmlns:dc="http://purl.org/dc/elements/1.1/" xmlns:dcterms="http://purl.org/dc/terms/" xmlns:vCard="http://www.w3.org/2001/vcard-rdf/3.0#" xmlns:bqbiol="http://biomodels.net/biology-qualifiers/" xmlns:bqmodel="http://biomodels.net/model-qualifiers/">
+        <rdf:Description rdf:about="#arFramework3">
+          <dc:creator>
+            <rdf:Bag>
+              <rdf:li rdf:parseType="Resource">
+                <vCard:N rdf:parseType="Resource">
+                  <vCard:Family>Ognissanti</vCard:Family>
+                  <vCard:Given>Damiano</vCard:Given>
+                </vCard:N>
+              </rdf:li>
+            </rdf:Bag>
+          </dc:creator>
+          <dcterms:created rdf:parseType="Resource">
+            <dcterms:W3CDTF>2022-08-19T11:46:48Z</dcterms:W3CDTF>
+          </dcterms:created>
+          <dcterms:modified rdf:parseType="Resource">
+            <dcterms:W3CDTF>2022-08-19T11:46:48Z</dcterms:W3CDTF>
+          </dcterms:modified>
+          <bqbiol:isDescribedBy>
+            <rdf:Bag>
+              <rdf:li rdf:resource=""/>
+            </rdf:Bag>
+          </bqbiol:isDescribedBy>
+        </rdf:Description>
+      </rdf:RDF>
+    </annotation>
+    <listOfUnitDefinitions>
+      <unitDefinition id="time" name="time">
+        <listOfUnits>
+          <unit kind="second" exponent="1" scale="0" multiplier="60"/>
+        </listOfUnits>
+      </unitDefinition>
+    </listOfUnitDefinitions>
+    <listOfCompartments>
+      <compartment metaid="metaid_0" id="default" size="1" constant="true"/>
+    </listOfCompartments>
+    <listOfSpecies>
+      <species id="prey" name="prey" compartment="default" initialConcentration="0.44249296" hasOnlySubstanceUnits="true" boundaryCondition="false" constant="false"/>
+      <species id="predator" name="predator" compartment="default" initialConcentration="4.6280594" hasOnlySubstanceUnits="true" boundaryCondition="false" constant="false"/>
+    </listOfSpecies>
+    <listOfParameters>
+      <parameter id="alpha" value="1.3" constant="true"/>
+      <parameter id="beta" value="0.9" constant="true"/>
+      <parameter id="gamma" value="0.8" constant="true"/>
+      <parameter id="delta" value="1.8" constant="true"/>
+    </listOfParameters>
+    <listOfReactions>
+      <reaction id="v1" name="v1" reversible="false">
+        <listOfProducts>
+          <speciesReference species="prey" stoichiometry="1"/>
+        </listOfProducts>
+        <kineticLaw>
+          <math xmlns="http://www.w3.org/1998/Math/MathML">
+            <apply>
+              <times/>
+              <ci> alpha </ci>
+              <ci> prey </ci>
+            </apply>
+          </math>
+        </kineticLaw>
+      </reaction>
+      <reaction id="v2" name="v2" reversible="false">
+        <listOfReactants>
+          <speciesReference species="predator" stoichiometry="1"/>
+        </listOfReactants>
+        <kineticLaw>
+          <math xmlns="http://www.w3.org/1998/Math/MathML">
+            <apply>
+              <times/>
+              <ci> delta </ci>
+              <ci> predator </ci>
+            </apply>
+          </math>
+        </kineticLaw>
+      </reaction>
+      <reaction id="v3" name="v3" reversible="false">
+        <listOfReactants>
+          <speciesReference species="prey" stoichiometry="1"/>
+        </listOfReactants>
+        <listOfModifiers>
+          <modifierSpeciesReference species="predator"/>
+        </listOfModifiers>
+        <kineticLaw>
+          <math xmlns="http://www.w3.org/1998/Math/MathML">
+            <apply>
+              <times/>
+              <ci> beta </ci>
+              <ci> prey </ci>
+              <ci> predator </ci>
+            </apply>
+          </math>
+        </kineticLaw>
+      </reaction>
+      <reaction id="v4" name="v4" reversible="false">
+        <listOfProducts>
+          <speciesReference species="predator" stoichiometry="1"/>
+        </listOfProducts>
+        <kineticLaw>
+          <math xmlns="http://www.w3.org/1998/Math/MathML">
+            <apply>
+              <times/>
+              <ci> gamma </ci>
+            </apply>
+          </math>
+        </kineticLaw>
+      </reaction>
+    </listOfReactions>
+  </model>
+</sbml>
@@ -0,0 +1,5 @@
+petabEntityId	modelEntityId
+net1_input1	net1.inputs[0][0]
+net1_input2	net1.inputs[0][1]
+net1_output1	net1.outputs[0][0]
+net1_ps	net1.parameters
@@ -0,0 +1,21 @@
+observableId	experimentId	measurement	time
+prey_o	e1	0.17301723066954827	1.0
+prey_o	e1	0.48917673783383914	2.0
+prey_o	e1	1.643995531803569	3.0
+prey_o	e1	5.45196278566626	4.0
+prey_o	e1	2.9775220192442062	5.0
+prey_o	e1	0.18166337927167636	6.0
+prey_o	e1	0.3481122259853288	7.0
+prey_o	e1	0.9379191088947554	8.0
+prey_o	e1	3.113240033804055	9.0
+prey_o	e1	8.863933242141234	10.0
+predator_o	e1	0.8474159434934865	1.0
+predator_o	e1	0.2111345157163205	2.0
+predator_o	e1	-0.025053892755649274	3.0
+predator_o	e1	0.12501049397609923	4.0
+predator_o	e1	6.700454554758639	5.0
+predator_o	e1	2.007158288516007	6.0
+predator_o	e1	0.42009248269510124	7.0
+predator_o	e1	0.04803185161440761	8.0
+predator_o	e1	0.12866939374360575	9.0
+predator_o	e1	1.192783778293036	10.0
@@ -0,0 +1,53 @@
+# ruff: noqa: F401, F821, F841
+import amici
+import equinox as eqx
+import jax
+import jax.random as jr
+
+
+class net1(eqx.Module):
+    layers: dict
+    inputs: list[str]
+    outputs: list[str]
+
+    def __init__(self, key):
+        super().__init__()
+        keys = jr.split(key, 3)
+        self.layers = {
+            "layer1": eqx.nn.Linear(
+                in_features=2, out_features=5, use_bias=True, key=keys[0]
+            ),
+            "layer2": eqx.nn.Linear(
+                in_features=5, out_features=5, use_bias=True, key=keys[1]
+            ),
+            "layer3": eqx.nn.Linear(
+                in_features=5, out_features=1, use_bias=True, key=keys[2]
+            ),
+        }
+        self.inputs = ["input0"]
+        self.outputs = ["layer3"]
+
+    def forward(self, input, key=None):
+        net_input = input
+        layer1 = (
+            jax.vmap(self.layers["layer1"])
+            if len(net_input.shape) == 2
+            else self.layers["layer1"]
+        )(net_input)
+        tanh = jax.nn.tanh(layer1)
+        layer2 = (
+            jax.vmap(self.layers["layer2"])
+            if len(tanh.shape) == 2
+            else self.layers["layer2"]
+        )(tanh)
+        tanh_1 = jax.nn.tanh(layer2)
+        layer3 = (
+            jax.vmap(self.layers["layer3"])
+            if len(tanh_1.shape) == 2
+            else self.layers["layer3"]
+        )(tanh_1)
+        output = layer3
+        return output
+
+
+net = net1
@@ -0,0 +1,3 @@
+observableId	observableFormula	noiseFormula	observableTransformation	noiseDistribution
+prey_o	prey	0.05	lin	normal
+predator_o	predator	0.05	lin	normal
@@ -0,0 +1,5 @@
+parameterId	parameterScale	lowerBound	upperBound	nominalValue	estimate
+alpha	lin	0.0	15.0	1.3	true
+delta	lin	0.0	15.0	1.8	true
+beta	lin	0.0	15.0	0.9	true
+net1_ps	lin	-inf	inf	array	true
@@ -0,0 +1,28 @@
+model_files:
+  lv:
+    location: "lv.xml"
+    language: "sbml"
+measurement_files:
+  - "measurements.tsv"
+observable_files:
+  - "observables.tsv"
+format_version: "2.0.0"
+experiment_files:
+  - "experiments.tsv"
+parameter_files:
+  - "parameters.tsv"
+mapping_files:
+  - "mapping.tsv"
+extensions:
+  sciml:
+    array_files:
+      - "net1_ps.hdf5"
+    version: "0.1.0"
+    hybridization_files:
+      - "hybridization.tsv"
+    required: true
+    neural_networks:
+      net1:
+        location: "net1.py"
+        pre_initialization: false
+        format: "equinox"
@@ -0,0 +1,9 @@
+simulation_files:
+  - "simulations.tsv"
+tol_grad: 0.1
+grad_files:
+  mech: "grad_mech.tsv"
+  net1: "grad_net1.hdf5"
+llh: 33.02909543616689
+tol_simulations: 0.001
+tol_llh: 0.001
@@ -3,6 +3,7 @@
 import amici
 import pytest
 
+pytest.importorskip("pysb")
 pytest.importorskip("jax")
 import diffrax
 import jax