Core

The differentiable heart: pure tensor-in/tensor-out kernels from structural parameters to simulated intensities and the R-loss. Nothing here imports a parser, an optimizer, or a vendor SDK.

ADP transforms used at the raw-parameter constraint boundary.

diffBloch.core.adp.cholesky_adp(raw_factor: Tensor) → Tensor

Map raw 3x3 factors to symmetric positive-semidefinite ADP matrices.

Anisotropic ADPs are stored as Cholesky factors and expanded as L @ L.T. Only the lower triangle is used, giving the six degrees of freedom of a symmetric PSD matrix and avoiding gauge-redundant upper-triangular parameters.

diffBloch.core.adp.cholesky_raw_from_adp(uij: Tensor) → Tensor

Return a Cholesky factor suitable for initializing cholesky_adp.

This initializer requires positive-definite ADPs, matching torch.linalg.cholesky. Singular positive-semidefinite matrices should be regularized before initialization.

diffBloch.core.adp.isotropic_adp(u_iso: Tensor) → Tensor

Expand isotropic Uiso values to 3x3 Cartesian ADP matrices.

diffBloch.core.adp.cif_adp_to_star(uij_cif: Tensor, reciprocal_lengths: Tensor) → Tensor

Convert CIF-frame anisotropic ADPs Uij_cif to the reciprocal U* frame.

U*_ij = d*_i d*_j Uij_cif with d* = (|a*|, |b*|, |c*|) (reciprocal_lengths, shape (3,)). This is the CIF->Cartesian->star transform with the orthogonalization matrix A cancelled algebraically (A^-1 A D* U D* A^T A^-T = D* U D*), so it depends only on reciprocal_cell. Differentiable in uij_cif; supports a batch axis.

diffBloch.core.adp.cartesian_adp_to_star(uij_cart: Tensor, reciprocal_basis: Tensor) → Tensor

Convert Cartesian-frame ADPs Uij_cart to the reciprocal U* frame.

U* = B Uij_cart B^T with B = reciprocal_cell (rows a*, b*, c*). For an isotropic Cartesian displacement Uij_cart = Uiso I this reduces to the textbook U* = Uiso G* (G* = B B^T the reciprocal metric), so DWF = exp(-2 pi^2 Uiso |g|^2). Building U* directly from B avoids the A^-1 (Uiso I) A^-T route (A per Trueblood et al. 1996, eq. 50), where a c* formed from cross(c, b) = |a*| rather than cross(a, b) = |c*| mislabels |a*| as |c*| and corrupts anisotropic cells. Differentiable in uij_cart; supports a leading batch axis.

diffBloch.core.adp.equivalent_isotropic_adp(uij: Tensor) → Tensor

Return the trace-equivalent isotropic ADP for 3x3 matrices.

diffBloch.core.adp.ueq_from_cif_uij(uij_cif: Tensor, reciprocal_lengths: Tensor, metric_tensor: Tensor) → Tensor

Return the crystallographic equivalent isotropic displacement Ueq for CIF-frame Uij.

Ueq = (1/3) sum_ij Uij_cif d*_i d*_j (a_i . a_j) (Fischer & Tillmanns 1988): the reciprocal U* tensor (cif_adp_to_star()) contracted with the direct-cell metric tensor metric_tensor = cell @ cell.T. Not trace(uij_cif) / 3 – that shortcut only equals Ueq in an orthonormal frame, and is wrong whenever the cell has non-90-degree angles or unequal axes.

Pure tensor constraints and bijectors for raw refinement parameters.

diffBloch.core.constraints.unit_interval(raw: Tensor) → Tensor

Map unconstrained values to the open interval (0, 1).

diffBloch.core.constraints.positive(raw: Tensor) → Tensor

Map unconstrained values to positive values.

diffBloch.core.constraints.apply_symmetry_projection(raw: Tensor, *, projection: Tensor, offset: Tensor) → Tensor

Project raw coordinates onto the site-symmetry-allowed subspace.

This is the public API for the operation; the tiny callable object is kept private so there is one advertised way to apply the projector.

type diffBloch.core.constraints.AdpConstraint = tuple[int, int, int, int, float]

type diffBloch.core.constraints.AdpConstraints = tuple[tuple[AdpConstraint, ...], ...]

diffBloch.core.constraints.apply_adp_constraints(uij: Tensor, constraints: AdpConstraints) → Tensor

Enforce site-symmetry ADP equalities Uij[i,j] = coeff * Uij[src_i, src_j] per atom.

uij is the (N, 3, 3) symmetric ADP tensor in the CIF frame; constraints holds, per atom (one entry each), the tuples (i, j, src_i, src_j, coeff) – meaning Uij[i,j] = coeff * Uij[src_i, src_j] – extracted from the space-group Ueqns (see diffBloch.io.symmetry_setup.symmetry_constraints()). Both are aligned by construction – the caller (diffBloch.params.constrain(), validated once at its boundary) guarantees len(constraints) == N – so this is a pure transform, not a validator. Each right-hand side reads the original (unconstrained) component, so the order of application does not matter, and both [i,j] and [j,i] are set to keep the tensor symmetric. Atoms with no constraints pass through unchanged; dependent raw components simply stop affecting the output (their gradient vanishes), which is the ADP analogue of freezing an over-parameterized coordinate.

diffBloch.core.constraints.diagonal_projection(mask: Tensor, fixed: Tensor) → tuple[Tensor, Tensor]

Build (projection, offset) for the axis-aligned special case from a 0/1 mask.

The diagonal projector P = diag(mask) with offset = fixed * (1 - mask) reproduces the per-coordinate freeze raw * mask + fixed * (1 - mask) – the special case of apply_symmetry_projection() in which every constrained coordinate is fixed to a constant (1 = free, 0 = held fixed). Coupled degrees of freedom (x = y) need the full off-diagonal projector and cannot be expressed this way.

Pure unit-cell and lattice-centering helpers.

diffBloch.core.crystal.cell_matrix_from_parameters(parameters: FloatArray) → FloatArray

Return the fractional-to-Cartesian cell matrix from six cell parameters.

parameters are ordered as (a, b, c, alpha, beta, gamma), with lengths in Angstrom and angles in degrees. Rows of the returned matrix are the real-space basis vectors.

diffBloch.core.crystal.reciprocal_cell(cell: FloatArray) → FloatArray

Return reciprocal lattice vectors as rows in Angstrom^-1.

Uses the ASE-compatible convention reciprocal_cell = pinv(cell).T.

diffBloch.core.crystal.cell_volume(cell: FloatArray) → float

Return the positive unit-cell volume in Angstrom^3.

diffBloch.core.crystal.orientation_basis(cell: FloatArray, orientation: FloatArray) → FloatArray

Lab-frame reciprocal basis for one orientation: reciprocal_cell(cell @ orientation.T).

The single home for the orientation convention. orientation is not guaranteed exactly orthonormal – it carries PETS’s own UB-vs-cell-parameters fit residual (see preprocess.orientation) – so this is NOT reciprocal_basis @ orientation.T. cell rows are the real-space basis; returns rows a*, b*, c* in inverse Angstrom. orientation = I reproduces reciprocal_cell(cell).

diffBloch.core.crystal.reflection_condition(hkl: IntArray, centering: str) → NDArray[bool]

Return the mask of reflections allowed by a lattice-centering rule.

Pure reciprocal-space helpers for Miller-index grids.

diffBloch.core.reciprocal.g_vectors(hkl: IntArray, reciprocal_basis: FloatArray) → FloatArray

Return reciprocal-space vectors for Miller indices and a reciprocal basis.

diffBloch.core.reciprocal.g_vector_lengths(hkl: IntArray, reciprocal_basis: FloatArray) → FloatArray

Return |g| for Miller indices and a reciprocal basis.

diffBloch.core.reciprocal.reciprocal_space_gpts(cell: FloatArray, g_max: float) → tuple[int, int, int]

Return symmetric reciprocal-grid dimensions needed to cover g_max.

diffBloch.core.reciprocal.make_hkl_grid(cell: FloatArray, g_max: float, axes: Sequence[int] = (0, 1, 2)) → IntArray

Return Miller indices whose reciprocal vectors satisfy |g| <= g_max.

axes may restrict the generated Miller dimensions while still filtering in full 3D with omitted axes fixed at zero.

diffBloch.core.reciprocal.gmax_mask(hkl: IntArray, reciprocal_basis: FloatArray, g_max: float) → NDArray[bool]

Return a boolean mask selecting reflections with |g| <= g_max.

diffBloch.core.reciprocal.ravel_hkl(hkl: IntArray, gpts: tuple[int, int, int]) → NDArray[int64]

Map signed Miller indices to flat indices for a centered reciprocal grid.

Symmetry expansion with precomputed ASU membership.

type diffBloch.core.symmetry.DuplicatePolicy = Literal['error', 'keep', 'replace']

class diffBloch.core.symmetry.DuplicateSite(existing_asu_index: int, candidate_asu_index: int, existing_symop_index: int, candidate_symop_index: int)

Bases: object

A candidate ASU expansion site equivalent to an already-planned site.

existing_asu_index: int

candidate_asu_index: int

existing_symop_index: int

candidate_symop_index: int

class diffBloch.core.symmetry.AsuExpansionPlan(asu_indices: Tensor, symop_indices: Tensor, rotations: Tensor, translations: Tensor, n_asu_sites: int, duplicate_sites: tuple[DuplicateSite, ...] = ())

Bases: object

Precomputed ASU membership for differentiable symmetry expansion.

asu_indices: Tensor

symop_indices: Tensor

rotations: Tensor

translations: Tensor

n_asu_sites: int

duplicate_sites: tuple[DuplicateSite, ...] = ()

property n_expanded_sites: int

Number of unique expanded sites in the plan.

class diffBloch.core.symmetry.ExpandedAsu(positions: Tensor, asu_indices: Tensor, symop_indices: Tensor, numbers: Tensor | None = None, uij: Tensor | None = None, occupancies: Tensor | None = None)

Bases: object

Expanded ASU tensors and the membership that produced them.

positions: Tensor

asu_indices: Tensor

symop_indices: Tensor

numbers: Tensor | None = None

uij: Tensor | None = None

occupancies: Tensor | None = None

diffBloch.core.symmetry.build_asu_expansion_plan(frac_positions: NDArray[float64], symops_R: NDArray[float64], symops_t: NDArray[float64], *, symprec: float = 1e-3, on_duplicates: DuplicatePolicy = 'error') → AsuExpansionPlan

Precompute unique ASU/symop memberships for later torch expansion.

Membership is ordered atom-major, symop-minor; duplicate detection is kept out of the differentiable path.

diffBloch.core.symmetry.expand_asu(plan: AsuExpansionPlan, positions: Tensor, *, numbers: Tensor | None = None, uij: Tensor | None = None, occupancies: Tensor | None = None) → ExpandedAsu

Expand ASU tensors using a precomputed membership plan.

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.

type diffBloch.core.scattering.StructureFactorCutoff = Literal['hard', 'taper']

diffBloch.core.scattering.lobato_form_factors(numbers: Tensor, g: Tensor) → Tensor

Electron scattering factor f_e(Z, |g|) for each atom, vectorised over unique Z.

f_e(s) = sum_i a_i (2 + b_i s^2) / (1 + b_i s^2)^2 with s^2 = |g|^2. Returns a real (N_atoms, N_g) tensor (in g’s dtype); a constant with respect to the refinement (depends only on Z and the fixed geometry).

diffBloch.core.scattering.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.

diffBloch.core.scattering.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.

diffBloch.core.scattering.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.

Dynamical-diffraction core: electron-optics primitives + structure-matrix assembly.

primitives holds the NumPy electron-optics constants (setup geometry, not refined); assembly holds the torch differentiable structure-matrix path built on them. Re-exported here so from diffBloch.core.dynamical import ... stays stable across the package split.

class diffBloch.core.dynamical.BeamPlan(gather: StructureFactorGather, mii: Tensor, prefactor: float, diagonal: Tensor, k_n: float, psi0: Tensor, mask: Tensor)

Bases: object

Geometry/numerics-only plan for a beam set (immutable across refinement).

Everything fixed by geometry and beam energy, with no dependence on the refined Fgb: the gather of structure factors onto beam pairs, the symmetrisation factors mii ((N,)), the scalar off-diagonal prefactor, the precomputed real structure-matrix diagonal ((N,) = 2 * k_n * Sg * Mii), the propagation constant k_n, the incident wavefunction psi0 ((N,), 1 at the 000 beam), and the active-beam mask ((N,)). structure_matrix() consumes it with Fgb to produce A; build_bloch_system() wraps that into a BlochSystem for the propagators.

gather: StructureFactorGather

mii: Tensor

prefactor: float

diagonal: Tensor

k_n: float

psi0: Tensor

mask: Tensor

class diffBloch.core.dynamical.BeamPlanBatch(gather: StructureFactorGather, mii: Tensor, prefactor: float, diagonal: Tensor, k_n: float, psi0: Tensor, mask: Tensor)

Bases: object

A stack of BeamPlans over one shared beam set (an orientation’s rocking tilts).

Invariant (enforced by stack_beam_plans()): every plan shares the F-gather, prefactor, k_n, psi0, and mask – only the geometry-per-tilt mii (B, N) and diagonal (B, N) are stacked. build_bloch_systems() consumes it with Fgb to produce a batched BlochSystem (operator (B, N, N)).

gather: StructureFactorGather

mii: Tensor

prefactor: float

diagonal: Tensor

k_n: float

psi0: Tensor

mask: Tensor

class diffBloch.core.dynamical.BlochSystem(a: Tensor, mii: Tensor, psi0: Tensor, k_n: float, mask: Tensor)

Bases: object

A fully specified, propagator-agnostic Bloch-wave system.

Reifies the coupled dynamical-diffraction equations dpsi/dz = (i pi / k_n) A psi with psi(0) = psi0: the structure-matrix operator A ((N, N) complex), its symmetrisation companion mii ((N,) – A is stored Hermitian-symmetrised so eigh applies; mii un-does that to recover physical amplitudes), the incident wavefunction psi0 ((N,)), the propagation constant k_n, and the active-beam mask ((N,)).

Defining invariant: a propagator consumes only a BlochSystem – no geometry, energy, or hkl – which is what makes it a closed system rather than a field bag. See core.solver.

a: Tensor

mii: Tensor

psi0: Tensor

k_n: float

mask: Tensor

class diffBloch.core.dynamical.StructureFactorGather(structure_factor_indices: Tensor, beam_difference_indices: Tensor, n_beams: int, buffer_size: int, gpts: tuple[int, int, int])

Bases: object

Precomputed indices mapping structure factors onto the (N, N) off-diagonal grid.

Geometry-only plan: structure_factor_indices ravel the Fgb support grid and beam_difference_indices ravel the pairwise beam differences hkl_j - hkl_i, both with the same gpts box (the shared-grid contract — one gpts keeps the two from drifting). Consumed by gather_structure_factors(), which scatters Fgb into a flat buffer and indexes it, preserving gradients.

structure_factor_indices: Tensor

beam_difference_indices: Tensor

n_beams: int

buffer_size: int

gpts: tuple[int, int, int]

diffBloch.core.dynamical.build_beam_plan(beam_hkl: IntArray, structure_factor_hkl: IntArray, reciprocal_basis: FloatArray, *, energy: float, gpts: tuple[int, int, int], u0: float = 0.0, gather: StructureFactorGather | None = None, validate: bool = True) → BeamPlan

Precompute the geometry/numerics for a beam set.

beam_hkl (N, 3) selects the beams; structure_factor_hkl (G, 3) / gpts define the Fgb support grid; reciprocal_basis (3, 3) gives g = beam_hkl @ reciprocal_basis. energy (eV) and u0 (mean-inner-potential) set the wavevector. Composes the native primitives into the off-diagonal scale (the no-absorption structure-matrix), the structure-matrix diagonal, and the propagation pieces (k_n, psi0, mask), with psi0 the (hkl == 000) incident-beam convention. mask is all-True here: the beams are the pre-selected active set (per-orientation sg_max selection is deferred).

gather may be a precomputed StructureFactorGather for this exact (structure_factor_hkl, beam_hkl, gpts) – the F-gather is basis-independent, so callers rebuilding many plans over a single beam set (rocking-curve tilts, orientation-search trials) build it once and pass it in, skipping the per-call index-map construction and its validation (the dominant cost). When None it is built here. A cheap shape guard rejects a gather that does not match this beam set / box.

validate is forwarded to build_structure_factor_gather() when building the gather here (default True); pass False on a hot rebuild loop whose grid coverage is guaranteed by an upstream g_max guard. Ignored when a precomputed gather is supplied.

diffBloch.core.dynamical.build_bloch_system(plan: BeamPlan, structure_factors: Tensor, absorption: Absorption = NO_ABSORPTION) → BlochSystem

Assemble the closed Bloch system for a beam plan and structure factors Fgb.

A is built from the differentiable Fgb (so the system is differentiable in Fgb); the symmetrisation factors, incident wavefunction, propagation constant, and active mask are carried straight from the geometry plan. The result is solver-agnostic – see core.solver.propagate().

diffBloch.core.dynamical.build_bloch_systems(batch: BeamPlanBatch, structure_factors: Tensor, absorption: Absorption = NO_ABSORPTION) → BlochSystem

Assemble the batched Bloch system for a beam-plan batch and structure factors Fgb.

Gathers Fgb once onto the shared (N, N) off-diagonal grid, then scales it per tilt with the stacked mii and fills the per-tilt diagonal – yielding a batched operator a (B, N, N). The result is a BlochSystem whose fields carry the batch dim (a (B, N, N), mii (B, N)); core.solver.propagate() is rank-polymorphic over it. Differentiable in Fgb (the shared gather feeds every tilt).

diffBloch.core.dynamical.build_structure_factor_gather(structure_factor_hkl: IntArray, beam_hkl: IntArray, gpts: tuple[int, int, int], *, validate: bool = True, structure_factor_indices: IntArray | None = None) → StructureFactorGather

Precompute the structure-factor gather for a beam set against an Fgb support grid.

structure_factor_hkl (G, 3) are the Miller indices the structure factors are tabulated on; beam_hkl (N, 3) are the selected beams. The pairwise differences hkl_j - hkl_i range to ~2x the beam g_max, so structure_factor_hkl must cover them (the difference-support constraint) — validated here rather than silently gathering zeros. Both sets ravel through the same gpts box (diffBloch.core.reciprocal.ravel_hkl(), which rejects indices outside the box).

validate (default True) runs three O(N^2 / N^2 log G) integrity checks: no duplicate structure_factor_hkl, every beam difference in-box, and structure_factor_hkl covers every difference. They are the dominant cost when rebuilding a gather per trial over a large beam union, and are pure checks – when they pass, the returned indices are identical to skipping them.

validate=False skips them for a hot loop whose grid coverage is guaranteed upstream (a g_max guard). It is not self-guarding: the two skipped classes fail differently. A genuinely out-of-box difference still raises (ravel_hkl, numpy’s terser message). But structure_factor_hkl is the |g| <= g_max sphere – a subset of the rectangular box – so an in-box difference outside that sphere is absent from it yet ravels to a valid box index the scatter never wrote: gather_structure_factors() reads a silent zero, with no runtime backstop. So validate=False is sound only under the upstream coverage guarantee; that guard (not the ravel_hkl check) is what keeps a mis-sized grid from silently gathering zeros.

structure_factor_indices optionally supplies the precomputed grid ravel (grid_source_indices()) so a hot rebuild loop skips re-raveling the (large) support grid every call – it is grid-constant, so one precompute serves every beam set. When None it is raveled here (also validating gpts + the grid box, which a supplied one is trusted to have satisfied).

diffBloch.core.dynamical.energy2sigma(energy: float) → float

Electron interaction parameter sigma in 1/(angstrom*eV) for a beam energy in eV.

sigma = 2 pi m e lambda / h^2 with the relativistic mass m = (1 + E e / (m_e c^2)) m_e (Spence & Zuo 1992). It reproduces the standard values 9.2440e-4 / 7.2884e-4 / 6.5262e-4 at 100 / 200 / 300 keV.

diffBloch.core.dynamical.energy2wavelength(energy: float) → float

Relativistic electron wavelength in angstrom for a beam energy in eV.

lambda = h c / sqrt(E e (2 m_e c^2 + E e)), the relativistic de Broglie wavelength. It reproduces the textbook values 0.03701 / 0.02508 / 0.01969 Å at 100 / 200 / 300 keV.

diffBloch.core.dynamical.excitation_errors(g: FloatArray, energy: float, *, u0: float = 0.0) → FloatArray

Excitation errors Sg (Å^-1) for reciprocal vectors g (Spence & Zuo method).

Sg = (|K|^2 - |K + g|^2) / (2 |K|) with the beam K along -z and magnitude wavevector_magnitude(energy, u0=u0). Measures each reflection’s distance from the Ewald sphere; Sg = 0 exactly at g = 0. g is (N, 3) in Å^-1; returns (N,).

diffBloch.core.dynamical.gather_structure_factors(gather: StructureFactorGather, structure_factors: Tensor) → Tensor

Gather structure factors onto the (N, N) off-diagonal grid, preserving gradients.

structure_factors (G,) is the Fgb tensor aligned with the plan’s structure_factor_hkl order. Scatters it into a flat reciprocal buffer (out-of-place index_add) and indexes the buffer at the beam differences, so out[i, j] = F(beam_j - beam_i). Differentiable in structure_factors.

diffBloch.core.dynamical.grid_source_indices(structure_factor_hkl: IntArray, gpts: tuple[int, int, int]) → IntArray

The raveled structure_factor_hkl source offsets – the grid-constant half of every gather index map.

build_structure_factor_gather ravels the support grid to these offsets on every call, but they depend only on (structure_factor_hkl, gpts) – invariant across beam sets, orientations, and trials. Precompute once and pass via structure_factor_indices= to skip re-raveling the (potentially large) support grid on every per-trial coupled rebuild. Same values build_structure_factor_gather would compute internally.

diffBloch.core.dynamical.mii_factors(g: FloatArray, energy: float, *, u0: float = 0.0) → FloatArray

Diagonal Mii factors that symmetrise the Bloch structure matrix.

Mii = 1 / sqrt(1 - g_z / K_n) with K_n = wavevector_magnitude(energy, u0=u0). The structure matrix uses them on both axes off-diagonal (Mii_i Mii_j) and once on the diagonal. Mii = 1 at g = 0; g is (N, 3) in Å^-1, returns (N,).

diffBloch.core.dynamical.snap_to_standard_energy(energy: float) → float

Snap energy (eV) onto the nearest standard TEM voltage if it’s close enough.

Returns energy unchanged when it is not within _STANDARD_ENERGY_RELATIVE_TOLERANCE of any STANDARD_MICROSCOPE_ENERGIES_EV entry (a genuinely non-standard voltage), so this only removes PETS’s wavelength-rounding noise, never silently reassigns a real, different accelerating voltage.

diffBloch.core.dynamical.stack_beam_plans(plans: Sequence[BeamPlan]) → BeamPlanBatch

Stack beam plans sharing one beam set into a BeamPlanBatch.

Validates the shared-beam-set invariant – identical gather indices, prefactor, k_n, psi0, and mask – then stacks the per-tilt mii / diagonal along a new leading batch axis. Raises if the plans do not share a beam set (the caller passed unrelated plans, not rocking-curve tilts of one orientation). Pure geometry: no dependence on Fgb.

diffBloch.core.dynamical.structure_matrix(plan: BeamPlan, structure_factors: Tensor, absorption: Absorption = NO_ABSORPTION) → Tensor

Assemble the Bloch structure matrix A from a plan and structure factors Fgb.

Off-diagonal A[i,j] = prefactor * Mii_i * Mii_j * F(g_j - g_i) (gathered, then broadcast-scaled); the diagonal is replaced by the precomputed 2 * k_n * Sg_i * Mii_i. Differentiable in structure_factors (the diagonal is a geometry constant). Returns a complex (N, N) tensor in the dtype of structure_factors.

diffBloch.core.dynamical.structure_matrix_prefactor(energy: float) → float

Off-diagonal structure-matrix prefactor sigma / (kappa * lambda * pi).

Scales structure factors into the Bloch structure matrix A. It is energy2sigma(energy) / (kappa * energy2wavelength(energy) * pi).

diffBloch.core.dynamical.wavelength2energy(wavelength: float) → float

Beam energy in eV for a relativistic electron wavelength in angstrom.

Exact algebraic inverse of energy2wavelength(): solving lambda = h c / sqrt(E e (2 m_e c^2 + E e)) for E e gives E e = sqrt((m_e c^2)^2 + (h c / lambda)^2) - m_e c^2. Lets the boundary derive the beam energy the dynamical path needs from the wavelength a PETS file records.

diffBloch.core.dynamical.wavevector_magnitude(energy: float, *, u0: float = 0.0) → float

Corrected wavevector magnitude K_n = sqrt(1/lambda^2 + U0) in Å^-1.

u0 is the mean-inner-potential correction term (added to 1/lambda^2 in Å^-2); u0=0 gives the vacuum wavevector 1/lambda.

Bloch-wave propagators: integrate a BlochSystem to the exit wavefunction.

Two first-class methods, selected by a SolverMethod value (strategy-as-value, not a stateful class):

• matrix_exp – the refine default. psi(t) = matrix_exp(A * i pi t / k_n) @ psi0; a single dense matrix exponential with stable autograd.

• bloch_eigen – eval-only. Diagonalise A once, then every thickness is a cheap phase multiply – fast for many thicknesses, but eigh’s backward is ill-conditioned near degenerate eigenvalues (which symmetric crystals routinely produce), so it is not the refine default.

Both are first-class and swappable off the same BlochSystem (no geometry/energy/hkl needed – the system is the closed problem). They differ in what they return – symmetrised vs physical amplitudes – coinciding only at Mii == 1; that distinction is a feature to experiment with, not a bug. The no-absorption path assumes A is Hermitian.

type diffBloch.core.solver.SolverMethod = Literal['matrix_exp', 'bloch_eigen']

type diffBloch.core.solver.Thicknesses = float | Sequence[float] | Tensor

diffBloch.core.solver.propagate(system: BlochSystem, thicknesses: Thicknesses, *, method: SolverMethod = 'matrix_exp', max_batch: int | None = None) → Tensor

Propagate system.psi0 to each thickness, returning the exit wavefunction.

thicknesses is a scalar or 1-D sequence/tensor (Å). Rank-polymorphic in the operator: a single system (a (N, N)) returns (T, N); a batched system (a (B, N, N), e.g. an orientation’s rocking-curve tilts stacked by core.dynamical.build_bloch_systems()) returns (B, T, N) – one batched eigh / matrix_exp over all tilts. The single-system path is exactly the un-batched computation (the batch axis is simply absent). Differentiable in A (hence in Fgb). method picks the propagator: matrix_exp (refine default, stable autograd) or bloch_eigen (eval-only). The no-absorption path assumes A is Hermitian. The whole propagation runs at float32/complex64.

max_batch (matrix_exp only) caps how many (N, N) operators are exponentiated in one torch.matrix_exp call. None (the default) builds the whole (..., T, N, N) transfer at once. A positive integer instead streams the flattened (B*T, N, N) operator stack in row-blocks, so that transfer is never materialized – the peak drops from ~K*B*T*N**2 to ~K*max_batch*N**2. torch.matrix_exp shares one scaling-and-squaring count across a batch (from its max norm), so regrouping shifts rounding by ~1 ulp: the result matches the unbounded solve to machine precision, never in accuracy. A memory knob, not a result knob. bloch_eigen ignores it (it diagonalises once, no (B, T, N, N) intermediate).

diffBloch.core.solver.memory_safe_max_batch(n_beams: int, *, budget_bytes: int = DEFAULT_MATRIX_EXP_BUDGET_BYTES) → int

The largest max_batch keeping one (max_batch, N, N) matrix_exp block under budget.

N (n_beams) is the solve’s beam count, so the bound adapts to cell size – a large-cell compound (bigger N, cubically heavier propagator) gets a proportionally smaller block, where a fixed block count would still blow up. Returns at least 1 (a single matrix is always attempted, even if it alone exceeds the budget). This is the safe default the engine applies when a caller does not pin max_batch itself; the result matches an unbounded solve to machine precision (a rounding-level ~1 ulp shift, see propagate()), never in accuracy.

Observables and typed product objects bridging propagation to losses.

intensities is the pure observable |psi|^2. On top of it sit three frozen, tensor-carrying product objects:

• BlochSolution – the calculated side: amplitudes/intensities per thickness over a beam set (built from a core.solver.propagate() output).

• PatternBatch – the observed side: measured intensities/sigmas per reflection (built from an io.ExperimentalRecord).

• AlignmentPlan – the precomputed hkl bridge between the two (mirrors BeamPlan: built once from geometry, reused every step), consumed by align() to put calculated and observed on a common reflection axis ready for core.losses.

hkl alignment here is exact (no symmetry merging); symmetry-equivalent merging is deferred.

class diffBloch.core.products.AlignedIntensities(calculated: Tensor, observed: Tensor, sigmas: Tensor)

Bases: object

Calculated/observed intensities and sigmas on a common (T, K) reflection axis.

observed/sigmas are broadcast across the T calculated thicknesses, so the trio drops straight into core.losses (e.g. rbragg(calculated, observed, sigmas)).

calculated: Tensor

observed: Tensor

sigmas: Tensor

class diffBloch.core.products.AlignmentPlan(hkl: Tensor, solution_index: Tensor, pattern_index: Tensor)

Bases: object

Precomputed hkl bridge between a BlochSolution and a PatternBatch.

hkl (K, 3) lists the shared reflections (those observed and calculated, in observed order); solution_index / pattern_index (K,) gather the matching rows from BlochSolution.beam_hkl and PatternBatch.hkl respectively. Geometry-only and reusable.

hkl: Tensor

solution_index: Tensor

pattern_index: Tensor

class diffBloch.core.products.BlochSolution(amplitudes: Tensor, intensities: Tensor, beam_hkl: Tensor, thicknesses: Tensor)

Bases: object

Calculated diffraction over a beam set: amplitudes and intensities per thickness.

amplitudes / intensities are (T, N) (T thicknesses, N beams); beam_hkl is (N, 3); thicknesses is (T,) (Å). Build with from_propagation() from a core.solver.propagate() output.

amplitudes: Tensor

intensities: Tensor

beam_hkl: Tensor

thicknesses: Tensor

classmethod from_propagation(amplitudes: Tensor, beam_hkl: Tensor, thicknesses: Tensor) → Self

Wrap a (T, N) propagated wavefunction, deriving intensities = |amplitudes|^2.

classmethod integrate(solutions: Sequence[BlochSolution], *, reduction: TiltReduction = PLAIN_SUM) → Self

Incoherently reduce tilt sub-solutions into one rocking-curve-integrated solution.

Rocking-curve integration samples N slightly-tilted sub-orientations sharing one beam set and reduces their intensities |psi|^2 over the tilt axis (an incoherent reduction, the physical rotation-frame integration – not their amplitudes). reduction selects the tilt-axis reduction: PlainSum (the default) sums the tilts; MosaicSmoothed applies a moving-average mosaicity broadening first. All sub-solutions must share the beam set (beam_hkl) and thicknesses: the tilts reuse the one nominal beam set, varying only geometry. The integrated observable has no single exit-wave, so amplitudes is stored as the real effective amplitude sqrt(total intensity) (phase is physically lost in an incoherent reduction); only intensities feeds alignment/losses (amplitudes has no downstream consumer). A single-element sequence returns an equivalent solution (the N=1 identity is handled by the caller returning the sub-solution directly).

classmethod integrate_batched(amplitudes: Tensor, beam_hkl: Tensor, thicknesses: Tensor, *, reduction: PlainSum | MosaicSmoothed = PLAIN_SUM) → Self

Reduce a batched (N_tilts, T, N) propagation into one integrated solution.

The batched-solver sibling of integrate(): instead of stacking per-tilt sub-solutions, it takes the stacked exit-wave amplitudes a single batched core.solver.propagate() returns for all tilts at once (leading tilt axis), derives their intensities = |psi|^2, and applies the same tilt-axis reduction. Byte-for-byte equivalent to integrate on the corresponding per-tilt sub-solutions (identical stack, identical _reduce_tilts); only the geometry of how the tilts were solved differs. amplitudes is (N_tilts, T, N) complex; beam_hkl (N, 3); thicknesses (T,) (the shared beam set / thicknesses the tilts co-vary over).

class diffBloch.core.products.MosaicSmoothed(samples: int)

Bases: object

Mosaicity: a sampled moving average over the tilt axis, applied before the sum.

Models crystal mosaic spread by broadening the rocking curve. samples consecutive tilts are averaged, then the smoothed curve is summed. Equivalently the integrated intensity is the sum of the N - samples + 1 local means: padding the smoothed curve back to length N with zeros before summing does not change the sum. samples must not exceed the tilt count N (checked at reduction time).

samples: int

class diffBloch.core.products.PatternBatch(hkl: Tensor, intensities: Tensor, sigmas: Tensor, rotation_index: int = 0)

Bases: object

Observed diffraction intensities: hkl (M, 3), intensities/sigmas (M,).

Build with from_experimental_record() from a validated io.ExperimentalRecord (optionally restricted to one PETS zone-axis row).

hkl: Tensor

intensities: Tensor

sigmas: Tensor

rotation_index: int = 0

classmethod from_experimental_record(record: ExperimentalRecord, *, zone_axis_id: int | None = None, rotation_index: int = 0) → Self

Tensorise observed reflections, optionally filtering to one zone_axis_id.

class diffBloch.core.products.PlainSum

Bases: object

Incoherent sum over the rocking-curve tilts – the default rotation-frame integration.

diffBloch.core.products.align(solution: BlochSolution, pattern: PatternBatch, plan: AlignmentPlan) → AlignedIntensities

Gather calculated and observed intensities onto the plan’s shared reflection axis.

Device-safe: the geometry-only plan indices may live on CPU, so (mirroring the gather/diagonal use sites in core.dynamical) they are moved to each tensor’s device, and observed/sigmas land on calculated.device so the trio is co-located for core.losses.

diffBloch.core.products.build_alignment_plan(solution_hkl: Tensor, pattern_hkl: Tensor, *, restrict_to: Tensor | None = None) → AlignmentPlan

Match observed reflections to calculated beams by exact hkl (observed-order intersection).

restrict_to (S, 3) optionally pins the scored reflection set: only pattern rows whose hkl is in restrict_to are eligible, so the result is pattern ∩ solution ∩ restrict_to. This is how couple_beams keeps scoring on the select_beams selection while the solve set expands to the coupling union – solution_hkl (the union) grows, but the scored axis stays the pre-couple set. It is an intersection, so a restrict_to reflection absent from solution_hkl is dropped: you can only score a reflection you solved, which is the scored ⊆ coupled invariant. None (the default) scores the whole pattern ∩ solution, keeping the tilt-independent path unchanged. pattern_index indexes the full pattern_hkl regardless, so align is untouched.

diffBloch.core.products.intensities(amplitudes: Tensor) → Tensor

Elastic diffracted intensity |psi|^2 of complex exit-wave amplitudes.

Shape-preserving; returns a real tensor (the real dtype matching the complex input). Differentiable in amplitudes (hence back through A / Fgb).

diffBloch.core.products.reduce_tilts(stacked: Tensor, reduction: PlainSum | MosaicSmoothed) → Tensor

Reduce stacked per-tilt intensities (N_tilts, ...) over the leading tilt axis.

The rocking-curve rotation-frame integration: PlainSum sums the tilts; MosaicSmoothed applies a sampled moving average first (the mosaicity broadening). Public because the tilt axis is reduced from two places – a single shared beam set (BlochSolution.integrate() / BlochSolution.integrate_batched()) and the segmented coupling path, which reassembles each reflection’s curve across per-chunk beam sets onto a shared union axis and reduces that (N_tilts, T, N_union) stack here.

Intensity-space loss/metric functions for refinement.

The intensity comparisons used to score and refine a structure. (Position-space metrics belong to the engine/eval layer and are out of scope here.) All functions compare a calculated intensity tensor against an observed one and reduce over the final (reflection) axis, so a (T, N) thickness/orientation batch yields a (T,) loss.

• mse / l1 – generic regression losses.

• rbragg – the crystallographic Bragg R(obs) factor over reflections with I_obs > 3*sigma.

• w_rbragg – the weighted R2 of Klar et al. 2023.

Differentiable in calculated.

diffBloch.core.losses.l1(calculated: Tensor, observed: Tensor) → Tensor

Mean absolute error over the reflection axis.

diffBloch.core.losses.mse(calculated: Tensor, observed: Tensor) → Tensor

Mean squared error over the reflection axis.

diffBloch.core.losses.optimal_scale(calculated: Tensor, observed: Tensor, sigmas: Tensor, *, metric: Callable[[Tensor, Tensor, Tensor], Tensor] = w_rbragg, num_points: int = 100, lo: float = 0.02, hi: float = 2.0) → tuple[Tensor, Tensor]

Grid-search the multiplicative scale on calculated that minimises metric.

The search runs num_points factors in [lo, hi] relative to the total ratio sum(observed)/sum(calculated) (so it is centred near scale 1), evaluates metric(scaled, observed, sigmas) at each, and returns the absolute scale applied to calculated and the minimum metric value. metric defaults to w_rbragg() (the wR2 used to score orientations); it must reduce over the final reflection axis so a (num_points, N) batch yields (num_points,).

diffBloch.core.losses.rbragg(calculated: Tensor, observed: Tensor, sigmas: Tensor) → Tensor

Bragg R(obs) factor sum|sqrt(I_obs) - sqrt(I_calc)| / sum sqrt(I_obs).

Restricted to observed reflections (I_obs > 3*sigma), the standard crystallographic significance cut. Reduces over the reflection axis.

The I_obs > 3*sigma cut is applied by selection (torch.where), not by multiplying a 0/1 mask: experimental intensities can be negative (background-subtracted), so sqrt of an excluded reflection is NaN, and NaN * 0 would poison the sum. Masked-in reflections have I_obs > 3*sigma > 0 (and calculated |psi|^2 >= 0), so their square roots are always finite; the clamps guard only against numerical noise. A 0/1 multiply-mask would be NaN-unsafe on the negative excluded intensities, which is why selection is used instead.

diffBloch.core.losses.w_rbragg(calculated: Tensor, observed: Tensor, sigmas: Tensor, *, mu: float = 0.01) → Tensor

Weighted R2 sqrt( sum w*(I_calc - I_obs)^2 / sum (w*I_obs)^2 ) (Klar et al. 2023).

w = 1 / sqrt( sigma(sqrt(I_obs))^2 + (mu*sqrt(I_obs))^2 ) with mu the instability factor; weak reflections (I_obs < 0.01*sigma) use the 5*sqrt(sigma) floor from the Klar et al. 2023 supplementary information.