mxlpy 0.17.0__py3-none-any.whl → 0.19.0__py3-none-any.whl

mxlpy/__init__.py

@@ -48,18 +48,15 @@ from . import (
     fns,
     mc,
     mca,
-    npe,
     plot,
     report,
     sbml,
-    surrogates,
 )
 from .integrators import DefaultIntegrator, Scipy
 from .label_map import LabelMapper
 from .linear_label_map import LinearLabelMapper
 from .mc import Cache
 from .model import Model
-from .nn import _torch
 from .scan import (
     steady_state,
     time_course,
@@ -75,6 +72,19 @@ with contextlib.suppress(ImportError):
 if TYPE_CHECKING:
     from mxlpy.types import ArrayLike
 
+    from . import (
+        nn,
+        npe,
+        surrogates,
+    )
+else:
+    from lazy_import import lazy_module
+
+    nn = lazy_module("mxlpy.nn")
+    npe = lazy_module("mxlpy.npe")
+    surrogates = lazy_module("mxlpy.surrogates")
+
+
 __all__ = [
     "Assimulo",
     "Cache",
@@ -87,7 +97,6 @@ __all__ = [
     "Scipy",
     "Simulator",
     "SymbolicModel",
-    "_torch",
     "cartesian_product",
     "distributions",
     "experimental",
@@ -96,6 +105,7 @@ __all__ = [
     "make_protocol",
     "mc",
     "mca",
+    "nn",
     "npe",
     "plot",
     "report",

mxlpy/experimental/diff.py

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

mxlpy/fit.py

@@ -28,12 +28,16 @@ from mxlpy.types import (
 
 __all__ = [
     "InitialGuess",
+    "LossFn",
     "MinimizeFn",
+    "ProtocolResidualFn",
     "ResidualFn",
     "SteadyStateResidualFn",
     "TimeSeriesResidualFn",
+    "rmse",
     "steady_state",
     "time_course",
+    "time_course_over_protocol",
 ]
 
 if TYPE_CHECKING:
@@ -44,6 +48,21 @@ if TYPE_CHECKING:
 type InitialGuess = dict[str, float]
 type ResidualFn = Callable[[Array], float]
 type MinimizeFn = Callable[[ResidualFn, InitialGuess], dict[str, float]]
+type LossFn = Callable[
+    [
+        pd.DataFrame | pd.Series,
+        pd.DataFrame | pd.Series,
+    ],
+    float,
+]
+
+
+def rmse(
+    y_pred: pd.DataFrame | pd.Series,
+    y_true: pd.DataFrame | pd.Series,
+) -> float:
+    """Calculate root mean square error between model and data."""
+    return cast(float, np.sqrt(np.mean(np.square(y_pred - y_true))))
 
 
 class SteadyStateResidualFn(Protocol):
@@ -58,6 +77,7 @@ class SteadyStateResidualFn(Protocol):
         model: Model,
         y0: dict[str, float],
         integrator: IntegratorType,
+        loss_fn: LossFn,
     ) -> float:
         """Calculate residual error between model steady state and experimental data."""
         ...
@@ -75,6 +95,27 @@ class TimeSeriesResidualFn(Protocol):
         model: Model,
         y0: dict[str, float],
         integrator: IntegratorType,
+        loss_fn: LossFn,
+    ) -> float:
+        """Calculate residual error between model time course and experimental data."""
+        ...
+
+
+class ProtocolResidualFn(Protocol):
+    """Protocol for time series residual functions."""
+
+    def __call__(
+        self,
+        par_values: Array,
+        # This will be filled out by partial
+        par_names: list[str],
+        data: pd.DataFrame,
+        model: Model,
+        y0: dict[str, float],
+        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."""
         ...
@@ -109,6 +150,7 @@ def _steady_state_residual(
     model: Model,
     y0: dict[str, float] | None,
     integrator: IntegratorType,
+    loss_fn: LossFn,
 ) -> float:
     """Calculate residual error between model steady state and experimental data.
 
@@ -119,6 +161,7 @@ def _steady_state_residual(
         y0: Initial conditions
         par_names: Names of parameters being fit
         integrator: ODE integrator class to use
+        loss_fn: Loss function to use for residual calculation
 
     Returns:
         float: Root mean square error between model and data
@@ -143,9 +186,11 @@ def _steady_state_residual(
     )
     if res is None:
         return cast(float, np.inf)
-    results_ss = res.get_combined()
-    diff = data - results_ss.loc[:, data.index]  # type: ignore
-    return cast(float, np.sqrt(np.mean(np.square(diff))))
+
+    return loss_fn(
+        res.get_combined().loc[:, cast(list, data.index)],
+        data,
+    )
 
 
 def _time_course_residual(
@@ -156,6 +201,53 @@ def _time_course_residual(
     model: Model,
     y0: dict[str, float] | None,
     integrator: IntegratorType,
+    loss_fn: LossFn,
+) -> float:
+    """Calculate residual error between model time course and experimental data.
+
+    Args:
+        par_values: Parameter values to test
+        data: Experimental time course data
+        model: Model instance to simulate
+        y0: Initial conditions
+        par_names: Names of parameters being fit
+        integrator: ODE integrator class to use
+        loss_fn: Loss function to use for residual calculation
+
+    Returns:
+        float: Root mean square error between model and data
+
+    """
+    res = (
+        Simulator(
+            model.update_parameters(dict(zip(par_names, par_values, strict=True))),
+            y0=y0,
+            integrator=integrator,
+        )
+        .simulate_time_course(cast(list, data.index))
+        .get_result()
+    )
+    if res is None:
+        return cast(float, np.inf)
+    results_ss = res.get_combined()
+
+    return loss_fn(
+        results_ss.loc[:, cast(list, data.columns)],
+        data,
+    )
+
+
+def _protocol_residual(
+    par_values: ArrayLike,
+    # This will be filled out by partial
+    par_names: list[str],
+    data: pd.DataFrame,
+    model: Model,
+    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.
 
@@ -166,6 +258,9 @@ def _time_course_residual(
         y0: Initial conditions
         par_names: Names of parameters being fit
         integrator: ODE integrator class to use
+        loss_fn: Loss function to use for residual calculation
+        protocol: Experimental protocol
+        time_points_per_step: Number of time points per step in the protocol
 
     Returns:
         float: Root mean square error between model and data
@@ -177,14 +272,20 @@ def _time_course_residual(
             y0=y0,
             integrator=integrator,
         )
-        .simulate_time_course(data.index)  # type: ignore
+        .simulate_over_protocol(
+            protocol=protocol,
+            time_points_per_step=time_points_per_step,
+        )
         .get_result()
     )
     if res is None:
         return cast(float, np.inf)
     results_ss = res.get_combined()
-    diff = data - results_ss.loc[:, data.columns]  # type: ignore
-    return cast(float, np.sqrt(np.mean(np.square(diff))))
+
+    return loss_fn(
+        results_ss.loc[:, cast(list, data.columns)],
+        data,
+    )
 
 
 def steady_state(
@@ -195,6 +296,7 @@ def steady_state(
     minimize_fn: MinimizeFn = _default_minimize_fn,
     residual_fn: SteadyStateResidualFn = _steady_state_residual,
     integrator: IntegratorType = DefaultIntegrator,
+    loss_fn: LossFn = rmse,
 ) -> dict[str, float]:
     """Fit model parameters to steady-state experimental data.
 
@@ -210,6 +312,7 @@ def steady_state(
         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
 
     Returns:
         dict[str, float]: Fitted parameters as {parameter_name: fitted_value}
@@ -232,6 +335,7 @@ def steady_state(
             y0=y0,
             par_names=par_names,
             integrator=integrator,
+            loss_fn=loss_fn,
         ),
     )
     res = minimize_fn(fn, p0)
@@ -249,6 +353,62 @@ def time_course(
     minimize_fn: MinimizeFn = _default_minimize_fn,
     residual_fn: TimeSeriesResidualFn = _time_course_residual,
     integrator: IntegratorType = DefaultIntegrator,
+    loss_fn: LossFn = rmse,
+) -> dict[str, float]:
+    """Fit model parameters to time course of experimental data.
+
+    Examples:
+        >>> time_course(model, p0, data)
+        {'k1': 0.1, 'k2': 0.2}
+
+    Args:
+        model: Model instance to fit
+        data: Experimental time course data
+        p0: Initial parameter guesses as {parameter_name: value}
+        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
+
+    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
+
+    """
+    par_names = list(p0.keys())
+    p_orig = model.parameters
+
+    fn = cast(
+        ResidualFn,
+        partial(
+            residual_fn,
+            data=data,
+            model=model,
+            y0=y0,
+            par_names=par_names,
+            integrator=integrator,
+            loss_fn=loss_fn,
+        ),
+    )
+    res = minimize_fn(fn, p0)
+    model.update_parameters(p_orig)
+    return res
+
+
+def time_course_over_protocol(
+    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,
+    loss_fn: LossFn = rmse,
+    time_points_per_step: int = 10,
 ) -> dict[str, float]:
     """Fit model parameters to time course of experimental data.
 
@@ -258,12 +418,15 @@ def time_course(
 
     Args:
         model: Model instance to fit
-        data: Experimental time course data as pandas DataFrame
         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
 
     Returns:
         dict[str, float]: Fitted parameters as {parameter_name: fitted_value}
@@ -284,6 +447,9 @@ def time_course(
             y0=y0,
             par_names=par_names,
             integrator=integrator,
+            loss_fn=loss_fn,
+            protocol=protocol,
+            time_points_per_step=time_points_per_step,
         ),
     )
     res = minimize_fn(fn, p0)

mxlpy/identify.py

@@ -19,6 +19,7 @@ def _mc_fit_time_course_worker(
     p0: pd.Series,
     model: Model,
     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
@@ -28,6 +29,7 @@ def _mc_fit_time_course_worker(
         model=model,
         y0=None,
         integrator=fit.DefaultIntegrator,
+        loss_fn=loss_fn,
     )
 
 
@@ -37,6 +39,7 @@ def profile_likelihood(
     parameter_name: str,
     parameter_values: Array,
     n_random: int = 10,
+    loss_fn: fit.LossFn = fit.rmse,
 ) -> pd.Series:
     """Estimate the profile likelihood of model parameters given data.
 
@@ -46,6 +49,7 @@ def profile_likelihood(
         parameter_name: The name of the parameter to profile.
         parameter_values: The values of the parameter to profile.
         n_random: Number of Monte Carlo samples.
+        loss_fn: Loss function to use for fitting.
 
     """
     parameter_distributions = sample(
@@ -57,7 +61,9 @@ def profile_likelihood(
     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),
+            partial(
+                _mc_fit_time_course_worker, model=model, data=data, loss_fn=loss_fn
+            ),
             inputs=list(
                 parameter_distributions.drop(columns=parameter_name).iterrows()
             ),

mxlpy/integrators/int_assimulo.py

@@ -97,7 +97,11 @@ class Assimulo:
         if steps is None:
             steps = 0
         try:
-            return self.integrator.simulate(t_end, steps)  # type: ignore
+            t, y = self.integrator.simulate(t_end, steps)
+            return (
+                np.atleast_1d(np.array(t, dtype=float)),
+                np.atleast_2d(np.array(y, dtype=float)),
+            )
         except CVodeError:
             return None, None
 
@@ -116,8 +120,11 @@ class Assimulo:
 
         """
         try:
-            t, y = self.integrator.simulate(time_points[-1], 0, time_points)  # type: ignore
-            return np.array(t, dtype=float), np.array(y, dtype=float)
+            t, y = self.integrator.simulate(time_points[-1], 0, time_points)
+            return (
+                np.atleast_1d(np.array(t, dtype=float)),
+                np.atleast_2d(np.array(y, dtype=float)),
+            )
         except CVodeError:
             return None, None
 

mxlpy/integrators/int_scipy.py

@@ -108,10 +108,14 @@ class Scipy:
             rtol=self.rtol,
             method="LSODA",
         )
+
         if res.success:
-            self.t0 = time_points[-1]
-            self.y0 = res.y[:, -1]
-            return np.array(time_points, dtype=float), res.y.T
+            t = np.atleast_1d(np.array(res.t, dtype=float))
+            y = np.atleast_2d(np.array(res.y, dtype=float).T)
+
+            self.t0 = t[-1]
+            self.y0 = y[-1]
+            return t, y
         return None, None
 
     def integrate_to_steady_state(

mxlpy/label_map.py

@@ -29,6 +29,8 @@ __all__ = ["LabelMapper"]
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping
 
+    from mxlpy.types import Derived
+
 
 def _total_concentration(*args: float) -> float:
     """Calculate sum of isotopomer concentrations.
@@ -552,7 +554,7 @@ class LabelMapper:
         for name, dp in self.model.derived_parameters.items():
             m.add_derived(name, fn=dp.fn, args=dp.args)
 
-        variables: dict[str, float] = {}
+        variables: dict[str, float | Derived] = {}
         for k, v in self.model.variables.items():
             if (isos := isotopomers.get(k)) is None:
                 variables[k] = v

mxlpy/meta/codegen_latex.py

@@ -694,7 +694,7 @@ class TexExport:
 def _to_tex_export(self: Model) -> TexExport:
     return TexExport(
         parameters=self.parameters,
-        variables=self.variables,
+        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()},
         stoichiometries={k: v.stoichiometry for k, v in self.reactions.items()},

mxlpy/meta/source_tools.py

@@ -434,5 +434,5 @@ def _handle_call(node: ast.Call, ctx: Context) -> sympy.Expr:
                 ctx=ctx.updated(parent_module=imports[module_name.id]),
             )
 
-    msg = f"Onsupported function type {node.func}"
+    msg = f"Unsupported function type {node.func}"
     raise NotImplementedError(msg)