torchspin 0.3.0__py3-none-any.whl

torchspin/__init__.py

@@ -0,0 +1,363 @@
+"""torchspin — PyTorch EPR/ESR spin Hamiltonian engine.
+
+A Python/PyTorch port of EasySpin's core physics layer.
+All energies are in **MHz**; magnetic fields in **mT**; Euler angles in **radians**.
+
+Quick start::
+
+    import torch
+    from torchspin import SpinSystem, ham
+
+    # S=1/2 nitroxide-like spin system
+    sys = SpinSystem(
+        S=[0.5],
+        g=[[2.009, 2.006, 2.002]],
+        Nucs=['14N'],
+        A=[[10.0, 10.0, 95.0]],   # MHz
+    )
+
+    B0 = [0.0, 0.0, 340.0]         # mT, along z
+    H = ham(sys, B0)
+    energies = torch.linalg.eigvalsh(H)   # MHz
+"""
+from importlib.metadata import version as _pkg_version, PackageNotFoundError as _PkgNotFound
+try:
+    __version__ = _pkg_version("torchspin")
+except _PkgNotFound:
+    # Editable install during development, or version metadata unavailable.
+    __version__ = "0.3.0+dev"
+
+from torchspin._compile import maybe_compile, set_compile_enabled, is_compile_available
+from torchspin.autograd import differentiable_spectrum
+from torchspin.batch import batch_pepper, batch_simulate
+from torchspin.cardamom import cardamom, CardamomPar, CardamomOptions, MDInput
+from torchspin.ctafft import ctafft
+from torchspin.ewrls import ewrls
+from torchspin.lpsvd import lpsvd, LPSVDResult
+from torchspin.rapidscan2spc import rapidscan2spc
+from torchspin.stochtraj_diffusion import stochtraj_diffusion, DiffusionPar
+from torchspin.stochtraj_jump import stochtraj_jump, JumpPar
+from torchspin.autoguess import (estimate_parameters, field_from_g, g_from_field,
+                                  print_parameter_report)
+from torchspin.lineshape import (gaussian, lorentzian, voigtian, apowin, addnoise, deriv, lshape)
+from torchspin.dataproc import basecorr, fieldmod, rescaledata
+from torchspin.utils import (datasmooth, equivcouple, equivsplit,
+                              hilberttrans, larmorfrq, rcfilt, unitconvert,
+                              mhz2mt, mt2mhz, hsdim, eprconvert)
+from torchspin.nucdata import nucabund, nucspin, nucgval, nucqmom
+from torchspin.blochsteady import blochsteady, BlochOptions
+from torchspin.chili import chili, ChiliOptions
+from torchspin.endorfrq import endorfrq
+from torchspin.salt import salt
+from torchspin.constants import (BMAGN, BOLTZMANN, CLIGHT, EVOLT, GFREE, NMAGN, PLANCK,
+                                   HBAR, AMU, ECHARGE, EMASS, PMASS, NMASS, EPS0, MU0,
+                                   FARADAY, AVOGADRO, BOHRRAD, HARTREE, RYDBERG, GAMMAE, GAMMAN,
+                                   ANGSTROM, DEGREE, BARN, MOLGAS)
+from torchspin.convspec import convspec
+from torchspin.eprload import eprload, eprload_info
+from torchspin.eprsave import eprsave, eprsave_info
+from torchspin.esfit import esfit, FitOptions, FitResult
+from torchspin.experiment import Experiment, Options
+from torchspin.curry import curry
+from torchspin.fastmotion import fastmotion
+from torchspin.garlic import garlic
+from torchspin.ham import ham
+from torchspin.levels import levels
+from torchspin.resfreqs_matrix import resfreqs_matrix
+from torchspin.ham_cf import ham_cf
+from torchspin.ham_ee import ham_ee
+from torchspin.ham_ez import ham_ez
+from torchspin.ham_ezho import ham_ezho
+from torchspin.ham_hf import ham_hf
+from torchspin.ham_nn import ham_nn
+from torchspin.ham_nq import ham_nq
+from torchspin.ham_nz import ham_nz
+from torchspin.ham_oz import ham_oz
+from torchspin.ham_so import ham_so
+from torchspin.ham_zf import ham_zf, zfsframes
+from torchspin.hamsymm import hamsymm
+from torchspin.makespec import makespec
+from torchspin.pepper import pepper
+from torchspin.pepper_autograd import pepper_autograd
+from torchspin.rotations import erot
+from torchspin.rotutils import (eulang, euler2quat, quat2euler, quat2rotmat,
+                                 rotmat2quat, rotaxi2mat, rotmat2axi,
+                                 quatinv, quatmult, quatvecmult,
+                                 vec2ang, ang2vec,
+                                 tensor_cart2sph, tensor_sph2cart, rotateframe)
+from torchspin.diptensor import diptensor
+from torchspin.spinladder import spinladder, SpinManifold
+from torchspin.orisel import orisel, OriselOptions
+from torchspin.spinops import sop, commute
+from torchspin.sphgrid import sphgrid, sphrand
+from torchspin.angmom import (wigner3j, wigner6j, wignerd, spherharm, clebschgordan, cgmatrix,
+                               isto, isto2stev)
+from torchspin.endorfrq_perturb import endorfrq_perturb
+from torchspin.resfreqs_perturb import resfreqs_perturb
+from torchspin.resfields_eig import resfields_eig, EigOptions
+from torchspin.spidyan import spidyan, SpidyanOptions, SpidyanInfo
+from torchspin.mdload import mdload, MDData, MDInfo, MDOptions
+from torchspin.mdhmm import mdhmm, HMMResult, HMMOptions
+from torchspin.mdtraj2oripot import mdtraj2oripot
+from torchspin.orca2torchspin import orca2torchspin, OrcaData
+from torchspin.mlpsvd import mlpsvd, MLPSVDResult
+from torchspin.oripotentialplot import oripotentialplot
+from torchspin.saffron import saffron, PulseExperiment, SaffronOptions
+from torchspin.signalprocessing import signalprocessing
+from torchspin.saffron_pathways import find_refocusing_pathways
+from torchspin.spinsystem import SpinSystem, nucspinrmv, nucspinkeep, nucspinadd, isotopologues, spinvec
+from torchspin.sigeq import sigeq
+from torchspin.fdaxis import fdaxis
+from torchspin.plegendre import plegendre
+from torchspin.propint import propint
+from torchspin.resonatorprofile import resonatorprofile
+from torchspin.rfmixer import rfmixer
+from torchspin.transmitter import transmitter
+from torchspin.pulse import pulse
+from torchspin.exciteprofile import exciteprofile
+from torchspin.resonator import resonator
+from torchspin.photoselect import photoselect
+from torchspin.levelsplot import levelsplot, LevelsPlotOptions
+from torchspin.stackplot import stackplot
+from torchspin.dipkernel import dipkernel
+from torchspin.dipbackground import dipbackground
+from torchspin.exponfit import exponfit
+from torchspin.nucfrq2d import nucfrq2d
+from torchspin.evolve import evolve
+
+__all__ = [
+    # Version
+    '__version__',
+    # Differentiable simulator
+    'differentiable_spectrum',
+    'pepper_autograd',
+    # Core
+    'SpinSystem',
+    'nucspinrmv',
+    'nucspinkeep',
+    'nucspinadd',
+    'isotopologues',
+    'spinvec',
+    'sop',
+    'erot',
+    # Hamiltonians
+    'ham',
+    'ham_cf',
+    'ham_ez',
+    'ham_ezho',
+    'ham_hf',
+    'ham_ee',
+    'ham_nz',
+    'ham_nq',
+    'ham_nn',
+    'ham_oz',
+    'ham_so',
+    'ham_zf',
+    # Spectrum simulation
+    'blochsteady',
+    'BlochOptions',
+    'pepper',
+    'garlic',
+    'fastmotion',
+    'chili',
+    'ChiliOptions',
+    'endorfrq',
+    'salt',
+    'levels',
+    'resfreqs_matrix',
+    'hamsymm',
+    'Experiment',
+    'Options',
+    'sphgrid',
+    'sphrand',
+    'makespec',
+    'convspec',
+    # Magnetometry
+    'curry',
+    # Data I/O (save)
+    'eprsave',
+    'eprsave_info',
+    # Signal processing utilities
+    'larmorfrq',
+    'equivsplit',
+    'equivcouple',
+    'datasmooth',
+    'rcfilt',
+    'hilberttrans',
+    'unitconvert',
+    'mhz2mt',
+    'mt2mhz',
+    'hsdim',
+    # Nuclear data
+    'nucabund',
+    'nucspin',
+    'nucgval',
+    'nucqmom',
+    # Lineshape functions
+    'gaussian',
+    'lorentzian',
+    'voigtian',
+    'apowin',
+    'addnoise',
+    'deriv',
+    'lshape',
+    # Data processing
+    'basecorr',
+    'fieldmod',
+    'rescaledata',
+    # Angular momentum functions
+    'wigner3j',
+    'wigner6j',
+    'wignerd',
+    'spherharm',
+    'clebschgordan',
+    'cgmatrix',
+    'isto',
+    'isto2stev',
+    # Perturbation-theory resonance/ENDOR
+    'endorfrq_perturb',
+    'resfreqs_perturb',
+    'resfields_eig',
+    'EigOptions',
+    # Spin dynamics (spidyan)
+    'spidyan',
+    'SpidyanOptions',
+    'SpidyanInfo',
+    # MD trajectory analysis (cardamom Phase 5)
+    'mdload',
+    'MDData',
+    'MDInfo',
+    'MDOptions',
+    'mdhmm',
+    'HMMResult',
+    'HMMOptions',
+    'mdtraj2oripot',
+    # Rotation utilities
+    'eulang',
+    'euler2quat',
+    'quat2euler',
+    'quat2rotmat',
+    'rotmat2quat',
+    'rotaxi2mat',
+    'rotmat2axi',
+    'quatinv',
+    'quatmult',
+    'quatvecmult',
+    'vec2ang',
+    'ang2vec',
+    'tensor_cart2sph',
+    'tensor_sph2cart',
+    'rotateframe',
+    # Spin operators
+    'commute',
+    # Dipolar coupling
+    'diptensor',
+    # ZFS analysis
+    'zfsframes',
+    # Spin ladder
+    'spinladder',
+    'SpinManifold',
+    # Orientation selectivity
+    'orisel',
+    'OriselOptions',
+    # Data I/O
+    'eprload',
+    'eprload_info',
+    'orca2torchspin',
+    'OrcaData',
+    # Pulse EPR
+    'saffron',
+    'PulseExperiment',
+    'SaffronOptions',
+    'find_refocusing_pathways',
+    # Parameter estimation
+    'estimate_parameters',
+    'field_from_g',
+    'g_from_field',
+    'print_parameter_report',
+    # Fitting
+    'esfit',
+    'FitOptions',
+    'FitResult',
+    # Batch simulation
+    'batch_pepper',
+    'batch_simulate',
+    # Trajectory-based simulation (cardamom)
+    'cardamom',
+    'CardamomPar',
+    'CardamomOptions',
+    'MDInput',
+    'stochtraj_diffusion',
+    'DiffusionPar',
+    'stochtraj_jump',
+    'JumpPar',
+    # Compilation
+    'maybe_compile',
+    'set_compile_enabled',
+    'is_compile_available',
+    # Constants
+    'BMAGN',
+    'PLANCK',
+    'GFREE',
+    'NMAGN',
+    'BOLTZMANN',
+    'CLIGHT',
+    'EVOLT',
+    'HBAR',
+    'AMU',
+    'ECHARGE',
+    'EMASS',
+    'PMASS',
+    'NMASS',
+    'EPS0',
+    'MU0',
+    'FARADAY',
+    'AVOGADRO',
+    'BOHRRAD',
+    'HARTREE',
+    'RYDBERG',
+    'GAMMAE',
+    'GAMMAN',
+    'ANGSTROM',
+    'DEGREE',
+    'BARN',
+    'MOLGAS',
+    # EPR conversion
+    'eprconvert',
+    # Tier 1 utilities
+    'sigeq',
+    'fdaxis',
+    'plegendre',
+    'photoselect',
+    'levelsplot',
+    'LevelsPlotOptions',
+    'stackplot',
+    # Tier 2 utilities
+    'dipkernel',
+    'dipbackground',
+    'exponfit',
+    'nucfrq2d',
+    'evolve',
+    # Tier 3b — Signal Processing
+    'ctafft',
+    'ewrls',
+    'lpsvd',
+    'LPSVDResult',
+    'mlpsvd',
+    'MLPSVDResult',
+    'rapidscan2spc',
+    'signalprocessing',
+    # Visualization
+    'oripotentialplot',
+    # Tier 3c — Higher-order Zeeman
+    'ham_ezho',
+    # Tier 3d-i — Pulse EPR primitives
+    'resonatorprofile',
+    'transmitter',
+    'rfmixer',
+    'propint',
+    # Tier 3d-ii — Pulse waveforms
+    'pulse',
+    'exciteprofile',
+    'resonator',
+]

torchspin/_cardamom_istos.py

@@ -0,0 +1,327 @@
+"""IST decomposition and Wigner D-matrices for cardamom ISTOs method.
+
+Implements:
+- ``istotensor``: irreducible spherical tensor operators from two vector operators
+- ``magint``: IST decomposition of spin Hamiltonian interactions
+- ``wigD``: rank-2 Wigner D-matrices from quaternion trajectories
+
+References
+----------
+[1] Oganesyan, Phys. Chem. Chem. Phys. 13, 4724 (2011)
+[2] Mehring, Principles of High Resolution NMR in Solids, Appendix A (1983)
+"""
+from __future__ import annotations
+
+import numpy as np
+from torchspin.constants import BMAGN, NMAGN, PLANCK
+from torchspin.spinops import sop
+
+
+# ---------------------------------------------------------------------------
+# istotensor — rank-0, rank-1, rank-2 IST operators from two vector operators
+# ---------------------------------------------------------------------------
+
+def istotensor(
+    a: list[np.ndarray] | np.ndarray,
+    b: list[np.ndarray] | np.ndarray,
+) -> tuple[np.ndarray, list[np.ndarray], list[np.ndarray]]:
+    """Build irreducible spherical tensor operators from two vector operators.
+
+    Port of EasySpin's ``private/istotensor.m``.
+
+    Parameters
+    ----------
+    a, b:
+        Each is either a list/tuple of 3 matrices [ax, ay, az] (spin operators)
+        or a length-3 array (e.g. magnetic field components).
+
+    Returns
+    -------
+    T0:
+        Rank-0 tensor (scalar or matrix).
+    T1:
+        Rank-1 tensor, list of 3 elements [T1(+1), T1(0), T1(-1)].
+    T2:
+        Rank-2 tensor, list of 5 elements [T2(+2), T2(+1), T2(0), T2(-1), T2(-2)].
+    """
+    if isinstance(a, (list, tuple)):
+        ax, ay, az = a[0], a[1], a[2]
+    else:
+        a = np.asarray(a)
+        ax, ay, az = a[0], a[1], a[2]
+
+    if isinstance(b, (list, tuple)):
+        bx, by, bz = b[0], b[1], b[2]
+    else:
+        b = np.asarray(b)
+        bx, by, bz = b[0], b[1], b[2]
+
+    # Rank 0: T0 = -(1/sqrt(3)) * (ax*bx + ay*by + az*bz)
+    T0 = -(1 / np.sqrt(3)) * (ax @ bx + ay @ by + az @ bz) \
+        if isinstance(ax, np.ndarray) and ax.ndim == 2 \
+        else -(1 / np.sqrt(3)) * (ax * bx + ay * by + az * bz)
+
+    # Helper for products (handles both matrix and scalar cases)
+    def _prod(x, y):
+        if isinstance(x, np.ndarray) and x.ndim == 2:
+            return x @ y
+        return x * y
+
+    # Rank 1
+    T1 = [
+        -0.5 * (_prod(ax, bz) - _prod(az, bx) + 1j * (_prod(ay, bz) - _prod(az, by))),  # (+1)
+        -(1j / np.sqrt(2)) * (_prod(ay, bx) - _prod(ax, by)),  # (0)
+        -0.5 * (_prod(ax, bz) - _prod(az, bx) - 1j * (_prod(ay, bz) - _prod(az, by))),  # (-1)
+    ]
+
+    # Rank 2
+    T2 = [
+        0.5 * (_prod(ax, bx) - _prod(ay, by) + 1j * (_prod(ax, by) + _prod(ay, bx))),  # (+2)
+        -0.5 * (_prod(ax, bz) + _prod(az, bx) + 1j * (_prod(ay, bz) + _prod(az, by))),  # (+1)
+        np.sqrt(2 / 3) * (_prod(az, bz) - 0.5 * (_prod(ax, bx) + _prod(ay, by))),  # (0)
+        0.5 * (_prod(ax, bz) + _prod(az, bx) - 1j * (_prod(ay, bz) + _prod(az, by))),  # (-1)
+        0.5 * (_prod(ax, bx) - _prod(ay, by) - 1j * (_prod(ax, by) + _prod(ay, bx))),  # (-2)
+    ]
+
+    return T0, T1, T2
+
+
+# ---------------------------------------------------------------------------
+# magint — IST decomposition of spin Hamiltonian
+# ---------------------------------------------------------------------------
+
+def magint(
+    spins: list[float],
+    g: np.ndarray,
+    CenterField: float,
+    A: np.ndarray | None = None,
+    D: np.ndarray | None = None,
+    gn: list[float] | None = None,
+    nuc_spins: list[float] | None = None,
+    include_nuc_zeeman: bool = False,
+) -> tuple[dict, dict]:
+    """Decompose spin Hamiltonian interactions into IST components.
+
+    Port of EasySpin's ``private/magint.m``.
+
+    Parameters
+    ----------
+    spins:
+        List of spin quantum numbers [S1, ..., I1, ...].
+    g:
+        g-tensor, shape (nElectrons, 3) or (3,) for single electron.
+    CenterField:
+        Center magnetic field in mT.
+    A:
+        Hyperfine tensor principal values in MHz, shape (nNuclei, 3) or (3,).
+    D:
+        ZFS tensor principal values in MHz, shape (nElectrons, 3).
+    gn:
+        Nuclear g-values.
+    nuc_spins:
+        Nuclear spin quantum numbers.
+    include_nuc_zeeman:
+        Whether to include nuclear Zeeman interaction.
+
+    Returns
+    -------
+    T:
+        dict with 'T0' (list of rank-0 ops), 'T2' (list of 5-element lists of rank-2 ops)
+    F:
+        dict with 'F0' (array of rank-0 spatial), 'F2' (array of rank-2 spatial, shape (nInt, 5))
+    """
+    g = np.atleast_2d(np.asarray(g, dtype=float))
+    nElSpins = g.shape[0]
+
+    # Count interactions
+    nNuclei = 0
+    if nuc_spins is not None:
+        nNuclei = len(nuc_spins)
+
+    # Build spin operators for all spins
+    SpinOps = []
+    for iSpin in range(len(spins)):
+        Sx = sop(spins, [iSpin + 1, 1]).numpy()
+        Sy = sop(spins, [iSpin + 1, 2]).numpy()
+        Sz = sop(spins, [iSpin + 1, 3]).numpy()
+        SpinOps.append([Sx, Sy, Sz])
+
+    # B0 field vector (z-direction)
+    B0 = [0.0, 0.0, CenterField / 1e3]  # mT → T
+
+    T0_list = []
+    T2_list = []
+    F0_list = []
+    F2_list = []
+
+    # --- Electron Zeeman: muB * B * g * S / h ---
+    for iEl in range(nElSpins):
+        g_mat = np.diag(g[iEl, :])
+        # IST operators: istotensor(B0, S)
+        T0, T1, T2 = istotensor(B0, SpinOps[iEl])
+        T0_list.append(T0)
+        T2_list.append(T2)
+        # Spatial components: tensor_cart2sph(g * bmagn / planck) in Hz
+        g_scaled = g_mat * (BMAGN / PLANCK)  # Hz/T
+        F0_scalar = _cart2sph_rank0(g_scaled)
+        F2_vec = _tensor_cart2sph_rank2(g_scaled)  # 5-element
+        F0_list.append(F0_scalar)
+        F2_list.append(F2_vec)
+
+    # --- Hyperfine: S * A * I ---
+    if A is not None and nNuclei > 0:
+        A = np.atleast_2d(np.asarray(A, dtype=float))
+        for iEl in range(nElSpins):
+            for iNuc in range(nNuclei):
+                A_mat = np.diag(A[iNuc, :]) * 1e6  # MHz → Hz
+                S_ops = SpinOps[iEl]
+                I_ops = SpinOps[nElSpins + iNuc]
+                T0, T1, T2 = istotensor(S_ops, I_ops)
+                T0_list.append(T0)
+                T2_list.append(T2)
+                F0_scalar = _cart2sph_rank0(A_mat)
+                F2_vec = _tensor_cart2sph_rank2(A_mat)
+                F0_list.append(F0_scalar)
+                F2_list.append(F2_vec)
+
+    # --- ZFS: S * D * S ---
+    if D is not None:
+        D = np.atleast_2d(np.asarray(D, dtype=float))
+        for iEl in range(nElSpins):
+            if spins[iEl] < 1:
+                continue
+            D_mat = np.diag(D[iEl, :]) * 1e6  # MHz → Hz
+            T0, T1, T2 = istotensor(SpinOps[iEl], SpinOps[iEl])
+            T0_list.append(T0)
+            T2_list.append(T2)
+            F0_scalar = _cart2sph_rank0(D_mat)
+            F2_vec = _tensor_cart2sph_rank2(D_mat)
+            F0_list.append(F0_scalar)
+            F2_list.append(F2_vec)
+
+    # --- Nuclear Zeeman ---
+    if include_nuc_zeeman and gn is not None:
+        for iNuc in range(nNuclei):
+            I_ops = SpinOps[nElSpins + iNuc]
+            T0, T1, T2 = istotensor(B0, I_ops)
+            T0_list.append(T0)
+            T2_list.append(T2)
+            gn_scaled = -gn[iNuc] * NMAGN / PLANCK  # Hz/T (scalar → isotropic tensor)
+            F0_list.append(_cart2sph_rank0(gn_scaled * np.eye(3)))
+            F2_list.append(_tensor_cart2sph_rank2(gn_scaled * np.eye(3)))
+
+    F0 = np.array(F0_list, dtype=complex)
+    F2 = np.array(F2_list, dtype=complex) if len(F2_list) > 0 else np.zeros((0, 5), dtype=complex)
+
+    return (
+        {'T0': T0_list, 'T2': T2_list},
+        {'F0': F0, 'F2': F2},
+    )
+
+
+def _tensor_cart2sph_rank2(T: np.ndarray) -> np.ndarray:
+    """Convert 3x3 Cartesian tensor to rank-2 spherical tensor components.
+
+    Uses the Mehring convention (same as EasySpin's tensor_cart2sph.m):
+    T2(0) = sqrt(2/3) * (Tzz - 0.5*(Txx + Tyy))
+
+    Returns [T2(+2), T2(+1), T2(0), T2(-1), T2(-2)].
+    """
+    T = np.asarray(T, dtype=complex)
+    Txx, Txy, Txz = T[0, 0], T[0, 1], T[0, 2]
+    Tyx, Tyy, Tyz = T[1, 0], T[1, 1], T[1, 2]
+    Tzx, Tzy, Tzz = T[2, 0], T[2, 1], T[2, 2]
+
+    T2p2 = 0.5 * ((Txx - Tyy) + 1j * (Txy + Tyx))
+    T2p1 = -0.5 * ((Txz + Tzx) + 1j * (Tyz + Tzy))
+    T20 = np.sqrt(2 / 3) * (Tzz - 0.5 * (Txx + Tyy))
+    T2m1 = 0.5 * ((Txz + Tzx) - 1j * (Tyz + Tzy))
+    T2m2 = 0.5 * ((Txx - Tyy) - 1j * (Txy + Tyx))
+
+    return np.array([T2p2, T2p1, T20, T2m1, T2m2])
+
+
+def _cart2sph_rank0(T: np.ndarray) -> complex:
+    """Extract rank-0 (isotropic) component from a 3x3 tensor.
+
+    F0 = -(1/sqrt(3)) * Tr(T)
+    """
+    T = np.asarray(T, dtype=complex)
+    return -(1 / np.sqrt(3)) * np.trace(T)
+
+
+# ---------------------------------------------------------------------------
+# wigD — rank-2 Wigner D-matrices from quaternion trajectories
+# ---------------------------------------------------------------------------
+
+def wigD(q: np.ndarray) -> np.ndarray:
+    """Compute rank-2 Wigner D-matrices from quaternion trajectories.
+
+    Port of the ``wigD`` subfunction in EasySpin's ``cardamom_propagatedm.m``.
+
+    Parameters
+    ----------
+    q:
+        Quaternion trajectory, shape ``(4, nSteps, nTraj)`` or ``(4, nSteps)``.
+        Convention: q = [q0, q1, q2, q3] where q0 is scalar part.
+
+    Returns
+    -------
+    D2:
+        Rank-2 Wigner D-matrix trajectory, shape ``(5, 5, nSteps, nTraj)``.
+        Indices correspond to m' = +2, +1, 0, -1, -2 (rows)
+        and m = +2, +1, 0, -1, -2 (columns).
+    """
+    if q.ndim == 2:
+        q = q[:, :, np.newaxis]
+
+    # Quaternion components → Cayley-Klein parameters
+    A = q[0] - 1j * q[3]
+    B = -q[2] - 1j * q[1]
+    Ast = q[0] + 1j * q[3]
+    Bst = -q[2] + 1j * q[1]
+
+    Z = A * Ast - B * Bst
+
+    ASq = A * A
+    BSq = B * B
+    AstSq = Ast * Ast
+    BstSq = Bst * Bst
+
+    nSteps = q.shape[1]
+    nTraj = q.shape[2]
+    D2 = np.zeros((5, 5, nSteps, nTraj), dtype=complex)
+
+    sqrt6 = np.sqrt(6)
+
+    D2[0, 0] = ASq * ASq
+    D2[0, 1] = 2 * A**3 * B
+    D2[0, 2] = sqrt6 * ASq * BSq
+    D2[0, 3] = 2 * A * B**3
+    D2[0, 4] = BSq * BSq
+
+    D2[1, 0] = -2 * A**3 * Bst
+    D2[1, 1] = ASq * (2 * Z - 1)
+    D2[1, 2] = sqrt6 * A * B * Z
+    D2[1, 3] = BSq * (2 * Z + 1)
+    D2[1, 4] = 2 * Ast * B**3
+
+    D2[2, 0] = sqrt6 * ASq * BstSq
+    D2[2, 1] = -sqrt6 * A * Bst * Z
+    D2[2, 2] = 0.5 * (3 * Z**2 - 1)
+    D2[2, 3] = sqrt6 * Ast * B * Z
+    D2[2, 4] = sqrt6 * AstSq * BSq
+
+    D2[3, 0] = -2 * A * Bst**3
+    D2[3, 1] = BstSq * (2 * Z + 1)
+    D2[3, 2] = -sqrt6 * Ast * Bst * Z
+    D2[3, 3] = AstSq * (2 * Z - 1)
+    D2[3, 4] = 2 * Ast**3 * B
+
+    D2[4, 0] = BstSq * BstSq
+    D2[4, 1] = -2 * Ast * Bst**3
+    D2[4, 2] = sqrt6 * AstSq * BstSq
+    D2[4, 3] = -2 * Ast**3 * Bst
+    D2[4, 4] = AstSq * AstSq
+
+    return D2