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

mxlpy/fit.py

@@ -10,39 +10,89 @@ Functions:
 
 from __future__ import annotations
 
+import logging
+from copy import deepcopy
+from dataclasses import dataclass
 from functools import partial
 from typing import TYPE_CHECKING, Protocol
 
 import numpy as np
 from scipy.optimize import minimize
 
-from mxlpy.integrators import DefaultIntegrator
+from mxlpy import parallel
 from mxlpy.simulator import Simulator
 from mxlpy.types import Array, ArrayLike, Callable, IntegratorType, cast
 
 if TYPE_CHECKING:
     import pandas as pd
 
+    from mxlpy.carousel import Carousel
     from mxlpy.model import Model
 
+LOGGER = logging.getLogger(__name__)
+
 __all__ = [
+    "Bounds",
+    "CarouselFit",
+    "FitResult",
     "InitialGuess",
+    "LOGGER",
     "LossFn",
+    "MinResult",
     "MinimizeFn",
     "ProtocolResidualFn",
     "ResidualFn",
     "SteadyStateResidualFn",
     "TimeSeriesResidualFn",
+    "carousel_protocol_time_course",
+    "carousel_steady_state",
+    "carousel_time_course",
+    "protocol_time_course",
     "rmse",
     "steady_state",
     "time_course",
-    "time_course_over_protocol",
 ]
 
 
+@dataclass
+class MinResult:
+    """Result of a minimization operation."""
+
+    parameters: dict[str, float]
+    residual: float
+
+
+@dataclass
+class FitResult:
+    """Result of a fit operation."""
+
+    model: Model
+    best_pars: dict[str, float]
+    loss: float
+
+
+@dataclass
+class CarouselFit:
+    """Result of a carousel fit operation."""
+
+    fits: list[FitResult]
+
+    def get_best_fit(self) -> FitResult:
+        """Get the best fit from the carousel."""
+        return min(self.fits, key=lambda x: x.loss)
+
+
 type InitialGuess = dict[str, float]
 type ResidualFn = Callable[[Array], float]
-type MinimizeFn = Callable[[ResidualFn, InitialGuess], dict[str, float]]
+type Bounds = dict[str, tuple[float | None, float | None]]
+type MinimizeFn = Callable[
+    [
+        ResidualFn,
+        InitialGuess,
+        Bounds,
+    ],
+    MinResult | None,
+]
 type LossFn = Callable[
     [
         pd.DataFrame | pd.Series,
@@ -70,7 +120,7 @@ class SteadyStateResidualFn(Protocol):
         par_names: list[str],
         data: pd.Series,
         model: Model,
-        y0: dict[str, float],
+        y0: dict[str, float] | None,
         integrator: IntegratorType,
         loss_fn: LossFn,
     ) -> float:
@@ -88,7 +138,7 @@ class TimeSeriesResidualFn(Protocol):
         par_names: list[str],
         data: pd.DataFrame,
         model: Model,
-        y0: dict[str, float],
+        y0: dict[str, float] | None,
         integrator: IntegratorType,
         loss_fn: LossFn,
     ) -> float:
@@ -106,11 +156,10 @@ class ProtocolResidualFn(Protocol):
         par_names: list[str],
         data: pd.DataFrame,
         model: Model,
-        y0: dict[str, float],
+        y0: dict[str, float] | None,
         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."""
         ...
@@ -119,22 +168,28 @@ class ProtocolResidualFn(Protocol):
 def _default_minimize_fn(
     residual_fn: ResidualFn,
     p0: dict[str, float],
-) -> dict[str, float]:
+    bounds: Bounds,
+) -> MinResult | None:
     res = minimize(
         residual_fn,
         x0=list(p0.values()),
-        bounds=[(1e-12, 1e6) for _ in range(len(p0))],
+        bounds=[bounds.get(name, (1e-6, 1e6)) for name in p0],
         method="L-BFGS-B",
     )
     if res.success:
-        return dict(
-            zip(
-                p0,
-                res.x,
-                strict=True,
-            )
+        return MinResult(
+            parameters=dict(
+                zip(
+                    p0,
+                    res.x,
+                    strict=True,
+                ),
+            ),
+            residual=res.fun,
         )
-    return dict(zip(p0, np.full(len(p0), np.nan, dtype=float), strict=True))
+
+    LOGGER.warning("Minimisation failed.")
+    return None
 
 
 def _steady_state_residual(
@@ -232,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],
@@ -242,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.
 
@@ -267,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()
     )
@@ -283,16 +337,96 @@ 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,
+    *,
     p0: dict[str, float],
     data: pd.Series,
     y0: dict[str, float] | None = None,
     minimize_fn: MinimizeFn = _default_minimize_fn,
     residual_fn: SteadyStateResidualFn = _steady_state_residual,
-    integrator: IntegratorType = DefaultIntegrator,
+    integrator: IntegratorType | None = None,
     loss_fn: LossFn = rmse,
-) -> dict[str, float]:
+    bounds: Bounds | None = None,
+) -> FitResult | None:
     """Fit model parameters to steady-state experimental data.
 
     Examples:
@@ -308,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,
@@ -333,23 +468,31 @@ def steady_state(
             loss_fn=loss_fn,
         ),
     )
-    res = minimize_fn(fn, p0)
-
-    # Restore
+    min_result = minimize_fn(fn, p0, {} if bounds is None else bounds)
+    # Restore original model
     model.update_parameters(p_orig)
-    return res
+    if min_result is None:
+        return min_result
+
+    return FitResult(
+        model=deepcopy(model).update_parameters(min_result.parameters),
+        best_pars=min_result.parameters,
+        loss=min_result.residual,
+    )
 
 
 def time_course(
     model: Model,
+    *,
     p0: dict[str, float],
     data: pd.DataFrame,
     y0: dict[str, float] | None = None,
     minimize_fn: MinimizeFn = _default_minimize_fn,
     residual_fn: TimeSeriesResidualFn = _time_course_residual,
-    integrator: IntegratorType = DefaultIntegrator,
+    integrator: IntegratorType | None = None,
     loss_fn: LossFn = rmse,
-) -> dict[str, float]:
+    bounds: Bounds | None = None,
+) -> FitResult | None:
     """Fit model parameters to time course of experimental data.
 
     Examples:
@@ -365,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,
@@ -388,23 +532,33 @@ def time_course(
             loss_fn=loss_fn,
         ),
     )
-    res = minimize_fn(fn, p0)
+
+    min_result = minimize_fn(fn, p0, {} if bounds is None else bounds)
+    # Restore original model
     model.update_parameters(p_orig)
-    return res
+    if min_result is None:
+        return min_result
 
+    return FitResult(
+        model=deepcopy(model).update_parameters(min_result.parameters),
+        best_pars=min_result.parameters,
+        loss=min_result.residual,
+    )
 
-def time_course_over_protocol(
+
+def protocol_time_course(
     model: Model,
+    *,
     p0: dict[str, float],
     data: pd.DataFrame,
     protocol: pd.DataFrame,
     y0: dict[str, float] | None = None,
     minimize_fn: MinimizeFn = _default_minimize_fn,
-    residual_fn: ProtocolResidualFn = _protocol_residual,
-    integrator: IntegratorType = DefaultIntegrator,
+    residual_fn: ProtocolResidualFn = _protocol_time_course_residual,
+    integrator: IntegratorType | None = None,
     loss_fn: LossFn = rmse,
-    time_points_per_step: int = 10,
-) -> dict[str, float]:
+    bounds: Bounds | None = None,
+) -> FitResult | None:
     """Fit model parameters to time course of experimental data.
 
     Examples:
@@ -422,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,
@@ -444,9 +599,196 @@ def time_course_over_protocol(
             integrator=integrator,
             loss_fn=loss_fn,
             protocol=protocol,
-            time_points_per_step=time_points_per_step,
         ),
     )
-    res = minimize_fn(fn, p0)
+
+    min_result = minimize_fn(fn, p0, {} if bounds is None else bounds)
+    # Restore original model
     model.update_parameters(p_orig)
-    return res
+    if min_result is None:
+        return min_result
+
+    return FitResult(
+        model=deepcopy(model).update_parameters(min_result.parameters),
+        best_pars=min_result.parameters,
+        loss=min_result.residual,
+    )
+
+
+def carousel_steady_state(
+    carousel: Carousel,
+    *,
+    p0: dict[str, float],
+    data: pd.Series,
+    y0: dict[str, float] | None = None,
+    minimize_fn: MinimizeFn = _default_minimize_fn,
+    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.
+
+    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
+            for i in parallel.parallelise(
+                partial(
+                    _carousel_steady_state_worker,
+                    p0=p0,
+                    data=data,
+                    y0=y0,
+                    integrator=integrator,
+                    loss_fn=loss_fn,
+                    minimize_fn=minimize_fn,
+                    residual_fn=residual_fn,
+                    bounds=bounds,
+                ),
+                inputs=list(enumerate(carousel.variants)),
+            )
+            if (fit := i[1]) is not None
+        ]
+    )
+
+
+def carousel_time_course(
+    carousel: Carousel,
+    *,
+    p0: dict[str, float],
+    data: pd.DataFrame,
+    y0: dict[str, float] | None = None,
+    minimize_fn: MinimizeFn = _default_minimize_fn,
+    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.
+
+    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
+            for i in parallel.parallelise(
+                partial(
+                    _carousel_time_course_worker,
+                    p0=p0,
+                    data=data,
+                    y0=y0,
+                    integrator=integrator,
+                    loss_fn=loss_fn,
+                    minimize_fn=minimize_fn,
+                    residual_fn=residual_fn,
+                    bounds=bounds,
+                ),
+                inputs=list(enumerate(carousel.variants)),
+            )
+            if (fit := i[1]) is not None
+        ]
+    )
+
+
+def carousel_protocol_time_course(
+    carousel: Carousel,
+    *,
+    p0: dict[str, float],
+    data: pd.DataFrame,
+    protocol: pd.DataFrame,
+    y0: dict[str, float] | None = None,
+    minimize_fn: MinimizeFn = _default_minimize_fn,
+    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.
+
+    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
+            for i in parallel.parallelise(
+                partial(
+                    _carousel_protocol_worker,
+                    p0=p0,
+                    data=data,
+                    protocol=protocol,
+                    y0=y0,
+                    integrator=integrator,
+                    loss_fn=loss_fn,
+                    minimize_fn=minimize_fn,
+                    residual_fn=residual_fn,
+                    bounds=bounds,
+                ),
+                inputs=list(enumerate(carousel.variants)),
+            )
+            if (fit := i[1]) is not None
+        ]
+    )

mxlpy/identify.py

@@ -28,16 +28,15 @@ def _mc_fit_time_course_worker(
     data: pd.DataFrame,
     loss_fn: fit.LossFn,
 ) -> float:
-    p_fit = fit.time_course(model=model, p0=p0.to_dict(), data=data)
-    return fit._time_course_residual(  # noqa: SLF001
-        par_values=list(p_fit.values()),
-        par_names=list(p_fit.keys()),
-        data=data,
+    fit_result = fit.time_course(
         model=model,
-        y0=None,
-        integrator=fit.DefaultIntegrator,
+        p0=p0.to_dict(),
+        data=data,
         loss_fn=loss_fn,
     )
+    if fit_result is None:
+        return np.inf
+    return fit_result.loss
 
 
 def profile_likelihood(
@@ -60,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/int_scipy.py

@@ -98,6 +98,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,