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[detector]
imported_measurement = Measurement.read(measurements[export_detector])
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_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_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_export_import_potential(tmp_path):
atoms = read(_set_path('orthogonal_graphene.cif'))
d = tmp_path / 'sub'
d.mkdir()
path = d / 'potential.hdf5'
potential = Potential(atoms, sampling=.05)
precalculated_potential = potential.build(pbar=False)
precalculated_potential.write(path)
imported_potential = PotentialArray.read(path)
assert np.allclose(imported_potential.array, precalculated_potential.array)
assert np.allclose(imported_potential.extent,
precalculated_potential.extent)
assert np.allclose(imported_potential._slice_thicknesses,
precalculated_potential._slice_thicknesses)
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_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_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_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_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_probe_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)
measurement = probe.scan(linescan, detector, potential, max_batch=50, pbar=False)
tds_measurement = probe.scan(linescan, detector, tds_potential, max_batch=50, pbar=False)
assert np.allclose(measurement.array, tds_measurement.array, atol=1e-6)
frozen_phonons = FrozenPhonons(atoms, 2, sigmas={'C': 0, 'O': 0.1})
tds_potential = Potential(frozen_phonons, sampling=.05)
tds_measurement = probe.scan(linescan, detector, tds_potential, max_batch=50, pbar=False)
assert not np.allclose(measurement.array, tds_measurement.array, atol=1e-6)
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_potential_storage():
atoms = Atoms('CO', positions=[(2, 3, 1), (3, 2, 3)], cell=(4, 6, 4.3))
potential = Potential(atoms=atoms, sampling=.1, device='gpu')
assert type(potential.build().array) is cp.ndarray
potential = Potential(atoms=atoms, sampling=.1, device='gpu', storage='cpu')
assert type(potential.build().array) is np.ndarray
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_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(30, 80e3, 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_potential_centered():
Lx = 5
Ly = 5
gpts_x = 40
gpts_y = 60
atoms1 = Atoms('C', positions=[(0, Ly / 2, 2)], cell=[Lx, Ly, 4])
atoms2 = Atoms('C', positions=[(Lx / 2, Ly / 2, 2)], cell=[Lx, Ly, 4])
potential1 = Potential(atoms1,
gpts=(gpts_x, gpts_y),
slice_thickness=4,
cutoff_tolerance=1e-2)
potential2 = Potential(atoms2,
gpts=(gpts_x, gpts_y),
slice_thickness=4,
cutoff_tolerance=1e-2)
assert np.allclose(
potential1[0].array[:, gpts_y // 2],
np.roll(potential2[0].array[:, gpts_y // 2], gpts_x // 2))
assert np.allclose(potential2[0].array[:, gpts_y // 2][1:],
(potential2[0].array[:, gpts_y // 2])[::-1][:-1])
assert np.allclose(potential2[0].array[gpts_x // 2][1:],
potential2[0].array[gpts_x // 2][::-1][:-1])
def multislice(self,
potential: Union[AbstractPotential, Atoms],
pbar: bool = True) -> Waves:
"""
Build plane wave function and propagate it through the potential.
The grid of the potential will be matched to the wave function.
:param potential: The potential through which to propagate the wave function.
:param pbar: If true, display a progress bar.
:return: Wave function at the exit plane of the potential.
"""
if isinstance(potential, Atoms):
potential = Potential(atoms=potential)
potential.grid.match(self)
return self.build().multislice(potential, pbar=pbar)
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_dft():
atoms = read('data/graphene.traj')
calc = GPAW(mode=PW(400), h=.1, txt=None, kpts=(4, 2, 2))
atoms.set_calculator(calc)
atoms.get_potential_energy()
potential_dft = GPAWPotential(calc, sampling=.05).build()
potential_iam = Potential(atoms, sampling=.05).build()
projected_dft = potential_dft.array.sum(0)
projected_dft -= projected_dft.min()
projected_iam = potential_iam.array.sum(0)
projected_iam -= projected_iam.min()
rel_diff = (projected_iam - projected_dft) / (projected_iam + 1e-16) * 100
rel_diff[projected_iam < 10] = np.nan
assert np.round(np.nanmax(rel_diff) / 10, 0) * 10 == 40
def test_fig_5_12():
atoms = Atoms('CSiCuAuU',
positions=[(x, 25, 4) for x in np.linspace(5, 45, 5)],
cell=(50, 50, 8))
potential = Potential(atoms=atoms,
gpts=512,
parametrization='kirkland',
cutoff_tolerance=1e-4)
waves = PlaneWave(energy=200e3)
waves = waves.multislice(potential, pbar=False)
waves = waves.apply_ctf(defocus=700, Cs=1.3e7, semiangle_cutoff=.01037)
intensity = np.abs(waves.array)**2
assert np.round(intensity.min(), 2) == np.float32(.72)
assert np.round(intensity.max(), 2) == np.float32(1.03)
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_dft():
from gpaw import GPAW
from abtem.dft import GPAWPotential
atoms = read(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
'data/hexagonal_graphene.cif'))
gpaw = GPAW(h=.1, txt=None, kpts=(3, 3, 1))
atoms.calc = gpaw
atoms.get_potential_energy()
dft_pot = GPAWPotential(gpaw, sampling=.02)
dft_array = dft_pot.build()
dft_potential = dft_array.tile((3, 2))
atoms = orthogonalize_cell(gpaw.atoms) * (3, 2, 1)
iam_potential = Potential(atoms,
gpts=dft_potential.gpts,
cutoff_tolerance=1e-4,
device='cpu').build()
projected_iam = iam_potential.array.sum(0)
projected_iam -= projected_iam.min()
projected_dft = dft_potential.array.sum(0)
projected_dft -= projected_dft.min()
absolute_difference = projected_iam - projected_dft
valid = np.abs(projected_iam) > 1
relative_difference = np.zeros_like(projected_iam)
relative_difference[:] = np.nan
relative_difference[valid] = 100 * (
projected_iam[valid] - projected_dft[valid]) / projected_iam[valid]
assert np.isclose(9.553661, absolute_difference.max(), atol=.1)
assert np.isclose(43.573837, relative_difference[valid].max(), atol=.1)
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_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_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)
cutatoms.cell.niggli_reduce()
cutatoms.center(vacuum=60.0)
# # Show sample after rotation and with vaccum added.
# # Uncomment only for a single example. Used for illustrating samples.
# view(cutatoms)
# # # # # # # # # # # # # # # # # # # # #
# # abtem - Simulating TEM Images: # # #
# # # # # # # # # # # # # # # # # # # # #
# # Building the potential
cutatoms = orthogonalize_cell(cutatoms)
potential = Potential(cutatoms,
gpts=512,
slice_thickness=1,
parametrization='kirkland',
projection='infinite')
wave = PlaneWave(energy=300e3 # acceleration voltage in eV
)
exit_wave = wave.multislice(potential)
# # Show exit_wave before adding noise and ctf
# # Uncomment only for a single example. Used for illustrating samples.
# exit_wave.intensity().mean(0).show();
ctf = CTF(
energy=wave.energy,
semiangle_cutoff=20, # mrad
cbar = plt.colorbar(sc,
cax=cax,
ticks=np.arange(1, m + 1, 1),
orientation='vertical')
#plt.colorbar(sc, cax=cax, orientation='vertical',ticks=[-5,-2.5,0,2.5,5],label='$\epsilon_p$ [\%]')
#plt.tight_layout()
plt.show()
for i, model in enumerate(models_list):
# # # Building the potential
potential = Potential(model,
gpts=L,
slice_thickness=1,
parametrization='kirkland',
projection='infinite')
wave = PlaneWave(energy=300e3 # acceleration voltage in eV
)
exit_wave = wave.multislice(potential)
np.savez('{0}/points/points_{1:04d}.npz'.format(dir_name,
first_number + i),
sites=sites_list[i],
classes=classes_list[i])
exit_wave.write('{0}/wave/wave_{1:04d}.hdf5'.format(
dir_name, first_number + i))
write('{0}/model/model_{1:04d}.cfg'.format(dir_name, first_number + i),
def test_potential_build():
atoms = Atoms('CO', positions=[(2, 3, 1), (3, 2, 3)], cell=(4, 6, 4.3))
potential = Potential(atoms=atoms, sampling=.1)
array_potential = potential.build()
assert np.all(array_potential[2].array == potential[2].array)
def test_potential_raises():
with pytest.raises(RuntimeError) as e:
Potential(Atoms('C', positions=[(2, 2, 2)], cell=(4, 4, 0)))
assert str(e.value) == 'atoms has no thickness'