modelbase2 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

modelbase2/simulator.py

@@ -10,7 +10,7 @@ Classes:
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Literal, Self, cast, overload
 
 import numpy as np
@@ -18,13 +18,13 @@ import pandas as pd
 
 from modelbase2.integrators import DefaultIntegrator
 
-__all__ = ["Simulator"]
+__all__ = ["Result", "Simulator"]
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Callable, Iterator
 
     from modelbase2.model import Model
-    from modelbase2.types import ArrayLike, IntegratorProtocol
+    from modelbase2.types import Array, ArrayLike, IntegratorProtocol
 
 
 def _normalise_split_results(
@@ -56,6 +56,240 @@ def _normalise_split_results(
     return results
 
 
+@dataclass(slots=True)
+class Result:
+    """Simulation results."""
+
+    model: Model
+    _raw_variables: list[pd.DataFrame]
+    _parameters: list[dict[str, float]]
+    _dependent: list[pd.DataFrame] = field(default_factory=list)
+
+    @property
+    def variables(self) -> pd.DataFrame:
+        """Simulation variables."""
+        return self.get_variables(
+            include_derived=True,
+            include_readouts=True,
+            concatenated=True,
+            normalise=None,
+        )
+
+    @property
+    def fluxes(self) -> pd.DataFrame:
+        """Simulation fluxes."""
+        return self.get_fluxes()
+
+    def __iter__(self) -> Iterator[pd.DataFrame]:
+        """Iterate over the concentration and flux response coefficients."""
+        return iter((self.variables, self.fluxes))
+
+    def _get_dependent(
+        self,
+        *,
+        include_readouts: bool = True,
+    ) -> list[pd.DataFrame]:
+        # Already computed
+        if len(self._dependent) > 0:
+            return self._dependent
+
+        # Compute new otherwise
+        for res, p in zip(self._raw_variables, self._parameters, strict=True):
+            self.model.update_parameters(p)
+            self._dependent.append(
+                self.model.get_dependent_time_course(
+                    variables=res,
+                    include_readouts=include_readouts,
+                )
+            )
+        return self._dependent
+
+    def _select_variables(
+        self,
+        dependent: list[pd.DataFrame],
+        *,
+        include_derived: bool,
+        include_readouts: bool,
+    ) -> list[pd.DataFrame]:
+        names = self.model.get_variable_names()
+        if include_derived:
+            names.extend(self.model.get_derived_variable_names())
+        if include_readouts:
+            names.extend(self.model.get_readout_names())
+        return [i.loc[:, names] for i in dependent]
+
+    def _select_fluxes(
+        self,
+        dependent: list[pd.DataFrame],
+        *,
+        include_surrogates: bool,
+    ) -> list[pd.DataFrame]:
+        names = self.model.get_reaction_names()
+        if include_surrogates:
+            names.extend(self.model.get_surrogate_reaction_names())
+        return [i.loc[:, names] for i in dependent]
+
+    def _adjust_data(
+        self,
+        data: list[pd.DataFrame],
+        normalise: float | ArrayLike | None = None,
+        *,
+        concatenated: bool = True,
+    ) -> pd.DataFrame | list[pd.DataFrame]:
+        if normalise is not None:
+            data = _normalise_split_results(data, normalise=normalise)
+        if concatenated:
+            return pd.concat(data, axis=0)
+        return data
+
+    @overload
+    def get_variables(  # type: ignore
+        self,
+        *,
+        include_derived: bool = True,
+        include_readouts: bool = True,
+        concatenated: Literal[False],
+        normalise: float | ArrayLike | None = None,
+    ) -> list[pd.DataFrame]: ...
+
+    @overload
+    def get_variables(
+        self,
+        *,
+        include_derived: bool = True,
+        include_readouts: bool = True,
+        concatenated: Literal[True],
+        normalise: float | ArrayLike | None = None,
+    ) -> pd.DataFrame: ...
+
+    @overload
+    def get_variables(
+        self,
+        *,
+        include_derived: bool = True,
+        include_readouts: bool = True,
+        concatenated: bool = True,
+        normalise: float | ArrayLike | None = None,
+    ) -> pd.DataFrame: ...
+
+    def get_variables(
+        self,
+        *,
+        include_derived: bool = True,
+        include_readouts: bool = True,
+        concatenated: bool = True,
+        normalise: float | ArrayLike | None = None,
+    ) -> pd.DataFrame | list[pd.DataFrame]:
+        """Get the variables over time.
+
+        Examples:
+            >>> Result().get_variables()
+            Time            ATP      NADPH
+            0.000000   1.000000   1.000000
+            0.000100   0.999900   0.999900
+            0.000200   0.999800   0.999800
+
+        """
+        if not include_derived and not include_readouts:
+            return self._adjust_data(
+                self._raw_variables,
+                normalise=normalise,
+                concatenated=concatenated,
+            )
+
+        variables = self._select_variables(
+            self._get_dependent(),
+            include_derived=include_derived,
+            include_readouts=include_readouts,
+        )
+        return self._adjust_data(
+            variables, normalise=normalise, concatenated=concatenated
+        )
+
+    @overload
+    def get_fluxes(  # type: ignore
+        self,
+        *,
+        include_surrogates: bool = True,
+        normalise: float | ArrayLike | None = None,
+        concatenated: Literal[False],
+    ) -> list[pd.DataFrame]: ...
+
+    @overload
+    def get_fluxes(
+        self,
+        *,
+        include_surrogates: bool = True,
+        normalise: float | ArrayLike | None = None,
+        concatenated: Literal[True],
+    ) -> pd.DataFrame: ...
+
+    @overload
+    def get_fluxes(
+        self,
+        *,
+        include_surrogates: bool = True,
+        normalise: float | ArrayLike | None = None,
+        concatenated: bool = True,
+    ) -> pd.DataFrame: ...
+
+    def get_fluxes(
+        self,
+        *,
+        include_surrogates: bool = True,
+        normalise: float | ArrayLike | None = None,
+        concatenated: bool = True,
+    ) -> pd.DataFrame | list[pd.DataFrame]:
+        """Get the flux results.
+
+        Examples:
+            >>> Result.get_fluxes()
+            Time             v1         v2
+            0.000000   1.000000   10.00000
+            0.000100   0.999900   9.999000
+            0.000200   0.999800   9.998000
+
+        Returns:
+            pd.DataFrame: DataFrame of fluxes.
+
+        """
+        fluxes = self._select_fluxes(
+            self._get_dependent(),
+            include_surrogates=include_surrogates,
+        )
+        return self._adjust_data(
+            fluxes,
+            normalise=normalise,
+            concatenated=concatenated,
+        )
+
+    def get_combined(self) -> pd.DataFrame:
+        """Get the variables and fluxes as a single pandas.DataFrame.
+
+        Examples:
+            >>> Result.get_combined()
+            Time            ATP      NADPH         v1         v2
+            0.000000   1.000000   1.000000   1.000000   10.00000
+            0.000100   0.999900   0.999900   0.999900   9.999000
+            0.000200   0.999800   0.999800   0.999800   9.998000
+
+        Returns:
+            pd.DataFrame: DataFrame of fluxes.
+
+        """
+        return pd.concat((self.variables, self.fluxes), axis=1)
+
+    def get_new_y0(self) -> dict[str, float]:
+        """Get the new initial conditions after the simulation.
+
+        Examples:
+            >>> Simulator(model).simulate_to_steady_state().get_new_y0()
+            {"ATP": 1.0, "NADPH": 1.0}
+
+        """
+        return dict(self.variables.iloc[-1])
+
+
 @dataclass(
     init=False,
     slots=True,
@@ -68,19 +302,23 @@ class Simulator:
         model: Model instance to simulate.
         y0: Initial conditions for the simulation.
         integrator: Integrator protocol to use for the simulation.
-        concs: List of DataFrames containing concentration results.
-        args: List of DataFrames containing argument values.
+        variables: List of DataFrames containing concentration results.
+        dependent: List of DataFrames containing argument values.
         simulation_parameters: List of dictionaries containing simulation parameters.
 
     """
 
     model: Model
-    y0: ArrayLike
+    y0: dict[str, float]
     integrator: IntegratorProtocol
-    concs: list[pd.DataFrame] | None
-    args: list[pd.DataFrame] | None
+    variables: list[pd.DataFrame] | None
+    dependent: list[pd.DataFrame] | None
     simulation_parameters: list[dict[str, float]] | None
 
+    # For resets (e.g. update variable)
+    _integrator_type: Callable[[Callable, ArrayLike], IntegratorProtocol]
+    _time_shift: float | None
+
     def __init__(
         self,
         model: Model,
@@ -94,66 +332,44 @@ class Simulator:
         """Initialize the Simulator.
 
         Args:
-            model (Model): The model to be simulated.
-            y0 (dict[str, float] | None, optional): Initial conditions for the model variables.
-                If None, the initial conditions are obtained from the model. Defaults to None.
-            integrator (Callable[[Callable, ArrayLike], IntegratorProtocol], optional): The integrator to use for the simulation.
-                Defaults to DefaultIntegrator.
-            test_run (bool, optional): If True, performs a test run to ensure the model's methods
-                (get_full_concs, get_fluxes, get_right_hand_side) work correctly with the initial conditions.
-                Defaults to True.
+            model: The model to be simulated.
+            y0: Initial conditions for the model variables.
+                If None, the initial conditions are obtained from the model.
+            integrator: The integrator to use for the simulation.
+            test_run (bool, optional): If True, performs a test run for better error messages
 
         """
         self.model = model
-        y0 = model.get_initial_conditions() if y0 is None else y0
-        self.y0 = [y0[k] for k in model.get_variable_names()]
+        self.y0 = model.get_initial_conditions() if y0 is None else y0
 
-        self.integrator = integrator(self.model, self.y0)
-        self.concs = None
-        self.args = None
+        self._integrator_type = integrator
+        self._time_shift = None
+        self.variables = None
         self.simulation_parameters = None
 
         if test_run:
-            y0 = dict(zip(model.get_variable_names(), self.y0, strict=True))
-            self.model.get_full_concs(y0, 0)
-            self.model.get_fluxes(y0, 0)
-            self.model.get_right_hand_side(y0, 0)
-
-    def _save_simulation_results(
-        self,
-        *,
-        results: pd.DataFrame,
-        skipfirst: bool,
-    ) -> None:
-        """Save simulation results.
-
-        Args:
-            results: DataFrame containing the simulation results.
-            skipfirst: Whether to skip the first row of results.
+            self.model.get_right_hand_side(self.y0, time=0)
 
-        """
-        if self.concs is None:
-            self.concs = [results]
-        elif skipfirst:
-            self.concs.append(results.iloc[1:, :])
-        else:
-            self.concs.append(results)
+        self._initialise_integrator()
 
-        if self.simulation_parameters is None:
-            self.simulation_parameters = []
-        self.simulation_parameters.append(self.model.parameters)
+    def _initialise_integrator(self) -> None:
+        y0 = self.y0
+        self.integrator = self._integrator_type(
+            self.model,
+            [y0[k] for k in self.model.get_variable_names()],
+        )
 
     def clear_results(self) -> None:
         """Clear simulation results."""
-        self.concs = None
-        self.args = None
+        self.variables = None
+        self.dependent = None
         self.simulation_parameters = None
-        if self.integrator is not None:
-            self.integrator.reset()
+        self._time_shift = None
+        self._initialise_integrator()
 
     def _handle_simulation_results(
         self,
-        time: ArrayLike | None,
+        time: Array | None,
         results: ArrayLike | None,
         *,
         skipfirst: bool,
@@ -172,15 +388,28 @@ class Simulator:
             self.clear_results()
             return
 
+        if self._time_shift is not None:
+            time += self._time_shift
+
         # NOTE: IMPORTANT!
-        # model._get_rhs sorts the return array by model.get_compounds()
+        # model._get_rhs sorts the return array by model.get_variable_names()
         # Do NOT change this ordering
         results_df = pd.DataFrame(
             results,
             index=time,
             columns=self.model.get_variable_names(),
         )
-        self._save_simulation_results(results=results_df, skipfirst=skipfirst)
+
+        if self.variables is None:
+            self.variables = [results_df]
+        elif skipfirst:
+            self.variables.append(results_df.iloc[1:, :])
+        else:
+            self.variables.append(results_df)
+
+        if self.simulation_parameters is None:
+            self.simulation_parameters = []
+        self.simulation_parameters.append(self.model.parameters)
 
     def simulate(
         self,
@@ -204,7 +433,11 @@ class Simulator:
             Self: The Simulator instance with updated results.
 
         """
+        if self._time_shift is not None:
+            t_end -= self._time_shift
+
         time, results = self.integrator.integrate(t_end=t_end, steps=steps)
+
         self._handle_simulation_results(time, results, skipfirst=True)
         return self
 
@@ -226,6 +459,10 @@ class Simulator:
             Self: The Simulator instance with updated results.
 
         """
+        if self._time_shift is not None:
+            time_points = np.array(time_points, dtype=float)
+            time_points -= self._time_shift
+
         time, results = self.integrator.integrate_time_course(time_points=time_points)
         self._handle_simulation_results(time, results, skipfirst=True)
         return self
@@ -259,7 +496,7 @@ class Simulator:
             rel_norm=rel_norm,
         )
         self._handle_simulation_results(
-            [time] if time is not None else None,
+            np.array([time], dtype=float) if time is not None else None,
             [results] if results is not None else None,  # type: ignore
             skipfirst=False,
         )
@@ -290,298 +527,37 @@ class Simulator:
             t_end = cast(pd.Timedelta, t_end)
             self.model.update_parameters(pars.to_dict())
             self.simulate(t_end.total_seconds(), steps=time_points_per_step)
-            if self.concs is None:
+            if self.variables is None:
                 break
         return self
 
-    def _get_args_vectorised(
-        self,
-        concs: list[pd.DataFrame],
-        params: list[dict[str, float]],
-        *,
-        include_readouts: bool = True,
-    ) -> list[pd.DataFrame]:
-        args: list[pd.DataFrame] = []
-
-        for res, p in zip(concs, params, strict=True):
-            self.model.update_parameters(p)
-            args.append(
-                self.model.get_args_time_course(
-                    concs=res,
-                    include_readouts=include_readouts,
-                )
-            )
-        return args
-
-    @overload
-    def get_concs(  # type: ignore
-        self,
-        *,
-        normalise: float | ArrayLike | None = None,
-        concatenated: Literal[False],
-    ) -> None | list[pd.DataFrame]: ...
-
-    @overload
-    def get_concs(
-        self,
-        *,
-        normalise: float | ArrayLike | None = None,
-        concatenated: Literal[True],
-    ) -> None | pd.DataFrame: ...
-
-    @overload
-    def get_concs(
-        self,
-        *,
-        normalise: float | ArrayLike | None = None,
-        concatenated: Literal[True] = True,
-    ) -> None | pd.DataFrame: ...
-
-    def get_concs(
-        self,
-        *,
-        normalise: float | ArrayLike | None = None,
-        concatenated: bool = True,
-    ) -> None | pd.DataFrame | list[pd.DataFrame]:
-        """Get the concentration results.
-
-        Examples:
-            >>> Simulator(model).get_concs()
-            Time            ATP      NADPH
-            0.000000   1.000000   1.000000
-            0.000100   0.999900   0.999900
-            0.000200   0.999800   0.999800
-
-        Returns:
-            pd.DataFrame: DataFrame of concentrations.
-
-        """
-        if self.concs is None:
-            return None
-
-        results = self.concs.copy()
-        if normalise is not None:
-            results = _normalise_split_results(results=results, normalise=normalise)
-        if concatenated:
-            return pd.concat(results, axis=0)
-
-        return results
-
-    @overload
-    def get_full_concs(  # type: ignore
-        self,
-        *,
-        normalise: float | ArrayLike | None = None,
-        concatenated: Literal[False],
-        include_readouts: bool = True,
-    ) -> list[pd.DataFrame] | None: ...
-
-    @overload
-    def get_full_concs(
-        self,
-        *,
-        normalise: float | ArrayLike | None = None,
-        concatenated: Literal[True],
-        include_readouts: bool = True,
-    ) -> pd.DataFrame | None: ...
-
-    @overload
-    def get_full_concs(
-        self,
-        *,
-        normalise: float | ArrayLike | None = None,
-        concatenated: bool = True,
-        include_readouts: bool = True,
-    ) -> pd.DataFrame | None: ...
-
-    def get_full_concs(
-        self,
-        *,
-        normalise: float | ArrayLike | None = None,
-        concatenated: bool = True,
-        include_readouts: bool = True,
-    ) -> pd.DataFrame | list[pd.DataFrame] | None:
-        """Get the full concentration results, including derived quantities.
+    def get_result(self) -> Result | None:
+        """Get result of the simulation.
 
         Examples:
-            >>> Simulator(model).get_full_concs()
+            >>> variables, fluxes = Simulator(model).simulate().get_result()
+            >>> variables
             Time            ATP      NADPH
             0.000000   1.000000   1.000000
             0.000100   0.999900   0.999900
             0.000200   0.999800   0.999800
-
-        Returns: DataFrame of full concentrations.
-
-        """
-        if (concs := self.concs) is None:
-            return None
-        if (params := self.simulation_parameters) is None:
-            return None
-        if (args := self.args) is None:
-            args = self._get_args_vectorised(concs, params)
-
-        names = (
-            self.model.get_variable_names() + self.model.get_derived_variable_names()
-        )
-        if include_readouts:
-            names.extend(self.model.get_readout_names())
-        full_concs = [i.loc[:, names] for i in args]
-        if normalise is not None:
-            full_concs = _normalise_split_results(
-                results=full_concs,
-                normalise=normalise,
-            )
-        if concatenated:
-            return pd.concat(full_concs, axis=0)
-        return full_concs
-
-    @overload
-    def get_fluxes(  # type: ignore
-        self,
-        *,
-        normalise: float | ArrayLike | None = None,
-        concatenated: Literal[False],
-    ) -> list[pd.DataFrame] | None: ...
-
-    @overload
-    def get_fluxes(
-        self,
-        *,
-        normalise: float | ArrayLike | None = None,
-        concatenated: Literal[True],
-    ) -> pd.DataFrame | None: ...
-
-    @overload
-    def get_fluxes(
-        self,
-        *,
-        normalise: float | ArrayLike | None = None,
-        concatenated: bool = True,
-    ) -> pd.DataFrame | None: ...
-
-    def get_fluxes(
-        self,
-        *,
-        normalise: float | ArrayLike | None = None,
-        concatenated: bool = True,
-    ) -> pd.DataFrame | list[pd.DataFrame] | None:
-        """Get the flux results.
-
-        Examples:
-            >>> Simulator(model).get_fluxes()
+            >>> fluxes
             Time             v1         v2
             0.000000   1.000000   10.00000
             0.000100   0.999900   9.999000
             0.000200   0.999800   9.998000
 
-        Returns:
-            pd.DataFrame: DataFrame of fluxes.
-
         """
-        if (concs := self.concs) is None:
+        if (variables := self.variables) is None:
             return None
-        if (params := self.simulation_parameters) is None:
+        if (parameters := self.simulation_parameters) is None:
             return None
-        if (args := self.args) is None:
-            args = self._get_args_vectorised(concs, params)
-
-        fluxes: list[pd.DataFrame] = []
-        for y, p in zip(args, params, strict=True):
-            self.model.update_parameters(p)
-            fluxes.append(self.model.get_fluxes_time_course(args=y))
-
-        if normalise is not None:
-            fluxes = _normalise_split_results(
-                results=fluxes,
-                normalise=normalise,
-            )
-        if concatenated:
-            return pd.concat(fluxes, axis=0)
-        return fluxes
-
-    def get_concs_and_fluxes(self) -> tuple[pd.DataFrame | None, pd.DataFrame | None]:
-        """Get the concentrations and fluxes.
-
-        Examples:
-            >>> Simulator(model).get_concs_and_fluxes()
-            (concs, fluxes)
-
-
-        Returns:
-            tuple[pd.DataFrame, pd.DataFrame]: Tuple of concentrations and fluxes.
-
-        """
-        return self.get_concs(), self.get_fluxes()
-
-    def get_full_concs_and_fluxes(
-        self,
-        *,
-        include_readouts: bool = True,
-    ) -> tuple[pd.DataFrame | None, pd.DataFrame | None]:
-        """Get the full concentrations and fluxes.
-
-        >>> Simulator(model).get_full_concs_and_fluxes()
-        (full_concs, full_fluxes)
-
-        Args:
-            include_readouts: Whether to include readouts in the results.
-
-        Returns:
-            Full concentrations and fluxes
-
-        """
-        return (
-            self.get_full_concs(include_readouts=include_readouts),
-            self.get_fluxes(),
+        return Result(
+            model=self.model,
+            _raw_variables=variables,
+            _parameters=parameters,
         )
 
-    def get_results(self) -> pd.DataFrame | None:
-        """Get the combined results of concentrations and fluxes.
-
-        Examples:
-            >>> Simulator(model).get_results()
-            Time            ATP      NADPH         v1         v2
-            0.000000   1.000000   1.000000   1.000000   1.000000
-            0.000100   0.999900   0.999900   0.999900   0.999900
-            0.000200   0.999800   0.999800   0.999800   0.999800
-
-        Returns:
-            pd.DataFrame: Combined DataFrame of concentrations and fluxes.
-
-        """
-        c, v = self.get_concs_and_fluxes()
-        if c is None or v is None:
-            return None
-        return pd.concat((c, v), axis=1)
-
-    def get_full_results(self) -> pd.DataFrame | None:
-        """Get the combined full results of concentrations and fluxes.
-
-        Examples:
-            >>> Simulator(model).get_full_results()
-            Time            ATP      NADPH         v1         v2
-            0.000000   1.000000   1.000000   1.000000   1.000000
-            0.000100   0.999900   0.999900   0.999900   0.999900
-            0.000200   0.999800   0.999800   0.999800   0.999800
-
-        """
-        c, v = self.get_full_concs_and_fluxes()
-        if c is None or v is None:
-            return None
-        return pd.concat((c, v), axis=1)
-
-    def get_new_y0(self) -> dict[str, float] | None:
-        """Get the new initial conditions after the simulation.
-
-        Examples:
-            >>> Simulator(model).get_new_y0()
-            {"ATP": 1.0, "NADPH": 1.0}
-
-        """
-        if (res := self.get_concs()) is None:
-            return None
-        return dict(res.iloc[-1])
-
     def update_parameter(self, parameter: str, value: float) -> Self:
         """Updates the value of a specified parameter in the model.
 
@@ -637,3 +613,38 @@ class Simulator:
         """
         self.model.scale_parameters(parameters)
         return self
+
+    def update_variable(self, variable: str, value: float) -> Self:
+        """Updates the value of a specified value in the simulation.
+
+        Examples:
+            >>> Simulator(model).update_variable("k1", 0.1)
+
+        Args:
+            variable: name of the model variable
+            value: new value
+
+        """
+        return self.update_variables({variable: value})
+
+    def update_variables(self, variables: dict[str, float]) -> Self:
+        """Updates the value of a specified value in the simulation.
+
+        Examples:
+            >>> Simulator(model).update_variables({"k1": 0.1})
+
+        Args:
+            variables: {variable: value} pairs
+
+        """
+        sim_variables = self.variables
+
+        # In case someone calls this before the first simulation
+        if sim_variables is None:
+            self.y0 = self.y0 | variables
+            return self
+
+        self.y0 = sim_variables[-1].iloc[-1, :].to_dict() | variables
+        self._time_shift = float(sim_variables[-1].index[-1])
+        self._initialise_integrator()
+        return self