├── GGA.py
├── LICENSE
├── LSDA.py
├── NNGGA
├── NNLSDA
├── NNNRA
├── NNmGGA
├── NRA.py
├── README.md
├── metaGGA.py
├── numint.py
└── numint_1.5.py


/GGA.py:
--------------------------------------------------------------------------------
  1 | # -*- coding: utf-8 -*-
  2 | #!/usr/bin/env python
  3 | import numpy as np
  4 | from pyscf import gto, dft
  5 | import math
  6 | import torch 
  7 | from torch.autograd import Variable 
  8 | import torch.nn as nn 
  9 | import torch.nn.functional as F 
 10 | 
 11 | 
 12 | # definition of target molecule #
 13 | mol = gto.Mole()
 14 | mol.verbose = 4
 15 | 
 16 | mol.atom="""8            .000000     .000000     .119262
 17 | 1            .000000     .763239    -.477047
 18 | 1            .000000    -.763239    -.477047"""
 19 | mol.charge=0
 20 | mol.spin  =0
 21 | mol.basis = "6-31G"
 22 | mol.build()
 23 | 
 24 | # definition of NN structure #
 25 | hidden=100
 26 | print("hidden nodes= "+str(hidden))
 27 | 
 28 | class Net(nn.Module):
 29 |     def __init__(self):
 30 |         super(Net, self).__init__()
 31 |         self.fc1 = nn.Linear(3, hidden)
 32 |         self.fc2 = nn.Linear(hidden, hidden)
 33 |         self.fc3 = nn.Linear(hidden, hidden)
 34 |         self.fc4 = nn.Linear(hidden, 1)
 35 | 
 36 | 
 37 |     def forward(self, x):
 38 |         t=torch.empty((x.shape[0],3))
 39 |         unif=(x[:,1]+x[:,0]+1e-7)**(1.0/3)
 40 |         t[:,0]=unif
 41 |         t[:,1]=((1+torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3)+(1-torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3))*0.5
 42 |         t[:,2]=((x[:,2]+x[:,4]+2*x[:,3])**0.5+1e-7)/unif**4
 43 |         ds=(1+torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(5.0/3)+(1-torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(5.0/3)
 44 | 
 45 |         #print(t)
 46 |         logt=torch.log(t)
 47 |         g1=F.elu(self.fc1(logt))
 48 |         g2=F.elu(self.fc2(g1))
 49 |         g3=F.elu(self.fc3(g2))
 50 | 
 51 |         eunif=(x[:,1]+x[:,0]).view(-1,1)**(1.0/3)
 52 |         spinscale=(((1+torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3)+(1-torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3))*0.5).view(-1,1)
 53 |         g4=-(F.elu(self.fc4(g3))+1)*eunif*spinscale
 54 |         return g4
 55 | 
 56 | #loading NN weights
 57 | model = Net()
 58 | model.load_state_dict(torch.load("NNGGA",map_location='cpu'))
 59 | 
 60 | # definition of functional #
 61 | def eval_xc(xc_code, rho, spin, relativity=0, deriv=2, verbose=None):
 62 |     if spin!=0:
 63 |         rho1=rho[0]
 64 |         rho2=rho[1]
 65 |         rho01, dx1, dy1, dz1 = rho1[:4]
 66 |         rho02, dx2, dy2, dz2 = rho2[:4]
 67 |         gamma1=dx1**2+dy1**2+dz1**2+1e-7
 68 |         gamma2=dx2**2+dy2**2+dz2**2+1e-7
 69 |         gamma12=dx1*dx2+dy1*dy2+dz1*dz2+1e-7
 70 |     else:
 71 |         rho0, dx, dy, dz= rho[:4]
 72 |         gamma1=gamma2=gamma12=(dx**2+dy**2+dz**2)*0.25+1e-7
 73 |         rho01=rho02=rho0*0.5
 74 | 
 75 |     rho0=np.concatenate((rho01.reshape((-1,1)),rho02.reshape((-1,1)),gamma1.reshape((-1,1)),gamma12.reshape((-1,1)),gamma2.reshape((-1,1))),axis=1)
 76 |     
 77 |     N=rho0.shape[0]
 78 | 
 79 |     x=Variable(torch.Tensor(rho0),requires_grad=True)
 80 |     pred_exc=model(x)
 81 |     
 82 |     exc=pred_exc.data[:,0].numpy()#+0.015471944-0.002533411+0.001574637
 83 |     eneden=torch.dot(pred_exc[:,0],x[:,0]+x[:,1])
 84 |     eneden.backward()
 85 |     grad=x.grad.data.numpy()
 86 | 
 87 |     if spin!=0:
 88 |         vlapl=np.zeros((N,2))
 89 |         vrho=np.hstack((grad[:,0].reshape((-1,1)),grad[:,1].reshape((-1,1))))
 90 |         vgamma=np.hstack((grad[:,2].reshape((-1,1)),grad[:,3].reshape((-1,1)),grad[:,4].reshape((-1,1))))
 91 | 
 92 |     else:
 93 |         vlapl=np.zeros(N)    
 94 |         vrho=(grad[:,0]+grad[:,1])/2
 95 |         vgamma=(grad[:,2]+grad[:,4]+grad[:,3])/4
 96 | 
 97 |     vxc=(vrho,vgamma,None,None)
 98 |     return exc, vxc, None, None
 99 | 
100 | # DFT calculation #
101 | if mol.spin==0:
102 |     mfl = dft.RKS(mol)
103 | else:
104 |     mfl = dft.UKS(mol)
105 | 
106 | mfl = mfl.define_xc_(eval_xc, 'GGA')
107 | mfl.kernel()


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2019
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/LSDA.py:
--------------------------------------------------------------------------------
 1 | # -*- coding: utf-8 -*-
 2 | #!/usr/bin/env python
 3 | import numpy as np
 4 | from pyscf import gto, dft
 5 | import math
 6 | import torch 
 7 | from torch.autograd import Variable 
 8 | import torch.nn as nn 
 9 | import torch.nn.functional as F 
10 | 
11 | 
12 | # definition of target molecule #
13 | mol = gto.Mole()
14 | mol.verbose = 4
15 | 
16 | mol.atom="""8            .000000     .000000     .119262
17 | 1            .000000     .763239    -.477047
18 | 1            .000000    -.763239    -.477047"""
19 | mol.charge=0
20 | mol.spin  =0
21 | mol.basis = "6-31G"
22 | mol.build()
23 | 
24 | # definition of NN structure #
25 | hidden=100
26 | print("hidden nodes= "+str(hidden))
27 | 
28 | class Net(nn.Module):
29 |     def __init__(self):
30 |         super(Net, self).__init__()
31 |         self.fc1 = nn.Linear(2, hidden)
32 |         self.fc2 = nn.Linear(hidden, hidden)
33 |         self.fc3 = nn.Linear(hidden, hidden)
34 |         self.fc4 = nn.Linear(hidden, 1)
35 | 
36 | 
37 |     def forward(self, x):
38 |         t=torch.empty((x.shape[0],2))
39 |         unif=(x[:,1]+x[:,0]+1e-7)**(1.0/3)
40 |         t[:,0]=unif
41 |         t[:,1]=((1+torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3)+(1-torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3))*0.5
42 | 
43 |         #print(t)
44 |         logt=torch.log(t)#/scale.float()
45 |         g1=F.elu(self.fc1(logt))
46 |         g2=F.elu(self.fc2(g1))
47 |         g3=F.elu(self.fc3(g2))
48 |         #g4=F.elu(self.fc4(g3))
49 |         eunif=(x[:,1]+x[:,0]).view(-1,1)**(1.0/3)
50 |         spinscale=(((1+torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3)+(1-torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3))*0.5).view(-1,1)
51 |         g4=-(F.elu(self.fc4(g3))+1)*eunif*spinscale
52 |         return g4
53 | 
54 | 
55 | #loading NN weights
56 | model = Net()
57 | model.load_state_dict(torch.load("NNLSDA",map_location='cpu'))
58 | 
59 | # definition of functional #
60 | 
61 | def eval_xc(xc_code, rho, spin, relativity=0, deriv=2, verbose=None):
62 |     if spin!=0:
63 |         rho01=rho[0]
64 |         rho02=rho[1]
65 |         
66 |     else:
67 |         rho01=rho02=rho*0.5
68 | 
69 |     rho0=np.concatenate((rho01.reshape((-1,1)),rho02.reshape((-1,1))),axis=1)
70 |     N=rho0.shape[0]
71 |     x=Variable(torch.Tensor(rho0),requires_grad=True)
72 |     pred_exc=model(x)
73 |     exc=pred_exc.data[:,0].numpy()
74 |     eneden=torch.dot(pred_exc[:,0],x[:,0]+x[:,1])
75 |     eneden.backward()
76 |     grad=x.grad.data.numpy()
77 | 
78 |     if spin!=0:
79 |         vrho=np.hstack((grad[:,0].reshape((-1,1)),grad[:,1].reshape((-1,1))))
80 | 
81 |     else:
82 |         vlapl=np.zeros(N)    
83 |         vrho=(grad[:,0]+grad[:,1])/2
84 |     vxc=(vrho, None, None, None)
85 |     return exc, vxc, None, None
86 | 
87 | # DFT calculation #
88 | if mol.spin==0:
89 |     mfl = dft.RKS(mol)
90 | else:
91 |     mfl = dft.UKS(mol)
92 | 
93 | mfl = mfl.define_xc_(eval_xc, 'LDA')
94 | mfl.kernel()
95 | 


--------------------------------------------------------------------------------
/NNGGA:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ml-electron-project/NNfunctional/3f4a97998747dc14a9296a88d16572601621efab/NNGGA


--------------------------------------------------------------------------------
/NNLSDA:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ml-electron-project/NNfunctional/3f4a97998747dc14a9296a88d16572601621efab/NNLSDA


--------------------------------------------------------------------------------
/NNNRA:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ml-electron-project/NNfunctional/3f4a97998747dc14a9296a88d16572601621efab/NNNRA


--------------------------------------------------------------------------------
/NNmGGA:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ml-electron-project/NNfunctional/3f4a97998747dc14a9296a88d16572601621efab/NNmGGA


--------------------------------------------------------------------------------
/NRA.py:
--------------------------------------------------------------------------------
  1 | # -*- coding: utf-8 -*-
  2 | #!/usr/bin/env python
  3 | import numpy as np
  4 | from pyscf import gto, dft
  5 | import math
  6 | import torch 
  7 | from torch.autograd import Variable 
  8 | import torch.nn as nn 
  9 | import torch.nn.functional as F 
 10 | 
 11 | ####################
 12 | # In order to use the NRA functional, please replace 
 13 | # "(where you installed PySCF)/pyscf/dft/numint.py"
 14 | #into numint.py in this folder.
 15 | ####################
 16 | 
 17 | # definition of target molecule #
 18 | mol = gto.Mole()
 19 | mol.verbose = 4
 20 | 
 21 | mol.atom="""8            .000000     .000000     .119262
 22 | 1            .000000     .763239    -.477047
 23 | 1            .000000    -.763239    -.477047"""
 24 | mol.charge=0
 25 | mol.spin  =0
 26 | mol.basis = "6-31G"
 27 | mol.build()
 28 | 
 29 | # definition of NN structure #
 30 | hidden=100
 31 | print("hidden nodes= "+str(hidden))
 32 | 
 33 | class Net(nn.Module):
 34 |     def __init__(self):
 35 |         super(Net, self).__init__()
 36 |         self.fc1 = nn.Linear(5, hidden)
 37 |         self.fc2 = nn.Linear(hidden, hidden)
 38 |         self.fc3 = nn.Linear(hidden, hidden)
 39 |         self.fc4 = nn.Linear(hidden, 1)
 40 | 
 41 | 
 42 |     def forward(self, x):
 43 |         t=torch.empty((x.shape[0],5))
 44 |         unif=(x[:,1]+x[:,0]+1e-7)**(1.0/3)
 45 |         t[:,0]=unif
 46 |         t[:,1]=((1+torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3)+(1-torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3))*0.5
 47 |         t[:,2]=((x[:,2]+x[:,4]+2*x[:,3])**0.5+1e-7)/unif**4
 48 |         ds=(1+torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(5.0/3)+(1-torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(5.0/3)
 49 |         t[:,3]=(x[:,5]+x[:,6]+1e-7)/(unif**5*ds)
 50 |         t[:,4]=x[:,7]
 51 |         logt=torch.log(t)
 52 |         g1=F.elu(self.fc1(logt))
 53 |         g2=F.elu(self.fc2(g1))
 54 |         g3=F.elu(self.fc3(g2))
 55 |         eunif=(x[:,1]+x[:,0]).view(-1,1)**(1.0/3)
 56 |         spinscale=(((1+torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3)+(1-torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3))*0.5).view(-1,1)
 57 |         g4=-(F.elu(self.fc4(g3))+1)*eunif*spinscale
 58 |         return g4
 59 | 
 60 | #loading NN weights
 61 | model = Net()
 62 | model.load_state_dict(torch.load("NNNRA",map_location='cpu'))
 63 | 
 64 | # definition of functional #
 65 | 
 66 | old_coords=np.empty(1)
 67 | cutoff=math.log(1e3)
 68 | intlist=[]
 69 | def eval_xc(xc_code, rho, spin, weight, coords, relativity=0, deriv=2, verbose=None):
 70 |     global old_coords
 71 |     global intlist
 72 |     if spin!=0:
 73 |         rho1=rho[0]
 74 |         rho2=rho[1]
 75 |         rho01, dx1, dy1, dz1, lapl1, tau1 = rho1[:6]
 76 |         rho02, dx2, dy2, dz2, lapl2, tau2 = rho2[:6]
 77 |         gamma1=dx1**2+dy1**2+dz1**2+1e-7
 78 |         gamma2=dx2**2+dy2**2+dz2**2+1e-7
 79 |         gamma12=dx1*dx2+dy1*dy2+dz1*dz2+1e-7
 80 |     else:
 81 |         rho0, dx, dy, dz, lapl, tau = rho[:6]
 82 |         gamma1=gamma2=gamma12=(dx**2+dy**2+dz**2)*0.25+1e-7
 83 |         rho01=rho02=rho0*0.5
 84 |         tau1=tau2=tau*0.5
 85 | 
 86 |     rhos=rho01+rho02
 87 | 
 88 |     N=rhos.shape[0]
 89 |     Kr=np.empty(N)
 90 |     dKr=np.empty(N)
 91 |     if np.array_equal(coords,old_coords) :
 92 |         for i in range (0,N):
 93 |             dis=((coords[i]-coords[intlist[i]])**2).sum(axis=1)**0.5/0.2
 94 |             Rw=weight[intlist[i]]*np.exp(-dis)
 95 |             Kr[i]=np.dot(Rw,rhos[intlist[i]])
 96 |         
 97 |     else:
 98 |         intlist=[]
 99 |         for i in range (0,N):
100 |             dis=((coords[i]-coords)**2).sum(axis=1)**0.5/0.2
101 |             intlist.append(np.where(dis<cutoff)[0])
102 |             Rw=weight[intlist[i]]*np.exp(-dis[intlist[i]])
103 |             Kr[i]=np.dot(Rw,rhos[intlist[i]])
104 | 
105 | 
106 |     Kr=np.clip(Kr,0,None)+1e-7
107 |     rho0=np.concatenate((rho01.reshape((-1,1)),rho02.reshape((-1,1)),gamma1.reshape((-1,1)),gamma12.reshape((-1,1)),gamma2.reshape((-1,1)),tau1.reshape((-1,1)),tau2.reshape((-1,1)),Kr.reshape((-1,1))),axis=1)
108 | 
109 |     x=Variable(torch.Tensor(rho0),requires_grad=True)
110 |     pred_exc=model(x)
111 |     
112 |     exc=pred_exc.data[:,0].numpy()
113 |     eneden=torch.dot(pred_exc[:,0],x[:,0]+x[:,1])
114 |     eneden.backward()
115 |     grad=x.grad.data.numpy()
116 |     Rgra=grad[:,7]
117 |     
118 |     for i in range (0,N):
119 |         Rw=weight[intlist[i]]*np.exp(-((coords[i]-coords[intlist[i]])**2).sum(axis=1)**0.5/0.2)
120 |         dKr[i]=np.dot(Rw,Rgra[intlist[i]])
121 | 
122 |     if spin!=0:
123 |         vlapl=np.zeros((N,2))
124 |         vrho=np.hstack(((grad[:,0]+dKr).reshape((-1,1)),(grad[:,1]+dKr).reshape((-1,1))))
125 |         vgamma=np.hstack((grad[:,2].reshape((-1,1)),grad[:,3].reshape((-1,1)),grad[:,4].reshape((-1,1))))
126 |         vtau=np.hstack((grad[:,5].reshape((-1,1)),grad[:,6].reshape((-1,1))))
127 |     else:
128 |         vlapl=np.zeros(N)    
129 |         vrho=(grad[:,0]+grad[:,1])/2+dKr
130 |         vgamma=(grad[:,2]+grad[:,4]+grad[:,3])/4
131 |         vtau=(grad[:,5]+grad[:,6])/2
132 |     vxc=(vrho,vgamma,vlapl,vtau)
133 |     old_coords=np.copy(coords)
134 |     return exc, vxc, None, None
135 | 
136 | 
137 | 
138 | 
139 | # DFT calculation #
140 | if mol.spin==0:
141 |     mfl = dft.RKS(mol)
142 | else:
143 |     mfl = dft.UKS(mol)
144 | 
145 | mfl = mfl.define_xc_(eval_xc, 'MGGA+')
146 | mfl.kernel()


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | NN-based functional
 2 | ====
 3 | 
 4 | Parameters of NN-based functionals and example codes to call them (.py).
 5 | 
 6 | PySCF 1.6.2 and Pytorch 1.1.0 packages are required to run the example codes.
 7 | 
 8 | To call the NRA functional, (Where you installed PySCF)/pyscf/dft/numint.py should be replaced to the one in this folder.
 9 | 
10 | Note that the numerical integration in the xc energy and potential of "NRA.py" is currently implemented in inefficient way. It will take a lot of computational time even for a small molecule. 
11 | 
12 | NOTICE:The distributed codes includes the work that is distributed in the Apache License 2.0.
13 | 
14 | 
15 | 21/07/2019 numint.py is modified for PySCF version 1.6.2.
16 | 
17 | ## Reference
18 | Cite the following paper in any related works.
19 | @misc{1903.00238,
20 | Author = {Ryo Nagai and Ryosuke Akashi and Osamu Sugino},
21 | Title = {Completing density functional theory by machine-learning hidden messages from molecules},
22 | Year = {2019},
23 | Eprint = {arXiv:1903.00238},
24 | }
25 | 


--------------------------------------------------------------------------------
/metaGGA.py:
--------------------------------------------------------------------------------
  1 | # -*- coding: utf-8 -*-
  2 | #!/usr/bin/env python
  3 | import numpy as np
  4 | from pyscf import gto, dft
  5 | import math
  6 | import torch 
  7 | from torch.autograd import Variable 
  8 | import torch.nn as nn 
  9 | import torch.nn.functional as F 
 10 | 
 11 | 
 12 | # definition of target molecule #
 13 | mol = gto.Mole()
 14 | mol.verbose = 4
 15 | 
 16 | mol.atom="""8            .000000     .000000     .119262
 17 | 1            .000000     .763239    -.477047
 18 | 1            .000000    -.763239    -.477047"""
 19 | mol.charge=0
 20 | mol.spin  =0
 21 | mol.basis = "6-31G"
 22 | mol.build()
 23 | 
 24 | # definition of NN structure #
 25 | hidden=100
 26 | print("hidden nodes= "+str(hidden))
 27 | 
 28 | class Net(nn.Module):
 29 |     def __init__(self):
 30 |         super(Net, self).__init__()
 31 |         self.fc1 = nn.Linear(4, hidden)
 32 |         self.fc2 = nn.Linear(hidden, hidden)
 33 |         self.fc3 = nn.Linear(hidden, hidden)
 34 |         self.fc4 = nn.Linear(hidden, 1)
 35 |         #self.fc5 = nn.Linear(hidden, 1)
 36 | 
 37 | 
 38 |     def forward(self, x):
 39 |         t=torch.empty((x.shape[0],4))
 40 |         unif=(x[:,1]+x[:,0]+1e-7)**(1.0/3)
 41 |         t[:,0]=unif
 42 |         t[:,1]=((1+torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3)+(1-torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3))*0.5
 43 |         t[:,2]=((x[:,2]+x[:,4]+2*x[:,3])**0.5+1e-7)/unif**4
 44 |         ds=(1+torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(5.0/3)+(1-torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(5.0/3)
 45 |         t[:,3]=(x[:,5]+x[:,6]+1e-7)/(unif**5*ds)
 46 |         logt=torch.log(t)
 47 |         g1=F.elu(self.fc1(logt))
 48 |         g2=F.elu(self.fc2(g1))
 49 |         g3=F.elu(self.fc3(g2))
 50 |         eunif=(x[:,1]+x[:,0]).view(-1,1)**(1.0/3)
 51 |         spinscale=(((1+torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3)+(1-torch.div((x[:,0]-x[:,1]),(x[:,1]+x[:,0]+1e-7)))**(4.0/3))*0.5).view(-1,1)
 52 |         g4=-(F.elu(self.fc4(g3))+1)*eunif*spinscale
 53 |         return g4
 54 | 
 55 | #loading NN weights
 56 | model = Net()
 57 | model.load_state_dict(torch.load("NNmGGA",map_location='cpu'))
 58 | 
 59 | # definition of functional #
 60 | def eval_xc(xc_code, rho, spin, relativity=0, deriv=2, verbose=None):
 61 |     if spin!=0:
 62 |         rho1=rho[0]
 63 |         rho2=rho[1]
 64 |         rho01, dx1, dy1, dz1, lapl1, tau1 = rho1[:6]
 65 |         rho02, dx2, dy2, dz2, lapl2, tau2 = rho2[:6]
 66 |         gamma1=dx1**2+dy1**2+dz1**2+1e-7
 67 |         gamma2=dx2**2+dy2**2+dz2**2+1e-7
 68 |         gamma12=dx1*dx2+dy1*dy2+dz1*dz2+1e-7
 69 |     else:
 70 |         rho0, dx, dy, dz, lapl, tau = rho[:6]
 71 |         gamma1=gamma2=gamma12=(dx**2+dy**2+dz**2)*0.25+1e-7
 72 |         rho01=rho02=rho0*0.5
 73 |         tau1=tau2=tau*0.5
 74 | 
 75 |     rho0=np.concatenate((rho01.reshape((-1,1)),rho02.reshape((-1,1)),gamma1.reshape((-1,1)),gamma12.reshape((-1,1)),gamma2.reshape((-1,1)),tau1.reshape((-1,1)),tau2.reshape((-1,1))),axis=1)
 76 |     
 77 |     N=rho0.shape[0]
 78 | 
 79 |     x=Variable(torch.Tensor(rho0),requires_grad=True)
 80 |     pred_exc=model(x)
 81 |     
 82 |     exc=pred_exc.data[:,0].numpy()#+0.015471944-0.002533411+0.001574637
 83 |     eneden=torch.dot(pred_exc[:,0],x[:,0]+x[:,1])
 84 |     eneden.backward()
 85 |     grad=x.grad.data.numpy()
 86 | 
 87 |     if spin!=0:
 88 |         vlapl=np.zeros((N,2))
 89 |         vrho=np.hstack((grad[:,0].reshape((-1,1)),grad[:,1].reshape((-1,1))))
 90 |         vgamma=np.hstack((grad[:,2].reshape((-1,1)),grad[:,3].reshape((-1,1)),grad[:,4].reshape((-1,1))))
 91 |         vtau=np.hstack((grad[:,5].reshape((-1,1)),grad[:,6].reshape((-1,1))))
 92 |     else:
 93 |         vlapl=np.zeros(N)    
 94 |         vrho=(grad[:,0]+grad[:,1])/2
 95 |         vgamma=(grad[:,2]+grad[:,4]+grad[:,3])/4
 96 |         vtau=(grad[:,5]+grad[:,6])/2
 97 |     vxc=(vrho,vgamma,vlapl,vtau)
 98 |     return exc, vxc, None, None
 99 | 
100 | 
101 | # DFT calculation #
102 | if mol.spin==0:
103 |     mfl = dft.RKS(mol)
104 | else:
105 |     mfl = dft.UKS(mol)
106 | 
107 | mfl = mfl.define_xc_(eval_xc, 'MGGA')
108 | mfl.kernel()
109 | 


--------------------------------------------------------------------------------
/numint.py:
--------------------------------------------------------------------------------
   1 | #!/usr/bin/env python
   2 | 
   3 | #######################
   4 | #Changed by Ryo Nagai, March 2019.
   5 | #For PySCF ver 1.6.2
   6 | #This software includes the work that is distributed in the Apache License 2.0.
   7 | #######################
   8 | 
   9 | # Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
  10 | #
  11 | # Licensed under the Apache License, Version 2.0 (the "License");
  12 | # you may not use this file except in compliance with the License.
  13 | # You may obtain a copy of the License at
  14 | #
  15 | #     http://www.apache.org/licenses/LICENSE-2.0
  16 | #
  17 | # Unless required by applicable law or agreed to in writing, software
  18 | # distributed under the License is distributed on an "AS IS" BASIS,
  19 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  20 | # See the License for the specific language governing permissions and
  21 | # limitations under the License.
  22 | #
  23 | # Author: Qiming Sun <osirpt.sun@gmail.com>
  24 | #
  25 | 
  26 | import warnings
  27 | import ctypes
  28 | import numpy
  29 | import scipy.linalg
  30 | from pyscf import lib
  31 | from pyscf.lib import logger
  32 | try:
  33 |     from pyscf.dft import libxc
  34 | except (ImportError, OSError):
  35 |     try:
  36 |         from pyscf.dft import xcfun
  37 |         libxc = xcfun
  38 |     except (ImportError, OSError):
  39 |         import warnings
  40 |         warnings.warn('XC functional libraries (libxc or XCfun) are not available.')
  41 |         from pyscf.dft import xc
  42 |         libxc = xc
  43 | 
  44 | from pyscf.dft.gen_grid import make_mask, BLKSIZE
  45 | from pyscf import __config__
  46 | 
  47 | libdft = lib.load_library('libdft')
  48 | OCCDROP = getattr(__config__, 'dft_numint_OCCDROP', 1e-12)
  49 | # The system size above which to consider the sparsity of the density matrix.
  50 | # If the number of AOs in the system is less than this value, all tensors are
  51 | # treated as dense quantities and contracted by dgemm directly.
  52 | SWITCH_SIZE = getattr(__config__, 'dft_numint_SWITCH_SIZE', 800)
  53 | 
  54 | def eval_ao(mol, coords, deriv=0, shls_slice=None,
  55 |             non0tab=None, out=None, verbose=None):
  56 |     '''Evaluate AO function value on the given grids.
  57 | 
  58 |     Args:
  59 |         mol : an instance of :class:`Mole`
  60 | 
  61 |         coords : 2D array, shape (N,3)
  62 |             The coordinates of the grids.
  63 | 
  64 |     Kwargs:
  65 |         deriv : int
  66 |             AO derivative order.  It affects the shape of the return array.
  67 |             If deriv=0, the returned AO values are stored in a (N,nao) array.
  68 |             Otherwise the AO values are stored in an array of shape (M,N,nao).
  69 |             Here N is the number of grids, nao is the number of AO functions,
  70 |             M is the size associated to the derivative deriv.
  71 |         relativity : bool
  72 |             No effects.
  73 |         shls_slice : 2-element list
  74 |             (shl_start, shl_end).
  75 |             If given, only part of AOs (shl_start <= shell_id < shl_end) are
  76 |             evaluated.  By default, all shells defined in mol will be evaluated.
  77 |         non0tab : 2D bool array
  78 |             mask array to indicate whether the AO values are zero.  The mask
  79 |             array can be obtained by calling :func:`make_mask`
  80 |         out : ndarray
  81 |             If provided, results are written into this array.
  82 |         verbose : int or object of :class:`Logger`
  83 |             No effects.
  84 | 
  85 |     Returns:
  86 |         2D array of shape (N,nao) for AO values if deriv = 0.
  87 |         Or 3D array of shape (:,N,nao) for AO values and AO derivatives if deriv > 0.
  88 |         In the 3D array, the first (N,nao) elements are the AO values,
  89 |         followed by (3,N,nao) for x,y,z compoents;
  90 |         Then 2nd derivatives (6,N,nao) for xx, xy, xz, yy, yz, zz;
  91 |         Then 3rd derivatives (10,N,nao) for xxx, xxy, xxz, xyy, xyz, xzz, yyy, yyz, yzz, zzz;
  92 |         ...
  93 | 
  94 |     Examples:
  95 | 
  96 |     >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz')
  97 |     >>> coords = numpy.random.random((100,3))  # 100 random points
  98 |     >>> ao_value = eval_ao(mol, coords)
  99 |     >>> print(ao_value.shape)
 100 |     (100, 24)
 101 |     >>> ao_value = eval_ao(mol, coords, deriv=1, shls_slice=(1,4))
 102 |     >>> print(ao_value.shape)
 103 |     (4, 100, 7)
 104 |     >>> ao_value = eval_ao(mol, coords, deriv=2, shls_slice=(1,4))
 105 |     >>> print(ao_value.shape)
 106 |     (10, 100, 7)
 107 |     '''
 108 |     comp = (deriv+1)*(deriv+2)*(deriv+3)//6
 109 |     if mol.cart:
 110 |         feval = 'GTOval_cart_deriv%d' % deriv
 111 |     else:
 112 |         feval = 'GTOval_sph_deriv%d' % deriv
 113 |     return mol.eval_gto(feval, coords, comp, shls_slice, non0tab, out=out)
 114 | 
 115 | #TODO: \nabla^2 rho and tau = 1/2 (\nabla f)^2
 116 | def eval_rho(mol, ao, dm, non0tab=None, xctype='LDA', hermi=0, verbose=None):
 117 |     r'''Calculate the electron density for LDA functional, and the density
 118 |     derivatives for GGA functional.
 119 | 
 120 |     Args:
 121 |         mol : an instance of :class:`Mole`
 122 | 
 123 |         ao : 2D array of shape (N,nao) for LDA, 3D array of shape (4,N,nao) for GGA
 124 |             or (5,N,nao) for meta-GGA.  N is the number of grids, nao is the
 125 |             number of AO functions.  If xctype is GGA, ao[0] is AO value
 126 |             and ao[1:3] are the AO gradients.  If xctype is meta-GGA, ao[4:10]
 127 |             are second derivatives of ao values.
 128 |         dm : 2D array
 129 |             Density matrix
 130 | 
 131 |     Kwargs:
 132 |         non0tab : 2D bool array
 133 |             mask array to indicate whether the AO values are zero.  The mask
 134 |             array can be obtained by calling :func:`make_mask`
 135 |         xctype : str
 136 |             LDA/GGA/mGGA.  It affects the shape of the return density.
 137 |         hermi : bool
 138 |             dm is hermitian or not
 139 |         verbose : int or object of :class:`Logger`
 140 |             No effects.
 141 | 
 142 |     Returns:
 143 |         1D array of size N to store electron density if xctype = LDA;  2D array
 144 |         of (4,N) to store density and "density derivatives" for x,y,z components
 145 |         if xctype = GGA;  (6,N) array for meta-GGA, where last two rows are
 146 |         \nabla^2 rho and tau = 1/2(\nabla f)^2
 147 | 
 148 |     Examples:
 149 | 
 150 |     >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz')
 151 |     >>> coords = numpy.random.random((100,3))  # 100 random points
 152 |     >>> ao_value = eval_ao(mol, coords, deriv=0)
 153 |     >>> dm = numpy.random.random((mol.nao_nr(),mol.nao_nr()))
 154 |     >>> dm = dm + dm.T
 155 |     >>> rho, dx_rho, dy_rho, dz_rho = eval_rho(mol, ao, dm, xctype='LDA')
 156 |     '''
 157 |     xctype = xctype.upper()
 158 |     if xctype == 'LDA' or xctype == 'HF':
 159 |         ngrids, nao = ao.shape
 160 |     else:
 161 |         ngrids, nao = ao[0].shape
 162 | 
 163 |     if non0tab is None:
 164 |         non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
 165 |                              dtype=numpy.uint8)
 166 |     if not hermi:
 167 |         # (D + D.T)/2 because eval_rho computes 2*(|\nabla i> D_ij <j|) instead of
 168 |         # |\nabla i> D_ij <j| + |i> D_ij <\nabla j| for efficiency
 169 |         dm = (dm + dm.conj().T) * .5
 170 | 
 171 |     shls_slice = (0, mol.nbas)
 172 |     ao_loc = mol.ao_loc_nr()
 173 |     if xctype == 'LDA' or xctype == 'HF':
 174 |         c0 = _dot_ao_dm(mol, ao, dm, non0tab, shls_slice, ao_loc)
 175 |         #:rho = numpy.einsum('pi,pi->p', ao, c0)
 176 |         rho = _contract_rho(ao, c0)
 177 |     elif xctype in ('GGA', 'NLC'):
 178 |         rho = numpy.empty((4,ngrids))
 179 |         c0 = _dot_ao_dm(mol, ao[0], dm, non0tab, shls_slice, ao_loc)
 180 |         #:rho[0] = numpy.einsum('pi,pi->p', c0, ao[0])
 181 |         rho[0] = _contract_rho(c0, ao[0])
 182 |         for i in range(1, 4):
 183 |             #:rho[i] = numpy.einsum('pi,pi->p', c0, ao[i])
 184 |             rho[i] = _contract_rho(c0, ao[i])
 185 |             rho[i] *= 2 # *2 for +c.c. in the next two lines
 186 |             #c1 = _dot_ao_dm(mol, ao[i], dm, non0tab, shls_slice, ao_loc)
 187 |             #rho[i] += numpy.einsum('pi,pi->p', c1, ao[0])
 188 |     else: # meta-GGA
 189 |         # rho[4] = \nabla^2 rho, rho[5] = 1/2 |nabla f|^2
 190 |         rho = numpy.empty((6,ngrids))
 191 |         c0 = _dot_ao_dm(mol, ao[0], dm, non0tab, shls_slice, ao_loc)
 192 |         #:rho[0] = numpy.einsum('pi,pi->p', ao[0], c0)
 193 |         rho[0] = _contract_rho(ao[0], c0)
 194 |         rho[5] = 0
 195 |         for i in range(1, 4):
 196 |             #:rho[i] = numpy.einsum('pi,pi->p', c0, ao[i]) * 2 # *2 for +c.c.
 197 |             rho[i] = _contract_rho(c0, ao[i]) * 2
 198 |             c1 = _dot_ao_dm(mol, ao[i], dm.T, non0tab, shls_slice, ao_loc)
 199 |             #:rho[5] += numpy.einsum('pi,pi->p', c1, ao[i])
 200 |             rho[5] += _contract_rho(c1, ao[i])
 201 |         XX, YY, ZZ = 4, 7, 9
 202 |         ao2 = ao[XX] + ao[YY] + ao[ZZ]
 203 |         #:rho[4] = numpy.einsum('pi,pi->p', c0, ao2)
 204 |         rho[4] = _contract_rho(c0, ao2)
 205 |         rho[4] += rho[5]
 206 |         rho[4] *= 2
 207 |         rho[5] *= .5
 208 |     return rho
 209 | 
 210 | def eval_rho2(mol, ao, mo_coeff, mo_occ, non0tab=None, xctype='LDA',
 211 |               verbose=None):
 212 |     r'''Calculate the electron density for LDA functional, and the density
 213 |     derivatives for GGA functional.  This function has the same functionality
 214 |     as :func:`eval_rho` except that the density are evaluated based on orbital
 215 |     coefficients and orbital occupancy.  It is more efficient than
 216 |     :func:`eval_rho` in most scenario.
 217 | 
 218 |     Args:
 219 |         mol : an instance of :class:`Mole`
 220 | 
 221 |         ao : 2D array of shape (N,nao) for LDA, 3D array of shape (4,N,nao) for GGA
 222 |             or (5,N,nao) for meta-GGA.  N is the number of grids, nao is the
 223 |             number of AO functions.  If xctype is GGA, ao[0] is AO value
 224 |             and ao[1:3] are the AO gradients.  If xctype is meta-GGA, ao[4:10]
 225 |             are second derivatives of ao values.
 226 |         dm : 2D array
 227 |             Density matrix
 228 | 
 229 |     Kwargs:
 230 |         non0tab : 2D bool array
 231 |             mask array to indicate whether the AO values are zero.  The mask
 232 |             array can be obtained by calling :func:`make_mask`
 233 |         xctype : str
 234 |             LDA/GGA/mGGA.  It affects the shape of the return density.
 235 |         verbose : int or object of :class:`Logger`
 236 |             No effects.
 237 | 
 238 |     Returns:
 239 |         1D array of size N to store electron density if xctype = LDA;  2D array
 240 |         of (4,N) to store density and "density derivatives" for x,y,z components
 241 |         if xctype = GGA;  (6,N) array for meta-GGA, where last two rows are
 242 |         \nabla^2 rho and tau = 1/2(\nabla f)^2
 243 |     '''
 244 |     xctype = xctype.upper()
 245 |     if xctype == 'LDA' or xctype == 'HF':
 246 |         ngrids, nao = ao.shape
 247 |     else:
 248 |         ngrids, nao = ao[0].shape
 249 | 
 250 |     if non0tab is None:
 251 |         non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
 252 |                              dtype=numpy.uint8)
 253 |     shls_slice = (0, mol.nbas)
 254 |     ao_loc = mol.ao_loc_nr()
 255 |     pos = mo_occ > OCCDROP
 256 |     if pos.sum() > 0:
 257 |         cpos = numpy.einsum('ij,j->ij', mo_coeff[:,pos], numpy.sqrt(mo_occ[pos]))
 258 |         if xctype == 'LDA' or xctype == 'HF':
 259 |             c0 = _dot_ao_dm(mol, ao, cpos, non0tab, shls_slice, ao_loc)
 260 |             #:rho = numpy.einsum('pi,pi->p', c0, c0)
 261 |             rho = _contract_rho(c0, c0)
 262 |         elif xctype in ('GGA', 'NLC'):
 263 |             rho = numpy.empty((4,ngrids))
 264 |             c0 = _dot_ao_dm(mol, ao[0], cpos, non0tab, shls_slice, ao_loc)
 265 |             #:rho[0] = numpy.einsum('pi,pi->p', c0, c0)
 266 |             rho[0] = _contract_rho(c0, c0)
 267 |             for i in range(1, 4):
 268 |                 c1 = _dot_ao_dm(mol, ao[i], cpos, non0tab, shls_slice, ao_loc)
 269 |                 #:rho[i] = numpy.einsum('pi,pi->p', c0, c1) * 2 # *2 for +c.c.
 270 |                 rho[i] = _contract_rho(c0, c1) * 2
 271 |         else: # meta-GGA
 272 |             # rho[4] = \nabla^2 rho, rho[5] = 1/2 |nabla f|^2
 273 |             rho = numpy.empty((6,ngrids))
 274 |             c0 = _dot_ao_dm(mol, ao[0], cpos, non0tab, shls_slice, ao_loc)
 275 |             #:rho[0] = numpy.einsum('pi,pi->p', c0, c0)
 276 |             rho[0] = _contract_rho(c0, c0)
 277 |             rho[5] = 0
 278 |             for i in range(1, 4):
 279 |                 c1 = _dot_ao_dm(mol, ao[i], cpos, non0tab, shls_slice, ao_loc)
 280 |                 #:rho[i] = numpy.einsum('pi,pi->p', c0, c1) * 2 # *2 for +c.c.
 281 |                 #:rho[5] += numpy.einsum('pi,pi->p', c1, c1)
 282 |                 rho[i] = _contract_rho(c0, c1) * 2
 283 |                 rho[5] += _contract_rho(c1, c1)
 284 |             XX, YY, ZZ = 4, 7, 9
 285 |             ao2 = ao[XX] + ao[YY] + ao[ZZ]
 286 |             c1 = _dot_ao_dm(mol, ao2, cpos, non0tab, shls_slice, ao_loc)
 287 |             #:rho[4] = numpy.einsum('pi,pi->p', c0, c1)
 288 |             rho[4] = _contract_rho(c0, c1)
 289 |             rho[4] += rho[5]
 290 |             rho[4] *= 2
 291 | 
 292 |             rho[5] *= .5
 293 |     else:
 294 |         if xctype == 'LDA' or xctype == 'HF':
 295 |             rho = numpy.zeros(ngrids)
 296 |         elif xctype in ('GGA', 'NLC'):
 297 |             rho = numpy.zeros((4,ngrids))
 298 |         else:
 299 |             rho = numpy.zeros((6,ngrids))
 300 | 
 301 |     neg = mo_occ < -OCCDROP
 302 |     if neg.sum() > 0:
 303 |         cneg = numpy.einsum('ij,j->ij', mo_coeff[:,neg], numpy.sqrt(-mo_occ[neg]))
 304 |         if xctype == 'LDA' or xctype == 'HF':
 305 |             c0 = _dot_ao_dm(mol, ao, cneg, non0tab, shls_slice, ao_loc)
 306 |             #:rho -= numpy.einsum('pi,pi->p', c0, c0)
 307 |             rho -= _contract_rho(c0, c0)
 308 |         elif xctype == 'GGA':
 309 |             c0 = _dot_ao_dm(mol, ao[0], cneg, non0tab, shls_slice, ao_loc)
 310 |             #:rho[0] -= numpy.einsum('pi,pi->p', c0, c0)
 311 |             rho[0] -= _contract_rho(c0, c0)
 312 |             for i in range(1, 4):
 313 |                 c1 = _dot_ao_dm(mol, ao[i], cneg, non0tab, shls_slice, ao_loc)
 314 |                 #:rho[i] -= numpy.einsum('pi,pi->p', c0, c1) * 2 # *2 for +c.c.
 315 |                 rho[i] -= _contract_rho(c0, c1) * 2 # *2 for +c.c.
 316 |         else:
 317 |             c0 = _dot_ao_dm(mol, ao[0], cneg, non0tab, shls_slice, ao_loc)
 318 |             #:rho[0] -= numpy.einsum('pi,pi->p', c0, c0)
 319 |             rho[0] -= _contract_rho(c0, c0)
 320 |             rho5 = 0
 321 |             for i in range(1, 4):
 322 |                 c1 = _dot_ao_dm(mol, ao[i], cneg, non0tab, shls_slice, ao_loc)
 323 |                 #:rho[i] -= numpy.einsum('pi,pi->p', c0, c1) * 2 # *2 for +c.c.
 324 |                 #:rho5 += numpy.einsum('pi,pi->p', c1, c1)
 325 |                 rho[i] -= _contract_rho(c0, c1) * 2 # *2 for +c.c.
 326 |                 rho5 += _contract_rho(c1, c1)
 327 |             XX, YY, ZZ = 4, 7, 9
 328 |             ao2 = ao[XX] + ao[YY] + ao[ZZ]
 329 |             c1 = _dot_ao_dm(mol, ao2, cneg, non0tab, shls_slice, ao_loc)
 330 |             #:rho[4] -= numpy.einsum('pi,pi->p', c0, c1) * 2
 331 |             rho[4] -= _contract_rho(c0, c1) * 2
 332 |             rho[4] -= rho5 * 2
 333 | 
 334 |             rho[5] -= rho5 * .5
 335 |     return rho
 336 | 
 337 | def _vv10nlc(rho,coords,vvrho,vvweight,vvcoords,nlc_pars):
 338 |     thresh=1e-8
 339 | 
 340 |     #output
 341 |     exc=numpy.zeros(rho[0,:].size)
 342 |     vxc=numpy.zeros([2,rho[0,:].size])
 343 | 
 344 |     #outer grid needs threshing
 345 |     threshind=rho[0,:]>=thresh
 346 |     coords=coords[threshind]
 347 |     R=rho[0,:][threshind]
 348 |     Gx=rho[1,:][threshind]
 349 |     Gy=rho[2,:][threshind]
 350 |     Gz=rho[3,:][threshind]
 351 |     G=Gx**2.+Gy**2.+Gz**2.
 352 | 
 353 |     #threshed output
 354 |     excthresh=numpy.zeros(R.size)
 355 |     vxcthresh=numpy.zeros([2,R.size])
 356 | 
 357 |     #inner grid needs threshing
 358 |     innerthreshind=vvrho[0,:]>=thresh
 359 |     vvcoords=vvcoords[innerthreshind]
 360 |     vvweight=vvweight[innerthreshind]
 361 |     Rp=vvrho[0,:][innerthreshind]
 362 |     RpW=Rp*vvweight
 363 |     Gxp=vvrho[1,:][innerthreshind]
 364 |     Gyp=vvrho[2,:][innerthreshind]
 365 |     Gzp=vvrho[3,:][innerthreshind]
 366 |     Gp=Gxp**2.+Gyp**2.+Gzp**2.
 367 | 
 368 |     #constants and parameters
 369 |     Pi=numpy.pi
 370 |     Pi43=4.*Pi/3.
 371 |     Bvv, Cvv = nlc_pars
 372 |     Kvv=Bvv*1.5*Pi*((9.*Pi)**(-1./6.))
 373 |     Beta=((3./(Bvv*Bvv))**(0.75))/32.
 374 | 
 375 |     #inner grid
 376 |     W0p=Gp/(Rp*Rp)
 377 |     W0p=Cvv*W0p*W0p
 378 |     W0p=(W0p+Pi43*Rp)**0.5
 379 |     Kp=Kvv*(Rp**(1./6.))
 380 | 
 381 |     #outer grid
 382 |     W0tmp=G/(R**2)
 383 |     W0tmp=Cvv*W0tmp*W0tmp
 384 |     W0=(W0tmp+Pi43*R)**0.5
 385 |     dW0dR=(0.5*Pi43*R-2.*W0tmp)/W0
 386 |     dW0dG=W0tmp*R/(G*W0)
 387 |     K=Kvv*(R**(1./6.))
 388 |     dKdR=(1./6.)*K
 389 | 
 390 |     for i in range(R.size):
 391 |         DX=vvcoords[:,0]-coords[i,0]
 392 |         DY=vvcoords[:,1]-coords[i,1]
 393 |         DZ=vvcoords[:,2]-coords[i,2]
 394 |         R2=DX*DX+DY*DY+DZ*DZ
 395 |         gp=R2*W0p+Kp
 396 |         g=R2*W0[i]+K[i]
 397 |         gt=g+gp
 398 |         T=RpW/(g*gp*gt)
 399 |         F=numpy.sum(T)
 400 |         T*=(1./g+1./gt)
 401 |         U=numpy.sum(T)
 402 |         W=numpy.sum(T*R2)
 403 |         F*=-1.5
 404 |         #excthresh is multiplied by Rho later
 405 |         excthresh[i]=Beta+0.5*F
 406 |         vxcthresh[0,i]=Beta+F+1.5*(U*dKdR[i]+W*dW0dR[i])
 407 |         vxcthresh[1,i]=1.5*W*dW0dG[i]
 408 |     exc[threshind]=excthresh
 409 |     vxc[0,threshind]=vxcthresh[0,:]
 410 |     vxc[1,threshind]=vxcthresh[1,:]
 411 | 
 412 |     return exc,vxc
 413 | 
 414 | def eval_mat(mol, ao, weight, rho, vxc,
 415 |              non0tab=None, xctype='LDA', spin=0, verbose=None):
 416 |     r'''Calculate XC potential matrix.
 417 | 
 418 |     Args:
 419 |         mol : an instance of :class:`Mole`
 420 | 
 421 |         ao : ([4/10,] ngrids, nao) ndarray
 422 |             2D array of shape (N,nao) for LDA,
 423 |             3D array of shape (4,N,nao) for GGA
 424 |             or (10,N,nao) for meta-GGA.
 425 |             N is the number of grids, nao is the number of AO functions.
 426 |             If xctype is GGA, ao[0] is AO value and ao[1:3] are the real space
 427 |             gradients.  If xctype is meta-GGA, ao[4:10] are second derivatives
 428 |             of ao values.
 429 |         weight : 1D array
 430 |             Integral weights on grids.
 431 |         rho : ([4/6,] ngrids) ndarray
 432 |             Shape of ((*,N)) for electron density (and derivatives) if spin = 0;
 433 |             Shape of ((*,N),(*,N)) for alpha/beta electron density (and derivatives) if spin > 0;
 434 |             where N is number of grids.
 435 |             rho (*,N) are ordered as (den,grad_x,grad_y,grad_z,laplacian,tau)
 436 |             where grad_x = d/dx den, laplacian = \nabla^2 den, tau = 1/2(\nabla f)^2
 437 |             In spin unrestricted case,
 438 |             rho is ((den_u,grad_xu,grad_yu,grad_zu,laplacian_u,tau_u)
 439 |                     (den_d,grad_xd,grad_yd,grad_zd,laplacian_d,tau_d))
 440 |         vxc : ([4,] ngrids) ndarray
 441 |             XC potential value on each grid = (vrho, vsigma, vlapl, vtau)
 442 |             vsigma is GGA potential value on each grid.
 443 |             If the kwarg spin != 0, a list [vsigma_uu,vsigma_ud] is required.
 444 | 
 445 |     Kwargs:
 446 |         xctype : str
 447 |             LDA/GGA/mGGA.  It affects the shape of `ao` and `rho`
 448 |         non0tab : 2D bool array
 449 |             mask array to indicate whether the AO values are zero.  The mask
 450 |             array can be obtained by calling :func:`make_mask`
 451 |         spin : int
 452 |             If not 0, the returned matrix is the Vxc matrix of alpha-spin.  It
 453 |             is computed with the spin non-degenerated UKS formula.
 454 | 
 455 |     Returns:
 456 |         XC potential matrix in 2D array of shape (nao,nao) where nao is the
 457 |         number of AO functions.
 458 |     '''
 459 |     xctype = xctype.upper()
 460 |     if xctype == 'LDA' or xctype == 'HF':
 461 |         ngrids, nao = ao.shape
 462 |     else:
 463 |         ngrids, nao = ao[0].shape
 464 | 
 465 |     if non0tab is None:
 466 |         non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
 467 |                              dtype=numpy.uint8)
 468 |     shls_slice = (0, mol.nbas)
 469 |     ao_loc = mol.ao_loc_nr()
 470 |     transpose_for_uks = False
 471 |     if xctype == 'LDA' or xctype == 'HF':
 472 |         if not isinstance(vxc, numpy.ndarray) or vxc.ndim == 2:
 473 |             vrho = vxc[0]
 474 |         else:
 475 |             vrho = vxc
 476 |         # *.5 because return mat + mat.T
 477 |         #:aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho)
 478 |         aow = _scale_ao(ao, .5*weight*vrho)
 479 |         mat = _dot_ao_ao(mol, ao, aow, non0tab, shls_slice, ao_loc)
 480 |     else:
 481 |         #wv = weight * vsigma * 2
 482 |         #aow  = numpy.einsum('pi,p->pi', ao[1], rho[1]*wv)
 483 |         #aow += numpy.einsum('pi,p->pi', ao[2], rho[2]*wv)
 484 |         #aow += numpy.einsum('pi,p->pi', ao[3], rho[3]*wv)
 485 |         #aow += numpy.einsum('pi,p->pi', ao[0], .5*weight*vrho)
 486 |         vrho, vsigma = vxc[:2]
 487 |         wv = numpy.empty((4,ngrids))
 488 |         if spin == 0:
 489 |             assert(vsigma is not None and rho.ndim==2)
 490 |             wv[0]  = weight * vrho * .5
 491 |             wv[1:4] = rho[1:4] * (weight * vsigma * 2)
 492 |         else:
 493 |             rho_a, rho_b = rho
 494 |             wv[0]  = weight * vrho * .5
 495 |             try:
 496 |                 wv[1:4] = rho_a[1:4] * (weight * vsigma[0] * 2)  # sigma_uu
 497 |                 wv[1:4]+= rho_b[1:4] * (weight * vsigma[1])      # sigma_ud
 498 |             except ValueError:
 499 |                 warnings.warn('Note the output of libxc.eval_xc cannot be '
 500 |                               'directly used in eval_mat.\nvsigma from eval_xc '
 501 |                               'should be restructured as '
 502 |                               '(vsigma[:,0],vsigma[:,1])\n')
 503 |                 transpose_for_uks = True
 504 |                 vsigma = vsigma.T
 505 |                 wv[1:4] = rho_a[1:4] * (weight * vsigma[0] * 2)  # sigma_uu
 506 |                 wv[1:4]+= rho_b[1:4] * (weight * vsigma[1])      # sigma_ud
 507 |         #:aow = numpy.einsum('npi,np->pi', ao[:4], wv)
 508 |         aow = _scale_ao(ao[:4], wv)
 509 |         mat = _dot_ao_ao(mol, ao[0], aow, non0tab, shls_slice, ao_loc)
 510 | 
 511 | # JCP, 138, 244108
 512 | # JCP, 112, 7002
 513 |     if xctype == 'MGGA':
 514 |         vlapl, vtau = vxc[2:]
 515 | 
 516 |         if vlapl is None:
 517 |             vlapl = 0
 518 |         else:
 519 |             if spin != 0:
 520 |                 if transpose_for_uks:
 521 |                     vlapl = vlapl.T
 522 |                 vlapl = vlapl[0]
 523 |             XX, YY, ZZ = 4, 7, 9
 524 |             ao2 = ao[XX] + ao[YY] + ao[ZZ]
 525 |             #:aow = numpy.einsum('pi,p->pi', ao2, .5 * weight * vlapl, out=aow)
 526 |             aow = _scale_ao(ao2, .5 * weight * vlapl, out=aow)
 527 |             mat += _dot_ao_ao(mol, ao[0], aow, non0tab, shls_slice, ao_loc)
 528 | 
 529 |         if spin != 0:
 530 |             if transpose_for_uks:
 531 |                 vtau = vtau.T
 532 |             vtau = vtau[0]
 533 |         wv = weight * (.25*vtau + vlapl)
 534 |         #:aow = numpy.einsum('pi,p->pi', ao[1], wv, out=aow)
 535 |         aow = _scale_ao(ao[1], wv, out=aow)
 536 |         mat += _dot_ao_ao(mol, ao[1], aow, non0tab, shls_slice, ao_loc)
 537 |         #:aow = numpy.einsum('pi,p->pi', ao[2], wv, out=aow)
 538 |         aow = _scale_ao(ao[2], wv, out=aow)
 539 |         mat += _dot_ao_ao(mol, ao[2], aow, non0tab, shls_slice, ao_loc)
 540 |         #:aow = numpy.einsum('pi,p->pi', ao[3], wv, out=aow)
 541 |         aow = _scale_ao(ao[3], wv, out=aow)
 542 |         mat += _dot_ao_ao(mol, ao[3], aow, non0tab, shls_slice, ao_loc)
 543 | 
 544 |     return mat + mat.T.conj()
 545 | 
 546 | 
 547 | def _dot_ao_ao(mol, ao1, ao2, non0tab, shls_slice, ao_loc, hermi=0):
 548 |     '''return numpy.dot(ao1.T, ao2)'''
 549 |     ngrids, nao = ao1.shape
 550 |     if nao < SWITCH_SIZE:
 551 |         return lib.dot(ao1.T.conj(), ao2)
 552 | 
 553 |     if not ao1.flags.f_contiguous:
 554 |         ao1 = lib.transpose(ao1)
 555 |     if not ao2.flags.f_contiguous:
 556 |         ao2 = lib.transpose(ao2)
 557 |     if ao1.dtype == ao2.dtype == numpy.double:
 558 |         fn = libdft.VXCdot_ao_ao
 559 |     else:
 560 |         fn = libdft.VXCzdot_ao_ao
 561 |         ao1 = numpy.asarray(ao1, numpy.complex128)
 562 |         ao2 = numpy.asarray(ao2, numpy.complex128)
 563 | 
 564 |     if non0tab is None or shls_slice is None or ao_loc is None:
 565 |         pnon0tab = pshls_slice = pao_loc = lib.c_null_ptr()
 566 |     else:
 567 |         pnon0tab    = non0tab.ctypes.data_as(ctypes.c_void_p)
 568 |         pshls_slice = (ctypes.c_int*2)(*shls_slice)
 569 |         pao_loc     = ao_loc.ctypes.data_as(ctypes.c_void_p)
 570 | 
 571 |     vv = numpy.empty((nao,nao), dtype=ao1.dtype)
 572 |     fn(vv.ctypes.data_as(ctypes.c_void_p),
 573 |        ao1.ctypes.data_as(ctypes.c_void_p),
 574 |        ao2.ctypes.data_as(ctypes.c_void_p),
 575 |        ctypes.c_int(nao), ctypes.c_int(ngrids),
 576 |        ctypes.c_int(mol.nbas), ctypes.c_int(hermi),
 577 |        pnon0tab, pshls_slice, pao_loc)
 578 |     return vv
 579 | 
 580 | def _dot_ao_dm(mol, ao, dm, non0tab, shls_slice, ao_loc, out=None):
 581 |     '''return numpy.dot(ao, dm)'''
 582 |     ngrids, nao = ao.shape
 583 |     if nao < SWITCH_SIZE:
 584 |         return lib.dot(dm.T, ao.T).T
 585 | 
 586 |     if not ao.flags.f_contiguous:
 587 |         ao = lib.transpose(ao)
 588 |     if ao.dtype == dm.dtype == numpy.double:
 589 |         fn = libdft.VXCdot_ao_dm
 590 |     else:
 591 |         fn = libdft.VXCzdot_ao_dm
 592 |         ao = numpy.asarray(ao, numpy.complex128)
 593 |         dm = numpy.asarray(dm, numpy.complex128)
 594 | 
 595 |     if non0tab is None or shls_slice is None or ao_loc is None:
 596 |         pnon0tab = pshls_slice = pao_loc = lib.c_null_ptr()
 597 |     else:
 598 |         pnon0tab    = non0tab.ctypes.data_as(ctypes.c_void_p)
 599 |         pshls_slice = (ctypes.c_int*2)(*shls_slice)
 600 |         pao_loc     = ao_loc.ctypes.data_as(ctypes.c_void_p)
 601 | 
 602 |     vm = numpy.ndarray((ngrids,dm.shape[1]), dtype=ao.dtype, order='F', buffer=out)
 603 |     dm = numpy.asarray(dm, order='C')
 604 |     fn(vm.ctypes.data_as(ctypes.c_void_p),
 605 |        ao.ctypes.data_as(ctypes.c_void_p),
 606 |        dm.ctypes.data_as(ctypes.c_void_p),
 607 |        ctypes.c_int(nao), ctypes.c_int(dm.shape[1]),
 608 |        ctypes.c_int(ngrids), ctypes.c_int(mol.nbas),
 609 |        pnon0tab, pshls_slice, pao_loc)
 610 |     return vm
 611 | 
 612 | def _scale_ao(ao, wv, out=None):
 613 |     #:aow = numpy.einsum('npi,np->pi', ao[:4], wv)
 614 |     if wv.ndim == 2:
 615 |         ao = ao.transpose(0,2,1)
 616 |     else:
 617 |         ngrids, nao = ao.shape
 618 |         ao = ao.T.reshape(1,nao,ngrids)
 619 |         wv = wv.reshape(1,ngrids)
 620 | 
 621 |     wv = numpy.asarray(wv, order='C')
 622 |     comp, nao, ngrids = ao.shape
 623 |     aow = numpy.ndarray((nao,ngrids), dtype=ao.dtype, buffer=out).T
 624 | 
 625 |     if not ao.flags.c_contiguous:
 626 |         aow = numpy.einsum('nip,np->pi', ao, wv)
 627 |     elif aow.dtype == numpy.double:
 628 |         libdft.VXC_dscale_ao(aow.ctypes.data_as(ctypes.c_void_p),
 629 |                              ao.ctypes.data_as(ctypes.c_void_p),
 630 |                              wv.ctypes.data_as(ctypes.c_void_p),
 631 |                              ctypes.c_int(comp), ctypes.c_int(nao),
 632 |                              ctypes.c_int(ngrids))
 633 |     elif aow.dtype == numpy.complex128:
 634 |         libdft.VXC_zscale_ao(aow.ctypes.data_as(ctypes.c_void_p),
 635 |                              ao.ctypes.data_as(ctypes.c_void_p),
 636 |                              wv.ctypes.data_as(ctypes.c_void_p),
 637 |                              ctypes.c_int(comp), ctypes.c_int(nao),
 638 |                              ctypes.c_int(ngrids))
 639 |     else:
 640 |         aow = numpy.einsum('nip,np->pi', ao, wv)
 641 |     return aow
 642 | 
 643 | def _contract_rho(bra, ket):
 644 |     #:rho  = numpy.einsum('pi,pi->p', bra.real, ket.real)
 645 |     #:rho += numpy.einsum('pi,pi->p', bra.imag, ket.imag)
 646 |     bra = bra.T
 647 |     ket = ket.T
 648 |     nao, ngrids = bra.shape
 649 |     rho = numpy.empty(ngrids)
 650 | 
 651 |     if not (bra.flags.c_contiguous and ket.flags.c_contiguous):
 652 |         rho  = numpy.einsum('ip,ip->p', bra.real, ket.real)
 653 |         rho += numpy.einsum('ip,ip->p', bra.imag, ket.imag)
 654 |     elif bra.dtype == numpy.double and ket.dtype == numpy.double:
 655 |         libdft.VXC_dcontract_rho(rho.ctypes.data_as(ctypes.c_void_p),
 656 |                                  bra.ctypes.data_as(ctypes.c_void_p),
 657 |                                  ket.ctypes.data_as(ctypes.c_void_p),
 658 |                                  ctypes.c_int(nao), ctypes.c_int(ngrids))
 659 |     elif bra.dtype == numpy.complex128 and ket.dtype == numpy.complex128:
 660 |         libdft.VXC_zcontract_rho(rho.ctypes.data_as(ctypes.c_void_p),
 661 |                                  bra.ctypes.data_as(ctypes.c_void_p),
 662 |                                  ket.ctypes.data_as(ctypes.c_void_p),
 663 |                                  ctypes.c_int(nao), ctypes.c_int(ngrids))
 664 |     else:
 665 |         rho  = numpy.einsum('ip,ip->p', bra.real, ket.real)
 666 |         rho += numpy.einsum('ip,ip->p', bra.imag, ket.imag)
 667 |     return rho
 668 | 
 669 | def nr_vxc(mol, grids, xc_code, dms, spin=0, relativity=0, hermi=0,
 670 |            max_memory=2000, verbose=None):
 671 |     '''
 672 |     Evaluate RKS/UKS XC functional and potential matrix on given meshgrids
 673 |     for a set of density matrices.  See :func:`nr_rks` and :func:`nr_uks`
 674 |     for more details.
 675 | 
 676 |     Args:
 677 |         mol : an instance of :class:`Mole`
 678 | 
 679 |         grids : an instance of :class:`Grids`
 680 |             grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
 681 |         xc_code : str
 682 |             XC functional description.
 683 |             See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
 684 |         dms : 2D array or a list of 2D arrays
 685 |             Density matrix or multiple density matrices
 686 | 
 687 |     Kwargs:
 688 |         hermi : int
 689 |             Input density matrices symmetric or not
 690 |         max_memory : int or float
 691 |             The maximum size of cache to use (in MB).
 692 | 
 693 |     Returns:
 694 |         nelec, excsum, vmat.
 695 |         nelec is the number of electrons generated by numerical integration.
 696 |         excsum is the XC functional value.  vmat is the XC potential matrix in
 697 |         2D array of shape (nao,nao) where nao is the number of AO functions.
 698 | 
 699 |     Examples:
 700 | 
 701 |     >>> from pyscf import gto, dft
 702 |     >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')
 703 |     >>> grids = dft.gen_grid.Grids(mol)
 704 |     >>> grids.coords = numpy.random.random((100,3))  # 100 random points
 705 |     >>> grids.weights = numpy.random.random(100)
 706 |     >>> nao = mol.nao_nr()
 707 |     >>> dm = numpy.random.random((2,nao,nao))
 708 |     >>> nelec, exc, vxc = dft.numint.nr_vxc(mol, grids, 'lda,vwn', dm, spin=1)
 709 |     '''
 710 |     ni = NumInt()
 711 |     return ni.nr_vxc(mol, grids, xc_code, dms, spin, relativity,
 712 |                      hermi, max_memory, verbose)
 713 | 
 714 | def nr_rks(ni, mol, grids, xc_code, dms, relativity=0, hermi=0,
 715 |            max_memory=2000, verbose=None):
 716 |     '''Calculate RKS XC functional and potential matrix on given meshgrids
 717 |     for a set of density matrices
 718 | 
 719 |     Args:
 720 |         ni : an instance of :class:`NumInt`
 721 | 
 722 |         mol : an instance of :class:`Mole`
 723 | 
 724 |         grids : an instance of :class:`Grids`
 725 |             grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
 726 |         xc_code : str
 727 |             XC functional description.
 728 |             See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
 729 |         dms : 2D array a list of 2D arrays
 730 |             Density matrix or multiple density matrices
 731 | 
 732 |     Kwargs:
 733 |         hermi : int
 734 |             Input density matrices symmetric or not
 735 |         max_memory : int or float
 736 |             The maximum size of cache to use (in MB).
 737 | 
 738 |     Returns:
 739 |         nelec, excsum, vmat.
 740 |         nelec is the number of electrons generated by numerical integration.
 741 |         excsum is the XC functional value.  vmat is the XC potential matrix in
 742 |         2D array of shape (nao,nao) where nao is the number of AO functions.
 743 | 
 744 |     Examples:
 745 | 
 746 |     >>> from pyscf import gto, dft
 747 |     >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')
 748 |     >>> grids = dft.gen_grid.Grids(mol)
 749 |     >>> grids.coords = numpy.random.random((100,3))  # 100 random points
 750 |     >>> grids.weights = numpy.random.random(100)
 751 |     >>> nao = mol.nao_nr()
 752 |     >>> dm = numpy.random.random((nao,nao))
 753 |     >>> ni = dft.numint.NumInt()
 754 |     >>> nelec, exc, vxc = ni.nr_rks(mol, grids, 'lda,vwn', dm)
 755 |     '''
 756 |     xctype = ni._xc_type(xc_code)
 757 |     make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
 758 | 
 759 |     shls_slice = (0, mol.nbas)
 760 |     ao_loc = mol.ao_loc_nr()
 761 | 
 762 |     nelec = numpy.zeros(nset)
 763 |     excsum = numpy.zeros(nset)
 764 |     vmat = numpy.zeros((nset,nao,nao))
 765 |     aow = None
 766 |     if xctype == 'LDA':
 767 |         ao_deriv = 0
 768 |         for ao, mask, weight, coords \
 769 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 770 |             aow = numpy.ndarray(ao.shape, order='F', buffer=aow)
 771 |             for idm in range(nset):
 772 |                 rho = make_rho(idm, ao, mask, 'LDA')
 773 |                 exc, vxc = ni.eval_xc(xc_code, rho, 0, relativity, 1, verbose)[:2]
 774 |                 vrho = vxc[0]
 775 |                 den = rho * weight
 776 |                 nelec[idm] += den.sum()
 777 |                 excsum[idm] += numpy.dot(den, exc)
 778 |                 # *.5 because vmat + vmat.T
 779 |                 #:aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho, out=aow)
 780 |                 aow = _scale_ao(ao, .5*weight*vrho, out=aow)
 781 |                 vmat[idm] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc)
 782 |                 rho = exc = vxc = vrho = None
 783 |     elif xctype == 'GGA':
 784 |         ao_deriv = 1
 785 |         for ao, mask, weight, coords \
 786 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 787 |             ngrid = weight.size
 788 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
 789 |             for idm in range(nset):
 790 |                 rho = make_rho(idm, ao, mask, 'GGA')
 791 |                 exc, vxc = ni.eval_xc(xc_code, rho, 0, relativity, 1, verbose)[:2]
 792 |                 den = rho[0] * weight
 793 |                 nelec[idm] += den.sum()
 794 |                 excsum[idm] += numpy.dot(den, exc)
 795 | # ref eval_mat function
 796 |                 wv = _rks_gga_wv0(rho, vxc, weight)
 797 |                 #:aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
 798 |                 aow = _scale_ao(ao, wv, out=aow)
 799 |                 vmat[idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 800 |                 rho = exc = vxc = wv = None
 801 |     elif xctype == 'NLC':
 802 |         nlc_pars = ni.nlc_coeff(xc_code[:-6])
 803 |         if nlc_pars == [0,0]:
 804 |             raise NotImplementedError('VV10 cannot be used with %s. '
 805 |                                       'The supported functionals are %s' %
 806 |                                       (xc_code[:-6], ni.libxc.VV10_XC))
 807 |         ao_deriv = 1
 808 |         vvrho=numpy.empty([nset,4,0])
 809 |         vvweight=numpy.empty([nset,0])
 810 |         vvcoords=numpy.empty([nset,0,3])
 811 |         for ao, mask, weight, coords \
 812 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 813 |             rhotmp = numpy.empty([0,4,weight.size])
 814 |             weighttmp = numpy.empty([0,weight.size])
 815 |             coordstmp = numpy.empty([0,weight.size,3])
 816 |             for idm in range(nset):
 817 |                 rho = make_rho(idm, ao, mask, 'GGA')
 818 |                 rho = numpy.expand_dims(rho,axis=0)
 819 |                 rhotmp = numpy.concatenate((rhotmp,rho),axis=0)
 820 |                 weighttmp = numpy.concatenate((weighttmp,numpy.expand_dims(weight,axis=0)),axis=0)
 821 |                 coordstmp = numpy.concatenate((coordstmp,numpy.expand_dims(coords,axis=0)),axis=0)
 822 |                 rho = None
 823 |             vvrho=numpy.concatenate((vvrho,rhotmp),axis=2)
 824 |             vvweight=numpy.concatenate((vvweight,weighttmp),axis=1)
 825 |             vvcoords=numpy.concatenate((vvcoords,coordstmp),axis=1)
 826 |             rhotmp = weighttmp = coordstmp = None
 827 |         for ao, mask, weight, coords \
 828 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 829 |             ngrid = weight.size
 830 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
 831 |             for idm in range(nset):
 832 |                 rho = make_rho(idm, ao, mask, 'GGA')
 833 |                 exc, vxc = _vv10nlc(rho,coords,vvrho[idm],vvweight[idm],vvcoords[idm],nlc_pars)
 834 |                 den = rho[0] * weight
 835 |                 nelec[idm] += den.sum()
 836 |                 excsum[idm] += numpy.dot(den, exc)
 837 | # ref eval_mat function
 838 |                 wv = _rks_gga_wv0(rho, vxc, weight)
 839 |                 #:aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
 840 |                 aow = _scale_ao(ao, wv, out=aow)
 841 |                 vmat[idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 842 |                 rho = exc = vxc = wv = None
 843 |         vvrho = vvweight = vvcoords = None
 844 |     elif xctype == 'MGGA':
 845 |         if (any(x in xc_code.upper() for x in ('CC06', 'CS', 'BR89', 'MK00'))):
 846 |             raise NotImplementedError('laplacian in meta-GGA method')
 847 |         ao_deriv = 2
 848 |         for ao, mask, weight, coords \
 849 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 850 |             ngrid = weight.size
 851 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
 852 |             for idm in range(nset):
 853 |                 rho = make_rho(idm, ao, mask, 'MGGA')
 854 |                 exc, vxc = ni.eval_xc(xc_code, rho, 0, relativity, 1, verbose)[:2]
 855 |                 vrho, vsigma, vlapl, vtau = vxc[:4]
 856 |                 den = rho[0] * weight
 857 |                 nelec[idm] += den.sum()
 858 |                 excsum[idm] += numpy.dot(den, exc)
 859 | 
 860 |                 wv = _rks_gga_wv0(rho, vxc, weight)
 861 |                 #:aow = numpy.einsum('npi,np->pi', ao[:4], wv, out=aow)
 862 |                 aow = _scale_ao(ao[:4], wv, out=aow)
 863 |                 vmat[idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 864 | 
 865 | # FIXME: .5 * .5   First 0.5 for v+v.T symmetrization.
 866 | # Second 0.5 is due to the Libxc convention tau = 1/2 \nabla\phi\dot\nabla\phi
 867 |                 wv = (.5 * .5 * weight * vtau).reshape(-1,1)
 868 |                 vmat[idm] += _dot_ao_ao(mol, ao[1], wv*ao[1], mask, shls_slice, ao_loc)
 869 |                 vmat[idm] += _dot_ao_ao(mol, ao[2], wv*ao[2], mask, shls_slice, ao_loc)
 870 |                 vmat[idm] += _dot_ao_ao(mol, ao[3], wv*ao[3], mask, shls_slice, ao_loc)
 871 | 
 872 |                 rho = exc = vxc = vrho = vsigma = wv = None
 873 | 
 874 |     elif xctype == 'MGGA+':
 875 |         if (any(x in xc_code.upper() for x in ('CC06', 'CS', 'BR89', 'MK00'))):
 876 |             raise NotImplementedError('laplacian in meta-GGA method')
 877 |         ao_deriv = 2
 878 |         for ao, mask, weight, coords \
 879 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 880 |             ngrid = weight.size
 881 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
 882 |             for idm in range(nset):
 883 |                 rho = make_rho(idm, ao, mask, 'MGGA')
 884 |                 exc, vxc = ni.eval_xc(xc_code,rho,0,weight,coords, relativity, 1, verbose)[:2]
 885 |                 vrho, vsigma, vlapl, vtau = vxc[:4]
 886 |                 den = rho[0] * weight
 887 |                 nelec[idm] += den.sum()
 888 |                 excsum[idm] += numpy.dot(den, exc)
 889 | 
 890 |                 wv = _rks_gga_wv0(rho, vxc, weight)
 891 |                 aow = numpy.einsum('npi,np->pi', ao[:4], wv, out=aow)
 892 |                 vmat[idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 893 | 
 894 | # FIXME: .5 * .5   First 0.5 for v+v.T symmetrization.
 895 | # Second 0.5 is due to the Libxc convention tau = 1/2 \nabla\phi\dot\nabla\phi
 896 |                 wv = (.5 * .5 * weight * vtau).reshape(-1,1)
 897 |                 vmat[idm] += _dot_ao_ao(mol, ao[1], wv*ao[1], mask, shls_slice, ao_loc)
 898 |                 vmat[idm] += _dot_ao_ao(mol, ao[2], wv*ao[2], mask, shls_slice, ao_loc)
 899 |                 vmat[idm] += _dot_ao_ao(mol, ao[3], wv*ao[3], mask, shls_slice, ao_loc)
 900 | 
 901 |                 rho = exc = vxc = vrho = vsigma = wv = None
 902 | 
 903 |     for i in range(nset):
 904 |         vmat[i] = vmat[i] + vmat[i].T
 905 |     if nset == 1:
 906 |         nelec = nelec[0]
 907 |         excsum = excsum[0]
 908 |         vmat = vmat.reshape(nao,nao)
 909 |     return nelec, excsum, vmat
 910 | 
 911 | def nr_uks(ni, mol, grids, xc_code, dms, relativity=0, hermi=0,
 912 |            max_memory=2000, verbose=None):
 913 |     '''Calculate UKS XC functional and potential matrix on given meshgrids
 914 |     for a set of density matrices
 915 | 
 916 |     Args:
 917 |         mol : an instance of :class:`Mole`
 918 | 
 919 |         grids : an instance of :class:`Grids`
 920 |             grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
 921 |         xc_code : str
 922 |             XC functional description.
 923 |             See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
 924 |         dms : a list of 2D arrays
 925 |             A list of density matrices, stored as (alpha,alpha,...,beta,beta,...)
 926 | 
 927 |     Kwargs:
 928 |         hermi : int
 929 |             Input density matrices symmetric or not
 930 |         max_memory : int or float
 931 |             The maximum size of cache to use (in MB).
 932 | 
 933 |     Returns:
 934 |         nelec, excsum, vmat.
 935 |         nelec is the number of (alpha,beta) electrons generated by numerical integration.
 936 |         excsum is the XC functional value.
 937 |         vmat is the XC potential matrix for (alpha,beta) spin.
 938 | 
 939 |     Examples:
 940 | 
 941 |     >>> from pyscf import gto, dft
 942 |     >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')
 943 |     >>> grids = dft.gen_grid.Grids(mol)
 944 |     >>> grids.coords = numpy.random.random((100,3))  # 100 random points
 945 |     >>> grids.weights = numpy.random.random(100)
 946 |     >>> nao = mol.nao_nr()
 947 |     >>> dm = numpy.random.random((2,nao,nao))
 948 |     >>> ni = dft.numint.NumInt()
 949 |     >>> nelec, exc, vxc = ni.nr_uks(mol, grids, 'lda,vwn', dm)
 950 |     '''
 951 |     xctype = ni._xc_type(xc_code)
 952 |     if xctype == 'NLC':
 953 |         dms_sf = dms[0] + dms[1]
 954 |         nelec, excsum, vmat = nr_rks(ni, mol, grids, xc_code, dms_sf, relativity, hermi,
 955 |                                      max_memory, verbose)
 956 |         return [nelec,nelec], excsum, numpy.asarray([vmat,vmat])
 957 | 
 958 |     shls_slice = (0, mol.nbas)
 959 |     ao_loc = mol.ao_loc_nr()
 960 | 
 961 |     dma, dmb = _format_uks_dm(dms)
 962 |     nao = dma.shape[-1]
 963 |     make_rhoa, nset = ni._gen_rho_evaluator(mol, dma, hermi)[:2]
 964 |     make_rhob       = ni._gen_rho_evaluator(mol, dmb, hermi)[0]
 965 | 
 966 |     nelec = numpy.zeros((2,nset))
 967 |     excsum = numpy.zeros(nset)
 968 |     vmat = numpy.zeros((2,nset,nao,nao))
 969 |     aow = None
 970 |     if xctype == 'LDA':
 971 |         ao_deriv = 0
 972 |         for ao, mask, weight, coords \
 973 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 974 |             aow = numpy.ndarray(ao.shape, order='F', buffer=aow)
 975 |             for idm in range(nset):
 976 |                 rho_a = make_rhoa(idm, ao, mask, xctype)
 977 |                 rho_b = make_rhob(idm, ao, mask, xctype)
 978 |                 exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
 979 |                                       1, relativity, 1, verbose)[:2]
 980 |                 vrho = vxc[0]
 981 |                 den = rho_a * weight
 982 |                 nelec[0,idm] += den.sum()
 983 |                 excsum[idm] += numpy.dot(den, exc)
 984 |                 den = rho_b * weight
 985 |                 nelec[1,idm] += den.sum()
 986 |                 excsum[idm] += numpy.dot(den, exc)
 987 | 
 988 |                 # *.5 due to +c.c. in the end
 989 |                 #:aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho[:,0], out=aow)
 990 |                 aow = _scale_ao(ao, .5*weight*vrho[:,0], out=aow)
 991 |                 vmat[0,idm] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc)
 992 |                 #:aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho[:,1], out=aow)
 993 |                 aow = _scale_ao(ao, .5*weight*vrho[:,1], out=aow)
 994 |                 vmat[1,idm] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc)
 995 |                 rho_a = rho_b = exc = vxc = vrho = None
 996 |     elif xctype == 'GGA':
 997 |         ao_deriv = 1
 998 |         for ao, mask, weight, coords \
 999 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1000 |             ngrid = weight.size
1001 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
1002 |             for idm in range(nset):
1003 |                 rho_a = make_rhoa(idm, ao, mask, xctype)
1004 |                 rho_b = make_rhob(idm, ao, mask, xctype)
1005 |                 exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
1006 |                                       1, relativity, 1, verbose)[:2]
1007 |                 den = rho_a[0]*weight
1008 |                 nelec[0,idm] += den.sum()
1009 |                 excsum[idm] += numpy.dot(den, exc)
1010 |                 den = rho_b[0]*weight
1011 |                 nelec[1,idm] += den.sum()
1012 |                 excsum[idm] += numpy.dot(den, exc)
1013 | 
1014 |                 wva, wvb = _uks_gga_wv0((rho_a,rho_b), vxc, weight)
1015 |                 #:aow = numpy.einsum('npi,np->pi', ao, wva, out=aow)
1016 |                 aow = _scale_ao(ao, wva, out=aow)
1017 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
1018 |                 #:aow = numpy.einsum('npi,np->pi', ao, wvb, out=aow)
1019 |                 aow = _scale_ao(ao, wvb, out=aow)
1020 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
1021 |                 rho_a = rho_b = exc = vxc = wva = wvb = None
1022 |     elif xctype == 'MGGA':
1023 |         if (any(x in xc_code.upper() for x in ('CC06', 'CS', 'BR89', 'MK00'))):
1024 |             raise NotImplementedError('laplacian in meta-GGA method')
1025 |         ao_deriv = 2
1026 |         for ao, mask, weight, coords \
1027 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1028 |             ngrid = weight.size
1029 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
1030 |             for idm in range(nset):
1031 |                 rho_a = make_rhoa(idm, ao, mask, xctype)
1032 |                 rho_b = make_rhob(idm, ao, mask, xctype)
1033 |                 exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
1034 |                                       1, relativity, 1, verbose)[:2]
1035 |                 vrho, vsigma, vlapl, vtau = vxc[:4]
1036 |                 den = rho_a[0]*weight
1037 |                 nelec[0,idm] += den.sum()
1038 |                 excsum[idm] += numpy.dot(den, exc)
1039 |                 den = rho_b[0]*weight
1040 |                 nelec[1,idm] += den.sum()
1041 |                 excsum[idm] += numpy.dot(den, exc)
1042 | 
1043 |                 wva, wvb = _uks_gga_wv0((rho_a,rho_b), vxc, weight)
1044 |                 #:aow = numpy.einsum('npi,np->pi', ao[:4], wva, out=aow)
1045 |                 aow = _scale_ao(ao[:4], wva, out=aow)
1046 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
1047 |                 #:aow = numpy.einsum('npi,np->pi', ao[:4], wvb, out=aow)
1048 |                 aow = _scale_ao(ao[:4], wvb, out=aow)
1049 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
1050 | 
1051 | # FIXME: .5 * .5   First 0.5 for v+v.T symmetrization.
1052 | # Second 0.5 is due to the Libxc convention tau = 1/2 \nabla\phi\dot\nabla\phi
1053 |                 wv = (.25 * weight * vtau[:,0]).reshape(-1,1)
1054 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[1], wv*ao[1], mask, shls_slice, ao_loc)
1055 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[2], wv*ao[2], mask, shls_slice, ao_loc)
1056 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[3], wv*ao[3], mask, shls_slice, ao_loc)
1057 |                 wv = (.25 * weight * vtau[:,1]).reshape(-1,1)
1058 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[1], wv*ao[1], mask, shls_slice, ao_loc)
1059 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[2], wv*ao[2], mask, shls_slice, ao_loc)
1060 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[3], wv*ao[3], mask, shls_slice, ao_loc)
1061 |                 rho_a = rho_b = exc = vxc = vrho = vsigma = wva = wvb = None
1062 |     elif xctype == 'MGGA+':
1063 |         if (any(x in xc_code.upper() for x in ('CC06', 'CS', 'BR89', 'MK00'))):
1064 |             raise NotImplementedError('laplacian in meta-GGA method')
1065 |         ao_deriv = 2
1066 |         for ao, mask, weight, coords \
1067 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1068 |             ngrid = weight.size
1069 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
1070 |             for idm in range(nset):
1071 |                 rho_a = make_rhoa(idm, ao, mask, xctype)
1072 |                 rho_b = make_rhob(idm, ao, mask, xctype)
1073 |                 exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
1074 |                                       1,weight,coords, relativity, 1, verbose)[:2]
1075 |                 vrho, vsigma, vlapl, vtau = vxc[:4]
1076 |                 den = rho_a[0]*weight
1077 |                 nelec[0,idm] += den.sum()
1078 |                 excsum[idm] += numpy.dot(den, exc)
1079 |                 den = rho_b[0]*weight
1080 |                 nelec[1,idm] += den.sum()
1081 |                 excsum[idm] += numpy.dot(den, exc)
1082 | 
1083 |                 wva, wvb = _uks_gga_wv0((rho_a,rho_b), vxc, weight)
1084 |                 aow = numpy.einsum('npi,np->pi', ao[:4], wva, out=aow)
1085 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
1086 |                 aow = numpy.einsum('npi,np->pi', ao[:4], wvb, out=aow)
1087 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
1088 | 
1089 | # FIXME: .5 * .5   First 0.5 for v+v.T symmetrization.
1090 | # Second 0.5 is due to the Libxc convention tau = 1/2 \nabla\phi\dot\nabla\phi
1091 |                 wv = (.25 * weight * vtau[:,0]).reshape(-1,1)
1092 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[1], wv*ao[1], mask, shls_slice, ao_loc)
1093 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[2], wv*ao[2], mask, shls_slice, ao_loc)
1094 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[3], wv*ao[3], mask, shls_slice, ao_loc)
1095 |                 wv = (.25 * weight * vtau[:,1]).reshape(-1,1)
1096 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[1], wv*ao[1], mask, shls_slice, ao_loc)
1097 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[2], wv*ao[2], mask, shls_slice, ao_loc)
1098 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[3], wv*ao[3], mask, shls_slice, ao_loc)
1099 |                 rho_a = rho_b = exc = vxc = vrho = vsigma = wva = wvb = None
1100 | 
1101 |     for i in range(nset):
1102 |         vmat[0,i] = vmat[0,i] + vmat[0,i].T
1103 |         vmat[1,i] = vmat[1,i] + vmat[1,i].T
1104 |     if isinstance(dma, numpy.ndarray) and dma.ndim == 2:
1105 |         vmat = vmat[:,0]
1106 |         nelec = nelec.reshape(2)
1107 |         excsum = excsum[0]
1108 |     return nelec, excsum, vmat
1109 | 
1110 | def _format_uks_dm(dms):
1111 |     if isinstance(dms, numpy.ndarray) and dms.ndim == 2:  # RHF DM
1112 |         dma = dmb = dms * .5
1113 |     else:
1114 |         dma, dmb = dms
1115 |     if getattr(dms, 'mo_coeff', None) is not None:
1116 |         mo_coeff = dms.mo_coeff
1117 |         mo_occ = dms.mo_occ
1118 |         if mo_coeff[0].ndim < dma.ndim: # handle ROKS
1119 |             mo_occa = numpy.array(mo_occ> 0, dtype=numpy.double)
1120 |             mo_occb = numpy.array(mo_occ==2, dtype=numpy.double)
1121 |             dma = lib.tag_array(dma, mo_coeff=mo_coeff, mo_occ=mo_occa)
1122 |             dmb = lib.tag_array(dmb, mo_coeff=mo_coeff, mo_occ=mo_occb)
1123 |         else:
1124 |             dma = lib.tag_array(dma, mo_coeff=mo_coeff[0], mo_occ=mo_occ[0])
1125 |             dmb = lib.tag_array(dmb, mo_coeff=mo_coeff[1], mo_occ=mo_occ[1])
1126 |     return dma, dmb
1127 | 
1128 | nr_rks_vxc = nr_rks
1129 | nr_uks_vxc = nr_uks
1130 | 
1131 | def nr_rks_fxc(ni, mol, grids, xc_code, dm0, dms, relativity=0, hermi=0,
1132 |                rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
1133 |     '''Contract RKS XC (singlet hessian) kernel matrix with given density matrices
1134 | 
1135 |     Args:
1136 |         ni : an instance of :class:`NumInt`
1137 | 
1138 |         mol : an instance of :class:`Mole`
1139 | 
1140 |         grids : an instance of :class:`Grids`
1141 |             grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
1142 |         xc_code : str
1143 |             XC functional description.
1144 |             See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
1145 |         dms : 2D array a list of 2D arrays
1146 |             Density matrix or multiple density matrices
1147 | 
1148 |     Kwargs:
1149 |         hermi : int
1150 |             Input density matrices symmetric or not
1151 |         max_memory : int or float
1152 |             The maximum size of cache to use (in MB).
1153 |         rho0 : float array
1154 |             Zero-order density (and density derivative for GGA).  Giving kwargs rho0,
1155 |             vxc and fxc to improve better performance.
1156 |         vxc : float array
1157 |             First order XC derivatives
1158 |         fxc : float array
1159 |             Second order XC derivatives
1160 | 
1161 |     Returns:
1162 |         nelec, excsum, vmat.
1163 |         nelec is the number of electrons generated by numerical integration.
1164 |         excsum is the XC functional value.  vmat is the XC potential matrix in
1165 |         2D array of shape (nao,nao) where nao is the number of AO functions.
1166 | 
1167 |     Examples:
1168 | 
1169 |     '''
1170 |     xctype = ni._xc_type(xc_code)
1171 | 
1172 |     make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
1173 |     if ((xctype == 'LDA' and fxc is None) or
1174 |         (xctype == 'GGA' and rho0 is None)):
1175 |         make_rho0 = ni._gen_rho_evaluator(mol, dm0, 1)[0]
1176 | 
1177 |     shls_slice = (0, mol.nbas)
1178 |     ao_loc = mol.ao_loc_nr()
1179 | 
1180 |     vmat = numpy.zeros((nset,nao,nao))
1181 |     aow = None
1182 |     if xctype == 'LDA':
1183 |         ao_deriv = 0
1184 |         ip = 0
1185 |         for ao, mask, weight, coords \
1186 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1187 |             ngrid = weight.size
1188 |             aow = numpy.ndarray(ao.shape, order='F', buffer=aow)
1189 |             if fxc is None:
1190 |                 rho = make_rho0(0, ao, mask, 'LDA')
1191 |                 fxc0 = ni.eval_xc(xc_code, rho, 0, relativity, 2, verbose)[2]
1192 |                 frr = fxc0[0]
1193 |             else:
1194 |                 frr = fxc[0][ip:ip+ngrid]
1195 |                 ip += ngrid
1196 | 
1197 |             for i in range(nset):
1198 |                 rho1 = make_rho(i, ao, mask, 'LDA')
1199 |                 #:aow = numpy.einsum('pi,p->pi', ao, weight*frr*rho1, out=aow)
1200 |                 aow = _scale_ao(ao, weight*frr*rho1, out=aow)
1201 |                 vmat[i] += _dot_ao_ao(mol, aow, ao, mask, shls_slice, ao_loc)
1202 |                 rho1 = None
1203 | 
1204 |     elif xctype == 'GGA':
1205 |         ao_deriv = 1
1206 |         ip = 0
1207 |         for ao, mask, weight, coords \
1208 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1209 |             ngrid = weight.size
1210 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
1211 |             if rho0 is None:
1212 |                 rho = make_rho0(0, ao, mask, 'GGA')
1213 |             else:
1214 |                 rho = numpy.asarray(rho0[:,ip:ip+ngrid], order='C')
1215 |             if vxc is None or fxc is None:
1216 |                 vxc0, fxc0 = ni.eval_xc(xc_code, rho, 0, relativity, 2, verbose)[1:3]
1217 |             else:
1218 |                 vxc0 = (None, vxc[1][ip:ip+ngrid])
1219 |                 fxc0 = (fxc[0][ip:ip+ngrid], fxc[1][ip:ip+ngrid], fxc[2][ip:ip+ngrid])
1220 |                 ip += ngrid
1221 | 
1222 |             for i in range(nset):
1223 |                 rho1 = make_rho(i, ao, mask, 'GGA')
1224 |                 wv = _rks_gga_wv1(rho, rho1, vxc0, fxc0, weight)
1225 |                 #:aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
1226 |                 aow = _scale_ao(ao, wv, out=aow)
1227 |                 vmat[i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
1228 |                 rho1 = sigma1 = None
1229 | 
1230 |         for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
1231 |             vmat[i] = vmat[i] + vmat[i].T.conj()
1232 | 
1233 |     elif xctype == 'NLC':
1234 |         raise NotImplementedError('NLC')
1235 |     elif xctype == 'MGGA':
1236 |         raise NotImplementedError('meta-GGA')
1237 | 
1238 |     if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
1239 |         vmat = vmat[0]
1240 |     return vmat
1241 | 
1242 | def nr_rks_fxc_st(ni, mol, grids, xc_code, dm0, dms_alpha, relativity=0, singlet=True,
1243 |                   rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
1244 |     '''Associated to singlet or triplet Hessian
1245 |     Note the difference to nr_rks_fxc, dms_alpha is the response density
1246 |     matrices of alpha spin, alpha+/-beta DM is applied due to singlet/triplet
1247 |     coupling
1248 | 
1249 |     Ref. CPL, 256, 454
1250 |     '''
1251 |     xctype = ni._xc_type(xc_code)
1252 | 
1253 |     make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms_alpha, hermi=0)
1254 |     if ((xctype == 'LDA' and fxc is None) or
1255 |         (xctype == 'GGA' and rho0 is None)):
1256 |         make_rho0 = ni._gen_rho_evaluator(mol, dm0, hermi=1)[0]
1257 | 
1258 |     shls_slice = (0, mol.nbas)
1259 |     ao_loc = mol.ao_loc_nr()
1260 | 
1261 |     vmat = numpy.zeros((nset,nao,nao))
1262 |     aow = None
1263 |     if xctype == 'LDA':
1264 |         ao_deriv = 0
1265 |         ip = 0
1266 |         for ao, mask, weight, coords \
1267 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1268 |             ngrid = weight.size
1269 |             aow = numpy.ndarray(ao.shape, order='F', buffer=aow)
1270 |             if fxc is None:
1271 |                 rho = make_rho0(0, ao, mask, 'LDA')
1272 |                 rho *= .5  # alpha density
1273 |                 fxc0 = ni.eval_xc(xc_code, (rho,rho), 1, deriv=2)[2]
1274 |                 u_u, u_d, d_d = fxc0[0].T
1275 |             else:
1276 |                 if fxc[0].ndim == 1:
1277 |                     raise RuntimeError('cached (rho, vxc, fxc) need to be '
1278 |                                        'generated by cache_xc_kernel with flag spin=1')
1279 |                 u_u, u_d, d_d = fxc[0][ip:ip+ngrid].T
1280 |                 ip += ngrid
1281 |             if singlet:
1282 |                 frho = u_u + u_d
1283 |                 if 0:
1284 |                     rho = ni.eval_rho2(mol, ao, mo_coeff, mo_occ, mask, 'LDA')
1285 |                     fxc_test = ni.eval_xc(xc_code, rho, 0, deriv=2)[2]
1286 |                     assert(numpy.linalg.norm(fxc_test[0]*2-frho) < 1e-4)
1287 |             else:
1288 |                 frho = u_u - u_d
1289 | 
1290 |             for i in range(nset):
1291 |                 rho1 = make_rho(i, ao, mask, 'LDA')
1292 |                 #:aow = numpy.einsum('pi,p->pi', ao, weight*frho*rho1, out=aow)
1293 |                 aow = _scale_ao(ao, weight*frho*rho1, out=aow)
1294 |                 vmat[i] += _dot_ao_ao(mol, aow, ao, mask, shls_slice, ao_loc)
1295 |                 rho1 = None
1296 | 
1297 |     elif xctype == 'GGA':
1298 |         ao_deriv = 1
1299 |         ip = 0
1300 |         for ao, mask, weight, coords \
1301 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1302 |             ngrid = weight.size
1303 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
1304 |             if vxc is None or fxc is None:
1305 |                 rho = make_rho0(0, ao, mask, 'GGA')
1306 |                 rho *= .5  # alpha density
1307 |                 vxc0, fxc0 = ni.eval_xc(xc_code, (rho,rho), 1, deriv=2)[1:3]
1308 | 
1309 |                 vsigma = vxc0[1].T
1310 |                 u_u, u_d, d_d = fxc0[0].T  # v2rho2
1311 |                 u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc0[1].T  # v2rhosigma
1312 |                 uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc0[2].T  # v2sigma2
1313 |             else:
1314 |                 if rho0[0].ndim == 1:
1315 |                     raise RuntimeError('cached (rho, vxc, fxc) need to be '
1316 |                                        'generated by cache_xc_kernel with flag spin=1')
1317 |                 rho = rho0[0][:,ip:ip+ngrid]
1318 |                 vsigma = vxc[1][ip:ip+ngrid].T
1319 |                 u_u, u_d, d_d = fxc[0][ip:ip+ngrid].T  # v2rho2
1320 |                 u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1][ip:ip+ngrid].T  # v2rhosigma
1321 |                 uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2][ip:ip+ngrid].T  # v2sigma2
1322 |                 ip += ngrid
1323 | 
1324 |             # Factorization differs to CPL, 256, 454, to use _rks_gga_wv1 function
1325 |             if singlet:
1326 |                 fgamma = vsigma[0] + vsigma[1] * .5
1327 |                 frho = u_u + u_d
1328 |                 fgg = uu_uu + .5*ud_ud + 2*uu_ud + uu_dd
1329 |                 frhogamma = u_uu + u_dd + u_ud
1330 |             else:
1331 |                 fgamma = vsigma[0] - vsigma[1] * .5
1332 |                 frho = u_u - u_d
1333 |                 fgg = uu_uu - uu_dd
1334 |                 frhogamma = u_uu - u_dd
1335 | 
1336 |             for i in range(nset):
1337 |                 # rho1[0 ] = |b><j| z_{bj}
1338 |                 # rho1[1:] = \nabla(|b><j|) z_{bj}
1339 |                 rho1 = make_rho(i, ao, mask, 'GGA')
1340 |                 wv = _rks_gga_wv1(rho, rho1, (None,fgamma), (frho,frhogamma,fgg), weight)
1341 |                 #:aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
1342 |                 aow = _scale_ao(ao, wv, out=aow)
1343 |                 vmat[i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
1344 |                 rho1 = sigma1 = None
1345 | 
1346 |         for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
1347 |             vmat[i] = vmat[i] + vmat[i].T.conj()
1348 | 
1349 |     elif xctype == 'NLC':
1350 |         raise NotImplementedError('NLC')
1351 |     elif xctype == 'MGGA':
1352 |         raise NotImplementedError('meta-GGA')
1353 | 
1354 |     if isinstance(dms_alpha, numpy.ndarray) and dms_alpha.ndim == 2:
1355 |         vmat = vmat[0]
1356 |     return vmat
1357 | 
1358 | def _rks_gga_wv0(rho, vxc, weight):
1359 |     vrho, vgamma = vxc[:2]
1360 |     ngrid = vrho.size
1361 |     wv = numpy.empty((4,ngrid))
1362 |     wv[0]  = weight * vrho
1363 |     wv[1:] = (weight * vgamma * 2) * rho[1:4]
1364 |     wv[0] *= .5  # v+v.T should be applied in the caller
1365 |     return wv
1366 | 
1367 | def _rks_gga_wv1(rho0, rho1, vxc, fxc, weight):
1368 |     vgamma = vxc[1]
1369 |     frho, frhogamma, fgg = fxc[:3]
1370 |     # sigma1 ~ \nabla(\rho_\alpha+\rho_\beta) dot \nabla(|b><j|) z_{bj}
1371 |     sigma1 = numpy.einsum('xi,xi->i', rho0[1:4], rho1[1:4])
1372 |     ngrid = vgamma.size
1373 |     wv = numpy.empty((4,ngrid))
1374 |     wv[0]  = frho * rho1[0]
1375 |     wv[0] += frhogamma * sigma1 * 2
1376 |     wv[1:] = (fgg * sigma1 * 4 + frhogamma * rho1[0] * 2) * rho0[1:4]
1377 |     wv[1:]+= vgamma * rho1[1:4] * 2
1378 |     wv *= weight
1379 |     wv[0] *= .5  # v+v.T should be applied in the caller
1380 |     return wv
1381 | 
1382 | def _rks_gga_wv2(rho0, rho1, fxc, kxc, weight):
1383 |     frr, frg, fgg = fxc[:3]
1384 |     frrr, frrg, frgg, fggg = kxc
1385 |     sigma1 = numpy.einsum('xi,xi->i', rho0[1:], rho1[1:])
1386 |     r1r1 = rho1[0]**2
1387 |     s1s1 = sigma1**2
1388 |     r1s1 = rho1[0] * sigma1
1389 |     sigma2 = numpy.einsum('xi,xi->i', rho1[1:], rho1[1:])
1390 |     ngrid = frrr.size
1391 |     wv = numpy.empty((4,ngrid))
1392 |     wv[0]  = frrr * r1r1
1393 |     wv[0] += 4 * frrg * r1s1
1394 |     wv[0] += 4 * frgg * s1s1
1395 |     wv[0] += 2 * frg * sigma2
1396 |     wv[1:]  = 2 * frrg * r1r1 * rho0[1:]
1397 |     wv[1:] += 8 * frgg * r1s1 * rho0[1:]
1398 |     wv[1:] += 4 * frg * rho1[0] * rho1[1:]
1399 |     wv[1:] += 4 * fgg * sigma2 * rho0[1:]
1400 |     wv[1:] += 8 * fgg * sigma1 * rho1[1:]
1401 |     wv[1:] += 8 * fggg * s1s1 * rho0[1:]
1402 |     wv *= weight
1403 |     wv[0]*=.5  # v+v.T should be applied in the caller
1404 |     return wv
1405 | 
1406 | def nr_uks_fxc(ni, mol, grids, xc_code, dm0, dms, relativity=0, hermi=0,
1407 |                rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
1408 |     '''Contract UKS XC kernel matrix with given density matrices
1409 | 
1410 |     Args:
1411 |         ni : an instance of :class:`NumInt`
1412 | 
1413 |         mol : an instance of :class:`Mole`
1414 | 
1415 |         grids : an instance of :class:`Grids`
1416 |             grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
1417 |         xc_code : str
1418 |             XC functional description.
1419 |             See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
1420 |         dms : 2D array a list of 2D arrays
1421 |             Density matrix or multiple density matrices
1422 | 
1423 |     Kwargs:
1424 |         hermi : int
1425 |             Input density matrices symmetric or not
1426 |         max_memory : int or float
1427 |             The maximum size of cache to use (in MB).
1428 |         rho0 : float array
1429 |             Zero-order density (and density derivative for GGA).  Giving kwargs rho0,
1430 |             vxc and fxc to improve better performance.
1431 |         vxc : float array
1432 |             First order XC derivatives
1433 |         fxc : float array
1434 |             Second order XC derivatives
1435 | 
1436 |     Returns:
1437 |         nelec, excsum, vmat.
1438 |         nelec is the number of electrons generated by numerical integration.
1439 |         excsum is the XC functional value.  vmat is the XC potential matrix in
1440 |         2D array of shape (nao,nao) where nao is the number of AO functions.
1441 | 
1442 |     Examples:
1443 | 
1444 |     '''
1445 |     xctype = ni._xc_type(xc_code)
1446 | 
1447 |     dma, dmb = _format_uks_dm(dms)
1448 |     nao = dms.shape[-1]
1449 |     make_rhoa, nset = ni._gen_rho_evaluator(mol, dma, hermi)[:2]
1450 |     make_rhob       = ni._gen_rho_evaluator(mol, dmb, hermi)[0]
1451 | 
1452 |     if ((xctype == 'LDA' and fxc is None) or
1453 |         (xctype == 'GGA' and rho0 is None)):
1454 |         make_rho0 = ni._gen_rho_evaluator(mol, _format_uks_dm(dm0), 1)[0]
1455 | 
1456 |     shls_slice = (0, mol.nbas)
1457 |     ao_loc = mol.ao_loc_nr()
1458 | 
1459 |     vmat = numpy.zeros((2,nset,nao,nao))
1460 |     aow = None
1461 |     if xctype == 'LDA':
1462 |         ao_deriv = 0
1463 |         ip = 0
1464 |         for ao, mask, weight, coords \
1465 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1466 |             ngrid = weight.size
1467 |             aow = numpy.ndarray(ao.shape, order='F', buffer=aow)
1468 |             if fxc is None:
1469 |                 rho0a = make_rho0(0, ao, mask, xctype)
1470 |                 rho0b = make_rho0(1, ao, mask, xctype)
1471 |                 fxc0 = ni.eval_xc(xc_code, (rho0a,rho0b), 1, relativity, 2, verbose)[2]
1472 |                 u_u, u_d, d_d = fxc0[0].T
1473 |             else:
1474 |                 u_u, u_d, d_d = fxc[0][ip:ip+ngrid].T
1475 |                 ip += ngrid
1476 | 
1477 |             for i in range(nset):
1478 |                 rho1a = make_rhoa(i, ao, mask, xctype)
1479 |                 rho1b = make_rhob(i, ao, mask, xctype)
1480 |                 wv = u_u * rho1a + u_d * rho1b
1481 |                 wv *= weight
1482 |                 #:aow = numpy.einsum('pi,p->pi', ao, wv, out=aow)
1483 |                 aow = _scale_ao(ao, wv, out=aow)
1484 |                 vmat[0,i] += _dot_ao_ao(mol, aow, ao, mask, shls_slice, ao_loc)
1485 |                 wv = u_d * rho1a + d_d * rho1b
1486 |                 wv *= weight
1487 |                 #:aow = numpy.einsum('pi,p->pi', ao, wv, out=aow)
1488 |                 aow = _scale_ao(ao, wv, out=aow)
1489 |                 vmat[1,i] += _dot_ao_ao(mol, aow, ao, mask, shls_slice, ao_loc)
1490 | 
1491 |     elif xctype == 'GGA':
1492 |         ao_deriv = 1
1493 |         ip = 0
1494 |         for ao, mask, weight, coords \
1495 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1496 |             ngrid = weight.size
1497 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
1498 |             if rho0 is None:
1499 |                 rho0a = make_rho0(0, ao, mask, xctype)
1500 |                 rho0b = make_rho0(1, ao, mask, xctype)
1501 |             else:
1502 |                 rho0a = rho0[0][:,ip:ip+ngrid]
1503 |                 rho0b = rho0[1][:,ip:ip+ngrid]
1504 |             if vxc is None or fxc is None:
1505 |                 vxc0, fxc0 = ni.eval_xc(xc_code, (rho0a,rho0b), 1, relativity, 2, verbose)[1:3]
1506 |             else:
1507 |                 vxc0 = (None, vxc[1][ip:ip+ngrid])
1508 |                 fxc0 = (fxc[0][ip:ip+ngrid], fxc[1][ip:ip+ngrid], fxc[2][ip:ip+ngrid])
1509 |                 ip += ngrid
1510 | 
1511 |             for i in range(nset):
1512 |                 rho1a = make_rhoa(i, ao, mask, xctype)
1513 |                 rho1b = make_rhob(i, ao, mask, xctype)
1514 |                 wva, wvb = _uks_gga_wv1((rho0a,rho0b), (rho1a,rho1b), vxc0, fxc0, weight)
1515 |                 #:aow = numpy.einsum('npi,np->pi', ao, wva, out=aow)
1516 |                 aow = _scale_ao(ao, wva, out=aow)
1517 |                 vmat[0,i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
1518 |                 #:aow = numpy.einsum('npi,np->pi', ao, wvb, out=aow)
1519 |                 aow = _scale_ao(ao, wvb, out=aow)
1520 |                 vmat[1,i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
1521 | 
1522 |         for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
1523 |             vmat[0,i] = vmat[0,i] + vmat[0,i].T.conj()
1524 |             vmat[1,i] = vmat[1,i] + vmat[1,i].T.conj()
1525 | 
1526 |     elif xctype == 'NLC':
1527 |         raise NotImplementedError('NLC')
1528 |     elif xctype == 'MGGA':
1529 |         raise NotImplementedError('meta-GGA')
1530 | 
1531 |     if isinstance(dma, numpy.ndarray) and dma.ndim == 2:
1532 |         vmat = vmat[:,0]
1533 |     return vmat
1534 | 
1535 | def _uks_gga_wv0(rho, vxc, weight):
1536 |     rhoa, rhob = rho
1537 |     vrho, vsigma = vxc[:2]
1538 |     ngrid = vrho.shape[0]
1539 |     wva = numpy.empty((4,ngrid))
1540 |     wva[0]  = weight * vrho[:,0] * .5  # v+v.T should be applied in the caller
1541 |     wva[1:] = rhoa[1:4] * (weight * vsigma[:,0] * 2)  # sigma_uu
1542 |     wva[1:]+= rhob[1:4] * (weight * vsigma[:,1])      # sigma_ud
1543 |     wvb = numpy.empty((4,ngrid))
1544 |     wvb[0]  = weight * vrho[:,1] * .5  # v+v.T should be applied in the caller
1545 |     wvb[1:] = rhob[1:4] * (weight * vsigma[:,2] * 2)  # sigma_dd
1546 |     wvb[1:]+= rhoa[1:4] * (weight * vsigma[:,1])      # sigma_ud
1547 |     return wva, wvb
1548 | 
1549 | def _uks_gga_wv1(rho0, rho1, vxc, fxc, weight):
1550 |     uu, ud, dd = vxc[1].T
1551 |     u_u, u_d, d_d = fxc[0].T
1552 |     u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1].T
1553 |     uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2].T
1554 |     ngrid = uu.size
1555 | 
1556 |     rho0a, rho0b = rho0
1557 |     rho1a, rho1b = rho1
1558 |     a0a1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1a[1:4])
1559 |     a0b1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1b[1:4])
1560 |     b0a1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1a[1:4])
1561 |     b0b1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1b[1:4])
1562 | 
1563 |     wva = numpy.empty((4,ngrid))
1564 |     wvb = numpy.empty((4,ngrid))
1565 |     # alpha = alpha-alpha * alpha
1566 |     wva[0]  = u_u * rho1a[0]
1567 |     wva[0] += u_uu * a0a1 * 2
1568 |     wva[0] += u_ud * b0a1
1569 |     wva[1:] = uu * rho1a[1:4] * 2
1570 |     wva[1:]+= u_uu * rho1a[0] * rho0a[1:4] * 2
1571 |     wva[1:]+= u_ud * rho1a[0] * rho0b[1:4]
1572 |     wva[1:]+= uu_uu * a0a1 * rho0a[1:4] * 4
1573 |     wva[1:]+= uu_ud * a0a1 * rho0b[1:4] * 2
1574 |     wva[1:]+= uu_ud * b0a1 * rho0a[1:4] * 2
1575 |     wva[1:]+= ud_ud * b0a1 * rho0b[1:4]
1576 | 
1577 |     # alpha = alpha-beta  * beta
1578 |     wva[0] += u_d * rho1b[0]
1579 |     wva[0] += u_ud * a0b1
1580 |     wva[0] += u_dd * b0b1 * 2
1581 |     wva[1:]+= ud * rho1b[1:4]
1582 |     wva[1:]+= d_uu * rho1b[0] * rho0a[1:4] * 2
1583 |     wva[1:]+= d_ud * rho1b[0] * rho0b[1:4]
1584 |     wva[1:]+= uu_ud * a0b1 * rho0a[1:4] * 2
1585 |     wva[1:]+= ud_ud * a0b1 * rho0b[1:4]
1586 |     wva[1:]+= uu_dd * b0b1 * rho0a[1:4] * 4
1587 |     wva[1:]+= ud_dd * b0b1 * rho0b[1:4] * 2
1588 |     wva *= weight
1589 |     wva[0] *= .5  # v+v.T should be applied in the caller
1590 | 
1591 |     # beta = beta-alpha * alpha
1592 |     wvb[0]  = u_d * rho1a[0]
1593 |     wvb[0] += d_ud * b0a1
1594 |     wvb[0] += d_uu * a0a1 * 2
1595 |     wvb[1:] = ud * rho1a[1:4]
1596 |     wvb[1:]+= u_dd * rho1a[0] * rho0b[1:4] * 2
1597 |     wvb[1:]+= u_ud * rho1a[0] * rho0a[1:4]
1598 |     wvb[1:]+= ud_dd * b0a1 * rho0b[1:4] * 2
1599 |     wvb[1:]+= ud_ud * b0a1 * rho0a[1:4]
1600 |     wvb[1:]+= uu_dd * a0a1 * rho0b[1:4] * 4
1601 |     wvb[1:]+= uu_ud * a0a1 * rho0a[1:4] * 2
1602 | 
1603 |     # beta = beta-beta  * beta
1604 |     wvb[0] += d_d * rho1b[0]
1605 |     wvb[0] += d_dd * b0b1 * 2
1606 |     wvb[0] += d_ud * a0b1
1607 |     wvb[1:]+= dd * rho1b[1:4] * 2
1608 |     wvb[1:]+= d_dd * rho1b[0] * rho0b[1:4] * 2
1609 |     wvb[1:]+= d_ud * rho1b[0] * rho0a[1:4]
1610 |     wvb[1:]+= dd_dd * b0b1 * rho0b[1:4] * 4
1611 |     wvb[1:]+= ud_dd * b0b1 * rho0a[1:4] * 2
1612 |     wvb[1:]+= ud_dd * a0b1 * rho0b[1:4] * 2
1613 |     wvb[1:]+= ud_ud * a0b1 * rho0a[1:4]
1614 |     wvb *= weight
1615 |     wvb[0] *= .5  # v+v.T should be applied in the caller
1616 |     return wva, wvb
1617 | 
1618 | def _uks_gga_wv2(rho0, rho1, fxc, kxc, weight):
1619 |     u_u, u_d, d_d = fxc[0].T
1620 |     u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1].T
1621 |     uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2].T
1622 |     u_u_u, u_u_d, u_d_d, d_d_d = kxc[0].T
1623 |     u_u_uu, u_u_ud, u_u_dd, u_d_uu, u_d_ud, u_d_dd, d_d_uu, \
1624 |             d_d_ud, d_d_dd = kxc[1].T
1625 |     u_uu_uu, u_uu_ud, u_uu_dd, u_ud_ud, u_ud_dd, u_dd_dd, d_uu_uu, d_uu_ud, \
1626 |             d_uu_dd, d_ud_ud, d_ud_dd, d_dd_dd = kxc[2].T
1627 |     uu_uu_uu, uu_uu_ud, uu_uu_dd, uu_ud_ud, uu_ud_dd, uu_dd_dd, ud_ud_ud, \
1628 |             ud_ud_dd, ud_dd_dd, dd_dd_dd = kxc[3].T
1629 |     ngrid = u_u.size
1630 | 
1631 |     rho0a, rho0b = rho0
1632 |     rho1a, rho1b = rho1
1633 |     a0a1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1a[1:4])
1634 |     a0b1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1b[1:4])
1635 |     b0a1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1a[1:4])
1636 |     b0b1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1b[1:4])
1637 |     a1a1 = numpy.einsum('xi,xi->i', rho1a[1:4], rho1a[1:4])
1638 |     a1b1 = numpy.einsum('xi,xi->i', rho1a[1:4], rho1b[1:4])
1639 |     b1a1 = a1b1
1640 |     b1b1 = numpy.einsum('xi,xi->i', rho1b[1:4], rho1b[1:4])
1641 |     a0a1_a0a1 = numpy.einsum('i,i->i', a0a1, a0a1)
1642 |     a0a1_a0b1 = numpy.einsum('i,i->i', a0a1, a0b1)
1643 |     a0a1_b0a1 = numpy.einsum('i,i->i', a0a1, b0a1)
1644 |     a0a1_b0b1 = numpy.einsum('i,i->i', a0a1, b0b1)
1645 |     a0b1_a0a1 = a0a1_a0b1
1646 |     a0b1_a0b1 = numpy.einsum('i,i->i', a0b1, a0b1)
1647 |     a0b1_b0a1 = numpy.einsum('i,i->i', a0b1, b0a1)
1648 |     a0b1_b0b1 = numpy.einsum('i,i->i', a0b1, b0b1)
1649 |     b0a1_a0a1 = a0a1_b0a1
1650 |     b0a1_a0b1 = a0b1_b0a1
1651 |     b0a1_b0a1 = numpy.einsum('i,i->i', b0a1, b0a1)
1652 |     b0a1_b0b1 = numpy.einsum('i,i->i', b0a1, b0b1)
1653 |     b0b1_a0a1 = a0a1_b0b1
1654 |     b0b1_a0b1 = a0b1_b0b1
1655 |     b0b1_b0a1 = b0a1_b0b1
1656 |     b0b1_b0b1 = numpy.einsum('i,i->i', b0b1, b0b1)
1657 | 
1658 |     wva = numpy.zeros((4,ngrid))
1659 |     wva[0] += numpy.einsum('i,i,i->i', u_u_u, rho1a[0], rho1a[0])
1660 |     wva[0] += numpy.einsum('i,i,i->i', u_u_d, rho1a[0], rho1b[0]) * 2
1661 |     wva[0] += numpy.einsum('i,i,i->i', u_d_d, rho1b[0], rho1b[0])
1662 |     wva[0] += numpy.einsum('i,i->i', u_uu, a1a1) * 2
1663 |     wva[0] += numpy.einsum('i,i->i', u_ud, a1b1) * 2
1664 |     wva[0] += numpy.einsum('i,i->i', u_dd, b1b1) * 2
1665 |     wva[1:] += numpy.einsum('i,i,xi->xi', u_uu, rho1a[0], rho1a[1:]) * 4
1666 |     wva[1:] += numpy.einsum('i,i,xi->xi', d_uu, rho1b[0], rho1a[1:]) * 4
1667 |     wva[1:] += numpy.einsum('i,i,xi->xi', u_ud, rho1a[0], rho1b[1:]) * 2
1668 |     wva[1:] += numpy.einsum('i,i,xi->xi', d_ud, rho1b[0], rho1b[1:]) * 2
1669 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu, a0a1, rho1a[1:]) * 8
1670 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud, a0a1, rho1b[1:]) * 4
1671 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud, a0b1, rho1a[1:]) * 4
1672 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud, a0b1, rho1b[1:]) * 2
1673 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud, b0a1, rho1a[1:]) * 4
1674 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_dd, b0b1, rho1a[1:]) * 8
1675 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud, b0a1, rho1b[1:]) * 2
1676 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_dd, b0b1, rho1b[1:]) * 4
1677 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu, a1a1, rho0a[1:]) * 4
1678 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud, a1b1, rho0a[1:]) * 4
1679 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_dd, b1b1, rho0a[1:]) * 4
1680 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud, a1a1, rho0b[1:]) * 2
1681 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud, a1b1, rho0b[1:]) * 2
1682 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_dd, b1b1, rho0b[1:]) * 2
1683 |     wva[0] += numpy.einsum('i,i,i->i', u_u_uu, rho1a[0], a0a1) * 4
1684 |     wva[0] += numpy.einsum('i,i,i->i', u_d_uu, rho1b[0], a0a1) * 4
1685 |     wva[0] += numpy.einsum('i,i,i->i', u_u_ud, rho1a[0], a0b1) * 2
1686 |     wva[0] += numpy.einsum('i,i,i->i', u_d_ud, rho1b[0], a0b1) * 2
1687 |     wva[0] += numpy.einsum('i,i,i->i', u_u_ud, rho1a[0], b0a1) * 2
1688 |     wva[0] += numpy.einsum('i,i,i->i', u_d_ud, rho1b[0], b0a1) * 2
1689 |     wva[0] += numpy.einsum('i,i,i->i', u_u_dd, rho1a[0], b0b1) * 4
1690 |     wva[0] += numpy.einsum('i,i,i->i', u_d_dd, rho1b[0], b0b1) * 4
1691 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_u_uu, rho1a[0], rho1a[0], rho0a[1:]) * 2
1692 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_d_uu, rho1a[0], rho1b[0], rho0a[1:]) * 4
1693 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_d_uu, rho1b[0], rho1b[0], rho0a[1:]) * 2
1694 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_u_ud, rho1a[0], rho1a[0], rho0a[1:])
1695 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_d_ud, rho1a[0], rho1b[0], rho0a[1:]) * 2
1696 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_d_ud, rho1b[0], rho1b[0], rho0a[1:])
1697 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_uu, rho1a[0], a0a1, rho0a[1:]) * 8
1698 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_uu, rho1b[0], a0a1, rho0a[1:]) * 8
1699 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_ud, rho1a[0], a0b1, rho0a[1:]) * 4
1700 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_ud, rho1b[0], a0b1, rho0a[1:]) * 4
1701 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_ud, rho1a[0], b0a1, rho0a[1:]) * 4
1702 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_ud, rho1b[0], b0a1, rho0a[1:]) * 4
1703 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_dd, rho1a[0], b0b1, rho0a[1:]) * 8
1704 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_dd, rho1b[0], b0b1, rho0a[1:]) * 8
1705 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_ud, rho1a[0], a0a1, rho0b[1:]) * 4
1706 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_ud, rho1b[0], a0a1, rho0b[1:]) * 4
1707 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_ud, rho1a[0], a0b1, rho0b[1:]) * 2
1708 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_ud, rho1b[0], a0b1, rho0b[1:]) * 2
1709 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_ud, rho1a[0], b0a1, rho0b[1:]) * 2
1710 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_ud, rho1b[0], b0a1, rho0b[1:]) * 2
1711 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_dd, rho1a[0], b0b1, rho0b[1:]) * 4
1712 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_dd, rho1b[0], b0b1, rho0b[1:]) * 4
1713 |     wva[0] += numpy.einsum('i,i->i', u_uu_uu, a0a1_a0a1) * 4
1714 |     wva[0] += numpy.einsum('i,i->i', u_uu_ud, a0a1_a0b1) * 2
1715 |     wva[0] += numpy.einsum('i,i->i', u_uu_ud, a0b1_a0a1) * 2
1716 |     wva[0] += numpy.einsum('i,i->i', u_ud_ud, a0b1_a0b1)
1717 |     wva[0] += numpy.einsum('i,i->i', u_uu_ud, a0a1_b0a1) * 4
1718 |     wva[0] += numpy.einsum('i,i->i', u_ud_ud, a0b1_b0a1) * 2
1719 |     wva[0] += numpy.einsum('i,i->i', u_uu_dd, a0a1_b0b1) * 8
1720 |     wva[0] += numpy.einsum('i,i->i', u_ud_dd, a0b1_b0b1) * 4
1721 |     wva[0] += numpy.einsum('i,i->i', u_ud_ud, b0a1_b0a1)
1722 |     wva[0] += numpy.einsum('i,i->i', u_ud_dd, b0a1_b0b1) * 2
1723 |     wva[0] += numpy.einsum('i,i->i', u_ud_dd, b0b1_b0a1) * 2
1724 |     wva[0] += numpy.einsum('i,i->i', u_dd_dd, b0b1_b0b1) * 4
1725 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu_uu, a0a1_a0a1, rho0a[1:]) * 8
1726 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu_ud, a0a1_a0b1, rho0a[1:]) * 4
1727 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu_ud, a0b1_a0a1, rho0a[1:]) * 4
1728 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0b1_a0b1, rho0a[1:]) * 2
1729 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu_ud, a0a1_a0a1, rho0b[1:]) * 4
1730 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0a1_a0b1, rho0b[1:]) * 2
1731 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0b1_a0a1, rho0b[1:]) * 2
1732 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud_ud, a0b1_a0b1, rho0b[1:])
1733 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu_ud, a0a1_b0a1, rho0a[1:]) * 8
1734 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0b1_b0a1, rho0a[1:]) * 4
1735 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu_dd, a0a1_b0b1, rho0a[1:]) * 16
1736 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, a0b1_b0b1, rho0a[1:]) * 8
1737 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0a1_b0a1, rho0b[1:]) * 4
1738 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud_ud, a0b1_b0a1, rho0b[1:]) * 2
1739 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, a0a1_b0b1, rho0b[1:]) * 8
1740 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, a0b1_b0b1, rho0b[1:]) * 4
1741 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, b0a1_b0a1, rho0a[1:]) * 2
1742 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, b0b1_b0a1, rho0a[1:]) * 4
1743 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, b0a1_b0b1, rho0a[1:]) * 4
1744 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_dd_dd, b0b1_b0b1, rho0a[1:]) * 8
1745 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud_ud, b0a1_b0a1, rho0b[1:])
1746 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0b1_b0a1, rho0b[1:]) * 2
1747 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0a1_b0b1, rho0b[1:]) * 2
1748 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_dd_dd, b0b1_b0b1, rho0b[1:]) * 4
1749 |     wva *= weight
1750 |     wva[0]*=.5  # v+v.T should be applied in the caller
1751 | 
1752 |     wvb = numpy.zeros((4,ngrid))
1753 |     wvb[0] += numpy.einsum('i,i,i->i', d_d_d, rho1b[0], rho1b[0])
1754 |     wvb[0] += numpy.einsum('i,i,i->i', u_d_d, rho1b[0], rho1a[0]) * 2
1755 |     wvb[0] += numpy.einsum('i,i,i->i', u_u_d, rho1a[0], rho1a[0])
1756 |     wvb[0] += numpy.einsum('i,i->i', d_dd, b1b1) * 2
1757 |     wvb[0] += numpy.einsum('i,i->i', d_ud, b1a1) * 2
1758 |     wvb[0] += numpy.einsum('i,i->i', d_uu, a1a1) * 2
1759 |     wvb[1:] += numpy.einsum('i,i,xi->xi', d_dd, rho1b[0], rho1b[1:]) * 4
1760 |     wvb[1:] += numpy.einsum('i,i,xi->xi', u_dd, rho1a[0], rho1b[1:]) * 4
1761 |     wvb[1:] += numpy.einsum('i,i,xi->xi', d_ud, rho1b[0], rho1a[1:]) * 2
1762 |     wvb[1:] += numpy.einsum('i,i,xi->xi', u_ud, rho1a[0], rho1a[1:]) * 2
1763 |     wvb[1:] += numpy.einsum('i,i,xi->xi', dd_dd, b0b1, rho1b[1:]) * 8
1764 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd, b0b1, rho1a[1:]) * 4
1765 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd, b0a1, rho1b[1:]) * 4
1766 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud, b0a1, rho1a[1:]) * 2
1767 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd, a0b1, rho1b[1:]) * 4
1768 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_dd, a0a1, rho1b[1:]) * 8
1769 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud, a0b1, rho1a[1:]) * 2
1770 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud, a0a1, rho1a[1:]) * 4
1771 |     wvb[1:] += numpy.einsum('i,i,xi->xi', dd_dd, b1b1, rho0b[1:]) * 4
1772 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd, b1a1, rho0b[1:]) * 4
1773 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_dd, a1a1, rho0b[1:]) * 4
1774 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd, b1b1, rho0a[1:]) * 2
1775 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud, b1a1, rho0a[1:]) * 2
1776 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud, a1a1, rho0a[1:]) * 2
1777 |     wvb[0] += numpy.einsum('i,i,i->i', d_d_dd, rho1b[0], b0b1) * 4
1778 |     wvb[0] += numpy.einsum('i,i,i->i', u_d_dd, rho1a[0], b0b1) * 4
1779 |     wvb[0] += numpy.einsum('i,i,i->i', d_d_ud, rho1b[0], b0a1) * 2
1780 |     wvb[0] += numpy.einsum('i,i,i->i', u_d_ud, rho1a[0], b0a1) * 2
1781 |     wvb[0] += numpy.einsum('i,i,i->i', d_d_ud, rho1b[0], a0b1) * 2
1782 |     wvb[0] += numpy.einsum('i,i,i->i', u_d_ud, rho1a[0], a0b1) * 2
1783 |     wvb[0] += numpy.einsum('i,i,i->i', d_d_uu, rho1b[0], a0a1) * 4
1784 |     wvb[0] += numpy.einsum('i,i,i->i', u_d_uu, rho1a[0], a0a1) * 4
1785 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_d_dd, rho1b[0], rho1b[0], rho0b[1:]) * 2
1786 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_d_dd, rho1b[0], rho1a[0], rho0b[1:]) * 4
1787 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_u_dd, rho1a[0], rho1a[0], rho0b[1:]) * 2
1788 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_d_ud, rho1b[0], rho1b[0], rho0b[1:])
1789 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_d_ud, rho1b[0], rho1a[0], rho0b[1:]) * 2
1790 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_u_ud, rho1a[0], rho1a[0], rho0b[1:])
1791 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_dd_dd, rho1b[0], b0b1, rho0b[1:]) * 8
1792 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_dd_dd, rho1a[0], b0b1, rho0b[1:]) * 8
1793 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_dd, rho1b[0], b0a1, rho0b[1:]) * 4
1794 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_dd, rho1a[0], b0a1, rho0b[1:]) * 4
1795 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_dd, rho1b[0], a0b1, rho0b[1:]) * 4
1796 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_dd, rho1a[0], a0b1, rho0b[1:]) * 4
1797 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_dd, rho1b[0], a0a1, rho0b[1:]) * 8
1798 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_dd, rho1a[0], a0a1, rho0b[1:]) * 8
1799 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_dd, rho1b[0], b0b1, rho0a[1:]) * 4
1800 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_dd, rho1a[0], b0b1, rho0a[1:]) * 4
1801 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_ud, rho1b[0], b0a1, rho0a[1:]) * 2
1802 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_ud, rho1a[0], b0a1, rho0a[1:]) * 2
1803 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_ud, rho1b[0], a0b1, rho0a[1:]) * 2
1804 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_ud, rho1a[0], a0b1, rho0a[1:]) * 2
1805 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_ud, rho1b[0], a0a1, rho0a[1:]) * 4
1806 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_ud, rho1a[0], a0a1, rho0a[1:]) * 4
1807 |     wvb[0] += numpy.einsum('i,i->i', d_dd_dd, b0b1_b0b1) * 4
1808 |     wvb[0] += numpy.einsum('i,i->i', d_ud_dd, b0b1_b0a1) * 2
1809 |     wvb[0] += numpy.einsum('i,i->i', d_ud_dd, b0a1_b0b1) * 2
1810 |     wvb[0] += numpy.einsum('i,i->i', d_ud_ud, b0a1_b0a1)
1811 |     wvb[0] += numpy.einsum('i,i->i', d_ud_dd, b0b1_a0b1) * 4
1812 |     wvb[0] += numpy.einsum('i,i->i', d_ud_ud, b0a1_a0b1) * 2
1813 |     wvb[0] += numpy.einsum('i,i->i', d_uu_dd, b0b1_a0a1) * 8
1814 |     wvb[0] += numpy.einsum('i,i->i', d_uu_ud, b0a1_a0a1) * 4
1815 |     wvb[0] += numpy.einsum('i,i->i', d_ud_ud, a0b1_a0b1)
1816 |     wvb[0] += numpy.einsum('i,i->i', d_uu_ud, a0b1_a0a1) * 2
1817 |     wvb[0] += numpy.einsum('i,i->i', d_uu_ud, a0a1_a0b1) * 2
1818 |     wvb[0] += numpy.einsum('i,i->i', d_uu_uu, a0a1_a0a1) * 4
1819 |     wvb[1:] += numpy.einsum('i,i,xi->xi', dd_dd_dd, b0b1_b0b1, rho0b[1:]) * 8
1820 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd_dd, b0b1_b0a1, rho0b[1:]) * 4
1821 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd_dd, b0a1_b0b1, rho0b[1:]) * 4
1822 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0a1_b0a1, rho0b[1:]) * 2
1823 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd_dd, b0b1_b0b1, rho0a[1:]) * 4
1824 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0b1_b0a1, rho0a[1:]) * 2
1825 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0a1_b0b1, rho0a[1:]) * 2
1826 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_ud, b0a1_b0a1, rho0a[1:])
1827 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd_dd, b0b1_a0b1, rho0b[1:]) * 8
1828 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0a1_a0b1, rho0b[1:]) * 4
1829 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_dd_dd, b0b1_a0a1, rho0b[1:]) * 16
1830 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, b0a1_a0a1, rho0b[1:]) * 8
1831 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0b1_a0b1, rho0a[1:]) * 4
1832 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_ud, b0a1_a0b1, rho0a[1:]) * 2
1833 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, b0b1_a0a1, rho0a[1:]) * 8
1834 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, b0a1_a0a1, rho0a[1:]) * 4
1835 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, a0b1_a0b1, rho0b[1:]) * 2
1836 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, a0a1_a0b1, rho0b[1:]) * 4
1837 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, a0b1_a0a1, rho0b[1:]) * 4
1838 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_uu_dd, a0a1_a0a1, rho0b[1:]) * 8
1839 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_ud, a0b1_a0b1, rho0a[1:])
1840 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0a1_a0b1, rho0a[1:]) * 2
1841 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0b1_a0a1, rho0a[1:]) * 2
1842 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_uu_ud, a0a1_a0a1, rho0a[1:]) * 4
1843 |     wvb *= weight
1844 |     wvb[0]*=.5
1845 | 
1846 |     return wva, wvb
1847 | 
1848 | def nr_fxc(mol, grids, xc_code, dm0, dms, spin=0, relativity=0, hermi=0,
1849 |            rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
1850 |     r'''Contract XC kernel matrix with given density matrices
1851 | 
1852 |     ... math::
1853 | 
1854 |             a_{pq} = f_{pq,rs} * x_{rs}
1855 | 
1856 |     '''
1857 |     ni = NumInt()
1858 |     return ni.nr_fxc(mol, grids, xc_code, dm0, dms, spin, relativity,
1859 |                      hermi, rho0, vxc, fxc, max_memory, verbose)
1860 | 
1861 | 
1862 | def cache_xc_kernel(ni, mol, grids, xc_code, mo_coeff, mo_occ, spin=0,
1863 |                     max_memory=2000):
1864 |     '''Compute the 0th order density, Vxc and fxc.  They can be used in TDDFT,
1865 |     DFT hessian module etc.
1866 |     '''
1867 |     xctype = ni._xc_type(xc_code)
1868 |     ao_deriv = 0
1869 |     if xctype == 'GGA':
1870 |         ao_deriv = 1
1871 |     elif xctype == 'NLC':
1872 |         raise NotImplementedError('NLC')
1873 |     elif xctype == 'MGGA':
1874 |         raise NotImplementedError('meta-GGA')
1875 | 
1876 |     if spin == 0:
1877 |         nao = mo_coeff.shape[0]
1878 |         rho = []
1879 |         for ao, mask, weight, coords \
1880 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory=max_memory):
1881 |             rho.append(ni.eval_rho2(mol, ao, mo_coeff, mo_occ, mask, xctype))
1882 |         rho = numpy.hstack(rho)
1883 |     else:
1884 |         nao = mo_coeff[0].shape[0]
1885 |         rhoa = []
1886 |         rhob = []
1887 |         for ao, mask, weight, coords \
1888 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1889 |             rhoa.append(ni.eval_rho2(mol, ao, mo_coeff[0], mo_occ[0], mask, xctype))
1890 |             rhob.append(ni.eval_rho2(mol, ao, mo_coeff[1], mo_occ[1], mask, xctype))
1891 |         rho = (numpy.hstack(rhoa), numpy.hstack(rhob))
1892 |     vxc, fxc = ni.eval_xc(xc_code, rho, spin, 0, 2, 0)[1:3]
1893 |     return rho, vxc, fxc
1894 | 
1895 | def get_rho(ni, mol, dm, grids, max_memory=2000):
1896 |     '''Density in real space
1897 |     '''
1898 |     make_rho, nset, nao = ni._gen_rho_evaluator(mol, dm, 1)
1899 |     assert(nset == 1)
1900 |     rho = numpy.empty(grids.weights.size)
1901 |     p1 = 0
1902 |     for ao, mask, weight, coords \
1903 |             in ni.block_loop(mol, grids, nao, 0, max_memory):
1904 |         p0, p1 = p1, p1 + weight.size
1905 |         rho[p0:p1] = make_rho(0, ao, mask, 'LDA')
1906 |     return rho
1907 | 
1908 | 
1909 | class NumInt(object):
1910 |     def __init__(self):
1911 |         self.libxc = libxc
1912 | 
1913 |     @lib.with_doc(nr_vxc.__doc__)
1914 |     def nr_vxc(self, mol, grids, xc_code, dms, spin=0, relativity=0, hermi=0,
1915 |                max_memory=2000, verbose=None):
1916 |         if spin == 0:
1917 |             return self.nr_rks(mol, grids, xc_code, dms, relativity, hermi,
1918 |                                max_memory, verbose)
1919 |         else:
1920 |             return self.nr_uks(mol, grids, xc_code, dms, relativity, hermi,
1921 |                                max_memory, verbose)
1922 | 
1923 |     @lib.with_doc(nr_fxc.__doc__)
1924 |     def nr_fxc(self, mol, grids, xc_code, dm0, dms, spin=0, relativity=0, hermi=0,
1925 |                rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
1926 |         if spin == 0:
1927 |             return self.nr_rks_fxc(mol, grids, xc_code, dm0, dms, relativity,
1928 |                                    hermi, rho0, vxc, fxc, max_memory, verbose)
1929 |         else:
1930 |             return self.nr_uks_fxc(mol, grids, xc_code, dm0, dms, relativity,
1931 |                                    hermi, rho0, vxc, fxc, max_memory, verbose)
1932 | 
1933 |     nr_rks = nr_rks
1934 |     nr_uks = nr_uks
1935 |     nr_rks_fxc = nr_rks_fxc
1936 |     nr_uks_fxc = nr_uks_fxc
1937 |     cache_xc_kernel  = cache_xc_kernel
1938 |     get_rho = get_rho
1939 | 
1940 |     @lib.with_doc(eval_ao.__doc__)
1941 |     def eval_ao(self, mol, coords, deriv=0, shls_slice=None,
1942 |                 non0tab=None, out=None, verbose=None):
1943 |         return eval_ao(mol, coords, deriv, shls_slice, non0tab, out, verbose)
1944 | 
1945 |     @lib.with_doc(make_mask.__doc__)
1946 |     def make_mask(self, mol, coords, relativity=0, shls_slice=None,
1947 |                   verbose=None):
1948 |         return make_mask(mol, coords, relativity, shls_slice, verbose)
1949 | 
1950 |     @lib.with_doc(eval_rho2.__doc__)
1951 |     def eval_rho2(self, mol, ao, mo_coeff, mo_occ, non0tab=None, xctype='LDA',
1952 |                   verbose=None):
1953 |         return eval_rho2(mol, ao, mo_coeff, mo_occ, non0tab, xctype, verbose)
1954 | 
1955 |     @lib.with_doc(eval_rho.__doc__)
1956 |     def eval_rho(self, mol, ao, dm, non0tab=None, xctype='LDA', hermi=0, verbose=None):
1957 |         return eval_rho(mol, ao, dm, non0tab, xctype, hermi, verbose)
1958 | 
1959 |     def block_loop(self, mol, grids, nao, deriv=0, max_memory=2000,
1960 |                    non0tab=None, blksize=None, buf=None):
1961 |         '''Define this macro to loop over grids by blocks.
1962 |         '''
1963 |         if grids.coords is None:
1964 |             grids.build(with_non0tab=True)
1965 |         ngrids = grids.coords.shape[0]
1966 |         comp = (deriv+1)*(deriv+2)*(deriv+3)//6
1967 | # NOTE to index grids.non0tab, the blksize needs to be the integer multiplier of BLKSIZE
1968 |         if blksize is None:
1969 |             blksize = int(max_memory*1e6/(comp*2*nao*8*BLKSIZE))*BLKSIZE
1970 |             blksize = max(BLKSIZE, min(blksize, ngrids, BLKSIZE*1200))
1971 |         if non0tab is None:
1972 |             non0tab = grids.non0tab
1973 |         if non0tab is None:
1974 |             non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
1975 |                                  dtype=numpy.uint8)
1976 |         if buf is None:
1977 |             buf = numpy.empty((comp,blksize,nao))
1978 |         for ip0 in range(0, ngrids, blksize):
1979 |             ip1 = min(ngrids, ip0+blksize)
1980 |             coords = grids.coords[ip0:ip1]
1981 |             weight = grids.weights[ip0:ip1]
1982 |             non0 = non0tab[ip0//BLKSIZE:]
1983 |             ao = self.eval_ao(mol, coords, deriv=deriv, non0tab=non0, out=buf)
1984 |             yield ao, non0, weight, coords
1985 | 
1986 |     def _gen_rho_evaluator(self, mol, dms, hermi=0):
1987 |         if getattr(dms, 'mo_coeff', None) is not None:
1988 | #TODO: test whether dm.mo_coeff matching dm
1989 |             mo_coeff = dms.mo_coeff
1990 |             mo_occ = dms.mo_occ
1991 |             if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
1992 |                 mo_coeff = [mo_coeff]
1993 |                 mo_occ = [mo_occ]
1994 |             nao = mo_coeff[0].shape[0]
1995 |             ndms = len(mo_occ)
1996 |             def make_rho(idm, ao, non0tab, xctype):
1997 |                 return self.eval_rho2(mol, ao, mo_coeff[idm], mo_occ[idm],
1998 |                                       non0tab, xctype)
1999 |         else:
2000 |             if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
2001 |                 dms = [dms]
2002 |             if not hermi:
2003 | # For eval_rho when xctype==GGA, which requires hermitian DMs
2004 |                 dms = [(dm+dm.conj().T)*.5 for dm in dms]
2005 |             nao = dms[0].shape[0]
2006 |             ndms = len(dms)
2007 |             def make_rho(idm, ao, non0tab, xctype):
2008 |                 return self.eval_rho(mol, ao, dms[idm], non0tab, xctype, hermi=1)
2009 |         return make_rho, ndms, nao
2010 | 
2011 | ####################
2012 | # Overwrite following functions to use custom XC functional
2013 | 
2014 |     def hybrid_coeff(self, xc_code, spin=0):
2015 |         return self.libxc.hybrid_coeff(xc_code, spin)
2016 | 
2017 |     def nlc_coeff(self, xc_code):
2018 |         return self.libxc.nlc_coeff(xc_code)
2019 | 
2020 |     def rsh_coeff(self, xc_code):
2021 |         return self.libxc.rsh_coeff(xc_code)
2022 | 
2023 |     def eval_xc(self, xc_code, rho, spin=0, relativity=0, deriv=1, verbose=None):
2024 |         return self.libxc.eval_xc(xc_code, rho, spin, relativity, deriv, verbose)
2025 |     eval_xc.__doc__ = libxc.eval_xc.__doc__
2026 | 
2027 |     def _xc_type(self, xc_code):
2028 |         return self.libxc.xc_type(xc_code)
2029 | 
2030 |     def rsh_and_hybrid_coeff(self, xc_code, spin=0):
2031 |         '''Range-separated parameter and HF exchange components: omega, alpha, beta
2032 | 
2033 |         Exc_RSH = c_SR * SR_HFX + c_LR * LR_HFX + (1-c_SR) * Ex_SR + (1-c_LR) * Ex_LR + Ec
2034 |                 = alpha * HFX + beta * SR_HFX + (1-c_SR) * Ex_SR + (1-c_LR) * Ex_LR + Ec
2035 |                 = alpha * LR_HFX + hyb * SR_HFX + (1-c_SR) * Ex_SR + (1-c_LR) * Ex_LR + Ec
2036 | 
2037 |         SR_HFX = < pi | e^{-omega r_{12}}/r_{12} | iq >
2038 |         LR_HFX = < pi | (1-e^{-omega r_{12}})/r_{12} | iq >
2039 |         alpha = c_LR
2040 |         beta = c_SR - c_LR
2041 |         '''
2042 |         omega, alpha, beta = self.rsh_coeff(xc_code)
2043 |         if abs(omega) > 1e-10:
2044 |             hyb = alpha + beta
2045 |         else:
2046 |             hyb = self.hybrid_coeff(xc_code, spin)
2047 |         return omega, alpha, hyb
2048 | _NumInt = NumInt
2049 | 
2050 | 
2051 | if __name__ == '__main__':
2052 |     import time
2053 |     from pyscf import gto
2054 |     from pyscf import dft
2055 | 
2056 |     mol = gto.M(
2057 |         atom = [
2058 |         ["O" , (0. , 0.     , 0.)],
2059 |         [1   , (0. , -0.757 , 0.587)],
2060 |         [1   , (0. , 0.757  , 0.587)] ],
2061 |         basis = '6311g**',)
2062 |     mf = dft.RKS(mol)
2063 |     mf.grids.atom_grid = {"H": (30, 194), "O": (30, 194),}
2064 |     mf.grids.prune = None
2065 |     mf.grids.build()
2066 |     dm = mf.get_init_guess(key='minao')
2067 | 
2068 |     numpy.random.seed(1)
2069 |     dm1 = numpy.random.random((dm.shape))
2070 |     dm1 = lib.hermi_triu(dm1)
2071 |     print(time.clock())
2072 |     res = mf._numint.nr_vxc(mol, mf.grids, mf.xc, dm1, spin=0)
2073 |     print(res[1] - -37.084047825971282)
2074 |     res = mf._numint.nr_vxc(mol, mf.grids, mf.xc, (dm1,dm1), spin=1)
2075 |     print(res[1] - -92.436362308687094)
2076 |     res = mf._numint.nr_vxc(mol, mf.grids, mf.xc, dm, spin=0)
2077 |     print(res[1] - -8.6313329288394947)
2078 |     res = mf._numint.nr_vxc(mol, mf.grids, mf.xc, (dm,dm), spin=1)
2079 |     print(res[1] - -21.520301399504582)
2080 |     print(time.clock())
2081 | 


--------------------------------------------------------------------------------
/numint_1.5.py:
--------------------------------------------------------------------------------
   1 | #!/usr/bin/env python
   2 | 
   3 | #######################
   4 | #Changed by Ryo Nagai, March 2019.
   5 | #For PySCF version 1.5
   6 | #This software includes the work that is distributed in the Apache License 2.0.
   7 | #######################
   8 | 
   9 | # Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
  10 | #
  11 | # Licensed under the Apache License, Version 2.0 (the "License");
  12 | # you may not use this file except in compliance with the License.
  13 | # You may obtain a copy of the License at
  14 | #
  15 | #     http://www.apache.org/licenses/LICENSE-2.0
  16 | #
  17 | # Unless required by applicable law or agreed to in writing, software
  18 | # distributed under the License is distributed on an "AS IS" BASIS,
  19 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  20 | # See the License for the specific language governing permissions and
  21 | # limitations under the License.
  22 | #
  23 | # Author: Qiming Sun <osirpt.sun@gmail.com>
  24 | #
  25 | 
  26 | import warnings
  27 | import ctypes
  28 | import numpy
  29 | import scipy.linalg
  30 | from pyscf import lib
  31 | from pyscf.lib import logger
  32 | try:
  33 |     from pyscf.dft import libxc
  34 | except (ImportError, OSError):
  35 |     try:
  36 |         from pyscf.dft import xcfun
  37 |         libxc = xcfun
  38 |     except (ImportError, OSError):
  39 |         import warnings
  40 |         warnings.warn('XC functional libraries (libxc or XCfun) are not available.')
  41 |         from pyscf.dft import xc
  42 |         libxc = xc
  43 | 
  44 | from pyscf.dft.gen_grid import make_mask, BLKSIZE
  45 | from pyscf import __config__
  46 | 
  47 | libdft = lib.load_library('libdft')
  48 | OCCDROP = getattr(__config__, 'dft_numint_OCCDROP', 1e-12)
  49 | # The system size above which to consider the sparsity of the density matrix.
  50 | # If the number of AOs in the system is less than this value, all tensors are
  51 | # treated as dense quantities and contracted by dgemm directly.
  52 | SWITCH_SIZE = getattr(__config__, 'dft_numint_SWITCH_SIZE', 800)
  53 | 
  54 | def eval_ao(mol, coords, deriv=0, shls_slice=None,
  55 |             non0tab=None, out=None, verbose=None):
  56 |     '''Evaluate AO function value on the given grids.
  57 |     Args:
  58 |         mol : an instance of :class:`Mole`
  59 |         coords : 2D array, shape (N,3)
  60 |             The coordinates of the grids.
  61 |     Kwargs:
  62 |         deriv : int
  63 |             AO derivative order.  It affects the shape of the return array.
  64 |             If deriv=0, the returned AO values are stored in a (N,nao) array.
  65 |             Otherwise the AO values are stored in an array of shape (M,N,nao).
  66 |             Here N is the number of grids, nao is the number of AO functions,
  67 |             M is the size associated to the derivative deriv.
  68 |         relativity : bool
  69 |             No effects.
  70 |         shls_slice : 2-element list
  71 |             (shl_start, shl_end).
  72 |             If given, only part of AOs (shl_start <= shell_id < shl_end) are
  73 |             evaluated.  By default, all shells defined in mol will be evaluated.
  74 |         non0tab : 2D bool array
  75 |             mask array to indicate whether the AO values are zero.  The mask
  76 |             array can be obtained by calling :func:`make_mask`
  77 |         out : ndarray
  78 |             If provided, results are written into this array.
  79 |         verbose : int or object of :class:`Logger`
  80 |             No effects.
  81 |     Returns:
  82 |         2D array of shape (N,nao) for AO values if deriv = 0.
  83 |         Or 3D array of shape (:,N,nao) for AO values and AO derivatives if deriv > 0.
  84 |         In the 3D array, the first (N,nao) elements are the AO values,
  85 |         followed by (3,N,nao) for x,y,z compoents;
  86 |         Then 2nd derivatives (6,N,nao) for xx, xy, xz, yy, yz, zz;
  87 |         Then 3rd derivatives (10,N,nao) for xxx, xxy, xxz, xyy, xyz, xzz, yyy, yyz, yzz, zzz;
  88 |         ...
  89 |     Examples:
  90 |     >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz')
  91 |     >>> coords = numpy.random.random((100,3))  # 100 random points
  92 |     >>> ao_value = eval_ao(mol, coords)
  93 |     >>> print(ao_value.shape)
  94 |     (100, 24)
  95 |     >>> ao_value = eval_ao(mol, coords, deriv=1, shls_slice=(1,4))
  96 |     >>> print(ao_value.shape)
  97 |     (4, 100, 7)
  98 |     >>> ao_value = eval_ao(mol, coords, deriv=2, shls_slice=(1,4))
  99 |     >>> print(ao_value.shape)
 100 |     (10, 100, 7)
 101 |     '''
 102 |     comp = (deriv+1)*(deriv+2)*(deriv+3)//6
 103 |     if mol.cart:
 104 |         feval = 'GTOval_cart_deriv%d' % deriv
 105 |     else:
 106 |         feval = 'GTOval_sph_deriv%d' % deriv
 107 |     return mol.eval_gto(feval, coords, comp, shls_slice, non0tab, out=out)
 108 | 
 109 | #TODO: \nabla^2 rho and tau = 1/2 (\nabla f)^2
 110 | def eval_rho(mol, ao, dm, non0tab=None, xctype='LDA', hermi=0, verbose=None):
 111 |     r'''Calculate the electron density for LDA functional, and the density
 112 |     derivatives for GGA functional.
 113 |     Args:
 114 |         mol : an instance of :class:`Mole`
 115 |         ao : 2D array of shape (N,nao) for LDA, 3D array of shape (4,N,nao) for GGA
 116 |             or (5,N,nao) for meta-GGA.  N is the number of grids, nao is the
 117 |             number of AO functions.  If xctype is GGA, ao[0] is AO value
 118 |             and ao[1:3] are the AO gradients.  If xctype is meta-GGA, ao[4:10]
 119 |             are second derivatives of ao values.
 120 |         dm : 2D array
 121 |             Density matrix
 122 |     Kwargs:
 123 |         non0tab : 2D bool array
 124 |             mask array to indicate whether the AO values are zero.  The mask
 125 |             array can be obtained by calling :func:`make_mask`
 126 |         xctype : str
 127 |             LDA/GGA/mGGA.  It affects the shape of the return density.
 128 |         hermi : bool
 129 |             dm is hermitian or not
 130 |         verbose : int or object of :class:`Logger`
 131 |             No effects.
 132 |     Returns:
 133 |         1D array of size N to store electron density if xctype = LDA;  2D array
 134 |         of (4,N) to store density and "density derivatives" for x,y,z components
 135 |         if xctype = GGA;  (6,N) array for meta-GGA, where last two rows are
 136 |         \nabla^2 rho and tau = 1/2(\nabla f)^2
 137 |     Examples:
 138 |     >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz')
 139 |     >>> coords = numpy.random.random((100,3))  # 100 random points
 140 |     >>> ao_value = eval_ao(mol, coords, deriv=0)
 141 |     >>> dm = numpy.random.random((mol.nao_nr(),mol.nao_nr()))
 142 |     >>> dm = dm + dm.T
 143 |     >>> rho, dx_rho, dy_rho, dz_rho = eval_rho(mol, ao, dm, xctype='LDA')
 144 |     '''
 145 |     xctype = xctype.upper()
 146 |     if xctype == 'LDA' or xctype == 'HF':
 147 |         ngrids, nao = ao.shape
 148 |     else:
 149 |         ngrids, nao = ao[0].shape
 150 | 
 151 |     if non0tab is None:
 152 |         non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
 153 |                              dtype=numpy.uint8)
 154 |     if not hermi:
 155 |         # (D + D.T)/2 because eval_rho computes 2*(|\nabla i> D_ij <j|) instead of
 156 |         # |\nabla i> D_ij <j| + |i> D_ij <\nabla j| for efficiency
 157 |         dm = (dm + dm.conj().T) * .5
 158 | 
 159 |     shls_slice = (0, mol.nbas)
 160 |     ao_loc = mol.ao_loc_nr()
 161 |     if xctype == 'LDA' or xctype == 'HF':
 162 |         c0 = _dot_ao_dm(mol, ao, dm, non0tab, shls_slice, ao_loc)
 163 |         rho = numpy.einsum('pi,pi->p', ao, c0)
 164 |     elif xctype in ('GGA', 'NLC'):
 165 |         rho = numpy.empty((4,ngrids))
 166 |         c0 = _dot_ao_dm(mol, ao[0], dm, non0tab, shls_slice, ao_loc)
 167 |         rho[0] = numpy.einsum('pi,pi->p', c0, ao[0])
 168 |         for i in range(1, 4):
 169 |             rho[i] = numpy.einsum('pi,pi->p', c0, ao[i])
 170 |             rho[i] *= 2 # *2 for +c.c. in the next two lines
 171 |             #c1 = _dot_ao_dm(mol, ao[i], dm, non0tab, shls_slice, ao_loc)
 172 |             #rho[i] += numpy.einsum('pi,pi->p', c1, ao[0])
 173 |     else: # meta-GGA
 174 |         # rho[4] = \nabla^2 rho, rho[5] = 1/2 |nabla f|^2
 175 |         rho = numpy.empty((6,ngrids))
 176 |         c0 = _dot_ao_dm(mol, ao[0], dm, non0tab, shls_slice, ao_loc)
 177 |         rho[0] = numpy.einsum('pi,pi->p', ao[0], c0)
 178 |         rho[5] = 0
 179 |         for i in range(1, 4):
 180 |             rho[i] = numpy.einsum('pi,pi->p', c0, ao[i]) * 2 # *2 for +c.c.
 181 |             c1 = _dot_ao_dm(mol, ao[i], dm.T, non0tab, shls_slice, ao_loc)
 182 |             rho[5] += numpy.einsum('pi,pi->p', c1, ao[i])
 183 |         XX, YY, ZZ = 4, 7, 9
 184 |         ao2 = ao[XX] + ao[YY] + ao[ZZ]
 185 |         rho[4] = numpy.einsum('pi,pi->p', c0, ao2)
 186 |         rho[4] += rho[5]
 187 |         rho[4] *= 2
 188 |         rho[5] *= .5
 189 |     return rho
 190 | 
 191 | def eval_rho2(mol, ao, mo_coeff, mo_occ, non0tab=None, xctype='LDA',
 192 |               verbose=None):
 193 |     r'''Calculate the electron density for LDA functional, and the density
 194 |     derivatives for GGA functional.  This function has the same functionality
 195 |     as :func:`eval_rho` except that the density are evaluated based on orbital
 196 |     coefficients and orbital occupancy.  It is more efficient than
 197 |     :func:`eval_rho` in most scenario.
 198 |     Args:
 199 |         mol : an instance of :class:`Mole`
 200 |         ao : 2D array of shape (N,nao) for LDA, 3D array of shape (4,N,nao) for GGA
 201 |             or (5,N,nao) for meta-GGA.  N is the number of grids, nao is the
 202 |             number of AO functions.  If xctype is GGA, ao[0] is AO value
 203 |             and ao[1:3] are the AO gradients.  If xctype is meta-GGA, ao[4:10]
 204 |             are second derivatives of ao values.
 205 |         dm : 2D array
 206 |             Density matrix
 207 |     Kwargs:
 208 |         non0tab : 2D bool array
 209 |             mask array to indicate whether the AO values are zero.  The mask
 210 |             array can be obtained by calling :func:`make_mask`
 211 |         xctype : str
 212 |             LDA/GGA/mGGA.  It affects the shape of the return density.
 213 |         verbose : int or object of :class:`Logger`
 214 |             No effects.
 215 |     Returns:
 216 |         1D array of size N to store electron density if xctype = LDA;  2D array
 217 |         of (4,N) to store density and "density derivatives" for x,y,z components
 218 |         if xctype = GGA;  (6,N) array for meta-GGA, where last two rows are
 219 |         \nabla^2 rho and tau = 1/2(\nabla f)^2
 220 |     '''
 221 |     xctype = xctype.upper()
 222 |     if xctype == 'LDA' or xctype == 'HF':
 223 |         ngrids, nao = ao.shape
 224 |     else:
 225 |         ngrids, nao = ao[0].shape
 226 | 
 227 |     if non0tab is None:
 228 |         non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
 229 |                              dtype=numpy.uint8)
 230 |     shls_slice = (0, mol.nbas)
 231 |     ao_loc = mol.ao_loc_nr()
 232 |     pos = mo_occ > OCCDROP
 233 |     if pos.sum() > 0:
 234 |         cpos = numpy.einsum('ij,j->ij', mo_coeff[:,pos], numpy.sqrt(mo_occ[pos]))
 235 |         if xctype == 'LDA' or xctype == 'HF':
 236 |             c0 = _dot_ao_dm(mol, ao, cpos, non0tab, shls_slice, ao_loc)
 237 |             rho = numpy.einsum('pi,pi->p', c0, c0)
 238 |         elif xctype in ('GGA', 'NLC'):
 239 |             rho = numpy.empty((4,ngrids))
 240 |             c0 = _dot_ao_dm(mol, ao[0], cpos, non0tab, shls_slice, ao_loc)
 241 |             rho[0] = numpy.einsum('pi,pi->p', c0, c0)
 242 |             for i in range(1, 4):
 243 |                 c1 = _dot_ao_dm(mol, ao[i], cpos, non0tab, shls_slice, ao_loc)
 244 |                 rho[i] = numpy.einsum('pi,pi->p', c0, c1) * 2 # *2 for +c.c.
 245 |         else: # meta-GGA
 246 |             # rho[4] = \nabla^2 rho, rho[5] = 1/2 |nabla f|^2
 247 |             rho = numpy.empty((6,ngrids))
 248 |             c0 = _dot_ao_dm(mol, ao[0], cpos, non0tab, shls_slice, ao_loc)
 249 |             rho[0] = numpy.einsum('pi,pi->p', c0, c0)
 250 |             rho[5] = 0
 251 |             for i in range(1, 4):
 252 |                 c1 = _dot_ao_dm(mol, ao[i], cpos, non0tab, shls_slice, ao_loc)
 253 |                 rho[i] = numpy.einsum('pi,pi->p', c0, c1) * 2 # *2 for +c.c.
 254 |                 rho[5] += numpy.einsum('pi,pi->p', c1, c1)
 255 |             XX, YY, ZZ = 4, 7, 9
 256 |             ao2 = ao[XX] + ao[YY] + ao[ZZ]
 257 |             c1 = _dot_ao_dm(mol, ao2, cpos, non0tab, shls_slice, ao_loc)
 258 |             rho[4] = numpy.einsum('pi,pi->p', c0, c1)
 259 |             rho[4] += rho[5]
 260 |             rho[4] *= 2
 261 | 
 262 |             rho[5] *= .5
 263 |     else:
 264 |         if xctype == 'LDA' or xctype == 'HF':
 265 |             rho = numpy.zeros(ngrids)
 266 |         elif xctype in ('GGA', 'NLC'):
 267 |             rho = numpy.zeros((4,ngrids))
 268 |         else:
 269 |             rho = numpy.zeros((6,ngrids))
 270 | 
 271 |     neg = mo_occ < -OCCDROP
 272 |     if neg.sum() > 0:
 273 |         cneg = numpy.einsum('ij,j->ij', mo_coeff[:,neg], numpy.sqrt(-mo_occ[neg]))
 274 |         if xctype == 'LDA' or xctype == 'HF':
 275 |             c0 = _dot_ao_dm(mol, ao, cneg, non0tab, shls_slice, ao_loc)
 276 |             rho -= numpy.einsum('pi,pi->p', c0, c0)
 277 |         elif xctype == 'GGA':
 278 |             c0 = _dot_ao_dm(mol, ao[0], cneg, non0tab, shls_slice, ao_loc)
 279 |             rho[0] -= numpy.einsum('pi,pi->p', c0, c0)
 280 |             for i in range(1, 4):
 281 |                 c1 = _dot_ao_dm(mol, ao[i], cneg, non0tab, shls_slice, ao_loc)
 282 |                 rho[i] -= numpy.einsum('pi,pi->p', c0, c1) * 2 # *2 for +c.c.
 283 |         else:
 284 |             c0 = _dot_ao_dm(mol, ao[0], cneg, non0tab, shls_slice, ao_loc)
 285 |             rho[0] -= numpy.einsum('pi,pi->p', c0, c0)
 286 |             rho5 = 0
 287 |             for i in range(1, 4):
 288 |                 c1 = _dot_ao_dm(mol, ao[i], cneg, non0tab, shls_slice, ao_loc)
 289 |                 rho[i] -= numpy.einsum('pi,pi->p', c0, c1) * 2 # *2 for +c.c.
 290 |                 rho5 += numpy.einsum('pi,pi->p', c1, c1)
 291 |             XX, YY, ZZ = 4, 7, 9
 292 |             ao2 = ao[XX] + ao[YY] + ao[ZZ]
 293 |             c1 = _dot_ao_dm(mol, ao2, cneg, non0tab, shls_slice, ao_loc)
 294 |             rho[4] -= numpy.einsum('pi,pi->p', c0, c1) * 2
 295 |             rho[4] -= rho5 * 2
 296 | 
 297 |             rho[5] -= rho5 * .5
 298 |     return rho
 299 | 
 300 | def _vv10nlc(rho,coords,vvrho,vvweight,vvcoords,nlc_pars):
 301 |     thresh=1e-8
 302 | 
 303 |     #output
 304 |     exc=numpy.zeros(rho[0,:].size)
 305 |     vxc=numpy.zeros([2,rho[0,:].size])
 306 | 
 307 |     #outer grid needs threshing
 308 |     threshind=rho[0,:]>=thresh
 309 |     coords=coords[threshind]
 310 |     R=rho[0,:][threshind]
 311 |     Gx=rho[1,:][threshind]
 312 |     Gy=rho[2,:][threshind]
 313 |     Gz=rho[3,:][threshind]
 314 |     G=Gx**2.+Gy**2.+Gz**2.
 315 | 
 316 |     #threshed output
 317 |     excthresh=numpy.zeros(R.size)
 318 |     vxcthresh=numpy.zeros([2,R.size])
 319 | 
 320 |     #inner grid needs threshing
 321 |     innerthreshind=vvrho[0,:]>=thresh
 322 |     vvcoords=vvcoords[innerthreshind]
 323 |     vvweight=vvweight[innerthreshind]
 324 |     Rp=vvrho[0,:][innerthreshind]
 325 |     RpW=Rp*vvweight
 326 |     Gxp=vvrho[1,:][innerthreshind]
 327 |     Gyp=vvrho[2,:][innerthreshind]
 328 |     Gzp=vvrho[3,:][innerthreshind]
 329 |     Gp=Gxp**2.+Gyp**2.+Gzp**2.
 330 | 
 331 |     #constants and parameters
 332 |     Pi=numpy.pi
 333 |     Pi43=4.*Pi/3.
 334 |     Bvv, Cvv = nlc_pars
 335 |     Kvv=Bvv*1.5*Pi*((9.*Pi)**(-1./6.))
 336 |     Beta=((3./(Bvv*Bvv))**(0.75))/32.
 337 | 
 338 |     #inner grid
 339 |     W0p=Gp/(Rp*Rp)
 340 |     W0p=Cvv*W0p*W0p
 341 |     W0p=(W0p+Pi43*Rp)**0.5
 342 |     Kp=Kvv*(Rp**(1./6.))
 343 | 
 344 |     #outer grid
 345 |     W0tmp=G/(R**2)
 346 |     W0tmp=Cvv*W0tmp*W0tmp
 347 |     W0=(W0tmp+Pi43*R)**0.5
 348 |     dW0dR=(0.5*Pi43*R-2.*W0tmp)/W0
 349 |     dW0dG=W0tmp*R/(G*W0)
 350 |     K=Kvv*(R**(1./6.))
 351 |     dKdR=(1./6.)*K
 352 | 
 353 |     for i in range(R.size):
 354 |         DX=vvcoords[:,0]-coords[i,0]
 355 |         DY=vvcoords[:,1]-coords[i,1]
 356 |         DZ=vvcoords[:,2]-coords[i,2]
 357 |         R2=DX*DX+DY*DY+DZ*DZ
 358 |         gp=R2*W0p+Kp
 359 |         g=R2*W0[i]+K[i]
 360 |         gt=g+gp
 361 |         T=RpW/(g*gp*gt)
 362 |         F=numpy.sum(T)
 363 |         T*=(1./g+1./gt)
 364 |         U=numpy.sum(T)
 365 |         W=numpy.sum(T*R2)
 366 |         F*=-1.5
 367 |         #excthresh is multiplied by Rho later
 368 |         excthresh[i]=Beta+0.5*F
 369 |         vxcthresh[0,i]=Beta+F+1.5*(U*dKdR[i]+W*dW0dR[i])
 370 |         vxcthresh[1,i]=1.5*W*dW0dG[i]
 371 |     exc[threshind]=excthresh
 372 |     vxc[0,threshind]=vxcthresh[0,:]
 373 |     vxc[1,threshind]=vxcthresh[1,:]
 374 | 
 375 |     return exc,vxc
 376 | 
 377 | def eval_mat(mol, ao, weight, rho, vxc,
 378 |              non0tab=None, xctype='LDA', spin=0, verbose=None):
 379 |     r'''Calculate XC potential matrix.
 380 |     Args:
 381 |         mol : an instance of :class:`Mole`
 382 |         ao : ([4/10,] ngrids, nao) ndarray
 383 |             2D array of shape (N,nao) for LDA,
 384 |             3D array of shape (4,N,nao) for GGA
 385 |             or (10,N,nao) for meta-GGA.
 386 |             N is the number of grids, nao is the number of AO functions.
 387 |             If xctype is GGA, ao[0] is AO value and ao[1:3] are the real space
 388 |             gradients.  If xctype is meta-GGA, ao[4:10] are second derivatives
 389 |             of ao values.
 390 |         weight : 1D array
 391 |             Integral weights on grids.
 392 |         rho : ([4/6,] ngrids) ndarray
 393 |             Shape of ((*,N)) for electron density (and derivatives) if spin = 0;
 394 |             Shape of ((*,N),(*,N)) for alpha/beta electron density (and derivatives) if spin > 0;
 395 |             where N is number of grids.
 396 |             rho (*,N) are ordered as (den,grad_x,grad_y,grad_z,laplacian,tau)
 397 |             where grad_x = d/dx den, laplacian = \nabla^2 den, tau = 1/2(\nabla f)^2
 398 |             In spin unrestricted case,
 399 |             rho is ((den_u,grad_xu,grad_yu,grad_zu,laplacian_u,tau_u)
 400 |                     (den_d,grad_xd,grad_yd,grad_zd,laplacian_d,tau_d))
 401 |         vxc : ([4,] ngrids) ndarray
 402 |             XC potential value on each grid = (vrho, vsigma, vlapl, vtau)
 403 |             vsigma is GGA potential value on each grid.
 404 |             If the kwarg spin != 0, a list [vsigma_uu,vsigma_ud] is required.
 405 |     Kwargs:
 406 |         xctype : str
 407 |             LDA/GGA/mGGA.  It affects the shape of `ao` and `rho`
 408 |         non0tab : 2D bool array
 409 |             mask array to indicate whether the AO values are zero.  The mask
 410 |             array can be obtained by calling :func:`make_mask`
 411 |         spin : int
 412 |             If not 0, the returned matrix is the Vxc matrix of alpha-spin.  It
 413 |             is computed with the spin non-degenerated UKS formula.
 414 |     Returns:
 415 |         XC potential matrix in 2D array of shape (nao,nao) where nao is the
 416 |         number of AO functions.
 417 |     '''
 418 |     xctype = xctype.upper()
 419 |     if xctype == 'LDA' or xctype == 'HF':
 420 |         ngrids, nao = ao.shape
 421 |     else:
 422 |         ngrids, nao = ao[0].shape
 423 | 
 424 |     if non0tab is None:
 425 |         non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
 426 |                              dtype=numpy.uint8)
 427 |     shls_slice = (0, mol.nbas)
 428 |     ao_loc = mol.ao_loc_nr()
 429 |     transpose_for_uks = False
 430 |     if xctype == 'LDA' or xctype == 'HF':
 431 |         if not isinstance(vxc, numpy.ndarray) or vxc.ndim == 2:
 432 |             vrho = vxc[0]
 433 |         else:
 434 |             vrho = vxc
 435 |         # *.5 because return mat + mat.T
 436 |         aow = numpy.empty_like(ao)
 437 |         aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho, out=aow)
 438 |         mat = _dot_ao_ao(mol, ao, aow, non0tab, shls_slice, ao_loc)
 439 |     else:
 440 |         #wv = weight * vsigma * 2
 441 |         #aow  = numpy.einsum('pi,p->pi', ao[1], rho[1]*wv)
 442 |         #aow += numpy.einsum('pi,p->pi', ao[2], rho[2]*wv)
 443 |         #aow += numpy.einsum('pi,p->pi', ao[3], rho[3]*wv)
 444 |         #aow += numpy.einsum('pi,p->pi', ao[0], .5*weight*vrho)
 445 |         vrho, vsigma = vxc[:2]
 446 |         wv = numpy.empty((4,ngrids))
 447 |         if spin == 0:
 448 |             assert(vsigma is not None and rho.ndim==2)
 449 |             wv[0]  = weight * vrho * .5
 450 |             wv[1:4] = rho[1:4] * (weight * vsigma * 2)
 451 |         else:
 452 |             rho_a, rho_b = rho
 453 |             wv[0]  = weight * vrho * .5
 454 |             try:
 455 |                 wv[1:4] = rho_a[1:4] * (weight * vsigma[0] * 2)  # sigma_uu
 456 |                 wv[1:4]+= rho_b[1:4] * (weight * vsigma[1])      # sigma_ud
 457 |             except ValueError:
 458 |                 warnings.warn('Note the output of libxc.eval_xc cannot be '
 459 |                               'directly used in eval_mat.\nvsigma from eval_xc '
 460 |                               'should be restructured as '
 461 |                               '(vsigma[:,0],vsigma[:,1])\n')
 462 |                 transpose_for_uks = True
 463 |                 vsigma = vsigma.T
 464 |                 wv[1:4] = rho_a[1:4] * (weight * vsigma[0] * 2)  # sigma_uu
 465 |                 wv[1:4]+= rho_b[1:4] * (weight * vsigma[1])      # sigma_ud
 466 |         aow = numpy.empty_like(ao[0])
 467 |         aow = numpy.einsum('npi,np->pi', ao[:4], wv, out=aow)
 468 |         mat = _dot_ao_ao(mol, ao[0], aow, non0tab, shls_slice, ao_loc)
 469 | 
 470 | # JCP, 138, 244108
 471 | # JCP, 112, 7002
 472 |     if xctype == 'MGGA':
 473 |         vlapl, vtau = vxc[2:]
 474 | 
 475 |         if vlapl is None:
 476 |             vlapl = 0
 477 |         else:
 478 |             if spin != 0:
 479 |                 if transpose_for_uks:
 480 |                     vlapl = vlapl.T
 481 |                 vlapl = vlapl[0]
 482 |             XX, YY, ZZ = 4, 7, 9
 483 |             ao2 = ao[XX] + ao[YY] + ao[ZZ]
 484 |             aow = numpy.einsum('pi,p->pi', ao2, .5 * weight * vlapl, out=aow)
 485 |             mat += _dot_ao_ao(mol, ao[0], aow, non0tab, shls_slice, ao_loc)
 486 | 
 487 |         if spin != 0:
 488 |             if transpose_for_uks:
 489 |                 vtau = vtau.T
 490 |             vtau = vtau[0]
 491 |         wv = weight * (.25*vtau + vlapl)
 492 |         aow = numpy.einsum('pi,p->pi', ao[1], wv, out=aow)
 493 |         mat += _dot_ao_ao(mol, ao[1], aow, non0tab, shls_slice, ao_loc)
 494 |         aow = numpy.einsum('pi,p->pi', ao[2], wv, out=aow)
 495 |         mat += _dot_ao_ao(mol, ao[2], aow, non0tab, shls_slice, ao_loc)
 496 |         aow = numpy.einsum('pi,p->pi', ao[3], wv, out=aow)
 497 |         mat += _dot_ao_ao(mol, ao[3], aow, non0tab, shls_slice, ao_loc)
 498 | 
 499 |     return mat + mat.T.conj()
 500 | 
 501 | 
 502 | def _dot_ao_ao(mol, ao1, ao2, non0tab, shls_slice, ao_loc, hermi=0):
 503 |     '''return numpy.dot(ao1.T, ao2)'''
 504 |     ngrids, nao = ao1.shape
 505 |     if nao < SWITCH_SIZE:
 506 |         return lib.dot(ao1.T.conj(), ao2)
 507 | 
 508 |     if not ao1.flags.f_contiguous:
 509 |         ao1 = lib.transpose(ao1)
 510 |     if not ao2.flags.f_contiguous:
 511 |         ao2 = lib.transpose(ao2)
 512 |     if ao1.dtype == ao2.dtype == numpy.double:
 513 |         fn = libdft.VXCdot_ao_ao
 514 |     else:
 515 |         fn = libdft.VXCzdot_ao_ao
 516 |         ao1 = numpy.asarray(ao1, numpy.complex128)
 517 |         ao2 = numpy.asarray(ao2, numpy.complex128)
 518 | 
 519 |     if non0tab is None or shls_slice is None or ao_loc is None:
 520 |         pnon0tab = pshls_slice = pao_loc = lib.c_null_ptr()
 521 |     else:
 522 |         pnon0tab    = non0tab.ctypes.data_as(ctypes.c_void_p)
 523 |         pshls_slice = (ctypes.c_int*2)(*shls_slice)
 524 |         pao_loc     = ao_loc.ctypes.data_as(ctypes.c_void_p)
 525 | 
 526 |     vv = numpy.empty((nao,nao), dtype=ao1.dtype)
 527 |     fn(vv.ctypes.data_as(ctypes.c_void_p),
 528 |        ao1.ctypes.data_as(ctypes.c_void_p),
 529 |        ao2.ctypes.data_as(ctypes.c_void_p),
 530 |        ctypes.c_int(nao), ctypes.c_int(ngrids),
 531 |        ctypes.c_int(mol.nbas), ctypes.c_int(hermi),
 532 |        pnon0tab, pshls_slice, pao_loc)
 533 |     return vv
 534 | 
 535 | def _dot_ao_dm(mol, ao, dm, non0tab, shls_slice, ao_loc, out=None):
 536 |     '''return numpy.dot(ao, dm)'''
 537 |     ngrids, nao = ao.shape
 538 |     if nao < SWITCH_SIZE:
 539 |         return lib.dot(ao, dm)
 540 | 
 541 |     if not ao.flags.f_contiguous:
 542 |         ao = lib.transpose(ao)
 543 |     if ao.dtype == dm.dtype == numpy.double:
 544 |         fn = libdft.VXCdot_ao_dm
 545 |     else:
 546 |         fn = libdft.VXCzdot_ao_dm
 547 |         ao = numpy.asarray(ao, numpy.complex128)
 548 |         dm = numpy.asarray(dm, numpy.complex128)
 549 | 
 550 |     if non0tab is None or shls_slice is None or ao_loc is None:
 551 |         pnon0tab = pshls_slice = pao_loc = lib.c_null_ptr()
 552 |     else:
 553 |         pnon0tab    = non0tab.ctypes.data_as(ctypes.c_void_p)
 554 |         pshls_slice = (ctypes.c_int*2)(*shls_slice)
 555 |         pao_loc     = ao_loc.ctypes.data_as(ctypes.c_void_p)
 556 | 
 557 |     vm = numpy.ndarray((ngrids,dm.shape[1]), dtype=ao.dtype, order='F', buffer=out)
 558 |     dm = numpy.asarray(dm, order='C')
 559 |     fn(vm.ctypes.data_as(ctypes.c_void_p),
 560 |        ao.ctypes.data_as(ctypes.c_void_p),
 561 |        dm.ctypes.data_as(ctypes.c_void_p),
 562 |        ctypes.c_int(nao), ctypes.c_int(dm.shape[1]),
 563 |        ctypes.c_int(ngrids), ctypes.c_int(mol.nbas),
 564 |        pnon0tab, pshls_slice, pao_loc)
 565 |     return vm
 566 | 
 567 | def nr_vxc(mol, grids, xc_code, dms, spin=0, relativity=0, hermi=0,
 568 |            max_memory=2000, verbose=None):
 569 |     '''
 570 |     Evaluate RKS/UKS XC functional and potential matrix on given meshgrids
 571 |     for a set of density matrices.  See :func:`nr_rks` and :func:`nr_uks`
 572 |     for more details.
 573 |     Args:
 574 |         mol : an instance of :class:`Mole`
 575 |         grids : an instance of :class:`Grids`
 576 |             grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
 577 |         xc_code : str
 578 |             XC functional description.
 579 |             See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
 580 |         dms : 2D array or a list of 2D arrays
 581 |             Density matrix or multiple density matrices
 582 |     Kwargs:
 583 |         hermi : int
 584 |             Input density matrices symmetric or not
 585 |         max_memory : int or float
 586 |             The maximum size of cache to use (in MB).
 587 |     Returns:
 588 |         nelec, excsum, vmat.
 589 |         nelec is the number of electrons generated by numerical integration.
 590 |         excsum is the XC functional value.  vmat is the XC potential matrix in
 591 |         2D array of shape (nao,nao) where nao is the number of AO functions.
 592 |     Examples:
 593 |     >>> from pyscf import gto, dft
 594 |     >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')
 595 |     >>> grids = dft.gen_grid.Grids(mol)
 596 |     >>> grids.coords = numpy.random.random((100,3))  # 100 random points
 597 |     >>> grids.weights = numpy.random.random(100)
 598 |     >>> nao = mol.nao_nr()
 599 |     >>> dm = numpy.random.random((2,nao,nao))
 600 |     >>> nelec, exc, vxc = dft.numint.nr_vxc(mol, grids, 'lda,vwn', dm, spin=1)
 601 |     '''
 602 |     ni = NumInt()
 603 |     return ni.nr_vxc(mol, grids, xc_code, dms, spin, relativity,
 604 |                      hermi, max_memory, verbose)
 605 | 
 606 | def nr_rks(ni, mol, grids, xc_code, dms, relativity=0, hermi=0,
 607 |            max_memory=2000, verbose=None):
 608 |     '''Calculate RKS XC functional and potential matrix on given meshgrids
 609 |     for a set of density matrices
 610 |     Args:
 611 |         ni : an instance of :class:`NumInt`
 612 |         mol : an instance of :class:`Mole`
 613 |         grids : an instance of :class:`Grids`
 614 |             grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
 615 |         xc_code : str
 616 |             XC functional description.
 617 |             See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
 618 |         dms : 2D array a list of 2D arrays
 619 |             Density matrix or multiple density matrices
 620 |     Kwargs:
 621 |         hermi : int
 622 |             Input density matrices symmetric or not
 623 |         max_memory : int or float
 624 |             The maximum size of cache to use (in MB).
 625 |     Returns:
 626 |         nelec, excsum, vmat.
 627 |         nelec is the number of electrons generated by numerical integration.
 628 |         excsum is the XC functional value.  vmat is the XC potential matrix in
 629 |         2D array of shape (nao,nao) where nao is the number of AO functions.
 630 |     Examples:
 631 |     >>> from pyscf import gto, dft
 632 |     >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')
 633 |     >>> grids = dft.gen_grid.Grids(mol)
 634 |     >>> grids.coords = numpy.random.random((100,3))  # 100 random points
 635 |     >>> grids.weights = numpy.random.random(100)
 636 |     >>> nao = mol.nao_nr()
 637 |     >>> dm = numpy.random.random((nao,nao))
 638 |     >>> ni = dft.numint.NumInt()
 639 |     >>> nelec, exc, vxc = ni.nr_rks(mol, grids, 'lda,vwn', dm)
 640 |     '''
 641 |     xctype = ni._xc_type(xc_code)
 642 |     make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
 643 | 
 644 |     shls_slice = (0, mol.nbas)
 645 |     ao_loc = mol.ao_loc_nr()
 646 | 
 647 |     nelec = numpy.zeros(nset)
 648 |     excsum = numpy.zeros(nset)
 649 |     vmat = numpy.zeros((nset,nao,nao))
 650 |     aow = None
 651 |     if xctype == 'LDA':
 652 |         ao_deriv = 0
 653 |         for ao, mask, weight, coords \
 654 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 655 |             aow = numpy.ndarray(ao.shape, order='F', buffer=aow)
 656 |             for idm in range(nset):
 657 |                 rho = make_rho(idm, ao, mask, 'LDA')
 658 |                 exc, vxc = ni.eval_xc(xc_code, rho, 0, relativity, 1, verbose)[:2]
 659 |                 vrho = vxc[0]
 660 |                 den = rho * weight
 661 |                 nelec[idm] += den.sum()
 662 |                 excsum[idm] += numpy.dot(den, exc)
 663 |                 # *.5 because vmat + vmat.T
 664 |                 aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho, out=aow)
 665 |                 vmat[idm] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc)
 666 |                 rho = exc = vxc = vrho = None
 667 |     elif xctype == 'GGA':
 668 |         ao_deriv = 1
 669 |         for ao, mask, weight, coords \
 670 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 671 |             ngrid = weight.size
 672 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
 673 |             for idm in range(nset):
 674 |                 rho = make_rho(idm, ao, mask, 'GGA')
 675 |                 exc, vxc = ni.eval_xc(xc_code, rho, 0, relativity, 1, verbose)[:2]
 676 |                 den = rho[0] * weight
 677 |                 nelec[idm] += den.sum()
 678 |                 excsum[idm] += numpy.dot(den, exc)
 679 | # ref eval_mat function
 680 |                 wv = _rks_gga_wv0(rho, vxc, weight)
 681 |                 aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
 682 |                 vmat[idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 683 |                 rho = exc = vxc = wv = None
 684 |     elif xctype == 'NLC':
 685 |         nlc_pars = ni.nlc_coeff(xc_code[:-6])
 686 |         if nlc_pars == [0,0]:
 687 |             raise NotImplementedError('VV10 cannot be used with %s. '
 688 |                                       'The supported functionals are %s' %
 689 |                                       (xc_code[:-6], ni.libxc.VV10_XC))
 690 |         ao_deriv = 1
 691 |         vvrho=numpy.empty([nset,4,0])
 692 |         vvweight=numpy.empty([nset,0])
 693 |         vvcoords=numpy.empty([nset,0,3])
 694 |         for ao, mask, weight, coords \
 695 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 696 |             rhotmp = numpy.empty([0,4,weight.size])
 697 |             weighttmp = numpy.empty([0,weight.size])
 698 |             coordstmp = numpy.empty([0,weight.size,3])
 699 |             for idm in range(nset):
 700 |                 rho = make_rho(idm, ao, mask, 'GGA')
 701 |                 rho = numpy.expand_dims(rho,axis=0)
 702 |                 rhotmp = numpy.concatenate((rhotmp,rho),axis=0)
 703 |                 weighttmp = numpy.concatenate((weighttmp,numpy.expand_dims(weight,axis=0)),axis=0)
 704 |                 coordstmp = numpy.concatenate((coordstmp,numpy.expand_dims(coords,axis=0)),axis=0)
 705 |                 rho = None
 706 |             vvrho=numpy.concatenate((vvrho,rhotmp),axis=2)
 707 |             vvweight=numpy.concatenate((vvweight,weighttmp),axis=1)
 708 |             vvcoords=numpy.concatenate((vvcoords,coordstmp),axis=1)
 709 |             rhotmp = weighttmp = coordstmp = None
 710 |         for ao, mask, weight, coords \
 711 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 712 |             ngrid = weight.size
 713 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
 714 |             for idm in range(nset):
 715 |                 rho = make_rho(idm, ao, mask, 'GGA')
 716 |                 exc, vxc = _vv10nlc(rho,coords,vvrho[idm],vvweight[idm],vvcoords[idm],nlc_pars)
 717 |                 den = rho[0] * weight
 718 |                 nelec[idm] += den.sum()
 719 |                 excsum[idm] += numpy.dot(den, exc)
 720 | # ref eval_mat function
 721 |                 wv = _rks_gga_wv0(rho, vxc, weight)
 722 |                 aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
 723 |                 vmat[idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 724 |                 rho = exc = vxc = wv = None
 725 |         vvrho = vvweight = vvcoords = None
 726 |     elif xctype == 'MGGA':
 727 |         if (any(x in xc_code.upper() for x in ('CC06', 'CS', 'BR89', 'MK00'))):
 728 |             raise NotImplementedError('laplacian in meta-GGA method')
 729 |         ao_deriv = 2
 730 |         for ao, mask, weight, coords \
 731 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 732 |             ngrid = weight.size
 733 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
 734 |             for idm in range(nset):
 735 |                 rho = make_rho(idm, ao, mask, 'MGGA')
 736 |                 exc, vxc = ni.eval_xc(xc_code, rho, 0, relativity, 1, verbose)[:2]
 737 |                 vrho, vsigma, vlapl, vtau = vxc[:4]
 738 |                 den = rho[0] * weight
 739 |                 nelec[idm] += den.sum()
 740 |                 excsum[idm] += numpy.dot(den, exc)
 741 | 
 742 |                 wv = _rks_gga_wv0(rho, vxc, weight)
 743 |                 aow = numpy.einsum('npi,np->pi', ao[:4], wv, out=aow)
 744 |                 vmat[idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 745 | 
 746 | # FIXME: .5 * .5   First 0.5 for v+v.T symmetrization.
 747 | # Second 0.5 is due to the Libxc convention tau = 1/2 \nabla\phi\dot\nabla\phi
 748 |                 wv = (.5 * .5 * weight * vtau).reshape(-1,1)
 749 |                 vmat[idm] += _dot_ao_ao(mol, ao[1], wv*ao[1], mask, shls_slice, ao_loc)
 750 |                 vmat[idm] += _dot_ao_ao(mol, ao[2], wv*ao[2], mask, shls_slice, ao_loc)
 751 |                 vmat[idm] += _dot_ao_ao(mol, ao[3], wv*ao[3], mask, shls_slice, ao_loc)
 752 | 
 753 |                 rho = exc = vxc = vrho = vsigma = wv = None
 754 | 
 755 |     elif xctype == 'MGGA+':
 756 |         if (any(x in xc_code.upper() for x in ('CC06', 'CS', 'BR89', 'MK00'))):
 757 |             raise NotImplementedError('laplacian in meta-GGA method')
 758 |         ao_deriv = 2
 759 |         for ao, mask, weight, coords \
 760 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 761 |             ngrid = weight.size
 762 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
 763 |             for idm in range(nset):
 764 |                 rho = make_rho(idm, ao, mask, 'MGGA')
 765 |                 exc, vxc = ni.eval_xc(xc_code,rho,0,weight,coords, relativity, 1, verbose)[:2]
 766 |                 vrho, vsigma, vlapl, vtau = vxc[:4]
 767 |                 den = rho[0] * weight
 768 |                 nelec[idm] += den.sum()
 769 |                 excsum[idm] += numpy.dot(den, exc)
 770 | 
 771 |                 wv = _rks_gga_wv0(rho, vxc, weight)
 772 |                 aow = numpy.einsum('npi,np->pi', ao[:4], wv, out=aow)
 773 |                 vmat[idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 774 | 
 775 | # FIXME: .5 * .5   First 0.5 for v+v.T symmetrization.
 776 | # Second 0.5 is due to the Libxc convention tau = 1/2 \nabla\phi\dot\nabla\phi
 777 |                 wv = (.5 * .5 * weight * vtau).reshape(-1,1)
 778 |                 vmat[idm] += _dot_ao_ao(mol, ao[1], wv*ao[1], mask, shls_slice, ao_loc)
 779 |                 vmat[idm] += _dot_ao_ao(mol, ao[2], wv*ao[2], mask, shls_slice, ao_loc)
 780 |                 vmat[idm] += _dot_ao_ao(mol, ao[3], wv*ao[3], mask, shls_slice, ao_loc)
 781 | 
 782 |                 rho = exc = vxc = vrho = vsigma = wv = None
 783 | 
 784 |     for i in range(nset):
 785 |         vmat[i] = vmat[i] + vmat[i].T
 786 |     if nset == 1:
 787 |         nelec = nelec[0]
 788 |         excsum = excsum[0]
 789 |         vmat = vmat.reshape(nao,nao)
 790 |     return nelec, excsum, vmat
 791 | 
 792 | def nr_uks(ni, mol, grids, xc_code, dms, relativity=0, hermi=0,
 793 |            max_memory=2000, verbose=None):
 794 |     '''Calculate UKS XC functional and potential matrix on given meshgrids
 795 |     for a set of density matrices
 796 |     Args:
 797 |         mol : an instance of :class:`Mole`
 798 |         grids : an instance of :class:`Grids`
 799 |             grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
 800 |         xc_code : str
 801 |             XC functional description.
 802 |             See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
 803 |         dms : a list of 2D arrays
 804 |             A list of density matrices, stored as (alpha,alpha,...,beta,beta,...)
 805 |     Kwargs:
 806 |         hermi : int
 807 |             Input density matrices symmetric or not
 808 |         max_memory : int or float
 809 |             The maximum size of cache to use (in MB).
 810 |     Returns:
 811 |         nelec, excsum, vmat.
 812 |         nelec is the number of (alpha,beta) electrons generated by numerical integration.
 813 |         excsum is the XC functional value.
 814 |         vmat is the XC potential matrix for (alpha,beta) spin.
 815 |     Examples:
 816 |     >>> from pyscf import gto, dft
 817 |     >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')
 818 |     >>> grids = dft.gen_grid.Grids(mol)
 819 |     >>> grids.coords = numpy.random.random((100,3))  # 100 random points
 820 |     >>> grids.weights = numpy.random.random(100)
 821 |     >>> nao = mol.nao_nr()
 822 |     >>> dm = numpy.random.random((2,nao,nao))
 823 |     >>> ni = dft.numint.NumInt()
 824 |     >>> nelec, exc, vxc = ni.nr_uks(mol, grids, 'lda,vwn', dm)
 825 |     '''
 826 |     xctype = ni._xc_type(xc_code)
 827 |     if xctype == 'NLC':
 828 |         dms_sf = dms[0] + dms[1]
 829 |         nelec, excsum, vmat = nr_rks(ni, mol, grids, xc_code, dms_sf, relativity, hermi,
 830 |                                      max_memory, verbose)
 831 |         return [nelec,nelec], excsum, numpy.asarray([vmat,vmat])
 832 | 
 833 |     shls_slice = (0, mol.nbas)
 834 |     ao_loc = mol.ao_loc_nr()
 835 | 
 836 |     dma, dmb = _format_uks_dm(dms)
 837 |     nao = dma.shape[-1]
 838 |     make_rhoa, nset = ni._gen_rho_evaluator(mol, dma, hermi)[:2]
 839 |     make_rhob       = ni._gen_rho_evaluator(mol, dmb, hermi)[0]
 840 | 
 841 |     nelec = numpy.zeros((2,nset))
 842 |     excsum = numpy.zeros(nset)
 843 |     vmat = numpy.zeros((2,nset,nao,nao))
 844 |     aow = None
 845 |     if xctype == 'LDA':
 846 |         ao_deriv = 0
 847 |         for ao, mask, weight, coords \
 848 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 849 |             aow = numpy.ndarray(ao.shape, order='F', buffer=aow)
 850 |             for idm in range(nset):
 851 |                 rho_a = make_rhoa(idm, ao, mask, xctype)
 852 |                 rho_b = make_rhob(idm, ao, mask, xctype)
 853 |                 exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
 854 |                                       1, relativity, 1, verbose)[:2]
 855 |                 vrho = vxc[0]
 856 |                 den = rho_a * weight
 857 |                 nelec[0,idm] += den.sum()
 858 |                 excsum[idm] += numpy.dot(den, exc)
 859 |                 den = rho_b * weight
 860 |                 nelec[1,idm] += den.sum()
 861 |                 excsum[idm] += numpy.dot(den, exc)
 862 | 
 863 |                 # *.5 due to +c.c. in the end
 864 |                 aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho[:,0], out=aow)
 865 |                 vmat[0,idm] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc)
 866 |                 aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho[:,1], out=aow)
 867 |                 vmat[1,idm] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc)
 868 |                 rho_a = rho_b = exc = vxc = vrho = None
 869 |     elif xctype == 'GGA':
 870 |         ao_deriv = 1
 871 |         for ao, mask, weight, coords \
 872 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 873 |             ngrid = weight.size
 874 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
 875 |             for idm in range(nset):
 876 |                 rho_a = make_rhoa(idm, ao, mask, xctype)
 877 |                 rho_b = make_rhob(idm, ao, mask, xctype)
 878 |                 exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
 879 |                                       1, relativity, 1, verbose)[:2]
 880 |                 den = rho_a[0]*weight
 881 |                 nelec[0,idm] += den.sum()
 882 |                 excsum[idm] += numpy.dot(den, exc)
 883 |                 den = rho_b[0]*weight
 884 |                 nelec[1,idm] += den.sum()
 885 |                 excsum[idm] += numpy.dot(den, exc)
 886 | 
 887 |                 wva, wvb = _uks_gga_wv0((rho_a,rho_b), vxc, weight)
 888 |                 aow = numpy.einsum('npi,np->pi', ao, wva, out=aow)
 889 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 890 |                 aow = numpy.einsum('npi,np->pi', ao, wvb, out=aow)
 891 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 892 |                 rho_a = rho_b = exc = vxc = wva = wvb = None
 893 |     elif xctype == 'MGGA':
 894 |         if (any(x in xc_code.upper() for x in ('CC06', 'CS', 'BR89', 'MK00'))):
 895 |             raise NotImplementedError('laplacian in meta-GGA method')
 896 |         ao_deriv = 2
 897 |         for ao, mask, weight, coords \
 898 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 899 |             ngrid = weight.size
 900 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
 901 |             for idm in range(nset):
 902 |                 rho_a = make_rhoa(idm, ao, mask, xctype)
 903 |                 rho_b = make_rhob(idm, ao, mask, xctype)
 904 |                 exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
 905 |                                       1, relativity, 1, verbose)[:2]
 906 |                 vrho, vsigma, vlapl, vtau = vxc[:4]
 907 |                 den = rho_a[0]*weight
 908 |                 nelec[0,idm] += den.sum()
 909 |                 excsum[idm] += numpy.dot(den, exc)
 910 |                 den = rho_b[0]*weight
 911 |                 nelec[1,idm] += den.sum()
 912 |                 excsum[idm] += numpy.dot(den, exc)
 913 | 
 914 |                 wva, wvb = _uks_gga_wv0((rho_a,rho_b), vxc, weight)
 915 |                 aow = numpy.einsum('npi,np->pi', ao[:4], wva, out=aow)
 916 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 917 |                 aow = numpy.einsum('npi,np->pi', ao[:4], wvb, out=aow)
 918 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 919 | 
 920 | # FIXME: .5 * .5   First 0.5 for v+v.T symmetrization.
 921 | # Second 0.5 is due to the Libxc convention tau = 1/2 \nabla\phi\dot\nabla\phi
 922 |                 wv = (.25 * weight * vtau[:,0]).reshape(-1,1)
 923 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[1], wv*ao[1], mask, shls_slice, ao_loc)
 924 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[2], wv*ao[2], mask, shls_slice, ao_loc)
 925 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[3], wv*ao[3], mask, shls_slice, ao_loc)
 926 |                 wv = (.25 * weight * vtau[:,1]).reshape(-1,1)
 927 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[1], wv*ao[1], mask, shls_slice, ao_loc)
 928 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[2], wv*ao[2], mask, shls_slice, ao_loc)
 929 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[3], wv*ao[3], mask, shls_slice, ao_loc)
 930 |                 rho_a = rho_b = exc = vxc = vrho = vsigma = wva = wvb = None
 931 |     elif xctype == 'MGGA+':
 932 |         if (any(x in xc_code.upper() for x in ('CC06', 'CS', 'BR89', 'MK00'))):
 933 |             raise NotImplementedError('laplacian in meta-GGA method')
 934 |         ao_deriv = 2
 935 |         for ao, mask, weight, coords \
 936 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
 937 |             ngrid = weight.size
 938 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
 939 |             for idm in range(nset):
 940 |                 rho_a = make_rhoa(idm, ao, mask, xctype)
 941 |                 rho_b = make_rhob(idm, ao, mask, xctype)
 942 |                 exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
 943 |                                       1,weight,coords, relativity, 1, verbose)[:2]
 944 |                 vrho, vsigma, vlapl, vtau = vxc[:4]
 945 |                 den = rho_a[0]*weight
 946 |                 nelec[0,idm] += den.sum()
 947 |                 excsum[idm] += numpy.dot(den, exc)
 948 |                 den = rho_b[0]*weight
 949 |                 nelec[1,idm] += den.sum()
 950 |                 excsum[idm] += numpy.dot(den, exc)
 951 | 
 952 |                 wva, wvb = _uks_gga_wv0((rho_a,rho_b), vxc, weight)
 953 |                 aow = numpy.einsum('npi,np->pi', ao[:4], wva, out=aow)
 954 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 955 |                 aow = numpy.einsum('npi,np->pi', ao[:4], wvb, out=aow)
 956 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
 957 | 
 958 | # FIXME: .5 * .5   First 0.5 for v+v.T symmetrization.
 959 | # Second 0.5 is due to the Libxc convention tau = 1/2 \nabla\phi\dot\nabla\phi
 960 |                 wv = (.25 * weight * vtau[:,0]).reshape(-1,1)
 961 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[1], wv*ao[1], mask, shls_slice, ao_loc)
 962 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[2], wv*ao[2], mask, shls_slice, ao_loc)
 963 |                 vmat[0,idm] += _dot_ao_ao(mol, ao[3], wv*ao[3], mask, shls_slice, ao_loc)
 964 |                 wv = (.25 * weight * vtau[:,1]).reshape(-1,1)
 965 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[1], wv*ao[1], mask, shls_slice, ao_loc)
 966 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[2], wv*ao[2], mask, shls_slice, ao_loc)
 967 |                 vmat[1,idm] += _dot_ao_ao(mol, ao[3], wv*ao[3], mask, shls_slice, ao_loc)
 968 |                 rho_a = rho_b = exc = vxc = vrho = vsigma = wva = wvb = None
 969 | 
 970 |     for i in range(nset):
 971 |         vmat[0,i] = vmat[0,i] + vmat[0,i].T
 972 |         vmat[1,i] = vmat[1,i] + vmat[1,i].T
 973 |     if isinstance(dma, numpy.ndarray) and dma.ndim == 2:
 974 |         vmat = vmat[:,0]
 975 |         nelec = nelec.reshape(2)
 976 |         excsum = excsum[0]
 977 |     return nelec, excsum, vmat
 978 | 
 979 | def _format_uks_dm(dms):
 980 |     if isinstance(dms, numpy.ndarray) and dms.ndim == 2:  # RHF DM
 981 |         dma = dmb = dms * .5
 982 |     else:
 983 |         dma, dmb = dms
 984 |     if hasattr(dms, 'mo_coeff'):
 985 |         mo_coeff = dms.mo_coeff
 986 |         mo_occ = dms.mo_occ
 987 |         if mo_coeff[0].ndim < dma.ndim: # handle ROKS
 988 |             mo_occa = numpy.array(mo_occ> 0, dtype=numpy.double)
 989 |             mo_occb = numpy.array(mo_occ==2, dtype=numpy.double)
 990 |             dma = lib.tag_array(dma, mo_coeff=mo_coeff, mo_occ=mo_occa)
 991 |             dmb = lib.tag_array(dmb, mo_coeff=mo_coeff, mo_occ=mo_occb)
 992 |         else:
 993 |             dma = lib.tag_array(dma, mo_coeff=mo_coeff[0], mo_occ=mo_occ[0])
 994 |             dmb = lib.tag_array(dmb, mo_coeff=mo_coeff[1], mo_occ=mo_occ[1])
 995 |     return dma, dmb
 996 | 
 997 | nr_rks_vxc = nr_rks
 998 | nr_uks_vxc = nr_uks
 999 | 
1000 | def nr_rks_fxc(ni, mol, grids, xc_code, dm0, dms, relativity=0, hermi=0,
1001 |                rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
1002 |     '''Contract RKS XC (singlet hessian) kernel matrix with given density matrices
1003 |     Args:
1004 |         ni : an instance of :class:`NumInt`
1005 |         mol : an instance of :class:`Mole`
1006 |         grids : an instance of :class:`Grids`
1007 |             grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
1008 |         xc_code : str
1009 |             XC functional description.
1010 |             See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
1011 |         dms : 2D array a list of 2D arrays
1012 |             Density matrix or multiple density matrices
1013 |     Kwargs:
1014 |         hermi : int
1015 |             Input density matrices symmetric or not
1016 |         max_memory : int or float
1017 |             The maximum size of cache to use (in MB).
1018 |         rho0 : float array
1019 |             Zero-order density (and density derivative for GGA).  Giving kwargs rho0,
1020 |             vxc and fxc to improve better performance.
1021 |         vxc : float array
1022 |             First order XC derivatives
1023 |         fxc : float array
1024 |             Second order XC derivatives
1025 |     Returns:
1026 |         nelec, excsum, vmat.
1027 |         nelec is the number of electrons generated by numerical integration.
1028 |         excsum is the XC functional value.  vmat is the XC potential matrix in
1029 |         2D array of shape (nao,nao) where nao is the number of AO functions.
1030 |     Examples:
1031 |     '''
1032 |     xctype = ni._xc_type(xc_code)
1033 | 
1034 |     make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
1035 |     if ((xctype == 'LDA' and fxc is None) or
1036 |         (xctype == 'GGA' and rho0 is None)):
1037 |         make_rho0 = ni._gen_rho_evaluator(mol, dm0, 1)[0]
1038 | 
1039 |     shls_slice = (0, mol.nbas)
1040 |     ao_loc = mol.ao_loc_nr()
1041 | 
1042 |     vmat = numpy.zeros((nset,nao,nao))
1043 |     aow = None
1044 |     if xctype == 'LDA':
1045 |         ao_deriv = 0
1046 |         ip = 0
1047 |         for ao, mask, weight, coords \
1048 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1049 |             ngrid = weight.size
1050 |             aow = numpy.ndarray(ao.shape, order='F', buffer=aow)
1051 |             if fxc is None:
1052 |                 rho = make_rho0(0, ao, mask, 'LDA')
1053 |                 fxc0 = ni.eval_xc(xc_code, rho, 0, relativity, 2, verbose)[2]
1054 |                 frr = fxc0[0]
1055 |             else:
1056 |                 frr = fxc[0][ip:ip+ngrid]
1057 |                 ip += ngrid
1058 | 
1059 |             for i in range(nset):
1060 |                 rho1 = make_rho(i, ao, mask, 'LDA')
1061 |                 aow = numpy.einsum('pi,p->pi', ao, weight*frr*rho1, out=aow)
1062 |                 vmat[i] += _dot_ao_ao(mol, aow, ao, mask, shls_slice, ao_loc)
1063 |                 rho1 = None
1064 | 
1065 |     elif xctype == 'GGA':
1066 |         ao_deriv = 1
1067 |         ip = 0
1068 |         for ao, mask, weight, coords \
1069 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1070 |             ngrid = weight.size
1071 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
1072 |             if rho0 is None:
1073 |                 rho = make_rho0(0, ao, mask, 'GGA')
1074 |             else:
1075 |                 rho = numpy.asarray(rho0[:,ip:ip+ngrid], order='C')
1076 |             if vxc is None or fxc is None:
1077 |                 vxc0, fxc0 = ni.eval_xc(xc_code, rho, 0, relativity, 2, verbose)[1:3]
1078 |             else:
1079 |                 vxc0 = (None, vxc[1][ip:ip+ngrid])
1080 |                 fxc0 = (fxc[0][ip:ip+ngrid], fxc[1][ip:ip+ngrid], fxc[2][ip:ip+ngrid])
1081 |                 ip += ngrid
1082 | 
1083 |             for i in range(nset):
1084 |                 rho1 = make_rho(i, ao, mask, 'GGA')
1085 |                 wv = _rks_gga_wv1(rho, rho1, vxc0, fxc0, weight)
1086 |                 aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
1087 |                 vmat[i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
1088 |                 rho1 = sigma1 = None
1089 | 
1090 |         for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
1091 |             vmat[i] = vmat[i] + vmat[i].T.conj()
1092 | 
1093 |     elif xctype == 'NLC':
1094 |         raise NotImplementedError('NLC')
1095 |     elif xctype == 'MGGA':
1096 |         raise NotImplementedError('meta-GGA')
1097 | 
1098 |     if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
1099 |         vmat = vmat[0]
1100 |     return vmat
1101 | 
1102 | def nr_rks_fxc_st(ni, mol, grids, xc_code, dm0, dms_alpha, relativity=0, singlet=True,
1103 |                   rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
1104 |     '''Associated to singlet or triplet Hessian
1105 |     Note the difference to nr_rks_fxc, dms_alpha is the response density
1106 |     matrices of alpha spin, alpha+/-beta DM is applied due to singlet/triplet
1107 |     coupling
1108 |     Ref. CPL, 256, 454
1109 |     '''
1110 |     xctype = ni._xc_type(xc_code)
1111 | 
1112 |     make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms_alpha, hermi=0)
1113 |     if ((xctype == 'LDA' and fxc is None) or
1114 |         (xctype == 'GGA' and rho0 is None)):
1115 |         make_rho0 = ni._gen_rho_evaluator(mol, dm0, hermi=1)[0]
1116 | 
1117 |     shls_slice = (0, mol.nbas)
1118 |     ao_loc = mol.ao_loc_nr()
1119 | 
1120 |     vmat = numpy.zeros((nset,nao,nao))
1121 |     aow = None
1122 |     if xctype == 'LDA':
1123 |         ao_deriv = 0
1124 |         ip = 0
1125 |         for ao, mask, weight, coords \
1126 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1127 |             ngrid = weight.size
1128 |             aow = numpy.ndarray(ao.shape, order='F', buffer=aow)
1129 |             if fxc is None:
1130 |                 rho = make_rho0(0, ao, mask, 'LDA')
1131 |                 rho *= .5  # alpha density
1132 |                 fxc0 = ni.eval_xc(xc_code, (rho,rho), 1, deriv=2)[2]
1133 |                 u_u, u_d, d_d = fxc0[0].T
1134 |             else:
1135 |                 if fxc[0].ndim == 1:
1136 |                     raise RuntimeError('cached (rho, vxc, fxc) need to be '
1137 |                                        'generated by cache_xc_kernel with flag spin=1')
1138 |                 u_u, u_d, d_d = fxc[0][ip:ip+ngrid].T
1139 |                 ip += ngrid
1140 |             if singlet:
1141 |                 frho = u_u + u_d
1142 |                 if 0:
1143 |                     rho = ni.eval_rho2(mol, ao, mo_coeff, mo_occ, mask, 'LDA')
1144 |                     fxc_test = ni.eval_xc(xc_code, rho, 0, deriv=2)[2]
1145 |                     assert(numpy.linalg.norm(fxc_test[0]*2-frho) < 1e-4)
1146 |             else:
1147 |                 frho = u_u - u_d
1148 | 
1149 |             for i in range(nset):
1150 |                 rho1 = make_rho(i, ao, mask, 'LDA')
1151 |                 aow = numpy.einsum('pi,p->pi', ao, weight*frho*rho1, out=aow)
1152 |                 vmat[i] += _dot_ao_ao(mol, aow, ao, mask, shls_slice, ao_loc)
1153 |                 rho1 = None
1154 | 
1155 |     elif xctype == 'GGA':
1156 |         ao_deriv = 1
1157 |         ip = 0
1158 |         for ao, mask, weight, coords \
1159 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1160 |             ngrid = weight.size
1161 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
1162 |             if vxc is None or fxc is None:
1163 |                 rho = make_rho0(0, ao, mask, 'GGA')
1164 |                 rho *= .5  # alpha density
1165 |                 vxc0, fxc0 = ni.eval_xc(xc_code, (rho,rho), 1, deriv=2)[1:3]
1166 | 
1167 |                 vsigma = vxc0[1].T
1168 |                 u_u, u_d, d_d = fxc0[0].T  # v2rho2
1169 |                 u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc0[1].T  # v2rhosigma
1170 |                 uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc0[2].T  # v2sigma2
1171 |             else:
1172 |                 if rho0[0].ndim == 1:
1173 |                     raise RuntimeError('cached (rho, vxc, fxc) need to be '
1174 |                                        'generated by cache_xc_kernel with flag spin=1')
1175 |                 rho = rho0[0][:,ip:ip+ngrid]
1176 |                 vsigma = vxc[1][ip:ip+ngrid].T
1177 |                 u_u, u_d, d_d = fxc[0][ip:ip+ngrid].T  # v2rho2
1178 |                 u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1][ip:ip+ngrid].T  # v2rhosigma
1179 |                 uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2][ip:ip+ngrid].T  # v2sigma2
1180 |                 ip += ngrid
1181 | 
1182 |             # Factorization differs to CPL, 256, 454, to use _rks_gga_wv1 function
1183 |             if singlet:
1184 |                 fgamma = vsigma[0] + vsigma[1] * .5
1185 |                 frho = u_u + u_d
1186 |                 fgg = uu_uu + .5*ud_ud + 2*uu_ud + uu_dd
1187 |                 frhogamma = u_uu + u_dd + u_ud
1188 |             else:
1189 |                 fgamma = vsigma[0] - vsigma[1] * .5
1190 |                 frho = u_u - u_d
1191 |                 fgg = uu_uu - uu_dd
1192 |                 frhogamma = u_uu - u_dd
1193 | 
1194 |             for i in range(nset):
1195 |                 # rho1[0 ] = |b><j| z_{bj}
1196 |                 # rho1[1:] = \nabla(|b><j|) z_{bj}
1197 |                 rho1 = make_rho(i, ao, mask, 'GGA')
1198 |                 wv = _rks_gga_wv1(rho, rho1, (None,fgamma), (frho,frhogamma,fgg), weight)
1199 |                 aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
1200 |                 vmat[i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
1201 |                 rho1 = sigma1 = None
1202 | 
1203 |         for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
1204 |             vmat[i] = vmat[i] + vmat[i].T.conj()
1205 | 
1206 |     elif xctype == 'NLC':
1207 |         raise NotImplementedError('NLC')
1208 |     elif xctype == 'MGGA':
1209 |         raise NotImplementedError('meta-GGA')
1210 | 
1211 |     if isinstance(dms_alpha, numpy.ndarray) and dms_alpha.ndim == 2:
1212 |         vmat = vmat[0]
1213 |     return vmat
1214 | 
1215 | def _rks_gga_wv0(rho, vxc, weight):
1216 |     vrho, vgamma = vxc[:2]
1217 |     ngrid = vrho.size
1218 |     wv = numpy.empty((4,ngrid))
1219 |     wv[0]  = weight * vrho
1220 |     wv[1:] = (weight * vgamma * 2) * rho[1:4]
1221 |     wv[0] *= .5  # v+v.T should be applied in the caller
1222 |     return wv
1223 | 
1224 | def _rks_gga_wv1(rho0, rho1, vxc, fxc, weight):
1225 |     vgamma = vxc[1]
1226 |     frho, frhogamma, fgg = fxc[:3]
1227 |     # sigma1 ~ \nabla(\rho_\alpha+\rho_\beta) dot \nabla(|b><j|) z_{bj}
1228 |     sigma1 = numpy.einsum('xi,xi->i', rho0[1:4], rho1[1:4])
1229 |     ngrid = vgamma.size
1230 |     wv = numpy.empty((4,ngrid))
1231 |     wv[0]  = frho * rho1[0]
1232 |     wv[0] += frhogamma * sigma1 * 2
1233 |     wv[1:] = (fgg * sigma1 * 4 + frhogamma * rho1[0] * 2) * rho0[1:4]
1234 |     wv[1:]+= vgamma * rho1[1:4] * 2
1235 |     wv *= weight
1236 |     wv[0] *= .5  # v+v.T should be applied in the caller
1237 |     return wv
1238 | 
1239 | def _rks_gga_wv2(rho0, rho1, fxc, kxc, weight):
1240 |     frr, frg, fgg = fxc[:3]
1241 |     frrr, frrg, frgg, fggg = kxc
1242 |     sigma1 = numpy.einsum('xi,xi->i', rho0[1:], rho1[1:])
1243 |     r1r1 = rho1[0]**2
1244 |     s1s1 = sigma1**2
1245 |     r1s1 = rho1[0] * sigma1
1246 |     sigma2 = numpy.einsum('xi,xi->i', rho1[1:], rho1[1:])
1247 |     ngrid = frrr.size
1248 |     wv = numpy.empty((4,ngrid))
1249 |     wv[0]  = frrr * r1r1
1250 |     wv[0] += 4 * frrg * r1s1
1251 |     wv[0] += 4 * frgg * s1s1
1252 |     wv[0] += 2 * frg * sigma2
1253 |     wv[1:]  = 2 * frrg * r1r1 * rho0[1:]
1254 |     wv[1:] += 8 * frgg * r1s1 * rho0[1:]
1255 |     wv[1:] += 4 * frg * rho1[0] * rho1[1:]
1256 |     wv[1:] += 4 * fgg * sigma2 * rho0[1:]
1257 |     wv[1:] += 8 * fgg * sigma1 * rho1[1:]
1258 |     wv[1:] += 8 * fggg * s1s1 * rho0[1:]
1259 |     wv *= weight
1260 |     wv[0]*=.5  # v+v.T should be applied in the caller
1261 |     return wv
1262 | 
1263 | def nr_uks_fxc(ni, mol, grids, xc_code, dm0, dms, relativity=0, hermi=0,
1264 |                rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
1265 |     '''Contract UKS XC kernel matrix with given density matrices
1266 |     Args:
1267 |         ni : an instance of :class:`NumInt`
1268 |         mol : an instance of :class:`Mole`
1269 |         grids : an instance of :class:`Grids`
1270 |             grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
1271 |         xc_code : str
1272 |             XC functional description.
1273 |             See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
1274 |         dms : 2D array a list of 2D arrays
1275 |             Density matrix or multiple density matrices
1276 |     Kwargs:
1277 |         hermi : int
1278 |             Input density matrices symmetric or not
1279 |         max_memory : int or float
1280 |             The maximum size of cache to use (in MB).
1281 |         rho0 : float array
1282 |             Zero-order density (and density derivative for GGA).  Giving kwargs rho0,
1283 |             vxc and fxc to improve better performance.
1284 |         vxc : float array
1285 |             First order XC derivatives
1286 |         fxc : float array
1287 |             Second order XC derivatives
1288 |     Returns:
1289 |         nelec, excsum, vmat.
1290 |         nelec is the number of electrons generated by numerical integration.
1291 |         excsum is the XC functional value.  vmat is the XC potential matrix in
1292 |         2D array of shape (nao,nao) where nao is the number of AO functions.
1293 |     Examples:
1294 |     '''
1295 |     xctype = ni._xc_type(xc_code)
1296 | 
1297 |     dma, dmb = _format_uks_dm(dms)
1298 |     nao = dms.shape[-1]
1299 |     make_rhoa, nset = ni._gen_rho_evaluator(mol, dma, hermi)[:2]
1300 |     make_rhob       = ni._gen_rho_evaluator(mol, dmb, hermi)[0]
1301 | 
1302 |     if ((xctype == 'LDA' and fxc is None) or
1303 |         (xctype == 'GGA' and rho0 is None)):
1304 |         make_rho0 = ni._gen_rho_evaluator(mol, _format_uks_dm(dm0), 1)[0]
1305 | 
1306 |     shls_slice = (0, mol.nbas)
1307 |     ao_loc = mol.ao_loc_nr()
1308 | 
1309 |     vmat = numpy.zeros((2,nset,nao,nao))
1310 |     aow = None
1311 |     if xctype == 'LDA':
1312 |         ao_deriv = 0
1313 |         ip = 0
1314 |         for ao, mask, weight, coords \
1315 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1316 |             ngrid = weight.size
1317 |             aow = numpy.ndarray(ao.shape, order='F', buffer=aow)
1318 |             if fxc is None:
1319 |                 rho0a = make_rho0(0, ao, mask, xctype)
1320 |                 rho0b = make_rho0(1, ao, mask, xctype)
1321 |                 fxc0 = ni.eval_xc(xc_code, (rho0a,rho0b), 1, relativity, 2, verbose)[2]
1322 |                 u_u, u_d, d_d = fxc0[0].T
1323 |             else:
1324 |                 u_u, u_d, d_d = fxc[0][ip:ip+ngrid].T
1325 |                 ip += ngrid
1326 | 
1327 |             for i in range(nset):
1328 |                 rho1a = make_rhoa(i, ao, mask, xctype)
1329 |                 rho1b = make_rhob(i, ao, mask, xctype)
1330 |                 wv = u_u * rho1a + u_d * rho1b
1331 |                 wv *= weight
1332 |                 aow = numpy.einsum('pi,p->pi', ao, wv, out=aow)
1333 |                 vmat[0,i] += _dot_ao_ao(mol, aow, ao, mask, shls_slice, ao_loc)
1334 |                 wv = u_d * rho1a + d_d * rho1b
1335 |                 wv *= weight
1336 |                 aow = numpy.einsum('pi,p->pi', ao, wv, out=aow)
1337 |                 vmat[1,i] += _dot_ao_ao(mol, aow, ao, mask, shls_slice, ao_loc)
1338 | 
1339 |     elif xctype == 'GGA':
1340 |         ao_deriv = 1
1341 |         ip = 0
1342 |         for ao, mask, weight, coords \
1343 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1344 |             ngrid = weight.size
1345 |             aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
1346 |             if rho0 is None:
1347 |                 rho0a = make_rho0(0, ao, mask, xctype)
1348 |                 rho0b = make_rho0(1, ao, mask, xctype)
1349 |             else:
1350 |                 rho0a = rho0[0][:,ip:ip+ngrid]
1351 |                 rho0b = rho0[1][:,ip:ip+ngrid]
1352 |             if vxc is None or fxc is None:
1353 |                 vxc0, fxc0 = ni.eval_xc(xc_code, (rho0a,rho0b), 1, relativity, 2, verbose)[1:3]
1354 |             else:
1355 |                 vxc0 = (None, vxc[1][ip:ip+ngrid])
1356 |                 fxc0 = (fxc[0][ip:ip+ngrid], fxc[1][ip:ip+ngrid], fxc[2][ip:ip+ngrid])
1357 |                 ip += ngrid
1358 | 
1359 |             for i in range(nset):
1360 |                 rho1a = make_rhoa(i, ao, mask, xctype)
1361 |                 rho1b = make_rhob(i, ao, mask, xctype)
1362 |                 wva, wvb = _uks_gga_wv1((rho0a,rho0b), (rho1a,rho1b), vxc0, fxc0, weight)
1363 |                 aow = numpy.einsum('npi,np->pi', ao, wva, out=aow)
1364 |                 vmat[0,i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
1365 |                 aow = numpy.einsum('npi,np->pi', ao, wvb, out=aow)
1366 |                 vmat[1,i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
1367 | 
1368 |         for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
1369 |             vmat[0,i] = vmat[0,i] + vmat[0,i].T.conj()
1370 |             vmat[1,i] = vmat[1,i] + vmat[1,i].T.conj()
1371 | 
1372 |     elif xctype == 'NLC':
1373 |         raise NotImplementedError('NLC')
1374 |     elif xctype == 'MGGA':
1375 |         raise NotImplementedError('meta-GGA')
1376 | 
1377 |     if isinstance(dma, numpy.ndarray) and dma.ndim == 2:
1378 |         vmat = vmat[:,0]
1379 |     return vmat
1380 | 
1381 | def _uks_gga_wv0(rho, vxc, weight):
1382 |     rhoa, rhob = rho
1383 |     vrho, vsigma = vxc[:2]
1384 |     ngrid = vrho.shape[0]
1385 |     wva = numpy.empty((4,ngrid))
1386 |     wva[0]  = weight * vrho[:,0] * .5  # v+v.T should be applied in the caller
1387 |     wva[1:] = rhoa[1:4] * (weight * vsigma[:,0] * 2)  # sigma_uu
1388 |     wva[1:]+= rhob[1:4] * (weight * vsigma[:,1])      # sigma_ud
1389 |     wvb = numpy.empty((4,ngrid))
1390 |     wvb[0]  = weight * vrho[:,1] * .5  # v+v.T should be applied in the caller
1391 |     wvb[1:] = rhob[1:4] * (weight * vsigma[:,2] * 2)  # sigma_dd
1392 |     wvb[1:]+= rhoa[1:4] * (weight * vsigma[:,1])      # sigma_ud
1393 |     return wva, wvb
1394 | 
1395 | def _uks_gga_wv1(rho0, rho1, vxc, fxc, weight):
1396 |     uu, ud, dd = vxc[1].T
1397 |     u_u, u_d, d_d = fxc[0].T
1398 |     u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1].T
1399 |     uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2].T
1400 |     ngrid = uu.size
1401 | 
1402 |     rho0a, rho0b = rho0
1403 |     rho1a, rho1b = rho1
1404 |     a0a1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1a[1:4])
1405 |     a0b1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1b[1:4])
1406 |     b0a1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1a[1:4])
1407 |     b0b1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1b[1:4])
1408 | 
1409 |     wva = numpy.empty((4,ngrid))
1410 |     wvb = numpy.empty((4,ngrid))
1411 |     # alpha = alpha-alpha * alpha
1412 |     wva[0]  = u_u * rho1a[0]
1413 |     wva[0] += u_uu * a0a1 * 2
1414 |     wva[0] += u_ud * b0a1
1415 |     wva[1:] = uu * rho1a[1:4] * 2
1416 |     wva[1:]+= u_uu * rho1a[0] * rho0a[1:4] * 2
1417 |     wva[1:]+= u_ud * rho1a[0] * rho0b[1:4]
1418 |     wva[1:]+= uu_uu * a0a1 * rho0a[1:4] * 4
1419 |     wva[1:]+= uu_ud * a0a1 * rho0b[1:4] * 2
1420 |     wva[1:]+= uu_ud * b0a1 * rho0a[1:4] * 2
1421 |     wva[1:]+= ud_ud * b0a1 * rho0b[1:4]
1422 | 
1423 |     # alpha = alpha-beta  * beta
1424 |     wva[0] += u_d * rho1b[0]
1425 |     wva[0] += u_ud * a0b1
1426 |     wva[0] += u_dd * b0b1 * 2
1427 |     wva[1:]+= ud * rho1b[1:4]
1428 |     wva[1:]+= d_uu * rho1b[0] * rho0a[1:4] * 2
1429 |     wva[1:]+= d_ud * rho1b[0] * rho0b[1:4]
1430 |     wva[1:]+= uu_ud * a0b1 * rho0a[1:4] * 2
1431 |     wva[1:]+= ud_ud * a0b1 * rho0b[1:4]
1432 |     wva[1:]+= uu_dd * b0b1 * rho0a[1:4] * 4
1433 |     wva[1:]+= ud_dd * b0b1 * rho0b[1:4] * 2
1434 |     wva *= weight
1435 |     wva[0] *= .5  # v+v.T should be applied in the caller
1436 | 
1437 |     # beta = beta-alpha * alpha
1438 |     wvb[0]  = u_d * rho1a[0]
1439 |     wvb[0] += d_ud * b0a1
1440 |     wvb[0] += d_uu * a0a1 * 2
1441 |     wvb[1:] = ud * rho1a[1:4]
1442 |     wvb[1:]+= u_dd * rho1a[0] * rho0b[1:4] * 2
1443 |     wvb[1:]+= u_ud * rho1a[0] * rho0a[1:4]
1444 |     wvb[1:]+= ud_dd * b0a1 * rho0b[1:4] * 2
1445 |     wvb[1:]+= ud_ud * b0a1 * rho0a[1:4]
1446 |     wvb[1:]+= uu_dd * a0a1 * rho0b[1:4] * 4
1447 |     wvb[1:]+= uu_ud * a0a1 * rho0a[1:4] * 2
1448 | 
1449 |     # beta = beta-beta  * beta
1450 |     wvb[0] += d_d * rho1b[0]
1451 |     wvb[0] += d_dd * b0b1 * 2
1452 |     wvb[0] += d_ud * a0b1
1453 |     wvb[1:]+= dd * rho1b[1:4] * 2
1454 |     wvb[1:]+= d_dd * rho1b[0] * rho0b[1:4] * 2
1455 |     wvb[1:]+= d_ud * rho1b[0] * rho0a[1:4]
1456 |     wvb[1:]+= dd_dd * b0b1 * rho0b[1:4] * 4
1457 |     wvb[1:]+= ud_dd * b0b1 * rho0a[1:4] * 2
1458 |     wvb[1:]+= ud_dd * a0b1 * rho0b[1:4] * 2
1459 |     wvb[1:]+= ud_ud * a0b1 * rho0a[1:4]
1460 |     wvb *= weight
1461 |     wvb[0] *= .5  # v+v.T should be applied in the caller
1462 |     return wva, wvb
1463 | 
1464 | def _uks_gga_wv2(rho0, rho1, fxc, kxc, weight):
1465 |     u_u, u_d, d_d = fxc[0].T
1466 |     u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1].T
1467 |     uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2].T
1468 |     u_u_u, u_u_d, u_d_d, d_d_d = kxc[0].T
1469 |     u_u_uu, u_u_ud, u_u_dd, u_d_uu, u_d_ud, u_d_dd, d_d_uu, \
1470 |             d_d_ud, d_d_dd = kxc[1].T
1471 |     u_uu_uu, u_uu_ud, u_uu_dd, u_ud_ud, u_ud_dd, u_dd_dd, d_uu_uu, d_uu_ud, \
1472 |             d_uu_dd, d_ud_ud, d_ud_dd, d_dd_dd = kxc[2].T
1473 |     uu_uu_uu, uu_uu_ud, uu_uu_dd, uu_ud_ud, uu_ud_dd, uu_dd_dd, ud_ud_ud, \
1474 |             ud_ud_dd, ud_dd_dd, dd_dd_dd = kxc[3].T
1475 |     ngrid = u_u.size
1476 | 
1477 |     rho0a, rho0b = rho0
1478 |     rho1a, rho1b = rho1
1479 |     a0a1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1a[1:4])
1480 |     a0b1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1b[1:4])
1481 |     b0a1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1a[1:4])
1482 |     b0b1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1b[1:4])
1483 |     a1a1 = numpy.einsum('xi,xi->i', rho1a[1:4], rho1a[1:4])
1484 |     a1b1 = numpy.einsum('xi,xi->i', rho1a[1:4], rho1b[1:4])
1485 |     b1a1 = a1b1
1486 |     b1b1 = numpy.einsum('xi,xi->i', rho1b[1:4], rho1b[1:4])
1487 |     a0a1_a0a1 = numpy.einsum('i,i->i', a0a1, a0a1)
1488 |     a0a1_a0b1 = numpy.einsum('i,i->i', a0a1, a0b1)
1489 |     a0a1_b0a1 = numpy.einsum('i,i->i', a0a1, b0a1)
1490 |     a0a1_b0b1 = numpy.einsum('i,i->i', a0a1, b0b1)
1491 |     a0b1_a0a1 = a0a1_a0b1
1492 |     a0b1_a0b1 = numpy.einsum('i,i->i', a0b1, a0b1)
1493 |     a0b1_b0a1 = numpy.einsum('i,i->i', a0b1, b0a1)
1494 |     a0b1_b0b1 = numpy.einsum('i,i->i', a0b1, b0b1)
1495 |     b0a1_a0a1 = a0a1_b0a1
1496 |     b0a1_a0b1 = a0b1_b0a1
1497 |     b0a1_b0a1 = numpy.einsum('i,i->i', b0a1, b0a1)
1498 |     b0a1_b0b1 = numpy.einsum('i,i->i', b0a1, b0b1)
1499 |     b0b1_a0a1 = a0a1_b0b1
1500 |     b0b1_a0b1 = a0b1_b0b1
1501 |     b0b1_b0a1 = b0a1_b0b1
1502 |     b0b1_b0b1 = numpy.einsum('i,i->i', b0b1, b0b1)
1503 | 
1504 |     wva = numpy.zeros((4,ngrid))
1505 |     wva[0] += numpy.einsum('i,i,i->i', u_u_u, rho1a[0], rho1a[0])
1506 |     wva[0] += numpy.einsum('i,i,i->i', u_u_d, rho1a[0], rho1b[0]) * 2
1507 |     wva[0] += numpy.einsum('i,i,i->i', u_d_d, rho1b[0], rho1b[0])
1508 |     wva[0] += numpy.einsum('i,i->i', u_uu, a1a1) * 2
1509 |     wva[0] += numpy.einsum('i,i->i', u_ud, a1b1) * 2
1510 |     wva[0] += numpy.einsum('i,i->i', u_dd, b1b1) * 2
1511 |     wva[1:] += numpy.einsum('i,i,xi->xi', u_uu, rho1a[0], rho1a[1:]) * 4
1512 |     wva[1:] += numpy.einsum('i,i,xi->xi', d_uu, rho1b[0], rho1a[1:]) * 4
1513 |     wva[1:] += numpy.einsum('i,i,xi->xi', u_ud, rho1a[0], rho1b[1:]) * 2
1514 |     wva[1:] += numpy.einsum('i,i,xi->xi', d_ud, rho1b[0], rho1b[1:]) * 2
1515 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu, a0a1, rho1a[1:]) * 8
1516 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud, a0a1, rho1b[1:]) * 4
1517 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud, a0b1, rho1a[1:]) * 4
1518 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud, a0b1, rho1b[1:]) * 2
1519 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud, b0a1, rho1a[1:]) * 4
1520 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_dd, b0b1, rho1a[1:]) * 8
1521 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud, b0a1, rho1b[1:]) * 2
1522 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_dd, b0b1, rho1b[1:]) * 4
1523 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu, a1a1, rho0a[1:]) * 4
1524 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud, a1b1, rho0a[1:]) * 4
1525 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_dd, b1b1, rho0a[1:]) * 4
1526 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud, a1a1, rho0b[1:]) * 2
1527 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud, a1b1, rho0b[1:]) * 2
1528 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_dd, b1b1, rho0b[1:]) * 2
1529 |     wva[0] += numpy.einsum('i,i,i->i', u_u_uu, rho1a[0], a0a1) * 4
1530 |     wva[0] += numpy.einsum('i,i,i->i', u_d_uu, rho1b[0], a0a1) * 4
1531 |     wva[0] += numpy.einsum('i,i,i->i', u_u_ud, rho1a[0], a0b1) * 2
1532 |     wva[0] += numpy.einsum('i,i,i->i', u_d_ud, rho1b[0], a0b1) * 2
1533 |     wva[0] += numpy.einsum('i,i,i->i', u_u_ud, rho1a[0], b0a1) * 2
1534 |     wva[0] += numpy.einsum('i,i,i->i', u_d_ud, rho1b[0], b0a1) * 2
1535 |     wva[0] += numpy.einsum('i,i,i->i', u_u_dd, rho1a[0], b0b1) * 4
1536 |     wva[0] += numpy.einsum('i,i,i->i', u_d_dd, rho1b[0], b0b1) * 4
1537 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_u_uu, rho1a[0], rho1a[0], rho0a[1:]) * 2
1538 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_d_uu, rho1a[0], rho1b[0], rho0a[1:]) * 4
1539 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_d_uu, rho1b[0], rho1b[0], rho0a[1:]) * 2
1540 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_u_ud, rho1a[0], rho1a[0], rho0a[1:])
1541 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_d_ud, rho1a[0], rho1b[0], rho0a[1:]) * 2
1542 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_d_ud, rho1b[0], rho1b[0], rho0a[1:])
1543 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_uu, rho1a[0], a0a1, rho0a[1:]) * 8
1544 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_uu, rho1b[0], a0a1, rho0a[1:]) * 8
1545 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_ud, rho1a[0], a0b1, rho0a[1:]) * 4
1546 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_ud, rho1b[0], a0b1, rho0a[1:]) * 4
1547 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_ud, rho1a[0], b0a1, rho0a[1:]) * 4
1548 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_ud, rho1b[0], b0a1, rho0a[1:]) * 4
1549 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_dd, rho1a[0], b0b1, rho0a[1:]) * 8
1550 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_dd, rho1b[0], b0b1, rho0a[1:]) * 8
1551 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_ud, rho1a[0], a0a1, rho0b[1:]) * 4
1552 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_ud, rho1b[0], a0a1, rho0b[1:]) * 4
1553 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_ud, rho1a[0], a0b1, rho0b[1:]) * 2
1554 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_ud, rho1b[0], a0b1, rho0b[1:]) * 2
1555 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_ud, rho1a[0], b0a1, rho0b[1:]) * 2
1556 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_ud, rho1b[0], b0a1, rho0b[1:]) * 2
1557 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_dd, rho1a[0], b0b1, rho0b[1:]) * 4
1558 |     wva[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_dd, rho1b[0], b0b1, rho0b[1:]) * 4
1559 |     wva[0] += numpy.einsum('i,i->i', u_uu_uu, a0a1_a0a1) * 4
1560 |     wva[0] += numpy.einsum('i,i->i', u_uu_ud, a0a1_a0b1) * 2
1561 |     wva[0] += numpy.einsum('i,i->i', u_uu_ud, a0b1_a0a1) * 2
1562 |     wva[0] += numpy.einsum('i,i->i', u_ud_ud, a0b1_a0b1)
1563 |     wva[0] += numpy.einsum('i,i->i', u_uu_ud, a0a1_b0a1) * 4
1564 |     wva[0] += numpy.einsum('i,i->i', u_ud_ud, a0b1_b0a1) * 2
1565 |     wva[0] += numpy.einsum('i,i->i', u_uu_dd, a0a1_b0b1) * 8
1566 |     wva[0] += numpy.einsum('i,i->i', u_ud_dd, a0b1_b0b1) * 4
1567 |     wva[0] += numpy.einsum('i,i->i', u_ud_ud, b0a1_b0a1)
1568 |     wva[0] += numpy.einsum('i,i->i', u_ud_dd, b0a1_b0b1) * 2
1569 |     wva[0] += numpy.einsum('i,i->i', u_ud_dd, b0b1_b0a1) * 2
1570 |     wva[0] += numpy.einsum('i,i->i', u_dd_dd, b0b1_b0b1) * 4
1571 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu_uu, a0a1_a0a1, rho0a[1:]) * 8
1572 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu_ud, a0a1_a0b1, rho0a[1:]) * 4
1573 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu_ud, a0b1_a0a1, rho0a[1:]) * 4
1574 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0b1_a0b1, rho0a[1:]) * 2
1575 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu_ud, a0a1_a0a1, rho0b[1:]) * 4
1576 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0a1_a0b1, rho0b[1:]) * 2
1577 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0b1_a0a1, rho0b[1:]) * 2
1578 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud_ud, a0b1_a0b1, rho0b[1:])
1579 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu_ud, a0a1_b0a1, rho0a[1:]) * 8
1580 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0b1_b0a1, rho0a[1:]) * 4
1581 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_uu_dd, a0a1_b0b1, rho0a[1:]) * 16
1582 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, a0b1_b0b1, rho0a[1:]) * 8
1583 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0a1_b0a1, rho0b[1:]) * 4
1584 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud_ud, a0b1_b0a1, rho0b[1:]) * 2
1585 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, a0a1_b0b1, rho0b[1:]) * 8
1586 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, a0b1_b0b1, rho0b[1:]) * 4
1587 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, b0a1_b0a1, rho0a[1:]) * 2
1588 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, b0b1_b0a1, rho0a[1:]) * 4
1589 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, b0a1_b0b1, rho0a[1:]) * 4
1590 |     wva[1:] += numpy.einsum('i,i,xi->xi', uu_dd_dd, b0b1_b0b1, rho0a[1:]) * 8
1591 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud_ud, b0a1_b0a1, rho0b[1:])
1592 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0b1_b0a1, rho0b[1:]) * 2
1593 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0a1_b0b1, rho0b[1:]) * 2
1594 |     wva[1:] += numpy.einsum('i,i,xi->xi', ud_dd_dd, b0b1_b0b1, rho0b[1:]) * 4
1595 |     wva *= weight
1596 |     wva[0]*=.5  # v+v.T should be applied in the caller
1597 | 
1598 |     wvb = numpy.zeros((4,ngrid))
1599 |     wvb[0] += numpy.einsum('i,i,i->i', d_d_d, rho1b[0], rho1b[0])
1600 |     wvb[0] += numpy.einsum('i,i,i->i', u_d_d, rho1b[0], rho1a[0]) * 2
1601 |     wvb[0] += numpy.einsum('i,i,i->i', u_u_d, rho1a[0], rho1a[0])
1602 |     wvb[0] += numpy.einsum('i,i->i', d_dd, b1b1) * 2
1603 |     wvb[0] += numpy.einsum('i,i->i', d_ud, b1a1) * 2
1604 |     wvb[0] += numpy.einsum('i,i->i', d_uu, a1a1) * 2
1605 |     wvb[1:] += numpy.einsum('i,i,xi->xi', d_dd, rho1b[0], rho1b[1:]) * 4
1606 |     wvb[1:] += numpy.einsum('i,i,xi->xi', u_dd, rho1a[0], rho1b[1:]) * 4
1607 |     wvb[1:] += numpy.einsum('i,i,xi->xi', d_ud, rho1b[0], rho1a[1:]) * 2
1608 |     wvb[1:] += numpy.einsum('i,i,xi->xi', u_ud, rho1a[0], rho1a[1:]) * 2
1609 |     wvb[1:] += numpy.einsum('i,i,xi->xi', dd_dd, b0b1, rho1b[1:]) * 8
1610 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd, b0b1, rho1a[1:]) * 4
1611 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd, b0a1, rho1b[1:]) * 4
1612 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud, b0a1, rho1a[1:]) * 2
1613 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd, a0b1, rho1b[1:]) * 4
1614 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_dd, a0a1, rho1b[1:]) * 8
1615 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud, a0b1, rho1a[1:]) * 2
1616 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud, a0a1, rho1a[1:]) * 4
1617 |     wvb[1:] += numpy.einsum('i,i,xi->xi', dd_dd, b1b1, rho0b[1:]) * 4
1618 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd, b1a1, rho0b[1:]) * 4
1619 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_dd, a1a1, rho0b[1:]) * 4
1620 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd, b1b1, rho0a[1:]) * 2
1621 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud, b1a1, rho0a[1:]) * 2
1622 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud, a1a1, rho0a[1:]) * 2
1623 |     wvb[0] += numpy.einsum('i,i,i->i', d_d_dd, rho1b[0], b0b1) * 4
1624 |     wvb[0] += numpy.einsum('i,i,i->i', u_d_dd, rho1a[0], b0b1) * 4
1625 |     wvb[0] += numpy.einsum('i,i,i->i', d_d_ud, rho1b[0], b0a1) * 2
1626 |     wvb[0] += numpy.einsum('i,i,i->i', u_d_ud, rho1a[0], b0a1) * 2
1627 |     wvb[0] += numpy.einsum('i,i,i->i', d_d_ud, rho1b[0], a0b1) * 2
1628 |     wvb[0] += numpy.einsum('i,i,i->i', u_d_ud, rho1a[0], a0b1) * 2
1629 |     wvb[0] += numpy.einsum('i,i,i->i', d_d_uu, rho1b[0], a0a1) * 4
1630 |     wvb[0] += numpy.einsum('i,i,i->i', u_d_uu, rho1a[0], a0a1) * 4
1631 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_d_dd, rho1b[0], rho1b[0], rho0b[1:]) * 2
1632 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_d_dd, rho1b[0], rho1a[0], rho0b[1:]) * 4
1633 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_u_dd, rho1a[0], rho1a[0], rho0b[1:]) * 2
1634 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_d_ud, rho1b[0], rho1b[0], rho0b[1:])
1635 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_d_ud, rho1b[0], rho1a[0], rho0b[1:]) * 2
1636 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_u_ud, rho1a[0], rho1a[0], rho0b[1:])
1637 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_dd_dd, rho1b[0], b0b1, rho0b[1:]) * 8
1638 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_dd_dd, rho1a[0], b0b1, rho0b[1:]) * 8
1639 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_dd, rho1b[0], b0a1, rho0b[1:]) * 4
1640 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_dd, rho1a[0], b0a1, rho0b[1:]) * 4
1641 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_dd, rho1b[0], a0b1, rho0b[1:]) * 4
1642 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_dd, rho1a[0], a0b1, rho0b[1:]) * 4
1643 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_dd, rho1b[0], a0a1, rho0b[1:]) * 8
1644 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_dd, rho1a[0], a0a1, rho0b[1:]) * 8
1645 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_dd, rho1b[0], b0b1, rho0a[1:]) * 4
1646 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_dd, rho1a[0], b0b1, rho0a[1:]) * 4
1647 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_ud, rho1b[0], b0a1, rho0a[1:]) * 2
1648 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_ud, rho1a[0], b0a1, rho0a[1:]) * 2
1649 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_ud_ud, rho1b[0], a0b1, rho0a[1:]) * 2
1650 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_ud_ud, rho1a[0], a0b1, rho0a[1:]) * 2
1651 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', d_uu_ud, rho1b[0], a0a1, rho0a[1:]) * 4
1652 |     wvb[1:] += numpy.einsum('i,i,i,xi->xi', u_uu_ud, rho1a[0], a0a1, rho0a[1:]) * 4
1653 |     wvb[0] += numpy.einsum('i,i->i', d_dd_dd, b0b1_b0b1) * 4
1654 |     wvb[0] += numpy.einsum('i,i->i', d_ud_dd, b0b1_b0a1) * 2
1655 |     wvb[0] += numpy.einsum('i,i->i', d_ud_dd, b0a1_b0b1) * 2
1656 |     wvb[0] += numpy.einsum('i,i->i', d_ud_ud, b0a1_b0a1)
1657 |     wvb[0] += numpy.einsum('i,i->i', d_ud_dd, b0b1_a0b1) * 4
1658 |     wvb[0] += numpy.einsum('i,i->i', d_ud_ud, b0a1_a0b1) * 2
1659 |     wvb[0] += numpy.einsum('i,i->i', d_uu_dd, b0b1_a0a1) * 8
1660 |     wvb[0] += numpy.einsum('i,i->i', d_uu_ud, b0a1_a0a1) * 4
1661 |     wvb[0] += numpy.einsum('i,i->i', d_ud_ud, a0b1_a0b1)
1662 |     wvb[0] += numpy.einsum('i,i->i', d_uu_ud, a0b1_a0a1) * 2
1663 |     wvb[0] += numpy.einsum('i,i->i', d_uu_ud, a0a1_a0b1) * 2
1664 |     wvb[0] += numpy.einsum('i,i->i', d_uu_uu, a0a1_a0a1) * 4
1665 |     wvb[1:] += numpy.einsum('i,i,xi->xi', dd_dd_dd, b0b1_b0b1, rho0b[1:]) * 8
1666 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd_dd, b0b1_b0a1, rho0b[1:]) * 4
1667 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd_dd, b0a1_b0b1, rho0b[1:]) * 4
1668 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0a1_b0a1, rho0b[1:]) * 2
1669 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd_dd, b0b1_b0b1, rho0a[1:]) * 4
1670 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0b1_b0a1, rho0a[1:]) * 2
1671 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0a1_b0b1, rho0a[1:]) * 2
1672 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_ud, b0a1_b0a1, rho0a[1:])
1673 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_dd_dd, b0b1_a0b1, rho0b[1:]) * 8
1674 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0a1_a0b1, rho0b[1:]) * 4
1675 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_dd_dd, b0b1_a0a1, rho0b[1:]) * 16
1676 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, b0a1_a0a1, rho0b[1:]) * 8
1677 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, b0b1_a0b1, rho0a[1:]) * 4
1678 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_ud, b0a1_a0b1, rho0a[1:]) * 2
1679 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, b0b1_a0a1, rho0a[1:]) * 8
1680 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, b0a1_a0a1, rho0a[1:]) * 4
1681 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_dd, a0b1_a0b1, rho0b[1:]) * 2
1682 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, a0a1_a0b1, rho0b[1:]) * 4
1683 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_dd, a0b1_a0a1, rho0b[1:]) * 4
1684 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_uu_dd, a0a1_a0a1, rho0b[1:]) * 8
1685 |     wvb[1:] += numpy.einsum('i,i,xi->xi', ud_ud_ud, a0b1_a0b1, rho0a[1:])
1686 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0a1_a0b1, rho0a[1:]) * 2
1687 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_ud_ud, a0b1_a0a1, rho0a[1:]) * 2
1688 |     wvb[1:] += numpy.einsum('i,i,xi->xi', uu_uu_ud, a0a1_a0a1, rho0a[1:]) * 4
1689 |     wvb *= weight
1690 |     wvb[0]*=.5
1691 | 
1692 |     return wva, wvb
1693 | 
1694 | def nr_fxc(mol, grids, xc_code, dm0, dms, spin=0, relativity=0, hermi=0,
1695 |            rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
1696 |     r'''Contract XC kernel matrix with given density matrices
1697 |     ... math::
1698 |             a_{pq} = f_{pq,rs} * x_{rs}
1699 |     '''
1700 |     ni = NumInt()
1701 |     return ni.nr_fxc(mol, grids, xc_code, dm0, dms, spin, relativity,
1702 |                      hermi, rho0, vxc, fxc, max_memory, verbose)
1703 | 
1704 | 
1705 | def cache_xc_kernel(ni, mol, grids, xc_code, mo_coeff, mo_occ, spin=0,
1706 |                     max_memory=2000):
1707 |     '''Compute the 0th order density, Vxc and fxc.  They can be used in TDDFT,
1708 |     DFT hessian module etc.
1709 |     '''
1710 |     xctype = ni._xc_type(xc_code)
1711 |     ao_deriv = 0
1712 |     if xctype == 'GGA':
1713 |         ao_deriv = 1
1714 |     elif xctype == 'NLC':
1715 |         raise NotImplementedError('NLC')
1716 |     elif xctype == 'MGGA':
1717 |         raise NotImplementedError('meta-GGA')
1718 | 
1719 |     if spin == 0:
1720 |         nao = mo_coeff.shape[0]
1721 |         rho = []
1722 |         for ao, mask, weight, coords \
1723 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory=max_memory):
1724 |             rho.append(ni.eval_rho2(mol, ao, mo_coeff, mo_occ, mask, xctype))
1725 |         rho = numpy.hstack(rho)
1726 |     else:
1727 |         nao = mo_coeff[0].shape[0]
1728 |         rhoa = []
1729 |         rhob = []
1730 |         for ao, mask, weight, coords \
1731 |                 in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
1732 |             rhoa.append(ni.eval_rho2(mol, ao, mo_coeff[0], mo_occ[0], mask, xctype))
1733 |             rhob.append(ni.eval_rho2(mol, ao, mo_coeff[1], mo_occ[1], mask, xctype))
1734 |         rho = (numpy.hstack(rhoa), numpy.hstack(rhob))
1735 |     vxc, fxc = ni.eval_xc(xc_code, rho, spin, 0, 2, 0)[1:3]
1736 |     return rho, vxc, fxc
1737 | 
1738 | def get_rho(ni, mol, dm, grids, max_memory=2000):
1739 |     '''Density in real space
1740 |     '''
1741 |     make_rho, nset, nao = ni._gen_rho_evaluator(mol, dm, 1)
1742 |     rho = numpy.empty(grids.weights.size)
1743 |     p1 = 0
1744 |     for ao, mask, weight, coords \
1745 |             in ni.block_loop(mol, grids, nao, 0, max_memory):
1746 |         p0, p1 = p1, p1 + weight.size
1747 |         rho[p0:p1] = make_rho(0, ao, mask, 'LDA')
1748 |     return rho
1749 | 
1750 | 
1751 | class NumInt(object):
1752 |     def __init__(self):
1753 |         self.libxc = libxc
1754 | 
1755 |     @lib.with_doc(nr_vxc.__doc__)
1756 |     def nr_vxc(self, mol, grids, xc_code, dms, spin=0, relativity=0, hermi=0,
1757 |                max_memory=2000, verbose=None):
1758 |         if spin == 0:
1759 |             return self.nr_rks(mol, grids, xc_code, dms, relativity, hermi,
1760 |                                max_memory, verbose)
1761 |         else:
1762 |             return self.nr_uks(mol, grids, xc_code, dms, relativity, hermi,
1763 |                                max_memory, verbose)
1764 | 
1765 |     @lib.with_doc(nr_fxc.__doc__)
1766 |     def nr_fxc(self, mol, grids, xc_code, dm0, dms, spin=0, relativity=0, hermi=0,
1767 |                rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
1768 |         if spin == 0:
1769 |             return self.nr_rks_fxc(mol, grids, xc_code, dm0, dms, relativity,
1770 |                                    hermi, rho0, vxc, fxc, max_memory, verbose)
1771 |         else:
1772 |             return self.nr_uks_fxc(mol, grids, xc_code, dm0, dms, relativity,
1773 |                                    hermi, rho0, vxc, fxc, max_memory, verbose)
1774 | 
1775 |     nr_rks = nr_rks
1776 |     nr_uks = nr_uks
1777 |     nr_rks_fxc = nr_rks_fxc
1778 |     nr_uks_fxc = nr_uks_fxc
1779 |     cache_xc_kernel  = cache_xc_kernel
1780 |     get_rho = get_rho
1781 | 
1782 |     @lib.with_doc(eval_ao.__doc__)
1783 |     def eval_ao(self, mol, coords, deriv=0, shls_slice=None,
1784 |                 non0tab=None, out=None, verbose=None):
1785 |         return eval_ao(mol, coords, deriv, shls_slice, non0tab, out, verbose)
1786 | 
1787 |     @lib.with_doc(make_mask.__doc__)
1788 |     def make_mask(self, mol, coords, relativity=0, shls_slice=None,
1789 |                   verbose=None):
1790 |         return make_mask(mol, coords, relativity, shls_slice, verbose)
1791 | 
1792 |     @lib.with_doc(eval_rho2.__doc__)
1793 |     def eval_rho2(self, mol, ao, mo_coeff, mo_occ, non0tab=None, xctype='LDA',
1794 |                   verbose=None):
1795 |         return eval_rho2(mol, ao, mo_coeff, mo_occ, non0tab, xctype, verbose)
1796 | 
1797 |     @lib.with_doc(eval_rho.__doc__)
1798 |     def eval_rho(self, mol, ao, dm, non0tab=None, xctype='LDA', hermi=0, verbose=None):
1799 |         return eval_rho(mol, ao, dm, non0tab, xctype, hermi, verbose)
1800 | 
1801 |     def block_loop(self, mol, grids, nao, deriv=0, max_memory=2000,
1802 |                    non0tab=None, blksize=None, buf=None):
1803 |         '''Define this macro to loop over grids by blocks.
1804 |         '''
1805 |         if grids.coords is None:
1806 |             grids.build(with_non0tab=True)
1807 |         ngrids = grids.coords.shape[0]
1808 |         comp = (deriv+1)*(deriv+2)*(deriv+3)//6
1809 | # NOTE to index grids.non0tab, the blksize needs to be the integer multiplier of BLKSIZE
1810 |         if blksize is None:
1811 |             blksize = int(max_memory*1e6/(comp*2*nao*8*BLKSIZE))*BLKSIZE
1812 |             blksize = max(BLKSIZE, min(blksize, ngrids, BLKSIZE*1200))
1813 |         if non0tab is None:
1814 |             non0tab = grids.non0tab
1815 |         if non0tab is None:
1816 |             non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
1817 |                                  dtype=numpy.uint8)
1818 |         if buf is None:
1819 |             buf = numpy.empty((comp,blksize,nao))
1820 |         for ip0 in range(0, ngrids, blksize):
1821 |             ip1 = min(ngrids, ip0+blksize)
1822 |             coords = grids.coords[ip0:ip1]
1823 |             weight = grids.weights[ip0:ip1]
1824 |             non0 = non0tab[ip0//BLKSIZE:]
1825 |             ao = self.eval_ao(mol, coords, deriv=deriv, non0tab=non0, out=buf)
1826 |             yield ao, non0, weight, coords
1827 | 
1828 |     def _gen_rho_evaluator(self, mol, dms, hermi=0):
1829 |         if hasattr(dms, 'mo_coeff'):
1830 | #TODO: test whether dm.mo_coeff matching dm
1831 |             mo_coeff = dms.mo_coeff
1832 |             mo_occ = dms.mo_occ
1833 |             if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
1834 |                 mo_coeff = [mo_coeff]
1835 |                 mo_occ = [mo_occ]
1836 |             nao = mo_coeff[0].shape[0]
1837 |             ndms = len(mo_occ)
1838 |             def make_rho(idm, ao, non0tab, xctype):
1839 |                 return self.eval_rho2(mol, ao, mo_coeff[idm], mo_occ[idm],
1840 |                                       non0tab, xctype)
1841 |         else:
1842 |             if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
1843 |                 dms = [dms]
1844 |             if not hermi:
1845 | # For eval_rho when xctype==GGA, which requires hermitian DMs
1846 |                 dms = [(dm+dm.conj().T)*.5 for dm in dms]
1847 |             nao = dms[0].shape[0]
1848 |             ndms = len(dms)
1849 |             def make_rho(idm, ao, non0tab, xctype):
1850 |                 return self.eval_rho(mol, ao, dms[idm], non0tab, xctype, hermi=1)
1851 |         return make_rho, ndms, nao
1852 | 
1853 | ####################
1854 | # Overwrite following functions to use custom XC functional
1855 | 
1856 |     def hybrid_coeff(self, xc_code, spin=0):
1857 |         return self.libxc.hybrid_coeff(xc_code, spin)
1858 | 
1859 |     def nlc_coeff(self, xc_code):
1860 |         return self.libxc.nlc_coeff(xc_code)
1861 | 
1862 |     def rsh_coeff(self, xc_code):
1863 |         return self.libxc.rsh_coeff(xc_code)
1864 | 
1865 |     def eval_xc(self, xc_code, rho, spin=0, relativity=0, deriv=1, verbose=None):
1866 |         return self.libxc.eval_xc(xc_code, rho, spin, relativity, deriv, verbose)
1867 |     eval_xc.__doc__ = libxc.eval_xc.__doc__
1868 | 
1869 |     def _xc_type(self, xc_code):
1870 |         return self.libxc.xc_type(xc_code)
1871 | 
1872 |     def rsh_and_hybrid_coeff(self, xc_code, spin=0):
1873 |         '''Range-separated parameter and HF exchange components: omega, alpha, beta
1874 |         Exc_RSH = c_SR * SR_HFX + c_LR * LR_HFX + (1-c_SR) * Ex_SR + (1-c_LR) * Ex_LR + Ec
1875 |                 = alpha * HFX + beta * SR_HFX + (1-c_SR) * Ex_SR + (1-c_LR) * Ex_LR + Ec
1876 |                 = alpha * LR_HFX + hyb * SR_HFX + (1-c_SR) * Ex_SR + (1-c_LR) * Ex_LR + Ec
1877 |         SR_HFX = < pi | e^{-omega r_{12}}/r_{12} | iq >
1878 |         LR_HFX = < pi | (1-e^{-omega r_{12}})/r_{12} | iq >
1879 |         alpha = c_LR
1880 |         beta = c_SR - c_LR
1881 |         '''
1882 |         omega, alpha, beta = self.rsh_coeff(xc_code)
1883 |         if abs(omega) > 1e-10:
1884 |             hyb = alpha + beta
1885 |         else:
1886 |             hyb = self.hybrid_coeff(xc_code, spin)
1887 |         return omega, alpha, hyb
1888 | _NumInt = NumInt
1889 | 
1890 | 
1891 | if __name__ == '__main__':
1892 |     import time
1893 |     from pyscf import gto
1894 |     from pyscf import dft
1895 | 
1896 |     mol = gto.M(
1897 |         atom = [
1898 |         ["O" , (0. , 0.     , 0.)],
1899 |         [1   , (0. , -0.757 , 0.587)],
1900 |         [1   , (0. , 0.757  , 0.587)] ],
1901 |         basis = '6311g**',)
1902 |     mf = dft.RKS(mol)
1903 |     mf.grids.atom_grid = {"H": (30, 194), "O": (30, 194),}
1904 |     mf.grids.prune = None
1905 |     mf.grids.build()
1906 |     dm = mf.get_init_guess(key='minao')
1907 | 
1908 |     numpy.random.seed(1)
1909 |     dm1 = numpy.random.random((dm.shape))
1910 |     dm1 = lib.hermi_triu(dm1)
1911 |     print(time.clock())
1912 |     res = mf._numint.nr_vxc(mol, mf.grids, mf.xc, dm1, spin=0)
1913 |     print(res[1] - -37.084047825971282)
1914 |     res = mf._numint.nr_vxc(mol, mf.grids, mf.xc, (dm1,dm1), spin=1)
1915 |     print(res[1] - -92.436362308687094)
1916 |     res = mf._numint.nr_vxc(mol, mf.grids, mf.xc, dm, spin=0)
1917 |     print(res[1] - -8.6313329288394947)
1918 |     res = mf._numint.nr_vxc(mol, mf.grids, mf.xc, (dm,dm), spin=1)
1919 |     print(res[1] - -21.520301399504582)
1920 |     print(time.clock())


--------------------------------------------------------------------------------