mxlpy 0.22.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

mxlpy/scan.py

@@ -15,16 +15,21 @@ Functions:
 
 from __future__ import annotations
 
-from dataclasses import dataclass
 from functools import partial
-from typing import TYPE_CHECKING, Protocol, Self, cast
+from typing import TYPE_CHECKING, Protocol
 
 import numpy as np
 import pandas as pd
 
 from mxlpy.parallel import Cache, parallelise
-from mxlpy.simulator import Result, Simulator
-from mxlpy.types import IntegratorType, ProtocolByPars, SteadyStates, TimeCourseByPars
+from mxlpy.simulator import Simulator
+from mxlpy.types import (
+    IntegratorType,
+    ProtocolScan,
+    Result,
+    SteadyStateScan,
+    TimeCourseScan,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -34,14 +39,14 @@ if TYPE_CHECKING:
 
 
 __all__ = [
+    "ProtocolTimeCourseWorker",
     "ProtocolWorker",
     "SteadyStateWorker",
-    "TimeCourse",
     "TimeCourseWorker",
-    "TimePoint",
+    "protocol",
+    "protocol_time_course",
     "steady_state",
     "time_course",
-    "time_course_over_protocol",
 ]
 
 
@@ -67,206 +72,6 @@ def _update_parameters_and_initial_conditions[T](
     return fn(model)
 
 
-def _empty_conc_series(model: Model) -> pd.Series:
-    """Create an empty concentration series for the model.
-
-    Args:
-        model: Model instance to generate the series for.
-
-    Returns:
-        pd.Series: Series with NaN values for each model variable.
-
-    """
-    return pd.Series(
-        data=np.full(shape=len(model.get_variable_names()), fill_value=np.nan),
-        index=model.get_variable_names(),
-    )
-
-
-def _empty_flux_series(model: Model) -> pd.Series:
-    """Create an empty flux series for the model.
-
-    Args:
-        model: Model instance to generate the series for.
-
-    Returns:
-        pd.Series: Series with NaN values for each model reaction.
-
-    """
-    return pd.Series(
-        data=np.full(shape=len(model.get_reaction_names()), fill_value=np.nan),
-        index=model.get_reaction_names(),
-    )
-
-
-def _empty_conc_df(model: Model, time_points: Array) -> pd.DataFrame:
-    """Create an empty concentration DataFrame for the model over given time points.
-
-    Args:
-        model: Model instance to generate the DataFrame for.
-        time_points: Array of time points.
-
-    Returns:
-        pd.DataFrame: DataFrame with NaN values for each model variable at each time point.
-
-    """
-    return pd.DataFrame(
-        data=np.full(
-            shape=(len(time_points), len(model.get_variable_names())),
-            fill_value=np.nan,
-        ),
-        index=time_points,
-        columns=model.get_variable_names(),
-    )
-
-
-def _empty_flux_df(model: Model, time_points: Array) -> pd.DataFrame:
-    """Create an empty concentration DataFrame for the model over given time points.
-
-    Args:
-        model: Model instance to generate the DataFrame for.
-        time_points: Array of time points.
-
-    Returns:
-        pd.DataFrame: DataFrame with NaN values for each model variable at each time point.
-
-    """
-    return pd.DataFrame(
-        data=np.full(
-            shape=(len(time_points), len(model.get_reaction_names())),
-            fill_value=np.nan,
-        ),
-        index=time_points,
-        columns=model.get_reaction_names(),
-    )
-
-
-###############################################################################
-# Single returns
-###############################################################################
-
-
-@dataclass(slots=True)
-class TimePoint:
-    """Represents a single time point in a simulation.
-
-    Attributes:
-        concs: Series of concentrations at the time point.
-        fluxes: Series of fluxes at the time point.
-
-    Args:
-        model: Model instance to generate the time point for.
-        concs: DataFrame of concentrations (default: None).
-        fluxes: DataFrame of fluxes (default: None).
-        idx: Index of the time point in the DataFrame (default: -1).
-
-    """
-
-    variables: pd.Series
-    fluxes: pd.Series
-
-    @classmethod
-    def from_result(
-        cls,
-        *,
-        model: Model,
-        result: Result | None,
-        idx: int = -1,
-    ) -> Self:
-        """Initialize the Scan object.
-
-        Args:
-            model: The model object.
-            result: Result of the simulation
-            idx: Index to select specific row from concs and fluxes DataFrames.
-
-        """
-        if result is None:
-            return cls(
-                variables=_empty_conc_series(model),
-                fluxes=_empty_flux_series(model),
-            )
-
-        return cls(
-            variables=result.variables.iloc[idx],
-            fluxes=result.fluxes.iloc[idx],
-        )
-
-    @property
-    def results(self) -> pd.Series:
-        """Get the combined results of concentrations and fluxes.
-
-        Example:
-            >>> time_point.results
-            x1    1.0
-            x2    0.5
-            v1    0.1
-            v2    0.2
-
-        Returns:
-            pd.Series: Combined series of concentrations and fluxes.
-
-        """
-        return pd.concat((self.variables, self.fluxes), axis=0)
-
-
-@dataclass(slots=True)
-class TimeCourse:
-    """Represents a time course in a simulation.
-
-    Attributes:
-        variables: DataFrame of concentrations over time.
-        fluxes: DataFrame of fluxes over time.
-
-    """
-
-    variables: pd.DataFrame
-    fluxes: pd.DataFrame
-
-    @classmethod
-    def from_scan(
-        cls,
-        *,
-        model: Model,
-        time_points: Array,
-        result: Result | None,
-    ) -> Self:
-        """Initialize the Scan object.
-
-        Args:
-            model (Model): The model object.
-            time_points (Array): Array of time points.
-            result: Result of the simulation
-
-        """
-        if result is None:
-            return cls(
-                _empty_conc_df(model, time_points),
-                _empty_flux_df(model, time_points),
-            )
-        return cls(
-            result.variables,
-            result.fluxes,
-        )
-
-    @property
-    def results(self) -> pd.DataFrame:
-        """Get the combined results of concentrations and fluxes over time.
-
-        Examples:
-            >>> time_course.results
-            Time   x1     x2     v1     v2
-            0.0   1.0   1.00   1.00   1.00
-            0.1   0.9   0.99   0.99   0.99
-            0.2   0.8   0.99   0.99   0.99
-
-        Returns:
-            pd.DataFrame: Combined DataFrame of concentrations and fluxes.
-
-        """
-        return pd.concat((self.variables, self.fluxes), axis=1)
-
-
 ###############################################################################
 # Workers
 ###############################################################################
@@ -282,7 +87,7 @@ class SteadyStateWorker(Protocol):
         rel_norm: bool,
         integrator: IntegratorType | None,
         y0: dict[str, float] | None,
-    ) -> TimePoint:
+    ) -> Result:
         """Call the worker function."""
         ...
 
@@ -297,7 +102,7 @@ class TimeCourseWorker(Protocol):
         *,
         integrator: IntegratorType | None,
         y0: dict[str, float] | None,
-    ) -> TimeCourse:
+    ) -> Result:
         """Call the worker function."""
         ...
 
@@ -313,7 +118,23 @@ class ProtocolWorker(Protocol):
         integrator: IntegratorType | None,
         y0: dict[str, float] | None,
         time_points_per_step: int = 10,
-    ) -> TimeCourse:
+    ) -> Result:
+        """Call the worker function."""
+        ...
+
+
+class ProtocolTimeCourseWorker(Protocol):
+    """Worker function for protocol-based simulations."""
+
+    def __call__(
+        self,
+        model: Model,
+        protocol: pd.DataFrame,
+        time_points: Array,
+        *,
+        integrator: IntegratorType | None,
+        y0: dict[str, float] | None,
+    ) -> Result:
         """Call the worker function."""
         ...
 
@@ -324,7 +145,7 @@ def _steady_state_worker(
     rel_norm: bool,
     integrator: IntegratorType | None,
     y0: dict[str, float] | None,
-) -> TimePoint:
+) -> Result:
     """Simulate the model to steady state and return concentrations and fluxes.
 
     Args:
@@ -345,7 +166,9 @@ def _steady_state_worker(
         )
     except ZeroDivisionError:
         res = None
-    return TimePoint.from_result(model=model, result=res)
+    return (
+        Result.default(model=model, time_points=np.array([0.0])) if res is None else res
+    )
 
 
 def _time_course_worker(
@@ -353,7 +176,7 @@ def _time_course_worker(
     time_points: Array,
     y0: dict[str, float] | None,
     integrator: IntegratorType | None,
-) -> TimeCourse:
+) -> Result:
     """Simulate the model to steady state and return concentrations and fluxes.
 
     Args:
@@ -374,11 +197,7 @@ def _time_course_worker(
         )
     except ZeroDivisionError:
         res = None
-    return TimeCourse.from_scan(
-        model=model,
-        time_points=time_points,
-        result=res,
-    )
+    return Result.default(model=model, time_points=time_points) if res is None else res
 
 
 def _protocol_worker(
@@ -388,7 +207,7 @@ def _protocol_worker(
     integrator: IntegratorType | None,
     y0: dict[str, float] | None,
     time_points_per_step: int = 10,
-) -> TimeCourse:
+) -> Result:
     """Simulate the model over a protocol and return concentrations and fluxes.
 
     Args:
@@ -419,11 +238,43 @@ def _protocol_worker(
         protocol.index[-1].total_seconds(),
         len(protocol) * time_points_per_step,
     )
-    return TimeCourse.from_scan(
-        model=model,
-        time_points=time_points,
-        result=res,
-    )
+    return Result.default(model=model, time_points=time_points) if res is None else res
+
+
+def _protocol_time_course_worker(
+    model: Model,
+    protocol: pd.DataFrame,
+    time_points: Array,
+    *,
+    integrator: IntegratorType | None,
+    y0: dict[str, float] | None,
+) -> Result:
+    """Simulate the model over a protocol and return concentrations and fluxes.
+
+    Args:
+        model: Model instance to simulate.
+        y0: Initial conditions as a dictionary {species: value}.
+        protocol: DataFrame containing the protocol steps.
+        time_points: Time points where to return the simulation
+        integrator: Integrator function to use for steady state calculation
+
+    Returns:
+        TimeCourse: Object containing protocol series concentrations and fluxes.
+
+    """
+    try:
+        res = (
+            Simulator(model, integrator=integrator, y0=y0)
+            .simulate_protocol_time_course(
+                protocol=protocol,
+                time_points=time_points,
+            )
+            .get_result()
+        )
+    except ZeroDivisionError:
+        res = None
+
+    return Result.default(model=model, time_points=time_points) if res is None else res
 
 
 def steady_state(
@@ -436,7 +287,7 @@ def steady_state(
     cache: Cache | None = None,
     worker: SteadyStateWorker = _steady_state_worker,
     integrator: IntegratorType | None = None,
-) -> SteadyStates:
+) -> SteadyStateScan:
     """Get steady-state results over supplied values.
 
     Args:
@@ -492,16 +343,16 @@ def steady_state(
         cache=cache,
         parallel=parallel,
     )
-    concs = pd.DataFrame({k: v.variables.T for k, v in res}).T
-    fluxes = pd.DataFrame({k: v.fluxes.T for k, v in res}).T
-    idx = (
-        pd.Index(to_scan.iloc[:, 0])
-        if to_scan.shape[1] == 1
-        else pd.MultiIndex.from_frame(to_scan)
+
+    return SteadyStateScan(
+        raw_index=(
+            pd.Index(to_scan.iloc[:, 0])
+            if to_scan.shape[1] == 1
+            else pd.MultiIndex.from_frame(to_scan)
+        ),
+        raw_results=[i[1] for i in res],
+        to_scan=to_scan,
     )
-    concs.index = idx
-    fluxes.index = idx
-    return SteadyStates(variables=concs, fluxes=fluxes, parameters=to_scan)
 
 
 def time_course(
@@ -514,7 +365,7 @@ def time_course(
     cache: Cache | None = None,
     integrator: IntegratorType | None = None,
     worker: TimeCourseWorker = _time_course_worker,
-) -> TimeCourseByPars:
+) -> TimeCourseScan:
     """Get time course for each supplied parameter.
 
     Examples:
@@ -585,16 +436,13 @@ def time_course(
         cache=cache,
         parallel=parallel,
     )
-    concs = cast(dict, {k: v.variables for k, v in res})
-    fluxes = cast(dict, {k: v.fluxes for k, v in res})
-    return TimeCourseByPars(
-        parameters=to_scan,
-        variables=pd.concat(concs, names=["n", "time"]),
-        fluxes=pd.concat(fluxes, names=["n", "time"]),
+    return TimeCourseScan(
+        to_scan=to_scan,
+        raw_results=dict(res),
     )
 
 
-def time_course_over_protocol(
+def protocol(
     model: Model,
     *,
     to_scan: pd.DataFrame,
@@ -605,7 +453,7 @@ def time_course_over_protocol(
     cache: Cache | None = None,
     worker: ProtocolWorker = _protocol_worker,
     integrator: IntegratorType | None = None,
-) -> ProtocolByPars:
+) -> ProtocolScan:
     """Get protocol series for each supplied parameter.
 
     Examples:
@@ -656,11 +504,77 @@ def time_course_over_protocol(
         cache=cache,
         parallel=parallel,
     )
-    concs = cast(dict, {k: v.variables for k, v in res})
-    fluxes = cast(dict, {k: v.fluxes for k, v in res})
-    return ProtocolByPars(
-        parameters=to_scan,
+    return ProtocolScan(
+        to_scan=to_scan,
+        protocol=protocol,
+        raw_results=dict(res),
+    )
+
+
+def protocol_time_course(
+    model: Model,
+    *,
+    to_scan: pd.DataFrame,
+    protocol: pd.DataFrame,
+    time_points: Array,
+    y0: dict[str, float] | None = None,
+    parallel: bool = True,
+    cache: Cache | None = None,
+    worker: ProtocolTimeCourseWorker = _protocol_time_course_worker,
+    integrator: IntegratorType | None = None,
+) -> ProtocolScan:
+    """Get protocol series for each supplied parameter.
+
+    Examples:
+        >>> scan.time_course_over_protocol(
+        ...     model,
+        ...     parameters=pd.DataFrame({"k2": np.linspace(1, 2, 11)}),
+        ...     protocol=make_protocol(
+        ...         {
+        ...             1: {"k1": 1},
+        ...             2: {"k1": 2},
+        ...         }
+        ...     ),
+        ... )
+
+    Args:
+        model: Model instance to simulate.
+        to_scan: DataFrame containing parameter or initial values to scan.
+        protocol: Protocol to follow for the simulation.
+        time_points: Time points where to return simulation results
+        y0: Initial conditions as a dictionary {variable: value}.
+        parallel: Whether to execute in parallel (default: True).
+        cache: Optional cache to store and retrieve results.
+        worker: Worker function to use for the simulation.
+        integrator: Integrator function to use for steady state calculation
+
+    Returns:
+        TimeCourseByPars: Protocol series results for each parameter set.
+
+    """
+    # We update the initial conditions separately here, because `to_scan` might also
+    # contain initial conditions.
+    if y0 is not None:
+        model.update_variables(y0)
+
+    res = parallelise(
+        partial(
+            _update_parameters_and_initial_conditions,
+            fn=partial(
+                worker,
+                protocol=protocol,
+                time_points=time_points,
+                integrator=integrator,
+                y0=None,
+            ),
+            model=model,
+        ),
+        inputs=list(to_scan.iterrows()),
+        cache=cache,
+        parallel=parallel,
+    )
+    return ProtocolScan(
+        to_scan=to_scan,
         protocol=protocol,
-        variables=pd.concat(concs, names=["n", "time"]),
-        fluxes=pd.concat(fluxes, names=["n", "time"]),
+        raw_results=dict(res),
     )