PetabSimulationResult for PetabSimulator.simulate

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",

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
 

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",
 ]

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]

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

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
 

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)