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

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,
@@ -98,6 +100,9 @@ class Scipy:
             tuple[ArrayLike, ArrayLike]: 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 = spi.solve_ivp(
             fun=self.rhs,
             y0=self.y0,
@@ -140,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",
@@ -551,13 +550,13 @@ class LabelMapper:
 
         m = Model()
 
-        m.add_parameters(self.model.parameters)
+        m.add_parameters(self.model.get_parameter_values())
 
-        for name, dp in self.model.derived_parameters.items():
+        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] = {}
-        for k, v in self.model.variables.items():
+        variables: dict[str, float] = {}
+        for k, v in self.model.get_initial_conditions().items():
             if (isos := isotopomers.get(k)) is None:
                 variables[k] = v
             else:
@@ -585,14 +584,14 @@ class LabelMapper:
                 args=label_names,
             )
 
-        for name, dv in self.model.derived_variables.items():
+        for name, dv in self.model.get_derived_variables().items():
             m.add_derived(
                 name,
                 fn=dv.fn,
                 args=[f"{i}__total" if i in isotopomers else i for i in dv.args],
             )
 
-        for rxn_name, rxn in self.model.reactions.items():
+        for rxn_name, rxn in self.model.get_raw_reactions().items():
             if (label_map := self.label_maps.get(rxn_name)) is None:
                 m.add_reaction(
                     rxn_name,

mxlpy/linear_label_map.py

@@ -272,8 +272,10 @@ class LinearLabelMapper:
         m = Model()
         m.add_variables(variables)
         m.add_parameters(concs.to_dict() | fluxes.to_dict() | {"EXT": external_label})
+
+        rxns = self.model.get_raw_reactions()
         for rxn_name, label_map in self.label_maps.items():
-            rxn = self.model.reactions[rxn_name]
+            rxn = rxns[rxn_name]
             subs, prods = _unpack_stoichiometries(rxn.stoichiometry)
 
             subs = _stoichiometry_to_duplicate_list(subs)

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:
@@ -66,9 +66,10 @@ class ParameterScanWorker(Protocol):
         model: Model,
         *,
         parameters: pd.DataFrame,
+        y0: dict[str, float] | None,
         rel_norm: bool,
         integrator: IntegratorType,
-    ) -> SteadyStates:
+    ) -> SteadyStateScan:
         """Call the worker function."""
         ...
 
@@ -77,9 +78,10 @@ def _parameter_scan_worker(
     model: Model,
     *,
     parameters: pd.DataFrame,
+    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
@@ -109,6 +111,7 @@ def _parameter_scan_worker(
         parallel=False,
         rel_norm=rel_norm,
         integrator=integrator,
+        y0=y0,
     )
 
 
@@ -122,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:
@@ -160,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,
     )
 
 
@@ -177,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:
@@ -216,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),
     )
 
 
@@ -234,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:
@@ -274,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

@@ -71,7 +71,7 @@ def _response_coefficient_worker(
             - Series of flux response coefficients
 
     """
-    old = model.parameters[parameter]
+    old = model.get_parameter_values()[parameter]
     if y0 is not None:
         model.update_variables(y0)
 
@@ -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
 
 
@@ -205,7 +209,7 @@ def parameter_elasticities(
 
     variables = model.get_initial_conditions() if variables is None else variables
     for par in to_scan:
-        old = model.parameters[par]
+        old = model.get_parameter_values()[par]
 
         model.update_parameters({par: old * (1 + displacement)})
         upper = model.get_fluxes(variables=variables, time=time)

mxlpy/meta/__init__.py

@@ -2,12 +2,14 @@
 
 from __future__ import annotations
 
-from .codegen_latex import generate_latex_code
-from .codegen_modebase import generate_mxlpy_code
-from .codegen_py import generate_model_code_py
+from .codegen_latex import generate_latex_code, to_tex_export
+from .codegen_model import generate_model_code_py, generate_model_code_rs
+from .codegen_mxlpy import generate_mxlpy_code
 
 __all__ = [
     "generate_latex_code",
     "generate_model_code_py",
+    "generate_model_code_rs",
     "generate_mxlpy_code",
+    "to_tex_export",
 ]

mxlpy/meta/codegen_latex.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING
 
 import sympy
 
-from mxlpy.meta.source_tools import fn_to_sympy
+from mxlpy.meta.sympy_tools import fn_to_sympy, list_of_symbols
 from mxlpy.types import Derived, RateFn
 
 if TYPE_CHECKING:
@@ -21,6 +21,7 @@ __all__ = [
     "default_init",
     "generate_latex_code",
     "get_model_tex_diff",
+    "to_tex_export",
 ]
 
 cdot = r"\cdot"
@@ -31,10 +32,6 @@ newline = r"\\" + "\n"
 floatbarrier = r"\FloatBarrier"
 
 
-def _list_of_symbols(args: list[str]) -> list[sympy.Symbol | sympy.Expr]:
-    return [sympy.Symbol(arg) for arg in args]
-
-
 def default_init[T1, T2](d: dict[T1, T2] | None) -> dict[T1, T2]:
     """Return empty dict if d is None.
 
@@ -63,10 +60,6 @@ def _gls(s: str) -> str:
     return rf"\gls{{{s}}}"
 
 
-def _abbrev_and_full(s: str) -> str:
-    return rf"\acrfull{{{s}}}"
-
-
 def _gls_short(s: str) -> str:
     return rf"\acrshort{{{s}}}"
 
@@ -75,6 +68,10 @@ def _gls_full(s: str) -> str:
     return rf"\acrlong{{{s}}}"
 
 
+def _gls_short_and_full(s: str) -> str:
+    return rf"\acrfull{{{s}}}"
+
+
 def _rename_latex(s: str) -> str:
     if s[0].isdigit():
         s = s[1:]
@@ -109,6 +106,8 @@ def _sympy_to_latex(expr: sympy.Expr) -> str:
 
 def _fn_to_latex(
     fn: Callable,
+    *,
+    origin: str,
     arg_names: list[str],
     long_name_cutoff: int,
 ) -> tuple[str, dict[str, str]]:
@@ -121,10 +120,13 @@ def _fn_to_latex(
     replacements = {k: _name_to_latex(f"_x{i}") for i, k in enumerate(long_names)}
 
     expr = fn_to_sympy(
-        fn, _list_of_symbols([replacements.get(k, k) for k in tex_names])
+        fn,
+        origin=origin,
+        model_args=list_of_symbols([replacements.get(k, k) for k in tex_names]),
     )
-    fn_str = _sympy_to_latex(expr)
-    return fn_str, replacements
+    if expr is None:
+        return rf"\textcolor{{red}}{{{origin}}}", replacements
+    return _sympy_to_latex(expr), replacements
 
 
 def _table(
@@ -323,7 +325,8 @@ def _stoichs_to_latex(
             )
             sympy_fn = fn_to_sympy(
                 rxn_stoich.fn,
-                _list_of_symbols([replacements.get(k, k) for k in arg_names]),
+                origin=rxn_name,
+                model_args=list_of_symbols([replacements.get(k, k) for k in arg_names]),
             )
             expr = expr + sympy_fn * sympy.Symbol(rxn_name)  # type: ignore
         else:
@@ -480,7 +483,7 @@ class TexExport:
 
         def _add_gls_if_found(k: str) -> str:
             if (new := gls.get(k)) is not None:
-                return _abbrev_and_full(new)
+                return _gls_short_and_full(new)
             return k
 
         return TexExport(
@@ -564,7 +567,7 @@ class TexExport:
 
     def export_derived(
         self,
-        long_name_cutoff: int,
+        long_name_cutoff: int = 10,
     ) -> str:
         """Export derived quantities as LaTeX equations.
 
@@ -587,16 +590,20 @@ class TexExport:
         for k, v in sorted(self.derived.items()):
             fn_str, repls = _fn_to_latex(
                 v.fn,
+                origin=k,
                 arg_names=v.args,
                 long_name_cutoff=long_name_cutoff,
             )
-            rows.append(f"{_mathrm(_name_to_latex(k))} &= {fn_str} \\\\")
+            rows.append(f"    {_mathrm(_name_to_latex(k))} &= {fn_str} \\\\")
             if repls:
                 rows.append(_replacements_in_align(repls))
 
         return _latex_align(rows)
 
-    def export_reactions(self, long_name_cutoff: int) -> str:
+    def export_reactions(
+        self,
+        long_name_cutoff: int = 10,
+    ) -> str:
         """Export reactions as LaTeX equations.
 
         Returns
@@ -618,17 +625,18 @@ class TexExport:
         for k, v in sorted(self.reactions.items()):
             fn_str, repls = _fn_to_latex(
                 v.fn,
+                origin=k,
                 arg_names=v.args,
                 long_name_cutoff=long_name_cutoff,
             )
-            rows.append(f"{_mathrm(_name_to_latex(k))} &= {fn_str} \\\\")
+            rows.append(f"    {_mathrm(_name_to_latex(k))} &= {fn_str} \\\\")
             if repls:
                 rows.append(_replacements_in_align(repls))
         return _latex_align(rows)
 
     def export_diff_eqs(
         self,
-        long_name_cutoff: int,
+        long_name_cutoff: int = 10,
     ) -> str:
         """Export stoichiometries as LaTeX table.
 
@@ -654,12 +662,15 @@ class TexExport:
                 long_name_cutoff=long_name_cutoff,
             )
 
-            rows.append(f"{dxdt} &= {stoich_str} \\\\")
+            rows.append(f"    {dxdt} &= {stoich_str} \\\\")
             if repls:
                 rows.append(_replacements_in_align(repls))
         return _latex_align(rows)
 
-    def export_all(self, long_name_cutoff: int = 10) -> str:
+    def export_all(
+        self,
+        long_name_cutoff: int = 10,
+    ) -> str:
         """Export all model parts as a complete LaTeX document section.
 
         Returns
@@ -754,7 +765,7 @@ class TexExport:
 \usepackage[a4paper,top=2cm,bottom=2cm,left=2cm,right=2cm,marginparwidth=1.75cm]{{geometry}}
 \usepackage{{amsmath, amssymb, array, booktabs,
             breqn, caption, longtable, mathtools, placeins,
-            ragged2e, tabularx, titlesec, titling}}
+            ragged2e, tabularx, titlesec, titling, xcolor}}
 \newcommand{{\sectionbreak}}{{\clearpage}}
 \setlength{{\parindent}}{{0pt}}
 \allowdisplaybreaks
@@ -769,17 +780,20 @@ class TexExport:
 """
 
 
-def _to_tex_export(self: Model) -> TexExport:
+def to_tex_export(model: Model) -> TexExport:
+    """Create TexExport object from a model."""
     diff_eqs = {}
-    for rxn_name, rxn in self.reactions.items():
+    for rxn_name, rxn in model.get_raw_reactions().items():
         for var_name, factor in rxn.stoichiometry.items():
             diff_eqs.setdefault(var_name, {})[rxn_name] = factor
 
     return TexExport(
-        parameters=self.parameters,
-        variables=self.get_initial_conditions(),  # FIXME: think about this later
-        derived=self.derived,
-        reactions={k: TexReaction(v.fn, v.args) for k, v in self.reactions.items()},
+        parameters=model.get_parameter_values(),
+        variables=model.get_initial_conditions(),  # FIXME: think about this later
+        derived=model.get_raw_derived(),
+        reactions={
+            k: TexReaction(v.fn, v.args) for k, v in model.get_raw_reactions().items()
+        },
         diff_eqs=diff_eqs,
     )
 
@@ -820,7 +834,7 @@ def generate_latex_code(
     """
     gls = default_init(gls)
     return (
-        _to_tex_export(model)
+        to_tex_export(model)
         .rename_with_glossary(gls)
         .export_document(long_name_cutoff=long_name_cutoff)
     )
@@ -866,7 +880,7 @@ def get_model_tex_diff(
     return f"""{" start autogenerated ":%^60}
 {_clearpage()}
 {_subsubsection("Model changes")}{_label(section_label)}
-{((_to_tex_export(m1) - _to_tex_export(m2)).rename_with_glossary(gls).export_all(long_name_cutoff=long_name_cutoff))}
+{((to_tex_export(m1) - to_tex_export(m2)).rename_with_glossary(gls).export_all(long_name_cutoff=long_name_cutoff))}
 {_clearpage()}
 {" end autogenerated ":%^60}
 """