torch-scattering 0.6.0__py3-none-any.whl

torch_scattering/__init__.py

@@ -0,0 +1,32 @@
+"""Multislice electron scattering simulation in PyTorch."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from torch_scattering._core import (
+    chunk_slices,
+    fresnel_propagator,
+    interaction_parameter,
+    multislice_step,
+    transmission_function,
+)
+from torch_scattering.firstborn import firstborn
+from torch_scattering.multislice import multislice
+from torch_scattering.projection import projection
+from torch_scattering.rytov import rytov
+
+try:
+    __version__ = version("torch-scattering")
+except PackageNotFoundError:
+    __version__ = "uninstalled"
+
+__all__ = [
+    "chunk_slices",
+    "firstborn",
+    "fresnel_propagator",
+    "interaction_parameter",
+    "multislice",
+    "multislice_step",
+    "projection",
+    "rytov",
+    "transmission_function",
+]

torch_scattering/_core.py

@@ -0,0 +1,309 @@
+"""Pure-math primitives for scattering algorithms."""
+
+import math
+from numbers import Real
+
+import torch
+from scipy import constants as C
+from torch_ctf import calculate_relativistic_electron_wavelength
+from torch_grid_utils import fftfreq_grid
+
+_SUPPORTED_POTENTIAL_DTYPES = {
+    torch.float32,
+    torch.float64,
+    torch.complex64,
+    torch.complex128,
+}
+
+
+def _complex_dtype(dtype: torch.dtype) -> torch.dtype:
+    """Return the complex dtype corresponding to a supported potential dtype."""
+    if dtype in {torch.float32, torch.complex64}:
+        return torch.complex64
+    return torch.complex128
+
+
+def _validate_scalar(
+    value: float | torch.Tensor,
+    name: str,
+) -> None:
+    """Validate a positive, finite scalar without replacing differentiable tensors."""
+    if isinstance(value, torch.Tensor):
+        if value.ndim != 0:
+            raise ValueError(f"{name} must be a scalar")
+        if value.is_complex():
+            raise ValueError(f"{name} must be real")
+        if not bool(torch.isfinite(value).item()) or not bool((value > 0).item()):
+            raise ValueError(f"{name} must be positive and finite")
+    elif (
+        not isinstance(value, Real)
+        or isinstance(value, bool)
+        or not math.isfinite(float(value))
+        or value <= 0
+    ):
+        raise ValueError(f"{name} must be a positive, finite scalar")
+
+
+def _validate_propagation_inputs(
+    potential: torch.Tensor,
+    pixel_size: float | torch.Tensor,
+    voltage: float | torch.Tensor,
+) -> None:
+    """Validate inputs shared by all high-level propagation modes."""
+    if not isinstance(potential, torch.Tensor):
+        raise TypeError("potential must be a torch.Tensor")
+    if potential.ndim < 3:
+        raise ValueError("potential must have shape (..., Z, H, W)")
+    if any(size == 0 for size in potential.shape[-3:]):
+        raise ValueError("potential spatial dimensions Z, H, and W must be nonempty")
+    if potential.dtype not in _SUPPORTED_POTENTIAL_DTYPES:
+        raise TypeError(
+            "potential dtype must be float32, float64, complex64, or complex128"
+        )
+    _validate_scalar(pixel_size, "pixel_size")
+    _validate_scalar(voltage, "voltage")
+
+
+def fresnel_propagator(
+    frequency_grid: torch.Tensor,
+    wavelength: float | torch.Tensor,
+    dz: float | torch.Tensor,
+) -> torch.Tensor:
+    """
+    Compute the Fresnel free-space propagator for one multislice step.
+
+    Parameters
+    ----------
+    frequency_grid : torch.Tensor
+        Real-valued grid of spatial frequency magnitudes (1/Angstrom), e.g.
+        from ``torch_grid_utils.fftfreq_grid(..., norm=True)``.
+    wavelength : float | torch.Tensor
+        Relativistic electron wavelength in Angstroms.
+    dz : float | torch.Tensor
+        Propagation distance (slice thickness) in Angstroms.
+
+    Returns
+    -------
+    torch.Tensor
+        Complex-valued Fresnel propagator with the same shape as
+        `frequency_grid`.
+
+    Notes
+    -----
+    ``H(k) = exp(i * pi * wavelength * dz * k^2)``. This is a pure phase
+    transfer function, so ``|H(k)| == 1`` for all `k`.
+
+    References
+    ----------
+    .. [1] E. J. Kirkland, Advanced Computing in Electron Microscopy,
+       Springer US, Boston, MA, 2010.
+    """
+    return torch.exp(1j * C.pi * wavelength * dz * frequency_grid**2)
+
+
+def transmission_function(
+    potential_slice: torch.Tensor,
+    sigma: float | torch.Tensor,
+    dz: float | torch.Tensor,
+) -> torch.Tensor:
+    """
+    Compute the transmission function for one electrostatic-potential slice.
+
+    Parameters
+    ----------
+    potential_slice : torch.Tensor
+        Electrostatic potential in volts for a single slice. Real-valued for
+        a non-absorbing specimen, or complex-valued with a positive imaginary
+        part to model absorption.
+    sigma : float | torch.Tensor
+        Electron-specimen interaction parameter, e.g. from
+        `interaction_parameter`, in rad/(V*Angstrom).
+    dz : float | torch.Tensor
+        Slice thickness in Angstroms.
+
+    Returns
+    -------
+    torch.Tensor
+        Complex-valued transmission function, same shape as
+        `potential_slice`.
+
+    Notes
+    -----
+    ``T = exp(i * sigma * dz * V)``. The phase shift imparted by a slice is
+    the line integral of `sigma * V` along the beam direction through the
+    slice; treating `V` as constant over the (thin) slice thickness reduces
+    that integral to `sigma * V * dz`. For real `V` this makes `T` a pure
+    phase mask (``|T| == 1``); a positive imaginary part of `V` attenuates
+    the wave.
+    """
+    return torch.exp(1j * sigma * dz * potential_slice)
+
+
+def multislice_step(
+    wave: torch.Tensor,
+    potential_slice: torch.Tensor,
+    propagator: torch.Tensor,
+    sigma: float | torch.Tensor,
+    dz: float | torch.Tensor,
+) -> torch.Tensor:
+    """
+    Advance the electron wave through one electrostatic-potential slice.
+
+    Parameters
+    ----------
+    wave : torch.Tensor
+        Complex-valued incident wave for this slice, shape (..., H, W).
+    potential_slice : torch.Tensor
+        Electrostatic potential in volts for this slice, shape (..., H, W).
+    propagator : torch.Tensor
+        Fresnel propagator for this slice thickness, e.g. from
+        `fresnel_propagator`, shape (H, W).
+    sigma : float | torch.Tensor
+        Electron-specimen interaction parameter, in rad/(V*Angstrom).
+    dz : float | torch.Tensor
+        Slice thickness in Angstroms.
+
+    Returns
+    -------
+    torch.Tensor
+        Complex-valued wave after transmission through and propagation
+        past this slice, same shape as `wave`.
+
+    Notes
+    -----
+    Each step alternates transmission through the slice (real space) with
+    Fresnel propagation to the next slice (Fourier space):
+    ``wave_out = IFFT(FFT(wave * T) * H)``
+    """
+    transmitted = wave * transmission_function(potential_slice, sigma, dz)
+    propagated_wave: torch.Tensor = torch.fft.ifft2(
+        torch.fft.fft2(transmitted) * propagator
+    )
+    return propagated_wave
+
+
+def chunk_slices(total_slices: int, n_chunks: int) -> list[int]:
+    """
+    Compute contiguous, back-loaded chunk sizes for coarsened multislice.
+
+    Parameters
+    ----------
+    total_slices : int
+        Total number of slices along the beam direction.
+    n_chunks : int
+        Desired number of chunks. Must satisfy ``0 < n_chunks <= total_slices``.
+
+    Returns
+    -------
+    list[int]
+        Chunk sizes summing to `total_slices`, length `n_chunks`, suitable
+        for `torch.split(..., split_size_or_sections=sizes, dim=...)`. Every
+        chunk has ``total_slices // n_chunks`` slices except the last,
+        which absorbs the remainder ``total_slices % n_chunks``. This is a
+        back-loaded remainder, unlike `torch.chunk`/`torch.tensor_split`,
+        which front-load it.
+
+    Raises
+    ------
+    ValueError
+        If `n_chunks` is not in ``(0, total_slices]``.
+
+    Examples
+    --------
+    >>> chunk_slices(10, 3)
+    [3, 3, 4]
+    """
+    if not (0 < n_chunks <= total_slices):
+        raise ValueError(
+            "n_chunks must satisfy 0 < n_chunks <= total_slices "
+            f"({total_slices}), got {n_chunks}"
+        )
+    base = total_slices // n_chunks
+    sizes = [base] * n_chunks
+    sizes[-1] += total_slices - base * n_chunks
+    return sizes
+
+
+def interaction_parameter(voltage: float | torch.Tensor) -> float | torch.Tensor:
+    """
+    Calculate the electron-specimen interaction parameter.
+
+    Computes the interaction constant sigma for electron scattering,
+    following Kirkland Eq. (5.6). Wavelength is computed internally from
+    `voltage` via `torch_ctf.calculate_relativistic_electron_wavelength`.
+
+    Parameters
+    ----------
+    voltage : float | torch.Tensor
+        Electron beam acceleration voltage in kilovolts.
+
+    Returns
+    -------
+    float | torch.Tensor
+        Interaction parameter sigma in units of rad/(V*Angstrom). A tensor
+        if `voltage` is a tensor, otherwise a plain float.
+
+    References
+    ----------
+    .. [1] E. J. Kirkland, Advanced Computing in Electron Microscopy,
+       Eq. (5.6), Springer US, Boston, MA, 2010.
+    """
+    wavelength_m = calculate_relativistic_electron_wavelength(voltage * 1.0e3)
+    wavelength = wavelength_m * 1.0e10  # meters -> Angstroms
+    # Electron rest mass energy, m0*c^2, in electronvolts.
+    rest_ev = C.electron_mass * C.speed_of_light**2 / C.elementary_charge
+    ev = voltage * 1.0e3  # [eV]
+    sigma: float | torch.Tensor = (
+        2.0 * C.pi / (wavelength * ev) * ((ev + rest_ev) / (ev + 2.0 * rest_ev))
+    )
+    return sigma
+
+
+def _prepare_propagation_parameters(
+    potential: torch.Tensor,
+    pixel_size: float | torch.Tensor,
+    voltage: float | torch.Tensor,
+) -> tuple[torch.Tensor, float | torch.Tensor, torch.Tensor]:
+    """
+    Compute the wavelength, interaction parameter, and frequency grid.
+
+    Shared setup for the high-level wrapper functions (`firstborn`,
+    `rytov`, `multislice`).
+
+    Parameters
+    ----------
+    potential : torch.Tensor
+        Real or complex 3D electrostatic potential in volts, shape
+        (..., Z, H, W).
+    pixel_size : float | torch.Tensor
+        Isotropic voxel spacing in Angstroms, used as both the in-plane pixel
+        spacing and slice thickness.
+    voltage : float | torch.Tensor
+        Electron beam acceleration voltage in kilovolts.
+
+    Returns
+    -------
+    tuple[torch.Tensor, float | torch.Tensor, torch.Tensor]
+        `wavelength` in Angstroms, `sigma` (interaction parameter) in
+        rad/(V*Angstrom), and `frequency_grid`, the real-valued grid of
+        spatial frequency magnitudes (1/Angstrom), shape (H, W).
+    """
+    _validate_propagation_inputs(potential, pixel_size, voltage)
+    wavelength_m = calculate_relativistic_electron_wavelength(voltage * 1.0e3)
+    wavelength = wavelength_m * 1.0e10  # meters -> Angstroms
+    sigma = interaction_parameter(voltage=voltage)
+
+    height, width = potential.shape[-2:]
+    spacing = (
+        float(pixel_size.detach().cpu())
+        if isinstance(pixel_size, torch.Tensor)
+        else float(pixel_size)
+    )
+    frequency_grid = fftfreq_grid(
+        image_shape=(height, width),
+        rfft=False,
+        spacing=spacing,
+        norm=True,
+        device=potential.device,
+    ).to(dtype=potential.real.dtype)
+    return wavelength, sigma, frequency_grid

torch_scattering/firstborn.py

@@ -0,0 +1,97 @@
+"""High-level first Born approximation wrapper.
+
+Wires up frequency grids, wavelength, and a per-slice Fresnel propagator
+stack on top of the pure-math primitives in `_core`.
+"""
+
+import torch
+
+from torch_scattering._core import (
+    _prepare_propagation_parameters,
+    chunk_slices,
+    fresnel_propagator,
+)
+
+
+def firstborn(
+    potential: torch.Tensor,
+    pixel_size: float | torch.Tensor,
+    voltage: float | torch.Tensor,
+    n_slices: int | None = None,
+) -> torch.Tensor:
+    """
+    Compute the exit wave using the first Born approximation.
+
+    Parameters
+    ----------
+    potential : torch.Tensor
+        Real-valued electrostatic potential in volts for a non-absorbing
+        specimen, or complex-valued potential when modelling absorption.
+        Shape (..., Z, H, W), where Z is the beam direction.
+    pixel_size : float | torch.Tensor
+        Positive finite isotropic voxel spacing in Angstroms. This is both
+        the in-plane pixel spacing and slice thickness.
+    voltage : float | torch.Tensor
+        Electron beam acceleration voltage in kilovolts (e.g. 300 for 300 kV).
+    n_slices : int | None
+        Number of terms to sum in the Born series. If `None` (default),
+        every slice of `potential` is treated individually (equivalent to
+        `n_slices=Z`). If smaller than `Z`, slices are grouped into
+        `n_slices` contiguous chunks via `chunk_slices` (all but the last
+        chunk have the same size; the last absorbs the remainder), and each
+        chunk's summed potential is propagated as one thicker slab. Must
+        satisfy ``0 < n_slices <= Z``.
+
+    Returns
+    -------
+    torch.Tensor
+        Complex-valued 2D exit wave, shape (..., H, W).
+
+    Notes
+    -----
+    Each chunk's scattered contribution is propagated directly to the exit
+    plane and summed, rather than being propagated incrementally
+    slice-by-slice as in `multislice`:
+    ``psi = 1 + sum_c IFFT(FFT(i * sigma * dz * V_c) * H_c)``, where `V_c`
+    is the summed potential of chunk `c` and `H_c` is the Fresnel
+    propagator for the distance from chunk `c` to the exit plane. This is
+    the linearised (first-order) term of the exponentiated Born series
+    used by `rytov`, valid only for weak potentials (``sigma * dz * V <<
+    1``). Faster than `multislice` (one FFT pass instead of one per chunk)
+    but less accurate for thick or strongly scattering specimens, and less
+    accurate than `rytov` for the same potential since it does not
+    exponentiate the scattered sum.
+
+    References
+    ----------
+    .. [1] E. J. Kirkland, Advanced Computing in Electron Microscopy,
+       Springer US, Boston, MA, 2010.
+    """
+    wavelength, sigma, frequency_grid = _prepare_propagation_parameters(
+        potential, pixel_size, voltage
+    )
+
+    total_slices = potential.shape[-3]
+    if n_slices is None:
+        n_slices = total_slices
+    sizes = chunk_slices(total_slices, n_slices)
+    chunk_potential = torch.stack(
+        [chunk.sum(dim=-3) for chunk in torch.split(potential, sizes, dim=-3)],
+        dim=-3,
+    )  # (..., n_slices, H, W)
+
+    # Chunk c is (sum of sizes[c:]) slice thicknesses from the exit plane:
+    # the wave still has to cross the remainder of chunk c itself plus
+    # every chunk after it.
+    size_tensor = torch.as_tensor(sizes, device=potential.device)
+    remaining_slices = size_tensor.flip(0).cumsum(0).flip(0)
+    distances = remaining_slices * pixel_size
+    propagators = fresnel_propagator(
+        frequency_grid, wavelength=wavelength, dz=distances[:, None, None]
+    )  # (n_slices, H, W)
+
+    scattered = torch.fft.ifft2(
+        torch.fft.fft2(1j * sigma * pixel_size * chunk_potential) * propagators
+    )
+    exit_wave: torch.Tensor = 1 + scattered.sum(dim=-3)
+    return exit_wave

torch_scattering/multislice.py

@@ -0,0 +1,95 @@
+"""High-level multislice wrapper.
+
+Wires up frequency grids, wavelength, and the per-slice loop on top of the
+pure-math primitives in `_core`.
+"""
+
+import torch
+
+from torch_scattering._core import (
+    _complex_dtype,
+    _prepare_propagation_parameters,
+    chunk_slices,
+    fresnel_propagator,
+    multislice_step,
+)
+
+
+def multislice(
+    potential: torch.Tensor,
+    pixel_size: float | torch.Tensor,
+    voltage: float | torch.Tensor,
+    n_slices: int | None = None,
+) -> torch.Tensor:
+    """
+    Compute the 2D exit wave from a 3D electrostatic potential.
+
+    Parameters
+    ----------
+    potential : torch.Tensor
+        Real-valued electrostatic potential in volts for a non-absorbing
+        specimen, or complex-valued potential when modelling absorption.
+        Shape (..., Z, H, W), where Z is the beam direction.
+    pixel_size : float | torch.Tensor
+        Positive finite isotropic voxel spacing in Angstroms. This is both
+        the in-plane pixel spacing and slice thickness.
+    voltage : float | torch.Tensor
+        Electron beam acceleration voltage in kilovolts (e.g. 300 for 300 kV).
+    n_slices : int | None
+        Number of multislice steps to take. If `None` (default), every
+        slice of `potential` is propagated individually (equivalent to
+        `n_slices=Z`). If smaller than `Z`, slices are grouped into
+        `n_slices` contiguous chunks via `chunk_slices` (all but the last
+        chunk have the same size; the last absorbs the remainder), each
+        chunk's potential is summed, and propagated as one thicker slab.
+        Must satisfy ``0 < n_slices <= Z``.
+
+    Returns
+    -------
+    torch.Tensor
+        Complex-valued 2D exit wave, shape (..., H, W).
+
+    Notes
+    -----
+    The incident wave entering the specimen is a uniform (unit-amplitude,
+    zero-phase) plane wave. Each chunk's summed potential is transmitted
+    with `dz=pixel_size` (the per-voxel Riemann-sum weight) and then
+    propagated with `dz=chunk_size * pixel_size` (the chunk's physical
+    thickness). `n_slices=None` propagates one true slice at a time and
+    is the most accurate setting; grouping trades accuracy for speed.
+
+    `torch_grid_utils.fftfreq_grid` only accepts a plain float for its
+    `spacing` argument, so if `pixel_size` is a tensor, the frequency grid
+    itself is built from a detached float copy; `pixel_size` still flows
+    as a tensor into the propagator and transmission function, so gradients
+    with respect to it are not lost, just not routed through the grid
+    spacing.
+    """
+    wavelength, sigma, frequency_grid = _prepare_propagation_parameters(
+        potential, pixel_size, voltage
+    )
+
+    height, width = potential.shape[-2:]
+    wave = torch.ones(
+        (*potential.shape[:-3], height, width),
+        dtype=_complex_dtype(potential.dtype),
+        device=potential.device,
+    )
+    total_slices = potential.shape[-3]
+    if n_slices is None:
+        n_slices = total_slices
+    sizes = chunk_slices(total_slices, n_slices)
+    # At most 2 distinct chunk sizes occur; build one propagator per size.
+    propagators = {
+        size: fresnel_propagator(
+            frequency_grid, wavelength=wavelength, dz=pixel_size * size
+        )
+        for size in set(sizes)
+    }
+    for chunk in torch.split(potential, sizes, dim=-3):
+        chunk_size = chunk.shape[-3]
+        potential_chunk = chunk.sum(dim=-3)
+        wave = multislice_step(
+            wave, potential_chunk, propagators[chunk_size], sigma, pixel_size
+        )
+    return wave

torch_scattering/projection.py

@@ -0,0 +1,55 @@
+"""High-level projection-approximation wrapper.
+
+Wires up the interaction parameter and slice summation on top of the
+pure-math primitives in `_core`.
+"""
+
+import torch
+
+from torch_scattering._core import (
+    _validate_propagation_inputs,
+    interaction_parameter,
+    transmission_function,
+)
+
+
+def projection(
+    potential: torch.Tensor,
+    pixel_size: float | torch.Tensor,
+    voltage: float | torch.Tensor,
+) -> torch.Tensor:
+    """
+    Compute the exit wave using the projection approximation.
+
+    Parameters
+    ----------
+    potential : torch.Tensor
+        Real-valued electrostatic potential in volts for a non-absorbing
+        specimen, or complex-valued potential when modelling absorption.
+        Shape (..., Z, H, W), where Z is the beam direction.
+    pixel_size : float | torch.Tensor
+        Positive finite isotropic voxel spacing in Angstroms. This is both
+        the in-plane pixel spacing and slice thickness.
+    voltage : float | torch.Tensor
+        Electron beam acceleration voltage in kilovolts (e.g. 300 for 300 kV).
+
+    Returns
+    -------
+    torch.Tensor
+        Complex-valued 2D exit wave, shape (..., H, W).
+
+    Notes
+    -----
+    This wave-propagation approximation numerically sums a sampled 3D
+    potential along the beam direction, then treats the specimen as
+    infinitely thin and ignores Fresnel propagation between slices:
+    ``psi = exp(i * sigma * dz * sum_z V(z))``. This is valid only when the
+    specimen is thin enough that propagation effects within it are
+    negligible. It is distinct from an analytic projected-potential
+    calculation and from projection alignment used when fitting structures
+    into maps.
+    """
+    _validate_propagation_inputs(potential, pixel_size, voltage)
+    sigma = interaction_parameter(voltage=voltage)
+    projected_potential = potential.sum(dim=-3)
+    return transmission_function(projected_potential, sigma=sigma, dz=pixel_size)

torch_scattering/rytov.py

@@ -0,0 +1,94 @@
+"""High-level Rytov-approximation wrapper.
+
+Wires up frequency grids, wavelength, and a per-slice Fresnel propagator
+stack on top of the pure-math primitives in `_core`.
+"""
+
+import torch
+
+from torch_scattering._core import (
+    _prepare_propagation_parameters,
+    chunk_slices,
+    fresnel_propagator,
+)
+
+
+def rytov(
+    potential: torch.Tensor,
+    pixel_size: float | torch.Tensor,
+    voltage: float | torch.Tensor,
+    n_slices: int | None = None,
+) -> torch.Tensor:
+    """
+    Compute the exit wave using the Rytov approximation.
+
+    Parameters
+    ----------
+    potential : torch.Tensor
+        Real-valued electrostatic potential in volts for a non-absorbing
+        specimen, or complex-valued potential when modelling absorption.
+        Shape (..., Z, H, W), where Z is the beam direction.
+    pixel_size : float | torch.Tensor
+        Positive finite isotropic voxel spacing in Angstroms. This is both
+        the in-plane pixel spacing and slice thickness.
+    voltage : float | torch.Tensor
+        Electron beam acceleration voltage in kilovolts (e.g. 300 for 300 kV).
+    n_slices : int | None
+        Number of terms to sum in the exponentiated Born series. If `None`
+        (default), every slice of `potential` is treated individually
+        (equivalent to `n_slices=Z`). If smaller than `Z`, slices are
+        grouped into `n_slices` contiguous chunks via `chunk_slices` (all
+        but the last chunk have the same size; the last absorbs the
+        remainder), and each chunk's summed potential is propagated as one
+        thicker slab. Must satisfy ``0 < n_slices <= Z``.
+
+    Returns
+    -------
+    torch.Tensor
+        Complex-valued 2D exit wave, shape (..., H, W).
+
+    Notes
+    -----
+    Each chunk's scattered contribution is propagated directly to the exit
+    plane and summed inside a single exponential, rather than being
+    propagated incrementally slice-by-slice as in `multislice`:
+    ``psi = exp(sum_c IFFT(FFT(i * sigma * dz * V_c) * H_c))``, where `V_c`
+    is the summed potential of chunk `c` and `H_c` is the Fresnel
+    propagator for the distance from chunk `c` to the exit plane. This is
+    the exponentiated Born series, reducing to the first Born approximation
+    for small `V`. Faster than `multislice` (one FFT pass instead of one
+    per chunk) but less accurate for thick specimens. `n_slices=None` is
+    the most accurate setting; grouping trades accuracy for speed.
+
+    References
+    ----------
+    .. [1] E. J. Kirkland, Advanced Computing in Electron Microscopy,
+       Springer US, Boston, MA, 2010.
+    """
+    wavelength, sigma, frequency_grid = _prepare_propagation_parameters(
+        potential, pixel_size, voltage
+    )
+
+    total_slices = potential.shape[-3]
+    if n_slices is None:
+        n_slices = total_slices
+    sizes = chunk_slices(total_slices, n_slices)
+    chunk_potential = torch.stack(
+        [chunk.sum(dim=-3) for chunk in torch.split(potential, sizes, dim=-3)],
+        dim=-3,
+    )  # (..., n_slices, H, W)
+
+    # Chunk c is (sum of sizes[c:]) slice thicknesses from the exit plane:
+    # the wave still has to cross the remainder of chunk c itself plus
+    # every chunk after it.
+    size_tensor = torch.as_tensor(sizes, device=potential.device)
+    remaining_slices = size_tensor.flip(0).cumsum(0).flip(0)
+    distances = remaining_slices * pixel_size
+    propagators = fresnel_propagator(
+        frequency_grid, wavelength=wavelength, dz=distances[:, None, None]
+    )  # (n_slices, H, W)
+
+    scattered = torch.fft.ifft2(
+        torch.fft.fft2(1j * sigma * pixel_size * chunk_potential) * propagators
+    )
+    return torch.exp(scattered.sum(dim=-3))

torch_scattering-0.6.0.dist-info/METADATA

@@ -0,0 +1,212 @@
+Metadata-Version: 2.5
+Name: torch-scattering
+Version: 0.6.0
+Summary: Multislice electron scattering simulation in PyTorch
+Project-URL: homepage, https://github.com/teamtomo/teamtomo
+Project-URL: repository, https://github.com/teamtomo/teamtomo
+Author-email: Joel Yeo <joelyeo.is@gmail.com>
+License: BSD-3-Clause
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Requires-Dist: scipy
+Requires-Dist: torch
+Requires-Dist: torch-ctf
+Requires-Dist: torch-grid-utils
+Description-Content-Type: text/markdown
+
+# torch-scattering
+
+[![License](https://img.shields.io/pypi/l/torch-scattering.svg?color=green)](https://github.com/joelyeois/torch-scattering/raw/main/LICENSE)
+[![PyPI](https://img.shields.io/pypi/v/torch-scattering.svg?color=green)](https://pypi.org/project/torch-scattering)
+[![Python Version](https://img.shields.io/pypi/pyversions/torch-scattering.svg?color=green)](https://python.org)
+[![CI](https://github.com/joelyeois/torch-scattering/actions/workflows/ci.yml/badge.svg)](https://github.com/joelyeois/torch-scattering/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/joelyeois/torch-scattering/branch/main/graph/badge.svg)](https://codecov.io/gh/joelyeois/torch-scattering)
+
+Multislice electron scattering simulation in PyTorch, for cryo-EM/cryo-ET forward modelling.
+
+## Overview
+
+`torch_scattering` computes the 2D exit wave produced by propagating an electron
+beam through a 3D electrostatic potential in volts. The potential has shape
+`(..., Z, H, W)`, where Z is the beam direction. `pixel_size` is the isotropic
+voxel spacing in Angstroms, so it specifies both the Y/X pixel spacing and the
+Z slice thickness. Every function returns a complex exit wave of shape
+`(..., H, W)`.
+
+Real `float32` and `float64` potentials model non-absorbing specimens and can be
+passed directly; callers do not need to cast them to complex. Complex potentials
+remain supported for modelling absorption.
+
+Four propagation modes are provided, trading physical accuracy for speed:
+
+* `multislice()` - full multislice propagation (Kirkland, *Advanced Computing in
+  Electron Microscopy*), alternating transmission through each slice with Fresnel
+  propagation to the next. The most accurate mode.
+* `rytov()` - Rytov approximation, accumulating phase in the exponent rather than
+  the wave itself.
+* `firstborn()` - first Born approximation, summing single-scattering
+  contributions from each slice.
+* `projection()` - projection approximation, treating the specimen as infinitely
+  thin and skipping inter-slice propagation entirely. The fastest and least
+  accurate mode.
+
+All four share the same required inputs and can be swapped in for one another.
+`multislice`, `rytov`, and `firstborn` also accept an `n_slices` argument to
+coarsen the potential into fewer, thicker slabs before propagating.
+
+Lower-level, pure-math primitives (`fresnel_propagator`, `transmission_function`,
+`multislice_step`, `chunk_slices`, `interaction_parameter`) are also exposed for
+building custom propagation schemes.
+
+## Installation
+
+```shell
+pip install torch-scattering
+```
+
+## Usage
+
+```python
+import torch
+from torch_scattering import multislice
+
+# A real electrostatic potential in volts, shape (Z, H, W).
+potential = torch.zeros((50, 64, 64), dtype=torch.float32)
+
+# propagate a plane wave through it
+exit_wave = multislice(
+    potential=potential,
+    pixel_size=1.0,   # Angstroms
+    voltage=300,      # kV
+)
+# exit_wave.shape is (64, 64)
+# exit_wave.dtype is torch.complex64
+```
+
+`rytov`, `firstborn`, and `projection` share the same call signature:
+
+```python
+from torch_scattering import firstborn, projection, rytov
+
+exit_wave = rytov(potential, pixel_size=1.0, voltage=300)
+exit_wave = firstborn(potential, pixel_size=1.0, voltage=300)
+exit_wave = projection(potential, pixel_size=1.0, voltage=300)  # n_slices not applicable
+```
+
+### Coarsening slices
+
+`n_slices` groups the potential into fewer, thicker slabs before propagating.
+By default (`n_slices=None`), every slice of the potential is propagated
+individually - the most accurate but slowest setting.
+
+```python
+# propagate as 10 chunks instead of all 50 slices individually
+exit_wave = multislice(potential, pixel_size=1.0, voltage=300, n_slices=10)
+```
+
+### Batching
+
+All functions accept arbitrary leading batch dimensions on `potential`:
+
+```python
+potential = torch.zeros((8, 50, 64, 64), dtype=torch.complex64)  # batch of 8
+exit_wave = multislice(potential, pixel_size=1.0, voltage=300)
+# exit_wave.shape is (8, 64, 64)
+```
+
+## Structure-to-wave pipeline
+
+Structure handling and potential generation are deliberately separate packages.
+They are not runtime dependencies of `torch-scattering`; their real tensor
+output is passed through the public tensor API:
+
+```python
+import pandas as pd
+from torch_calculate_electrostatic_potential import (
+    GridConfig,
+    potential_from_structure_3d,
+)
+from torch_scattering import multislice
+from torch_structure_manipulation import (
+    AtomicStructure,
+    annotate_bonding_environments,
+)
+
+# mmdf-compatible coordinates are in Angstroms.
+atoms = pd.DataFrame(
+    [
+        ("A", 1, "ALA", "C", "C", 0.0, 0.0, 0.0),
+        ("A", 1, "ALA", "O", "O", 1.2, 0.0, 0.0),
+        ("A", 1, "ALA", "CA", "C", -1.2, 0.0, 0.0),
+        ("A", 2, "GLY", "N", "N", 2.4, 0.0, 0.0),
+    ],
+    columns=[
+        "chain", "residue_id", "residue", "atom", "element", "x", "y", "z"
+    ],
+)
+atoms["b_isotropic"] = 10.0  # Angstrom squared
+atoms["occupancy"] = 1.0
+
+# Annotate a complete local residue context, then build the desired structure.
+annotated = annotate_bonding_environments(atoms, include_hydrogens=False)
+structure = AtomicStructure.from_dataframe(annotated.iloc[[0]])
+
+grid = GridConfig.from_grid_shape_and_voxel_size(
+    grid_shape=(9, 9, 9),       # Z, Y, X
+    voxel_size=(1.0, 1.0, 1.0), # Angstroms; isotropic for scattering
+    center_zyx=(0.0, 0.0, 0.0),
+    sublattice_radius=4.0,
+)
+elemental_volts = potential_from_structure_3d(structure, grid)
+bonded_volts = potential_from_structure_3d(
+    structure,
+    grid,
+    scattering_factors="peng_bonded",
+    bonded_fallback="error",
+)
+
+# Both volumes are real tensors in volts and are accepted directly.
+elemental_wave = multislice(elemental_volts, pixel_size=1.0, voltage=300.0)
+bonded_wave = multislice(bonded_volts, pixel_size=1.0, voltage=300.0)
+# Both waves are complex tensors; voltage is in kV.
+```
+
+`projection()` is a wave-propagation approximation that numerically sums this
+sampled 3D volume along Z. It is distinct from the electrostatic package's
+analytic 2D projected-potential calculation and from projection alignment in
+`torch-fit-in-map`.
+
+## Low-level primitives
+
+For building custom propagation schemes directly on top of the multislice
+recurrence:
+
+```python
+import torch
+from torch_grid_utils import fftfreq_grid
+from torch_scattering import (
+    fresnel_propagator,
+    interaction_parameter,
+    multislice_step,
+)
+
+frequency_grid = fftfreq_grid(image_shape=(64, 64), rfft=False, spacing=1.0, norm=True)
+propagator = fresnel_propagator(frequency_grid, wavelength=0.01969, dz=1.0)
+sigma = interaction_parameter(voltage=300)
+
+wave = torch.ones((64, 64), dtype=torch.complex64)
+potential_slice = torch.zeros((64, 64), dtype=torch.complex64)
+wave = multislice_step(wave, potential_slice, propagator, sigma, dz=1.0)
+```
+
+## License
+
+This project is licensed under the BSD 3-Clause License - see the LICENSE file for details.

torch_scattering-0.6.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+torch_scattering/__init__.py,sha256=SCLWYPfsGT1d8e4aggxkk01ZW97ewWaG5eEpV0NzUSE,792
+torch_scattering/_core.py,sha256=i_Y6wUcB2yIGmQoUY-VhOUUVaCsVz4ValZFJgvh4mCw,10345
+torch_scattering/firstborn.py,sha256=dnizhedOFoek5ZrtPnk7N4q3zlfGCbi1WyRAo0Mld7A,3772
+torch_scattering/multislice.py,sha256=NjfqqhogZoLxd34kFVMZ21sLmdrNvyuaSfr28DtiPpQ,3532
+torch_scattering/projection.py,sha256=Z29xlUCA8wfYNBvubmTw70hkn83HrrrzyLCyNQdxSqw,1936
+torch_scattering/rytov.py,sha256=l2m7oL8ZuIWuGElVtZbKyLp07pvhp5EcFmhAtE9pcQY,3660
+torch_scattering-0.6.0.dist-info/METADATA,sha256=0cqRCdkmIldmKjSvkoKkOMQx22Xas0hiv0Z6BjR3698,7786
+torch_scattering-0.6.0.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
+torch_scattering-0.6.0.dist-info/licenses/LICENSE,sha256=Nt6shgig6_4SzvDw6xDx2NTmbmbwnqylMO5P54YuT3Y,1520
+torch_scattering-0.6.0.dist-info/RECORD,,

torch_scattering-0.6.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.0
+Root-Is-Purelib: true
+Tag: py3-none-any

torch_scattering-0.6.0.dist-info/licenses/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2020-2026, TeamTomo
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.