Pair-angle distribution function (PADF) and 2D-PADF

src/quantem/core/visualization/__init__.py

@@ -3,6 +3,7 @@
     NormalizationConfig as NormalizationConfig,
 )
 from quantem.core.visualization.show_params import ShowParams as ShowParams
+from quantem.core.visualization.visualization import show_1d as show_1d
 from quantem.core.visualization.visualization import show_2d as show_2d
 from quantem.core.visualization.visualization_utils import (
     ScalebarConfig as ScalebarConfig,

src/quantem/core/visualization/visualization.py

@@ -669,3 +669,172 @@ def show_2d(
         axs_out = axs
 
     return fig, axs_out
+
+
+def show_1d(
+    arrays: ArrayLike | Sequence[ArrayLike] | Sequence[Sequence[ArrayLike]],
+    *,
+    x: ArrayLike | Sequence[ArrayLike] | None = None,
+    labels: str | Sequence[str] | Sequence[Sequence[str]] | None = None,
+    title: str | Sequence[str] | None = None,
+    xlabel: str | Sequence[str] | None = None,
+    ylabel: str | Sequence[str] | None = None,
+    log_y: bool | Sequence[bool] = False,
+    xlim: tuple[float, float] | None = None,
+    xticks: ArrayLike | None = None,
+    colors: Sequence[str] | None = None,
+    figax: tuple[Any, Any] | None = None,
+    axsize: tuple[float, float] = (5, 3),
+    save: os.PathLike | str | None = None,
+    **kwargs: Any,
+) -> tuple[Any, Any]:
+    """Plot one or more 1D arrays as line profiles.
+
+    A single array or a flat sequence of arrays is drawn in one panel, so that the curves can
+    be compared directly. A nested sequence draws one panel per inner sequence, in one row.
+    ``title``, ``xlabel``, ``ylabel``, and ``log_y`` may be scalars or one value per panel.
+
+    Parameters
+    ----------
+    arrays : ndarray or sequence of ndarray or sequence of sequences of ndarray
+        The 1D arrays to plot.
+    x : ndarray or sequence of ndarray, optional
+        Horizontal coordinates, either one array shared by all curves or one array per panel.
+        Defaults to the array indices.
+    labels : str or sequence of str or sequence of sequences of str, optional
+        Legend labels, one per curve. A legend is drawn in each panel that has labels.
+    title : str or sequence of str, optional
+        Panel titles.
+    xlabel, ylabel : str or sequence of str, optional
+        Axis labels.
+    log_y : bool or sequence of bool, default=False
+        Whether to use a logarithmic vertical axis.
+    xlim : tuple, optional
+        Horizontal axis limits, shared by all panels.
+    xticks : array-like, optional
+        Horizontal tick positions, shared by all panels.
+    colors : sequence of str, optional
+        Line colors, cycled over the curves of each panel. Defaults to ``viz.colors.set``.
+    figax : tuple, optional
+        (fig, axs) tuple to plot on. If None, a new figure and axes are created.
+    axsize : tuple, default=(5, 3)
+        Size of each panel in inches.
+    save : os.PathLike or str, optional
+        Path to save the figure to.
+    **kwargs : dict
+        Additional keyword arguments passed to ``matplotlib.axes.Axes.plot``, for example
+        ``marker``, ``linestyle``, or ``linewidth`` (default 2).
+
+    Returns
+    -------
+    fig : Figure
+        The matplotlib figure object.
+    axs : Axes | ndarray of Axes
+        A single Axes for one panel, otherwise a 1D ndarray of Axes.
+
+    Examples
+    --------
+    >>> depth = np.arange(9) * 17.5
+    >>> fig, ax = show_1d(
+    ...     [mean_potential, std_potential],
+    ...     x=depth,
+    ...     labels=["mean", "std"],
+    ...     xlabel="depth (Å)",
+    ...     ylabel="potential (rad)",
+    ...     marker="o",
+    ... )
+    """
+
+    def as_numpy(value: Any) -> Any:
+        if isinstance(value, (list, tuple)) and len(value) > 0 and np.isscalar(value[0]):
+            return np.asarray(value)  # a plain list of numbers is one curve
+        return to_cpu(value)
+
+    dpi = kwargs.pop("dpi", 300)
+    arrays = as_numpy(arrays)
+    if isinstance(arrays, np.ndarray):
+        arrays = [arrays]
+    panels = (
+        [[as_numpy(a) for a in p] for p in arrays]
+        if not isinstance(arrays[0], np.ndarray)
+        else [list(arrays)]
+    )
+    panels = [[np.asarray(a).ravel() for a in panel] for panel in panels]
+    num_panels = len(panels)
+
+    def per_panel(value: Any, name: str) -> list[Any]:
+        if value is None or isinstance(value, (str, bool)):
+            return [value] * num_panels
+        value = list(value)
+        if len(value) != num_panels:
+            raise ValueError(f"{name} has {len(value)} entries for {num_panels} panels")
+        return value
+
+    if x is None:
+        xs = [None] * num_panels
+    else:
+        x = as_numpy(x)
+        if isinstance(x, np.ndarray) and x.ndim == 1:
+            xs = [x] * num_panels
+        else:
+            xs = [np.asarray(v).ravel() for v in per_panel(x, "x")]
+
+    if labels is None or isinstance(labels, str):
+        panel_labels = [[labels] for _ in range(num_panels)]
+    elif num_panels == 1 and all(isinstance(lab, str) for lab in labels):
+        panel_labels = [list(labels)]
+    else:
+        panel_labels = [
+            [lab] if isinstance(lab, str) else list(lab) for lab in per_panel(labels, "labels")
+        ]
+
+    titles = per_panel(title, "title")
+    xlabels = per_panel(xlabel, "xlabel")
+    ylabels = per_panel(ylabel, "ylabel")
+    log_ys = per_panel(log_y, "log_y")
+    colors = list(config.get("viz.colors.set")) if colors is None else list(colors)
+    kwargs.setdefault("linewidth", 2)
+
+    if figax is not None:
+        fig, axs = figax
+        axs = np.atleast_1d(axs).ravel()
+        if len(axs) != num_panels:
+            raise ValueError(f"figax has {len(axs)} axes for {num_panels} panels")
+    else:
+        fig, axs = plt.subplots(
+            1, num_panels, figsize=(axsize[0] * num_panels, axsize[1]), squeeze=False
+        )
+        axs = axs.ravel()
+
+    for ax, panel, xv, labs, ttl, xl, yl, logy in zip(
+        axs, panels, xs, panel_labels, titles, xlabels, ylabels, log_ys
+    ):
+        for k, y in enumerate(panel):
+            coords = np.arange(len(y)) if xv is None else xv
+            if len(coords) != len(y):
+                raise ValueError(f"x has length {len(coords)} but the curve has length {len(y)}")
+            label = labs[k] if k < len(labs) else None
+            ax.plot(coords, y, color=colors[k % len(colors)], label=label, **kwargs)
+        if logy:
+            ax.set_yscale("log")
+        if xlim is not None:
+            ax.set_xlim(xlim)
+        if xticks is not None:
+            ax.set_xticks(xticks)
+        if ttl is not None:
+            ax.set_title(ttl)
+        if xl is not None:
+            ax.set_xlabel(xl)
+        if yl is not None:
+            ax.set_ylabel(yl)
+        if any(lab is not None for lab in labs):
+            ax.legend()
+
+    if figax is None:
+        fig.tight_layout()
+
+    if save is not None:
+        print(f"Saving figure to {save}")
+        fig.savefig(save, bbox_inches="tight", dpi=dpi)
+
+    return fig, (axs[0] if num_panels == 1 else axs)

src/quantem/diffraction/bispectrum.py

@@ -0,0 +1,161 @@
+"""
+Angular bispectrum (triple correlation) of polar diffraction data.
+
+The pair angular correlation used by the PADF keeps only |I_hat(q, n)|^2 -
+the angular power spectrum - discarding the relative phases between angular
+harmonics. The bispectrum
+
+    B_{q1 q2 q3}(n1, n2) = < I_hat_{q1}(n1) I_hat_{q2}(n2)
+                             conj( I_hat_{q3}(n1 + n2) ) >_patterns
+
+is the lowest-order rotation-invariant statistic that retains those phases
+(the harmonic closure n3 = n1 + n2 cancels the e^{i n alpha} rotation
+phases). Its 2D inverse Fourier transform is the angular triple correlation
+
+    C3(delta1, delta2) = < mean_phi I(phi) I(phi + delta1) I(phi + delta2) >.
+
+Working in the harmonic domain with a cutoff n_max keeps the object small:
+(2 n_max + 1)^2 complex values per ring triple, instead of the
+(N_phi)^2-per-triple maps of a direct triple loop.
+
+Note: for single flat-Ewald diffraction patterns, Friedel symmetry makes
+I(q, phi) pi-periodic, so odd harmonics vanish and C3 is 180-deg periodic
+in each angle; only the joint reflection (d1, d2) -> (pi - d1, pi - d2)
+remains degenerate - the relative fold between the two angles is broken.
+
+Class: AngularBispectrum
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+
+class AngularBispectrum:
+    """
+    Angular bispectrum / triple correlation of polar diffraction data.
+
+    Parameters
+    ----------
+    polar : Polar4dstem or torch.Tensor or ndarray
+        Polar data with the angular axis second-to-last and the radial axis
+        last: (..., N_phi, N_q). Leading dimensions (scan positions,
+        pattern index) are flattened into an ensemble of patterns.
+    n_max : int
+        Maximum angular harmonic retained. Governs angular resolution
+        (~180 / n_max degrees) and cost; the bispectrum is
+        (2 n_max + 1)^2 per ring triple.
+    subtract_mean : bool
+        Subtract each ring's angular mean per pattern (removes the n = 0
+        term, which otherwise dominates C3 with structureless offsets).
+    batch_size : int
+        Patterns per chunk when accumulating ensemble averages.
+    """
+
+    def __init__(self, polar, n_max: int = 60, subtract_mean: bool = True,
+                 batch_size: int = 4096):
+        tensor = polar.tensor if hasattr(polar, "tensor") else torch.as_tensor(np.asarray(polar))
+        tensor = tensor.to(torch.float64)
+        n_phi, n_q = tensor.shape[-2], tensor.shape[-1]
+        if n_max >= n_phi // 2:
+            n_max = n_phi // 2 - 1
+        self.n_max = int(n_max)
+        self.n_phi = n_phi
+        self.n_q = n_q
+        self.batch_size = int(batch_size)
+
+        data = tensor.reshape(-1, n_phi, n_q)
+        if subtract_mean:
+            data = data - data.mean(dim=1, keepdim=True)
+        self.n_patterns = data.shape[0]
+
+        # angular FFT per ring; normalized so coefficients are per-radian
+        # Fourier amplitudes independent of the phi sampling
+        F = torch.fft.rfft(data, dim=1) / n_phi  # (Npat, Nphi//2+1, Nq)
+        F = F[:, : self.n_max + 1, :]
+        # signed-harmonic spectrum from Hermitian symmetry of real data:
+        # index i <-> n = i - n_max, so F_signed[..., 0] is n = -n_max
+        neg = torch.conj(torch.flip(F[:, 1:, :], dims=[1]))
+        self.F = torch.cat([neg, F], dim=1).permute(0, 2, 1).contiguous()
+        # self.F: (Npat, Nq, 2*n_max+1), harmonic n at index n + n_max
+        self.n_values = np.arange(-self.n_max, self.n_max + 1)
+
+    def bispectrum(self, q1: int, q2: int | None = None, q3: int | None = None):
+        """
+        Ensemble-averaged bispectrum B(n1, n2) for the ring triple
+        (q1, q2, q3) (radial bin indices; q2, q3 default to q1).
+
+        Returns
+        -------
+        B : ndarray, complex, shape (2 n_max + 1, 2 n_max + 1)
+            B[i, j] with n1 = i - n_max, n2 = j - n_max. Entries with
+            |n1 + n2| > n_max are zero (outside the harmonic cutoff).
+        """
+        q2 = q1 if q2 is None else q2
+        q3 = q1 if q3 is None else q3
+        m = self.n_max
+        n = self.n_values
+        idx = n[:, None] + n[None, :]              # n1 + n2
+        valid = np.abs(idx) <= m
+        idx_c = torch.from_numpy(np.where(valid, idx + m, 0))
+        valid_t = torch.from_numpy(valid)
+
+        B = torch.zeros((2 * m + 1, 2 * m + 1), dtype=torch.complex128)
+        for start in range(0, self.n_patterns, self.batch_size):
+            F1 = self.F[start:start + self.batch_size, q1]  # (b, 2m+1)
+            F2 = self.F[start:start + self.batch_size, q2]
+            F3 = self.F[start:start + self.batch_size, q3]
+            F3g = F3[:, idx_c]                              # (b, 2m+1, 2m+1)
+            B += (F1[:, :, None] * F2[:, None, :] * torch.conj(F3g)).sum(dim=0)
+        B = B / self.n_patterns
+        B[~valid_t] = 0
+        return B.numpy()
+
+    def triple_correlation(self, q1: int, q2: int | None = None,
+                           q3: int | None = None, n_delta: int = 181):
+        """
+        Angular triple correlation C3(delta1, delta2) for the ring triple,
+        evaluated on an n_delta x n_delta grid over [0, 360) degrees.
+
+        Returns
+        -------
+        delta_deg : ndarray (n_delta,)
+        C3 : ndarray (n_delta, n_delta), real
+        """
+        B = self.bispectrum(q1, q2, q3)
+        delta = 2 * np.pi * np.arange(n_delta) / n_delta
+        E = np.exp(1j * np.outer(self.n_values, delta))  # (2m+1, n_delta)
+        C3 = np.real(E.T @ B @ E)
+        return np.degrees(delta), C3
+
+    def plot_triple_correlation(self, q1: int, q2: int | None = None,
+                                q3: int | None = None, n_delta: int = 181,
+                                delta_max: float = 180.0,
+                                title: str | None = None,
+                                figsize: tuple[float, float] = (4.6, 4.0),
+                                returnfig: bool = False):
+        """
+        Plot C3(delta1, delta2) as a zero-centered map. For diffraction
+        data C3 is 180-deg periodic in each angle, so the default view is
+        [0, 180]^2.
+        """
+        delta_deg, C3 = self.triple_correlation(q1, q2, q3, n_delta=2 * n_delta)
+        sel = delta_deg <= delta_max + 1e-9
+        d = delta_deg[sel]
+        M = C3[np.ix_(sel, sel)]
+        vmax = np.quantile(np.abs(M), 0.999)
+
+        fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
+        im = ax.pcolormesh(d, d, M.T, cmap="RdBu_r", vmin=-vmax, vmax=vmax,
+                           shading="auto")
+        ax.set_xlabel("$\\Delta_1$ (deg)")
+        ax.set_ylabel("$\\Delta_2$ (deg)")
+        ax.set_xticks(np.arange(0, delta_max + 1, 45))
+        ax.set_yticks(np.arange(0, delta_max + 1, 45))
+        ax.set_aspect("equal")
+        fig.colorbar(im, ax=ax, label="3-point correlation")
+        ax.set_title(title if title is not None else f"C3, ring {q1}",
+                     fontsize=10)
+        if returnfig:
+            return fig, ax
+        plt.show()