Examples

Included in the KS-pies download is an example folder including several scripts which demonstrate possible uses of KS-pies. Each script intends to showcase the use of a KS-pies feature.

1. Test Script

A complete example using ZMP, WY and the util functions. A Be atom is used for the test density. Prints to terminal results and expected results, that can be compared to confirm that the code is working correctly.

import numpy as np
from pyscf import gto, scf, dft
from kspies import wy, zmp, util
'''Blackbox test script'''

#Reference simple HF run
mol = gto.M(atom = 'Be',
            basis = 'cc-pVTZ',
            spin = 0,
            verbose = 0 )
mf = scf.RHF(mol).run()
dm_rhf = mf.make_rdm1()
mf = scf.UHF(mol).run()
dm_uhf = mf.make_rdm1()

print("#####Check WY#####")
print("!Expected -Ws = -14.57272496 !")
#Check wy giving correct answers
mw1 = wy.RWY(mol, dm_rhf)
mw1.reg = 1e-5
mw1.run()
mw1.info()

mw2 = wy.UWY(mol, dm_uhf)
mw2.reg = 1e-5
mw2.run()
mw2.info()

print()
print("#####Check ZMP#####")
print("!Expected :                 gap=  0.1262392 dN=   78.65 C= 6.75e-04 !")
#Check zmp giving correct answers
mz1 = zmp.RZMP(mol, dm_rhf)
mz1.zscf(16)

mz2 = zmp.UZMP(mol, dm_uhf)
mz2.zscf(16)

#Check util giving correct answers
grids = dft.gen_grid.Grids(mol)
grids.build()
coords = grids.coords
weights = grids.weights
ao = dft.numint.eval_ao(mol, coords, deriv=1)
rho = dft.numint.eval_rho(mol, ao, dm_rhf, xctype='GGA')

vj_ana = scf.hf.get_jk(mol, dm_rhf)[0]
Eh_ana = .5*np.einsum('ij,ji', vj_ana, dm_rhf)

print()
print("#####Check util.eval_vh#####")
print("Eh(numerical) - Eh(analytic) in Hartrees")
saved = [-0.000269882, -0.000132191, -0.000072399, 
         -0.000042809, -0.000026879, -0.000017730] #Precomputed values
print("Lvl  computed      saved")
for i, Lvl in enumerate([ 3, 4, 5, 6, 7, 8]):
    vj_num = util.eval_vh(mol, coords, dm_rhf, Lvl = Lvl)
    Eh_num = .5*np.einsum('r,r,r', vj_num, rho[0,:], weights)
    print('%i  %.7f %.7f'%(Lvl,(Eh_num-Eh_ana),saved[i]))

print()
print("#####Check util.eval_vxc#####")
print("Check with vxc*n differences on same density")

xc_code = 'pbe'
_, vxc, _, _ = dft.libxc.eval_xc(xc_code, rho)
Vxc = dft.numint.eval_mat(mol, ao, weights, rho, vxc, xctype='GGA' ) #PySCF Vxc matrix
Exc = np.einsum('ij,ji', Vxc, dm_rhf) #analytic Vxc*n 
vxcr = util.eval_vxc(mol, dm_rhf, xc_code, coords) 
Exc_num = np.einsum('r,r,r', vxcr, rho[0,:], weights) #numerical Vxc*n
print('PySCF - util  : %.7f'%(Exc-Exc_num))

xc_code2 = '0.999*pbe,pbe'
_, vxc, _, _ = dft.libxc.eval_xc(xc_code2, rho)
Vxc2 = dft.numint.eval_mat(mol, ao, weights, rho, vxc, xctype='GGA' )
Exc2 = np.einsum('ij,ji', Vxc2, dm_rhf)
print('PBE - 0.999PBE: %.7f'%(Exc-Exc2))

2. Benzene

First perform DF-RHF calculation of benzene, and then perform RWY, UWY, DF-RZMP, DF-UZMP, RZMP and DF-UZMP calculations with HF target density. The example also plot density differences (requires matplotlib) Note that RZMP and UZMP with DF approximation require a substantial amount of time.

import numpy as np
from pyscf import gto, scf, lib, dft
from kspies import zmp, wy
lib.num_threads(8)

#Define system
mol=gto.M(atom= '''
C    0.0000    1.3936     0.0000
C    1.2069    0.6968     0.0000
C    1.2069   -0.6968     0.0000
C    0.0000   -1.3936     0.0000
C   -1.2069   -0.6968     0.0000
C   -1.2069    0.6968     0.0000
H    0.0000    2.4787     0.0000
H    2.1467    1.2394     0.0000
H    2.1467   -1.2394     0.0000
H    0.0000   -2.4787     0.0000
H   -2.1467   -1.2394     0.0000
H   -2.1467    1.2394     0.0000
''',
basis = 'cc-pVTZ', verbose=3)

#Reference HF calculation
mf=scf.RHF(mol).density_fit()
mf.kernel()
dm_tar = mf.make_rdm1()
udm_tar = np.array((dm_tar*0.5,dm_tar*0.5)) #to pass to UZMP and UWY

#Grid generation for density accuracy check
grids=dft.gen_grid.Grids(mol)
grids.build()
coords=grids.coords
weights=grids.weights
ao=dft.numint.eval_ao(mol, coords)
def ndiff(P):
    if P.ndim==3: P=P[0]+P[1]
    rho=dft.numint.eval_rho(mol, ao, P)
    return np.einsum('r,r',abs(rho),weights)

#RWY calculation
mw1 = wy.RWY(mol, dm_tar)
mw1.run()
mw1.info()
dm_wy1 = mw1.dm
print("RWY density difference: ",ndiff(dm_wy1-dm_tar))

#UWY calculation
mw2 = wy.UWY(mol, udm_tar)
mw2.run()
mw2.info()
dm_wy2 = mw2.dm
print("UWY density difference: ",ndiff(dm_wy2-udm_tar))

#DF-RZMP
mz1=zmp.RZMP(mol, dm_tar)
mz1.with_df = True
for l in [ 8, 16, 32, 64, 128, 256, 512 ]:
    mz1.level_shift = l*0.1
    mz1.zscf(l)
    #np.save('P'+str(l), mz1.dm)

#DF-UZMP
mz2=zmp.UZMP(mol, udm_tar)
mz2.with_df = True
for l in [ 8, 16, 32, 64, 128, 256, 512 ]:
    mz2.level_shift = l*0.1
    mz2.zscf(l)

#RZMP (slow)
mz3=zmp.RZMP(mol, dm_tar)
for l in [ 8, 16, 32, 64, 128, 256, 512 ]:
    mz3.level_shift = l*0.1
    mz3.zscf(l)

#UZMP (slow)
mz4=zmp.RZMP(mol, udm_tar)
for l in [ 8, 16, 32, 64, 128, 256, 512 ]:
    mz4.level_shift = l*0.1
    mz4.zscf(l)

#Setting for plotting density difference
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

def make_colormap(seq):
    """Return a LinearSegmentedColormap
    seq: a sequence of floats and RGB-tuples. The floats should be increasing
    and in the interval (0,1).
    """
    seq = [(None,) * 3, 0.0] + list(seq) + [1.0, (None,) * 3]
    cdict = {'red': [], 'green': [], 'blue': []}
    for i, item in enumerate(seq):
        if isinstance(item, float):
            r1, g1, b1 = seq[i - 1]
            r2, g2, b2 = seq[i + 1]
            cdict['red'].append([item, r1, r2])
            cdict['green'].append([item, g1, g2])
            cdict['blue'].append([item, b1, b2])
    return mcolors.LinearSegmentedColormap('CustomMap', cdict)
c = mcolors.ColorConverter().to_rgb
rvb = make_colormap(
      [c('#E94B00'), c('#FF9021'), 0.30,
       c('#FF9021'), c('#FFFBEF'), 0.49,
       c('white')  , 0.51,
       c('#F4FFDD'), c('#7BFF0D'), 0.80,
       c('#7BFF0D'), c('#0AE000')])

#Define drawing region
unit = 1.88973
nx = 101 
ny = 101
xs = unit*np.linspace(-3, 3, nx)
ys = unit*np.linspace(-3, 3, ny)
iz = 0.
coords = []
for ix in xs:
    for iy in ys:
        coords.append((ix, iy, iz))
coords=np.array(coords)
ao=dft.numint.eval_ao(mol,coords)

ls = [8, 32, 128, 512]
fig = plt.figure(figsize = (8, 7))
ax = fig.subplots(2, 2, sharex=True, sharey=True, gridspec_kw={'hspace': 0.06, 'wspace': 0.06})
for i in range(4):
    P=np.load('P'+str(ls[i])+'.npy')
    drho=dft.numint.eval_rho(mol,ao,P-P_tar)
    drho=np.reshape(drho, (nx, ny))
    c = ax[i//2,i%2].contourf(xs/unit, ys/unit, drho.T, levels=np.linspace(-0.05,0.05,101), cmap=rvb, extend='both')
    ax[i//2,i%2].set_aspect('equal', adjustable='box')

plt.subplots_adjust(top=0.95,right=0.84,bottom=0.06,left=0.06)
cbar_ax = fig.add_axes([0.88, 0.06, 0.03, 0.88])
fig.colorbar(c, cax=cbar_ax,ticks=np.linspace(-0.05,0.05,11))
#plt.savefig('bz_dendiff.pdf', format='pdf')
plt.show()

3. Oxygen

Similar with benzene example but CCSD density is used as a target

import numpy as np
from pyscf import gto, scf, lib, dft, cc
from kspies import zmp, wy, util
lib.num_threads(8)

#Define system
mol=gto.M(atom= '''
O 1.208 0.0 0.0
O 0.0   0.0 0.0
''',
basis = 'cc-pVQZ',spin=2,  verbose=3)

#Reference calculation
mf=scf.UHF(mol)
mf.kernel()
mc=cc.CCSD(mf)
mc.kernel()
dm_tar = util.mo2ao(mf.mo_coeff, mc.make_rdm1())
np.save('P_tar',dm_tar)

dm_tar=np.load('P_tar.npy')

#Grid generation for density accuracy check
grids=dft.gen_grid.Grids(mol)
grids.build()
coords=grids.coords
weights=grids.weights
ao=dft.numint.eval_ao(mol, coords)
def ndiff(P):
    if P.ndim==3: P=P[0]+P[1]
    rho=dft.numint.eval_rho(mol, ao, P)
    return np.einsum('r,r',abs(rho),weights)

#DF-UZMP
mz =zmp.UZMP(mol, dm_tar)
mz.with_df = True
for l in [ 8, 16, 32, 64, 128, 256, 512 ]:
    mz.level_shift = l*0.1
    mz.zscf(l)

#UWY
mw = wy.UWY(mol, dm_tar)
mw.run()
mw.info()
dm_wy = mw.dm
print("RWY density difference: ",ndiff(dm_wy-dm_tar))

4. Regularized WY

Perform restricted WY on molecular nitrogen with regularization

import numpy as np
from pyscf import gto, scf
from kspies import wy

mol = gto.M(atom = 'N 0 0 0 ; N 1.1 0 0', 
            basis = 'cc-pVDZ')
mf = scf.RHF(mol).run()
dm_tar = mf.make_rdm1()

PBS = gto.expand_etbs([(0, 13, 2**-4 , 2),
                       (1, 3 , 2**-2 , 2)])
mw = wy.RWY(mol, dm_tar, pbas=PBS)
#Note that for this designed-to-be ill-conditioned problem,
#Hessian-based optimization algorithms are problematic.
mw.method = 'bfgs'
mw.tol = 2e-7
mw.run()
mw.info()
Ws_fin = mw.Ws

etas = [ 2.**(-a) for a in np.linspace(5., 27., 45) ]
v = np.zeros(len(etas))
W = np.zeros(len(etas))
for i, eta in enumerate(etas):
    mw.reg=eta
    mw.run()
    v[i] = mw.Dvb()
    W[i] = mw.Ws
    mw.info()

import matplotlib.pyplot as plt
fig,ax = plt.subplots(2)
ax[0].scatter(np.log10(Ws_fin-W), np.log10(v))
ax[1].scatter(np.log10(etas), v*etas/(Ws_fin-W))

plt.tight_layout()
#plt.savefig('L_curves.pdf', format='pdf')
#plt.savefig('L_curves.eps', format='eps')
plt.show()

5. User Defined System

Perform restricted WY on a user defined harmonic oscillator. This example shows generation of finite-difference Hamiltonian, solving HF equation with it, and perform WY calculation with that HF density as a target To create a user defined instance, a number of settings must be specified:

import numpy as np
from pyscf import lib, gto
from kspies import wy
import matplotlib.pyplot as plt
from scipy.linalg import toeplitz
from scipy.linalg import eigh

#Define system
x = np.linspace(-10, 10, 201) #Domain
h = (x[-1]-x[0])/(len(x)-1)   #grid spacing
n = len(x)                    #Dimension of basis
a = np.zeros(len(x))
a[0] = 2.
a[1] = -1.
T = toeplitz(a,a)/(2*h**2)    #Kinetic energy matrix by 2nd order FD

S = np.identity(len(x))       #Overlap matrix
k = 0.25
V = np.diag(0.5*k*x**2)       #Harmonic potential matrix

l = 0.5                       #1D e-e soft repulsion parameter
def deno(l):
    b = np.expand_dims(x, axis=0)
    dist = abs(b-b.T)
    return 1./np.sqrt(dist**2+l**2)

def get_J(dm):
    J = np.diag(np.einsum('ii,ik->k', dm, deno(l)))
    return J

def get_K(dm):
    K = np.einsum('il,il->il', dm, deno(l))
    return K

#Pass to mole object
mol = gto.M()
mol.nelectron = 4
mol.verbose = 0
mol.incore_anyway = True

#Solve HF equation
F = T+V
for i in range(15):
    e,C = eigh(F,S)
    dm = 2*np.einsum('ik,jk->ij', C[:,:mol.nelectron//2], C[:,:mol.nelectron//2])
    J = get_J(dm)
    K = get_K(dm)
    F = T+V+J-0.5*K
    print("EHF = ",np.einsum('ij,ji', T+V+0.5*J-0.25*K, dm))
dm_tar = dm

plt.plot(x, 10*np.diag(dm_tar)/h, label='den(HF)', color='black') # x10 scaled density

#Three-center overlap integral
Sijt = np.zeros((n,n,n))
for i in range(n):
  Sijt[i,i,i] = 1.

#Run WY
mw = wy.RWY(mol, dm_tar, Sijt=Sijt)
mw.tol = 1e-7
mw.method = 'bfgs'
mw.T = T    #Kinetic energy matrix - finite difference in this example
mw.Tp = T   #Kinetic energy matrix in potential basis 
mw.V = V    #External potential matrix
mw.S = S    #Overlap matrix
mw.guide = None
mw.run()
mw.info()
mw.time_profile()

#Plotting
Vb = np.diag(mw.b) #-mw.b[50]) #KS potential is unique up to a constant. 
plt.plot(x, 10*np.diag(mw.dm)/h, label='den(WY)', color='red', linestyle='--') # x10 scaled density
plt.plot(x, np.diag(V), label=r'$v_{ext}$(r)')
plt.plot(x, np.diag(V+Vb), label=r'$v_{S}$(r)')
plt.plot(x, 1e+6*np.diag(mw.dm-dm_tar)/h,label='den(WY-HF)', color='blue', linestyle='--') # x10^6 scaled diff
plt.xlim(-10, 10)
plt.ylim(-0.5, 10)
plt.tight_layout()

plt.legend()
plt.show()

6. Plot XC ZMP FA

Calculate and plot regularized ZMP using the util functions with exchange-correlation aspects of the Fermi-Amaldi potential. Demonstrating the guide, level_shift, dm, and zscf routines.

import numpy as np
from pyscf import gto, scf
from kspies import zmp, util
import matplotlib.pyplot as plt

#Target density from HF
mol = gto.M(atom='Ne',
            basis='aug-cc-pVTZ')
mf=scf.RHF(mol).run()
dm_tar = mf.make_rdm1()

#Define plotting domain
coords = []
for x in np.linspace(0, 3, 1001):
    coords.append((x, 0., 0.))
coords = np.array(coords)

plt.figure(figsize=(3,4))

#ZMP calculations
mz = zmp.RZMP(mol, dm_tar)
mz.guide='faxc'
for l in [ 16, 128, 1024 ]:
    mz.level_shift = 0.1*l
    mz.zscf(l) 
    dmxc = l*mz.dm-(l+1./mol.nelectron)*dm_tar
    vxc = util.eval_vh(mol, coords, dmxc )
    plt.plot(coords[:, 0], vxc, label = r'$\lambda$='+str(l))

plt.xlabel("x")
plt.ylabel("vx(r)")
plt.xlim(0, 3)
plt.ylim(-5, 0)
plt.legend()
plt.subplots_adjust(left=0.2, right=0.93, bottom=0.12, top=0.97)
plt.show()

7. Plot XC WY PBE

Calculate and plot regularized WY using the util functions with a PBE guiding potential.

import numpy as np
from pyscf import gto, scf, dft
from kspies import wy, util
import matplotlib.pyplot as plt

mol = gto.M(atom='Ne',
            basis='aug-cc-pVTZ')
mf=scf.RHF(mol).run()
dm_tar = mf.make_rdm1()

#Define plotting domain
coords = []
for x in np.linspace(0, 3, 1001):
    coords.append((x, 0., 0.))
coords = np.array(coords)

plt.figure(figsize=(3,4))

#WY calculations
mw = wy.RWY(mol, dm_tar, pbas='cc-pVQZ')
mw.guide = 'pbe'
mw.tol = 1e-7
pb = dft.numint.eval_ao(mw.pmol, coords) #potential basis values on grid
vg = util.eval_vxc(mol, dm_tar, mw.guide, coords) #guiding potential on grid

for expo in np.arange(2,6):
    mw.reg = 10.**(-expo)
    mw.run()
    vC = np.einsum('t,rt->r', mw.b, pb)
    mw.info()
    plt.plot(coords[:,0], vg+vC, label=r'$\eta$ = 10^'+str(-expo))

plt.xlabel("x")
plt.ylabel("vx(r)")
plt.xlim(0, 3)
plt.ylim(-5, 0)
plt.legend()
plt.subplots_adjust(left=0.2, right=0.93, bottom=0.12, top=0.97)
plt.show()

8. User Defined Potential Basis

Use of user-defined potential basis. Slater-type basis is given as an example

import numpy as np
from pyscf import gto, scf, dft 
from kspies import wy

mol_1 = gto.M(atom='Ne', basis='aug-cc-pVTZ')
mf = scf.UHF(mol_1).run()
P_tar_1 = mf.make_rdm1()

def make_sto(zeta):
    def sto(coords):
        dist = np.sum(coords**2, axis=1)**.5
        return np.exp(-zeta*dist)
    return sto

pbas = []
for zeta in [ 0.25, 0.5, 1., 2., 4., 8. ]:
    pbas.append(make_sto(zeta))

mol_3=gto.M(atom='Ne', basis='aug-cc-pVTZ')
def make_gto(i): #make aug-cc-pVTZ as user-defined basis
    def gto(coords):
        return dft.numint.eval_ao(mol_3, coords)[:,i]
    return gto

#If this pbas is used, result is same with setting
#pbas='aug-cc-pVTZ'
#pbas = []
#for i in range(mol_3.nao_nr()):
#  pbas.append(make_gto(i))

mywy = wy.UWY(mol_1, P_tar_1, pbas=pbas)
mywy.run()
mywy.info()