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)[source]#

Bases: abtem.prism.s_matrix.BaseSMatrix, abtem.array.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)[source]#

Methods

`__init__`(array, wave_vectors, semiangle_cutoff)
`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_items`(items[, keepdims])	Index the array and the corresponding axes metadata.
`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 `Optional`[`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_gpu`([device])	Move the array from the host memory to a gpu.
`to_hyperspy`()	Convert measurement 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: abtem.core.energy.Accelerator#

Accelerator object describing the acceleration energy.

Return type: Accelerator

property angular_sampling: tuple[float, float]#

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

Return type: tuple[float, float]

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.

Return type: tuple[int, int]

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.

Return type: tuple[int, int]

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: np.ndarray | da.core.Array#

Underlying array describing the array object.

Return type: np.ndarray | da.core.Array

property axes_metadata: abtem.core.axes.AxesMetadataList#

List of AxisMetadata.

Return type: AxesMetadataList

property base_axes_metadata: list[abtem.core.axes.AxisMetadata]#

List of AxisMetadata for the base axes in real space.

Return type: list[AxisMetadata]

property base_dims#: Number of base dimensions.

property base_shape: tuple[int, ...]#

Shape of the base axes of the underlying array.

Return type: tuple[int, ...]

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)[source]#

Copy SMatrixArray to specified device.

Return type: SMatrixArray

property cutoff_angles: tuple[float, float]#

Scattering angles at the antialias cutoff [mrad].

Return type: tuple[float, float]

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[abtem.core.axes.AxisMetadata]#

Axis metadata for each ensemble axis.

Return type: list[AxisMetadata]

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.

Return type: tuple[int, int]

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.

Return type: tuple[float] | tuple[float, float] | tuple[float, …]

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.

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.

Return type: TypeVar(T, bound= ArrayObject)

property full_cutoff_angles: tuple[float, float]#

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

Return type: tuple[float, float]

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.

Return type: tuple[int] | tuple[int, int] | tuple[int, …]

property grid: abtem.core.grid.Grid#

Simulation grid.

Return type: Grid

property interpolation: tuple[int, int]#

Interpolation factor in the x and y directions

Return type: tuple[int, int]

property is_complex: bool#

True if array is complex.

Return type: bool

property is_lazy: bool#

True if array is lazy.

Return type: bool

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.

Return type: dict

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)[source]#

Parameters: potential (Optional[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.

Return type: tuple[bool, bool]

rechunk(chunks, in_place=True)[source]#

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[abtem.core.axes.AxisMetadata]#

List of AxisMetadata for base axes in reciprocal space.

Return type: list[AxisMetadata]

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

Reciprocal-space sampling in reciprocal Ångstrom.

Return type: tuple[float] | tuple[float, float] | tuple[float, …]

property rectangle_cutoff_angles: tuple[float, float]#

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

Return type: tuple[float, float]

reduce(scan=None, ctf=None, detectors=None, max_batch_reduction='auto', reduction_scheme='auto')[source]#

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.

Return type: tuple[float] | tuple[float, float] | tuple[float, …]

scan(scan=None, detectors=None, ctf=None, max_batch_reduction='auto', rechunk='auto')[source]#

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.

Return type: float

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.

Return type: tuple[int, ...]

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_gpu(device='gpu')#

Move the array from the host memory to a gpu.

Return type: TypeVar(T, bound= ArrayObject)

to_hyperspy()#: Convert measurement 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: numpy.ndarray#

The wave vectors corresponding to each plane wave.

Return type: ndarray

property wavelength#: Relativistic wavelength in Ångstrom.

property waves: abtem.waves.Waves#

The wave vectors describing each plane wave.

Return type: Waves

property window_extent: tuple[float, float]#

The cropping window extent of the waves.

Return type: tuple[float, float]

property window_gpts: tuple[int, int]#

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

Return type: tuple[int, int]

property window_offset: tuple[float, float]#

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

Return type: tuple[float, float]