def detect(self, waves, normalize: bool = True) -> np.ndarray:
"""
Integrate the intensity of a the wave functions over the detector range.
Parameters
----------
waves : Waves object
The batch of wave functions to detect.
normalize : bool
Normalize output by the total intensity of the wave function.
Returns
-------
1d array
Detected values as a 1D array. The array has the same length as the batch size of the wave functions.
"""
xp = get_array_module(waves.array)
fft2 = get_device_function(xp, 'fft2')
abs2 = get_device_function(xp, 'abs2')
intensity = abs2(fft2(waves.array, overwrite_x=False))
return self._integrate_array(intensity, waves.angular_sampling,
min(waves.cutoff_scattering_angles),
normalize)
def detect(self, waves) -> np.ndarray:
"""
Integrate the intensity of a the wave functions over the detector range.
Parameters
----------
waves: Waves object
The batch of wave functions to detect.
Returns
-------
2d array
Detected values. The array has shape of (batch size, number of bins).
"""
xp = get_array_module(waves.array)
fft2 = get_device_function(xp, 'fft2')
abs2 = get_device_function(xp, 'abs2')
sum_run_length_encoded = get_device_function(xp, 'sum_run_length_encoded')
intensity = abs2(fft2(waves.array, overwrite_x=False))
indices = self._get_regions(waves.gpts, waves.angular_sampling, min(waves.cutoff_scattering_angles), xp)
total = xp.sum(intensity, axis=(-2, -1))
separators = xp.concatenate((xp.array([0]), xp.cumsum(xp.array([len(ring) for ring in indices]))))
intensity = intensity.reshape((intensity.shape[0], -1))[:, xp.concatenate(indices)]
result = xp.zeros((len(intensity), len(separators) - 1), dtype=xp.float32)
sum_run_length_encoded(intensity, result, separators)
return result
def fft_interpolate_2d(array,
new_shape,
normalization='values',
overwrite_x=False):
xp = get_array_module(array)
fft2 = get_device_function(xp, 'fft2')
ifft2 = get_device_function(xp, 'ifft2')
old_size = array.shape[-2] * array.shape[-1]
if np.iscomplexobj(array):
cropped = fft_crop(fft2(array), new_shape)
array = ifft2(cropped, overwrite_x=overwrite_x)
else:
array = xp.complex64(array)
array = ifft2(fft_crop(fft2(array), new_shape),
overwrite_x=overwrite_x).real
if normalization == 'values':
array *= array.shape[-1] * array.shape[-2] / old_size
elif normalization == 'norm':
array *= array.shape[-1] * array.shape[-2] / old_size
elif (normalization != False) and (normalization != None):
raise RuntimeError()
return array
def calculate_far_field_intensity(waves, overwrite: bool = False):
xp = get_array_module(waves.array)
fft2 = get_device_function(xp, 'fft2')
abs2 = get_device_function(xp, 'abs2')
array = fft2(waves.array, overwrite_x=overwrite)
intensity = crop_to_center(xp.fft.fftshift(array, axes=(-2, -1)))
return abs2(intensity)
def detect(self, waves) -> np.ndarray:
xp = get_array_module(waves.array)
abs2 = get_device_function(xp, 'abs2')
fft2 = get_device_function(xp, 'fft2')
intensity = abs2(fft2(waves.array, overwrite_x=False))
intensity = xp.fft.fftshift(intensity, axes=(-1, -2))
intensity = crop_to_center(intensity)
return intensity
def apply_ctf(self, ctf: CTF = None, **kwargs):
"""
Apply the aberrations defined by a CTF object to wave function.
:param ctf: Contrast Transfer Function object to be applied.
:param kwargs: Provide the aberration coefficients as keyword arguments.
:return: The wave functions with aberrations applied.
"""
xp = get_array_module(self.array)
fft2_convolve = get_device_function(get_array_module(self.array),
'fft2_convolve')
if ctf is None:
ctf = CTF(**kwargs)
ctf.accelerator.match(self.accelerator)
kx, ky = spatial_frequencies(self.grid.gpts, self.grid.sampling)
alpha, phi = polargrid(xp.asarray(kx * self.wavelength),
xp.asarray(ky * self.wavelength))
kernel = ctf.evaluate(alpha, phi)
return self.__class__(fft2_convolve(self.array,
kernel,
overwrite_x=False),
extent=self.extent,
energy=self.energy)
def _bandlimit(self, array):
xp = get_array_module(array)
fft2_convolve = get_device_function(xp, 'fft2_convolve')
fft2_convolve(array,
self.get_mask(array.shape[-2:], (1, 1), xp),
overwrite_x=True)
return array
def fourier_translation_operator(positions: np.ndarray, shape: tuple):
positions_shape = positions.shape
if len(positions_shape) == 1:
positions = positions[None]
xp = get_array_module(positions)
complex_exponential = get_device_function(xp, 'complex_exponential')
kx, ky = spatial_frequencies(shape, (1., 1.))
kx = kx.reshape((1, -1, 1))
ky = ky.reshape((1, 1, -1))
kx = xp.asarray(kx)
ky = xp.asarray(ky)
positions = xp.asarray(positions)
x = positions[:, 0].reshape((-1, ) + (1, 1))
y = positions[:, 1].reshape((-1, ) + (1, 1))
result = complex_exponential(-2 * np.pi * kx * x) * complex_exponential(
-2 * np.pi * ky * y)
if len(positions_shape) == 1:
return result[0]
else:
return result
def detect(self, waves) -> np.ndarray:
"""
Calculate the far field intensity of the wave functions. The output is cropped to include the non-suppressed
frequencies from the antialiased 2D fourier spectrum.
Parameters
----------
waves: Waves object
The batch of wave functions to detect.
Returns
-------
Detected values. The first dimension indexes the batch size, the second and third indexes the two components
of the spatial frequency.
"""
xp = get_array_module(waves.array)
abs2 = get_device_function(xp, 'abs2')
waves = waves.far_field(max_angle=self.max_angle)
intensity = abs2(waves.array)
intensity = xp.fft.fftshift(intensity, axes=(-2, -1))
intensity = self._interpolate(intensity, waves.angular_sampling)
return intensity
def generate_slices(self, start=0, end=None):
if end is None:
end = len(self)
self.grid.check_is_defined()
xp = get_array_module(self._device)
interpolate_radial_functions = get_device_function(
xp, 'interpolate_radial_functions')
atoms = self.atoms.copy()
indices_by_number = {
number: np.where(atoms.numbers == number)[0]
for number in np.unique(atoms.numbers)
}
array = xp.zeros(self.gpts, dtype=xp.float32)
a = np.sum([self.get_slice_thickness(i) for i in range(0, start)])
for i in range(start, end):
array[:] = 0.
b = a + self.get_slice_thickness(i)
for number, indices in indices_by_number.items():
slice_atoms = atoms[indices]
integrator, disc_indices = self.get_integrator(number)
disc_indices = xp.asarray(disc_indices)
slice_atoms = slice_atoms[
(slice_atoms.positions[:, 2] > a - integrator.cutoff) *
(slice_atoms.positions[:, 2] < b + integrator.cutoff)]
slice_atoms = pad_atoms(slice_atoms, integrator.cutoff)
if len(slice_atoms) == 0:
continue
vr = np.zeros((len(slice_atoms), len(integrator.r)),
np.float32)
dvdr = np.zeros((len(slice_atoms), len(integrator.r)),
np.float32)
for j, atom in enumerate(slice_atoms):
am, bm = a - atom.z, b - atom.z
vr[j], dvdr[j, :-1] = integrator.integrate(am, bm)
vr = xp.asarray(vr, dtype=xp.float32)
dvdr = xp.asarray(dvdr, dtype=xp.float32)
r = xp.asarray(integrator.r, dtype=xp.float32)
slice_positions = xp.asarray(slice_atoms.positions[:, :2],
dtype=xp.float32)
sampling = xp.asarray(self.sampling, dtype=xp.float32)
interpolate_radial_functions(array, disc_indices,
slice_positions, vr, r, dvdr,
sampling)
a = b
yield ProjectedPotential(array / kappa,
self.get_slice_thickness(i),
extent=self.extent)
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 _evaluate_propagator_array(self, gpts, sampling, wavelength, dz, xp):
complex_exponential = get_device_function(xp, 'complex_exponential')
kx = xp.fft.fftfreq(gpts[0], sampling[0]).astype(xp.float32)
ky = xp.fft.fftfreq(gpts[1], sampling[1]).astype(xp.float32)
f = (complex_exponential(-(kx**2)[:, None] * np.pi * wavelength * dz) *
complex_exponential(-(ky**2)[None] * np.pi * wavelength * dz))
f *= xp.asarray(self._antialiasing_aperture(gpts))
return f
def as_transmission_functions(self, energy):
xp = get_array_module(self.array)
complex_exponential = get_device_function(xp, 'complex_exponential')
array = complex_exponential(energy2sigma(energy) * self._array)
return TransmissionFunctions(array,
slice_thicknesses=self._slice_thicknesses,
extent=self.extent,
energy=energy)
def detect(self, waves) -> np.ndarray:
"""
Integrate the intensity of a the wave functions over the detector range.
Parameters
----------
waves: Waves object
The batch of wave functions to detect.
Returns
-------
3d array
Detected values. The first dimension indexes the batch size, the second and third indexes the radial and
angular bins, respectively.
"""
xp = get_array_module(waves.array)
fft2 = get_device_function(xp, 'fft2')
abs2 = get_device_function(xp, 'abs2')
sum_run_length_encoded = get_device_function(xp,
'sum_run_length_encoded')
intensity = abs2(fft2(waves.array, overwrite_x=False))
indices = self._get_regions(waves.gpts, waves.angular_sampling,
min(waves.cutoff_scattering_angles), xp)
separators = xp.concatenate(
(xp.array([0]), xp.cumsum(xp.array([len(ring)
for ring in indices]))))
intensity = intensity.reshape(
(intensity.shape[0], -1))[:, xp.concatenate(indices)]
result = xp.zeros((len(intensity), len(separators) - 1),
dtype=xp.float32)
sum_run_length_encoded(intensity, result, separators)
shape = (-1, )
if self.nbins_radial > 1:
shape += (self.nbins_radial, )
if self.nbins_angular > 1:
shape += (self.nbins_angular, )
return result.reshape(shape)
def intensity(self) -> Measurement:
"""
:return: The intensity of the wave functions at the image plane.
"""
calibrations = calibrations_from_grid(self.grid.gpts,
self.grid.sampling, ['x', 'y'])
calibrations = (None, ) * (len(self.array.shape) - 2) + calibrations
abs2 = get_device_function(get_array_module(self.array), 'abs2')
return Measurement(abs2(self.array), calibrations)
def _interpolate(self, array, angular_sampling):
xp = get_array_module(array)
interpolate_bilinear = get_device_function(xp, 'interpolate_bilinear')
new_gpts, new_angular_sampling = self._resampled_gpts(array.shape[-2:], angular_sampling)
v, u, vw, uw = self._bilinear_nodes_and_weight(array.shape[-2:],
new_gpts,
angular_sampling,
new_angular_sampling,
xp)
return interpolate_bilinear(array, v, u, vw, uw)
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 diffraction_pattern(self) -> Measurement:
"""
:return: The intensity of the wave functions at the diffraction plane.
"""
calibrations = calibrations_from_grid(self.grid.antialiased_gpts,
self.grid.antialiased_sampling,
names=['alpha_x', 'alpha_y'],
units='mrad',
scale_factor=self.wavelength *
1000,
fourier_space=True)
calibrations = (None, ) * (len(self.array.shape) - 2) + calibrations
xp = get_array_module(self.array)
abs2 = get_device_function(xp, 'abs2')
fft2 = get_device_function(xp, 'fft2')
pattern = asnumpy(
abs2(
crop_to_center(
xp.fft.fftshift(fft2(self.array, overwrite_x=False)))))
return Measurement(pattern, calibrations)
def _fourier_translation_operator(self, positions):
xp = get_array_module(positions)
complex_exponential = get_device_function(xp, 'complex_exponential')
kx, ky = spatial_frequencies(self.grid.gpts, self.grid.sampling)
kx = kx.reshape((1, -1, 1))
ky = ky.reshape((1, 1, -1))
kx = xp.asarray(kx)
ky = xp.asarray(ky)
positions = xp.asarray(positions)
x = positions[:, 0].reshape((-1, ) + (1, 1))
y = positions[:, 1].reshape((-1, ) + (1, 1))
return complex_exponential(2 * np.pi * kx * x) * complex_exponential(
2 * np.pi * ky * y)
def propagate(self, waves: Union['Waves', 'SMatrix', 'PartialSMatrix'],
dz: float):
"""
Propgate wave function or scattering matrix.
:param waves: Wave function or scattering matrix to propagate.
:param dz: Propagation distance [Å].
"""
propagator_array = self._evaluate_propagator_array(
waves.grid.gpts, waves.grid.sampling, dz, waves.wavelength,
get_array_module(waves.array))
fft2_convolve = get_device_function(get_array_module(waves.array),
'fft2_convolve')
fft2_convolve(waves._array, propagator_array)
def bandlimit(self, waves):
"""
Parameters
----------
waves
Returns
-------
"""
xp = get_array_module(waves.array)
fft2_convolve = get_device_function(xp, 'fft2_convolve')
fft2_convolve(waves.array,
self.get_mask(waves.gpts, waves.sampling, xp),
overwrite_x=True)
return waves
def transmit(waves: Union['Waves', 'SMatrix', 'PartialSMatrix'],
potential_slice: ProjectedPotential):
"""
Transmit wave function or scattering matrix.
:param waves: Wave function or scattering matrix to propagate.
:param potential_slice: Projected potential to transmit the wave function through.
"""
xp = get_array_module(waves.array)
complex_exponential = get_device_function(xp, 'complex_exponential')
dim_padding = len(waves._array.shape) - len(potential_slice.array.shape)
slice_array = potential_slice.array.reshape((1, ) * dim_padding +
potential_slice.array.shape)
if np.iscomplexobj(slice_array):
waves._array *= copy_to_device(slice_array, xp)
else:
waves._array *= complex_exponential(
copy_to_device(waves.accelerator.sigma * slice_array, xp))
def as_transmission_function(self,
energy: float,
in_place: bool = True,
max_batch: int = 1,
antialias_filter: AntialiasFilter = None):
"""
Calculate the transmission functions for a specific energy.
Parameters
----------
energy: float
Electron energy [eV].
Returns
-------
TransmissionFunction object
"""
xp = get_array_module(self.array)
complex_exponential = get_device_function(xp, 'complex_exponential')
array = self._array
if not in_place:
array = array.copy()
array = complex_exponential(energy2sigma(energy) * array)
t = TransmissionFunction(
array,
slice_thicknesses=self._slice_thicknesses.copy(),
extent=self.extent,
energy=energy)
if antialias_filter is None:
antialias_filter = AntialiasFilter()
for start, end, potential_slices in t.generate_slices(
max_batch=max_batch):
antialias_filter.bandlimit(potential_slices)
return t
def fourier_translation_operator(positions: np.ndarray,
shape: tuple) -> np.ndarray:
"""
Create an array representing one or more phase ramp(s) for shifting another array.
Parameters
----------
positions : array of xy-positions
shape : two int
Returns
-------
"""
positions_shape = positions.shape
if len(positions_shape) == 1:
positions = positions[None]
xp = get_array_module(positions)
complex_exponential = get_device_function(xp, 'complex_exponential')
kx, ky = spatial_frequencies(shape, (1., 1.))
kx = kx.reshape((1, -1, 1))
ky = ky.reshape((1, 1, -1))
kx = xp.asarray(kx)
ky = xp.asarray(ky)
positions = xp.asarray(positions)
x = positions[:, 0].reshape((-1, ) + (1, 1))
y = positions[:, 1].reshape((-1, ) + (1, 1))
result = complex_exponential(-2 * np.pi * kx * x) * complex_exponential(
-2 * np.pi * ky * y)
if len(positions_shape) == 1:
return result[0]
else:
return result
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
def evaluate_aberrations(self, alpha: Union[float, np.ndarray], phi: Union[float, np.ndarray]) -> \
Union[float, np.ndarray]:
xp = get_array_module(alpha)
complex_exponential = get_device_function(xp, 'complex_exponential')
return complex_exponential(-self.evaluate_chi(alpha, phi))
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 measure(self):
array = np.fft.fftshift(self.build())[0]
calibrations = calibrations_from_grid(self.gpts, self.sampling,
['x', 'y'])
abs2 = get_device_function(get_array_module(array), 'abs2')
return Measurement(array, calibrations, name=str(self))
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:
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)
