"""Electron structure factors (Lobato parametrization), vectorised.
The elastic structure-factor path -- form factors, Debye-Waller factors, and the phase sum --
vectorised over unique-Z groups into a single batched phase sum.
Form factors use the Lobato–Van Dyck (2014) parametrization (coefficients vendored in
``core/data/lobato.json``), so no external scattering library is needed at runtime. The form factor
is a setup constant — only ``positions``, ``uij_star`` and ``occupancies`` carry gradients.
Absorption (the imaginary ``U0'`` path) is intentionally deferred (runs set ``absorption: false``).
``structure_factors`` consumes ADPs already in the U* (reciprocal) frame; the Cartesian→U*
conversion belongs to the ADP/spec layer that wires this into the engine.
"""
from __future__ import annotations
import json
from functools import lru_cache
from importlib import resources
from typing import Literal
import torch
from torch import Tensor
from diffBloch.core.absorption import absorptive_form_factors, equivalent_isotropic_b
from diffBloch.specs import NO_ABSORPTION, Absorption
type StructureFactorCutoff = Literal["hard", "taper"]
def _g_vector_lengths(hkl: Tensor, reciprocal_basis: Tensor) -> Tensor:
"""``|g|`` per reflection from Miller indices and a reciprocal basis (rows = a*, b*, c*).
The torch counterpart of the NumPy ``core.reciprocal.g_vector_lengths`` (which serves the
planning path); kept here so the differentiable structure-factor path stays in torch.
"""
if hkl.ndim != 2 or hkl.shape[1] != 3:
raise ValueError("hkl must have shape (M, 3)")
if reciprocal_basis.shape != (3, 3):
raise ValueError("reciprocal_basis must have shape (3, 3)")
g_vectors = hkl.to(reciprocal_basis.dtype) @ reciprocal_basis
lengths: Tensor = torch.linalg.vector_norm(g_vectors, dim=1)
return lengths
@lru_cache(maxsize=1)
def _lobato_table() -> dict[int, tuple[tuple[float, ...], tuple[float, ...]]]:
"""Load the vendored Lobato coefficients keyed by atomic number: ``{Z: ((a1..a5), (b1..b5))}``.
The table is keyed by Z, not element symbol, so the core needs no symbol authority — element
identity comes from the parsed atomic ``numbers`` plus the vendored data alone, keeping
``core/`` free of any parser/periodic-table dependency.
"""
text = (resources.files("diffBloch.core") / "data" / "lobato.json").read_text()
raw = json.loads(text)
return {int(z): (tuple(ab[0]), tuple(ab[1])) for z, ab in raw.items()}
[docs]
def debye_waller_factor(hkl: Tensor, uij_star: Tensor) -> Tensor:
"""Anisotropic Debye–Waller factor ``exp(-2 pi^2 h^T U* h)`` per (atom, reflection).
``hkl`` is ``(M, 3)``; ``uij_star`` is ``(N, 3, 3)`` in the U* (reciprocal) frame. Returns
``(N, M)``, differentiable in ``uij_star``.
"""
if hkl.ndim != 2 or hkl.shape[1] != 3:
raise ValueError("hkl must have shape (M, 3)")
if uij_star.ndim != 3 or tuple(uij_star.shape[1:]) != (3, 3):
raise ValueError("uij_star must have shape (N, 3, 3)")
h = hkl.to(uij_star.dtype)
quadratic = torch.einsum("rx,axy,ry->ar", h, uij_star, h)
return torch.exp(-2.0 * torch.pi**2 * quadratic)
[docs]
def structure_factor_cutoff(
g: Tensor, g_max: float, *, mode: StructureFactorCutoff = "hard"
) -> Tensor:
"""Reflection resolution cutoff: a hard ``|g| <= g_max`` mask or a logistic taper window."""
if g_max <= 0.0:
raise ValueError("g_max must be positive")
if mode == "hard":
return (g <= g_max).to(g.dtype)
if mode == "taper":
# Logistic taper window: roll-off centred at TAPER_ALPHA * g_max with
# logistic width TAPER_WIDTH. TAPER_ALPHA = 1 - 0.05 starts the roll-off ~5% below g_max.
taper_width, taper_alpha = 0.005, 1.0 - 0.05
return 1.0 / (1.0 + torch.exp((g / g_max - taper_alpha) / taper_width))
raise ValueError(f"unsupported cutoff mode: {mode!r}")
[docs]
def structure_factors(
positions: Tensor,
numbers: Tensor,
occupancies: Tensor,
uij_star: Tensor,
hkl: Tensor,
reciprocal_basis: Tensor,
cell_volume: float,
*,
g_max: float,
cutoff: StructureFactorCutoff = "hard",
zero_threshold: float = 1e-12,
absorption: Absorption = NO_ABSORPTION,
energy: float | None = None,
) -> Tensor:
"""Vectorised electron structure factors ``Fgb``, optionally including absorption.
``Fgb(h) = (1/V) sum_atoms f_e * DWF * occ * cutoff * exp(2 pi i r . h)``. ``|g|`` is derived
internally from ``hkl`` and ``reciprocal_basis`` (no separate ``g`` argument to keep in sync).
Differentiable in ``positions``, ``uij_star``, ``occupancies``; ``f_e`` is a constant form
factor. Symmetry-related atoms can sum a component to exactly zero mathematically (systematic
absences); floating-point roundoff lands near but not at zero, so components below
``zero_threshold`` are snapped to a clean ``0.0``. Returns a complex ``(M,)`` tensor.
"""
n_atoms = positions.shape[0]
if positions.ndim != 2 or positions.shape[1] != 3:
raise ValueError("positions must have shape (N, 3)")
if numbers.shape != (n_atoms,) or occupancies.shape != (n_atoms,):
raise ValueError("numbers and occupancies must have shape (N,) matching positions")
if uij_star.shape != (n_atoms, 3, 3):
raise ValueError("uij_star must have shape (N, 3, 3) matching positions")
if cell_volume <= 0.0:
raise ValueError("cell_volume must be positive")
g = _g_vector_lengths(hkl, reciprocal_basis)
form_factors = lobato_form_factors(numbers, g)
dwf = debye_waller_factor(hkl, uij_star)
cutoff_window = structure_factor_cutoff(g, g_max, mode=cutoff)
atomic_factors = form_factors
if absorption.enabled:
if energy is None:
raise ValueError("energy is required for parameterized absorption")
b_iso = equivalent_isotropic_b(uij_star, reciprocal_basis)
atomic_factors = torch.complex(
form_factors,
absorptive_form_factors(numbers, g / 2.0, b_iso, energy=energy),
)
per_atom = (atomic_factors * dwf * occupancies[:, None]) * cutoff_window[None, :]
phase = torch.exp(2.0j * torch.pi * (positions @ hkl.to(positions.dtype).transpose(0, 1)))
unmasked = (per_atom.to(phase.dtype) * phase).sum(dim=0)
real = torch.where(
unmasked.real.abs() >= zero_threshold, unmasked.real, torch.zeros_like(unmasked.real)
)
imag = torch.where(
unmasked.imag.abs() >= zero_threshold, unmasked.imag, torch.zeros_like(unmasked.imag)
)
return torch.complex(real, imag) / cell_volume