from __future__ import annotations
import math
from typing import TYPE_CHECKING, Callable, Sequence
import numpy as np
import scipy.ndimage # type: ignore
from numba import cuda, njit # type: ignore
from abtem.antialias import AntialiasAperture
from abtem.core.backend import get_array_module
from abtem.core.energy import energy2sigma, energy2wavelength
from abtem.core.utils import get_dtype
if TYPE_CHECKING:
from abtem.potentials.iam import PotentialArray
from abtem.waves import Waves
import dask.array as da
# 1D second-derivative centered stencils
fd_coefficients = {
2: [1.0, -2.0, 1.0],
4: [
-0.08333333333333333,
1.3333333333333333,
-2.5,
1.3333333333333333,
-0.08333333333333333,
],
6: [
0.011111111111111112,
-0.15,
1.5,
-2.7222222222222223,
1.5,
-0.15,
0.011111111111111112,
],
8: [
-0.0017857142857142857,
0.025396825396825397,
-0.2,
1.6,
-2.8472222222222223,
1.6,
-0.2,
0.025396825396825397,
-0.0017857142857142857,
],
10: [
0.00031746031746031746,
-0.00496031746031746,
0.03968253968253968,
-0.23809523809523808,
1.6666666666666667,
-2.9272222222222224,
1.6666666666666667,
-0.23809523809523808,
0.03968253968253968,
-0.00496031746031746,
0.00031746031746031746,
],
12: [
-6.012506012506013e-05,
0.001038961038961039,
-0.008928571428571428,
0.05291005291005291,
-0.26785714285714285,
1.7142857142857142,
-2.9827777777777778,
1.7142857142857142,
-0.26785714285714285,
0.05291005291005291,
-0.008928571428571428,
0.001038961038961039,
-6.012506012506013e-05,
],
14: [
1.1892869035726179e-05,
-0.00022662522662522663,
0.0021212121212121214,
-0.013257575757575758,
0.06481481481481481,
-0.2916666666666667,
1.75,
-3.02359410430839,
1.75,
-0.2916666666666667,
0.06481481481481481,
-0.013257575757575758,
0.0021212121212121214,
-0.00022662522662522663,
1.1892869035726179e-05,
],
16: [
-2.428127428127428e-06,
5.074290788576503e-05,
-0.000518000518000518,
0.003480963480963481,
-0.017676767676767676,
0.07542087542087542,
-0.3111111111111111,
1.7777777777777777,
-3.05484410430839,
1.7777777777777777,
-0.3111111111111111,
0.07542087542087542,
-0.017676767676767676,
0.003480963480963481,
-0.000518000518000518,
5.074290788576503e-05,
-2.428127428127428e-06,
],
18: [
5.078436450985471e-07,
-1.1569313039901276e-05,
0.00012844298558584272,
-0.0009324009324009324,
0.005034965034965035,
-0.022027972027972027,
0.08484848484848485,
-0.32727272727272727,
1.8,
-3.0795354623330815,
1.8,
-0.32727272727272727,
0.08484848484848485,
-0.022027972027972027,
0.005034965034965035,
-0.0009324009324009324,
0.00012844298558584272,
-1.1569313039901276e-05,
5.078436450985471e-07,
],
}
def _build_matrix(offsets: list[int]):
import sympy # type: ignore
"""Constructs the equation system matrix for the finite difference coefficients"""
A = [([1 for _ in offsets])]
for i in range(1, len(offsets)):
A.append([j**i for j in offsets])
return sympy.Matrix(A)
def _build_rhs(offsets: list[int], deriv: int):
import sympy # type: ignore
"""The right hand side of the equation system matrix"""
b = [0 for _ in offsets]
b[deriv] = math.factorial(deriv)
return sympy.Matrix(b)
def _calculate_finite_difference_coefficient(derivative: int, accuracy: int = 2):
import sympy # type: ignore
num_central = 2 * math.floor((derivative + 1) / 2) - 1 + accuracy
num_side = num_central // 2
offsets = list(range(-num_side, num_side + 1))
matrix = _build_matrix(offsets)
rhs = _build_rhs(offsets, derivative)
coefs = sympy.linsolve((matrix, rhs))
coefs = np.array([float(coef) for coef in tuple(coefs)[0]])
return coefs
[docs]
def finite_difference_coefficients(derivative: int, accuracy: int = 2):
if accuracy % 2 == 1 or accuracy <= 0:
raise ValueError("accuracy order must be a positive even integer")
if derivative < 0:
raise ValueError("derivative degree must be a positive integer")
if accuracy <= 18:
return np.array(fd_coefficients[accuracy])
return _calculate_finite_difference_coefficient(derivative, accuracy)
def _laplace_stencil_array(accuracy):
coefficients = finite_difference_coefficients(2, accuracy)
stencil = np.zeros((len(coefficients),) * 2)
stencil[len(coefficients) // 2, :] = coefficients
stencil[:, len(coefficients) // 2] += coefficients
return stencil
def _laplace_operator_stencil(
accuracy, prefactor, mode: str = "wrap", dtype=None, device: str = "cpu"
):
if dtype is None:
dtype = get_dtype(complex=True)
c = finite_difference_coefficients(2, accuracy)
c = c * prefactor
c = c.astype(dtype)
c = np.roll(c, -(len(c) // 2))
n = len(c) // 2
padding = n + 1
from numba import prange # type: ignore
@njit(parallel=True, fastmath=True)
def _laplace_stencil_cpu_batch(a):
M, H, W = a.shape
out = a.copy()
out[:] = 0
for m in prange(M):
for i in range(n, H - n):
for j in range(n, W - n):
cumul = dtype(0.0)
for k in range(-n, n + 1):
cumul += c[k] * a[m, i + k, j] + c[k] * a[m, i, j + k]
out[m, i, j] = cumul
return out
@cuda.jit
def stencil_func_gpu_batch(a, out):
m, i, j = cuda.grid(3)
M, H, W = a.shape
if m < M and n <= i < H - n and n <= j < W - n:
cumul = dtype(0.0)
for k in range(-n, n + 1):
cumul += c[k] * a[m, i + k, j] + c[k] * a[m, i, j + k]
out[m, i, j] = cumul
def _laplace_stencil_gpu(a):
xp = get_array_module(a)
out = xp.zeros_like(a)
M, H, W = a.shape
threads_x = 8
threads_y = 8
target_threads = 256
threads_m = max(1, target_threads // (threads_x * threads_y))
threadsperblock = (threads_m, threads_x, threads_y)
blockspergrid_m = math.ceil(a.shape[0] / threadsperblock[0])
blockspergrid_x = math.ceil(a.shape[1] / threadsperblock[1])
blockspergrid_y = math.ceil(a.shape[2] / threadsperblock[2])
blockspergrid = (blockspergrid_m, blockspergrid_x, blockspergrid_y)
stencil_func_gpu_batch[blockspergrid, threadsperblock](a, out)
return out
def _laplace_stencil(a):
# Store original shape and reshape to 3D
original_shape = a.shape
if a.ndim == 2:
a = a.reshape(1, *a.shape)
elif a.ndim > 3:
a = a.reshape(-1, *a.shape[-2:])
elif a.ndim != 3:
raise ValueError(f"Array must have at least 2 dimensions, got {a.ndim}")
# Apply stencil
if device == "cpu":
result = _laplace_stencil_cpu_batch(a)
elif device == "gpu":
result = _laplace_stencil_gpu(a)
else:
raise ValueError(f"Unsupported device: {device}")
# Reshape back to original shape
return result.reshape(original_shape)
def _apply_boundary(mode, padding):
def stencil_with_boundary(func):
def func_wrapper(a):
xp = get_array_module(a)
# Build pad_width for arbitrary dimensions
# Only pad the last two spatial dimensions
pad_width = [(0, 0)] * (a.ndim - 2) + [(padding,) * 2, (padding,) * 2]
# Build slicing to remove padding from last two dimensions
slicing = tuple(
[slice(None)] * (a.ndim - 2)
+ [slice(padding, -padding), slice(padding, -padding)]
)
a = xp.pad(a, pad_width=pad_width, mode=mode)
res = func(a)
return res[slicing]
return func_wrapper
return stencil_with_boundary
if mode != "none":
return _apply_boundary(mode="wrap", padding=padding)(_laplace_stencil)
else:
return _laplace_stencil
def _laplace_operator_func_slow(accuracy, prefactor):
stencil = _laplace_stencil_array(accuracy) * prefactor
def func(array):
return scipy.ndimage.convolve(array, stencil, mode="wrap")
return func
[docs]
class LaplaceOperator:
def __init__(self, accuracy):
"""
Centered finite-difference laplacian operator.
Parameters
----------
accuracy: int
centered finite-difference stencil accuracy
"""
self._accuracy = accuracy
self._key = None
self._stencil = None
def _get_new_stencil(self, key, device: str = "cpu"):
wavelength, sampling = key
prefactor = 1 / np.prod(np.array(sampling, dtype=float))
return _laplace_operator_stencil(
self._accuracy, prefactor, mode="wrap",
dtype=get_dtype(complex=True), device=device
)
[docs]
def get_stencil(self, waves: Waves, device: str = "cpu") -> Callable:
"""
Cached method to return finite-difference stencil using specified
waves parameters, namely wavelength and sampling for the prefactor.
Parameters
----------
waves: Waves
Waves object stencil will ultimately be applied to
device: str, optional
Device to evaluate stencil on, "cpu" or "gpu"
Returns
----------
stencil: Callable
Cached stencil function with waves prefactor, on specified device
"""
key = (
energy2wavelength(waves._valid_energy),
waves.sampling,
)
if key == self._key:
return self._stencil
self._stencil = self._get_new_stencil(key, device=device)
self._key = key
return self._stencil
[docs]
def apply(self, waves):
laplace_stencil = self.get_stencil(waves, device=waves.device)
waves._array = laplace_stencil(waves._array)
return waves
[docs]
class DivergedError(Exception):
def __init__(self, message="the multislice exponential series diverged"):
super().__init__(message)
[docs]
class NotConvergedError(Exception):
def __init__(self, message="the series did not converge"):
super().__init__(message)
def _multislice_exponential_series(
waves: np.ndarray | da.core.Array,
transmission_function: np.ndarray,
laplace: Callable,
wavelength: float,
thickness: float,
tolerance: float = 1e-16,
max_terms: int = 300,
order: int = 1,
fully_corrected: bool = False,
):
xp = get_array_module(waves)
initial_amplitude = xp.abs(waves).sum()
if fully_corrected:
temp = full_series(
waves, laplace, transmission_function, order, wavelength, thickness
)
else:
temp = propagator_taylor_series(
waves,
order=order,
laplace=laplace,
transmission_function=transmission_function,
wavelength=wavelength,
thickness=thickness,
)
waves += temp
for i in range(2, max_terms + 1):
if fully_corrected:
temp = (
full_series(
temp, laplace, transmission_function, order, wavelength, thickness
)
/ i
)
else:
temp = (
propagator_taylor_series(
temp,
order=order,
laplace=laplace,
transmission_function=transmission_function,
wavelength=wavelength,
thickness=thickness,
)
) / i
waves += temp
temp_amplitude = xp.abs(temp).sum()
if temp_amplitude / initial_amplitude <= tolerance:
break
if not xp.all(xp.isfinite(temp)) or temp_amplitude > initial_amplitude:
raise DivergedError()
else:
raise NotConvergedError(
f"series did not converge to a tolerance of {tolerance} in {max_terms}terms"
)
return waves
[docs]
def conventional_operator(
waves: np.ndarray | da.core.Array,
laplace: Callable,
transmission_function: np.ndarray,
wavelength: float,
):
"""
Split-step real-space multislice operator used in all higher-order expansions.
Parameters
----------
waves: Waves
Waves object to apply multislice operator on
laplace: Callable
Fast laplace operator stencil function
transmission_function: numpy.ndarray,
Scaled potential slice to multiply incoming waves with
wavelength: float
Waves wavelength
"""
K0 = 1 / wavelength
return laplace(waves) / (4 * np.pi * K0) + transmission_function * waves
[docs]
def propagator_taylor_series(
waves: np.ndarray | da.core.Array,
order: int,
laplace: Callable,
transmission_function: np.ndarray,
wavelength: float,
thickness: float,
):
"""
Taylor series expansion of the propagator term in the MS equation.
Eq.(8) in Ultramicroscopy 134 (2013) 135-143.
"""
if order < 1:
raise ValueError("order must be a positive integer and at least 1")
if order == 1:
return (
conventional_operator(waves, laplace, transmission_function, wavelength)
* 1.0j
* thickness
)
K0 = 1 / wavelength
laplace_waves = laplace(waves) / (4 * np.pi * K0)
series = laplace_waves.copy()
temp = laplace_waves.copy()
for i in range(2, order + 1):
prefactor = (wavelength / (-2.0 * np.pi)) ** (i - 1) * 0.5
temp = laplace(temp) / (4 * np.pi * K0)
series += temp * prefactor
return (series + waves * transmission_function) * 1.0j * thickness
[docs]
def full_series(
waves: np.ndarray | da.core.Array,
laplace: Callable,
transmission_function: np.ndarray,
order: int,
wavelength: float,
thickness: float,
override_prefactor: list[float] = [],
):
"""
Full Taylor series expansion of the MS Eq.(14) in Ultramicrscopy 134 (2013) 135-143.
override_prefactor used in backscatter call, Eq. (13) in Micron 190 (2025) 103778.
"""
series = conventional_operator(waves, laplace, transmission_function, wavelength)
temp = series.copy()
for i in range(2, order + 1):
if override_prefactor:
prefactor = override_prefactor[
i - 1
] # Note that the first prefactor always gets skipped and is always 1
else:
prefactor = (wavelength / (-2.0 * np.pi)) ** (i - 1) * 0.5
temp = conventional_operator(temp, laplace, transmission_function, wavelength)
series += temp * prefactor
return series * 1.0j * thickness
[docs]
def multislice_step(
waves: Waves,
potential_slice: PotentialArray,
next_slice: PotentialArray | None,
laplace: LaplaceOperator,
tolerance: float = 1e-16,
max_terms: int = 300,
order: int = 1,
fully_corrected: bool = False,
) -> Waves | Sequence[Waves]:
"""
Performs a single multislice step.
Parameters
----------
waves: Waves
Waves object to apply multislice operator on
potential_slice: PotentialArray
PotentialArray of slice to be used in transmission operator
next_slice: PotentialArray | None
PotentialArray of next slice (used to 'look ahead' for backscattered waves)
laplace: Callable
Fast laplace operator stencil function
tolerance: float
Convergence tolerance for exponent Taylor series amplitude
max_terms: int
Maximum terms in exponent Taylor series
order: int
Multislice operator expansion order
fully_corrected: bool
If True, transmission and propagator operators are expanded to specified order.
If next_slice is not None, the backscattered wave component is also computed.
"""
xp = get_array_module(waves.array)
if max_terms < 1:
raise ValueError()
if waves.device != potential_slice.device:
potential_slice = potential_slice.copy_to_device(device=waves.device)
if next_slice is not None:
next_slice = next_slice.copy_to_device(device=waves.device)
thickness = potential_slice.thickness
transmission_function_array = (
potential_slice.array[0] * energy2sigma(waves._valid_energy) / thickness
)
if fully_corrected and next_slice is not None:
transmission_function_array_next_slice = (
next_slice.array[0] * energy2sigma(waves._valid_energy) / thickness
)
laplace_stencil = laplace.get_stencil(waves, device=waves.device)
wavelength = energy2wavelength(waves._valid_energy)
# Forward scattering term
waves._array = _multislice_exponential_series(
waves._array,
transmission_function_array,
laplace_stencil,
wavelength,
thickness,
tolerance,
max_terms,
order,
fully_corrected,
)
# Bandlimit before the correction term to suppress high-k content that
# would be amplified by repeated conventional_operator applications at
# low energies (large 1/K0) or after many slices (channeling build-up).
aperture = AntialiasAperture()
waves = aperture.bandlimit(waves)
# Correction-term contributions
backscatter = xp.zeros_like(waves._array)
if fully_corrected and next_slice is not None:
# constants and prefactors
K0 = 1 / wavelength
# Eq. 7 in Micron 190 (2025) 103778.
backscatter = (
1
/ (2 * np.pi * 1.0j * thickness)
* (
full_series(
waves._array,
laplace_stencil,
transmission_function_array_next_slice,
order,
wavelength,
thickness,
)
- full_series(
waves._array,
laplace_stencil,
transmission_function_array,
order,
wavelength,
thickness,
)
)
)
# 1/k series with custom prefactors
prefactors = [1]
for i in range(1, order + 1):
prefactors.append(prefactors[-1] * (1 - 2 * i) / (2 * i))
for i in range(len(prefactors)):
prefactors[i] = prefactors[i] / (1.0j * thickness) / (np.pi * K0) ** i
backscatter *= (
1
/ (2 * K0)
* (
1
+ full_series(
waves._array,
laplace_stencil,
transmission_function_array_next_slice,
order,
wavelength,
thickness,
override_prefactor=prefactors,
)
)
)
if not xp.all(xp.isfinite(backscatter)):
raise DivergedError(
"overflow in correction-term computation; try using a smaller "
"slice_thickness or a lower expansion order"
)
# Eq.10 in Micron 190 (2025) 103778.
waves._array = waves._array - backscatter
if fully_corrected:
kwargs = waves._copy_kwargs(exclude=("array",))
backscatter_waves = waves.__class__(backscatter, **kwargs)
return waves, aperture.bandlimit(backscatter_waves)
return waves
