"""Module for describing abTEM measurement objects."""
from __future__ import annotations
import copy
import functools
import itertools
import warnings
from abc import ABCMeta, abstractmethod
from collections import defaultdict
from numbers import Number
from types import ModuleType
from typing import (
TYPE_CHECKING,
Callable,
Dict,
Optional,
Self,
Sequence,
SupportsFloat,
Type,
TypeVar,
)
import dask.array as da
import numpy as np
from ase import Atom
from ase.cell import Cell
from matplotlib.axes import Axes
from numba import jit # type: ignore
from abtem.array import ArrayObject, _validate_array_items, stack
from abtem.core import config
from abtem.core.axes import (
AxisMetadata,
LinearAxis,
NonLinearAxis,
RealSpaceAxis,
ReciprocalSpaceAxis,
ScaleAxis,
ScanAxis,
)
from abtem.core.backend import (
asnumpy,
cp,
get_array_module,
get_ndimage_module,
get_scipy_module,
)
from abtem.core.complex import abs2
from abtem.core.energy import energy2wavelength
from abtem.core.fft import fft_crop, fft_interpolate
from abtem.core.grid import (
adjusted_gpts,
polar_spatial_frequencies,
spatial_frequencies,
)
from abtem.core.units import get_conversion_factor
from abtem.core.utils import (
CopyMixin,
EqualityMixin,
cos_sin_deg,
get_dtype,
is_broadcastable,
label_to_index,
)
# from abtem.distributions import BaseDistribution
from abtem.noise import NoiseTransform, ScanNoiseTransform
from abtem.visualize.artists import LinesArtist
from abtem.visualize.visualizations import Visualization
from abtem.visualize.widgets import ImageGUI, LinesGUI, ScatterGUI
interpolate_bilinear_cuda: Optional[Callable] = None
sum_run_length_encoded: Optional[Callable] = None
sum_run_length_encoded_cuda: Optional[Callable] = None
if cp is not None:
from abtem.core._cuda import interpolate_bilinear as interpolate_bilinear_cuda
from abtem.core._cuda import sum_run_length_encoded as sum_run_length_encoded_cuda
xr: Optional[ModuleType] = None
try:
import xarray as xr
except ImportError:
xr = None
pd: Optional[ModuleType] = None
try:
import pandas as pd # type: ignore
except ImportError:
pd = None
if TYPE_CHECKING:
from abtem.waves import BaseWaves
BaseMeasurementsSubclass = TypeVar("BaseMeasurementsSubclass", bound="BaseMeasurements")
def _scanned_measurement_type(
measurement: BaseMeasurements | BaseWaves,
) -> Type[RealSpaceLineProfiles | Images | MeasurementsEnsemble]:
if len(_scan_shape(measurement)) == 1:
return RealSpaceLineProfiles
elif len(_scan_shape(measurement)) == 2:
return Images
else:
return MeasurementsEnsemble
def _bin_extent(n: int) -> tuple[float, float]:
if n % 2 == 0:
return -n // 2 - 0.5, n // 2 - 0.5
else:
return -n // 2 + 0.5, n // 2 + 0.5
def _reduced_scanned_images_or_line_profiles(
new_array: np.ndarray,
old_measurement: BaseMeasurements,
metadata: dict | None = None,
) -> RealSpaceLineProfiles | Images | MeasurementsEnsemble | np.ndarray:
if metadata is None:
metadata = {}
metadata = {**old_measurement.metadata, **metadata}
ensemble_axes = tuple(range(len(old_measurement.ensemble_shape)))
source = _scan_axes(old_measurement)
destination = tuple(range(len(ensemble_axes) - len(source), len(ensemble_axes)))
scan_axes_metadata = [old_measurement.ensemble_axes_metadata[i] for i in source]
ensemble_axes_metadata = [
m
for i, m in enumerate(old_measurement.ensemble_axes_metadata)
if i not in source
]
if source != destination:
xp = get_array_module(new_array)
if old_measurement.is_lazy:
new_array = da.moveaxis(new_array, source, destination)
else:
new_array = xp.moveaxis(new_array, source, destination)
sampling: float | tuple[float, ...]
if len(scan_axes_metadata) == 1:
sampling = _scan_sampling(old_measurement)[-1]
return RealSpaceLineProfiles(
new_array,
sampling=sampling,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
elif len(scan_axes_metadata) == 2:
sampling = _scan_sampling(old_measurement)
sampling = (sampling[-2], sampling[-1])
images = Images(
new_array,
sampling=sampling,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
return images
elif _scanned_measurement_type(old_measurement) is MeasurementsEnsemble:
ensemble_axes_metadata = old_measurement.ensemble_axes_metadata
measurement_ensemble = MeasurementsEnsemble(
new_array,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
return measurement_ensemble
else:
return new_array
def _scan_axes(measurement):
num_scan_axes = 0
scan_axes = ()
for i, axis in enumerate(measurement.ensemble_axes_metadata):
if num_scan_axes == 2:
break
if isinstance(axis, ScanAxis) and axis._main is True:
scan_axes += (i,)
num_scan_axes += 1
scan_axes = scan_axes[-2:]
return scan_axes
def _scan_sampling(measurements: BaseMeasurements | BaseWaves) -> tuple[float, ...]:
return tuple(
measurements.axes_metadata[i].sampling for i in _scan_axes(measurements)
)
def _scan_axes_metadata(
measurements: BaseMeasurements | BaseWaves,
) -> list[AxisMetadata]:
return [measurements.axes_metadata[i] for i in _scan_axes(measurements)]
def _scan_shape(measurements: BaseMeasurements | BaseWaves) -> tuple[int, ...]:
return tuple(measurements.shape[i] for i in _scan_axes(measurements))
def _scan_area_per_pixel(measurements):
if len(_scan_sampling(measurements)) == 2:
return np.prod(_scan_sampling(measurements))
else:
raise RuntimeError("Cannot infer pixel area from axes metadata.")
def _scan_extent(measurement):
extent = ()
for n, metadata in zip(_scan_shape(measurement), _scan_axes_metadata(measurement)):
extent += (metadata.sampling * n,)
return extent
def _array_module_fn(array, xp: ModuleType, name: str):
"""Return the array-module function ``name`` (e.g. ``"concatenate"``),
routed through dask's own implementation whenever ``array`` is still a
lazy dask array, regardless of device.
NumPy dispatches a top-level call like ``numpy.concatenate([dask_array])``
to dask automatically via ``__array_function__``, but CuPy does not: its
functions raise ``TypeError`` when given a ``dask.array.core.Array``
rather than a genuine ``cupy.ndarray``. Any code that resolves ``xp`` via
:func:`get_array_module` and then calls ``xp.<name>(...)`` directly on an
array that might still be lazy (e.g. because it was built from
``exit_waves.array``/``dp.array`` without forcing computation) must go
through this helper instead of calling ``xp.<name>`` unconditionally.
"""
if isinstance(array, da.core.Array):
return getattr(da, name)
return getattr(xp, name)
def _spatial_frequency_squared(
gpts: tuple[int, int],
sampling: tuple[float, float],
xp: ModuleType = np,
) -> np.ndarray:
"""Squared spatial frequency grid, the shared (and expensive) part of
:func:`_annular_detector_mask` — factored out so callers building many
masks for the same ``gpts``/``sampling`` (e.g. a q-sweep) can compute it
once and reuse it."""
kx, ky = spatial_frequencies(
gpts, (1 / sampling[0] / gpts[0], 1 / sampling[1] / gpts[1]), False, xp
)
return kx[:, None] ** 2 + ky[None] ** 2
def _annular_detector_mask(
gpts: tuple[int, int],
sampling: tuple[float, float],
inner: float,
outer: float,
offset: tuple[float, float] = (0.0, 0.0),
fftshift: bool = False,
xp: ModuleType = np,
k2: Optional[np.ndarray] = None,
) -> np.ndarray | list[np.ndarray]:
if k2 is None:
k2 = _spatial_frequency_squared(gpts, sampling, xp)
bins = (k2 >= inner**2) & (k2 < outer**2)
if np.any(np.array(offset) != 0.0):
offset = (
int(round(offset[0] / sampling[0])),
int(round(offset[1] / sampling[1])),
)
# if (abs(offset[0]) > bins[0]) or (abs(offset[1]) > bins[1]):
# raise RuntimeError("Detector offset exceeds maximum detected angle.")
bins = np.roll(bins, offset, (0, 1))
if fftshift:
bins = xp.fft.fftshift(bins)
return bins
def _polar_detector_bins(
gpts: tuple[int, int],
sampling: tuple[float, float],
inner: float,
outer: float,
nbins_radial: int,
nbins_azimuthal: int,
rotation: float = 0.0,
offset: tuple[float, float] = (0.0, 0.0),
fftshift: bool = False,
return_indices: bool = False,
) -> np.ndarray | list[np.ndarray]:
alpha, phi = polar_spatial_frequencies(
gpts, sampling=(1 / sampling[0] / gpts[0], 1 / sampling[1] / gpts[1])
)
phi = (phi - rotation) % (2 * np.pi)
radial_bins = -np.ones(gpts, dtype=int)
valid = (alpha >= inner) & (alpha < outer)
radial_bins[valid] = nbins_radial * (alpha[valid] - inner) / (outer - inner)
angular_bins = np.floor(nbins_azimuthal * (phi / (2 * np.pi)))
angular_bins = np.clip(angular_bins, 0, nbins_azimuthal - 1).astype(int)
bins = -np.ones(gpts, dtype=int)
bins[valid] = angular_bins[valid] + radial_bins[valid] * nbins_azimuthal
if np.any(np.array(offset) != 0.0):
offset = (
int(round(offset[0] / sampling[0])),
int(round(offset[1] / sampling[1])),
)
# if (abs(offset[0]) > bins[0]) or (abs(offset[1]) > bins[1]):
# raise RuntimeError("Detector offset exceeds maximum detected angle.")
bins = np.roll(bins, offset, (0, 1))
if fftshift:
bins = np.fft.fftshift(bins)
if return_indices:
indices = []
for i in label_to_index(bins, nbins_radial * nbins_azimuthal - 1):
indices.append(i)
return indices
else:
return bins
@jit(nopython=True, nogil=True, fastmath=True)
def _sum_run_length_encoded(array, result, separators):
for x in range(result.shape[1]):
for i in range(result.shape[0]):
for j in range(separators[x], separators[x + 1]):
result[i, x] += array[i, j]
def _interpolate_stack(
array: np.ndarray, positions: np.ndarray, mode: str, order: int, **kwargs
):
map_coordinates = get_ndimage_module(array).map_coordinates
xp = get_array_module(array)
positions_shape = positions.shape
positions = positions.reshape((-1, 2))
old_shape = array.shape
array = array.reshape((-1,) + array.shape[-2:])
array = xp.pad(array, ((0, 0), (2 * order,) * 2, (2 * order,) * 2), mode=mode)
positions = positions + 2 * order
output = xp.zeros((array.shape[0], positions.shape[0]), dtype=array.dtype)
for i in range(array.shape[0]):
map_coordinates(array[i], positions.T, output=output[i], order=order, **kwargs)
output = output.reshape(old_shape[:-2] + positions_shape[:-1])
return output
[docs]
class BaseMeasurements(ArrayObject, EqualityMixin, CopyMixin, metaclass=ABCMeta):
"""
Base class for all measurement types.
Parameters
----------
array : ndarray
Array containing data of type `float` or `complex`.
ensemble_axes_metadata : list of AxisMetadata, optional
Metadata associated with an ensemble axis.
metadata : dict, optional
A dictionary defining simulation metadata.
"""
def __init__(
self,
array: np.ndarray | da.core.Array,
ensemble_axes_metadata: list[AxisMetadata],
metadata: dict,
):
super().__init__(
array=array,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
@property
@abstractmethod
def base_axes_metadata(self) -> list:
"""List of AxisMetadata of the base axes."""
@property
def metadata(self) -> dict:
"""Metadata describing the measurement."""
return self._metadata
def _get_energy(self) -> float:
"""Return the electron energy [eV].
Shares its first two resolution steps with
:func:`abtem.core.energy.resolve_energy` (``metadata["energy"]``,
then a single-element ``EnergyAxis``). For a full, un-indexed
energy-ensemble measurement that resolver returns ``None``, so this
falls back further to the maximum value of an ``EnergyAxis`` in
``ensemble_axes_metadata`` — the highest energy (shortest
wavelength), mirroring :attr:`Waves.angular_sampling`'s conservative
convention. This fallback is an imprecise convenience; callers that
need the exact per-member energy should index the ensemble first.
Raises
------
RuntimeError
If no energy can be found in metadata or ensemble axes metadata.
"""
from abtem.core.axes import EnergyAxis
from abtem.core.energy import resolve_energy
energy = resolve_energy(None, self.metadata, self.ensemble_axes_metadata)
if energy is not None:
return energy
for axis in self.ensemble_axes_metadata:
if isinstance(axis, EnergyAxis):
return float(max(axis.values))
raise RuntimeError("energy not in measurement metadata.")
def _check_is_complex(self):
if not np.iscomplexobj(self.array):
raise RuntimeError("Function not implemented for non-complex measurements.")
[docs]
def real(self) -> Self:
"""Returns the real part of a complex-valued measurement."""
self._check_is_complex()
self.metadata["label"] = "real"
self.metadata["units"] = "arb. unit"
return self._apply_element_wise_func(get_array_module(self.array).real)
[docs]
def imag(self) -> Self:
"""Returns the imaginary part of a complex-valued measurement."""
self._check_is_complex()
self.metadata["label"] = "imaginary"
self.metadata["units"] = "arb. unit"
return self._apply_element_wise_func(get_array_module(self.array).imag)
[docs]
def phase(self) -> Self:
"""Calculates the phase of a complex-valued measurement."""
self._check_is_complex()
self.metadata["label"] = "phase"
self.metadata["units"] = "rad."
return self._apply_element_wise_func(get_array_module(self.array).angle)
[docs]
def abs(self) -> Self:
"""Calculates the absolute value of a complex-valued measurement."""
# self._check_is_complex()
self.metadata["label"] = "amplitude"
self.metadata["units"] = "arb. unit"
return self._apply_element_wise_func(get_array_module(self.array).abs)
[docs]
def intensity(self) -> Self:
"""Calculates the squared norm of a complex-valued measurement."""
self._check_is_complex()
self.metadata["label"] = "intensity"
self.metadata["units"] = "arb. unit"
return self._apply_element_wise_func(abs2)
[docs]
def relative_difference(
self, other: BaseMeasurements, min_relative_tol: float = 0.0
) -> Self:
"""
Calculates the relative difference with respect to another compatible
measurement.
Parameters
----------
other : BaseMeasurements
Measurement to which the difference is calculated.
min_relative_tol : float
Avoids division by zero errors by defining a minimum value of the divisor in
the relative difference.
Returns
-------
difference : BaseMeasurements
The relative difference as a measurement of the same type.
"""
if not isinstance(other, self.__class__):
raise RuntimeError("Measurements are not of the same type.")
difference = self - other
xp = get_array_module(self.array)
valid = xp.abs(self.array) >= min_relative_tol * self.array.max()
difference._array[valid] /= self.array[valid]
difference._array[valid == 0] = np.nan
difference._array *= 100.0
difference.metadata["label"] = "Relative difference"
difference.metadata["units"] = "%"
difference.metadata["tex_units"] = r"$\%$"
return difference
[docs]
def normalize_ensemble(self, scale: str = "max", shift: str = "mean"):
"""
Normalize the ensemble by shifting ad scaling each member.
Parameters
----------
scale : {'max', 'min', 'sum', 'mean', 'ptp'}
shift : {'max', 'min', 'sum', 'mean', 'ptp'}
Returns
-------
normalized_measurements : BaseMeasurements or subclass of _BaseMeasurement
"""
if shift != "none":
array = self.array - getattr(np, shift)(self.array, axis=-1, keepdims=True)
else:
array = self.array
array = array / getattr(np, scale)(self.array, axis=-1, keepdims=True)
kwargs = self._copy_kwargs(exclude=("array",))
return self.__class__(array, **kwargs)
[docs]
def reduce_ensemble(self) -> Self:
"""
Calculates the mean of an ensemble measurement (e.g. of frozen phonon
configurations).
"""
axis = tuple(
i
for i, axis in enumerate(self.axes_metadata)
if hasattr(axis, "_ensemble_mean") and axis._ensemble_mean
)
if len(axis) == 0:
return self
if np.iscomplexobj(self.array):
warnings.warn("a complex reducing a complex measurement")
return self.mean(axis=axis)
def _apply_element_wise_func(self, func: Callable) -> Self:
d = self._copy_kwargs(exclude=("array",))
d["array"] = func(self.array)
return self.__class__(**d)
@property
@abstractmethod
def _area_per_pixel(self):
if len(_scan_sampling(self)) == 2:
return np.prod(_scan_sampling(self))
else:
raise RuntimeError("Cannot infer pixel area from axes metadata.")
[docs]
def poisson_noise(
self,
dose_per_area: SupportsFloat | Sequence[SupportsFloat] | None = None,
total_dose: SupportsFloat | Sequence[SupportsFloat] | None = None,
samples: int = 1,
seed: Optional[int] = None,
) -> Self:
"""
Add Poisson noise (i.e. shot noise) to a measurement corresponding to the
provided 'total_dose' (per measurement if applied to an ensemble) or
'dose_per_area' (not applicable for single measurements).
Parameters
----------
dose_per_area : float, sequence of float, optional
The irradiation dose [electrons per Å:sup:`2`]. May be given as a single
value or as a sequence of values for each ensemble member.
total_dose : float, optional
The irradiation dose per diffraction pattern.
samples : int, optional
The number of samples to draw from a Poisson distribution. If this is
greater than 1, an additional ensemble axis will be added to the
measurement.
seed : int, optional
Seed the random number generator.
Returns
-------
noisy_measurement : BaseMeasurements or subclass of _BaseMeasurement
The noisy measurement.
"""
dtype = get_dtype(complex=False)
wrong_dose_error = RuntimeError(
"Provide one of 'dose_per_area' or 'total_dose'."
)
if (dose_per_area is not None) and (total_dose is not None):
raise wrong_dose_error
if dose_per_area is not None and total_dose is None:
total_dose = self._area_per_pixel * dose_per_area
elif total_dose is None:
pass
elif total_dose is not None:
if dose_per_area is not None:
raise wrong_dose_error
else:
raise wrong_dose_error
total_dose = np.array(total_dose, dtype=dtype)
# print(seed, samples)
transform = NoiseTransform(dose=total_dose, samples=samples, seeds=seed)
measurement = transform.apply(self)
return measurement
def _scale_axis_from_metadata(self):
return ScaleAxis(
label=self.metadata.get("label", ""),
units=self.metadata.get("units", ""),
tex_label=None,
)
[docs]
def to_measurement_ensemble(self):
return MeasurementsEnsemble(
array=self.array,
ensemble_axes_metadata=self.axes_metadata,
metadata=self.metadata,
)
[docs]
@abstractmethod
def show(self, *args, **kwargs):
"""Documented in subclasses"""
[docs]
def periodic_crop(
array: np.ndarray, corner: tuple[float, float], new_shape: tuple[int, int]
) -> np.ndarray:
"""
Crop an array with periodic boundary conditions. The cropping region is wrapped
around the array.
Parameters
----------
array : ndarray
The array to crop.
corner : two floats
The corner of the cropping region.
new_shape : two ints
The shape of the cropping region.
Returns
-------
cropped_array : ndarray
The cropped array.
"""
xp = get_array_module(array)
if (
(corner[0] > 0)
& (corner[1] > 0)
& (corner[0] + new_shape[0] < array.shape[-2])
& (corner[1] + new_shape[1] < array.shape[-1])
):
array = array[
...,
corner[0] : corner[0] + new_shape[0],
corner[1] : corner[1] + new_shape[1],
]
return array
x = xp.arange(corner[0], corner[0] + new_shape[0], dtype=xp.int64) % array.shape[-2]
y = xp.arange(corner[1], corner[1] + new_shape[1], dtype=xp.int64) % array.shape[-1]
x, y = xp.meshgrid(x, y, indexing="ij")
array = array[..., x.ravel(), y.ravel()].reshape(array.shape[:-2] + new_shape)
return array
[docs]
def integrate_disc(
measurement: Images | DiffractionPatterns,
position: np.ndarray,
radius: float,
return_mean: bool = False,
border: str = "wrap",
interpolate: Optional[tuple[float, bool]] = None,
) -> float:
"""
Integrate the values of a 2d measurement on a disc-shaped region.
Parameters
----------
measurement : 2d measurement
The measurement to integrate
position : two floats
Center of disc-shaped integration region
radius : float
Radius of disc-shaped integration region
return_mean : bool
If true return the mean, otherwise return the sum.
border : str
Specify how to treat integration regions that cross the image border. The valid
values and their behaviours are:
'wrap'
The measurement is extended by wrapping around to the opposite edge.
'raise'
Raise an error if the integration region crosses the measurement border.
interpolate : float or False
The image will be interpolated to this sampling. Units of Angstrom.
Returns
-------
float
Integral value
"""
if interpolate is not None:
measurement = measurement.interpolate(interpolate)
x_axis = measurement.base_axes_metadata[-2]
y_axis = measurement.base_axes_metadata[-1]
assert isinstance(x_axis, RealSpaceAxis) and isinstance(y_axis, RealSpaceAxis)
offset = (x_axis.offset, y_axis.offset)
position = np.array(position) - offset
integration_shape = (
int(np.ceil(2 * radius / x_axis.sampling)),
int(np.ceil(2 * radius / y_axis.sampling)),
)
corner = (
int(np.floor(position[-2] / x_axis.sampling)) - integration_shape[0] // 2,
int(np.floor(position[-1] / y_axis.sampling)) - integration_shape[1] // 2,
)
if border == "wrap":
cropped = periodic_crop(measurement.array, corner, integration_shape)
elif border == "raise":
if (
(np.any(np.array(corner) < 0))
| (corner[0] + integration_shape[0] > measurement.array.shape[0])
| (corner[1] + integration_shape[1] > measurement.array.shape[1])
):
raise RuntimeError("The integration region is outside the image.")
cropped = periodic_crop(measurement.array, corner, integration_shape)
else:
raise RuntimeError('border must be one of "wrap" or "raise"')
x = np.linspace(
0.0,
cropped.shape[-2] * x_axis.sampling,
cropped.shape[-2],
endpoint=x_axis.endpoint,
)
y = np.linspace(
0.0,
cropped.shape[-1] * measurement.sampling[-1],
cropped.shape[-1],
endpoint=y_axis.endpoint,
)
x, y = np.meshgrid(x, y, indexing="ij")
cropped_position = np.array(position)[:2] - (
corner[-2] * x_axis.sampling,
corner[-1] * y_axis.sampling,
)
r = np.sqrt((x - cropped_position[-2]) ** 2 + (y - cropped_position[-1]) ** 2)
mean_sampling = (x_axis.sampling + y_axis.sampling) / 2
mask = 1 - np.clip((r - radius) / mean_sampling, 0, 1)
if return_mean:
return (cropped * mask).sum((-2, -1)) / mask.sum((-2, -1))
else:
return (cropped * mask).sum((-2, -1))
[docs]
class MeasurementsEnsemble(BaseMeasurements):
_base_dims = 0
def __init__(
self,
array: np.ndarray,
ensemble_axes_metadata: list[AxisMetadata],
metadata: dict | None = None,
):
super().__init__(
array=array,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
@property
def _area_per_pixel(self):
raise RuntimeError("Cannot infer pixel area from metadata.")
@property
def base_axes_metadata(self):
return []
[docs]
def show(
self,
type: str = "lines",
ax: Optional[Axes] = None,
power: float = 1.0,
common_scale: bool = False,
explode: bool | Sequence[int] = (),
overlay: bool | Sequence[int] = (),
figsize: Optional[tuple[int, int]] = None,
title: bool | str = True,
units: Optional[str] = None,
interact: bool = False,
display: bool = True,
**kwargs,
) -> Visualization:
"""
Show the image(s) using matplotlib.
Parameters
----------
ax : matplotlib.axes.Axes, optional
If given the plots are added to the axis. This is not available for exploded
plots.
power : float
Show image on a power scale.
explode : bool, optional
If True, a grid of images is created for all the items of the last two
ensemble axes. If False, the first ensemble item is shown. May be given as a
sequence of axis indices to create a grid of images from the specified axes.
The default is determined by the axis metadata.
overlay : bool or sequence of int, optional
If True, all line profiles in the ensemble are shown in a single plot.
If False, only the first ensemble item is shown. May be given as a sequence
of axis indices to specify which line profiles in the ensemble to show
together. The default is determined by the axis metadata.
figsize : two int, optional
The figure size given as width and height in inches, passed to
`matplotlib.pyplot.figure`.
title : bool or str, optional
Set the column title of the images. If True is given instead of a string the
title will be given by the value corresponding to the "name" key of the axes
metadata dictionary, if this item exists.
units : str
The units used for the x and y axes. The given units must be compatible with
the axes of the images.
interact : bool
If True, create an interactive visualization. This requires enabling the
ipympl Matplotlib backend.
display : bool, optional
If True (default) the figure is displayed immediately.
Returns
-------
measurement_visualization_2d : VisualizationImshow
"""
if not interact:
self.compute()
# scale_axis = self._scale_axis_from_metadata()
# base_axes_metadata = self._plot_base_axes_metadata(units)
# array = self.array
# raise RuntimeError("Cannot infer pixel area from metadata.")
# if display_axes != (-2, -1):
# array = np.moveaxis(self.array, source=display_axes, destination=(-2, -1))
# display_axes = normalize_axes(display_axes, self.shape)
# base_axes_metadata = [self.axes_metadata[i] for i in display_axes]
# ensemble_axes_metadata = [
# self.axes_metadata[i]
# for i in range(len(self.shape))
# if i not in display_axes
# ]
artist_type = LinesArtist
visualization = Visualization(
measurement=self,
ax=ax,
artist_type=artist_type,
# power=power,
aspect=False,
share_x=True,
share_y=common_scale,
common_scale=common_scale,
explode=explode,
overlay=overlay,
figsize=figsize,
# interact=interact,
title=title,
**kwargs,
)
if common_scale is False and visualization._explode:
visualization.axes.set_sizes(padding=0.8)
return visualization
class _BaseMeasurement2D(BaseMeasurements):
_base_dims = 2
@property
def base_shape(self) -> tuple[int, int]:
return super().base_shape[-2], super().base_shape[-1]
@abstractmethod
def _get_1d_equivalent(self):
pass
@property
@abstractmethod
def sampling(self) -> tuple[float, float]:
"""
Sampling of the measurements in `x` and `y` [Å] or [1/Å].
"""
@property
@abstractmethod
def extent(self) -> tuple[float, float]:
"""
Extent of measurements in `x` and `y` [Å] or [1/Å].
"""
@property
@abstractmethod
def offset(self) -> tuple[float, float]:
"""
The offset of the origin of the measurement coordinates [Å] or [1/Å].
"""
def interpolate_line(
self,
start: tuple[float, float] | Atom | None = None,
end: tuple[float, float] | Atom | None = None,
sampling: float | None = None,
gpts: int | None = None,
width: float = 0.0,
margin: float = 0.0,
order: int = 3,
endpoint: bool = False,
fractional: bool = False,
):
"""
Interpolate image(s) along a given line. Either 'sampling' or 'gpts' must be
provided.
Parameters
----------
start : two float, Atom, optional
Starting position of the line [Å] (alternatively taken from a selected
atom).
end : two float, Atom, optional
Ending position of the line [Å] (alternatively taken from a selected atom).
sampling : float
Sampling of grid points along the line [1 / Å].
gpts : int
Number of grid points along the line.
width : float, optional
The interpolation will be averaged across a perpendicular distance equal to
this width.
margin : float or tuple of float, optional
Add margin [Å] to the start and end interpolated line.
order : int, optional
The spline interpolation order.
endpoint : bool
Sets whether the ending position is included or not.
fractional : bool
If True, use fractional coordinates with respect to the extent of the
measurement.
Returns
-------
line_profiles : RealSpaceLineProfiles
The interpolated line(s).
"""
from abtem.scan import LineScan
# if self.is_complex:
# raise NotImplementedError
if (sampling is None) and (gpts is None):
sampling = min(self.sampling)
xp = get_array_module(self.array)
if start is None:
start = (0.0, 0.0)
if end is None and fractional:
end = (0.0, 1.0)
elif end is None:
end = (0.0, self.extent[0])
if fractional:
extent = self.extent
else:
extent = None
scan = LineScan(
start=start,
end=end,
gpts=gpts,
sampling=sampling,
endpoint=endpoint,
potential=extent,
fractional=fractional,
)
if margin != 0.0:
scan.add_margin(margin)
positions = xp.asarray(
(scan.get_positions(lazy=False) - self.offset) / self.sampling
)
if width:
direction = xp.array(scan.end) - xp.array(scan.start)
direction = direction / xp.linalg.norm(direction)
perpendicular_direction = xp.array([-direction[1], direction[0]])
n = xp.floor(width / min(self.sampling) / 2) * 2 + 1
perpendicular_positions = (
xp.linspace(-n / 2, n / 2, int(n))[:, None]
* perpendicular_direction[None]
)
positions = perpendicular_positions[None, :] + positions[:, None]
if self.is_lazy:
# raise NotImplementedError("Lazy interpolation not implemented.")
# TDOO: Implement lazy interpolation
base_axes = tuple(range(len(self.base_shape)))
chunks = self.array.chunks[:-2] + (positions.shape[0],)
new_axis = (base_axes[0],)
if width:
new_axis = new_axis + (new_axis[0] + 1,)
chunks = chunks + (1,)
array = da.map_blocks(
_interpolate_stack,
self.array,
positions=positions,
mode="wrap",
order=order,
drop_axis=base_axes,
new_axis=new_axis,
chunks=chunks,
meta=xp.array((), dtype=get_dtype(complex=False)),
)
else:
array = _interpolate_stack(self.array, positions, mode="wrap", order=order)
metadata = copy.copy(self.metadata)
metadata.update(scan.metadata)
metadata["label"] = "intensity"
metadata["units"] = "arb. unit"
if width:
array = array.mean(-1)
metadata["width"] = width
ensemble_axes_metadata = self.ensemble_axes_metadata
return self._get_1d_equivalent()(
array=array,
sampling=scan.sampling,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
def gaussian_filter(
self,
sigma: float | tuple[float, float],
boundary: str = "periodic",
cval: float = 0.0,
):
"""
Apply 2D gaussian filter to measurements.
Parameters
----------
sigma : float or two float
Standard deviation for the Gaussian kernel in the `x` and `y`-direction
given in physical units (Å for real-space images, 1/Å for diffraction
patterns). If given as a single number, the standard deviation is equal
for both axes.
boundary : {'periodic', 'reflect', 'constant'}
The boundary parameter determines how the images are extended beyond their
boundaries when the filter overlaps with a border.
``periodic`` :
The images are extended by wrapping around to the opposite edge. Use
this mode for periodic (default).
``reflect`` :
The images are extended by reflecting about the edge of the last
pixel.
``constant`` :
The images are extended by filling all values beyond the edge with
the same constant value, defined by the 'cval' parameter.
cval : scalar, optional
Value to fill past edges in spline interpolation input if boundary is
'constant' (default is 0.0).
Returns
-------
filtered_images : Images
The filtered image(s).
Notes
-----
The Gaussian kernel is parameterized by its standard deviation σ
(``sigma``). The corresponding full-width at half-maximum is
FWHM_G = 2√(2 ln 2)·σ ≈ 2.3548·σ.
The Lorentzian and Voigtian filters use the half-width at half-maximum
(HWHM) γ as their width parameter: for the same FWHM, γ = FWHM / 2
whereas σ = FWHM / (2√(2 ln 2)) ≈ FWHM / 2.3548.
"""
xp = get_array_module(self.array)
gaussian_filter = get_ndimage_module(self.array).gaussian_filter
if boundary == "periodic":
mode = "wrap"
elif boundary in ("reflect", "constant"):
mode = boundary
else:
raise ValueError()
if np.isscalar(sigma):
sigma = (sigma,) * 2
sigma = (0,) * (len(self.shape) - 2) + tuple(
s / d for s, d in zip(sigma, self.sampling)
)
if self.is_lazy:
depth = tuple(
min(int(np.ceil(4.0 * s)), n) for s, n in zip(sigma, self.shape)
)
array = da.map_overlap(
gaussian_filter,
self.array,
sigma=sigma,
boundary=boundary,
mode=mode,
cval=cval,
depth=depth,
meta=xp.array((), dtype=get_dtype(complex=False)),
)
else:
array = gaussian_filter(self.array, sigma=sigma, mode=mode, cval=cval)
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = array
return self.__class__(**kwargs)
def lorentzian_filter(
self,
half_width: float | tuple[float, float],
truncate: float = 10.0,
boundary: str = "periodic",
cval: float = 0.0,
):
"""
Apply a 2D Lorentzian (Cauchy) filter to measurements.
The filter convolves the image with the **true** 2-D Lorentzian profile
L(x, y) = 1 / (1 + (x/γ_x)² + (y/γ_y)²)
which is **not** separable into a product of 1-D Lorentzians (unlike a
Gaussian). A separable implementation would produce a cross-shaped
impulse response — several times brighter along the x and y axes than
along the diagonal at the same radius — which manifests as visible
halos/lines around sharp features.
Parameters
----------
half_width : float or two float
Half-width at half-maximum (HWHM) of the Lorentzian kernel in the ``x``
and ``y``-direction given in physical units (Å for real-space images,
1/Å for diffraction patterns). If given as a single number, the
half-width is equal for both axes (isotropic kernel); otherwise the
kernel is elliptical.
truncate : float, optional
Truncate the kernel at this many half-widths (default is 10.0). The
Lorentzian has heavier tails than the Gaussian, so a larger truncation
is recommended.
boundary : {'periodic', 'reflect', 'constant'}
How the image is extended beyond its boundary (default is 'periodic').
cval : float, optional
Constant value used when ``boundary`` is ``'constant'`` (default 0.0).
Returns
-------
filtered : same type as input
The filtered measurement(s).
Notes
-----
The Lorentzian kernel is parameterized by its half-width at half-maximum
(HWHM) γ (``half_width``). The corresponding full-width at half-maximum
is FWHM_L = 2γ.
Note that :meth:`gaussian_filter` uses the standard deviation σ as its
width parameter. For the same FWHM one needs γ = FWHM / 2 whereas
σ = FWHM / (2√(2 ln 2)) ≈ FWHM / 2.3548.
"""
xp = get_array_module(self.array)
if boundary == "periodic":
mode = "wrap"
elif boundary in ("reflect", "constant"):
mode = boundary
else:
raise ValueError(
f"boundary must be 'periodic', 'reflect', or 'constant', got {boundary!r}"
)
if np.isscalar(half_width):
half_width = (half_width,) * 2
hw_pixels = tuple(hw / d for hw, d in zip(half_width, self.sampling))
# Two spatial axes are the trailing dims; build a true 2-D kernel.
axes = (self.array.ndim - 2, self.array.ndim - 1)
kernel_2d = _lorentzian_kernel_2d(hw_pixels, truncate)
if self.is_lazy:
depth = [0] * self.array.ndim
depth[axes[0]] = min(kernel_2d.shape[0] // 2, self.shape[axes[0]])
depth[axes[1]] = min(kernel_2d.shape[1] // 2, self.shape[axes[1]])
array = da.map_overlap(
functools.partial(
_apply_convolve_2d_on_axes,
kernel_2d=kernel_2d,
axes=axes,
mode=mode,
cval=cval,
),
self.array,
depth=tuple(depth),
boundary=boundary,
meta=xp.array((), dtype=self.array.dtype),
)
else:
# Use FFT-based convolution: cost is independent of kernel size,
# which matters because the heavy-tailed Lorentzian kernel can
# easily be comparable in size to the image itself. The helper
# dispatches to scipy.signal or cupyx.scipy.signal based on the
# array's backend.
array = _apply_convolve_2d_on_axes(
self.array, kernel_2d, axes=axes, mode=mode, cval=cval
)
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = array
return self.__class__(**kwargs)
def voigtian_filter(
self,
gaussian_sigma: float | tuple[float, float],
lorentzian_gamma: float | tuple[float, float],
truncate: float = 10.0,
boundary: str = "periodic",
cval: float = 0.0,
):
"""
Apply a 2D Voigtian filter to measurements.
The Voigt profile is the convolution of a Gaussian and a Lorentzian.
Exploiting the associativity of convolution, the filter is implemented
as :meth:`gaussian_filter` followed by :meth:`lorentzian_filter`, which
uses the true 2-D (non-separable) Lorentzian kernel and therefore
avoids the cross-shaped artefacts a naively separable Voigt filter
would produce.
Parameters
----------
gaussian_sigma : float or two float
Standard deviation (σ) of the Gaussian component in physical units
(Å for real-space images, 1/Å for diffraction patterns). If given
as a single number it is equal for both axes.
lorentzian_gamma : float or two float
Half-width at half-maximum (HWHM, γ) of the Lorentzian component in
physical units. If given as a single number it is equal for both axes.
truncate : float, optional
Truncate the Lorentzian component at this many half-widths
(default is 10.0). The Gaussian component uses the SciPy default
(4·σ).
boundary : {'periodic', 'reflect', 'constant'}
How the image is extended beyond its boundary (default is 'periodic').
cval : float, optional
Constant value used when ``boundary`` is ``'constant'`` (default 0.0).
Returns
-------
filtered : same type as input
The filtered measurement(s).
Notes
-----
The Voigt kernel has two independent width parameters:
* ``gaussian_sigma`` (σ): standard deviation of the Gaussian component,
FWHM_G = 2√(2 ln 2)·σ ≈ 2.3548·σ.
* ``lorentzian_gamma`` (γ): half-width at half-maximum (HWHM) of the
Lorentzian component, FWHM_L = 2γ.
Because the two components use different parameterizations, σ and γ are
*not* directly comparable: for the same FWHM, γ = FWHM / 2 whereas
σ = FWHM / (2√(2 ln 2)) ≈ FWHM / 2.3548.
"""
return self.gaussian_filter(
gaussian_sigma, boundary=boundary, cval=cval
).lorentzian_filter(
lorentzian_gamma, truncate=truncate, boundary=boundary, cval=cval
)
def pseudo_voigtian_filter(
self,
gaussian_sigma: float | tuple[float, float],
lorentzian_gamma: float | tuple[float, float],
eta: float,
truncate: float = 10.0,
boundary: str = "periodic",
cval: float = 0.0,
):
"""
Apply a 2D pseudo-Voigtian filter to measurements.
The pseudo-Voigt profile is a weighted linear sum of a Gaussian and a
Lorentzian — PV = (1-η)·G + η·L — as proposed by Nguyen et al. (2014).
Exploiting the linearity of convolution, the filter is implemented as
the linear combination of :meth:`gaussian_filter` and
:meth:`lorentzian_filter`, which uses the true 2-D (non-separable)
Lorentzian kernel and therefore avoids the cross-shaped artefacts a
naively separable filter would produce.
Parameters
----------
gaussian_sigma : float or two float
Standard deviation (σ) of the Gaussian component in physical units
(Å for real-space images, 1/Å for diffraction patterns). If given
as a single number it is equal for both axes.
lorentzian_gamma : float or two float
Half-width at half-maximum (HWHM, γ) of the Lorentzian component in
physical units. If given as a single number it is equal for both axes.
eta : float
Lorentzian mixing fraction η ∈ [0, 1]. η = 0 gives a pure Gaussian
(equivalent to :meth:`gaussian_filter`); η = 1 gives a pure
Lorentzian (equivalent to :meth:`lorentzian_filter`).
truncate : float, optional
Truncate the Lorentzian component at this many half-widths
(default 10.0). The Gaussian component uses the SciPy default (4·σ).
boundary : {'periodic', 'reflect', 'constant'}
How the image is extended beyond its boundary (default is 'periodic').
cval : float, optional
Constant value used when ``boundary`` is ``'constant'`` (default 0.0).
Returns
-------
filtered : same type as input
The filtered measurement(s).
Notes
-----
Width parameterization:
* ``gaussian_sigma`` (σ): standard deviation; FWHM_G = 2√(2 ln 2)·σ ≈ 2.3548·σ.
* ``lorentzian_gamma`` (γ): HWHM; FWHM_L = 2γ.
For the same FWHM, γ = FWHM / 2 whereas σ = FWHM / (2√(2 ln 2)) ≈ FWHM / 2.3548.
"""
if eta == 0.0:
return self.gaussian_filter(gaussian_sigma, boundary=boundary, cval=cval)
if eta == 1.0:
return self.lorentzian_filter(
lorentzian_gamma, truncate=truncate, boundary=boundary, cval=cval
)
g = self.gaussian_filter(gaussian_sigma, boundary=boundary, cval=cval)
l = self.lorentzian_filter(
lorentzian_gamma, truncate=truncate, boundary=boundary, cval=cval
)
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = (1.0 - eta) * g.array + eta * l.array
return self.__class__(**kwargs)
def interpolate_line_at_position(
self,
center: tuple[float, float] | Atom,
angle: float,
extent: float,
gpts: Optional[int] = None,
sampling: Optional[float] = None,
width: float = 0.0,
order: int = 3,
endpoint: bool = True,
):
"""
Interpolate image(s) along a line centered at a specified position.
Parameters
----------
center : two float
Center position of the line [Å]. May be given as an Atom.
angle : float
Angle of the line [deg.].
extent : float
Extent of the line [Å].
gpts : int
Number of grid points along the line.
sampling : float
Sampling of grid points along the line [Å].
width : float, optional
The interpolation will be averaged across a perpendicular distance equal to
this width.
order : int, optional
The spline interpolation order.
endpoint : bool
Sets whether the ending position is included or not.
Returns
-------
line_profiles : RealSpaceLineProfiles or ReciprocalSpaceProfiles
The interpolated line(s).
"""
from abtem.scan import LineScan
scan = LineScan.at_position(center=center, extent=extent, angle=angle)
return self.interpolate_line(
scan.start,
scan.end,
gpts=gpts,
sampling=sampling,
width=width,
order=order,
endpoint=endpoint,
)
def show(
self,
ax: Optional[Axes] = None,
cbar: bool = False,
cmap: Optional[str] = None,
vmin: Optional[float] = None,
vmax: Optional[float] = None,
power: float = 1.0,
logscale: bool = False,
common_color_scale: bool = False,
explode: bool | Sequence[int] = (),
overlay: bool | Sequence[int] = (),
figsize: Optional[tuple[int, int]] = None,
title: bool | str = True,
units: Optional[str] = None,
interact: bool = False,
display: bool = True,
**kwargs,
) -> Visualization:
"""
Show the image(s) using matplotlib.
Parameters
----------
ax : matplotlib.axes.Axes, optional
If given the plots are added to the axis. This is not available for exploded
plots.
cbar : bool, optional
Add colorbar(s) to the image(s). The size and padding of the colorbars may
be adjusted using the `set_cbar_size` and `set_cbar_padding` methods.
cmap : str, optional
Matplotlib colormap name used to map scalar data to colors. If the
measurement is complex the colormap must be one of 'hsv' or 'hsluv'.
vmin : float, optional
Minimum of the intensity color scale. Default is the minimum of the array
values.
vmax : float, optional
Maximum of the intensity color scale. Default is the maximum of the array
values.
power : float
Show image on a power scale. Cannot be used together with ``logscale``.
logscale : bool
If True, show image on a logarithmic intensity scale. Cannot be used
together with ``power != 1.0``.
common_color_scale : bool, optional
If True all images in an image grid are shown on the same colorscale, and a
single colorbar is created (if it is requested). Default is False.
explode : bool, optional
If True, a grid of images is created for all the items of the last two
ensemble axes. If False, the first ensemble item is shown. May be given as a
sequence of axis indices to create a grid of images from the specified axes.
The default is determined by the axis metadata.
overlay : bool or sequence of int, optional
If True, all line profiles in the ensemble are shown in a single plot.
If False, only the first ensemble item is shown. May be given as a sequence
of axis indices to specify which line profiles in the ensemble to show
together. The default is determined by the axis metadata.
figsize : two int, optional
The figure size given as width and height in inches, passed to
`matplotlib.pyplot.figure`.
title : bool or str, optional
Set the column title of the images. If True is given instead of a string the
title will be given by the
value corresponding to the "name" key of the axes metadata dictionary, if
this item exists.
units : str
The units used for the x and y axes. The given units must be compatible with
the axes of the images.
interact : bool
If True, create an interactive visualization. This requires enabling the
ipympl Matplotlib backend.
display : bool, optional
If True (default) the figure is displayed immediately.
Returns
-------
measurement_visualization_2d : VisualizationImshow
"""
visualization = Visualization(
measurement=self,
ax=ax,
common_scale=common_color_scale,
figsize=figsize,
title=title,
aspect=True,
share_x=True,
share_y=True,
explode=explode,
overlay=overlay,
interactive=not interact and display,
value_limits=(vmin, vmax),
power=power,
logscale=logscale,
cmap=cmap,
cbar=cbar,
units=units,
**kwargs,
)
if interact:
visualization.interact(ImageGUI, display=display)
return visualization
[docs]
class Images(_BaseMeasurement2D):
"""
A collection of 2D measurements such as HRTEM or STEM-ADF images. May be used to
represent a reconstructed phase.
Parameters
----------
array : numpy.ndarray
2D or greater array containing data of type `float` or `complex`. The
second-to-last and last
dimensions are the image `y`- and `x`-axis, respectively.
sampling : two float
Lateral sampling of images in `x` and `y` [Å].
ensemble_axes_metadata : list of AxisMetadata, optional
List of metadata associated with the ensemble axes. The length and item order
must match the ensemble axes.
metadata : dict, optional
A dictionary defining measurement metadata.
"""
def __init__(
self,
array: da.core.Array | np.array,
sampling: float | tuple[float, float],
ensemble_axes_metadata: Optional[list[AxisMetadata]] = None,
metadata: Optional[Dict] = None,
):
if np.isscalar(sampling):
sampling = (float(sampling),) * 2
else:
sampling = float(sampling[0]), float(sampling[1])
self._sampling = sampling
super().__init__(
array=array,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
def _get_1d_equivalent(self):
return RealSpaceLineProfiles
@property
def _area_per_pixel(self):
return np.prod(self.sampling)
@property
def sampling(self) -> tuple[float, float]:
return self._sampling
@property
def offset(self) -> tuple[float, float]:
return 0.0, 0.0
@property
def extent(self) -> tuple[float, float]:
return (
self.sampling[0] * self.base_shape[0],
self.sampling[1] * self.base_shape[1],
)
@property
def coordinates(self) -> tuple[np.ndarray, np.ndarray]:
"""Coordinates of pixels in `x` and `y` [Å]."""
x = np.linspace(0.0, self.shape[-2] * self.sampling[0], self.shape[-2])
y = np.linspace(0.0, self.shape[-1] * self.sampling[1], self.shape[-1])
return x, y
@property
def base_axes_metadata(self) -> list[AxisMetadata]:
return [
RealSpaceAxis(
label="x", sampling=self.sampling[0], units="Å", tex_label="$x$"
),
RealSpaceAxis(
label="y", sampling=self.sampling[1], units="Å", tex_label="$y$"
),
]
[docs]
def integrate_disc(self, position: np.ndarray, radius: float) -> float:
"""
Integrate the values of the images on a disc-shaped region.
Parameters
----------
position : two floats
Center of disc-shaped integration region [Å].
radius : float
Radius of disc-shaped integration region [Å].
Returns
-------
float
Integral value.
"""
return integrate_disc(self, position=position, radius=radius)
[docs]
def integrate_gradient(self):
"""
Calculate integrated gradients. Requires complex images whose real and imaginary
parts represent the `x` and `y` components of a gradient.
Returns
-------
integrated_gradient : Images
The integrated gradient.
"""
self._check_is_complex()
if self.is_lazy:
xp = get_array_module(self.array)
array = self.array.rechunk(
self.array.chunks[:-2] + ((self.shape[-2],), (self.shape[-1],))
)
array = array.map_blocks(
_integrate_gradient_2d,
sampling=self.sampling,
meta=xp.array((), dtype=get_dtype(complex=False)),
)
else:
array = _integrate_gradient_2d(self.array, sampling=self.sampling)
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = array
kwargs["metadata"].update({"label": "iCOM", "units": "arb. unit"})
return self.__class__(**kwargs)
[docs]
def crop(
self,
extent: tuple[float, float],
offset: tuple[float, float] = (0.0, 0.0),
centered: bool = False,
):
"""
Crop images to a smaller extent.
Parameters
----------
extent : tuple of float
Extent of rectangular cropping region in `x` and `y` [Å].
offset : tuple of float
Lower corner of cropping region in `x` and `y` [Å] (default is (0,0)).
Returns
-------
cropped_images : Images
The cropped images.
"""
if centered and offset == (0.0, 0.0):
offset = (
self.extent[0] / 2 - extent[0] / 2,
self.extent[1] / 2 - extent[1] / 2,
)
elif centered:
raise ValueError("Offset is not used when centered is True.")
if extent[0] > self.extent[0] or extent[1] > self.extent[1]:
raise ValueError("Extent must be smaller than the original extent.")
offset = (
int(np.round(self.base_shape[0] * offset[0] / self.extent[0])),
int(np.round(self.base_shape[1] * offset[1] / self.extent[1])),
)
new_shape = (
int(np.round(self.base_shape[0] * extent[0] / self.extent[0])),
int(np.round(self.base_shape[1] * extent[1] / self.extent[1])),
)
array = self.array[
...,
offset[0] : offset[0] + new_shape[0],
offset[1] : offset[1] + new_shape[1],
]
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = array
return self.__class__(**kwargs)
@staticmethod
def _interpolate_spline(array, old_gpts, new_gpts, pad_mode, order, cval):
xp = get_array_module(array)
x = xp.linspace(0.0, old_gpts[0], new_gpts[0], endpoint=False)
y = xp.linspace(0.0, old_gpts[1], new_gpts[1], endpoint=False)
positions = xp.meshgrid(x, y, indexing="ij")
positions = xp.stack(positions, axis=-1)
return _interpolate_stack(array, positions, pad_mode, order=order, cval=cval)
[docs]
def interpolate(
self,
sampling: Optional[float | tuple[float, float]] = None,
gpts: Optional[int | tuple[int, int]] = None,
method: str = "fft",
boundary: str = "periodic",
order: int = 3,
normalization: str = "values",
cval: float = 0.0,
) -> Images:
"""
Interpolate images producing equivalent images with a different sampling.
Either 'sampling' or 'gpts' must be provided (but not both).
Parameters
----------
sampling : float or two float
Sampling of images after interpolation in `x` and `y` [Å].
gpts : int or two int
Number of grid points of images after interpolation in `x` and `y`.
Do not use if 'sampling' is used.
method : {'fft', 'spline'}
The interpolation method.
``fft`` :
Interpolate by cropping or zero-padding in reciprocal space.
This method should be preferred for periodic images.
``spline`` :
Interpolate using spline interpolation. This method should be
preferred for non-periodic images.
boundary : {'periodic', 'reflect', 'constant'}
The boundary parameter determines how the input array is extended beyond its
boundaries for spline interpolation.
``periodic`` :
The images are extended by wrapping around to the opposite edge.
Use this mode for periodic images (default).
``reflect`` :
The images are extended by reflecting about the edge of the last
pixel.
``constant`` :
The images are extended by filling all values beyond the edge with
the same constant value, defined by the 'cval' parameter.
order : int
The order of the spline interpolation (default is 3). The order has to be in
the range 0-5.
normalization : {'values', 'amplitude'}
The normalization parameter determines which quantity is preserved after
normalization.
``values`` :
The pixel-wise values of the images are preserved.
``intensity`` :
The total intensity of the images is preserved.
cval : scalar, optional
Value to fill past edges in spline interpolation input if boundary is
'constant' (default is 0.0).
Returns
-------
interpolated_images : Images
The interpolated images.
"""
if method == "fft" and boundary != "periodic":
raise ValueError(
"Only periodic boundaries available for FFT interpolation."
)
if sampling is None and gpts is None:
raise ValueError()
if gpts is None and sampling is not None:
if np.isscalar(sampling):
sampling = (sampling,) * 2
gpts = tuple(int(np.ceil(l / d)) for d, l in zip(sampling, self.extent))
elif gpts is not None:
if np.isscalar(gpts):
gpts = (gpts,) * 2
else:
raise ValueError()
xp = get_array_module(self.array)
sampling = (self.extent[0] / gpts[0], self.extent[1] / gpts[1])
if boundary == "periodic":
boundary = "wrap"
array = None
if self.is_lazy:
array = self.array.rechunk(
chunks=self.array.chunks[:-2] + ((self.shape[-2],), (self.shape[-1],))
)
if method == "fft":
array = array.map_blocks(
fft_interpolate,
new_shape=gpts,
normalization=normalization,
chunks=self.array.chunks[:-2] + ((gpts[0],), (gpts[1],)),
meta=xp.array((), dtype=self.array.dtype),
)
elif method == "spline":
array = array.map_blocks(
self._interpolate_spline,
old_gpts=self.shape[-2:],
new_gpts=gpts,
order=order,
cval=cval,
pad_mode=boundary,
chunks=self.array.chunks[:-2] + ((gpts[0],), (gpts[1],)),
meta=xp.array((), dtype=self.array.dtype),
)
else:
if method == "fft":
array = fft_interpolate(self.array, gpts, normalization=normalization)
elif method == "spline":
array = self._interpolate_spline(
self.array,
old_gpts=self.shape[-2:],
new_gpts=gpts,
pad_mode=boundary,
order=order,
cval=cval,
)
if array is None:
raise RuntimeError()
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["sampling"] = sampling
kwargs["array"] = array
return self.__class__(**kwargs)
[docs]
def tile(self, repetitions: tuple[int, int]) -> Images:
"""
Tile image(s).
Parameters
----------
repetitions : tuple of int
The number of repetitions of the images along the `x`- and `y`-axis,
respectively.
Returns
-------
tiled_images : Images
The tiled image(s).
"""
if len(repetitions) != 2:
raise RuntimeError()
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = np.tile(
self.array, (1,) * (len(self.array.shape) - 2) + repetitions
)
return self.__class__(**kwargs)
[docs]
def scan_noise(
self,
dwell_time: float,
flyback_time: float,
rms_power: float,
max_frequency: float = 500.0,
num_components: int = 200,
seed: Optional[int] = None,
):
"""
Apply scan noise to images.
Parameters
----------
dwell_time : float
Dwell time of the beam [s].
flyback_time : float
Flyback time of the beam [s].
rms_power : float
RMS power of the scan noise [V].
max_frequency : float
Maximum frequency of the scan noise [1/Å].
"""
transform = ScanNoiseTransform(
dwell_time=dwell_time,
flyback_time=flyback_time,
rms_power=rms_power,
max_frequency=max_frequency,
num_components=num_components,
seeds=seed,
)
return self.apply_transform(transform)
@staticmethod
def _diffractograms(array):
xp = get_array_module(array)
array = xp.fft.fft2(array)
return xp.fft.fftshift(xp.abs(array), axes=(-2, -1))
[docs]
def diffractograms(self) -> DiffractionPatterns:
"""
Calculate diffractograms (i.e. power spectra) from image(s).
Returns
-------
diffractograms : DiffractionPatterns
Diffractograms of image(s).
"""
xp = get_array_module(self.array)
if self.is_lazy:
array = self.array.rechunk(
chunks=self.array.chunks[:-2] + ((self.shape[-2],), (self.shape[-1],))
)
array = array.map_blocks(
self._diffractograms, meta=xp.array((), dtype=get_dtype(complex=False))
)
else:
array = self._diffractograms(self.array)
sampling = 1 / self.extent[0], 1 / self.extent[1]
return DiffractionPatterns(
array=array,
sampling=sampling,
ensemble_axes_metadata=self.ensemble_axes_metadata,
metadata=self.metadata,
)
# def _plot_base_axes_metadata(self, units: str = None):
# return self.base_axes_metadata
class _BaseMeasurement1D(BaseMeasurements):
_base_dims = 1
def __init__(
self,
array: np.ndarray,
sampling: Optional[float] = None,
ensemble_axes_metadata: Optional[list[AxisMetadata]] = None,
metadata: Optional[dict] = None,
):
self._sampling = sampling
super().__init__(
array=array,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
@property
def _area_per_pixel(self):
raise RuntimeError("Cannot infer pixel area from metadata.")
@classmethod
def from_array_and_metadata(
cls,
array: np.ndarray,
axes_metadata: list[AxisMetadata],
metadata: Optional[dict] = None,
) -> Self:
"""
Creates line profile(s) from a given array and metadata.
Parameters
----------
array : array
Complex array defining one or more 1D line profiles.
axes_metadata : list of AxesMetadata
Axis metadata for each axis. The axis metadata must be compatible with the
shape of the array. The last two axes must be RealSpaceAxis.
metadata : dict, optional
A dictionary defining the measurement metadata.
Returns
-------
line_profiles : RealSpaceLineProfiles
Line profiles from the array and metadata.
"""
x_axis = axes_metadata[-1]
if isinstance(x_axis, LinearAxis):
sampling = x_axis.sampling
else:
raise RuntimeError()
axes_metadata = axes_metadata[:-1]
return cls(
array,
sampling=sampling,
ensemble_axes_metadata=axes_metadata,
metadata=metadata,
)
@property
def extent(self) -> float:
"""
Extent of measurements [Å] or [1/Å].
"""
return self.sampling * self.shape[-1]
@property
def sampling(self) -> float:
"""
Extent of measurements [Å] or [1/Å].
"""
return self._sampling
@property
@abstractmethod
def base_axes_metadata(self) -> list[RealSpaceAxis | ReciprocalSpaceAxis]:
pass
def _line_scan(self, sampling=None):
start, end = self.metadata["start"], self.metadata["end"]
from abtem.scan import LineScan
return LineScan(start=start, end=end, sampling=sampling)
def _add_to_plot(self, *args, **kwargs):
if not all(key in self.metadata for key in ("start", "end")):
raise RuntimeError(
"The metadata does not contain the keys 'start' and 'end'"
)
if "width" in self.metadata:
kwargs["width"] = self.metadata["width"]
self._line_scan().add_to_plot(*args, **kwargs)
@staticmethod
def _calculate_widths(array, sampling, height):
xp = get_array_module(array)
array = array - xp.max(array, axis=-1, keepdims=True) * height
widths = xp.zeros(array.shape[:-1], dtype=np.float32)
for i in np.ndindex(array.shape[:-1]):
zero_crossings = xp.where(xp.diff(xp.sign(array[i]), axis=-1))[0]
left, right = zero_crossings[0], zero_crossings[-1]
widths[i] = (right - left) * sampling
return widths
def width(self, height: float = 0.5):
"""
Calculate the width of line(s) at a given height, e.g. full width at half
maximum (the default).
Parameters
----------
height : float
Fractional height at which the width is calculated.
Returns
-------
width : float
The calculated width.
"""
if self.is_lazy:
return self.array.map_blocks(
self._calculate_widths,
drop_axis=(len(self.array.shape) - 1,),
dtype=np.float32,
sampling=self.sampling,
height=height,
)
else:
return self._calculate_widths(self.array, self.sampling, height)
@staticmethod
def _interpolate(array, gpts, endpoint, order):
xp = get_array_module(array)
map_coordinates = get_ndimage_module(array).map_coordinates
old_shape = array.shape
array = array.reshape((-1, array.shape[-1]))
array = xp.pad(array, ((0,) * 2, (3,) * 2), mode="wrap")
new_points = xp.linspace(3.0, array.shape[-1] - 3.0, gpts, endpoint=endpoint)[
None
]
new_array = xp.zeros(array.shape[:-1] + (gpts,), dtype=xp.float32)
for i in range(len(array)):
map_coordinates(array[i], new_points, new_array[i], order=order)
return new_array.reshape(old_shape[:-1] + (gpts,))
def interpolate(
self,
sampling: Optional[float] = None,
gpts: Optional[int] = None,
order: int = 3,
endpoint: bool = False,
) -> BaseMeasurementsSubclass:
"""
Interpolate line profile(s) producing equivalent line profile(s) with a
different sampling. Either 'sampling' or 'gpts' must be provided (but not both).
Parameters
----------
sampling : float, optional
Sampling of line profiles after interpolation [Å].
gpts : int, optional
Number of grid points of line profiles after interpolation. Do not use if
'sampling' is used.
order : int, optional
The order of the spline interpolation (default is 3). The order has to be in
the range 0-5.
endpoint : bool, optional
If True, end is the last position. Otherwise, it is not included.
Default is False.
Returns
-------
interpolated_profiles : RealSpaceLineProfiles
The interpolated line profile(s).
"""
xp = get_array_module(self.array)
if (gpts is not None) and (sampling is not None):
raise RuntimeError()
if sampling is None and gpts is None:
sampling = self.sampling
if gpts is None:
gpts = int(np.ceil(self.extent / sampling))
if sampling is None:
sampling = self.extent / gpts
if self.is_lazy:
array = self.array.rechunk(self.array.chunks[:-1] + ((self.shape[-1],),))
array = array.map_blocks(
self._interpolate,
gpts=gpts,
endpoint=endpoint,
order=order,
chunks=self.array.chunks[:-1] + (gpts,),
meta=xp.array((), dtype=get_dtype(complex=False)),
)
else:
array = self._interpolate(self.array, gpts, endpoint, order)
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = array
kwargs["sampling"] = sampling
return self.__class__(**kwargs)
def show(
self,
ax: Optional[Axes] = None,
common_scale: bool = True,
explode: bool | Sequence[int] = False,
overlay: Optional[bool | Sequence[int]] = None,
figsize: Optional[tuple[int, int]] = None,
title: str | bool = True,
units: Optional[str] = None,
legend: bool = False,
interact: bool = False,
display: bool = True,
**kwargs,
) -> Visualization:
"""
Show the reciprocal-space line profile(s) using matplotlib.
Parameters
----------
ax : matplotlib Axes, optional
If given the plots are added to the Axes. This is not available for image
grids.
common_scale : bool
If True all plots are shown with a common y-axis. Default is False.
explode : bool or sequence of bool, optional
If True, a grid of plots is created for all the items of the last two
ensemble axes. If False, only the one plot is created. May be given as a
sequence of axis indices to create a grid of plots from the specified
axes. The default is determined by the axis metadata.
overlay : bool or sequence of int, optional
If True, all line profiles in the ensemble are shown in a single plot.
If False, only the first ensemble item is shown. May be given as a sequence
of axis indices to specify which line profiles in the ensemble to show
together. The default is determined by the axis metadata.
figsize : two int, optional
The figure size given as width and height in inches, passed to
`matplotlib.pyplot.figure`.
title : bool or str, optional
Set the column title of the plots. If True is given instead of a string the
title will be given by the value corresponding to the "name" key of the axes
metadata dictionary, if this item exists.
legend : bool
Add a legend to the plot. The labels will be derived from
units : str, optional
The units used for the x-axis. The given units must be compatible.
interact : bool
If True, create an interactive visualization. This requires enabling the
`ipympl` Matplotlib backend.
display : bool, optional
If True (default) the figure is displayed immediately.
Returns
-------
visualization : Visualization
"""
if overlay is None and explode is False:
overlay = True
elif overlay is False or overlay is None:
overlay = ()
visualization = Visualization(
measurement=self,
ax=ax,
figsize=figsize,
title=title,
aspect=False,
share_x=True,
share_y=common_scale,
explode=explode,
overlay=overlay,
interactive=not interact and display,
legend=legend,
common_scale=common_scale,
**kwargs,
)
if interact:
visualization.interact(LinesGUI, display=display)
if common_scale is False and visualization._explode:
visualization.axes.set_sizes(padding=0.8)
return visualization
[docs]
class RealSpaceLineProfiles(_BaseMeasurement1D):
"""
A collection of real-space line profile(s).
Parameters
----------
array : numpy.ndarray
1D or greater array containing data of type `float` or `complex`.
sampling : float
Sampling of line profiles [Å].
ensemble_axes_metadata : list of AxisMetadata, optional
List of metadata associated with the ensemble axes. The length and item order
must match the ensemble axes.
metadata : dict, optional
A dictionary defining measurement metadata.
"""
def __init__(
self,
array: np.ndarray,
sampling: Optional[float] = None,
ensemble_axes_metadata: Optional[list[AxisMetadata]] = None,
metadata: Optional[dict] = None,
):
super().__init__(
array=array,
sampling=sampling,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
@property
def base_axes_metadata(self) -> list[RealSpaceAxis]:
return [
RealSpaceAxis(label="r", sampling=self.sampling, units="Å", tex_label="$r$")
]
[docs]
def tile(self, repetitions: int) -> "RealSpaceLineProfiles":
"""
Tile line profiles(s).
Parameters
----------
repetitions : int
The number of repetitions of the line profiles.
Returns
-------
tiled_line_profiles : RealSpaceLineProfiles
The tiled line profiles(s).
"""
kwargs = self._copy_kwargs(exclude=("array",))
xp = get_array_module(self.array)
reps = (1,) * (len(self.array.shape) - 1) + (repetitions,)
if self.is_lazy:
kwargs["array"] = da.tile(self.array, reps)
else:
kwargs["array"] = xp.tile(self.array, reps)
return self.__class__(**kwargs)
# def _plot_extent(self, units=None):
# scale = get_conversion_factor(units, "Å")
# return [0, self.extent * scale]
#
# # def _plot_extent(self, units=None):
# # scale = {"Å": 1, "nm": 0.1}[_validate_real_space_units(units)]
# # return [0, self.extent * scale]
#
# def _plot_x_label(self, units=None):
# return f"x [{_validate_units(units, 'Å')}]"
#
# def _plot_y_label(self, units=None):
# return f"y [{_validate_units(units, 'Å')}]"
[docs]
class ReciprocalSpaceLineProfiles(_BaseMeasurement1D):
"""
A collection of reciprocal-space line profile(s).
Parameters
----------
array : numpy.ndarray
1D or greater array containing data of type `float` or `complex`.
sampling : float
Sampling of line profiles [1 / Å].
ensemble_axes_metadata : list of AxisMetadata, optional
List of metadata associated with the ensemble axes. The length and item order
must match the ensemble axes.
metadata : dict, optional
A dictionary defining measurement metadata.
"""
def __init__(
self,
array: np.ndarray,
sampling: Optional[float] = None,
ensemble_axes_metadata: Optional[list[AxisMetadata]] = None,
metadata: Optional[dict] = None,
):
super().__init__(
array=array,
sampling=sampling,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
@property
def base_axes_metadata(self) -> list[AxisMetadata]:
return [
ReciprocalSpaceAxis(
label="k", sampling=self.sampling, units="1/Å", tex_label="$k$"
)
]
@property
def angular_extent(self):
"""Extent of line profiles given as scattering angels [mrad]."""
wavelength = energy2wavelength(self._get_energy())
return self.extent * wavelength * 1e3
# def _plot_x_label(self, units=None):
# return f"x [{_validate_units(units, '1/Å')}]"
#
# def _plot_y_label(self, units=None):
# return f"y [{_validate_units(units, '1/Å')}]"
#
# def _plot_extent(self, units=None):
# if units is None:
# units = "1/Å"
#
# if units == "mrad":
# return [0, self.angular_extent]
# elif units == "1/Å":
# return [0, self.extent]
def _integrate_gradient_2d(gradient, sampling):
xp = get_array_module(gradient)
gx, gy = gradient.real, gradient.imag
(nx, ny) = gx.shape[-2:]
ikx = xp.fft.fftfreq(nx, d=sampling[0])
iky = xp.fft.fftfreq(ny, d=sampling[1])
grid_ikx, grid_iky = xp.meshgrid(ikx, iky, indexing="ij")
k = grid_ikx**2 + grid_iky**2
k[k == 0] = 1e-12
That = (xp.fft.fft2(gx) * grid_ikx + xp.fft.fft2(gy) * grid_iky) / (2j * np.pi * k)
T = xp.real(xp.fft.ifft2(That))
T -= xp.min(T)
return T
def _fourier_space_bilinear_nodes_and_weight(
old_shape: tuple[int, int],
new_shape: tuple[int, int],
old_angular_sampling: tuple[float, float],
new_angular_sampling: tuple[float, float],
xp,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
nodes = []
weights = []
old_sampling = (
1 / old_angular_sampling[0] / old_shape[0],
1 / old_angular_sampling[1] / old_shape[1],
)
new_sampling = (
1 / new_angular_sampling[0] / new_shape[0],
1 / new_angular_sampling[1] / new_shape[1],
)
for n, m, r, d in zip(old_shape, new_shape, old_sampling, new_sampling):
k = xp.fft.fftshift(xp.fft.fftfreq(n, r).astype(xp.float32))
k_new = xp.fft.fftshift(xp.fft.fftfreq(m, d).astype(xp.float32))
distances = k_new[None] - k[:, None]
distances[distances < 0.0] = np.inf
w = distances.min(0) / (k[1] - k[0])
w[w == np.inf] = 0.0
nodes.append(distances.argmin(0))
weights.append(w)
v, u = nodes
vw, uw = weights
v, u, vw, uw = xp.broadcast_arrays(v[:, None], u[None, :], vw[:, None], uw[None, :])
return v, u, vw, uw
# Threshold below which (1/hw)^2 would overflow float64.
# For the Lorentzian kernel the minimum radius is 1, so the maximum operand
# is x/hw = 1/hw; squaring overflows when hw < 1/sqrt(float_max).
_LORENTZIAN_SAFE_HW = 1.0 / np.sqrt(np.finfo(np.float64).max) # ≈ 7.5e-155
def _lorentzian_kernel_2d(
hw_pixels: tuple[float, float],
truncate: float = 10.0,
) -> np.ndarray:
"""Build a normalized 2-D Lorentzian (Cauchy) convolution kernel in pixel units.
The kernel is the **true** 2-D Lorentzian profile
L(x, y) = 1 / (1 + (x/γ_x)² + (y/γ_y)²)
which is **not** separable into a product of 1-D Lorentzians. Filtering with
a separable Lorentzian (outer product of two 1-D kernels) produces a
non-rotationally-symmetric impulse response that is several times brighter
along the axes than along the diagonal at the same radius, manifesting as
cross-shaped halos around sharp features. This 2-D kernel avoids that
artefact (it is elliptical when γ_x ≠ γ_y, isotropic when γ_x = γ_y).
The returned kernel uses the dtype configured via ``abtem.config['precision']``
(``float32`` or ``float64``), so the filter respects the global precision
setting.
Parameters
----------
hw_pixels : tuple of two float
Half-width at half-maximum along each axis (axis 0, axis 1) in pixel
units. An axis with HWHM below ``_LORENTZIAN_SAFE_HW`` is treated as
zero and contributes a delta along that axis.
truncate : float
Truncate the kernel at this many half-widths along each axis.
"""
out_dtype = get_dtype(complex=False)
hw0, hw1 = hw_pixels
zero0 = hw0 < _LORENTZIAN_SAFE_HW
zero1 = hw1 < _LORENTZIAN_SAFE_HW
if zero0 and zero1:
return np.array([[1.0]], dtype=out_dtype)
radius0 = 0 if zero0 else int(np.ceil(truncate * hw0))
radius1 = 0 if zero1 else int(np.ceil(truncate * hw1))
# Always compute the kernel in float64 to keep the small-tail values
# accurate, then cast down to the configured precision at the end.
y = np.arange(-radius0, radius0 + 1, dtype=np.float64)[:, None]
x = np.arange(-radius1, radius1 + 1, dtype=np.float64)[None, :]
denom = np.broadcast_to(np.ones(()), (y.shape[0], x.shape[1])).copy()
if not zero0:
denom = denom + (y / hw0) ** 2
if not zero1:
denom = denom + (x / hw1) ** 2
kernel = 1.0 / denom
return (kernel / kernel.sum()).astype(out_dtype)
def _apply_convolve_2d_on_axes(array, kernel_2d, axes, mode, cval=0.0):
"""Apply a 2-D convolution on a specified pair of axes of an n-D array.
Uses FFT-based convolution (``scipy.signal.fftconvolve`` on CPU,
``cupyx.scipy.signal.fftconvolve`` on GPU) so the cost is O(N log N) per
spatial slice — independent of kernel size. The heavy tails of the
Lorentzian and Voigtian kernels can make them comparable in size to the
image itself, which would be prohibitively slow with direct convolution.
The image is first padded according to ``mode`` so that the valid-mode
convolution returns the same shape as the input.
The kernel is broadcast to ``array.ndim`` with size-1 dimensions on all
axes other than ``axes``. Suitable for ``dask.array.map_overlap``.
"""
xp = get_array_module(array)
scipy_signal = get_scipy_module(array).signal
kh, kw = kernel_2d.shape
# Delta kernel → no-op shortcut.
if kh == 1 and kw == 1:
return (array * float(kernel_2d[0, 0])).astype(array.dtype, copy=False)
# Kernels from _lorentzian_kernel_2d always have odd sizes; `valid`-mode
# convolution on a (kh//2, kw//2)-padded array then returns the original
# shape exactly.
ph = kh // 2
pw = kw // 2
pad_widths = [(0, 0)] * array.ndim
pad_widths[axes[0]] = (ph, ph)
pad_widths[axes[1]] = (pw, pw)
if mode == "wrap":
padded = xp.pad(array, pad_widths, mode="wrap")
elif mode == "reflect":
padded = xp.pad(array, pad_widths, mode="reflect")
elif mode == "constant":
padded = xp.pad(array, pad_widths, mode="constant", constant_values=cval)
else:
raise ValueError(f"Unknown convolution mode: {mode!r}")
shape = [1] * array.ndim
shape[axes[0]] = kh
shape[axes[1]] = kw
kernel_nd = xp.asarray(kernel_2d).reshape(shape)
# Convolve over all axes (no ``axes=`` kwarg). The size-1 dims of the
# kernel make the FFT along non-target axes a length-1 no-op, so the
# result is numerically identical to passing ``axes=axes`` — verified
# against scipy.signal.fftconvolve with explicit axes.
#
# cupyx.scipy.signal.fftconvolve internally uses cupyx.jit.rawkernel,
# which emits a FutureWarning at import / decorator-application time
# ("cupyx.jit.rawkernel is experimental ..."). That warning is purely
# an upstream API-stability notice, so it is silenced project-wide via
# the ``filterwarnings`` block in pyproject.toml. If CuPy changes the
# warning text, the category, or the fftconvolve interface itself, the
# filter stops matching and the GPU tests catch the change.
result = scipy_signal.fftconvolve(padded, kernel_nd, mode="valid")
return result.astype(array.dtype, copy=False)
return array
def _lorentzian_source_size(
measurements,
half_width: float | tuple[float, float],
truncate: float = 10.0,
):
"""Apply a true 2-D Lorentzian convolution across the two scan axes.
Mirrors :func:`_gaussian_source_size` but uses the non-separable 2-D
Lorentzian kernel from :func:`_lorentzian_kernel_2d`, avoiding the
cross-shaped artefacts that a separable 1-D × 1-D Lorentzian would
produce around sharp features.
"""
scan_axes = _scan_axes(measurements)
if len(scan_axes) < 2:
raise RuntimeError(
"Lorentzian source size not implemented for diffraction patterns with less"
" than two scan axes."
)
if np.isscalar(half_width):
half_width = (half_width,) * 2
xp = get_array_module(measurements.array)
scan_sampling = _scan_sampling(measurements)
hw_pixels = (
half_width[0] / scan_sampling[0],
half_width[1] / scan_sampling[1],
)
kernel_2d = _lorentzian_kernel_2d(hw_pixels, truncate)
axes = (scan_axes[0], scan_axes[1])
if measurements.is_lazy:
depth = [0] * measurements.array.ndim
depth[axes[0]] = min(kernel_2d.shape[0] // 2, measurements.shape[axes[0]])
depth[axes[1]] = min(kernel_2d.shape[1] // 2, measurements.shape[axes[1]])
array = measurements.array.map_overlap(
functools.partial(
_apply_convolve_2d_on_axes,
kernel_2d=kernel_2d,
axes=axes,
mode="wrap",
),
depth=tuple(depth),
boundary="periodic",
meta=xp.array((), dtype=measurements.array.dtype),
)
else:
array = _apply_convolve_2d_on_axes(
measurements.array, kernel_2d, axes=axes, mode="wrap"
)
kwargs = measurements._copy_kwargs(exclude=("array",))
return measurements.__class__(array, **kwargs)
def _voigtian_source_size(
measurements,
gaussian_sigma: float | tuple[float, float],
lorentzian_gamma: float | tuple[float, float],
truncate: float = 10.0,
):
"""Apply a Voigtian source-size convolution via Gaussian then Lorentzian.
Exploits the associativity of convolution: convolving with a Voigt profile
(Gaussian ∗ Lorentzian) is equivalent to convolving sequentially with a
Gaussian and then with a Lorentzian. The Lorentzian step uses the true
2-D non-separable kernel.
"""
if len(_scan_axes(measurements)) < 2:
raise RuntimeError(
"Voigtian source size not implemented for diffraction patterns with less"
" than two scan axes."
)
g = _gaussian_source_size(measurements, gaussian_sigma)
return _lorentzian_source_size(g, lorentzian_gamma, truncate)
def _pseudo_voigtian_source_size(
measurements,
gaussian_sigma: float | tuple[float, float],
lorentzian_gamma: float | tuple[float, float],
eta: float,
truncate: float = 10.0,
):
"""Apply a pseudo-Voigtian source-size convolution via a linear combination.
Exploits the linearity of convolution: PV = (1−η)·G + η·L, so the
pseudo-Voigt-filtered image is the same weighted combination of the
Gaussian- and Lorentzian-filtered images. Each component is computed
with its proper (separable Gaussian / true-2-D Lorentzian) kernel.
"""
if len(_scan_axes(measurements)) < 2:
raise RuntimeError(
"Pseudo-Voigtian source size not implemented for diffraction patterns"
" with less than two scan axes."
)
if eta == 0.0:
return _gaussian_source_size(measurements, gaussian_sigma)
if eta == 1.0:
return _lorentzian_source_size(measurements, lorentzian_gamma, truncate)
g = _gaussian_source_size(measurements, gaussian_sigma)
l = _lorentzian_source_size(measurements, lorentzian_gamma, truncate)
kwargs = measurements._copy_kwargs(exclude=("array",))
return measurements.__class__(
(1.0 - eta) * g.array + eta * l.array, **kwargs
)
def _gaussian_source_size(measurements, sigma: float | tuple[float, float]):
if len(_scan_axes(measurements)) < 2:
raise RuntimeError(
"Gaussian source size not implemented for diffraction patterns with less"
"than two scan axes."
)
if np.isscalar(sigma):
sigma = (sigma,) * 2
xp = get_array_module(measurements.array)
gaussian_filter = get_ndimage_module(measurements._array).gaussian_filter
ensemble_axes = tuple(range(len(measurements.ensemble_shape)))
padded_sigma = ()
depth = ()
i = 0
for axis, n in zip(ensemble_axes, measurements.ensemble_shape):
if axis in _scan_axes(measurements):
scan_sampling = _scan_sampling(measurements)[i]
padded_sigma += (sigma[i] / scan_sampling,)
depth += (min(int(np.ceil(4.0 * sigma[i] / scan_sampling)), n),)
i += 1
else:
padded_sigma += (0.0,)
depth += (0,)
padded_sigma += (0.0,) * 2
depth += (0,) * 2
if measurements.is_lazy:
array = measurements.array.map_overlap(
gaussian_filter,
sigma=padded_sigma,
mode="wrap",
depth=depth,
meta=xp.array((), dtype=get_dtype(complex=False)),
)
else:
array = gaussian_filter(measurements.array, sigma=padded_sigma, mode="wrap")
kwargs = measurements._copy_kwargs(exclude=("array",))
return measurements.__class__(array, **kwargs)
def _infer_lines(b, H, W, out_H, out_W, kH, kW):
target_size = 2**17
line_size = b * (H * W // out_H + kH * kW * out_W)
target_lines = target_size // line_size
if target_lines < out_H:
lines = 1
while True:
next_lines = lines * 2
if next_lines > target_lines:
break
lines = next_lines
else:
lines = out_H
return lines
def _interpolate_bilinear(x, v, u, vw, uw):
B, H, W = x.shape
out_H, out_W = v.shape
# Interpolation is done by each output panel (i.e. multi lines)
# in order to better utilize CPU cache memory.
lines = _infer_lines(B, H, W, out_H, out_W, 2, 2)
vcol = np.empty((2, lines, out_W), dtype=v.dtype)
ucol = np.empty((2, lines, out_W), dtype=u.dtype)
wcol = np.empty((2, 2, lines, out_W), dtype=x.dtype)
y = np.empty((B, out_H * out_W), dtype=x.dtype)
for i in range(0, out_H, lines):
n = min(lines, out_H - i)
vcol = vcol[:, :n]
ucol = ucol[:, :n]
wcol = wcol[:, :, :n]
i_end = i + n
# indices
vcol[0] = v[i:i_end]
ucol[0] = u[i:i_end]
np.add(vcol[0], 1, out=vcol[1])
np.add(ucol[0], 1, out=ucol[1])
np.minimum(vcol[1], H - 1, out=vcol[1])
np.minimum(ucol[1], W - 1, out=ucol[1])
wcol[0, 1] = uw[i:i_end]
np.subtract(1, wcol[0, 1], out=wcol[0, 0])
np.multiply(wcol[0], vw[i:i_end], out=wcol[1])
wcol[0] -= wcol[1]
# packing to the panel whose shape is (B, C, 2, 2, l, out_W)
panel = x[:, vcol[:, None], ucol[None, :]]
# interpolation
panel = panel.reshape((B, 4, n * out_W))
weights = wcol.reshape((4, n * out_W))
iout = i * out_W
iout_end = i_end * out_W
np.einsum("ijk,jk->ik", panel, weights, out=y[:, iout:iout_end])
del panel, weights
return y.reshape((B, out_H, out_W))
def _diffraction_pattern_resampling_gpts(
old_sampling: tuple[float, float],
old_gpts: tuple[int, int],
sampling: str | float | tuple[float, float],
gpts: Optional[tuple[int, int]] = None,
adjust_sampling: bool = True,
):
if gpts is None:
if sampling == "uniform":
validated_sampling = (max(old_sampling),) * 2
elif isinstance(sampling, float):
validated_sampling = (sampling,) * 2
else:
raise ValueError("Invalid sampling.")
assert isinstance(validated_sampling, tuple)
adjusted_sampling, gpts = adjusted_gpts(
validated_sampling, old_sampling, old_gpts
)
if adjust_sampling:
validated_sampling = adjusted_sampling
else:
if np.isscalar(gpts):
gpts = (gpts,) * 2
validated_sampling = tuple(
d * old_n / new_n for d, old_n, new_n in zip(sampling, old_gpts, gpts)
)
return gpts, validated_sampling
[docs]
class DiffractionPatterns(_BaseMeasurement2D):
"""
One or more diffraction patterns.
Parameters
----------
array : numpy.ndarray
2D or greater array containing data with `float` type. The second-to-last and
last dimensions are the reciprocal space `y`- and `x`-axis of the diffraction
pattern.
sampling : float or two float
The reciprocal-space sampling of the diffraction patterns [1 / Å].
fftshift : bool, optional
If True, the diffraction patterns are assumed to have the zero-frequency
component to the center of the spectrum, otherwise the center(s) are assumed to
be at `(0, 0)`.
ensemble_axes_metadata : list of AxisMetadata, optional
List of metadata associated with the ensemble axes. The length and item order
must match the ensemble axes.
metadata : dict, optional
A dictionary defining measurement metadata.
"""
def __init__(
self,
array: np.ndarray | da.core.Array,
sampling: float | tuple[float, float],
fftshift: bool = False,
ensemble_axes_metadata: Optional[list[AxisMetadata]] = None,
metadata: Optional[dict] = None,
):
if np.isscalar(sampling):
sampling = (float(sampling),) * 2
else:
sampling = float(sampling[0]), float(sampling[1])
self._fftshift = fftshift
self._sampling = sampling
self._base_axes = (-2, -1)
super().__init__(
array=array,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
@property
def _area_per_pixel(self):
return _scan_area_per_pixel(self)
def _get_1d_equivalent(self):
return ReciprocalSpaceLineProfiles
@property
def base_axes_metadata(self):
limits = self.limits
return [
ReciprocalSpaceAxis(
sampling=self.sampling[0],
offset=limits[0][0],
label="kx",
units="1/Å",
fftshift=self.fftshift,
tex_label="$k_x$",
),
ReciprocalSpaceAxis(
sampling=self.sampling[1],
offset=limits[1][0],
label="ky",
units="1/Å",
fftshift=self.fftshift,
tex_label="$k_y$",
),
]
[docs]
def tile_scan(self, repetitions: tuple[int, int]) -> DiffractionPatterns:
"""
Tile the scan axes of the diffraction patterns. The diffraction patterns must
have one or more scan axes.
Parameters
----------
repetitions : two int
The number of repetitions of the scan positions along the `x`- and `y`-axis.
Returns
-------
tiled_diffraction_patterns : DiffractionPatterns
The tiled diffraction patterns.
"""
scan_axes = _scan_axes(self)
if len(scan_axes) != 2:
raise NotImplementedError
xp = get_array_module(self.array)
tiling = ()
j = 0
for i in range(len(self.shape)):
if i in scan_axes:
tiling += (repetitions[j],)
j += 1
else:
tiling += (1,)
array = xp.tile(self.array, tiling)
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = array
return self.__class__(**kwargs)
@staticmethod
def _index_diffraction_spots(
array,
orientation_matrices,
hkl,
mask_all,
sampling,
cell,
sg_max,
g_max,
centering,
energy,
radius,
):
from abtem.bloch.indexing import index_diffraction_spots
from abtem.bloch.utils import filter_reciprocal_space_vectors
mask = filter_reciprocal_space_vectors(
hkl,
cell,
energy=energy,
sg_max=sg_max,
g_max=g_max,
centering=centering,
orientation_matrices=orientation_matrices,
)
array = index_diffraction_spots(
array=array,
hkl=hkl[mask],
sampling=sampling,
cell=cell,
energy=energy,
radius=radius,
orientation_matrices=orientation_matrices,
)
array_all = np.zeros(array.shape[:-1] + (mask_all.sum(),), dtype=array.dtype)
array_all[..., mask[mask_all]] = array
return array_all
[docs]
def index_diffraction_spots(
self,
cell: Cell | float | tuple[float, float, float],
sg_max: Optional[float] = None,
g_max: Optional[float] = None,
orientation_matrices: Optional[np.ndarray] = None,
radius: Optional[float] = None,
centering: str = "P",
energy: Optional[float] = None,
) -> IndexedDiffractionPatterns:
"""
Indexes the Bragg reflections (diffraction spots) by their Miller indices.
Parameters
----------
cell : ase.cell.Cell or float or tuple of float
The assumed unit cell with respect to the diffraction pattern should be
indexed. Must be one of ASE `Cell` object, float (for a cubic unit cell) or
three floats (for orthorhombic unit cells).
orientation_matrices : numpy.ndarray, optional
Orientation matrices used for indexing the diffraction spots. The shape of
the orientation matrices must be broadcastable with the ensemble shape of
the diffraction patterns.
sg_max : float, optional
Maximum excitation error [1/Å] of the indexed diffraction spots.
The default is estimated from the energy and `g_max`.
g_max : float, optional
Maximum scattering vector [1/Å] of the indexed diffraction spots.
The default is the maximum frequency of the diffraction patterns.
radius : float, optional
Integration Radius of the diffraction spots [1/Å]. The default is the
reciprocal-space sampling of the diffraction patterns.
centering : {'P', 'F', 'I', 'A', 'B', 'C'}
Assumed lattice centering used for determining the reflection conditions.
energy : float, optional
The energy of the electrons [keV]. The default is the energy stored in the
metadata.
Return
-------
indexed_patterns : IndexedDiffractionPatterns
The indexed diffraction pattern(s).
"""
from abtem.bloch.indexing import (
estimate_necessary_excitation_error,
validate_cell,
)
from abtem.bloch.utils import filter_reciprocal_space_vectors, make_hkl_grid
if orientation_matrices is not None and not is_broadcastable(
self.ensemble_shape, orientation_matrices.shape[:-2]
):
raise ValueError(
"The ensemble shape and the shape of the orientation matrices must be"
" broadcastable."
)
if energy is None:
energy = self._get_energy()
if g_max is None:
g_max = max(self.max_frequency)
if sg_max is None:
sg_max = estimate_necessary_excitation_error(energy, g_max)
if orientation_matrices is None:
orientation_matrices = np.eye(3)[(None,) * len(self.array.shape[:-2])]
cell = validate_cell(cell)
hkl = make_hkl_grid(cell, g_max)
mask = filter_reciprocal_space_vectors(
hkl,
cell,
energy=energy,
sg_max=sg_max,
g_max=g_max,
centering=centering,
orientation_matrices=orientation_matrices,
)
if self.is_lazy:
orientation_matrices = orientation_matrices[
(None,)
* (len(self.array.shape[:-2]) - len(orientation_matrices.shape[:-2]))
]
chunks = tuple(
c if n == sum(c) else 1
for n, c in zip(orientation_matrices.shape, self.array.chunks[:-2])
)
lazy_orientation_matrices = da.from_array(
orientation_matrices, chunks=chunks + (3, 3)
)
intensities = da.map_blocks(
self._index_diffraction_spots,
self.array,
lazy_orientation_matrices,
hkl=hkl,
mask_all=mask,
sampling=self.sampling,
cell=cell,
sg_max=sg_max,
g_max=g_max,
centering=centering,
energy=energy,
radius=radius,
drop_axis=len(self.array.shape) - 1,
chunks=self.array.chunks[:-2] + (mask.sum(),),
meta=np.array((), dtype=self.dtype),
)
else:
intensities = self._index_diffraction_spots(
array=self.array,
orientation_matrices=orientation_matrices,
hkl=hkl,
mask_all=mask,
sampling=self.sampling,
cell=cell,
sg_max=sg_max,
g_max=g_max,
centering=centering,
energy=energy,
radius=radius,
)
reciprocal_lattice_vectors = np.matmul(
cell.reciprocal(),
np.swapaxes(orientation_matrices, -2, -1),
)
return IndexedDiffractionPatterns(
intensities,
hkl[mask],
reciprocal_lattice_vectors=reciprocal_lattice_vectors,
ensemble_axes_metadata=self.ensemble_axes_metadata,
metadata=self.metadata,
)
@property
def fftshift(self) -> bool:
"""
True if the zero-frequency is shifted to the center of the array.
"""
return self._fftshift
@property
def sampling(self) -> tuple[float, float]:
"""Sampling of diffraction patterns in `x` and `y` [1 / Å]."""
return self._sampling
@property
def angular_sampling(self) -> tuple[float, float]:
"""
Angular sampling of diffraction patterns in `x` and `y` [mrad].
"""
wavelength = energy2wavelength(self._get_energy())
return (
self.sampling[0] * wavelength * 1e3,
self.sampling[1] * wavelength * 1e3,
)
@property
def max_angles(self) -> tuple[float, float]:
"""Maximum scattering angle in `x` and `y` [mrad]."""
return (
self.shape[-2] // 2 * self.angular_sampling[0],
self.shape[-1] // 2 * self.angular_sampling[1],
)
@property
def max_frequency(self):
"""Maximum spatial frequency in `x` and `y` [1 / Å]."""
return abs(self.limits[0][1]), abs(self.limits[1][1])
@property
def limits(self) -> list[tuple[float, float]]:
"""Lowest and highest spatial frequency in `x` and `y` [1 / Å]."""
limits = []
for i in (-2, -1):
if self.shape[i] % 2:
limits += [
(
-(self.shape[i] - 1) // 2 * self.sampling[i],
(self.shape[i] - 1) // 2 * self.sampling[i],
)
]
else:
limits += [
(
-self.shape[i] // 2 * self.sampling[i],
(self.shape[i] // 2 - 1) * self.sampling[i],
)
]
return limits
@property
def angular_limits(self) -> list[tuple[float, float]]:
"""Lowest and highest scattering angle in `x` and `y` [mrad]."""
limits = self.limits
wavelength = energy2wavelength(self._get_energy())
limits[0] = (
limits[0][0] * wavelength * 1e3,
limits[0][1] * wavelength * 1e3,
)
limits[1] = (
limits[1][0] * wavelength * 1e3,
limits[1][1] * wavelength * 1e3,
)
return limits
@property
def offset(self) -> tuple[float, float]:
limits = self.limits
return limits[0][0], limits[1][0]
@property
def extent(self) -> tuple[float, float]:
limits = self.limits
return limits[0][0] - limits[0][1], limits[1][0] - limits[1][1]
@property
def coordinates(self) -> tuple[np.ndarray, np.ndarray]:
"""Reciprocal-space frequency coordinates [1 / Å]."""
return (
self.axes_metadata[-2].coordinates(self.base_shape[-2]),
self.axes_metadata[-1].coordinates(self.base_shape[-1]),
)
@property
def angular_coordinates(self) -> tuple[np.ndarray, np.ndarray]:
"""Scattering angle coordinates [mrad]."""
xp = get_array_module(self.array)
limits = self.angular_limits
alpha_x = xp.linspace(
limits[0][0], limits[0][1], self.shape[-2], dtype=xp.float32
)
alpha_y = xp.linspace(
limits[1][0], limits[1][1], self.shape[-1], dtype=xp.float32
)
if self.fftshift:
return alpha_x, alpha_y
else:
return np.fft.fftshift(alpha_x), np.fft.fftshift(alpha_y)
@staticmethod
def _batch_interpolate_bilinear(array, new_sampling, sampling, new_gpts):
xp = get_array_module(array)
v, u, vw, uw = _fourier_space_bilinear_nodes_and_weight(
array.shape[-2:], new_gpts, sampling, new_sampling, xp
)
old_shape = array.shape
array = array.reshape((-1,) + array.shape[-2:])
old_sums = array.sum((-2, -1), keepdims=True)
if xp is cp:
array = interpolate_bilinear_cuda(array, v, u, vw, uw)
else:
array = _interpolate_bilinear(array, v, u, vw, uw)
array = array / array.sum((-2, -1), keepdims=True) * old_sums
return array.reshape(old_shape[:-2] + array.shape[-2:])
[docs]
def interpolate(
self,
sampling: Optional[str | float | tuple[float, float]] = None,
gpts: Optional[tuple[int, int]] = None,
):
"""
Interpolate diffraction pattern(s) producing equivalent pattern(s) with a
different sampling.
Parameters
----------
sampling : 'uniform' or float or two floats
Sampling of diffraction patterns after interpolation in `x` and `y` [1 / Å].
If a single value, the same sampling is used for both axes. If 'uniform',
the diffraction patterns are down-sampled along the axis with the smallest
pixel size such that the sampling is uniform.
gpts : tuple of int
Number of grid points of the diffraction patterns after interpolation in
`x` and `y`.
Do not use if 'sampling' is used.
Returns
-------
interpolated_diffraction_patterns : DiffractionPatterns
The interpolated diffraction pattern(s).
"""
gpts, sampling = _diffraction_pattern_resampling_gpts(
self.sampling, self.base_shape, sampling, gpts, adjust_sampling=False
)
if self.is_lazy:
array = self.array.map_blocks(
self._batch_interpolate_bilinear,
sampling=self.sampling,
new_sampling=sampling,
new_gpts=gpts,
chunks=self.array.chunks[:-2] + ((gpts[0],), (gpts[1],)),
dtype=get_dtype(complex=False),
)
else:
array = self._batch_interpolate_bilinear(
self.array, sampling=self.sampling, new_sampling=sampling, new_gpts=gpts
)
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["sampling"] = sampling
kwargs["array"] = array
return self.__class__(**kwargs)
def _check_integration_limits(self, inner: float, outer: float):
if inner > outer:
raise RuntimeError(
f"Inner detection ({inner} mrad) angle cannot exceed the outer"
f" detection angle ({outer} mrad)."
)
if (outer > self.max_angles[0]) or (outer > self.max_angles[1]):
if not np.isclose(min(self.max_angles), outer, atol=1e-5):
raise RuntimeError(
f"Outer integration limit cannot exceed the maximum simulated angle"
f" ({outer} mrad > {min(self.max_angles)} mrad), increase the"
" number of grid points."
)
[docs]
def gaussian_source_size(
self, sigma: float | tuple[float, float]
) -> DiffractionPatterns:
"""
Simulate the effect of a finite source size on diffraction pattern(s) using a
Gaussian filter.
The filter is not applied to diffraction pattern individually, but the intensity
of diffraction patterns are mixed across scan axes. Applying this filter
requires two linear scan axes.
Applying this filter before integrating the diffraction patterns will produce
the same image as integrating the diffraction patterns first then applying a
Gaussian filter.
Parameters
----------
sigma : float or two float
Standard deviation of Gaussian kernel in the `x` and `y`-direction. If given
as a single number, the standard deviation is equal for both axes.
Returns
-------
filtered_diffraction_patterns : DiffractionPatterns
The filtered diffraction pattern(s).
"""
return _gaussian_source_size(self, sigma)
[docs]
def lorentzian_source_size(
self,
half_width: float | tuple[float, float],
truncate: float = 10.0,
) -> DiffractionPatterns:
"""
Simulate the effect of a finite source size on diffraction pattern(s) using a
Lorentzian (Cauchy) filter.
The filter mixes intensities across scan axes (not within each diffraction
pattern) and requires two linear scan axes. Applying this filter before
integrating gives the same result as integrating first and then applying a
:meth:`lorentzian_filter` to the resulting images.
Parameters
----------
half_width : float or two float
Half-width at half-maximum (HWHM) of the Lorentzian kernel in the ``x``
and ``y``-direction [Å]. If given as a single number it is equal for
both axes.
truncate : float, optional
Truncate the kernel at this many half-widths (default is 10.0).
Returns
-------
filtered_diffraction_patterns : DiffractionPatterns
The filtered diffraction pattern(s).
Notes
-----
The Lorentzian kernel is parameterized by its HWHM γ (``half_width``),
with FWHM_L = 2γ. This differs from :meth:`gaussian_source_size`,
which uses the standard deviation σ (FWHM_G ≈ 2.3548·σ).
"""
return _lorentzian_source_size(self, half_width, truncate)
[docs]
def voigtian_source_size(
self,
gaussian_sigma: float | tuple[float, float],
lorentzian_gamma: float | tuple[float, float],
truncate: float = 10.0,
) -> DiffractionPatterns:
"""
Simulate the effect of a finite source size on diffraction pattern(s) using a
Voigt filter (convolution of Gaussian and Lorentzian).
The filter mixes intensities across scan axes (not within each diffraction
pattern) and requires two linear scan axes. Applying this filter before
integrating gives the same result as integrating first and then applying a
:meth:`voigtian_filter` to the resulting images.
Parameters
----------
gaussian_sigma : float or two float
Standard deviation (σ) of the Gaussian component [Å]. If given as a
single number it is equal for both axes.
lorentzian_gamma : float or two float
Half-width at half-maximum (HWHM, γ) of the Lorentzian component [Å].
If given as a single number it is equal for both axes.
truncate : float, optional
Truncate the Lorentzian component at this many half-widths
(default 10.0). The Gaussian component uses the SciPy default (4·σ).
Returns
-------
filtered_diffraction_patterns : DiffractionPatterns
The filtered diffraction pattern(s).
Notes
-----
The Voigt profile combines two width parameters:
* ``gaussian_sigma`` (σ): standard deviation of the Gaussian component,
FWHM_G = 2√(2 ln 2)·σ ≈ 2.3548·σ.
* ``lorentzian_gamma`` (γ): HWHM of the Lorentzian component,
FWHM_L = 2γ.
For the same FWHM, γ = FWHM / 2 but σ = FWHM / (2√(2 ln 2)) ≈ FWHM / 2.3548.
"""
return _voigtian_source_size(self, gaussian_sigma, lorentzian_gamma, truncate)
[docs]
def pseudo_voigtian_source_size(
self,
gaussian_sigma: float | tuple[float, float],
lorentzian_gamma: float | tuple[float, float],
eta: float,
truncate: float = 10.0,
) -> DiffractionPatterns:
"""
Simulate the effect of a finite source size on diffraction pattern(s) using a
pseudo-Voigtian filter (weighted sum of Gaussian and Lorentzian).
The filter mixes intensities across scan axes (not within each diffraction
pattern) and requires two linear scan axes. Applying this filter before
integrating gives the same result as integrating first and then applying a
:meth:`pseudo_voigtian_filter` to the resulting images.
Parameters
----------
gaussian_sigma : float or two float
Standard deviation (σ) of the Gaussian component [Å]. If given as a
single number it is equal for both axes.
lorentzian_gamma : float or two float
Half-width at half-maximum (HWHM, γ) of the Lorentzian component [Å].
If given as a single number it is equal for both axes.
eta : float
Lorentzian mixing fraction η ∈ [0, 1]. η = 0 gives a pure Gaussian;
η = 1 gives a pure Lorentzian.
truncate : float, optional
Truncate the kernel at this many effective half-widths (default 10.0).
Returns
-------
filtered_diffraction_patterns : DiffractionPatterns
The filtered diffraction pattern(s).
Notes
-----
The pseudo-Voigt profile is PV = (1-η)·G + η·L, where G and L are
Gaussian and Lorentzian profiles with widths σ and γ respectively. Width
parameterization: FWHM_G = 2√(2 ln 2)·σ ≈ 2.3548·σ; FWHM_L = 2γ.
"""
return _pseudo_voigtian_source_size(
self, gaussian_sigma, lorentzian_gamma, eta, truncate
)
[docs]
def poisson_noise(
self,
dose_per_area: Optional[float] = None,
total_dose: Optional[float] = None,
samples: int = 1,
seed: Optional[int] = None,
):
"""
Add Poisson noise (i.e. shot noise) to a measurement corresponding to the
provided `total_dose` (per measurement if applied to an ensemble) or
`dose_per_area` (not applicable for single measurements).
Parameters
----------
dose_per_area : float, optional
The irradiation dose per unit of scan area [electrons per Å:sup:`2`].
This is only valid if the diffraction patterns has two scan axes.
total_dose : float, optional
The irradiation dose per diffraction pattern.
samples : int, optional
The number of samples to draw from a Poisson distribution. If this is
greater than 1, an additional ensemble axis will be added to the
measurement.
seed : int, optional
Seed the random number generator.
Returns
-------
noisy_measurement : BaseMeasurements
The noisy measurement.
"""
if len(_scan_shape(self)) < 2 and dose_per_area is not None:
raise ValueError(
"diffraction patterns has less than two scan axes, provide `total_dose`"
" not `dose_per_area`"
)
# TODO: normalization
return super().poisson_noise(
dose_per_area=dose_per_area,
total_dose=total_dose,
samples=samples,
seed=seed,
)
@staticmethod
def _radial_binning(
array,
nbins_radial,
nbins_azimuthal,
sampling,
inner,
outer,
fftshift,
rotation,
offset,
):
xp = get_array_module(array)
indices = _polar_detector_bins(
gpts=array.shape[-2:],
sampling=sampling,
inner=inner,
outer=outer,
nbins_radial=nbins_radial,
nbins_azimuthal=nbins_azimuthal,
fftshift=fftshift,
rotation=rotation,
offset=offset,
return_indices=True,
)
separators = xp.concatenate(
(xp.array([0]), xp.cumsum(xp.array([len(i) for i in indices])))
)
new_shape = array.shape[:-2] + (nbins_radial, nbins_azimuthal)
array = array.reshape(
(
-1,
array.shape[-2] * array.shape[-1],
)
)[..., np.concatenate(indices)]
# Use the configured floating-point precision, not a hardcoded float32.
# _AbstractRadialDetector._out_dtype returns get_dtype(complex=False), so
# the result dtype must match to avoid a silent precision downgrade.
fp_dtype = array.dtype
result = xp.zeros(
(
array.shape[0],
len(indices),
),
dtype=fp_dtype,
)
if xp is cp:
sum_run_length_encoded_cuda(array, result, separators)
else:
_sum_run_length_encoded(array, result, separators)
return result.reshape(new_shape)
[docs]
def polar_binning(
self,
nbins_radial: int,
nbins_azimuthal: int,
inner: float = 0.0,
outer: Optional[float] = None,
rotation: float = 0.0,
offset: tuple[float, float] = (0.0, 0.0),
):
"""
Create polar measurements from the diffraction patterns by binning the
measurements on a polar grid. This method may be used to simulate a segmented
detector with a specified number of radial and azimuthal bins.
Each bin is a segment of an annulus and the bins are spaced equally in the
radial and azimuthal directions.
The bins fit between a given inner and outer integration limit, they may be
rotated around the origin, and their center may be shifted from the origin.
Parameters
----------
nbins_radial : int
Number of radial bins.
nbins_azimuthal : int
Number of angular bins.
inner : float
Inner integration limit of the bins [mrad] (default is 0.0).
outer : float
Outer integration limit of the bins [mrad]. If not specified, this is set to
be the maximum detected angle of the diffraction pattern.
rotation : float
Rotation of the bins around the origin [mrad] (default is 0.0).
offset : two float
Offset of the bins from the origin in `x` and `y` [mrad].
Default is (0.0, 0.0).
Returns
-------
polar_measurements : PolarMeasurements
The polar measurements.
"""
if nbins_radial <= 0 or nbins_azimuthal <= 0:
raise RuntimeError("number of bins must be greater than zero")
if outer is None:
outer = min(self.max_angles)
self._check_integration_limits(inner, outer)
xp = get_array_module(self.array)
if self.is_lazy:
array = self.array.map_blocks(
self._radial_binning,
nbins_radial=nbins_radial,
nbins_azimuthal=nbins_azimuthal,
sampling=self.angular_sampling,
inner=inner,
outer=outer,
fftshift=self.fftshift,
rotation=rotation,
offset=offset,
drop_axis=(len(self.shape) - 2, len(self.shape) - 1),
chunks=self.array.chunks[:-2]
+ (
(nbins_radial,),
(nbins_azimuthal,),
),
new_axis=(
len(self.shape) - 2,
len(self.shape) - 1,
),
meta=xp.array((), dtype=get_dtype(complex=False)),
)
else:
array = self._radial_binning(
self.array,
nbins_radial=nbins_radial,
nbins_azimuthal=nbins_azimuthal,
sampling=self.angular_sampling,
inner=inner,
outer=outer,
fftshift=self.fftshift,
rotation=rotation,
offset=offset,
)
radial_sampling = (outer - inner) / nbins_radial
azimuthal_sampling = 2 * np.pi / nbins_azimuthal
return PolarMeasurements(
array,
radial_sampling=radial_sampling,
azimuthal_sampling=azimuthal_sampling,
radial_offset=inner,
azimuthal_offset=rotation,
ensemble_axes_metadata=self.ensemble_axes_metadata,
metadata=self.metadata,
)
[docs]
def radial_binning(
self, step_size: float = 1.0, inner: float = 0.0, outer: Optional[float] = None
) -> PolarMeasurements:
"""
Create polar measurement(s) from the diffraction pattern(s) by binning the
measurements in annular regions. This method may be used to simulate a segmented
detector with a specified number of radial bins.
This is equivalent to detecting a wave function using the
`FlexibleAnnularDetector`.
Parameters
----------
step_size : float, optional
Radial extent of the bins [mrad]. Default is 1.0.
inner : float, optional
Inner integration limit of the bins [mrad]. Default is 0.0.
outer : float, optional
Outer integration limit of the bins [mrad]. If not specified, this is set to
be the maximum detected angle of the diffraction pattern.
Returns
-------
radially_binned_measurement : PolarMeasurements
Radially binned polar measurement(s).
"""
if outer is None:
outer = min(self.max_angles)
nbins_radial = int((outer - inner) / step_size)
return self.polar_binning(nbins_radial, 1, inner, outer)
@staticmethod
def _integrate_fourier_space(array, sampling, inner, outer, fftshift, offset):
xp = get_array_module(array)
bins = _annular_detector_mask(
gpts=array.shape[-2:],
sampling=sampling,
inner=inner,
outer=outer,
fftshift=fftshift,
offset=offset,
xp=xp,
)
return xp.sum(array * bins, axis=(-2, -1))
[docs]
def integrate_radial(
self,
inner: float,
outer: float = None,
offset: tuple[float, float] = (0.0, 0.0),
) -> Images:
"""
Create images by integrating the diffraction patterns over an annulus defined by
an inner and outer integration angle.
Parameters
----------
inner : float
Inner integration limit [mrad].
outer : float
Outer integration limit [mrad].
offset : tuple of float
Offset of center of annular integration region [mrad].
Returns
-------
integrated_images : Images
The integrated images.
"""
if isinstance(inner, Sequence) or isinstance(outer, Sequence):
if isinstance(inner, Number):
inners = (inner,) * len(outer)
outers = outer
else:
outers = (outer,) * len(inner)
inners = inner
measurements = [
self.integrate_radial(inner=inner, outer=outer)
for inner, outer in zip(inners, outers)
]
measurements = stack(
measurements,
axis_metadata=NonLinearAxis(
label="Limits", values=tuple(zip(inners, outers)), units="mrad"
),
)
return measurements
if outer is None:
outer = min(self.max_angles)
self._check_integration_limits(inner, outer)
xp = get_array_module(self.array)
if self.is_lazy:
integrated_intensity = self.array.map_blocks(
self._integrate_fourier_space,
sampling=self.angular_sampling,
inner=inner,
outer=outer,
fftshift=self.fftshift,
offset=offset,
drop_axis=(len(self.shape) - 2, len(self.shape) - 1),
meta=xp.array((), dtype=get_dtype(complex=False)),
)
else:
integrated_intensity = self._integrate_fourier_space(
self.array,
sampling=self.angular_sampling,
inner=inner,
outer=outer,
fftshift=self.fftshift,
offset=offset,
)
return _reduced_scanned_images_or_line_profiles(integrated_intensity, self)
[docs]
def integrated_center_of_mass(self) -> Images:
"""
Calculate integrated center-of-mass (iCOM) images from diffraction patterns.
This method is only implemented for diffraction patterns with exactly two scan
axes.
Returns
-------
icom_images : Images
The iCOM images.
"""
com = self.center_of_mass()
if isinstance(com, Images):
return com.integrate_gradient()
else:
raise RuntimeError(
f"Integrated center-of-mass not implemented for DiffractionPatterns"
f" with {len(_scan_shape(self))} scan axes."
)
@staticmethod
def _com(array: np.ndarray, x: np.ndarray, y: np.ndarray):
com_x = (array * x[:, None]).sum(axis=(-2, -1))
com_y = (array * y[None]).sum(axis=(-2, -1))
com = com_x + 1.0j * com_y
return com
[docs]
def center_of_mass(self, units: str = "1/Å") -> Images | RealSpaceLineProfiles:
"""
Calculate center-of-mass images or line profiles from diffraction patterns.
The results are of type `complex` where the real and imaginary part represents
the `x` and `y` component.
Parameters
----------
units : {'1/Å', 'mrad'}
Units of the center-of-mass values of the output images or line profiles.
Default is '1/Å'.
Returns
-------
com_images : Images
Center-of-mass images.
com_line_profiles : RealSpaceLineProfiles
Center-of-mass line profiles (returned if there is only one scan axis).
"""
if units == "mrad":
x, y = self.angular_coordinates
elif units == "1/Å":
x, y = self.coordinates
else:
raise ValueError("units must be '1/Å' or 'mrad'")
xp = get_array_module(self.array)
x, y = xp.asarray(x), xp.asarray(y)
if self.is_lazy:
base_axes = tuple(
range(
len(self.ensemble_shape),
len(self.base_shape) + len(self.ensemble_shape),
)
)
array = self.array.map_blocks(
self._com, x=x, y=y, drop_axis=base_axes, dtype=get_dtype(complex=True)
)
else:
array = self._com(self.array, x=x, y=y)
return _reduced_scanned_images_or_line_profiles(array, self)
@staticmethod
def _bandlimit(
array: np.ndarray,
inner: float,
outer: float,
angular_coordinates: tuple[np.ndarray, np.ndarray],
):
alpha_x, alpha_y = angular_coordinates
alpha = np.sqrt(alpha_x[:, None] ** 2 + alpha_y[None] ** 2)
block = alpha > inner
if outer != np.inf:
block *= alpha < outer
return array * block
[docs]
def bandlimit(
self, inner: float = 0.0, outer: float = np.inf
) -> DiffractionPatterns:
"""
Bandlimit diffraction pattern(s) by setting everything outside an annulus
defined by two radial angles to zero.
Parameters
----------
inner : float
Inner limit of zero region [mrad].
outer : float
Outer limit of zero region [mrad].
Returns
-------
band-limited_diffraction_patterns : DiffractionPatterns
The band-limited diffraction pattern(s).
"""
xp = get_array_module(self.array)
if self.is_lazy:
array = self.array.map_blocks(
self._bandlimit,
inner=inner,
outer=outer,
angular_coordinates=self.angular_coordinates,
meta=xp.array((), dtype=get_dtype(complex=False)),
)
else:
array = self._bandlimit(self.array, inner, outer, self.angular_coordinates)
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = array
return self.__class__(**kwargs)
@staticmethod
def _crop(array: np.ndarray, gpts: tuple[int, int]):
xp = get_array_module(array)
array = xp.fft.fftshift(
fft_crop(xp.fft.ifftshift(array, axes=(-2, -1)), new_shape=gpts),
axes=(-2, -1),
)
return array
[docs]
def crop(
self,
max_angle: Optional[float] = None,
max_frequency: Optional[float] = None,
gpts: Optional[tuple[int, int]] = None,
) -> DiffractionPatterns:
"""
Crop the diffraction patterns such that they only include spatial frequencies
(scattering angles) up to a given limit.
Parameters
----------
max_angle : float, optional
The maximum included scattering angle in the cropped diffraction patterns.
max_frequency : float, optional
The maximum included spatial frequency in the cropped diffraction patterns.
gpts : tuple of int
The number of gpts in the cropped diffraction patterns.
Returns
-------
cropped_diffraction_patterns : DiffractionPatterns
The cropped diffraction pattern(s).
"""
none_args = (max_angle is None) + (max_frequency is None) + (gpts is None)
if none_args != 2:
raise ValueError(
"provide exactly one of 'max_angle', 'max_frequency' or 'gpts'"
)
if gpts is None and max_angle is not None:
gpts = (
int(2 * np.round(max_angle / self.angular_sampling[0])) + 1,
int(2 * np.round(max_angle / self.angular_sampling[1])) + 1,
)
elif gpts is None and max_frequency is not None:
gpts = (
int(2 * np.round(max_frequency / self.sampling[0])) + 1,
int(2 * np.round(max_frequency / self.sampling[1])) + 1,
)
if gpts is None:
raise ValueError()
if self.is_lazy:
xp = get_array_module(self.array)
array = self.array.map_blocks(
self._crop,
gpts=gpts,
chunks=self.array.chunks[:-2] + gpts,
meta=xp.array((), dtype=self.dtype),
)
else:
array = self._crop(self.array, gpts=gpts)
kwargs = self._copy_kwargs(exclude=("array",))
kwargs["array"] = array
return self.__class__(**kwargs)
@staticmethod
def _azimuthal_average(
array: np.ndarray,
angular_coordinates: tuple[np.ndarray, np.ndarray],
max_angle: float,
radial_sampling: float,
weighting_function: str,
width: float,
):
x, y = np.meshgrid(*angular_coordinates, indexing="ij")
r = np.sqrt(x**2 + y**2)
centers = np.arange(0, max_angle, radial_sampling)
values = np.zeros(array.shape[:-2] + centers.shape)
for i, center in enumerate(centers):
if weighting_function == "step":
mask = np.abs(r - center) < width
elif weighting_function == "gaussian":
mask = np.exp(-((r - center) ** 2) / (width**2 / 2))
else:
raise ValueError("weighting function must be 'step' or 'gaussian'")
weight = np.sum(mask)
if weight > 0:
values[..., i] = np.sum(array * mask, axis=(-2, -1)) / weight
else:
values[..., i] = 0.0
return values
[docs]
def azimuthal_average(
self,
max_angle: Optional[float] = None,
radial_sampling: float = 1.0,
weighting_function: str = "step",
width: float = 1.0,
) -> ReciprocalSpaceLineProfiles:
"""
Calculate the azimuthal averages of the diffraction patterns.
Parameters
----------
max_angle : float, optional
The maximum included scattering angle in the azimuthal averages [mrad].
radial_sampling : float, optional
The radial sampling of the azimuthal averages [mrad]. Default is equal to
the smallest value of the x and y component of the angular sampling.
weighting_function : str
The weighting function to determining how to average the diffraction
patterns. The weighting function determines the shape of the mask that is
applied to the diffraction patterns before averaging. The options are 'step'
and 'gaussian'. Default is 'step'.
width : float, optional
The width of the weighting function [mrad]. Default is 1.0.
For the 'step' weighting function, this is the width of the step function.
For the 'gaussian' weighting function, this is the standard deviation of the
Gaussian function.
Returns
-------
azimuthal_averages : ReciprocalSpaceLineProfiles
The azimuthal averages of the diffraction patterns.
"""
if max_angle is None:
max_angle = -min(min(self.angular_limits))
radial_sampling = radial_sampling * min(self.angular_sampling)
width = width * min(self.angular_sampling)
if self.is_lazy:
xp = get_array_module(self.array)
n = int(max_angle / radial_sampling)
base_axes = tuple(
range(
len(self.ensemble_shape),
len(self.ensemble_shape) + len(self.base_shape),
)
)
array = self.array.map_blocks(
self._azimuthal_average,
angular_coordinates=self.angular_coordinates,
max_angle=max_angle,
radial_sampling=radial_sampling,
weighting_function=weighting_function,
width=width,
drop_axis=base_axes,
new_axis=base_axes[0],
chunks=self.array.chunks[:-2] + (n,),
meta=xp.array((), dtype=get_dtype(complex=False)),
)
else:
array = self._azimuthal_average(
self.array,
angular_coordinates=self.angular_coordinates,
max_angle=max_angle,
radial_sampling=radial_sampling,
weighting_function=weighting_function,
width=width,
)
wavelength = energy2wavelength(self._get_energy())
return ReciprocalSpaceLineProfiles(
array,
sampling=radial_sampling / (wavelength * 1e3),
ensemble_axes_metadata=self.ensemble_axes_metadata,
metadata=self.metadata,
)
# def fourier_shell_correlation(
# self,
# other: DiffractionPatterns,
# radial_sampling: float = 1.0,
# width: float = 1.0,
# weighting_function: str = "step",
# ):
# fsc = (self**0.5 * other**0.5).azimuthal_average(
# radial_sampling=radial_sampling,
# width=width,
# weighting_function=weighting_function,
# ) / (
# self.azimuthal_average(
# radial_sampling=radial_sampling,
# width=width,
# weighting_function=weighting_function,
# )
# * other.azimuthal_average(
# radial_sampling=radial_sampling,
# width=width,
# weighting_function=weighting_function,
# )
# )
# return fsc
[docs]
def block_direct(
self, radius: Optional[float] = None, margin: Optional[bool] = None
) -> DiffractionPatterns:
"""
Block the direct beam by setting the pixels of the zeroth-order Bragg reflection
(non-scattered beam) to zero.
Parameters
----------
radius : float, optional
The radius of the zeroth-order reflection to block [mrad]. If not given this
will be inferred from the metadata, if available.
margin : bool, optional
If True adds a margin to the blocking radius to fully block soft apertures.
Margin is true by default for diffraction patterns with `semiangle_cutoff`
in metadata.
Returns
-------
diffraction_patterns : DiffractionPatterns
The diffraction pattern(s) with the direct beam removed.
"""
if radius is None:
if "semiangle_cutoff" in self.metadata.keys():
radius = self.metadata["semiangle_cutoff"]
else:
radius = max(self.angular_sampling) * 1.0001
if "semiangle_cutoff" in self.metadata.keys() and margin is None:
margin = True
if margin:
radius += max(self.angular_sampling)
return self.bandlimit(radius, outer=np.inf)
[docs]
class PolarMeasurements(BaseMeasurements):
"""
Class describing polar measurements with a specified number of radial and azimuthal
bins.
Each bin is a segment of an annulus and the bins are spaced equally in the radial
and azimuthal directions. The bins may be rotated around the origin, and their
center may be shifted from the origin.
Parameters
----------
array : numpy.ndarray
Array containing the measurement.
radial_sampling : float
Sampling of the radial bins [mrad].
azimuthal_sampling : int
Sampling of the azimuthal bins [rad].
radial_offset : float, optional
Offset of the bins from the origin [mrad] (default is 0.0).
azimuthal_offset : float, optional
Rotation of the bins around the origin [rad] (default is 0.0).
ensemble_axes_metadata : list of AxisMetadata, optional
List of metadata associated with the ensemble axes. The length and item order
must match the ensemble axes.
metadata : dict, optional
A dictionary defining measurement metadata.
Returns
-------
polar_measurements : PolarMeasurements
The polar measurements.
"""
_base_dims = 2
def __init__(
self,
array: np.ndarray | da.core.Array,
radial_sampling: float,
azimuthal_sampling: float,
radial_offset: float = 0.0,
azimuthal_offset: float = 0.0,
ensemble_axes_metadata: Optional[list[AxisMetadata]] = None,
metadata: Optional[dict] = None,
):
self._radial_sampling = radial_sampling
self._azimuthal_sampling = azimuthal_sampling
self._radial_offset = radial_offset
self._azimuthal_offset = azimuthal_offset
super().__init__(
array=array,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
@property
def _area_per_pixel(self):
return _scan_area_per_pixel(self)
@property
def sampling(self):
raise RuntimeError("Sampling not defined for polar measurement.")
@property
def offset(self):
raise RuntimeError("Offset not defined for polar measurement.")
@property
def extent(self):
raise RuntimeError("Extent not defined for polar measurement.")
def _get_1d_equivalent(self):
raise RuntimeError("Not defined for polar measurement.")
@property
def base_axes_metadata(self) -> list[AxisMetadata]:
return [
LinearAxis(
label="Radial scattering angle",
offset=self.radial_offset,
sampling=self.radial_sampling,
_concatenate=False,
units="mrad",
),
LinearAxis(
label="Azimuthal scattering angle",
offset=self.azimuthal_offset,
sampling=self.azimuthal_sampling,
_concatenate=False,
units="rad",
),
]
@property
def radial_offset(self) -> float:
"""Offset of the bins from the origin [mrad]."""
return self._radial_offset
@property
def outer_angle(self) -> float:
"""The outer angle of the outermost radial bin [mrad]."""
return self._radial_offset + self.radial_sampling * self.shape[-2]
@property
def radial_sampling(self) -> float:
"""Sampling of the radial bins [mrad]."""
return self._radial_sampling
@property
def azimuthal_sampling(self) -> float:
"""Sampling of the azimuthal bins [rad]."""
return self._azimuthal_sampling
@property
def azimuthal_offset(self) -> float:
"""Rotation of the bins around the origin [rad]."""
return self._azimuthal_offset
[docs]
def integrate_radial(
self, inner: float, outer: float
) -> Images | RealSpaceLineProfiles:
"""
Create images by integrating the polar measurements over an annulus defined by
an inner and outer integration angle.
Parameters
----------
inner : float
Inner integration limit [mrad].
outer : float
Outer integration limit [mrad].
Returns
-------
integrated_images : Images
The integrated images.
real_space_line_profiles : RealSpaceLineProfiles
Integrated line profiles (returned if there is only one scan axis).
"""
return self.integrate(radial_limits=(inner, outer))
[docs]
def integrate(
self,
radial_limits: Optional[tuple[float, float]] = None,
azimuthal_limits: Optional[tuple[float, float]] = None,
detector_regions: Optional[int | Sequence[int]] = None,
) -> Images | RealSpaceLineProfiles:
"""
Integrate polar regions to produce an image or line profiles.
Parameters
----------
radial_limits : tuple of float
Inner and outer radial angles of the integration limits [mrad].
azimuthal_limits : tuple of float
Lower and upper azimuthal angles of the integration limits [rad].
detector_regions : int or sequence of int
The explicit detector regions to integrate over.
Returns
-------
integrated_images : Images or RealSpaceLineProfiles
"""
if detector_regions is not None:
if (radial_limits is not None) or (azimuthal_limits is not None):
raise ValueError()
if np.isscalar(detector_regions):
detector_regions = [detector_regions]
array = self.array.reshape(self.shape[:-2] + (-1,))[
..., list(detector_regions)
].sum(axis=-1)
else:
if radial_limits is None:
radial_slice = slice(None)
else:
inner_index = int(
(radial_limits[0] - self.radial_offset) / self.radial_sampling
)
outer_index = int(
(radial_limits[1] - self.radial_offset) / self.radial_sampling
)
radial_slice = slice(inner_index, outer_index)
if outer_index > self.shape[-2]:
raise RuntimeError("Integration limit exceeded.")
if azimuthal_limits is None:
azimuthal_slice = slice(None)
else:
left_index = int(azimuthal_limits[0] / self.radial_sampling)
right_index = int(azimuthal_limits[1] / self.radial_sampling)
azimuthal_slice = slice(left_index, right_index)
array = self.array[..., radial_slice, azimuthal_slice].sum(axis=(-2, -1))
return _reduced_scanned_images_or_line_profiles(array, self)
[docs]
def gaussian_source_size(
self, sigma: float | tuple[float, float]
) -> PolarMeasurements:
"""
Simulate the effect of a finite source size on diffraction pattern(s) using a
Gaussian filter.
The filter is not applied to diffraction pattern individually, but the intensity
of diffraction patterns are mixed across scan axes. Applying this filter
requires two linear scan axes.
Applying this filter before integrating the diffraction patterns will produce
the same image as integrating the diffraction patterns first then applying a
Gaussian filter.
Parameters
----------
sigma : float or two float
Standard deviation of Gaussian kernel in the `x` and `y`-direction. If given
as a single number, the standard deviation is equal for both axes.
Returns
-------
filtered_diffraction_patterns : DiffractionPatterns
The filtered diffraction pattern(s).
"""
return _gaussian_source_size(self, sigma)
[docs]
def lorentzian_source_size(
self,
half_width: float | tuple[float, float],
truncate: float = 10.0,
) -> PolarMeasurements:
"""
Simulate the effect of a finite source size on polar measurement(s) using a
Lorentzian (Cauchy) filter.
The filter mixes intensities across scan axes (not within each measurement)
and requires two linear scan axes.
Parameters
----------
half_width : float or two float
Half-width at half-maximum (HWHM) of the Lorentzian kernel in the ``x``
and ``y``-direction [Å]. If given as a single number it is equal for
both axes.
truncate : float, optional
Truncate the kernel at this many half-widths (default is 10.0).
Returns
-------
filtered_measurements : PolarMeasurements
The filtered measurement(s).
Notes
-----
The Lorentzian kernel is parameterized by its HWHM γ (``half_width``),
with FWHM_L = 2γ. This differs from :meth:`gaussian_source_size`,
which uses the standard deviation σ (FWHM_G ≈ 2.3548·σ).
"""
return _lorentzian_source_size(self, half_width, truncate)
[docs]
def voigtian_source_size(
self,
gaussian_sigma: float | tuple[float, float],
lorentzian_gamma: float | tuple[float, float],
truncate: float = 10.0,
) -> PolarMeasurements:
"""
Simulate the effect of a finite source size on polar measurement(s) using a
Voigt filter (convolution of Gaussian and Lorentzian).
The filter mixes intensities across scan axes (not within each measurement)
and requires two linear scan axes.
Parameters
----------
gaussian_sigma : float or two float
Standard deviation (σ) of the Gaussian component [Å]. If given as a
single number it is equal for both axes.
lorentzian_gamma : float or two float
Half-width at half-maximum (HWHM, γ) of the Lorentzian component [Å].
If given as a single number it is equal for both axes.
truncate : float, optional
Truncate the Lorentzian component at this many half-widths
(default 10.0). The Gaussian component uses the SciPy default (4·σ).
Returns
-------
filtered_measurements : PolarMeasurements
The filtered measurement(s).
Notes
-----
The Voigt profile combines two width parameters:
* ``gaussian_sigma`` (σ): standard deviation of the Gaussian component,
FWHM_G = 2√(2 ln 2)·σ ≈ 2.3548·σ.
* ``lorentzian_gamma`` (γ): HWHM of the Lorentzian component,
FWHM_L = 2γ.
For the same FWHM, γ = FWHM / 2 but σ = FWHM / (2√(2 ln 2)) ≈ FWHM / 2.3548.
"""
return _voigtian_source_size(self, gaussian_sigma, lorentzian_gamma, truncate)
[docs]
def pseudo_voigtian_source_size(
self,
gaussian_sigma: float | tuple[float, float],
lorentzian_gamma: float | tuple[float, float],
eta: float,
truncate: float = 10.0,
) -> PolarMeasurements:
"""
Simulate the effect of a finite source size on polar measurement(s) using a
pseudo-Voigtian filter (weighted sum of Gaussian and Lorentzian).
The filter mixes intensities across scan axes (not within each measurement)
and requires two linear scan axes.
Parameters
----------
gaussian_sigma : float or two float
Standard deviation (σ) of the Gaussian component [Å]. If given as a
single number it is equal for both axes.
lorentzian_gamma : float or two float
Half-width at half-maximum (HWHM, γ) of the Lorentzian component [Å].
If given as a single number it is equal for both axes.
eta : float
Lorentzian mixing fraction η ∈ [0, 1]. η = 0 gives a pure Gaussian;
η = 1 gives a pure Lorentzian.
truncate : float, optional
Truncate the kernel at this many effective half-widths (default 10.0).
Returns
-------
filtered_measurements : PolarMeasurements
The filtered measurement(s).
Notes
-----
The pseudo-Voigt profile is PV = (1-η)·G + η·L, where G and L are
Gaussian and Lorentzian profiles with widths σ and γ respectively. Width
parameterization: FWHM_G = 2√(2 ln 2)·σ ≈ 2.3548·σ; FWHM_L = 2γ.
"""
return _pseudo_voigtian_source_size(
self, gaussian_sigma, lorentzian_gamma, eta, truncate
)
[docs]
def poisson_noise(
self,
dose_per_area: Optional[float] = None,
total_dose: Optional[float] = None,
samples: int = 1,
seed: Optional[int] = None,
):
"""
Add Poisson noise (i.e. shot noise) to a measurement corresponding to the provided 'total_dose' (per measurement
if applied to an ensemble) or 'dose_per_area' (not applicable for single measurements).
Parameters
----------
dose_per_area : float, optional
The irradiation dose per unit of scan area [electrons per Å:sup:`2`]. This is only valid if the diffraction
patterns has two scan axes.
total_dose : float, optional
The irradiation dose per diffraction pattern.
samples : int, optional
The number of samples to draw from a Poisson distribution. If this is greater than 1, an additional
ensemble axis will be added to the measurement.
seed : int, optional
Seed the random number generator.
Returns
-------
noisy_measurement : BaseMeasurements
The noisy measurement.
"""
if len(_scan_shape(self)) < 2 and dose_per_area is not None:
raise ValueError(
"Polar measurement has less than two scan axes, provide 'total_dose' not 'dose_per_area'."
)
return super().poisson_noise(
dose_per_area=dose_per_area,
total_dose=total_dose,
samples=samples,
seed=seed,
)
[docs]
def to_diffraction_patterns(
self, gpts: int | tuple[int, int], margin: float | tuple[float, float] = 0.1
):
"""
Convert the polar measurements to diffraction patterns by discretizing the polar
bins on a regular grid.
Parameters
----------
gpts : int or two int
Number of grid points describing the diffraction patterns.
margin : float or two float, optional
The margin as a fraction of the outer angle of the polar measurements to add
to the maximum angle of the diffraction patterns.
Returns
-------
diffraction_patterns : DiffractionPatterns
The diffraction patterns discretized from the polar measurements.
"""
if np.isscalar(gpts):
gpts = (gpts,) * 2
if np.isscalar(margin):
margin = (margin,) * 2
angular_sampling = (
(1 + margin[0]) * self.outer_angle / gpts[0] * 2,
(1 + margin[1]) * self.outer_angle / gpts[1] * 2,
)
nbins_radial, nbins_azimuthal = self.base_shape
regions = _polar_detector_bins(
gpts=gpts,
sampling=angular_sampling,
inner=self.radial_offset,
outer=self.outer_angle,
nbins_radial=nbins_radial,
nbins_azimuthal=nbins_azimuthal,
fftshift=True,
rotation=self.azimuthal_offset,
offset=(0.0, 0.0),
return_indices=False,
)
new_array = np.zeros(self.ensemble_shape + regions.shape, dtype=np.float32)
for i, indices in enumerate(label_to_index(regions)):
x, y = np.unravel_index(indices, regions.shape)
radial, azimuthal = np.unravel_index(i, (nbins_radial, nbins_azimuthal))
new_array[..., x, y] = self.array[..., radial, azimuthal][..., None]
new_array[..., regions < 0] = np.nan
wavelength = energy2wavelength(self._get_energy())
sampling = (
angular_sampling[0] / (wavelength * 1e3),
angular_sampling[1] / (wavelength * 1e3),
)
return DiffractionPatterns(
new_array,
sampling=sampling,
ensemble_axes_metadata=self.ensemble_axes_metadata,
metadata=self.metadata,
)
[docs]
def differentials(
self,
direction_1: tuple[int | tuple[int, ...], int | tuple[int, ...]],
direction_2: tuple[int | tuple[int, ...], int | tuple[int, ...]],
return_complex: bool = True,
) -> Images:
"""
Calculate the differential signal by subtracting the intensity of specified
detector regions.
Parameters
----------
direction_1 : tuple of int or tuple of tuple of int
The detector regions used for calculating the differential signal for the
first direction. The first item is the detector region(s) contributing to
the positive term and the second item is the detector region(s) contributing
to the negative terms.
direction_2 : tuple of int or tuple of tuple of int
The detector regions used for calculating the differential signal for the
second direction. The first item is the detector region(s) contributing to
the positive term and the second item is the detector region(s) contributing
to the negative terms.
return_complex : bool, optional
If True, return a complex image where the real and imaginary part represents
`direction_1` and `direction_2`. If False, return images with an ensemble
dimension for the directions.
Returns
-------
differential_image : Images
The (complex) differential image(s).
"""
differential_1 = self.integrate(
detector_regions=direction_1[1]
) - self.integrate(detector_regions=direction_1[0])
differential_2 = self.integrate(
detector_regions=direction_2[1]
) - self.integrate(detector_regions=direction_2[0])
if not return_complex:
stacked = stack(
(differential_1, differential_2), ("direction_1", "direction_2")
)
return stacked
xp = get_array_module(self.array)
array = xp.zeros_like(xp.array(differential_1.array), dtype=get_dtype(complex=True))
array.real = differential_1.array
array.imag = differential_2.array
return Images(array, **differential_1._copy_kwargs(exclude=("array",)))
[docs]
def to_image_ensemble(self):
"""
Convert the polar measurements to an ensemble of images, where the radial and
azimuthal angles becomes ensemble axes.
Returns
-------
image_ensemble : Images
"""
image_axes = _scan_axes(self)
xp = get_array_module(self.array)
array = xp.moveaxis(self.array, image_axes, (-2, -1))[..., 0, :, :]
ensemble_axes_metadata = [
axis.copy()
for i, axis in enumerate(self.axes_metadata[:-1])
if i not in image_axes
]
ensemble_axes_metadata[-1]._default_type = "range"
sampling = _scan_sampling(self)
return Images(
array,
sampling=(sampling[0], sampling[1]),
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=self.metadata,
)
[docs]
def show(
self,
ax: Optional[Axes] = None,
gpts: int | tuple[int, int] = (512, 512),
cbar: bool = False,
cmap: Optional[str] = None,
vmin: Optional[float] = None,
vmax: Optional[float] = None,
power: float = 1.0,
logscale: bool = False,
common_color_scale: bool = False,
explode: bool | Sequence[bool] = (),
overlay: bool | Sequence[int] = (),
figsize: Optional[tuple[int, int]] = None,
title: bool | str = True,
units: Optional[str] = None,
interact: bool = False,
display: bool = True,
) -> Visualization:
"""
Show the polar measurements as an image or a grid of images. The images are
shown by discretizing the polar bins on a regular grid.
Parameters
----------
ax : matplotlib.axes.Axes, optional
If given the plots are added to the axis. This is not available for exploded
plots.
gpts : int or tuple of int, optional
Number of grid points in the image(s).
cbar : bool, optional
Add colorbar(s) to the image(s). The size and padding of the colorbars may
be adjusted using the `set_cbar_size` and `set_cbar_padding` methods.
cmap : str, optional
Matplotlib colormap name used to map scalar data to colors. If the
measurement is complex the colormap must be one of 'hsv' or 'hsluv'.
vmin : float, optional
Minimum of the intensity color scale. Default is the minimum of the array
values.
vmax : float, optional
Maximum of the intensity color scale. Default is the maximum of the array
values.
power : float
Show image on a power scale. Cannot be used together with ``logscale``.
logscale : bool
If True, show image on a logarithmic intensity scale. Cannot be used
together with ``power != 1.0``.
common_color_scale : bool, optional
If True all images in an image grid are shown on the same colorscale, and a
single colorbar is created (if it is requested). Default is False.
explode : bool, optional
If True, a grid of images is created for all the items of the last two
ensemble axes. If False, the first ensemble item is shown. May be given as a
sequence of axis indices to create a grid of images from the specified axes.
The default is determined by the axis metadata.
overlay : bool or sequence of int, optional
If True, all line profiles in the ensemble are shown in a single plot.
If False, only the first ensemble item is shown. May be given as a sequence
of axis indices to specify which line profiles in the ensemble to show
together. The default is determined by the axis metadata.
figsize : two int, optional
The figure size given as width and height in inches, passed to
`matplotlib.pyplot.figure`.
title : bool or str, optional
Set the column title of the images. If True is given instead of a string the
title will be given by the value corresponding to the "name" key of the axes
metadata dictionary, if this item exists.
units : str
The units used for the x and y axes. The given units must be compatible with
the axes of the images.
interact : bool
If True, create an interactive visualization. This requires enabling the
`ipympl` Matplotlib backend.
display : bool, optional
If True (default) the figure is displayed immediately.
Returns
-------
measurement_visualization_2d : MeasurementVisualizationImshow
"""
diffraction_patterns = self.to_diffraction_patterns(gpts=gpts)
if not interact:
diffraction_patterns.compute()
return diffraction_patterns.show(
ax=ax,
cbar=cbar,
cmap=cmap,
vmin=vmin,
vmax=vmax,
power=power,
logscale=logscale,
common_color_scale=common_color_scale,
explode=explode,
overlay=overlay,
figsize=figsize,
title=title,
units=units,
interact=interact,
display=display,
)
[docs]
@jit(nopython=True, nogil=True, fastmath=True)
def calculate_max_reciprocal_space_vector(hkl, reciprocal_lattice_vectors):
k_max = 0.0
for i in range(len(hkl)):
lengths = (
(
hkl[i, 0] * reciprocal_lattice_vectors[..., 0, :]
+ hkl[i, 1] * reciprocal_lattice_vectors[..., 1, :]
+ hkl[i, 2] * reciprocal_lattice_vectors[..., 2, :]
)
** 2
).sum(-1)
if hasattr(lengths, "max"):
lengths = lengths.max()
k_max = max(k_max, lengths)
return np.sqrt(k_max)
# @jit(nopython=True, nogil=True)
def _reciprocal_lattice_vector_lengths(max_lengths, hkl, reciprocal_lattice_vectors):
for i in range(len(hkl)): # pylint: disable=not-an-iterable
lengths = np.sqrt(
(
(
hkl[i, 0] * reciprocal_lattice_vectors[..., 0, :]
+ hkl[i, 1] * reciprocal_lattice_vectors[..., 1, :]
+ hkl[i, 2] * reciprocal_lattice_vectors[..., 2, :]
)
** 2
).sum(-1)
)
max_lengths[i] = lengths.max()
return max_lengths
[docs]
def reciprocal_lattice_vector_lengths(hkl, reciprocal_lattice_vectors):
max_lengths = np.zeros(len(hkl))
return _reciprocal_lattice_vector_lengths(
max_lengths, hkl, reciprocal_lattice_vectors
)
@jit(nopython=True, nogil=True)
def _reciprocal_lattice_vector_mask(mask, hkl, reciprocal_lattice_vectors, k_max):
for i in range(len(hkl)): # pylint: disable=not-an-iterable
lengths = (
(
hkl[i, 0] * reciprocal_lattice_vectors[..., 0, :]
+ hkl[i, 1] * reciprocal_lattice_vectors[..., 1, :]
+ hkl[i, 2] * reciprocal_lattice_vectors[..., 2, :]
)
** 2
).sum(-1)
new_mask = lengths < k_max**2
if isinstance(new_mask, bool):
mask[i] = new_mask
else:
mask[i] = new_mask.any()
[docs]
class IndexedDiffractionPatterns(BaseMeasurements):
"""
Diffraction patterns indexed by their Miller indices.
Parameters
----------
array : numpy.ndarray
1D or greater array of type `float` or `complex`. The last axis represents the
diffraction spots and should have the same length as the number of miller
indices, any preceding axis represents an ensemble axis.
miller_indices : numpy.ndarray
The miller indices of the diffraction spots as an N x 3 array where N is the
number of miller indices. The order of the miller indices must correspond to the
array of intensities. The second axis represents each hkl miller index.
reciprocal_lattice_vectors : numpy.ndarray
The reciprocal lattice vectors of the crystal as a 3 x 3 array. The first axis
represents miller indices and the order of the items must correspond to the
array of intensities. The second axis represents the reciprocal space positions
in x, y and z [1/Å].
ensemble_axes_metadata : list of AxisMetadata, optional
List of metadata associated with the ensemble axes. The length and item order
must match the ensemble axes.
metadata : dict, optional
A dictionary defining measurement metadata.
"""
_base_dims = 1
def __init__(
self,
array: da.core.Array | np.ndarray,
miller_indices: np.ndarray,
reciprocal_lattice_vectors: np.ndarray,
ensemble_axes_metadata: Optional[list[AxisMetadata]] = None,
metadata: Optional[dict] = None,
):
if isinstance(reciprocal_lattice_vectors, Cell):
reciprocal_lattice_vectors = np.array(reciprocal_lattice_vectors)
if len(reciprocal_lattice_vectors.shape) <= len(array.shape):
reciprocal_lattice_vectors = reciprocal_lattice_vectors[
(None,) * (len(reciprocal_lattice_vectors.shape) - len(array.shape) + 1)
]
if not is_broadcastable(
array.shape[:-1], reciprocal_lattice_vectors.shape[:-2]
):
raise ValueError()
if not len(miller_indices) == array.shape[-1]:
raise ValueError(
"The number of miller indices must be equal to the number of"
" diffraction spots."
)
super().__init__(
array=array,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
self._miller_indices = miller_indices
self._intensities = array
self._reciprocal_lattice_vectors = reciprocal_lattice_vectors
def _area_per_pixel(self):
raise NotImplementedError
@property
def base_axes_metadata(self) -> list[AxisMetadata]:
return [AxisMetadata(label="hkl")]
@property
def intensities(self) -> np.ndarray:
"""
Intensities of the diffraction spots.
"""
return self._array
@property
def reciprocal_lattice_vectors(self) -> np.ndarray:
"""
Reciprocal lattice vectors of the diffraction spots.
"""
return self._reciprocal_lattice_vectors
@property
def positions(self) -> np.ndarray:
"""
Reciprocal space positions of the diffraction spots.
"""
positions = self.miller_indices @ self.reciprocal_lattice_vectors
return positions
@property
def all_positions(self) -> np.ndarray:
"""
Reciprocal space positions of the diffraction spots.
"""
repeats = ()
for n, m in zip(self.shape[:-1], self.reciprocal_lattice_vectors.shape[:-2]):
if n == m:
repeats += (1,)
elif m == 1:
repeats += (n,)
else:
raise RuntimeError("Incompatible shapes.")
positions = np.tile(self.positions, repeats + (1, 1))
return positions
@property
def miller_indices(self) -> np.ndarray:
"""
Miller indices of the diffraction spots.
"""
return self._miller_indices
@property
def angular_positions(self):
"""
Scattering angles of the diffraction spots.
"""
wavelength = energy2wavelength(self._get_energy())
return self.positions * wavelength * 1e3
@property
def ensemble_shape(self) -> tuple:
return self.intensities.shape[:-1]
@classmethod
def _pack_kwargs(cls, kwargs):
kwargs["miller_indices"] = kwargs["miller_indices"].tolist()
# kwargs["reciprocal_lattice_vectors"] = [
# (float(position[0]), float(position[1]), float(position[2]))
# for position in kwargs["reciprocal_lattice_vectors"]
# ]
kwargs["reciprocal_lattice_vectors"] = kwargs[
"reciprocal_lattice_vectors"
].tolist()
packed_kwargs = super()._pack_kwargs(kwargs)
return packed_kwargs
@classmethod
def _unpack_kwargs(cls, attrs):
kwargs = super()._unpack_kwargs(attrs)
kwargs["miller_indices"] = np.array(kwargs["miller_indices"], dtype=int)
kwargs["reciprocal_lattice_vectors"] = np.array(
kwargs["reciprocal_lattice_vectors"], dtype=float
)
return kwargs
def __getitem__(self, items):
items = _validate_array_items(items, self.shape)
kwargs = self.get_items(items)
new_items = ()
for i, n in zip(items, self.reciprocal_lattice_vectors.shape):
if n == 1:
if isinstance(i, int):
new_items += (0,)
else:
new_items += (slice(None),)
else:
new_items += (i,)
# items = tuple(
# 0 if (isinstance(i, int) and (n == 1)) else slice(None)
# for i, n in zip(items, self.reciprocal_lattice_vectors.shape)
# )
kwargs["reciprocal_lattice_vectors"] = self.reciprocal_lattice_vectors[
new_items
]
return self.__class__(**kwargs)
def _reduction(
self,
reduction_func,
axes,
keepdims: bool = False,
split_every: int = 2,
**kwargs,
):
reciprocal_lattice_vectors = getattr(np, reduction_func)(
self.reciprocal_lattice_vectors, axis=axes, keepdims=keepdims
)
kwargs["reciprocal_lattice_vectors"] = reciprocal_lattice_vectors
return super()._reduction(reduction_func, axes, keepdims, split_every, **kwargs)
[docs]
def remove_low_intensity(self, threshold: float = 1e-3):
"""
Remove diffraction spots with intensity below a threshold for all ensemble
dimensions.
Parameters
----------
threshold : float
Intensity threshold for removing diffraction spots.
Returns
-------
thresholded_spots : IndexedDiffractionPatterns
The indexed diffraction spots with an intensity above the given threshold.
"""
if self.is_lazy:
raise RuntimeError("Cannot threshold lazy IndexedDiffractionPatterns.")
ensemble_axes = tuple(range(self.ensemble_dims))
xp = get_array_module(self.intensities)
mask = xp.max(self.intensities, axis=ensemble_axes) > threshold
mask = asnumpy(mask)
miller_indices = self.miller_indices[mask]
intensities = self.intensities[..., mask]
return self.__class__(
intensities,
miller_indices,
reciprocal_lattice_vectors=self.reciprocal_lattice_vectors,
ensemble_axes_metadata=self.ensemble_axes_metadata,
metadata=self._metadata,
)
[docs]
def sort(self, criterion: str = "distance"):
"""
Sort the diffraction spots according to a given criterion.
Parameters
----------
criterion : {'distance', 'intensity'}
The boundary parameter determines how the images are extended beyond their
boundaries when the filter overlaps with a border.
``distance`` :
Sort according to the distance in reciprocal space from the zero
frequency.
``intensity`` :
Sort according to the intensity of the diffraction spots.
Returns
-------
sorted_spots : IndexedDiffractionPatterns
The indexed diffraction spots sorted according to the given criterion.
"""
if self.lazy:
raise RuntimeError("Cannot sort lazy IndexedDiffractionPatterns.")
if criterion == "distance":
criterion = -np.linalg.norm(self.positions, axis=1)
elif criterion == "intensity":
ensemble_axes = tuple(range(len(self.ensemble_shape)))
criterion = -np.max(self.intensities, axis=ensemble_axes)
else:
raise ValueError()
order = np.argsort(criterion)
array = self.array[..., order]
miller_indices = self.miller_indices[order]
reciprocal_lattice_vectors = self.reciprocal_lattice_vectors[..., order, :, :]
return self.__class__(
array,
miller_indices,
reciprocal_lattice_vectors=reciprocal_lattice_vectors,
ensemble_axes_metadata=self.ensemble_axes_metadata,
metadata=self._metadata,
)
[docs]
def crop(self, max_angle: Optional[float] = None, k_max: Optional[float] = None):
"""
Crop the indexed diffraction patterns such that they only include spots with
spatial frequencies (scattering angles) up to a given limit.
Parameters
----------
max_angle : float, optional
The maximum included scattering angle in the cropped diffraction patterns.
k_max : float, optional
The maximum included reciprocal lattice vector in the cropped diffraction
spots.
Returns
-------
cropped : IndexedDiffractionPatterns
The cropped indexed diffraction spots.
"""
if max_angle is not None and k_max is None:
wavelength = energy2wavelength(self._get_energy())
k_max = max_angle / wavelength / 1e3
elif not k_max or max_angle:
raise ValueError("Either 'max_angle' or 'k_max' must be given.")
mask = np.zeros(len(self.miller_indices), dtype=bool)
_reciprocal_lattice_vector_mask(
mask,
self.miller_indices.astype(self.reciprocal_lattice_vectors.dtype),
self.reciprocal_lattice_vectors,
k_max,
)
miller_indices = self.miller_indices[mask]
array = self.array[..., mask]
return self.__class__(
array,
miller_indices,
reciprocal_lattice_vectors=self.reciprocal_lattice_vectors,
ensemble_axes_metadata=self.ensemble_axes_metadata,
metadata=self._metadata,
)
[docs]
def normalize_to_spot(self, spot: Optional[tuple[int, int, int]] = None):
"""
Normalize the intensity of the diffraction spots.
Parameters
----------
spot : tuple of three int
The intensities will be normalized with respect to the intensity of this
spot. Defaults to the most intense spot.
Returns
-------
normalized_indexed_diffraction_patterns : IndexedDiffractionPatterns
"""
intensities_dict = self.intensities_dict
if spot is None:
c = np.max(self.intensities)
else:
c = intensities_dict[spot]
intensities = self.intensities / c
return self.__class__(
intensities,
self.miller_indices.copy(),
self.positions.copy(),
ensemble_axes_metadata=self.ensemble_axes_metadata,
metadata=self._metadata,
)
[docs]
def to_data_array(self):
"""
Convert the indexed diffraction patterns to xarray DataArray.
Returns
-------
data_array_of_indexed_spots : xarray.DataArray
"""
if xr is None:
raise RuntimeError("xarray is required to convert to DataArray.")
coords = []
for axes_metadata, n in zip(self.ensemble_axes_metadata, self.ensemble_shape):
new_coords = list(axes_metadata.coordinates(n))
new_coords = xr.DataArray(
new_coords,
dims=[axes_metadata.label],
attrs={"units": axes_metadata.units},
)
coords.append(new_coords)
hkl = list([f"{h} {k} {l}" for h, k, l in self.miller_indices])
hkl = xr.DataArray(hkl, name="hkl", dims=("hkl",))
coords.append(hkl)
dims = [axes_metadata.label for axes_metadata in self.ensemble_axes_metadata]
dims = dims + ["hkl"]
data = xr.DataArray(
self.array,
dims=dims,
coords=coords,
attrs={
"long_name": "intensity",
"units": self.metadata.get("units", "arb. unit"),
},
)
k = reciprocal_lattice_vector_lengths(
self.miller_indices, self.reciprocal_lattice_vectors
)
data = data.assign_coords(k=("hkl", k))
return data
def _miller_indices_to_string(self):
return [f"{h} {k} {l}" for h, k, l in self.miller_indices]
[docs]
def to_dataframe(self):
"""
Convert the indexed diffraction patterns to pandas DataFrame.
Returns
-------
data_frame : pd.DataFrame
"""
if pd is None:
raise RuntimeError("pandas is required to convert to DataFrame.")
if self.ensemble_shape:
if len(self.ensemble_shape) > 1:
raise RuntimeError(
"cannot convert indexed diffraction patterns with more than one"
" ensemble axis to dataframe"
)
intensities = {
hkl: self.intensities[..., i]
for i, hkl in enumerate(self._miller_indices_to_string())
}
axes_metadata = self.ensemble_axes_metadata[0]
if hasattr(axes_metadata, "values"):
index = axes_metadata.values
else:
index = list(range(len(self.intensities)))
df = pd.DataFrame(intensities, index=index)
with config.set({"visualize.use_tex": False}):
df.index.name = self.axes_metadata[0].format_label()
df.columns.name = self.axes_metadata[1].format_label()
return df
else:
intensities = {
hkl: intensity
for hkl, intensity in zip(
self._miller_indices_to_string(), self.intensities
)
}
return pd.DataFrame(intensities, index=[0])
[docs]
def block_direct(self):
"""
Remove the zero-order spot.
Returns
-------
blocked : IndexedDiffractionPatterns
The indexed diffraction spots without the zero-order spot.
"""
to_delete = np.where(np.all(self.miller_indices == 0, axis=1))[0]
miller_indices = np.delete(self.miller_indices, to_delete, axis=0)
intensities = np.delete(self.intensities, to_delete, axis=-1)
return self.__class__(
intensities,
miller_indices,
reciprocal_lattice_vectors=self.reciprocal_lattice_vectors,
ensemble_axes_metadata=self.ensemble_axes_metadata,
metadata=self._metadata,
)
[docs]
def max_reciprocal_space_vector_length(self):
return calculate_max_reciprocal_space_vector(
self.miller_indices, self.reciprocal_lattice_vectors
)
[docs]
def show(
self,
ax: Optional[Axes] = None,
cbar: bool = False,
cmap: Optional[str] = None,
vmin: Optional[float] = None,
vmax: Optional[float] = None,
power: float = 1.0,
logscale: bool = False,
common_color_scale: bool = False,
scale: float = 0.5,
explode: bool | Sequence[bool] = (),
overlay: bool | Sequence[bool] = (),
figsize: Optional[tuple[int, int]] = None,
title: bool | str = True,
units: Optional[str] = None,
interact: bool = False,
display: bool = True,
**kwargs,
):
"""
Show the diffraction spots as an EllipseCollection using matplotlib.
Parameters
----------
ax : matplotlib.axes.Axes, optional
If given the plots are added to the axis. This is not available for exploded
plots.
cbar : bool, optional
Add colorbar(s) to the image(s). The size and padding of the colorbars may
be adjusted using the `set_cbar_size` and `set_cbar_padding` methods.
cmap : str, optional
Matplotlib colormap name used to map scalar data to colors. If the
measurement is complex the colormap must be one of 'hsv' or 'hsluv'.
vmin : float, optional
Minimum of the intensity color scale. Default is the minimum of the array
values.
vmax : float, optional
Maximum of the intensity color scale. Default is the maximum of the array
values.
power : float
Show diffraction spots intensities on a power scale. Cannot be used
together with ``logscale``.
logscale : bool
If True, show diffraction spots on a logarithmic intensity scale.
Cannot be used together with ``power != 1.0``.
common_color_scale : bool, optional
If True all images in an image grid are shown on the same colorscale, and a
single colorbar is created (if it is requested). Default is False.
scale : float, optional
Scale the radii of the circles representing the diffraction spots.
explode : bool or sequence of bool, optional
If True, a grid of plots is created for all the items of the last two
ensemble axes. If False, the first ensemble item is shown. May be given as a
sequence of axis indices to create a grid of plots from the specified axes.
The default is determined by the axis metadata.
overlay : bool or sequence of int, optional
If True, all line profiles in the ensemble are shown in a single plot.
If False, only the first ensemble item is shown. May be given as a sequence
of axis indices to specify which line profiles in the ensemble to show
together. The default is determined by the axis metadata.
figsize : two int, optional
The figure size given as width and height in inches, passed to
`matplotlib.pyplot.figure`.
title : bool or str, optional
Set the column title of the plots. If True is given instead of a string the
title will be given by the value corresponding to the "name" key of the axes
metadata dictionary, if this item exists.
units : str
The units used for the x and y axes. The given units must be compatible with
the axes of the plots.
interact : bool
If True, create an interactive visualization. This requires enabling the
`ipympl` Matplotlib backend.
display : bool, optional
If True (default) the figure is displayed immediately.
Returns
-------
visualization : Visualization
"""
k_max = self.max_reciprocal_space_vector_length() * get_conversion_factor(
units, "1/Å", self.metadata.get("energy", None)
)
xlim = (-k_max, k_max)
ylim = (-k_max, k_max)
visualization = Visualization(
measurement=self,
ax=ax,
common_scale=common_color_scale,
figsize=figsize,
title=title,
aspect=True,
share_x=True,
share_y=True,
explode=explode,
overlay=overlay,
interactive=not interact and display,
value_limits=(vmin, vmax),
power=power,
logscale=logscale,
cmap=cmap,
cbar=cbar,
scale=scale,
units=units,
xlim=xlim,
ylim=ylim,
**kwargs,
)
if interact:
visualization.interact(ScatterGUI, display=display)
return visualization
@property
def intensities_dict(self) -> dict[tuple, np.ndarray]:
"""
A dictionary mapping miller indices to intensities.
"""
intensities = {
tuple(hkl): intensity
for hkl, intensity in zip(
self.miller_indices,
np.moveaxis(self.intensities, -1, 0),
)
}
values = np.zeros(self.shape[:-1], dtype=np.float32)
intensities = defaultdict(lambda: values, intensities)
return intensities
@property
def positions_dict(self) -> dict[tuple, np.ndarray]:
"""
A dictionary mapping miller indices to reciprocal space positions [1/Å].
"""
positions = {
tuple(hkl): position
for hkl, position in zip(
self.miller_indices,
np.moveaxis(self.reciprocal_lattice_vectors, -2, 0),
)
}
return positions
@classmethod
def _stack(
cls,
diffraction_spots: IndexedDiffractionPatterns,
axis_metadata: list[AxisMetadata],
axis: int,
):
intensities = [spots.intensities_dict for spots in diffraction_spots]
# positions = [spots.positions_dict for spots in diffraction_spots]
# def merge_dicts_no_overwrite(dict1, dict2):
# return {**dict1, **{k: v for k, v in dict2.items() if k not in dict1}}
# merged = {}
# for positions1 in positions:
# merged = merge_dicts_no_overwrite(merged, positions1)
# positions = [
# merge_dicts_no_overwrite(positions1, merged) for positions1 in positions
# ]
miller_indices = list(
set(itertools.chain(*[intensities1.keys() for intensities1 in intensities]))
)
new_intensities = {}
for hkl in miller_indices:
new_intensities[hkl] = []
for intensities1 in intensities:
new_intensities[hkl].append(intensities1[hkl])
new_intensities[hkl] = np.stack(new_intensities[hkl], axis=axis)
miller_indices = np.stack(list(new_intensities.keys()), axis=0)
intensities = np.stack(list(new_intensities.values()), axis=-1)
positions = np.stack(
[spots.reciprocal_lattice_vectors for spots in diffraction_spots], axis=0
)
ensemble_axes_metadata = [
axis_metadata.copy()
for axis_metadata in diffraction_spots[0].ensemble_axes_metadata
]
ensemble_axes_metadata.insert(axis, axis_metadata)
metadata = diffraction_spots[0].metadata
return IndexedDiffractionPatterns(
intensities, miller_indices, positions, ensemble_axes_metadata, metadata
)
def _axis_range_mask(
values: np.ndarray, value_range: Optional[tuple[float, float]], param_name: str
) -> np.ndarray:
if value_range is None:
return np.ones(len(values), dtype=bool)
lo, hi = value_range
if lo > hi:
raise ValueError(f"{param_name}={value_range}: min must not exceed max")
mask = (values >= lo) & (values <= hi)
if not mask.any():
raise ValueError(
f"{param_name}={value_range} selects no data -- values span "
f"[{values.min()}, {values.max()}]"
)
return mask
[docs]
class MomentumResolvedSpectrum(BaseMeasurements):
"""
Momentum-resolved energy-loss spectrum S(q, E).
Stores a 2D intensity array whose last two dimensions correspond to
scattering-vector bins (q) and energy bins (E). Additional leading
dimensions are ensemble axes.
Use :meth:`crop` to restrict the q/energy-loss range before plotting
with :meth:`show` -- besides zooming in, this also rescales the
colour scale (and, for an exploded grid, every panel's shared colour
scale) to the cropped region instead of the full data.
Parameters
----------
array : numpy.ndarray or dask array
Array of shape ``(..., n_q, n_E)``.
q_values : sequence of float
Scattering-vector values [mrad] for each q bin.
e_values : sequence of float
Energy-loss values [eV] for each energy bin.
ensemble_axes_metadata : list of AxisMetadata, optional
metadata : dict, optional
"""
_base_dims = 2
def __init__(
self,
array: np.ndarray | da.core.Array,
q_values: Sequence[float] | np.ndarray,
e_values: Sequence[float] | np.ndarray,
ensemble_axes_metadata: Optional[list[AxisMetadata]] = None,
metadata: Optional[dict] = None,
):
self._q_values = tuple(float(v) for v in q_values)
self._e_values = tuple(float(v) for v in e_values)
if metadata is None:
metadata = {}
metadata.setdefault("label", "intensity")
metadata.setdefault("units", "arb. unit")
super().__init__(
array=array,
ensemble_axes_metadata=ensemble_axes_metadata,
metadata=metadata,
)
@property
def q_values(self) -> tuple[float, ...]:
"""Scattering-vector values [mrad]."""
return self._q_values
@property
def e_values(self) -> tuple[float, ...]:
"""Energy-loss values [eV]."""
return self._e_values
@property
def _area_per_pixel(self):
raise RuntimeError("Cannot infer pixel area for a spectrum.")
@property
def base_axes_metadata(self) -> list[AxisMetadata]:
return [
NonLinearAxis(
label="q",
values=self._q_values,
units="mrad",
),
NonLinearAxis(
label="energy loss",
values=self._e_values,
units="eV",
),
]
[docs]
def crop(
self,
q_range: Optional[tuple[float, float]] = None,
e_range: Optional[tuple[float, float]] = None,
) -> "MomentumResolvedSpectrum":
"""
Crop the spectrum to a q and/or energy-loss range.
Parameters
----------
q_range : tuple of float, optional
Inclusive ``(min, max)`` scattering-vector range [mrad] to keep.
If None (default), the full q range is kept.
e_range : tuple of float, optional
Inclusive ``(min, max)`` energy-loss range [eV] to keep -- note
this is in eV (the unit ``e_values`` is stored in) regardless of
the ``e_units`` display option of :meth:`show`. If None
(default), the full energy range is kept.
Returns
-------
cropped : MomentumResolvedSpectrum
The cropped spectrum.
"""
q = np.array(self._q_values)
e = np.array(self._e_values)
q_mask = _axis_range_mask(q, q_range, "q_range")
e_mask = _axis_range_mask(e, e_range, "e_range")
array = self.array[..., q_mask, :][..., :, e_mask]
kwargs = self._copy_kwargs(exclude=("array", "q_values", "e_values"))
kwargs["array"] = array
kwargs["q_values"] = q[q_mask]
kwargs["e_values"] = e[e_mask]
return self.__class__(**kwargs)
[docs]
def show(
self,
ax: Optional[Axes] = None,
cbar: bool = True,
cmap: Optional[str] = "viridis",
vmin: Optional[float] = None,
vmax: Optional[float] = None,
power: float = 1.0,
explode: bool | Sequence[int] = (),
figsize: Optional[tuple[int, int]] = None,
title: bool | str = True,
e_units: str = "meV",
**kwargs,
) -> tuple:
"""
Show the spectrum as a 2D heatmap with q on the x-axis and energy
on the y-axis.
Unlike other abTEM measurements this method plots directly with
matplotlib (the q and energy axes are non-linear, so the shared
:class:`.Visualization` imshow machinery does not apply) and therefore
returns the raw matplotlib objects rather than a :class:`.Visualization`.
Parameters
----------
ax : matplotlib Axes, optional
Only used for the single-panel case (ignored when exploding).
cbar : bool
Show colorbar. Default True.
cmap : str, optional
Colormap name. Default 'viridis'.
vmin, vmax : float, optional
Colour-scale limits. When exploding, a value left as None is filled
from the global min/max across all panels so the shared colorbar is
meaningful.
power : float
Display on a power scale.
explode : bool or sequence of int
If True, explode all ensemble axes into a panel grid. If a sequence
of ints, explode only those ensemble-axis indices (the remaining
ensemble axes collapse to their first element). If falsy (default),
a single panel is shown.
figsize : (width, height), optional
title : bool or str
If a string, used as the (base) title. If True, a default title is
generated. If False, no title is set.
e_units : str
Units for the energy axis ('meV' or 'eV'). Default 'meV'.
kwargs
Forwarded to :meth:`matplotlib.axes.Axes.pcolormesh`.
Returns
-------
fig, ax : matplotlib Figure and Axes (single panel) or ndarray of Axes
(exploded grid).
"""
import itertools
import warnings
import matplotlib.pyplot as plt
from matplotlib.colors import PowerNorm
array = self.array
if hasattr(array, "compute"):
array = array.compute()
q = np.array(self._q_values)
e = np.array(self._e_values)
if e_units == "meV":
e = e * 1000.0
e_label = "Energy loss [meV]"
else:
e_label = "Energy loss [eV]"
n_q, n_e = len(q), len(e)
ensemble_shape = array.shape[: -self._base_dims]
n_ensemble = len(ensemble_shape)
# Resolve which ensemble axes to explode into a panel grid.
if explode is True:
explode_axes: tuple[int, ...] = tuple(range(n_ensemble))
elif not explode:
explode_axes = ()
else:
explode_axes = tuple(int(a) % n_ensemble for a in explode)
cbar_label = (
f"{self.metadata.get('label', '')} [{self.metadata.get('units', '')}]"
)
def panel_data(grid_index: tuple[int, ...]) -> np.ndarray:
# Exploded axes take their grid value; other ensemble axes collapse
# to their first element. grid_index is positional in explode_axes
# order (matching how indices was built below), not ascending axis
# order, so map by axis identity rather than consuming positionally.
axis_to_value = dict(zip(explode_axes, grid_index))
full = tuple(
axis_to_value.get(d, 0) for d in range(n_ensemble)
)
data = array[full] if full else array
# array may be GPU-resident (cupy); np.asarray() cannot convert a
# cupy array implicitly (cupy deliberately blocks it). Duck-type
# on .get() rather than using asnumpy(), which is a no-op when
# the cupy package itself isn't installed and so would not
# convert a genuinely GPU-resident array in that environment.
data = data.get() if hasattr(data, "get") else np.asarray(data)
if data.shape != (n_q, n_e):
raise ValueError(
f"Shape mismatch in MomentumResolvedSpectrum.show(): "
f"data.shape={data.shape} but expected ({n_q}, {n_e}) from "
f"q_values (len {n_q}) and e_values (len {n_e}). "
f"Full array shape: {array.shape}"
)
return data
if explode_axes:
grid_sizes = [ensemble_shape[a] for a in explode_axes]
indices = list(itertools.product(*[range(s) for s in grid_sizes]))
n = len(indices)
ncols = min(n, 4)
nrows = (n + ncols - 1) // ncols
if figsize is None:
figsize = (4 * ncols, 3.5 * nrows)
# Axes are linked explicitly below rather than via plt.subplots'
# own sharex=True/sharey=True: that path reads back the object
# array slot it just allocated (matplotlib/gridspec.py's
# `axarr[0, 0]`) before every axes is created, and on some numpy
# builds that slot is not reliably None on first read.
fig, axes_arr = plt.subplots(
nrows,
ncols,
figsize=figsize,
squeeze=False,
)
axes_flat = axes_arr.flatten()
for ax in axes_flat[1:]:
ax.sharex(axes_flat[0])
ax.sharey(axes_flat[0])
# Shared colour scale across panels so the single colorbar applies to
# every panel (otherwise the norm autoscales to the first panel only).
panels = [panel_data(idx).T for idx in indices]
_vmin = min(float(p.min()) for p in panels) if vmin is None else vmin
_vmax = max(float(p.max()) for p in panels) if vmax is None else vmax
norm = PowerNorm(gamma=power, vmin=_vmin, vmax=_vmax)
im = None
for k, (idx, data_t) in enumerate(zip(indices, panels)):
a = axes_flat[k]
im = a.pcolormesh(
q, e, data_t, shading="nearest", cmap=cmap, norm=norm, **kwargs
)
a.set_xlabel("q [mrad]")
a.set_ylabel(e_label)
parts = [
self.ensemble_axes_metadata[a_idx][idx[j]].format_title()
for j, a_idx in enumerate(explode_axes)
]
panel_title = ", ".join(parts)
if isinstance(title, str):
panel_title = f"{title} — {panel_title}" if panel_title else title
if title:
a.set_title(panel_title)
for k in range(len(indices), len(axes_flat)):
axes_flat[k].set_visible(False)
# Hide the x/y tick labels and axis labels of interior panels
# (kept only on the bottom row / left column), since sharex/sharey
# above already makes them redundant on every other panel.
for a in axes_flat[: len(indices)]:
a.label_outer()
# tight_layout() must run before fig.colorbar(): colorbar()
# shrinks the given panel axes via their gridspec to make room for
# the new colorbar axes, and a tight_layout() call afterwards
# would re-layout every axes from scratch -- including the
# colorbar axes, which it has no special handling for -- and can
# end up overlapping it with the last panel instead of leaving it
# in its own strip.
fig.tight_layout()
if cbar and im is not None:
fig.colorbar(
im,
ax=axes_flat[: len(indices)].tolist(),
label=cbar_label,
pad=0.015,
)
return fig, axes_arr
# Single panel — collapse any ensemble axes to their first element.
if n_ensemble > 0:
warnings.warn(
"MomentumResolvedSpectrum.show(): array has ensemble axes "
f"{ensemble_shape}; showing member {(0,) * n_ensemble} only. Pass "
"explode=True (or a sequence of ensemble-axis indices) to see all.",
stacklevel=2,
)
data = panel_data(())
if ax is None:
if figsize is None:
figsize = (6, 4)
fig, ax = plt.subplots(figsize=figsize)
else:
fig = ax.get_figure()
norm = PowerNorm(gamma=power, vmin=vmin, vmax=vmax)
im = ax.pcolormesh(
q, e, data.T, shading="nearest", cmap=cmap, norm=norm, **kwargs
)
ax.set_xlabel("q [mrad]")
ax.set_ylabel(e_label)
if isinstance(title, str):
ax.set_title(title)
elif title is True:
ax.set_title("S(q, E)")
if cbar:
fig.colorbar(im, ax=ax, label=cbar_label)
fig.tight_layout()
return fig, ax
def _thermal_weight_tds(
I_tds: np.ndarray,
e_values: np.ndarray,
energy_axis_idx: int,
temperature: float,
) -> tuple[np.ndarray, np.ndarray]:
"""Unfold a classical (loss/gain-symmetric) TDS intensity, computed at
energy magnitudes ``|E|`` only, into quantum loss (``+E``) and gain
(``-E``) sides via detailed balance.
A single frozen-phonon multislice run per energy magnitude already
contains the *combined* loss + gain intensity — the split between the
two sides is a quantum-statistical effect governed by the Bose-Einstein
phonon occupation ``n(E)``, not something the classical thermal sampling
distinguishes. The zero-energy bin is passed through unweighted: there is
no loss/gain asymmetry at zero energy transfer, and the classical
sampling already accounts for the thermal occupation there.
Loss and gain weights, ``(n+1)/(2n+1)`` and ``n/(2n+1)``, sum to 1, so
splitting preserves the total (loss + gain) spectral weight of the
unweighted input at each energy magnitude.
"""
from ase import units
if e_values[0] != 0.0 or np.any(np.diff(e_values) <= 0):
raise ValueError(
"temperature-based loss/gain unfolding requires energies "
"starting at 0 and strictly increasing (one multislice run per "
"|energy loss|, unfolded into +/- sides here)."
)
xp = get_array_module(I_tds)
is_lazy = isinstance(I_tds, da.core.Array)
concatenate = _array_module_fn(I_tds, xp, "concatenate")
flip = _array_module_fn(I_tds, xp, "flip")
nonzero_e = e_values[1:]
beta = 1.0 / (units.kB * temperature)
n_occ = 1.0 / (np.exp(nonzero_e * beta) - 1.0)
loss_weight = (n_occ + 1.0) / (2.0 * n_occ + 1.0)
gain_weight = n_occ / (2.0 * n_occ + 1.0)
def _broadcast(weight):
shape = [1] * I_tds.ndim
shape[energy_axis_idx] = len(weight)
return xp.asarray(weight.reshape(shape))
zero_slice = tuple(
slice(0, 1) if i == energy_axis_idx else slice(None)
for i in range(I_tds.ndim)
)
nonzero_slice = tuple(
slice(1, None) if i == energy_axis_idx else slice(None)
for i in range(I_tds.ndim)
)
I_zero = I_tds[zero_slice]
I_loss = I_tds[nonzero_slice] * _broadcast(loss_weight)
# Reverse along the energy axis: ascending |E| -> descending (most
# negative first), so the final signed axis comes out ascending overall.
I_gain = flip(I_tds[nonzero_slice] * _broadcast(gain_weight), axis=energy_axis_idx)
result_array = concatenate([I_gain, I_zero, I_loss], axis=energy_axis_idx)
if is_lazy:
# I_gain/I_zero/I_loss each inherit exit_waves' original per-config
# chunking along the energy axis (e.g. one chunk per energy value),
# so the concatenated result ends up with many tiny chunks there.
# Downstream dask operations on this axis (e.g. interpolate_line in
# momentum_resolved_spectrum) mishandle that level of fragmentation,
# so consolidate it into a single chunk here at the source rather
# than relying on every consumer to work around it.
result_array = result_array.rechunk({energy_axis_idx: -1})
e_values_signed = np.concatenate([-nonzero_e[::-1], [0.0], nonzero_e])
return result_array, e_values_signed
[docs]
def phonon_loss_diffraction_patterns(
exit_waves,
component: str = "tds",
max_angle: str | float = "cutoff",
parity: str = "odd",
block_direct: bool | float = False,
temperature: Optional[float] = None,
) -> "DiffractionPatterns":
"""
Compute inelastic (TDS) diffraction patterns from energy-resolved
frozen-phonon exit waves.
The thermal diffuse scattering signal is obtained per energy bin as::
I_coherent = |FT(Σ_j psi_j)|² / N² (elastic)
I_incoherent = Σ_j |FT(psi_j)|² / N (total)
I_tds = I_incoherent - I_coherent (inelastic / phonon loss)
The returned ``DiffractionPatterns`` retain the ``EnergyLossAxis`` but the
``FrozenPhononsAxis`` is collapsed. Apply an offset ``AnnularDetector``
or ``SlitDetector`` to integrate over desired q points.
Note that ``I_tds`` is the variance of the diffracted amplitude across
frozen-phonon configurations: with a single configuration per energy,
``I_incoherent`` and ``I_coherent`` are identical by construction and
``I_tds`` is exactly zero everywhere, not a numerical artifact. At least
2 configurations per energy are required for ``component="tds"``/``"all"``
(this is enforced with a ``ValueError``); in practice many more are
needed for good statistics.
Parameters
----------
exit_waves : Waves
Complex exit waves from a multislice simulation with an
``EnergyResolvedAtomsEnsemble`` (``ensemble_mean=False``).
Must contain both a ``FrozenPhononsAxis`` and an ``EnergyLossAxis``
in its ensemble axes. If the potential used more than one slice,
prefer building it with ``projection="finite"`` — see the
:class:`~abtem.inelastic.phonons.EnergyResolvedAtomsEnsemble` notes
on slice-boundary artifacts with out-of-plane displacement.
component : {'tds', 'coherent', 'incoherent', 'all'}
Which component to return. ``'all'`` stacks the three along a
new leading ``OrdinalAxis(label='component')``.
max_angle : str or float
Passed to ``Waves.diffraction_patterns``.
parity : str
Passed to ``Waves.diffraction_patterns``.
block_direct : bool or float, optional
If True, the direct beam is blocked in the resulting diffraction
patterns. If given as a float, masks up to that scattering angle
[mrad]. Default is False.
temperature : float, optional
Sample temperature [K]. If given, unfolds the TDS signal — computed
from a single frozen-phonon run per energy *magnitude* — into signed
quantum loss (``+E``) and gain (``-E``) sides using Bose-Einstein
detailed balance, following P. Zeiger's approach. Requires
``component="tds"`` and an ``EnergyLossAxis`` whose values start at 0
and strictly increase (the classical/incoherent-minus-coherent
signal is symmetric in loss/gain; only their *split* is a quantum
effect). The zero-energy bin is unweighted. Default is None (no
unfolding — the returned energies are the ones in ``exit_waves``).
Returns
-------
DiffractionPatterns
Intensity patterns with the ``FrozenPhononsAxis`` removed and the
``EnergyLossAxis`` preserved (or replaced by its signed loss/gain
unfolding if ``temperature`` is given).
"""
from abtem.core.axes import EnergyLossAxis, FrozenPhononsAxis, OrdinalAxis
# --- validate ensemble axes ---
fp_axis_idx = None
energy_axis_idx = None
for i, ax in enumerate(exit_waves.ensemble_axes_metadata):
if isinstance(ax, FrozenPhononsAxis):
fp_axis_idx = i
if isinstance(ax, EnergyLossAxis):
energy_axis_idx = i
if fp_axis_idx is None:
raise ValueError(
"exit_waves must have a FrozenPhononsAxis in ensemble_axes_metadata. "
"Did you create the EnergyResolvedAtomsEnsemble with ensemble_mean=False?"
)
if energy_axis_idx is None:
raise ValueError(
"exit_waves must have an EnergyLossAxis in ensemble_axes_metadata."
)
if not np.iscomplexobj(exit_waves.array):
raise ValueError(
"exit_waves must contain complex wave functions (not intensities). "
"Pass the Waves object directly, not DiffractionPatterns."
)
# --- number of frozen-phonon configurations ---
N = exit_waves.shape[fp_axis_idx]
if N < 2 and component in ("tds", "all"):
raise ValueError(
f"component={component!r} requires at least 2 frozen-phonon "
f"configurations per energy, got N={N}. TDS/phonon-loss is "
"I_incoherent - I_coherent, the variance of the diffracted "
"amplitude across configurations -- for a single configuration "
"I_incoherent and I_coherent are identical by construction, so "
"the result is exactly zero everywhere, not a numerical fluke. "
"Build the EnergyResolvedAtomsEnsemble with more than one "
"configuration per energy to get a non-trivial signal."
)
dp_kwargs = dict(max_angle=max_angle, parity=parity, fftshift=True)
# Coherent: sum complex exit waves first, then compute diffraction pattern
# I_coh = |FT(Σ_j psi_j)|² / N²
coherent_waves = exit_waves.sum(axis=fp_axis_idx)
dp_coherent = coherent_waves.diffraction_patterns(**dp_kwargs)
I_coherent = dp_coherent.array / N**2
# Incoherent: compute diffraction patterns first, then sum intensities
# I_inc = Σ_j |FT(psi_j)|² / N
dp_all = exit_waves.diffraction_patterns(**dp_kwargs)
I_incoherent = dp_all.array.sum(axis=fp_axis_idx) / N
# TDS = incoherent - coherent
I_tds = I_incoherent - I_coherent
# --- select component ---
valid_components = ("tds", "coherent", "incoherent", "all")
if component not in valid_components:
raise ValueError(f"component must be one of {valid_components}")
if temperature is not None and component != "tds":
raise ValueError(
"temperature-based loss/gain unfolding requires component='tds'."
)
remaining_axes = [
ax
for i, ax in enumerate(exit_waves.ensemble_axes_metadata)
if i != fp_axis_idx
]
if temperature is not None:
remaining_energy_axis_idx = next(
i for i, ax in enumerate(remaining_axes) if isinstance(ax, EnergyLossAxis)
)
energy_axis = remaining_axes[remaining_energy_axis_idx]
e_values = np.asarray(energy_axis.values, dtype=float)
I_tds, e_values_signed = _thermal_weight_tds(
I_tds, e_values, remaining_energy_axis_idx, temperature
)
remaining_axes[remaining_energy_axis_idx] = EnergyLossAxis(
values=tuple(e_values_signed), units=energy_axis.units
)
if component == "all":
xp = get_array_module(I_tds)
stack_fn = _array_module_fn(I_tds, xp, "stack")
result_array = stack_fn([I_coherent, I_incoherent, I_tds], axis=0)
component_axis = OrdinalAxis(
label="component",
values=("coherent", "incoherent", "tds"),
)
remaining_axes = [component_axis] + remaining_axes
elif component == "tds":
result_array = I_tds
elif component == "coherent":
result_array = I_coherent
else:
result_array = I_incoherent
metadata = dict(dp_coherent.metadata)
metadata["phonon_loss_component"] = component
result = DiffractionPatterns(
result_array,
sampling=dp_coherent.sampling,
fftshift=dp_coherent.fftshift,
ensemble_axes_metadata=remaining_axes or None,
metadata=metadata,
)
if block_direct:
radius = block_direct if isinstance(block_direct, (int, float)) else None
result = result.block_direct(radius=radius)
return result
[docs]
def momentum_resolved_spectrum(
tds_diffraction_patterns: "DiffractionPatterns",
detector,
) -> "MomentumResolvedSpectrum":
"""
Build S(q, E) from energy-resolved TDS diffraction patterns.
Dispatches on the type of *detector*:
* :class:`~abtem.detectors.SpectralAnnularDetector` — sweeps an offset
circular acceptance disk (acceptance **radius** ``outer``, inner always 0)
along a radial direction at each q-step. The q-axis runs from
``q_min`` to ``q_max`` in steps of ``outer`` (one disk-radius per step).
* :class:`~abtem.detectors.SpectralSlitDetector` — uses
:meth:`DiffractionPatterns.interpolate_line` (spline interpolation) to
integrate strips along the slit's long axis. Samples directly from
``q_min`` to ``q_max`` (default ``q_min=0`` includes the direct beam).
The integration aperture perpendicular to q has **full width** ``width``
(= ``2 * outer`` of an equivalent annular detector).
Both detector types share the same ``q_min`` / ``q_max`` convention — the
same numerical values yield the same q-range in the output spectrum::
SpectralAnnularDetector(outer=r, q_min=Q0, q_max=Q)
SpectralSlitDetector(q_min=Q0, q_max=Q, width=2*r)
Parameters
----------
tds_diffraction_patterns : DiffractionPatterns
Energy-resolved TDS diffraction patterns as returned by
:func:`phonon_loss_diffraction_patterns`. Must have an
``EnergyLossAxis`` in its ensemble axes.
detector : SpectralAnnularDetector or SpectralSlitDetector
Detector that controls the integration strategy and q-range.
Returns
-------
MomentumResolvedSpectrum
"""
from abtem.core.axes import EnergyLossAxis, OrdinalAxis
from abtem.core.energy import energy2wavelength
from abtem.detectors import SpectralAnnularDetector, SpectralSlitDetector
dp = tds_diffraction_patterns
# --- auto-select TDS slice when component="all" was used ---
if dp.metadata.get("phonon_loss_component") == "all":
for i, ax in enumerate(dp.ensemble_axes_metadata):
if isinstance(ax, OrdinalAxis) and ax.label == "component":
tds_idx = list(ax.values).index("tds")
slicing = tuple(
tds_idx if j == i else slice(None)
for j in range(len(dp.ensemble_axes_metadata))
)
dp = dp[slicing]
break
# --- find the energy axis ---
energy_axis_idx = None
for i, ax in enumerate(dp.ensemble_axes_metadata):
if isinstance(ax, EnergyLossAxis):
energy_axis_idx = i
break
if energy_axis_idx is None:
raise ValueError(
"tds_diffraction_patterns must have an EnergyLossAxis in "
"ensemble_axes_metadata. Use phonon_loss_diffraction_patterns first."
)
energy_axis = dp.ensemble_axes_metadata[energy_axis_idx]
e_values = np.array(energy_axis.values)
# Guard against a mismatch between axis metadata and actual array size
# (e.g. a frozen-phonon axis that was not fully collapsed).
n_e_array = dp.array.shape[energy_axis_idx]
n_e_axis = len(e_values)
if n_e_array != n_e_axis:
raise ValueError(
f"EnergyLossAxis at position {energy_axis_idx} has {n_e_axis} values "
f"but dp.array has size {n_e_array} along that axis. "
f"Make sure the input was produced by phonon_loss_diffraction_patterns "
f"with all FrozenPhonons axes collapsed."
)
remaining_ensemble = [
ax for i, ax in enumerate(dp.ensemble_axes_metadata) if i != energy_axis_idx
]
# ---- SpectralSlitDetector: interpolated strip integration ----
if isinstance(detector, SpectralSlitDetector):
# interpolate_line works in the DP's internal coordinate system (Å⁻¹).
# Convert mrad → Å⁻¹ using the electron wavelength.
energy = dp.metadata.get("energy")
if energy is None:
raise ValueError(
"DiffractionPatterns metadata must contain 'energy' [eV] for "
"mrad → Å⁻¹ conversion in SpectralSlitDetector mode."
)
wavelength = energy2wavelength(energy) # Å
mrad_to_inv_ang = 1.0 / (wavelength * 1e3)
offset_inv = (
detector.offset[0] * mrad_to_inv_ang,
detector.offset[1] * mrad_to_inv_ang,
)
q_min_inv = detector.q_min * mrad_to_inv_ang
q_max_inv = detector.q_max * mrad_to_inv_ang
width_inv = detector.width * mrad_to_inv_ang
sampling_inv = min(dp.sampling)
cos_a, sin_a = cos_sin_deg(detector.angle)
# Sample directly from q_min to q_max along the slit direction.
# For q_min=0 this naturally includes q=0 as the first point.
N = max(2, int(round((q_max_inv - q_min_inv) / sampling_inv)) + 1)
start_inv = (offset_inv[0] + q_min_inv * cos_a, offset_inv[1] + q_min_inv * sin_a)
end_inv = (offset_inv[0] + q_max_inv * cos_a, offset_inv[1] + q_max_inv * sin_a)
# A fragmented energy axis (e.g. one chunk per energy value, as
# EnergyResolvedAtomsEnsemble produces by default, or as
# _thermal_weight_tds's concatenate leaves it) causes dask's
# interpolate_line + the q-binning reshape below to compute an
# incorrect result size. Consolidate it into a single chunk first.
dp_for_interp = dp.rechunk({energy_axis_idx: -1}) if dp.is_lazy else dp
line_profiles = dp_for_interp.interpolate_line(
start=start_inv,
end=end_inv,
gpts=N,
width=width_inv,
endpoint=True,
order=1, # linear interp — cubic splines ring around sharp peaks
)
# interpolate_line with width>0 returns the AVERAGE across the
# perpendicular direction. Convert to a SUM to be consistent with the
# annular detector path (which sums all pixels inside the mask).
# n_perp is the number of perpendicular samples used internally.
n_perp = int(np.floor(width_inv / sampling_inv / 2) * 2 + 1)
# energy dim is at energy_axis_idx; move it last → (...other, n_q, n_E)
xp = get_array_module(line_profiles.array)
moveaxis = _array_module_fn(line_profiles.array, xp, "moveaxis")
spectrum_array = moveaxis(
line_profiles.array * n_perp, energy_axis_idx, -1
)
q_values_fine = np.linspace(detector.q_min, detector.q_max, N)
# Optionally bin along q to reduce the number of points
if detector.q_sampling is not None:
q_sampling_mrad = detector.q_sampling
native_step = (detector.q_max - detector.q_min) / max(N - 1, 1)
bin_size = max(1, int(round(q_sampling_mrad / native_step)))
n_q_fine = spectrum_array.shape[-2]
n_q_binned = n_q_fine // bin_size
# Trim to a multiple of bin_size, then reshape and sum
trimmed = spectrum_array[..., : n_q_binned * bin_size, :]
new_shape = trimmed.shape[:-2] + (n_q_binned, bin_size, trimmed.shape[-1])
spectrum_array = trimmed.reshape(new_shape).sum(axis=-2)
# Bin centres for q_values
q_trimmed = q_values_fine[: n_q_binned * bin_size]
q_values = q_trimmed.reshape(n_q_binned, bin_size).mean(axis=1)
else:
q_values = q_values_fine
# ---- SpectralAnnularDetector: offset-disk sweep ----
elif isinstance(detector, SpectralAnnularDetector):
outer = detector.outer
q_max = detector.q_max if detector.q_max is not None else min(dp.max_angles)
q_step = detector.q_sampling if detector.q_sampling is not None else outer
n_steps = max(2, round((q_max - detector.q_min) / q_step) + 1)
q_values = np.linspace(detector.q_min, q_max, n_steps)
cos_a, sin_a = cos_sin_deg(detector.sweep_angle)
xp = get_array_module(dp.array)
gpts = dp.shape[-2:]
sampling = dp.angular_sampling
# The spatial-frequency grid is the same for every q-step (only the
# offset/roll differs), so compute it once rather than inside the
# per-q loop below.
k2 = _spatial_frequency_squared(gpts, sampling, xp)
# Build all masks at once and stack: (n_q, gy, gx)
masks = xp.stack(
[
_annular_detector_mask(
gpts=gpts,
sampling=sampling,
inner=0.0,
outer=outer,
offset=(float(q * cos_a), float(q * sin_a)),
fftshift=dp.fftshift,
xp=xp,
k2=k2,
)
for q in q_values
],
axis=0,
)
# Single batched contraction: (...ens, gy, gx) × (n_q, gy, gx) → (...ens, n_q)
tensordot = _array_module_fn(dp.array, xp, "tensordot")
# dask (as of 2025.11.0) silently computes the wrong result for
# da.tensordot with negative axis indices -- declared shape is
# correct but the actual data is not. Positive indices are
# unaffected and work identically for numpy/cupy's own tensordot,
# so always use them here rather than only when routed through dask.
dp_axes = (dp.array.ndim - 2, dp.array.ndim - 1)
mask_axes = (masks.ndim - 2, masks.ndim - 1)
spectrum_array = tensordot(dp.array, masks, axes=(dp_axes, mask_axes))
moveaxis = _array_module_fn(spectrum_array, xp, "moveaxis")
spectrum_array = moveaxis(spectrum_array, energy_axis_idx, -1) # (...ens, n_q, n_E)
else:
raise TypeError(
f"detector must be a SpectralAnnularDetector or SpectralSlitDetector, "
f"got {type(detector).__name__}"
)
# Validate: last two dims must be (n_q, n_e)
expected_tail = (len(q_values), len(e_values))
if spectrum_array.shape[-2:] != expected_tail:
raise RuntimeError(
f"momentum_resolved_spectrum internal error: spectrum_array.shape="
f"{spectrum_array.shape} but expected last two dims {expected_tail} "
f"(n_q={len(q_values)}, n_e={len(e_values)}). "
f"dp.array.shape={dp.array.shape}, energy_axis_idx={energy_axis_idx}"
)
return MomentumResolvedSpectrum(
array=spectrum_array,
q_values=q_values,
e_values=e_values,
ensemble_axes_metadata=remaining_ensemble or None,
metadata={"label": "intensity", "units": "arb. unit"},
)
