def allocate_measurement(self, grid: Grid, wavelength: float,
scan: AbstractScan) -> Measurement:
inner, outer, nbins_radial, nbins_azimuthal = self._get_bins(
grid.antialiased_sampling, wavelength)
shape = scan.shape
calibrations = scan.calibrations
if nbins_radial > 1:
shape += (nbins_radial, )
calibrations += (Calibration(offset=inner,
sampling=self._radial_steps,
units='mrad'), )
if nbins_azimuthal > 1:
shape += (nbins_azimuthal, )
calibrations += (Calibration(offset=0,
sampling=self._azimuthal_steps,
units='rad'), )
array = np.zeros(shape, dtype=np.float32)
measurement = Measurement(array, calibrations=calibrations)
if isinstance(self.save_file, str):
measurement = measurement.write(self.save_file)
return measurement
def poisson_noise(measurement: Measurement, dose: float):
"""
Add Poisson noise to a measurment.
Parameters
----------
measurement: Measurement object
The measurement to add noise to.
dose: float
The irradiation dose in electrons per Å^2.
Returns
-------
measurement: Measurement object
The noisy measurement.
"""
pixel_area = np.product([calibration.sampling for calibration in measurement.calibrations[:2]])
measurement = measurement.copy()
array = measurement.array
electrons_per_pixel = dose * pixel_area
measurement.array[:] = array * electrons_per_pixel
measurement.array[:] = np.random.poisson(array).astype(np.float)
return measurement
def allocate_measurement(self, waves, scan: AbstractScan) -> Measurement:
"""
Allocate a Measurement object or an hdf5 file.
Parameters
----------
waves : Waves or SMatrix object
The wave function that will define the shape of the diffraction patterns.
scan: Scan object
The scan object that will define the scan dimensions the measurement.
Returns
-------
Measurement object or str
The allocated measurement or path to hdf5 file with the measurement data.
"""
waves.grid.check_is_defined()
calibrations = calibrations_from_grid(waves.gpts,
waves.sampling,
names=['x', 'y'],
units='Å')
array = np.zeros(scan.shape + waves.gpts, dtype=np.complex64)
measurement = Measurement(array,
calibrations=scan.calibrations +
calibrations)
if isinstance(self.save_file, str):
measurement = measurement.write(self.save_file)
return measurement
def allocate_measurement(self,
waves,
scan: AbstractScan = None) -> Measurement:
"""
Allocate a Measurement object or an hdf5 file.
Parameters
----------
waves : Waves or SMatrix object
The wave function that will define the shape of the diffraction patterns.
scan: Scan object
The scan object that will define the scan dimensions the measurement.
Returns
-------
Measurement object or str
The allocated measurement or path to hdf5 file with the measurement data.
"""
waves.grid.check_is_defined()
waves.accelerator.check_is_defined()
check_max_angle_exceeded(waves, self.max_angle)
gpts = waves.downsampled_gpts(self.max_angle)
gpts, new_angular_sampling = self._resampled_gpts(
gpts, angular_sampling=waves.angular_sampling)
sampling = (1 / new_angular_sampling[0] / gpts[0] * waves.wavelength *
1000, 1 / new_angular_sampling[1] / gpts[1] *
waves.wavelength * 1000)
calibrations = calibrations_from_grid(gpts,
sampling,
names=['alpha_x', 'alpha_y'],
units='mrad',
scale_factor=waves.wavelength *
1000,
fourier_space=True)
if scan is None:
scan_shape = ()
scan_calibrations = ()
elif isinstance(scan, tuple):
scan_shape = scan
scan_calibrations = (None, ) * len(scan)
else:
scan_shape = scan.shape
scan_calibrations = scan.calibrations
array = np.zeros(scan_shape + gpts)
measurement = Measurement(array,
calibrations=scan_calibrations +
calibrations)
if isinstance(self.save_file, str):
measurement = measurement.write(self.save_file)
return measurement
def allocate_measurement(self,
waves,
scan: AbstractScan = None) -> Measurement:
"""
Allocate a Measurement object or an hdf5 file.
Parameters
----------
waves : Waves object
An example of the
scan : Scan object
The scan object that will define the scan dimensions the measurement.
Returns
-------
Measurement object or str
The allocated measurement or path to hdf5 file with the measurement data.
"""
waves.grid.check_is_defined()
waves.accelerator.check_is_defined()
if scan is None:
shape = ()
calibrations = ()
else:
shape = scan.shape
calibrations = scan.calibrations
nbins_radial, nbins_azimuthal, inner, _ = self._get_bins(
min(waves.cutoff_scattering_angles))
if nbins_radial > 1:
shape += (nbins_radial, )
calibrations += (Calibration(offset=inner,
sampling=self._radial_steps,
units='mrad'), )
if nbins_azimuthal > 1:
shape += (nbins_azimuthal, )
calibrations += (Calibration(offset=0,
sampling=self._azimuthal_steps,
units='rad'), )
array = np.zeros(shape, dtype=np.float32)
measurement = Measurement(array, calibrations=calibrations)
if isinstance(self.save_file, str):
measurement = measurement.write(self.save_file)
return measurement
def test_export_import_measurement(tmp_path):
d = tmp_path / 'sub'
d.mkdir()
path = d / 'measurement.hdf5'
calibrations = calibrations_from_grid((512, 256), (.1, .3), ['x', 'y'],
'Å')
measurement = Measurement(np.random.rand(512, 256), calibrations)
measurement.write(path)
imported_measurement = Measurement.read(path)
assert np.allclose(measurement.array, imported_measurement.array)
assert measurement.calibrations[0] == imported_measurement.calibrations[0]
assert measurement.calibrations[1] == imported_measurement.calibrations[1]
def allocate_measurement(self, grid: Grid, wavelength: float,
scan: AbstractScan) -> Measurement:
grid.check_is_defined()
calibrations = calibrations_from_grid(grid.gpts,
grid.sampling,
names=['x', 'y'],
units='Å')
array = np.zeros(scan.shape + grid.gpts, dtype=np.complex64)
measurement = Measurement(array,
calibrations=scan.calibrations +
calibrations)
if isinstance(self.save_file, str):
measurement = measurement.write(self.save_file)
return measurement
def poisson_noise(measurement: Measurement,
dose: float,
pixel_area: float = None,
negative_values='clip'):
"""
Add Poisson noise to a measurment.
Parameters
----------
measurement: Measurement object
The measurement to add noise to.
dose: float
The irradiation dose in electrons per Å^2.
pixel_area: float, optional
Pixel area in Å^2. If not provided this will be calculated from the calibraions when possible.
Returns
-------
measurement: Measurement object
The noisy measurement.
"""
if pixel_area is None:
pixel_areas = []
for calibration in measurement.calibrations:
if calibration is not None:
if calibration.units.lower() in ('angstrom', 'å'):
pixel_areas.append(calibration.sampling)
if len(pixel_areas) != 2:
raise RuntimeError(
'Real space pixel size not recognized from calibrations.')
pixel_area = np.product(pixel_areas)
measurement = measurement.copy()
array = measurement.array
if negative_values == 'clip':
array = np.clip(array, a_min=1e-12, a_max=None)
elif negative_values != 'raise':
if np.any(array < 0.):
raise ValueError('Measurement values must be positive.')
electrons_per_pixel = dose * pixel_area
array = array * electrons_per_pixel
measurement.array[:] = np.random.poisson(array).astype(np.float)
return measurement
def test_gridscan_to_file(tmp_path):
d = tmp_path / 'sub'
d.mkdir()
path = d / 'measurement2.hdf5'
atoms = read(_set_path('orthogonal_graphene.cif'))
potential = Potential(atoms=atoms, sampling=.05)
probe = Probe(energy=200e3, semiangle_cutoff=30)
probe.grid.match(potential)
scan = GridScan(start=[0, 0], end=[0, potential.extent[1]], gpts=(10, 9))
detector = PixelatedDetector()
export_detector = PixelatedDetector(save_file=path)
measurements = probe.scan(scan, [detector, export_detector],
potential,
pbar=False)
measurement = measurements[0]
imported_measurement = Measurement.read(measurements[1])
assert np.allclose(measurement.array, imported_measurement.array)
assert measurement.calibrations[0] == imported_measurement.calibrations[0]
assert measurement.calibrations[1] == imported_measurement.calibrations[1]
def integrate(self, diffraction_patterns: Measurement) -> Measurement:
"""
Integrate diffraction pattern measurements on the detector region.
Parameters
----------
diffraction_patterns : 2d, 3d or 4d Measurement object
The collection diffraction patterns to be integrated.
Returns
-------
Measurement
"""
if diffraction_patterns.dimensions < 2:
raise ValueError()
if not (diffraction_patterns.calibrations[-1].units == diffraction_patterns.calibrations[-2].units):
raise ValueError()
sampling = (diffraction_patterns.calibrations[-2].sampling, diffraction_patterns.calibrations[-1].sampling)
calibrations = diffraction_patterns.calibrations[:-2]
array = np.fft.ifftshift(diffraction_patterns.array, axes=(-2, -1))
cutoff_scattering_angle = min(diffraction_patterns.calibrations[-2].sampling *
diffraction_patterns.array.shape[-2],
diffraction_patterns.calibrations[-1].sampling *
diffraction_patterns.array.shape[-1], )
return Measurement(self._integrate_array(array, sampling, cutoff_scattering_angle), calibrations=calibrations)
def show(self, waves, **kwargs):
"""
Visualize the detector region(s) of the detector as applied to a specified wave function.
Parameters
----------
waves : Waves or SMatrix object
The wave function the visualization will be created to match
kwargs :
Additional keyword arguments for abtem.visualize.mpl.show_measurement_2d.
"""
waves.grid.check_is_defined()
array = np.full(waves.gpts, -1, dtype=np.int)
for i, indices in enumerate(
self._get_regions(waves.gpts, waves.angular_sampling,
min(waves.cutoff_scattering_angles))):
array.ravel()[indices] = i
calibrations = calibrations_from_grid(waves.gpts,
waves.sampling,
names=['alpha_x', 'alpha_y'],
units='mrad',
scale_factor=waves.wavelength *
1e3,
fourier_space=True)
array = np.fft.fftshift(array, axes=(-1, -2))
measurement = Measurement(array,
calibrations=calibrations,
name='Detector regions')
return show_measurement_2d(measurement, discrete_cmap=True, **kwargs)
def allocate_measurement(self, grid: Grid, wavelength: float,
scan: AbstractScan) -> Measurement:
grid.check_is_defined()
shape = (grid.gpts[0] // 2, grid.gpts[1] // 2)
calibrations = calibrations_from_grid(grid.antialiased_gpts,
grid.antialiased_sampling,
names=['alpha_x', 'alpha_y'],
units='mrad',
scale_factor=wavelength * 1000,
fourier_space=True)
array = np.zeros(scan.shape + shape)
measurement = Measurement(array,
calibrations=scan.calibrations +
calibrations)
if isinstance(self.save_file, str):
measurement = measurement.write(self.save_file)
return measurement
def intensity(self) -> Measurement:
"""
:return: The intensity of the wave functions at the image plane.
"""
calibrations = calibrations_from_grid(self.grid.gpts,
self.grid.sampling, ['x', 'y'])
calibrations = (None, ) * (len(self.array.shape) - 2) + calibrations
abs2 = get_device_function(get_array_module(self.array), 'abs2')
return Measurement(abs2(self.array), calibrations)
def project(self):
"""
Create a 2d measurement of the projected potential.
Returns
-------
Measurement
"""
calibrations = calibrations_from_grid(self.grid.gpts,
self.grid.sampling,
names=['x', 'y'])
array = asnumpy(self.array.sum(0))
array -= array.min()
return Measurement(array, calibrations)
def add_scan_noise(measurement: Measurement,
dwell_time: float,
flyback_time: float,
max_frequency: float,
rms_power: float,
num_components: int = 200):
"""
Add scan noise to a measurement.
Parameters
----------
measurement: Measurement object or 2d array
The measurement to add noise to.
dwell_time: float
Dwell time on a single pixel in s.
flyback_time: float
Flyback time for the scanning probe at the end of each scan line in s.
max_frequency: float
Maximum noise frequency in 1 / s.
rms_power: float
Root-mean-square power of the distortion in unit of percent.
num_components: int, optional
Number of frequency components. More components will be more 'white' but will take longer.
Returns
-------
measurement: Measurement object
The noisy measurement.
"""
measurement = measurement.copy()
if isinstance(measurement, Measurement):
array = measurement.array
else:
array = measurement
time = _pixel_times(dwell_time, flyback_time, array.T.shape)
displacement_x, displacement_y = _make_displacement_field(
time, max_frequency, num_components, rms_power)
array[:] = _apply_displacement_field(array.T, displacement_x,
displacement_y).T
return measurement
def diffraction_pattern(self) -> Measurement:
"""
:return: The intensity of the wave functions at the diffraction plane.
"""
calibrations = calibrations_from_grid(self.grid.antialiased_gpts,
self.grid.antialiased_sampling,
names=['alpha_x', 'alpha_y'],
units='mrad',
scale_factor=self.wavelength *
1000,
fourier_space=True)
calibrations = (None, ) * (len(self.array.shape) - 2) + calibrations
xp = get_array_module(self.array)
abs2 = get_device_function(xp, 'abs2')
fft2 = get_device_function(xp, 'fft2')
pattern = asnumpy(
abs2(
crop_to_center(
xp.fft.fftshift(fft2(self.array, overwrite_x=False)))))
return Measurement(pattern, calibrations)
def show(self, transitions_idx=0):
intensity = None
if self._sliced_atoms.slice_thicknesses is None:
none_slice_thickess = True
self._sliced_atoms.slice_thicknesses = self._sliced_atoms.atoms.cell[
2, 2]
else:
none_slice_thickess = False
for slice_idx in range(self.num_slices):
for t in self._generate_slice_transition_potentials(
slice_idx, transitions_idx):
if intensity is None:
intensity = np.abs(t)**2
else:
intensity += np.abs(t)**2
if none_slice_thickess:
self._sliced_atoms.slice_thicknesses = None
calibrations = calibrations_from_grid(self.gpts, self.sampling,
['x', 'y'])
Measurement(intensity[0], calibrations, name=str(self)).show()
def profiles(self, max_semiangle: float = None, phi: float = 0.):
if max_semiangle is None:
if self._semiangle_cutoff == np.inf:
max_semiangle = 50
else:
max_semiangle = self._semiangle_cutoff * 1.6
alpha = np.linspace(0, max_semiangle / 1000., 500)
aberrations = self.evaluate_aberrations(alpha, phi)
aperture = self.evaluate_aperture(alpha)
temporal_envelope = self.evaluate_temporal_envelope(alpha)
spatial_envelope = self.evaluate_spatial_envelope(alpha, phi)
gaussian_envelope = self.evaluate_gaussian_envelope(alpha)
envelope = aperture * temporal_envelope * spatial_envelope * gaussian_envelope
calibration = Calibration(offset=0.,
sampling=(alpha[1] - alpha[0]) * 1000.,
units='mrad',
name='alpha')
profiles = {}
profiles['ctf'] = Measurement(aberrations.imag * envelope,
calibrations=[calibration],
name='CTF')
profiles['aperture'] = Measurement(aperture,
calibrations=[calibration],
name='Aperture')
profiles['temporal_envelope'] = Measurement(temporal_envelope,
calibrations=[calibration],
name='Temporal')
profiles['spatial_envelope'] = Measurement(spatial_envelope,
calibrations=[calibration],
name='Spatial')
profiles['gaussian_spread'] = Measurement(gaussian_envelope,
calibrations=[calibration],
name='Gaussian')
profiles['envelope'] = Measurement(envelope,
calibrations=[calibration],
name='Envelope')
return profiles
def allocate_measurement(self, waves, scan):
array = np.zeros(scan.shape, dtype=np.float32)
return Measurement(array, calibrations=scan.calibrations)
from abtem.measure import Measurement
import matplotlib.pyplot as plt
measurement = Measurement.read('STEM_MoS2.hdf5')
measurement.tile((5, 3)).interpolate(.01).show()
plt.show()
def measure(self):
array = np.fft.fftshift(self.build())[0]
calibrations = calibrations_from_grid(self.gpts, self.sampling,
['x', 'y'])
abs2 = get_device_function(get_array_module(array), 'abs2')
return Measurement(array, calibrations, name=str(self))
def epie(
measurement: Measurement,
probe_guess: Probe,
maxiter: int = 5,
alpha: float = 1.,
beta: float = 1.,
fix_probe: bool = False,
fix_com: bool = False,
return_iterations: bool = False,
max_angle=None,
seed=None,
device='cpu',
):
"""
Reconstruct the phase of a 4D-STEM measurement using the extended Ptychographical Iterative Engine.
See https://doi.org/10.1016/j.ultramic.2009.05.012
Parameters
----------
measurement : Measurement object
4D-STEM measurement.
probe_guess : Probe object
The initial guess for the probe.
maxiter : int
Run the algorithm for this many iterations.
alpha : float
Controls the size of the iterative updates for the object. See reference.
beta : float
Controls the size of the iterative updates for the probe. See reference.
fix_probe : bool
If True, the probe will not be updated by the algorithm. Default is False.
fix_com : bool
If True, the center of mass of the probe will be centered. Default is True.
return_iterations : bool
If True, return the reconstruction after every iteration. Default is False.
max_angle : float, optional
The maximum reconstructed scattering angle. If this is larger than the input data, the data will be zero-padded.
seed : int, optional
Seed the random number generator.
device : str
Set the calculation device.
Returns
-------
List of Measurement objects
"""
diffraction_patterns = measurement.array.reshape(
(-1, ) + measurement.array.shape[2:])
if max_angle:
padding_x = int((max_angle / abs(measurement.calibrations[-2].offset) *
diffraction_patterns.shape[-2]) //
2) - diffraction_patterns.shape[-2] // 2
padding_y = int((max_angle / abs(measurement.calibrations[-1].offset) *
diffraction_patterns.shape[-1]) //
2) - diffraction_patterns.shape[-1] // 2
diffraction_patterns = np.pad(diffraction_patterns,
((0, ) * 2, (padding_x, ) * 2,
(padding_y, ) * 2))
extent = (probe_guess.wavelength * 1e3 /
measurement.calibrations[2].sampling, probe_guess.wavelength *
1e3 / measurement.calibrations[3].sampling)
sampling = (extent[0] / diffraction_patterns.shape[-2],
extent[1] / diffraction_patterns.shape[-1])
x = measurement.calibrations[0].coordinates(
measurement.shape[0]) / sampling[0]
y = measurement.calibrations[1].coordinates(
measurement.shape[1]) / sampling[1]
x, y = np.meshgrid(x, y, indexing='ij')
positions = np.array([x.ravel(), y.ravel()]).T
probe_guess.extent = extent
probe_guess.gpts = diffraction_patterns.shape[-2:]
calibrations = calibrations_from_grid(probe_guess.gpts,
probe_guess.sampling,
names=['x', 'y'],
units='Å')
probe_guess._device = device
probe_guess = probe_guess.build(np.array([0, 0])).array[0]
result = _run_epie(diffraction_patterns.shape[-2:],
probe_guess,
diffraction_patterns,
positions,
maxiter=maxiter,
alpha=alpha,
beta=beta,
return_iterations=return_iterations,
fix_probe=fix_probe,
fix_com=fix_com,
seed=seed)
if return_iterations:
object_iterations = [
Measurement(object, calibrations=calibrations)
for object in result[0]
]
probe_iterations = [
Measurement(np.fft.fftshift(probe), calibrations=calibrations)
for probe in result[1]
]
return object_iterations, probe_iterations, result[2]
else:
return (Measurement(result[0], calibrations=calibrations),
Measurement(np.fft.fftshift(result[1]),
calibrations=calibrations), result[2])
