Python numpy.complex128() Examples

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 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 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 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_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_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 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 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 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_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 _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 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 __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):
        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 qubit_generator_matrix(self):
        generator = np.zeros((8, 8), dtype=np.complex128)
        # w0 |110><101| + h.c.
        generator[6, 5] = self.weights[0]
        generator[5, 6] = self.weights[0].conjugate()
        # w1 |110><011| + h.c.
        generator[6, 3] = self.weights[1]
        generator[3, 6] = self.weights[1].conjugate()
        # w2 |101><011| + h.c.
        generator[5, 3] = self.weights[2]
        generator[3, 5] = self.weights[2].conjugate()
        return generator 

def test_three_qubit_rotation_gates_on_simulator(gate: cirq.Gate,
                                                 initial_state: np.ndarray,
                                                 correct_state: np.ndarray):
    op = gate(*cirq.LineQubit.range(3))
    result = cirq.Circuit(op).final_wavefunction(
        initial_state, dtype=np.complex128)
    cirq.testing.assert_allclose_up_to_global_phase(result,
                                                    correct_state,
                                                    atol=1e-8) 

def generate_fswap_unitaries(swap_pairs: List[List[Tuple]], dimension: int):
    swap_unitaries = []
    for swap_tuples in swap_pairs:
        generator = np.zeros((dimension, dimension), dtype=np.complex128)
        for i, j in swap_tuples:
            generator[i, i] = -1
            generator[j, j] = -1
            generator[i, j] = 1
            generator[j, i] = 1
        swap_unitaries.append(expm(-1j * np.pi * generator / 2))
    return swap_unitaries 

def test_opdm_func_vals():
    # coverage: ignore
    rhf_objective, molecule, parameters, obi, tbi = make_h6_1_3()
    qubits = [cirq.GridQubit(0, x) for x in range(molecule.n_orbitals)]
    np.random.seed(43)
    sampler = cirq.Simulator(dtype=np.complex128)
    opdm_func = OpdmFunctional(qubits=qubits,
                               sampler=sampler,
                               constant=molecule.nuclear_repulsion,
                               one_body_integrals=obi,
                               two_body_integrals=tbi,
                               num_electrons=molecule.n_electrons // 2)

    assert isinstance(opdm_func, OpdmFunctional)

    data = opdm_func.calculate_data(parameters)
    assert isinstance(data, dict)
    assert list(data.keys()) == ['z', 'xy_even', 'xy_odd', 'qubits',
                           'qubit_permutations', 'circuits',
                           'circuits_with_measurement']

    opdm_from_data = compute_opdm(data, return_variance=False)

    opdm_from_obj, var_dict = opdm_func.calculate_rdm(parameters)
    assert isinstance(var_dict, dict)
    assert np.linalg.norm(opdm_from_data - opdm_from_obj) < 1.0E-2

    assert np.isclose(opdm_func.energy_from_opdm(opdm_from_data),
                      rhf_objective.energy_from_opdm(opdm_from_data))

    rdm_gen = RDMGenerator(opdm_func)
    rdm_gen.opdm_generator(parameters)
    assert len(rdm_gen.noisy_opdms) == 1
    assert len(rdm_gen.variance_dicts) == 1 

def get_opdm(wf, num_orbitals, transform=of.jordan_wigner):
    opdm_hw = np.zeros((num_orbitals, num_orbitals),
                       dtype=np.complex128)
    creation_ops = [
        of.get_sparse_operator(transform(of.FermionOperator(((p, 1)))),
                               n_qubits=num_orbitals)
        for p in range(num_orbitals)
    ]
    # not using display style objects
    for p, q in product(range(num_orbitals), repeat=2):
        operator = creation_ops[p] @ creation_ops[q].conj().transpose()
        opdm_hw[p, q] = wf.conj().T @ operator @ wf

    return opdm_hw 

def test_gen_fswap_unitaries():
    fswapu = generate_fswap_unitaries([((0, 1), (2, 3))], 4)
    true_generator = np.zeros((4, 4), dtype=np.complex128)
    true_generator[0, 0], true_generator[1, 1] = -1, -1
    true_generator[0, 1], true_generator[1, 0] = 1, 1
    true_generator[2, 2], true_generator[3, 3] = -1, -1
    true_generator[2, 3], true_generator[3, 2] = 1, 1
    true_u = sp.linalg.expm(-1j * np.pi * true_generator / 2)
    assert np.allclose(true_u, fswapu[0]) 

def eval_mat(self, cell, ao, weight, rho, vxc,
                 non0tab=None, xctype='LDA', spin=0, verbose=None):
        # Guess whether ao is evaluated for kpts_band.  When xctype is LDA, ao on grids
        # should be a 2D array.  For other xc functional, ao should be a 3D array.
        if ao.ndim == 2 or (xctype != 'LDA' and ao.ndim == 3):
            mat = eval_mat(cell, ao, weight, rho, vxc, non0tab, xctype, spin, verbose)
        else:
            nkpts = len(ao)
            nao = ao[0].shape[-1]
            mat = numpy.empty((nkpts,nao,nao), dtype=numpy.complex128)
            for k in range(nkpts):
                mat[k] = eval_mat(cell, ao[k], weight, rho, vxc,
                                  non0tab, xctype, spin, verbose)
        return mat 

def test_eval_rhoG_orth_kpts(self):
        numpy.random.seed(9)
        dm = numpy.random.random(dm1.shape) + numpy.random.random(dm1.shape) * 1j
        mydf = multigrid.MultiGridFFTDF(cell_orth)
        rhoG = multigrid._eval_rhoG(mydf, dm, hermi=0, kpts=kpts, deriv=0,
                                    rhog_high_order=True)
        self.assertTrue(rhoG.dtype == numpy.complex128)

        mydf = df.FFTDF(cell_orth)
        ni = dft.numint.KNumInt()
        ao_kpts = ni.eval_ao(cell_orth, mydf.grids.coords, kpts, deriv=0)
        ref = ni.eval_rho(cell_orth, ao_kpts, dm, hermi=0, xctype='LDA')
        rhoR = tools.ifft(rhoG[0], cell_orth.mesh).real
        rhoR *= numpy.prod(cell_orth.mesh)/cell_orth.vol
        self.assertAlmostEqual(abs(rhoR-ref).max(), 0, 7) 

def test_eval_rho(self):
        cell, grids = make_grids([61]*3)
        numpy.random.seed(10)
        nao = 10
        ngrids = 500
        dm = numpy.random.random((nao,nao))
        dm = dm + dm.T
        ao =(numpy.random.random((10,ngrids,nao)) +
             numpy.random.random((10,ngrids,nao))*1j)
        ao = ao.transpose(0,2,1).copy().transpose(0,2,1)

        rho0 = numpy.zeros((6,ngrids), dtype=numpy.complex128)
        rho0[0] = numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[0].conj())
        rho0[1] = numpy.einsum('pi,ij,pj->p', ao[1], dm, ao[0].conj()) + numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[1].conj())
        rho0[2] = numpy.einsum('pi,ij,pj->p', ao[2], dm, ao[0].conj()) + numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[2].conj())
        rho0[3] = numpy.einsum('pi,ij,pj->p', ao[3], dm, ao[0].conj()) + numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[3].conj())
        rho0[4]+= numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[4].conj()) + numpy.einsum('pi,ij,pj->p', ao[4], dm, ao[0].conj())
        rho0[4]+= numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[7].conj()) + numpy.einsum('pi,ij,pj->p', ao[7], dm, ao[0].conj())
        rho0[4]+= numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[9].conj()) + numpy.einsum('pi,ij,pj->p', ao[9], dm, ao[0].conj())
        rho0[5]+= numpy.einsum('pi,ij,pj->p', ao[1], dm, ao[1].conj())
        rho0[5]+= numpy.einsum('pi,ij,pj->p', ao[2], dm, ao[2].conj())
        rho0[5]+= numpy.einsum('pi,ij,pj->p', ao[3], dm, ao[3].conj())
        rho0[4]+= rho0[5]*2
        rho0[5] *= .5

        rho1 = numint.eval_rho(cell, ao, dm, xctype='MGGA')
        self.assertTrue(numpy.allclose(rho0, rho1))

        rho1 = numint.eval_rho(cell, ao, dm, xctype='GGA')
        self.assertAlmostEqual(lib.fp(rho1), -255.45150185669198, 7)

        rho1 = numint.eval_rho(cell, ao[0], dm, xctype='LDA')
        self.assertAlmostEqual(lib.fp(rho1), -17.198879910245601, 7) 

def test_get_veff(self):
        mf = pscf.RHF(cell)
        numpy.random.seed(1)
        nao = cell.nao_nr()
        dm = numpy.random.random((nao,nao)) + numpy.random.random((nao,nao))*1j
        dm = dm + dm.conj().T
        v11 = mf.get_veff(cell, dm, kpt=cell.get_abs_kpts([.25,.25,.25]))
        v12 = mf.get_veff(cell, dm, kpts_band=cell.get_abs_kpts([.25,.25,.25]))
        v13 = mf.get_veff(cell, dm, kpt=cell.get_abs_kpts([-1./3,1./3,.25]),
                          kpts_band=cell.get_abs_kpts([.25,.25,.25]))
        v14 = mf.get_veff(cell, dm, kpt=cell.get_abs_kpts([-1./3,1./3,.25]),
                          kpts_band=cell.make_kpts([2,1,1]))
        self.assertTrue(v11.dtype == numpy.complex128)
        self.assertTrue(v12.dtype == numpy.complex128)

        mf = pscf.UHF(cell)
        v21 = mf.get_veff(cell, dm, kpt=cell.get_abs_kpts([.25,.25,.25]))
        dm = [dm*.5,dm*.5]
        v22 = mf.get_veff(cell, dm, kpts_band=cell.get_abs_kpts([.25,.25,.25]))
        v23 = mf.get_veff(cell, dm, kpt=cell.get_abs_kpts([-1./3,1./3,.25]),
                          kpts_band=cell.get_abs_kpts([.25,.25,.25]))
        v24 = mf.get_veff(cell, dm, kpt=cell.get_abs_kpts([-1./3,1./3,.25]),
                          kpts_band=cell.make_kpts([2,1,1]))
        self.assertAlmostEqual(abs(v11-v21).max(), 0, 9)
        self.assertAlmostEqual(abs(v12-v22).max(), 0, 9)
        self.assertAlmostEqual(abs(v13-v23).max(), 0, 9)
        self.assertAlmostEqual(abs(v14-v24).max(), 0, 9)
        self.assertAlmostEqual(lib.finger(v11), -0.30110964334164825+0.81409418199767414j, 9)
        self.assertAlmostEqual(lib.finger(v12), -2.1601376488983997-9.4070613374115908j, 9) 

def run_kcell_complex(cell, nk):
    abs_kpts = cell.make_kpts(nk, wrap_around=True)
    kmf = pbcscf.KRHF(cell, abs_kpts)
    kmf.conv_tol = 1e-12
    ekpt = kmf.scf()
    kmf.mo_coeff = [kmf.mo_coeff[i].astype(np.complex128) for i in range(np.prod(nk))]
    mp = pyscf.pbc.mp.kmp2.KMP2(kmf).run()
    return ekpt, mp.e_corr 

def get_SI(cell, Gv=None):
    '''Calculate the structure factor (0D, 1D, 2D, 3D) for all atoms; see MH (3.34).

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

        Gv : (N,3) array
            G vectors

    Returns:
        SI : (natm, ngrids) ndarray, dtype=np.complex128
            The structure factor for each atom at each G-vector.
    '''
    coords = cell.atom_coords()
    ngrids = np.prod(cell.mesh)
    if Gv is None or Gv.shape[0] == ngrids:
        basex, basey, basez = cell.get_Gv_weights(cell.mesh)[1]
        b = cell.reciprocal_vectors()
        rb = np.dot(coords, b.T)
        SIx = np.exp(-1j*np.einsum('z,g->zg', rb[:,0], basex))
        SIy = np.exp(-1j*np.einsum('z,g->zg', rb[:,1], basey))
        SIz = np.exp(-1j*np.einsum('z,g->zg', rb[:,2], basez))
        SI = SIx[:,:,None,None] * SIy[:,None,:,None] * SIz[:,None,None,:]
        SI = SI.reshape(-1,ngrids)
    else:
        SI = np.exp(-1j*np.dot(coords, Gv.T))
    return SI 

def get_A_mat(self):
        '''
        Construct the projection matrix: A_{m,n}^{\mathbf{k}}
        Equation (62) in MV, Phys. Rev. B 56, 12847 or equation (22) in SMV, Phys. Rev. B 65, 035109
        '''

        A_matrix_loc = np.empty([self.num_kpts_loc, self.num_wann_loc, self.num_bands_loc], dtype = np.complex128)

        if self.use_bloch_phases == True:
            Amn = np.zeros([self.num_wann_loc, self.num_bands_loc])
            np.fill_diagonal(Amn, 1)
            A_matrix_loc[:,:,:] = Amn
        else:
            from pyscf.dft import numint
            grids = gen_grid.Grids(self.cell).build()
            coords = grids.coords
            weights = grids.weights
            for ith_wann in range(self.num_wann_loc):
                frac_site = self.proj_site[ith_wann]
                abs_site = frac_site.dot(self.real_lattice_loc) / param.BOHR
                l = self.proj_l[ith_wann]
                mr = self.proj_m[ith_wann]
                r = self.proj_radial[ith_wann]
                zona = self.proj_zona[ith_wann]
                x_axis = self.proj_x[ith_wann]
                z_axis = self.proj_z[ith_wann]
                gr = g_r(coords, abs_site, l, mr, r, zona, x_axis, z_axis, unit = 'B')
                ao_L0 = numint.eval_ao(self.cell, coords)
                s_aoL0_g = np.einsum('i,i,iv->v', weights, gr, ao_L0, optimize = True)
                for k_id in range(self.num_kpts_loc):
                    kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
                    mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
                    s_kpt = self.cell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpt, pbcopt=lib.c_null_ptr())
                    A_matrix_loc[k_id,ith_wann,:] = np.einsum('v,vu,um->m', s_aoL0_g, s_kpt, mo_included, optimize = True).conj()

        return A_matrix_loc 

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 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