From 48f65c211a58d431b89466c5971d3bb4b10f186d Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Mon, 8 Jun 2026 15:08:45 +0100 Subject: [PATCH 01/16] add equinox module import support --- python/sdist/amici/exporters/jax/nn.py | 7 +- .../amici/importers/petab/_petab_importer.py | 10 +- python/tests/test_sciml.py | 162 +++++++++++++++++- tests/sciml/sciml_helpers.py | 104 +++++++++++ tests/sciml/test_sciml.py | 105 +----------- 5 files changed, 276 insertions(+), 112 deletions(-) create mode 100644 tests/sciml/sciml_helpers.py diff --git a/python/sdist/amici/exporters/jax/nn.py b/python/sdist/amici/exporters/jax/nn.py index 576cb6265d..b51522111f 100644 --- a/python/sdist/amici/exporters/jax/nn.py +++ b/python/sdist/amici/exporters/jax/nn.py @@ -1,4 +1,5 @@ from pathlib import Path +from shutil import copyfile import equinox as eqx import jax @@ -374,7 +375,7 @@ def cat(tensors, axis: int = 0): def generate_equinox( - nn_model: "NNModel", # noqa: F821 + nn_model: "NNModel | str", # noqa: F821 filename: Path | str, frozen_layers: dict[str, bool] | None = None, ) -> None: @@ -391,6 +392,10 @@ def generate_equinox( # TODO: move to top level import and replace forward type definitions from petab_sciml import Layer + if isinstance(nn_model, Path): + copyfile(nn_model, filename) + return + if frozen_layers is None: frozen_layers = {} diff --git a/python/sdist/amici/importers/petab/_petab_importer.py b/python/sdist/amici/importers/petab/_petab_importer.py index a9e1c20c4c..b78905f794 100644 --- a/python/sdist/amici/importers/petab/_petab_importer.py +++ b/python/sdist/amici/importers/petab/_petab_importer.py @@ -685,10 +685,14 @@ def _build_hybridization(self) -> dict[str, dict]: if mid.split(".")[1].startswith("output") ] + model = ( + Path(net_config["location"]) + 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 diff --git a/python/tests/test_sciml.py b/python/tests/test_sciml.py index e9f97aaec2..a44d82dfa4 100644 --- a/python/tests/test_sciml.py +++ b/python/tests/test_sciml.py @@ -1,12 +1,16 @@ """Tests for SBML/SciML functionality, including JAX neural network code generation.""" +import importlib.util +import os +from contextlib import contextmanager +from pathlib import Path from unittest.mock import Mock -import pytest - -pytest.importorskip("jax") -pytest.importorskip("equinox") - +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy as np +import pandas as pd import pytest from amici.exporters.jax.nn import ( _format_function_call, @@ -15,7 +19,37 @@ _process_activation_call, _process_layer_call, ) +from amici.importers.petab import * +from amici.importers.petab import PetabImporter from amici.importers.utils import symbol_with_assumptions +from amici.sim.jax import petab_simulate, run_simulations +from petab import v2 +from yaml import safe_load + +# TODO: remove once sciml linter is released in libpetab +_sciml_helpers_spec = importlib.util.spec_from_file_location( + "sciml_helpers", + Path(__file__).parents[2] / "tests" / "sciml" / "sciml_helpers.py", +) +_sciml_helpers_mod = importlib.util.module_from_spec(_sciml_helpers_spec) +_sciml_helpers_spec.loader.exec_module(_sciml_helpers_mod) +_v2_sciml_problem_helper = _sciml_helpers_mod._v2_sciml_problem_helper + + +@contextmanager +def change_directory(destination): + # Save the current working directory + original_directory = os.getcwd() + try: + # Change to the new directory + os.chdir(destination) + yield + finally: + # Change back to the original directory + os.chdir(original_directory) + + +jax.config.update("jax_enable_x64", True) class TestFormatFunctionCall: @@ -789,3 +823,121 @@ def test_valid_hybridization_no_error(self, mock_de_model): # Should not raise any errors mock_de_model._process_hybridization(hybridization) + + +class TestEquinoxImport: + """Test that an Equinox model can be imported and used in a PEtab problem.""" + + def test_equinox_model_import(self): + """Test that the Equinox model is correctly imported and can be called.""" + + test_dir = ( + Path(__file__).parent / "sciml_test_problems" / "equinox_import" + ) + with open(test_dir / "petab" / "problem.yaml") as f: + petab_yaml = safe_load(f) + + with open(test_dir / "solutions.yaml") as f: + solutions = safe_load(f) + + with change_directory(test_dir / "petab"): + petab_problem = _v2_sciml_problem_helper( + petab_yaml, str(test_dir / "petab") + ) + + pi = PetabImporter( + petab_problem=petab_problem, + module_name="sciml_test", + compile_=True, + jax=True, + validate=False, # And again...around "array" in parameters table + ) + + jax_problem = pi.create_simulator( + force_import=True, + ) + + llh, _ = run_simulations(jax_problem) + + np.testing.assert_allclose( + llh, + solutions["llh"], + atol=solutions["tol_llh"], + rtol=solutions["tol_llh"], + ) + simulations = pd.concat( + [ + pd.read_csv(test_dir / simulation, sep="\t") + for simulation in solutions["simulation_files"] + ] + ) + + # simulations + sort_by = [v2.C.OBSERVABLE_ID, v2.C.TIME, v2.C.EXPERIMENT_ID] + actual = petab_simulate(jax_problem).sort_values(by=sort_by) + expected = simulations.sort_values(by=sort_by) + np.testing.assert_allclose( + actual[v2.C.SIMULATION].values, + expected[v2.C.SIMULATION].values, + atol=solutions["tol_simulations"], + rtol=solutions["tol_simulations"], + ) + + +def _normal_logpdf(x: jnp.ndarray, mean: float, std: float) -> jnp.ndarray: + var = std**2 + return jnp.sum( + -0.5 * jnp.log(2.0 * jnp.pi * var) - 0.5 * ((x - mean) ** 2) / var + ) + + +def _uniform_logpdf(x: jnp.ndarray, low: float, high: float) -> jnp.ndarray: + return jnp.sum( + jnp.where( + (x >= low) & (x <= high), + -jnp.log(high - low), + -jnp.inf, + ) + ) + + +def _tree_array_lognormprior(tree, mean: float, std: float) -> jnp.ndarray: + arrays, _ = eqx.partition(tree, eqx.is_inexact_array) + leaves = jax.tree_util.tree_leaves(arrays) + + total = jnp.array(0.0) + for leaf in leaves: + if leaf is not None: + total = total + _normal_logpdf(leaf, mean, std) + return total + + +def _tree_array_loguniformprior(tree, low: float, high: float) -> jnp.ndarray: + arrays, _ = eqx.partition(tree, eqx.is_inexact_array) + leaves = jax.tree_util.tree_leaves(arrays) + + total = jnp.array(0.0) + for leaf in leaves: + if leaf is not None: + total = total + _uniform_logpdf(leaf, low, high) + return total + + +def _model_logprior( + model, layer1_bias_std=1.0, layer1_weight_std=1.0 +) -> jnp.ndarray: + mech = model.parameters + layer1_bias = model.model.nns["net1"].layers["layer1"].bias + layer1_weight = model.model.nns["net1"].layers["layer1"].weight + rest = eqx.tree_at( + lambda m: m["net1"].layers["layer1"], model.model.nns, replace=None + ) + + return ( + _tree_array_loguniformprior(mech, low=0.0, high=15.0) + + _tree_array_lognormprior(layer1_bias, mean=0.0, std=layer1_bias_std) + + _tree_array_lognormprior( + layer1_weight, mean=0.0, std=layer1_weight_std + ) + + _tree_array_lognormprior(rest, mean=0.0, std=1.0) + ) diff --git a/tests/sciml/sciml_helpers.py b/tests/sciml/sciml_helpers.py new file mode 100644 index 0000000000..c15fc34545 --- /dev/null +++ b/tests/sciml/sciml_helpers.py @@ -0,0 +1,104 @@ +import pandas as pd +from amici.sim.jax.petab import _try_float +from petab import v1, v2 + + +def _process_prior_params(prior_params): + if isinstance(prior_params, float): + return prior_params + else: + return [float(param) for param in prior_params.split(";")] + + +def _v2_sciml_problem_helper(yaml_config, base_path): + config = v2.ProblemConfig(**yaml_config) + + parameter_tables = [] + for f in config.parameter_files: + df = pd.read_csv(f, sep="\t") + df.nominalValue = df.nominalValue.apply(_try_float) + if "priorParameters" in df.columns: + df.priorParameters = df.priorParameters.apply( + _process_prior_params + ) + parameters = [ + v2.Parameter.model_construct(**row.to_dict()) + for _, row in df.reset_index().iterrows() + ] + parameter_tables.append(v2.ParameterTable(elements=parameters)) + + models = [ + v1.models.model.model_factory( + model_info.location, + base_path=base_path, + model_language=model_info.language, + model_id=model_id, + ) + for model_id, model_info in (config.model_files or {}).items() + ] + + measurement_tables = ( + [ + v2.MeasurementTable.from_tsv(f, base_path) + for f in config.measurement_files + ] + if config.measurement_files + else None + ) + + experiment_tables = ( + [ + v2.ExperimentTable.from_tsv(f, base_path) + for f in config.experiment_files + ] + if config.experiment_files + else None + ) + + condition_tables = ( + [ + v2.ConditionTable.from_tsv(f, base_path) + for f in config.condition_files + ] + if config.condition_files + else None + ) + + if condition_tables is None: + cond_ids = [ + cid + for exp_table in experiment_tables + for exp in exp_table.elements + for p in exp.periods + for cid in p.condition_ids + ] + condition_tables = [ + v2.ConditionTable(elements=[v2.Condition(id=cid, changes=[])]) + for cid in set(cond_ids) + ] + + observable_tables = ( + [ + v2.ObservableTable.from_tsv(f, base_path) + for f in config.observable_files + ] + if config.observable_files + else None + ) + + mapping_tables = ( + [v2.MappingTable.from_tsv(f, base_path) for f in config.mapping_files] + if config.mapping_files + else None + ) + + return v2.Problem( + config=config, + models=models, + condition_tables=condition_tables, + experiment_tables=experiment_tables, + observable_tables=observable_tables, + measurement_tables=measurement_tables, + parameter_tables=parameter_tables, + mapping_tables=mapping_tables, + ) diff --git a/tests/sciml/test_sciml.py b/tests/sciml/test_sciml.py index e3666fd79a..1691dcc6eb 100644 --- a/tests/sciml/test_sciml.py +++ b/tests/sciml/test_sciml.py @@ -15,9 +15,9 @@ from amici.exporters.jax import generate_equinox from amici.importers.petab import * from amici.sim.jax import petab_simulate, run_simulations -from amici.sim.jax.petab import _try_float -from petab import v1, v2 +from petab import v2 from petab_sciml import NNModelStandard +from sciml_helpers import _v2_sciml_problem_helper from yaml import safe_load @@ -404,107 +404,6 @@ def test_sciml_problem_import(test): ) -def _v2_sciml_problem_helper(yaml_config, base_path): - config = v2.ProblemConfig(**yaml_config) - - parameter_tables = [] - for f in config.parameter_files: - df = pd.read_csv(f, sep="\t") - df.nominalValue = df.nominalValue.apply(_try_float) - if "priorParameters" in df.columns: - df.priorParameters = df.priorParameters.apply( - _process_prior_params - ) - parameters = [ - v2.Parameter.model_construct(**row.to_dict()) - for _, row in df.reset_index().iterrows() - ] - parameter_tables.append(v2.ParameterTable(elements=parameters)) - - models = [ - v1.models.model.model_factory( - model_info.location, - base_path=base_path, - model_language=model_info.language, - model_id=model_id, - ) - for model_id, model_info in (config.model_files or {}).items() - ] - - measurement_tables = ( - [ - v2.MeasurementTable.from_tsv(f, base_path) - for f in config.measurement_files - ] - if config.measurement_files - else None - ) - - experiment_tables = ( - [ - v2.ExperimentTable.from_tsv(f, base_path) - for f in config.experiment_files - ] - if config.experiment_files - else None - ) - - condition_tables = ( - [ - v2.ConditionTable.from_tsv(f, base_path) - for f in config.condition_files - ] - if config.condition_files - else None - ) - - if condition_tables is None: - cond_ids = [ - cid - for exp_table in experiment_tables - for exp in exp_table.elements - for p in exp.periods - for cid in p.condition_ids - ] - condition_tables = [ - v2.ConditionTable(elements=[v2.Condition(id=cid, changes=[])]) - for cid in set(cond_ids) - ] - - observable_tables = ( - [ - v2.ObservableTable.from_tsv(f, base_path) - for f in config.observable_files - ] - if config.observable_files - else None - ) - - mapping_tables = ( - [v2.MappingTable.from_tsv(f, base_path) for f in config.mapping_files] - if config.mapping_files - else None - ) - - return v2.Problem( - config=config, - models=models, - condition_tables=condition_tables, - experiment_tables=experiment_tables, - observable_tables=observable_tables, - measurement_tables=measurement_tables, - parameter_tables=parameter_tables, - mapping_tables=mapping_tables, - ) - - -def _process_prior_params(prior_params): - if isinstance(prior_params, float): - return prior_params - else: - return [float(param) for param in prior_params.split(";")] - - def _normal_logpdf(x: jnp.ndarray, mean: float, std: float) -> jnp.ndarray: var = std**2 return jnp.sum( From 811ab4b1689d33cab1a795f3af99bee9d9815034 Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Mon, 8 Jun 2026 15:45:59 +0100 Subject: [PATCH 02/16] make tests more unit-ful --- python/tests/test_sciml.py | 112 +++++-------------------------------- 1 file changed, 13 insertions(+), 99 deletions(-) diff --git a/python/tests/test_sciml.py b/python/tests/test_sciml.py index a44d82dfa4..0e6e9f93c9 100644 --- a/python/tests/test_sciml.py +++ b/python/tests/test_sciml.py @@ -6,24 +6,22 @@ from pathlib import Path from unittest.mock import Mock -import equinox as eqx -import jax -import jax.numpy as jnp -import numpy as np -import pandas as pd import pytest + +pytest.importorskip("jax") +pytest.importorskip("equinox") + from amici.exporters.jax.nn import ( _format_function_call, _generate_forward, _generate_layer, _process_activation_call, _process_layer_call, + generate_equinox, ) from amici.importers.petab import * from amici.importers.petab import PetabImporter from amici.importers.utils import symbol_with_assumptions -from amici.sim.jax import petab_simulate, run_simulations -from petab import v2 from yaml import safe_load # TODO: remove once sciml linter is released in libpetab @@ -49,9 +47,6 @@ def change_directory(destination): os.chdir(original_directory) -jax.config.update("jax_enable_x64", True) - - class TestFormatFunctionCall: """Test the utility function for formatting function calls.""" @@ -837,9 +832,6 @@ def test_equinox_model_import(self): with open(test_dir / "petab" / "problem.yaml") as f: petab_yaml = safe_load(f) - with open(test_dir / "solutions.yaml") as f: - solutions = safe_load(f) - with change_directory(test_dir / "petab"): petab_problem = _v2_sciml_problem_helper( petab_yaml, str(test_dir / "petab") @@ -852,92 +844,14 @@ def test_equinox_model_import(self): jax=True, validate=False, # And again...around "array" in parameters table ) + hybridization = pi._build_hybridization() - jax_problem = pi.create_simulator( - force_import=True, - ) - - llh, _ = run_simulations(jax_problem) - - np.testing.assert_allclose( - llh, - solutions["llh"], - atol=solutions["tol_llh"], - rtol=solutions["tol_llh"], - ) - simulations = pd.concat( - [ - pd.read_csv(test_dir / simulation, sep="\t") - for simulation in solutions["simulation_files"] - ] - ) - - # simulations - sort_by = [v2.C.OBSERVABLE_ID, v2.C.TIME, v2.C.EXPERIMENT_ID] - actual = petab_simulate(jax_problem).sort_values(by=sort_by) - expected = simulations.sort_values(by=sort_by) - np.testing.assert_allclose( - actual[v2.C.SIMULATION].values, - expected[v2.C.SIMULATION].values, - atol=solutions["tol_simulations"], - rtol=solutions["tol_simulations"], - ) - - -def _normal_logpdf(x: jnp.ndarray, mean: float, std: float) -> jnp.ndarray: - var = std**2 - return jnp.sum( - -0.5 * jnp.log(2.0 * jnp.pi * var) - 0.5 * ((x - mean) ** 2) / var - ) + assert isinstance(hybridization["net1"]["model"], Path) + generate_equinox( + hybridization["net1"]["model"], + pi.output_dir / "net1.py", + hybridization["net1"]["frozen_layers"], + ) -def _uniform_logpdf(x: jnp.ndarray, low: float, high: float) -> jnp.ndarray: - return jnp.sum( - jnp.where( - (x >= low) & (x <= high), - -jnp.log(high - low), - -jnp.inf, - ) - ) - - -def _tree_array_lognormprior(tree, mean: float, std: float) -> jnp.ndarray: - arrays, _ = eqx.partition(tree, eqx.is_inexact_array) - leaves = jax.tree_util.tree_leaves(arrays) - - total = jnp.array(0.0) - for leaf in leaves: - if leaf is not None: - total = total + _normal_logpdf(leaf, mean, std) - return total - - -def _tree_array_loguniformprior(tree, low: float, high: float) -> jnp.ndarray: - arrays, _ = eqx.partition(tree, eqx.is_inexact_array) - leaves = jax.tree_util.tree_leaves(arrays) - - total = jnp.array(0.0) - for leaf in leaves: - if leaf is not None: - total = total + _uniform_logpdf(leaf, low, high) - return total - - -def _model_logprior( - model, layer1_bias_std=1.0, layer1_weight_std=1.0 -) -> jnp.ndarray: - mech = model.parameters - layer1_bias = model.model.nns["net1"].layers["layer1"].bias - layer1_weight = model.model.nns["net1"].layers["layer1"].weight - rest = eqx.tree_at( - lambda m: m["net1"].layers["layer1"], model.model.nns, replace=None - ) - - return ( - _tree_array_loguniformprior(mech, low=0.0, high=15.0) - + _tree_array_lognormprior(layer1_bias, mean=0.0, std=layer1_bias_std) - + _tree_array_lognormprior( - layer1_weight, mean=0.0, std=layer1_weight_std - ) - + _tree_array_lognormprior(rest, mean=0.0, std=1.0) - ) + assert os.path.isfile(pi.output_dir / "net1.py") From b09d74fa70d7150b0aa2e690dc50516a60051738 Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Mon, 8 Jun 2026 16:13:27 +0100 Subject: [PATCH 03/16] add test equinox import problem --- .../equinox_import/petab/experiments.tsv | 2 + .../equinox_import/petab/hybridization.tsv | 4 + .../equinox_import/petab/lv.xml | 118 ++++++++++++++++++ .../equinox_import/petab/mapping.tsv | 5 + .../equinox_import/petab/measurements.tsv | 21 ++++ .../equinox_import/petab/net1.py | 53 ++++++++ .../equinox_import/petab/net1_ps.hdf5 | Bin 0 -> 12868 bytes .../equinox_import/petab/observables.tsv | 3 + .../equinox_import/petab/parameters.tsv | 5 + .../equinox_import/petab/problem.yaml | 28 +++++ 10 files changed, 239 insertions(+) create mode 100644 python/tests/sciml_test_problems/equinox_import/petab/experiments.tsv create mode 100644 python/tests/sciml_test_problems/equinox_import/petab/hybridization.tsv create mode 100644 python/tests/sciml_test_problems/equinox_import/petab/lv.xml create mode 100644 python/tests/sciml_test_problems/equinox_import/petab/mapping.tsv create mode 100644 python/tests/sciml_test_problems/equinox_import/petab/measurements.tsv create mode 100644 python/tests/sciml_test_problems/equinox_import/petab/net1.py create mode 100644 python/tests/sciml_test_problems/equinox_import/petab/net1_ps.hdf5 create mode 100644 python/tests/sciml_test_problems/equinox_import/petab/observables.tsv create mode 100644 python/tests/sciml_test_problems/equinox_import/petab/parameters.tsv create mode 100644 python/tests/sciml_test_problems/equinox_import/petab/problem.yaml diff --git a/python/tests/sciml_test_problems/equinox_import/petab/experiments.tsv b/python/tests/sciml_test_problems/equinox_import/petab/experiments.tsv new file mode 100644 index 0000000000..68854689dd --- /dev/null +++ b/python/tests/sciml_test_problems/equinox_import/petab/experiments.tsv @@ -0,0 +1,2 @@ +experimentId time conditionId +e1 0.0 diff --git a/python/tests/sciml_test_problems/equinox_import/petab/hybridization.tsv b/python/tests/sciml_test_problems/equinox_import/petab/hybridization.tsv new file mode 100644 index 0000000000..d9d3e42d60 --- /dev/null +++ b/python/tests/sciml_test_problems/equinox_import/petab/hybridization.tsv @@ -0,0 +1,4 @@ +targetId targetValue +net1_input1 prey +net1_input2 predator +gamma net1_output1 diff --git a/python/tests/sciml_test_problems/equinox_import/petab/lv.xml b/python/tests/sciml_test_problems/equinox_import/petab/lv.xml new file mode 100644 index 0000000000..f573753073 --- /dev/null +++ b/python/tests/sciml_test_problems/equinox_import/petab/lv.xml @@ -0,0 +1,118 @@ + + + + + + 
PEtab implementation of the simple model
+ + + + + + + + + + Ognissanti + Damiano + + + + + + 2022-08-19T11:46:48Z + + + 2022-08-19T11:46:48Z + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + alpha + prey + + + + + + + + + + + + + delta + predator + + + + + + + + + + + + + + + + beta + prey + predator + + + + + + + + + + + + + gamma + + + + + + + diff --git a/python/tests/sciml_test_problems/equinox_import/petab/mapping.tsv b/python/tests/sciml_test_problems/equinox_import/petab/mapping.tsv new file mode 100644 index 0000000000..767a48a23a --- /dev/null +++ b/python/tests/sciml_test_problems/equinox_import/petab/mapping.tsv @@ -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 diff --git a/python/tests/sciml_test_problems/equinox_import/petab/measurements.tsv b/python/tests/sciml_test_problems/equinox_import/petab/measurements.tsv new file mode 100644 index 0000000000..509250fdec --- /dev/null +++ b/python/tests/sciml_test_problems/equinox_import/petab/measurements.tsv @@ -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 diff --git a/python/tests/sciml_test_problems/equinox_import/petab/net1.py b/python/tests/sciml_test_problems/equinox_import/petab/net1.py new file mode 100644 index 0000000000..40b473ba85 --- /dev/null +++ b/python/tests/sciml_test_problems/equinox_import/petab/net1.py @@ -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 diff --git a/python/tests/sciml_test_problems/equinox_import/petab/net1_ps.hdf5 b/python/tests/sciml_test_problems/equinox_import/petab/net1_ps.hdf5 new file mode 100644 index 0000000000000000000000000000000000000000..102fedf43f416d58ed7bdcef20492cdcaf8806bb GIT binary patch literal 12868 zcmeHNdrVVT7(bU{EfwTikPS4U*~CnwOoYWcmz>LZm6$tH_H2`hM@8oE*<1^beA%nRvhvFy&-=r;3^4prP_gm_IO{x#qcGUXfDpch2n* zf#$D^8epPbv-ul5#t%YErTKf5@*rRk>iG+Ni;~C)0v$+;G~Rg)8ZI+kKK{Ua|tM84a!)c5oi^SS7BkKLk;y~ z71t&4IB#R5qrG_gynMa393l7Le|RmGX8_-eal}ZY#MYA|Gm*o|H*nH5OT?ca{Dfbu zK9d)Qg-P_8E27VzaLi#<4DLbv`vWT zcx~?1)1%__3}g7V?e)V)h9+>_vE-Jbh9B`*=Ck9i zXNN^srLk$W@mH~PY;r^K-m79s)8-Ge_sxiZr$u*9zB4bHs`B-{!Zc3KKN)VC`WfrD z8LBe)IsEbb^j|Swj^dkBlfSh;%{g{wb;hS>I`O-+uFs!Jyn)TehB@`F%lJ4y@ZO1$ zK|FDBME}|0NwFcizNqB<5H5QmqFBfq#|!hNy6@|zaeG;3cPf7qKlg22Yj(3s96wOr zkz9BOH}5OdS!0L9#5Ys-N3>rP-%r?;a^%*8SQAyH>WC$<>7y=9dhCoOWkmpxqN0sCJ^DD@! zg3qUBX+A~rH~0v7K2-;*|3&)jSEy^P6|a=rO_9EasOdhlXWdp>VL8~fo@kZUZD<40 zvvr#jVq2}qu_J3qNo!&+ng;*r#5r3Wrdx4*&<=}amk0ur@#w#m> zr!|j|uURTj1D+s*$J0ph)WB5pN_pH{ZuxYZpNB|!dR9+9N9laLU^CI1h(PQ-kOIL( zKm?R{x@(hvb?`@s|Nh#gmzKcO8uyvFrOaZtWF_AnN6yp4&%sU=J-3MA5%L^mD4$~p zv{(NlcVC{UZ$0HE-L|`fFT Date: Tue, 9 Jun 2026 09:37:18 +0100 Subject: [PATCH 04/16] add sciml dep to install script --- scripts/installAmiciSource.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/scripts/installAmiciSource.sh b/scripts/installAmiciSource.sh index 4c416a5788..5e5bd02349 100755 --- a/scripts/installAmiciSource.sh +++ b/scripts/installAmiciSource.sh @@ -40,6 +40,7 @@ python -m pip install --upgrade pip setuptools cmake_build_extension==0.6.0 nump python -m pip install git+https://github.com/pysb/pysb@master # for SPM with compartments python -m pip install git+https://github.com/patrick-kidger/diffrax@main # for events with direction python -m pip install optax # for jax petab notebook +python -m pip install git+https://github.com/petab-dev/petab_sciml.git@main AMICI_BUILD_TEMP="${AMICI_PATH}/python/sdist/build/temp" \ python -m pip install --verbose -e "${AMICI_PATH}/python/sdist[petab,test,vis,jax]" --no-build-isolation deactivate From 92b06757bb8dee5e0e86127c0605a1ee32971840 Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Tue, 9 Jun 2026 11:10:40 +0100 Subject: [PATCH 05/16] use abs path --- python/sdist/amici/importers/petab/_petab_importer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sdist/amici/importers/petab/_petab_importer.py b/python/sdist/amici/importers/petab/_petab_importer.py index b78905f794..bc0c991962 100644 --- a/python/sdist/amici/importers/petab/_petab_importer.py +++ b/python/sdist/amici/importers/petab/_petab_importer.py @@ -686,7 +686,7 @@ def _build_hybridization(self) -> dict[str, dict]: ] model = ( - Path(net_config["location"]) + Path(net_config["location"]).resolve() if net_config["format"] == "equinox" else NNModelStandard.load_data(Path(net_config["location"])) ) From b2b4e4cc4aeaf76597f8ddfda79c8f5fb1dca7d7 Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Tue, 9 Jun 2026 11:46:12 +0100 Subject: [PATCH 06/16] create dir earlier in nn.py --- python/sdist/amici/exporters/jax/nn.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sdist/amici/exporters/jax/nn.py b/python/sdist/amici/exporters/jax/nn.py index b51522111f..0a79a6103a 100644 --- a/python/sdist/amici/exporters/jax/nn.py +++ b/python/sdist/amici/exporters/jax/nn.py @@ -392,6 +392,8 @@ 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, Path): copyfile(nn_model, filename) return @@ -458,8 +460,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, From e7e682f703fc7f44c24acc892221e64d25d76ead Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Mon, 15 Jun 2026 11:29:35 +0100 Subject: [PATCH 07/16] handle strings --- python/sdist/amici/exporters/jax/nn.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/sdist/amici/exporters/jax/nn.py b/python/sdist/amici/exporters/jax/nn.py index 0a79a6103a..1395b66f8b 100644 --- a/python/sdist/amici/exporters/jax/nn.py +++ b/python/sdist/amici/exporters/jax/nn.py @@ -375,7 +375,7 @@ def cat(tensors, axis: int = 0): def generate_equinox( - nn_model: "NNModel | str", # noqa: F821 + nn_model: "NNModel | Path | str", # noqa: F821 filename: Path | str, frozen_layers: dict[str, bool] | None = None, ) -> None: @@ -394,6 +394,9 @@ def generate_equinox( 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 From cd8b2c18464cc973cf4880cf2157fdaca298fb53 Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Mon, 15 Jun 2026 11:45:46 +0100 Subject: [PATCH 08/16] add solutions test --- python/tests/test_sciml.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/python/tests/test_sciml.py b/python/tests/test_sciml.py index 0e6e9f93c9..9bd5620943 100644 --- a/python/tests/test_sciml.py +++ b/python/tests/test_sciml.py @@ -6,9 +6,10 @@ from pathlib import Path from unittest.mock import Mock +import jax +import numpy as np import pytest -pytest.importorskip("jax") pytest.importorskip("equinox") from amici.exporters.jax.nn import ( @@ -22,8 +23,11 @@ from amici.importers.petab import * from amici.importers.petab import PetabImporter from amici.importers.utils import symbol_with_assumptions +from amici.sim.jax import run_simulations from yaml import safe_load +jax.config.update("jax_enable_x64", True) + # TODO: remove once sciml linter is released in libpetab _sciml_helpers_spec = importlib.util.spec_from_file_location( "sciml_helpers", @@ -832,6 +836,9 @@ def test_equinox_model_import(self): with open(test_dir / "petab" / "problem.yaml") as f: petab_yaml = safe_load(f) + with open(test_dir / "solutions.yaml") as f: + solutions = safe_load(f) + with change_directory(test_dir / "petab"): petab_problem = _v2_sciml_problem_helper( petab_yaml, str(test_dir / "petab") @@ -855,3 +862,16 @@ def test_equinox_model_import(self): ) assert os.path.isfile(pi.output_dir / "net1.py") + + jax_problem = pi.create_simulator( + force_import=True, + ) + + llh, _ = run_simulations(jax_problem) + + np.testing.assert_allclose( + llh, + solutions["llh"], + atol=solutions["tol_llh"], + rtol=solutions["tol_llh"], + ) From b55058266acf7c1dc6977e3b091f8d6d1c844981 Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Mon, 15 Jun 2026 13:31:48 +0100 Subject: [PATCH 09/16] change generator to list comp in test --- python/tests/test_conserved_quantities_rref.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/tests/test_conserved_quantities_rref.py b/python/tests/test_conserved_quantities_rref.py index 4a0f4a74db..2d87c68c49 100644 --- a/python/tests/test_conserved_quantities_rref.py +++ b/python/tests/test_conserved_quantities_rref.py @@ -12,8 +12,10 @@ def random_matrix_generator(min_dim, max_dim, count): """Generate random 2D square matrix""" rng = np.random.default_rng(12345) - for rows, cols in rng.integers(min_dim, max_dim, [count, 2]): - yield np.random.rand(rows, cols) + return [ + np.random.rand(rows, cols) + for rows, cols in rng.integers(min_dim, max_dim, [count, 2]) + ] @skip_on_valgrind From a07b6e2843c03b41d28de611884236a8a210ac00 Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Mon, 15 Jun 2026 14:05:36 +0100 Subject: [PATCH 10/16] add solutions file for equinox import test --- .../sciml_test_problems/equinox_import/solutions.yaml | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 python/tests/sciml_test_problems/equinox_import/solutions.yaml diff --git a/python/tests/sciml_test_problems/equinox_import/solutions.yaml b/python/tests/sciml_test_problems/equinox_import/solutions.yaml new file mode 100644 index 0000000000..b1960f75c2 --- /dev/null +++ b/python/tests/sciml_test_problems/equinox_import/solutions.yaml @@ -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 From f3ea906ebcd342745619b31f799e2233da022522 Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Mon, 15 Jun 2026 14:06:03 +0100 Subject: [PATCH 11/16] add jax dep to windows workflow install --- .github/workflows/test_windows.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test_windows.yml b/.github/workflows/test_windows.yml index bb3efb8ce2..f3148dcbdb 100644 --- a/.github/workflows/test_windows.yml +++ b/.github/workflows/test_windows.yml @@ -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 From 08d013580f6071073d29bb5f24acfe674293c349 Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Mon, 15 Jun 2026 14:35:20 +0100 Subject: [PATCH 12/16] skip pysb tests for windows --- python/tests/test_jax.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tests/test_jax.py b/python/tests/test_jax.py index 52426835c3..238e2bc89c 100644 --- a/python/tests/test_jax.py +++ b/python/tests/test_jax.py @@ -3,6 +3,7 @@ import amici import pytest +pytest.importorskip("pysb") pytest.importorskip("jax") import diffrax import jax From 762052ec64d7ce174eb0b3f7693bcd962b8b794f Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Mon, 15 Jun 2026 14:36:23 +0100 Subject: [PATCH 13/16] update generate equinox docstring --- python/sdist/amici/exporters/jax/nn.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sdist/amici/exporters/jax/nn.py b/python/sdist/amici/exporters/jax/nn.py index 1395b66f8b..1370738c03 100644 --- a/python/sdist/amici/exporters/jax/nn.py +++ b/python/sdist/amici/exporters/jax/nn.py @@ -380,10 +380,11 @@ def generate_equinox( 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: From c4ed8be3b41af199638f5a86fa59399a571ddbaa Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Mon, 15 Jun 2026 15:18:55 +0100 Subject: [PATCH 14/16] try to fix the windows workflow again --- python/tests/test_sciml.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tests/test_sciml.py b/python/tests/test_sciml.py index 9bd5620943..fc89177556 100644 --- a/python/tests/test_sciml.py +++ b/python/tests/test_sciml.py @@ -11,6 +11,7 @@ import pytest pytest.importorskip("equinox") +pytest.importorskip("petab_sciml") from amici.exporters.jax.nn import ( _format_function_call, From 22373463f1e252a48101d21abb6c65a5faf3fe15 Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Mon, 15 Jun 2026 16:40:35 +0100 Subject: [PATCH 15/16] increase test coverage --- python/tests/test_sciml.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tests/test_sciml.py b/python/tests/test_sciml.py index fc89177556..1fbf3924ef 100644 --- a/python/tests/test_sciml.py +++ b/python/tests/test_sciml.py @@ -856,8 +856,9 @@ def test_equinox_model_import(self): assert isinstance(hybridization["net1"]["model"], Path) + model_file_str = str(hybridization["net1"]["model"]) generate_equinox( - hybridization["net1"]["model"], + model_file_str, pi.output_dir / "net1.py", hybridization["net1"]["frozen_layers"], ) From 020c7639699b86d2ad5c81212e4b1e5121f2ebea Mon Sep 17 00:00:00 2001 From: Branwen Snelling Date: Tue, 16 Jun 2026 09:48:21 +0100 Subject: [PATCH 16/16] fixup rebase --- python/tests/test_conserved_quantities_rref.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/python/tests/test_conserved_quantities_rref.py b/python/tests/test_conserved_quantities_rref.py index 2d87c68c49..4a0f4a74db 100644 --- a/python/tests/test_conserved_quantities_rref.py +++ b/python/tests/test_conserved_quantities_rref.py @@ -12,10 +12,8 @@ def random_matrix_generator(min_dim, max_dim, count): """Generate random 2D square matrix""" rng = np.random.default_rng(12345) - return [ - np.random.rand(rows, cols) - for rows, cols in rng.integers(min_dim, max_dim, [count, 2]) - ] + for rows, cols in rng.integers(min_dim, max_dim, [count, 2]): + yield np.random.rand(rows, cols) @skip_on_valgrind