def build(self, positions: Sequence[Sequence[float]] = None) -> Waves:
"""
Build probe wave functions at the provided positions.
:param positions: Positions of the probe wave functions
:return: Probe wave functions as a Waves object.
"""
self.grid.check_is_defined()
self.accelerator.check_is_defined()
xp = get_array_module_from_device(self._device)
fft2 = get_device_function(xp, 'fft2')
if positions is None:
positions = xp.zeros((1, 2), dtype=xp.float32)
else:
positions = xp.array(positions, dtype=xp.float32)
if len(positions.shape) == 1:
positions = xp.expand_dims(positions, axis=0)
array = fft2(self._evaluate_ctf(xp) *
self._fourier_translation_operator(positions),
overwrite_x=True)
return Waves(array, extent=self.extent, energy=self.energy)
def build(self) -> Waves:
"""
:return: Wave function as a Waves object.
"""
xp = get_array_module_from_device(self._device)
self.grid.check_is_defined()
array = xp.ones((self.gpts[0], self.gpts[1]), dtype=xp.complex64)
return Waves(array, extent=self.extent, energy=self.energy)
def build(self) -> SMatrix:
self.grid.check_is_defined()
self.accelerator.check_is_defined()
xp = get_array_module(self._device)
storage_xp = get_array_module_from_device(self._storage)
complex_exponential = get_device_function(xp, 'complex_exponential')
n_max = int(
xp.ceil(self.expansion_cutoff / 1000. /
(self.wavelength / self.extent[0] * self.interpolation)))
m_max = int(
xp.ceil(self.expansion_cutoff / 1000. /
(self.wavelength / self.extent[1] * self.interpolation)))
n = xp.arange(-n_max, n_max + 1, dtype=xp.float32)
w = xp.asarray(self.extent[0], dtype=xp.float32)
m = xp.arange(-m_max, m_max + 1, dtype=xp.float32)
h = xp.asarray(self.extent[1], dtype=xp.float32)
kx = n / w * xp.float32(self.interpolation)
ky = m / h * xp.float32(self.interpolation)
mask = kx[:, None]**2 + ky[None, :]**2 < (self.expansion_cutoff /
1000. / self.wavelength)**2
kx, ky = xp.meshgrid(kx, ky, indexing='ij')
kx = kx[mask]
ky = ky[mask]
x, y = coordinates(extent=self.extent,
gpts=self.gpts,
endpoint=self.grid.endpoint)
x = xp.asarray(x)
y = xp.asarray(y)
array = storage_xp.zeros((len(kx), ) + (self.gpts[0], self.gpts[1]),
dtype=np.complex64)
for i in range(len(kx)):
array[i] = copy_to_device(
complex_exponential(
-2 * np.pi * kx[i, None, None] * x[:, None]) *
complex_exponential(
-2 * np.pi * ky[i, None, None] * y[None, :]),
self._storage)
return SMatrix(array,
expansion_cutoff=self.expansion_cutoff,
interpolation=self.interpolation,
extent=self.extent,
energy=self.energy,
k=(kx, ky))
def build(self,
pbar: Union[bool, ProgressBar] = False) -> 'ArrayPotential':
self.grid.check_is_defined()
storage_xp = get_array_module_from_device(self._storage)
array = storage_xp.zeros(
(self.num_slices, ) + (self.gpts[0], self.gpts[1]),
dtype=np.float32)
slice_thicknesses = np.zeros(self.num_slices)
if isinstance(pbar, bool):
pbar = ProgressBar(total=len(self),
desc='Potential',
disable=not pbar)
pbar.reset()
for i, potential_slice in enumerate(self.generate_slices()):
array[i] = copy_to_device(potential_slice.array, self._storage)
slice_thicknesses[i] = potential_slice.thickness
pbar.update(1)
pbar.refresh()
return ArrayPotential(array, slice_thicknesses, self.extent)
def _generate_slices_finite(self,
first_slice=0,
last_slice=None,
max_batch=1) -> Generator:
xp = get_array_module_from_device(self._device)
interpolate_radial_functions = get_device_function(
xp, 'interpolate_radial_functions')
atoms = self.atoms.copy()
atoms.wrap()
indices_by_number = {
number: np.where(atoms.numbers == number)[0]
for number in np.unique(atoms.numbers)
}
start, end = next(
generate_batches(last_slice - first_slice,
max_batch=max_batch,
start=first_slice))
array = xp.zeros((end - start, ) + self.gpts, dtype=xp.float32)
slice_edges = np.linspace(0, self.atoms.cell[2, 2],
self.num_slices + 1)
for start, end in generate_batches(last_slice - first_slice,
max_batch=max_batch,
start=first_slice):
array[:] = 0.
for number, indices in indices_by_number.items():
species_atoms = atoms[indices]
integrator = self.get_integrator(number)
disc_indices = xp.asarray(
self._get_radial_interpolation_points(number))
a = slice_edges[start]
b = slice_edges[end]
chunk_atoms = species_atoms[
(species_atoms.positions[:, 2] > a - integrator.cutoff) *
(species_atoms.positions[:, 2] < b + integrator.cutoff)]
chunk_atoms = pad_atoms(chunk_atoms, integrator.cutoff)
chunk_positions = chunk_atoms.positions
if len(chunk_atoms) == 0:
continue
positions = np.zeros((0, 3), dtype=xp.float32)
A = np.zeros((0, ), dtype=xp.float32)
B = np.zeros((0, ), dtype=xp.float32)
run_length_enconding = np.zeros((end - start + 1, ),
dtype=xp.int32)
for i, j in enumerate(range(start, end)):
a = slice_edges[j]
b = slice_edges[j + 1]
slice_positions = chunk_positions[
(chunk_positions[:, 2] > a - integrator.cutoff) *
(chunk_positions[:, 2] < b + integrator.cutoff)]
positions = np.vstack((positions, slice_positions))
A = np.concatenate((A, [a] * len(slice_positions)))
B = np.concatenate((B, [b] * len(slice_positions)))
run_length_enconding[
i + 1] = run_length_enconding[i] + len(slice_positions)
vr, dvdr = integrator.integrate(positions[:, 2], A, B, xp=xp)
vr = xp.asarray(vr, dtype=xp.float32)
dvdr = xp.asarray(dvdr, dtype=xp.float32)
r = xp.asarray(integrator.r, dtype=xp.float32)
sampling = xp.asarray(self.sampling, dtype=xp.float32)
interpolate_radial_functions(array, run_length_enconding,
disc_indices, positions, vr, r,
dvdr, sampling)
slice_thicknesses = [
self.get_slice_thickness(i) for i in range(start, end)
]
yield start, end, PotentialArray(array[:end - start] / kappa,
slice_thicknesses,
extent=self.extent)
def _generate_slices_infinite(self,
first_slice=0,
last_slice=None,
max_batch=1) -> Generator:
xp = get_array_module_from_device(self._device)
fft2_convolve = get_device_function(xp, 'fft2_convolve')
atoms = self.atoms.copy()
atoms.wrap()
positions = atoms.get_positions().astype(np.float32)
numbers = atoms.get_atomic_numbers()
unique = np.unique(numbers)
order = np.argsort(positions[:, 2])
positions = positions[order]
numbers = numbers[order]
kx = xp.fft.fftfreq(self.gpts[0], self.sampling[0])
ky = xp.fft.fftfreq(self.gpts[1], self.sampling[1])
kx, ky = xp.meshgrid(kx, ky, indexing='ij')
k = xp.sqrt(kx**2 + ky**2)
sinc = xp.sinc(
xp.sqrt((kx * self.sampling[0])**2 + (kx * self.sampling[1])**2))
scattering_factors = {}
for atomic_number in unique:
f = kirkland_projected_fourier(k, self.parameters[atomic_number])
scattering_factors[atomic_number] = (
f /
(sinc * self.sampling[0] * self.sampling[1] * kappa)).astype(
xp.complex64)
slice_idx = np.floor(positions[:, 2] / atoms.cell[2, 2] *
self.num_slices).astype(np.int)
start, end = next(
generate_batches(last_slice - first_slice,
max_batch=max_batch,
start=first_slice))
array = xp.zeros((end - start, ) + self.gpts, dtype=xp.complex64)
temp = xp.zeros((end - start, ) + self.gpts, dtype=xp.complex64)
for start, end in generate_batches(last_slice - first_slice,
max_batch=max_batch,
start=first_slice):
array[:] = 0.
start_idx = np.searchsorted(slice_idx, start)
end_idx = np.searchsorted(slice_idx, end)
if start_idx != end_idx:
for j, number in enumerate(unique):
temp[:] = 0.
chunk_positions = positions[start_idx:end_idx]
chunk_slice_idx = slice_idx[start_idx:end_idx] - start
if len(unique) > 1:
chunk_positions = chunk_positions[
numbers[start_idx:end_idx] == number]
chunk_slice_idx = chunk_slice_idx[
numbers[start_idx:end_idx] == number]
chunk_positions = xp.asarray(chunk_positions[:, :2] /
self.sampling)
superpose_deltas(chunk_positions, chunk_slice_idx, temp)
fft2_convolve(temp, scattering_factors[number])
array += temp
slice_thicknesses = [
self.get_slice_thickness(i) for i in range(start, end)
]
yield start, end, PotentialArray(array.real[:end - start],
slice_thicknesses,
extent=self.extent)
def build(
self,
first_slice: int = 0,
last_slice: int = None,
energy: float = None,
max_batch: int = None,
pbar: Union[bool, ProgressBar] = False,
) -> 'PotentialArray':
"""
Precalcaulate the potential as a potential array.
Parameters
----------
first_slice: int
First potential slice to generate.
last_slice: int, optional
Last potential slice generate.
energy: float
Electron energy [eV]. If given, the transmission functions will be returned.
max_batch: int
Maximum number of potential slices calculated in parallel.
pbar: bool
If true, show progress bar.
Returns
-------
PotentialArray object
"""
self.grid.check_is_defined()
if last_slice is None:
last_slice = len(self)
if max_batch is None:
max_batch = self._estimate_max_batch()
storage_xp = get_array_module_from_device(self._storage)
if energy is None:
array = storage_xp.zeros(
(last_slice - first_slice, ) + (self.gpts[0], self.gpts[1]),
dtype=np.float32)
generator = self.generate_slices(max_batch=max_batch,
first_slice=first_slice,
last_slice=last_slice)
else:
array = storage_xp.zeros(
(last_slice - first_slice, ) + (self.gpts[0], self.gpts[1]),
dtype=np.complex64)
generator = self.generate_transmission_functions(
energy=energy,
max_batch=max_batch,
first_slice=first_slice,
last_slice=last_slice)
slice_thicknesses = np.zeros(last_slice - first_slice)
if isinstance(pbar, bool):
pbar = ProgressBar(total=len(self),
desc='Potential',
disable=not pbar)
close_pbar = True
else:
close_pbar = False
pbar.reset()
for start, end, potential_slice in generator:
array[start:end] = copy_to_device(potential_slice.array,
self._storage)
slice_thicknesses[start:end] = potential_slice.slice_thicknesses
pbar.update(end - start)
pbar.refresh()
if close_pbar:
pbar.close()
if energy is None:
return PotentialArray(array,
slice_thicknesses=slice_thicknesses,
extent=self.extent)
else:
return TransmissionFunction(array,
slice_thicknesses=slice_thicknesses,
extent=self.extent,
energy=energy)
def _generate_slices_finite(self,
first_slice=0,
last_slice=None,
max_batch=1) -> Generator:
sliced_atoms = SlicedAtoms(self.atoms, self.slice_thickness)
xp = get_array_module_from_device(self._device)
interpolate_radial_functions = get_device_function(
xp, 'interpolate_radial_functions')
array = None
unique = np.unique(self.atoms.numbers)
for start, end in generate_batches(last_slice - first_slice,
max_batch=max_batch,
start=first_slice):
if array is None:
array = xp.zeros((end - start, ) + self.gpts, dtype=xp.float32)
else:
array[:] = 0.
for number in unique:
integrator = self.get_integrator(number)
disc_indices = xp.asarray(
self._get_radial_interpolation_points(number))
chunk_atoms = sliced_atoms.get_subsliced_atoms(
start, end, number, z_margin=integrator.cutoff)
if len(chunk_atoms) == 0:
continue
positions = np.zeros((0, 3), dtype=xp.float32)
slice_entrances = np.zeros((0, ), dtype=xp.float32)
slice_exits = np.zeros((0, ), dtype=xp.float32)
run_length_enconding = np.zeros((end - start + 1, ),
dtype=xp.int32)
for i, slice_idx in enumerate(range(start, end)):
slice_atoms = chunk_atoms.get_subsliced_atoms(
slice_idx,
padding=integrator.cutoff,
z_margin=integrator.cutoff)
slice_positions = slice_atoms.positions
slice_entrance = slice_atoms.get_slice_entrance(slice_idx)
slice_exit = slice_atoms.get_slice_exit(slice_idx)
positions = np.vstack((positions, slice_positions))
slice_entrances = np.concatenate(
(slice_entrances,
[slice_entrance] * len(slice_positions)))
slice_exits = np.concatenate(
(slice_exits, [slice_exit] * len(slice_positions)))
run_length_enconding[
i + 1] = run_length_enconding[i] + len(slice_positions)
vr, dvdr = integrator.integrate(positions[:, 2],
slice_entrances, slice_exits)
vr = xp.asarray(vr, dtype=xp.float32)
dvdr = xp.asarray(dvdr, dtype=xp.float32)
r = xp.asarray(integrator.r, dtype=xp.float32)
sampling = xp.asarray(self.sampling, dtype=xp.float32)
interpolate_radial_functions(array, run_length_enconding,
disc_indices, positions, vr, r,
dvdr, sampling)
slice_thicknesses = [
self.get_slice_thickness(i) for i in range(start, end)
]
yield start, end, PotentialArray(array[:end - start] / kappa,
slice_thicknesses,
extent=self.extent)
def _generate_slices_infinite(self,
first_slice=0,
last_slice=None,
max_batch=1) -> Generator:
xp = get_array_module_from_device(self._device)
fft2_convolve = get_device_function(xp, 'fft2_convolve')
atoms = self.atoms.copy()
atoms.wrap()
indices_by_number = {
number: np.where(atoms.numbers == number)[0]
for number in np.unique(atoms.numbers)
}
kx = xp.fft.fftfreq(self.gpts[0], self.sampling[0])
ky = xp.fft.fftfreq(self.gpts[1], self.sampling[1])
kx, ky = xp.meshgrid(kx, ky, indexing='ij')
k = xp.sqrt(kx**2 + ky**2)
sinc = xp.sinc(
xp.sqrt((kx * self.sampling[0])**2 + (kx * self.sampling[1])**2))
scattering_factors = {}
for atomic_number in indices_by_number.keys():
f = kirkland_projected_fourier(k, self.parameters[atomic_number])
scattering_factors[atomic_number] = (
f /
(sinc * self.sampling[0] * self.sampling[1] * kappa)).astype(
xp.complex64)
slice_idx = np.floor(atoms.positions[:, 2] / atoms.cell[2, 2] *
self.num_slices).astype(np.int)
start, end = next(
generate_batches(last_slice - first_slice,
max_batch=max_batch,
start=first_slice))
array = xp.zeros((end - start, ) + self.gpts, dtype=xp.complex64)
temp = xp.zeros((end - start, ) + self.gpts, dtype=xp.complex64)
for start, end in generate_batches(last_slice - first_slice,
max_batch=max_batch,
start=first_slice):
array[:] = 0.
for j, (number, indices) in enumerate(indices_by_number.items()):
temp[:] = 0.
in_slice = (slice_idx >= start) * (slice_idx < end) * (
atoms.numbers == number)
slice_atoms = atoms[in_slice]
positions = xp.asarray(slice_atoms.positions[:, :2] /
self.sampling)
superpose_deltas(positions, slice_idx[in_slice] - start, temp)
fft2_convolve(temp, scattering_factors[number])
array += temp
slice_thicknesses = [
self.get_slice_thickness(i) for i in range(start, end)
]
yield start, end, PotentialArray(array.real[:end - start],
slice_thicknesses,
extent=self.extent)