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