sindy-exp 0.2.0__py3-none-any.whl

sindy_exp/_plotting.py

@@ -0,0 +1,544 @@
+from dataclasses import dataclass
+from typing import Annotated, Optional, Sequence
+from warnings import warn
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+import scipy
+import seaborn as sns
+import sympy as sp
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from matplotlib.typing import ColorType
+
+PAL = sns.color_palette("Set1")
+PLOT_KWS = {"alpha": 0.7, "linewidth": 3}
+
+
+@dataclass
+class _ColorConstants:
+    color_sequence: list[ColorType]
+
+    def set_sequence(self, color_sequence: list[ColorType]):
+        self.color_sequence = color_sequence
+
+    @property
+    def TRUE(self):
+        return self.color_sequence[0]
+
+    @property
+    def MEAS(self):
+        return self.color_sequence[1]
+
+    @property
+    def EST(self):
+        return self.color_sequence[2]
+
+    @property
+    def SIM(self):
+        return self.color_sequence[3]
+
+    @property
+    def TRAIN(self):
+        return self.color_sequence[4]
+
+    @property
+    def TEST(self):
+        return self.color_sequence[5]
+
+
+COLOR = _ColorConstants(mpl.color_sequences["tab10"])
+
+
+def plot_coefficients(
+    coefficients: Annotated[np.ndarray, "(n_coord, n_features)"],
+    input_features: Sequence[str],
+    feature_names: Sequence[str],
+    ax: Axes,
+    **heatmap_kws,
+) -> None:
+    """Plot a set of dynamical system coefficients in a heatmap.
+
+    Args:
+        coefficients: A 2D array holding the coefficients of different
+            library functions.  System dimension is rows, function index
+            is columns
+        input_features: system coordinate names, e.g. "x","y","z" or "u","v"
+        feature_names: the names of the functions in the library.
+        ax: the matplotlib axis to plot on
+        **heatmap_kws: additional kwargs to seaborn's styling
+    """
+
+    def detex(input: str) -> str:
+        if input[0] == "$":
+            input = input[1:]
+        if input[-1] == "$":
+            input = input[:-1]
+        return input
+
+    if input_features is None:
+        input_features = [r"$\dot x_" + f"{k}$" for k in range(coefficients.shape[0])]
+    else:
+        input_features = [r"$\dot " + f"{detex(fi)}$" for fi in input_features]
+
+    if feature_names is None:
+        feature_names = [f"f{k}" for k in range(coefficients.shape[1])]
+
+    with sns.axes_style(style="white", rc={"axes.facecolor": (0, 0, 0, 0)}):
+        heatmap_args = {
+            "xticklabels": input_features,
+            "yticklabels": feature_names,
+            "center": 0.0,
+            "cmap": sns.color_palette("vlag", n_colors=20, as_cmap=True),
+            "ax": ax,
+            "linewidths": 0.1,
+            "linecolor": "whitesmoke",
+        }
+        heatmap_args.update(**heatmap_kws)
+        coefficients = np.where(
+            coefficients == 0, np.nan * np.empty_like(coefficients), coefficients
+        )
+        sns.heatmap(coefficients.T, **heatmap_args)
+
+        ax.tick_params(axis="y", rotation=0)
+
+    return ax
+
+
+def _coeff_dicts_to_matrix(
+    coeffs: Sequence[dict[sp.Expr, float]],
+) -> tuple[np.ndarray, list[str]]:
+    """Convert a list of coefficient dictionaries to a dense matrix.
+
+    The input is a list of dictionaries mapping feature identifiers (
+    SymPy expressions) to coefficients. All dictionaries are assumed to
+    correspond to different coordinates of the same system. This helper builds
+    a consistent feature ordering across coordinates and returns a numeric
+    matrix along with the stringified feature names.
+    """
+
+    if not coeffs:
+        raise ValueError("No coefficient dictionaries provided.")
+
+    features: list[sp.Expr] = sorted({key for d in coeffs for key in d.keys()}, key=str)
+
+    mat = np.zeros((len(coeffs), len(features)), dtype=float)
+    for row, d in enumerate(coeffs):
+        for col, feat in enumerate(features):
+            mat[row, col] = d[feat]
+
+    feature_names = [str(f).replace("**", "^") for f in features]
+    return mat, feature_names
+
+
+def _compare_coefficient_plots_impl(
+    coefficients_est: Annotated[np.ndarray, "(n_coord, n_feat)"],
+    coefficients_true: Annotated[np.ndarray, "(n_coord, n_feat)"],
+    input_features: Sequence[str],
+    feature_names: Sequence[str],
+    scaling: bool = True,
+    axs: Optional[Sequence[Axes]] = None,
+) -> None:
+    """Internal implementation for coefficient comparison heatmaps."""
+    n_cols = len(coefficients_est)
+
+    # helps boost the color of small coefficients.  Maybe log is better?
+    all_vals = np.hstack((coefficients_est.flatten(), coefficients_true.flatten()))
+    nzs = all_vals[all_vals.nonzero()]
+    max_val = np.max(np.abs(nzs), initial=0.0)
+    min_val = np.min(np.abs(nzs), initial=np.inf)
+    if scaling and np.isfinite(min_val) and max_val / min_val > 10:
+        pwr_ratio = 1.0 / np.log10(max_val / min_val)
+    else:
+        pwr_ratio = 1
+
+    def signed_root(x):
+        return np.sign(x) * np.power(np.abs(x), pwr_ratio)
+
+    with sns.axes_style(style="white", rc={"axes.facecolor": (0, 0, 0, 0)}):
+        if axs is None:
+            fig, axs = plt.subplots(
+                1, 2, figsize=(1.9 * n_cols, 8), sharey=True, sharex=True
+            )
+            fig.tight_layout()
+
+        vmax = signed_root(max_val)
+
+        plot_coefficients(
+            signed_root(coefficients_true),
+            input_features=input_features,
+            feature_names=feature_names,
+            ax=axs[0],
+            cbar=False,
+            vmax=vmax,
+            vmin=-vmax,
+        )
+
+        plot_coefficients(
+            signed_root(coefficients_est),
+            input_features=input_features,
+            feature_names=feature_names,
+            ax=axs[1],
+            cbar=False,
+            vmax=vmax,
+            vmin=-vmax,
+        )
+
+        axs[0].set_title("True Coefficients", rotation=45)
+        axs[1].set_title("Est. Coefficients", rotation=45)
+
+
+def compare_coefficient_plots(
+    coefficients_est: Annotated[np.ndarray, "(n_coord, n_feat)"],
+    coefficients_true: Annotated[np.ndarray, "(n_coord, n_feat)"],
+    input_features: Sequence[str],
+    feature_names: Sequence[str],
+    scaling: bool = True,
+    axs: Optional[Sequence[Axes]] = None,
+) -> None:
+    """Create plots of true and estimated coefficients.
+
+    Deprecated:
+        Use :func:`compare_coefficient_plots_from_dicts` with coefficient
+        dictionaries instead. This function will be removed in a future
+        release.
+    """
+
+    warn(
+        "compare_coefficient_plots is deprecated; use "
+        "compare_coefficient_plots_from_dicts with coefficient dictionaries instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+
+    _compare_coefficient_plots_impl(
+        coefficients_est,
+        coefficients_true,
+        input_features=input_features,
+        feature_names=feature_names,
+        scaling=scaling,
+        axs=axs,
+    )
+
+
+def compare_coefficient_plots_from_dicts(
+    coefficients_est: Sequence[dict[sp.Expr, float]],
+    coefficients_true: Sequence[dict[sp.Expr, float]],
+    input_features: Sequence[str],
+    feature_names: Sequence[str] | None = None,
+    scaling: bool = True,
+    axs: Optional[Sequence[Axes]] = None,
+):
+    """Wrapper to compare coefficients given as dictionaries.
+
+    Converts aligned coefficient dictionaries into dense matrices and then
+    delegates to :func:`compare_coefficient_plots` for plotting.
+
+    This assumes that the coefficient dictionaries are aligned, i.e., that they
+    contain the same keys across all coordinates, as produced by
+    ``unionize_coeff_dicts()``.
+    """
+
+    true_mat, inferred_feature_names = _coeff_dicts_to_matrix(coefficients_true)
+    est_mat, est_feature_names = _coeff_dicts_to_matrix(coefficients_est)
+
+    if true_mat.shape != est_mat.shape:
+        raise ValueError("True and estimated coefficient shapes do not match")
+
+    if inferred_feature_names != est_feature_names:
+        raise ValueError(
+            "Feature names inferred from true and estimated coefficients do not match"
+        )
+
+    _compare_coefficient_plots_impl(
+        est_mat,
+        true_mat,
+        input_features=input_features,
+        feature_names=inferred_feature_names,
+        scaling=scaling,
+        axs=axs,
+    )
+
+
+def _plot_training_trajectory(
+    ax: Axes,
+    x_train: np.ndarray,
+    x_true: np.ndarray,
+    x_smooth: np.ndarray | None,
+    labels: bool = True,
+) -> None:
+    """Plot a single training trajectory
+
+    If x_smooth is provided, it is only plotted if sufficiently different
+    from x_train.
+    """
+    if x_train.shape[1] == 2:
+        ax.plot(
+            x_true[:, 0], x_true[:, 1], ".", label="True", color=COLOR.TRUE, **PLOT_KWS
+        )
+        ax.plot(
+            x_train[:, 0],
+            x_train[:, 1],
+            ".",
+            label="Measured",
+            color=COLOR.MEAS,
+            **PLOT_KWS,
+        )
+        if (
+            x_smooth is not None
+            and np.linalg.norm(x_smooth - x_train) / np.linalg.norm(x_train) > 1e-12
+        ):
+            ax.plot(
+                x_smooth[:, 0],
+                x_smooth[:, 1],
+                ".",
+                label="Smoothed",
+                color=COLOR.EST,
+                **PLOT_KWS,
+            )
+        if labels:
+            ax.set(xlabel="$x_0$", ylabel="$x_1$")
+        else:
+            ax.set(xticks=[], yticks=[])
+    elif x_train.shape[1] == 3:
+        ax.plot(
+            x_true[:, 0],
+            x_true[:, 1],
+            x_true[:, 2],
+            color=COLOR.TRUE,
+            label="True values",
+            **PLOT_KWS,
+        )
+
+        ax.plot(
+            x_train[:, 0],
+            x_train[:, 1],
+            x_train[:, 2],
+            ".",
+            color=COLOR.MEAS,
+            label="Measured values",
+            alpha=0.3,
+        )
+        if (
+            x_smooth is not None
+            and np.linalg.norm(x_smooth - x_train) / x_smooth.size > 1e-12
+        ):
+            ax.plot(
+                x_smooth[:, 0],
+                x_smooth[:, 1],
+                x_smooth[:, 2],
+                ".",
+                color=COLOR.EST,
+                label="Smoothed values",
+                alpha=0.3,
+            )
+        if labels:
+            ax.set(xlabel="$x$", ylabel="$y$", zlabel="$z$")
+        else:
+            ax.set(xticks=[], yticks=[], zticks=[])
+    else:
+        raise ValueError("Can only plot 2d or 3d data.")
+
+
+def plot_training_data(
+    t_train: np.ndarray,
+    x_train: np.ndarray,
+    x_true: np.ndarray,
+    x_smooth: np.ndarray | None = None,
+    coord_names: Optional[Sequence[str]] = None,
+) -> tuple[Figure, Figure]:
+    """Plot training data (and smoothed training data, if different)."""
+    if coord_names is None:
+        coord_names = [f"$x_{i}$" for i in range(x_true.shape[1])]
+
+    fig_composite = plt.figure(figsize=(12, 6))
+    if x_train.shape[-1] == 2:
+        ax_traj = fig_composite.add_subplot(1, 2, 1)
+    elif x_train.shape[-1] == 3:
+        ax_traj = fig_composite.add_subplot(1, 2, 1, projection="3d")
+    else:
+        raise ValueError("Too many or too few coordinates to plot")
+    _plot_training_trajectory(ax_traj, x_train, x_true, x_smooth)
+    ax_traj.legend()
+    ax_traj.set(title="Trajectory Plot")
+    ax_psd = fig_composite.add_subplot(1, 2, 2)
+    _plot_data_psd(ax_psd, x_train, coord_names, traj_type="train")
+    _plot_data_psd(ax_psd, x_true, coord_names, traj_type="true")
+    ax_psd.set(title="Absolute Spectral Density")
+
+    n_coord = x_true.shape[-1]
+    fig_by_coord_1d = plt.figure(figsize=(n_coord * 4, 6))
+    for coord_ind, cname in enumerate(coord_names):
+        ax = fig_by_coord_1d.add_subplot(n_coord, 1, coord_ind + 1)
+        _plot_training_1d(ax, coord_ind, t_train, x_train, x_true, x_smooth, cname)
+
+    fig_by_coord_1d.axes[-1].legend()
+
+    return fig_composite, fig_by_coord_1d
+
+
+def _plot_data_psd(
+    ax: Axes,
+    x: np.ndarray,
+    coord_names: Sequence[str],
+    traj_type: str = "train",
+):
+    """Plot the power spectral density of training data."""
+    if traj_type == "train":
+        color = COLOR.MEAS
+    elif traj_type == "sim":
+        color = COLOR.SIM
+    elif traj_type == "true":
+        color = COLOR.TRUE
+    elif traj_type == "smooth":
+        color = COLOR.EST
+    else:
+        raise ValueError(f"Unknown traj_type '{traj_type}'")
+    coord_names = [name + f" {traj_type}" for name in coord_names]
+    for coord, series in zip(coord_names, x.T):
+        ax.loglog(
+            np.abs(scipy.fft.rfft(series)) / np.sqrt(len(series)),
+            color=color,
+            label=coord,
+        )
+    ax.legend()
+    ax.set(xlabel="Wavenumber")
+    ax.set(ylabel="Magnitude")
+
+
+def _plot_training_1d(
+    ax: Axes,
+    coord_ind: int,
+    t_train: np.ndarray,
+    x_train: np.ndarray,
+    x_true: np.ndarray,
+    x_smooth: Optional[np.ndarray],
+    coord_name: str,
+):
+    ax.plot(t_train, x_train[..., coord_ind], ".", color=COLOR.MEAS, label="measured")
+    ax.plot(t_train, x_true[..., coord_ind], "-", color=COLOR.TRUE, label="true")
+    if x_smooth is not None:
+        ax.plot(t_train, x_smooth[..., coord_ind], color=COLOR.EST, label="smoothed")
+    ax.set(xlabel="t", ylabel=coord_name)
+
+
+def _plot_pde_training_data(last_train, last_train_true, smoothed_last_train):
+    """Plot training data (and smoothed training data, if different)."""
+    # 1D:
+    if len(last_train.shape) == 3:
+        fig, axs = plt.subplots(1, 3, figsize=(18, 6))
+        axs[0].imshow(last_train_true, vmin=0, vmax=last_train_true.max())
+        axs[0].set(title="True Data")
+        axs[1].imshow(last_train, vmin=0, vmax=last_train_true.max())
+        axs[1].set(title="Noisy Data")
+        axs[2].imshow(smoothed_last_train, vmin=0, vmax=last_train_true.max())
+        axs[2].set(title="Smoothed Data")
+        return plt.show()
+
+
+def _plot_test_sim_data_1d_panel(
+    axs: Sequence[Axes],
+    x_true: Optional[np.ndarray],
+    x_sim: np.ndarray,
+    t_test: np.ndarray,
+    t_sim: np.ndarray,
+    coord_names: Sequence[str],
+) -> None:
+    for ordinate, ax in enumerate(axs):
+        if x_true is not None:
+            ax.plot(t_test, x_true[:, ordinate], color=COLOR.TRUE, label="True")
+        axs[ordinate].plot(
+            t_sim, x_sim[:, ordinate], "--", color=COLOR.SIM, label="Simulation"
+        )
+        axs[ordinate].set(xlabel="t", ylabel=coord_names[ordinate])
+
+
+def _plot_test_sim_data_2d(
+    ax: Axes,
+    x_true: Optional[np.ndarray],
+    x_sim: np.ndarray,
+    labels: bool,
+    coord_names: Sequence[str],
+) -> None:
+    if x_true is not None:
+        ax.plot(x_true[:, 0], x_true[:, 1], color=COLOR.TRUE, label="True Values")
+    ax.plot(x_sim[:, 0], x_sim[:, 1], "--", color=COLOR.SIM, label="Simulation")
+    if labels:
+        ax.set(xlabel=coord_names[0], ylabel=coord_names[1])
+    else:
+        ax.set(xticks=[], yticks=[])
+
+
+def _plot_test_sim_data_3d(
+    ax: Axes, x_vals: np.ndarray, label: Optional[str], coord_names: Sequence[str]
+):
+    if label == "True":
+        color = COLOR.TRUE
+    elif label == "Simulation":
+        color = COLOR.SIM
+    else:
+        color = None
+    ax.plot(x_vals[:, 0], x_vals[:, 1], x_vals[:, 2], color=color, label=label)
+    if label:
+        ax.set(xlabel=coord_names[0], ylabel=coord_names[1], zlabel=coord_names[2])
+    else:
+        ax.set(xticks=[], yticks=[], zticks=[])
+
+
+def plot_test_trajectory(
+    x_true: np.ndarray,
+    x_sim: np.ndarray,
+    t_test: np.ndarray,
+    t_sim: np.ndarray,
+    figs: Optional[tuple[Figure, Figure]] = None,
+    coord_names: Optional[Sequence[str]] = None,
+) -> tuple[Figure, Figure]:
+    """Plot a test trajectory
+
+    Args:
+        last_test: a single trajectory of the system
+        model: a trained model to simulate and compare to test data
+        dt: the time interval in test data
+
+    Returns:
+        The sequence of axes used for the single-dimension time-series plots.
+        If ``axs`` is provided, the same sequence is returned.
+    """
+    if coord_names is None:
+        coord_names = [f"$x_{i}$" for i in range(x_true.shape[1])]
+    if not figs:
+        fig_by_coord_1d, axs_by_coord = plt.subplots(
+            x_true.shape[1], 1, sharex=True, figsize=(7, 9)
+        )
+        if x_true.shape[1] == 2:
+            fig_composite, axs_composite = plt.subplots(1, 2, figsize=(10, 4.5))
+        elif x_true.shape[1] == 3:
+            fig_composite, axs_composite = plt.subplots(
+                1, 2, figsize=(10, 4.5), subplot_kw={"projection": "3d"}
+            )
+        else:
+            raise ValueError("Can only plot 2d or 3d data.")
+    else:
+        fig_composite, fig_by_coord_1d = figs
+        axs_composite = fig_composite.axes
+        axs_by_coord = fig_by_coord_1d.axes
+
+    assert isinstance(axs_composite, list)
+    assert isinstance(axs_by_coord, list)
+    _plot_test_sim_data_1d_panel(axs_by_coord, None, x_sim, t_test, t_sim, coord_names)
+    axs_by_coord[-1].legend()
+    if x_true.shape[1] == 2:
+        _plot_test_sim_data_2d(
+            axs_composite[0], None, x_sim, labels=True, coord_names=coord_names
+        )
+    elif x_true.shape[1] == 3:
+        _plot_test_sim_data_3d(axs_composite[0], x_sim, "Simulation", coord_names)
+    axs_composite[0].legend()
+    _plot_data_psd(axs_composite[1], x_sim, coord_names, traj_type="sim")
+    if not figs:
+        fig_by_coord_1d.suptitle("Test Trajectories by Dimension")
+        fig_composite.suptitle("Full Test Trajectories")
+        axs_composite[0].set(title="true trajectory")
+        axs_composite[0].set(title="model simulation")
+    return fig_composite, fig_by_coord_1d

sindy_exp/_typing.py

@@ -0,0 +1,158 @@
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import (
+    Any,
+    Callable,
+    Literal,
+    NamedTuple,
+    Optional,
+    Protocol,
+    TypeVar,
+    overload,
+)
+
+import numpy as np
+import sympy as sp
+from numpy.typing import NBitBase
+from sympy import Expr
+from typing_extensions import Self
+
+NpFlt = np.dtype[np.floating]
+Float1D = np.ndarray[tuple[int], NpFlt]
+Float2D = np.ndarray[tuple[int, int], NpFlt]
+Shape = TypeVar("Shape", bound=tuple[int, ...])
+FloatND = np.ndarray[Shape, np.dtype[np.floating[NBitBase]]]
+
+
+TrajectoryType = TypeVar("TrajectoryType", list[np.ndarray], np.ndarray)
+
+
+class _BaseSINDy(Protocol):
+    optimizer: Any
+    feature_library: Any
+    feature_names: list[str]
+
+    def fit(self, x: TrajectoryType, t: TrajectoryType, *args, **kwargs) -> Self: ...
+
+    def simulate(self, x0: np.ndarray, t: np.ndarray, **kwargs) -> np.ndarray: ...
+
+    def score(
+        self,
+        x: TrajectoryType,
+        t: TrajectoryType,
+        x_dot: TrajectoryType,
+        metric: Callable,
+    ) -> float: ...
+
+    def predict(self, x: np.ndarray, u: None | np.ndarray = None) -> np.ndarray: ...
+
+    def coefficients(self): ...
+
+    @overload
+    def equations(self) -> list[str]: ...
+
+    @overload
+    def equations(self, precision: int) -> list[str]: ...
+
+    @overload
+    def equations(self, precision: int, fmt: Literal["str"] | None) -> list[str]: ...
+
+    @overload
+    def equations(
+        self, precision: int, fmt: Literal["sympy"]
+    ) -> list[dict[Expr, float]]: ...
+
+    def print(self, precision: int, **kwargs) -> None: ...
+
+    def get_feature_names(self) -> list[str]: ...
+
+
+class ProbData(NamedTuple):
+    """Data bundle for a single trajectory.
+
+    Represents a trajectory's training data and associated metadata.
+    """
+
+    dt: float
+    t_train: Float1D
+    x_train: Float2D
+    x_train_true: Float2D
+    x_train_true_dot: Float2D
+    input_features: list[str]
+    integrator: Optional[Any] = None  # diffrax.Solution
+
+
+class NestedDict(defaultdict):
+    """A dictionary that splits all keys by ".", creating a sub-dict.
+
+    Args: see superclass
+
+    Example:
+
+        >>> foo = NestedDict("a.b"=1)
+        >>> foo["a.c"] = 2
+        >>> foo["a"]["b"]
+        1
+    """
+
+    def __missing__(self, key):
+        try:
+            prefix, subkey = key.split(".", 1)
+        except ValueError:
+            raise KeyError(key)
+        return self[prefix][subkey]
+
+    def __setitem__(self, key, value):
+        if "." in key:
+            prefix, suffix = key.split(".", 1)
+            if self.get(prefix) is None:
+                self[prefix] = NestedDict()
+            return self[prefix].__setitem__(suffix, value)
+        else:
+            return super().__setitem__(key, value)
+
+    def update(self, other: dict):  # type: ignore
+        try:
+            for k, v in other.items():
+                self.__setitem__(k, v)
+        except:  # noqa: E722
+            super().update(other)
+
+    def flatten(self):
+        """Flattens a nested dictionary without mutating.  Returns new dict"""
+
+        def _flatten(nested_d: dict) -> dict:
+            new = {}
+            for key, value in nested_d.items():
+                if not isinstance(key, str):
+                    raise TypeError("Only string keys allowed in flattening")
+                if not isinstance(value, dict):
+                    new[key] = value
+                    continue
+                for sub_key, sub_value in _flatten(value).items():
+                    new[key + "." + sub_key] = sub_value
+            return new
+
+        return _flatten(self)
+
+
+@dataclass
+class DynamicsTrialData:
+    trajectories: list[ProbData]
+    true_equations: list[dict[sp.Expr, float]]
+    sindy_equations: list[dict[sp.Expr, float]]
+    model: _BaseSINDy
+    input_features: list[str]
+    smooth_train: list[np.ndarray]
+
+
+@dataclass
+class SINDyTrialUpdate:
+    t_sim: Float1D
+    t_test: Float1D
+    x_sim: FloatND
+
+
+@dataclass
+class FullDynamicsTrialData(DynamicsTrialData):
+    sims: list[SINDyTrialUpdate]