import hfscf
import numpy as np

SCF

def SCF (r = 1.4632, Z=[1,1], b1 = hfscf.GTO["H"], b2 = hfscf.GTO["H"], b = 3, vbs=False):
    R = [0, r]
    # Задали геометрию, заряды ядер и описание атомов
    # vbs = verbose, опция чтобы следить за выполнением кода
    # с помощью комментариев на испанском
    if vbs: print("*) Generando matriz de traslape S.")
    # Считаем матрицу перекрывания 
    s_scf = hfscf.S(R, b1, b2, b)
    # Считаем одноэлектронный гамильтониан или H_core
    if vbs: print("\n*) Generando hamiltoniano H.")
    h_scf = hfscf.H(R, Z, b1, b2, b)
        
    # Diagonalizar matriz S y hallar matriz X
    if vbs: print("\n*) Diagonalizando matriz S y hallando matriz diagonal X.")
    # Считаем матрицу трансформации Х и сопряженно-транспонированную к ней Xa
    X = hfscf.diagon(m=s_scf)
    Xa = X.getH()
    
    # Estimar matriz de densidad P
    # Делаем inital guess матрицы плотности
    if vbs: print("\n*) Creando matriz de densidad P.")
    p_scf = np.matrix([[0,0],[0,0]], dtype=np.float64)  # Referencia (7) p. 148
    
    # Comenzar proceso iterativo
    # Начинаем наши SCF итерации
    if vbs: print("\n*) Comenzando con el SCF.")
    for iteracion in range(50):
        # Construir matriz de Fock F
        # F = H + G
        # Считаем фокиан. Для этого к H_core прибавляем
        # двуэлектронный вклад в энергию G
        if vbs: print("\n**) Generando la matriz de Fock: calculando \
integrales de dos electrones.")
        g_scf = hfscf.G(r, p_scf, b1, b2, b)
        f_scf = h_scf + g_scf
        
        # Construir matriz F'
        # F' = X_adj * F * X
        # Считаем F'
        if vbs: print("**) Cambiando la base de F.")
        f_tra = Xa * f_scf * X
        
        # Diagonalizar matriz F' y constuir matriz C'
        # Диагонализируя F', получаем C'
        if vbs: print("**) Diagonalizando F' y generando C'.")
        c_tra = hfscf.diagon2(m=f_tra)
        
        # Construir matriz C
        # C = X * C'
        # Считаем матрицу коэффициентов перед базисными функциями
        if vbs: print("**) Construyendo matriz de coeficientes C.")
        c_scf = X * c_tra
        
        # Construir matriz P a partir de matriz C
        # Обновляем матрицу плотности, используя матрицу С
        if vbs: print("**) Recalculando matriz de densidad P.")
        p_temp = hfscf.P(C=c_scf)
        
        print("\nConcluida la " + str(iteracion + 1) + ". iteracion.\n")
        
        # Revisar convergencia
        #Проверяем на сходимость
        if np.linalg.norm(p_temp - p_scf) < 1E-4: # Referencia (7) p. 148
            print("\n\n-->El campo autoconsistente SI ha convergido!")
            return {"S":s_scf,"H":h_scf,"X": X,"F":f_scf,"C":c_scf,"P":p_temp}
        else:
            p_scf = p_temp
    print("\n\n-->El campo autoconsistente NO ha convergido!\nRevisar supuestos.")
    return {"S":s_scf,"H":h_scf,"X": X,"F":f_scf,"C":c_scf,"P":p_temp}

scf_mtxs = SCF(vbs=True)

*) Generando matriz de traslape S.

*) Generando hamiltoniano H.

*) Diagonalizando matriz S y hallando matriz diagonal X.

*) Creando matriz de densidad P.

*) Comenzando con el SCF.

**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 1. iteracion.


**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 2. iteracion.


**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 3. iteracion.


**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 4. iteracion.


**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 5. iteracion.


**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 6. iteracion.


**) Generando la matriz de Fock: calculando integrales de dos electrones.
**) Cambiando la base de F.
**) Diagonalizando F' y generando C'.
**) Construyendo matriz de coeficientes C.
**) Recalculando matriz de densidad P.

Concluida la 7. iteracion.



-->El campo autoconsistente SI ha convergido!

S - матрица перекрывания, содержит интегралы перекрывания базисных функций.
H - одноэлектронный гамильтониан системы. Содержит одноэлектронные интегралы.
X - матрица перехода между разными базисами, в которых записаны F и F'
F - фокиан системы. Равен сумме матриц H и G
C - матрица МО-ЛКАО коэффициентов.

Графики

hfscf.orbital(scf_mtxs['C'])

hfscf.orbital2D(scf_mtxs['C'], scf_mtxs['X'], scf_mtxs['F'])

Две орбитали: связывающая (1) и разрыхляющая (2).

Расчеты энергий

orbital_en = hfscf.ener_orbs(scf_mtxs['X'], scf_mtxs['F'])
elec_en = hfscf.ener_elec(scf_mtxs['P'], scf_mtxs['H'], scf_mtxs['F'])
total_en = hfscf.ener_tot(elec=elec_en)

print(f"Энергии орбиталей:\t{orbital_en[0]} и {orbital_en[1]}")
print(f"Электронная энергия:\t{elec_en}")
print(f"Полная энергия:\t\t{total_en}")

Энергии орбиталей:	-0.5646153594583119 и 0.6368673980935621
Электронная энергия:	-1.7974485548087507
Полная энергия:		-1.1140149845517797