chemparseplot 1.9.17__tar.gz → 1.10.0__tar.gz

{chemparseplot-1.9.17 → chemparseplot-1.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chemparseplot
-Version: 1.9.17
+Version: 1.10.0
 Summary: Parsers and plotting tools for computational chemistry
 Project-URL: Documentation, https://chemparseplot.rgoswami.me
 Project-URL: Issues, https://github.com/HaoZeke/chemparseplot/issues

{chemparseplot-1.9.17 → chemparseplot-1.10.0}/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '1.9.17'
-__version_tuple__ = version_tuple = (1, 9, 17)
+__version__ = version = '1.10.0'
+__version_tuple__ = version_tuple = (1, 10, 0)
 
 __commit_id__ = commit_id = None

{chemparseplot-1.9.17 → chemparseplot-1.10.0}/chemparseplot/plot/__init__.py

@@ -25,7 +25,14 @@ __all__ = [
 def __getattr__(name):
     import importlib
 
-    lazy_submodules = {"geomscan", "structs", "chemgp", "optimization"}
+    lazy_submodules = {
+        "geomscan",
+        "structs",
+        "chemgp",
+        "optimization",
+        "landscape",
+        "disconnectivity",
+    }
     if name in lazy_submodules:
         return importlib.import_module(f".{name}", __name__)
     if name == "ureg":

chemparseplot-1.10.0/chemparseplot/plot/disconnectivity.py

@@ -0,0 +1,323 @@
+# SPDX-FileCopyrightText: 2023-present Rohit Goswami <rog32@hi.is>
+#
+# SPDX-License-Identifier: MIT
+"""Approximate disconnectivity graphs from minima databases.
+
+The classic disconnectivity graph needs transition states; a hopping
+campaign records only quenched minima. These helpers draw the superbasin
+merge tree of a descriptor-space k-nearest-neighbour graph instead: at each
+energy level the minima below it split into connected components, and
+components merge as the level rises. That approximation is a rendering of
+the database made for reading the figure; it should never feed a reported
+statistic.
+
+Basins can be lumped for display by spectral clustering on the same graph.
+Everything is iterative: with a thousand minima the merge chains grow
+deeper than any recursion limit.
+
+```{versionadded} 1.10.0
+```
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy.sparse import coo_matrix
+from scipy.sparse.csgraph import connected_components, laplacian
+
+from chemparseplot.plot.theme import RUHI_COLORS
+
+BASIN_PALETTE = [
+    RUHI_COLORS[c] for c in ["teal", "sky", "magenta", "coral", "sunshine"]
+]
+
+
+def dedup(matrix, energies, tol):
+    """Merge near-identical minima, keeping the lowest-energy representative.
+
+    Returns the kept row indices, lowest energy first.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    order = np.argsort(energies)
+    rep_idx = []
+    for i in order:
+        if any(np.linalg.norm(matrix[i] - matrix[j]) < tol for j in rep_idx):
+            continue
+        rep_idx.append(i)
+    return np.array(rep_idx)
+
+
+def knn_graph(matrix, k=6):
+    """Symmetric k-nearest-neighbour adjacency in descriptor space.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    d = np.linalg.norm(matrix[:, None, :] - matrix[None, :, :], axis=-1)
+    np.fill_diagonal(d, np.inf)
+    n = len(matrix)
+    rows, cols = [], []
+    for i in range(n):
+        for j in np.argsort(d[i])[:k]:
+            rows += [i, j]
+            cols += [j, i]
+    data = np.ones(len(rows))
+    return coo_matrix((data, (rows, cols)), shape=(n, n)).tocsr()
+
+
+def bridge_components(adjacency, matrix, weight=0.1):
+    """Weakly connect graph components through their closest point pairs.
+
+    Spectral clustering assumes a connected graph; a k-NN graph over a
+    minima database often fragments. One weak edge per component pair
+    (closest points in descriptor space) restores the geometry without
+    drowning the strong intra-basin structure.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    n_comp, labels = connected_components(adjacency, directed=False)
+    if n_comp == 1:
+        return adjacency
+    adjacency = adjacency.tolil()
+    for a in range(n_comp):
+        for b in range(a + 1, n_comp):
+            ia = np.where(labels == a)[0]
+            ib = np.where(labels == b)[0]
+            d = np.linalg.norm(
+                matrix[ia][:, None, :] - matrix[ib][None, :, :], axis=-1
+            )
+            i, j = np.unravel_index(int(np.argmin(d)), d.shape)
+            adjacency[ia[i], ib[j]] = weight
+            adjacency[ib[j], ia[i]] = weight
+    return adjacency.tocsr()
+
+
+def spectral_basins(adjacency, n_basins, rng, matrix=None):
+    """Display-only basin assignment by spectral clustering on the graph.
+
+    With ``matrix`` given, disconnected components are weakly bridged
+    through their closest descriptor pairs first.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    if matrix is not None:
+        adjacency = bridge_components(adjacency, matrix)
+    lap = laplacian(adjacency.astype(float), normed=True)
+    k = min(n_basins, adjacency.shape[0] - 1)
+    # Dense eigendecomposition: minima databases stay around a thousand
+    # nodes, and iterative smallest-magnitude solvers are unreliable on
+    # the near-singular Laplacians of disconnected graphs.
+    _, vecs = np.linalg.eigh(np.asarray(lap.todense()))
+    x = vecs[:, :k]
+    x = x / (np.linalg.norm(x, axis=1, keepdims=True) + 1e-12)
+    # Farthest-point seeding: deterministic given the first pick, and it
+    # cannot drop both centers into one embedded cloud.
+    first = int(rng.integers(len(x)))
+    centers = [x[first]]
+    while len(centers) < k:
+        d = np.min(
+            [np.linalg.norm(x - c, axis=1) for c in centers], axis=0
+        )
+        centers.append(x[int(np.argmax(d))])
+    centers = np.array(centers)
+    for _ in range(60):
+        assign = np.argmin(
+            np.linalg.norm(x[:, None, :] - centers[None, :, :], axis=-1), axis=1
+        )
+        new = np.array(
+            [
+                x[assign == c].mean(axis=0) if (assign == c).any() else centers[c]
+                for c in range(k)
+            ]
+        )
+        if np.allclose(new, centers):
+            break
+        centers = new
+    return assign
+
+
+class Node:
+    """A merge-tree node; leaves carry their minimum's index.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+
+    __slots__ = ("level", "children", "leaf", "x", "emin", "argmin")
+
+    def __init__(self, level, children=None, leaf=None):
+        self.level = level
+        self.children = children or []
+        self.leaf = leaf
+        self.x = None
+        if self.leaf is not None:
+            self.emin = level
+            self.argmin = leaf
+        else:
+            self.emin = min(c.emin for c in self.children)
+            self.argmin = min(self.children, key=lambda c: c.emin).argmin
+
+
+def merge_tree(adjacency, energies, levels):
+    """Superbasin merge tree over an energy-threshold sweep.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    n = len(energies)
+    leaves = [Node(energies[i], leaf=i) for i in range(n)]
+    active = {}
+    prev_sets = []
+    for lev in levels:
+        mask = energies <= lev
+        if not mask.any():
+            continue
+        sub = np.where(mask)[0]
+        _, labels = connected_components(
+            adjacency[sub][:, sub], directed=False
+        )
+        comps = {}
+        for local, g in enumerate(sub):
+            comps.setdefault(labels[local], []).append(g)
+        new_sets = [frozenset(v) for v in comps.values()]
+        for s in new_sets:
+            if s in active:
+                continue
+            parts = [p for p in prev_sets if p <= s and p in active]
+            merged = [active[p] for p in parts]
+            covered = frozenset().union(*parts) if parts else frozenset()
+            for p in parts:
+                del active[p]
+            extra = sorted(s - covered, key=lambda i: energies[i])
+            if not merged:
+                node = leaves[extra[0]]
+                for m in extra[1:]:
+                    node = Node(
+                        max(energies[m], node.level), children=[node, leaves[m]]
+                    )
+            else:
+                node = merged[0] if len(merged) == 1 else Node(lev, children=merged)
+                for m in extra:
+                    node = Node(lev, children=[node, leaves[m]])
+            active[s] = node
+        prev_sets = new_sets
+    roots = list(active.values())
+    if len(roots) > 1:
+        return Node(levels[-1], children=roots)
+    return roots[0]
+
+
+def iter_nodes(root):
+    """Post-order traversal without recursion.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    out, stack = [], [(root, False)]
+    while stack:
+        node, done = stack.pop()
+        if done or node.leaf is not None:
+            out.append(node)
+            continue
+        stack.append((node, True))
+        for c in reversed(node.children):
+            stack.append((c, False))
+    return out
+
+
+def layout(root):
+    """Leaf x positions in tree order, children sorted deepest-first.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    for node in iter_nodes(root):
+        if node.leaf is None:
+            node.children.sort(key=lambda c: c.emin)
+    x = 0.0
+    for node in iter_nodes(root):
+        if node.leaf is not None:
+            node.x = x
+            x += 1.0
+        else:
+            node.x = float(np.mean([c.x for c in node.children]))
+
+
+def paint(root, assign, palette=None):
+    """Colour each subtree by the spectral basin of its deepest minimum.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    palette = palette or BASIN_PALETTE
+    colors = {}
+    for node in iter_nodes(root):
+        colors[id(node)] = palette[int(assign[node.argmin]) % len(palette)]
+    return colors
+
+
+def draw(ax, root, colors, lw=1.1):
+    """Vertical stems to the merge level, horizontal bars at merges.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    for node in iter_nodes(root):
+        if node.leaf is not None:
+            continue
+        for c in node.children:
+            ax.plot(
+                [c.x, c.x],
+                [c.level, node.level],
+                color=colors.get(id(c), "#444444"),
+                lw=lw,
+                solid_capstyle="round",
+                zorder=2,
+            )
+        xs = [c.x for c in node.children]
+        ax.plot(
+            [min(xs), max(xs)],
+            [node.level, node.level],
+            color="#444444",
+            lw=0.9,
+            zorder=1,
+        )
+
+
+def plot_disconnectivity(
+    ax,
+    matrix,
+    energies,
+    *,
+    dedup_tol=0.15,
+    knn=6,
+    n_levels=48,
+    n_basins=5,
+    rng=None,
+):
+    """Full pipeline: dedup, graph, merge tree, spectral paint, draw.
+
+    Returns ``(kept_indices, root)`` so callers can label leaves.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    rng = rng or np.random.default_rng(11)
+    reps = dedup(matrix, energies, dedup_tol)
+    m = matrix[reps]
+    en = energies[reps]
+    adjacency = knn_graph(m, knn)
+    levels = np.linspace(en.min(), en.max(), n_levels)[1:]
+    root = merge_tree(adjacency, en, levels)
+    layout(root)
+    assign = spectral_basins(adjacency, n_basins, rng, matrix=m)
+    colors = paint(root, assign)
+    draw(ax, root, colors)
+    ax.set_xticks([])
+    ax.set_ylabel("quenched energy")
+    ax.spines["bottom"].set_visible(False)
+    return reps, root

chemparseplot-1.10.0/chemparseplot/plot/landscape.py

@@ -0,0 +1,315 @@
+# SPDX-FileCopyrightText: 2023-present Rohit Goswami <rog32@hi.is>
+#
+# SPDX-License-Identifier: MIT
+"""Sketch-map landscape figures over minima databases.
+
+A campaign of quenched minima becomes one figure: a sketch-map plane
+(lab-cosmo ``dimlandmark``/``dimred`` binaries) over a permutation-invariant
+descriptor, a filled-contour energy surface with thin contour lines, per-arm
+scatter overlays, pinned reference markers, and rendered structure insets in
+the figure margins tied to their map points.
+
+The default descriptor is the sorted smoothed-coordination-number vector,
+the discriminant of the sketch-map literature for Lennard-Jones clusters:
+it separates the fcc funnel from the icosahedral one. SOAP power spectra
+(through ``featomic``) remain available for molecular systems where
+coordination alone degenerates.
+
+```{versionadded} 1.10.0
+```
+"""
+
+from __future__ import annotations
+
+import subprocess
+from pathlib import Path
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+
+from chemparseplot.plot.theme import RUHI_COLORS, RUHI_THEME
+
+ARM_MARKERS = ["o", "^", "s", "P", "X"]
+
+SOAP_HYPERS = {
+    "cutoff": {"radius": 4.0, "smoothing": {"type": "ShiftedCosine", "width": 0.5}},
+    "density": {"type": "Gaussian", "width": 0.35},
+    "basis": {
+        "type": "TensorProduct",
+        "max_angular": 4,
+        "radial": {"type": "Gto", "max_radial": 5},
+    },
+}
+
+
+def cn_matrix(structures, symbols, cutoff=1.5, cn_species=None):
+    """Sorted smoothed coordination numbers, permutation invariant.
+
+    A Fermi switching function at the cutoff keeps the vector continuous,
+    so near-degenerate minima do not collapse onto identical rows.
+
+    Parameters
+    ----------
+    structures : sequence of array-like
+        Flat ``3N`` coordinate rows.
+    symbols : list of str
+        One chemical symbol per atom.
+    cutoff : float
+        Coordination cutoff in the coordinate units.
+    cn_species : str, optional
+        Restrict the descriptor to atoms of one species.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    width = 0.1 * cutoff
+    keep = None
+    if cn_species:
+        keep = np.array([s == cn_species for s in symbols])
+    rows = []
+    for x in structures:
+        pos = np.asarray(x, dtype=float).reshape(-1, 3)
+        if keep is not None:
+            pos = pos[keep]
+        d = np.linalg.norm(pos[:, None, :] - pos[None, :, :], axis=-1)
+        np.fill_diagonal(d, np.inf)
+        cn = (1.0 / (1.0 + np.exp((d - cutoff) / width))).sum(axis=1)
+        rows.append(np.sort(cn))
+    return np.vstack(rows)
+
+
+def soap_matrix(structures, symbols, cutoff=4.0, hypers=None):
+    """Structure-averaged SOAP power spectra through featomic.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    import ase
+    from featomic import SoapPowerSpectrum
+
+    hy = dict(hypers or SOAP_HYPERS)
+    hy["cutoff"] = dict(hy["cutoff"], radius=cutoff)
+    calc = SoapPowerSpectrum(**hy)
+    rows = []
+    for x in structures:
+        pos = np.asarray(x, dtype=float).reshape(-1, 3)
+        pos = pos - pos.mean(axis=0)
+        atoms = ase.Atoms(symbols=symbols, positions=pos)
+        ps = calc.compute(atoms)
+        ps = ps.keys_to_samples("center_type").keys_to_properties(
+            ["neighbor_1_type", "neighbor_2_type"]
+        )
+        rows.append(ps.block().values.mean(axis=0))
+    return np.vstack(rows)
+
+
+def median_sigma(matrix, rng, sample=400):
+    """The sketch-map switching scale: median pairwise descriptor distance.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    sub = matrix[rng.choice(len(matrix), size=min(len(matrix), sample), replace=False)]
+    d = np.linalg.norm(sub[:, None, :] - sub[None, :, :], axis=-1)
+    return float(np.median(d[np.triu_indices(len(sub), 1)]))
+
+
+def run_sketchmap(
+    matrix,
+    smbin,
+    lapack_lib,
+    n_landmark,
+    sigma,
+    workdir,
+    n_pinned=0,
+):
+    """Landmark selection and the published sketch-map projection schedule.
+
+    ``dimlandmark`` picks minmax landmarks (the first ``n_pinned`` rows are
+    pinned); ``dimred`` runs sigmoid exponents 8,8 against 2,8, conjugate
+    gradient pre-optimization, then the pointwise global grid stage. No
+    ``-pi``: descriptor space is not periodic.
+
+    Returns
+    -------
+    coords, idx : ndarray
+        Projected landmark coordinates and their row indices in ``matrix``.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    workdir = Path(workdir)
+    workdir.mkdir(parents=True, exist_ok=True)
+    smbin = Path(smbin)
+    mat = workdir / "desc.dat"
+    np.savetxt(mat, matrix, fmt="%.8f")
+    env = {"LD_LIBRARY_PATH": str(lapack_lib), "PATH": "/usr/bin:/bin"}
+    dim = matrix.shape[1]
+    lm_raw = subprocess.run(  # noqa: S603
+        [
+            str(smbin / "dimlandmark"),
+            "-D",
+            str(dim),
+            "-n",
+            str(n_landmark),
+            "-mode",
+            "minmax",
+            "-i",
+            "-ifirst",
+            str(n_pinned),
+        ],
+        stdin=mat.open(),
+        capture_output=True,
+        text=True,
+        env=env,
+        check=True,
+    ).stdout
+    rows = [
+        line.split()
+        for line in lm_raw.splitlines()
+        if line.strip() and not line.startswith("#")
+    ]
+    idx = np.array([int(r[0]) for r in rows])
+    (workdir / "landmarks.dat").write_text(
+        "\n".join(" ".join(r[1:]) for r in rows) + "\n"
+    )
+    proj = subprocess.run(  # noqa: S603
+        [
+            str(smbin / "dimred"),
+            "-D",
+            str(dim),
+            "-d",
+            "2",
+            "-center",
+            "-fun-hd",
+            f"{sigma},8,8",
+            "-fun-ld",
+            f"{sigma},2,8",
+            "-preopt",
+            "500",
+            "-grid",
+            f"{10 * sigma:.4f},51,501",
+            "-gopt",
+            "10",
+        ],
+        stdin=(workdir / "landmarks.dat").open(),
+        capture_output=True,
+        text=True,
+        env=env,
+        check=True,
+    ).stdout
+    coords = np.array(
+        [
+            [float(v) for v in line.split()[:2]]
+            for line in proj.splitlines()
+            if line.strip() and not line.startswith("#")
+        ]
+    )
+    if len(idx) != len(coords):
+        msg = (
+            f"dimlandmark gave {len(idx)} indices, dimred {len(coords)} "
+            "points; refusing to plot with mismatched labels"
+        )
+        raise RuntimeError(msg)
+    return coords, idx
+
+
+def plot_landscape_surface(ax, coords, energies, cmap=None, levels=20):
+    """Filled-contour energy surface with thin contour lines.
+
+    Returns the contour set for colorbar wiring.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    from scipy.interpolate import griddata
+
+    cmap = cmap or mpl.colormaps["ruhi_diverging"]
+    vmin = float(np.nanmin(energies))
+    vmax = float(np.nanpercentile(energies, 97))
+    gx, gy = np.meshgrid(
+        np.linspace(coords[:, 0].min(), coords[:, 0].max(), 240),
+        np.linspace(coords[:, 1].min(), coords[:, 1].max(), 240),
+    )
+    surface = griddata(
+        coords, np.clip(energies, vmin, vmax), (gx, gy), method="linear"
+    )
+    lv = np.linspace(vmin, vmax, levels)
+    cf = ax.contourf(gx, gy, surface, levels=lv, cmap=cmap, vmin=vmin, vmax=vmax,
+                     alpha=0.85)
+    ax.contour(gx, gy, surface, levels=lv, colors="black", linewidths=0.3,
+               alpha=0.45)
+    ax.set_facecolor(RUHI_THEME.gridcolor)
+    return cf
+
+
+def margin_inset_slots(n):
+    """Axes-fraction slots down the left and right figure margins.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    slots = [(-0.26, 0.82), (-0.26, 0.22), (1.30, 0.82), (1.30, 0.22),
+             (-0.26, 0.52), (1.30, 0.52)]
+    return slots[:n]
+
+
+def place_margin_insets(fig, ax, entries, zoom=0.14, renderer="xyzrender"):
+    """Structure insets at margin slots, arrowed to their map points.
+
+    Parameters
+    ----------
+    entries : list of dict
+        Each carries ``atoms`` (ASE Atoms), ``xy`` (map point), and
+        optionally ``label`` drawn under the slot.
+
+    ```{versionadded} 1.10.0
+    ```
+    """
+    from chemparseplot.plot.neb import plot_structure_inset
+
+    slots = margin_inset_slots(len(entries))
+    fig.canvas.draw()
+    remaining = list(range(len(slots)))
+    for entry in entries:
+        x, y = entry["xy"]
+        pt_disp = ax.transData.transform((x, y))
+        if not remaining:
+            break
+        best = min(
+            remaining,
+            key=lambda s: np.linalg.norm(ax.transAxes.transform(slots[s]) - pt_disp),
+        )
+        remaining.remove(best)
+        slot_disp = ax.transAxes.transform(slots[best])
+        offset_pts = (slot_disp - pt_disp) * 72.0 / fig.dpi
+        plot_structure_inset(
+            ax,
+            entry["atoms"],
+            x,
+            y,
+            xybox=tuple(offset_pts),
+            rad=0.08,
+            zoom=zoom,
+            rotation="auto",
+            renderer=renderer,
+            xyzrender_config="paton",
+        )
+        ax.scatter([x], [y], marker="o", s=120, facecolors="none",
+                   edgecolors="black", linewidths=1.4, zorder=7)
+        if entry.get("label"):
+            ax.annotate(
+                entry["label"],
+                xy=slots[best],
+                xycoords="axes fraction",
+                xytext=(0, -58),
+                textcoords="offset points",
+                ha="center",
+                va="top",
+                fontsize=plt.rcParams["font.size"] - 3,
+                color="black",
+                weight="bold",
+                zorder=81,
+                annotation_clip=False,
+            )

chemparseplot-1.10.0/tests/plot/test_landscape_disconnectivity.py

@@ -0,0 +1,140 @@
+# SPDX-FileCopyrightText: 2023-present Rohit Goswami <rog32@hi.is>
+#
+# SPDX-License-Identifier: MIT
+"""Pure-python coverage of the landscape descriptors and the merge tree.
+
+The sketch-map binaries and featomic are external; everything here
+exercises the descriptor math and the disconnectivity pipeline only.
+"""
+
+import matplotlib
+
+matplotlib.use("Agg")
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+
+from chemparseplot.plot.disconnectivity import (
+    Node,
+    dedup,
+    iter_nodes,
+    knn_graph,
+    layout,
+    merge_tree,
+    plot_disconnectivity,
+    spectral_basins,
+)
+from chemparseplot.plot.landscape import cn_matrix, median_sigma
+
+
+def dimer_chain(n, spacing):
+    """A 1D chain of n atoms as a flat coordinate row."""
+    pos = np.zeros((n, 3))
+    pos[:, 0] = np.arange(n) * spacing
+    return pos.ravel()
+
+
+def test_cn_matrix_permutation_invariant():
+    rng = np.random.default_rng(3)
+    pos = rng.normal(size=(8, 3))
+    perm = rng.permutation(8)
+    a = cn_matrix([pos.ravel()], ["X"] * 8, cutoff=1.5)
+    b = cn_matrix([pos[perm].ravel()], ["X"] * 8, cutoff=1.5)
+    np.testing.assert_allclose(a, b, atol=1e-10)
+
+
+def test_cn_matrix_species_restriction():
+    pos = np.zeros((4, 3))
+    pos[:, 0] = [0.0, 1.0, 5.0, 6.0]
+    full = cn_matrix([pos.ravel()], ["O", "H", "O", "H"], cutoff=1.5)
+    only_o = cn_matrix(
+        [pos.ravel()], ["O", "H", "O", "H"], cutoff=1.5, cn_species="O"
+    )
+    assert full.shape[1] == 4
+    assert only_o.shape[1] == 2
+
+
+def test_cn_matrix_distinguishes_compact_from_stretched():
+    tight = dimer_chain(6, 1.1)
+    loose = dimer_chain(6, 3.0)
+    m = cn_matrix([tight, loose], ["X"] * 6, cutoff=1.5)
+    assert m[0].sum() > m[1].sum()
+
+
+def test_median_sigma_positive():
+    rng = np.random.default_rng(5)
+    m = rng.normal(size=(50, 4))
+    sigma = median_sigma(m, rng)
+    assert sigma > 0
+
+
+def test_dedup_keeps_lowest_energy_representative():
+    m = np.array([[0.0], [0.01], [5.0]])
+    es = np.array([2.0, 1.0, 3.0])
+    reps = dedup(m, es, tol=0.5)
+    assert 1 in reps  # the lower-energy twin survives
+    assert 0 not in reps
+    assert 2 in reps
+
+
+def test_merge_tree_two_basins_merge_once():
+    # Two well-separated pairs: components must merge only at the top.
+    m = np.array([[0.0], [0.1], [10.0], [10.1]])
+    es = np.array([-4.0, -3.0, -3.9, -2.9])
+    adjacency = knn_graph(m, k=1)
+    levels = np.linspace(es.min(), es.max() + 1.0, 20)[1:]
+    root = merge_tree(adjacency, es, levels)
+    members = sorted(n.leaf for n in iter_nodes(root) if n.leaf is not None)
+    assert members == [0, 1, 2, 3]
+    assert root.emin == -4.0
+    layout(root)
+    xs = [n.x for n in iter_nodes(root) if n.leaf is not None]
+    assert sorted(xs) == [0.0, 1.0, 2.0, 3.0]
+
+
+def test_merge_tree_deep_chain_no_recursion():
+    n = 3000
+    m = np.arange(n, dtype=float).reshape(-1, 1)
+    es = np.linspace(-10.0, -1.0, n)
+    adjacency = knn_graph(m, k=2)
+    levels = np.linspace(es.min(), es.max(), 40)[1:]
+    root = merge_tree(adjacency, es, levels)
+    layout(root)
+    assert len([x for x in iter_nodes(root) if x.leaf is not None]) == n
+
+
+def test_spectral_basins_shapes():
+    rng = np.random.default_rng(7)
+    m = np.vstack(
+        [rng.normal(0, 0.1, size=(10, 2)), rng.normal(5, 0.1, size=(10, 2))]
+    )
+    adjacency = knn_graph(m, k=3)
+    assign = spectral_basins(adjacency, 2, rng, matrix=m)
+    assert assign.shape == (20,)
+    # The two clouds separate into distinct basins.
+    assert len({int(a) for a in assign[:10]}) == 1
+    assert assign[0] != assign[-1]
+
+
+def test_plot_disconnectivity_end_to_end():
+    rng = np.random.default_rng(9)
+    m = np.vstack(
+        [rng.normal(0, 0.2, size=(15, 3)), rng.normal(4, 0.2, size=(15, 3))]
+    )
+    es = np.concatenate(
+        [rng.uniform(-10, -8, size=15), rng.uniform(-9, -7, size=15)]
+    )
+    fig, ax = plt.subplots()
+    reps, root = plot_disconnectivity(ax, m, es, dedup_tol=0.01, knn=3)
+    assert len(reps) <= 30
+    assert root.emin == pytest.approx(es.min())
+    plt.close(fig)
+
+
+def test_node_tracks_argmin():
+    a = Node(-5.0, leaf=0)
+    b = Node(-7.0, leaf=1)
+    parent = Node(-4.0, children=[a, b])
+    assert parent.emin == -7.0
+    assert parent.argmin == 1