Python numpy.complex128() Examples

def get_M_mat(self):
        '''
        Construct the ovelap matrix: M_{m,n}^{(\mathbf{k,b})}
        Equation (25) in MV, Phys. Rev. B 56, 12847
        '''

        M_matrix_loc = np.empty([self.num_kpts_loc, self.nntot_loc, self.num_bands_loc, self.num_bands_loc], dtype = np.complex128)

        for k_id in range(self.num_kpts_loc):
            for nn in range(self.nntot_loc):
                    k1 = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
                    k_id2 = self.nn_list[nn, k_id, 0] - 1
                    k2_ = self.kpt_latt_loc[k_id2]
                    k2_scaled = k2_ + self.nn_list[nn, k_id, 1:4]
                    k2 = self.cell.get_abs_kpts(k2_scaled)
                    s_AO = df.ft_ao.ft_aopair(self.cell, -k2+k1, kpti_kptj=[k2,k1], q = np.zeros(3))[0]
                    Cm = self.mo_coeff_kpts[k_id][:,self.band_included_list]
                    Cn = self.mo_coeff_kpts[k_id2][:,self.band_included_list]
                    M_matrix_loc[k_id, nn,:,:] = np.einsum('nu,vm,uv->nm', Cn.T.conj(), Cm, s_AO, optimize = True).conj()

        return M_matrix_loc 

def test_quartic_fermionic_simulation_unitary(weights, exponent):
    generator = np.zeros((1 << 4,) * 2, dtype=np.complex128)

    # w0 |1001><0110| + h.c.
    generator[9, 6] = weights[0]
    generator[6, 9] = weights[0].conjugate()
    # w1 |1010><0101| + h.c.
    generator[10, 5] = weights[1]
    generator[5, 10] = weights[1].conjugate()
    # w2 |1100><0011| + h.c.
    generator[12, 3] = weights[2]
    generator[3, 12] = weights[2].conjugate()
    expected_unitary = la.expm(-1j * exponent * generator)

    gate = ofc.QuarticFermionicSimulationGate(weights, exponent=exponent)
    actual_unitary = cirq.unitary(gate)

    assert np.allclose(expected_unitary, actual_unitary) 

def test_rhf_exx_ewald_with_kpt(self):
        numpy.random.seed(1)
        k = numpy.random.random(3)
        mf = pbchf.RHF(cell, k, exxdiv='ewald')
        e1 = mf.kernel()
        self.assertAlmostEqual(e1, -4.2048655827967139, 8)
        self.assertTrue(mf.mo_coeff.dtype == numpy.complex128)

        kmf = pscf.KRHF(cell, k, exxdiv='ewald')
        e0 = kmf.kernel()
        self.assertTrue(numpy.allclose(e0,e1))

        # test bands
        numpy.random.seed(1)
        kpt_band = numpy.random.random(3)
        e1, c1 = mf.get_bands(kpt_band)
        e0, c0 = kmf.get_bands(kpt_band)
        self.assertAlmostEqual(abs(e0-e1).max(), 0, 7)
        self.assertAlmostEqual(lib.finger(e1), -6.8312867098806249, 7) 

def test_cubic_fermionic_simulation_gate_consistency_docstring(
        weights, exponent):
    generator = np.zeros((8, 8), dtype=np.complex128)
    # w0 |110><101| + h.c.
    generator[6, 5] = weights[0]
    generator[5, 6] = weights[0].conjugate()
    # w1 |110><011| + h.c.
    generator[6, 3] = weights[1]
    generator[3, 6] = weights[1].conjugate()
    # w2 |101><011| + h.c.
    generator[5, 3] = weights[2]
    generator[3, 5] = weights[2].conjugate()
    expected_unitary = la.expm(-1j * exponent * generator)

    gate = ofc.CubicFermionicSimulationGate(weights, exponent=exponent)
    actual_unitary = cirq.unitary(gate)

    assert np.allclose(expected_unitary, actual_unitary) 

def so_by_shell(mol, shls):
    '''Spin-orbit coupling ECP in spinor basis
    i/2 <Pauli_matrix dot l U(r)>
    '''
    li = mol.bas_angular(shls[0])
    lj = mol.bas_angular(shls[1])
    di = (li*4+2) * mol.bas_nctr(shls[0])
    dj = (lj*4+2) * mol.bas_nctr(shls[1])
    bas = numpy.vstack((mol._bas, mol._ecpbas))
    mol._env[AS_ECPBAS_OFFSET] = len(mol._bas)
    mol._env[AS_NECPBAS] = len(mol._ecpbas)
    buf = numpy.empty((di,dj), order='F', dtype=numpy.complex128)
    cache = numpy.empty(buf.size*48)
    fn = libecp.ECPso_spinor
    fn(buf.ctypes.data_as(ctypes.c_void_p),
       (ctypes.c_int*2)(di, dj),
       (ctypes.c_int*2)(*shls),
       mol._atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
       bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
       mol._env.ctypes.data_as(ctypes.c_void_p), lib.c_null_ptr(),
       cache.ctypes.data_as(ctypes.c_void_p))
    return buf 

def _eigen_components(self):
        components = [(0, np.diag([1, 1, 1, 0, 1, 0, 0, 1]))]
        nontrivial_part = np.zeros((3, 3), dtype=np.complex128)
        for ij, w in zip([(1, 2), (0, 2), (0, 1)], self.weights):
            nontrivial_part[ij] = w
            nontrivial_part[ij[::-1]] = w.conjugate()
        assert np.allclose(nontrivial_part, nontrivial_part.conj().T)
        eig_vals, eig_vecs = np.linalg.eigh(nontrivial_part)
        for eig_val, eig_vec in zip(eig_vals, eig_vecs.T):
            exp_factor = -eig_val / np.pi
            proj = np.zeros((8, 8), dtype=np.complex128)
            nontrivial_indices = np.array([3, 5, 6], dtype=np.intp)
            proj[nontrivial_indices[:, np.newaxis], nontrivial_indices] = (
                np.outer(eig_vec.conjugate(), eig_vec))
            components.append((exp_factor, proj))
        return components 

def __init__(self, j):
    self._j = j
    self._c2r = np.zeros( (2*j+1, 2*j+1), dtype=np.complex128)
    self._c2r[j,j]=1.0
    for m in range(1,j+1):
      self._c2r[m+j, m+j] = sgn[m] * np.sqrt(0.5) 
      self._c2r[m+j,-m+j] = np.sqrt(0.5) 
      self._c2r[-m+j,-m+j]= 1j*np.sqrt(0.5)
      self._c2r[-m+j, m+j]= -sgn[m] * 1j * np.sqrt(0.5)
    
    self._hc_c2r = np.conj(self._c2r).transpose()
    self._conj_c2r = np.conjugate(self._c2r) # what is the difference ? conj and conjugate
    self._tr_c2r = np.transpose(self._c2r)
    
    #print(abs(self._hc_c2r.conj().transpose()-self._c2r).sum())

  #
  #
  # 

def qubit_generator_matrix(self):
        """The (Hermitian) matrix G such that the gate's unitary is
        exp(-i * G).
        """

        generator = np.zeros((1 << 4,) * 2, dtype=np.complex128)

        # w0 |1001><0110| + h.c.
        generator[9, 6] = self.weights[0]
        generator[6, 9] = self.weights[0].conjugate()
        # w1 |1010><0101| + h.c.
        generator[10, 5] = self.weights[1]
        generator[5, 10] = self.weights[1].conjugate()
        # w2 |1100><0011| + h.c.
        generator[12, 3] = self.weights[2]
        generator[3, 12] = self.weights[2].conjugate()
        return generator 

def _contract_rho(bra, ket):
    #:rho  = numpy.einsum('pi,pi->p', bra.real, ket.real)
    #:rho += numpy.einsum('pi,pi->p', bra.imag, ket.imag)
    bra = bra.T
    ket = ket.T
    nao, ngrids = bra.shape
    rho = numpy.empty(ngrids)

    if not (bra.flags.c_contiguous and ket.flags.c_contiguous):
        rho  = numpy.einsum('ip,ip->p', bra.real, ket.real)
        rho += numpy.einsum('ip,ip->p', bra.imag, ket.imag)
    elif bra.dtype == numpy.double and ket.dtype == numpy.double:
        libdft.VXC_dcontract_rho(rho.ctypes.data_as(ctypes.c_void_p),
                                 bra.ctypes.data_as(ctypes.c_void_p),
                                 ket.ctypes.data_as(ctypes.c_void_p),
                                 ctypes.c_int(nao), ctypes.c_int(ngrids))
    elif bra.dtype == numpy.complex128 and ket.dtype == numpy.complex128:
        libdft.VXC_zcontract_rho(rho.ctypes.data_as(ctypes.c_void_p),
                                 bra.ctypes.data_as(ctypes.c_void_p),
                                 ket.ctypes.data_as(ctypes.c_void_p),
                                 ctypes.c_int(nao), ctypes.c_int(ngrids))
    else:
        rho  = numpy.einsum('ip,ip->p', bra.real, ket.real)
        rho += numpy.einsum('ip,ip->p', bra.imag, ket.imag)
    return rho 

def test_get_matrix_of_eigs():
    """
    Generate the matrix of [exp(i (li - lj)) - 1] / (i(li - lj)
    :return:
    """
    lam_vals = np.random.randn(4) + 1j * np.random.randn(4)
    mat_eigs = np.zeros((lam_vals.shape[0],
                         lam_vals.shape[0]),
                         dtype=np.complex128)
    for i, j in product(range(lam_vals.shape[0]), repeat=2):
        if np.isclose(abs(lam_vals[i] - lam_vals[j]), 0):
            mat_eigs[i, j] = 1
        else:
            mat_eigs[i, j] = (np.exp(1j * (lam_vals[i] - lam_vals[j])) - 1) / (
                        1j * (lam_vals[i] - lam_vals[j]))

    test_mat_eigs = get_matrix_of_eigs(lam_vals)
    assert np.allclose(test_mat_eigs, mat_eigs) 

def test_rhf_vcut_sph(self):
        mf = pbchf.RHF(cell, exxdiv='vcut_sph')
        e1 = mf.kernel()
        self.assertAlmostEqual(e1, -4.29190260870812, 8)
        self.assertTrue(mf.mo_coeff.dtype == numpy.double)

        mf = pscf.KRHF(cell, [[0,0,0]], exxdiv='vcut_sph')
        e0 = mf.kernel()
        self.assertTrue(numpy.allclose(e0,e1))

        numpy.random.seed(1)
        k = numpy.random.random(3)
        mf = pbchf.RHF(cell, k, exxdiv='vcut_sph')
        e1 = mf.kernel()
        self.assertAlmostEqual(e1, -4.1379172088570595, 8)
        self.assertTrue(mf.mo_coeff.dtype == numpy.complex128)

        mf = pscf.KRHF(cell, k, exxdiv='vcut_sph')
        e0 = mf.kernel()
        self.assertTrue(numpy.allclose(e0,e1)) 

def test_get_eri_gamma(self):
        odf = df.DF(cell)
        odf.linear_dep_threshold = 1e-7
        odf.auxbasis = 'weigend'
        odf.mesh = (6,)*3
        eri0000 = odf.get_eri()
        self.assertTrue(eri0000.dtype == numpy.double)
        self.assertAlmostEqual(eri0000.real.sum(), 41.61280626625331, 9)
        self.assertAlmostEqual(finger(eri0000), 1.9981472468639465, 9)

        eri1111 = kmdf.get_eri((kpts[0],kpts[0],kpts[0],kpts[0]))
        self.assertTrue(eri1111.dtype == numpy.double)
        self.assertAlmostEqual(eri1111.real.sum(), 41.61280626625331, 9)
        self.assertAlmostEqual(eri1111.imag.sum(), 0, 9)
        self.assertAlmostEqual(finger(eri1111), 1.9981472468639465, 9)
        self.assertAlmostEqual(abs(eri1111-eri0000).max(), 0, 9)

        eri4444 = kmdf.get_eri((kpts[4],kpts[4],kpts[4],kpts[4]))
        self.assertTrue(eri4444.dtype == numpy.complex128)
        self.assertAlmostEqual(eri4444.real.sum(), 62.55120674032798, 9)
        self.assertAlmostEqual(abs(eri4444.imag).sum(), 0.0016507912195378644, 7)
        self.assertAlmostEqual(finger(eri4444), 0.6206014899350296-7.413680313987067e-05j, 8)
        eri0000 = ao2mo.restore(1, eri0000, cell.nao_nr()).reshape(eri4444.shape)
        self.assertAlmostEqual(abs(eri0000-eri4444).max(), 0, 4) 

def cc_Fov(cc, t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nkpts, nocc_a, nvir_a = t1a.shape
    nocc_b, nvir_b = t1b.shape[1:]

    fov = eris.fock[0][:,:nocc_a,nocc_a:]
    fOV = eris.fock[1][:,:nocc_b,nocc_b:]

    fa = np.zeros((nkpts,nocc_a,nvir_a), dtype=np.complex128)
    fb = np.zeros((nkpts,nocc_b,nvir_b), dtype=np.complex128)

    for km in range(nkpts):
        fa[km]+=fov[km]
        fb[km]+=fOV[km]
        for kn in range(nkpts):
            fa[km]+=einsum('nf,menf->me',t1a[kn],eris.ovov[km,km,kn])
            fa[km]+=einsum('nf,menf->me',t1b[kn],eris.ovOV[km,km,kn])
            fa[km]-=einsum('nf,mfne->me',t1a[kn],eris.ovov[km,kn,kn])
            fb[km]+=einsum('nf,menf->me',t1b[kn],eris.OVOV[km,km,kn])
            fb[km]+=einsum('nf,nfme->me',t1a[kn],eris.ovOV[kn,kn,km])
            fb[km]-=einsum('nf,mfne->me',t1b[kn],eris.OVOV[km,kn,kn])

    return fa,fb 

def export_unk(self, grid = [50,50,50]):
        '''
        Export the periodic part of BF in a real space grid for plotting with wannier90
        '''

        from scipy.io import FortranFile
        grids_coor, weights = periodic_grid(self.cell, grid, order = 'F')

        for k_id in range(self.num_kpts_loc):
            spin = '.1'
            if self.spin_up != None and self.spin_up == False : spin = '.2'
            kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
            ao = numint.eval_ao(self.cell, grids_coor, kpt = kpt)
            u_ao = np.einsum('x,xi->xi', np.exp(-1j*np.dot(grids_coor, kpt)), ao, optimize = True)
            unk_file = FortranFile('UNK' + "%05d" % (k_id + 1) + spin, 'w')
            unk_file.write_record(np.asarray([grid[0], grid[1], grid[2], k_id + 1, self.num_bands_loc], dtype = np.int32))
            mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
            u_mo = np.einsum('xi,in->xn', u_ao, mo_included, optimize = True)
            for band in range(len(self.band_included_list)):
                unk_file.write_record(np.asarray(u_mo[:,band], dtype = np.complex128))
            unk_file.close() 

def test_get_eri_0110(self):
        eri0110 = kmdf.get_eri((kpts[0],kpts[1],kpts[1],kpts[0]))
        self.assertTrue(eri0110.dtype == numpy.complex128)
        self.assertAlmostEqual(eri0110.real.sum(), 83.11379784456005, 9)
        self.assertAlmostEqual(abs(eri0110.imag).sum(), 5.083432749354445, 9)
        self.assertAlmostEqual(finger(eri0110), 0.9700096733262158-0.3318635561567801j, 9)
        check2 = kmdf.get_eri((kpts[0]+5e-8,kpts[1]+5e-8,kpts[1],kpts[0]))
        self.assertTrue(numpy.allclose(eri0110, check2, atol=1e-7)) 

def qubit_generator_matrix(self):
        generator = np.zeros((4, 4), dtype=np.complex128)
        # w0 |10><01| + h.c.
        generator[2, 1] = self.weights[0]
        generator[1, 2] = self.weights[0].conjugate()
        # w1 |11><11|
        generator[3, 3] = self.weights[1]
        return generator 

def test_get_eri_0123(self):
        eri0123 = kmdf.get_eri(kpts[:4])
        self.assertTrue(eri0123.dtype == numpy.complex128)
        self.assertAlmostEqual(eri0123.real.sum(), 83.11713402928555, 9)
        self.assertAlmostEqual(abs(eri0123.imag.sum()), 0.0006633634465733618, 9)
        self.assertAlmostEqual(finger(eri0123), 0.9693314315640165-0.33152709566516436j, 9) 

def eval_ao(mol, coords, deriv=0, with_s=True, shls_slice=None,
            non0tab=None, out=None, verbose=None):
    comp = (deriv+1)*(deriv+2)*(deriv+3)//6
    feval = 'GTOval_spinor_deriv%d' % deriv
    aoLa, aoLb = mol.eval_gto(feval, coords, comp, shls_slice, non0tab, out=out)
    if with_s:
        assert(deriv <= 1)  # only GTOval_ipsp_spinor
        ngrid, nao = aoLa.shape[-2:]
        if out is not None:
            aoSa = numpy.empty((comp,nao,ngrid), dtype=numpy.complex128)
            aoSb = numpy.empty((comp,nao,ngrid), dtype=numpy.complex128)
        else:
            out = numpy.ndarray((4,comp,nao,ngrid), dtype=numpy.complex128, buffer=out)
            aoSa, aoSb = out[2:]
        comp = 1
        ao = mol.eval_gto('GTOval_sp_spinor', coords, comp, shls_slice, non0tab)
        aoSa[0] = ao[0].T
        aoSb[0] = ao[1].T
        fevals = ['GTOval_sp_spinor', 'GTOval_ipsp_spinor']
        p1 = 1
        for n in range(1, deriv+1):
            comp = (n+1)*(n+2)//2
            ao = mol.eval_gto(fevals[n], coords, comp, shls_slice, non0tab)
            p0, p1 = p1, p1 + comp
            for k in range(comp):
                aoSa[p0:p1] = ao[0].transpose(0,2,1)
                aoSb[p0:p1] = ao[1].transpose(0,2,1)
        aoSa = aoSa.transpose(0,2,1)
        aoSb = aoSb.transpose(0,2,1)
        if deriv == 0:
            aoSa = aoSa[0]
            aoSb = aoSb[0]
        return aoLa, aoLb, aoSa, aoSb
    else:
        return aoLa, aoLb 

def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None):
    if kpts_band is not None:
        return get_j_for_bands(mydf, dm_kpts, hermi, kpts, kpts_band)

    dm_kpts = lib.asarray(dm_kpts, order='C')
    dms = _format_dms(dm_kpts, kpts)
    nset, nkpts, nao = dms.shape[:3]
    vj_kpts = numpy.zeros((nset,nkpts,nao,nao), dtype=numpy.complex128)
    kpt_allow = numpy.zeros(3)
    mesh = mydf.mesh
    coulG = mydf.weighted_coulG(kpt_allow, False, mesh)
    max_memory = (mydf.max_memory - lib.current_memory()[0]) * .8
    weight = 1./len(kpts)
    dmsC = dms.conj()
    for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts, max_memory=max_memory):
        vG = [0] * nset
        #:rho = numpy.einsum('lkL,lk->L', pqk.conj(), dm)
        for k, aoao in enumerate(aoaoks):
            for i in range(nset):
                rho = numpy.einsum('ij,Lij->L', dmsC[i,k],
                                   aoao.reshape(-1,nao,nao)).conj()
                vG[i] += rho * coulG[p0:p1]
        for i in range(nset):
            vG[i] *= weight
        for k, aoao in enumerate(aoaoks):
            for i in range(nset):
                vj_kpts[i,k] += numpy.einsum('L,Lij->ij', vG[i],
                                             aoao.reshape(-1,nao,nao))
    aoao = aoaoks = p0 = p1 = None

    if gamma_point(kpts):
        vj_kpts = vj_kpts.real.copy()
    return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts) 

def calc_prefactor_2D(x, y, f, wmin, tmin, dt, sign=1.0):

    if abs(sign) != 1.0:
        raise ValueError("sign must be 1.0 or -1.0")
    fmod = np.zeros(f.shape, dtype=np.complex128)

    for i in range(f.shape[0]):
        for j in range(f.shape[1]):
            fmod[i, j] = f[i, j]*dt[0]*dt[1]*np.exp(sign*1.0j*(wmin[0]*(x[i]-tmin[0]) +
                                    wmin[1]*(y[j]-tmin[1])))

    return fmod 

def test_get_eri_0110_high_cost(self):
        eri0110 = kmdf.get_eri((kpts[0],kpts[1],kpts[1],kpts[0]))
        self.assertTrue(eri0110.dtype == numpy.complex128)
        self.assertAlmostEqual(eri0110.real.sum(), 411.86033298299179, 6)
        self.assertAlmostEqual(abs(eri0110.imag).sum(), 136.58633427242452, 6)
        self.assertAlmostEqual(finger(eri0110), 1.3767132918850329+0.12378724026874122j, 6)
        check2 = kmdf.get_eri((kpts[0]+5e-9,kpts[1]+5e-9,kpts[1],kpts[0]))
        self.assertTrue(numpy.allclose(eri0110, check2, atol=1e-7))

#    def test_get_eri_0123_high_cost(self):
#        eri0123 = kmdf.get_eri(kpts[:4])
#        self.assertTrue(eri0123.dtype == numpy.complex128)
#        self.assertAlmostEqual(eri0123.real.sum(), 410.38308763371651, 6)
#        self.assertAlmostEqual(abs(eri0123.imag.sum()), 0.18510527268199378, 6)
#        self.assertAlmostEqual(finger(eri0123), 1.7644500565943559+0.30677193151572507j, 6) 

def test_get_eri_0011_high_cost(self):
        eri0011 = kmdf.get_eri((kpts[0],kpts[0],kpts[1],kpts[1]))
        self.assertTrue(eri0011.dtype == numpy.complex128)
        self.assertAlmostEqual(eri0011.real.sum(), 259.13073670793142, 6)
        self.assertAlmostEqual(abs(eri0011.imag).sum(), 8.4042424538275426, 6)
        self.assertAlmostEqual(finger(eri0011), 2.1374953278715552+0.12350314965485282j, 6) 

def test_get_eri_1111_high_cost(self):
        eri1111 = kmdf.get_eri((kpts[1],kpts[1],kpts[1],kpts[1]))
        self.assertTrue(eri1111.dtype == numpy.complex128)
        self.assertAlmostEqual(eri1111.real.sum(), 258.81872464108312, 6)
        self.assertAlmostEqual(abs(eri1111.imag).sum(), 16.275864968641145, 6)
        self.assertAlmostEqual(finger(eri1111), 2.2339104732873363+0.10954687420327755j, 7)
        check2 = kmdf.get_eri((kpts[1]+5e-9,kpts[1]+5e-9,kpts[1],kpts[1]))
        self.assertTrue(numpy.allclose(eri1111, check2, atol=1e-7)) 

def test_get_jk_kpts(self):
        df = fft.FFTDF(cell)
        dm = mf0.get_init_guess()
        nkpts = len(kpts)
        dms = [dm] * nkpts
        vj0, vk0 = get_jk_kpts(mf0, cell, dms, kpts=kpts)
        vj1, vk1 = df.get_jk(dms, kpts=kpts, exxdiv=None)
        self.assertTrue(vj1.dtype == numpy.complex128)
        self.assertTrue(vk1.dtype == numpy.complex128)
        self.assertTrue(np.allclose(vj0, vj1, atol=1e-9, rtol=1e-9))
        self.assertTrue(np.allclose(vk0, vk1, atol=1e-9, rtol=1e-9))

        ej1 = numpy.einsum('xij,xji->', vj1, dms) / len(kpts)
        ek1 = numpy.einsum('xij,xji->', vk1, dms) / len(kpts)
        self.assertAlmostEqual(ej1, 2.3163352969873445*(6/6.82991739766009)**2, 9)
        self.assertAlmostEqual(ek1, 7.7311228144548600*(6/6.82991739766009)**2, 9)

        numpy.random.seed(1)
        kpts_band = numpy.random.random((2,3))
        vj1, vk1 = df.get_jk(dms, kpts=kpts, kpts_band=kpts_band, exxdiv=None)
        self.assertAlmostEqual(lib.finger(vj1), 6/6.82991739766009*(3.437188138446714+0.1360466492092307j), 9)
        self.assertAlmostEqual(lib.finger(vk1), 6/6.82991739766009*(7.479986541097368+1.1980593415201204j), 9)

        nao = dm.shape[0]
        mo_coeff = numpy.random.random((nkpts,nao,nao))
        mo_occ = numpy.array(numpy.random.random((nkpts,nao))>.6, dtype=numpy.double)
        dms = numpy.einsum('kpi,ki,kqi->kpq', mo_coeff, mo_occ, mo_coeff)
        dms = lib.tag_array(lib.asarray(dms), mo_coeff=mo_coeff, mo_occ=mo_occ)
        vk1 = df.get_jk(dms, kpts=kpts, kpts_band=kpts_band, exxdiv=None)[1]
        self.assertAlmostEqual(lib.finger(vk1), 10.239828255099447+2.1190549216896182j, 9) 

def test_get_jk(self):
        df = fft.FFTDF(cell)
        dm = mf0.get_init_guess()
        vj0, vk0 = get_jk(mf0, cell, dm, kpt=kpts[0])
        vj1, vk1 = df.get_jk(dm, kpts=kpts[0], exxdiv=None)
        self.assertTrue(vj1.dtype == numpy.complex128)
        self.assertTrue(vk1.dtype == numpy.complex128)
        self.assertTrue(np.allclose(vj0, vj1, atol=1e-9, rtol=1e-9))
        self.assertTrue(np.allclose(vk0, vk1, atol=1e-9, rtol=1e-9))

        ej1 = numpy.einsum('ij,ji->', vj1, dm)
        ek1 = numpy.einsum('ij,ji->', vk1, dm)
        self.assertAlmostEqual(ej1, 2.3002596914518700*(6/6.82991739766009)**2, 9)
        self.assertAlmostEqual(ek1, 3.3165691757797346*(6/6.82991739766009)**2, 9)

        dm = mf0.get_init_guess()
        vj0, vk0 = get_jk(mf0, cell, dm)
        vj1, vk1 = df.get_jk(dm, exxdiv=None)
        self.assertTrue(vj1.dtype == numpy.float64)
        self.assertTrue(vk1.dtype == numpy.float64)
        self.assertTrue(np.allclose(vj0, vj1, atol=1e-9, rtol=1e-9))
        self.assertTrue(np.allclose(vk0, vk1, atol=1e-9, rtol=1e-9))

        ej1 = numpy.einsum('ij,ji->', vj1, dm)
        ek1 = numpy.einsum('ij,ji->', vk1, dm)
        self.assertAlmostEqual(ej1, 2.4673139106639925*(6/6.82991739766009)**2, 9)
        self.assertAlmostEqual(ek1, 3.6886674521354221*(6/6.82991739766009)**2, 9) 

def get_ao_pairs_G(cell, kpt=np.zeros(3)):
    '''Calculate forward (G|ij) and "inverse" (ij|G) FFT of all AO pairs.

    Args:
        cell : instance of :class:`Cell`

    Returns:
        ao_pairs_G, ao_pairs_invG : (ngrids, nao*(nao+1)/2) ndarray
            The FFTs of the real-space AO pairs.

    '''
    coords = gen_grid.gen_uniform_grids(cell)
    aoR = numint.eval_ao(cell, coords, kpt) # shape = (coords, nao)
    ngrids, nao = aoR.shape
    gamma_point = abs(kpt).sum() < 1e-9
    if gamma_point:
        npair = nao*(nao+1)//2
        ao_pairs_G = np.empty([ngrids, npair], np.complex128)

        ij = 0
        for i in range(nao):
            for j in range(i+1):
                ao_ij_R = np.conj(aoR[:,i]) * aoR[:,j]
                ao_pairs_G[:,ij] = tools.fft(ao_ij_R, cell.mesh)
                #ao_pairs_invG[:,ij] = ngrids*tools.ifft(ao_ij_R, cell.mesh)
                ij += 1
        ao_pairs_invG = ao_pairs_G.conj()
    else:
        ao_pairs_G = np.zeros([ngrids, nao,nao], np.complex128)
        for i in range(nao):
            for j in range(nao):
                ao_ij_R = np.conj(aoR[:,i]) * aoR[:,j]
                ao_pairs_G[:,i,j] = tools.fft(ao_ij_R, cell.mesh)
        ao_pairs_invG = ao_pairs_G.transpose(0,2,1).conj().reshape(-1,nao**2)
        ao_pairs_G = ao_pairs_G.reshape(-1,nao**2)
    return ao_pairs_G, ao_pairs_invG 

def _fftn_blas(f, mesh):
    Gx = np.fft.fftfreq(mesh[0])
    Gy = np.fft.fftfreq(mesh[1])
    Gz = np.fft.fftfreq(mesh[2])
    expRGx = np.exp(np.einsum('x,k->xk', -2j*np.pi*np.arange(mesh[0]), Gx))
    expRGy = np.exp(np.einsum('x,k->xk', -2j*np.pi*np.arange(mesh[1]), Gy))
    expRGz = np.exp(np.einsum('x,k->xk', -2j*np.pi*np.arange(mesh[2]), Gz))
    out = np.empty(f.shape, dtype=np.complex128)
    buf = np.empty(mesh, dtype=np.complex128)
    for i, fi in enumerate(f):
        buf[:] = fi.reshape(mesh)
        g = lib.dot(buf.reshape(mesh[0],-1).T, expRGx, c=out[i].reshape(-1,mesh[0]))
        g = lib.dot(g.reshape(mesh[1],-1).T, expRGy, c=buf.reshape(-1,mesh[1]))
        g = lib.dot(g.reshape(mesh[2],-1).T, expRGz, c=out[i].reshape(-1,mesh[2]))
    return out.reshape(-1, *mesh) 

def test_symm_orb_serialization(self):
        '''Handle the complex symmetry-adapted orbitals'''
        mol = gto.M(atom='He', basis='ccpvdz', symmetry=True)
        mol.loads(mol.dumps())

        lz_minus = numpy.sqrt(.5) * (mol.symm_orb[3] - mol.symm_orb[2] * 1j)
        lz_plus = -numpy.sqrt(.5) * (mol.symm_orb[3] + mol.symm_orb[2] * 1j)
        mol.symm_orb[2] = lz_minus
        mol.symm_orb[3] = lz_plus
        mol.loads(mol.dumps())
        self.assertTrue(mol.symm_orb[0].dtype == numpy.double)
        self.assertTrue(mol.symm_orb[2].dtype == numpy.complex128)
        self.assertTrue(mol.symm_orb[3].dtype == numpy.complex128) 

def spin_square(mol, mo):
    s = mol.intor('int1e_ovlp')
    nao = s.shape[0]
    sx = numpy.zeros((nao*2,nao*2))
    sy = numpy.zeros((nao*2,nao*2), dtype=numpy.complex128)
    sz = numpy.zeros((nao*2,nao*2))
    s1 = numpy.zeros((nao*2,nao*2))
    sx[:nao,nao:] = .5 * s
    sx[nao:,:nao] = .5 * s
    sy[:nao,nao:] =-.5j* s
    sy[nao:,:nao] = .5j* s
    sz[:nao,:nao] = .5 * s
    sz[nao:,nao:] =-.5 * s
    sx = reduce(numpy.dot, (mo.T.conj(), sx, mo))
    sy = reduce(numpy.dot, (mo.T.conj(), sy, mo))
    sz = reduce(numpy.dot, (mo.T.conj(), sz, mo))
    ss = numpy.einsum('ij,kl->ijkl', sx, sx) + 0j
    ss+= numpy.einsum('ij,kl->ijkl', sy, sy)
    ss+= numpy.einsum('ij,kl->ijkl', sz, sz)
    nmo = mo.shape[1]
    dm2 = numpy.einsum('ij,kl->ijkl', numpy.eye(nmo), numpy.eye(nmo))
    dm2-= numpy.einsum('jk,il->ijkl', numpy.eye(nmo), numpy.eye(nmo))
    ss = numpy.einsum('ijkl,ijkl', ss, dm2).real

    s1[:nao,:nao] = s
    s1[nao:,nao:] = s
    s1 = reduce(numpy.dot, (mo.T.conj(), s1, mo))
    ss+= s1.trace().real * .75
    return ss 

def test_get_eri_0011_1(self):
        eri0011 = kmdf1.get_eri((kpts[0],kpts[0],kpts[1],kpts[1]))
        self.assertTrue(eri0011.dtype == numpy.complex128)
        self.assertAlmostEqual(eri0011.real.sum(), 44.247892059548818, 6)
        self.assertAlmostEqual(abs(eri0011.imag).sum(), 6.5167066853467244, 6)
        self.assertAlmostEqual(finger(eri0011), (6.2170722376745422-0.10266245792326825j), 6)