1#!/usr/bin/env python
2# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16from functools import reduce
17import unittest
18import numpy
19import scipy.linalg
20from pyscf import lib
21from pyscf import gto
22from pyscf.x2c import sfx2c1e
23from pyscf.x2c import sfx2c1e_grad
24from pyscf.x2c import sfx2c1e_hess
25
26def _sqrt0(a):
27    w, v = scipy.linalg.eigh(a)
28    return numpy.dot(v*numpy.sqrt(w), v.conj().T)
29
30def _invsqrt0(a):
31    w, v = scipy.linalg.eigh(a)
32    return numpy.dot(v/numpy.sqrt(w), v.conj().T)
33
34def _sqrt1(a0, a1):
35    '''Solving first order derivative of x^2 = a'''
36    w, v = scipy.linalg.eigh(a0)
37    w = numpy.sqrt(w)
38    a1 = reduce(numpy.dot, (v.conj().T, a1, v))
39    x1 = a1 / (w[:,None] + w)
40    x1 = reduce(numpy.dot, (v, x1, v.conj().T))
41    return x1
42
43def _sqrt2(a0, a1i, a1j, a2ij):
44    '''Solving second order derivative of x^2 = a'''
45    w, v = scipy.linalg.eigh(a0)
46    w = numpy.sqrt(w)
47    a1i = reduce(numpy.dot, (v.conj().T, a1i, v))
48    x1i = a1i / (w[:,None] + w)
49
50    a1j = reduce(numpy.dot, (v.conj().T, a1j, v))
51    x1j = a1j / (w[:,None] + w)
52
53    a2ij = reduce(numpy.dot, (v.conj().T, a2ij, v))
54    tmp = x1i.dot(x1j)
55    a2ij -= tmp + tmp.conj().T
56    x2 = a2ij / (w[:,None] + w)
57    x2 = reduce(numpy.dot, (v, x2, v.conj().T))
58    return x2
59
60def _invsqrt1(a0, a1):
61    '''Solving first order derivative of x^2 = a^{-1}'''
62    w, v = scipy.linalg.eigh(a0)
63    w = 1./numpy.sqrt(w)
64    a1 = -reduce(numpy.dot, (v.conj().T, a1, v))
65    x1 = numpy.einsum('i,ij,j->ij', w**2, a1, w**2) / (w[:,None] + w)
66    x1 = reduce(numpy.dot, (v, x1, v.conj().T))
67    return x1
68
69def _invsqrt2(a0, a1i, a1j, a2ij):
70    '''Solving first order derivative of x^2 = a^{-1}'''
71    w, v = scipy.linalg.eigh(a0)
72    w = 1./numpy.sqrt(w)
73    a1i = reduce(numpy.dot, (v.conj().T, a1i, v))
74    x1i = numpy.einsum('i,ij,j->ij', w**2, a1i, w**2) / (w[:,None] + w)
75
76    a1j = reduce(numpy.dot, (v.conj().T, a1j, v))
77    x1j = numpy.einsum('i,ij,j->ij', w**2, a1j, w**2) / (w[:,None] + w)
78
79    a2ij = reduce(numpy.dot, (v.conj().T, a2ij, v))
80    tmp = (a1i*w**2).dot(a1j)
81    a2ij -= tmp + tmp.conj().T
82    a2ij = -numpy.einsum('i,ij,j->ij', w**2, a2ij, w**2)
83    tmp = x1i.dot(x1j)
84    a2ij -= tmp + tmp.conj().T
85    x2 = a2ij / (w[:,None] + w)
86    x2 = reduce(numpy.dot, (v, x2, v.conj().T))
87    return x2
88
89def get_h0_s0(mol):
90    s = mol.intor_symmetric('int1e_ovlp')
91    t = mol.intor_symmetric('int1e_kin')
92    v = mol.intor_symmetric('int1e_nuc')
93    w = mol.intor_symmetric('int1e_pnucp')
94    nao = s.shape[0]
95    n2 = nao * 2
96    h = numpy.zeros((n2,n2), dtype=v.dtype)
97    m = numpy.zeros((n2,n2), dtype=v.dtype)
98    c = lib.param.LIGHT_SPEED
99    h[:nao,:nao] = v
100    h[:nao,nao:] = t
101    h[nao:,:nao] = t
102    h[nao:,nao:] = w * (.25/c**2) - t
103    m[:nao,:nao] = s
104    m[nao:,nao:] = t * (.5/c**2)
105    return h, m
106
107def get_h1_s1(mol, ia):
108    aoslices = mol.aoslice_by_atom()
109    ish0, ish1, p0, p1 = aoslices[ia]
110    nao = mol.nao_nr()
111    s1 = mol.intor('int1e_ipovlp', comp=3)
112    t1 = mol.intor('int1e_ipkin', comp=3)
113    v1 = mol.intor('int1e_ipnuc', comp=3)
114    w1 = mol.intor('int1e_ipspnucsp', comp=12).reshape(3,4,nao,nao)[:,3]
115    with mol.with_rinv_origin(mol.atom_coord(ia)):
116        z = -mol.atom_charge(ia)
117        rinv1 = z*mol.intor('int1e_iprinv', comp=3)
118        prinvp1 = z*mol.intor('int1e_ipsprinvsp', comp=12).reshape(3,4,nao,nao)[:,3]
119    n2 = nao * 2
120    h = numpy.zeros((3,n2,n2), dtype=v1.dtype)
121    m = numpy.zeros((3,n2,n2), dtype=v1.dtype)
122    rinv1[:,p0:p1,:] -= v1[:,p0:p1]
123    rinv1 = rinv1 + rinv1.transpose(0,2,1).conj()
124    prinvp1[:,p0:p1,:] -= w1[:,p0:p1]
125    prinvp1 = prinvp1 + prinvp1.transpose(0,2,1).conj()
126
127    s1ao = numpy.zeros_like(s1)
128    t1ao = numpy.zeros_like(t1)
129    s1ao[:,p0:p1,:] = -s1[:,p0:p1]
130    s1ao[:,:,p0:p1]+= -s1[:,p0:p1].transpose(0,2,1)
131    t1ao[:,p0:p1,:] = -t1[:,p0:p1]
132    t1ao[:,:,p0:p1]+= -t1[:,p0:p1].transpose(0,2,1)
133
134    c = lib.param.LIGHT_SPEED
135    h[:,:nao,:nao] = rinv1
136    h[:,:nao,nao:] = t1ao
137    h[:,nao:,:nao] = t1ao
138    h[:,nao:,nao:] = prinvp1 * (.25/c**2) - t1ao
139    m[:,:nao,:nao] = s1ao
140    m[:,nao:,nao:] = t1ao * (.5/c**2)
141    return h, m
142
143def get_h2_s2(mol, ia, ja):
144    aoslices = mol.aoslice_by_atom()
145    ish0, ish1, i0, i1 = aoslices[ia]
146    jsh0, jsh1, j0, j1 = aoslices[ja]
147    nao = mol.nao_nr()
148    s2aa = mol.intor('int1e_ipipovlp', comp=9).reshape(3,3,nao,nao)
149    t2aa = mol.intor('int1e_ipipkin', comp=9).reshape(3,3,nao,nao)
150    v2aa = mol.intor('int1e_ipipnuc', comp=9).reshape(3,3,nao,nao)
151    w2aa = mol.intor('int1e_ipippnucp', comp=9).reshape(3,3,nao,nao)
152    s2ab = mol.intor('int1e_ipovlpip', comp=9).reshape(3,3,nao,nao)
153    t2ab = mol.intor('int1e_ipkinip', comp=9).reshape(3,3,nao,nao)
154    v2ab = mol.intor('int1e_ipnucip', comp=9).reshape(3,3,nao,nao)
155    w2ab = mol.intor('int1e_ippnucpip', comp=9).reshape(3,3,nao,nao)
156
157    n2 = nao * 2
158    h2 = numpy.zeros((3,3,n2,n2), dtype=v2aa.dtype)
159    m2 = numpy.zeros((3,3,n2,n2), dtype=v2aa.dtype)
160    s2ao = numpy.zeros_like(s2aa)
161    t2ao = numpy.zeros_like(s2aa)
162    v2ao = numpy.zeros_like(s2aa)
163    w2ao = numpy.zeros_like(s2aa)
164    if ia == ja:
165        with mol.with_rinv_origin(mol.atom_coord(ia)):
166            z = mol.atom_charge(ia)
167            rinv2aa = z*mol.intor('int1e_ipiprinv', comp=9).reshape(3,3,nao,nao)
168            rinv2ab = z*mol.intor('int1e_iprinvip', comp=9).reshape(3,3,nao,nao)
169            prinvp2aa = z*mol.intor('int1e_ipipprinvp', comp=9).reshape(3,3,nao,nao)
170            prinvp2ab = z*mol.intor('int1e_ipprinvpip', comp=9).reshape(3,3,nao,nao)
171        s2ao[:,:,i0:i1      ] = s2aa[:,:,i0:i1      ]
172        s2ao[:,:,i0:i1,j0:j1]+= s2ab[:,:,i0:i1,j0:j1]
173        t2ao[:,:,i0:i1      ] = t2aa[:,:,i0:i1      ]
174        t2ao[:,:,i0:i1,j0:j1]+= t2ab[:,:,i0:i1,j0:j1]
175        v2ao -= rinv2aa + rinv2ab
176        v2ao[:,:,i0:i1      ]+= v2aa[:,:,i0:i1      ]
177        v2ao[:,:,i0:i1,j0:j1]+= v2ab[:,:,i0:i1,j0:j1]
178        v2ao[:,:,i0:i1      ]+= rinv2aa[:,:,i0:i1] * 2
179        v2ao[:,:,i0:i1      ]+= rinv2ab[:,:,i0:i1] * 2
180        w2ao -= prinvp2aa + prinvp2ab
181        w2ao[:,:,i0:i1      ]+= w2aa[:,:,i0:i1      ]
182        w2ao[:,:,i0:i1,j0:j1]+= w2ab[:,:,i0:i1,j0:j1]
183        w2ao[:,:,i0:i1      ]+= prinvp2aa[:,:,i0:i1] * 2
184        w2ao[:,:,i0:i1      ]+= prinvp2ab[:,:,i0:i1] * 2
185        s2ao = s2ao + s2ao.transpose(0,1,3,2)
186        t2ao = t2ao + t2ao.transpose(0,1,3,2)
187        v2ao = v2ao + v2ao.transpose(0,1,3,2)
188        w2ao = w2ao + w2ao.transpose(0,1,3,2)
189    else:
190        s2ao[:,:,i0:i1,j0:j1] = s2ab[:,:,i0:i1,j0:j1]
191        t2ao[:,:,i0:i1,j0:j1] = t2ab[:,:,i0:i1,j0:j1]
192        v2ao[:,:,i0:i1,j0:j1] = v2ab[:,:,i0:i1,j0:j1]
193        w2ao[:,:,i0:i1,j0:j1] = w2ab[:,:,i0:i1,j0:j1]
194        zi = mol.atom_charge(ia)
195        zj = mol.atom_charge(ja)
196        with mol.with_rinv_at_nucleus(ia):
197            shls_slice = (jsh0, jsh1, 0, mol.nbas)
198            rinv2aa = mol.intor('int1e_ipiprinv', comp=9, shls_slice=shls_slice)
199            rinv2ab = mol.intor('int1e_iprinvip', comp=9, shls_slice=shls_slice)
200            prinvp2aa = mol.intor('int1e_ipipprinvp', comp=9, shls_slice=shls_slice)
201            prinvp2ab = mol.intor('int1e_ipprinvpip', comp=9, shls_slice=shls_slice)
202            rinv2aa = zi * rinv2aa.reshape(3,3,j1-j0,nao)
203            rinv2ab = zi * rinv2ab.reshape(3,3,j1-j0,nao)
204            prinvp2aa = zi * prinvp2aa.reshape(3,3,j1-j0,nao)
205            prinvp2ab = zi * prinvp2ab.reshape(3,3,j1-j0,nao)
206            v2ao[:,:,j0:j1] += rinv2aa
207            v2ao[:,:,j0:j1] += rinv2ab.transpose(1,0,2,3)
208            w2ao[:,:,j0:j1] += prinvp2aa
209            w2ao[:,:,j0:j1] += prinvp2ab.transpose(1,0,2,3)
210
211        with mol.with_rinv_at_nucleus(ja):
212            shls_slice = (ish0, ish1, 0, mol.nbas)
213            rinv2aa = mol.intor('int1e_ipiprinv', comp=9, shls_slice=shls_slice)
214            rinv2ab = mol.intor('int1e_iprinvip', comp=9, shls_slice=shls_slice)
215            prinvp2aa = mol.intor('int1e_ipipprinvp', comp=9, shls_slice=shls_slice)
216            prinvp2ab = mol.intor('int1e_ipprinvpip', comp=9, shls_slice=shls_slice)
217            rinv2aa = zj * rinv2aa.reshape(3,3,i1-i0,nao)
218            rinv2ab = zj * rinv2ab.reshape(3,3,i1-i0,nao)
219            prinvp2aa = zj * prinvp2aa.reshape(3,3,i1-i0,nao)
220            prinvp2ab = zj * prinvp2ab.reshape(3,3,i1-i0,nao)
221            v2ao[:,:,i0:i1] += rinv2aa
222            v2ao[:,:,i0:i1] += rinv2ab
223            w2ao[:,:,i0:i1] += prinvp2aa
224            w2ao[:,:,i0:i1] += prinvp2ab
225        s2ao = s2ao + s2ao.transpose(0,1,3,2)
226        t2ao = t2ao + t2ao.transpose(0,1,3,2)
227        v2ao = v2ao + v2ao.transpose(0,1,3,2)
228        w2ao = w2ao + w2ao.transpose(0,1,3,2)
229
230    c = lib.param.LIGHT_SPEED
231    h2[:,:,:nao,:nao] = v2ao
232    h2[:,:,:nao,nao:] = t2ao
233    h2[:,:,nao:,:nao] = t2ao
234    h2[:,:,nao:,nao:] = w2ao * (.25/c**2) - t2ao
235    m2[:,:,:nao,:nao] = s2ao
236    m2[:,:,nao:,nao:] = t2ao * (.5/c**2)
237    return h2, m2
238
239def get_x0(mol):
240    c = lib.param.LIGHT_SPEED
241    h0, s0 = get_h0_s0(mol)
242    e, c = scipy.linalg.eigh(h0, s0)
243    nao = mol.nao_nr()
244    cl = c[:nao,nao:]
245    cs = c[nao:,nao:]
246    x0 = scipy.linalg.solve(cl.T, cs.T).T
247    return x0
248
249def get_x1(mol, ia):
250    h0, s0 = get_h0_s0(mol)
251    h1, s1 = get_h1_s1(mol, ia)
252    e0, c0 = scipy.linalg.eigh(h0, s0)
253    c0[:,c0[1]<0] *= -1
254    nao = mol.nao_nr()
255    cl0 = c0[:nao,nao:]
256    cs0 = c0[nao:,nao:]
257    x0 = scipy.linalg.solve(cl0.T, cs0.T).T
258    h1 = numpy.einsum('pi,xpq,qj->xij', c0.conj(), h1, c0[:,nao:])
259    s1 = numpy.einsum('pi,xpq,qj->xij', c0.conj(), s1, c0[:,nao:])
260    epi = e0[:,None] - e0[nao:]
261    degen_mask = abs(epi) < 1e-7
262    epi[degen_mask] = 1e200
263    c1 = (h1 - s1 * e0[nao:]) / -epi
264#    c1[:,degen_mask] = -.5 * s1[:,degen_mask]
265    c1 = numpy.einsum('pq,xqi->xpi', c0, c1)
266    cl1 = c1[:,:nao]
267    cs1 = c1[:,nao:]
268    x1 = [scipy.linalg.solve(cl0.T, (cs1[i] - x0.dot(cl1[i])).T).T
269          for i in range(3)]
270    return numpy.asarray(x1)
271
272def get_x2(mol, ia, ja):
273    h0, s0 = get_h0_s0(mol)
274    e0, c0 = scipy.linalg.eigh(h0, s0)
275    c0[:,c0[1]<0] *= -1
276    nao = mol.nao_nr()
277    cl0 = c0[:nao,nao:]
278    cs0 = c0[nao:,nao:]
279    x0 = scipy.linalg.solve(cl0.T, cs0.T).T
280    epq = e0[:,None] - e0
281    degen_mask = abs(epq) < 1e-7
282    epq[degen_mask] = 1e200
283
284    h1i, s1i = get_h1_s1(mol, ia)
285    h1i = numpy.einsum('pi,xpq,qj->xij', c0.conj(), h1i, c0)
286    s1i = numpy.einsum('pi,xpq,qj->xij', c0.conj(), s1i, c0)
287    c1i = (h1i - s1i * e0) / -epq
288    c1i[:,degen_mask] = -.5 * s1i[:,degen_mask]
289    e1i = h1i - s1i * e0
290    e1i[:,~degen_mask] = 0
291    c1i_ao = numpy.einsum('pq,xqi->xpi', c0, c1i[:,:,nao:])
292    cl1i = c1i_ao[:,:nao]
293    cs1i = c1i_ao[:,nao:]
294    x1i = [scipy.linalg.solve(cl0.T, (cs1i[i] - x0.dot(cl1i[i])).T).T
295           for i in range(3)]
296
297    h1j, s1j = get_h1_s1(mol, ja)
298    h1j = numpy.einsum('pi,xpq,qj->xij', c0.conj(), h1j, c0)
299    s1j = numpy.einsum('pi,xpq,qj->xij', c0.conj(), s1j, c0)
300    c1j = (h1j - s1j * e0) / -epq
301    c1j[:,degen_mask] = -.5 * s1j[:,degen_mask]
302    e1j = h1j - s1j * e0
303    e1j[:,~degen_mask] = 0
304    c1j_ao = numpy.einsum('pq,xqi->xpi', c0, c1j[:,:,nao:])
305    cl1j = c1j_ao[:,:nao]
306    cs1j = c1j_ao[:,nao:]
307    x1j = [scipy.linalg.solve(cl0.T, (cs1j[i] - x0.dot(cl1j[i])).T).T
308           for i in range(3)]
309
310    h2, s2 = get_h2_s2(mol, ia, ja)
311    h2 = numpy.einsum('pi,xypq,qj->xyij', c0.conj(), h2, c0[:,nao:])
312    s2 = numpy.einsum('pi,xypq,qj->xyij', c0.conj(), s2, c0[:,nao:])
313    f2 = h2 + numpy.einsum('xip,ypj->xyij', h1i, c1j[:,:,nao:])
314    f2+= numpy.einsum('yip,xpj->xyij', h1j, c1i[:,:,nao:])
315    f2-=(s2 + numpy.einsum('xip,ypj->xyij', s1i, c1j[:,:,nao:]) +
316         numpy.einsum('yip,xpj->xyij', s1j, c1i[:,:,nao:])) * e0[nao:]
317    f2-= numpy.einsum('xip,ypj->xyij', s1i + c1i, e1j[:,:,nao:])
318    f2-= numpy.einsum('yip,xpj->xyij', s1j + c1j, e1i[:,:,nao:])
319    c2 = f2 / -epq[:,nao:]
320    s2pp = numpy.einsum('xip,ypj->xyij', s1i, c1j)
321    s2pp+= numpy.einsum('yip,xpj->xyij', s1j, c1i)
322    s2pp+= numpy.einsum('xpi,ypj->xyij', c1i, c1j)
323    s2pp = s2pp + s2pp.transpose(0,1,3,2)
324    s2pp = s2pp[:,:,:,nao:] + s2
325    c2[:,:,degen_mask[:,nao:]] = -.5 * s2pp[:,:,degen_mask[:,nao:]]
326
327    c2_ao = numpy.einsum('pq,xyqi->xypi', c0, c2)
328    cl2 = c2_ao[:,:,:nao]
329    cs2 = c2_ao[:,:,nao:]
330    x2 = numpy.zeros((3,3,nao,nao))
331    for i in range(3):
332        for j in range(3):
333            tmp = cs2[i,j] - x0.dot(cl2[i,j])
334            tmp -= x1i[i].dot(cl1j[j])
335            tmp -= x1j[j].dot(cl1i[i])
336            x2[i,j] = scipy.linalg.solve(cl0.T, tmp.T).T
337    return numpy.asarray(x2).reshape(3,3,nao,nao)
338
339
340def get_r1(mol, atm_id, pos):
341# See JCP 135 084114, Eq (34)
342    c = lib.param.LIGHT_SPEED
343    aoslices = mol.aoslice_by_atom()
344    ish0, ish1, p0, p1 = aoslices[atm_id]
345    s0 = mol.intor('int1e_ovlp')
346    t0 = mol.intor('int1e_kin')
347    s1all = mol.intor('int1e_ipovlp', comp=3)
348    t1all = mol.intor('int1e_ipkin', comp=3)
349    s1 = numpy.zeros_like(s0)
350    t1 = numpy.zeros_like(t0)
351    s1[p0:p1,:]  =-s1all[pos][p0:p1]
352    s1[:,p0:p1] -= s1all[pos][p0:p1].T
353    t1[p0:p1,:]  =-t1all[pos][p0:p1]
354    t1[:,p0:p1] -= t1all[pos][p0:p1].T
355    x0 = get_x0(mol)
356    x1 = get_x1(mol, atm_id)[pos]
357    sa0 = s0 + reduce(numpy.dot, (x0.T, t0*(.5/c**2), x0))
358    sa1 = s1 + reduce(numpy.dot, (x0.T, t1*(.5/c**2), x0))
359    sa1+= reduce(numpy.dot, (x1.T, t0*(.5/c**2), x0))
360    sa1+= reduce(numpy.dot, (x0.T, t0*(.5/c**2), x1))
361
362    s0_sqrt = _sqrt0(s0)
363    s0_invsqrt = _invsqrt0(s0)
364    s1_sqrt = _sqrt1(s0, s1)
365    s1_invsqrt = _invsqrt1(s0, s1)
366    R0_part = reduce(numpy.dot, (s0_invsqrt, sa0, s0_invsqrt))
367    R1_part = (reduce(numpy.dot, (s0_invsqrt, sa1, s0_invsqrt)) +
368               reduce(numpy.dot, (s1_invsqrt, sa0, s0_invsqrt)) +
369               reduce(numpy.dot, (s0_invsqrt, sa0, s1_invsqrt)))
370    R1  = reduce(numpy.dot, (s0_invsqrt, _invsqrt1(R0_part, R1_part), s0_sqrt))
371    R1 += reduce(numpy.dot, (s1_invsqrt, _invsqrt0(R0_part), s0_sqrt))
372    R1 += reduce(numpy.dot, (s0_invsqrt, _invsqrt0(R0_part), s1_sqrt))
373    return R1
374
375def get_r2(mol, ia, ja, ipos, jpos):
376# See JCP 135 084114, Eq (34)
377    c = lib.param.LIGHT_SPEED
378    aoslices = mol.aoslice_by_atom()
379    ish0, ish1, i0, i1 = aoslices[ia]
380    jsh0, jsh1, j0, j1 = aoslices[ja]
381
382    s0 = mol.intor('int1e_ovlp')
383    t0 = mol.intor('int1e_kin')
384    s1all = mol.intor('int1e_ipovlp', comp=3)
385    t1all = mol.intor('int1e_ipkin', comp=3)
386    s1i = numpy.zeros_like(s0)
387    t1i = numpy.zeros_like(t0)
388    s1i[i0:i1,:]  =-s1all[ipos][i0:i1]
389    s1i[:,i0:i1] -= s1all[ipos][i0:i1].T
390    t1i[i0:i1,:]  =-t1all[ipos][i0:i1]
391    t1i[:,i0:i1] -= t1all[ipos][i0:i1].T
392
393    s1j = numpy.zeros_like(s0)
394    t1j = numpy.zeros_like(t0)
395    s1j[j0:j1,:]  =-s1all[jpos][j0:j1]
396    s1j[:,j0:j1] -= s1all[jpos][j0:j1].T
397    t1j[j0:j1,:]  =-t1all[jpos][j0:j1]
398    t1j[:,j0:j1] -= t1all[jpos][j0:j1].T
399
400    x0 = get_x0(mol)
401    x1i = get_x1(mol, ia)[ipos]
402    x1j = get_x1(mol, ja)[jpos]
403    x2 = get_x2(mol, ia, ja)[ipos,jpos]
404    sa0 = s0 + reduce(numpy.dot, (x0.T, t0*(.5/c**2), x0))
405    sa1i = s1i + reduce(numpy.dot, (x0.T, t1i*(.5/c**2), x0))
406    sa1i+= reduce(numpy.dot, (x1i.T, t0*(.5/c**2), x0))
407    sa1i+= reduce(numpy.dot, (x0.T, t0*(.5/c**2), x1i))
408    sa1j = s1j + reduce(numpy.dot, (x0.T, t1j*(.5/c**2), x0))
409    sa1j+= reduce(numpy.dot, (x1j.T, t0*(.5/c**2), x0))
410    sa1j+= reduce(numpy.dot, (x0.T, t0*(.5/c**2), x1j))
411
412    nao = mol.nao_nr()
413    s2aa = mol.intor('int1e_ipipovlp', comp=9).reshape(3,3,nao,nao)
414    t2aa = mol.intor('int1e_ipipkin', comp=9).reshape(3,3,nao,nao)
415    s2ab = mol.intor('int1e_ipovlpip', comp=9).reshape(3,3,nao,nao)
416    t2ab = mol.intor('int1e_ipkinip', comp=9).reshape(3,3,nao,nao)
417    s2ao = numpy.zeros_like(s0)
418    t2ao = numpy.zeros_like(t0)
419    if ia == ja:
420        s2ao[i0:i1      ] = s2aa[ipos,jpos,i0:i1      ]
421        s2ao[i0:i1,j0:j1]+= s2ab[ipos,jpos,i0:i1,j0:j1]
422        t2ao[i0:i1      ] = t2aa[ipos,jpos,i0:i1      ]
423        t2ao[i0:i1,j0:j1]+= t2ab[ipos,jpos,i0:i1,j0:j1]
424    else:
425        s2ao[i0:i1,j0:j1] = s2ab[ipos,jpos,i0:i1,j0:j1]
426        t2ao[i0:i1,j0:j1] = t2ab[ipos,jpos,i0:i1,j0:j1]
427    s2ao = s2ao + s2ao.T
428    t2ao = t2ao + t2ao.T
429    sa2  = reduce(numpy.dot, (x2.T, t0*(.5/c**2), x0))
430    sa2 += reduce(numpy.dot, (x1i.T, t1j*(.5/c**2), x0))
431    sa2 += reduce(numpy.dot, (x0.T, t1i*(.5/c**2), x1j))
432    sa2 += reduce(numpy.dot, (x1i.T, t0*(.5/c**2), x1j))
433    sa2  = sa2 + sa2.T
434    sa2 += s2ao + reduce(numpy.dot, (x0.T, t2ao*(.5/c**2), x0))
435
436    s0_sqrt = _sqrt0(s0)
437    s0_invsqrt = _invsqrt0(s0)
438    s1i_sqrt = _sqrt1(s0, s1i)
439    s1i_invsqrt = _invsqrt1(s0, s1i)
440    s1j_sqrt = _sqrt1(s0, s1j)
441    s1j_invsqrt = _invsqrt1(s0, s1j)
442    s2_sqrt = _sqrt2(s0, s1i, s1j, s2ao)
443    s2_invsqrt = _invsqrt2(s0, s1i, s1j, s2ao)
444
445    R0_mid = reduce(numpy.dot, (s0_invsqrt, sa0, s0_invsqrt))
446    R1i_mid = (reduce(numpy.dot, (s0_invsqrt, sa1i, s0_invsqrt)) +
447               reduce(numpy.dot, (s1i_invsqrt, sa0, s0_invsqrt)) +
448               reduce(numpy.dot, (s0_invsqrt, sa0, s1i_invsqrt)))
449    R1j_mid = (reduce(numpy.dot, (s0_invsqrt, sa1j, s0_invsqrt)) +
450               reduce(numpy.dot, (s1j_invsqrt, sa0, s0_invsqrt)) +
451               reduce(numpy.dot, (s0_invsqrt, sa0, s1j_invsqrt)))
452# second derivative of (s_invsqrt * sa * s_invsqrt), 9 terms
453    R2_mid = (reduce(numpy.dot, (s0_invsqrt, sa0, s2_invsqrt)) +
454              reduce(numpy.dot, (s1i_invsqrt, sa1j, s0_invsqrt)) +
455              reduce(numpy.dot, (s0_invsqrt, sa1i, s1j_invsqrt)) +
456              reduce(numpy.dot, (s1i_invsqrt, sa0, s1j_invsqrt)))
457    R2_mid  = R2_mid + R2_mid.T
458    R2_mid += reduce(numpy.dot, (s0_invsqrt, sa2, s0_invsqrt))
459    R2_mid = _invsqrt2(R0_mid, R1i_mid, R1j_mid, R2_mid)
460    R1i_mid = _invsqrt1(R0_mid, R1i_mid)
461    R1j_mid = _invsqrt1(R0_mid, R1j_mid)
462    R0_mid = _invsqrt0(R0_mid)
463
464    R2  = reduce(numpy.dot, (s2_invsqrt, R0_mid, s0_sqrt))
465    R2 += reduce(numpy.dot, (s1i_invsqrt, R1j_mid, s0_sqrt))
466    R2 += reduce(numpy.dot, (s1i_invsqrt, R0_mid, s1j_sqrt))
467    R2 += reduce(numpy.dot, (s1j_invsqrt, R1i_mid, s0_sqrt))
468    R2 += reduce(numpy.dot, (s0_invsqrt, R2_mid, s0_sqrt))
469    R2 += reduce(numpy.dot, (s0_invsqrt, R1i_mid, s1j_sqrt))
470    R2 += reduce(numpy.dot, (s1j_invsqrt, R0_mid, s1i_sqrt))
471    R2 += reduce(numpy.dot, (s0_invsqrt, R1j_mid, s1i_sqrt))
472    R2 += reduce(numpy.dot, (s0_invsqrt, R0_mid, s2_sqrt))
473    return R2
474
475mol1 = gto.M(
476    verbose = 0,
477    atom = [["He" , (0. , 0.     , 0.0001)],
478            [1   , (0. , -0.757 , 0.587)],
479            [1   , (0. , 0.757  , 0.587)]],
480    basis = '3-21g',
481)
482
483mol2 = gto.M(
484    verbose = 0,
485    atom = [["He" , (0. , 0.     ,-0.0001)],
486            [1   , (0. , -0.757 , 0.587)],
487            [1   , (0. , 0.757  , 0.587)]],
488    basis = '3-21g',
489)
490
491mol = gto.M(
492    verbose = 0,
493    atom = [["He" , (0. , 0.     , 0.   )],
494            [1   , (0. , -0.757 , 0.587)],
495            [1   , (0. , 0.757  , 0.587)]],
496    basis = '3-21g',
497)
498
499class KnownValues(unittest.TestCase):
500    def test_sqrt_second_order(self):
501        with lib.light_speed(10) as c:
502            nao = mol.nao_nr()
503            aoslices = mol.aoslice_by_atom()
504            p0, p1 = aoslices[0][2:]
505            s1p1 = mol1.intor('int1e_ipovlp', comp=3)
506            s1p2 = mol2.intor('int1e_ipovlp', comp=3)
507            s1_1 = numpy.zeros((3,nao,nao))
508            s1_1[:,p0:p1] = -s1p1[:,p0:p1]
509            s1_1 = s1_1 + s1_1.transpose(0,2,1)
510            s1_2 = numpy.zeros((3,nao,nao))
511            s1_2[:,p0:p1] = -s1p2[:,p0:p1]
512            s1_2 = s1_2 + s1_2.transpose(0,2,1)
513            s2sqrt_ref = (_sqrt1(mol1.intor('int1e_ovlp'), s1_1[2]) -
514                          _sqrt1(mol2.intor('int1e_ovlp'), s1_2[2])) / 0.0002 * lib.param.BOHR
515            s2invsqrt_ref = (_invsqrt1(mol1.intor('int1e_ovlp'), s1_1[2]) -
516                             _invsqrt1(mol2.intor('int1e_ovlp'), s1_2[2])) / 0.0002 * lib.param.BOHR
517
518            s1p = mol.intor('int1e_ipovlp', comp=3)
519            s1i = numpy.zeros((3,nao,nao))
520            s1i[:,p0:p1] = -s1p[:,p0:p1]
521            s1i = s1i + s1i.transpose(0,2,1)
522            s2aap = mol.intor('int1e_ipipovlp', comp=9).reshape(3,3,nao,nao)
523            s2abp = mol.intor('int1e_ipovlpip', comp=9).reshape(3,3,nao,nao)
524            s2 = numpy.zeros((3,3,nao,nao))
525            s2[:,:,p0:p1]        = s2aap[:,:,p0:p1]
526            s2[:,:,p0:p1,p0:p1] += s2abp[:,:,p0:p1,p0:p1]
527            s2 = s2 + s2.transpose(0,1,3,2)
528            s2sqrt = _sqrt2(mol.intor('int1e_ovlp'), s1i[2], s1i[2], s2[2,2])
529            s2invsqrt = _invsqrt2(mol.intor('int1e_ovlp'), s1i[2], s1i[2], s2[2,2])
530
531            self.assertAlmostEqual(abs(s2sqrt-s2sqrt_ref).max(), 0, 7)
532            self.assertAlmostEqual(abs(s2invsqrt-s2invsqrt_ref).max(), 0, 7)
533
534            p0, p1 = aoslices[1][2:]
535            s1_1 = numpy.zeros((3,nao,nao))
536            s1_1[:,p0:p1] = -s1p1[:,p0:p1]
537            s1_1 = s1_1 + s1_1.transpose(0,2,1)
538            s1_2 = numpy.zeros((3,nao,nao))
539            s1_2[:,p0:p1] = -s1p2[:,p0:p1]
540            s1_2 = s1_2 + s1_2.transpose(0,2,1)
541            s2sqrt_ref = (_sqrt1(mol1.intor('int1e_ovlp'), s1_1[2]) -
542                          _sqrt1(mol2.intor('int1e_ovlp'), s1_2[2])) / 0.0002 * lib.param.BOHR
543            s2invsqrt_ref = (_invsqrt1(mol1.intor('int1e_ovlp'), s1_1[2]) -
544                             _invsqrt1(mol2.intor('int1e_ovlp'), s1_2[2])) / 0.0002 * lib.param.BOHR
545            q0, q1 = aoslices[0][2:]
546            s1i = numpy.zeros((3,nao,nao))
547            s1i[:,p0:p1] = -s1p[:,p0:p1]
548            s1i = s1i + s1i.transpose(0,2,1)
549            s1j = numpy.zeros((3,nao,nao))
550            s1j[:,q0:q1] = -s1p[:,q0:q1]
551            s1j = s1j + s1j.transpose(0,2,1)
552            s2 = numpy.zeros((3,3,nao,nao))
553            s2[:,:,p0:p1,q0:q1] = s2abp[:,:,p0:p1,q0:q1]
554            s2 = s2 + s2.transpose(0,1,3,2)
555            s2sqrt = _sqrt2(mol.intor('int1e_ovlp'), s1i[2], s1j[2], s2[2,2])
556            s2invsqrt = _invsqrt2(mol.intor('int1e_ovlp'), s1i[2], s1j[2], s2[2,2])
557
558            self.assertAlmostEqual(abs(s2sqrt-s2sqrt_ref).max(), 0, 7)
559            self.assertAlmostEqual(abs(s2invsqrt-s2invsqrt_ref).max(), 0, 7)
560
561    def test_h2(self):
562        with lib.light_speed(10) as c:
563            h1_1, s1_1 = get_h1_s1(mol1, 0)
564            h1_2, s1_2 = get_h1_s1(mol2, 0)
565            h2_ref = (h1_1[2] - h1_2[2]) / 0.0002 * lib.param.BOHR
566            s2_ref = (s1_1[2] - s1_2[2]) / 0.0002 * lib.param.BOHR
567
568            h2, s2 = get_h2_s2(mol, 0, 0)
569            self.assertAlmostEqual(abs(h2[2,2]-h2_ref).max(), 0, 7)
570            self.assertAlmostEqual(abs(s2[2,2]-s2_ref).max(), 0, 7)
571
572            h1_1, s1_1 = get_h1_s1(mol1, 1)
573            h1_2, s1_2 = get_h1_s1(mol2, 1)
574            h2_ref = (h1_1[2] - h1_2[2]) / 0.0002 * lib.param.BOHR
575            s2_ref = (s1_1[2] - s1_2[2]) / 0.0002 * lib.param.BOHR
576
577            h2, s2 = get_h2_s2(mol, 1, 0)
578            self.assertAlmostEqual(abs(h2[2,2]-h2_ref).max(), 0, 7)
579            self.assertAlmostEqual(abs(s2[2,2]-s2_ref).max(), 0, 7)
580
581    def test_x2(self):
582        with lib.light_speed(10) as c:
583            x1_1 = get_x1(mol1, 0)
584            x1_2 = get_x1(mol2, 0)
585            x2_ref = (x1_1[2] - x1_2[2]) / 0.0002 * lib.param.BOHR
586            x2 = get_x2(mol, 0, 0)
587            self.assertAlmostEqual(abs(x2[2,2]-x2_ref).max(), 0, 7)
588
589            x1_1 = get_x1(mol1, 1)
590            x1_2 = get_x1(mol2, 1)
591            x2_ref = (x1_1[2] - x1_2[2]) / 0.0002 * lib.param.BOHR
592            x2 = get_x2(mol, 1, 0)
593            self.assertAlmostEqual(abs(x2[2,2]-x2_ref).max(), 0, 7)
594
595    def test_r2(self):
596        with lib.light_speed(10) as c:
597            r1_1 = get_r1(mol1, 0, 2)
598            r1_2 = get_r1(mol2, 0, 2)
599            r2_ref = (r1_1 - r1_2) / 0.0002 * lib.param.BOHR
600            r2 = get_r2(mol, 0, 0, 2, 2)
601            self.assertAlmostEqual(abs(r2-r2_ref).max(), 0, 7)
602
603            r1_1 = get_r1(mol1, 1, 2)
604            r1_2 = get_r1(mol2, 1, 2)
605            r2_ref = (r1_1 - r1_2) / 0.0002 * lib.param.BOHR
606            r2 = get_r2(mol, 1, 0, 2, 2)
607            self.assertAlmostEqual(abs(r2-r2_ref).max(), 0, 7)
608
609    def test_hfw2(self):
610        h1_deriv_1 = sfx2c1e_grad.gen_sf_hfw(mol1, approx='1E')
611        h1_deriv_2 = sfx2c1e_grad.gen_sf_hfw(mol2, approx='1E')
612        h2_deriv = sfx2c1e_hess.gen_sf_hfw(mol, approx='1E')
613
614        h2 = h2_deriv(0,0)
615        h2_ref = (h1_deriv_1(0)[2] - h1_deriv_2(0)[2]) / 0.0002 * lib.param.BOHR
616        self.assertAlmostEqual(abs(h2[2,2]-h2_ref).max(), 0, 7)
617
618        h2 = h2_deriv(1,0)
619        h2_ref = (h1_deriv_1(1)[2] - h1_deriv_2(1)[2]) / 0.0002 * lib.param.BOHR
620        self.assertAlmostEqual(abs(h2[2,2]-h2_ref).max(), 0, 7)
621
622        h1_deriv_1 = sfx2c1e_grad.gen_sf_hfw(mol1, approx='ATOM1E')
623        h1_deriv_2 = sfx2c1e_grad.gen_sf_hfw(mol2, approx='ATOM1E')
624        h2_deriv = sfx2c1e_hess.gen_sf_hfw(mol, approx='ATOM1E')
625
626        h2 = h2_deriv(0,0)
627        h2_ref = (h1_deriv_1(0)[2] - h1_deriv_2(0)[2]) / 0.0002 * lib.param.BOHR
628        self.assertAlmostEqual(abs(h2[2,2]-h2_ref).max(), 0, 7)
629
630        h2 = h2_deriv(1,0)
631        h2_ref = (h1_deriv_1(1)[2] - h1_deriv_2(1)[2]) / 0.0002 * lib.param.BOHR
632        self.assertAlmostEqual(abs(h2[2,2]-h2_ref).max(), 0, 7)
633
634
635if __name__ == "__main__":
636    print("Full Tests for sfx2c1e gradients")
637    unittest.main()
638