bugfix and improve prop calculations

gpu4pyscf/properties/ir.py

@@ -27,15 +27,16 @@
 
 
 def eval_ir_freq_intensity(mf, hessian_obj):
-    """main function to calculate the polarizability
+    '''Main driver of infrared spectra intensity
 
     Args:
         mf: mean field object
-        unit (str, optional): the unit of the polarizability. Defaults to 'au'.
+        hessian_obj: hessian object
 
     Returns:
-        polarizability (numpy.array): polarizability
-    """
+        frequency mode: in cm^-1
+        infrared spectra intensity: in km/mol
+    '''
     log = logger.new_logger(hessian_obj, mf.mol.verbose)
     assert isinstance(mf, RHF)
     hessian = hessian_obj.kernel()
@@ -56,7 +57,7 @@ def eval_ir_freq_intensity(mf, hessian_obj):
     TR = thermo._get_TR(mass.get(), mf.mol.atom_coords())
     TRspace = []
     TRspace.append(TR[:3])
-    
+
     rot_const = thermo.rotation_const(mass.get(), mf.mol.atom_coords())
     rotor_type = thermo._get_rotor_type(rot_const)
     if rotor_type == 'ATOM':
@@ -75,7 +76,7 @@ def eval_ir_freq_intensity(mf, hessian_obj):
         h = reduce(cupy.dot, (bvec.T, hessian_mass.transpose(0,2,1,3).reshape(3*natm,3*natm), bvec))
         e, mode = cupy.linalg.eigh(h)
         mode = bvec.dot(mode)
-    
+
     c = contract('ixn,i->ixn', mode.reshape(natm, 3, -1),
                   1/np.sqrt(mass)).reshape(3*natm, -1)
     freq = cupy.sign(e)*cupy.sqrt(cupy.abs(e))*unit2cm
@@ -94,10 +95,10 @@ def eval_ir_freq_intensity(mf, hessian_obj):
     # TODO: compact with hessian method, which can save one time cphf solve.
     # ! Different from PySCF, mo1 is all in mo!
     mo1, mo_e1 = hessian_obj.solve_mo1(mo_energy, mo_coeff, mo_occ, h1ao,
-                                       None, atmlst, hessian_obj.max_memory, log)  
+                                       None, atmlst, hessian_obj.max_memory, log)
     mo1 = cupy.asarray(mo1)
     mo_e1 = cupy.asarray(mo_e1)
-    
+
     tmp = cupy.empty((3, 3, natm))  # dipole moment, x,y,z
     aoslices = mf.mol.aoslice_by_atom()
     with mf.mol.with_common_orig((0, 0, 0)):

gpu4pyscf/properties/polarizability.py

@@ -40,7 +40,7 @@ def gen_vind(mf, mo_coeff, mo_occ):
     def fx(mo1):
         mo1 = mo1.reshape(-1, nvir, nocc)  # * the saving pattern
         mo1_mo_real = contract('nai,ua->nui', mo1, mvir)
-        dm1 = 2*contract('nui,vi->nuv', mo1_mo_real, mocc.conj()) 
+        dm1 = 2*contract('nui,vi->nuv', mo1_mo_real, mocc.conj())
         dm1+= dm1.transpose(0,2,1)
 
         v1 = vresp(dm1)  # (nset, nao, nao)
@@ -51,15 +51,14 @@ def fx(mo1):
     return fx
 
 
-def eval_polarizability(mf, unit='au'):
+def eval_polarizability(mf):
     """main function to calculate the polarizability
 
     Args:
         mf: mean field object
-        unit (str, optional): the unit of the polarizability. Defaults to 'au'.
 
     Returns:
-        polarizability (numpy.array): polarizability
+        polarizability (numpy.array): polarizability in au
     """
 
     polarizability = np.empty((3, 3))

gpu4pyscf/properties/shielding.py

@@ -16,10 +16,10 @@
 import numpy as np
 import cupy
 from pyscf.data import nist
-from pyscf.scf import _vhf, jk
+from pyscf.scf import _vhf
 from gpu4pyscf.dft import numint
 from gpu4pyscf.lib.cupy_helper import contract, sandwich_dot, add_sparse
-from gpu4pyscf.scf import cphf
+from gpu4pyscf.scf import cphf, jk
 
 def gen_vind(mf, mo_coeff, mo_occ):
     """get the induced potential. This is the same as contract the mo1 with the kernel.
@@ -38,8 +38,7 @@ def gen_vind(mf, mo_coeff, mo_occ):
     nocc = mocc.shape[1]
     nvir = nmo - nocc
     omega, alpha, hyb = mf._numint.rsh_and_hybrid_coeff(mf.xc, spin=mf.mol.spin)
-    # FIXME: check if hybrid
-    # FIXME: handle rsh
+    assert omega == 0, "The module of NMR shielding does not support range-separated functional yet."
 
     def fx(mo1):
         mo1 = mo1.reshape(-1, nvir, nocc)  # * the saving pattern
@@ -49,7 +48,8 @@ def fx(mo1):
         if hasattr(mf,'with_df'):
             vk = mf.get_jk(mf.mol, dm1, hermi=2, with_j=False)[1]
         else:
-            vk = cupy.array(jk.get_jk(mf.mol, dm1.get(), ['ijkl,jk->il']*3))
+            #vk = cupy.array(jk.get_jk(mf.mol, dm1.get(), ['ijkl,jk->il']*3))
+            vk = jk.get_jk(mf.mol, dm1, with_j=False)[1]
         v1 = -.5*hyb * vk
         tmp = contract('nuv,vi->nui', v1, mocc)
         v1vo = contract('nui,ua->nai', tmp, mvir.conj())
@@ -161,8 +161,7 @@ def get_vxc(mf, dm0):
         vxc += vj
     else:
         omega, alpha, hyb = mf._numint.rsh_and_hybrid_coeff(mf.xc, spin=mf.mol.spin)
-        # FIXME: check if hybrid
-        # FIXME: handle rsh
+        assert omega == 0, "The module of NMR shielding does not support range-separated functional yet."
         vxc += vj - vk*hyb*0.5
     return vxc
 
@@ -180,7 +179,15 @@ def get_h1ao(mf):
 
 
 def eval_shielding(mf):
+    ''' Main driver of NMR shielding
 
+    Args:
+        mf: mean field object
+
+    Returns:
+        dia-magnetic of NMR shielding: in PPM
+        para-magnetic of NMR shielding: in PPM
+    '''
     mo_coeff = mf.mo_coeff
     mo_occ = mf.mo_occ
     mo_energy = mf.mo_energy
@@ -210,13 +217,13 @@ def eval_shielding(mf):
     s1jkdm1 = contract('nui,vi->nuv', tmp, mocc.conj())*2
     s1jkdm1 = s1jkdm1 - s1jkdm1.transpose(0, 2, 1)
     omega, alpha, hyb = mf._numint.rsh_and_hybrid_coeff(mf.xc, spin=mf.mol.spin)
-    # FIXME: check if hybrid
-    # FIXME: handle rsh
+    assert omega == 0.0, "The module of NMR shielding does not support range-separated functional yet."
 
     if hasattr(mf,'with_df'):
         vk = mf.get_jk(mf.mol, s1jkdm1, hermi=2, with_j=False)[1]
     else:
-        vk = cupy.array(jk.get_jk(mf.mol, s1jkdm1.get(), ['ijkl,jk->il']*3))
+        #vk = cupy.array(jk.get_jk(mf.mol, s1jkdm1.get(), ['ijkl,jk->il']*3))
+        vk = jk.get_jk(mf.mol, s1jkdm1, with_j=False)[1]
     vk2 = -.5*hyb * vk
     h1ao += vk2
     tmp = contract('xuv,ua->xav', h1ao, mvir)