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

mxlpy/__init__.py

@@ -54,7 +54,7 @@ from . import (
     sbml,
     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
@@ -62,7 +62,14 @@ 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",

mxlpy/carousel.py

@@ -18,7 +18,7 @@ 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
@@ -35,11 +35,13 @@ class CarouselSteadyState:
     """Time course of a carousel simulation."""
 
     carousel: list[Model]
-    results: list[scan.TimePoint]
+    results: list[Result]
 
     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 +49,7 @@ class CarouselTimeCourse:
     """Time course of a carousel simulation."""
 
     carousel: list[Model]
-    results: list[scan.TimeCourse]
+    results: list[Result]
 
     def get_variables_by_model(self) -> pd.DataFrame:
         """Get the variables of the time course results, indexed by model."""

mxlpy/compare.py

@@ -8,8 +8,8 @@ from typing import TYPE_CHECKING, cast
 import pandas as pd
 
 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

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

mxlpy/mc.py

@@ -35,10 +35,10 @@ from mxlpy.scan import (
 from mxlpy.types import (
     IntegratorType,
     McSteadyStates,
-    ProtocolByPars,
+    ProtocolScan,
     ResponseCoefficientsByPars,
-    SteadyStates,
-    TimeCourseByPars,
+    SteadyStateScan,
+    TimeCourseScan,
 )
 
 if TYPE_CHECKING:
@@ -69,7 +69,7 @@ class ParameterScanWorker(Protocol):
         y0: dict[str, float] | None,
         rel_norm: bool,
         integrator: IntegratorType,
-    ) -> SteadyStates:
+    ) -> SteadyStateScan:
         """Call the worker function."""
         ...
 
@@ -81,7 +81,7 @@ def _parameter_scan_worker(
     y0: dict[str, float] | None,
     rel_norm: bool,
     integrator: IntegratorType,
-) -> SteadyStates:
+) -> SteadyStateScan:
     """Worker function for parallel steady state scanning across parameter sets.
 
     This function executes a parameter scan for steady state solutions for a
@@ -125,7 +125,7 @@ def steady_state(
     rel_norm: bool = False,
     worker: SteadyStateWorker = _steady_state_worker,
     integrator: IntegratorType | None = None,
-) -> SteadyStates:
+) -> SteadyStateScan:
     """Monte-carlo scan of steady states.
 
     Examples:
@@ -163,10 +163,14 @@ def steady_state(
         max_workers=max_workers,
         cache=cache,
     )
-    return SteadyStates(
-        variables=pd.concat({k: v.variables for k, v in res}, axis=1).T,
-        fluxes=pd.concat({k: v.fluxes for k, v in res}, axis=1).T,
-        parameters=mc_to_scan,
+    return SteadyStateScan(
+        raw_index=(
+            pd.Index(mc_to_scan.iloc[:, 0])
+            if mc_to_scan.shape[1] == 1
+            else pd.MultiIndex.from_frame(mc_to_scan)
+        ),
+        raw_results=[i[1] for i in res],
+        to_scan=mc_to_scan,
     )
 
 
@@ -180,7 +184,7 @@ def time_course(
     cache: Cache | None = None,
     worker: TimeCourseWorker = _time_course_worker,
     integrator: IntegratorType | None = None,
-) -> TimeCourseByPars:
+) -> TimeCourseScan:
     """MC time course.
 
     Examples:
@@ -219,10 +223,9 @@ def time_course(
         cache=cache,
     )
 
-    return TimeCourseByPars(
-        parameters=mc_to_scan,
-        variables=pd.concat({k: v.variables.T for k, v in res}, axis=1).T,
-        fluxes=pd.concat({k: v.fluxes.T for k, v in res}, axis=1).T,
+    return TimeCourseScan(
+        to_scan=mc_to_scan,
+        raw_results=dict(res),
     )
 
 
@@ -237,7 +240,7 @@ def time_course_over_protocol(
     cache: Cache | None = None,
     worker: ProtocolWorker = _protocol_worker,
     integrator: IntegratorType | None = None,
-) -> ProtocolByPars:
+) -> ProtocolScan:
     """MC time course.
 
     Examples:
@@ -277,13 +280,10 @@ def time_course_over_protocol(
         max_workers=max_workers,
         cache=cache,
     )
-    concs = {k: v.variables.T for k, v in res}
-    fluxes = {k: v.fluxes.T for k, v in res}
-    return ProtocolByPars(
-        variables=pd.concat(concs, axis=1).T,
-        fluxes=pd.concat(fluxes, axis=1).T,
-        parameters=mc_to_scan,
+    return ProtocolScan(
+        to_scan=mc_to_scan,
         protocol=protocol,
+        raw_results=dict(res),
     )
 
 

mxlpy/mca.py

@@ -91,8 +91,12 @@ def _response_coefficient_worker(
         y0=None,
     )
 
-    conc_resp = (upper.variables - lower.variables) / (2 * displacement * old)
-    flux_resp = (upper.fluxes - lower.fluxes) / (2 * displacement * old)
+    conc_resp = (upper.variables.iloc[-1] - lower.variables.iloc[-1]) / (
+        2 * displacement * old
+    )
+    flux_resp = (upper.fluxes.iloc[-1] - lower.fluxes.iloc[-1]) / (
+        2 * displacement * old
+    )
     # Reset
     model.update_parameters({parameter: old})
     if normalized:
@@ -102,8 +106,8 @@ def _response_coefficient_worker(
             integrator=integrator,
             y0=None,
         )
-        conc_resp *= old / norm.variables
-        flux_resp *= old / norm.fluxes
+        conc_resp *= old / norm.variables.iloc[-1]
+        flux_resp *= old / norm.fluxes.iloc[-1]
     return conc_resp, flux_resp
 
 

mxlpy/meta/codegen_model.py

@@ -110,7 +110,8 @@ def _generate_model_code(
         _LOGGER.warning(msg)
 
     # Return
-    ret = ", ".join(f"d{i}dt" for i in diff_eqs) if len(diff_eqs) > 0 else "()"
+    ret_order = [i for i in variables if i in diff_eqs]
+    ret = ", ".join(f"d{i}dt" for i in ret_order) if len(diff_eqs) > 0 else "()"
     source.append(return_template.format(ret))
 
     if end is not None: