def __init__(self, structure, kpoints, coefficients): logger.info("Initializing wavefunction overlap calculator") self.structure = structure # k-points has to cover the full BZ kpoints = kpoints_to_first_bz(kpoints) mesh_dim = get_mesh_dim_from_kpoints(kpoints, tol=1e-4) round_dp = int(np.log10(1 / 1e-6)) kpoints = np.round(kpoints, round_dp) # get the indices to sort the k-points on the Z, then Y, then X columns sort_idx = np.lexsort((kpoints[:, 2], kpoints[:, 1], kpoints[:, 0])) # put the kpoints into a 3D grid so that they can be indexed as # kpoints[ikx][iky][ikz] = [kx, ky, kz] grid_kpoints = kpoints[sort_idx].reshape(mesh_dim + (3, )) x = grid_kpoints[:, 0, 0, 0] y = grid_kpoints[0, :, 0, 1] z = grid_kpoints[0, 0, :, 2] self.nbands = {s: c.shape[0] for s, c in coefficients.items()} # TODO: Expand the k-point mesh to account for periodic boundary conditions self.interpolators = {} for spin, spin_coefficients in coefficients.items(): nbands = spin_coefficients.shape[0] ncoefficients = spin_coefficients.shape[-1] # sort the coefficients then reshape them into the grid. The coefficients # can now be indexed as coefficients[iband][ikx][iky][ikz] sorted_coefficients = spin_coefficients[:, sort_idx] grid_shape = (nbands, ) + mesh_dim + (ncoefficients, ) grid_coefficients = sorted_coefficients.reshape(grid_shape) if nbands == 1: # this can cause a bug in RegularGridInterpolator. Have to fake # having at least two bands nbands = 2 grid_coefficients = np.tile(grid_coefficients, (2, 1, 1, 1, 1)) if eval_linear: grid = UCGrid( (0, nbands - 1, nbands), (x[0], x[-1], len(x)), (y[0], y[-1], len(y)), (z[0], z[-1], len(z)), ) self.interpolators[spin] = (grid, grid_coefficients) else: interp_range = (np.arange(nbands), x, y, z) self.interpolators[spin] = RegularGridInterpolator( interp_range, grid_coefficients, bounds_error=False, fill_value=None)
def __init__( self, structure, kpoints, projections, band_centers, symprec=defaults["symprec"], kpoint_symmetry_mapping=None, ): logger.info("Initializing orbital overlap calculator") # k-points have to be on a regular grid, even if only the irreducible part of # the grid is used. If the irreducible part is given, we have to expand it # to the full BZ. Also need to expand the projections to the full BZ using # the rotation mapping if kpoint_symmetry_mapping: full_kpoints, ir_to_full_idx, rot_mapping = kpoint_symmetry_mapping else: full_kpoints, ir_to_full_idx, rot_mapping = expand_kpoints( structure, kpoints, symprec=symprec) mesh_dim = get_mesh_dim_from_kpoints(full_kpoints) round_dp = int(np.log10(1 / 1e-6)) full_kpoints = np.round(full_kpoints, round_dp) # get the indices to sort the k-points on the Z, then Y, then X columns sort_idx = np.lexsort( (full_kpoints[:, 2], full_kpoints[:, 1], full_kpoints[:, 0])) # put the kpoints into a 3D grid so that they can be indexed as # kpoints[ikx][iky][ikz] = [kx, ky, kz] grid_kpoints = full_kpoints[sort_idx].reshape(mesh_dim + (3, )) x = grid_kpoints[:, 0, 0, 0] y = grid_kpoints[0, :, 0, 1] z = grid_kpoints[0, 0, :, 2] self.nbands = {s: p.shape[0] for s, p in projections.items()} # TODO: Expand the k-point mesh to account for periodic boundary conditions self.interpolators = {} for spin, spin_projections in projections.items(): nbands = spin_projections.shape[0] nkpoints = len(ir_to_full_idx) nprojections = np.product(spin_projections.shape[2:]) expand_projections = spin_projections[:, ir_to_full_idx] flat_projections = expand_projections.reshape( (nbands, nkpoints, -1), order="F") # aim is to get the wavefunction coefficients norm_projection = (flat_projections / np.sqrt( (flat_projections**2).sum(axis=2))[..., None]) norm_projection[np.isnan(norm_projection)] = 0 coefficients = norm_projection # sort the coefficients then reshape them into the grid. The coefficients # can now be indexed as coefficients[iband][ikx][iky][ikz] sorted_coefficients = coefficients[:, sort_idx] grid_shape = (nbands, ) + mesh_dim + (nprojections, ) grid_coefficients = sorted_coefficients.reshape(grid_shape) if nbands == 1: # this can cause a bug in RegularGridInterpolator. Have to fake # having at least two bands nbands = 2 grid_coefficients = np.tile(grid_coefficients, (2, 1, 1, 1, 1)) interp_range = (np.arange(nbands), x, y, z) self.interpolators[spin] = RegularGridInterpolator( interp_range, grid_coefficients, bounds_error=False, fill_value=None) self.rotation_masks = get_rotation_masks(projections) self.band_centers = band_centers