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