text: #!/usr/bin/env python
  # Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  #
  # Author: Qiming Sun <osirpt.sun@gmail.com>
  #
  
  '''
  Non-relativistic unrestricted Kohn-Sham g-tensor
  (In testing)
  
  Refs:
      JPC, 101, 3388
      JCP, 115, 11080
      JCP, 119, 10489
  '''
  
  from functools import reduce
  import numpy
  from pyscf import lib
  from pyscf.lib import logger
  from pyscf.dft import numint
  from pyscf.prop.gtensor import uhf as uhf_g
  from pyscf.prop.gtensor.uhf import _write, align
  from pyscf.data import nist
  from pyscf.grad import rks as rks_grad
  
  
  # Note mo10 is the imaginary part of MO^1
  def para(gobj, mo10, mo_coeff, mo_occ, qed_fac=1):
      mol = gobj.mol
      effspin = mol.spin * .5
      muB = .5  # Bohr magneton
      #qed_fac = (nist.G_ELECTRON - 1)
  
      orboa = mo_coeff[0][:,mo_occ[0]>0]
      orbob = mo_coeff[1][:,mo_occ[1]>0]
      dm0a = numpy.dot(orboa, orboa.T)
      dm0b = numpy.dot(orbob, orbob.T)
      dm10a = [reduce(numpy.dot, (mo_coeff[0], x, orboa.T)) for x in mo10[0]]
      dm10b = [reduce(numpy.dot, (mo_coeff[1], x, orbob.T)) for x in mo10[1]]
      dm10a = numpy.asarray([x-x.T for x in dm10a])
      dm10b = numpy.asarray([x-x.T for x in dm10b])
  
      hso1e = uhf_g.make_h01_soc1e(gobj, mo_coeff, mo_occ, qed_fac)
      gpara1e =-numpy.einsum('xji,yij->xy', dm10a, hso1e)
      gpara1e+= numpy.einsum('xji,yij->xy', dm10b, hso1e)
      gpara1e *= 1./effspin / muB
      _write(gobj, align(gpara1e)[0], 'SOC(1e)/OZ')
  
      if gobj.para_soc2e:
          gpara2e = gobj.make_para_soc2e((dm0a,dm0b), (dm10a,dm10b), qed_fac)
          _write(gobj, align(gpara2e)[0], 'SOC(2e)/OZ')
      else:
          gpara2e = 0
  
      gpara = gpara1e + gpara2e
      return gpara
  
  
  def make_para_soc2e(gobj, dm0, dm10, sso_qed_fac=1):
      if isinstance(gobj.para_soc2e, str):
          with_sso = 'SSO' in gobj.para_soc2e.upper()
          with_soo = 'SOO' in gobj.para_soc2e.upper()
          with_somf = 'SOMF' in gobj.para_soc2e.upper()
          with_amfi = 'AMFI' in gobj.para_soc2e.upper()
          assert(not (with_somf and (with_sso or with_soo)))
      elif gobj.para_soc2e:
          with_sso = with_soo = True
          with_somf = with_amfi = False
      else:
          with_sso = with_soo = with_somf = False
          with_amfi = True
  
      mol = gobj.mol
      alpha2 = nist.ALPHA ** 2
      effspin = mol.spin * .5
      muB = .5  # Bohr magneton
      #sso_qed_fac = (nist.G_ELECTRON - 1)
  
      mf = gobj._scf
      ni = mf._numint
      omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)
      if abs(omega) > 1e-10:
          raise NotImplementedError
      mem_now = lib.current_memory()[0]
      max_memory = max(2000, mf.max_memory*.9-mem_now)
      v1 = get_vxc_soc(ni, mol, mf.grids, mf.xc, dm0,
                       max_memory=max_memory, verbose=gobj.verbose)
      dm10a, dm10b = dm10
      if with_somf:
          ej = numpy.einsum('yil,xli->xy', v1[0]+v1[1], dm10a-dm10b)
      else:
          ej  = numpy.einsum('yil,xli->xy', v1[0], dm10a)
          ej -= numpy.einsum('yil,xli->xy', v1[1], dm10b)
      #ej *= -2  #Veff(-2X) approximation of JCP 122 034107
  
      gpara2e = 0
      if abs(hyb) > 1e-10:
          if with_amfi:
              vj, vk = uhf_g.get_jk_amfi(mol, dm0)
          else:
              vj, vk = uhf_g.get_jk(mol, dm0)
          if with_sso or with_soo:
              ek  = numpy.einsum('yil,xli->xy', vk[0], dm10a)
              ek -= numpy.einsum('yil,xli->xy', vk[1], dm10b)
              if with_sso:
                  ej += numpy.einsum('yij,xji->xy', vj[0]+vj[1], dm10a-dm10b)
                  gpara2e -= sso_qed_fac * (ej - hyb * ek)
              if with_soo:
                  ej += numpy.einsum('yij,xji->xy', vj[0]-vj[1], dm10a+dm10b)
                  gpara2e -= 2 * (ej - hyb * ek)
          else:  # SOMF, see JCP 122, 034107 Eq (19)
              ej += numpy.einsum('yij,xji->xy', vj[0]+vj[1], dm10a-dm10b)
              ek  = numpy.einsum('yil,xli->xy', vk[0]+vk[1], dm10a-dm10b)
              gpara2e -= ej - 1.5 * hyb * ek
      else:
          if with_amfi:
              vj = uhf_g.get_j_amfi(mol, dm0)
          else:
              vj = uhf_g.get_j(mol, dm0)
          if with_sso or with_soo:
              if with_sso:
                  ej += numpy.einsum('yij,xji->xy', vj[0]+vj[1], dm10a-dm10b)
                  gpara2e -= sso_qed_fac * ej
              if with_soo:
                  ej += numpy.einsum('yij,xji->xy', vj[0]-vj[1], dm10a+dm10b)
                  gpara2e -= 2 * ej
          else:  # SOMF, see JCP 122, 034107 Eq (19)
              ej += numpy.einsum('yij,xji->xy', vj[0]+vj[1], dm10a-dm10b)
              gpara2e -= ej
  
      gpara2e *= (alpha2/4) / effspin / muB
      if gobj.verbose >= logger.INFO:
          _write(gobj, align(gpara2e)[0], 'SOC(2e)/OZ')
      return gpara2e
  
  
  # Treat Vxc as one-particle operator Vnuc
  def get_vxc_soc(ni, mol, grids, xc_code, dms, max_memory=2000, verbose=None):
      xctype = ni._xc_type(xc_code)
      make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi=1)
      ngrids = len(grids.weights)
      BLKSIZE = numint.BLKSIZE
      blksize = min(int(max_memory/12*1e6/8/nao/BLKSIZE)*BLKSIZE, ngrids)
      shls_slice = (0, mol.nbas)
      ao_loc = mol.ao_loc_nr()
  
      vmat = numpy.zeros((2,3,nao,nao))
      if xctype == 'LDA':
          buf = numpy.empty((4,blksize,nao))
          ao_deriv = 1
          for ao, mask, weight, coords \
                  in ni.block_loop(mol, grids, nao, ao_deriv, max_memory,
                                   blksize=blksize, buf=buf):
              rho_a = make_rho(0, ao[0], mask, 'LDA')
              rho_b = make_rho(1, ao[0], mask, 'LDA')
              vxc = ni.eval_xc(xc_code, (rho_a, rho_b), 1, deriv=1)[1]
              vrho = vxc[0]
              aow = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho[:,0])
              _cross3x3_(vmat[0], mol, aow, ao[1:], mask, shls_slice, ao_loc)
              aow = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho[:,1])
              _cross3x3_(vmat[1], mol, aow, ao[1:], mask, shls_slice, ao_loc)
              rho = vxc = vrho = aow = None
  
      elif xctype == 'GGA':
          buf = numpy.empty((10,blksize,nao))
          ao_deriv = 2
          for ao, mask, weight, coords \
                  in ni.block_loop(mol, grids, nao, ao_deriv, max_memory,
                                   blksize=blksize, buf=buf):
              rho_a = make_rho(0, ao, mask, 'GGA')
              rho_b = make_rho(1, ao, mask, 'GGA')
              vxc = ni.eval_xc(xc_code, (rho_a,rho_b), 1, deriv=1)[1]
              wva, wvb = numint._uks_gga_wv0((rho_a, rho_b), vxc, weight)
  
              ip_ao = ao[1:4]
              ipip_ao = ao[4:]
              aow = rks_grad._make_dR_dao_w(ao, wva)
              _cross3x3_(vmat[0], mol, aow, ip_ao, mask, shls_slice, ao_loc)
              aow = rks_grad._make_dR_dao_w(ao, wvb)
              _cross3x3_(vmat[1], mol, aow, ip_ao, mask, shls_slice, ao_loc)
              rho = vxc = vrho = vsigma = wv = aow = None
          vmat = vmat - vmat.transpose(0,1,3,2)
  
      else:
          raise NotImplementedError('meta-GGA')
  
      return vmat
  
  
  def _cross3x3_(out, mol, ao1, ao2, mask, shls_slice, ao_loc):
      out[0] += numint._dot_ao_ao(mol, ao1[1], ao2[2], mask, shls_slice, ao_loc)
      out[0] -= numint._dot_ao_ao(mol, ao1[2], ao2[1], mask, shls_slice, ao_loc)
      out[1] += numint._dot_ao_ao(mol, ao1[2], ao2[0], mask, shls_slice, ao_loc)
      out[1] -= numint._dot_ao_ao(mol, ao1[0], ao2[2], mask, shls_slice, ao_loc)
      out[2] += numint._dot_ao_ao(mol, ao1[0], ao2[1], mask, shls_slice, ao_loc)
      out[2] -= numint._dot_ao_ao(mol, ao1[1], ao2[0], mask, shls_slice, ao_loc)
      return out
  
  
  class GTensor(uhf_g.GTensor):
      '''dE = B dot gtensor dot s'''
      def __init__(self, scf_method):
          uhf_g.GTensor.__init__(self, scf_method)
          self.dia_soc2e = False
          self.para_soc2e = 'AMFI+SSO+SOO'
  
      def para(self, mo10=None, mo_coeff=None, mo_occ=None):
          if mo_coeff is None: mo_coeff = self._scf.mo_coeff
          if mo_occ is None:   mo_occ = self._scf.mo_occ
          if mo10 is None:
              self.mo10, self.mo_e10 = self.solve_mo1()
              mo10 = self.mo10
          return para(self, mo10, mo_coeff, mo_occ)
  
      make_para_soc2e = make_para_soc2e
  
  
  if __name__ == '__main__':
      from pyscf import gto, scf
      mol = gto.M(atom='H 0 0 0; H 0 0 1.',
                  basis='ccpvdz', spin=1, charge=-1, verbose=3)
      mf = scf.UKS(mol).set(xc='bp86').run()
      gobj = GTensor(mf)
      gobj.gauge_orig = (0,0,0)
      gobj.para_soc2e = False
      gobj.so_eff_charge = True
      print(gobj.align(gobj.kernel())[0])
  
      mol = gto.M(atom='''
                  H 0   0   1
                  H 1.2 0   1
                  H .1  1.1 0.3
                  H .8  .7  .6
                  ''',
                  basis='ccpvdz', spin=1, charge=1, verbose=3)
      mf = scf.UKS(mol).set(xc='bp86').run()
      gobj = GTensor(mf)
      #print(gobj.kernel())
      gobj.para_soc2e = 'SSO'
      gobj.dia_soc2e = None
      gobj.so_eff_charge = False
      nao, nmo = mf.mo_coeff[0].shape
      nelec = mol.nelec
      numpy.random.seed(1)
      mo10 =[numpy.random.random((3,nmo,nelec[0])),
             numpy.random.random((3,nmo,nelec[1]))]
      print(lib.finger(para(gobj, mo10, mf.mo_coeff, mf.mo_occ)) - -2.1813250579863279e-05)
      numpy.random.seed(1)
      dm0 = numpy.random.random((2,nao,nao))
      dm0 = dm0 + dm0.transpose(0,2,1)
      dm10 = numpy.random.random((2,3,nao,nao))
      dm10 = dm10 - dm10.transpose(0,1,3,2)
      print(lib.finger(make_para_soc2e(gobj, dm0, dm10)) - 0.0036073897889263721)
  
  