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