def test_downsample_max_angle(): S = SMatrix(30, 80e3, 1, gpts=500) potential = Potential(Atoms('C', positions=[(2.5, 2.5, 2)], cell=(5, 5, 4))) S = S.multislice(potential, pbar=False) pattern1 = S.collapse((0, 0)).diffraction_pattern(max_angle=64) S = S.downsample(max_angle=64) pattern2 = S.collapse((0, 0)).diffraction_pattern(None) pattern3 = S.collapse((0, 0)).diffraction_pattern(max_angle=64) assert np.allclose(pattern1.array, pattern2.array, atol=1e-6 * pattern1.array.max(), rtol=1e-6) assert np.allclose(pattern3.array, pattern2.array, atol=1e-6 * pattern1.array.max(), rtol=1e-6)
def test_cropped_scan(): atoms = read( os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data/amorphous_carbon.cif')) potential = Potential(atoms, gpts=512, slice_thickness=1, device='gpu', projection='infinite', parametrization='kirkland', storage='gpu').build(pbar=True) detector = AnnularDetector(inner=40, outer=60) gridscan = GridScan(start=[0, 0], end=potential.extent, gpts=16) S = SMatrix(expansion_cutoff=20, interpolation=4, energy=300e3, device='gpu', storage='cpu') # .build() S = S.multislice(potential, pbar=True) S = S.downsample('limit') measurements = S.scan(gridscan, [detector], max_batch_probes=64) scans = gridscan.partition_scan((2, 2)) cropped_measurements = { detector: detector.allocate_measurement(S.collapse((0, 0)), gridscan) } for scan in scans: cropped = S.crop_to_scan(scan) cropped = cropped.transfer('gpu') cropped_measurements = cropped.scan(scan, cropped_measurements, pbar=False) assert np.allclose(cropped_measurements[detector].array, measurements[detector].array)