def generate_slices(self,
first_slice: int = 0,
last_slice: int = None,
max_batch: int = 1):
if last_slice is None:
last_slice = len(self)
for start, end in generate_batches(last_slice - first_slice,
max_batch=max_batch,
start=first_slice):
slice_thicknesses = np.array(
[self.get_slice_thickness(i) for i in range(start, end)])
yield start, end, self.__class__(
self.array[start:end],
slice_thicknesses=slice_thicknesses,
extent=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 _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)