def generate_slices(self, start=0, end=None, max_batch=1):
for i in range(self.num_slices):
array = np.zeros(self.gpts, dtype=np.float32)
array[:self.gpts[0] // 2] = 1
yield start, end, PotentialArray(array[None],
slice_thicknesses=np.array([self.get_slice_thickness(i)]),
extent=self.extent)
def test_export_import_potential(tmp_path, graphene_atoms):
d = tmp_path / 'sub'
d.mkdir()
path = d / 'potential.hdf5'
potential = Potential(graphene_atoms, sampling=.05)
precalculated_potential = potential.build(pbar=False)
precalculated_potential.write(path)
imported_potential = PotentialArray.read(path)
assert np.allclose(imported_potential.array, precalculated_potential.array)
assert np.allclose(imported_potential.extent,
precalculated_potential.extent)
assert np.allclose(imported_potential._slice_thicknesses,
precalculated_potential._slice_thicknesses)
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