SMatrixArray

class abtem.prism.s_matrix.SMatrixArray(array, wave_vectors, semiangle_cutoff, energy=None, interpolation=(1, 1), sampling=None, extent=None, window_gpts=(0, 0), window_offset=(0, 0), periodic=(True, True), device=None, ensemble_axes_metadata=None, metadata=None)

Bases: BaseSMatrix, ArrayObject

A scattering matrix defined by a given array of dimension 3, where the first indexes the probe plane waves and the latter two are the y and x scan directions.

Parameters:

• array (np.ndarray) – Array defining the scattering matrix. Must be 3D or higher, dimensions before the last three dimensions should represent ensemble dimensions, the next dimension indexes the plane waves and the last two dimensions represent the spatial extent of the plane waves.

• wave_vectors (np.ndarray) – Array defining the wave vectors corresponding to each plane wave. Must have shape Nx2, where N is equal to the number of plane waves.

• semiangle_cutoff (float) – The radial cutoff of the plane-wave expansion [mrad].

• energy (float) – Electron energy [eV].

• sampling (one or two float, optional) – Lateral sampling of wave functions [1 / Å]. Provide only if potential is not given. Will be ignored if ‘gpts’ is also provided.

• extent (one or two float, optional) – Lateral extent of wave functions [Å]. Provide only if potential is not given.

• interpolation (one or two int, optional) – Interpolation factor in the x and y directions (default is 1, ie. no interpolation). If a single value is provided, assumed to be the same for both directions.

• window_gpts (tuple of int) – The number of grid points describing the cropping window of the wave functions.

• window_offset (tuple of int) – The number of grid points from the origin the cropping windows of the wave functions is displaced.

• periodic (tuple of bool) – Specifies whether the SMatrix should be assumed to be periodic along the x and y-axis.

• device (str, optional) – The calculations will be carried out on this device (‘cpu’ or ‘gpu’). Default is ‘cpu’. The default is determined by the user configuration.

• ensemble_axes_metadata (list of AxesMetadata) – Axis metadata for each ensemble axis. The axis metadata must be compatible with the shape of the array.

• metadata (dict) – A dictionary defining wave function metadata. All items will be added to the metadata of measurements derived from the waves.

__init__(array, wave_vectors, semiangle_cutoff, energy=None, interpolation=(1, 1), sampling=None, extent=None, window_gpts=(0, 0), window_offset=(0, 0), periodic=(True, True), device=None, ensemble_axes_metadata=None, metadata=None)

Methods

__init__(array, wave_vectors, semiangle_cutoff)
apply_func(func, **kwargs)	rtype: TypeVar(T, bound= ArrayObject)
apply_transform(transform[, max_batch])	Transform the wave functions by a given transformation.
compute([progress_bar, profiler, ...])	Turn a lazy abTEM object into its in-memory equivalent.
copy()	Make a copy.
copy_to_device(device)	Copy SMatrixArray to specified device.
dummy_probes([scan, ctf, plane])	A probe or an ensemble of probes equivalent reducing the SMatrix at a single position.
ensemble_blocks([chunks])	Split the ensemble into an array of smaller ensembles.
ensure_lazy([chunks])	Creates an equivalent lazy version of the array object.
expand_dims([axis, axis_metadata])	Expand the shape of the array object.
from_array_and_metadata(array, ...)	Creates array object from a given array and metadata.
from_zarr(url[, chunks])	Read wave functions from a hdf5 file.
generate_blocks([chunks])	Generate chunks of the ensemble.
generate_ensemble([keepdims])	Generate every member of the ensemble.
get_from_metadata(name[, broadcastable])
get_items(items[, keepdims])	Index the array and the corresponding axes metadata.
lazy([chunks])	rtype: TypeVar(T, bound= ArrayObject)
match_grid(other[, check_match])	Match the grid to another object with a Grid.
max([axis, keepdims, split_every])	Maximum of array object over one or more axes.
mean([axis, keepdims, split_every])	Mean of array object over one or more axes.
min([axis, keepdims, split_every])	Minmimum of array object over one or more axes.
multislice([potential])	type potential: BasePotential
no_base_chunks()	Rechunk to remove chunks across the base dimensions.
rechunk(chunks[, in_place])	Rechunk dask array.
reduce([scan, ctf, detectors, ...])	Scan the probe across the potential and record a measurement for each detector.
scan([scan, detectors, ctf, ...])	Reduce the SMatrix using coefficients calculated by a BaseScan and a CTF, to obtain the exit wave functions at given initial probe positions and aberrations.
select_block(index, chunks)	Select a block from the ensemble.
set_ensemble_axes_metadata(axes_metadata, axis)	Sets the axes metadata of an ensemble axis.
squeeze([axis])	Remove axes of length one from array object.
std([axis, keepdims, split_every])	Standard deviation of array object over one or more axes.
sum([axis, keepdims, split_every])	Sum of array object over one or more axes.
to_cpu()	Move the array to the host memory from an arbitrary source array.
to_data_array()	Convert ArrayObject to a xarray DataArray.
to_gpu([device])	Move the array from the host memory to a gpu.
to_hyperspy()	Convert ArrayObject to a Hyperspy signal.
to_tiff(filename, **kwargs)	Write data to a tiff file.
to_zarr(url[, compute, overwrite])	Write data to a zarr file.

Attributes

accelerator	Accelerator object describing the acceleration energy.
angular_sampling	Reciprocal-space sampling in units of scattering angles [mrad].
antialias_cutoff_gpts	The number of grid points along the x and y direction in the simulation grid at the antialiasing cutoff scattering angle.
antialias_valid_gpts	The number of grid points along the x and y direction in the simulation grid for the largest rectangle that fits within antialiasing cutoff scattering angle.
array	Underlying array describing the array object.
axes_metadata	List of AxisMetadata.
base_axes_metadata	List of AxisMetadata for the base axes in real space.
base_dims	Number of base dimensions.
base_shape	Shape of the base axes of the underlying array.
cutoff_angles	Scattering angles at the antialias cutoff [mrad].
device	The device on which the SMatrixArray is reduced.
dtype	The datatype of waves.
energy	Electron acceleration energy in electron volts.
ensemble_axes_metadata	Axis metadata for each ensemble axis.
ensemble_dims	Number of ensemble dimensions.
ensemble_shape	Shape of the ensemble axes of the underlying array.
extent	Extent of grid for each dimension in Ångstrom.
full_cutoff_angles	Scattering angles corresponding to the full wave function size [mrad].
gpts	Number of grid points for each dimension.
grid	Simulation grid.
interpolation	Interpolation factor in the x and y directions
is_complex	True if array is complex.
is_lazy	True if array is lazy.
metadata	Metadata stored as a dictionary.
periodic	If True the SMatrix is assumed to be periodic along corresponding axis.
reciprocal_space_axes_metadata	List of AxisMetadata for base axes in reciprocal space.
reciprocal_space_sampling	Reciprocal-space sampling in reciprocal Ångstrom.
rectangle_cutoff_angles	Scattering angles corresponding to the sides of the largest rectangle within the antialias cutoff [mrad].
sampling	Grid sampling for each dimension in Ångstrom per grid point.
semiangle_cutoff	The cutoff semiangle of the plane wave expansion.
shape	Shape of the underlying array.
storage_device	The device on which the SMatrixArray is stored.
wave_vectors	The wave vectors corresponding to each plane wave.
wavelength	Relativistic wavelength in Ångstrom.
waves	The wave vectors describing each plane wave.
window_extent	The cropping window extent of the waves.
window_gpts	The number of grid points describing the cropping window of the wave functions.
window_offset	The number of grid points from the origin the cropping windows of the wave functions is displaced.

property accelerator: Accelerator

Accelerator object describing the acceleration energy.

property angular_sampling: tuple[float, float]

Reciprocal-space sampling in units of scattering angles [mrad].

property antialias_cutoff_gpts: tuple[int, int]

The number of grid points along the x and y direction in the simulation grid at the antialiasing cutoff scattering angle.

property antialias_valid_gpts: tuple[int, int]

The number of grid points along the x and y direction in the simulation grid for the largest rectangle that fits within antialiasing cutoff scattering angle.

apply_transform(transform, max_batch='auto')

Transform the wave functions by a given transformation.

Parameters:

• transform (ArrayObjectTransform) – The array object transformation to apply.

• max_batch (int, optional) – The number of wave functions in each chunk of the Dask array. If ‘auto’ (default), the batch size is automatically chosen based on the abtem user configuration settings “dask.chunk-size” and “dask.chunk-size-gpu”.

Returns:

transformed_array_object – The transformed array object.

Return type:

ArrayObjectTransform

property array: ndarray | Array

Underlying array describing the array object.

property axes_metadata: AxesMetadataList

List of AxisMetadata.

property base_axes_metadata: list[AxisMetadata]

List of AxisMetadata for the base axes in real space.

property base_dims

Number of base dimensions.

property base_shape: tuple[int, ...]

Shape of the base axes of the underlying array.

compute(progress_bar=None, profiler=False, resource_profiler=False, **kwargs)

Turn a lazy abTEM object into its in-memory equivalent.

Parameters:

• progress_bar (bool) – Display a progress bar in the terminal or notebook during computation. The progress bar is only displayed with a local scheduler.

• profiler (bool) – Return Profiler class used to profile Dask’s execution at the task level. Only execution with a local is profiled.

• resource_profiler (bool) – Return ResourceProfiler class is used to profile Dask’s execution at the resource level.

• kwargs – Additional keyword arguments passes to dask.compute.

copy()

Make a copy.

copy_to_device(device)

Copy SMatrixArray to specified device.

Return type:

SMatrixArray

property cutoff_angles: tuple[float, float]

Scattering angles at the antialias cutoff [mrad].

property device

The device on which the SMatrixArray is reduced.

property dtype

The datatype of waves.

dummy_probes(scan=None, ctf=None, plane='entrance', **kwargs)

A probe or an ensemble of probes equivalent reducing the SMatrix at a single position.

Parameters:

• scan (BaseScan) –

• ctf (CTF) –

• plane (str) –

Returns:

dummy_probes

Return type:

Probes

property energy

Electron acceleration energy in electron volts.

property ensemble_axes_metadata: list[AxisMetadata]

Axis metadata for each ensemble axis.

ensemble_blocks(chunks=None)

Split the ensemble into an array of smaller ensembles.

Parameters:

chunks (iterable of tuples) – Block sizes along each dimension.

Return type:

Array

property ensemble_dims

Number of ensemble dimensions.

property ensemble_shape: tuple[int, int]

Shape of the ensemble axes of the underlying array.

ensure_lazy(chunks='auto')

Creates an equivalent lazy version of the array object.

Parameters:

chunks (int or tuple or str) – How to chunk the array. See dask.array.from_array.

Returns:

lazy_array_object – Lazy version of the array object.

Return type:

ArrayObject or subclass of ArrayObject

expand_dims(axis=None, axis_metadata=None)

Expand the shape of the array object.

Parameters:

• axis (int or tuple of ints) – Position in the expanded axes where the new axis (or axes) is placed.

• axis_metadata (AxisMetadata or List of AxisMetadata, optional) – The axis metadata describing the expanded axes. Default is UnknownAxis.

Returns:

expanded – View of array object with the number of dimensions increased.

Return type:

ArrayObject or subclass of ArrayObject

property extent: tuple[float] | tuple[float, float] | tuple[float, ...]

Extent of grid for each dimension in Ångstrom.

classmethod from_array_and_metadata(array, axes_metadata, metadata)

Creates array object from a given array and metadata.

Parameters:

• array (array) – Complex array defining one or more 2D wave functions. The second-to-last and last dimensions are the wave function y- and x-axis, respectively.

• 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) – A dictionary defining wave function metadata. All items will be added to the metadata of measurements derived from the waves. The metadata must contain the electron energy [eV].

Returns:

wave_functions – The created wave functions.

Return type:

Waves

classmethod from_zarr(url, chunks='auto')

Read wave functions from a hdf5 file.

Return type:

TypeVar(T, bound= ArrayObject)

urlstr

Location of the data, typically a path to a local file. A URL can also include a protocol specifier like s3:// for remote data.

chunkstuple of ints or tuples of ints

Passed to dask.array.from_array(), allows setting the chunks on initialisation, if the chunking scheme in the on-disc dataset is not optimal for the calculations to follow.

property full_cutoff_angles: tuple[float, float]

Scattering angles corresponding to the full wave function size [mrad].

generate_blocks(chunks=1)

Generate chunks of the ensemble.

Parameters:

chunks (iterable of tuples) – Block sizes along each dimension.

generate_ensemble(keepdims=False)

Generate every member of the ensemble.

Parameters:

keepdims (bool, opptional) – If True, all ensemble axes are left in the result as dimensions with size one. Default is False.

Yields:

ArrayObject or subclass of ArrayObject – Member of the ensemble.

get_items(items, keepdims=False)

Index the array and the corresponding axes metadata. Only ensemble axes can be indexed.

Parameters:

• items (int or tuple of int or slice) – The array is indexed according to this.

• keepdims (bool, optional) – If True, all ensemble axes are left in the result as dimensions with size one. Default is False.

Returns:

indexed_array – The indexed array object.

Return type:

ArrayObject or subclass of ArrayObject

property gpts: tuple[int] | tuple[int, int] | tuple[int, ...]

Number of grid points for each dimension.

property grid: Grid

Simulation grid.

property interpolation: tuple[int, int]

Interpolation factor in the x and y directions

property is_complex: bool

True if array is complex.

property is_lazy: bool

True if array is lazy.

match_grid(other, check_match=False)

Match the grid to another object with a Grid.

max(axis=None, keepdims=False, split_every=2)

Maximum of array object over one or more axes. Only ensemble axes can be reduced.

Parameters:

• axis (int or tuple of ints, optional) – Axis or axes along which a maxima are calculated. The default is to compute the mean of the flattened array. If this is a tuple of ints, the maxima are calculated over multiple axes. The indicated axes must be ensemble axes.

• keepdims (bool, optional) – If True, the reduced axes are left in the result as dimensions with size one. Default is False.

• split_every (int) – Only used for lazy arrays. See dask.array.reductions.

Returns:

reduced_array – The reduced array object.

Return type:

ArrayObject or subclass of ArrayObject

mean(axis=None, keepdims=False, split_every=2)

Mean of array object over one or more axes. Only ensemble axes can be reduced.

Parameters:

• axis (int or tuple of ints, optional) – Axis or axes along which a means are calculated. The default is to compute the mean of the flattened array. If this is a tuple of ints, the mean is calculated over multiple axes. The indicated axes must be ensemble axes.

• keepdims (bool, optional) – If True, the reduced axes are left in the result as dimensions with size one. Default is False.

• split_every (int) – Only used for lazy arrays. See dask.array.reductions.

Returns:

reduced_array – The reduced array object.

Return type:

ArrayObject or subclass of ArrayObject

property metadata: dict

Metadata stored as a dictionary.

min(axis=None, keepdims=False, split_every=2)

Minmimum of array object over one or more axes. Only ensemble axes can be reduced.

Parameters:

• axis (int or tuple of ints, optional) – Axis or axes along which a minima are calculated. The default is to compute the mean of the flattened array. If this is a tuple of ints, the minima are calculated over multiple axes. The indicated axes must be ensemble axes.

• keepdims (bool, optional) – If True, the reduced axes are left in the result as dimensions with size one. Default is False.

• split_every (int) – Only used for lazy arrays. See dask.array.reductions.

Returns:

reduced_array – The reduced array object.

Return type:

ArrayObject or subclass of ArrayObject

multislice(potential=None)

Parameters:

potential (BasePotential) –

Return type:

SMatrixArray

no_base_chunks()

Rechunk to remove chunks across the base dimensions.

property periodic: tuple[bool, bool]

If True the SMatrix is assumed to be periodic along corresponding axis.

rechunk(chunks, in_place=True)

Rechunk dask array.

chunksint or tuple or str

How to rechunk the array. See dask.array.rechunk.

kwargs :

Additional keyword arguments passes to dask.array.rechunk.

property reciprocal_space_axes_metadata: list[AxisMetadata]

List of AxisMetadata for base axes in reciprocal space.

property reciprocal_space_sampling: tuple[float] | tuple[float, float] | tuple[float, ...]

Reciprocal-space sampling in reciprocal Ångstrom.

property rectangle_cutoff_angles: tuple[float, float]

Scattering angles corresponding to the sides of the largest rectangle within the antialias cutoff [mrad].

reduce(scan=None, ctf=None, detectors=None, max_batch_reduction='auto', reduction_scheme='auto')

Scan the probe across the potential and record a measurement for each detector.

Parameters:

• detectors (list of Detector objects) – The detectors recording the measurements.

• scan (Scan object) – Scan defining the positions of the probe wave functions.

• ctf (CTF object, optional) – The probe contrast transfer function. Default is None (aperture is set by the planewave cutoff).

• max_batch_reduction (int or str, optional) – Number of positions per reduction operation. A large number of positions better utilize thread parallelization, but requires more memory and floating point operations. If ‘auto’ (default), the batch size is automatically chosen based on the abtem user configuration settings “dask.chunk-size” and “dask.chunk-size-gpu”.

• rechunk (two int or str, optional) – Partitioning of the scan. The scattering matrix will be reduced in similarly partitioned chunks. Should be equal to or greater than the interpolation.

Return type:

BaseMeasurements | Waves | list[BaseMeasurements | Waves]

property sampling: tuple[float] | tuple[float, float] | tuple[float, ...]

Grid sampling for each dimension in Ångstrom per grid point.

scan(scan=None, detectors=None, ctf=None, max_batch_reduction='auto', rechunk='auto')

Reduce the SMatrix using coefficients calculated by a BaseScan and a CTF, to obtain the exit wave functions at given initial probe positions and aberrations.

Parameters:

• scan (BaseScan) – Positions of the probe wave functions. If not given, scans across the entire potential at Nyquist sampling.

• detectors (BaseDetector, list of BaseDetector, optional) – A detector or a list of detectors defining how the wave functions should be converted to measurements after running the multislice algorithm. See abtem.measurements.detect for a list of implemented detectors.

• ctf (CTF) – Contrast transfer function from used for calculating the expansion coefficients in the reduction of the SMatrix.

• max_batch_reduction (int or str, optional) – Number of positions per reduction operation. A large number of positions better utilize thread parallelization, but requires more memory and floating point operations. If ‘auto’ (default), the batch size is automatically chosen based on the abtem user configuration settings “dask.chunk-size” and “dask.chunk-size-gpu”.

• rechunk (str or tuple of int, optional) – Parallel reduction of the SMatrix requires rechunking the Dask array from chunking along the expansion axis to chunking over the spatial axes. If given as a tuple of int of length the SMatrix is rechunked to have those chunks. If ‘auto’ (default) the chunks are taken to be identical to the interpolation factor.

Returns:

• detected_waves (BaseMeasurements or list of BaseMeasurement) – The detected measurement (if detector(s) given).

• exit_waves (Waves) – Wave functions at the exit plane(s) of the potential (if no detector(s) given).

select_block(index, chunks)

Select a block from the ensemble.

Parameters:

• index (tuple of ints) – Index of selected block.

• chunks (iterable of tuples) – Block sizes along each dimension.

property semiangle_cutoff: float

The cutoff semiangle of the plane wave expansion.

set_ensemble_axes_metadata(axes_metadata, axis)

Sets the axes metadata of an ensemble axis.

Parameters:

• axes_metadata (AxisMetadata) – The new axis metadata.

• axis (int) – The axis to set.

property shape: tuple[int, ...]

Shape of the underlying array.

squeeze(axis=None)

Remove axes of length one from array object.

Parameters:

axis (int or tuple of ints, optional) – Selects a subset of the entries of length one in the shape.

Returns:

squeezed – The input array object, but with all or a subset of the dimensions of length 1 removed.

Return type:

ArrayObject or subclass of ArrayObject

std(axis=None, keepdims=False, split_every=2)

Standard deviation of array object over one or more axes. Only ensemble axes can be reduced.

Parameters:

• axis (int or tuple of ints, optional) – Axis or axes along which a standard deviations are calculated. The default is to compute the mean of the flattened array. If this is a tuple of ints, the standard deviations are calculated over multiple axes. The indicated axes must be ensemble axes.

• keepdims (bool, optional) – If True, the reduced axes are left in the result as dimensions with size one. Default is False.

• split_every (int) – Only used for lazy arrays. See dask.array.reductions.

Returns:

reduced_array – The reduced array object.

Return type:

ArrayObject or subclass of ArrayObject

property storage_device

The device on which the SMatrixArray is stored.

sum(axis=None, keepdims=False, split_every=2)

Sum of array object over one or more axes. Only ensemble axes can be reduced.

Parameters:

• axis (int or tuple of ints, optional) – Axis or axes along which a sums are performed. The default is to compute the mean of the flattened array. If this is a tuple of ints, the sum is performed over multiple axes. The indicated axes must be ensemble axes.

• keepdims (bool, optional) – If True, the reduced axes are left in the result as dimensions with size one. Default is False.

• split_every (int) – Only used for lazy arrays. See dask.array.reductions.

Returns:

reduced_array – The reduced array object.

Return type:

ArrayObject or subclass of ArrayObject

to_cpu()

Move the array to the host memory from an arbitrary source array.

Return type:

TypeVar(T, bound= ArrayObject)

to_data_array()

Convert ArrayObject to a xarray DataArray.

to_gpu(device='gpu')

Move the array from the host memory to a gpu.

Return type:

TypeVar(T, bound= ArrayObject)

to_hyperspy()

Convert ArrayObject to a Hyperspy signal.

to_tiff(filename, **kwargs)

Write data to a tiff file.

Parameters:

• filename (str) – The filename of the file to write.

• kwargs – Keyword arguments passed to tifffile.imwrite.

to_zarr(url, compute=True, overwrite=False, **kwargs)

Write data to a zarr file.

Parameters:

• url (str) – Location of the data, typically a path to a local file. A URL can also include a protocol specifier like s3:// for remote data.

• compute (bool) – If true compute immediately; return dask.delayed.Delayed otherwise.

• overwrite (bool) – If given array already exists, overwrite=False will cause an error, where overwrite=True will replace the existing data.

• kwargs – Keyword arguments passed to dask.array.to_zarr.

property wave_vectors: ndarray

The wave vectors corresponding to each plane wave.

property wavelength

Relativistic wavelength in Ångstrom.

property waves: Waves

The wave vectors describing each plane wave.

property window_extent: tuple[float, float]

The cropping window extent of the waves.

property window_gpts: tuple[int, int]

The number of grid points describing the cropping window of the wave functions.

property window_offset: tuple[float, float]

The number of grid points from the origin the cropping windows of the wave functions is displaced.