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

mxlpy/__init__.py

@@ -52,8 +52,9 @@ from . import (
     plot,
     report,
     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
@@ -61,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, unwrap
+from .types import (
+    Derived,
+    InitialAssignment,
+    IntegratorProtocol,
+    Parameter,
+    Variable,
+    unwrap,
+)
 
 with contextlib.suppress(ImportError):
     from .integrators import Assimulo
@@ -83,13 +91,17 @@ __all__ = [
     "Cache",
     "DefaultIntegrator",
     "Derived",
+    "Diffrax",
+    "InitialAssignment",
     "IntegratorProtocol",
     "LabelMapper",
     "LinearLabelMapper",
     "Model",
+    "Parameter",
     "Scipy",
     "Simulator",
     "SymbolicModel",
+    "Variable",
     "cartesian_product",
     "compare",
     "distributions",
@@ -110,6 +122,7 @@ __all__ = [
     "time_course",
     "time_course_over_protocol",
     "to_symbolic_model",
+    "units",
     "unwrap",
 ]
 

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
@@ -230,11 +230,7 @@ def protocol_time_courses(
 ) -> ProtocolComparison:
     """Compare the time courses of two models."""
     return ProtocolComparison(
-        res1=unwrap(
-            Simulator(m1).simulate_over_protocol(protocol=protocol).get_result()
-        ),
-        res2=unwrap(
-            Simulator(m2).simulate_over_protocol(protocol=protocol).get_result()
-        ),
+        res1=unwrap(Simulator(m1).simulate_protocol(protocol=protocol).get_result()),
+        res2=unwrap(Simulator(m2).simulate_protocol(protocol=protocol).get_result()),
         protocol=protocol,
     )

mxlpy/experimental/diff.py

@@ -187,7 +187,7 @@ def model_diff(m1: Model, m2: Model) -> ModelDiff:
         if (v2 := m2._parameters.get(k)) is None:  # noqa: SLF001
             diff.missing_parameters.add(k)
         elif v1 != v2:
-            diff.different_parameters[k] = (v1, v2)
+            diff.different_parameters[k] = (v1, v2)  # type: ignore
 
     for k, v1 in m1._variables.items():  # noqa: SLF001
         if (v2 := m2._variables.get(k)) is None:  # noqa: SLF001

mxlpy/fit.py

@@ -32,6 +32,7 @@ if TYPE_CHECKING:
 LOGGER = logging.getLogger(__name__)
 
 __all__ = [
+    "Bounds",
     "CarouselFit",
     "FitResult",
     "InitialGuess",
@@ -43,13 +44,13 @@ __all__ = [
     "ResidualFn",
     "SteadyStateResidualFn",
     "TimeSeriesResidualFn",
+    "carousel_protocol_time_course",
     "carousel_steady_state",
     "carousel_time_course",
-    "carousel_time_course_over_protocol",
+    "protocol_time_course",
     "rmse",
     "steady_state",
     "time_course",
-    "time_course_over_protocol",
 ]
 
 
@@ -83,7 +84,15 @@ class CarouselFit:
 
 type InitialGuess = dict[str, float]
 type ResidualFn = Callable[[Array], float]
-type MinimizeFn = Callable[[ResidualFn, InitialGuess], MinResult | None]
+type Bounds = dict[str, tuple[float | None, float | None]]
+type MinimizeFn = Callable[
+    [
+        ResidualFn,
+        InitialGuess,
+        Bounds,
+    ],
+    MinResult | None,
+]
 type LossFn = Callable[
     [
         pd.DataFrame | pd.Series,
@@ -151,7 +160,6 @@ class ProtocolResidualFn(Protocol):
         integrator: IntegratorType,
         loss_fn: LossFn,
         protocol: pd.DataFrame,
-        time_points_per_step: int = 10,
     ) -> float:
         """Calculate residual error between model time course and experimental data."""
         ...
@@ -160,11 +168,12 @@ class ProtocolResidualFn(Protocol):
 def _default_minimize_fn(
     residual_fn: ResidualFn,
     p0: dict[str, float],
+    bounds: Bounds,
 ) -> MinResult | None:
     res = minimize(
         residual_fn,
         x0=list(p0.values()),
-        bounds=[(0, None) for _ in range(len(p0))],
+        bounds=[bounds.get(name, (1e-6, 1e6)) for name in p0],
         method="L-BFGS-B",
     )
     if res.success:
@@ -278,7 +287,7 @@ def _time_course_residual(
     )
 
 
-def _protocol_residual(
+def _protocol_time_course_residual(
     par_values: ArrayLike,
     # This will be filled out by partial
     par_names: list[str],
@@ -288,7 +297,6 @@ def _protocol_residual(
     integrator: IntegratorType,
     loss_fn: LossFn,
     protocol: pd.DataFrame,
-    time_points_per_step: int = 10,
 ) -> float:
     """Calculate residual error between model time course and experimental data.
 
@@ -313,9 +321,9 @@ def _protocol_residual(
             y0=y0,
             integrator=integrator,
         )
-        .simulate_over_protocol(
+        .simulate_protocol_time_course(
             protocol=protocol,
-            time_points_per_step=time_points_per_step,
+            time_points=data.index,
         )
         .get_result()
     )
@@ -329,6 +337,84 @@ def _protocol_residual(
     )
 
 
+def _carousel_steady_state_worker(
+    model: Model,
+    p0: dict[str, float],
+    data: pd.Series,
+    y0: dict[str, float] | None,
+    integrator: IntegratorType | None,
+    loss_fn: LossFn,
+    minimize_fn: MinimizeFn,
+    residual_fn: SteadyStateResidualFn,
+    bounds: Bounds | None,
+) -> FitResult | None:
+    model_pars = model.get_parameter_values()
+
+    return steady_state(
+        model,
+        p0={k: v for k, v in p0.items() if k in model_pars},
+        y0=y0,
+        data=data,
+        minimize_fn=minimize_fn,
+        residual_fn=residual_fn,
+        integrator=integrator,
+        loss_fn=loss_fn,
+        bounds=bounds,
+    )
+
+
+def _carousel_time_course_worker(
+    model: Model,
+    p0: dict[str, float],
+    data: pd.DataFrame,
+    y0: dict[str, float] | None,
+    integrator: IntegratorType | None,
+    loss_fn: LossFn,
+    minimize_fn: MinimizeFn,
+    residual_fn: TimeSeriesResidualFn,
+    bounds: Bounds | None,
+) -> FitResult | None:
+    model_pars = model.get_parameter_values()
+    return time_course(
+        model,
+        p0={k: v for k, v in p0.items() if k in model_pars},
+        y0=y0,
+        data=data,
+        minimize_fn=minimize_fn,
+        residual_fn=residual_fn,
+        integrator=integrator,
+        loss_fn=loss_fn,
+        bounds=bounds,
+    )
+
+
+def _carousel_protocol_worker(
+    model: Model,
+    p0: dict[str, float],
+    data: pd.DataFrame,
+    protocol: pd.DataFrame,
+    y0: dict[str, float] | None,
+    integrator: IntegratorType | None,
+    loss_fn: LossFn,
+    minimize_fn: MinimizeFn,
+    residual_fn: ProtocolResidualFn,
+    bounds: Bounds | None,
+) -> FitResult | None:
+    model_pars = model.get_parameter_values()
+    return protocol_time_course(
+        model,
+        p0={k: v for k, v in p0.items() if k in model_pars},
+        y0=y0,
+        protocol=protocol,
+        data=data,
+        minimize_fn=minimize_fn,
+        residual_fn=residual_fn,
+        integrator=integrator,
+        loss_fn=loss_fn,
+        bounds=bounds,
+    )
+
+
 def steady_state(
     model: Model,
     *,
@@ -339,6 +425,7 @@ def steady_state(
     residual_fn: SteadyStateResidualFn = _steady_state_residual,
     integrator: IntegratorType | None = None,
     loss_fn: LossFn = rmse,
+    bounds: Bounds | None = None,
 ) -> FitResult | None:
     """Fit model parameters to steady-state experimental data.
 
@@ -355,18 +442,19 @@ def steady_state(
         residual_fn: Function to calculate fitting error
         integrator: ODE integrator class
         loss_fn: Loss function to use for residual calculation
+        bounds: Mapping of bounds per parameter
 
     Returns:
         dict[str, float]: Fitted parameters as {parameter_name: fitted_value}
 
     Note:
-        Uses L-BFGS-B optimization with bounds [1e-12, 1e6] for all parameters
+        Uses L-BFGS-B optimization with bounds [1e-6, 1e6] for all parameters
 
     """
     par_names = list(p0.keys())
 
     # Copy to restore
-    p_orig = model.parameters
+    p_orig = model.get_parameter_values()
 
     fn = cast(
         ResidualFn,
@@ -380,7 +468,7 @@ def steady_state(
             loss_fn=loss_fn,
         ),
     )
-    min_result = minimize_fn(fn, p0)
+    min_result = minimize_fn(fn, p0, {} if bounds is None else bounds)
     # Restore original model
     model.update_parameters(p_orig)
     if min_result is None:
@@ -403,6 +491,7 @@ def time_course(
     residual_fn: TimeSeriesResidualFn = _time_course_residual,
     integrator: IntegratorType | None = None,
     loss_fn: LossFn = rmse,
+    bounds: Bounds | None = None,
 ) -> FitResult | None:
     """Fit model parameters to time course of experimental data.
 
@@ -419,16 +508,17 @@ def time_course(
         residual_fn: Function to calculate fitting error
         integrator: ODE integrator class
         loss_fn: Loss function to use for residual calculation
+        bounds: Mapping of bounds per parameter
 
     Returns:
         dict[str, float]: Fitted parameters as {parameter_name: fitted_value}
 
     Note:
-        Uses L-BFGS-B optimization with bounds [1e-12, 1e6] for all parameters
+        Uses L-BFGS-B optimization with bounds [1e-6, 1e6] for all parameters
 
     """
     par_names = list(p0.keys())
-    p_orig = model.parameters
+    p_orig = model.get_parameter_values()
 
     fn = cast(
         ResidualFn,
@@ -443,7 +533,7 @@ def time_course(
         ),
     )
 
-    min_result = minimize_fn(fn, p0)
+    min_result = minimize_fn(fn, p0, {} if bounds is None else bounds)
     # Restore original model
     model.update_parameters(p_orig)
     if min_result is None:
@@ -456,7 +546,7 @@ def time_course(
     )
 
 
-def time_course_over_protocol(
+def protocol_time_course(
     model: Model,
     *,
     p0: dict[str, float],
@@ -464,10 +554,10 @@ def time_course_over_protocol(
     protocol: pd.DataFrame,
     y0: dict[str, float] | None = None,
     minimize_fn: MinimizeFn = _default_minimize_fn,
-    residual_fn: ProtocolResidualFn = _protocol_residual,
+    residual_fn: ProtocolResidualFn = _protocol_time_course_residual,
     integrator: IntegratorType | None = None,
     loss_fn: LossFn = rmse,
-    time_points_per_step: int = 10,
+    bounds: Bounds | None = None,
 ) -> FitResult | None:
     """Fit model parameters to time course of experimental data.
 
@@ -486,16 +576,17 @@ def time_course_over_protocol(
         integrator: ODE integrator class
         loss_fn: Loss function to use for residual calculation
         time_points_per_step: Number of time points per step in the protocol
+        bounds: Mapping of bounds per parameter
 
     Returns:
         dict[str, float]: Fitted parameters as {parameter_name: fitted_value}
 
     Note:
-        Uses L-BFGS-B optimization with bounds [1e-12, 1e6] for all parameters
+        Uses L-BFGS-B optimization with bounds [1e-6, 1e6] for all parameters
 
     """
     par_names = list(p0.keys())
-    p_orig = model.parameters
+    p_orig = model.get_parameter_values()
 
     fn = cast(
         ResidualFn,
@@ -508,11 +599,10 @@ def time_course_over_protocol(
             integrator=integrator,
             loss_fn=loss_fn,
             protocol=protocol,
-            time_points_per_step=time_points_per_step,
         ),
     )
 
-    min_result = minimize_fn(fn, p0)
+    min_result = minimize_fn(fn, p0, {} if bounds is None else bounds)
     # Restore original model
     model.update_parameters(p_orig)
     if min_result is None:
@@ -525,78 +615,6 @@ def time_course_over_protocol(
     )
 
 
-def _carousel_steady_state_worker(
-    model: Model,
-    p0: dict[str, float],
-    data: pd.Series,
-    y0: dict[str, float] | None,
-    integrator: IntegratorType | None,
-    loss_fn: LossFn,
-    minimize_fn: MinimizeFn,
-    residual_fn: SteadyStateResidualFn,
-) -> FitResult | None:
-    model_pars = model.parameters
-
-    return steady_state(
-        model,
-        p0={k: v for k, v in p0.items() if k in model_pars},
-        y0=y0,
-        data=data,
-        minimize_fn=minimize_fn,
-        residual_fn=residual_fn,
-        integrator=integrator,
-        loss_fn=loss_fn,
-    )
-
-
-def _carousel_time_course_worker(
-    model: Model,
-    p0: dict[str, float],
-    data: pd.DataFrame,
-    y0: dict[str, float] | None,
-    integrator: IntegratorType | None,
-    loss_fn: LossFn,
-    minimize_fn: MinimizeFn,
-    residual_fn: TimeSeriesResidualFn,
-) -> FitResult | None:
-    model_pars = model.parameters
-    return time_course(
-        model,
-        p0={k: v for k, v in p0.items() if k in model_pars},
-        y0=y0,
-        data=data,
-        minimize_fn=minimize_fn,
-        residual_fn=residual_fn,
-        integrator=integrator,
-        loss_fn=loss_fn,
-    )
-
-
-def _carousel_protocol_worker(
-    model: Model,
-    p0: dict[str, float],
-    data: pd.DataFrame,
-    protocol: pd.DataFrame,
-    y0: dict[str, float] | None,
-    integrator: IntegratorType | None,
-    loss_fn: LossFn,
-    minimize_fn: MinimizeFn,
-    residual_fn: ProtocolResidualFn,
-) -> FitResult | None:
-    model_pars = model.parameters
-    return time_course_over_protocol(
-        model,
-        p0={k: v for k, v in p0.items() if k in model_pars},
-        y0=y0,
-        protocol=protocol,
-        data=data,
-        minimize_fn=minimize_fn,
-        residual_fn=residual_fn,
-        integrator=integrator,
-        loss_fn=loss_fn,
-    )
-
-
 def carousel_steady_state(
     carousel: Carousel,
     *,
@@ -607,8 +625,33 @@ def carousel_steady_state(
     residual_fn: SteadyStateResidualFn = _steady_state_residual,
     integrator: IntegratorType | None = None,
     loss_fn: LossFn = rmse,
+    bounds: Bounds | None = None,
 ) -> CarouselFit:
-    """Fit model parameters to steady-state experimental data over a carousel."""
+    """Fit model parameters to steady-state experimental data over a carousel.
+
+    Examples:
+        >>> carousel_steady_state(carousel, p0=p0, data=data)
+
+    Args:
+        carousel: Model carousel to fit
+        p0: Initial parameter guesses as {parameter_name: value}
+        data: Experimental time course data
+        protocol: Experimental protocol
+        y0: Initial conditions as {species_name: value}
+        minimize_fn: Function to minimize fitting error
+        residual_fn: Function to calculate fitting error
+        integrator: ODE integrator class
+        loss_fn: Loss function to use for residual calculation
+        time_points_per_step: Number of time points per step in the protocol
+        bounds: Mapping of bounds per parameter
+
+    Returns:
+        dict[str, float]: Fitted parameters as {parameter_name: fitted_value}
+
+    Note:
+        Uses L-BFGS-B optimization with bounds [1e-6, 1e6] for all parameters
+
+    """
     return CarouselFit(
         [
             fit
@@ -622,6 +665,7 @@ def carousel_steady_state(
                     loss_fn=loss_fn,
                     minimize_fn=minimize_fn,
                     residual_fn=residual_fn,
+                    bounds=bounds,
                 ),
                 inputs=list(enumerate(carousel.variants)),
             )
@@ -640,8 +684,33 @@ def carousel_time_course(
     residual_fn: TimeSeriesResidualFn = _time_course_residual,
     integrator: IntegratorType | None = None,
     loss_fn: LossFn = rmse,
+    bounds: Bounds | None = None,
 ) -> CarouselFit:
-    """Fit model parameters to time course of experimental data over a carousel."""
+    """Fit model parameters to time course of experimental data over a carousel.
+
+    Examples:
+        >>> carousel_steady_state(carousel, p0=p0, data=data)
+
+    Args:
+        carousel: Model carousel to fit
+        p0: Initial parameter guesses as {parameter_name: value}
+        data: Experimental time course data
+        protocol: Experimental protocol
+        y0: Initial conditions as {species_name: value}
+        minimize_fn: Function to minimize fitting error
+        residual_fn: Function to calculate fitting error
+        integrator: ODE integrator class
+        loss_fn: Loss function to use for residual calculation
+        time_points_per_step: Number of time points per step in the protocol
+        bounds: Mapping of bounds per parameter
+
+    Returns:
+        dict[str, float]: Fitted parameters as {parameter_name: fitted_value}
+
+    Note:
+        Uses L-BFGS-B optimization with bounds [1e-6, 1e6] for all parameters
+
+    """
     return CarouselFit(
         [
             fit
@@ -655,6 +724,7 @@ def carousel_time_course(
                     loss_fn=loss_fn,
                     minimize_fn=minimize_fn,
                     residual_fn=residual_fn,
+                    bounds=bounds,
                 ),
                 inputs=list(enumerate(carousel.variants)),
             )
@@ -663,7 +733,7 @@ def carousel_time_course(
     )
 
 
-def carousel_time_course_over_protocol(
+def carousel_protocol_time_course(
     carousel: Carousel,
     *,
     p0: dict[str, float],
@@ -671,11 +741,36 @@ def carousel_time_course_over_protocol(
     protocol: pd.DataFrame,
     y0: dict[str, float] | None = None,
     minimize_fn: MinimizeFn = _default_minimize_fn,
-    residual_fn: ProtocolResidualFn = _protocol_residual,
+    residual_fn: ProtocolResidualFn = _protocol_time_course_residual,
     integrator: IntegratorType | None = None,
     loss_fn: LossFn = rmse,
+    bounds: Bounds | None = None,
 ) -> CarouselFit:
-    """Fit model parameters to time course of experimental data over a protocol."""
+    """Fit model parameters to time course of experimental data over a protocol.
+
+    Examples:
+        >>> carousel_steady_state(carousel, p0=p0, data=data)
+
+    Args:
+        carousel: Model carousel to fit
+        p0: Initial parameter guesses as {parameter_name: value}
+        data: Experimental time course data
+        protocol: Experimental protocol
+        y0: Initial conditions as {species_name: value}
+        minimize_fn: Function to minimize fitting error
+        residual_fn: Function to calculate fitting error
+        integrator: ODE integrator class
+        loss_fn: Loss function to use for residual calculation
+        time_points_per_step: Number of time points per step in the protocol
+        bounds: Mapping of bounds per parameter
+
+    Returns:
+        dict[str, float]: Fitted parameters as {parameter_name: fitted_value}
+
+    Note:
+        Uses L-BFGS-B optimization with bounds [1e-6, 1e6] for all parameters
+
+    """
     return CarouselFit(
         [
             fit
@@ -690,6 +785,7 @@ def carousel_time_course_over_protocol(
                     loss_fn=loss_fn,
                     minimize_fn=minimize_fn,
                     residual_fn=residual_fn,
+                    bounds=bounds,
                 ),
                 inputs=list(enumerate(carousel.variants)),
             )

mxlpy/identify.py

@@ -59,21 +59,26 @@ def profile_likelihood(
 
     """
     parameter_distributions = sample(
-        {k: LogNormal(np.log(v), sigma=1) for k, v in model.parameters.items()},
+        {
+            k: LogNormal(np.log(v), sigma=1)
+            for k, v in model.get_parameter_values().items()
+        },
         n=n_random,
     )
 
     res = {}
     for value in tqdm(parameter_values, desc=parameter_name):
         model.update_parameter(parameter_name, value)
-        res[value] = parallelise(
-            partial(
-                _mc_fit_time_course_worker, model=model, data=data, loss_fn=loss_fn
-            ),
-            inputs=list(
-                parameter_distributions.drop(columns=parameter_name).iterrows()
-            ),
-            disable_tqdm=True,
+        res[value] = dict(
+            parallelise(
+                partial(
+                    _mc_fit_time_course_worker, model=model, data=data, loss_fn=loss_fn
+                ),
+                inputs=list(
+                    parameter_distributions.drop(columns=parameter_name).iterrows()
+                ),
+                disable_tqdm=True,
+            )
         )
     errors = pd.DataFrame(res, dtype=float).T.abs().mean(axis=1)
     errors.index.name = "fitting error"

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",
 ]