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_downsample_detect():
atoms = read('data/srtio3_100.cif')
atoms *= (4, 4, 1)
potential = Potential(atoms,
gpts=256,
projection='infinite',
slice_thickness=.5,
parametrization='kirkland').build(pbar=False)
detector = FlexibleAnnularDetector()
end = (potential.extent[0] / 4, potential.extent[1] / 4)
gridscan = GridScan(start=[0, 0], end=end, sampling=.2)
S = SMatrix(energy=300e3,
semiangle_cutoff=9.4,
rolloff=0.05,
expansion_cutoff=10)
S_exit = S.multislice(potential, pbar=False)
measurement = S_exit.scan(gridscan, detector, pbar=False)
S_downsampled = S_exit.downsample()
downsampled_measurement = S_downsampled.scan(gridscan,
detector,
pbar=False)
assert S_downsampled.array.shape != S_exit.array.shape
assert not np.all(measurement.array == downsampled_measurement.array)
assert np.allclose(measurement.array, downsampled_measurement.array)
def test_prism_storage():
potential = Potential(Atoms('C', positions=[(2.5, 2.5, 2)],
cell=(5, 5, 4)))
S_builder = SMatrix(30, 80e3, 1, gpts=500, device='gpu', storage='cpu')
S_gpu = S_builder.multislice(potential, pbar=False)
assert type(S_gpu.array) is np.ndarray
def test_prism_batch():
potential = Potential(Atoms('C', positions=[(2.5, 2.5, 2)],
cell=(5, 5, 4)))
S_builder = SMatrix(30, 80e3, 1, gpts=500)
S1 = S_builder.multislice(potential, pbar=False)
S2 = S_builder.multislice(potential, max_batch=5, pbar=False)
assert np.allclose(S1.array, S2.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_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_tilt():
S = SMatrix(semiangle_cutoff=30.,
energy=60e3,
interpolation=1,
extent=5,
gpts=50,
tilt=(1, 2))
S_array = S.build()
wave = S_array.collapse()
assert S_array.tilt == S.tilt == wave.tilt
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_raises():
with pytest.raises(ValueError) as e:
SMatrix(.01, 80e3, .5)
assert str(e.value) == 'Interpolation factor must be int'
with pytest.raises(RuntimeError) as e:
prism = SMatrix(10, 80e3, 1)
prism.build()
assert str(e.value) == 'Grid extent is not defined'
def test_prism_batch():
potential = Potential(Atoms('C', positions=[(2.5, 2.5, 2)],
cell=(5, 5, 4)))
S_builder = SMatrix(semiangle_cutoff=30.,
energy=80e3,
interpolation=1,
gpts=500)
S1 = S_builder.multislice(potential, pbar=False)
S2 = S_builder.multislice(potential, max_batch=5, pbar=False)
assert np.allclose(S1.array, S2.array)
def test_prism_storage():
potential = Potential(Atoms('C', positions=[(2.5, 2.5, 2)],
cell=(5, 5, 4)))
S_builder = SMatrix(semiangle_cutoff=30.,
energy=80e3,
interpolation=1,
gpts=500,
device='gpu',
storage='cpu')
S_gpu = S_builder.multislice(potential, pbar=False)
assert type(S_gpu.array) is np.ndarray
def test_prism_multislice():
potential = Potential(Atoms('C', positions=[(0, 0, 2)], cell=(5, 5, 4)))
S = SMatrix(30, 60e3, 1, extent=5, gpts=500)
probe = Probe(extent=5,
gpts=500,
energy=60e3,
semiangle_cutoff=30,
rolloff=0.)
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_crop():
S = SMatrix(expansion_cutoff=30,
interpolation=3,
energy=300e3,
extent=10,
sampling=.02).build()
gridscan = GridScan(start=[0, 0], end=S.extent, gpts=64)
scans = gridscan.partition_scan((2, 2))
cropped = S.crop_to_scan(scans[0])
assert cropped.gpts != S.gpts
position = (4.9, 0.)
assert np.allclose(
S.collapse(position).array - cropped.collapse(position).array, 0.)
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_unequal_interpolation():
from ase.build import nanotube
def make_atoms(l):
atoms = nanotube(6, 0, length=l)
atoms.rotate(-90, 'x', rotate_cell=True)
atoms.center(vacuum=5, axis=(0, 1))
cell = atoms.cell.copy()
cell[1] = np.abs(atoms.cell[2])
cell[2] = np.abs(atoms.cell[1])
atoms.set_cell(cell)
return atoms
detector = AnnularDetector(50, 150)
atoms1 = make_atoms(8)
S1 = SMatrix(energy=80e3,
expansion_cutoff=10,
semiangle_cutoff=10,
interpolation=(2, 4),
device='cpu',
gpts=(256, 512))
S1.extent = np.diag(atoms1.cell)[:2]
scan = GridScan((0, 0), (atoms1.cell[0, 0], atoms1.cell[1, 1] / 2),
sampling=S1.ctf.nyquist_sampling * .9)
measurement1 = S1.scan(scan, detector, atoms1, pbar=False)
atoms2 = make_atoms(4)
S2 = SMatrix(energy=80e3,
expansion_cutoff=10,
semiangle_cutoff=10,
interpolation=(2, 2),
device='cpu',
gpts=(256, 256))
S2.extent = np.diag(atoms2.cell)[:2]
scan = GridScan((0, 0), (atoms2.cell[0, 0], atoms2.cell[1, 1]),
sampling=S2.ctf.nyquist_sampling * .9)
measurement2 = S2.scan(scan, detector, atoms2, pbar=False)
assert len(S1) == len(S2)
assert np.allclose((S1.build().collapse().intensity() -
S2.build().collapse().intensity()).array, 0)
assert np.allclose((measurement2 - measurement1).array, 0, 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_partition_measurement():
atoms = read(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
'data/amorphous_carbon.cif'))
potential = Potential(atoms,
gpts=512,
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=16)
S = SMatrix(expansion_cutoff=15, interpolation=1, energy=300e3)
S = S.multislice(potential, pbar=False)
measurements = S.scan(gridscan, [detector], pbar=False)
scans = gridscan.partition_scan((2, 2))
partitioned_measurements = {
detector: detector.allocate_measurement(S.collapse((0, 0)), gridscan)
}
for scan in scans:
partitioned_measurements = S.scan(scan, measurements, pbar=False)
assert np.allclose(partitioned_measurements[detector].array,
measurements[detector].array)
def test_prism_gpu():
potential = Potential(Atoms('C', positions=[(2.5, 2.5, 2)],
cell=(5, 5, 4)))
S_builder = SMatrix(30, 80e3, 1, gpts=500, device='cpu')
S_cpu = S_builder.multislice(potential, pbar=False)
assert type(S_cpu.array) is np.ndarray
S_builder = SMatrix(30, 80e3, 1, gpts=500, device='gpu')
S_gpu = S_builder.multislice(potential, pbar=False)
assert type(S_gpu.array) is cp.ndarray
assert np.allclose(S_cpu.array, asnumpy(S_gpu.array))
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_downsample_smatrix():
S = SMatrix(expansion_cutoff=10,
interpolation=2,
energy=300e3,
extent=10,
sampling=.05)
S = S.build().downsample()
S2 = SMatrix(expansion_cutoff=10,
interpolation=2,
energy=300e3,
extent=10,
gpts=S.gpts)
S2 = S2.build()
assert np.allclose(S.array - S2.array, 0., atol=5e-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)
def test_prism_tilt():
S = SMatrix(30., 60e3, 1, extent=5, gpts=50, tilt=(1, 2))
S_array = S.build()
wave = S_array.collapse()
assert S_array.tilt == S.tilt == wave.tilt
