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 diff --git a/python/sdist/amici/exporters/jax/nn.py b/python/sdist/amici/exporters/jax/nn.py index 576cb6265d..1370738c03 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,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, diff --git a/python/sdist/amici/importers/petab/_petab_importer.py b/python/sdist/amici/importers/petab/_petab_importer.py index a9e1c20c4c..bc0c991962 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"]).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 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 0000000000..102fedf43f Binary files /dev/null and b/python/tests/sciml_test_problems/equinox_import/petab/net1_ps.hdf5 differ diff --git a/python/tests/sciml_test_problems/equinox_import/petab/observables.tsv b/python/tests/sciml_test_problems/equinox_import/petab/observables.tsv new file mode 100644 index 0000000000..fc5f605b2c --- /dev/null +++ b/python/tests/sciml_test_problems/equinox_import/petab/observables.tsv @@ -0,0 +1,3 @@ +observableId observableFormula noiseFormula observableTransformation noiseDistribution +prey_o prey 0.05 lin normal +predator_o predator 0.05 lin normal diff --git a/python/tests/sciml_test_problems/equinox_import/petab/parameters.tsv b/python/tests/sciml_test_problems/equinox_import/petab/parameters.tsv new file mode 100644 index 0000000000..dcaf68d365 --- /dev/null +++ b/python/tests/sciml_test_problems/equinox_import/petab/parameters.tsv @@ -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 diff --git a/python/tests/sciml_test_problems/equinox_import/petab/problem.yaml b/python/tests/sciml_test_problems/equinox_import/petab/problem.yaml new file mode 100644 index 0000000000..0ea6784db6 --- /dev/null +++ b/python/tests/sciml_test_problems/equinox_import/petab/problem.yaml @@ -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" 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 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 diff --git a/python/tests/test_sciml.py b/python/tests/test_sciml.py index e9f97aaec2..1fbf3924ef 100644 --- a/python/tests/test_sciml.py +++ b/python/tests/test_sciml.py @@ -1,21 +1,55 @@ """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 jax +import numpy as np import pytest -pytest.importorskip("jax") pytest.importorskip("equinox") +pytest.importorskip("petab_sciml") -import pytest 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 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", + 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) class TestFormatFunctionCall: @@ -789,3 +823,57 @@ 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 + ) + hybridization = pi._build_hybridization() + + assert isinstance(hybridization["net1"]["model"], Path) + + model_file_str = str(hybridization["net1"]["model"]) + generate_equinox( + model_file_str, + pi.output_dir / "net1.py", + hybridization["net1"]["frozen_layers"], + ) + + 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"], + ) 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 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(