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,