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

mxlpy/plot.py

@@ -23,9 +23,12 @@ import contextlib
 from cycler import cycler
 
 __all__ = [
+    "Color",
     "FigAx",
     "FigAxs",
     "Linestyle",
+    "RGB",
+    "RGBA",
     "add_grid",
     "bars",
     "context",
@@ -71,7 +74,7 @@ from mpl_toolkits.mplot3d import Axes3D
 from mxlpy.label_map import LabelMapper
 
 if TYPE_CHECKING:
-    from collections.abc import Generator
+    from collections.abc import Generator, Iterable
 
     from matplotlib.collections import QuadMesh
 
@@ -89,6 +92,11 @@ type Linestyle = Literal[
     "dashdot",
 ]
 
+
+type RGB = tuple[float, float, float]
+type RGBA = tuple[float, float, float, float]
+type Color = str | RGB | RGBA
+
 ##########################################################################
 # Helpers
 ##########################################################################
@@ -133,6 +141,13 @@ def _get_norm(vmin: float, vmax: float) -> Normalize:
     return norm
 
 
+def _norm(df: pd.DataFrame) -> Normalize:
+    """Get a normalization object for the given data."""
+    vmin = df.min().min()
+    vmax = df.max().max()
+    return _get_norm(vmin, vmax)
+
+
 def _norm_with_zero_center(df: pd.DataFrame) -> Normalize:
     """Get a normalization object with zero-centered values for the given data."""
     v = max(abs(df.min().min()), abs(df.max().max()))
@@ -166,7 +181,7 @@ def _split_large_groups[T](groups: list[list[T]], max_size: int) -> list[list[T]
     )  # type: ignore
 
 
-def _default_color(ax: Axes, color: str | None) -> str:
+def _default_color(ax: Axes, color: Color | None) -> Color:
     """Get a default color for the given axis."""
     return f"C{len(ax.lines)}" if color is None else color
 
@@ -291,16 +306,16 @@ def reset_prop_cycle(ax: Axes) -> None:
 @contextlib.contextmanager
 def context(
     colors: list[str] | None = None,
-    line_width: float | None = None,
-    line_style: Linestyle | None = None,
+    linewidth: float | None = None,
+    linestyle: Linestyle | None = None,
     rc: dict[str, Any] | None = None,
 ) -> Generator[None, None, None]:
     """Context manager to set the defaults for plots.
 
     Args:
         colors: colors to use for the plot.
-        line_width: line width to use for the plot.
-        line_style: line style to use for the plot.
+        linewidth: line width to use for the plot.
+        linestyle: line style to use for the plot.
         rc: additional keyword arguments to pass to the rc context.
 
     """
@@ -309,11 +324,11 @@ def context(
     if colors is not None:
         rc["axes.prop_cycle"] = cycler(color=colors)
 
-    if line_width is not None:
-        rc["lines.linewidth"] = line_width
+    if linewidth is not None:
+        rc["lines.linewidth"] = linewidth
 
-    if line_style is not None:
-        rc["lines.linestyle"] = line_style
+    if linestyle is not None:
+        rc["lines.linestyle"] = linestyle
 
     with plt.rc_context(rc):
         yield
@@ -477,9 +492,12 @@ def lines(
     x: pd.DataFrame | pd.Series,
     *,
     ax: Axes | None = None,
-    grid: bool = True,
     alpha: float = 1.0,
+    color: Color | list[Color] | None = None,
+    grid: bool = True,
     legend: bool = True,
+    linewidth: float | None = None,
+    linestyle: Linestyle | None = None,
 ) -> FigAx:
     """Plot multiple lines on the same axis."""
     fig, ax = _default_fig_ax(ax=ax, grid=grid)
@@ -487,6 +505,9 @@ def lines(
         x.index,
         x,
         alpha=alpha,
+        linewidth=linewidth,
+        linestyle=linestyle,
+        color=color,
     )
     _default_labels(ax, xlabel=x.index.name, ylabel=None)
     if legend:
@@ -497,6 +518,12 @@ def lines(
     return fig, ax
 
 
+def _repeat_color_if_necessary(
+    color: list[list[Color]] | Color | None, n: int
+) -> Iterable[list[Color] | Color | None]:
+    return [color] * n if not isinstance(color, list) else color
+
+
 def lines_grouped(
     groups: list[pd.DataFrame] | list[pd.Series],
     *,
@@ -506,6 +533,9 @@ def lines_grouped(
     sharex: bool = True,
     sharey: bool = False,
     grid: bool = True,
+    color: Color | list[list[Color]] | None = None,
+    linewidth: float | None = None,
+    linestyle: Linestyle | None = None,
 ) -> FigAxs:
     """Plot multiple groups of lines on separate axes."""
     fig, axs = grid_layout(
@@ -518,8 +548,20 @@ def lines_grouped(
         grid=grid,
     )
 
-    for group, ax in zip(groups, axs, strict=False):
-        lines(group, ax=ax, grid=grid)
+    for group, ax, color_ in zip(
+        groups,
+        axs,
+        _repeat_color_if_necessary(color, n=len(groups)),
+        strict=False,
+    ):
+        lines(
+            group,
+            ax=ax,
+            grid=grid,
+            color=color_,
+            linewidth=linewidth,
+            linestyle=linestyle,
+        )
 
     for i in range(len(groups), len(axs)):
         axs[i].set_visible(False)
@@ -535,6 +577,9 @@ def line_autogrouped(
     row_height: float = 3,
     max_group_size: int = 6,
     grid: bool = True,
+    color: Color | list[list[Color]] | None = None,
+    linewidth: float | None = None,
+    linestyle: Linestyle | None = None,
 ) -> FigAxs:
     """Plot a series or dataframe with lines grouped by order of magnitude."""
     group_names = _split_large_groups(
@@ -556,6 +601,9 @@ def line_autogrouped(
         col_width=col_width,
         row_height=row_height,
         grid=grid,
+        color=color,
+        linestyle=linestyle,
+        linewidth=linewidth,
     )
 
 
@@ -564,7 +612,9 @@ def line_mean_std(
     *,
     label: str | None = None,
     ax: Axes | None = None,
-    color: str | None = None,
+    color: Color | None = None,
+    linewidth: float | None = None,
+    linestyle: Linestyle | None = None,
     alpha: float = 0.2,
     grid: bool = True,
 ) -> FigAx:
@@ -579,6 +629,8 @@ def line_mean_std(
         mean,
         color=color,
         label=label,
+        linewidth=linewidth,
+        linestyle=linestyle,
     )
     ax.fill_between(
         df.index,
@@ -598,6 +650,9 @@ def lines_mean_std_from_2d_idx(
     ax: Axes | None = None,
     alpha: float = 0.2,
     grid: bool = True,
+    color: Color | None = None,
+    linewidth: float | None = None,
+    linestyle: Linestyle | None = None,
 ) -> FigAx:
     """Plot the mean and standard deviation of a 2D indexed dataframe."""
     if len(cast(pd.MultiIndex, df.index).levels) != 2:  # noqa: PLR2004
@@ -612,25 +667,32 @@ def lines_mean_std_from_2d_idx(
             label=name,
             alpha=alpha,
             ax=ax,
+            color=color,
+            linestyle=linestyle,
+            linewidth=linewidth,
         )
     ax.legend()
     return fig, ax
 
 
-def heatmap(
-    df: pd.DataFrame,
+def _create_heatmap(
     *,
+    ax: Axes | None,
+    df: pd.DataFrame,
+    title: str | None = None,
+    xlabel: str,
+    ylabel: str,
+    xticklabels: list[str],
+    yticklabels: list[str],
+    norm: Normalize,
     annotate: bool = False,
     colorbar: bool = True,
     invert_yaxis: bool = True,
     cmap: str = "RdBu_r",
-    norm: Normalize | None = None,
-    ax: Axes | None = None,
     cax: Axes | None = None,
     sci_annotation_bounds: tuple[float, float] = (0.01, 100),
     annotation_style: str = "2g",
 ) -> tuple[Figure, Axes, QuadMesh]:
-    """Plot a heatmap of the given data."""
     fig, ax = _default_fig_ax(
         ax=ax,
         figsize=(
@@ -639,17 +701,23 @@ def heatmap(
         ),
         grid=False,
     )
-    if norm is None:
-        norm = _norm_with_zero_center(df)
 
+    # Note: pcolormesh swaps index/columns
     hm = ax.pcolormesh(df, norm=norm, cmap=cmap)
+
+    if xlabel is not None:
+        ax.set_xlabel(xlabel)
+    if ylabel is not None:
+        ax.set_ylabel(ylabel)
+    if title is not None:
+        ax.set_title(title)
     ax.set_xticks(
         np.arange(0, len(df.columns), 1) + 0.5,
-        labels=df.columns,
+        labels=xticklabels,
     )
     ax.set_yticks(
         np.arange(0, len(df.index), 1) + 0.5,
-        labels=df.index,
+        labels=yticklabels,
     )
 
     if annotate:
@@ -666,39 +734,77 @@ def heatmap(
     return fig, ax, hm
 
 
+def heatmap(
+    df: pd.DataFrame,
+    *,
+    ax: Axes | None = None,
+    title: str | None = None,
+    annotate: bool = False,
+    colorbar: bool = True,
+    invert_yaxis: bool = True,
+    cmap: str = "RdBu_r",
+    norm: Normalize | None = None,
+    cax: Axes | None = None,
+    sci_annotation_bounds: tuple[float, float] = (0.01, 100),
+    annotation_style: str = "2g",
+) -> tuple[Figure, Axes, QuadMesh]:
+    """Plot a heatmap of the given data."""
+    return _create_heatmap(
+        ax=ax,
+        df=df,
+        title=title,
+        xlabel=df.index.name,
+        ylabel=df.columns.name,
+        xticklabels=cast(list, df.columns),
+        yticklabels=cast(list, df.index),
+        annotate=annotate,
+        colorbar=colorbar,
+        invert_yaxis=invert_yaxis,
+        cmap=cmap,
+        norm=_norm_with_zero_center(df) if norm is None else norm,
+        cax=cax,
+        sci_annotation_bounds=sci_annotation_bounds,
+        annotation_style=annotation_style,
+    )
+
+
 def heatmap_from_2d_idx(
     df: pd.DataFrame,
     variable: str,
+    *,
     ax: Axes | None = None,
-) -> FigAx:
+    annotate: bool = False,
+    colorbar: bool = True,
+    invert_yaxis: bool = False,
+    cmap: str = "viridis",
+    norm: Normalize | None = None,
+    cax: Axes | None = None,
+    sci_annotation_bounds: tuple[float, float] = (0.01, 100),
+    annotation_style: str = "2g",
+) -> tuple[Figure, Axes, QuadMesh]:
     """Plot a heatmap of a 2D indexed dataframe."""
     if len(cast(pd.MultiIndex, df.index).levels) != 2:  # noqa: PLR2004
         msg = "MultiIndex must have exactly two levels"
         raise ValueError(msg)
-
-    fig, ax = _default_fig_ax(ax=ax, grid=False)
-    df2d = df[variable].unstack()
-
-    ax.set_title(variable)
-    # Note: pcolormesh swaps index/columns
-    hm = ax.pcolormesh(df2d.T)
-    ax.set_xlabel(df2d.index.name)
-    ax.set_ylabel(df2d.columns.name)
-    ax.set_xticks(
-        np.arange(0, len(df2d.index), 1) + 0.5,
-        labels=[f"{i:.2f}" for i in df2d.index],
-    )
-    ax.set_yticks(
-        np.arange(0, len(df2d.columns), 1) + 0.5,
-        labels=[f"{i:.2f}" for i in df2d.columns],
+    df2d = df[variable].unstack().T
+
+    return _create_heatmap(
+        df=df2d,
+        xlabel=df2d.index.name,
+        ylabel=df2d.columns.name,
+        xticklabels=[f"{i:.2f}" for i in df2d.columns],
+        yticklabels=[f"{i:.2f}" for i in df2d.index],
+        ax=ax,
+        cax=cax,
+        annotate=annotate,
+        colorbar=colorbar,
+        invert_yaxis=invert_yaxis,
+        cmap=cmap,
+        norm=_norm(df2d) if norm is None else norm,
+        sci_annotation_bounds=sci_annotation_bounds,
+        annotation_style=annotation_style,
     )
 
-    rotate_xlabels(ax, rotation=45, ha="right")
-
-    # Add colorbar
-    fig.colorbar(hm, ax=ax)
-    return fig, ax
-
 
 def heatmaps_from_2d_idx(
     df: pd.DataFrame,
@@ -708,6 +814,13 @@ def heatmaps_from_2d_idx(
     row_height_factor: float = 0.6,
     sharex: bool = True,
     sharey: bool = False,
+    annotate: bool = False,
+    colorbar: bool = True,
+    invert_yaxis: bool = False,
+    cmap: str = "viridis",
+    norm: Normalize | None = None,
+    sci_annotation_bounds: tuple[float, float] = (0.01, 100),
+    annotation_style: str = "2g",
 ) -> FigAxs:
     """Plot multiple heatmaps of a 2D indexed dataframe."""
     idx = cast(pd.MultiIndex, df.index)
@@ -722,7 +835,18 @@ def heatmaps_from_2d_idx(
         grid=False,
     )
     for ax, var in zip(axs, df.columns, strict=False):
-        heatmap_from_2d_idx(df, var, ax=ax)
+        heatmap_from_2d_idx(
+            df,
+            var,
+            ax=ax,
+            annotate=annotate,
+            colorbar=colorbar,
+            invert_yaxis=invert_yaxis,
+            cmap=cmap,
+            norm=norm,
+            sci_annotation_bounds=sci_annotation_bounds,
+            annotation_style=annotation_style,
+        )
     return fig, axs
 
 
@@ -892,6 +1016,9 @@ def relative_label_distribution(
     row_height: float = 3,
     sharey: bool = False,
     grid: bool = True,
+    color: Color | None = None,
+    linewidth: float | None = None,
+    linestyle: Linestyle | None = None,
 ) -> FigAxs:
     """Plot the relative distribution of labels in the given data."""
     variables = list(mapper.label_variables) if subset is None else subset
@@ -903,20 +1030,37 @@ def relative_label_distribution(
         sharey=sharey,
         grid=grid,
     )
+    # FIXME: rewrite as building a dict of dataframes
+    # and passing it to lines_grouped
     if isinstance(mapper, LabelMapper):
         for ax, name in zip(axs, variables, strict=False):
             for i in range(mapper.label_variables[name]):
                 isos = mapper.get_isotopomers_of_at_position(name, i)
                 labels = cast(pd.DataFrame, concs.loc[:, isos])
                 total = concs.loc[:, f"{name}__total"]
-                ax.plot(labels.index, (labels.sum(axis=1) / total), label=f"C{i + 1}")
+                ax.plot(
+                    labels.index,
+                    (labels.sum(axis=1) / total),
+                    label=f"C{i + 1}",
+                    linewidth=linewidth,
+                    linestyle=linestyle,
+                    color=color,
+                )
             ax.set_title(name)
             ax.legend()
     else:
         for ax, (name, isos) in zip(
-            axs, mapper.get_isotopomers(variables).items(), strict=False
+            axs,
+            mapper.get_isotopomers(variables).items(),
+            strict=False,
         ):
-            ax.plot(concs.index, concs.loc[:, isos])
+            ax.plot(
+                concs.index,
+                concs.loc[:, isos],
+                linewidth=linewidth,
+                linestyle=linestyle,
+                color=color,
+            )
             ax.set_title(name)
             ax.legend([f"C{i + 1}" for i in range(len(isos))])
 

mxlpy/surrogates/_torch.py

@@ -5,12 +5,11 @@ from typing import Self
 import numpy as np
 import pandas as pd
 import torch
-import tqdm
 from torch import nn
 from torch.optim.adam import Adam
 from torch.optim.optimizer import ParamsT
 
-from mxlpy.nn._torch import MLP, DefaultDevice
+from mxlpy.nn._torch import MLP, DefaultDevice, train
 from mxlpy.types import AbstractSurrogate
 
 type LossFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
@@ -106,28 +105,16 @@ class TorchTrainer:
         epochs: int,
         batch_size: int | None = None,
     ) -> Self:
-        if batch_size is None:
-            losses = _train_full(
-                aprox=self.approximator,
-                features=self.features,
-                targets=self.targets,
-                epochs=epochs,
-                optimizer=self.optimizer,
-                device=self.device,
-                loss_fn=self.loss_fn,
-            )
-        else:
-            losses = _train_batched(
-                aprox=self.approximator,
-                features=self.features,
-                targets=self.targets,
-                epochs=epochs,
-                optimizer=self.optimizer,
-                device=self.device,
-                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,
+            device=self.device,
+            loss_fn=self.loss_fn,
+        )
         if len(self.losses) > 0:
             losses.index += self.losses[-1].index[-1]
         self.losses.append(losses)
@@ -152,83 +139,6 @@ class TorchTrainer:
         )
 
 
-def _train_batched(
-    aprox: nn.Module,
-    features: pd.DataFrame,
-    targets: pd.DataFrame,
-    epochs: int,
-    optimizer: Adam,
-    device: torch.device,
-    batch_size: int,
-    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.
-
-    """
-    rng = np.random.default_rng()
-    losses = {}
-    for i in tqdm.trange(epochs):
-        idxs = rng.choice(features.index, size=batch_size)
-        X = torch.Tensor(features.iloc[idxs].to_numpy(), device=device)
-        Y = torch.Tensor(targets.iloc[idxs].to_numpy(), device=device)
-        optimizer.zero_grad()
-        loss = loss_fn(aprox(X), Y)
-        loss.backward()
-        optimizer.step()
-        losses[i] = loss.detach().numpy()
-    return pd.Series(losses, dtype=float)
-
-
-def _train_full(
-    aprox: nn.Module,
-    features: pd.DataFrame,
-    targets: pd.DataFrame,
-    epochs: int,
-    optimizer: Adam,
-    device: torch.device,
-    loss_fn: Callable,
-) -> pd.Series:
-    """Train the neural network using full-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
-        loss_fn: Loss function
-
-    Returns:
-        pd.Series: Series containing the training loss history.
-
-    """
-    X = torch.Tensor(features.to_numpy(), device=device)
-    Y = torch.Tensor(targets.to_numpy(), device=device)
-
-    losses = {}
-    for i in tqdm.trange(epochs):
-        optimizer.zero_grad()
-        loss = loss_fn(aprox(X), Y)
-        loss.backward()
-        optimizer.step()
-        losses[i] = loss.detach().numpy()
-    return pd.Series(losses, dtype=float)
-
-
 def train_torch(
     features: pd.DataFrame,
     targets: pd.DataFrame,

mxlpy/types.py

@@ -46,9 +46,6 @@ __all__ = [
     "unwrap2",
 ]
 
-# Re-exporting some types here, because their imports have
-# changed between Python versions and I have no interest in
-# fixing it in every file
 from collections.abc import Callable, Iterator, Mapping
 from typing import TYPE_CHECKING, Any, ParamSpec, Protocol, TypeVar, cast
 

{mxlpy-0.17.0.dist-info → mxlpy-0.18.0.dist-info}/METADATA

@@ -1,17 +1,17 @@
 Metadata-Version: 2.4
 Name: mxlpy
-Version: 0.17.0
+Version: 0.18.0
 Summary: A package to build metabolic models
 Author-email: Marvin van Aalst <marvin.vanaalst@gmail.com>
 Maintainer-email: Marvin van Aalst <marvin.vanaalst@gmail.com>
-License-Expression: GPL-3.0-or-later
+License-Expression: MIT
 License-File: LICENSE
 Keywords: metabolic,modelling,ode
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: Console
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)
+Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: MacOS
 Classifier: Operating System :: Microsoft :: Windows
 Classifier: Operating System :: OS Independent
@@ -67,7 +67,7 @@ Description-Content-Type: text/markdown
 
 [![pypi](https://img.shields.io/pypi/v/mxlpy.svg)](https://pypi.python.org/pypi/mxlpy)
 [![docs][docs-badge]][docs]
-![License](https://img.shields.io/badge/license-GPL--3.0-blue?style=flat-square)
+![License](https://img.shields.io/badge/license-MIT-blue?style=flat-square)
 ![Coverage](https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fgist.github.com%2Fmarvinvanaalst%2F98ab3ce1db511de42f9871e91d85e4cd%2Fraw%2Fcoverage.json&query=%24.message&label=Coverage&color=%24.color&suffix=%20%25)
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 [![security: bandit](https://img.shields.io/badge/security-bandit-yellow.svg)](https://github.com/PyCQA/bandit)
@@ -90,10 +90,10 @@ pixi add --pypi mxlpy[torch]
 
 ## How to cite
 
-If you use this software in your scientific work, please cite [this article](...):
+If you use this software in your scientific work, please cite [this article](https://doi.org/10.1101/2025.05.06.652335):
 
-- [doi](https://doi.org/)
-- [bibtex file](https://fillme.out)
+- [doi](https://doi.org/10.1101/2025.05.06.652335)
+- [bibtex file](https://github.com/Computational-Biology-Aachen/MxlPy/citation.bibtex)
 
 
 ## Development setup