Python numpy.complex() Examples

def diagonalize_asymm(H):
    """
    Diagonalize a real, *asymmetric* matrix and return sorted results.

    Return the eigenvalues and eigenvectors (column matrix)
    sorted from lowest to highest eigenvalue.
    """
    E,C = np.linalg.eig(H)
    #if np.allclose(E.imag, 0*E.imag):
    #    E = np.real(E)
    #else:
    #    print "WARNING: Eigenvalues are complex, will be returned as such."

    idx = E.real.argsort()
    E = E[idx]
    C = C[:,idx]

    return E,C 

def get_init_guess(self, kshift, nroots=1, koopmans=False, diag=None):
        size = self.vector_size()
        dtype = getattr(diag, 'dtype', np.complex)
        nroots = min(nroots, size)
        guess = []
        if koopmans:
            for n in self.nonzero_opadding[kshift][::-1][:nroots]:
                g = np.zeros(int(size), dtype=dtype)
                g[n] = 1.0
                g = self.mask_frozen(g, kshift, const=0.0)
                guess.append(g)
        else:
            idx = diag.argsort()[:nroots]
            for i in idx:
                g = np.zeros(int(size), dtype=dtype)
                g[i] = 1.0
                g = self.mask_frozen(g, kshift, const=0.0)
                guess.append(g)
        return guess 

def get_init_guess(self, kshift, nroots=1, koopmans=False, diag=None, **kwargs):
        """Initial guess vectors of R coefficients"""
        size = self.vector_size(kshift)
        dtype = getattr(diag, 'dtype', np.complex)
        nroots = min(nroots, size)
        guess = []
        # TODO do Koopmans later
        if koopmans:
            raise NotImplementedError
        else:
            idx = diag.argsort()[:nroots]
            for i in idx:
                g = np.zeros(int(size), dtype=dtype)
                g[i] = 1.0
                # TODO do mask_frozen later
                guess.append(g)
        return guess 

def allocateVecs(self):
        self.subH = np.zeros( shape=(self.maxM,self.maxM), dtype=complex )
        self.sol = np.zeros( shape=(self.maxM), dtype=complex )
        self.dgks = np.zeros( shape=(self.maxM), dtype=complex )
        self.nConv = np.zeros( shape=(self.maxM), dtype=int )
        self.eigs = np.zeros( shape=(self.maxM), dtype=complex )
        self.evecs = np.zeros( shape=(self.maxM,self.maxM), dtype=complex )
        self.oldeigs = np.zeros( shape=(self.maxM), dtype=complex )
        self.deigs = np.zeros( shape=(self.maxM), dtype=complex )
        self.outeigs = np.zeros( shape=(self.nEigen), dtype=complex )
        self.outevecs = np.zeros( shape=(self.size,self.nEigen), dtype=complex)
        self.currentSize = 0

        self.Ax = np.zeros( shape=(self.size), dtype=complex )
        self.res = np.zeros( shape=(self.size), dtype=complex )
        self.vlist = np.zeros( shape=(self.maxM,self.size), dtype=complex )
        self.cv = np.zeros( shape = (self.size), dtype = complex )
        self.cAv = np.zeros( shape = (self.size), dtype = complex )
        self.Avlist = np.zeros( shape=(self.maxM,self.size), dtype=complex )
        self.dres = 999.9
        self.resnorm = 999.9
        self.cvEig = 0.1
        self.ciEig = 0
        self.deflated = 0 

def cc_Wvvvv(t1,t2,eris):
    tau = make_tau(t2,t1,t1)
    #eris_vovv = np.array(eris.ovvv).transpose(1,0,3,2)
    #tmp = einsum('mb,amef->abef',t1,eris_vovv)
    #Wabef = eris.vvvv - tmp + tmp.transpose(1,0,2,3)
    #Wabef += 0.25*einsum('mnab,mnef->abef',tau,eris.oovv)
    if t1.dtype == np.complex: ds_type = 'c16'
    else: ds_type = 'f8'
    _tmpfile1 = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
    fimd = h5py.File(_tmpfile1.name)
    nocc, nvir = t1.shape
    Wabef = fimd.create_dataset('vvvv', (nvir,nvir,nvir,nvir), ds_type)
    for a in range(nvir):
        Wabef[a] = eris.vvvv[a] 
        Wabef[a] -= einsum('mb,mfe->bef',t1,eris.ovvv[:,a,:,:]) 
        Wabef[a] += einsum('m,mbfe->bef',t1[:,a],eris.ovvv) 
        Wabef[a] += 0.25*einsum('mnb,mnef->bef',tau[:,:,a,:],eris.oovv)
    return Wabef 

def Wvvvv(t1,t2,eris):
    tau = make_tau(t2,t1,t1)
    #Wabef = cc_Wvvvv(t1,t2,eris) + 0.25*einsum('mnab,mnef->abef',tau,eris.oovv)
    if t1.dtype == np.complex: ds_type = 'c16'
    else: ds_type = 'f8'
    _tmpfile1 = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
    fimd = h5py.File(_tmpfile1.name)
    nocc, nvir = t1.shape
    Wabef = fimd.create_dataset('vvvv', (nvir,nvir,nvir,nvir), ds_type)
    #_cc_Wvvvv = cc_Wvvvv(t1,t2,eris)
    for a in range(nvir):
        #Wabef[a] = _cc_Wvvvv[a]
        Wabef[a] = eris.vvvv[a] 
        Wabef[a] -= einsum('mb,mfe->bef',t1,eris.ovvv[:,a,:,:]) 
        Wabef[a] += einsum('m,mbfe->bef',t1[:,a],eris.ovvv) 
        #Wabef[a] += 0.25*einsum('mnb,mnef->bef',tau[:,:,a,:],eris.oovv)

        #Wabef[a] += 0.25*einsum('mnb,mnef->bef',tau[:,:,a,:],eris.oovv) 
        Wabef[a] += 0.5*einsum('mnb,mnef->bef',tau[:,:,a,:],eris.oovv) 
    return Wabef 

def test_intor_r2e(self):
        mol1 = gto.M(atom=[[1   , (0. , -0.7 , 0.)],
                           [1   , (0. ,  0.7 , 0.)]],
                     basis = '631g')
        nao = mol1.nao_2c()
        eri0 = numpy.empty((3,nao,nao,nao,nao), dtype=numpy.complex)
        ip = 0
        for i in range(mol1.nbas):
            jp = 0
            for j in range(mol1.nbas):
                kp = 0
                for k in range(mol1.nbas):
                    lp = 0
                    for l in range(mol1.nbas):
                        buf = mol1.intor_by_shell('int2e_ip1_spinor', (i,j,k,l), comp=3)
                        di,dj,dk,dl = buf.shape[1:]
                        eri0[:,ip:ip+di,jp:jp+dj,kp:kp+dk,lp:lp+dl] = buf
                        lp += dl
                    kp += dk
                jp += dj
            ip += di 

def hcore_generator(self, mol):
        aoslices = mol.aoslice_2c_by_atom()
        h1 = self.get_hcore(mol)
        n2c = mol.nao_2c()
        n4c = n2c * 2
        c = lib.param.LIGHT_SPEED
        def hcore_deriv(atm_id):
            shl0, shl1, p0, p1 = aoslices[atm_id]
            with mol.with_rinv_at_nucleus(atm_id):
                z = -mol.atom_charge(atm_id)
                vn = z * mol.intor('int1e_iprinv_spinor', comp=3)
                wn = z * mol.intor('int1e_ipsprinvsp_spinor', comp=3)

            v = numpy.zeros((3,n4c,n4c), numpy.complex)
            v[:,:n2c,:n2c] = vn
            v[:,n2c:,n2c:] = wn * (.25/c**2)
            v[:,p0:p1]         += h1[:,p0:p1]
            v[:,n2c+p0:n2c+p1] += h1[:,n2c+p0:n2c+p1]
            return v + v.conj().transpose(0,2,1)
        return hcore_deriv 

def make_s10(mol, gauge_orig=None, mb='RMB'):
    '''First order overlap matrix wrt external magnetic field.
    Note the side effects of set_common_origin.
    '''
    if gauge_orig is not None:
        mol.set_common_origin(gauge_orig)
    n2c = mol.nao_2c()
    n4c = n2c * 2
    c = lib.param.LIGHT_SPEED
    s1 = numpy.zeros((3, n4c, n4c), complex)
    if mb.upper() == 'RMB':
        if gauge_orig is None:
            t1 = mol.intor('int1e_giao_sa10sp_spinor', 3)
        else:
            t1 = mol.intor('int1e_cg_sa10sp_spinor', 3)
        for i in range(3):
            t1cc = t1[i] + t1[i].conj().T
            s1[i,n2c:,n2c:] = t1cc * (.25/c**2)

    if gauge_orig is None:
        sg = mol.intor('int1e_govlp_spinor', 3)
        tg = mol.intor('int1e_spgsp_spinor', 3)
        s1[:,:n2c,:n2c] += sg
        s1[:,n2c:,n2c:] += tg * (.25/c**2)
    return s1 

def test_better_float():
    # Better float function
    def check_against(f1, f2):
        return f1 if FLOAT_TYPES.index(f1) >= FLOAT_TYPES.index(f2) else f2
    for first in FLOAT_TYPES:
        for other in IUINT_TYPES + np.sctypes['complex']:
            assert_equal(better_float_of(first, other), first)
            assert_equal(better_float_of(other, first), first)
            for other2 in IUINT_TYPES + np.sctypes['complex']:
                assert_equal(better_float_of(other, other2), np.float32)
                assert_equal(better_float_of(other, other2, np.float64),
                             np.float64)
        for second in FLOAT_TYPES:
            assert_equal(better_float_of(first, second),
                         check_against(first, second))
    # Check codes and dtypes work
    assert_equal(better_float_of('f4', 'f8', 'f4'), np.float64)
    assert_equal(better_float_of('i4', 'i8', 'f8'), np.float64) 

def _num_to_rqc_str(num):
    """Convert float to string to be included in RQC quil DEFGATE block
    (as written by _np_to_quil_def_str)."""
    num = _np.complex(_np.real_if_close(num))
    if _np.imag(num) == 0:
        output = str(_np.real(num))
        return output
    else:
        real_part = _np.real(num)
        imag_part = _np.imag(num)
        if imag_part < 0:
            sgn = '-'
            imag_part = imag_part * -1
        elif imag_part > 0:
            sgn = '+'
        else:
            assert False
        return '{}{}{}i'.format(real_part, sgn, imag_part) 

def _patch_cython():
    # "monkey patch" some objects to avoid cyclic import structure
    from . import _npc_helper
    _npc_helper._charges = charges
    _npc_helper._np_conserved = np_conserved
    assert _npc_helper.QTYPE == charges.QTYPE
    # check types
    warn = False
    import numpy as np
    check_types = [
        (np.float64, np.complex128),
        (np.ones([1]).dtype, (1.j * np.ones([1])).dtype),
        (np.array(1.).dtype, np.array(1.j).dtype),
        (np.array(1., dtype=np.float).dtype, np.array(1., dtype=np.complex).dtype),
    ]
    types_ok = [
        _npc_helper._float_complex_are_64_bit(dt_float, dt_real)
        for dt_float, dt_real in check_types
    ]
    if not np.all(types_ok):
        import warnings
        warnings.warn("(Some of) the default dtypes are not 64-bit. "
                      "Using the compiled cython code (as you do) might make it slower.")
    # done 

def _griffin_lim(S):
    angles = np.exp(2j * np.pi * np.random.rand(*S.shape))
    S_complex = np.abs(S).astype(np.complex)
    for i in range(hparams.griffin_lim_iters):
        if i > 0:
            angles = np.exp(1j * np.angle(_stft(y)))
        y = _istft(S_complex * angles)
    return y 

def angle_to_cortex(self, theta, rho):
        'See help(neuropythy.registration.RetinotopyModel.angle_to_cortex).'
        #TODO: This should be made to work correctly with visual area boundaries: this could be done
        # by, for each area (e.g., V2) looking at its boundaries (with V1 and V3) and flipping the
        # adjacent triangles so that there is complete coverage of each hemifield, guaranteed.
        if not pimms.is_vector(theta): return self.angle_to_cortex([theta], [rho])[0]
        theta = np.asarray(theta)
        rho = np.asarray(rho)
        zs = np.asarray(
            rho * np.exp([np.complex(z) for z in 1j * ((90.0 - theta)/180.0*np.pi)]),
            dtype=np.complex)
        coords = np.asarray([zs.real, zs.imag]).T
        if coords.shape[0] == 0: return np.zeros((0, len(self.visual_meshes), 2))
        # we step through each area in the forward model and return the appropriate values
        tx = self.transform
        res = np.transpose(
            [self.visual_meshes[area].interpolate(coords, 'cortical_coordinates', method='linear')
             for area in sorted(self.visual_meshes.keys())],
            (1,0,2))
        if tx is not None:
            res = np.asarray(
                [np.dot(tx, np.vstack((area_xy.T, np.ones(len(area_xy)))))[0:2].T
                 for area_xy in res])
        return res 

def gen_H_tb(t,Nx,Ny,kvec):
    H = np.zeros((Nx,Ny,Nx,Ny),dtype=np.complex)
    for i in range(Nx):
        for j in range(Ny):
            if i == Nx-1:
                H[i,j,0   ,j] += np.exp(-1j*np.dot(np.array(kvec),np.array([Nx,0])))
            else:
                H[i,j,i+1 ,j] += 1

            if i == 0:
                H[i,j,Nx-1,j] += np.exp(-1j*np.dot(np.array(kvec),np.array([-Nx,0])))
            else:
                H[i,j,i-1 ,j] += 1

            if j == Ny-1:
                H[i,j,i,0   ] += np.exp(-1j*np.dot(np.array(kvec),np.array([0,Ny])))
            else:
                H[i,j,i,j+1] += 1

            if j == 0:
                H[i,j,i,Ny-1] += np.exp(-1j*np.dot(np.array(kvec),np.array([0,-Ny])))
            else:
                H[i,j,i,j-1] += 1
    return -t*H.reshape(Nx*Ny,Nx*Ny)

### get H_tb at a series of kpoints. 

def get_xmat(self, mol=None):
        if mol is None:
            xmol = self.get_xmol(mol)[0]
        else:
            xmol = mol
        c = lib.param.LIGHT_SPEED
        assert('1E' in self.approx.upper())

        if 'ATOM' in self.approx.upper():
            atom_slices = xmol.offset_2c_by_atom()
            n2c = xmol.nao_2c()
            x = numpy.zeros((n2c,n2c), dtype=numpy.complex)
            for ia in range(xmol.natm):
                ish0, ish1, p0, p1 = atom_slices[ia]
                shls_slice = (ish0, ish1, ish0, ish1)
                s1 = xmol.intor('int1e_ovlp_spinor', shls_slice=shls_slice)
                t1 = xmol.intor('int1e_spsp_spinor', shls_slice=shls_slice) * .5
                with xmol.with_rinv_at_nucleus(ia):
                    z = -xmol.atom_charge(ia)
                    v1 = z*xmol.intor('int1e_rinv_spinor', shls_slice=shls_slice)
                    w1 = z*xmol.intor('int1e_sprinvsp_spinor', shls_slice=shls_slice)
                x[p0:p1,p0:p1] = _x2c1e_xmatrix(t1, v1, w1, s1, c)
        else:
            s = xmol.intor_symmetric('int1e_ovlp_spinor')
            t = xmol.intor_symmetric('int1e_spsp_spinor') * .5
            v = xmol.intor_symmetric('int1e_nuc_spinor')
            w = xmol.intor_symmetric('int1e_spnucsp_spinor')
            x = _x2c1e_xmatrix(t, v, w, s, c)
        return x 

def get_jk(mol, dm, hermi=1, mf_opt=None, with_j=True, with_k=True, omega=None):
    '''non-relativistic J/K matrices (without SSO,SOO etc) in the j-adapted
    spinor basis.
    '''
    n2c = dm.shape[0]
    dd = numpy.zeros((n2c*2,)*2, dtype=numpy.complex)
    dd[:n2c,:n2c] = dm
    dhf._call_veff_llll(mol, dd, hermi, None)
    vj, vk = _vhf.rdirect_mapdm('int2e_spinor', 's8',
                                ('ji->s2kl', 'jk->s1il'), dm, 1,
                                mol._atm, mol._bas, mol._env, mf_opt)
    return dhf._jk_triu_(vj, vk, hermi) 

def get_init_guess(self, kshift, nroots=1, koopmans=False, diag=None):
        size = self.vector_size()
        dtype = getattr(diag, 'dtype', np.complex)
        nroots = min(nroots, size)
        nocca, noccb = self.nocc
        guess = []
        if koopmans:
            idx = np.zeros(nroots, dtype=np.int)
            tmp_oalpha, tmp_obeta = self.nonzero_opadding[kshift]
            tmp_oalpha = list(tmp_oalpha)
            tmp_obeta = list(tmp_obeta)
            if len(tmp_obeta) + len(tmp_oalpha) < nroots:
                raise ValueError("Max number of roots for k-point (idx=%3d) for koopmans "
                                 "is %3d.\nRequested %3d." %
                                 (kshift, len(tmp_obeta)+len(tmp_oalpha), nroots))

            total_count = 0
            while(total_count < nroots):
                if total_count % 2 == 0 and len(tmp_oalpha) > 0:
                    idx[total_count] = tmp_oalpha.pop()
                else:
                    # Careful! index depends on how we create vector
                    # (here the first elements are r1a, then r1b)
                    idx[total_count] = nocca + tmp_obeta.pop()
                total_count += 1
        else:
            idx = diag.argsort()

        for i in idx[:nroots]:
            g = np.zeros(size, dtype)
            g[i] = 1.0
            g = self.mask_frozen(g, kshift, const=0.0)
            guess.append(g)
        return guess 

def get_init_guess(self, kshift, nroots=1, koopmans=False, diag=None):
        size = self.vector_size()
        dtype = getattr(diag, 'dtype', np.complex)
        nroots = min(nroots, size)
        nocca, noccb = self.nocc
        nmoa, nmob = self.nmo
        nvira, nvirb = nmoa-nocca, nmob-noccb
        guess = []
        if koopmans:
            idx = np.zeros(nroots, dtype=np.int)
            tmp_valpha, tmp_vbeta = self.nonzero_vpadding[kshift]
            tmp_valpha = list(tmp_valpha)
            tmp_vbeta = list(tmp_vbeta)
            if len(tmp_vbeta) + len(tmp_valpha) < nroots:
                raise ValueError("Max number of roots for k-point (idx=%3d) for koopmans "
                                 "is %3d.\nRequested %3d." %
                                 (kshift, len(tmp_vbeta)+len(tmp_valpha), nroots))

            total_count = 0
            while(total_count < nroots):
                if total_count % 2 == 0 and len(tmp_valpha) > 0:
                    idx[total_count] = tmp_valpha.pop(0)
                else:
                    # Careful! index depends on how we create vector
                    # (here the first elements are r1a, then r1b)
                    idx[total_count] = nvira + tmp_vbeta.pop(0)
                total_count += 1
        else:
            idx = diag.argsort()

        for i in idx[:nroots]:
            g = np.zeros(size, dtype)
            g[i] = 1.0
            g = self.mask_frozen(g, kshift, const=0.0)
            guess.append(g)
        return guess 

def lipccsd(self, nroots=2*4, kptlist=None):
        time0 = time.clock(), time.time()
        log = logger.Logger(self.stdout, self.verbose)
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        if kptlist is None:
            kptlist = range(nkpts)
        size = self.vector_size_ip()
        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            nfrozen = np.sum(self.mask_frozen_ip(np.zeros(size, dtype=int), kshift, const=1))
            nroots = min(nroots, size - nfrozen)
        evals = np.zeros((len(kptlist),nroots), np.float)
        evecs = np.zeros((len(kptlist),nroots,size), np.complex)

        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            diag = self.ipccsd_diag()
            precond = lambda dx, e, x0: dx/(diag-e)
            # Initial guess from file
            amplitude_filename = "__lipccsd" + str(kshift) + "__.hdf5"
            rsuccess, x0 = read_eom_amplitudes((nroots,size),filename=amplitude_filename)
            if x0 is not None:
                x0 = x0.T
            #if not rsuccess:
            #    x0 = np.zeros_like(diag)
            #    x0[np.argmin(diag)] = 1.0

            conv, evals_k, evecs_k = eigs(self.lipccsd_matvec, size, nroots, x0=x0, Adiag=diag, verbose=self.verbose)

            evals_k = evals_k.real
            evals[k] = evals_k
            evecs[k] = evecs_k.T

            write_eom_amplitudes(evecs[k],filename=amplitude_filename)
        time0 = log.timer_debug1('converge ip-ccsd', *time0)
        comm.Barrier()
        return evals.real, evecs 

def eaccsd(self, nroots=2*4, kptlist=None):
        time0 = time.clock(), time.time()
        log = logger.Logger(self.stdout, self.verbose)
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        if kptlist is None:
            kptlist = range(nkpts)
        size = self.vector_size_ea()
        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            nfrozen = np.sum(self.mask_frozen_ea(np.zeros(size, dtype=int), kshift, const=1))
            nroots = min(nroots, size - nfrozen)
        evals = np.zeros((len(kptlist),nroots), np.float)
        evecs = np.zeros((len(kptlist),nroots,size), np.complex)

        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            diag = self.eaccsd_diag()
            diag = self.mask_frozen_ea(diag, kshift, const=LARGE_DENOM)
            precond = lambda dx, e, x0: dx/(diag-e)
            # Initial guess from file
            amplitude_filename = "__reaccsd" + str(kshift) + "__.hdf5"
            rsuccess, x0 = read_eom_amplitudes((nroots,size),filename=amplitude_filename)
            if x0 is not None:
                x0 = x0.T
            #if not rsuccess:
            #    x0 = np.zeros_like(diag)
            #    x0[np.argmin(diag)] = 1.0

            conv, evals_k, evecs_k = eigs(self.eaccsd_matvec, size, nroots, x0=x0, Adiag=diag, verbose=self.verbose)

            evals_k = evals_k.real
            evals[k] = evals_k
            evecs[k] = evecs_k.T

            write_eom_amplitudes(evecs[k],filename=amplitude_filename)
        time0 = log.timer_debug1('converge ea-ccsd', *time0)
        comm.Barrier()
        return evals.real, evecs 

def leaccsd(self, nroots=2*4, kptlist=None):
        time0 = time.clock(), time.time()
        log = logger.Logger(self.stdout, self.verbose)
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        if kptlist is None:
            kptlist = range(nkpts)
        size = self.vector_size_ea()
        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            nfrozen = np.sum(self.mask_frozen_ea(np.zeros(size, dtype=int), kshift, const=1))
            nroots = min(nroots, size - nfrozen)
        evals = np.zeros((len(kptlist),nroots), np.float)
        evecs = np.zeros((len(kptlist),nroots,size), np.complex)

        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            diag = self.eaccsd_diag()
            precond = lambda dx, e, x0: dx/(diag-e)
            # Initial guess from file
            amplitude_filename = "__leaccsd" + str(kshift) + "__.hdf5"
            rsuccess, x0 = read_eom_amplitudes((nroots,size),filename=amplitude_filename)
            if x0 is not None:
                x0 = x0.T
            #if not rsuccess:
            #    x0 = np.zeros_like(diag)
            #    x0[np.argmin(diag)] = 1.0

            conv, evals_k, evecs_k = eigs(self.leaccsd_matvec, size, nroots, x0=x0, Adiag=diag, verbose=self.verbose)

            evals_k = evals_k.real
            evals[k] = evals_k
            evecs[k] = evecs_k.T

            write_eom_amplitudes(evecs[k],amplitude_filename)
        time0 = log.timer_debug1('converge lea-ccsd', *time0)
        comm.Barrier()
        return evals.real, evecs 

def get_ao_pairs_G(mydf, kpts=numpy.zeros((2,3)), q=None, shls_slice=None,
                   compact=getattr(__config__, 'pbc_df_ao_pairs_compact', False)):
    '''Calculate forward Fourier tranform (G|ij) of all AO pairs.

    Returns:
        ao_pairs_G : 2D complex array
            For gamma point, the shape is (ngrids, nao*(nao+1)/2); otherwise the
            shape is (ngrids, nao*nao)
    '''
    if kpts is None: kpts = numpy.zeros((2,3))
    cell = mydf.cell
    kpts = numpy.asarray(kpts)
    q = kpts[1] - kpts[0]
    coords = cell.gen_uniform_grids(mydf.mesh)
    ngrids = len(coords)
    max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]) * .5)

    if shls_slice is None:
        shls_slice = (0, cell.nbas, 0, cell.nbas)
    ish0, ish1, jsh0, jsh1 = shls_slice
    ao_loc = cell.ao_loc_nr()
    i0 = ao_loc[ish0]
    i1 = ao_loc[ish1]
    j0 = ao_loc[jsh0]
    j1 = ao_loc[jsh1]
    compact = compact and (i0 == j0) and (i1 == j1)

    if compact and gamma_point(kpts):  # gamma point
        aosym = 's2'
    else:
        aosym = 's1'

    ao_pairs_G = numpy.empty((ngrids,(i1-i0)*(j1-j0)), dtype=numpy.complex)
    for pqkR, pqkI, p0, p1 \
            in mydf.pw_loop(mydf.mesh, kpts, q, shls_slice,
                            max_memory=max_memory, aosym=aosym):
        ao_pairs_G[p0:p1] = pqkR.T + pqkI.T * 1j
    return ao_pairs_G 

def get_mo_pairs_G(mydf, mo_coeffs, kpts=numpy.zeros((2,3)), q=None,
                   compact=getattr(__config__, 'pbc_df_mo_pairs_compact', False)):
    '''Calculate forward fourier transform (G|ij) of all MO pairs.

    Args:
        mo_coeff: length-2 list of (nao,nmo) ndarrays
            The two sets of MO coefficients to use in calculating the
            product |ij).

    Returns:
        mo_pairs_G : (ngrids, nmoi*nmoj) ndarray
            The FFT of the real-space MO pairs.
    '''
    if kpts is None: kpts = numpy.zeros((2,3))
    cell = mydf.cell
    kpts = numpy.asarray(kpts)
    q = kpts[1] - kpts[0]
    coords = cell.gen_uniform_grids(mydf.mesh)
    nmoi = mo_coeffs[0].shape[1]
    nmoj = mo_coeffs[1].shape[1]
    ngrids = len(coords)
    max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0]) * .5)

    mo_pairs_G = numpy.empty((ngrids,nmoi,nmoj), dtype=numpy.complex)
    for pqkR, pqkI, p0, p1 \
            in mydf.pw_loop(mydf.mesh, kpts, q,
                            max_memory=max_memory, aosym='s2'):
        pqk = (pqkR + pqkI*1j).reshape(nao,nao,-1)
        mo_pairs_G[p0:p1] = lib.einsum('pqk,pi,qj->kij', pqk, *mo_coeffs[:2])
    return mo_pairs_G.reshape(ngrids,nmoi*nmoj) 

def get_ovlp(mol):
    n2c = mol.nao_2c()
    n4c = n2c * 2
    c = lib.param.LIGHT_SPEED

    s = mol.intor_symmetric('int1e_ovlp_spinor')
    t = mol.intor_symmetric('int1e_spsp_spinor')
    s1e = numpy.zeros((n4c, n4c), numpy.complex)
    s1e[:n2c,:n2c] = s
    s1e[n2c:,n2c:] = t * (.5/c)**2
    return s1e 

def _call_veff_ssll(mol, dm, hermi=1, mf_opt=None):
    if isinstance(dm, numpy.ndarray) and dm.ndim == 2:
        n_dm = 1
        n2c = dm.shape[0] // 2
        dmll = dm[:n2c,:n2c].copy()
        dmsl = dm[n2c:,:n2c].copy()
        dmss = dm[n2c:,n2c:].copy()
        dms = (dmll, dmss, dmsl)
    else:
        n_dm = len(dm)
        n2c = dm[0].shape[0] // 2
        dms = [dmi[:n2c,:n2c].copy() for dmi in dm] \
            + [dmi[n2c:,n2c:].copy() for dmi in dm] \
            + [dmi[n2c:,:n2c].copy() for dmi in dm]
    jks = ('lk->s2ij',) * n_dm \
        + ('ji->s2kl',) * n_dm \
        + ('jk->s1il',) * n_dm
    c1 = .5 / lib.param.LIGHT_SPEED
    vx = _vhf.rdirect_bindm('int2e_spsp1_spinor', 's4', jks, dms, 1,
                            mol._atm, mol._bas, mol._env, mf_opt) * c1**2
    vj = numpy.zeros((n_dm,n2c*2,n2c*2), dtype=numpy.complex)
    vk = numpy.zeros((n_dm,n2c*2,n2c*2), dtype=numpy.complex)
    vj[:,n2c:,n2c:] = vx[      :n_dm  ,:,:]
    vj[:,:n2c,:n2c] = vx[n_dm  :n_dm*2,:,:]
    vk[:,n2c:,:n2c] = vx[n_dm*2:      ,:,:]
    if n_dm == 1:
        vj = vj.reshape(vj.shape[1:])
        vk = vk.reshape(vk.shape[1:])
    return _jk_triu_(vj, vk, hermi) 

def _proj_dmll(mol_nr, dm_nr, mol):
    '''Project non-relativistic atomic density matrix to large component spinor
    representation
    '''
    from pyscf.scf import addons
    proj = addons.project_mo_nr2r(mol_nr, numpy.eye(mol_nr.nao_nr()), mol)

    n2c = proj.shape[0]
    n4c = n2c * 2
    dm = numpy.zeros((n4c,n4c), dtype=complex)
    # *.5 because alpha and beta are summed in project_mo_nr2r
    dm_ll = reduce(numpy.dot, (proj, dm_nr*.5, proj.T.conj()))
    dm[:n2c,:n2c] = (dm_ll + time_reversal_matrix(mol, dm_ll)) * .5
    return dm 

def test_get_jk(self):
        n2c = h4.nao_2c()
        n4c = n2c * 2
        c1 = .5 / lib.param.LIGHT_SPEED
        eri0 = numpy.zeros((n4c,n4c,n4c,n4c), dtype=numpy.complex)
        eri0[:n2c,:n2c,:n2c,:n2c] = h4.intor('int2e_spinor')
        eri0[n2c:,n2c:,:n2c,:n2c] = h4.intor('int2e_spsp1_spinor') * c1**2
        eri0[:n2c,:n2c,n2c:,n2c:] = eri0[n2c:,n2c:,:n2c,:n2c].transpose(2,3,0,1)
        ssss = h4.intor('int2e_spsp1spsp2_spinor') * c1**4
        eri0[n2c:,n2c:,n2c:,n2c:] = ssss

        numpy.random.seed(1)
        dm = numpy.random.random((2,n4c,n4c))+numpy.random.random((2,n4c,n4c))*1j
        dm = dm + dm.transpose(0,2,1).conj()
        vj0 = numpy.einsum('ijkl,lk->ij', eri0, dm[0])
        vk0 = numpy.einsum('ijkl,jk->il', eri0, dm[0])
        vj, vk = scf.dhf.get_jk(h4, dm[0], hermi=1, coulomb_allow='SSSS')
        self.assertTrue(numpy.allclose(vj0, vj))
        self.assertTrue(numpy.allclose(vk0, vk))

        vj0 = numpy.einsum('ijkl,xlk->xij', ssss, dm[:,n2c:,n2c:])
        vk0 = numpy.einsum('ijkl,xjk->xil', ssss, dm[:,n2c:,n2c:])
        vj, vk = scf.dhf._call_veff_ssss(h4, dm, hermi=0)
        self.assertTrue(numpy.allclose(vj0, vj))
        self.assertTrue(numpy.allclose(vk0, vk))

        eri0[n2c:,n2c:,n2c:,n2c:] = 0
        vj0 = numpy.einsum('ijkl,xlk->xij', eri0, dm)
        vk0 = numpy.einsum('ijkl,xjk->xil', eri0, dm)
        vj, vk = scf.dhf.get_jk(h4, dm, hermi=1, coulomb_allow='SSLL')
        self.assertTrue(numpy.allclose(vj0, vj))
        self.assertTrue(numpy.allclose(vk0, vk))

        eri0[n2c:,n2c:,:n2c,:n2c] = 0
        eri0[:n2c,:n2c,n2c:,n2c:] = 0
        vj0 = numpy.einsum('ijkl,lk->ij', eri0, dm[0])
        vk0 = numpy.einsum('ijkl,jk->il', eri0, dm[0])
        vj, vk = scf.dhf.get_jk(h4, dm[0], hermi=0, coulomb_allow='LLLL')
        self.assertTrue(numpy.allclose(vj0, vj))
        self.assertTrue(numpy.allclose(vk0, vk)) 

def _fill_gaunt(mol, erig):
    n2c = erig.shape[0]
    n4c = n2c * 2

    tao = numpy.asarray(mol.time_reversal_map())
    idx = abs(tao)-1 # -1 for C indexing convention
    sign_mask = tao<0

    eri0 = numpy.zeros((n4c,n4c,n4c,n4c), dtype=numpy.complex)
    eri0[:n2c,n2c:,:n2c,n2c:] = erig # ssp1ssp2

    eri2 = erig.take(idx,axis=0).take(idx,axis=1) # sps1ssp2
    eri2[sign_mask,:] *= -1
    eri2[:,sign_mask] *= -1
    eri2 = -eri2.transpose(1,0,2,3)
    eri0[n2c:,:n2c,:n2c,n2c:] = eri2

    eri2 = erig.take(idx,axis=2).take(idx,axis=3) # ssp1sps2
    eri2[:,:,sign_mask,:] *= -1
    eri2[:,:,:,sign_mask] *= -1
    eri2 = -eri2.transpose(0,1,3,2)
    #self.assertTrue(numpy.allclose(eri0, eri2))
    eri0[:n2c,n2c:,n2c:,:n2c] = eri2

    eri2 = erig.take(idx,axis=0).take(idx,axis=1)
    eri2 = eri2.take(idx,axis=2).take(idx,axis=3) # sps1sps2
    eri2 = eri2.transpose(1,0,2,3)
    eri2 = eri2.transpose(0,1,3,2)
    eri2[sign_mask,:] *= -1
    eri2[:,sign_mask] *= -1
    eri2[:,:,sign_mask,:] *= -1
    eri2[:,:,:,sign_mask] *= -1
    eri0[n2c:,:n2c,n2c:,:n2c] = eri2
    return eri0 

def make_hdiag(h1e, eri, norb, nelec):
    '''Diagonal Hamiltonian for Davidson preconditioner
    '''
    if h1e.dtype == numpy.complex or eri.dtype == numpy.complex:
        raise NotImplementedError('Complex Hamiltonian')

    neleca, nelecb = _unpack_nelec(nelec)
    h1e = numpy.asarray(h1e, order='C')
    eri = ao2mo.restore(1, eri, norb)
    occslsta = occslstb = cistring._gen_occslst(range(norb), neleca)
    if neleca != nelecb:
        occslstb = cistring._gen_occslst(range(norb), nelecb)
    na = len(occslsta)
    nb = len(occslstb)

    hdiag = numpy.empty(na*nb)
    jdiag = numpy.asarray(numpy.einsum('iijj->ij',eri), order='C')
    kdiag = numpy.asarray(numpy.einsum('ijji->ij',eri), order='C')
    c_h1e = h1e.ctypes.data_as(ctypes.c_void_p)
    c_jdiag = jdiag.ctypes.data_as(ctypes.c_void_p)
    c_kdiag = kdiag.ctypes.data_as(ctypes.c_void_p)
    libfci.FCImake_hdiag_uhf(hdiag.ctypes.data_as(ctypes.c_void_p),
                             c_h1e, c_h1e, c_jdiag, c_jdiag, c_jdiag, c_kdiag, c_kdiag,
                             ctypes.c_int(norb),
                             ctypes.c_int(na), ctypes.c_int(nb),
                             ctypes.c_int(neleca), ctypes.c_int(nelecb),
                             occslsta.ctypes.data_as(ctypes.c_void_p),
                             occslstb.ctypes.data_as(ctypes.c_void_p))
    return hdiag