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