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

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/nn/_torch.py

@@ -8,17 +8,77 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, cast
 
+import numpy as np
+import pandas as pd
 import torch
+import tqdm
 from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+type LossFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-__all__ = ["DefaultDevice", "LSTM", "MLP"]
+    from torch.optim.adam import Adam
+
+    from mxlpy.types import Array
+
+__all__ = ["DefaultDevice", "LSTM", "LossFn", "MLP", "train"]
 
 DefaultDevice = torch.device("cpu")
 
 
+def train(
+    aprox: nn.Module,
+    features: Array,
+    targets: Array,
+    epochs: int,
+    optimizer: Adam,
+    device: torch.device,
+    batch_size: int | None,
+    loss_fn: LossFn,
+) -> pd.Series:
+    """Train the neural network using mini-batch gradient descent.
+
+    Args:
+        aprox: Neural network model to train.
+        features: Input features as a tensor.
+        targets: Target values as a tensor.
+        epochs: Number of training epochs.
+        optimizer: Optimizer for training.
+        device: torch device
+        batch_size: Size of mini-batches for training.
+        loss_fn: Loss function
+
+    Returns:
+        pd.Series: Series containing the training loss history.
+
+    """
+    losses = {}
+
+    data = TensorDataset(
+        torch.tensor(features.astype(np.float32), dtype=torch.float32, device=device),
+        torch.tensor(targets.astype(np.float32), dtype=torch.float32, device=device),
+    )
+    data_loader = DataLoader(
+        data,
+        batch_size=len(features) if batch_size is None else batch_size,
+        shuffle=True,
+    )
+
+    for i in tqdm.trange(epochs):
+        epoch_loss = 0
+        for xb, yb in data_loader:
+            optimizer.zero_grad()
+            loss = loss_fn(aprox(xb), yb)
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item() * xb.size(0)
+        losses[i] = epoch_loss / len(data_loader.dataset)  # type: ignore
+    return pd.Series(losses, dtype=float)
+
+
 class MLP(nn.Module):
     """Multilayer Perceptron (MLP) for surrogate modeling and neural posterior estimation.
 

mxlpy/npe/_torch.py

@@ -18,11 +18,10 @@ from typing import TYPE_CHECKING, Self, cast
 import numpy as np
 import pandas as pd
 import torch
-import tqdm
 from torch import nn
 from torch.optim.adam import Adam
 
-from mxlpy.nn._torch import LSTM, MLP, DefaultDevice
+from mxlpy.nn._torch import LSTM, MLP, DefaultDevice, train
 from mxlpy.parallel import Cache
 from mxlpy.types import AbstractEstimator
 
@@ -161,28 +160,16 @@ class TorchSteadyStateTrainer:
             batch_size: Size of mini-batches for training (None for full-batch)
 
         """
-        features = torch.Tensor(self.features.to_numpy(), device=self.device)
-        targets = torch.Tensor(self.targets.to_numpy(), device=self.device)
-
-        if batch_size is None:
-            losses = _train_full(
-                approximator=self.approximator,
-                features=features,
-                targets=targets,
-                epochs=epochs,
-                optimizer=self.optimizer,
-                loss_fn=self.loss_fn,
-            )
-        else:
-            losses = _train_batched(
-                approximator=self.approximator,
-                features=features,
-                targets=targets,
-                epochs=epochs,
-                optimizer=self.optimizer,
-                batch_size=batch_size,
-                loss_fn=self.loss_fn,
-            )
+        losses = train(
+            aprox=self.approximator,
+            features=self.features.to_numpy(),
+            targets=self.targets.to_numpy(),
+            epochs=epochs,
+            optimizer=self.optimizer,
+            batch_size=batch_size,
+            loss_fn=self.loss_fn,
+            device=self.device,
+        )
 
         if len(self.losses) > 0:
             losses.index += self.losses[-1].index[-1]
@@ -260,37 +247,22 @@ class TorchTimeCourseTrainer:
             batch_size: Size of mini-batches for training (None for full-batch)
 
         """
-        features = torch.Tensor(
-            np.swapaxes(
+        losses = train(
+            aprox=self.approximator,
+            features=np.swapaxes(
                 self.features.to_numpy().reshape(
                     (len(self.targets), -1, len(self.features.columns))
                 ),
                 axis1=0,
                 axis2=1,
             ),
+            targets=self.targets.to_numpy(),
+            epochs=epochs,
+            optimizer=self.optimizer,
+            batch_size=batch_size,
+            loss_fn=self.loss_fn,
             device=self.device,
         )
-        targets = torch.Tensor(self.targets.to_numpy(), device=self.device)
-
-        if batch_size is None:
-            losses = _train_full(
-                approximator=self.approximator,
-                features=features,
-                targets=targets,
-                epochs=epochs,
-                optimizer=self.optimizer,
-                loss_fn=self.loss_fn,
-            )
-        else:
-            losses = _train_batched(
-                approximator=self.approximator,
-                features=features,
-                targets=targets,
-                epochs=epochs,
-                optimizer=self.optimizer,
-                batch_size=batch_size,
-                loss_fn=self.loss_fn,
-            )
 
         if len(self.losses) > 0:
             losses.index += self.losses[-1].index[-1]
@@ -309,49 +281,6 @@ class TorchTimeCourseTrainer:
         )
 
 
-def _train_batched(
-    approximator: nn.Module,
-    features: torch.Tensor,
-    targets: torch.Tensor,
-    epochs: int,
-    optimizer: Adam,
-    batch_size: int,
-    loss_fn: LossFn,
-) -> pd.Series:
-    losses = {}
-    for epoch in tqdm.trange(epochs):
-        permutation = torch.randperm(features.size()[0])
-        epoch_loss = 0
-        for i in range(0, features.size()[0], batch_size):
-            optimizer.zero_grad()
-            indices = permutation[i : i + batch_size]
-            loss = loss_fn(approximator(features[indices]), targets[indices])
-            loss.backward()
-            optimizer.step()
-            epoch_loss += loss.detach().numpy()
-
-        losses[epoch] = epoch_loss / (features.size()[0] / batch_size)
-    return pd.Series(losses, dtype=float)
-
-
-def _train_full(
-    approximator: nn.Module,
-    features: torch.Tensor,
-    targets: torch.Tensor,
-    epochs: int,
-    optimizer: Adam,
-    loss_fn: LossFn,
-) -> pd.Series:
-    losses = {}
-    for i in tqdm.trange(epochs):
-        optimizer.zero_grad()
-        loss = loss_fn(approximator(features), targets)
-        loss.backward()
-        optimizer.step()
-        losses[i] = loss.detach().numpy()
-    return pd.Series(losses, dtype=float)
-
-
 def train_torch_steady_state(
     features: pd.DataFrame,
     targets: pd.DataFrame,