""" grider.py Generates grid for plotting """ import numpy as np import warnings from opt_einsum import contract import psi4 psi4.core.be_quiet() try: from pylibxc import LibXCFunctional as Functional except: pass # from .cubeprop import Cubeprop from .basis_set_artifact_correction import basis_set_correction, invert_kohn_sham_equations class Grider(): def grid_to_blocks(self, grid, basis=None): """ Generate list of blocks to allocate given grid Parameters ---------- grid: np.ndarray Grid to be distributed into blocks Size: (3, npoints) for homogeneous grid (4, npoints) for inhomogenous grid to account for weights basis: psi4.core.BasisSet; optional The basis set. If not given, it will use target wfn.basisset(). Returns ------- blocks: list List with psi4.core.BlockOPoints npoints: int Total number of points (for one dimension) points: psi4.core.{RKS, UKS} Points function to set matrices. """ assert (grid.shape[0] == 3) or (grid.shape[0] == 4), """Grid does not have the correct dimensions. \n Array must be of size (3, npoints) or (4, npoints)""" if_w = grid.shape[0] == 4 if basis is None: #basis = self.basis #added by Ehsan basis = psi4.core.BasisSet.build(self.molecule, "ORBITAL", self.basis) epsilon = psi4.core.get_global_option("CUBIC_BASIS_TOLERANCE") # added by Ehsan #print("\nThis is the epsilon: ", float(epsilon)) # basis_set = psi4.core.BasisSet.build(self.molecule, "ORBITAL", basis) extens = psi4.core.BasisExtents(basis, epsilon) #extens = psi4.core.BasisExtents(psi4.core.BasisExtents.basis(), 0.004) max_points = psi4.core.get_global_option("DFT_BLOCK_MAX_POINTS") npoints = grid.shape[1] nblocks = int(np.floor(npoints/max_points)) blocks = [] max_functions = 0 #Run through full blocks idx = 0 for nb in range(nblocks): x = psi4.core.Vector.from_array(grid[0][idx : idx + max_points]) y = psi4.core.Vector.from_array(grid[1][idx : idx + max_points]) z = psi4.core.Vector.from_array(grid[2][idx : idx + max_points]) if if_w: w = psi4.core.Vector.from_array(grid[3][idx : idx + max_points]) else: w = psi4.core.Vector.from_array(np.zeros(max_points)) # When w is not necessary and not given blocks.append(psi4.core.BlockOPoints(x, y, z, w, extens)) idx += max_points max_functions = max_functions if max_functions > len(blocks[-1].functions_local_to_global()) \ else len(blocks[-1].functions_local_to_global()) #Run through remaining points if idx < npoints: x = psi4.core.Vector.from_array(grid[0][idx:]) y = psi4.core.Vector.from_array(grid[1][idx:]) z = psi4.core.Vector.from_array(grid[2][idx:]) if if_w: w = psi4.core.Vector.from_array(grid[3][idx:]) else: w = psi4.core.Vector.from_array(np.zeros_like(grid[2][idx:])) # When w is not necessary and not given blocks.append(psi4.core.BlockOPoints(x, y, z, w, extens)) max_functions = max_functions if max_functions > len(blocks[-1].functions_local_to_global()) \ else len(blocks[-1].functions_local_to_global()) zero_matrix = psi4.core.Matrix(basis.nbf(), basis.nbf()) if self.ref == 1: point_func = psi4.core.RKSFunctions(basis, max_points, max_functions) point_func.set_pointers(zero_matrix) else: point_func = psi4.core.UKSFunctions(basis, max_points, max_functions) point_func.set_pointers(zero_matrix, zero_matrix) return blocks, npoints, point_func def generate_grids(self, x, y, z): """ Genrates Mesh from 3 separate linear spaces and flatten, needed for cubic grid. Parameters ---------- grid: tuple of three np.ndarray (x, y, z) Returns ------- grid: np.ndarray shape (3, len(x)*len(y)*len(z)). """ # x,y,z, = grid shape = (len(x), len(y), len(z)) X,Y,Z = np.meshgrid(x, y, z, indexing='ij') X = X.reshape((X.shape[0] * X.shape[1] * X.shape[2], 1)) Y = Y.reshape((Y.shape[0] * Y.shape[1] * Y.shape[2], 1)) Z = Z.reshape((Z.shape[0] * Z.shape[1] * Z.shape[2], 1)) grid = np.concatenate((X,Y,Z), axis=1).T return grid, shape def generate_dft_grid(self, Vpot): """ Extracts DFT spherical grid and weights from wfn object Parameters ---------- Vpot: psi4.core.VBase Vpot object with dft grid data Returns ------- dft_grid: list Numpy arrays corresponding to x,y,z, and w. Shape: (4, npoints) """ nblocks = Vpot.nblocks() blocks = [Vpot.get_block(i) for i in range(nblocks)] npoints = Vpot.grid().npoints() dft_grid = np.zeros((4, npoints)) offset = 0 for i_block in blocks: b_points = i_block.npoints() offset += b_points dft_grid[0, offset - b_points : offset] = i_block.x().np dft_grid[1, offset - b_points : offset] = i_block.y().np dft_grid[2, offset - b_points : offset] = i_block.z().np dft_grid[3, offset - b_points : offset] = i_block.w().np return dft_grid #Quantities on Grid def on_grid_ao(self, coeff, grid=None, basis=None, Vpot=None): """ Generates a quantity on the grid given its ao representation. *This is the most general function for basis to grid transformation. Parameters ---------- coeff: np.ndarray Vector/Matrix of quantity on ao basis. Shape: {(num_ao_basis, ), (num_ao_basis, num_ao_basis)} grid: np.ndarray Shape: (3, npoints) or (4, npoints) or tuple for block_handler (return of grid_to_blocks) grid where density will be computed. basis: psi4.core.BasisSet, optional The basis set. If not given it will use target wfn.basisset(). Vpot: psi4.core.VBase Vpotential object with info about grid. Provides DFT spherical grid. Only comes to play if no grid is given. Returns ------- coeff_r: np.ndarray Shape: (npoints, ) Quantity expressed by the coefficient on the given grid """ if grid is not None: if type(grid) is np.ndarray: if grid.shape[0] != 3 and grid.shape[0] != 4: raise ValueError("The shape of grid should be (3, npoints) " "or (4, npoints) but got (%i, %i)" % (grid.shape[0], grid.shape[1])) blocks, npoints, points_function = self.grid_to_blocks(grid, basis=basis) else: blocks, npoints, points_function = grid elif grid is None and Vpot is not None: nblocks = Vpot.nblocks() blocks = [Vpot.get_block(i) for i in range(nblocks)] npoints = Vpot.grid().npoints() points_function = Vpot.properties()[0] else: raise ValueError("A grid or a V_potential (DFT grid) must be given.") coeff_r = np.zeros((npoints)) offset = 0 for i_block in blocks: points_function.compute_points(i_block) b_points = i_block.npoints() offset += b_points lpos = np.array(i_block.functions_local_to_global()) if len(lpos)==0: continue phi = np.array(points_function.basis_values()["PHI"])[:b_points, :lpos.shape[0]] if coeff.ndim == 1: l_mat = coeff[(lpos[:])] coeff_r[offset - b_points : offset] = contract('pm,m->p', phi, l_mat) elif coeff.ndim == 2: l_mat = coeff[(lpos[:, None], lpos)] coeff_r[offset - b_points : offset] = contract('pm,mn,pn->p', phi, l_mat, phi) return coeff_r def on_grid_density(self, grid=None, Da=None, Db=None, Vpot=None): """ Generates Density given grid Parameters ---------- Da, Db: np.ndarray Alpha, Beta densities. Shape: (num_ao_basis, num_ao_basis) grid: np.ndarray Shape: (3, npoints) or (4, npoints) or tuple for block_handler (return of grid_to_blocks) grid where density will be computed. Vpot: psi4.core.VBase Vpotential object with info about grid. Provides DFT spherical grid. Only comes to play if no grid is given. Returns ------- density: np.ndarray Shape: (ref, npoints) Density on the given grid. """ if Da is None and Db is None: Da = psi4.core.Matrix.from_array(self.Dt[0]) Db = psi4.core.Matrix.from_array(self.Dt[1]) else: Da = psi4.core.Matrix.from_array(Da) Db = psi4.core.Matrix.from_array(Db) if self.ref == 2 and Db is None: raise ValueError("Db is required for an unrestricted system") if grid is not None: if type(grid) is np.ndarray: if grid.shape[0] != 3 and grid.shape[0] != 4: raise ValueError("The shape of grid should be (3, npoints) " "or (4, npoints) but got (%i, %i)" % (grid.shape[0], grid.shape[1])) blocks, npoints, points_function = self.grid_to_blocks(grid) else: blocks, npoints, points_function = grid elif grid is None and Vpot is not None: nblocks = Vpot.nblocks() blocks = [Vpot.get_block(i) for i in range(nblocks)] npoints = Vpot.grid().npoints() points_function = Vpot.properties()[0] else: raise ValueError("A grid or a V_potential (DFT grid) must be given.") if self.ref == 1: points_function.set_pointers(Da) rho_a = points_function.point_values()["RHO_A"] density = np.zeros((npoints)) if self.ref == 2: points_function.set_pointers(Da, Db) rho_a = points_function.point_values()["RHO_A"] rho_b = points_function.point_values()["RHO_B"] density = np.zeros((npoints, self.ref)) offset = 0 for i_block in blocks: points_function.compute_points(i_block) b_points = i_block.npoints() offset += b_points if self.ref == 1: density[offset - b_points : offset] = rho_a.np[ :b_points] else: density[offset - b_points : offset, 0] = rho_a.np[ :b_points] density[offset - b_points : offset, 1] = rho_b.np[ :b_points] return density def on_grid_orbitals(self, Ca=None, Cb=None, grid=None, Vpot=None): """ Generates orbitals given grid Parameters ---------- Ca, Cb: np.ndarray Alpha, Beta Orbital Coefficient Matrix. Shape: (num_ao_basis, num_ao_basis) grid: np.ndarray Shape: (3, npoints) or (4, npoints) or tuple for block_handler (return of grid_to_blocks) grid where density will be computed Vpot: psi4.core.VBase Vpotential object with info about grid. Provides DFT spherical grid. Only comes to play if no grid is given. Returns ------- orbitals: np.ndarray Orbitals on the given grid of size . Shape: (nbasis, npoints, ref) """ if Ca is None and Cb is None: Ca = psi4.core.Matrix.from_array(self.Ca) Cb = psi4.core.Matrix.from_array(self.Cb) else: Ca = psi4.core.Matrix.from_array(Ca) Cb = psi4.core.Matrix.from_array(Cb) if self.ref == 2 and Cb is None: raise ValueError("Db is required for an unrestricted system") if grid is not None: if type(grid) is np.ndarray: if grid.shape[0] != 3 and grid.shape[0] != 4: raise ValueError("The shape of grid should be (3, npoints) " "or (4, npoints) but got (%i, %i)" % (grid.shape[0], grid.shape[1])) blocks, npoints, points_function = self.grid_to_blocks(grid) else: blocks, npoints, points_function = grid elif grid is None and Vpot is not None: nblocks = Vpot.nblocks() blocks = [Vpot.get_block(i) for i in range(nblocks)] npoints = Vpot.grid().npoints() points_function = Vpot.properties()[0] else: raise ValueError("A grid or a V_potential (DFT grid) must be given.") if self.ref == 1: orbitals_r = [np.zeros((npoints)) for i_orb in range(self.nbf)] points_function.set_pointers(Ca) Ca_np = Ca.np if self.ref == 2: orbitals_r = [np.zeros((npoints, 2)) for i_orb in range(self.nbf)] points_function.set_pointers(Ca, Cb) Ca_np = Ca.np Cb_np = Cb.np offset = 0 for i_block in blocks: points_function.compute_points(i_block) b_points = i_block.npoints() offset += b_points lpos = np.array(i_block.functions_local_to_global()) if len(lpos)==0: continue phi = np.array(points_function.basis_values()["PHI"])[:b_points, :lpos.shape[0]] for i_orb in range(self.nbf): Ca_local = Ca_np[lpos, i_orb] if self.ref == 1: orbitals_r[i_orb][offset - b_points : offset] = contract('m, pm -> p', Ca_local, phi) else: Cb_local = Cb_np[lpos, i_orb] orbitals_r[i_orb][offset - b_points : offset,0] = contract('m, pm -> p', Ca_local, phi) orbitals_r[i_orb][offset - b_points : offset,1] = contract('m, pm -> p', Cb_local, phi) return orbitals_r def on_grid_esp(self, Da=None, Db=None, grid=None, Vpot=None, wfn=None): """ Generates EXTERNAL/ESP/HARTREE and Fermi Amaldi Potential on given grid Parameters ---------- Da,Db: np.ndarray, opt, shape (nbf, nbf) The electron density in the denominator of Hartee potential. If None, the original density matrix will be used. grid: np.ndarray Shape: (3, npoints) or (4, npoints) or tuple for block_handler (return of grid_to_blocks) grid where density will be computed. Vpot: psi4.core.VBase Vpotential object with info about grid. Provides DFT spherical grid. Only comes to play if no grid is given. Returns ------- vext, hartree, esp, v_fa: np.ndarray External, Hartree, ESP, and Fermi Amaldi potential on the given grid Shape: (npoints, ) """ if wfn is None: wfn = self.wfn if Da is not None or Db is not None: Da_temp = np.copy(self.wfn.Da().np) Db_temp = np.copy(self.wfn.Db().np) if Da is not None: wfn.Da().np[:] = Da if Db is not None: wfn.Db().np[:] = Db nthreads = psi4.get_num_threads() psi4.set_num_threads(1) if grid is not None: if type(grid) is np.ndarray: blocks, npoints, points_function = self.grid_to_blocks(grid) else: blocks, npoints, points_function = grid elif grid is None and Vpot is not None: nblocks = Vpot.nblocks() blocks = [Vpot.get_block(i) for i in range(nblocks)] npoints = Vpot.grid().npoints() else: raise ValueError("A grid or a V_potential (DFT grid) must be given.") #Initialize Arrays vext = np.zeros(npoints) esp = np.zeros(npoints) #Get Atomic Information mol_dict = self.mol.to_schema(dtype='psi4') natoms = len(mol_dict["elem"]) indx = [i for i in range(natoms) if self.mol.charge(i) != 0.0] natoms = len(indx) #Atomic numbers and Atomic positions zs = [mol_dict["elez"][i] for i in indx] rs = [self.mol.geometry().np[i] for i in indx] esp_wfn = psi4.core.ESPPropCalc(wfn) #Loop Through blocks offset = 0 with np.errstate(divide='ignore'): for i_block in blocks: b_points = i_block.npoints() offset += b_points x = i_block.x().np y = i_block.y().np z = i_block.z().np #EXTERNAL for atom in range(natoms): r = np.sqrt((x-rs[atom][0])**2 + (y-rs[atom][1])**2 + (z-rs[atom][2])**2) vext_temp = - 1.0 * zs[atom] / r vext_temp[np.isinf(vext_temp)] = 0.0 vext[offset - b_points : offset] += vext_temp #ESP xyz = np.concatenate((x[:,None],y[:,None],z[:,None]), axis=1) grid_block = psi4.core.Matrix.from_array(xyz) esp[offset - b_points : offset] = esp_wfn.compute_esp_over_grid_in_memory(grid_block).np #Hartree hartree = - 1.0 * (vext + esp) v_fa = (1 - 1.0 / (self.nalpha + self.nbeta)) * hartree if Da is not None: wfn.Da().np[:] = Da_temp if Db is not None: wfn.Db().np[:] = Db_temp psi4.set_num_threads(nthreads) return vext, hartree, v_fa, esp def on_grid_vxc(self, func_id=1, grid=None, Da=None, Db=None, Vpot=None): """ Generates Vxc given grid Parameters ---------- Da, Db: np.ndarray Alpha, Beta densities. Shape: (num_ao_basis, num_ao_basis) func_id: int Functional ID associated with Density Functional Approximationl. Full list of functionals: https://www.tddft.org/programs/libxc/functionals/ grid: np.ndarray Shape: (3, npoints) or (4, npoints) or tuple for block_handler (return of grid_to_blocks) grid where density will be computed. Vpot: psi4.core.VBase Vpotential object with info about grid. Provides DFT spherical grid. Only comes to play if no grid is given. Returns ------- VXC: np.ndarray Exchange correlation potential on the given grid Shape: (npoints, ) """ local_functionals = [1,546,549,532,692,641,552,287,307,578,5,24,4,579,308,289,551, 22,23,14,11,574,573,554,5900,12,13,25,9,10,27,3,684,683,17,7, 28,29,30,31,8,317,2,6,536,537,538,318,577,259,547,548,20,599,43, 51,580,50,550 ] if func_id not in local_functionals: raise ValueError("Only LDA fucntionals are supported on the grid") if Da is None: Da = self.Dt[0] if Db is None: Db = self.Dt[0] if grid is not None: if type(grid) is np.ndarray: blocks, npoints, points_function = self.grid_to_blocks(grid) else: blocks, npoints, points_function = grid density = self.on_grid_density(Da=Da, Db=Db, grid=grid) elif grid is None and Vpot is not None: nblocks = Vpot.nblocks() blocks = [Vpot.get_block(i) for i in range(nblocks)] npoints = Vpot.grid().npoints() density = self.on_grid_density(Da=Da, Db=Db, Vpot=Vpot) else: raise ValueError("A grid or a V_potential (DFT grid) must be given.") vxc = np.zeros((npoints, self.ref)) ingredients = {} offset = 0 for i_block in blocks: b_points = i_block.npoints() offset += b_points if self.ref == 1: ingredients["rho"] = density[offset - b_points : offset] else: ingredients["rho"] = density[offset - b_points : offset, :] if self.ref == 1: functional = Functional(func_id, 1) else: functional = Functional(func_id, 2) xc_dictionary = functional.compute(ingredients) vxc[offset - b_points : offset, :] = xc_dictionary['vrho'] return np.squeeze(vxc) def on_grid_lap_phi(self, Ca=None, Cb=None, grid=None, Vpot=None): """ Generates laplacian of molecular orbitals Parameters ---------- Ca, Cb: np.ndarray Alpha, Beta Orbital Coefficient Matrix. Shape: (num_ao_basis, num_ao_basis) grid: np.ndarray Shape: (3, npoints) or (4, npoints) or tuple for block_handler (return of grid_to_blocks) grid where density will be computed. Vpot: psi4.core.VBase Vpotential object with info about grid. Provides DFT spherical grid. Only comes to play if no grid is given. Returns ------- lap_phi: List[np.ndarray]. Where array is of shape (npoints, ref) Laplacian of molecular orbitals on the grid """ if Ca is None and Cb is None: Ca = psi4.core.Matrix.from_array(self.Ca) Cb = psi4.core.Matrix.from_array(self.Cb) else: Ca = psi4.core.Matrix.from_array(Ca) Cb = psi4.core.Matrix.from_array(Cb) if self.ref == 2 and Cb is None: raise ValueError("Db is required for an unrestricted system") if grid is not None: if type(grid) is np.ndarray: if grid.shape[0] != 3 and grid.shape[0] != 4: raise ValueError("The shape of grid should be (3, npoints) " "or (4, npoints) but got (%i, %i)" % (grid.shape[0], grid.shape[1])) blocks, npoints, points_function = self.grid_to_blocks(grid) else: blocks, npoints, points_function = grid elif grid is None and Vpot is not None: nblocks = Vpot.nblocks() blocks = [Vpot.get_block(i) for i in range(nblocks)] npoints = Vpot.grid().npoints() points_function = Vpot.properties()[0] else: raise ValueError("A grid or a V_potential (DFT grid) must be given.") points_function.set_ansatz(2) if self.ref == 1: points_function.set_pointers(Ca) lap_phi = [np.zeros((npoints)) for i_orb in range(self.nbf)] else: points_function.set_pointers(Ca, Cb) lap_phi = [np.zeros((npoints, 2)) for i_orb in range(self.nbf)] offset = 0 for i_block in blocks: points_function.compute_points(i_block) b_points = i_block.npoints() offset += b_points lpos = np.array(i_block.functions_local_to_global()) if len(lpos)==0: continue #Obtain subset of phi_@@ matrices lx = np.array(points_function.basis_values()["PHI_XX"])[:b_points, :lpos.shape[0]] ly = np.array(points_function.basis_values()["PHI_YY"])[:b_points, :lpos.shape[0]] lz = np.array(points_function.basis_values()["PHI_ZZ"])[:b_points, :lpos.shape[0]] for i_orb in range(self.nbf): Ca_local = Ca.np[lpos, i_orb][:,None] if self.ref ==1: lap_phi[i_orb][offset - b_points : offset] += ((lx + ly + lz) @ Ca_local)[:,0] else: Cb_local = Cb.np[lpos, i_orb][:,None] lap_phi[i_orb][offset - b_points : offset, 0] += ((lx + ly + lz) @ Ca_local)[:,0] lap_phi[i_orb][offset - b_points : offset, 1] += ((lx + ly + lz) @ Cb_local)[:,0] return lap_phi def on_grid_grad_phi(self, Ca=None, Cb=None, grid=None, Vpot=None): """ Generates laplacian of molecular orbitals Parameters ---------- Ca, Cb: np.ndarray Alpha, Beta Orbital Coefficient Matrix. Shape: (num_ao_basis, num_ao_basis) grid: np.ndarray Shape: (3, npoints) or (4, npoints) or tuple for block_handler (return of grid_to_blocks) grid where density will be computed. Vpot: psi4.core.VBase Vpotential object with info about grid. Provides DFT spherical grid. Only comes to play if no grid is given. Returns ------- grad_phi: List[np.ndarray]. Where array is of shape (npoints, ref) Gradient of molecular orbitals on the grid """ if Ca is None and Cb is None: Ca = psi4.core.Matrix.from_array(self.Ca) Cb = psi4.core.Matrix.from_array(self.Cb) else: Ca = psi4.core.Matrix.from_array(Ca) Cb = psi4.core.Matrix.from_array(Cb) if self.ref == 2 and Cb is None: raise ValueError("Db is required for an unrestricted system") if grid is not None: if type(grid) is np.ndarray: if grid.shape[0] != 3 and grid.shape[0] != 4: raise ValueError("The shape of grid should be (3, npoints) " "or (4, npoints) but got (%i, %i)" % (grid.shape[0], grid.shape[1])) blocks, npoints, points_function = self.grid_to_blocks(grid) else: blocks, npoints, points_function = grid elif grid is None and Vpot is not None: nblocks = Vpot.nblocks() blocks = [Vpot.get_block(i) for i in range(nblocks)] npoints = Vpot.grid().npoints() points_function = Vpot.properties()[0] else: raise ValueError("A grid or a V_potential (DFT grid) must be given.") points_function.set_ansatz(2) if self.ref == 1: points_function.set_pointers(Ca) grad_phi = [np.zeros((npoints)) for i_orb in range(self.nbf)] else: points_function.set_pointers(Ca, Cb) grad_phi = [np.zeros((npoints, 2)) for i_orb in range(self.nbf)] offset = 0 for i_block in blocks: points_function.compute_points(i_block) b_points = i_block.npoints() offset += b_points lpos = np.array(i_block.functions_local_to_global()) if len(lpos)==0: continue #Obtain subset of phi_@ matrix gx = np.array(points_function.basis_values()["PHI_X"])[:b_points, :lpos.shape[0]] gy = np.array(points_function.basis_values()["PHI_Y"])[:b_points, :lpos.shape[0]] gz = np.array(points_function.basis_values()["PHI_Z"])[:b_points, :lpos.shape[0]] for i_orb in range(self.nbf): Ca_local = Ca.np[lpos, i_orb][:,None] if self.ref == 1: grad_phi[i_orb][offset - b_points : offset] += ((gx + gy + gz) @ Ca_local)[:,0] if self.ref == 2: Cb_local = Cb.np[lpos, i_orb][:,None] grad_phi[i_orb][offset - b_points : offset, 0] += ((gx + gy + gz) @ Ca_local)[:,0] grad_phi[i_orb][offset - b_points : offset, 1] += ((gx + gy + gz) @ Cb_local)[:,0] return grad_phi def dft_grid_to_fock(self, value, Vpot): """For value on DFT spherical grid, Fock matrix is returned. VFock_ij = \int dx \phi_i(x) \phi_j(x) value(x) Parameters: ----------- value: np.ndarray of shape (npoint, ). Vpot:psi4.core.VBase Vpotential object with info about grid. Provides DFT spherical grid. Only comes to play if no grid is given. Returns: --------- VFock: np.ndarray of shape (nbasis, nbasis) """ VFock = np.zeros((self.nbf, self.nbf)) points_func = Vpot.properties()[0] i = 0 # Loop over the blocks for b in range(Vpot.nblocks()): # Obtain block information block = Vpot.get_block(b) points_func.compute_points(block) npoints = block.npoints() lpos = np.array(block.functions_local_to_global()) if len(lpos) == 0: i += npoints continue # Obtain the grid weight w = np.array(block.w()) # Compute phi! phi = np.array(points_func.basis_values()["PHI"])[:npoints, :lpos.shape[0]] Vtmp = np.einsum('pb,p,p,pa->ab', phi, value[i:i+npoints], w, phi, optimize=True) # Add the temporary back to the larger array by indexing, ensure it is symmetric VFock[(lpos[:, None], lpos)] += 0.5 * (Vtmp + Vtmp.T) i += npoints assert i == value.shape[0], "Did not run through all the points. %i %i" %(i, value.shape[0]) return VFock #Miscellaneous def get_basis_set_correction(self, grid): return basis_set_correction(self, grid)