Extended Amici history

pypesto/C.py

@@ -231,6 +231,12 @@ class InnerParameterType(str, Enum):
 SUFFIXES_HDF5 = ["hdf5", "h5"]
 SUFFIXES = SUFFIXES_CSV + SUFFIXES_HDF5
 
+CPU_TIME_TOTAL = 'cpu_time_total'
+PREEQ_CPU_TIME = 'preeq_cpu_time'
+PREEQ_CPU_TIME_B = 'preeq_cpu_timeB'
+POSTEQ_CPU_TIME = 'posteq_cpu_time'
+POSTEQ_CPU_TIME_B = 'posteq_cpu_timeB'
+
 
 ###############################################################################
 # PRIOR

pypesto/__init__.py

@@ -17,6 +17,7 @@
     CountHistoryBase,
     CsvHistory,
     Hdf5History,
+    Hdf5AmiciHistory,
     NoHistory,
     HistoryBase,
     HistoryOptions,

pypesto/history/__init__.py

@@ -8,6 +8,7 @@
 and evaluate performance.
 """
 
+from .amici import Hdf5AmiciHistory
 from .base import CountHistory, CountHistoryBase, HistoryBase, NoHistory
 from .csv import CsvHistory
 from .generate import create_history

pypesto/history/amici.py

@@ -0,0 +1,160 @@
+import time
+from pathlib import Path
+from typing import Sequence, Union
+
+import numpy as np
+
+from ..C import (
+    FVAL,
+    GRAD,
+    HESS,
+    N_ITERATIONS,
+    RES,
+    SRES,
+    RDATAS,
+    TIME,
+    ModeType,
+    X,
+    CPU_TIME_TOTAL,
+    PREEQ_CPU_TIME,
+    POSTEQ_CPU_TIME_B,
+    POSTEQ_CPU_TIME,
+    PREEQ_CPU_TIME_B
+)
+from .base import add_fun_from_res, reduce_result_via_options
+from .hdf5 import Hdf5History, with_h5_file
+from .options import HistoryOptions
+from .util import ResultDict, trace_wrap
+
+
+class Hdf5AmiciHistory(Hdf5History):
+    """
+    Stores a representation of the history in an HDF5 file, extended with
+    AMICI-specific traces of total simulation time, pre-equilibration time and
+    post-equilibration time
+
+    Parameters
+    ----------
+    id:
+        Id of the history
+    file:
+        HDF5 file name.
+    options:
+        History options. Defaults to ``None``.
+    """
+
+    def __init__(
+        self,
+        id: str,
+        file: Union[str, Path],
+        options: Union[HistoryOptions, dict, None] = None,
+    ):
+        super().__init__(id, file, options=options)
+
+    @with_h5_file("a")
+    def _update_trace(
+        self,
+        x: np.ndarray,
+        sensi_orders: tuple[int],
+        mode: ModeType,
+        result: ResultDict,
+    ) -> None:
+        """Update and possibly store the trace."""
+        if not self.options.trace_record:
+            return
+
+        # calculating function values from residuals
+        #  and reduce via requested history options
+        result = reduce_result_via_options(
+            add_fun_from_res(result), self.options
+        )
+
+        used_time = time.time() - self.start_time
+
+        values = {
+            X: x,
+            FVAL: result[FVAL],
+            GRAD: result[GRAD],
+            RES: result[RES],
+            SRES: result[SRES],
+            HESS: result[HESS],
+            TIME: used_time,
+            CPU_TIME_TOTAL: sum([rdata[CPU_TIME_TOTAL] for rdata in
+                                 result[RDATAS]]),
+            PREEQ_CPU_TIME: sum([rdata[PREEQ_CPU_TIME] for rdata in
+                                 result[RDATAS]]),
+            PREEQ_CPU_TIME_B: sum([rdata[PREEQ_CPU_TIME_B] for rdata in
+                                   result[RDATAS]]),
+            POSTEQ_CPU_TIME: sum([rdata[POSTEQ_CPU_TIME] for rdata in
+                                  result[RDATAS]]),
+            POSTEQ_CPU_TIME_B: sum([rdata[POSTEQ_CPU_TIME_B] for rdata in
+                                    result[RDATAS]])
+        }
+
+        iteration = self._require_group().attrs[N_ITERATIONS]
+
+        for key in values.keys():
+            if values[key] is not None:
+                self._require_group()[f'{iteration}/{key}'] = values[key]
+
+        self._require_group().attrs[N_ITERATIONS] += 1
+
+    @trace_wrap
+    def get_cpu_time_total_trace(
+        self, ix: Union[int, Sequence[int], None] = None, trim: bool = False
+    ) -> Union[Sequence[float], float]:
+        """
+        Cumulative simulation CPU time [ms].
+        Takes as parameter an index or indices and returns corresponding trace
+        values. If only a single value is requested, the list is flattened.
+        """
+        return self._get_hdf5_entries(CPU_TIME_TOTAL, ix)
+
+    @trace_wrap
+    def get_preeq_time_trace(
+        self, ix: Union[int, Sequence[int], None] = None, trim: bool = False
+    ) -> Union[Sequence[float], float]:
+        """
+        Cumulative computation time of the steady state solver [ms].
+        (preequilibration)
+        Takes as parameter an index or indices and returns corresponding trace
+        values. If only a single value is requested, the list is flattened.
+        """
+        return self._get_hdf5_entries(PREEQ_CPU_TIME, ix)
+
+    @trace_wrap
+    def get_preeq_timeB_trace(
+        self, ix: Union[int, Sequence[int], None] = None, trim: bool = False
+    ) -> Union[Sequence[float], float]:
+        """
+        Cumulative computation time of the steady state solver of the backward
+        problem [ms] (preequilibration).
+        Takes as parameter an index or indices and returns corresponding trace
+        values. If only a single value is requested, the list is flattened.
+        """
+        return self._get_hdf5_entries(PREEQ_CPU_TIME_B, ix)
+
+    @trace_wrap
+    def get_posteq_time_trace(
+        self, ix: Union[int, Sequence[int], None] = None, trim: bool = False
+    ) -> Union[Sequence[float], float]:
+        """
+        Cumulative computation time of the steady state solver [ms]
+        (postequilibration).
+        Takes as parameter an index or indices and returns corresponding trace
+        values. If only a single value is requested, the list is flattened.
+        """
+        return self._get_hdf5_entries(POSTEQ_CPU_TIME, ix)
+
+    @trace_wrap
+    def get_posteq_timeB_trace(
+        self, ix: Union[int, Sequence[int], None] = None, trim: bool = False
+    ) -> Union[Sequence[float], float]:
+        """
+        Cumulative computation time of the steady state solver of the backward
+        problem [ms] (postequilibration).
+        Takes as parameter an index or indices and returns corresponding trace
+        values. If only a single value is requested, the list is flattened.
+        """
+        return self._get_hdf5_entries(POSTEQ_CPU_TIME_B, ix)
+

pypesto/history/base.py

@@ -249,7 +249,7 @@ def get_time_trace(
         trim: bool = False,
     ) -> Union[Sequence[float], float]:
         """
-        Cumulative execution times.
+        Cumulative execution times [s].
 
         Takes as parameter an index or indices and returns corresponding trace
         values. If only a single value is requested, the list is flattened.

pypesto/history/generate.py

@@ -7,6 +7,7 @@
 from .base import CountHistory, HistoryBase
 from .csv import CsvHistory
 from .hdf5 import Hdf5History
+from .amici import Hdf5AmiciHistory
 from .memory import MemoryHistory
 from .options import HistoryOptions
 from .util import HistoryTypeError
@@ -16,6 +17,7 @@ def create_history(
     id: str,
     x_names: Sequence[str],
     options: HistoryOptions,
+    amici_objective: bool
 ) -> HistoryBase:
     """Create a :class:`HistoryBase` object; Factory method.
 
@@ -27,6 +29,8 @@ def create_history(
         Parameter names.
     options:
         History options.
+    amici_objective:
+        Indicates if AmiciObjective was used
 
     Returns
     -------
@@ -49,6 +53,9 @@ def create_history(
     if suffix in SUFFIXES_CSV:
         return CsvHistory(x_names=x_names, file=storage_file, options=options)
     elif suffix in SUFFIXES_HDF5:
-        return Hdf5History(id=id, file=storage_file, options=options)
+        if amici_objective:
+            return Hdf5AmiciHistory(id=id, file=storage_file, options=options)
+        else:
+            return Hdf5History(id=id, file=storage_file, options=options)
     else:
         raise HistoryTypeError(suffix)

pypesto/optimize/optimizer.py

@@ -18,7 +18,7 @@
     OptimizerHistory,
     create_history,
 )
-from ..objective import Objective
+from ..objective import AmiciObjective, Objective
 from ..problem import Problem
 from ..result import OptimizerResult
 from .load import fill_result_from_history
@@ -106,6 +106,7 @@ def wrapped_minimize(
             id=id,
             x_names=[problem.x_names[ix] for ix in problem.x_free_indices],
             options=history_options,
+            amici_objective=isinstance(objective, AmiciObjective)
         )
         optimizer_history = OptimizerHistory(
             history=history,

test/base/test_history.py

@@ -16,6 +16,7 @@
 from pypesto import (
     CsvHistory,
     Hdf5History,
+    Hdf5AmiciHistory,
     HistoryOptions,
     MemoryHistory,
     ObjectiveBase,
@@ -689,6 +690,56 @@ def test_hdf5_history_mp():
                             )
 
 
+def test_hdf5_amici_history():
+    objective1 = pypesto.Objective(
+        fun=so.rosen, grad=so.rosen_der, hess=so.rosen_hess
+    )
+    objective2 = load_amici_objective('conversion_reaction')[0]
+    lb = -2 * np.ones((1, 2))
+    ub = 2 * np.ones((1, 2))
+    problem1 = pypesto.Problem(
+        objective=objective1, lb=lb, ub=ub
+    )
+    problem2 = pypesto.Problem(
+        objective=objective2, lb=lb, ub=ub
+    )
+
+    optimizer = pypesto.optimize.ScipyOptimizer(options={'maxiter': 10})
+
+    with tempfile.TemporaryDirectory(dir=".") as tmpdirname:
+        _, fn = tempfile.mkstemp(".hdf5", dir=f"{tmpdirname}")
+
+        history_options_mp = pypesto.HistoryOptions(
+            trace_record=True, storage_file=fn
+        )
+
+        result1 = pypesto.optimize.minimize(
+            problem=problem1,
+            optimizer=optimizer,
+            n_starts=1,
+            history_options=history_options_mp,
+            progress_bar=False,
+        )
+        assert not isinstance(result1.optimize_result.list[0].history,
+                              Hdf5AmiciHistory)
+
+        # optimizing with amici history saved in hdf5
+        result2 = pypesto.optimize.minimize(
+            problem=problem2,
+            optimizer=optimizer,
+            n_starts=1,
+            history_options=history_options_mp,
+            progress_bar=False,
+        )
+        assert isinstance(result2.optimize_result.list[0].history,
+                          Hdf5AmiciHistory)
+        result2.optimize_result.list[0].history.get_cpu_time_total_trace()
+        result2.optimize_result.list[0].history.get_preeq_time_trace()
+        result2.optimize_result.list[0].history.get_preeq_timeB_trace()
+        result2.optimize_result.list[0].history.get_posteq_time_trace()
+        result2.optimize_result.list[0].history.get_posteq_timeB_trace()
+
+
 def test_trim_history():
     """
     Test whether the history gets correctly trimmed to be monotonically