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])