From 9bb778a508efb22c6024861bec939290b9fb9362 Mon Sep 17 00:00:00 2001 From: Daniel Weindl Date: Fri, 17 Oct 2025 13:51:45 +0200 Subject: [PATCH] Move fiddy-amici interface to amici Fiddy (https://github.com/ICB-DCM/fiddy/), a package for finite difference methods, is used here for testing. So far, the amici-fiddy interface code was in the fiddy repository. To resolve this cyclic dependency, let's move that interface code here. Co-authored-by: Dilan Pathirana <59329744+dilpath@users.noreply.github.com> --- doc/python_modules.rst | 1 + doc/rtd_requirements.txt | 1 + python/sdist/amici/adapters/__init__.py | 3 + python/sdist/amici/adapters/fiddy.py | 393 ++++++++++++++++++ python/sdist/pyproject.toml | 1 + python/tests/adapters/test_fiddy.py | 170 ++++++++ .../benchmark_models/test_petab_benchmark.py | 2 +- 7 files changed, 570 insertions(+), 1 deletion(-) create mode 100644 python/sdist/amici/adapters/__init__.py create mode 100644 python/sdist/amici/adapters/fiddy.py create mode 100644 python/tests/adapters/test_fiddy.py diff --git a/doc/python_modules.rst b/doc/python_modules.rst index 3408ef3e50..e9c5affdc4 100644 --- a/doc/python_modules.rst +++ b/doc/python_modules.rst @@ -34,3 +34,4 @@ AMICI Python API amici.numpy amici.sbml_utils amici.splines + amici.adapters.fiddy diff --git a/doc/rtd_requirements.txt b/doc/rtd_requirements.txt index cbb21058c2..d014b27755 100644 --- a/doc/rtd_requirements.txt +++ b/doc/rtd_requirements.txt @@ -26,3 +26,4 @@ ipykernel -e git+https://github.com/Benchmarking-Initiative/Benchmark-Models-PEtab.git@master#subdirectory=src/python&egg=benchmark_models_petab -e python/sdist/[jax] antimony>=2.13 +fiddy>=0.0.3 diff --git a/python/sdist/amici/adapters/__init__.py b/python/sdist/amici/adapters/__init__.py new file mode 100644 index 0000000000..ce7f80c730 --- /dev/null +++ b/python/sdist/amici/adapters/__init__.py @@ -0,0 +1,3 @@ +""" +Functionality for interfacing AMICI with other tools and frameworks. +""" diff --git a/python/sdist/amici/adapters/fiddy.py b/python/sdist/amici/adapters/fiddy.py new file mode 100644 index 0000000000..25687a396b --- /dev/null +++ b/python/sdist/amici/adapters/fiddy.py @@ -0,0 +1,393 @@ +""" +Adapters for using AMICI with the `fiddy `__ +package for finite difference checks. + +NOTE: Like fiddy, this module is experimental and subject to change. +""" + +from collections.abc import Callable +from functools import partial +from inspect import signature +from typing import Any + +import numpy as np +import petab.v1 as petab +from fiddy import CachedFunction, Type, fiddy_array +from petab.v1.C import LIN, LOG, LOG10 + +import amici +import amici.petab.simulations +from amici.petab.conditions import create_edatas +from amici.petab.parameter_mapping import create_parameter_mapping +from amici.petab.simulations import LLH, SLLH + +from .. import ReturnData, SensitivityOrder + +__all__ = [ + "run_simulation_to_cached_functions", + "simulate_petab_to_cached_functions", +] + +LOG_E_10 = np.log(10) + + +def transform_gradient_lin_to_lin(gradient_value, _): + return gradient_value + + +def transform_gradient_lin_to_log(gradient_value, parameter_value): + return gradient_value * parameter_value + + +def transform_gradient_lin_to_log10(gradient_value, parameter_value): + return gradient_value * (parameter_value * LOG_E_10) + + +transforms = { + LIN: transform_gradient_lin_to_lin, + LOG: transform_gradient_lin_to_log, + LOG10: transform_gradient_lin_to_log10, +} + + +all_rdata_derivatives = { + "x": "sx", + "x0": "sx0", + "x_ss": "sx_ss", + "y": "sy", + "sigmay": "ssigmay", + "z": "sz", + "rz": "srz", + "sigmaz": "ssigmaz", + "llh": "sllh", + "sllh": "s2llh", + "res": "sres", +} + +# The dimension of the AMICI ReturnData that contains parameters. +# Should be shifted to the last dimension to be compatible with fiddy. +derivative_parameter_dimension = { + "sx": 1, + "sx0": 0, + "sx_ss": 0, + "sy": 1, + "ssigmay": 1, + # 'sz' : ???, + "srz": 2, + # 'ssigmaz' : ???, + "sllh": 0, + "s2llh": 1, + "sres": 1, +} + + +def rdata_array_transpose(array: np.ndarray, variable: str) -> tuple[int]: + if array.size == 0: + return array + original_parameter_dimension = derivative_parameter_dimension[variable] + return np.moveaxis(array, original_parameter_dimension, -1) + + +def fiddy_array_transpose(array: np.ndarray, variable: str) -> tuple[int]: + if array.size == 0: + return array + original_parameter_dimension = derivative_parameter_dimension[variable] + return np.moveaxis(array, -1, original_parameter_dimension) + + +default_derivatives = { + k: v + for k, v in all_rdata_derivatives.items() + if v not in ["sz", "srz", "ssigmaz", "s2llh"] +} + + +def run_simulation_to_cached_functions( + amici_model: amici.AmiciModel, + *, + cache: bool = True, + parameter_ids: list[str] = None, + amici_solver: amici.AmiciSolver = None, + amici_edata: amici.AmiciExpData = None, + derivative_variables: list[str] = None, +): + """Convert `amici.run_simulation` to fiddy functions. + + :param amici_model: + The AMICI model to simulate. + :param amici_solver: + The AMICI solver to use. If `None`, a new solver will be created from + the model. + :param amici_edata: + The AMICI ExpData to use. If `None`, no data will be used. + :param derivative_variables: + The variables that derivatives will be computed or approximated for. + See the keys of `all_rdata_derivatives` for options. + :param parameter_ids: + The IDs that correspond to the values in the parameter vector that is + simulated. + :param cache: + Whether to cache the function calls. + :returns: function, derivatives and structure + """ + if amici_solver is None: + amici_solver = amici_model.create_solver() + if parameter_ids is None: + parameter_ids = amici_model.get_parameter_ids() + if amici_edata is not None and amici_edata.parameters is not None: + raise NotImplementedError( + "Customization of parameter values inside AMICI ExpData." + ) + chosen_derivatives = default_derivatives + if derivative_variables is not None: + chosen_derivatives = { + k: all_rdata_derivatives[k] for k in derivative_variables + } + + def run_amici_simulation( + point: Type.POINT, order: SensitivityOrder + ) -> ReturnData: + problem_parameters = dict(zip(parameter_ids, point, strict=True)) + amici_model.set_parameter_by_id(problem_parameters) + amici_solver.set_sensitivity_order(order) + rdata = amici.run_simulation( + model=amici_model, solver=amici_solver, edata=amici_edata + ) + return rdata + + def function(point: Type.POINT): + rdata = run_amici_simulation(point=point, order=SensitivityOrder.none) + outputs = { + variable: fiddy_array(getattr(rdata, variable)) + for variable in chosen_derivatives + } + rdata_flat = np.concatenate( + [output.flat for output in outputs.values()] + ) + return rdata_flat + + def derivative(point: Type.POINT, return_dict: bool = False): + rdata = run_amici_simulation(point=point, order=SensitivityOrder.first) + outputs = { + variable: rdata_array_transpose( + array=fiddy_array(getattr(rdata, derivative_variable)), + variable=derivative_variable, + ) + for variable, derivative_variable in chosen_derivatives.items() + } + rdata_flat = np.concatenate( + [ + output_array.reshape(-1, output_array.shape[-1]) + for output_array in outputs.values() + ], + axis=0, + ) + if return_dict: + return outputs + return rdata_flat + + if cache: + function = CachedFunction(function) + derivative = CachedFunction(derivative) + + # Get structure + dummy_point = fiddy_array( + [amici_model.get_parameter_by_id(par_id) for par_id in parameter_ids] + ) + dummy_rdata = run_amici_simulation( + point=dummy_point, order=SensitivityOrder.first + ) + + structures = { + "function": {variable: None for variable in chosen_derivatives}, + "derivative": {variable: None for variable in chosen_derivatives}, + } + function_position = 0 + derivative_position = 0 + for variable, derivative_variable in chosen_derivatives.items(): + function_array = fiddy_array(getattr(dummy_rdata, variable)) + derivative_array = fiddy_array( + getattr(dummy_rdata, derivative_variable) + ) + structures["function"][variable] = ( + function_position, + function_position + function_array.size, + function_array.shape, + ) + structures["derivative"][variable] = ( + derivative_position, + derivative_position + derivative_array.size, + derivative_array.shape, + ) + function_position += function_array.size + derivative_position += derivative_array.size + + return function, derivative, structures + + +# (start, stop, shape) +TYPE_STRUCTURE = tuple[int, int, tuple[int, ...]] + + +def flatten(arrays: dict[str, Type.ARRAY]) -> Type.ARRAY: + flattened_value = np.concatenate([array.flat for array in arrays.values()]) + return flattened_value + + +def reshape( + array: Type.ARRAY, + structure: TYPE_STRUCTURE, + sensitivities: bool = False, +) -> dict[str, Type.ARRAY]: + reshaped = {} + for variable, (start, stop, shape) in structure.items(): + # array is currently "flattened" w.r.t. fiddy dimensions + # hence, if sensis, reshape w.r.t. fiddy dimensions + if sensitivities and ( + dimension0 := derivative_parameter_dimension.get( + "s" + variable, False + ) + ): + shape = [ + size + for dimension, size in enumerate(shape) + if dimension != dimension0 + ] + [shape[dimension0]] + + array = array[start:stop] + if array.size != 0: + array = array.reshape(shape) + + # now reshape to AMICI dimensions + if sensitivities and ( + derivative_parameter_dimension.get(f"s{variable}", False) + ): + array = fiddy_array_transpose( + array=array, + variable=f"s{variable}", + ) + reshaped[variable] = array + + return reshaped + + +def simulate_petab_to_cached_functions( + petab_problem: petab.Problem, + amici_model: amici.Model, + parameter_ids: list[str] = None, + cache: bool = True, + precreate_edatas: bool = True, + precreate_parameter_mapping: bool = True, + simulate_petab: Callable[[Any], str] = None, + **kwargs, +) -> tuple[Type.FUNCTION, Type.FUNCTION]: + """ + Convert :func:`amici.petab.simulations.simulate_petab` + (PEtab v1 simulations) to fiddy functions. + + Note that all gradients are provided on linear scale. The correction from + `'log10'` scale is automatically done. + + :param amici_model: + The AMICI model to simulate. + :param simulate_petab: + A method to simulate PEtab problems with AMICI, e.g. + `amici.petab_objective.simulate_petab`. + :param parameter_ids: + The IDs of the parameters, in the order that parameter values will + be supplied. Defaults to `petab_problem.parameter_df.index`. + :param petab_problem: + The PEtab problem. + :param cache: + Whether to cache the function call. + :param precreate_edatas: + Whether to create the AMICI measurements object in advance, to save + time. + :param precreate_parameter_mapping: + Whether to create the AMICI parameter mapping object in advance, to + save time. + :param kwargs: + Passed to `simulate_petab`. + :returns: + A tuple of: + + * 1: A method to compute the function at a point. + * 2: A method to compute the gradient at a point. + """ + if parameter_ids is None: + parameter_ids = list(petab_problem.parameter_df.index) + + if simulate_petab is None: + simulate_petab = amici.petab.simulations.simulate_petab + + edatas = None + if precreate_edatas: + edatas = create_edatas( + amici_model=amici_model, + petab_problem=petab_problem, + simulation_conditions=petab_problem.get_simulation_conditions_from_measurement_df(), + ) + + parameter_mapping = None + if precreate_parameter_mapping: + parameter_mapping = create_parameter_mapping( + petab_problem=petab_problem, + simulation_conditions=petab_problem.get_simulation_conditions_from_measurement_df(), + scaled_parameters=kwargs.get( + "scaled_parameters", + ( + signature(simulate_petab) + .parameters["scaled_parameters"] + .default + ), + ), + amici_model=amici_model, + ) + + precreated_kwargs = { + "edatas": edatas, + "parameter_mapping": parameter_mapping, + "petab_problem": petab_problem, + } + precreated_kwargs = { + k: v for k, v in precreated_kwargs.items() if v is not None + } + + amici_solver = kwargs.pop("solver", amici_model.create_solver()) + + simulate_petab_partial = partial( + simulate_petab, + amici_model=amici_model, + **precreated_kwargs, + **kwargs, + ) + + def simulate_petab_full(point: Type.POINT, order: SensitivityOrder): + problem_parameters = dict(zip(parameter_ids, point, strict=True)) + amici_solver.set_sensitivity_order(order) + result = simulate_petab_partial( + problem_parameters=problem_parameters, + solver=amici_solver, + ) + return result + + def function(point: Type.POINT): + output = simulate_petab_full(point, order=SensitivityOrder.none) + result = output[LLH] + return np.array(result) + + def derivative(point: Type.POINT) -> Type.POINT: + result = simulate_petab_full(point, order=SensitivityOrder.first) + if result[SLLH] is None: + raise RuntimeError("Simulation failed.") + + sllh = np.array( + [result[SLLH][parameter_id] for parameter_id in parameter_ids] + ) + return sllh + + if cache: + function = CachedFunction(function) + derivative = CachedFunction(derivative) + + return function, derivative diff --git a/python/sdist/pyproject.toml b/python/sdist/pyproject.toml index 736bf0ca51..b6daa81c70 100644 --- a/python/sdist/pyproject.toml +++ b/python/sdist/pyproject.toml @@ -63,6 +63,7 @@ classifiers = [ petab = ["petab>=0.4.0"] pysb = ["pysb>=1.13.1"] test = [ + "fiddy>=0.0.3", "h5py!=3.15", "pytest", "pytest-cov", diff --git a/python/tests/adapters/test_fiddy.py b/python/tests/adapters/test_fiddy.py new file mode 100644 index 0000000000..2ae5c88f96 --- /dev/null +++ b/python/tests/adapters/test_fiddy.py @@ -0,0 +1,170 @@ +"""Tests for `amici.adapters.fiddy`.""" + +from pathlib import Path + +import amici +import amici.petab.petab_import +import amici.petab.simulations +import numpy as np +import pytest +from amici import SensitivityOrder, SteadyStateSensitivityMode +from amici.adapters.fiddy import ( + run_simulation_to_cached_functions, + simulate_petab_to_cached_functions, +) +from fiddy import MethodId, Type, get_derivative +from fiddy.derivative_check import NumpyIsCloseDerivativeCheck +from fiddy.success import Consistency +from petab import v1 + +# Absolute and relative tolerances for finite difference gradient checks. +ATOL: float = 1e-3 +RTOL: float = 1e-3 + + +def lotka_volterra() -> tuple[v1.Problem, np.ndarray]: + petab_problem = v1.Problem.from_yaml( + str( + Path(__file__).parents[1] + / "petab_test_problems" + / "lotka_volterra" + / "petab" + / "problem.yaml" + ) + ) + point = np.array([2, 3], dtype=Type.SCALAR) + return petab_problem, point + + +@pytest.mark.parametrize("problem_generator", [lotka_volterra]) +def test_run_amici_simulation_to_functions(problem_generator): + petab_problem, point = problem_generator() + timepoints = sorted(set(petab_problem.measurement_df.time)) + amici_model = amici.petab.petab_import.import_petab_problem(petab_problem) + amici_model.set_timepoints(timepoints) + amici_solver = amici_model.create_solver() + + amici_solver.set_sensitivity_order(SensitivityOrder.first) + + parameter_ids = list( + petab_problem.parameter_df[ + petab_problem.parameter_df.estimate == 1 + ].index + ) + parameter_indices = [ + amici_model.get_parameter_ids().index(parameter_id) + for parameter_id in parameter_ids + ] + + ( + amici_function, + amici_derivative, + structures, + ) = run_simulation_to_cached_functions( + parameter_ids=parameter_ids, + amici_model=amici_model, + amici_solver=amici_solver, + ) + + expected_derivative = amici_derivative(point)[..., parameter_indices] + + derivative = get_derivative( + function=amici_function, + point=point, + sizes=[1e-10, 1e-5], + direction_ids=parameter_ids, + method_ids=[MethodId.FORWARD, MethodId.BACKWARD, MethodId.CENTRAL], + # analysis_classes=[], + # analysis_classes=[ + # lambda: TransformByDirectionScale(scales=parameter_scales), + # ], + success_checker=Consistency(atol=1e-2), + ) + test_derivative = derivative.value + + # The test derivative is close to the expected derivative. + assert np.isclose( + test_derivative, + expected_derivative, + rtol=1e-1, + atol=1e-1, + equal_nan=True, + ).all() + + # Same as above assert. + check = NumpyIsCloseDerivativeCheck( + derivative=derivative, + expectation=expected_derivative, + point=point, + ) + result = check(rtol=1e-1, atol=1e-1, equal_nan=True) + assert result.success + + +@pytest.mark.parametrize("problem_generator", [lotka_volterra]) +@pytest.mark.parametrize("scaled_parameters", (False, True)) +def test_simulate_petab_to_functions(problem_generator, scaled_parameters): + petab_problem, point = problem_generator() + amici_model = amici.petab.petab_import.import_petab_problem(petab_problem) + amici_solver = amici_model.create_solver() + + if amici_model.get_name() == "simple": + amici_model.set_steady_state_sensitivity_mode( + SteadyStateSensitivityMode.integrationOnly + ) + + amici_solver.set_sensitivity_order(SensitivityOrder.first) + + if scaled_parameters: + point = np.asarray( + list( + petab_problem.scale_parameters( + dict( + zip( + petab_problem.parameter_df.index, + point, + strict=True, + ) + ) + ).values() + ) + ) + + amici_function, amici_derivative = simulate_petab_to_cached_functions( + parameter_ids=petab_problem.parameter_df.index, + petab_problem=petab_problem, + amici_model=amici_model, + solver=amici_solver, + scaled_gradients=scaled_parameters, + scaled_parameters=scaled_parameters, + ) + + expected_derivative = amici_derivative(point) + + parameter_ids = list( + petab_problem.parameter_df[ + petab_problem.parameter_df.estimate == 1 + ].index + ) + # parameter_scales = dict( + # petab_problem.parameter_df[ + # petab_problem.parameter_df.estimate == 1 + # ].parameterScale + # ) + + derivative = get_derivative( + function=amici_function, + point=point, + sizes=[1e-10, 1e-5, 1e-3, 1e-1], + direction_ids=parameter_ids, + method_ids=[MethodId.FORWARD, MethodId.BACKWARD, MethodId.CENTRAL], + success_checker=Consistency(), + ) + + check = NumpyIsCloseDerivativeCheck( + derivative=derivative, + expectation=expected_derivative, + point=point, + ) + result = check(rtol=1e-2) + assert result.success diff --git a/tests/benchmark_models/test_petab_benchmark.py b/tests/benchmark_models/test_petab_benchmark.py index d9466ebd4d..242b45a5ec 100644 --- a/tests/benchmark_models/test_petab_benchmark.py +++ b/tests/benchmark_models/test_petab_benchmark.py @@ -21,6 +21,7 @@ import pytest import yaml from amici import SensitivityMethod +from amici.adapters.fiddy import simulate_petab_to_cached_functions from amici.logging import get_logger from amici.petab.petab_import import import_petab_problem from amici.petab.simulations import ( @@ -31,7 +32,6 @@ ) from fiddy import MethodId, get_derivative from fiddy.derivative_check import NumpyIsCloseDerivativeCheck -from fiddy.extensions.amici import simulate_petab_to_cached_functions from fiddy.success import Consistency from petab.v1.lint import measurement_table_has_timepoint_specific_mappings from petab.v1.visualize import plot_problem