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_BACKWARD = 'preeq_cpu_timeB'
+POSTEQ_CPU_TIME = 'posteq_cpu_time'
+POSTEQ_CPU_TIME_BACKWARD = 'posteq_cpu_timeB'
+
 
 ###############################################################################
 # PRIOR

pypesto/__init__.py

@@ -16,7 +16,9 @@
     CountHistory,
     CountHistoryBase,
     CsvHistory,
+    CsvAmiciHistory,
     Hdf5History,
+    Hdf5AmiciHistory,
     NoHistory,
     HistoryBase,
     HistoryOptions,

pypesto/history/__init__.py

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

pypesto/history/amici.py

@@ -0,0 +1,250 @@
+from pathlib import Path
+from typing import Sequence, Union
+
+import numpy as np
+
+from ..C import (
+    CPU_TIME_TOTAL,
+    POSTEQ_CPU_TIME,
+    POSTEQ_CPU_TIME_BACKWARD,
+    PREEQ_CPU_TIME,
+    PREEQ_CPU_TIME_BACKWARD,
+    RDATAS,
+)
+from .csv import CsvHistory
+from .hdf5 import Hdf5History
+from .options import HistoryOptions
+from .util import trace_wrap
+
+
+class Hdf5AmiciHistory(Hdf5History):
+    """
+    Stores history extended by AMICI-specific time traces in an HDF5 file.
+
+    Stores 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)
+
+    @staticmethod
+    def _simulation_to_values(x, result, used_time):
+        values = Hdf5History._simulation_to_values(x, result, used_time)
+        # default unit for time in amici is [ms], converted to [s]
+        values |= {
+            CPU_TIME_TOTAL: sum(
+                [rdata[CPU_TIME_TOTAL] for rdata in result[RDATAS]]
+            )*0.001,
+            PREEQ_CPU_TIME: sum(
+                [rdata[PREEQ_CPU_TIME] for rdata in result[RDATAS]]
+            )*0.001,
+            PREEQ_CPU_TIME_BACKWARD: sum(
+                [rdata[PREEQ_CPU_TIME_BACKWARD] for rdata in result[RDATAS]]
+            )*0.001,
+            POSTEQ_CPU_TIME: sum(
+                [rdata[POSTEQ_CPU_TIME] for rdata in result[RDATAS]]
+            )*0.001,
+            POSTEQ_CPU_TIME_BACKWARD: sum(
+                [rdata[POSTEQ_CPU_TIME_BACKWARD] for rdata in result[RDATAS]]
+            )*0.001,
+        }
+        return values
+
+    @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 [s].
+
+        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 pre-equilibration time, [s].
+
+        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 pre-equilibration time of the backward problem, [s].
+
+        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_BACKWARD, ix)
+
+    @trace_wrap
+    def get_posteq_time_trace(
+        self, ix: Union[int, Sequence[int], None] = None, trim: bool = False
+    ) -> Union[Sequence[float], float]:
+        """
+        Cumulative post-equilibration time [s].
+
+        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 post-equilibration time of the backward problem [s].
+
+        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_BACKWARD, ix)
+
+
+class CsvAmiciHistory(CsvHistory):
+    """
+    Stores history extended by AMICI-specific time traces in a CSV file.
+
+    Stores AMICI-specific traces of total simulation time, pre-equilibration
+    time and post-equilibration time.
+
+    Parameters
+    ----------
+    file:
+        CSV file name.
+    x_names:
+        Parameter names.
+    options:
+        History options.
+    load_from_file:
+        If True, history will be initialized from data in the specified file.
+    """
+
+    def __init__(
+        self,
+        file: str,
+        x_names: Sequence[str] = None,
+        options: Union[HistoryOptions, dict] = None,
+        load_from_file: bool = False,
+    ):
+        super().__init__(file, x_names, options, load_from_file=load_from_file)
+
+    def _trace_columns(self) -> list[tuple]:
+        columns = super()._trace_columns()
+        return columns + [
+            (c, np.nan)
+            for c in [
+                CPU_TIME_TOTAL,
+                PREEQ_CPU_TIME,
+                PREEQ_CPU_TIME_BACKWARD,
+                POSTEQ_CPU_TIME,
+                POSTEQ_CPU_TIME_BACKWARD,
+            ]
+        ]
+
+    def _simulation_to_values(self, result, used_time):
+        values = super()._simulation_to_values(result, used_time)
+        # default unit for time in amici is [ms], converted to [s]
+        values |= {
+            CPU_TIME_TOTAL: sum(
+                [rdata[CPU_TIME_TOTAL] for rdata in result[RDATAS]]
+            )*0.001,
+            PREEQ_CPU_TIME: sum(
+                [rdata[PREEQ_CPU_TIME] for rdata in result[RDATAS]]
+            )*0.001,
+            PREEQ_CPU_TIME_BACKWARD: sum(
+                [rdata[PREEQ_CPU_TIME_BACKWARD] for rdata in result[RDATAS]]
+            )*0.001,
+            POSTEQ_CPU_TIME: sum(
+                [rdata[POSTEQ_CPU_TIME] for rdata in result[RDATAS]]
+            )*0.001,
+            POSTEQ_CPU_TIME_BACKWARD: sum(
+                [rdata[POSTEQ_CPU_TIME_BACKWARD] for rdata in result[RDATAS]]
+            )*0.001,
+        }
+        return values
+
+    @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 [s].
+
+        Takes as parameter an index or indices and returns corresponding trace
+        values. If only a single value is requested, the list is flattened.
+        """
+        return list(self._trace[(CPU_TIME_TOTAL, np.nan)].values[ix])
+
+    @trace_wrap
+    def get_preeq_time_trace(
+        self, ix: Union[int, Sequence[int], None] = None, trim: bool = False
+    ) -> Union[Sequence[float], float]:
+        """
+        Cumulative pre-equilibration time [s].
+
+        Takes as parameter an index or indices and returns corresponding trace
+        values. If only a single value is requested, the list is flattened.
+        """
+        return list(self._trace[(PREEQ_CPU_TIME, np.nan)].values[ix])
+
+    @trace_wrap
+    def get_preeq_timeB_trace(
+        self, ix: Union[int, Sequence[int], None] = None, trim: bool = False
+    ) -> Union[Sequence[float], float]:
+        """
+        Cumulative pre-equilibration time of the backward problem [s].
+
+        Takes as parameter an index or indices and returns corresponding trace
+        values. If only a single value is requested, the list is flattened.
+        """
+        return list(self._trace[(PREEQ_CPU_TIME_BACKWARD, np.nan)].values[ix])
+
+    @trace_wrap
+    def get_posteq_time_trace(
+        self, ix: Union[int, Sequence[int], None] = None, trim: bool = False
+    ) -> Union[Sequence[float], float]:
+        """
+        Cumulative post-equilibration time [s].
+
+        Takes as parameter an index or indices and returns corresponding trace
+        values. If only a single value is requested, the list is flattened.
+        """
+        return list(self._trace[(POSTEQ_CPU_TIME, np.nan)].values[ix])
+
+    @trace_wrap
+    def get_posteq_timeB_trace(
+        self, ix: Union[int, Sequence[int], None] = None, trim: bool = False
+    ) -> Union[Sequence[float], float]:
+        """
+        Cumulative post-equilibration time of the backward problem [s].
+
+        Takes as parameter an index or indices and returns corresponding trace
+        values. If only a single value is requested, the list is flattened.
+        """
+        return list(self._trace[(POSTEQ_CPU_TIME_BACKWARD, np.nan)].values[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/csv.py

@@ -100,6 +100,20 @@ def finalize(self, message: str = None, exitflag: str = None):
         super().finalize(message=message, exitflag=exitflag)
         self._save_trace(finalize=True)
 
+    def _simulation_to_values(self, result, used_time):
+        values = {
+            TIME: used_time,
+            N_FVAL: self._n_fval,
+            N_GRAD: self._n_grad,
+            N_HESS: self._n_hess,
+            N_RES: self._n_res,
+            N_SRES: self._n_sres,
+            FVAL: result[FVAL],
+            RES: result[RES],
+            HESS: result[HESS],
+        }
+        return values
+
     def _update_trace(
         self,
         x: np.ndarray,
@@ -127,17 +141,7 @@ def _update_trace(
             name=len(self._trace), index=self._trace.columns, dtype='object'
         )
 
-        values = {
-            TIME: used_time,
-            N_FVAL: self._n_fval,
-            N_GRAD: self._n_grad,
-            N_HESS: self._n_hess,
-            N_RES: self._n_res,
-            N_SRES: self._n_sres,
-            FVAL: result[FVAL],
-            RES: result[RES],
-            HESS: result[HESS],
-        }
+        values = self._simulation_to_values(result, used_time)
 
         for var, val in values.items():
             row[(var, np.nan)] = val
@@ -162,12 +166,8 @@ def _update_trace(
         # save trace to file
         self._save_trace()
 
-    def _init_trace(self, x: np.ndarray):
-        """Initialize the trace."""
-        if self.x_names is None:
-            self.x_names = [f'x{i}' for i, _ in enumerate(x)]
-
-        columns: list[tuple] = [
+    def _trace_columns(self) -> list[tuple]:
+        return [
             (c, np.nan)
             for c in [
                 TIME,
@@ -183,6 +183,13 @@ def _init_trace(self, x: np.ndarray):
             ]
         ]
 
+    def _init_trace(self, x: np.ndarray):
+        """Initialize the trace."""
+        if self.x_names is None:
+            self.x_names = [f'x{i}' for i, _ in enumerate(x)]
+
+        columns = self._trace_columns()
+
         for var in [X, GRAD]:
             if var == X or self.options[f'trace_record_{var}']:
                 columns.extend([(var, x_name) for x_name in self.x_names])

pypesto/history/generate.py

@@ -15,7 +15,7 @@
 def create_history(
     id: str,
     x_names: Sequence[str],
-    options: HistoryOptions,
+    options: HistoryOptions
 ) -> HistoryBase:
     """Create a :class:`HistoryBase` object; Factory method.
 
@@ -47,7 +47,9 @@ def create_history(
 
     # create history type based on storage type
     if suffix in SUFFIXES_CSV:
-        return CsvHistory(x_names=x_names, file=storage_file, options=options)
+        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)
     else: