diff --git a/src/quantem/core/visualization/__init__.py b/src/quantem/core/visualization/__init__.py index 125fb02e..76082a11 100644 --- a/src/quantem/core/visualization/__init__.py +++ b/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, diff --git a/src/quantem/core/visualization/visualization.py b/src/quantem/core/visualization/visualization.py index e252caa6..252c5f0d 100644 --- a/src/quantem/core/visualization/visualization.py +++ b/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) diff --git a/src/quantem/diffraction/bispectrum.py b/src/quantem/diffraction/bispectrum.py new file mode 100644 index 00000000..a36c3a13 --- /dev/null +++ b/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() diff --git a/src/quantem/diffraction/gr_truth.py b/src/quantem/diffraction/gr_truth.py new file mode 100644 index 00000000..a77a1262 --- /dev/null +++ b/src/quantem/diffraction/gr_truth.py @@ -0,0 +1,400 @@ +""" +Adapted from PDFs_3batch by Arthur McCray (renamed 9/23/2026) + +Three-body g3(r, r', theta) histograms of atomic structures (PDF3B), for 3D +cells and for 2D (monolayer / coplanar) structures. + +Differences from PDFs_3batch.py +------------------------------- +1. Flat-axis detection: a non-periodic axis along which every atom shares the + same coordinate is "flat". Centers are no longer trimmed along such an axis + (a monolayer loses no neighbors above/below the plane), which previously + emptied `center_inds` and crashed on `max([len(i) for i in inds_all])`. + For genuinely thick non-periodic axes the trim margin is now measured from + the atomic extent rather than the cell, so a thin film in a tall vacuum cell + is handled correctly. +2. Areal density (atoms / A^2) is used when exactly one axis is flat, instead of + a volumetric density taken from an arbitrary vacuum-padded cell. +3. 2D normalization: the coplanar ideal-gas pair measure is + 4 * pi * sigma^2 * r1 * r2 * dr^2 * dtheta + i.e. linear in r and with NO sin(theta) factor. Dividing a 2D PADF by + sin(theta) manufactures spurious intensity at 0 and 180 degrees. +4. `good_vals` now compares the theta bin index against the number of theta + bins rather than against the literal value 180 (the original was only + correct for dtheta == 1). + +The 3D code path is unchanged; a fully periodic or thick non-periodic cell +behaves exactly as before. +""" + +import numpy as np +import time + +from scipy.spatial import KDTree +from datetime import timedelta + +try: + import cupy as cp +except (ImportError, ModuleNotFoundError): + cp = np + + +class PDF3B(object): + + def __init__(self, device="cpu", v=1): + if isinstance(device, int): + self._xp = cp + cp.cuda.runtime.setDevice(device) + elif device.lower() == "cpu": + self._xp = np + self.device = device + self._v = v + return + + def pdf3B( + self, + atoms, + dr=0.1, + dtheta=1, + hist_r_max=20, + skip=1, + batch_size=1, + v=None, + pbcs=None, + subpixel_weights=True, + flat_tol=1e-6, + ): + xp = self._xp + stime = time.perf_counter() + v = self._v if v is None else v + vprint = print if v >= 1 else lambda *a, **k: None + if skip == None: + skip = 1 + if pbcs is None: + pbcs = atoms.get_pbc() + self.pbcs = pbcs + + bbox_np = np.array(atoms.cell[(0, 1, 2), (0, 1, 2)]) + bbox = xp.array(bbox_np) + pbcs_np = np.copy(pbcs) + pbcs = xp.array(pbcs) + positions_np = np.array(atoms.positions) # for KD tree + + ### wrap into the cell along periodic dims (KDTree requires 0 <= x < L) + for i, p in enumerate(pbcs_np): + if p: + positions_np[:, i] = np.mod(positions_np[:, i], bbox_np[i]) + positions_np[positions_np[:, i] >= bbox_np[i], i] = 0.0 # fp edge case + positions = xp.array(positions_np) + + vprint(f"Cell size (A): {bbox_np}") + vprint(f"Using a max radius of {hist_r_max} A") + + ### FIX 1a: identify flat axes -- non-periodic with zero atomic extent. + ### A monolayer has no neighbors to lose along such an axis, so every + ### atom already has a complete neighborhood and must not be trimmed. + flat_axes = [ + i + for i in range(3) + if not pbcs_np[i] and np.ptp(positions_np[:, i]) < flat_tol + ] + self.flat_axes = flat_axes + if len(flat_axes) == 1: + vprint( + f"Detected flat (2D) axis {flat_axes[0]}: using areal density " + "and 2D normalization" + ) + elif len(flat_axes) > 1: + vprint( + f"WARNING: {len(flat_axes)} flat axes {flat_axes} -> structure is " + "1D or a single point. Falling back to volumetric normalization, " + "which is not correct for this geometry." + ) + + ### warn if the minimum-image convention is violated in a periodic dim. + ### _get_vecs_pbcs and query_ball_tree both require L > 2 * hist_r_max + for i, p in enumerate(pbcs_np): + if p and bbox_np[i] <= 2 * hist_r_max: + vprint( + f"WARNING: periodic axis {i} has L = {bbox_np[i]:.2f} A " + f"<= 2 * hist_r_max = {2 * hist_r_max} A. Minimum-image " + "distances and neighbor counts will be wrong -- tile the " + "cell first, e.g. atoms * (n, m, 1)." + ) + + ### FIX 1b: get center positions that aren't within r_max of a boundary + ### without pbcs. Margins are measured from the atomic extent, not the + ### cell, and flat axes are skipped entirely. + goods = np.ones(len(positions_np)).astype("bool") + for i, p in enumerate(pbcs_np): + if p or i in flat_axes: + continue + lo = positions_np[:, i].min() + hi = positions_np[:, i].max() + goods *= (positions_np[:, i] > lo + hist_r_max) & ( + positions_np[:, i] < hi - hist_r_max + ) + center_inds = np.where(goods)[0] + if len(center_inds) == 0: + raise ValueError( + "No valid center atoms: every atom lies within hist_r_max of a " + "non-periodic boundary. The atomic extent along each finite, " + f"non-flat axis must exceed 2 * hist_r_max = {2 * hist_r_max} A. " + "Reduce hist_r_max, use a thicker sample, or enable pbcs." + ) + + ### KD tree for getting neighbors, cpu cuz faster + start_KD = time.perf_counter() + bbox_np_pbcs = bbox_np * pbcs_np + bbox_np_pbcs[pbcs_np.astype(bool)] += 1e-9 + tree_all = KDTree(positions_np, copy_data=True, boxsize=bbox_np_pbcs) + center_inds_skip = center_inds[::skip] + tree_centers = KDTree( + positions_np[center_inds_skip], copy_data=True, boxsize=bbox_np_pbcs + ) + inds_all = tree_centers.query_ball_tree(tree_all, hist_r_max) + inds_all = [xp.array(i) for i in inds_all] + end_KD = time.perf_counter() + vprint(f"KDTree calc time: {end_KD - start_KD:.1e} s") + + ### Histogram coordinates + hist_theta = xp.arange(0.0, 180 + dtheta, dtheta) + hist_r = xp.arange(0.0, hist_r_max + dr, dr) + num_bins_r = hist_r.shape[0] + num_bins_theta = hist_theta.shape[0] + hist_sig = xp.zeros((num_bins_theta, num_bins_r, num_bins_r)) + + ### FIX 2: Density -- areal for a 2D structure, volumetric otherwise + cell = xp.array(atoms.cell) + if len(flat_axes) == 1: + i, j = [k for k in range(3) if k != flat_axes[0]] + area = xp.linalg.norm(xp.cross(cell[i], cell[j])) + assert area > 0, "in-plane cell area is zero" + dens = positions.shape[0] / area # atoms / A^2 + vprint(f"Areal density = {float(dens):.4g} atoms / A^2") + else: + volume = xp.abs(xp.sum(cell[:, 0] * xp.cross(cell[:, 1], cell[:, 2]))) + assert volume > 0, "cell volume is zero" + dens = positions.shape[0] / volume # atoms / A^3 + + ### making big arrays that will be filled + maxinds = max([len(i) for i in inds_all]) + inds2 = xp.zeros((batch_size, maxinds)).astype("int") + + if hist_r_max > 15: + split_bincount = True + off_inds = xp.triu_indices(maxinds, k=1) + else: + split_bincount = False + off_inds = xp.where(~xp.eye(maxinds, dtype=bool)) + + ### exact upper bound: one entry per (center, pair) in the batch + maxval_n = int(off_inds[0].shape[0]) + vals = xp.zeros((3, maxval_n * batch_size)).astype("int") + if subpixel_weights: + weights_full = xp.empty(maxval_n * batch_size) + weights = weights_full + else: + weights_full = None + weights = None + + vprint(f"Num total atoms in sim = {len(positions)}") + if not np.all(pbcs_np) and len(flat_axes) < 3: + vprint(f"Num valid center atoms = {len(center_inds)}") + if skip != 1: + vprint( + f"Skip = {skip}, so calculating using {len(center_inds_skip)} atoms as centers" + ) + vprint("Final shape will be: ", hist_sig.shape) + + for a0 in range(0, len(center_inds_skip), batch_size): + b0 = min(a0 + batch_size, len(center_inds_skip)) + batch_center_inds = center_inds_skip[a0:b0] + bsize = len(batch_center_inds) + bnum_inds = [len(inds) for inds in inds_all[a0:b0]] + for i in range(bsize): + inds2[i, : bnum_inds[i]] = inds_all[a0 + i] + + vecs = self._get_vecs_pbcs( + positions[batch_center_inds], + positions[inds2[:bsize]], + bbox=bbox, + pbcs=pbcs, + ) + for i in range(bsize): + vecs[i, bnum_inds[i] :] = 0 + + drs = xp.sqrt(xp.sum(vecs**2, axis=-1)) + + thetas_full = self._get_thetas_all(vecs, vecs, deg=True) + thetas = thetas_full[:, off_inds[0], off_inds[1]] + thetas = thetas[:bsize].ravel() + + dr1s = drs[:, off_inds[0]] + dr2s = drs[:, off_inds[1]] + dr1s = dr1s[:bsize].ravel() + dr2s = dr2s[:bsize].ravel() + + r_ind1 = dr1s / dr + r_ind2 = dr2s / dr + theta_ind = thetas / dtheta + + if subpixel_weights: + ### this isn't perfect, doesn't do corners + r_floor1 = xp.floor(r_ind1).astype("int") + dr_ind1 = r_ind1 - r_floor1 + + r_floor2 = xp.floor(r_ind2).astype("int") + dr_ind2 = r_ind2 - r_floor2 + + theta_floor = xp.floor(theta_ind).astype("int") + dtheta_ind = theta_ind - theta_floor + + else: + r_floor1 = xp.round(r_ind1).astype("int") + r_floor2 = xp.round(r_ind2).astype("int") + theta_floor = xp.round(theta_ind).astype("int") + + ### FIX 4: theta_floor is a BIN INDEX, so bound it by the number of + ### theta bins. The original compared it against the literal 180, + ### which was only correct for dtheta == 1. + good_vals = ( + (r_floor1 > 1e-9) + & (r_floor2 > 1e-9) + & (theta_floor > 0) + & (theta_floor < num_bins_theta - 1) + ) + good_inds = xp.where(good_vals)[0] + nvals = len(good_inds) + vals[0, :nvals] = theta_floor[good_inds] + vals[1, :nvals] = r_floor2[good_inds] + vals[2, :nvals] = r_floor1[good_inds] + if subpixel_weights: + weights_full[:nvals] = ( + dtheta_ind[good_inds] + dr_ind2[good_inds] + dr_ind1[good_inds] + ) / 3 + weights = weights_full[:nvals] + + ### Because taking triu, have to do this twice + ### and doing bincount twice with half indices seems to be ~10% faster than + ### once with twice as many (for a 20A cutoff), but slower if < 15 A + hist_sig += self._bincountdd( + vals[:, :nvals], + (num_bins_theta, num_bins_r, num_bins_r), + weights=weights, + ) + if split_bincount: + vals[[1, 2]] = vals[[2, 1]] + hist_sig += self._bincountdd( + vals[:, :nvals], + (num_bins_theta, num_bins_r, num_bins_r), + weights=weights, + ) + + hist_sig /= len(center_inds_skip) + + ### FIX 3: normalization + gr = hist_sig.copy() + r1 = hist_r[1:][None, None, ...] + r2 = hist_r[1:][None, ..., None] + th = hist_theta[1:-1, None, None] + if len(flat_axes) == 1: + ### 2D ideal-gas pair measure: + ### sigma^2 * (2 pi r1 dr) * (2 r2 dr dtheta) + ### -> linear in r, and no sin(theta): the out-of-plane solid-angle + ### band that produces sin(theta) in 3D does not exist for coplanar + ### vectors. The factor 2 folds +/- azimuth into [0, 180]. + gr[1:-1, 1:, 1:] /= ( + dens**2 + * dr**2 + * 4 # 2 * 2pi, for all off-diagonal (ordered) pairs; use 2 for triu only + * xp.pi + * r1 + * r2 + * np.deg2rad(dtheta) + ) + else: + gr[1:-1, 1:, 1:] /= ( + dens**2 + * dr**2 + * 8 # if all off diag indices + # * 4 # if triu + * xp.pi**2 + * r1**2 + * r2**2 + * xp.sin(xp.deg2rad(th)) + * np.deg2rad(dtheta) + ) + + ### edge bins are half-width, and bin 0 is excluded by good_vals + gr[-1] = 0 + gr[0] = 0 + gr[:, 0] = 0 + gr[:, :, 0] = 0 + + vprint("-- done --") + ttime = time.perf_counter() - stime + vprint( + f"Total time (h:m:s) {str(timedelta(seconds=round(ttime,3))).rstrip('0')}" + ) + vprint(f"Center atoms per sec: {len(center_inds_skip) / ttime:.1f}\n") + return gr + + def _get_vecs_pbcs(self, points, pointslists, bbox, pbcs=[1, 1, 1]): + xp = self._xp + if np.any(pbcs): + assert xp.all(xp.min(pointslists, axis=1) >= 0) + assert xp.all(xp.max(pointslists, axis=1) <= bbox) + dif = pointslists - points[:, None] + for i, ind in enumerate(pbcs): + if ind: + dif[:, :, i] = ( + xp.mod(dif[:, :, i] + bbox[i] * 0.5, bbox[i]) - bbox[i] * 0.5 + ) + return dif + + def _bincountdd(self, vals, nbins, weights=None): + # based on xp.histogramdd + """ + vals: array shape (N, D) + weights: array shape N + nbins: (D,) tuple (nbinsz, nbinsy, nbinsx) + based on xp.histogramdd + """ + xp = self._xp + vals = xp.floor(vals).astype("int") + xy = xp.ravel_multi_index(vals, nbins, mode="wrap") + minlength = int(np.prod(nbins)) + if weights is not None: + hist = xp.bincount(xy, weights=1 - weights, minlength=minlength).reshape( + nbins + ) + xy = xp.ravel_multi_index(vals + 1, nbins, mode="wrap") + hist += xp.bincount(xy, weights=weights, minlength=minlength).reshape(nbins) + else: + hist = xp.bincount(xy, minlength=minlength).reshape(nbins) + return hist + + def _get_thetas_all(self, v1, v2, deg=True): + """ + get angles between all vectors, outer product style + """ + xp = self._xp + v11 = self._norm_vecs(v1) + v21 = self._norm_vecs(v2) + mmul = xp.transpose(xp.matmul(v11, xp.transpose(v21, (0, 2, 1))), (0, 2, 1)) + theta = xp.arccos(xp.clip(mmul, -1.0, 1.0)) + if deg: + theta = xp.rad2deg(theta) + return theta + + def _norm_vecs(self, v): + xp = self._xp + if xp.ndim(v) == 1: + return v / xp.sqrt(xp.sum(v**2, axis=-1))[..., None] + mags = xp.sqrt(xp.sum(v**2, axis=-1)) + bads = xp.where(mags == 0) + mags[bads] = 1 + normed = v / mags[..., None] + normed[bads] = 0 + return normed diff --git a/src/quantem/diffraction/lattices.py b/src/quantem/diffraction/lattices.py new file mode 100644 index 00000000..e7276f1a --- /dev/null +++ b/src/quantem/diffraction/lattices.py @@ -0,0 +1,582 @@ +""" +Lattice test-image generation for PADF testing. + +Generates 2D "atom" images for lattices with different symmetries and +near-neighbor spacings: + - triangular (hexagonal Bravais lattice, 6-fold coordination) + - square + - honeycomb (graphene-like, 2-atom basis; called "hexagonal" informally) + - decagonal (10-fold quasilattice via the de Bruijn pentagrid method) + +plus kinematic diffraction patterns of random clusters cut from 3D structures +(random_clusters, cluster_patterns). + +The 2D lattice positions are computed in pixels with (row, col) ordering; bond +lengths and atom widths are specified in Angstroms and converted with +pixel_size_ang. The cluster functions work in Angstroms. +""" + +import matplotlib.pyplot as plt +import numpy as np +from scipy.ndimage import gaussian_filter +import torch + +GOLDEN = (1 + np.sqrt(5)) / 2 + + +def gaussian_form_factor(q, atom_sigma): + """ + Scattering factor f(q) of the Gaussian atoms drawn by render_atoms: + a real-space Gaussian of width sigma transforms to + f(q) = exp(-2 pi^2 sigma^2 q^2). + + Divide diffraction intensities by f(q)^2 (e.g. via the PADF module's + `fq` argument) to undo the atom-shape damping of the ring intensities; + floor the returned values before dividing so the high-q noise floor is + not amplified. + """ + q = np.asarray(q, dtype=float) + return np.exp(-2 * np.pi**2 * atom_sigma**2 * q**2) + + +def lattice_vectors(lattice: str, bond_length: float): + """ + Return (lat, basis) for a periodic lattice with the requested bond length. + + lat : (2, 2) array, rows are the lattice vectors in Angstroms. + basis : (n_atoms, 2) array of fractional coordinates in the cell. + + The triangular and square lattices have lattice parameter equal to the + bond length. The honeycomb cell is the triangular cell scaled by sqrt(3) + because its bond length is a / sqrt(3). + """ + if lattice == "boron_nitride": + # honeycomb geometry with two species; weights looked up separately + lattice = "honeycomb" + if lattice == "triangular": + lat = np.array([ + [1.0, 0.0], + [0.5, np.sqrt(3) / 2], + ]) * bond_length + basis = np.array([[0.0, 0.0]]) + elif lattice == "square": + lat = np.array([ + [1.0, 0.0], + [0.0, 1.0], + ]) * bond_length + basis = np.array([[0.0, 0.0]]) + elif lattice == "honeycomb": + lat = np.array([ + [1.0, 0.0], + [0.5, np.sqrt(3) / 2], + ]) * (bond_length * np.sqrt(3)) + basis = np.array([ + [1 / 3, 1 / 3], + [2 / 3, 2 / 3], + ]) + else: + raise ValueError(f"Unknown lattice '{lattice}'") + return lat, basis + + +def tile_lattice(lat_px, basis, im_shape, pad_px=2.0): + """ + Tile a periodic lattice so it fully covers a rectangular image. + + Parameters + ---------- + lat_px : (2, 2) array + Lattice vectors in pixels (rows are vectors, (row, col) components). + basis : (n, 2) array + Fractional coordinates of the atoms in the cell. + im_shape : tuple + Image shape (rows, cols) in pixels. + pad_px : float + Extra margin (pixels) kept around the image so atoms straddling the + border still render correctly. + + Returns + ------- + points_px : (m, 2) array of atom positions in pixels. + + Notes + ----- + The (a, b) index range needed to cover the image is found by expressing + the four image corners in lattice coordinates: solve lat_px.T @ (a, b) = + corner for each corner. Because the image is convex, the extreme lattice + coordinates occur at the corners, so tiling the bounding index box and + cropping afterwards is guaranteed to cover the image. + """ + n_row, n_col = im_shape + corners = np.array([ + [-pad_px, -pad_px], + [-pad_px, n_col + pad_px], + [n_row + pad_px, -pad_px], + [n_row + pad_px, n_col + pad_px], + ]) + # corners and lat_px are both in PIXELS - mixing units here breaks the + # tiling range. lat_px.T is 2x2 and invertible, so solve exactly. + ab = np.linalg.solve(lat_px.T, corners.T) # (2, 4) lattice coords + # widen by 1 for the basis offsets (fractional coords in [0, 1)) + a_lo, a_hi = int(np.floor(ab[0].min())) - 1, int(np.ceil(ab[0].max())) + 1 + b_lo, b_hi = int(np.floor(ab[1].min())) - 1, int(np.ceil(ab[1].max())) + 1 + + a, b = np.meshgrid( + np.arange(a_lo, a_hi + 1), + np.arange(b_lo, b_hi + 1), # b gets its own range - a and b differ + indexing="ij", + ) + cells = np.stack([a.ravel(), b.ravel()], axis=1).astype(float) # (m, 2) + # combine every cell with every basis atom + frac = (cells[:, None, :] + basis[None, :, :]).reshape(-1, 2) + points_px = frac @ lat_px + + keep = ( + (points_px[:, 0] >= -pad_px) + & (points_px[:, 0] <= n_row + pad_px) + & (points_px[:, 1] >= -pad_px) + & (points_px[:, 1] <= n_col + pad_px) + ) + return points_px[keep] + + +def decagonal_points(radius, edge_length, gamma=None, rng=None): + """ + Vertices of a decagonal (Penrose-like) quasilattice via the de Bruijn + pentagrid dual method, out to the given radius (both in Angstroms). + + Five families of parallel grid lines with normals e_j at angles 2*pi*j/5 + and offsets gamma_j are intersected pairwise; each intersection maps to a + rhombic tile of the dual tiling whose four vertices are sums of integer + multiples of the e_j. Edges of the tiling all have length `edge_length`. + NOTE: the SHORTEST pair distance in the vertex set is the thin-rhomb + short diagonal, edge_length / GOLDEN (= 0.618 * edge_length) - scale + edge_length by GOLDEN if the shortest distance should equal a target + bond length (make_images does this). + + gamma : (5,) grid offsets. Random (generalized Penrose) by default, + which avoids the degenerate >2-line intersections of the singular + sum(gamma) = 0 case. + """ + if rng is None: + rng = np.random.default_rng(0) + if gamma is None: + gamma = rng.uniform(0.05, 0.95, 5) + angles = 2 * np.pi * np.arange(5) / 5 + e = np.stack([np.cos(angles), np.sin(angles)], axis=1) # (5, 2) + + r_lat = radius / edge_length # work in units of the edge length + nmax = int(np.ceil(r_lat)) + 2 + ns = np.arange(-nmax, nmax + 1) + + all_verts = [] + for j in range(5): + for k in range(j + 1, 5): + A = np.array([e[j], e[k]]) # (2, 2) + nj, nk = np.meshgrid(ns, ns, indexing="ij") + nj, nk = nj.ravel(), nk.ravel() + rhs = np.stack([nj + gamma[j], nk + gamma[k]]) # (2, m) + X = np.linalg.solve(A, rhs) # (2, m) intersection points + near = np.hypot(X[0], X[1]) <= r_lat + 1.5 + X, nj_s, nk_s = X[:, near], nj[near], nk[near] + proj = e @ X - gamma[:, None] # (5, m) + K = np.ceil(proj) + for da in (0.0, 1.0): + for db in (0.0, 1.0): + Kv = K.copy() + # on the two defining grids the projection is exactly + # integer; set those indices explicitly (ceil is + # numerically unstable there) + Kv[j] = nj_s + da + Kv[k] = nk_s + db + all_verts.append(Kv.T @ e) # (m, 2) + + verts = np.vstack(all_verts) + verts = np.unique(np.round(verts, 6), axis=0) + verts = verts[np.hypot(verts[:, 0], verts[:, 1]) <= r_lat] + return verts * edge_length + + +# per-basis-atom scattering weights (relative); default is 1 for every atom. +# boron_nitride: honeycomb with B (Z=5) on one sublattice, N (Z=7) on the +# other - the projection loses its inversion center, so dynamical scattering +# can violate Friedel symmetry (unlike every other lattice here). +BASIS_WEIGHTS = { + "boron_nitride": (5.0 / 7.0, 1.0), +} + +SHELL_MULTS = { + "triangular": (1.0, np.sqrt(3), 2.0), + "square": (1.0, np.sqrt(2), 2.0), + "honeycomb": (1.0, np.sqrt(3), 2.0), + "boron_nitride": (1.0, np.sqrt(3), 2.0), + # decagonal: shortest-distance pairs (thin-rhomb diagonals) are rare; + # the strong shells are tile edges (phi d), the thick-rhomb short + # diagonal (2 sin36 phi d = 1.902 d), and phi^2 d + "decagonal": (GOLDEN, 2 * np.sin(np.radians(36)) * GOLDEN, GOLDEN**2), +} + + +def bond_angle_targets(lattice, bond_length, tol_frac=0.03): + """ + Expected (r', theta) maxima of the 3-body (PADF) correlation when one + arm is fixed on a 1st-shell bond: for every central atom, one arm goes + to a 1st-shell neighbor and the other to a 1st / 2nd / 3rd shell + neighbor; collect the shell distance r' and the angle between the arms. + + 1st: O--O 2nd: O----O 3rd: O------O + | | | + O O O + + Computed numerically by enumerating a patch of the actual lattice + (shell distances from SHELL_MULTS; for the decagonal quasilattice the + "1st shell" arm is the tile edge at GOLDEN * bond_length, its dominant + pair distance). + + Returns + ------- + targets : list of (r_shell, angles_deg, weights) + One entry per shell. angles_deg are the distinct arm-arm angles in + [0, 180] (degrees); weights are the relative multiplicities of each + angle (normalized to the most frequent). + """ + shells = [m * bond_length for m in SHELL_MULTS[lattice]] + span = shells[-1] * 2.5 + if lattice == "decagonal": + pts = decagonal_points(span + 2 * bond_length, bond_length * GOLDEN, + rng=np.random.default_rng(0)) + else: + lat, basis = lattice_vectors(lattice, bond_length) + n_cell = int(np.ceil(span / bond_length)) + 2 + a, b = np.meshgrid(np.arange(-n_cell, n_cell + 1), + np.arange(-n_cell, n_cell + 1), indexing="ij") + cells = np.stack([a.ravel(), b.ravel()], axis=1).astype(float) + frac = (cells[:, None, :] + basis[None, :, :]).reshape(-1, 2) + pts = frac @ lat + pts = pts[np.hypot(pts[:, 0], pts[:, 1]) <= span] + + tol = tol_frac * bond_length + 1e-9 + radii = np.hypot(pts[:, 0], pts[:, 1]) + # central atoms far enough from the patch edge to see all shells + centers = pts[radii <= span - shells[-1] - tol] + + targets = [] + diff_all = pts[None, :, :] - centers[:, None, :] # (Nc, Np, 2) + dist_all = np.hypot(diff_all[..., 0], diff_all[..., 1]) + arm_mask = np.abs(dist_all - shells[0]) < tol + for r_shell in shells: + partner_mask = np.abs(dist_all - r_shell) < tol + angles = [] + for i in range(centers.shape[0]): + arms = diff_all[i][arm_mask[i]] + partners = diff_all[i][partner_mask[i]] + if arms.size == 0 or partners.size == 0: + continue + cosang = (arms @ partners.T) / ( + np.hypot(arms[:, 0], arms[:, 1])[:, None] + * np.hypot(partners[:, 0], partners[:, 1])[None, :] + ) + ang = np.degrees(np.arccos(np.clip(cosang, -1, 1))).ravel() + # drop the self-pair (same vector for both arms -> angle 0) + if abs(r_shell - shells[0]) < tol: + ang = ang[ang > 0.5] + angles.append(ang) + ang = np.concatenate(angles) + # cluster to distinct angles (1 degree bins) + vals, counts = np.unique(np.round(ang).astype(int), return_counts=True) + keep = counts > 0.05 * counts.max() + vals, counts = vals[keep], counts[keep] + # merge adjacent-degree bins + merged_vals, merged_wts = [], [] + for v, c in zip(vals, counts): + if merged_vals and v - merged_vals[-1][-1] <= 1: + merged_vals[-1].append(v) + merged_wts[-1] += c + else: + merged_vals.append([v]) + merged_wts.append(c) + angs = np.array([float(np.mean(g)) for g in merged_vals]) + wts = np.array(merged_wts, dtype=float) + targets.append((float(r_shell), angs, wts / wts.max())) + return targets + + +def render_atoms(points_px, im_shape, atom_sigma_px, weights=None): + """ + Render atoms as Gaussians of width atom_sigma_px (pixels): bilinear + deposition of one weight per atom onto a padded grid, then a Gaussian + blur, then crop. Atoms slightly outside the image contribute their + in-image tails, as they would in a real micrograph. + + weights : per-atom deposit weights (e.g. relative scattering strengths + of different species); defaults to 1 for every atom. + """ + n_row, n_col = im_shape + p = int(np.ceil(4 * atom_sigma_px + 2)) # padding so border atoms blur in + im = np.zeros((n_row + 2 * p, n_col + 2 * p)) + + pts = points_px + p + w = np.ones(pts.shape[0]) if weights is None else np.asarray(weights, dtype=float) + keep = ( + (pts[:, 0] >= 0) & (pts[:, 0] < im.shape[0] - 1) + & (pts[:, 1] >= 0) & (pts[:, 1] < im.shape[1] - 1) + ) + pts, w = pts[keep], w[keep] + r0 = np.floor(pts[:, 0]).astype(int) + c0 = np.floor(pts[:, 1]).astype(int) + dr = pts[:, 0] - r0 + dc = pts[:, 1] - c0 + np.add.at(im, (r0, c0), w * (1 - dr) * (1 - dc)) + np.add.at(im, (r0, c0 + 1), w * (1 - dr) * dc) + np.add.at(im, (r0 + 1, c0), w * dr * (1 - dc)) + np.add.at(im, (r0 + 1, c0 + 1), w * dr * dc) + + im = gaussian_filter(im, atom_sigma_px) + return im[p:-p, p:-p] + + +def make_images( + im_shape=(256, 256), + pixel_size_ang=0.2, + atom_sigma_ang=0.4, + bond_length_ang=(2.0, 2.5, 3.0), + lattices=("triangular", "square", "honeycomb", "decagonal"), + rotation_deg=None, + seed=0, + plot_images=True, +): + """ + Generate lattice test images for every (lattice, bond length) pair. + + Parameters + ---------- + im_shape : tuple + Image shape (rows, cols) in pixels. + pixel_size_ang : float + Pixel size in Angstroms. + atom_sigma_ang : float + Gaussian atom width (sigma) in Angstroms. + bond_length_ang : sequence of float + Near-neighbor bond lengths to generate. For every lattice type + (including decagonal) this is the SHORTEST pair distance in the + generated points. + lattices : sequence of str + Any of "triangular", "square", "honeycomb", "decagonal". + rotation_deg : float or None + In-plane lattice rotation. None draws a random angle per image so no + lattice aligns with the pixel grid (avoids FFT axis artifacts). + seed : int + Seed for the rotation angles and quasilattice offsets. + plot_images : bool + Show a grid of the generated images. + + Returns + ------- + data : list of (rows, cols) float arrays + info : list of dicts with keys + lattice, bond_length_ang, rotation_deg, pixel_size_ang, num_atoms, + positions_ang (the (row, col) atom coordinates in Angstroms that + were rendered into the image, e.g. for building an ase.Atoms + ground truth structure) + """ + rng = np.random.default_rng(seed) + n_row, n_col = im_shape + atom_sigma_px = atom_sigma_ang / pixel_size_ang + center_px = np.array([(n_row - 1) / 2, (n_col - 1) / 2]) + + data = [] + info = [] + for lattice in lattices: + for bond in bond_length_ang: + theta = ( + rng.uniform(0, 360) if rotation_deg is None else float(rotation_deg) + ) + t = np.deg2rad(theta) + R = np.array([[np.cos(t), -np.sin(t)], [np.sin(t), np.cos(t)]]) + + if lattice == "decagonal": + # Tile edges are scaled by the golden ratio so the SHORTEST + # pair distance in the vertex set (the thin-rhomb diagonal, + # edge / GOLDEN) equals the requested bond length. Pair + # distances are then bond, GOLDEN * bond (tile edges), ... + radius = 0.5 * np.hypot(n_row, n_col) * pixel_size_ang + 2 * bond + pts_ang = decagonal_points(radius, bond * GOLDEN, rng=rng) + pts_px = (pts_ang @ R.T) / pixel_size_ang + center_px + pad = 2.0 + keep = ( + (pts_px[:, 0] >= -pad) & (pts_px[:, 0] <= n_row + pad) + & (pts_px[:, 1] >= -pad) & (pts_px[:, 1] <= n_col + pad) + ) + pts_px = pts_px[keep] + weights = None + else: + lat, basis = lattice_vectors(lattice, bond) + lat_px = (lat @ R.T) / pixel_size_ang + basis_weights = BASIS_WEIGHTS.get(lattice) + if basis_weights is None: + pts_px = tile_lattice(lat_px, basis, im_shape) + weights = None + else: + # tile each basis atom separately so it keeps its weight + pts_list, wt_list = [], [] + for b_frac, b_w in zip(basis, basis_weights): + pts_b = tile_lattice(lat_px, b_frac[None, :], im_shape) + pts_list.append(pts_b) + wt_list.append(np.full(pts_b.shape[0], b_w)) + pts_px = np.vstack(pts_list) + weights = np.concatenate(wt_list) + + data.append(render_atoms(pts_px, im_shape, atom_sigma_px, weights=weights)) + info.append({ + "lattice": lattice, + "bond_length_ang": bond, + "rotation_deg": theta, + "pixel_size_ang": pixel_size_ang, + "atom_sigma_ang": atom_sigma_ang, + "num_atoms": pts_px.shape[0], + "basis_weights": BASIS_WEIGHTS.get(lattice), + "positions_ang": pts_px * pixel_size_ang, + }) + + if plot_images: + n_bond = len(bond_length_ang) + n_lat = len(lattices) + fig, axs = plt.subplots( + n_lat, n_bond, + figsize=(3 * n_bond, 3 * n_lat), + squeeze=False, + ) + for i, (im, meta) in enumerate(zip(data, info)): + ax = axs[i // n_bond, i % n_bond] + ax.imshow(im, cmap="gray") + ax.set_title( + f"{meta['lattice']}, d = {meta['bond_length_ang']:.2f} $\\AA$", + fontsize=9, + ) + ax.set_xticks([]) + ax.set_yticks([]) + fig.tight_layout() + plt.show() + + return data, info + + +def cluster_patterns(atoms, n_patterns, n_px=256, q_max=2.5, r_cluster=None, + voltage=200e3, seed=0, random_orientation=True, blur_px=0.0): + """ + Kinematic diffraction patterns of spherical clusters cut from a periodic + structure, each with a random centre and a random 3D orientation + (random_clusters, same seed). + + Projected atoms are rendered as sampled Gaussians (sigma 1.5 px, which is + band-limited well inside Nyquist, so crystal reflections do not alias) on + a grid at 2x Nyquist, Fourier transformed and cropped to q_max. The + Gaussian is divided out and replaced by the Weickenmeier-Kohl electron + scattering factor of each element, so patterns follow f(q). blur_px is the sigma (pixels) of a + Gaussian blur of each intensity pattern, which approximates the partial + coherence of the illumination. + + Returns + ------- + Dataset4dstem of shape (n_patterns, 1, n_px, n_px), calibrated in 1/A, + with q = 0 at pixel (n_px // 2, n_px // 2). + """ + from quantem.core.datastructures import Dataset4dstem + from quantem.diffraction.wk_scattering_factors import compute_WK_factor + + pix = 1.0 / (4.0 * q_max) + n_big = 2 * n_px + dq = 1.0 / (n_big * pix) + sigma, half = 1.5, 6 # Gaussian atom width and patch half-width (px) + r_cluster = max_cluster_radius(atoms) if r_cluster is None else min(r_cluster, max_cluster_radius(atoms)) + if r_cluster + (half + 1) * pix > 0.5 * n_big * pix: + raise ValueError( + f"r_cluster = {r_cluster:.1f} A does not fit the render grid of half-width " + f"{0.5 * n_big * pix:.1f} A = n_px / (4 q_max); increase n_px or reduce r_cluster." + ) + qx = (torch.arange(n_px, dtype=torch.float64) - n_px // 2) * dq + QX, QY = torch.meshgrid(qx, qx, indexing="ij") + Q = torch.hypot(QX, QY) + gaussian = torch.exp(-2 * np.pi**2 * (sigma * pix) ** 2 * Q**2) + crop = slice(n_big // 2 - n_px // 2, n_big // 2 + n_px // 2) + species = np.unique(atoms.numbers) + filters = { + z: torch.from_numpy(compute_WK_factor(Q.numpy().ravel(), int(z), voltage).real).reshape(Q.shape) / gaussian + for z in species + } + + offsets = torch.arange(-half, half + 1) + n_pad = n_big + 2 * half + 2 + if blur_px > 0: # real-space Gaussian kernel, positive, so no ringing from the bright central beam + x = torch.arange(-int(np.ceil(4 * blur_px)), int(np.ceil(4 * blur_px)) + 1, dtype=torch.float64) + kernel = torch.exp(-(x**2) / (2 * blur_px**2)) + kernel = (kernel / kernel.sum())[None, None, :] + patterns = torch.empty((n_patterns, n_px, n_px), dtype=torch.float32) + for k, (pos, z_atoms) in enumerate(random_clusters(atoms, n_patterns, r_cluster, seed, random_orientation)): + xy = torch.as_tensor(pos[:, :2] / pix + n_big / 2) + psi = 0 + for z in species: + p = xy[torch.as_tensor(z_atoms == z)] + i0 = torch.round(p).long() + rows = i0[:, 0, None] + offsets[None, :] # (Natoms, patch) + cols = i0[:, 1, None] + offsets[None, :] + wr = torch.exp(-((rows - p[:, 0, None]) ** 2) / (2 * sigma**2)) + wc = torch.exp(-((cols - p[:, 1, None]) ** 2) / (2 * sigma**2)) + w = (wr[:, :, None] * wc[:, None, :]).reshape(-1) / (2 * np.pi * sigma**2) + idx = ((rows[:, :, None] + half + 1) * n_pad + cols[:, None, :] + half + 1).reshape(-1) + im = torch.zeros(n_pad * n_pad, dtype=torch.float64).index_add_(0, idx, w) + im = im.reshape(n_pad, n_pad)[half + 1:half + 1 + n_big, half + 1:half + 1 + n_big] + psi = psi + torch.fft.fftshift(torch.fft.fft2(im))[crop, crop] * filters[z] + intensity = torch.abs(psi) ** 2 + if blur_px > 0: + pad = kernel.shape[-1] // 2 + intensity = torch.nn.functional.conv1d(intensity[:, None, :], kernel, padding=pad)[:, 0, :] + intensity = torch.nn.functional.conv1d(intensity.T[:, None, :], kernel, padding=pad)[:, 0, :].T + patterns[k] = intensity + return Dataset4dstem.from_array( + patterns.numpy()[:, None], + origin=(0, 0, -(n_px // 2) * dq, -(n_px // 2) * dq), + sampling=(1, 1, dq, dq), + units=["pixels", "pixels", "A^-1", "A^-1"], + signal_units="intensity", + ) + + +def max_cluster_radius(atoms): + """Half the smallest distance between opposite faces of the periodic cell: + the largest sphere that contains no periodic image twice.""" + cell = np.asarray(atoms.cell) + volume = abs(np.linalg.det(cell)) + heights = [volume / np.linalg.norm(np.cross(cell[(i + 1) % 3], cell[(i + 2) % 3])) for i in range(3)] + return 0.5 * min(heights) + + +def random_clusters(atoms, n_clusters, r_cluster=None, seed=0, random_orientation=True): + """ + Yield (positions, atomic numbers) of spherical clusters cut from a periodic + structure, each with a random centre and rotated to a random orientation + (the beam is along z). With random_orientation=False the cell axes are + kept, so the beam is along the cell's z axis. The same seed gives the + clusters of cluster_patterns. r_cluster is capped at max_cluster_radius, + so no atom appears twice; any cell shape works. + """ + cell = torch.as_tensor(np.asarray(atoms.cell), dtype=torch.float64) # rows are the cell vectors + frac0 = torch.as_tensor(atoms.get_scaled_positions(wrap=True), dtype=torch.float64) + r_max = max_cluster_radius(atoms) + r_cluster = r_max if r_cluster is None else min(r_cluster, r_max) + gen = torch.Generator().manual_seed(seed) + for _ in range(n_clusters): + # uniform random rotation from a normalized Gaussian quaternion + w, x, y, z = torch.nn.functional.normalize(torch.randn(4, generator=gen, dtype=torch.float64), dim=0) + rot = torch.stack([ + torch.stack([1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)]), + torch.stack([2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)]), + torch.stack([2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)]), + ]) + # minimum image about a random centre, in fractional coordinates: within r_max of + # the centre every fractional offset is at most 1/2, so rounding picks the right image + frac = frac0 - torch.rand(3, generator=gen, dtype=torch.float64) + d = (frac - torch.round(frac)) @ cell + keep = torch.linalg.norm(d, dim=1) <= r_cluster + pos = d[keep] @ rot.T if random_orientation else d[keep] + yield pos.numpy(), atoms.numbers[keep.numpy()] diff --git a/src/quantem/diffraction/multislice.py b/src/quantem/diffraction/multislice.py new file mode 100644 index 00000000..2a60cf32 --- /dev/null +++ b/src/quantem/diffraction/multislice.py @@ -0,0 +1,109 @@ +""" +Repeat-slice multislice propagation for the PADF test lattices. + +Treats a 2D lattice image as the projected potential of one slice and +propagates a plane wave through n_slices identical copies (AA stacking) - +a minimal dynamical-diffraction model of a column structure. With a +NON-centrosymmetric projection (e.g. the boron_nitride lattice), multiple +scattering violates Friedel symmetry: odd angular harmonics appear in the +diffraction rings and grow with thickness. Centrosymmetric projections +(triangular, honeycomb, square) keep I(q) = I(-q) at any thickness. + +For quantitative work use abTEM with a real 3D structure; this module is +for fast, controlled tests of how dynamical scattering perturbs +correlation / PADF / bispectrum observables. +""" + +import numpy as np + + +def circular_window(shape): + """ + Radially symmetric Hann window: w = cos^2(pi rho / 2R) with R half the + smaller image dimension, zero outside. Unlike the separable 2D Hann + (outer product of 1D windows), its FFT skirt is isotropic, so the + windowed central beam carries no spurious 4-fold angular structure + into the correlations. + """ + n_row, n_col = shape + row = np.arange(n_row) - (n_row - 1) / 2 + col = np.arange(n_col) - (n_col - 1) / 2 + rho = np.hypot(row[:, None], col[None, :]) + R = min(n_row, n_col) / 2.0 + w = np.cos(0.5 * np.pi * np.minimum(rho / R, 1.0)) ** 2 + return w + + +def electron_wavelength(voltage_kv: float) -> float: + """Relativistic electron wavelength in Angstroms.""" + v = voltage_kv * 1e3 + h = 6.62607015e-34 + m0 = 9.1093837015e-31 + e = 1.602176634e-19 + c = 2.99792458e8 + lam = h / np.sqrt(2 * m0 * e * v * (1 + e * v / (2 * m0 * c**2))) + return lam * 1e10 + + +def multislice( + potential, + pixel_size_ang: float, + n_slices: int, + slice_thickness_ang: float = 3.33, + voltage_kv: float = 80.0, + phase_per_slice: float = 0.15, +): + """ + Propagate a plane wave through n_slices copies of the same 2D potential. + + Parameters + ---------- + potential : (N, N) array + Projected potential of one slice (arbitrary units; normalized + internally so its maximum produces `phase_per_slice` radians). + pixel_size_ang : float + Real-space pixel size in Angstroms. + n_slices : int + Number of identical slices (thickness = n_slices * slice_thickness). + slice_thickness_ang : float + Slice spacing in Angstroms (3.33 = graphite/h-BN interlayer). + voltage_kv : float + Accelerating voltage in kV (sets the Fresnel propagator). + phase_per_slice : float + Peak phase shift (radians) of one slice at the strongest atom. + n_slices = 1 with a small value approaches the kinematic limit. + + Returns + ------- + psi : (N, N) complex exit wave. + """ + pot = np.asarray(potential, dtype=np.float64) + n_row, n_col = pot.shape + lam = electron_wavelength(voltage_kv) + + transmission = np.exp(1j * phase_per_slice * pot / pot.max()) + + q_row = np.fft.fftfreq(n_row, d=pixel_size_ang) + q_col = np.fft.fftfreq(n_col, d=pixel_size_ang) + q2 = q_row[:, None] ** 2 + q_col[None, :] ** 2 + propagator = np.exp(-1j * np.pi * lam * slice_thickness_ang * q2) + # standard 2/3 bandwidth limit to prevent aliasing of scattered orders + q_max = min(np.abs(q_row).max(), np.abs(q_col).max()) + propagator *= q2 <= (2.0 / 3.0 * q_max) ** 2 + + psi = np.ones((n_row, n_col), dtype=np.complex128) + for _ in range(n_slices): + psi = np.fft.ifft2(np.fft.fft2(psi * transmission) * propagator) + return psi + + +def diffraction_intensity(psi, window: bool = True): + """ + Diffraction intensity of an exit wave: |FFT(psi * w)|^2, fftshifted, + with a circular (radial Hann) window so the finite-field skirt is + isotropic. The DC (unscattered beam) still dominates; downstream + analysis should treat it like the experimental central beam. + """ + if window: + psi = psi * circular_window(psi.shape) + return np.abs(np.fft.fftshift(np.fft.fft2(psi))) ** 2 diff --git a/src/quantem/diffraction/padf.py b/src/quantem/diffraction/padf.py new file mode 100644 index 00000000..0875bf98 --- /dev/null +++ b/src/quantem/diffraction/padf.py @@ -0,0 +1,728 @@ +import matplotlib.pyplot as plt +import torch +import numpy as np +from numpy.typing import NDArray +from scipy.special import spherical_jn + +from quantem.core.datastructures.dataset4dstem import Dataset4dstem +from quantem.core.datastructures.polar4dstem import Polar4dstem +from quantem.core.io.serialize import AutoSerialize +from quantem.diffraction.polar_transform import ( + as_dataset4dstem, + find_origin_angular_grid, + polar_transform, +) + +class PairAngleDistributionFunction(AutoSerialize): + """ + Compute the pair-angle distribution function from diffraction data: + a single 2D pattern (Dataset2d), a 3D stack of patterns, or a 4D-STEM + dataset. Non-4D inputs are wrapped as (1, 1, H, W) / (N, 1, H, W). + + Run the following pipeline: + padf = PairAngleDistributionFunction(ds) + 1. find_origin_angular_grid(ds) + 2. polar_transform(ds, origin_array=self.origin) + 3. rescale_intensity(self.polar, rho=1, fq=1) + 4. compute_avg_angular_correlation(rescaled intensity) + 5. subtract_angular_median(self.ang_corr) + 6. extract_Bl_matrices(self.ang_corr_corrected) + 7. transform_to_real_space(self.Bl_mats[0], self.Bl_mats[1]) + 8. reconstruct_PADF(self.real_Bl_mats, self.Bl_mats[1], Na=1); the overall scale is arbitrary + + polar_ds can instead be a polar dataset, or a list of them with the same + q axis, in which case steps 1-2 are skipped and the angular correlation is + averaged over the list (useful for more patterns than fit in memory). + fq may be a function of q (1/A), evaluated on the polar q axis self.q. + The real-space grid r = r_min, r_min + r_step, ... includes r_max. + """ + def __init__(self, + ds=None, + origin: NDArray | None = None, + polar_ds: Polar4dstem | list[Polar4dstem] | None = None, + ang_corr: torch.Tensor | None = None, + Bl_mats = None, + rho: float = 1.0, + fq=None, + l_max: int = 40, + r_min: float = 0.0, + r_max: float = 20.0, + r_step: float = 0.02, + l_taper: bool = False, + ): + super().__init__() + + # accept a single 2D pattern, a 3D stack, or a 4D-STEM dataset + self.ds = as_dataset4dstem(ds) if ds is not None else None + if polar_ds is None: + self.origin = origin if origin is not None else find_origin_angular_grid(self.ds) + polar_ds = polar_transform(self.ds, origin_array=self.origin) + else: + self.origin = origin + polars = list(polar_ds) if isinstance(polar_ds, (list, tuple)) else [polar_ds] + self.polar = polars[0] + # radial q calibration of the polar dataset: bin i sits at q_min + i * dq + dq = float(np.asarray(self.polar.sampling)[3]) + q_min = float(np.asarray(self.polar.origin)[3]) + self.q = q_min + np.arange(self.polar.shape[-1]) * dq + if callable(fq): + fq = fq(self.q) + + # TODO: look up rho and fq from sample identity when not provided + if ang_corr is None: + ang_corr = sum( + self.compute_avg_angular_correlation(self.rescale_intensity(p, rho=rho, fq=fq)) for p in polars + ) / len(polars) + self.ang_corr = ang_corr + self.ang_corr_corrected = self.subtract_angular_median(self.ang_corr) + self.Bl_mats = Bl_mats if Bl_mats is not None else self.extract_Bl_matrices(self.ang_corr_corrected, l_max=l_max) + self.real_Bl_mats = self.transform_to_real_space(self.Bl_mats[0], self.Bl_mats[1], dq=dq, q_min=q_min, + r_min=r_min, r_max=r_max + 0.5 * r_step, r_step=r_step) + self.padf = self.reconstruct_PADF(self.real_Bl_mats, self.Bl_mats[1], 1, l_taper=l_taper) + + @staticmethod + def rescale_intensity(data: Polar4dstem | Dataset4dstem | None = None, + rho: float | int = 1.0, + fq: NDArray | torch.Tensor | float | int | None = None): + """ + Equation 7: Take the raw intensity and divide by the following factors + - Mean number density (rho) + - Mean atomic scattering factor squared (f(q)^2) + - Number of atoms in the beam (N_a) + - phi_0 (dependent on experimental parameters) + + fq may be a scalar or a per-radial-bin array f(q) of length Nq + (broadcast along the last axis of the polar data). Dividing by + f(q)^2 removes the atom-shape damping of the diffraction rings; + leaving it in biases the real-space peak positions of the PADF. + Floor small f(q) values before passing to avoid amplifying the + high-q noise floor. + + Finally the intensity is normalized to unit mean. This is the + analogue of pypadf's scalar 1/beamnorm^2 correlation scaling: it + absorbs the unmeasurable overall factors (phi_0, N_a, detector + gain) so the PADF magnitude is independent of the raw intensity + scale. + + Handles dtype conversion to float64 + """ + + intensity = data.tensor.to(dtype=torch.float64) + if fq is None: + fq2 = 1.0 + elif np.ndim(fq) == 0: + fq2 = float(fq) ** 2 + else: + fq_t = torch.as_tensor(np.asarray(fq), dtype=torch.float64) + if fq_t.shape[-1] != intensity.shape[-1]: + raise ValueError( + f"fq has length {fq_t.shape[-1]} but the polar data has " + f"{intensity.shape[-1]} radial bins." + ) + fq2 = fq_t ** 2 # broadcasts along the radial (last) axis + rescaled_intensity = intensity / (rho * fq2) + rescaled_intensity /= rescaled_intensity.mean() + return rescaled_intensity + + @staticmethod + def compute_avg_angular_correlation(data): + """ + Equation 8: angular cross correlation implemented via Fourier correlation theorem + """ + + scan_row, scan_col, phi, r = data.shape + g_all = data.reshape(scan_row * scan_col, phi, r) # Flatten grid of dps to one dimension, makes for loop simpler + N_alpha = g_all.shape[0] # Number of dps we have = number of configurations (N_alpha) + + dphi = 2.0 * torch.pi / phi # A small change in angle + + G = torch.fft.rfft(g_all, dim=1) # Fourier transform along phi axis + P = torch.einsum('afi,afj->fij', G.conj(), G) # Outer product of r between G and its conjugate, summed over all N_alpha + C_sum = torch.fft.irfft(P, n=phi, dim=0) + + C_avg = (C_sum / N_alpha) * dphi + return C_avg + + def subtract_angular_mean(self, C_avg: torch.Tensor): + """ + Subtract the theta-average from each (q, q') ring of the angular + correlation, removing the uncorrelated background. + + Matches the subtraction in pypadf's maskcorr.py. We do not apply their |sin(theta)| + factor here because it changes the function being expanded and mixes Legendre orders + before the order-dependent Bessel transforms. Instead, we sin-weight the Legendre + fit in extract_Bl_matrices() to match the d(cos theta) measure of the Legendre + polynomials. + """ + return C_avg - C_avg.mean(dim=0, keepdim=True) + + @staticmethod + def subtract_angular_median(C_avg: torch.Tensor): + """ + Subtract the theta-median from each (q, q') ring of the angular + correlation. Unlike the mean, the median is not pulled up by the + self-correlation peak at theta = 0. The offset only changes B_0, so + the PADF (l >= 2) is the same as with subtract_angular_mean. + """ + return C_avg - C_avg.median(dim=0, keepdim=True).values + + def extract_Bl_matrices(self, C_avg, l_max=40, sv_cutoff=0.05, sin_weight=True, + theta_exclude=0.0): + """ + sin_weight : bool + Weight each angle in the least-squares fit by |sin(theta)|. C is + sampled evenly in theta, but Legendre polynomials are orthogonal in + cos(theta), and near 0 and 180 deg many samples crowd into a small + range of cos(theta). Since |sin(theta)| d(theta) = |d cos(theta)|, + this weight counts each sample by the cos(theta) range it covers. + theta_exclude : float + In degrees, correlation bins within theta_exclude of 0 and 360 deg get zero + weight in the fit. Use to exclude self correlations and keeps them out of B_l. + """ + # Equation 9, 10, 11 + # for each l: + # build the linear system relating C to B_l(q, q') + # solve it by SVD + # even l only: a single kinematic pattern is Friedel symmetric + # 5% SVD cutoff + Nphi, Nq, Nqp = C_avg.shape + dphi = torch.linspace(0, 2 * torch.pi, Nphi + 1, dtype=torch.float64)[:-1] # Ranges from [0, 2pi) + l_values = torch.arange(0, l_max, 2) # Even only + Nl = len(l_values) + cos_dphi = torch.cos(dphi) + + C_flat = C_avg.reshape(Nphi, Nq * Nqp) + + # Build matrix + L = l_values.reshape(1, Nl) + X = cos_dphi.reshape(Nphi, 1) + leg_matrix = torch.special.legendre_polynomial_p(X, L.to(torch.float64)) # (Nphi, Nl) by broadcasting + + # per-angle fit weight: sqrt|sin| (sin_weight), zero within theta_exclude of 0/360 deg + w = torch.sqrt(torch.abs(torch.sin(dphi))) if sin_weight else torch.ones_like(dphi) + if theta_exclude > 0: + te = np.deg2rad(theta_exclude) + w[(dphi < te) | (dphi > 2 * torch.pi - te)] = 0 + leg_matrix = leg_matrix * w[:, None] + C_flat = C_flat * w[:, None] + + U, S, Vh = torch.linalg.svd(leg_matrix, full_matrices=False) + S_max = S.max() + S_inv = torch.where(S > sv_cutoff * S_max, 1.0 / S, torch.zeros_like(S)) + + tmp = U.T @ C_flat + tmp = S_inv.unsqueeze(1) * tmp + B_flat = Vh.T @ tmp + + return B_flat.reshape(Nl, Nq, Nqp), l_values + + def transform_to_real_space(self, Bl_mats, l_values, dq=0.01, q_min=0.0, + r_min=0.0, r_max=20.0, r_step=0.02): + """ + Transform B_l(q, q') to B_l(r, r') by applying the spherical Bessel + transform (eq 8) along each q axis via direct quadrature. + + dq : radial q step of the polar dataset (e.g. angstrom^-1 / bin). + q_min : q value of the first radial bin (nonzero if the polar + transform used radial_min > 0). + r_min, r_max, r_step : real-space grid, in the reciprocal units of q. + + The transform kernel is nondimensionalized by q_max (the quadrature + uses (q / q_max)^2 d(q / q_max) per axis), so the output B_l(r, r') + is dimensionless rather than carrying units of q^6. Combined with + the unit-mean intensity normalization in rescale_intensity, this + keeps PADF values of order 1. + """ + # Equation 12, 13 + # apply bessel transform twice for each l + # → shape = (l, r, r') + + # Step one is to define q + q = q_min + torch.arange(0, Bl_mats.shape[1], dtype=torch.float64) * dq + q_max = float(q[-1]) + r = torch.arange(r_min, r_max, r_step, dtype=torch.float64) + self.r = r.numpy() + real_Bl = torch.zeros((Bl_mats.shape[0], r.shape[0], r.shape[0]), dtype=torch.float64) + + arg = 2 * torch.pi * torch.outer(r, q) # Shape len(r) x len(q) + for l in range(len(l_values)): + Bl = Bl_mats[l].to(dtype=torch.float64) # The corresponding q x q' matrix + jl = torch.from_numpy(spherical_jn(l_values[l], arg)).to(dtype=torch.float64) + + # Representing DSBT (eq 12) as a transformation matrix "sbessel", + # with q scaled by q_max so the quadrature is dimensionless + sbessel = 4 * torch.pi * jl * (q**2) * dq / q_max**3 + # Applying it twice (once along each axis) + real_Bl[l] = sbessel @ Bl @ sbessel.T * (-1)**l_values[l] + + return real_Bl + + def reconstruct_PADF(self, real_Bl, l_values, Na, per_theta=True, l_taper=False): + """ + Reconstruct the pair-angle distribution function from B_l(r, r') + For each theta: + - Sum Pl(cos theta) times B_l() over all l + - Multiply by n_alpha * 2 pi + + Theta goes from 0 to pi in 180 steps. Na scales the result; with the + unit-mean intensity normalization the overall scale is arbitrary. + + per_theta : bool + The Legendre sum is the real-space correlation C(r, r', cos theta) + (pypadf paper eq. 10). If True, multiply by |sin theta| to return a + quantity proportional to Theta(r, r', theta) (pypadf paper eq. 11). + l_taper : bool + Weight each order by order_taper(l, l_max), which suppresses the + ringing in theta from truncating the Legendre sum at l_max. + """ + padf = torch.zeros((real_Bl.shape[1], real_Bl.shape[2], 180), dtype=torch.float64) + l_values = np.asarray(l_values).astype(int) + weights = order_taper(l_values, l_values[-1] + 2) if l_taper else np.ones(len(l_values)) + theta = np.linspace(0, np.pi, 180) + self.theta = theta + self.theta_deg = np.degrees(theta) + cos_theta = torch.cos(torch.as_tensor(theta, dtype=torch.float64)) + chosen_l_values = l_values[1:] + + # Sum over all l + for l in range(len(chosen_l_values)): # skipping l = 0 + Pl = torch.special.legendre_polynomial_p(cos_theta, int(chosen_l_values[l])) + Bl = real_Bl[l + 1] * weights[l + 1] # skipping l = 0 + padf += Bl[:, :, np.newaxis] * Pl[np.newaxis, np.newaxis, :] + + padf *= 2 * torch.pi * Na + if per_theta: + padf *= torch.as_tensor(np.abs(np.sin(theta)), dtype=torch.float64) + return padf + + def plot_g2_g3( + self, + r_min: float | None = None, + r_max: float | None = None, + r_display_power: int = 1, + r_max_display: float | None = None, + r_search_min: float = 0.5, + r_marks=None, + markers=None, + title: str | None = None, + figsize: tuple[float, float] = (4.8, 4.4), + returnfig: bool = False, + ): + """ + Stacked 2-body / 3-body correlation summary of the PADF, after + atomode's G3 explorer. The two panels share the radial axis. + + Top panel: 3-body correlation map - Theta(r, r', theta) integrated + over r in [r_min, r_max] (symmetrized over both radial axes), + plotted as angle theta vs r'. + + Bottom panel: 2-body correlation profile - the |sin(theta)|-weighted + RMS of the r = r' diagonal over the interior angular range + (15-165 deg), + sqrt( sum_theta Theta(r, r, theta)^2 sin(theta) + / sum_theta sin(theta) ). + This stands in for the isotropic pair correlation (whose true l = 0 + component is removed by the correlation mean subtraction; a plain + sin-weighted integral of Theta cancels to ~zero by Legendre + orthogonality). It peaks at the radial shells where the PADF has + angular structure and is used to choose the near-neighbor shell; + the integration shell [r_min, r_max] is shaded. + + Both panels are weighted by r^(2 * r_display_power) for display, and + the map's color limits are taken from the interior angular range + (15-165 deg) so the collinear theta = 0 / 180 bands saturate rather + than compressing the color scale. + + Parameters + ---------- + r_min, r_max : float or None + Integration shell bounds in the r units of the PADF (Angstroms + for calibrated data). If None, the first peak of g2 is selected + automatically: from the last point below 5% of the peak height + up to the first local minimum past the peak (atomode's default + shell heuristic). + r_display_power : int + Display weighting exponent; each panel is multiplied by + r^(2 * r_display_power). 0 plots the raw values. + r_max_display : float or None + Upper limit of the radial axes; defaults to r.max(). + r_search_min : float + Ignore r below this value when auto-locating the g2 peak + (excludes the r ~ 0 reconstruction artifacts). + r_marks : sequence of float or None + Radii (e.g. the 1st/2nd/3rd neighbor shell distances) marked + with dotted vertical lines on both panels. + markers : sequence of (r, theta_deg) or None + Expected 3-body maxima overlaid on the map as open circles - + e.g. the (shell distance, arm-arm angle) targets from a known + structure. A third element per tuple, if present, scales the + marker size (relative multiplicity). + title : str or None + Figure title; a shell summary is used if None. + returnfig : bool + Return (fig, (ax_g2, ax_g3)) instead of calling plt.show(). + """ + r = np.asarray(self.r) + theta = np.asarray(self.theta) # radians + theta_deg = np.asarray(self.theta_deg) + arr = self.padf.numpy() if hasattr(self.padf, "numpy") else np.asarray(self.padf) + + # pair profile: sin(theta)-weighted RMS of the diagonal over the + # interior angles (a plain integral cancels by Legendre orthogonality + # since the l = 0 term is removed by the mean subtraction) + diag = np.einsum("iik->ik", arr) + interior_t = (theta_deg > 15) & (theta_deg < 165) + w_t = np.abs(np.sin(theta[interior_t])) + g2 = np.sqrt((diag[:, interior_t] ** 2 * w_t[None, :]).sum(axis=1) / w_t.sum()) + + # auto-select the first g2 peak if no shell was given + if r_min is None or r_max is None: + kernel = np.array([1.0, 2.0, 3.0, 2.0, 1.0]) + kernel /= kernel.sum() + smooth = np.convolve(np.where(r < r_search_min, 0.0, g2), kernel, mode="same") + i_pk = None + for idx in range(1, smooth.size - 1): + if smooth[idx] > 0 and smooth[idx] >= smooth[idx - 1] and smooth[idx] > smooth[idx + 1]: + i_pk = idx + break + if i_pk is None: + i_pk = int(np.argmax(smooth)) + below = np.flatnonzero(smooth[:i_pk] < 0.05 * smooth[i_pk]) + i_lo = int(below[-1]) if below.size else int(np.searchsorted(r, r_search_min)) + i_hi = min(smooth.size - 1, 2 * i_pk - i_lo) + for idx in range(i_pk + 1, smooth.size - 1): + if smooth[idx] <= smooth[idx - 1] and smooth[idx] <= smooth[idx + 1]: + i_hi = idx + break + if r_min is None: + r_min = float(r[i_lo]) + if r_max is None: + r_max = float(r[i_hi]) + + shell = (r >= r_min) & (r <= r_max) + if not np.any(shell): + shell[int(np.argmin(np.abs(r - 0.5 * (r_min + r_max))))] = True + + # g3 map: integrate over the shell, symmetrized over both r axes + g3_map = 0.5 * (arr[shell].sum(axis=0) + arr[:, shell].sum(axis=1)) # (Nr', Ntheta) + + # display weighting and color limits + if r_max_display is None: + r_max_display = float(r[-1]) + weight = r ** (2 * r_display_power) + g2_disp = g2 * weight + g3_disp = (g3_map * weight[:, None]).T # (Ntheta, Nr') + rs = r <= r_max_display + interior = (theta_deg > 15) & (theta_deg < 165) + vmax = np.quantile(np.abs(g3_disp[np.ix_(interior, rs)]), 0.999) + + fig, (ax_g3, ax_g2) = plt.subplots( + 2, 1, sharex=True, figsize=figsize, + gridspec_kw={"height_ratios": [2, 1]}, + constrained_layout=True, + ) + im = ax_g3.pcolormesh( + r[rs], theta_deg, g3_disp[:, rs], + cmap="RdBu_r", vmin=-vmax, vmax=vmax, shading="auto", + ) + ax_g3.set_ylabel("$\\theta$ (deg)") + ax_g3.set_ylim(0, 180) + ax_g3.set_yticks(np.arange(0, 181, 30)) + fig.colorbar(im, ax=[ax_g3, ax_g2], label="3-body correlation", pad=0.02) + + if markers is not None: + for mk in markers: + r_mk, t_mk = mk[0], mk[1] + size = 12.0 * (mk[2] ** 0.5) if len(mk) > 2 else 10.0 + if r_mk <= r_max_display: + ax_g3.plot(r_mk, t_mk, "o", ms=max(size, 5.0), mfc="none", + mec="k", mew=1.0) + + ax_g2.plot(r[rs], g2_disp[rs], "k-", lw=1) + ax_g2.axvspan(r_min, r_max, color="C1", alpha=0.25) + ax_g2.axhline(0, color="0.85", lw=0.8, zorder=0) + + if r_marks is not None: + for r_mk in r_marks: + for ax in (ax_g3, ax_g2): + ax.axvline(r_mk, color="k", ls=":", lw=0.8) + ax_g2.set_ylabel("2-body corr.") + ax_g2.set_xlabel("r, r' (Å)") + + ax_g3.set_title( + title if title is not None + else f"shell r = {r_min:.2f} - {r_max:.2f}", + fontsize=10, + ) + + if returnfig: + return fig, (ax_g3, ax_g2) + plt.show() + + def simple_plot( + self, + padf=None, + r_display_power: int = 1, + r_max_display: float | None = None, + r_fixed=None, + theta_marks=None, + r_marks=None, + title=None, + axsize: tuple[float, float] = (4.5, 3.5), + returnfig: bool = False, + ): + """ + Maps of the PADF against theta (y-axis, 0 to 180 deg) and r (x-axis), + one panel per volume in a single row: the r = r' diagonal + Theta(r, r, theta), or Theta(r_fixed, r', theta) when r_fixed is given. + The display is weighted by (r r')^r_display_power. + + padf : (Nr, Nr, Ntheta) array, or a list of them, or None + Volumes on the same (r, theta) grid, e.g. a model PADF; defaults to self.padf. + r_fixed : float, list of float (one per panel) or None + Fixed first radius (Angstroms) for each panel. + theta_marks, r_marks : sequence of float, list of them (one per panel) or None + Angles (degrees) and radii (Angstroms) marked with dotted lines. + title : str or list of str (one per panel) + Each panel has symmetric color limits from its largest |value| at + 15-165 deg and r <= r_max_display. + """ + vols = [self.padf] if padf is None else (list(padf) if isinstance(padf, (list, tuple)) else [padf]) + n = len(vols) + titles = _per_panel(title, n) + r_fixeds = list(r_fixed) if isinstance(r_fixed, (list, tuple)) else [r_fixed] * n + theta_markss = _per_panel(theta_marks, n, nested=True) + r_markss = _per_panel(r_marks, n, nested=True) + r = np.asarray(self.r) + theta_deg = np.asarray(self.theta_deg) + rs = r <= (r.max() if r_max_display is None else r_max_display) + interior = (theta_deg > 15) & (theta_deg < 165) + + fig, axs = plt.subplots(1, n, figsize=(axsize[0] * n, axsize[1]), squeeze=False, constrained_layout=True) + for ax, vol, t, r0, t_marks, rm_marks in zip(axs[0], vols, titles, r_fixeds, theta_markss, r_markss): + vol = vol.numpy() if isinstance(vol, torch.Tensor) else np.asarray(vol) + if r0 is None: + m = np.einsum("iik->ik", vol)[rs] * (r[rs] ** (2 * r_display_power))[:, None] + xlabel, label = "r = r' (Å)", rf"$\Theta(r, r, \theta) \cdot r^{{{2 * r_display_power}}}$" + else: + i0 = int(np.argmin(np.abs(r - r0))) + m = vol[i0][rs] * ((r[i0] * r[rs]) ** r_display_power)[:, None] + xlabel, label = f"r' (Å), r = {r[i0]:.2f} Å", rf"$\Theta(r, r', \theta) \cdot (r r')^{{{r_display_power}}}$" + vmax = np.abs(m[:, interior]).max() + im = ax.pcolormesh(r[rs], theta_deg, m.T, cmap="RdBu_r", vmin=-vmax, vmax=vmax, + shading="auto") + for th in t_marks or []: + ax.axhline(th, color="k", ls=":", lw=1) + for rm in rm_marks or []: + ax.axvline(rm, color="k", ls=":", lw=1) + ax.set_xlabel(xlabel) + ax.set_ylabel(r"$\theta$ (deg)") + ax.set_ylim(0, 180) + ax.set_yticks(np.arange(0, 181, 30)) + fig.colorbar(im, ax=ax, label=label) + if t is not None: + ax.set_title(t) + if returnfig: + return fig, axs[0] + plt.show() + + def plot_angular_correlation( + self, + q: float, + figsize: tuple[float, float] = (5, 3.5), + returnfig: bool = False, + ): + """ + Median-subtracted angular correlation C(dphi, q, q') at the radial bin + nearest q (1/Angstroms), against q' (x-axis) and dphi (y-axis). + """ + C = self.ang_corr_corrected.numpy() + q_axis = self.q + iq = int(np.argmin(np.abs(q_axis - q))) + dphi = np.arange(C.shape[0]) * 360.0 / C.shape[0] + M = C[:, iq, :] + vmax = np.quantile(np.abs(M), 0.995) + fig, ax = plt.subplots(figsize=figsize, constrained_layout=True) + im = ax.pcolormesh(q_axis, dphi, M, cmap="RdBu_r", vmin=-vmax, vmax=vmax, shading="auto") + ax.set_xlabel("q' (1/Å)") + ax.set_ylabel(r"$\Delta\phi$ (deg)") + ax.set_ylim(0, 360) + ax.set_yticks(np.arange(0, 361, 60)) + ax.set_title(f"q = {q_axis[iq]:.3f} 1/Å") + fig.colorbar(im, ax=ax, label=r"$C(\Delta\phi, q, q')$") + if returnfig: + return fig, ax + plt.show() + + +# Band-limit a model g3(r, r', theta) to what a PADF measurement can resolve +# (pypadf blfilter.py approach): expand in P_l(cos theta) for even l < l_max, +# take each order to the measured q points with the order-l spherical Bessel +# transform, and transform back. +def bandlimit_volume( + G: NDArray, + r: NDArray, + theta_deg: NDArray, + q: NDArray, + l_max: int = 40, +) -> NDArray: + """ + Parameters + ---------- + G : (Nr, Nr, Ntheta) array + Volume on a uniform r grid, theta in [0, 180] degrees. + r : (Nr,) array + Radial grid (Angstroms), uniformly spaced. + theta_deg : (Ntheta,) array + Angles in degrees. + q : (Nq,) array + Measured q values (1/Angstroms, no 2 pi), uniformly spaced. + l_max : int + Even orders l = 0, 2, ..., < l_max are kept (same as extract_Bl_matrices). + + Returns + ------- + (Nr, Nr, Ntheta) array, the band-limited volume. + """ + G = torch.as_tensor(np.asarray(G), dtype=torch.float64) + theta = torch.deg2rad(torch.as_tensor(np.asarray(theta_deg), dtype=torch.float64)) + l_values = torch.arange(0, l_max, 2, dtype=torch.float64) + + # Legendre projection over theta with the sin(theta) d(theta) measure + w = torch.abs(torch.sin(theta)) * torch.gradient(theta)[0] + P = torch.special.legendre_polynomial_p(torch.cos(theta)[None, :], l_values[:, None]) # (Nl, Ntheta) + Gl = torch.einsum("ijt,lt->lij", G, P * w[None, :]) * (2 * l_values[:, None, None] + 1) / 2 + + out = torch.zeros_like(G) + for k, l in enumerate(l_values.int().tolist()): + out += _bandlimit_order(Gl[k], l, r, q)[:, :, None] * P[k][None, None, :] + return out.numpy() + + +def order_taper(orders, order_max): + """cos^2 taper of harmonic orders, 1 at order 0 and 0 at order_max.""" + return np.cos(0.5 * np.pi * np.asarray(orders, dtype=np.float64) / order_max) ** 2 + + +def _bandlimit_order(Gl: torch.Tensor, l: int, r, q) -> torch.Tensor: + """Order-l spherical Bessel transform of Gl(r, r') to the q points and back.""" + r = np.asarray(r, dtype=np.float64) + q = np.asarray(q, dtype=np.float64) + jl = torch.from_numpy(spherical_jn(l, 2 * np.pi * np.outer(q, r))) # (Nq, Nr), not in torch + fwd = 4 * np.pi * jl * torch.from_numpy(r**2 * (r[1] - r[0]))[None, :] # r -> q + inv = 4 * np.pi * jl.T * torch.from_numpy(q**2 * (q[1] - q[0]))[None, :] # q -> r + return inv @ (fwd @ Gl @ fwd.T) @ inv.T + + +def _real_sph_harm_power(cos_t, sin_t, phi, l_values, weights, n_r, index): + """ + sum_m A_lm(r) conj(A_lm(r')) for each l in l_values, where + A_lm(r) = sum_v weights[v] Y_lm(v) binned to r by index (Nvec,) -> [0, n_r]. + Orthonormal associated Legendre functions by the standard stable + recurrence in l at fixed m; m < 0 enters through |A_l,-m| = |A_lm|. + """ + l_max = int(l_values.max()) + C = torch.zeros((len(l_values), n_r, n_r), dtype=torch.float64) + row = {int(l): k for k, l in enumerate(l_values.tolist())} + p_mm = torch.full_like(cos_t, (1 / (4 * np.pi)) ** 0.5) + for m in range(l_max + 1): + if m > 0: + p_mm = -((2 * m + 1) / (2 * m)) ** 0.5 * sin_t * p_mm + e = torch.polar(torch.ones_like(phi), m * phi) * weights + p_prev, p_l = None, p_mm + for l in range(m, l_max + 1): + if l == m + 1: + p_prev, p_l = p_l, (2 * m + 3) ** 0.5 * cos_t * p_l + elif l > m + 1: + a = ((4 * l * l - 1) / (l * l - m * m)) ** 0.5 + b = (((l - 1) ** 2 - m * m) / (4 * (l - 1) ** 2 - 1)) ** 0.5 + p_prev, p_l = p_l, a * (cos_t * p_l - b * p_prev) + if l in row: + A = torch.zeros(n_r + 1, dtype=torch.complex128).index_add_(0, index, p_l * e)[:n_r] + C[row[l]] += (1.0 if m == 0 else 2.0) * torch.real(A.conj()[:, None] * A[None, :]) + return C + + +def padf_from_clusters( + clusters, + r: NDArray, + theta_deg: NDArray, + q: NDArray | None = None, + l_max: int = 40, + l_taper: bool = False, +) -> NDArray: + """ + PADF of atomic models: the angular correlation of all interatomic vectors + within each cluster, summed over clusters. Like the diffraction PADF it + holds two-, three- and four-body terms (Martin et al. IUCrJ 2017), not + only triplets sharing an atom. + + Uses the addition theorem: for each even l < l_max, + C_l(r, r') = 2 pi sum_m A_lm(r) conj(A_lm(r')), + A_lm(r) = sum over vectors at |d| = r of Y_lm(d / |d|), per unit volume, + then Theta = |sin theta| sum_{l >= 2} C_l(r, r') P_l(cos theta), matching + reconstruct_PADF (l = 0 is dropped, as the PADF's mean subtraction does). + + Parameters + ---------- + clusters : iterable of (N, 3) arrays + Atom positions (Angstroms), one array per cluster. + r : (Nr,) array + Uniform radial grid (Angstroms); vectors are binned with linear weights. + theta_deg : (Ntheta,) array + Angles in degrees. + q : (Nq,) array or None + Measured q values (1/Angstroms, no 2 pi). If given, each C_l is + band-limited to them, as bandlimit_volume does. + l_max : int + Even orders l = 0, 2, ..., < l_max. + l_taper : bool + Weight each order by order_taper(l, l_max), as reconstruct_PADF does. + + Returns + ------- + (Nr, Nr, Ntheta) array. + """ + r_t = torch.as_tensor(np.asarray(r), dtype=torch.float64) + dr = float(r_t[1] - r_t[0]) + n_r = r_t.numel() + l_values = torch.arange(2, l_max, 2) # l = 0 is dropped, as the PADF's median subtraction does + shell = 4 * np.pi * torch.clamp(r_t, min=0.5 * dr) ** 2 * dr + Cl = torch.zeros((len(l_values), n_r, n_r), dtype=torch.float64) + for pos in clusters: + pos = torch.as_tensor(np.asarray(pos), dtype=torch.float64) + d = (pos[None, :, :] - pos[:, None, :]).reshape(-1, 3) + dist = torch.linalg.norm(d, dim=1) + keep = (dist > 0) & (dist >= float(r_t[0])) & (dist < float(r_t[-1])) + d, dist = d[keep], dist[keep] + cos_t = torch.clamp(d[:, 2] / dist, -1, 1) + sin_t = torch.sqrt(1 - cos_t**2) + phi = torch.atan2(d[:, 1], d[:, 0]) + # linear binning onto the r grid: each vector split between bins i0 and i0 + 1 + x = (dist - float(r_t[0])) / dr + i0 = torch.floor(x).long() + f = x - i0 + index = torch.cat([i0, i0 + 1]) + weights = torch.cat([1 - f, f]) / torch.cat([shell, shell[-1:]])[index] # per unit volume + cos_t, sin_t, phi = (torch.cat([t, t]) for t in (cos_t, sin_t, phi)) + Cl += _real_sph_harm_power(cos_t, sin_t, phi, l_values, weights, n_r, index) + Cl *= 2 * np.pi + + theta = torch.deg2rad(torch.as_tensor(np.asarray(theta_deg), dtype=torch.float64)) + out = torch.zeros((n_r, n_r, theta.numel()), dtype=torch.float64) + for k, l in enumerate(l_values.tolist()): + C = Cl[k] if q is None else _bandlimit_order(Cl[k], l, r, q) + if l_taper: + C = C * float(order_taper(l, int(l_values[-1]) + 2)) # same as reconstruct_PADF + out += C[:, :, None] * torch.special.legendre_polynomial_p(torch.cos(theta), l)[None, None, :] + return (out * torch.abs(torch.sin(theta))[None, None, :]).numpy() + + +def _per_panel(value, n, nested=False): + """One entry per panel: a list with one item per panel is used as is, + anything else is repeated. With nested=True only a list of lists (or + Nones) counts as per panel, so a flat list of marks is shared.""" + if isinstance(value, (list, tuple)) and len(value) == n: + if not nested or all(v is None or isinstance(v, (list, tuple, np.ndarray)) for v in value): + return list(value) + return [value] * n diff --git a/src/quantem/diffraction/padf2d.py b/src/quantem/diffraction/padf2d.py new file mode 100644 index 00000000..3b9f7a63 --- /dev/null +++ b/src/quantem/diffraction/padf2d.py @@ -0,0 +1,297 @@ +""" +2D pair-angle distribution function for in-plane (non-isotropic) samples. + +The standard PADF assumes 3D orientation averaging and expands the +correlation in Legendre polynomials / spherical Bessel functions. For +samples that only rotate in the plane - 2D materials at zone axis, surface +layers, the FFT-of-image test geometry - the correct decomposition is +circular harmonics and Hankel (J_n Bessel) transforms: + + C_n(q, q') = < conj(I_hat_n(q)) I_hat_n(q') >_patterns (exact, + read directly from the angular FFT - no basis inversion needed) + + C_n(r, r') = (2 pi)^2 int int q q' J_n(2 pi q r) J_n(2 pi q' r') + C_n(q, q') dq dq' / q_max^4 + + Theta(r, r', dtheta) = sum_n C_n(r, r') e^{i n dtheta} + +Compared with pushing 2D data through the 3D machinery this removes the +~2% systematic peak-position bias of the mismatched j_l kernel, needs no +SVD regularization, and has no sin(theta) sampling corrections (the 2D +angular measure is uniform). + +Friedel symmetry of single kinematic patterns makes odd n vanish, and the +relative angle folds to [0, 180]. When Friedel is broken (dynamical +scattering from a non-centrosymmetric projection), the odd harmonics are +retained naturally and the full 0-360 degree relative angle is recovered - +the 3D pipeline discards odd terms by construction. + +Class: PairAngleDistributionFunction2D +""" + +import matplotlib.pyplot as plt +import numpy as np +import torch +from scipy.special import jv + +from quantem.diffraction.padf import _per_panel, order_taper + + +class PairAngleDistributionFunction2D: + """ + 2D (in-plane) pair-angle distribution function. + + Parameters + ---------- + polar : Polar4dstem or torch.Tensor or ndarray + Polar diffraction data, angular axis second-to-last, radial axis + last: (..., N_phi, N_q). Leading dims are flattened into patterns. + n_max : int + Maximum circular harmonic order retained (angular resolution + ~180 / n_max degrees). + dq : float or None + Radial calibration (q units per bin). Read from the polar dataset's + sampling when available; required otherwise. + q_min : float + q of the first radial bin (from the polar dataset when available). + r_min, r_max, r_step : float + Real-space grid (including r_max), in the reciprocal units of q. + fq : scalar, (N_q,) array, callable or None + Atomic form factor f(q) per radial bin, or a function of q (1/A); the intensity is divided by + f(q)^2 (same eq-7 rescaling as the 3D pipeline). Leaving the atom- + shape damping in shrinks the effective q range and makes the + real-space maps broad and oscillatory. Floor small values before + passing to avoid amplifying the high-q noise floor. + q_taper : float + Fraction of the radial range cosine-tapered to zero at q_max + (default 0.25). Softens the hard-cutoff ringing of the Hankel + transform; 0 disables. + normalize_mean : bool + Normalize the intensity to unit mean (magnitude convention shared + with the 3D PADF pipeline). + batch_size : int + Patterns per chunk when accumulating C_n(q, q'). + """ + + def __init__(self, polar, n_max: int = 60, dq: float | None = None, + q_min: float | None = None, r_min: float = 0.0, + r_max: float = 20.0, r_step: float = 0.02, + fq=None, q_taper: float = 0.25, + normalize_mean: bool = True, batch_size: int = 4096): + if dq is None and hasattr(polar, "sampling"): + dq = float(np.asarray(polar.sampling)[3]) + if q_min is None: + q_min = float(np.asarray(polar.origin)[3]) if hasattr(polar, "origin") else 0.0 + if dq is None: + raise ValueError("dq must be given when polar carries no calibration.") + + 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.dq = float(dq) + self.q = q_min + np.arange(n_q) * self.dq + if callable(fq): + fq = fq(self.q) + + data = tensor.reshape(-1, n_phi, n_q) + if fq is not None: + fq_t = torch.as_tensor(np.broadcast_to(np.asarray(fq, dtype=np.float64), (n_q,)).copy()) + data = data / fq_t[None, None, :] ** 2 + if q_taper > 0: + data = data * _q_taper(n_q, q_taper)[None, None, :] + if normalize_mean: + data = data / data.mean() + n_patterns = data.shape[0] + + # circular harmonic correlation matrices C_n(q, q'), accumulated + # directly from the angular FFT (Fourier-series coefficients) + Cn = torch.zeros((self.n_max + 1, n_q, n_q), dtype=torch.complex128) + for start in range(0, n_patterns, batch_size): + F = torch.fft.rfft(data[start:start + batch_size], dim=1) / n_phi + F = F[:, : self.n_max + 1, :] # (b, n, q) + Cn += torch.einsum("bnq,bnp->nqp", torch.conj(F), F) + Cn = Cn / n_patterns + self.Cn_q = Cn.numpy() + + # Hankel transform of order n on both q axes, nondimensionalized by + # q_max^2 per axis (2D quadrature measure q dq / q_max^2) + r = np.arange(r_min, r_max + 0.5 * r_step, r_step) # includes r_max + self.r = r + q_max = self.q[-1] + w = 2.0 * np.pi * self.q * self.dq / q_max**2 # (Nq,) quadrature weight + arg = 2.0 * np.pi * np.outer(r, self.q) # (Nr, Nq) + Cn_r = torch.zeros((self.n_max + 1, r.size, r.size), dtype=torch.complex128) + for n in range(self.n_max + 1): + K = torch.from_numpy(jv(n, arg) * w[None, :]).to(torch.complex128) # jv is not in torch + Cn_r[n] = K @ Cn[n] @ K.T + self.Cn_r = Cn_r.numpy() + + def reconstruct(self, n_theta: int = 360, include_n0: bool = False, n_taper: bool = False): + """ + Reconstruct Theta(r, r', dtheta) on n_theta angles over [0, 360). + + include_n0 : keep the isotropic n = 0 term (a dtheta-independent + offset containing the uncorrelated background); excluded by + default, matching the 3D pipeline's l = 0 handling. + n_taper : weight each harmonic by order_taper(n, n_max + 1), which + suppresses the ringing in dtheta from truncating the sum at n_max. + + Returns + ------- + theta_deg : (n_theta,) angles in degrees. + padf : (Nr, Nr, n_theta) real array. + """ + theta = 2.0 * np.pi * np.arange(n_theta) / n_theta + n0 = 0 if include_n0 else 1 + n_vals = np.arange(n0, self.n_max + 1) + w = order_taper(n_vals, self.n_max + 1) if n_taper else np.ones(n_vals.size) + # Hermitian sum: C_{-n} = conj(C_n) -> 2 Re for n > 0 + E = torch.from_numpy(np.exp(1j * np.outer(n_vals, theta)) * w[:, None]) # (Nn, n_theta) + padf = 2.0 * torch.real(torch.einsum("nij,nt->ijt", torch.from_numpy(self.Cn_r[n0:]), E)).numpy() + if include_n0: + padf -= np.real(self.Cn_r[0])[:, :, None] # counted twice above + self.theta_deg = np.degrees(theta) + return self.theta_deg, padf + + def plot_diagonal(self, r_max_display: float | None = None, n_theta: int = 360, + theta_max: float = 360.0, r_display_power: float = 0.5, + title=None, axsize: tuple[float, float] = (4.5, 3.5), + returnfig: bool = False, padf=None, n_taper: bool = False, + subtract_median: bool = True, r_fixed=None): + """ + Maps against relative angle (y-axis, 0 to theta_max) and r (x-axis), + one panel per volume in a single row: the r = r' diagonal, or + Theta(r_fixed, r', dtheta) when r_fixed is given (one value or one per + panel). The display is weighted by (r r')^r_display_power, as in + PairAngleDistributionFunction.simple_plot (r^1 on the diagonal by + default), zero-centered RdBu_r. theta_max = 180 shows the folded view + (sufficient for Friedel-symmetric data); 360 shows the full relative + angle available when odd harmonics are present. Color limits are set + from interior angles and r <= 0.92 * r_max_display, so edge artifacts + do not compress the scale. + + padf: an (Nr, Nr, Ntheta) volume on this r grid and angles 360 k / Ntheta + (e.g. a model from padf2d_from_clusters), or a list of them, to plot + instead of this reconstruction. subtract_median subtracts, for display, + the median over the interior angles at each r: the self-correlation + bands at 0 and 180 deg otherwise offset every interior angle (the 2D + measure has no sin(theta) factor to suppress them). + """ + if padf is None: + vols = [self.reconstruct(n_theta=n_theta, n_taper=n_taper)[1]] + else: + vols = list(padf) if isinstance(padf, (list, tuple)) else [padf] + n = len(vols) + titles = _per_panel(title, n) + r_fixeds = list(r_fixed) if isinstance(r_fixed, (list, tuple)) else [r_fixed] * n + r = self.r + r_max_display = r.max() if r_max_display is None else r_max_display + rs = r <= r_max_display + r_inner = r[rs] <= 0.92 * r_max_display + + fig, axs = plt.subplots(1, n, figsize=(axsize[0] * n, axsize[1]), + squeeze=False, constrained_layout=True) + for ax, vol, t, r0 in zip(axs[0], vols, titles, r_fixeds): + theta_deg = 360.0 * np.arange(vol.shape[2]) / vol.shape[2] + ts = theta_deg <= theta_max + 1e-9 + if r0 is None: + m = np.einsum("iik->ik", vol)[rs][:, ts] * (r[rs] ** (2 * r_display_power))[:, None] + xlabel, label = "r = r' (Å)", f"$\\Theta_{{2D}} \\cdot r^{{{2 * r_display_power:g}}}$" + else: + i0 = int(np.argmin(np.abs(r - r0))) + m = vol[i0][rs][:, ts] * ((r[i0] * r[rs]) ** r_display_power)[:, None] + xlabel = f"r' (Å), r = {r[i0]:.2f} Å" + label = f"$\\Theta_{{2D}} \\cdot (r r')^{{{r_display_power:g}}}$" + disp = m.T + interior = (theta_deg[ts] % 180 > 15) & (theta_deg[ts] % 180 < 165) + if subtract_median: + disp = disp - np.median(disp[interior], axis=0, keepdims=True) + vmax = np.quantile(np.abs(disp[np.ix_(interior, r_inner)]), 0.999) + im = ax.pcolormesh(r[rs], theta_deg[ts], disp, cmap="RdBu_r", + vmin=-vmax, vmax=vmax, shading="auto") + ax.set_xlabel(xlabel) + ax.set_ylabel("$\\Delta\\theta$ (deg)") + ax.set_ylim(0, theta_max) + ax.set_yticks(np.arange(0, theta_max + 1, 30 if theta_max <= 180 else 60)) + fig.colorbar(im, ax=ax, label=label) + if t is not None: + ax.set_title(t, fontsize=10) + if returnfig: + return fig, axs[0] + plt.show() + + +def padf2d_from_clusters(clusters_xy, r, n_max: int = 60, n_theta: int = 360, + q=None, q_taper: float = 0.25, n_taper: bool = False): + """ + 2D PADF of atomic models: the angular correlation of all in-plane + interatomic vectors within each cluster (two-, three- and four-body + terms), summed over clusters, on the angle grid of + PairAngleDistributionFunction2D.reconstruct. + + A_n(r) = sum over vectors at |d| = r of exp(-i n phi), per unit area + C_n(r, r') = conj(A_n(r)) A_n(r') + + q : (Nq,) array or None + Measured q values. If given, each C_n is band-limited to them with + the same q_taper as the class (order-n Hankel transform and back). + n_taper : bool + Weight each harmonic by order_taper(n, n_max + 1), as reconstruct does. + + Returns + ------- + theta_deg : (n_theta,) angles in degrees. + padf : (Nr, Nr, n_theta) real array, n = 0 excluded. + """ + r = np.asarray(r, dtype=np.float64) + dr = r[1] - r[0] + n_r = r.size + ring = torch.from_numpy(2 * np.pi * np.maximum(r, 0.5 * dr) * dr) + n_vals = torch.arange(1, n_max + 1, dtype=torch.float64) + Cn = torch.zeros((n_max, n_r, n_r), dtype=torch.complex128) + for xy in clusters_xy: + xy = torch.as_tensor(np.asarray(xy)[:, :2], dtype=torch.float64) + d = (xy[None, :, :] - xy[:, None, :]).reshape(-1, 2) + dist = torch.hypot(d[:, 0], d[:, 1]) + keep = (dist > 0) & (dist >= r[0]) & (dist < r[-1]) + d, dist = d[keep], dist[keep] + phi = torch.atan2(d[:, 1], d[:, 0]) + # linear binning onto the r grid + x = (dist - r[0]) / dr + i0 = torch.floor(x).long() + f = (x - i0)[:, None] + E = torch.polar(torch.ones(phi.numel(), n_max, dtype=torch.float64), -phi[:, None] * n_vals[None, :]) + A = torch.zeros((n_r + 1, n_max), dtype=torch.complex128) + A.index_add_(0, i0, (1 - f) * E).index_add_(0, i0 + 1, f * E) + A = A[:n_r] / ring[:, None] + Cn += torch.einsum("in,jn->nij", A.conj(), A) + + if q is not None: + q = np.asarray(q, dtype=np.float64) + dq = q[1] - q[0] + taper = _q_taper(q.size, q_taper).numpy() + for k in range(n_max): + Jn = torch.from_numpy(jv(k + 1, 2 * np.pi * np.outer(q, r))) # (Nq, Nr), not in torch + fwd = (2 * np.pi * Jn * torch.from_numpy(r * dr)[None, :] + * torch.from_numpy(taper)[:, None]).to(torch.complex128) # r -> q, tapered + inv = (2 * np.pi * Jn.T * torch.from_numpy(q * dq)[None, :]).to(torch.complex128) # q -> r + Cn[k] = inv @ (fwd @ Cn[k] @ fwd.T) @ inv.T + + theta = 2.0 * np.pi * torch.arange(n_theta, dtype=torch.float64) / n_theta + E = torch.polar(torch.ones(n_max, n_theta, dtype=torch.float64), n_vals[:, None] * theta[None, :]) + if n_taper: + E = E * torch.from_numpy(order_taper(n_vals.numpy(), n_max + 1))[:, None] + padf = 2.0 * torch.real(torch.einsum("nij,nt->ijt", Cn, E)) + return np.degrees(theta.numpy()), padf.numpy() + + +def _q_taper(n_q: int, q_taper: float) -> torch.Tensor: + """Cosine taper of the top q_taper fraction of n_q radial bins to zero.""" + taper = torch.ones(n_q, dtype=torch.float64) + i0 = int((1.0 - q_taper) * n_q) + ramp = torch.arange(n_q - i0, dtype=torch.float64) / max(n_q - i0, 1) + taper[i0:] = 0.5 * (1 + torch.cos(torch.pi * ramp)) + return taper diff --git a/src/quantem/diffraction/polar.py b/src/quantem/diffraction/polar.py index c600f7a6..96898a95 100644 --- a/src/quantem/diffraction/polar.py +++ b/src/quantem/diffraction/polar.py @@ -35,7 +35,7 @@ class PairDistributionFunction(AutoSerialize): - polar transform of the diffraction patterns - azimuthal averaging to obtain I(k) - - parametric background fit B(k) (Gaussian model in k² and k⁴) + - parametric background fit B(k) (by default a constant plus a power law) - reduced structure factor F(k) = 2π · k · [S(k) − 1] - windowed sine transform of F(k) to recover the reduced PDF G(r) - optional low-r correction of S(k) for a known density (Kaplow 1965 / @@ -430,25 +430,29 @@ def fit_bg( Ik: torch.Tensor, kmin: float | None = None, kmax: float | None = None, + bg_model: str = "power", + num_power_laws: int = 1, + fk: NDArray | None = None, + bg_weight: str = "k/I", ) -> tuple[torch.Tensor, torch.Tensor]: """ Fit a smooth background B(k) to a radial intensity curve I(k). - The background model is a constant plus two monotonically decaying - terms (adopted from py4DSTEM): + The default background model is a constant plus a power law, - B(k) = c - + i0 * exp(-k^2 / (2 s0^2)) - + i1 * exp(-k^4 / (2 s1^4)) + B(k) = c + a * (1 + k^2 / s^2)^(-p), - B(k) is later subtracted from I(k) to isolate the diffuse signal, and - f(k) = B(k) - c is used as the denominator in the structure factor - S(k) = 1 + [I(k) − B(k)] / f(k). + with every parameter positive, so f(k) = B(k) - c is positive and + decreases monotonically. B(k) is later subtracted from I(k) to isolate + the diffuse signal, and f(k) is used as the denominator in the + structure factor S(k) = 1 + [I(k) − B(k)] / f(k). A second power law, + the py4DSTEM Gaussian model, or tabulated scattering factors can be + chosen with bg_model (see :meth:`fit_bg_batched`). - The five parameters are fit by weighted least squares over k in [kmin, kmax]. - Each residual is weighted by a sin² low-k taper that downweights the central - beam, divided by a linear factor that falls with k so that higher k counts - more (see :meth:`fit_bg_batched`). This is the single-curve convenience over + The parameters are fit by weighted least squares over k in [kmin, kmax]. + By default each residual is weighted by k / I(k), the error in F(k), times + a sin² low-k taper that downweights the central beam. This is the + single-curve convenience over :meth:`fit_bg_batched`, which holds the torch-native Levenberg-Marquardt solver. Parameters @@ -460,6 +464,8 @@ def fit_bg( Restrict the fit to k in [kmin, kmax]. The returned B(k) is still evaluated over the full k axis. If None, the full k range is used. + bg_model, num_power_laws, fk, bg_weight + Background model and fit weighting, see :meth:`fit_bg_batched`. Returns ------- @@ -471,7 +477,9 @@ def fit_bg( """ # A single curve is just a batch of one Ik_row = torch.as_tensor(Ik)[None, :] - bg_stack, f_stack = self.fit_bg_batched(Ik_row, kmin=kmin, kmax=kmax) + bg_stack, f_stack = self.fit_bg_batched( + Ik_row, kmin=kmin, kmax=kmax, bg_model=bg_model, num_power_laws=num_power_laws, fk=fk, bg_weight=bg_weight + ) self.bg, self.f = bg_stack[0], f_stack[0] return self.bg, self.f @@ -480,19 +488,17 @@ def fit_bg_batched( Ik_stack: torch.Tensor, kmin: float | None = None, kmax: float | None = None, + bg_model: str = "power", + num_power_laws: int = 1, + fk: NDArray | None = None, + bg_weight: str = "k/I", ) -> tuple[torch.Tensor, torch.Tensor]: - """Fit B(k) for a stack of N radial curves at once (vectorized LM). + """Fit B(k) for a stack of N radial curves at once. The fit minimises ``sum_k [w(k) * (B(k) - I(k))]^2`` over k in - [kmin, kmax], with the residual weight - - w(k) = taper(k) / (k_end - 0.9 * k + dk) - - Here ``taper`` rises as sin² from 0 at kmin to 1 at kmin + 0.25 1/A - (points where it is below 1e-4 get zero weight), and ``k_end`` is the - last k of the fit range. The weighting scheme follows the background fit - in py4DSTEM. All N curves are fit together in a single solve rather - than one at a time. + [kmin, kmax], by batched Levenberg-Marquardt on the log of the + (positive) parameters, or by linear least squares for "known". All N + curves are fit together. Parameters ---------- @@ -500,6 +506,29 @@ def fit_bg_batched( (N, Nk) radial means, one row per curve. kmin, kmax Fit range. + bg_model + "power" (default): B(k) = c + sum_i a_i (1 + k^2 / s_i^2)^(-p_i), with + num_power_laws (1 or 2) terms. Every parameter is positive, so the + single-scattering term f(k) = B(k) - c is positive and decreases + monotonically, and it falls off as a power of k like the atomic + scattering factors rather than exponentially. + "known": B(k) = N fk(k)^2 + c, the standard normalization by the + composition-averaged scattering factor, for a known composition. + Only N >= 0 and c >= 0 are fit. + "gaussian": B(k) = c + i0 exp(-k^2 / 2 s0^2) + i1 exp(-k^4 / 2 s1^4), + the py4DSTEM model, which falls off too fast to follow the atomic + scattering factors at high k. + fk + Mean atomic scattering factor on the k axis of the polar + data, or a function of k (1/A), required for bg_model="known". + For several elements pass sqrt(). + bg_weight + Residual weight w(k), always times a sin² taper that rises from 0 at + kmin to 1 at kmin + 0.25 1/A to suppress the central-beam tail. + "k/I" (default): k / I(k), the error in F(k) = k [I(k) - B(k)] / f(k) + for B close to I; I(k) is floored at 1e-3 of its median so empty + bins (e.g. a beam stop) cannot dominate. "k": k. + "py4dstem": 1 / (k_end - 0.9 k + dk), with k_end the end of the fit range. Returns ------- @@ -507,131 +536,83 @@ def fit_bg_batched( Each (N, Nk) float32: fitted background and f(k) = B(k) - c per row. """ k = torch.as_tensor(self.qq, dtype=torch.float64, device=self.device) - if kmin is None: - kmin = float(k.min()) - if kmax is None: - kmax = float(k.max()) - k2 = k**2 + kmin = float(k.min()) if kmin is None else kmin + kmax = float(k.max()) if kmax is None else kmax fit_mask = (k >= kmin) & (k <= kmax) - k2_fit = k2[fit_mask] # (Nf,) k_fit = k[fit_mask] - dk = k[1] - k[0] - # Residual weights over the fit window (length Nf): a sin^2 taper over the - # first 0.25 1/A above kmin suppresses the central-beam tail, and dividing - # by a line that falls with k favours high k. - taper = torch.sin(0.5 * torch.pi * torch.clamp((k_fit - kmin) / 0.25, 0.0, 1.0)) ** 2 - resid_weights = torch.where(taper > 1e-4, taper, 0.0) / (k_fit[-1] - 0.9 * k_fit + dk) Ik = torch.as_tensor(Ik_stack, dtype=torch.float64, device=self.device) Ik_fit = Ik[:, fit_mask].clamp(min=1e-10) # (N, Nf) - n = Ik_fit.shape[0] # number of Ik in batch - - # Build Jacobian - # Initial guesses: constant = min(I), amplitudes = median(I) - min(I), widths = mean(k). - c0 = Ik_fit.amin(dim=1) # (N,) - amp = (Ik_fit.median(dim=1).values - c0).clamp(min=1e-10) # (N,) - sig = torch.full((n,), max(float(k.mean()), 1e-3), dtype=torch.float64, device=self.device) - init = torch.stack([c0.clamp(min=1e-10), amp, sig, amp, sig], dim=1) # (N, 5) - # exp(theta) = parameters ensures that they remain positive - theta = torch.log(init) - - def _rj(theta: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Batched residual (N, Nf) and analytic Jacobian (N, Nf, 5). - - Model B(k): - B(k) = c + i0 * E0 + i1 * E1 - where: - E0 = exp(-k^2 / (2 s0^2)) - E1 = exp(-k^4 / (2 s1^4)) - """ - c, i0, s0, i1, s1 = torch.exp(theta).unbind(dim=1) # each (N,) - E0 = torch.exp( - torch.clamp(-k2_fit[None, :] / (2.0 * s0[:, None] ** 2), min=-100.0, max=0.0) - ) - E1 = torch.exp( - torch.clamp( - -(k2_fit[None, :] ** 2) / (2.0 * s1[:, None] ** 4), min=-100.0, max=0.0 - ) - ) - B = c[:, None] + i0[:, None] * E0 + i1[:, None] * E1 - w = resid_weights[None, :] # residual weights, broadcast over the batch - r = (B - Ik_fit) * w - t0 = i0[:, None] * E0 - t1 = i1[:, None] * E1 - jac = torch.stack( - [ - c[:, None] * w, - t0 * w, - t0 * (k2_fit[None, :] / s0[:, None] ** 2) * w, - t1 * w, - t1 * (2.0 * k2_fit[None, :] ** 2 / s1[:, None] ** 4) * w, - ], - dim=2, - ) - return r, jac - - def _loss(theta: torch.Tensor) -> torch.Tensor: - r, _ = _rj(theta) - return (r * r).sum(dim=1) # (N,) - - # Batched Levenberg-Marquardt, replaces scipy.curve_fit. - # recompute Jacobian every iteration and take one step, accepting - # per-curve where the loss improved. lamba is adapted per-curve - # over the course of fitting (decrease on accept toward Gauss-Newton, - # increase on reject toward gradient descent). Converged curves simply stop being updated. - lam = torch.ones(n, dtype=torch.float64, device=self.device) - loss = _loss(theta) - eye = 1e-12 * torch.eye(5, dtype=torch.float64, device=self.device) - max_log_step = 1.0 - for _ in range(200): # outer loop: recompute the Jacobian at the current theta - r, jac = _rj(theta) - JtJ = torch.einsum("nki,nkj->nij", jac, jac) # (N, 5, 5) - Jtr = torch.einsum("nki,nk->ni", jac, r) # (N, 5) - diag = torch.diagonal(JtJ, dim1=1, dim2=2).clamp(min=1e-12) # (N, 5) - theta_base = theta # all trial steps start from here this iter - prev_loss = loss - accepted = torch.zeros(n, dtype=torch.bool, device=self.device) - # inner loop: per-curve lambda search. increase damping and retry until - # each curve finds a downhill step - for _ in range(30): - A = JtJ + lam[:, None, None] * torch.diag_embed(diag) + eye - delta = torch.linalg.solve(A, Jtr) # (N, 5) - # shrink the step if largest component exceeds max_log_step - step_scale = (max_log_step / delta.abs().amax(dim=1).clamp(min=1e-30)).clamp( - max=1.0 - ) - theta_trial = theta_base - delta * step_scale[:, None] - loss_trial = _loss(theta_trial) - # accept improvements per-curve - improved = torch.isfinite(loss_trial) & (loss_trial < loss) & (~accepted) - theta = torch.where(improved[:, None], theta_trial, theta) - loss = torch.where(improved, loss_trial, loss) - accepted = accepted | improved - # Accepted curves relax lambda (toward Gauss-Newton), curves not accepted tighten it (toward gradient descent). - lam = torch.where( - improved, - (lam * 0.3).clamp(min=1e-12), - torch.where(~accepted, (lam * 3.0).clamp(max=1e12), lam), - ) - if bool(accepted.all()): - break - if not bool(accepted.any()): - break # no curve found a downhill step -> all converged / stuck - # Stop once every curve's relative improvement is negligible. - rel = (prev_loss - loss) / (prev_loss + 1e-30) - if float(rel.max()) < 1e-9: - break - # Evaluate B(k) over the full k-range (not just k_fit) for every curve + taper = torch.sin(0.5 * torch.pi * torch.clamp((k_fit - kmin) / 0.25, 0.0, 1.0)) ** 2 + taper = torch.where(taper > 1e-4, taper, 0.0) + if bg_weight == "k/I": + floor = 1e-3 * Ik_fit.median(dim=1, keepdim=True).values + weights = taper * k_fit / torch.maximum(Ik_fit, floor) + elif bg_weight == "k": + weights = (taper * k_fit).expand_as(Ik_fit) + elif bg_weight == "py4dstem": + weights = (taper / (k_fit[-1] - 0.9 * k_fit + (k[1] - k[0]))).expand_as(Ik_fit) + else: + raise ValueError(f"bg_weight must be 'k/I', 'k' or 'py4dstem', got {bg_weight!r}") + weights = weights / weights.amax(dim=1, keepdim=True) # (N, Nf), one row per curve + + if bg_model == "known": + if fk is None: + raise ValueError("bg_model='known' needs fk, the scattering factor on the k axis") + fk = fk(np.asarray(self.qq, dtype=float)) if callable(fk) else fk + f2 = torch.as_tensor(np.asarray(fk), dtype=torch.float64, device=self.device) ** 2 + basis = torch.stack([f2[fit_mask], torch.ones_like(k_fit)], dim=1) # (Nf, 2) + coef = torch.linalg.lstsq( + basis[None] * weights[:, :, None], (Ik_fit * weights)[:, :, None] + ).solution[:, :, 0] # (N, 2) + c = coef[:, 1].clamp(min=0.0) + # refit the scale alone where the constant was clamped to zero + w2f = weights**2 * f2[fit_mask] + scale = torch.where( + coef[:, 1] < 0, (Ik_fit * w2f).sum(1) / (w2f * f2[fit_mask]).sum(1), coef[:, 0] + ).clamp(min=1e-30) + f = scale[:, None] * f2[None, :] + return (f + c[:, None]).to(dtype=torch.float32), f.to(dtype=torch.float32) + + c0 = Ik_fit.amin(dim=1).clamp(min=1e-10) + if bg_model == "power": + if num_power_laws not in (1, 2): + raise ValueError("num_power_laws must be 1 or 2") + amp = (Ik_fit[:, 0] - c0).clamp(min=1e-10) + cols = [c0] # a steep term near the central beam and a slower tail + for s0, p0 in [(0.3, 2.0), (1.0, 1.0)][:num_power_laws]: + cols += [amp / num_power_laws, torch.full_like(c0, s0), torch.full_like(c0, p0)] + + def model(theta, kk): + par = torch.exp(theta) # [c, a_1, s_1, p_1, (a_2, s_2, p_2)] + B = par[0] + 0 * kk + for i in range(num_power_laws): + B = B + par[1 + 3 * i] * (1.0 + kk**2 / par[2 + 3 * i] ** 2) ** (-par[3 + 3 * i]) + return B + + elif bg_model == "gaussian": + amp = (Ik_fit.median(dim=1).values - c0).clamp(min=1e-10) + sig = torch.full_like(c0, max(float(k.mean()), 1e-3)) + cols = [c0, amp, sig, amp, sig] + + def model(theta, kk): + c, i0, s0, i1, s1 = torch.exp(theta) + E0 = torch.exp(torch.clamp(-(kk**2) / (2.0 * s0**2), min=-100.0, max=0.0)) + E1 = torch.exp(torch.clamp(-(kk**4) / (2.0 * s1**4), min=-100.0, max=0.0)) + return c + i0 * E0 + i1 * E1 + + else: + raise ValueError(f"bg_model must be 'power', 'known' or 'gaussian', got {bg_model!r}") + + theta = _levenberg_marquardt( + lambda th, I_row, w_row: (model(th, k_fit) - I_row) * w_row, + torch.log(torch.stack(cols, dim=1)), + Ik_fit, + weights, + ) with torch.no_grad(): - c, i0, s0, i1, s1 = torch.exp(theta).unbind(dim=1) - E0 = torch.exp( - torch.clamp(-k2[None, :] / (2.0 * s0[:, None] ** 2), min=-100.0, max=0.0) - ) - E1 = torch.exp( - torch.clamp(-(k2[None, :] ** 2) / (2.0 * s1[:, None] ** 4), min=-100.0, max=0.0) - ) - bg = c[:, None] + i0[:, None] * E0 + i1[:, None] * E1 # (N, Nk) - f = (bg - c[:, None]).clamp(min=1e-10 * bg.amax(dim=1, keepdim=True)) + bg = torch.stack([model(th, k) for th in theta]) # (N, Nk) + f = (bg - torch.exp(theta[:, :1])).clamp(min=1e-10 * bg.amax(dim=1, keepdim=True)) return bg.to(dtype=torch.float32), f.to(dtype=torch.float32) def calculate_Gr( @@ -650,6 +631,10 @@ def calculate_Gr( r_cut: float = 0.8, max_iter: int = 6, returnval: bool = False, + bg_model: str = "power", + num_power_laws: int = 1, + fk: NDArray | None = None, + bg_weight: str = "k/I", ) -> list[NDArray] | None: """ Calculate the reduced pair distribution function G(r) from a 4D-STEM dataset. @@ -710,6 +695,10 @@ def calculate_Gr( :meth:`damp_origin_oscillations`. returnval : bool, optional If True, return ``[r, G(r)]`` as numpy arrays. + bg_model, num_power_laws, fk, bg_weight : optional + Background model: "power" with 1 (default) or 2 power laws, + "gaussian", or "known" with the scattering factor fk; see + :meth:`fit_bg_batched`. bg_weight sets the fit weighting. Returns ------- @@ -758,7 +747,8 @@ def calculate_Gr( else: Ik = self.calculate_radial_mean(mask_realspace=mask_bool, returnval=True) # always refit the background fit bc it is cheap - bg, f = self.fit_bg(Ik, self.kmin_fit, self.kmax_fit) + bg, f = self.fit_bg(Ik, self.kmin_fit, self.kmax_fit, bg_model=bg_model, num_power_laws=num_power_laws, fk=fk, + bg_weight=bg_weight) # prevent division by near-zero values which cause NaNs at high k f_safe = torch.clamp(f, min=1e-10 * f.max()) @@ -952,15 +942,16 @@ def damp_origin_oscillations( k = torch.from_numpy(np.asarray(self.qq).astype(np.float32)).to(device=self.device) dk = k[1] - k[0] - k_fit_mask = (k >= self.kmin_fit) & (k <= self.kmax_window) - k_fit = k[k_fit_mask] + # correct S(k) over the transformed (window) range, which can be wider than the bg fit range + k_win_mask = (k >= self.kmin_window) & (k <= self.kmax_window) + k_win = k[k_win_mask] ka, ra = torch.meshgrid(k, self._r, indexing="ij") # Restrict r to [0, r_cut] for the correction r_mask = (self._r >= 0.0) & (self._r <= r_cut) r_short = self._r[r_mask] - k_fit_scaled = k_fit * 2 * torch.pi - k2d_fit, r2d_fit = torch.meshgrid(k_fit_scaled, r_short, indexing="ij") + k_win_scaled = k_win * 2 * torch.pi + k2d_win, r2d_win = torch.meshgrid(k_win_scaled, r_short, indexing="ij") Sk_cor = self.Sk.clone() # the window does not change during the iteration @@ -979,7 +970,7 @@ def damp_origin_oscillations( for _ in range(max_iter): # alpha and beta are the ideal and actual contributions to G(r) in the short-r range # from the current S(k) and G(r) - alpha, beta = self._compute_alpha_beta(k2d_fit, r2d_fit, G_iter[r_mask], r_short) + alpha, beta = self._compute_alpha_beta(k2d_win, r2d_win, G_iter[r_mask], r_short) # converge on the S(k) correction magnitude if beta_prev is not None: delta = float(torch.max(torch.abs(beta - beta_prev))) @@ -987,7 +978,7 @@ def damp_origin_oscillations( break beta_prev = beta.clone() # Update S_cor(k) and G_cor - Sk_cor[k_fit_mask] = Sk_cor[k_fit_mask] - beta + rho0 * alpha + Sk_cor[k_win_mask] = Sk_cor[k_win_mask] - beta + rho0 * alpha Fk_win = k * (Sk_cor - 1.0) * wk * 2 * torch.pi G_iter = ( (2.0 / torch.pi) @@ -1421,3 +1412,55 @@ def _apply_xrange( if not np.any(m): raise ValueError("Requested plot range contains no data.") return x[m], y[m] + + +def _levenberg_marquardt(residual, theta, data, weights, max_iter: int = 200, tol: float = 1e-9): + """ + Batched Levenberg-Marquardt. residual(theta_row, data_row, weight_row) + returns one curve's weighted residual; theta is (N, P), one parameter row + per curve. The Jacobian comes from autodiff. Each curve keeps its own + damping, which falls toward Gauss-Newton when a step lowers its loss and + rises toward gradient descent when it does not; converged curves stop + moving. + """ + res_fn = torch.func.vmap(residual) + jac_fn = torch.func.vmap(torch.func.jacrev(residual)) + n, n_par = theta.shape + + def loss_fn(th): + r = res_fn(th, data, weights) + return (r * r).sum(dim=1) + + lam = torch.ones(n, dtype=theta.dtype, device=theta.device) + loss = loss_fn(theta) + eye = 1e-12 * torch.eye(n_par, dtype=theta.dtype, device=theta.device) + for _ in range(max_iter): + r = res_fn(theta, data, weights) + jac = jac_fn(theta, data, weights) # (N, Nf, P) + JtJ = torch.einsum("nki,nkj->nij", jac, jac) + Jtr = torch.einsum("nki,nk->ni", jac, r) + diag = torch.diagonal(JtJ, dim1=1, dim2=2).clamp(min=1e-12) + theta_base, prev_loss = theta, loss + accepted = torch.zeros(n, dtype=torch.bool, device=theta.device) + for _ in range(30): # per-curve damping search + delta = torch.linalg.solve(JtJ + lam[:, None, None] * torch.diag_embed(diag) + eye, Jtr) + # limit the largest log-parameter step to 1 + step = (1.0 / delta.abs().amax(dim=1).clamp(min=1e-30)).clamp(max=1.0) + theta_trial = theta_base - delta * step[:, None] + loss_trial = loss_fn(theta_trial) + improved = torch.isfinite(loss_trial) & (loss_trial < loss) & (~accepted) + theta = torch.where(improved[:, None], theta_trial, theta) + loss = torch.where(improved, loss_trial, loss) + accepted = accepted | improved + lam = torch.where( + improved, + (lam * 0.3).clamp(min=1e-12), + torch.where(~accepted, (lam * 3.0).clamp(max=1e12), lam), + ) + if bool(accepted.all()): + break + if not bool(accepted.any()): + break # no curve found a downhill step + if float(((prev_loss - loss) / (prev_loss + 1e-30)).max()) < tol: + break + return theta diff --git a/src/quantem/diffraction/polar_transform.py b/src/quantem/diffraction/polar_transform.py index bac53001..95b5cbc9 100644 --- a/src/quantem/diffraction/polar_transform.py +++ b/src/quantem/diffraction/polar_transform.py @@ -21,6 +21,45 @@ # but is noted where the call occures +def as_dataset4dstem(data) -> Dataset4dstem: + """ + Normalize a dataset to 4D-STEM form so the polar/PADF pipeline can treat + all inputs uniformly: + + - 2D (single diffraction pattern) -> (1, 1, H, W) + - 3D (stack of patterns) -> (N, 1, H, W) + - 4D -> unchanged + + The detector calibration (last two dims) is preserved; wrapped scan dims + get sampling 1 pixel. + """ + if data.ndim == 4: + return data + arr = data.array if data.array is not None else to_numpy(data.tensor) + sampling = np.asarray(data.sampling, dtype=float) + origin = np.asarray(data.origin, dtype=float) + units = list(data.units) + if data.ndim == 2: + arr4 = arr[None, None] + scan_sampling = [1.0, 1.0] + elif data.ndim == 3: + arr4 = arr[:, None] + scan_sampling = [float(sampling[0]), 1.0] + else: + raise ValueError( + f"Got array with shape {data.shape}. " + "Expected a 2D pattern, 3D stack of patterns, or 4D-STEM dataset." + ) + return Dataset4dstem.from_array( + arr4, + name=data.name, + origin=np.concatenate([np.zeros(2), origin[-2:]]), + sampling=np.concatenate([scan_sampling, sampling[-2:]]), + units=["pixels", "pixels"] + units[-2:], + signal_units=data.signal_units, + ) + + def find_origin_angular_grid( data: Dataset4dstem, *, @@ -82,16 +121,8 @@ def find_origin_angular_grid( Array of shape (scan_row, scan_col, 2) containing (row, col) origin estimates in pixels. """ - if data.ndim == 2: - n_row, n_col = data.shape - scan_row, scan_col = 1, 1 - elif data.ndim == 4: - scan_row, scan_col, n_row, n_col = data.shape - else: - raise ValueError( - f" Got array with shape {data.shape}." - "To use find_origin_angular_grid, pass a 2D or 4DSTEM dataset." - ) + data = as_dataset4dstem(data) + scan_row, scan_col, n_row, n_col = data.shape # Move the full dataset to the chosen device once dps = ( @@ -396,11 +427,7 @@ def polar_transform( device: str = "cpu", batch_size: int = 128, ) -> Polar4dstem | torch.Tensor: - if data.ndim != 4: - raise ValueError( - f"Found array with shape: {data.shape}. " - "polar_transform requires a 4D-STEM dataset (ndim=4)." - ) + data = as_dataset4dstem(data) scan_row, scan_col, n_row, n_col = data.shape # Standardize origin_array input diff --git a/tests/diffraction/test_polar.py b/tests/diffraction/test_polar.py index 7f8e7716..bae2382b 100644 --- a/tests/diffraction/test_polar.py +++ b/tests/diffraction/test_polar.py @@ -315,29 +315,38 @@ def test_fit_bg_basic(self, synthetic_dataset2d): # Check that background is positive assert (bg >= 0).all() - def test_fit_bg_recovers_known_background(self, synthetic_4dstem_dataset): - """The torch fit recovers a known smooth background under a structure - modulation, on a realistic high-dynamic-range I(k) / calibrated axis.""" + @pytest.mark.parametrize("bg_model", ["power", "gaussian"]) + def test_fit_bg_recovers_known_background(self, synthetic_4dstem_dataset, bg_model): + """Each background model recovers a known smooth background of its own + form under a structure modulation, on a realistic high-dynamic-range + I(k) / calibrated axis.""" pdf = PairDistributionFunction.from_data(synthetic_4dstem_dataset, find_origin=False) nb = len(np.asarray(pdf.qq)) pdf.polar.sampling[3] = 2.3 / (nb - 1) # realistic 1/A q-axis k = np.asarray(pdf.qq) rng = np.random.RandomState(0) - bg_true = ( - 0.5 + 900.0 * np.exp(-(k**2) / (2 * 0.09**2)) + 6.0 * np.exp(-(k**4) / (2 * 0.8**4)) - ) + if bg_model == "gaussian": + bg_true = 0.5 + 900.0 * np.exp(-(k**2) / (2 * 0.09**2)) + 6.0 * np.exp(-(k**4) / (2 * 0.8**4)) + else: + bg_true = 0.5 + 900.0 * (1 + k**2 / 0.15**2) ** -1.8 # add 35% amplitude ripple to the true bg to simulate scattering modulation = 1.0 + 0.35 * np.sin(2 * np.pi * k / 0.42 + 0.5) * np.exp(-k / 1.2) Ik = np.clip(bg_true * modulation + np.abs(0.5 * rng.randn(k.size)), 1e-6, None).astype( np.float32 ) # try fitting to this simulated Ik - bg, _ = pdf.fit_bg(Ik.copy()) - bg = bg.cpu().numpy() - assert (bg >= 0).all() + weight = "py4dstem" if bg_model == "gaussian" else "k/I" + bg, f = pdf.fit_bg(Ik.copy(), bg_model=bg_model, bg_weight=weight) + bg, f = bg.cpu().numpy(), f.cpu().numpy() + assert (bg >= 0).all() and (f > 0).all() + assert (np.diff(f) <= 0).all(), "single scattering f(k) must decrease monotonically" # The fit tracks the modulated data, so pointwise error reflects the # ~35% structure modulation # the recovered shape of the smooth background should still correlate strongly with the truth. + # The power model is weighted toward the whole dynamic range (k / I), so + # compare it in log space; the Gaussian model keeps the linear comparison. + if bg_model == "power": + bg, bg_true = np.log(bg), np.log(bg_true) corr = float(np.corrcoef(bg, bg_true)[0, 1]) assert corr > 0.99, f"bg shape far from truth: corr {corr:.4g}" @@ -681,7 +690,7 @@ def _pdf_with_injected_S(cls): pdf.Ik = torch.tensor(S - 1, dtype=torch.float32) pdf.bg = torch.zeros(len(k)) pdf.f = torch.ones(len(k)) - pdf.fit_bg = lambda Ik, kmin=None, kmax=None: (pdf.bg, pdf.f) + pdf.fit_bg = lambda Ik, kmin=None, kmax=None, **kwargs: (pdf.bg, pdf.f) return pdf, k, S def test_G_matches_windowed_transform_of_true_S(self):