def scan(self,
scan: AbstractScan,
detectors: Union[AbstractDetector, Sequence[AbstractDetector]],
potential: Union[Atoms, AbstractPotential],
max_batch: int = 1,
pbar: bool = True) -> dict:
"""
Raster scan the probe across the potential and record a measurement for each detector.
:param scan: Scan object defining the positions of the probe wave functions.
:param detectors: The detectors recording the measurements.
:param potential: The potential across which to scan the probe .
:param max_batch: The probe batch size. Larger batches are faster, but require more memory.
:param pbar: If true, display progress bars.
:return: Dictionary of measurements with keys given by the detector.
"""
self.grid.match(potential.grid)
self.grid.check_is_defined()
if isinstance(detectors, AbstractDetector):
detectors = [detectors]
measurements = {}
for detector in detectors:
measurements[detector] = detector.allocate_measurement(
self.grid, self.wavelength, scan)
scan_bar = ProgressBar(total=len(scan), desc='Scan', disable=not pbar)
if isinstance(potential, AbstractTDSPotentialBuilder):
probe_generators = self._generate_tds_probes(
scan, potential, max_batch, pbar)
else:
if isinstance(potential, AbstractPotentialBuilder):
potential = potential.build(pbar=True)
probe_generators = [
self._generate_probes(scan, potential, max_batch)
]
for probe_generator in probe_generators:
scan_bar.reset()
for start, end, exit_probes in probe_generator:
for detector, measurement in measurements.items():
scan.insert_new_measurement(measurement, start, end,
detector.detect(exit_probes))
scan_bar.update(end - start)
scan_bar.refresh()
scan_bar.close()
return measurements
def multislice(self,
potential: AbstractPotential,
pbar: Union[ProgressBar, bool] = True) -> 'Waves':
"""
Propagate and transmit wave function through the provided potential.
:param potential: The potential through which to propagate the wave function.
:param pbar: If true, display a progress bar.
:return: Wave function at the exit plane of the potential.
"""
self.grid.match(potential)
propagator = FresnelPropagator()
if isinstance(potential, AbstractTDSPotentialBuilder):
xp = get_array_module(self.array)
N = len(potential.frozen_phonons)
out_array = xp.zeros((N, ) + self.array.shape, dtype=xp.complex64)
tds_waves = self.__class__(out_array,
extent=self.extent,
energy=self.energy)
tds_pbar = ProgressBar(total=N,
desc='TDS',
disable=(not pbar) or (N == 1))
multislice_pbar = ProgressBar(total=len(potential),
desc='Multislice',
disable=not pbar)
for i, potential_config in enumerate(
potential.generate_frozen_phonon_potentials(pbar=pbar)):
multislice_pbar.reset()
exit_waves = _multislice(copy(self),
potential_config,
propagator=propagator,
pbar=multislice_pbar)
tds_waves.array[i] = exit_waves.array
tds_pbar.update(1)
multislice_pbar.close()
tds_pbar.close()
return tds_waves
else:
return _multislice(self, potential, propagator, pbar)
def _run_epie(object,
probe: np.ndarray,
diffraction_patterns: np.ndarray,
positions: np.ndarray,
maxiter: int,
alpha: float = 1.,
beta: float = 1.,
fix_probe: bool = False,
fix_com: bool = False,
return_iterations: bool = False,
seed=None):
xp = get_array_module(probe)
object = xp.array(object)
probe = xp.array(probe)
if len(diffraction_patterns.shape) != 3:
raise ValueError()
if len(diffraction_patterns) != len(positions):
raise ValueError()
if object.shape == (2, ):
object = xp.ones((int(object[0]), int(object[1])), dtype=xp.complex64)
elif len(object.shape) != 2:
raise ValueError()
if probe.shape != diffraction_patterns.shape[1:]:
raise ValueError()
if probe.shape != object.shape:
raise ValueError()
if return_iterations:
object_iterations = []
probe_iterations = []
SSE_iterations = []
if seed is not None:
np.random.seed(seed)
diffraction_patterns = np.fft.ifftshift(np.sqrt(diffraction_patterns),
axes=(-2, -1))
SSE = 0.
k = 0
outer_pbar = ProgressBar(total=maxiter)
inner_pbar = ProgressBar(total=len(positions))
while k < maxiter:
indices = np.arange(len(positions))
np.random.shuffle(indices)
old_position = xp.array((0., 0.))
inner_pbar.reset()
SSE = 0.
for j in indices:
position = xp.array(positions[j])
diffraction_pattern = xp.array(diffraction_patterns[j])
illuminated_object = fft_shift(object, old_position - position)
g = illuminated_object * probe
gprime = xp.fft.ifft2(diffraction_pattern *
xp.exp(1j * xp.angle(xp.fft.fft2(g))))
object = illuminated_object + alpha * (
gprime - g) * xp.conj(probe) / (xp.max(xp.abs(probe))**2)
old_position = position
if not fix_probe:
probe = probe + beta * (
gprime - g) * xp.conj(illuminated_object) / (xp.max(
xp.abs(illuminated_object))**2)
# SSE += xp.sum(xp.abs(G) ** 2 - diffraction_pattern) ** 2
inner_pbar.update(1)
object = fft_shift(object, position)
if fix_com:
com = center_of_mass(xp.fft.fftshift(xp.abs(probe)**2))
probe = xp.fft.ifftshift(fft_shift(probe, -xp.array(com)))
# SSE = SSE / np.prod(diffraction_patterns.shape)
if return_iterations:
object_iterations.append(object)
probe_iterations.append(probe)
SSE_iterations.append(SSE)
outer_pbar.update(1)
# if verbose:
# print(f'Iteration {k:<{len(str(maxiter))}}, SSE = {float(SSE):.3e}')
k += 1
inner_pbar.close()
outer_pbar.close()
if return_iterations:
return object_iterations, probe_iterations, SSE_iterations
else:
return object, probe, SSE
def scan(self,
potential: Union[Atoms, AbstractPotential],
scan: AbstractScan,
detectors: Sequence[AbstractDetector],
max_batch_probes: int = 1,
max_batch_expansion: int = None,
pbar: bool = True):
self.grid.match(potential.grid)
self.grid.check_is_defined()
measurements = {}
for detector in detectors:
measurements[detector] = detector.allocate_measurement(
self.interpolated_grid, self.wavelength, scan)
if isinstance(potential, AbstractTDSPotentialBuilder):
probe_generators = self._generate_tds_probes(
scan,
potential,
max_batch_probes=max_batch_probes,
max_batch_expansion=max_batch_expansion,
potential_pbar=pbar,
multislice_pbar=pbar)
else:
if isinstance(potential, AbstractPotentialBuilder):
potential = potential.build(pbar=True)
S = self.multislice(potential,
max_batch=max_batch_probes,
pbar=pbar)
probe_generators = [
S._generate_probes(scan, max_batch_probes, max_batch_expansion)
]
tds_bar = ProgressBar(total=len(potential.frozen_phonons),
desc='TDS',
disable=(not pbar)
or (len(potential.frozen_phonons) == 1))
scan_bar = ProgressBar(total=len(scan), desc='Scan', disable=not pbar)
for probe_generator in probe_generators:
scan_bar.reset()
for start, end, exit_probes in probe_generator:
for detector, measurement in measurements.items():
scan.insert_new_measurement(measurement, start, end,
detector.detect(exit_probes))
scan_bar.update(end - start)
scan_bar.refresh()
tds_bar.update(1)
scan_bar.close()
tds_bar.refresh()
tds_bar.close()
return measurements