import

n2v.patched/engines/psi4.py

@@ -0,0 +1,193 @@
+"""
+Provides interface n2v interface to Psi4
+"""
+
+
+from .engine import Engine
+import numpy as np
+from opt_einsum import contract
+
+try:
+    import psi4
+    psi4.set_options({"save_jk" : True})
+    has_psi4 = True
+except ImportError:
+    has_psi4 = False
+    
+if has_psi4:
+    from ..grid import Psi4Grider
+    class Psi4Engine(Engine):
+        """
+        Psi4 Engine Class
+        """
+
+        def set_system(self, molecule, basis, ref='1', pbs='same', wfn=None):
+            """
+            Initializes geometry and basis infromation
+
+            Parameters
+            ----------
+            molecule: psi4.core.Molecule
+                Molecule of the system  used
+            basis: str
+                Basis set of calculation
+            ref: int
+                Reference: Restricted   -> 1
+                           Unrestricted -> 2
+            pbs: str
+                Basis set of potential used
+            wfn : psi4.core.{RHF, UHF, RKS, UKS, Wavefunction, CCWavefuncion...}
+                Psi4 wavefunction object
+            """
+            self.mol = molecule
+            
+            #Assert units are in bohr
+            # units = self.mol.to_schema(dtype='psi4')['units']
+            # if units != "Bohr":
+            #     raise ValueError("Units need to be set in Bohr")
+            self.basis_str = basis
+            self.ref = ref 
+            self.pbs = pbs
+            self.pbs_str   = basis if pbs == 'same' else pbs
+
+            self.nalpha = wfn.nalpha()
+            self.nbeta = wfn.nbeta()
+
+            self.wfn = wfn
+
+        def initialize(self):
+            """
+            Initializes basic objects required for the Psi4Engine
+            """
+            self.basis = psi4.core.BasisSet.build( self.mol, key='BASIS', target=self.basis_str)
+            self.pbs   = psi4.core.BasisSet.build( self.mol, key='BASIS', target=self.pbs_str)
+
+            self.nbf   = self.basis.nbf()
+            self.npbs  = self.pbs.nbf()
+
+            self.mints = psi4.core.MintsHelper( self.basis )
+            self.jk    = self.generate_jk()
+
+            self.grid = Psi4Grider(self.mol, self.basis, self.ref)
+
+        def get_T(self):
+            """Kinetic Potential in ao basis"""
+            return np.array( self.mints.ao_kinetic() )
+
+        def get_Tpbas(self):
+            """Kinetic Potential in pbs"""
+            return np.array( self.mints.ao_kinetic(self.pbs, self.pbs) )
+
+        def get_V(self):
+            """External potential in ao basis"""
+            return np.array( self.mints.ao_potential() )
+
+        def get_A(self):
+            """Inverse squared root of S matrix"""
+            A = self.mints.ao_overlap()
+            A.power( -0.5, 1e-16 )
+            return np.array( A )
+
+        def get_S(self):
+            """Overlap matrix in AO basis"""
+            return np.array( self.mints.ao_overlap() )
+
+        def get_S3(self):
+            """3 Orbitals Overlap matrix in AO basis"""
+            return np.array( self.mints.ao_3coverlap(self.basis,self.basis,self.pbs) )
+
+        def get_S4(self):
+            """
+            Calculates four overlap integral with Density Fitting method.
+            S4_{ijkl} = \int dr \phi_i(r)*\phi_j(r)*\phi_k(r)*\phi_l(r)
+            Parameters
+            ----------
+            wfn: psi4.core.Wavefunction
+                Wavefunction object of moleculep
+            Return
+            ------
+            S4
+            """
+
+            print(f"4-AO-Overlap tensor will take about {self.nbf **4 / 8 * 1e-9:f} GB.")
+
+            aux_basis = psi4.core.BasisSet.build(self.mol, "DF_BASIS_SCF", "", "JKFIT", self.basis_str)
+            S_Pmn = np.squeeze(self.mints.ao_3coverlap(aux_basis, self.basis, self.basis ))
+            S_PQ = np.array(self.mints.ao_overlap(aux_basis, aux_basis))
+            S_PQinv = np.linalg.pinv(S_PQ, rcond=1e-9)
+            S4 = contract('Pmn,PQ,Qrs->mnrs', S_Pmn, S_PQinv, S_Pmn)
+            return S4
+
+        def generate_jk(self, gen_K=False):
+            """
+            Creates jk object for generation of Coulomb and Exchange matrices
+            1.0e9 B -> 1.0 GB
+            """
+            jk = psi4.core.JK.build(self.basis)
+            memory = int(jk.memory_estimate() * 1.1)
+            jk.set_memory(int(memory)) 
+            # added by Ehsan
+            # .set_do_K(gen_K = False) determines if exchane matrices should be calculated or not
+            jk.set_do_K(gen_K)
+            jk.initialize()
+            #print("jk: ", jk)
+            return jk
+
+        def compute_hartree(self, Cocc_a, Cocc_b):
+            """
+            Generates Coulomb and Exchange matrices from occupied orbitals
+            """
+            Cocc_a = psi4.core.Matrix.from_array(Cocc_a)
+            Cocc_b = psi4.core.Matrix.from_array(Cocc_b)
+            self.jk.C_left_add(Cocc_a)
+            self.jk.C_left_add(Cocc_b) 
+            self.jk.compute()
+            self.jk.C_clear()
+            J = (np.array(self.jk.J()[0]), np.array(self.jk.J()[1]))
+            return J
+
+        def hartree_NO(self, Dta):
+            """
+            Computes Hartree potential in AO basis from Natural Orbitals
+            """
+
+            if self.wfn is None:
+                raise ValueError('Please provide a wfn object to the Inverter, i.e., Inverter.eng = wfn')
+
+            if type(self.wfn) == psi4.core.CCWavefunction:
+                C_NO    = psi4.core.Matrix(self.nbf, self.nbf)
+                eigs_NO = psi4.core.Vector(self.nbf)
+                self.wfn.Da().diagonalize( C_NO, eigs_NO, psi4.core.DiagonalizeOrder.Descending )
+                occ = np.sqrt( np.array(eigs_NO) )
+                new_CA = occ * np.array(C_NO)
+                assert np.allclose( new_CA @ new_CA.T, Dta )
+                if self.ref == 1:
+                    new_CB = np.copy( new_CA )
+                else:
+                    self.wfn.Db().diagonalize( C_NO, eigs_NO, psi4.core.DiagonalizeOrder.Descending )
+                    occ_b = np.sqrt( np.array( eigs_NO ) )
+                    new_CB = occ_b * np.array( C_NO )
+                J0 = self.compute_hartree(new_CA, new_CB)
+                return J0
+
+        def run_single_point(self, mol, basis, method):
+            """
+            Run a standard energy calculation
+            """
+
+            wfn_temp = psi4.energy(init+"/" + self.basis_str, 
+                                molecule=self.mol, 
+                                return_wfn=True)[1]
+
+            if self.ref == 1:
+                D = np.array(wfn_temp.Da()) + np.array(wfn_temp.Db())
+                C = np.array(wfn_temp.Ca())
+                e = np.array(wfn_temp.epsilon_a())
+                
+            else:
+                D = np.stack( (np.array(wfn_temp.Da()), np.array(wfn_temp.Db())) )
+                C = np.stack( (np.array(wfn_temp.Ca()), np.array(wfn_temp.Cb())) )
+                e = np.stack( (np.array(wfn_temp.epsilon_a()), np.array(wfn_temp.epsilon_b())) )
+
+            return D, C, e
+            

n2v.patched/grid/grider.py

@@ -0,0 +1,775 @@
+"""
+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)    

n2v.patched/inverter.py

@@ -0,0 +1,542 @@
+"""
+Inverter.py
+"""
+from warnings import warn
+from dataclasses import dataclass
+import numpy as np
+from opt_einsum import contract
+
+from .methods.zmp import ZMP
+from .methods.wuyang import WuYang
+from .methods.pdeco import PDECO
+from .methods.oucarter import OC
+from .methods.mrks import MRKS
+from .methods.direct import Direct
+
+#Grider was imported by Ehsan
+from .grid.grider import Grider
+
+@dataclass
+class V:
+    """Stores Potentials on AO"""
+    T : np.ndarray
+
+class E:
+    """Stores Energies"""
+
+# Grider was added by Ehsan 
+class Inverter(Direct, ZMP, WuYang, PDECO, OC, MRKS, Grider):
+    """
+    Attributes:
+    ----------
+
+    mol : Engine.molecule
+        Molecule class of engine used
+    basis : Engine.basis
+        Basis class of engine used
+    basis_str : str
+        Basis set
+    nbf : int
+        Number of basis functions for main calculation
+    nalpha : int
+        Number of alpha electrons
+    nbeta : int
+        Number of beta electrons
+    ref : {1,2}
+        Reference calculation
+        1 -> Restricted
+        2 -> Unrestricted
+    Dt : List
+        List of np.ndarray for target density matrices (on AO).
+    ct : List
+        List of np.ndarray for input occupied orbitals. This might not be correct for post-HartreeFock methods.
+    pbs_str: string
+        name of Potential basis set
+    pbs : Engine.basis
+        Basis class for Potential basis set of the engine used. 
+    npbs : int
+        the length of pbs
+    v_pbs : np.ndarray shape (npbs, ) for ref==1 and (2*npbs, ) for ref==2.
+        potential vector on the Potential Baiss Set.
+        If the potential is not represented on the basis set, this should
+        remain 0. It will be initialized to a 0 array. One can set this
+        value for initial guesses before Wu-Yang method (WY) or PDE-Constrained
+        Optimization method (PDE-CO). For example, if PDE-CO is ran after
+        a WY calculation, the initial for PDE-CO will be the result of WY
+        if v_pbs is not zeroed.
+    S2  : np.ndarray
+        The ao overlap matrix (i.e. S matrix)
+    S3  : np.ndarray
+        The three ao overlap matrix (ao, ao, pbs)
+    S4  : np.ndarray
+        The four ao overlap matrix, the size should be (ao, ao, ao, ao)
+    jk  : Engine.jk
+        Engine jk object.
+    T   : np.ndarray
+        kinetic matrix on ao
+    V   : np.ndarray
+        external potential matrix on ao
+    T_pbs: np.ndarray
+        kinetic matrix on pbs. Useful for regularization.
+    guide_potential_components: list of string
+        guide potential components name
+    va, vb: np.ndarray of shape (nbasis, nbasis)
+        guide potential Fock matrix.
+    """
+
+    def __init__( self, engine='psi4' ):
+        engine = 'psi4'
+        self.eng_str = engine.lower()
+        if engine.lower() == 'psi4':
+            from .engines import Psi4Engine
+            self.eng = Psi4Engine()
+        elif engine.lower() == 'pyscf':
+            from .engines import PySCFEngine
+            self.eng = PySCFEngine()
+        else:
+            raise ValueError("Engine name is incorrect. The availiable engines are: {psi4, pyscf}")
+            
+    def __repr__( self ):
+        return "n2v.Inverter"
+
+    def set_system( self, molecule, basis, ref=1, pbs='same' , **kwargs):
+        """
+        Stores relevant information and intitializes Engine
+
+        Parameters
+        ----------
+        molecule: Engine.molecule
+            Molecule object of selected engine
+        basis: str
+            Basis set of the main calculation
+        ref: int
+            reference for system. Restricted   -> 1
+                                  Unrestricted -> 2
+        pbs: str, default='same'
+            Basis set for the potential
+        **kwargs:
+            Optional Parameters for different Engiens 
+            Psi4 Engine:
+                wfn : psi4.core.{RHF, UHF, RKS, UKS, Wavefunction, CCWavefuncion...}
+                Psi4 wavefunction object
+            PySCF Engine:
+                None
+
+        """
+        # Communicate TO engine
+
+        self.eng.set_system(molecule, basis, ref, pbs, **kwargs)
+        self.ref = ref
+        
+        #added by Ehsan
+        self.basis = basis
+
+        #added by Ehsan
+        self.molecule = molecule
+
+        self.nalpha = self.eng.nalpha
+        self.nbeta = self.eng.nbeta
+
+        # Initialize ecompasses everything the engine builds with basis set 
+        self.eng.initialize()
+        self.set_basis_matrices()
+
+        # Receive FROM engine
+        self.nbf  = self.eng.nbf
+        self.npbs = self.eng.npbs
+        self.v_pbs = np.zeros( (self.npbs) ) if self.ref == 1 \
+                                             else np.zeros( 2 * self.npbs )
+
+    @classmethod
+    def from_wfn( self, wfn, pbs='same' ):
+        """
+        Generates Inverter directly from wavefunction. 
+        
+        Parameters
+        ----------
+        wfn: Psi4.Core.{RHF, RKS, ROHF, CCWavefunction, UHF, UKS, CUHF}
+            Wavefunction Object
+        Returns
+        -------
+        inv: n2v.Inverter
+            Inverter Object. 
+        """
+        from .engines import Psi4Engine
+        inv = self( engine='psi4' )
+        inv.eng = Psi4Engine()
+        ref = 1 if wfn.to_file()['boolean']['same_a_b_dens'] else 2
+        inv.set_system( wfn.molecule(), wfn.basisset().name(), pbs=pbs, ref=ref, wfn=wfn )
+        # done by Ehsan
+        #inv.Dt = [ np.array(wfn.Da()), np.array(wfn.Db()) ]
+        self.Dt = [ np.array(wfn.Da()), np.array(wfn.Db()) ]
+        
+        # done by Ehsan
+        #inv.ct = [ np.array(wfn.Ca_subset("AO", "OCC")), np.array(wfn.Cb_subset("AO", "OCC")) ]
+        # ct contains matrices of occupied orbitals alpah and betta (n x m)
+        self.ct = [ np.array(wfn.Ca_subset("AO", "OCC")), np.array(wfn.Cb_subset("AO", "OCC")) ]
+        
+        inv.et = [ np.array(wfn.epsilon_a_subset("AO", "OCC")), np.array(wfn.epsilon_b_subset("AO", "OCC")) ]
+        inv.eng_str = 'psi4'
+        inv.eng.wfn = wfn
+
+        return inv
+
+    def set_basis_matrices( self ):
+        """
+        Generate basis dependant matrices
+        """
+        self.T  = self.eng.get_T()
+        self.V  = self.eng.get_V()
+        self.A  = self.eng.get_A()
+        self.S2 = self.eng.get_S()
+        self.S3      = self.eng.get_S3()
+
+        if self.eng.pbs_str != 'same':  
+            self.T_pbs  = self.eng.get_Tpbas()
+
+        self.S4 = None
+
+    def compute_hartree( self, Cocc_a, Cocc_b ):
+        """
+        Computes Hartree Potential on AO basis set. 
+
+        Parameters
+        ----------
+        Cocc_a, Cocc_b: np.ndarray (nbf, nbf)
+            Occupied orbitals in ao basis
+
+        Returns
+        -------
+        J: List of np.ndarray
+            Hartree potential due to density from Cocc_a and Cocc_b
+        """
+        return self.eng.compute_hartree(Cocc_a, Cocc_b )
+
+    def diagonalize( self, matrix, ndocc ):
+        """
+        Diagonalizes Fock Matrix
+
+        Parameters
+        ----------
+        marrix: np.ndarray
+            Matrix to be diagonalized
+        ndocc: int
+            Number of occupied orbitals
+
+        Returns
+        -------
+        C: np.ndarray
+            Orbital Matrix
+        Cocc: np.ndarray
+            Occupied Orbital Matrix
+        D: np.ndarray
+            Density Matrix
+        eigves: np.ndarray
+            Eigenvalues
+        """
+
+        # np.linalg.eigh() gives eigenvalues and eigenvectors for a symmetric matrix of choice
+        Fp = self.A.dot(matrix).dot(self.A)
+        # eigvecs must be eigenvalues or energies here!
+        eigvecs, Cp = np.linalg.eigh(Fp)
+        C = self.A.dot(Cp)
+        Cocc = C[:, :ndocc]
+        # contract converts pi and qi to pq . here two matrices with n x m dimension
+        #are converted to one matrix with n x n shape,
+        #In fact it gives the product of Cocc matrix and its transpose matrix
+        D = contract('pi,qi->pq', Cocc, Cocc)
+        
+
+        return C, Cocc, D, eigvecs
+
+    def diagonalize_with_potential_vFock(self, v=None):
+        """
+        Diagonalize Fock matrix with additional external potential
+        Stores values in object. 
+
+        Parameters
+        ----------
+        v: np.ndarray
+            Additional external potential to be added to hamiltonian along with:
+            Kinetic_nm
+            External_nm
+            Guide_Potential_nm
+        """
+
+        if v is None:
+            fock_a = self.V + self.T + self.va
+        else:
+            if self.ref == 1:
+                fock_a = self.V + self.T + self.va + v
+            else:
+                valpha, vbeta = v
+                fock_a = self.V + self.T + self.va + valpha
+                fock_b = self.V + self.T + self.vb + vbeta
+
+
+        self.Ca, self.Coca, self.Da, self.eigvecs_a = self.diagonalize( fock_a, self.nalpha )
+
+        if self.ref == 1:
+            self.Cb, self.Cocb, self.Db, self.eigvecs_b = self.Ca.copy(), self.Coca.copy(), self.Da.copy(), self.eigvecs_a.copy()
+        else:
+            self.Cb, self.Cocb, self.Db, self.eigvecs_b = self.diagonalize( fock_b, self.nbeta )    
+
+    # Actual Methods
+    def generate_components(self, guide_components, **keywords):
+        """
+        Generates exact potential components to be added to
+        the Hamiltonian to aide in the inversion procedure. 
+        Parameters:
+        -----------
+        guide_potential_components: list
+            Components added as to guide inversion. 
+            Can be chosen from ["hartree", "fermi_amandi", "svwn"]
+        """
+
+        self.guide_components = guide_components
+        self.va = np.zeros( (self.nbf, self.nbf) )
+        self.vb = np.zeros( (self.nbf, self.nbf) )
+        self.J0 = self.compute_hartree(self.ct[0], self.ct[1])
+        N       = self.nalpha + self.nbeta
+
+        if self.eng_str == 'psi4':
+            J0_NO = self.eng.hartree_NO(self.Dt[0])
+            self.J0 = J0_NO if J0_NO is not None else self.J0
+
+        if guide_components == 'none':
+            warn("No guide potential was provided. Convergence may not be achieved")
+        elif guide_components == 'hartree':
+            self.va += self.J0[0] + self.J0[1]
+            self.vb += self.J0[0] + self.J0[1]
+        elif guide_components == 'fermi_amaldi':
+            v_fa = (1-1/N) * (self.J0[0] + self.J0[1])
+            self.va += v_fa
+            self.vb += v_fa
+        else:
+            raise ValueError("Guide component not recognized")
+
+    def invert(self, method, 
+                     guide_components = 'hartree',
+                     opt_max_iter = 50,
+                     **keywords):
+        """
+        Handler to all available inversion methods
+        Parameters
+        ----------
+        method: str
+            Method used to invert density. 
+            Can be chosen from {wuyang, zmp, mrks, oc}. 
+            See documentation below for each method. 
+        guide_components: list, opt
+            Components added as to guide inversion. 
+            Can be chosen from {"fermi_amandi", "svwn"}
+            Default: ["fermi_amaldi"]
+        opt_max_iter: int, opt
+            Maximum number of iterations inside the chosen inversion.
+            Default: 50
+        direct
+        ------
+        Direct inversion of a set of Kohn-Sham equations. 
+        $$v_{xc}(r) = \frac{1}{n(r)} \sum_i^N [\phi_i^{*} (r) \nabla^2 \phi_i(r) + \varepsilon_i | \phi_i(r)|^2] $$
+            Parameters:
+            -----------
+                grid: np.ndarray, opt
+                    Grid where result will be expressed in.
+                    If not provided, dft grid will be used instead. 
+                
+        wuyang
+        ------
+        the Wu-Yang method:
+        The Journal of chemical physics 118.6 (2003): 2498-2509.
+            Parameters:
+            ----------
+                opt_max_iter: int
+                    maximum iteration
+                opt_method: string, opt
+                    Method for scipy optimizer
+                    Currently only used by wuyang and pdeco method.
+                    Defaul: 'trust-krylov'
+                    https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html
+                reg : float, opt
+                    Regularization constant for Wuyant Inversion.
+                    Default: None -> No regularization is added.
+                    Becomes attribute of inverter -> inverter.lambda_reg
+                tol: float
+                    tol for scipy.optimize.minimize
+                gtol: float
+                    gtol for scipy.optimize.minimize: the gradient norm for
+                    convergence
+                opt: dict
+                    options for scipy.optimize.minimize
+                                Notice that opt has lower priorities than opt_max_iter and gtol.
+            return:
+                the result are stored in self.v_pbs
+            
+        zmp
+        ---
+        The Zhao-Morrison-Parr Method:
+        Phys. Rev. A 50, 2138 
+            Parameters:
+            ----------
+                lambda_list: list
+                    List of Lamda parameters used as a coefficient for Hartree 
+                    difference in SCF cycle.
+                zmp_mixing: float, optional
+                    mixing \in [0,1]. How much of the new potential is added in.
+                    For example, zmp_mixing = 0 means the traditional ZMP, i.e. all the potentials from previous
+                    smaller lambda are ignored.
+                    Zmp_mixing = 1 means that all the potentials of previous lambdas are accumulated, the larger lambda
+                    potential are meant to fix the wrong/inaccurate region of the potential of the sum of the previous
+                    potentials instead of providing an entire new potentials.
+                    default: 1
+                opt_max_iter: float
+                    Maximum number of iterations for scf cycle
+                opt_tol: float
+                    Convergence criteria set for Density Difference and DIIS error. 
+                return:
+                    The result will be stored in self.proto_density_a and self.proto_density_b
+                    For zmp_mixing==1, restricted (ref==1):
+                        self.proto_density_a = \sum_i lambda_i * (Da_i - Dt[0]) - 1/N * (Dt[0])
+                        self.proto_density_b = \sum_i lambda_i * (Db_i - Dt[1]) - 1/N * (Dt[1]);
+                    unrestricted (ref==1):
+                        self.proto_density_a = \sum_i lambda_i * (Da_i - Dt[0]) - 1/N * (Dt[0] + Dt[1])
+                        self.proto_density_b = \sum_i lambda_i * (Db_i - Dt[1]) - 1/N * (Dt[0] + Dt[1]);
+                    For restricted (ref==1):
+                        vxc = \int dr' \frac{self.proto_density_a + self.proto_density_b}{|r-r'|}
+                            = 2 * \int dr' \frac{self.proto_density_a}{|r-r'|};
+                    for unrestricted (ref==2):
+                        vxc_up = \int dr' \frac{self.proto_density_a}{|r-r'|}
+                        vxc_down = \int dr' \frac{self.proto_density_b}{|r-r'|}.
+                    To get potential on grid, one needs to do
+                        vxc = self.on_grid_esp(Da=self.proto_density_a, Db=self.proto_density_b, grid=grid) for restricted;
+                        vxc_up = self.on_grid_esp(Da=self.proto_density_a, Db=np.zeros_like(self.proto_density_a),
+                                  grid=grid) for unrestricted;
+        mRKS
+        ----
+        the modified Ryabinkin-Kohut-Staroverov method:
+        Phys. Rev. Lett. 115, 083001 
+        J. Chem. Phys. 146, 084103p
+            Parameters:
+            -----------
+                maxiter: int
+                    same as opt_max_iter
+                vxc_grid: np.ndarray of shape (3, num_grid_points), opt
+                    When this is given, the final result will be represented
+                v_tol: float, opt
+                    convergence criteria for vxc Fock matrices.
+                    default: 1e-4
+                D_tol: float, opt
+                    convergence criteria for density matrices.
+                    default: 1e-7
+                eig_tol: float, opt
+                    convergence criteria for occupied eigenvalue spectrum.
+                    default: 1e-4
+                frac_old: float, opt
+                    Linear mixing parameter for current vxc and old vxc.
+                    If 0, no old vxc is mixed in.
+                    Should be in [0,1)
+                    default: 0.5.
+                init: string or psi4.core.Wavefunction, opt
+                    Initial guess method.
+                    default: "SCAN"
+                    1) If None, input wfn info will be used as initial guess.
+                    2) If "continue" is given, then it will not initialize
+                    but use the densities and orbitals stored. Meaningly,
+                    one can run a quick WY calculation as the initial
+                    guess. This can also be used to user speficified
+                    initial guess by setting Da, Coca, eigvec_a.
+                    3) If it's not continue, it would be expecting a
+                    method name string that works for psi4. A separate psi4 calculation
+                    would be performed.
+                sing: tuple of float of length 4, opt.
+                    Singularity parameter for _vxc_hole_quadrature()
+                    default: (1e-5, 1e-4, 1e-5, 1e-4)
+                    [0]: atol, [1]: atol1 for dft_spherical grid calculation.
+                    [2]: atol, [3]: atol1 for vxc_grid calculation.
+            return:
+                The result will be stored in self.grid.vxc
+        oc
+        --
+        Ou-Carter method
+        J. Chem. Theory Comput. 2018, 14, 5680−5689
+            Parameters:
+            -----------
+                maxiter: int
+                    same as opt_max_iter
+                vxc_grid: np.ndarray of shape (3, num_grid_points)
+                    The final result will be represented on this grid
+                    default: 1e-4
+                D_tol: float, opt
+                    convergence criteria for density matrices.
+                    default: 1e-7
+                eig_tol: float, opt
+                    convergence criteria for occupied eigenvalue spectrum.
+                    default: 1e-4
+                frac_old: float, opt
+                    Linear mixing parameter for current vxc and old vxc.
+                    If 0, no old vxc is mixed in.
+                    Should be in [0,1)
+                    default: 0.5.
+                init: string, opt
+                    Initial guess method.
+                    default: "SCAN"
+                    1) If None, input wfn info will be used as initial guess.
+                    2) If "continue" is given, then it will not initialize
+                    but use the densities and orbitals stored. Meaningly,
+                    one can run a quick WY calculation as the initial
+                    guess. This can also be used to user speficified
+                    initial guess by setting Da, Coca, eigvec_a.
+                    3) If it's not continue, it would be expecting a
+                    method name string that works for psi4. A separate psi4 calculation
+                    would be performed.
+                    wuyang
+        pdeco
+        -----
+        the PDE-Constrained Optimization method:
+        Int J Quantum Chem. 2018;118:e25425;
+        Nat Commun 10, 4497 (2019).
+            Parameters:
+            ----------
+                opt_max_iter: int
+                    maximum iteration
+                opt_method: string, opt
+                    Method for scipy optimizer
+                    Currently only used by wuyang and pdeco method.
+                    Defaul: 'L-BFGS-B'
+                    Options: ['L-BFGS-B', 'BFGS']
+                    https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html
+                reg : float, opt
+                    Regularization constant for Wuyant Inversion.
+                    Default: None -> No regularization is added.
+                    Becomes attribute of inverter -> inverter.lambda_reg
+                gtol: float
+                    gtol for scipy.optimize.minimize: the gradient norm for
+                    convergence
+                opt: dict
+                    options for scipy.optimize.minimize
+                    Notice that opt has lower priorities than opt_max_iter and gtol.
+            return:
+                the result are stored in self.v_pbs
+        """
+
+        self.generate_components(guide_components)
+
+        if method.lower() == "direct":
+            return self.direct_inversion(**keywords)
+        elif method.lower() == "wuyang":
+            self.wuyang(opt_max_iter, **keywords)
+        elif method.lower() == "zmp":
+            self.zmp(opt_max_iter, **keywords)
+        elif method.lower() == "mrks":
+            if self.eng_str == 'pyscf':
+                raise ValueError("mRKS method not yet available with the PySCF engine. Try another method or another engine.")
+            return self.mRKS(opt_max_iter, **keywords)
+        elif method.lower() == 'oc':
+            if self.eng_str == 'pyscf':
+                raise ValueError("OuCarter method not yet available with the PySCF engine. Try another method or another engine.")
+            return self.oucarter(opt_max_iter, **keywords)
+        elif method.lower() == 'pdeco':
+            return self.pdeco(opt_max_iter, **keywords)
+        else:
+            raise ValueError(f"Inversion method not available. Methods available: {['wuyang', 'zmp', 'mrks', 'oc', 'pdeco']}")
+

n2v.patched/methods/zmp.py

@@ -0,0 +1,394 @@
+"""
+zmp.py
+
+Functions associated with zmp inversion
+"""
+
+
+import psi4
+psi4.core.be_quiet()
+import numpy as np
+from functools import reduce
+
+
+eps = np.finfo(float).eps
+
+#added by Ehsan
+import matplotlib.pyplot as plt
+
+class ZMP():
+    """
+    ZMP Class
+    Performs ZMP optimization according to: 
+    
+    1) 'From electron densities to Kohn-Sham kinetic energies, orbital energies,
+    exchange-correlation potentials, and exchange-correlation energies' by
+    Zhao + Morrison + Parr. 
+    https://doi.org/10.1103/PhysRevA.50.2138
+    """
+    def zmp(self, 
+            opt_max_iter=100, 
+            opt_tol= psi4.core.get_option("SCF", "D_CONVERGENCE"), 
+            lambda_list=[70],
+            zmp_mixing = 1,
+            print_scf = False, 
+            ):
+
+        """
+        Performs ZMP optimization according to: 
+        
+        1) 'From electron densities to Kohn-Sham kinetic energies, orbital energies,
+        exchange-correlation potentials, and exchange-correlation energies' by
+        Zhao + Morrison + Parr. 
+        https://doi.org/10.1103/PhysRevA.50.2138
+
+        Additional DIIS algorithms obtained from:
+        2) 'Psi4NumPy: An interactive quantum chemistry programming environment 
+        for reference implementations and rapid development.' by 
+        Daniel G.A. Smith and others. 
+        https://doi.org/10.1021/acs.jctc.8b00286
+ 
+        Functionals that drive the SCF procedure are obtained from:
+        https://doi.org/10.1002/qua.26400
+
+        Parameters:
+        -----------
+        lambda_list: list
+            List of Lamda parameters used as a coefficient for Hartree
+            difference in SCF cycle.
+        zmp_mixing: float, optional
+            mixing \in [0,1]. How much of the new potential is added in.
+            For example, zmp_mixing = 0 means the traditional ZMP, i.e. all the potentials from previous
+            smaller lambda are ignored.
+            Zmp_mixing = 1 means that all the potentials of previous lambdas are accumulated, the larger lambda
+            potential are meant to fix the wrong/inaccurate region of the potential of the sum of the previous
+            potentials instead of providing an entire new potentials.
+            default: 1
+        opt_max_iter: float
+            Maximum number of iterations for scf cycle
+        opt_tol: float
+            Convergence criteria set for Density Difference and DIIS error.
+        return:
+            The result will be stored in self.proto_density_a and self.proto_density_b
+            For zmp_mixing==1, restricted (ref==1):
+                self.proto_density_a = \sum_i lambda_i * (Da_i - Dt[0]) - 1/N * (Dt[0])
+                self.proto_density_b = \sum_i lambda_i * (Db_i - Dt[1]) - 1/N * (Dt[1]);
+            unrestricted (ref==1):
+                self.proto_density_a = \sum_i lambda_i * (Da_i - Dt[0]) - 1/N * (Dt[0] + Dt[1])
+                self.proto_density_b = \sum_i lambda_i * (Db_i - Dt[1]) - 1/N * (Dt[0] + Dt[1]);
+            For restricted (ref==1):
+                vxc = \int dr' \frac{self.proto_density_a + self.proto_density_b}{|r-r'|}
+                    = 2 * \int dr' \frac{self.proto_density_a}{|r-r'|};
+            for unrestricted (ref==2):
+                vxc_up = \int dr' \frac{self.proto_density_a}{|r-r'|}
+                vxc_down = \int dr' \frac{self.proto_density_b}{|r-r'|}.
+            To get potential on grid, one needs to do
+                vxc = self.on_grid_esp(Da=self.proto_density_a, Db=self.proto_density_b, grid=grid) for restricted;
+                vxc_up = self.on_grid_esp(Da=self.proto_density_a, Db=np.zeros_like(self.proto_density_a),
+                          grid=grid) for unrestricted;
+        """
+        self.diis_space = 100
+        self.mixing = zmp_mixing
+
+        print("\nRunning ZMP:")
+        self.zmp_scf(lambda_list, opt_max_iter, print_scf, D_conv=opt_tol)
+
+    def zmp_scf(self, 
+            lambda_list,
+            maxiter, 
+            print_scf,
+            D_conv):
+        """
+        Performs scf cycle
+        Parameters:
+            zmp_functional: options the penalty term.
+            But others are not currently working except for Hartree penalty (original ZMP).
+        ----------
+        """       
+        # Target density on grid
+        if self.ref == 1:
+            # density() is a method in the class Psi4Grider() in module  psi4grider.py (added by Ehsan)
+            D0 = self.eng.grid.density(Da=self.Dt[0])
+            
+            
+        else:
+            D0 = self.eng.grid.density(Da=self.Dt[0], Db=self.Dt[1])
+        
+        # Initialize Stuff
+        vc_previous_a = np.zeros((self.nbf, self.nbf))
+         
+        vc_previous_b = np.zeros((self.nbf, self.nbf))
+        self.Da = self.Dt[0]
+        self.Db = self.Dt[1]
+        Da = self.Dt[0]
+        Db = self.Dt[1]
+        Cocca = self.ct[0]
+        #print("Cocca: ", Cocca, type(Cocca))
+        Coccb = self.ct[1]
+
+        grid_diff_old = 1/np.finfo(float).eps
+        
+        self.proto_density_a =  np.zeros_like(Da)
+        
+        self.proto_density_b =  np.zeros_like(Db)
+
+#------------->  Iterating over lambdas:
+        Ddif = []
+        L = []
+        for lam_i in lambda_list:
+            self.shift = 0.1 * lam_i
+            D_old = self.Dt[0]
+
+            # Trial & Residual Vector Lists
+            state_vectors_a, state_vectors_b = [], []
+            error_vectors_a, error_vectors_b = [], []
+
+            for SCF_ITER in range(1,maxiter):
+
+#------------->  Generate Fock Matrix:
+                vc = self.generate_s_functional(lam_i,
+                                                Cocca, Coccb, 
+                                                Da, Db)
+
+                #Equation 10 of Reference (1). Level shift. 
+                Fa = self.T + self.V + self.va + vc[0] + vc_previous_a
+                Fa += (self.S2 - reduce(np.dot, (self.S2, Da, self.S2))) * self.shift
+                #added by Ehsan: a function (np.dot : dot product) applies on an iterable (self.S2, Da, self.S2) and gives one output (a new matrix)
+                
+                if self.ref == 2:
+                    Fb = self.T + self.V + self.vb + vc[1] + vc_previous_b
+                    Fb += (self.S2 - reduce(np.dot, (self.S2, Db, self.S2))) * self.shift
+                
+
+    #------------->  DIIS:
+                if SCF_ITER > 1:
+                    #Construct the AO gradient
+                    # r = (A(FDS - SDF)A)_mu_nu
+                    # added by Ehsan (self.A: Inverse squared root of S matrix), grad_a is a matrix showing the gradients
+                    grad_a = self.A.dot(Fa.dot(Da).dot(self.S2) - self.S2.dot(Da).dot(Fa)).dot(self.A)
+                    grad_a[np.abs(grad_a) < eps] = 0.0
+
+                    if SCF_ITER < self.diis_space:
+                        state_vectors_a.append(Fa.copy())
+                        error_vectors_a.append(grad_a.copy())
+                        
+                    else:
+                        state_vectors_a.append(Fa.copy())
+                        error_vectors_a.append(grad_a.copy())
+                        del state_vectors_a[0]
+                        del error_vectors_a[0]
+
+                    #Build inner product of error vectors
+                    #dimension of the DIIS subspace = Bdim (by Ehsan)
+                    Bdim = len(state_vectors_a) 
+                    # one column and one row are added to the matrix to be fiiled with -1 (this is part of the DIIS prodecure)
+                    #np.empty: the elements of the array will initially contain whatever data was already in the memory allocated for the array
+                    Ba = np.empty((Bdim + 1, Bdim + 1))  
+                    
+                    # sets the last row and the last column of the matrix Ba to -1
+                    Ba[-1, :] = -1 
+                    Ba[:, -1] = -1
+                    Ba[-1, -1] = 0
+                    
+                    Bb = Ba.copy()
+                    
+
+                    for i in range(len(state_vectors_a)):
+                        for j in range(len(state_vectors_a)):
+                            # Ba[i,j] will be a number made of the sum of inner products of corresponding elements in matrices 
+                            Ba[i,j] = np.einsum('ij,ij->', error_vectors_a[i], error_vectors_a[j])
+                            
+                    
+                    #Build almost zeros matrix to generate inverse. 
+                    RHS = np.zeros(Bdim + 1)
+                    RHS[-1] = -1
+                    
+
+                    #Find coefficient matrix:
+                    
+                    # x = np.linalg.solve(A, b) Solve the linear system A*x = b
+
+                    Ca = np.linalg.solve(Ba, RHS.copy())
+                    
+                    Ca[np.abs(Ca) < eps] = 0.0 
+
+                    #Generate new fock Matrix:
+                    Fa = np.zeros_like(Fa)
+                    # .shape[0] gives the number of rows in a 2D array
+                    for i in range(Ca.shape[0] - 1):
+                        Fa += Ca[i] * state_vectors_a[i]
+                    
+                    #diis_error_a is the maximum error element in the last error vectors matrix
+                    diis_error_a = np.max(np.abs(error_vectors_a[-1]))
+                    if self.ref ==  1:
+                        Fb = Fa.copy()
+                        diis_error = 2 * diis_error_a
+
+                    else:
+                        grad_b = self.A.dot(Fb.dot(Db).dot(self.S2) - self.S2.dot(Db).dot(Fb)).dot(self.A)
+                        grad_b[np.abs(grad_b) < eps] = 0.0
+                        
+                        if SCF_ITER < self.diis_space:
+                            state_vectors_b.append(Fb.copy())
+                            error_vectors_b.append(grad_b.copy())
+                        else:
+                            state_vectors_b.append(Fb.copy())
+                            error_vectors_b.append(grad_b.copy())
+                            del state_vectors_b[0]
+                            del error_vectors_b[0]
+
+                        for i in range(len(state_vectors_b)):
+                            for j in range(len(state_vectors_b)):
+                                Bb[i,j] = np.einsum('ij,ij->', error_vectors_b[i], error_vectors_b[j])
+
+                        diis_error_b = np.max(np.abs(error_vectors_b[-1]))
+                        diis_error = diis_error_a + diis_error_b
+
+                        Cb = np.linalg.solve(Bb, RHS.copy())
+                        Cb[np.abs(Cb) < eps] = 0.0 
+
+                        Fb = np.zeros_like(Fb)
+                        for i in range(Cb.shape[0] - 1):
+                            Fb += Cb[i] * state_vectors_b[i]
+                # for the first iteration the error is set to 1.0
+                else:
+                    diis_error = 1.0
+                
+    #------------->  Diagonalization | Check convergence:
+                # diagonalize() method has been defined in inventer.py . the inputs are fock matrix and number of occupied orbitals
+                # this is to find the new density matrix. Here the eigenfunction is solved to get eigenvalues and coefficients
+                Ca, Cocca, Da, eigs_a = self.diagonalize(Fa, self.nalpha)
+                # eigs_a is a one dimensioanl matrix (size = nbf) of eigenvalues or enrgies 
+                
+                if self.ref == 2:
+                    Cb, Coccb, Db, eigs_b = self.diagonalize(Fb, self.nbeta)
+                else: 
+                    Cb, Coccb, Db, eigs_b = Ca.copy(), Cocca.copy(), Da.copy(), eigs_a.copy()
+
+                #difference of the new and old density matrices
+                ddm = D_old - Da
+                D_old = Da
+                # the maximum element in the differnce denstiy matrix is taken as the density error value
+                derror = np.max(np.abs(ddm))
+
+                if print_scf:
+                    if np.mod(SCF_ITER,5) == 0.0:
+                        print(f"Iteration: {SCF_ITER:3d} | Self Convergence Error: {derror:10.5e} | DIIS Error: {diis_error:10.5e}")
+                
+                #DIIS error may improve as fast as the D_conv. Relax the criteria an order of magnitude. 
+                if abs(derror) < D_conv and abs(diis_error) < D_conv*10: 
+                    # here SCF convergence is reached
+                    break
+                if SCF_ITER == maxiter - 1:
+                    raise ValueError("ZMP Error: Maximum Number of SCF cycles reached. Try different settings.")
+
+            if self.ref == 1:
+                # map the current density on grid with n points depending on the size of basis set
+                density_current = self.eng.grid.density(Da=Da)
+                
+            else:
+                density_current_a = self.eng.grid.density(Da=Da, Db=Db)
+            #the difference between the current and exact density is evaluated on grid
+            grid_diff = np.max(np.abs(D0 - density_current))
+            if np.abs(grid_diff_old) < np.abs(grid_diff):
+                # This is a greedy algorithm: if the density error stopped improving for this lambda, we will stop here.
+                print(f"\nZMP halted at lambda={lam_i}. Density Error Stops Updating: old: {grid_diff_old}, current: {grid_diff}.")
+                break
+
+            grid_diff_old = grid_diff
+            print(f"SCF Converged for lambda:{int(lam_i):5d}. Max density difference: {grid_diff}")
+            #added by Ehsan
+            Ddif.append(grid_diff)
+            L.append(lam_i)
+            # D0 is on grid. density_current is also on grid.
+            # Dt[Dta,Dtb] and Da or Db are just arrays or matrices
+            self.proto_density_a += lam_i * (Da - self.Dt[0]) * self.mixing
+            if self.ref == 2:
+                self.proto_density_b += lam_i * (Db - self.Dt[1]) * self.mixing
+            else:
+                self.proto_density_b = self.proto_density_a.copy()
+
+            vc_previous_a += vc[0] * self.mixing
+            if self.ref == 2:
+                #add a portion of previous vc to the new one
+                vc_previous_b += vc[1] * self.mixing
+
+            # this is the lambda that is already proven to be improving the density, i.e. the corresponding
+            # potential has updated to proto_density
+            successful_lam = lam_i
+            # The proto_density corresponds to successful_lam
+            successful_proto_density = [(Da - self.Dt[0]), (Db - self.Dt[1])]
+# -------------> END Iterating over lambdas:
+        
+        #added by Ehsan
+        plt.plot(L, Ddif)
+        plt.xlabel('Lambda')
+        plt.ylabel('Delta-Density')
+        plt.title(f"basis set: {self.basis}")
+        plt.savefig('Lam_D_' + self.basis+ '.pdf')
+        plt.close()
+
+        self.proto_density_a += successful_lam * successful_proto_density[0] * (1 - self.mixing)
+        if self.guide_components.lower() == "fermi_amaldi":
+            # for ref==1, vxc = \int dr (proto_density_a + proto_density_b)/|r-r'| - 1/N*vH
+            if self.ref == 1:
+                self.proto_density_a -= (1 / (self.nalpha + self.nbeta)) * (self.Dt[0])
+            # for ref==1, vxc = \int dr (proto_density_a)/|r-r'| - 1/N*vH
+            else:
+                self.proto_density_a -= (1 / (self.nalpha + self.nbeta)) * (self.Dt[0] + self.Dt[1])
+
+        self.Da = Da
+        self.Ca = Ca
+        self.Coca = Cocca
+        self.eigvecs_a = eigs_a
+
+        if self.ref == 2:
+            self.proto_density_b += successful_lam * successful_proto_density[1] * (1 - self.mixing)
+            if self.guide_components.lower() == "fermi_amaldi":
+                # for ref==1, vxc = \int dr (proto_density_a + proto_density_b)/|r-r'| - 1/N*vH
+                # an inner if caluse with an opposite condition!!!
+                if self.ref == 1:
+                    self.proto_density_b -= (1 / (self.nalpha + self.nbeta)) * (self.Dt[1])
+                # for ref==1, vxc = \int dr (proto_density_a)/|r-r'| - 1/N*vH
+                else:
+                    self.proto_density_b -= (1 / (self.nalpha + self.nbeta)) * (self.Dt[0] + self.Dt[1])
+            self.Db = Db
+            self.Cb = Cb
+            self.Cocb = Coccb
+            self.eigvecs_b = eigs_b
+
+        else:
+            self.proto_density_b = self.proto_density_a.copy()
+            self.Db = self.Da.copy()
+            self.Cb = self.Ca.copy()
+            self.Cocb = self.Coca.copy()
+            self.eigvecs_b = self.eigvecs_a.copy()
+
+
+    
+    def generate_s_functional(self, lam, Cocca, Coccb, Da, Db):
+        """
+        Generates S_n Functional as described in:
+        https://doi.org/10.1002/qua.26400
+        """
+
+        # J is the Coulomb Matrix (note added by Ehsan)
+        J = self.eng.compute_hartree(Cocca, Coccb)
+        # J is computed from the occupied orbitals with the compoute_hartree() method in engine/psi4.py
+        # the matrix Cocc is typically an n×m matrix, where n is the total number of basis functions and m is the number of occupied orbitals.
+        # J[0] corresponds to the Coulomb Matrix based on alpha occupied orbitals (note added by Ehsan)
+        # Here density (D0) is not directly used!
+        #D0 = Cocc x Cocc*+, Cocc*+ is the conjugate transpose of the coefficient matrix of occupied orbitals
+
+        #Equation 7 of Reference (1), which gives Vc(r) for each lambda (original ZMP paper)
+        if self.ref == 1:
+            vc_a = 2 * lam * ( J[0] - self.J0[0] ) 
+            self.vca = J[0] - self.J0[0]
+            vc = [vc_a]
+        else:
+            vc_a = lam * ( J[0] - self.J0[0] ) 
+            vc_b = lam * ( J[1] - self.J0[1] )
+            vc = [vc_a, vc_b]            
+        # in practice Vc(r) will be a matrix (nbf x nbf) obtained from the difference between two Coulomb matrices
+        return vc
+        

start.py

@@ -0,0 +1,48 @@
+import n2v
+import psi4
+
+H2O = psi4.geometry(
+"""
+0 1
+O   0.000000   0.000000   0.000000
+H   0.757459   0.586790   0.000000
+H  -0.757459   0.586790   0.000000
+noreorient
+nocom
+units bohr
+symmetry c1
+""" )
+
+
+#n2v is driven by psi4's reference option. Make sure you set it accordingly.
+psi4.set_options({"reference" : "rhf"})
+
+#Perform a calculation for a target density.
+#Remember that for post scf calculations, Psi4 does not update the density.
+#Thus make sure you obtain something like a dipole in order to do so.
+e, wfn = psi4.properties("ccsd/cc-pvtz", return_wfn=True, properties=["dipole"], molecule=H2O)
+
+#Define inverter objects for each molcule. Simply use the wnf object from psi4 as an argument.
+ine = n2v.Inverter()
+ine.set_system(H2O, "cc-pvtz",wfn=wfn)
+ine.from_wfn(wfn)
+
+
+# how to change the increase lambda
+start = 1
+stop = 1000
+step = 2
+lam_list = []
+
+lam = start
+for i in range(int(start), int(stop), int(step)):
+    lam_list.append(i)
+
+
+# do zmp 
+ine.invert("zmp", guide_components='fermi_amaldi', opt_max_iter=2000, opt_tol=1e-7, zmp_mixing=0, print_scf=False,
+        lambda_list=lam_list)
+
+print(ine)
+
+

zmp_xc.py

@@ -0,0 +1,106 @@
+import os
+import psi4
+import matplotlib.pyplot as plt
+import numpy as np
+# import numpy_html
+psi4.set_options({"save_jk" : True})
+psi4.set_memory(int(2.50e9))
+psi4.core.clean()
+
+import n2v
+
+import matplotlib as mpl
+mpl.rcParams["font.size"] = 11
+mpl.rcParams["font.family"] = "sans-serif"
+mpl.rcParams["axes.edgecolor"] = "#eae8e9" 
+
+#Define Psi4 geometries. Symmetries need to be set to C1.
+
+
+Ne = psi4.geometry(
+"""
+0 1
+Ne 0.0 0.0 0.0
+noreorient
+nocom
+units bohr
+symmetry c1
+""" )
+
+#n2v is driven by psi4's reference option. Make sure you set it accordingly.
+psi4.set_options({"reference" : "rhf"})
+
+#Perform a calculation for a target density.
+#Remember that for post scf calculations, Psi4 does not update the density.
+#Thus make sure you obtain something like a dipole in order to do so.
+e, wfn = psi4.properties("CCSD/aug-cc-pvtz", return_wfn=True, properties=["dipole"], molecule=Ne)
+
+arrDa = wfn.Da().to_array()
+arrDb = wfn.Db().to_array()
+
+diff = arrDa - arrDb
+sum_difference = sum(sum(row) for row in diff)
+ 
+#Define inverter objects for each molcule. Simply use the wnf object from psi4 as an argument.
+inv = n2v.Inverter('psi4')
+inv.set_system(Ne, 'aug-cc-pvtz', wfn=wfn)
+inv.Dt = [ np.array(wfn.Da()), np.array(wfn.Db()) ]
+inv.ct = [ np.array(wfn.Ca_subset("AO", "OCC")), np.array(wfn.Cb_subset("AO", "OCC")) ]
+inv.et = [ np.array(wfn.epsilon_a_subset("AO", "OCC")), np.array(wfn.epsilon_b_subset("AO", "OCC"))]
+
+# Additionally one can simply initialize an Inverter using the wavefunction.
+inv = n2v.Inverter.from_wfn(wfn)
+
+# Let us define a plotting grid:
+
+npoints=1001
+x = np.linspace(-5,5,npoints)[:,None]
+y = np.zeros_like(x)
+z = y
+grid = np.concatenate((x,y,z), axis=1).T
+
+mix = [0.0, 0.1, 0.5, 1.0]
+vxc_lab = ['Vxc_mix0', 'Vxc_mix0.5', 'Vxc_mix_1.0']
+vxc_dic = {}
+for m in mix:
+    inv.invert("zmp", opt_max_iter=200, opt_tol=1e-7, zmp_mixing=m,
+           lambda_list=np.linspace(10, 1000, 20), guide_components="fermi_amaldi")
+    inv.eigvecs_a[:inv.nalpha]
+    np.diag(inv.Da)[:inv.nalpha]
+    results = inv.eng.grid.esp(Da=inv.proto_density_a, Db=inv.proto_density_b, grid=grid, )
+    vxc_dic[m] = results[1]
+    
+
+
+for k,v in vxc_dic.items():
+    plt.plot(x, v, label="Vxc_mix_"+str(k))
+
+plt.legend()
+plt.xlim(-5,5)
+
+fig, ax = plt.subplots()
+ls = ["solid","--", "-.", "-."]
+i = 0
+for k,v in vxc_dic.items():
+    ax.plot(x,  v, label="vxc_mix_"+str(k), ls=ls[i])
+    i += 1
+
+ax.set_title("Neon Exchange Correlation Potenial")
+ax.legend()
+ax.set_xlim(1e-5,5)
+ax.set_xscale("log")
+
+
+
+#ax.set_title("Neon Exchange Correlation Potential")
+
+#ax.legend()
+#ax.set_xlim(-5,5)
+
+pltname = '_Vxc_'  + '.pdf'
+plt.savefig(pltname)
+plt.close()
+
+
+plt.show()
+plt.close()