mxlpy 0.16.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/report.py

@@ -48,12 +48,39 @@ def markdown(
 ) -> str:
     """Generate a markdown report comparing two models.
 
-    Args:
-        m1: The first model to compare.
-        m2: The second model to compare.
-        analyses: A list of functions that take a Path and return a tuple of a string and a Path. Defaults to None.
-        rel_change: The relative change threshold for numerical differences. Defaults to 1e-2.
-        img_path: The path to save images. Defaults to Path().
+    Parameters
+    ----------
+    m1
+        The first model to compare
+    m2
+        The second model to compare
+    analyses
+        A list of functions that analyze both models and return a report section with image
+    rel_change
+        The relative change threshold for numerical differences
+    img_path
+        The path to save images
+
+    Returns
+    -------
+    str
+        Markdown formatted report comparing the two models
+
+    Examples
+    --------
+    >>> from mxlpy import Model
+    >>> m1 = Model().add_parameter("k1", 0.1).add_variable("S", 1.0)
+    >>> m2 = Model().add_parameter("k1", 0.2).add_variable("S", 1.0)
+    >>> report = markdown(m1, m2)
+    >>> "Parameters" in report and "k1" in report
+    True
+
+    >>> # With custom analysis function
+    >>> def custom_analysis(m1, m2, path):
+    ...     return "## Custom analysis", path / "image.png"
+    >>> report = markdown(m1, m2, analyses=[custom_analysis])
+    >>> "Custom analysis" in report
+    True
 
     """
     content: list[str] = [

mxlpy/sbml/_import.py

@@ -12,7 +12,7 @@ import libsbml
 import numpy as np  # noqa: F401  # models might need it
 import sympy
 
-from mxlpy.model import Model, _sort_dependencies
+from mxlpy.model import Dependency, Model, _sort_dependencies
 from mxlpy.paths import default_tmp_dir
 from mxlpy.sbml._data import (
     AtomicUnit,
@@ -522,7 +522,10 @@ def _codgen(name: str, sbml: Parser) -> Path:
         ^ set(variables)
         ^ set(sbml.derived)
         | {"time"},
-        elements=[(k, set(v.args)) for k, v in sbml.initial_assignment.items()],
+        elements=[
+            Dependency(name=k, required=set(v.args), provided={k})
+            for k, v in sbml.initial_assignment.items()
+        ],
     )
 
     if len(initial_assignment_order) > 0:

mxlpy/surrogates/__init__.py

@@ -19,13 +19,14 @@ from __future__ import annotations
 import contextlib
 
 with contextlib.suppress(ImportError):
-    from ._torch import TorchSurrogate, train_torch_surrogate
+    from ._torch import Torch, TorchTrainer, train_torch
 
-from ._poly import PolySurrogate, train_polynomial_surrogate
+from ._poly import Polynomial, train_polynomial
 
 __all__ = [
-    "PolySurrogate",
-    "TorchSurrogate",
-    "train_polynomial_surrogate",
-    "train_torch_surrogate",
+    "Polynomial",
+    "Torch",
+    "TorchTrainer",
+    "train_polynomial",
+    "train_torch",
 ]

mxlpy/surrogates/_poly.py

@@ -9,9 +9,9 @@ from numpy import polynomial
 from mxlpy.types import AbstractSurrogate, ArrayLike
 
 __all__ = [
-    "PolySurrogate",
+    "Polynomial",
     "PolynomialExpansion",
-    "train_polynomial_surrogate",
+    "train_polynomial",
 ]
 
 # define custom type
@@ -26,23 +26,24 @@ PolynomialExpansion = (
 
 
 @dataclass(kw_only=True)
-class PolySurrogate(AbstractSurrogate):
+class Polynomial(AbstractSurrogate):
     model: PolynomialExpansion
 
     def predict_raw(self, y: np.ndarray) -> np.ndarray:
         return self.model(y)
 
 
-def train_polynomial_surrogate(
-    feature: ArrayLike,
-    target: ArrayLike,
+def train_polynomial(
+    feature: ArrayLike | pd.Series,
+    target: ArrayLike | pd.Series,
     series: Literal[
         "Power", "Chebyshev", "Legendre", "Laguerre", "Hermite", "HermiteE"
     ] = "Power",
     degrees: Iterable[int] = (1, 2, 3, 4, 5, 6, 7),
     surrogate_args: list[str] | None = None,
+    surrogate_outputs: list[str] | None = None,
     surrogate_stoichiometries: dict[str, dict[str, float]] | None = None,
-) -> tuple[PolySurrogate, pd.DataFrame]:
+) -> tuple[Polynomial, pd.DataFrame]:
     """Train a surrogate model based on function series expansion.
 
     Args:
@@ -51,7 +52,8 @@ def train_polynomial_surrogate(
         series: Base functions for the surrogate model
         degrees: Degrees of the polynomial to fit to the data.
         surrogate_args: Additional arguments for the surrogate model.
-        surrogate_stoichiometries: Stoichiometries for the surrogate model.
+        surrogate_outputs: Names of the surrogate model outputs.
+        surrogate_stoichiometries: Mapping of variables to their stoichiometries
 
     Returns:
         PolySurrogate: Polynomial surrogate model.
@@ -83,9 +85,10 @@ def train_polynomial_surrogate(
     # Choose the model with the lowest AIC
     model = models[np.argmin(score)]
     return (
-        PolySurrogate(
+        Polynomial(
             model=model,
             args=surrogate_args if surrogate_args is not None else [],
+            outputs=surrogate_outputs if surrogate_outputs is not None else [],
             stoichiometries=surrogate_stoichiometries
             if surrogate_stoichiometries is not None
             else {},