def test_probe_waves_line_scan():
probe = Probe(energy=60e3)
detector = DummyDetector()
potential = DummyPotential(extent=5, sampling=.1)
scan = LineScan((0, 0), (1, 1), gpts=10)
measurement = probe.scan(scan,
detector,
potential,
max_batch=1,
pbar=False)
assert detector._detect_count == 10
assert np.all(measurement.array == 1.)
measurement = probe.scan(scan, [detector], potential, pbar=False)
assert detector._detect_count == 11
assert np.all(measurement.array == 1.)
measurement = probe.scan(scan,
detector,
potential,
max_batch=3,
pbar=False)
assert detector._detect_count == 15
assert np.all(measurement.array == 1.)
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 test_fig():
atoms = Atoms('CSiCuAuU',
positions=[(x, 25, 4) for x in np.linspace(5, 45, 5)],
cell=(50, 50, 8))
gpts = 2048
potential = Potential(atoms=atoms,
gpts=gpts,
parametrization='kirkland',
slice_thickness=8)
probe = Probe(energy=200e3, defocus=700, Cs=1.3e7, semiangle_cutoff=10.37)
probe.grid.match(potential)
scan = LineScan(start=[5, 25], end=[45, 25], gpts=5)
detector = AnnularDetector(inner=40, outer=200)
measurements = probe.scan(scan, [detector], potential, pbar=False)
#assert np.allclose(measurements[detector].array, [0.00010976, 0.00054356, 0.00198158, 0.00997221, 0.01098883])
assert np.allclose(
measurements[detector].array,
[0.0001168, 0.00059303, 0.00214667, 0.00977803, 0.01167613],
atol=1e-5)
def test_interpolation_scan():
atoms = Atoms('C', positions=[(2.5, 2.5, 2)], cell=(5, 5, 4))
potential = Potential(atoms)
linescan = LineScan(start=[0, 0], end=[2.5, 2.5], gpts=10)
detector = AnnularDetector(inner=80, outer=200)
probe = Probe(semiangle_cutoff=30, energy=80e3, gpts=250)
measurements = probe.scan(linescan, [detector],
potential,
max_batch=50,
pbar=False)
S_builder = SMatrix(semiangle_cutoff=30.,
energy=80e3,
interpolation=2,
gpts=500)
atoms = Atoms('C', positions=[(2.5, 2.5, 2)], cell=(5, 5, 4))
atoms *= (2, 2, 1)
potential = Potential(atoms)
S = S_builder.multislice(potential, pbar=False)
prism_measurements = S.scan(linescan,
detector,
max_batch_probes=10,
pbar=False)
assert np.allclose(measurements.array, prism_measurements.array, atol=1e-6)
def test_probe_waves_line_scan():
atoms = Atoms('CO',
positions=[(2.5, 2.5, 2), (2.5, 2.5, 3)],
cell=(5, 5, 4))
frozen_phonons = FrozenPhonons(atoms, 2, sigmas={'C': 0, 'O': 0.})
potential = Potential(atoms, sampling=.05)
tds_potential = Potential(frozen_phonons, sampling=.05)
linescan = LineScan(start=[0, 0], end=[2.5, 2.5], gpts=10)
detector = AnnularDetector(inner=80, outer=200)
probe = Probe(semiangle_cutoff=30, energy=80e3, gpts=500)
measurements = probe.scan(linescan, [detector],
potential,
max_batch=50,
pbar=False)
tds_measurements = probe.scan(linescan, [detector],
tds_potential,
max_batch=50,
pbar=False)
assert np.allclose(measurements[detector].array,
tds_measurements[detector].array,
atol=1e-6)
frozen_phonons = FrozenPhonons(atoms, 2, sigmas={'C': 0, 'O': 0.1})
tds_potential = Potential(frozen_phonons, sampling=.05)
tds_measurements = probe.scan(linescan, [detector],
tds_potential,
max_batch=50,
pbar=False)
assert not np.allclose(measurements[detector].array,
tds_measurements[detector].array,
atol=1e-6)
def test_flexible_annular_detector():
atoms = read('data/srtio3_100.cif')
atoms *= (4, 4, 1)
potential = Potential(
atoms,
gpts=512,
projection='infinite',
slice_thickness=.5,
parametrization='kirkland',
).build(pbar=False)
probe = Probe(energy=300e3, semiangle_cutoff=9.4, rolloff=0.05)
flexible_detector = FlexibleAnnularDetector()
annular_detector = AnnularDetector(inner=30, outer=80)
end = (potential.extent[0] / 4, potential.extent[1] / 4)
gridscan = GridScan(start=[0, 0], end=end, sampling=.2)
measurements = probe.scan(gridscan, [flexible_detector, annular_detector],
potential,
pbar=False)
assert np.allclose(measurements[flexible_detector].integrate(30, 80).array,
measurements[annular_detector].array)
def test_partition_measurement():
atoms = read(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
'data/amorphous_carbon.cif'))
potential = Potential(atoms,
gpts=256,
slice_thickness=1,
projection='infinite',
parametrization='kirkland').build(pbar=False)
detector = AnnularDetector(inner=70, outer=100)
gridscan = GridScan(start=[0, 0], end=potential.extent, gpts=4)
probe = Probe(semiangle_cutoff=15, energy=300e3)
measurements = probe.scan(gridscan, detector, potential, pbar=False)
scans = gridscan.partition_scan((2, 2))
partitioned_measurements = detector.allocate_measurement(probe, gridscan)
for scan in scans:
probe.scan(scan,
detector,
potential,
measurements=partitioned_measurements,
pbar=False)
assert np.allclose(partitioned_measurements.array, measurements.array)
def test_prism_translate():
S = SMatrix(30, 60e3, 1, extent=5, gpts=50)
probe = Probe(extent=5,
gpts=50,
energy=60e3,
semiangle_cutoff=30,
rolloff=0)
assert np.allclose(probe.build(np.array([(2.5, 2.5)])).array,
S.build().collapse([(2.5, 2.5)]).array,
atol=1e-5)
def test_create_probe_waves():
probe_waves = Probe(extent=10, gpts=10, energy=60e3, defocus=10, C32=30, **{'C10': 100})
waves = probe_waves.build()
assert waves.array.shape == (10, 10)
waves = probe_waves.build([[0, 0], [1, 1]])
assert waves.array.shape == (2, 10, 10)
def test_prism_multislice():
potential = Potential(Atoms('C', positions=[(0, 0, 2)], cell=(5, 5, 4)))
S = SMatrixBuilder(30, 1, extent=5, gpts=500, energy=60e3).build()
probe = Probe(extent=5, gpts=500, energy=60e3, semiangle_cutoff=30)
assert np.allclose(probe.build(np.array([[2.5, 2.5]
])).multislice(potential,
pbar=False).array,
S.multislice(potential,
pbar=False).collapse([(2.5, 2.5)]).array,
atol=2e-5)
def test_prism_match_probe():
S = SMatrix(30., 60e3, 1, extent=5, gpts=50)
probe = Probe(extent=5,
gpts=50,
energy=60e3,
semiangle_cutoff=30.,
rolloff=0.)
assert np.allclose(probe.build([(0., 0.)]).array,
S.build().collapse([(0., 0.)]).array,
atol=2e-5)
def test_prism_interpolation():
S_builder = SMatrix(30, 60e3, 2, extent=10, gpts=100)
probe = Probe(extent=5,
gpts=50,
energy=60e3,
semiangle_cutoff=30,
rolloff=0)
assert np.allclose(probe.build(np.array([(2.5, 2.5)])).array,
S_builder.build().collapse([(2.5, 2.5)]).array,
atol=1e-5)
def test_probe_waves_grid_scan():
probe = Probe(energy=60e3)
detector = DummyDetector()
potential = DummyPotential(extent=5, sampling=.1)
scan = GridScan((0, 0), (1, 1), gpts=10)
measurements = probe.scan(scan, [detector], potential, pbar=False)
assert detector._detect_count == 100
assert np.all(measurements[detector].array == 1.)
def test_export_import_waves(tmp_path):
d = tmp_path / 'sub'
d.mkdir()
path = d / 'waves.hdf5'
waves = Probe(semiangle_cutoff=30, sampling=.05, extent=10,
energy=80e3).build()
waves.write(path)
imported_waves = Waves.read(path)
assert np.allclose(waves.array, imported_waves.array)
assert np.allclose(waves.extent, imported_waves.extent)
assert np.allclose(waves.energy, imported_waves.energy)
def test_prism_interpolation():
S = SMatrix(semiangle_cutoff=30.,
energy=60e3,
interpolation=2,
extent=10,
gpts=100,
rolloff=0.)
probe = Probe(extent=5,
gpts=50,
energy=60e3,
semiangle_cutoff=30,
rolloff=0)
probe_array = probe.build(np.array([(2.5, 2.5)])).array
S_array = S.build().collapse([(2.5, 2.5)]).array
assert np.allclose(probe_array, S_array, atol=1e-5)
def test_probe_waves_line_scan():
potential = Potential(Atoms('C', positions=[(2.5, 2.5, 2)],
cell=(5, 5, 4)))
linescan = LineScan(start=[0, 0], end=[2.5, 2.5], gpts=10)
detector = AnnularDetector(inner=80, outer=200)
S = SMatrix(30, 80e3, 1, gpts=500).multislice(potential, pbar=False)
probe = Probe(semiangle_cutoff=30, energy=80e3, gpts=500)
prism_measurement = S.scan(linescan,
detector,
max_batch_probes=10,
pbar=False)
measurement = probe.scan(linescan,
detector,
potential,
max_batch=50,
pbar=False)
assert np.allclose(measurement.array, prism_measurement.array, atol=1e-6)
def test_fig_5_22():
atoms = Atoms('CSiCuAuU',
positions=[(x, 25, 4) for x in np.linspace(5, 45, 5)],
cell=(50, 50, 8))
gpts = 2048
potential = Potential(atoms=atoms,
gpts=gpts,
parametrization='kirkland',
slice_thickness=8)
#probe = Probe(energy=200e3, defocus=700, Cs=1.3e7, semiangle_cutoff=10.37, rolloff=.1)
probe = Probe(energy=200e3, defocus=700, Cs=1.3e7, semiangle_cutoff=10.37)
probe.grid.match(potential)
scan = LineScan(start=[5, 25], end=[45, 25], gpts=5)
detector = AnnularDetector(inner=40, outer=200)
measurement = probe.scan(scan, detector, potential, pbar=False)
#correct_values = np.array([0.0001168, 0.00059303, 0.00214667, 0.00977803, 0.01167613])
correct_values = np.array(
[0.00010675, 0.00055145, 0.00199743, 0.00911063, 0.01087296])
assert np.allclose(measurement.array, correct_values, atol=1e-5)
def test_resample_diffraction_patterns():
for extent in [(5, 10), (5, 8)]:
for gpts in [(256, 256), (256, 296)]:
probe = Probe(energy=60e3,
extent=extent,
gpts=gpts,
semiangle_cutoff=80,
rolloff=0.2)
detector = PixelatedDetector(max_angle='valid', resample='uniform')
wave = probe.build()
measurement = detector.detect(wave)
measurement /= measurement.max()
probe = Probe(energy=60e3,
extent=(5, 5),
gpts=(256, 256),
semiangle_cutoff=80,
rolloff=0.2)
wave = probe.build()
measurement2 = detector.detect(wave)
measurement2 /= measurement2.max()
s1 = (measurement2.shape[-2] - measurement.shape[-2]) // 2
s2 = (measurement2.shape[-1] - measurement.shape[-1]) // 2
measurement2 = measurement2[..., s1:s1 + measurement.shape[-2],
s2:s2 + measurement.shape[-1]]
assert np.all(np.abs(measurement2[0] - measurement[0]) < .5)
def test_pixelated_detector():
gpts = (512, 512)
extent = (12, 12)
probe = Probe(energy=60e3, extent=extent, gpts=gpts, semiangle_cutoff=80, rolloff=0.2)
detector = PixelatedDetector(max_angle=30, resample='uniform')
wave = probe.build().downsample(max_angle=30)
measurement = detector.detect(wave)
assert measurement.shape == wave.array.shape
detector = PixelatedDetector(max_angle='valid', resample='uniform')
measurement = detector.detect(wave)
assert measurement.shape == wave.array.shape
detector = PixelatedDetector(max_angle='limit', resample='uniform')
measurement = detector.detect(wave)
assert measurement.shape == wave.array.shape
gpts = (512, 512)
extent = (10, 12)
probe = Probe(energy=60e3, extent=extent, gpts=gpts, semiangle_cutoff=80, rolloff=0.2)
detector = PixelatedDetector(max_angle='valid', resample='uniform')
wave = probe.build()
measurement = detector.allocate_measurement(wave)
assert measurement.array.shape[0] == measurement.array.shape[1]
def test_propagator_cache():
wave = Probe(extent=10, gpts=100, energy=60e3, semiangle_cutoff=30).build()
initial_wave_array = wave.array.copy()
propagator = FresnelPropagator()
propagator.propagate(wave, .5)
propagator.propagate(wave, .5)
propagator.propagate(wave, -1)
assert np.allclose(wave.array.real[0], initial_wave_array.real[0], atol=5e-5)
assert np.allclose(wave.array.imag[0], initial_wave_array.imag[0], atol=5e-5)
assert propagator._cache._hits == 1
assert propagator._cache._misses == 2
def test_preallocated_measurement():
atoms = read(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
'data/amorphous_carbon.cif'))
potential = Potential(atoms,
gpts=256,
slice_thickness=1,
projection='infinite',
parametrization='kirkland').build(pbar=False)
scan = GridScan(start=[0, 0], end=potential.extent, gpts=4)
detector1 = AnnularDetector(inner=70, outer=100)
probe = Probe(semiangle_cutoff=15,
energy=300e3,
extent=potential.extent,
gpts=512)
measurement = detector1.allocate_measurement(probe, scan)
probe.scan(scan, detector1, potential, measurement, pbar=False)
assert np.any(measurement.array > 0)
detector2 = PixelatedDetector()
measurement1 = detector1.allocate_measurement(probe, scan)
measurement2 = detector2.allocate_measurement(probe, scan)
with pytest.raises(ValueError) as e:
probe.scan(scan, [detector1, detector2],
potential,
measurement1,
pbar=False)
probe.scan(scan, [detector1, detector2],
potential, {
detector1: measurement1,
detector2: measurement2
},
pbar=False)
def test_bandlimit():
atoms = Atoms(cell=(10, 10, 2))
probe = Probe(energy=300e3, semiangle_cutoff=30, rolloff=0.0, gpts=100)
waves = probe.multislice((0, 0), atoms, pbar=False)
diffraction_pattern = waves.diffraction_pattern(max_angle=30)
assert diffraction_pattern.shape == (33, 33)
probe = Probe(energy=300e3, semiangle_cutoff=1e3, rolloff=0.0, gpts=100)
waves = probe.multislice((0, 0), atoms, pbar=False)
diffraction_pattern = waves.diffraction_pattern('limit')
assert not np.allclose(diffraction_pattern.array / diffraction_pattern.array.max(), 1.)
assert diffraction_pattern.shape == (66, 66)
diffraction_pattern = waves.diffraction_pattern('valid')
assert diffraction_pattern.shape == (45, 45)
assert np.allclose(diffraction_pattern.array / diffraction_pattern.array.max(), 1.)
def test_probe_waves_raises():
with pytest.raises(ValueError) as e:
Probe(not_a_parameter=10)
assert str(e.value) == 'not_a_parameter not a recognized parameter'
probe = Probe()
with pytest.raises(RuntimeError) as e:
probe.build()
assert str(e.value) == 'Grid extent is not defined'
probe.extent = 10
probe.gpts = 100
with pytest.raises(RuntimeError) as e:
probe.build()
assert str(e.value) == 'Energy is not defined'
probe.energy = 60e3
probe.build()
def test_prism_match_probe():
S_builder = SMatrixBuilder(30., 1, extent=5, gpts=101, energy=60e3)
probe = Probe(extent=5, gpts=101, energy=60e3, semiangle_cutoff=30.)
assert np.allclose(probe.build([(0., 0.)]).array,
S_builder.build().collapse([(0., 0.)]).array,
atol=1e-5)
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])
