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

mxlpy/__init__.py

@@ -52,17 +52,24 @@ from . import (
     plot,
     report,
     sbml,
+    scan,
     units,
 )
-from .integrators import DefaultIntegrator, Scipy
+from .integrators import DefaultIntegrator, Diffrax, Scipy
 from .label_map import LabelMapper
 from .linear_label_map import LinearLabelMapper
 from .mc import Cache
 from .model import Model
-from .scan import steady_state, time_course, time_course_over_protocol
 from .simulator import Simulator
 from .symbolic import SymbolicModel, to_symbolic_model
-from .types import Derived, IntegratorProtocol, Parameter, Variable, unwrap
+from .types import (
+    Derived,
+    InitialAssignment,
+    IntegratorProtocol,
+    Parameter,
+    Variable,
+    unwrap,
+)
 
 with contextlib.suppress(ImportError):
     from .integrators import Assimulo
@@ -84,6 +91,8 @@ __all__ = [
     "Cache",
     "DefaultIntegrator",
     "Derived",
+    "Diffrax",
+    "InitialAssignment",
     "IntegratorProtocol",
     "LabelMapper",
     "LinearLabelMapper",
@@ -107,11 +116,9 @@ __all__ = [
     "plot",
     "report",
     "sbml",
-    "steady_state",
+    "scan",
     "surrogates",
     "symbolic",
-    "time_course",
-    "time_course_over_protocol",
     "to_symbolic_model",
     "units",
     "unwrap",

mxlpy/carousel.py

@@ -9,6 +9,7 @@ from functools import partial
 from typing import TYPE_CHECKING
 
 import pandas as pd
+from wadler_lindig import pformat
 
 from mxlpy import parallel, scan
 
@@ -18,13 +19,17 @@ if TYPE_CHECKING:
     from collections.abc import Iterable, Mapping
 
     from mxlpy import Model
-    from mxlpy.types import Array, IntegratorType, RateFn
+    from mxlpy.types import Array, IntegratorType, RateFn, Result
 
 
 @dataclass
 class ReactionTemplate:
     """Template for a reaction in a model."""
 
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
     fn: RateFn
     args: list[str]
     additional_parameters: dict[str, float] = field(default_factory=dict)
@@ -35,11 +40,17 @@ class CarouselSteadyState:
     """Time course of a carousel simulation."""
 
     carousel: list[Model]
-    results: list[scan.TimePoint]
+    results: list[Result]
+
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
 
     def get_variables_by_model(self) -> pd.DataFrame:
         """Get the variables of the time course results, indexed by model."""
-        return pd.DataFrame({i: r.variables for i, r in enumerate(self.results)}).T
+        return pd.DataFrame(
+            {i: r.variables.iloc[-1] for i, r in enumerate(self.results)}
+        ).T
 
 
 @dataclass
@@ -47,7 +58,11 @@ class CarouselTimeCourse:
     """Time course of a carousel simulation."""
 
     carousel: list[Model]
-    results: list[scan.TimeCourse]
+    results: list[Result]
+
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
 
     def get_variables_by_model(self) -> pd.DataFrame:
         """Get the variables of the time course results, indexed by model."""
@@ -74,6 +89,10 @@ class Carousel:
 
     variants: list[Model]
 
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
     def __init__(
         self,
         model: Model,
@@ -113,7 +132,7 @@ class Carousel:
             results=results,
         )
 
-    def protocol_time_course(
+    def protocol(
         self,
         protocol: pd.DataFrame,
         *,
@@ -139,6 +158,34 @@ class Carousel:
             results=results,
         )
 
+    def protocol_time_course(
+        self,
+        protocol: pd.DataFrame,
+        time_points: Array,
+        *,
+        y0: dict[str, float] | None = None,
+        integrator: IntegratorType | None = None,
+    ) -> CarouselTimeCourse:
+        """Simulate the carousel of models over a protocol time course."""
+        results = [
+            i[1]
+            for i in parallel.parallelise(
+                partial(
+                    scan._protocol_time_course_worker,  # noqa: SLF001
+                    protocol=protocol,
+                    integrator=integrator,
+                    time_points=time_points,
+                    y0=y0,
+                ),
+                list(enumerate(self.variants)),
+            )
+        ]
+
+        return CarouselTimeCourse(
+            carousel=self.variants,
+            results=results,
+        )
+
     def steady_state(
         self,
         *,

mxlpy/compare.py

@@ -6,10 +6,11 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING, cast
 
 import pandas as pd
+from wadler_lindig import pformat
 
 from mxlpy import plot
-from mxlpy.simulator import Result, Simulator
-from mxlpy.types import unwrap
+from mxlpy.simulator import Simulator
+from mxlpy.types import Result, unwrap
 
 if TYPE_CHECKING:
     from mxlpy.model import Model
@@ -32,6 +33,10 @@ class SteadyStateComparison:
     res1: Result
     res2: Result
 
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
     @property
     def variables(self) -> pd.DataFrame:
         """Compare the steady state variables."""
@@ -93,6 +98,10 @@ class TimeCourseComparison:
     res1: Result
     res2: Result
 
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
     # @property
     # def variables(self) -> pd.DataFrame:
     #     """Compare the steady state variables."""
@@ -144,6 +153,10 @@ class ProtocolComparison:
     res2: Result
     protocol: pd.DataFrame
 
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
     # @property
     # def variables(self) -> pd.DataFrame:
     #     """Compare the steady state variables."""

mxlpy/experimental/diff.py

@@ -5,6 +5,8 @@ from __future__ import annotations
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
+from wadler_lindig import pformat
+
 from mxlpy.types import Derived
 
 if TYPE_CHECKING:
@@ -28,6 +30,10 @@ class DerivedDiff:
     args1: list[str] = field(default_factory=list)
     args2: list[str] = field(default_factory=list)
 
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
 
 @dataclass
 class ReactionDiff:
@@ -38,6 +44,10 @@ class ReactionDiff:
     stoichiometry1: dict[str, float | Derived] = field(default_factory=dict)
     stoichiometry2: dict[str, float | Derived] = field(default_factory=dict)
 
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
 
 @dataclass
 class ModelDiff:
@@ -56,6 +66,10 @@ class ModelDiff:
     different_readouts: dict[str, DerivedDiff] = field(default_factory=dict)
     different_derived: dict[str, DerivedDiff] = field(default_factory=dict)
 
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
     def __str__(self) -> str:
         """Return a human-readable string representation of the diff."""
         content = ["Model Diff", "----------"]

mxlpy/fit.py

@@ -18,6 +18,7 @@ from typing import TYPE_CHECKING, Protocol
 
 import numpy as np
 from scipy.optimize import minimize
+from wadler_lindig import pformat
 
 from mxlpy import parallel
 from mxlpy.simulator import Simulator
@@ -61,6 +62,10 @@ class MinResult:
     parameters: dict[str, float]
     residual: float
 
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
 
 @dataclass
 class FitResult:
@@ -70,6 +75,10 @@ class FitResult:
     best_pars: dict[str, float]
     loss: float
 
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
 
 @dataclass
 class CarouselFit:
@@ -77,6 +86,10 @@ class CarouselFit:
 
     fits: list[FitResult]
 
+    def __repr__(self) -> str:
+        """Return default representation."""
+        return pformat(self)
+
     def get_best_fit(self) -> FitResult:
         """Get the best fit from the carousel."""
         return min(self.fits, key=lambda x: x.loss)
@@ -561,6 +574,8 @@ def protocol_time_course(
 ) -> FitResult | None:
     """Fit model parameters to time course of experimental data.
 
+    Time points of protocol time course are taken from the data.
+
     Examples:
         >>> time_course(model, p0, data)
         {'k1': 0.1, 'k2': 0.2}
@@ -688,8 +703,10 @@ def carousel_time_course(
 ) -> CarouselFit:
     """Fit model parameters to time course of experimental data over a carousel.
 
+    Time points are taken from the data.
+
     Examples:
-        >>> carousel_steady_state(carousel, p0=p0, data=data)
+        >>> carousel_time_course(carousel, p0=p0, data=data)
 
     Args:
         carousel: Model carousel to fit
@@ -748,6 +765,8 @@ def carousel_protocol_time_course(
 ) -> CarouselFit:
     """Fit model parameters to time course of experimental data over a protocol.
 
+    Time points of protocol time course are taken from the data.
+
     Examples:
         >>> carousel_steady_state(carousel, p0=p0, data=data)
 

mxlpy/integrators/__init__.py

@@ -6,6 +6,7 @@ It includes support for both Assimulo and Scipy integrators, with Assimulo being
 
 from __future__ import annotations
 
+from .int_diffrax import Diffrax
 from .int_scipy import Scipy
 
 try:
@@ -16,5 +17,8 @@ except ImportError:
     DefaultIntegrator = Scipy
 
 __all__ = [
+    "Assimulo",
     "DefaultIntegrator",
+    "Diffrax",
+    "Scipy",
 ]

mxlpy/integrators/int_assimulo.py

@@ -17,7 +17,7 @@ with contextlib.redirect_stderr(open(os.devnull, "w")):  # noqa: PTH123
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-    from mxlpy.types import Array, ArrayLike
+    from mxlpy.types import Array, ArrayLike, Rhs
 
 
 __all__ = [
@@ -43,8 +43,8 @@ class Assimulo:
 
     """
 
-    rhs: Callable
-    y0: ArrayLike
+    rhs: Rhs
+    y0: tuple[float, ...]
     jacobian: Callable | None = None
     atol: float = 1e-8
     rtol: float = 1e-8

mxlpy/integrators/int_diffrax.py

@@ -0,0 +1,119 @@
+"""Diffrax integrator for solving ODEs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+import numpy as np
+from diffrax import (
+    AbstractSolver,
+    AbstractStepSizeController,
+    Kvaerno5,
+    ODETerm,
+    PIDController,
+    SaveAt,
+    diffeqsolve,
+)
+
+__all__ = ["Diffrax"]
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from mxlpy.types import Array, Rhs
+
+
+@dataclass
+class Diffrax:
+    """Diffrax integrator for solving ODEs."""
+
+    rhs: Rhs
+    y0: tuple[float, ...]
+    jac: Callable | None = None
+    solver: AbstractSolver = field(default=Kvaerno5())
+    stepsize_controller: AbstractStepSizeController = field(
+        default=PIDController(rtol=1e-8, atol=1e-8)
+    )
+    t0: float = 0.0
+
+    def __post_init__(self) -> None:
+        """Create copy of initial state.
+
+        This method creates a copy of the initial state `y0` and stores it in the `_y0_orig` attribute.
+        This is useful for preserving the original initial state for future reference or reset operations.
+
+        """
+        self._y0_orig = self.y0
+
+    def reset(self) -> None:
+        """Reset the integrator."""
+        self.t0 = 0
+        self.y0 = self._y0_orig
+
+    def integrate_time_course(
+        self, *, time_points: Array
+    ) -> tuple[Array | None, Array | None]:
+        """Integrate the ODE system over a time course.
+
+        Args:
+            time_points: Time points for the integration.
+
+        Returns:
+            tuple[Array, Array]: Tuple containing the time points and the integrated values.
+
+        """
+        if time_points[0] != self.t0:
+            time_points = np.insert(time_points, 0, self.t0)
+
+        res = diffeqsolve(
+            ODETerm(lambda t, y, _: self.rhs(t, y)),  # type: ignore
+            solver=self.solver,
+            t0=time_points[0],
+            t1=time_points[-1],
+            dt0=None,
+            y0=self.y0,
+            max_steps=None,
+            saveat=SaveAt(ts=time_points),  # type: ignore
+            stepsize_controller=self.stepsize_controller,
+        )
+
+        t = np.atleast_1d(np.array(res.ts, dtype=float))
+        y = np.atleast_2d(np.array(res.ys, dtype=float).T)
+
+        self.t0 = t[-1]
+        self.y0 = y[-1]
+        return t, y
+
+    def integrate(
+        self,
+        *,
+        t_end: float,
+        steps: int | None = None,
+    ) -> tuple[Array | None, Array | None]:
+        """Integrate the ODE system over a time course."""
+        steps = 100 if steps is None else steps
+
+        return self.integrate_time_course(
+            time_points=np.linspace(self.t0, t_end, steps, dtype=float)
+        )
+
+    def integrate_to_steady_state(
+        self,
+        *,
+        tolerance: float,
+        rel_norm: bool,
+        t_max: float = 1_000_000_000,
+    ) -> tuple[float | None, Array | None]:
+        """Integrate the ODE system to steady state.
+
+        Args:
+            tolerance: Tolerance for determining steady state.
+            rel_norm: Whether to use relative normalization.
+            t_max: Maximum time point for the integration (default: 1,000,000,000).
+
+        Returns:
+            tuple[float | None, Array | None]: Tuple containing the final time point and the integrated values at steady state.
+
+        """
+        raise NotImplementedError

mxlpy/integrators/int_scipy.py

@@ -14,6 +14,8 @@ from mxlpy.types import Array, ArrayLike
 if TYPE_CHECKING:
     from collections.abc import Callable
 
+    from mxlpy.types import Rhs
+
 
 __all__ = [
     "Scipy",
@@ -40,13 +42,13 @@ class Scipy:
 
     """
 
-    rhs: Callable
-    y0: ArrayLike
+    rhs: Rhs
+    y0: tuple[float, ...]
     jacobian: Callable | None = None
     atol: float = 1e-8
     rtol: float = 1e-8
     t0: float = 0.0
-    _y0_orig: ArrayLike = field(default_factory=list)
+    _y0_orig: tuple[float, ...] = field(default_factory=tuple)
 
     def __post_init__(self) -> None:
         """Create copy of initial state.
@@ -55,12 +57,12 @@ class Scipy:
         This is useful for preserving the original initial state for future reference or reset operations.
 
         """
-        self._y0_orig = self.y0.copy()
+        self._y0_orig = self.y0
 
     def reset(self) -> None:
         """Reset the integrator."""
         self.t0 = 0
-        self.y0 = self._y0_orig.copy()
+        self.y0 = self._y0_orig
 
     def integrate(
         self,
@@ -143,9 +145,13 @@ class Scipy:
 
         """
         self.reset()
-        integ = spi.ode(self.rhs, jac=self.jacobian)
+
+        # If rhs returns a tuple, we get weird errors, so we need
+        # to wrap this in a list for some reason
+        integ = spi.ode(lambda t, x: list(self.rhs(t, x)), jac=self.jacobian)
         integ.set_integrator(name="lsoda")
         integ.set_initial_value(self.y0)
+
         t = self.t0 + step_size
         y1 = copy.deepcopy(self.y0)
         for _ in range(max_steps):

mxlpy/label_map.py

@@ -27,7 +27,6 @@ from mxlpy.model import Model
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping
 
-    from mxlpy.types import Derived
 
 __all__ = [
     "LabelMapper",
@@ -556,7 +555,7 @@ class LabelMapper:
         for name, dp in self.model.get_derived_parameters().items():
             m.add_derived(name, fn=dp.fn, args=dp.args)
 
-        variables: dict[str, float | Derived] = {}
+        variables: dict[str, float] = {}
         for k, v in self.model.get_initial_conditions().items():
             if (isos := isotopomers.get(k)) is None:
                 variables[k] = v