Put scf code into python script

02-scf/scf.py

+#!/usr/bin/env python
+# coding: utf-8
+
+from pyscf import gto
+import numpy as np
+
+
+delta_E = 1.0e-8            # Energy convergence criterion
+delta_P = 1.0e-8            # Density matrix convergence criterion
+MAXITER = 100               # Maximum number of SCF Iterations
+
+
+def build_hc(T_e: np.array, V_ne: np.array) -> np.array:
+    return T_e+V_ne
+
+
+def build_ortho_mat(S: np.array) -> np.array:
+    s, U =  np.linalg.eig(S)
+    L    =  np.diag(np.power(s, -0.5))
+    X = np.dot(U,(np.dot(L,np.transpose(U))))
+    return X
+
+
+def build_fock_prime(F: np.array, X: np.array) -> np.array:
+    Fp=np.dot(np.transpose(X),np.dot(F,X))
+    return Fp
+
+
+def eigensort(eigenvalues: np.array, eigenvectors: np.array) -> (np.array, np.array):
+    eigenvectors = eigenvectors[:,eigenvalues.argsort()]
+    eigenvalues = np.sort(eigenvalues)
+    return eigenvalues, eigenvectors
+
+
+def solve_eigenvalue(Fp: np.array, X: np.array) -> (np.array, np.array):
+    e, Cp   =  np.linalg.eigh(Fp)
+    e, Cp   =  eigensort(e, Cp)
+    C       =  np.dot(np.transpose(X),Cp)    
+    return e, C
+
+
+def build_density(C: np.array, NAO: int, NEL: int) -> np.array:
+    P =  np.zeros((NAO, NAO))
+    
+    for i in range(len(C)):
+        for j in range(len(C)):
+            for m in range(int(NEL/2)): #Number of occupied orbitals
+                P[i][j] += C[i][m] * C[j][m]
+    return P
+
+
+def electronic_energy(P: np.array, Hc: np.array, F: np.array, NAO: int) -> float:
+    E = 0
+    for u in range(NAO):
+        for v in range(NAO):
+            E += P[u][v] * (Hc[u][v] + F[u][v])
+    return E
+
+
+def build_new_fock(Hc: np.array, P: np.array, J: np.array, NAO: int) -> np.array:
+    F = np.copy(Hc)
+    for u in range(NAO):
+        for v in range(NAO):
+            for l in range(NAO):
+                for s in range(NAO):
+                    F[u][v] += P[l][s] * ((2 * J[u][v][l][s]) - J[u][l][v][s])
+    return F
+
+
+def compute_rmsd(P0: np.array, P1: np.array) -> float:
+    rmsd = 0
+    for u in range(len(P0)):
+        for v in range(len(P0[u])):
+            rmsd += (P1[u][v] - P0[u][v])**2
+    return np.sqrt(rmsd)
+
+
+def scf(Hc: np.array, X: np.array, J: np.array, E_nuc: float, NAO: int, NEL: int) -> dict:
+    
+    Fp_0 = build_fock_prime(Hc, X)
+    
+    e_0,C_0 = solve_eigenvalue(Fp_0, X) 
+    P_0 = build_density(C_0, NAO, NEL)
+    
+    E_el_0 = electronic_energy(P_0, Hc, Hc, NAO)
+
+    E_tot_0 = E_el_0 + E_nuc
+
+    this_scf = {
+        "tot_energies": [E_tot_0],
+        "orb_eigenvalues": [e_0],
+    }
+
+    # SCF loop
+    for i in range(1,MAXITER+1):
+        
+        F = build_new_fock(Hc, P_0, J, NAO)
+        
+        Fp = build_fock_prime(F, X)
+        
+        e, C = solve_eigenvalue(Fp, X)
+        
+        P = build_density(C, NAO, NEL)
+    
+        E_el = electronic_energy(P, Hc, F, NAO)
+        
+        E_tot = E_el + E_nuc
+
+        this_scf["tot_energies"].append(E_tot)
+        this_scf["orb_eigenvalues"].append(e)
+            
+        if np.abs(E_el-E_el_0)<delta_E:
+            print("E_el converged after: "+str(i)+" Iterations.")
+            print()
+            this_scf["iterations"] = i
+            break
+        
+        E_el_0 = E_el
+                
+        if compute_rmsd(P_0,P)<delta_P:
+            print("P converged after: "+str(i)+" Iterations.")
+            print()
+            this_scf["iterations"] = i
+            break
+        
+        P_0    = np.copy(P)
+    
+    print("   Final RHF results")
+    print("   -----------------")
+    print()
+    print("Total SCF energy = "+str(E_tot))
+    print()
+    print("   Final eigenvalues")
+    print("   -----------------")
+    print()
+    for i in range(len(e)):
+        print(str(i+1)+"  "+str(e[i]))
+
+    this_scf["tot_energies"] = np.array(this_scf["tot_energies"])
+    this_scf["orb_eigenvalues"] = np.array(this_scf["orb_eigenvalues"])
+    
+    return this_scf
+
+
+