bfflearn 0.1.0.dev1__py3-none-any.whl

bff/plotting.py

@@ -0,0 +1,362 @@
+from pathlib import Path
+from typing import Any, Mapping, Optional, Sequence, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.colors import ListedColormap
+from matplotlib.ticker import MaxNLocator
+from scipy.stats import gaussian_kde
+
+from .bayes.results import InferenceResults
+from .domain.specs import Specs
+
+PathLike = Union[str, Path]
+ArrayLike = Union[np.ndarray, torch.Tensor]
+
+
+def _wrap_label(text: str, max_per_line: int = 4) -> str:
+    words = text.split()
+    return "\n".join(
+        " ".join(words[i:i + max_per_line])
+        for i in range(0, len(words), max_per_line)
+    )
+
+
+def _coerce_specs(specs: Specs | PathLike) -> Specs:
+    return specs if isinstance(specs, Specs) else Specs(specs)
+
+
+def _coerce_samples(
+    samples: InferenceResults | ArrayLike,
+) -> np.ndarray:
+    if isinstance(samples, InferenceResults):
+        return np.asarray(samples.prepared_samples, dtype=float)
+    if isinstance(samples, torch.Tensor):
+        return samples.detach().cpu().numpy()
+    return np.asarray(samples, dtype=float)
+
+
+def _parameter_labels(
+    names: Sequence[str],
+    labels: Optional[Sequence[str] | Mapping[str, str]] = None,
+) -> list[str]:
+    if labels is None:
+        return list(names)
+    if isinstance(labels, Mapping):
+        return [labels.get(name, name) for name in names]
+    if len(labels) != len(names):
+        raise ValueError(
+            "parameter_labels must match the number of plotted parameters.")
+    return list(labels)
+
+
+def _expand_short_labels(
+    labels: Sequence[str],
+    full_labels: Sequence[str],
+) -> list[str]:
+    """Expand shortened labels like ``C1`` back to ``charge C1`` when possible."""
+    if len(labels) != len(full_labels):
+        return list(labels)
+
+    lookup: dict[str, str] = {}
+    for full_label in full_labels:
+        if full_label.startswith("$"):
+            lookup.setdefault(full_label, full_label)
+            continue
+        lookup.setdefault(full_label, full_label)
+        lookup.setdefault(full_label.split(maxsplit=1)[-1], full_label)
+
+    expanded: list[str] = []
+    changed = False
+    for label in labels:
+        replacement = lookup.get(label, label)
+        expanded.append(replacement)
+        changed |= replacement != label
+
+    return expanded if changed else list(labels)
+
+
+def _axis_labels(kind: str) -> tuple[str, str]:
+    if kind == "charge":
+        return "Atom", "Charge [e]"
+    if kind == "sigma":
+        return "Atom type", "$\\sigma$ [nm]"
+    return kind.capitalize(), kind.capitalize()
+
+
+def plot_marginals(
+    results: InferenceResults,
+    specs: Specs | PathLike,
+    *,
+    parameter_labels: Optional[Sequence[str] | Mapping[str, str]] = None,
+    color_prior: str = "gray",
+    color_posterior: str = "tab:red",
+    fn_out: Optional[PathLike] = None,
+) -> None:
+    specs = _coerce_specs(specs)
+    posterior = specs.with_implicit_charge(results.prepared_samples)
+    param_names = specs.bounds.names.tolist()
+    tick_labels = _parameter_labels(param_names, parameter_labels)
+
+    param_groups: dict[str, list[int]] = {}
+    for idx, name in enumerate(param_names):
+        kind = name.split()[0]
+        param_groups.setdefault(kind, []).append(idx)
+
+    fig, axes = plt.subplots(
+        1,
+        len(param_groups),
+        figsize=(4 * len(param_groups), 3.2),
+        gridspec_kw={"wspace": 0.3},
+    )
+    axes = np.atleast_1d(axes)
+
+    explicit_names = specs.bounds.without(specs.implicit_param).names.tolist()
+    prior_index = {name: i for i, name in enumerate(explicit_names)}
+    show_prior = results.priors is not None
+    legend_used = {"prior": False, "posterior": False, "bounds": False}
+
+    for ax, (kind, indices) in zip(axes, param_groups.items()):
+        bounds_block = np.asarray(
+            [specs.bounds.by_name[param_names[i]] for i in indices],
+            dtype=float,
+        )
+        y_min = bounds_block[:, 0].min()
+        y_max = bounds_block[:, 1].max()
+        y_pad = max(0.05, 0.18 * (y_max - y_min))
+        label_y = y_min - 0.80 * y_pad
+        posterior_peaks: list[float] = []
+        prior_peaks: list[float] = []
+        curves: dict[
+            int, tuple[np.ndarray, np.ndarray, Optional[np.ndarray], float]
+        ] = {}
+
+        for idx in indices:
+            name = param_names[idx]
+            lower, upper = specs.bounds.by_name[name]
+            y = np.linspace(lower - y_pad, upper + y_pad, 400)
+
+            prior_density = None
+            if show_prior and name in prior_index:
+                prior = results.priors.distributions[prior_index[name]]
+                prior_density = (
+                    prior.log_prob(torch.as_tensor(y, dtype=torch.float32))
+                    .exp()
+                    .detach()
+                    .cpu()
+                    .numpy()
+                )
+                prior_peaks.append(float(np.max(prior_density)))
+
+            posterior_density = gaussian_kde(posterior[:, idx])(y)
+            posterior_peaks.append(float(np.max(posterior_density)))
+            mode = float(y[np.argmax(posterior_density)])
+            curves[idx] = (y, posterior_density, prior_density, mode)
+
+        max_posterior_peak = max(posterior_peaks, default=1.0)
+        max_prior_peak = max(prior_peaks, default=max_posterior_peak)
+        posterior_width = 1.2
+        prior_width = 0.7
+        posterior_scale = posterior_width / max(max_posterior_peak, 1e-12)
+        prior_scale = prior_width / max(max_prior_peak, 1e-12)
+
+        for xpos, idx in enumerate(indices):
+            name = param_names[idx]
+            lower, upper = specs.bounds.by_name[name]
+            y, posterior_density, prior_density, mode = curves[idx]
+
+            if prior_density is not None:
+                ax.fill_betweenx(
+                    y,
+                    xpos - prior_scale * prior_density,
+                    xpos,
+                    color=color_prior,
+                    lw=0,
+                    label="prior" if not legend_used["prior"] else None,
+                )
+                legend_used["prior"] = True
+
+            ax.fill_betweenx(
+                y,
+                xpos,
+                xpos + posterior_scale * posterior_density,
+                color=color_posterior,
+                lw=0,
+                label="posterior" if not legend_used["posterior"] else None,
+            )
+            legend_used["posterior"] = True
+
+            center = 0.5 * (lower + upper)
+            yerr = np.array([[center - lower], [upper - center]])
+            ax.errorbar(
+                [xpos],
+                [center],
+                yerr=yerr,
+                lw=2,
+                ls="",
+                capsize=4,
+                capthick=2,
+                markeredgewidth=2,
+                color="k",
+                label="bounds" if not legend_used["bounds"] else None,
+            )
+            legend_used["bounds"] = True
+
+            ax.text(
+                xpos,
+                label_y,
+                f"{mode:.3f}",
+                color=color_posterior,
+                fontweight="bold",
+                ha="center",
+                va="top",
+            )
+
+        ax.set_xlim(-prior_width - 0.25, len(indices) - 1 + posterior_width + 0.25)
+        ax.set_ylim(label_y - 0.6 * y_pad, y_max + y_pad)
+        ax.set_xticks(range(len(indices)))
+        ax.set_xticklabels(
+            [_wrap_label(tick_labels[i]) for i in indices],
+            rotation=30,
+            ha="center",
+        )
+        xlabel, ylabel = _axis_labels(kind)
+        ax.set_xlabel(xlabel)
+        ax.set_ylabel(ylabel)
+        ax.tick_params(direction="in")
+
+    if len(axes):
+        axes[0].legend(
+            loc="upper center",
+            bbox_to_anchor=(0.5, 1.18),
+            ncol=3,
+            frameon=False,
+        )
+
+    if fn_out is not None:
+        plt.savefig(fn_out, bbox_inches="tight")
+        plt.close(fig)
+    else:
+        plt.show()
+
+
+def plot_corner(
+    samples: InferenceResults | ArrayLike,
+    labels: Optional[Sequence[str]] = None,
+    *,
+    quantiles: Sequence[float] = (0.16, 0.5, 0.84),
+    figsize: float = 2.6,
+    cmap: Any = "Reds",
+    levels: int = 5,
+    scatter_alpha: float = 0.15,
+    fn_out: Optional[PathLike] = None,
+) -> None:
+    sample_source = samples
+    samples = _coerce_samples(samples)
+    if samples.ndim != 2:
+        raise ValueError("plot_corner expects samples with shape (n_samples, n_dim).")
+
+    n_dim = samples.shape[1]
+    if labels is None:
+        if isinstance(sample_source, InferenceResults):
+            labels = list(sample_source.labels)
+        else:
+            labels = [f"theta_{i}" for i in range(n_dim)]
+    elif len(labels) != n_dim:
+        raise ValueError("labels must match the posterior sample dimension.")
+    elif isinstance(sample_source, InferenceResults):
+        labels = _expand_short_labels(labels, sample_source.labels)
+
+    labels = [_wrap_label(label) for label in labels]
+    base_cmap = plt.get_cmap(cmap)
+    colors = base_cmap(np.linspace(0, 1, max(levels, 2)))
+    colors[0] = np.array([1.0, 1.0, 1.0, 0.0])
+    contour_cmap = ListedColormap(colors)
+    fig, axes = plt.subplots(
+        n_dim,
+        n_dim,
+        figsize=(figsize * n_dim, figsize * n_dim),
+        gridspec_kw={"wspace": 0.05, "hspace": 0.05},
+    )
+
+    limits = [(samples[:, i].min(), samples[:, i].max()) for i in range(n_dim)]
+
+    for i in range(n_dim):
+        for j in range(n_dim):
+            ax = axes[i, j]
+            if i < j:
+                ax.axis("off")
+                continue
+
+            if i == j:
+                x = np.linspace(*limits[i], 400)
+                kde = gaussian_kde(samples[:, i])
+                density = kde(x)
+                ax.plot(x, density, color="k", lw=2.5)
+                ax.fill_between(x, 0, density, color="0.75", alpha=0.7)
+                if quantiles:
+                    q_values = np.quantile(samples[:, i], quantiles)
+                    for q in q_values:
+                        ax.axvline(q, color="k", ls="--", lw=1.3)
+                    if len(q_values) == 3:
+                        median = q_values[1]
+                        lower = median - q_values[0]
+                        upper = q_values[2] - median
+                        ax.set_title(
+                            (
+                                f"{labels[i]}\n"
+                                f"{median:.3f}\n"
+                                f"(+{upper:.3f} / -{lower:.3f})"
+                            ),
+                            fontsize=12,
+                        )
+                ax.set_xlim(*limits[i])
+                ax.set_yticks([])
+                ax.tick_params(axis="y", left=False, labelleft=False)
+            else:
+                x = samples[::10, j]
+                y = samples[::10, i]
+                ax.scatter(
+                    x, y, s=5, lw=0, alpha=scatter_alpha, color="k", rasterized=True
+                )
+                try:
+                    kde = gaussian_kde(np.vstack([x, y]))
+                    xi, yi = np.mgrid[
+                        limits[j][0]:limits[j][1]:100j,
+                        limits[i][0]:limits[i][1]:100j,
+                    ]
+                    zi = kde(np.vstack([xi.ravel(), yi.ravel()])).reshape(xi.shape)
+                    contour_levels = np.linspace(zi.min(), zi.max(), levels)
+                    ax.contourf(
+                        xi, yi, zi, levels=contour_levels, cmap=contour_cmap
+                    )
+                    ax.contour(
+                        xi, yi, zi, levels=contour_levels, colors="k", linewidths=0.8
+                    )
+                except np.linalg.LinAlgError:
+                    pass
+                ax.set_xlim(*limits[j])
+                ax.set_ylim(*limits[i])
+
+            ax.xaxis.set_major_locator(MaxNLocator(nbins=4))
+            ax.yaxis.set_major_locator(MaxNLocator(nbins=4))
+            ax.tick_params(direction="in", top=True, right=True, labelsize=11)
+
+            if i == n_dim - 1:
+                ax.set_xlabel(labels[j], fontsize=12)
+            else:
+                ax.set_xticklabels([])
+
+            if j == 0 and i > 0:
+                ax.set_ylabel(labels[i], fontsize=12)
+            elif i != j:
+                ax.set_yticklabels([])
+
+    fig.align_labels()
+
+    if fn_out is not None:
+        plt.savefig(fn_out, bbox_inches="tight")
+        plt.close(fig)
+    else:
+        plt.show()

bff/qoi/__init__.py

@@ -0,0 +1,9 @@
+from . import analysis, data, hbonds, rdf, routines
+
+__all__ = [
+    "data",
+    "hbonds",
+    "rdf",
+    "routines",
+    "analysis",
+]

bff/qoi/analysis.py

@@ -0,0 +1,209 @@
+import gc
+import multiprocessing as mp
+import warnings
+from functools import partial
+from pathlib import Path
+from typing import Any, Iterable, Sequence
+
+import MDAnalysis as mda
+
+from ..domain.trainset import TrajectorySet
+from ..io.logs import Logger
+from ..io.progress import iter_progress
+from ..tools import _normalized_dimensions
+from ..topology import prepare_universe
+from .data import QoI
+from .routines import RuntimeRoutine, run_analysis_routines
+
+PathLike = str | Path
+
+
+warnings.filterwarnings(
+    "ignore", category=DeprecationWarning, module="MDAnalysis.coordinates.DCD"
+)
+warnings.filterwarnings(
+    "ignore", category=UserWarning, module="MDAnalysis.coordinates.XDR"
+)
+
+
+def _cleanup_universe(universe: mda.Universe | None) -> None:
+    """Release trajectory resources after one analysis task."""
+    if universe is not None:
+        trajectory = getattr(universe, "trajectory", None)
+        if trajectory is not None:
+            close = getattr(trajectory, "close", None)
+            if callable(close):
+                close()
+        raw_trajectory = getattr(universe, "_trajectory", None)
+        if raw_trajectory is not None and raw_trajectory is not trajectory:
+            close = getattr(raw_trajectory, "close", None)
+            if callable(close):
+                close()
+
+
+def _prepare_universe(
+    fn_topol: PathLike,
+    fn_coord: PathLike,
+    fn_trj: PathLike,
+    *,
+    start: int,
+    stop: int | None,
+    step: int,
+    in_memory: bool,
+) -> mda.Universe:
+    """Prepare one MDAnalysis universe for trajectory analysis."""
+    universe = prepare_universe(str(fn_topol), str(fn_coord), dt=1)
+    default_dimensions = _normalized_dimensions(universe.dimensions)
+    universe.load_new(str(fn_trj))
+    universe._bff_default_dimensions = default_dimensions
+    if in_memory:
+        universe.transfer_to_memory(start=start, stop=stop, step=step)
+        if default_dimensions is not None:
+            for ts in universe.trajectory:
+                if ts.dimensions is None:
+                    ts.dimensions = default_dimensions
+    return universe
+
+
+def analyze_trajectory_set(
+    trajectory_set: TrajectorySet,
+    *,
+    routines_by_system: Sequence[tuple[RuntimeRoutine, ...]],
+    start: int = 0,
+    stop: int | None = None,
+    step: int = 1,
+    in_memory: bool = False,
+) -> list[dict[str, QoI]]:
+    """Analyze all trajectories that belong to one sample or reference set."""
+    if len(routines_by_system) != len(trajectory_set.fn_trj):
+        raise ValueError(
+            "Analysis routine count must match the number of trajectories in the set."
+        )
+
+    results: list[dict[str, QoI]] = []
+    routine_start = 0 if in_memory else start
+    routine_stop = None if in_memory else stop
+    routine_step = 1 if in_memory else step
+    for fn_topol, fn_coord, fn_trj, routines in zip(
+        trajectory_set.fn_topol,
+        trajectory_set.fn_coord,
+        trajectory_set.fn_trj,
+        routines_by_system,
+    ):
+        universe = None
+        try:
+            universe = _prepare_universe(
+                fn_topol,
+                fn_coord,
+                fn_trj,
+                start=start,
+                stop=stop,
+                step=step,
+                in_memory=in_memory,
+            )
+            result = run_analysis_routines(
+                routines,
+                universe=universe,
+                start=routine_start,
+                stop=routine_stop,
+                step=routine_step,
+            )
+            if not result:
+                raise ValueError("Analysis routine returned no QoI outputs.")
+            results.append(result)
+        finally:
+            _cleanup_universe(universe)
+    return results
+
+
+def _maybe_collect_garbage(
+    completed: int,
+    total: int,
+    *,
+    gc_collect: bool,
+    progress_stride: int,
+) -> None:
+    """Run the garbage collector at coarse progress intervals when requested."""
+    if not gc_collect:
+        return
+    if (completed % progress_stride == 0) or (completed == total):
+        gc.collect()
+
+
+def _iter_analyzed_sets(
+    trajectory_sets: Sequence[TrajectorySet],
+    *,
+    analyze_one: Any,
+    workers: int,
+    maxtasksperchild: int,
+) -> Iterable[list[dict[str, QoI]]]:
+    """Yield analyzed trajectory sets from the serial or multiprocessing path."""
+    if workers <= 1:
+        yield from (analyze_one(trajectory_set) for trajectory_set in trajectory_sets)
+        return
+
+    context = mp.get_context()
+    with context.Pool(workers, maxtasksperchild=maxtasksperchild) as pool:
+        yield from pool.imap(
+            analyze_one,
+            trajectory_sets,
+            chunksize=1,
+        )
+
+
+def analyze_trajectory_sets(
+    trajectory_sets: Sequence[TrajectorySet],
+    *,
+    routines_by_system: Sequence[tuple[RuntimeRoutine, ...]],
+    start: int = 0,
+    stop: int | None = None,
+    step: int = 1,
+    workers: int = 1,
+    progress_stride: int = 10,
+    progress_label: str = "Trajectory QoI",
+    logger: Logger | None = None,
+    in_memory: bool = False,
+    gc_collect: bool = False,
+    maxtasksperchild: int = 100,
+) -> list[list[dict[str, QoI]]]:
+    """Analyze multiple trajectory sets with one shared routine setup."""
+    logger = logger or Logger(progress_label)
+    n_sets = len(trajectory_sets)
+    if n_sets == 0:
+        return []
+
+    analyze_one = partial(
+        analyze_trajectory_set,
+        routines_by_system=routines_by_system,
+        start=start,
+        stop=stop,
+        step=step,
+        in_memory=in_memory,
+    )
+
+    workers = mp.cpu_count() if workers == -1 else workers
+    completed_sets = _iter_analyzed_sets(
+        trajectory_sets,
+        analyze_one=analyze_one,
+        workers=workers,
+        maxtasksperchild=maxtasksperchild,
+    )
+    qoi = []
+    for completed, result in enumerate(
+        iter_progress(
+            completed_sets,
+            total=n_sets,
+            stride=progress_stride,
+            logger=logger,
+            label=progress_label,
+        ),
+        start=1,
+    ):
+        qoi.append(result)
+        _maybe_collect_garbage(
+            completed,
+            n_sets,
+            gc_collect=gc_collect,
+            progress_stride=progress_stride,
+        )
+    return qoi