diff --git a/doc/examples/example_petab/petab_v2.ipynb b/doc/examples/example_petab/petab_v2.ipynb index 69ce771af1..b3b90b5aff 100644 --- a/doc/examples/example_petab/petab_v2.ipynb +++ b/doc/examples/example_petab/petab_v2.ipynb @@ -76,46 +76,46 @@ "source": "Now let's run the simulations:" }, { - "cell_type": "code", - "execution_count": null, - "id": "cee93ebff9477aca", "metadata": {}, + "cell_type": "code", "outputs": [], + "execution_count": null, "source": [ "# simulate all conditions encoded in the PEtab problem for which there are measurements\n", "# using the nominal parameter values from the PEtab problem\n", "result = simulator.simulate(problem.get_x_nominal_dict())\n", - "assert all(r.status == AMICI_SUCCESS for r in result[RDATAS]), (\n", + "assert all(r.status == AMICI_SUCCESS for r in result.rdatas), (\n", " \"Simulation failed.\"\n", ")\n", "result" - ] + ], + "id": "4eb108e1b01edfc0" }, { "cell_type": "markdown", "id": "3d73841ed61412ff", "metadata": {}, "source": [ - "The returned dictionary contains the simulation results for all experimental conditions encoded in the PEtab problem in `result['rdatas']`.\n", + "The returned object contains the simulation results for all experimental conditions encoded in the PEtab problem in `result.rdatas`.\n", "Those are the same objects as for any other AMICI simulation using `run_simulation`.\n", - "Additionally, the dictionary contains the `ExpData` instances used for the simulations in `result['edatas']`, which we will use below to visualize the PEtab-encoded measurements.\n", - "`result['llh']` and `result['sllh']` contain the aggregated log-likelihood value and its gradient, respectively, over all experimental conditions.\n", + "Additionally, the dictionary contains the `ExpData` instances used for the simulations in `result.edatas`, which we will use below to visualize the PEtab-encoded measurements.\n", + "`result.llh` and `result.sllh` contain the aggregated log-likelihood value and its gradient, respectively, over all experimental conditions.\n", "These can be used directly for parameter estimation. However, for parameter estimation, it is recommended to use the `pypesto` package that provides a full parameter estimation framework on top of AMICI and PEtab.\n", "\n", "Now, let's have a look at the results of the first experimental condition:" ] }, { - "cell_type": "code", - "execution_count": null, - "id": "e532cbc3ebb361ef", "metadata": {}, + "cell_type": "code", "outputs": [], + "execution_count": null, "source": [ - "rdata = result[\"rdatas\"][0]\n", - "edata = result[\"edatas\"][0]\n", + "rdata = result.rdatas[0]\n", + "edata = result.edatas[0]\n", "plot_observable_trajectories(rdata, model=simulator.model, edata=edata)" - ] + ], + "id": "d6bbefc2e3d3ccbc" }, { "cell_type": "markdown", diff --git a/python/sdist/amici/adapters/fiddy.py b/python/sdist/amici/adapters/fiddy.py index 026100679c..573361bbba 100644 --- a/python/sdist/amici/adapters/fiddy.py +++ b/python/sdist/amici/adapters/fiddy.py @@ -33,7 +33,7 @@ from amici.sim.sundials.petab.v1 import LLH, SLLH, create_edatas if TYPE_CHECKING: - from amici.sim.sundials.petab import PetabSimulator + from amici.sim.sundials.petab import PetabSimulationResult, PetabSimulator __all__ = [ "run_simulation_to_cached_functions", @@ -386,7 +386,9 @@ def simulate_petab_v2_to_cached_functions( if free_parameter_ids is None: free_parameter_ids = list(petab_simulator._petab_problem.x_free_ids) - def simulate(point: Type.POINT, order: SensitivityOrder) -> dict: + def simulate( + point: Type.POINT, order: SensitivityOrder + ) -> PetabSimulationResult: problem_parameters = dict(zip(free_parameter_ids, point, strict=True)) petab_simulator.solver.set_sensitivity_order(order) @@ -397,17 +399,17 @@ def simulate(point: Type.POINT, order: SensitivityOrder) -> dict: def function(point: Type.POINT) -> np.ndarray: output = simulate(point, order=SensitivityOrder.none) - result = output[LLH] + result = output.llh return np.array(result) def derivative(point: Type.POINT) -> Type.POINT: result = simulate(point, order=SensitivityOrder.first) - if result[SLLH] is None: + if result.sllh is None: raise RuntimeError("Simulation failed.") sllh = np.array( - [result[SLLH][parameter_id] for parameter_id in free_parameter_ids] + [result.sllh[parameter_id] for parameter_id in free_parameter_ids] ) return sllh diff --git a/python/sdist/amici/sim/sundials/petab/__init__.py b/python/sdist/amici/sim/sundials/petab/__init__.py index a3d3daa208..559c8781b3 100644 --- a/python/sdist/amici/sim/sundials/petab/__init__.py +++ b/python/sdist/amici/sim/sundials/petab/__init__.py @@ -23,25 +23,10 @@ stable. """ -from ._v2 import ExperimentManager, PetabSimulator -from .v1 import ( - EDATAS, - LLH, - RDATAS, - RES, - S2LLH, - SLLH, - SRES, -) +from ._v2 import ExperimentManager, PetabSimulationResult, PetabSimulator __all__ = [ "PetabSimulator", "ExperimentManager", - "EDATAS", - "LLH", - "RDATAS", - "RES", - "S2LLH", - "SLLH", - "SRES", + "PetabSimulationResult", ] diff --git a/python/sdist/amici/sim/sundials/petab/_v2.py b/python/sdist/amici/sim/sundials/petab/_v2.py index e96256c098..40b8d782b8 100644 --- a/python/sdist/amici/sim/sundials/petab/_v2.py +++ b/python/sdist/amici/sim/sundials/petab/_v2.py @@ -6,7 +6,7 @@ import numbers from collections import Counter from collections.abc import Sequence -from typing import Any +from dataclasses import dataclass, field import numpy as np import sympy as sp @@ -16,21 +16,13 @@ import amici from amici.logging import get_logger from amici.sim.sundials import SensitivityOrder -from amici.sim.sundials.petab.v1._simulations import ( - EDATAS, - LLH, - RDATAS, - RES, - S2LLH, - SLLH, - SRES, -) logger = get_logger(__name__, log_level=logging.INFO) __all__ = [ "PetabSimulator", "ExperimentManager", + "PetabSimulationResult", ] @@ -508,6 +500,55 @@ def _get_placeholder_mapping( return mapping +@dataclass +class PetabSimulationResult: + """ + Container for results of a PEtab simulation. + + Holds the per-experiment AMICI data objects and aggregated metrics + produced by :class:`PetabSimulator.simulate`. + """ + + #: List of :class:`amici.sim.sundials.ExpData` instances, one per + #: simulated experiment. These objects may be modified by subsequent + #: operations. + edatas: list[amici.sim.sundials.ExpData] = field(default_factory=list) + #: List of :class:`amici.sim.sundials.ReturnDataView` instances, + #: one per simulated experiment containing simulation outputs. + rdatas: list[amici.sim.sundials.ReturnDataView] = field( + default_factory=list + ) + #: Aggregated first-order sensitivities of the log-likelihood with + #: respect to PEtab problem parameters. Mapping from parameter ID + #: to sensitivity value, or ``None`` if sensitivities were not computed. + sllh: dict[str, float] | None = None + #: Aggregated second-order sensitivities (Hessian or FIM-based) + #: as a 2D :class:`numpy.ndarray` in the order of + #: ``Problem.x_free_ids``. ``None`` if second-order sensitivities + #: were not computed. + s2llh: np.ndarray | None = None + #: Sensitivities of the residuals (if computed) as a + #: :class:`numpy.ndarray`, or ``None`` when not computed. + sres: np.ndarray | None = None + + @property + def llh(self) -> float: + """The total log-likelihood across all experiments.""" + return sum(rdata.llh for rdata in self.rdatas) + + def res(self) -> np.ndarray | None: + """ + Concatenated residuals. + + :returns: Concatenated residuals from all experiments as a 1D + :class:`numpy.ndarray`, or ``None`` if not available. + """ + if any(rdata.res is None for rdata in self.rdatas): + return None + + return np.hstack([rdata.res for rdata in self.rdatas]) + + class PetabSimulator: """ Simulator for PEtab2 problems. @@ -563,22 +604,19 @@ def exp_man(self) -> ExperimentManager: def simulate( self, problem_parameters: dict[str, float] = None - ) -> dict[str, Any]: + ) -> PetabSimulationResult: # TODO params: dict|np.ndarray|None? """Simulate all experiments of the given PEtab problem. :return: - Dictionary of + A :class:`PetabSimulationResult` instance containing the + per-experiment data objects and aggregated results. - * the summed cost function value (``LLH``), - * list of :class:`amici.sim.sundials.ReturnData` (``RDATAS``) - for each experiment, - * list of :class:`amici.sim.sundials.ExpData` (``EDATAS``) - for each experiment + Note that the returned :class:`amici.sim.sundials.ExpData` + instances may be changed by subsequent calls to this function. + Create a copy if needed. - Note that the returned :class:`amici.amiciamici.sim.sundials.ExpData` instances - may be changed by subsequent calls to this function. - Create a copy if needed. + Aggregated residual sensitivities are not implemented yet. """ if problem_parameters is None: # use default parameters, i.e., nominal values for all parameters @@ -601,16 +639,14 @@ def simulate( self._model, self._solver, edatas, num_threads=self.num_threads ) - return { - EDATAS: edatas, - RDATAS: rdatas, - LLH: sum(rdata.llh for rdata in rdatas), - SLLH: self._aggregate_sllh(rdatas), - S2LLH: self._aggregate_s2llh(rdatas, use_fim=True), - RES: np.hstack([rdata.res for rdata in rdatas]), + return PetabSimulationResult( + edatas=edatas, + rdatas=rdatas, + sllh=self._aggregate_sllh(rdatas), + s2llh=self._aggregate_s2llh(rdatas, use_fim=True), # TODO: implement residual sensitivity aggregation - SRES: None, - } + sres=None, + ) def _aggregate_sllh( self, rdatas: Sequence[amici.sim.sundials.ReturnDataView] diff --git a/python/tests/petab_/test_petab_v2.py b/python/tests/petab_/test_petab_v2.py index e3b0038d5d..dcd862c935 100644 --- a/python/tests/petab_/test_petab_v2.py +++ b/python/tests/petab_/test_petab_v2.py @@ -312,7 +312,7 @@ def test_petab_simulator_deepcopy_and_pickle(): ps_copy = copy.deepcopy(ps) - assert ps.simulate({"kk": 2})["llh"] == ps_copy.simulate({"kk": 2})["llh"] + assert ps.simulate({"kk": 2}).llh == ps_copy.simulate({"kk": 2}).llh ps.solver.set_sensitivity_order(SensitivityOrder.first) assert ( @@ -323,6 +323,4 @@ def test_petab_simulator_deepcopy_and_pickle(): import pickle ps_pickle = pickle.loads(pickle.dumps(ps)) - assert ( - ps.simulate({"kk": 2})["llh"] == ps_pickle.simulate({"kk": 2})["llh"] - ) + assert ps.simulate({"kk": 2}).llh == ps_pickle.simulate({"kk": 2}).llh diff --git a/tests/benchmark_models/test_petab_benchmark.py b/tests/benchmark_models/test_petab_benchmark.py index a3428ebd93..42ce893921 100644 --- a/tests/benchmark_models/test_petab_benchmark.py +++ b/tests/benchmark_models/test_petab_benchmark.py @@ -38,7 +38,6 @@ from amici.sim.sundials.petab.v1 import ( LLH, RDATAS, - SLLH, rdatas_to_measurement_df, simulate_petab, ) @@ -680,13 +679,11 @@ def test_nominal_parameters_llh_v2(problem_id): SteadyStateSensitivityMode.integrationOnly ) problem_parameters = problem.get_x_nominal_dict(free=True, fixed=True) - ret = ps.simulate(problem_parameters=problem_parameters) + res = ps.simulate(problem_parameters=problem_parameters) - rdatas = ret[RDATAS] - # chi2 = sum(rdata["chi2"] for rdata in rdatas) - llh = ret[LLH] - - simulation_df = rdatas_to_simulation_df(rdatas, ps.model, pi.petab_problem) + simulation_df = rdatas_to_simulation_df( + res.rdatas, ps.model, pi.petab_problem + ) if was_flattened: simulation_df = unflatten_simulation_df(simulation_df, problem) print("v2 simulations:") @@ -729,7 +726,7 @@ def test_nominal_parameters_llh_v2(problem_id): print("max rel diff:", simulation_df_bm["rel_diff"].max()) # check llh - compare_to_reference(problem_id, llh) + compare_to_reference(problem_id, res.llh) # check gradient if problem_id not in problems_for_gradient_check: @@ -740,11 +737,9 @@ def test_nominal_parameters_llh_v2(problem_id): ps.solver.set_sensitivity_order(SensitivityOrder.first) ps.solver.set_sensitivity_method(SensitivityMethod.forward) ps.model.set_always_check_finite(True) - result = ps.simulate( - problem_parameters=problem_parameters, - ) - assert result[SLLH] is not None - actual_sens_pars = set(result[SLLH].keys()) + result = ps.simulate(problem_parameters=problem_parameters) + assert result.sllh is not None + actual_sens_pars = set(result.sllh.keys()) expected_sens_pars = set(problem.x_free_ids) assert actual_sens_pars == expected_sens_pars diff --git a/tests/petab_test_suite/test_petab_v2_suite.py b/tests/petab_test_suite/test_petab_v2_suite.py index 35c74bfb87..ba98eb3931 100755 --- a/tests/petab_test_suite/test_petab_v2_suite.py +++ b/tests/petab_test_suite/test_petab_v2_suite.py @@ -18,7 +18,7 @@ from amici.sim.sundials.gradient_check import ( check_derivatives as amici_check_derivatives, ) -from amici.sim.sundials.petab import LLH, RDATAS, PetabSimulator +from amici.sim.sundials.petab import PetabSimulator from petab import v2 logger = get_logger(__name__, logging.DEBUG) @@ -79,15 +79,14 @@ def _test_case(case, model_type, version, jax): ps.solver.set_steady_state_tolerance_factor(1.0) problem_parameters = problem.get_x_nominal_dict(free=True, fixed=True) - ret = ps.simulate(problem_parameters=problem_parameters) - - rdatas = ret[RDATAS] + res = ps.simulate(problem_parameters=problem_parameters) + rdatas = res.rdatas for rdata in rdatas: assert rdata.status == AMICI_SUCCESS, ( f"Simulation failed for {rdata.id}" ) chi2 = sum(rdata.chi2 for rdata in rdatas) - llh = ret[LLH] + llh = res.llh simulation_df = rdatas_to_simulation_df(rdatas, ps.model, pi.petab_problem) solution = petabtests.load_solution(case, model_type, version=version)