Python numpy.conj() Examples

def test_givens_inverse():
    r"""
    The Givens rotation in OpenFermion is defined as

    .. math::

        \begin{pmatrix}
            \cos(\theta) & -e^{i \varphi} \sin(\theta) \\
            \sin(\theta) &     e^{i \varphi} \cos(\theta)
        \end{pmatrix}.

    confirm numerically its hermitian conjugate is it's inverse
    """
    a = numpy.random.random() + 1j * numpy.random.random()
    b = numpy.random.random() + 1j * numpy.random.random()
    ab_rotation = givens_matrix_elements(a, b, which='right')

    assert numpy.allclose(ab_rotation.dot(numpy.conj(ab_rotation).T),
                          numpy.eye(2))
    assert numpy.allclose(numpy.conj(ab_rotation).T.dot(ab_rotation),
                          numpy.eye(2)) 

def test_circuit_generation_and_accuracy():
    for dim in range(2, 10):
        qubits = cirq.LineQubit.range(dim)
        u_generator = numpy.random.random(
            (dim, dim)) + 1j * numpy.random.random((dim, dim))
        u_generator = u_generator - numpy.conj(u_generator).T
        assert numpy.allclose(-1 * u_generator, numpy.conj(u_generator).T)

        unitary = scipy.linalg.expm(u_generator)
        circuit = cirq.Circuit()
        circuit.append(optimal_givens_decomposition(qubits, unitary))

        fermion_generator = QubitOperator(()) * 0.0
        for i, j in product(range(dim), repeat=2):
            fermion_generator += jordan_wigner(
                FermionOperator(((i, 1), (j, 0)), u_generator[i, j]))

        true_unitary = scipy.linalg.expm(
            get_sparse_operator(fermion_generator).toarray())
        assert numpy.allclose(true_unitary.conj().T.dot(true_unitary),
                              numpy.eye(2 ** dim, dtype=complex))

        test_unitary = cirq.unitary(circuit)
        assert numpy.isclose(
            abs(numpy.trace(true_unitary.conj().T.dot(test_unitary))), 2 ** dim) 

def _gamma1_intermediates(mp, t2=None):
    # Memory optimization should be here
    if t2 is None:
        t2 = mp.t2
    if t2 is None:
        raise NotImplementedError("Run kmp2.kernel with `with_t2=True`")
    nmo = mp.nmo
    nocc = mp.nocc
    nvir = nmo - nocc
    nkpts = mp.nkpts
    dtype = t2.dtype

    dm1occ = np.zeros((nkpts, nocc, nocc), dtype=dtype)
    dm1vir = np.zeros((nkpts, nvir, nvir), dtype=dtype)

    for ki in range(nkpts):
        for kj in range(nkpts):
            for ka in range(nkpts):
                kb = mp.khelper.kconserv[ki, ka, kj]

                dm1vir[kb] += einsum('ijax,ijay->yx', t2[ki][kj][ka].conj(), t2[ki][kj][ka]) * 2 -\
                              einsum('ijax,ijya->yx', t2[ki][kj][ka].conj(), t2[ki][kj][kb])
                dm1occ[kj] += einsum('ixab,iyab->xy', t2[ki][kj][ka].conj(), t2[ki][kj][ka]) * 2 -\
                              einsum('ixab,iyba->xy', t2[ki][kj][ka].conj(), t2[ki][kj][kb])
    return -dm1occ, dm1vir 

def get_j(cell, dm, hermi=1, vhfopt=None, kpt=np.zeros(3), kpts_band=None):
    dm = np.asarray(dm)
    nao = dm.shape[-1]

    coords = gen_grid.gen_uniform_grids(cell)
    if kpts_band is None:
        kpt1 = kpt2 = kpt
        aoR_k1 = aoR_k2 = numint.eval_ao(cell, coords, kpt)
    else:
        kpt1 = kpts_band
        kpt2 = kpt
        aoR_k1 = numint.eval_ao(cell, coords, kpt1)
        aoR_k2 = numint.eval_ao(cell, coords, kpt2)
    ngrids, nao = aoR_k1.shape

    def contract(dm):
        vjR_k2 = get_vjR(cell, dm, aoR_k2)
        vj = (cell.vol/ngrids) * np.dot(aoR_k1.T.conj(), vjR_k2.reshape(-1,1)*aoR_k1)
        return vj

    if dm.ndim == 2:
        vj = contract(dm)
    else:
        vj = lib.asarray([contract(x) for x in dm.reshape(-1,nao,nao)])
    return vj.reshape(dm.shape) 

def get_vkR(mf, cell, aoR_k1, aoR_k2, kpt1, kpt2):
    '''Get the real-space 2-index "exchange" potential V_{i,k1; j,k2}(r)
    where {i,k1} = exp^{i k1 r) |i> , {j,k2} = exp^{-i k2 r) <j|
    '''
    coords = gen_grid.gen_uniform_grids(cell)
    ngrids, nao = aoR_k1.shape

    expmikr = np.exp(-1j*np.dot(kpt1-kpt2,coords.T))
    coulG = tools.get_coulG(cell, kpt1-kpt2, exx=True, mf=mf)
    def prod(ij):
        i, j = divmod(ij, nao)
        rhoR = aoR_k1[:,i] * aoR_k2[:,j].conj()
        rhoG = tools.fftk(rhoR, cell.mesh, expmikr)
        vG = coulG*rhoG
        vR = tools.ifftk(vG, cell.mesh, expmikr.conj())
        return vR

    if aoR_k1.dtype == np.double and aoR_k2.dtype == np.double:
        vR = numpy.asarray([prod(ij).real for ij in range(nao**2)])
    else:
        vR = numpy.asarray([prod(ij) for ij in range(nao**2)])
    return vR.reshape(nao,nao,-1).transpose(2,0,1) 

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 get_k(self,dm):
        """
        Update Exact Exchange Matrix

        Args:
            dm: float or complex
                AO density matrix.
        Returns:
            kmat: float or complex
                Exact Exchange in AO basis
        """
        naux = self.auxmol.nao_nr()
        nao = self.ks.mol.nao_nr()
        kpj = np.einsum("ijp,jk->ikp", self.eri3c, dm)
        pik = np.linalg.solve(self.eri2c, kpj.reshape(-1,naux).T.conj())
        kmat = np.einsum("pik,kjp->ij", pik.reshape(naux,nao,nao), self.eri3c)
        return kmat 

def unitary_to_process_mx(U):
    """
    Compute the super-operator which acts on (row)-vectorized
    density matrices from a unitary operator (matrix) U which
    acts on state vectors.  This super-operator is given by
    the tensor product of U and conjugate(U), i.e. kron(U,U.conj).

    Parameters
    ----------
    U : numpy array
        The unitary matrix which acts on state vectors.

    Returns
    -------
    numpy array
       The super-operator process matrix.
    """
    # U -> kron(U,Uc) since U rho U_dag -> kron(U,Uc)
    #  since AXB --row-vectorize--> kron(A,B.T)*vec(X)
    return _np.kron(U, _np.conjugate(U)) 

def svd_dotH(self,s1,q1):
        """
        Performs diagonal matrix multiplication conjugate
        
        Implements :math:`q_0 = \\mathrm{diag}(s_1)^* q_1`.
        
        :param s1: diagonal parameters
        :param q1: input to the diagonal multiplication
        :returns: :code:`q0` diagonal multiplication output
        """
        srep = repeat_axes(np.conj(s1),self.sshape,self.srep_axes,rep=False)
        q0 = srep*q1
        return q0
        
    #def __str__(self):
    #    string = str(self.name) + '\n'\
    #              + 'Input: ' + str(self.shape0) + ',' + str(self.dtype0)
    #    return string 

def azimuthal_symmetry(X, Y, n, m):
    p1 = 2
    p2 = 1
    p = np.sqrt(p1**2 + p2**2)

    detX = np.real(X.hvhv*(X.hhhh*X.vvvv - X.hhvv*np.conj(X.hhvv)))
    detY = np.real(Y.hvhv*(Y.hhhh*Y.vvvv - Y.hhvv*np.conj(Y.hhvv)))
    detXY = np.real((X.hvhv+Y.hvhv) * ((X.hhhh+Y.hhhh)*(X.vvvv+Y.vvvv) - (X.hhvv+Y.hhvv)*(np.conj(X.hhvv)+np.conj(Y.hhvv))))

    lnq = (p*(n+m)*np.log(n+m) - p*n*np.log(n) - p*m*np.log(m)
            + n*np.log(detX) + m*np.log(detY) - (n+m)*np.log(detXY))

    rho1 = 1 - (2*p1**2 - 1)/(6*p1) * (1/n + 1/m - 1/(n+m))
    rho2 = 1 - (2*p2**2 - 1)/(6*p2) * (1/n + 1/m - 1/(n+m))
    rho = 1/p**2 * (p1**2 * rho1 + p2**2 * rho2)

    w2 = - p**2/4 * (1-1/rho)**2 + (p1**2*(p1**2-1) + p2**2*(p2**2-1))/24 * (1/n**2 + 1/m**2 - 1/(n+m)**2) * 1/rho**2

    return lnq, rho, w2 

def fexpand(x, ns=1, axis=None):
    """
    Reconstructs full spectrum from positive frequencies
    Works on the last dimension (contiguous in c-stored array)

    :param x: numpy.ndarray
    :param axis: axis along which to perform reduction (last axis by default)
    :return: numpy.ndarray
    """
    if axis is None:
        axis = x.ndim - 1
    # dec = int(ns % 2) * 2 - 1
    # xcomp = np.conj(np.flip(x[..., 1:x.shape[-1] + dec], axis=axis))
    ilast = int((ns + (ns % 2)) / 2)
    xcomp = np.conj(np.flip(np.take(x, np.arange(1, ilast), axis=axis), axis=axis))
    return np.concatenate((x, xcomp), axis=axis) 

def _power_spectrum(daft, dim, N, density):

    ps = (daft * np.conj(daft)).real

    if density:
        ps /= (np.asarray(N).prod()) ** 2
        for i in dim:
            ps /= daft['freq_' + i + '_spacing']

    return ps 

def _cross_spectrum(daft1, daft2, dim, N, density):
    cs = (daft1 * np.conj(daft2)).real

    if density:
        cs /= (np.asarray(N).prod())**2
        for i in dim:
            cs /= daft1['freq_' + i + '_spacing']

    return cs 

def test_power_spectrum(self, dask):
        """Test the power spectrum function"""
        N = 16
        da = xr.DataArray(np.random.rand(2,N,N), dims=['time','y','x'],
                         coords={'time':np.array(['2019-04-18', '2019-04-19'],
                                                dtype='datetime64'),
                                'y':range(N),'x':range(N)}
                         )
        if dask:
            da = da.chunk({'time': 1})
        ps = xrft.power_spectrum(da, dim=['y','x'], window=True, density=False,
                                detrend='constant')
        daft = xrft.dft(da, dim=['y','x'], detrend='constant', window=True)
        npt.assert_almost_equal(ps.values, np.real(daft*np.conj(daft)))
        npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)

        ps = xrft.power_spectrum(da, dim=['y'], real='x', window=True,
                                density=False, detrend='constant')
        daft = xrft.dft(da, dim=['y'], real='x', detrend='constant', window=True)
        npt.assert_almost_equal(ps.values, np.real(daft*np.conj(daft)))

        ### Normalized
        ps = xrft.power_spectrum(da, dim=['y','x'], window=True, detrend='constant')
        daft = xrft.dft(da, dim=['y','x'], window=True, detrend='constant')
        test = np.real(daft*np.conj(daft))/N**4
        dk = np.diff(np.fft.fftfreq(N, 1.))[0]
        test /= dk**2
        npt.assert_almost_equal(ps.values, test)
        npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.)

        ### Remove least-square fit
        ps = xrft.power_spectrum(da, dim=['y','x'],
                                window=True, density=False, detrend='linear'
                                )
        daft = xrft.dft(da, dim=['y','x'], window=True, detrend='linear')
        npt.assert_almost_equal(ps.values, np.real(daft*np.conj(daft)))
        npt.assert_almost_equal(np.ma.masked_invalid(ps).mask.sum(), 0.) 

def test_cross_spectrum(self, dask):
        """Test the cross spectrum function"""
        N = 16
        dim = ['x','y']
        da = xr.DataArray(np.random.rand(2,N,N), dims=['time','x','y'],
                          coords={'time':np.array(['2019-04-18', '2019-04-19'],
                                                 dtype='datetime64'),
                                 'x':range(N), 'y':range(N)})
        da2 = xr.DataArray(np.random.rand(2,N,N), dims=['time','x','y'],
                          coords={'time':np.array(['2019-04-18', '2019-04-19'],
                                                 dtype='datetime64'),
                                  'x':range(N), 'y':range(N)})
        if dask:
            da = da.chunk({'time': 1})
            da2 = da2.chunk({'time': 1})

        daft = xrft.dft(da, dim=dim, shift=True, detrend='constant',
                        window=True)
        daft2 = xrft.dft(da2, dim=dim, shift=True, detrend='constant',
                        window=True)
        cs = xrft.cross_spectrum(da, da2, dim=dim, window=True, density=False,
                                detrend='constant')
        npt.assert_almost_equal(cs.values, np.real(daft*np.conj(daft2)))
        npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.)

        cs = xrft.cross_spectrum(da, da2, dim=dim, shift=True, window=True,
                                detrend='constant')
        test = (daft * np.conj(daft2)).real.values/N**4

        dk = np.diff(np.fft.fftfreq(N, 1.))[0]
        test /= dk**2
        npt.assert_almost_equal(cs.values, test)
        npt.assert_almost_equal(np.ma.masked_invalid(cs).mask.sum(), 0.) 

def test_row_eliminate():
    """
    Test elemination of element in U[i, j] by rotating in i-1 and i.
    """
    dim = 3
    u_generator = numpy.random.random((dim, dim)) + 1j * numpy.random.random(
        (dim, dim))
    u_generator = u_generator - numpy.conj(u_generator).T

    # make sure the generator is actually antihermitian
    assert numpy.allclose(-1 * u_generator, numpy.conj(u_generator).T)

    unitary = scipy.linalg.expm(u_generator)

    # eliminate U[2, 0] by rotating in 1, 2
    gmat = givens_matrix_elements(unitary[1, 0], unitary[2, 0], which='right')
    givens_rotate(unitary, gmat, 1, 2, which='row')
    assert numpy.isclose(unitary[2, 0], 0.0)

    # eliminate U[1, 0] by rotating in 0, 1
    gmat = givens_matrix_elements(unitary[0, 0], unitary[1, 0], which='right')
    givens_rotate(unitary, gmat, 0, 1, which='row')
    assert numpy.isclose(unitary[1, 0], 0.0)

    # eliminate U[2, 1] by rotating in 1, 2
    gmat = givens_matrix_elements(unitary[1, 1], unitary[2, 1], which='right')
    givens_rotate(unitary, gmat, 1, 2, which='row')
    assert numpy.isclose(unitary[2, 1], 0.0) 

def test_circuit_generation_state():
    """
    Determine if we rotate the Hartree-Fock state correctly
    """
    simulator = cirq.Simulator()
    circuit = cirq.Circuit()
    qubits = cirq.LineQubit.range(4)
    circuit.append([cirq.X(qubits[0]), cirq.X(qubits[1]), cirq.X(qubits[1]),
                    cirq.X(qubits[2]), cirq.X(qubits[3]),
                    cirq.X(qubits[3])])  # alpha-spins are first then beta spins

    wavefunction = numpy.zeros((2 ** 4, 1), dtype=complex)
    wavefunction[10, 0] = 1.0

    dim = 2
    u_generator = numpy.random.random((dim, dim)) + 1j * numpy.random.random(
        (dim, dim))
    u_generator = u_generator - numpy.conj(u_generator).T
    unitary = scipy.linalg.expm(u_generator)

    circuit.append(optimal_givens_decomposition(qubits[:2], unitary))

    fermion_generator = QubitOperator(()) * 0.0
    for i, j in product(range(dim), repeat=2):
        fermion_generator += jordan_wigner(
            FermionOperator(((i, 1), (j, 0)), u_generator[i, j]))

    test_unitary = scipy.linalg.expm(
        get_sparse_operator(fermion_generator, 4).toarray())
    test_final_state = test_unitary.dot(wavefunction)
    cirq_wf = simulator.simulate(circuit).final_state
    assert numpy.allclose(cirq_wf, test_final_state.flatten()) 

def projection(a, b):
    return np.abs(np.dot(np.conj(a), b))**2 

def unobservable_modes(C, A, returnEigenValues = False):
    '''Returns all the unobservable modes of the pair A,C.
    
    Does the PBH test for observability for the system:
     dx = A*x
     y  = C*x
    
    Returns a list of the unobservable modes, and (optionally) 
    the corresponding eigenvalues.
    
    See Callier & Desoer "Linear System Theory", P. 253
    '''

    return uncontrollable_modes(A.conj().T, C.conj().T, returnEigenValues) 

def is_observable(C, A):
    '''Compute whether the pair (C,A) is observable.
    
    Returns True if observable, False otherwise.
    '''
    
    return is_controllable(A.conj().T, C.conj().T) 

def make_rdm1(mp, t2=None, kind="compact"):
    r"""
    Spin-traced one-particle density matrix in the MO basis representation.
    The occupied-virtual orbital response is not included.

    dm1[p,q] = <q_alpha^\dagger p_alpha> + <q_beta^\dagger p_beta>

    The convention of 1-pdm is based on McWeeney's book, Eq (5.4.20).
    The contraction between 1-particle Hamiltonian and rdm1 is
    E = einsum('pq,qp', h1, rdm1)

    Args:
        mp (KMP2): a KMP2 kernel object;
        t2 (ndarray): a t2 MP2 tensor;
        kind (str): either 'compact' or 'padded' - defines behavior for k-dependent MO basis sizes;

    Returns:
        A k-dependent single-particle density matrix.
    """
    if kind not in ("compact", "padded"):
        raise ValueError("The 'kind' argument should be either 'compact' or 'padded'")
    d_imds = _gamma1_intermediates(mp, t2=t2)
    result = []
    padding_idxs = padding_k_idx(mp, kind="joint")
    for (oo, vv), idxs in zip(zip(*d_imds), padding_idxs):
        oo += np.eye(*oo.shape)
        d = block_diag(oo, vv)
        d += d.conj().T
        if kind == "padded":
            result.append(d)
        else:
            result.append(d[np.ix_(idxs, idxs)])
    return result 

def init_amps(self, eris):
        time0 = time.clock(), time.time()
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        mo_e_o = [eris.mo_energy[k][:nocc] for k in range(nkpts)]
        mo_e_v = [eris.mo_energy[k][nocc:] for k in range(nkpts)]
        t1 = numpy.zeros((nkpts, nocc, nvir), dtype=numpy.complex128)
        t2 = numpy.zeros((nkpts, nkpts, nkpts, nocc, nocc, nvir, nvir), dtype=numpy.complex128)
        self.emp2 = 0
        eris_oovv = eris.oovv.copy()

        # Get location of padded elements in occupied and virtual space
        nonzero_opadding, nonzero_vpadding = padding_k_idx(self, kind="split")

        kconserv = kpts_helper.get_kconserv(self._scf.cell, self.kpts)
        for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
            kb = kconserv[ki, ka, kj]
            # For LARGE_DENOM, see t1new update above
            eia = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
            n0_ovp_ia = numpy.ix_(nonzero_opadding[ki], nonzero_vpadding[ka])
            eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]

            ejb = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
            n0_ovp_jb = numpy.ix_(nonzero_opadding[kj], nonzero_vpadding[kb])
            ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb]
            eijab = eia[:, None, :, None] + ejb[:, None, :]

            t2[ki, kj, ka] = eris_oovv[ki, kj, ka] / eijab

        t2 = numpy.conj(t2)
        self.emp2 = 0.25 * numpy.einsum('pqrijab,pqrijab', t2, eris_oovv).real
        self.emp2 /= nkpts

        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2.real)
        logger.timer(self, 'init mp2', *time0)
        return self.emp2, t1, t2 

def transform_irr2full(self,invec,kp,kq,kr):
        operation = self.get_transformation(kp,kq,kr)
        if operation == 0:
            return invec
        if operation == 1:
            return invec.transpose(2,3,0,1)
        if operation == 2:
            return numpy.conj(invec.transpose(1,0,3,2))
        if operation == 3:
            return numpy.conj(invec.transpose(3,2,1,0)) 

def get_jk(mf, cell, dm, hermi=1, vhfopt=None, kpt=np.zeros(3), kpts_band=None):
    dm = np.asarray(dm)
    nao = dm.shape[-1]

    coords = gen_grid.gen_uniform_grids(cell)
    if kpts_band is None:
        kpt1 = kpt2 = kpt
        aoR_k1 = aoR_k2 = numint.eval_ao(cell, coords, kpt)
    else:
        kpt1 = kpts_band
        kpt2 = kpt
        aoR_k1 = numint.eval_ao(cell, coords, kpt1)
        aoR_k2 = numint.eval_ao(cell, coords, kpt2)

    vkR_k1k2 = get_vkR(mf, cell, aoR_k1, aoR_k2, kpt1, kpt2)

    ngrids, nao = aoR_k1.shape
    def contract(dm):
        vjR_k2 = get_vjR(cell, dm, aoR_k2)
        vj = (cell.vol/ngrids) * np.dot(aoR_k1.T.conj(), vjR_k2.reshape(-1,1)*aoR_k1)

        #:vk = (cell.vol/ngrids) * np.einsum('rs,Rp,Rqs,Rr->pq', dm, aoR_k1.conj(),
        #:                                vkR_k1k2, aoR_k2)
        aoR_dm_k2 = np.dot(aoR_k2, dm)
        tmp_Rq = np.einsum('Rqs,Rs->Rq', vkR_k1k2, aoR_dm_k2)
        vk = (cell.vol/ngrids) * np.dot(aoR_k1.T.conj(), tmp_Rq)
        return vj, vk

    if dm.ndim == 2:
        vj, vk = contract(dm)
    else:
        jk = [contract(x) for x in dm.reshape(-1,nao,nao)]
        vj = lib.asarray([x[0] for x in jk])
        vk = lib.asarray([x[1] for x in jk])
    return vj.reshape(dm.shape), vk.reshape(dm.shape) 

def get_j_kpts(mf, cell, dm_kpts, kpts, kpts_band=None):
    coords = gen_grid.gen_uniform_grids(cell)
    nkpts = len(kpts)
    ngrids = len(coords)
    dm_kpts = np.asarray(dm_kpts)
    nao = dm_kpts.shape[-1]

    ni = numint.KNumInt(kpts)
    aoR_kpts = ni.eval_ao(cell, coords, kpts)
    if kpts_band is not None:
        aoR_kband = numint.eval_ao(cell, coords, kpts_band)

    dms = dm_kpts.reshape(-1,nkpts,nao,nao)
    nset = dms.shape[0]

    vjR = [get_vjR(cell, dms[i], aoR_kpts) for i in range(nset)]
    if kpts_band is not None:
        vj_kpts = [cell.vol/ngrids * lib.dot(aoR_kband.T.conj()*vjR[i], aoR_kband)
                   for i in range(nset)]
        if dm_kpts.ndim == 3:  # One set of dm_kpts for KRHF
            vj_kpts = vj_kpts[0]
        return lib.asarray(vj_kpts)
    else:
        vj_kpts = []
        for i in range(nset):
            vj = [cell.vol/ngrids * lib.dot(aoR_k.T.conj()*vjR[i], aoR_k)
                  for aoR_k in aoR_kpts]
            vj_kpts.append(lib.asarray(vj))
        return lib.asarray(vj_kpts).reshape(dm_kpts.shape) 

def test_get_j_non_hermitian(self):
        kpt = kpts[0]
        numpy.random.seed(2)
        nao = cell2.nao
        dm = numpy.random.random((nao,nao))
        mydf = fft.FFTDF(cell2)
        v1 = mydf.get_jk(dm, hermi=0, kpts=kpts[1], with_k=False)[0]
        eri = mydf.get_eri([kpts[1]]*4).reshape(nao,nao,nao,nao)
        ref = numpy.einsum('ijkl,ji->kl', eri, dm)
        self.assertAlmostEqual(abs(ref - v1).max(), 0, 12)
        self.assertTrue(abs(ref-ref.T.conj()).max() > 1e-5) 

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_uniform_grids(cell)
    aoR = 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 comp_tem_spectrum_nonin(self, tmp_fname = None):
        """ 
        Compute non-interacting polarizability

        Inputs:
        -------
            comegas (complex 1D array): frequency range (in Hartree) for which the polarizability is computed.
                                     The imaginary part control the width of the signal.
                                     For example, 
                                     td = tddft_iter_c(...)
                                     comegas = np.arange(0.0, 10.05, 0.05) + 1j*td.eps
        Output:
        -------
            gamma (complex 1D array): computed non-interacting eels spectrum
            self.dn0 (complex 2D array): computed non-interacting density change in prod basis
        
        """
        comegas = self.freq + 1.0j*self.eps

        gamma = np.zeros(comegas.shape, dtype=np.complex64)
        self.dn0 = np.zeros((comegas.shape[0], self.nprod), dtype=np.complex64)
        
        for iw, comega in enumerate(comegas):
            self.dn0[iw, :] = self.apply_rf0(self.V_freq[iw, :], comega) 
            gamma[iw] = np.dot(self.dn0[iw, :], np.conj(self.V_freq[iw, :]))
            if tmp_fname is not None:
                tmp = open(tmp_fname, "a")
                tmp.write("{0}   {1}   {2}\n".format(comega.real, -gamma[iw].real/np.pi, 
                                                                  -gamma[iw].imag/np.pi))
                tmp.close() # Need to open and close the file at every freq, otherwise
                            # tmp is written only at the end of the calculations, therefore,
                            # it is useless
 
        return -gamma/np.pi 

def llc_imag(x, w1, w2, e1, e2):
  """ Product of two Lorentzians one of which is conjugated """
  res = (lorentzian(x, w1, e1)*np.conj(lorentzian(x, w2, e2))).imag
  return res 

def transmat(M,U,inv = 1):
    if inv == 1:
        # U.t() * M * U
        Mtilde = np.dot(np.dot(U.T.conj(),M),U)
    elif inv == -1:
        # U * M * U.t()
        Mtilde = np.dot(np.dot(U,M),U.T.conj())
    return Mtilde