midas-2d 0.1.0__py3-none-any.whl

midas_2d/__init__.py

@@ -0,0 +1,211 @@
+"""midas-2d -- differentiable diffraction for 2D / few-layer & ultrafast work.
+
+Phase 1 (finite-size forward core):
+    from midas_2d import (
+        interference_factor, truncation_rod, nanoplatelet_rod,
+        rod_intensity, structure_factor_q,
+        cdse_zincblende, build_crystal_tensor,
+        make_spectrum, monochromatic,
+    )
+
+Builds on midas-hkls (structure factors, form factors, lattice, hkl gen) and
+midas-pink (energy spectrum).  See MIDAS_2D_ULTRAFAST_DIFFRACTION_PLAN.md.
+"""
+from .coherent import bcdi_forward, coherent_speckle, phase_retrieval
+from .debye import (
+    atomic_form_factors,
+    coherent_amplitude,
+    coherent_intensity,
+    debye_intensity,
+    ensemble_intensity,
+)
+from .disorder import (
+    AnisotropicMSD,
+    TransientMSD,
+    dwf_amplitude,
+    msd_tensor_from_frames,
+)
+from .dynamics import (
+    msd_from_stiffness,
+    recover_stiffness,
+    stiffness_from_msd,
+    thermal_ensemble,
+)
+from .energy import make_spectrum, monochromatic
+from .instrument import (
+    add_poisson_noise,
+    poisson_nll,
+    project_to_detector,
+    resolution_convolve,
+    solid_angle_polarization,
+)
+from .inverse import cosine_loss, fit, laplace_uncertainty, relative_l2_loss
+from .ml import ParameterMLP, make_dataset, train_surrogate
+from .phonon import (
+    apply_out_of_plane_strain,
+    bragg_timeseries,
+    fit_coherent_phonon,
+    strain_wave,
+)
+from .realdata import debye_reference_numpy, load_profile
+from .active_learning import (
+    fisher_information,
+    next_best_measurement,
+    rank_measurements,
+    sensitivity,
+)
+from .ensemble import polydisperse_rod, recover_thickness_distribution
+from .latent_dynamics import discover_eom, integrate_latent_ode, library_terms
+from .md_integrator import (
+    bragg_from_trajectory,
+    coherent_mode_kick,
+    harmonic_force,
+    recover_potential_from_movie,
+    velocity_verlet,
+)
+from .multimodal import (
+    carrier_density,
+    fit_multimodal,
+    lattice_temperature_from_carriers,
+    optical_signal,
+    strain_two_channel,
+    xray_only_degeneracy,
+)
+from .strain_profile import (
+    acoustic_pulse,
+    apply_depth_displacement,
+    depth_resolved_amplitude,
+    depth_resolved_intensity,
+    exponential_strain,
+    linear_strain,
+    recover_depth_strain,
+    strain_to_displacement,
+    temperature_to_displacement,
+    temperature_to_msd,
+    thermal_rod,
+)
+from .thermal_transport import (
+    fit_electron_phonon_coupling,
+    fit_thermal_diffusivity,
+    heat_diffusion_1d,
+    lattice_T_to_intensity_ratio,
+    two_temperature_model,
+)
+from .rocking import (
+    analytic_rod_model,
+    fwhm,
+    md_rod_model,
+    reciprocal_space_map,
+    rocking_curve,
+    thickness_from_fwhm,
+    thickness_loss_scan,
+)
+from .forward import rod_intensity, structure_factor_q
+from .io import (
+    build_crystal_tensor,
+    cdse_supercell,
+    cdse_zincblende,
+    load_xyz_frames,
+    supercell_from_crystal,
+)
+from .shape_factor import (
+    interference_factor,
+    laue_interference_1d,
+    nanoplatelet_rod,
+    truncation_rod,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "interference_factor",
+    "laue_interference_1d",
+    "nanoplatelet_rod",
+    "truncation_rod",
+    "rod_intensity",
+    "structure_factor_q",
+    "cdse_zincblende",
+    "build_crystal_tensor",
+    "cdse_supercell",
+    "supercell_from_crystal",
+    "load_xyz_frames",
+    "atomic_form_factors",
+    "coherent_amplitude",
+    "coherent_intensity",
+    "debye_intensity",
+    "ensemble_intensity",
+    "AnisotropicMSD",
+    "TransientMSD",
+    "dwf_amplitude",
+    "msd_tensor_from_frames",
+    "analytic_rod_model",
+    "md_rod_model",
+    "rocking_curve",
+    "fwhm",
+    "thickness_from_fwhm",
+    "reciprocal_space_map",
+    "thickness_loss_scan",
+    "fit",
+    "laplace_uncertainty",
+    "relative_l2_loss",
+    "cosine_loss",
+    "coherent_speckle",
+    "bcdi_forward",
+    "phase_retrieval",
+    "stiffness_from_msd",
+    "msd_from_stiffness",
+    "thermal_ensemble",
+    "recover_stiffness",
+    "strain_wave",
+    "apply_out_of_plane_strain",
+    "bragg_timeseries",
+    "fit_coherent_phonon",
+    "make_dataset",
+    "ParameterMLP",
+    "train_surrogate",
+    "project_to_detector",
+    "solid_angle_polarization",
+    "poisson_nll",
+    "add_poisson_noise",
+    "resolution_convolve",
+    "load_profile",
+    "debye_reference_numpy",
+    "depth_resolved_amplitude",
+    "depth_resolved_intensity",
+    "apply_depth_displacement",
+    "linear_strain",
+    "exponential_strain",
+    "acoustic_pulse",
+    "strain_to_displacement",
+    "temperature_to_displacement",
+    "temperature_to_msd",
+    "thermal_rod",
+    "recover_depth_strain",
+    "two_temperature_model",
+    "lattice_T_to_intensity_ratio",
+    "fit_electron_phonon_coupling",
+    "heat_diffusion_1d",
+    "fit_thermal_diffusivity",
+    "harmonic_force",
+    "velocity_verlet",
+    "bragg_from_trajectory",
+    "coherent_mode_kick",
+    "recover_potential_from_movie",
+    "carrier_density",
+    "lattice_temperature_from_carriers",
+    "strain_two_channel",
+    "optical_signal",
+    "fit_multimodal",
+    "xray_only_degeneracy",
+    "library_terms",
+    "integrate_latent_ode",
+    "discover_eom",
+    "polydisperse_rod",
+    "recover_thickness_distribution",
+    "sensitivity",
+    "fisher_information",
+    "rank_measurements",
+    "next_best_measurement",
+    "make_spectrum",
+    "monochromatic",
+]

midas_2d/active_learning.py

@@ -0,0 +1,13 @@
+"""Experiment design -- now provided by the shared ``midas_invert`` leaf.
+
+Thin re-export so existing ``midas_2d`` imports are unchanged.
+"""
+from midas_invert.design import (
+    fisher_information,
+    next_best_measurement,
+    rank_measurements,
+    sensitivity,
+)
+
+__all__ = ["sensitivity", "fisher_information", "rank_measurements",
+           "next_best_measurement"]

midas_2d/coherent.py

@@ -0,0 +1,120 @@
+"""Coherent diffraction & differentiable phase retrieval (BCDI thrust).
+
+Two faces of coherence:
+
+* :func:`coherent_speckle` -- finite-size *speckle* straight from atomic
+  coordinates: ``|A(Q)|^2`` on a 2-D/3-D Q grid (reuses
+  :func:`midas_2d.debye.coherent_intensity`).  This is the coordinate-level
+  coherent forward, differentiable w.r.t. positions.
+
+* :func:`bcdi_forward` + :func:`phase_retrieval` -- the object-level BCDI model
+  used by the Hruszkewycz/Cherukara thrust: a complex object
+  ``psi(r) = rho(r) exp(i Q0 . u(r))`` (amplitude = electron density / support,
+  phase = lattice displacement projected on the Bragg vector) produces speckle
+  ``|FFT(psi)|^2``.  :func:`phase_retrieval` inverts it by autograd -- the
+  differentiable alternative to iterative ER/HIO, and the natural place to drop
+  in a learned (Cherukara-style) prior.
+
+All complex-valued and torch-differentiable; CPU/CUDA/MPS portable.
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:  # pragma: no cover
+    import torch
+
+__all__ = [
+    "coherent_speckle",
+    "bcdi_forward",
+    "phase_retrieval",
+]
+
+
+def coherent_speckle(coords, elements, q_grid):
+    """``|A(Q)|^2`` over an arbitrary grid of Q-vectors (shape (..., 3))."""
+    from .debye import coherent_intensity
+    return coherent_intensity(coords, elements, q_grid)
+
+
+def bcdi_forward(psi, *, shift=True):
+    """Far-field coherent intensity ``|FFT(psi)|^2`` of a complex object.
+
+    Parameters
+    ----------
+    psi : complex tensor
+        The complex object (1-, 2-, or 3-D).
+    shift : bool
+        Apply ``fftshift`` so the zero frequency is centred.
+
+    Returns
+    -------
+    real tensor, same shape as ``psi``.
+    """
+    import torch
+    F = torch.fft.fftn(psi)
+    if shift:
+        F = torch.fft.fftshift(F)
+    return F.real * F.real + F.imag * F.imag
+
+
+def phase_retrieval(measured_intensity, support, *, init=None, steps=600,
+                    lr=0.05, beta_tv=0.0):
+    """Differentiable (autograd) phase retrieval of a complex object.
+
+    Minimises ``|| |FFT(psi)|^2 - measured ||^2`` over a complex object
+    constrained to live inside ``support`` (a 0/1 mask, applied every forward
+    so the constraint is exact and differentiable).
+
+    Parameters
+    ----------
+    measured_intensity : real tensor
+        Measured speckle ``|FFT(psi_true)|^2`` (fftshifted), same shape as the
+        object.
+    support : real tensor (0/1)
+        Object-domain support mask.
+    init : complex tensor, optional
+        Starting object; defaults to a random object inside the support seeded
+        from the measured magnitude.
+    steps, lr : optimisation controls.
+    beta_tv : float
+        Optional total-variation weight on |psi| (smoothness prior).
+
+    Returns
+    -------
+    dict: ``psi`` (recovered complex object), ``history`` (loss list).
+    """
+    import torch
+
+    measured = torch.as_tensor(measured_intensity)
+    support = torch.as_tensor(support, dtype=measured.dtype, device=measured.device)
+
+    if init is None:
+        # Random phase on a magnitude guess derived from the measured speckle.
+        mag = torch.sqrt(torch.clamp(measured.mean()) + 0.0) * support
+        phase = 2 * torch.pi * torch.rand(measured.shape, device=measured.device,
+                                          dtype=measured.dtype)
+        init = torch.complex(mag * torch.cos(phase), mag * torch.sin(phase))
+
+    re = (init.real * support).clone().detach().requires_grad_(True)
+    im = (init.imag * support).clone().detach().requires_grad_(True)
+    opt = torch.optim.Adam([re, im], lr=lr)
+
+    history: list[float] = []
+    obs_norm = measured.pow(2).sum() + 1e-12
+    for _ in range(steps):
+        opt.zero_grad()
+        psi = torch.complex(re * support, im * support)   # hard support constraint
+        pred = bcdi_forward(psi)
+        loss = ((pred - measured) ** 2).sum() / obs_norm
+        if beta_tv > 0:
+            amp = psi.abs()
+            tv = sum((torch.diff(amp, dim=d).abs().sum()
+                      for d in range(amp.dim())))
+            loss = loss + beta_tv * tv
+        loss.backward()
+        opt.step()
+        history.append(float(loss))
+
+    psi = torch.complex(re * support, im * support).detach()
+    return {"psi": psi, "history": history}

midas_2d/debye.py

@@ -0,0 +1,166 @@
+"""MD-coupled differentiable diffraction: scatter straight from atomic coordinates.
+
+Two forward paths, both autograd-differentiable **w.r.t. the atomic coordinates
+themselves** (this is the novel hook -- gradients flow from a measured pattern
+back to per-atom positions / displacement statistics, and onward to an MD force
+field):
+
+* :func:`coherent_amplitude` -- oriented, complex
+      A(Q) = sum_i  f_i(|Q|) exp(i Q . r_i),     I = |A|^2
+  Sampling ``Q`` along c* reproduces the analytic shape-function fringes of
+  Phase 1 (see the cross-check test) -- but here the fringes, the thermal
+  Debye-Waller falloff, and any *anisotropic* disordering all emerge from the
+  coordinates rather than from a phenomenological factor.
+
+* :func:`debye_intensity` -- orientationally averaged (the colloidal-solution
+  case for the Schaller/Flanders nanoplatelets)
+      I(q) = sum_i sum_j  f_i(q) f_j(q) sinc(q r_ij)
+
+:func:`ensemble_intensity` averages the incoherent intensity over MD frames;
+the spread of atomic positions across the trajectory *is* the disorder, so the
+anisotropic out-of-plane vs in-plane Debye-Waller falloff is reproduced with no
+DWF assumption at all.
+
+Form factors come from ``midas_hkls`` (Cromer-Mann), evaluated at the physical
+``s = |Q| / (4 pi)``.
+"""
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, Sequence
+
+from midas_hkls.form_factors import form_factor_batch
+
+if TYPE_CHECKING:  # pragma: no cover
+    import torch
+
+__all__ = [
+    "atomic_form_factors",
+    "coherent_amplitude",
+    "coherent_intensity",
+    "debye_intensity",
+    "ensemble_intensity",
+]
+
+_FOUR_PI = 4.0 * math.pi
+
+
+def atomic_form_factors(elements, q_mag):
+    """Per-atom Cromer-Mann form factors at scattering magnitude ``q_mag`` (1/A).
+
+    ``q_mag`` may have any shape; returns shape ``(*q_mag.shape, M)`` with M the
+    number of atoms in ``elements``.
+    """
+    import torch
+    q_mag = torch.as_tensor(q_mag)
+    s2 = (q_mag / _FOUR_PI) ** 2                       # s = |Q|/(4 pi)
+    return form_factor_batch(s2, list(elements))       # (..., M)
+
+
+def coherent_amplitude(coords, elements, q_vec):
+    """Complex coherent amplitude ``A(Q) = sum_i f_i(|Q|) exp(i Q . r_i)``.
+
+    Parameters
+    ----------
+    coords : tensor (M, 3)
+        Cartesian atomic coordinates (Angstrom).  Differentiable.
+    elements : sequence[str], length M
+    q_vec : tensor (..., 3)
+        Scattering vectors ``Q = 2 pi (h a* + k b* + l c*)`` in 1/A.
+
+    Returns
+    -------
+    tensor (...) complex
+    """
+    import torch
+    coords = torch.as_tensor(coords)
+    q_vec = torch.as_tensor(q_vec, dtype=coords.dtype, device=coords.device)
+    qmag = torch.linalg.vector_norm(q_vec, dim=-1)             # (...)
+    f = atomic_form_factors(elements, qmag)                    # (..., M)
+    phase = q_vec @ coords.T                                   # (..., M)
+    re = (f * torch.cos(phase)).sum(dim=-1)
+    im = (f * torch.sin(phase)).sum(dim=-1)
+    return torch.complex(re, im)
+
+
+def coherent_intensity(coords, elements, q_vec):
+    """``|A(Q)|^2`` for the coherent (oriented) path."""
+    A = coherent_amplitude(coords, elements, q_vec)
+    return A.real * A.real + A.imag * A.imag
+
+
+def debye_intensity(coords, elements, q_mag, *, tile=2048):
+    """Orientationally-averaged intensity via the Debye scattering equation.
+
+        I(q) = sum_i sum_j f_i(q) f_j(q) sinc(q r_ij),   sinc(x)=sin(x)/x
+
+    Parameters
+    ----------
+    coords : tensor (M, 3)
+        Cartesian coordinates (A).  Differentiable.
+    elements : sequence[str], length M
+    q_mag : tensor (Q,)
+        Scattering-vector magnitudes |Q| (1/A).
+    tile : int
+        Row-tile size over atoms to bound memory (the pair matrix is M x M).
+
+    Returns
+    -------
+    tensor (Q,)
+    """
+    import torch
+    coords = torch.as_tensor(coords)
+    q_mag = torch.as_tensor(q_mag, dtype=coords.dtype, device=coords.device)
+    M = coords.shape[0]
+    f = atomic_form_factors(elements, q_mag)                   # (Q, M)
+
+    out = torch.zeros(q_mag.shape, dtype=coords.dtype, device=coords.device)
+    for start in range(0, M, tile):
+        stop = min(start + tile, M)
+        ci = coords[start:stop]                                # (b, 3)
+        # pairwise distances block (b, M)
+        rij = torch.cdist(ci, coords)                          # (b, M)
+        qr = q_mag[:, None, None] * rij[None, :, :]            # (Q, b, M)
+        sinc = torch.sinc(qr / math.pi)                        # sin(qr)/(qr), =1 at 0
+        fi = f[:, start:stop]                                  # (Q, b)
+        # sum_j f_j sinc  -> (Q, b); then weight by f_i and sum over i
+        contrib = (fi * torch.einsum('qm,qbm->qb', f, sinc)).sum(dim=-1)
+        out = out + contrib
+    return out
+
+
+def ensemble_intensity(frames, elements, q, *, coherent=False, q_is_vector=None,
+                       tile=2048):
+    """Average intensity over a stack of MD frames (the disorder = the spread).
+
+    Parameters
+    ----------
+    frames : tensor (F, M, 3)
+        F snapshots of M atoms (Angstrom).
+    elements : sequence[str], length M
+    q : tensor
+        ``(Q, 3)`` Q-vectors if ``coherent`` (or ``q_is_vector``), else ``(Q,)``
+        magnitudes for the Debye path.
+    coherent : bool
+        If True, average ``|A(Q)|^2`` over frames (oriented).  Else average the
+        Debye powder intensity.
+
+    Returns
+    -------
+    tensor (Q,)
+        Frame-averaged intensity.
+    """
+    import torch
+    frames = torch.as_tensor(frames)
+    if frames.dim() == 2:
+        frames = frames.unsqueeze(0)
+    use_vec = coherent if q_is_vector is None else q_is_vector
+
+    acc = None
+    for fr in frames:
+        if use_vec:
+            I = coherent_intensity(fr, elements, q)
+        else:
+            I = debye_intensity(fr, elements, q, tile=tile)
+        acc = I if acc is None else acc + I
+    return acc / frames.shape[0]

midas_2d/disorder.py

@@ -0,0 +1,166 @@
+"""Transient, anisotropic disorder for few-layer materials.
+
+Schaller/Flanders observe that optical excitation disorders a CdSe nanoplatelet
+*anisotropically* -- the out-of-plane mean-square displacement grows differently
+from the in-plane one, and it evolves with pump-probe delay.  This module gives
+that signature two equivalent faces:
+
+1. **Analytic Debye-Waller** (:class:`AnisotropicMSD`, :class:`TransientMSD`):
+   a differentiable temperature factor ``T(Q) = exp(-1/2 Q . U . Q)`` with a
+   diagonal Cartesian MSD tensor ``U = diag(u_par, u_par, u_perp)`` (Angstrom^2),
+   optionally a function of time.  ``u_par``/``u_perp`` are fittable.
+
+2. **MD-derived** (:func:`msd_tensor_from_frames`): read the *same* tensor
+   straight off an MD trajectory as the covariance of atomic displacements.
+   This closes the loop with :mod:`midas_2d.debye` -- the ensemble spread of
+   coordinates *is* the disorder, and this function recovers the effective ``U``
+   that an analytic DWF would need to reproduce it.
+
+Convention: ``Q`` in 1/A, ``U`` in A^2, platelet normal = +z (the c axis).
+``T`` multiplies the structure-factor *amplitude*; intensity carries ``T^2``.
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:  # pragma: no cover
+    import torch
+
+__all__ = [
+    "dwf_amplitude",
+    "AnisotropicMSD",
+    "TransientMSD",
+    "msd_tensor_from_frames",
+]
+
+
+def dwf_amplitude(q_vec, u_par, u_perp):
+    """Anisotropic Debye-Waller *amplitude* factor ``exp(-1/2 Q.U.Q)``.
+
+    Parameters
+    ----------
+    q_vec : tensor (..., 3)
+        Scattering vectors (1/A).
+    u_par, u_perp : float or tensor
+        In-plane (x, y) and out-of-plane (z) mean-square displacements (A^2).
+
+    Returns
+    -------
+    tensor (...)
+        The amplitude multiplier; square it for the intensity multiplier.
+    """
+    import torch
+    q_vec = torch.as_tensor(q_vec)
+    u_par = torch.as_tensor(u_par, dtype=q_vec.dtype, device=q_vec.device)
+    u_perp = torch.as_tensor(u_perp, dtype=q_vec.dtype, device=q_vec.device)
+    qx, qy, qz = q_vec[..., 0], q_vec[..., 1], q_vec[..., 2]
+    quad = u_par * (qx * qx + qy * qy) + u_perp * (qz * qz)
+    return torch.exp(-0.5 * quad)
+
+
+class AnisotropicMSD:
+    """Fittable static anisotropic MSD tensor (positivity via softplus).
+
+    ``u_par``/``u_perp`` are stored as unconstrained ``raw_*`` parameters; the
+    physical values are ``softplus(raw)`` so optimisation stays unconstrained.
+    """
+    def __init__(self, u_par=1e-3, u_perp=1e-3, *, dtype=None, device=None,
+                 requires_grad=True):
+        import torch
+        dtype = dtype or torch.float64
+        self.raw_par = _inv_softplus(torch.as_tensor(u_par, dtype=dtype, device=device))
+        self.raw_perp = _inv_softplus(torch.as_tensor(u_perp, dtype=dtype, device=device))
+        if requires_grad:
+            self.raw_par = self.raw_par.clone().detach().requires_grad_(True)
+            self.raw_perp = self.raw_perp.clone().detach().requires_grad_(True)
+
+    @property
+    def u_par(self):
+        import torch
+        return torch.nn.functional.softplus(self.raw_par)
+
+    @property
+    def u_perp(self):
+        import torch
+        return torch.nn.functional.softplus(self.raw_perp)
+
+    def parameters(self):
+        return [self.raw_par, self.raw_perp]
+
+    def amplitude(self, q_vec):
+        return dwf_amplitude(q_vec, self.u_par, self.u_perp)
+
+
+class TransientMSD:
+    """Per-delay anisotropic MSD: ``u_par(t)``, ``u_perp(t)`` over T pump-probe
+    delays.  Each delay has its own fittable (softplus-positive) pair, so a full
+    time series of patterns is inverted jointly with shared structure.
+    """
+    def __init__(self, n_delays, u_par0=1e-3, u_perp0=1e-3, *, dtype=None,
+                 device=None, requires_grad=True):
+        import torch
+        dtype = dtype or torch.float64
+        rp = _inv_softplus(torch.full((n_delays,), float(u_par0), dtype=dtype, device=device))
+        rq = _inv_softplus(torch.full((n_delays,), float(u_perp0), dtype=dtype, device=device))
+        if requires_grad:
+            rp = rp.clone().detach().requires_grad_(True)
+            rq = rq.clone().detach().requires_grad_(True)
+        self.raw_par = rp
+        self.raw_perp = rq
+        self.n_delays = int(n_delays)
+
+    @property
+    def u_par(self):
+        import torch
+        return torch.nn.functional.softplus(self.raw_par)
+
+    @property
+    def u_perp(self):
+        import torch
+        return torch.nn.functional.softplus(self.raw_perp)
+
+    def parameters(self):
+        return [self.raw_par, self.raw_perp]
+
+    def amplitude(self, q_vec, delay_index):
+        return dwf_amplitude(q_vec, self.u_par[delay_index], self.u_perp[delay_index])
+
+    def anisotropy(self):
+        """``u_perp / u_par`` per delay -- the out-of-plane/in-plane ratio that
+        is the experimental disordering signature."""
+        return self.u_perp / self.u_par
+
+
+def msd_tensor_from_frames(frames, *, reference=None):
+    """Effective MSD tensor ``U`` (A^2) from an MD trajectory.
+
+    ``U = < (r - <r>) (r - <r>)^T >`` averaged over atoms and frames.  The
+    diagonal entries are ``(u_xx, u_yy, u_zz)``; for a platelet,
+    ``u_zz`` is ``u_perp`` and ``(u_xx + u_yy)/2`` is ``u_par``.
+
+    Parameters
+    ----------
+    frames : tensor (F, M, 3)
+    reference : tensor (M, 3), optional
+        Mean positions; defaults to the per-atom time average.
+
+    Returns
+    -------
+    U : tensor (3, 3)
+    """
+    import torch
+    frames = torch.as_tensor(frames)
+    if frames.dim() == 2:
+        frames = frames.unsqueeze(0)
+    ref = reference if reference is not None else frames.mean(dim=0)  # (M,3)
+    disp = (frames - ref[None]).reshape(-1, 3)                        # (F*M, 3)
+    U = (disp.T @ disp) / disp.shape[0]                               # (3,3)
+    return U
+
+
+def _inv_softplus(y):
+    """Inverse of softplus so that softplus(raw) == y for the init value."""
+    import torch
+    y = torch.as_tensor(y)
+    # raw = log(exp(y) - 1); stable for small/large y.
+    return torch.where(y > 20, y, torch.log(torch.expm1(y.clamp(min=1e-12))))