SMatrix#

class abtem.prism.s_matrix.SMatrix(semiangle_cutoff, energy, potential=None, gpts=None, sampling=None, extent=None, interpolation=1, downsample='cutoff', device=None, store_on_host=False)[source]#

Bases: BaseSMatrix, Ensemble, CopyMixin, EqualityMixin

The scattering matrix is used for simulating STEM experiments using the PRISM algorithm.

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

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

  • potential (Atoms or AbstractPotential, optional) – Atoms or a potential that the scattering matrix represents. If given as atoms, a default potential will be created. If nothing is provided the scattering matrix will represent a vacuum potential, in which case the sampling and extent must be provided.

  • gpts (one or two int, optional) – Number of grid points describing the scattering matrix. Provide only if potential is not given.

  • sampling (one or two float, optional) – Lateral sampling of scattering matrix [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 scattering matrix [Å]. 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.

  • downsample ({'cutoff', 'valid'} or float or bool) –

    Controls whether to downsample the scattering matrix after running the multislice algorithm.

    cutoff :

    Downsample to the antialias cutoff scattering angle (default).

    valid :

    Downsample to the largest rectangle that fits inside the circle with a radius defined by the antialias cutoff scattering angle.

    float :

    Downsample to a specified maximum scattering angle [mrad].

  • 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.

  • store_on_host (bool, optional) – If True, store the scattering matrix in host (cpu) memory so that the necessary memory is transferred as chunks to the device to run calculations (default is False).

__init__(semiangle_cutoff, energy, potential=None, gpts=None, sampling=None, extent=None, interpolation=1, downsample='cutoff', device=None, store_on_host=False)[source]#

Methods

__init__(semiangle_cutoff, energy[, ...])

build([lazy, max_batch, bound])

Build the plane waves of the scattering matrix and propagate them through the potential using the multislice algorithm.

copy()

Make a copy.

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.

generate_blocks([chunks])

Generate chunks of the ensemble.

match_grid(other[, check_match])

Match the grid to another object with a Grid.

multislice([potential, lazy, max_batch])

type potential:

reduce([scan, detectors, ctf, ...])

Run the multislice algorithm, then 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.

round_gpts_to_interpolation()

Round the gpts of the SMatrix to the closest multiple of the interpolation factor.

scan([scan, detectors, ctf, ...])

Run the multislice algorithm, then 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.

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.

axes_metadata

List of AxisMetadata.

base_axes_metadata

List of AxisMetadata for the base axes in real space.

base_shape

Shape of the base axes of the SMatrix.

cutoff_angles

Scattering angles at the antialias cutoff [mrad].

device

The device where the S-Matrix is created and reduced.

downsample

How to downsample the scattering matrix after running the multislice algorithm.

downsampled_gpts

The gpts of the SMatrix after downsampling.

dtype

The datatype of waves.

energy

Electron acceleration energy in electron volts.

ensemble_axes_metadata

Axis metadata for each ensemble axis.

ensemble_shape

Shape of the SMatrix ensemble axes.

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

metadata

Metadata stored as a dictionary.

potential

The potential described by the SMatrix.

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

Plane-wave expansion cutoff.

shape

Shape of the SMatrix.

store_on_host

Store the SMatrix in host memory.

tilt

The small-angle tilt of applied to the Fresnel propagator [mrad].

wave_vectors

The wave vectors corresponding to each plane wave.

wavelength

Relativistic wavelength in Ångstrom.

window_extent

The cropping window extent of the waves.

window_gpts

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

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.

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_shape: tuple[int, int, int]#

Shape of the base axes of the SMatrix.

build(lazy=None, max_batch='auto', bound=None)[source]#

Build the plane waves of the scattering matrix and propagate them through the potential using the multislice algorithm.

Parameters:
  • lazy (bool, optional) – If True, create the wave functions lazily, otherwise, calculate instantly. If not given, defaults to the setting in the user configuration file.

  • max_batch (int or str, optional) – The number of expansion plane waves in each run of the multislice algorithm.

Returns:

s_matrix_array – The built scattering matrix.

Return type:

SMatrixArray

copy()#

Make a copy.

property cutoff_angles: tuple[float, float]#

Scattering angles at the antialias cutoff [mrad].

property device#

The device where the S-Matrix is created and reduced.

property downsample: str | bool#

How to downsample the scattering matrix after running the multislice algorithm.

property downsampled_gpts: tuple[int, int]#

The gpts of the SMatrix after downsampling.

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:
Returns:

dummy_probes

Return type:

Probes

property energy#

Electron acceleration energy in electron volts.

property ensemble_axes_metadata#

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_shape: tuple[int, ...]#

Shape of the SMatrix ensemble axes.

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

Extent of grid for each dimension in Ångstrom.

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.

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

match_grid(other, check_match=False)#

Match the grid to another object with a Grid.

property metadata#

Metadata stored as a dictionary.

multislice(potential=None, lazy=None, max_batch='auto')[source]#
Parameters:
  • potential

  • lazy (bool, optional) – If True, create the wave functions lazily, otherwise, calculate instantly. If not given, defaults to the setting in the user configuration file.

  • max_batch (int or str, optional) – The number of expansion plane waves in each run of the multislice algorithm.

property potential: BasePotential#

The potential described by the SMatrix.

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, detectors=None, ctf=None, reduction_scheme='auto', max_batch_multislice='auto', max_batch_reduction='auto', disable_s_matrix_chunks='auto', lazy=None)[source]#

Run the multislice algorithm, then 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_multislice (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”.

  • 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”.

  • reduction_scheme (str, 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.

  • disable_s_matrix_chunks (bool, optional) – If True, each S-Matrix is kept as a single chunk, thus lowering the communication overhead, but providing fewer opportunities for parallelization.

  • lazy (bool, optional) – If True, create the measurements lazily, otherwise, calculate instantly. If None, this defaults to the value set in the configuration file.

Returns:

measurements – The detected measurement (if detector(s) given).

Return type:

BaseMeasurements or Waves or list of BaseMeasurements or list of Waves

round_gpts_to_interpolation()[source]#

Round the gpts of the SMatrix to the closest multiple of the interpolation factor.

Returns:

s_matrix_with_rounded_gpts

Return type:

SMatrix

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_multislice='auto', max_batch_reduction='auto', reduction_scheme='auto', disable_s_matrix_chunks='auto', lazy=None)[source]#

Run the multislice algorithm, then 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_multislice (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”.

  • 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”.

  • reduction_scheme (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.

  • disable_s_matrix_chunks (bool, optional) – If True, each S-Matrix is kept as a single chunk, thus lowering the communication overhead, but providing fewer opportunities for parallelization.

  • lazy (bool, optional) – If True, create the measurements lazily, otherwise, calculate instantly. If None, this defaults to the value set in the configuration file.

Return type:

BaseMeasurements | Waves | list[BaseMeasurements | Waves]

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#

Plane-wave expansion cutoff.

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

Shape of the SMatrix.

property store_on_host: bool#

Store the SMatrix in host memory. The reduction may still be calculated on the device.

property tilt#

The small-angle tilt of applied to the Fresnel propagator [mrad].

property wave_vectors: ndarray#

The wave vectors corresponding to each plane wave.

property wavelength#

Relativistic wavelength in Ångstrom.

property window_extent#

The cropping window extent of the waves.

property window_gpts#

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