def __init__(self,
calculator,
gpts: Union[int, Sequence[int]] = None,
sampling: Union[float, Sequence[float]] = None,
# origin: Union[float, Sequence[float]] = None,
slice_thickness=.5,
core_size=.005,
storage='cpu',
precalculate=True):
self._calculator = calculator
self._core_size = core_size
thickness = calculator.atoms.cell[2, 2]
nz = calculator.hamiltonian.finegd.N_c[2]
num_slices = int(np.ceil(nz / np.floor(slice_thickness / (thickness / nz))))
self._voxel_height = thickness / nz
self._slice_vertical_voxels = subdivide_into_batches(nz, num_slices)
# TODO: implement support for non-periodic extent
self._origin = (0., 0.)
extent = np.diag(orthogonalize_cell(calculator.atoms.copy()).cell)[:2]
self._grid = Grid(extent=extent, gpts=gpts, sampling=sampling, lock_extent=True)
super().__init__(precalculate=precalculate, storage=storage)
def __init__(self,
calculator,
gpts=None,
sampling=None,
slice_thickness=.5,
core_size=.005,
storage='cpu'):
self._calculator = calculator
self._core_size = core_size
thickness = calculator.atoms.cell[2, 2]
nz = calculator.hamiltonian.finegd.N_c[2]
num_slices = int(
np.ceil(nz / np.floor(slice_thickness / (thickness / nz))))
self._voxel_height = thickness / nz
self._slice_vertical_voxels = split_integer(nz, num_slices)
# TODO: implement support for non-periodic extent
self._origin = (0., 0.)
extent = np.diag(
orthogonalize_cell(calculator.atoms.copy(),
strain_error=True).cell)[:2]
self._grid = Grid(extent=extent,
gpts=gpts,
sampling=sampling,
lock_extent=True)
super().__init__(storage)
def __init__(self,
calculator,
gpts: Union[int, Sequence[int]] = None,
sampling: Union[float, Sequence[float]] = None,
origin: Union[float, Sequence[float]] = None,
orthogonal_cell: Sequence[float] = None,
periodic_z: bool = True,
slice_thickness=.5,
core_size=.005,
plane='xy',
storage='cpu',
precalculate=True):
self._calculator = calculator
self._core_size = core_size
self._plane = plane
if orthogonal_cell is None:
atoms = rotate_atoms_to_plane(calculator.atoms, plane)
thickness = atoms.cell[2, 2]
nz = calculator.hamiltonian.finegd.N_c[plane_to_axes(plane)[-1]]
extent = np.diag(orthogonalize_cell(atoms.copy()).cell)[:2]
else:
if plane != 'xy':
raise NotImplementedError()
thickness = orthogonal_cell[2]
nz = calculator.hamiltonian.finegd.N_c / np.linalg.norm(calculator.atoms.cell, axis=0) * orthogonal_cell[2]
nz = int(np.ceil(np.max(nz)))
extent = orthogonal_cell[:2]
num_slices = int(np.ceil(nz / np.floor(slice_thickness / (thickness / nz))))
self._orthogonal_cell = orthogonal_cell
self._voxel_height = thickness / nz
self._slice_vertical_voxels = subdivide_into_batches(nz, num_slices)
self._origin = (0., 0., 0.)
self._periodic_z = periodic_z
self._grid = Grid(extent=extent, gpts=gpts, sampling=sampling, lock_extent=True)
super().__init__(precalculate=precalculate, storage=storage)
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)
cutatoms.rotate(zrotation, 'z', center=(0, 0, 0), rotate_cell=False)
# # Adding vacuum around sample
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();
def generate_slices(self, first_slice=0, last_slice=None, max_batch=1):
interpolate_radial_functions = get_device_function(np, 'interpolate_radial_functions')
if last_slice is None:
last_slice = len(self)
valence = self._calculator.get_electrostatic_potential()
cell = self._calculator.atoms.cell[:2, :2]
atoms = self._calculator.atoms.copy()
atoms.set_tags(range(len(atoms)))
atoms = orthogonalize_cell(atoms)
indices_by_number = {number: np.where(atoms.numbers == number)[0] for number in np.unique(atoms.numbers)}
na = sum(self._slice_vertical_voxels[:first_slice])
a = na * self._voxel_height
for i in range(first_slice, last_slice):
nb = na + self._slice_vertical_voxels[i]
b = a + self._slice_vertical_voxels[i] * self._voxel_height
projected_valence = valence[..., na:nb].sum(axis=-1) * self._voxel_height
projected_valence = interpolate_rectangle(projected_valence, cell, self.extent, self.gpts, self._origin)
array = np.zeros((1,) + self.gpts, dtype=np.float32)
for number, indices in indices_by_number.items():
slice_atoms = atoms[indices]
if len(slice_atoms) == 0:
continue
r = self._calculator.density.setups[indices[0]].xc_correction.rgd.r_g[1:] * units.Bohr
cutoff = r[-1]
margin = np.int(np.ceil(cutoff / np.min(self.sampling)))
rows, cols = _disc_meshgrid(margin)
disc_indices = np.hstack((rows[:, None], cols[:, None]))
slice_atoms = slice_atoms[(slice_atoms.positions[:, 2] > a - cutoff) *
(slice_atoms.positions[:, 2] < b + cutoff)]
slice_atoms = pad_atoms(slice_atoms, cutoff)
R = np.geomspace(np.min(self.sampling) / 2, cutoff, int(np.ceil(cutoff / np.min(self.sampling))) * 10)
vr = np.zeros((len(slice_atoms), len(R)), np.float32)
dvdr = np.zeros((len(slice_atoms), len(R)), np.float32)
for j, atom in enumerate(slice_atoms):
r, v = get_paw_corrections(atom.tag, self._calculator, self._core_size)
f = interp1d(r * units.Bohr, v, fill_value=(v[0], 0), bounds_error=False, kind='linear')
integrator = PotentialIntegrator(f, R, self.get_slice_thickness(i), tolerance=1e-6)
vr[j], dvdr[j] = integrator.integrate(np.array([atom.z]), a, b)
sampling = np.asarray(self.sampling, dtype=np.float32)
run_length_enconding = np.zeros((2,), dtype=np.int32)
run_length_enconding[1] = len(slice_atoms)
interpolate_radial_functions(array,
run_length_enconding,
disc_indices,
slice_atoms.positions,
vr,
R,
dvdr,
sampling)
array = -(projected_valence + array / np.sqrt(4 * np.pi) * units.Ha)
yield i, i + 1, PotentialArray(array, np.array([self.get_slice_thickness(i)]), extent=self.extent)
a = b
na = nb
def generate_slices(self, first_slice=0, last_slice=None, max_batch=1):
interpolate_radial_functions = get_device_function(np, 'interpolate_radial_functions')
if last_slice is None:
last_slice = len(self)
if self._plane != 'xy':
atoms = rotate_atoms_to_plane(self._calculator.atoms.copy(), self._plane)
else:
atoms = self._calculator.atoms.copy()
old_cell = atoms.cell
atoms.set_tags(range(len(atoms)))
if self._orthogonal_cell is None:
atoms = orthogonalize_cell(atoms)
else:
scaled = atoms.cell.scaled_positions(np.diag(self._orthogonal_cell))
atoms = cut(atoms, a=scaled[0], b=scaled[1], c=scaled[2])
valence = self._calculator.get_electrostatic_potential()
new_gpts = self.gpts + (sum(self._slice_vertical_voxels),)
axes = plane_to_axes(self._plane)
if self._plane != 'xy':
array = np.moveaxis(valence, axes[:2], (0, 1))
else:
array = valence
from scipy.interpolate import RegularGridInterpolator
origin = (0., 0., 0.)
padded_array = np.zeros((array.shape[0] + 1, array.shape[1] + 1, array.shape[2] + 1))
padded_array[:-1, :-1, :-1] = array
padded_array[-1] = padded_array[0]
padded_array[:, -1] = padded_array[:, 0]
padded_array[:, :, -1] = padded_array[:, :, 0]
x = np.linspace(0, 1, padded_array.shape[0], endpoint=True)
y = np.linspace(0, 1, padded_array.shape[1], endpoint=True)
z = np.linspace(0, 1, padded_array.shape[2], endpoint=True)
interpolator = RegularGridInterpolator((x, y, z), padded_array)
new_cell = np.diag(atoms.cell)
x = np.linspace(origin[0], origin[0] + new_cell[0], new_gpts[0], endpoint=False)
y = np.linspace(origin[1], origin[1] + new_cell[1], new_gpts[1], endpoint=False)
z = np.linspace(origin[2], origin[2] + new_cell[2], new_gpts[2], endpoint=False)
P = np.array(old_cell)
P_inv = np.linalg.inv(P)
cutoffs = {}
for number in np.unique(atoms.numbers):
indices = np.where(atoms.numbers == number)[0]
r = self._calculator.density.setups[indices[0]].xc_correction.rgd.r_g[1:] * units.Bohr
cutoffs[number] = r[-1]
if self._periodic_z:
atoms = pad_atoms(atoms, margin=max(cutoffs.values()), directions='z', in_place=True)
indices_by_number = {number: np.where(atoms.numbers == number)[0] for number in np.unique(atoms.numbers)}
na = sum(self._slice_vertical_voxels[:first_slice])
a = na * self._voxel_height
for i in range(first_slice, last_slice):
nb = na + self._slice_vertical_voxels[i]
b = a + self._slice_vertical_voxels[i] * self._voxel_height
X, Y, Z = np.meshgrid(x, y, z[na:nb], indexing='ij')
points = np.array([X.ravel(), Y.ravel(), Z.ravel()]).T
scaled_points = np.dot(points, P_inv) % 1.0
projected_valence = interpolator(scaled_points).reshape(self.gpts + (nb - na,)).sum(
axis=-1) * self._voxel_height
array = np.zeros((1,) + self.gpts, dtype=np.float32)
for number, indices in indices_by_number.items():
slice_atoms = atoms[indices]
if len(slice_atoms) == 0:
continue
cutoff = cutoffs[number]
margin = np.int(np.ceil(cutoff / np.min(self.sampling)))
rows, cols = _disc_meshgrid(margin)
disc_indices = np.hstack((rows[:, None], cols[:, None]))
slice_atoms = slice_atoms[(slice_atoms.positions[:, 2] > a - cutoff) *
(slice_atoms.positions[:, 2] < b + cutoff)]
slice_atoms = pad_atoms(slice_atoms, margin=cutoff, directions='xy', )
R = np.geomspace(np.min(self.sampling) / 2, cutoff, int(np.ceil(cutoff / np.min(self.sampling))) * 10)
vr = np.zeros((len(slice_atoms), len(R)), np.float32)
dvdr = np.zeros((len(slice_atoms), len(R)), np.float32)
# TODO : improve speed of this
for j, atom in enumerate(slice_atoms):
r, v = get_paw_corrections(atom.tag, self._calculator, self._core_size)
f = interp1d(r * units.Bohr, v, fill_value=(v[0], 0), bounds_error=False, kind='linear')
integrator = PotentialIntegrator(f, R, self.get_slice_thickness(i), tolerance=1e-6)
vr[j], dvdr[j] = integrator.integrate(np.array([atom.z]), a, b)
sampling = np.asarray(self.sampling, dtype=np.float32)
run_length_enconding = np.zeros((2,), dtype=np.int32)
run_length_enconding[1] = len(slice_atoms)
interpolate_radial_functions(array,
run_length_enconding,
disc_indices,
slice_atoms.positions,
vr,
R,
dvdr,
sampling)
array = -(projected_valence + array / np.sqrt(4 * np.pi) * units.Ha)
yield i, i + 1, PotentialArray(array, np.array([self.get_slice_thickness(i)]), extent=self.extent)
a = b
na = nb
from abtem import *
from abtem.structures import orthogonalize_cell
device = 'cpu'
##### Set up atomic structure #####
atoms = mx2(formula='MoS2',
kind='2H',
a=3.18,
thickness=3.19,
size=(1, 1, 1),
vacuum=None)
repetitions = (3, 2, 1)
atoms = orthogonalize_cell(atoms)
atoms *= repetitions
atoms.center(vacuum=2, axis=2)
##### Define potential #####
potential = Potential(atoms,
gpts=256,
projection='finite',
slice_thickness=1,
parametrization='kirkland')
##### Define probe #####
probe = SMatrix(
energy=80e3,
