【わずか200行】pythonで量子化学のプログラミングをしよう

import math
import numpy as np
import numpy.linalg as LA
np.set_printoptions(precision=6,floatmode='fixed')

# 原子数、基底関数数、電子数を入力
n_atm = 2           # HeH+の原子数は2つ
n_sto = 3           # STO-3Gは3つのガウス軌道の和で1つの原子軌道を表す
n_basis = 2         # Heの1s軌道、Hの1s軌道で基底関数数は2つとする
n_elec = 2          # HeH+の電子数は2
n_dim = 3           # x,y,zの3次元
dis_HeH = 1.4       # He-H+の結合長(bohr)

# 配列宣言 
s_matrix = np.zeros((n_basis,n_basis))
t_matrix = np.zeros((n_basis,n_basis))
v_matrix = np.zeros((n_basis,n_basis)) 
e_matrix = np.zeros((n_basis,n_basis,n_basis,n_basis))
m_matrix = np.zeros((n_basis,n_basis)) 
pc=np.zeros((n_dim))
pp=np.zeros((n_dim))
pq=np.zeros((n_dim))
PI = 3.1415926

# 原子核の電荷を入力
CHG = [2,1]

# 原子座標を入力
C = [[0.0000 ,0.0000,0.0000],
     [dis_HeH,0.0000,0.0000]]

# STO-3Gの係数を入力
EXX = [[ 6.362421394 ,1.1589229990,0.3136497915 ],
       [ 3.425250914 ,0.6239137298,0.1688554040 ]]

CXX = [[ 0.4406303066,0.4261584492,0.1328149745],
       [ 0.2769343551,0.2678388516,0.0834736711]]

# STO-3G基底関数の中心座標を入力
CO = [[0.0000 ,0.0000,0.0000],
      [dis_HeH,0.0000,0.0000]]

# 重なり行列を計算 s_t[n][m] = (n|m) 
for n in range(n_basis):
    for m in range(n_basis):
      d2 = (CO[n][0]-CO[m][0])**2+(CO[n][1]-CO[m][1])**2+(CO[n][2]-CO[m][2])**2
      for i in range(n_sto):
          for j in range(n_sto):
              zeta = EXX[n][i]+EXX[m][j]
              gzai = EXX[n][i]*EXX[m][j]/zeta
              sss  = math.exp(-gzai*d2)*((PI/zeta)**1.5)
              cxij = CXX[n][i]*CXX[m][j]
              s_matrix[n,m] +=  cxij * sss
print("s_matrix")                    
print(s_matrix)                    

# 運動エネルギー積分を計算 t_t[n][m] = (n|-1/2 nabla |m) 
for n in range(n_basis):
    for m in range(n_basis):
      d2 = (CO[n][0]-CO[m][0])**2+(CO[n][1]-CO[m][1])**2+(CO[n][2]-CO[m][2])**2
      for i in range(n_sto):
          for j in range(n_sto):
              zeta = EXX[n][i]+EXX[m][j]
              gzai = EXX[n][i]*EXX[m][j]/zeta
              sss  = math.exp(-gzai*d2)*((PI/zeta)**1.5)
              tss  = gzai*(3-2*gzai*d2)*sss
              cxij = CXX[n][i]*CXX[m][j]
              t_matrix[n,m] +=  cxij * tss
print("t_matrix")                    
print(t_matrix)                    

# 核引力積分を計算 v_t[n][m] = (n|1/r|m) 
for n in range(n_basis):
    for m in range(n_basis):
        for l in range(n_atm):
            for ndim in range(n_dim):
                pc[ndim] =C[l][ndim]
            for i in range(n_sto):
                for j in range(n_sto):
                    zeta1 = EXX[n][i]+EXX[m][j]
                    gzai1 = EXX[n][i]*EXX[m][j]/zeta1
                    cxij  = CXX[n][i]*CXX[m][j]
                    for ndim in range(n_dim):
                        pp[ndim] =(EXX[n][i]*CO[n][ndim]+EXX[m][j]*CO[m][ndim])/zeta1
                    d2_ab = (CO[n][0]-CO[m][0])**2+(CO[n][1]-CO[m][1])**2+(CO[n][2]-CO[m][2])**2
                    d2_pc = (pc[0]-pp[0])**2+(pc[1]-pp[1])**2+(pc[2]-pp[2])**2
                    if(d2_pc<0.0001):
                        f0 = 1
                    else:
                        f0 = 0.5*math.sqrt(PI/(zeta1*d2_pc))*math.erf(math.sqrt(zeta1*d2_pc))
                    vss = 2*PI/zeta1*math.exp(-gzai1*d2_ab)*f0
                    v_matrix[n,m] +=  cxij * vss * (-CHG[l])
print("v_matrix")
print(v_matrix)                    

# 2電子積分を計算 e_t[l][m][n][o] = (lm|no) 
for l in range(n_basis):
    for m in range(n_basis):
        for n in range(n_basis):
            for o in range(n_basis):
                d2_ab = (CO[l][0]-CO[m][0])**2+(CO[l][1]-CO[m][1])**2+(CO[l][2]-CO[m][2])**2
                d2_cd = (CO[n][0]-CO[o][0])**2+(CO[n][1]-CO[o][1])**2+(CO[n][2]-CO[o][2])**2
                for h in range(n_sto):
                    for i in range(n_sto):
                        for j in range(n_sto):
                            for k in range(n_sto):
                                zeta1 = EXX[l][h]+EXX[m][i]
                                gzai1 = EXX[l][h]*EXX[m][i]/zeta1
                                zeta2 = EXX[n][j]+EXX[o][k]
                                gzai2 = EXX[n][j]*EXX[o][k]/zeta2
                                rho   = zeta1*zeta2/(zeta1+zeta2)
                                cxijkl= CXX[l][h]*CXX[m][i]*CXX[n][j]*CXX[o][k]
                             
                                for ndim in range(n_dim):
                                    pp[ndim] =(EXX[l][h]*CO[l][ndim]+EXX[m][i]*CO[m][ndim])/zeta1
                                    pq[ndim] =(EXX[n][j]*CO[n][ndim]+EXX[o][k]*CO[o][ndim])/zeta2 

                                d2_pq = (pp[0]-pq[0])**2+(pp[1]-pq[1])**2+(pp[2]-pq[2])**2
                                z_ab = math.sqrt(2)*(PI**1.25)/zeta1*math.exp(-gzai1*d2_ab)
                                z_cd = math.sqrt(2)*(PI**1.25)/zeta2*math.exp(-gzai2*d2_cd)

                                if(d2_pq<0.0001):
                                    f0 = 1
                                else:
                                    f0 = 0.5*math.sqrt(PI/(rho*d2_pq))*math.erf(math.sqrt(rho*d2_pq))

                                e_ssss = 1.0/math.sqrt(zeta1+zeta2)*z_ab*z_cd*f0
                                e_matrix[l,m,n,o] +=  cxijkl * e_ssss
print()                                
print("START SCF CALCULATION")
# S行列を対角化  U' * S * U = M 
w_matrix, u_matrix = LA.eig(s_matrix)
u_d_matrix = u_matrix.T

# 変換行列XとX'をつくる、  X = U * M^(-0.5)  
for i in range(n_basis):
    m_matrix[i][i] = 1/math.sqrt(w_matrix[i])

x_matrix = u_matrix @ m_matrix

# ハミルトニアン行列hをつくり、左右からX'とXをかけてa行列を作る a = X' * h * X  
h_matrix = t_matrix + v_matrix
a_matrix = x_matrix.T @ h_matrix @ x_matrix

# a行列を対角化する固有ベクトルwをつくる w' * a * w = eps 固有値の小さい順番で並び変える
eps_matrix,w_matrix = LA.eig(a_matrix)
index = np.argsort(eps_matrix)[::1]
w_matrix = w_matrix[:,index]

# 変換行列xと固有ベクトルwから初期分子軌道係数coffが得られる coff = x * w
coff_matrix = x_matrix @ w_matrix

# SCF計算開始、電子エネルギーeneが収束するまで行う
ene_save = 0
for n_scf in range(10):
    p_matrix = np.zeros((n_basis,n_basis))
    for n in range(n_basis):
        for m in range(n_basis):
            for i in range( int(n_elec/2) ):
                p_matrix[n][m] +=  2* coff_matrix[n][i] * coff_matrix[m][i] 
    
    # フォック行列hをつくる
    f_matrix =  t_matrix + v_matrix
    for n in range(n_basis):
        for m in range(n_basis):
            for i in range(n_basis):
                for j in range(n_basis):
                    f_matrix[n][m] +=  p_matrix[i][j]*(e_matrix[n][m][i][j]-0.5*e_matrix[n][i][m][j])

    # 電子エネルギーeneを計算し、 ene = 0.5 * p * ( h + f )　出力する
    ene = 0
    for n in range(n_basis):
        for m in range(n_basis):
            ene +=  0.5*p_matrix[n][m]*(h_matrix[m][n]+f_matrix[m][n])
    print("scf_roop",n_scf+1,"elec_ene",round(ene,7))

    # 収束判定を行う
    if(abs(ene - ene_save))<10 ** (-7):
        print("SCF CONVERSION")
        print()
        break

    # 次回roopの時に収束判定を行えるようにするため、エネルギ値をene_saveとして保存する
    ene_save = ene        

    # フォック行列hに、左右からX'とXをかけてa行列を作る a = X' * h * X  
    a_matrix = x_matrix.T @ f_matrix @ x_matrix
    
    # a行列を対角化する固有ベクトルwをつくる w' * a * w = eps 固有値の小さい順番で並び変える
    eps_matrix,w_matrix = LA.eig(a_matrix)
    index = np.argsort(eps_matrix)[::1]
    w_matrix = w_matrix[:,index]

    # 変換行列xと固有ベクトルwから分子軌道係数coffが得られる coff = x * w
    coff_matrix = x_matrix @ w_matrix

# 得られた分子軌道を出力する
print("coff_matrix")
print(coff_matrix)