Python numpy.require() Examples

def gw_comp_veff(self, vext, comega=1j*0.0):
    """
    This computes an effective field (scalar potential) given the external
    scalar potential as follows:
        (1-v\chi_{0})V_{eff} = V_{ext} = X_{a}^{n}V_{\mu}^{ab}X_{b}^{m} * 
                                         v\chi_{0}v * X_{a}^{n}V_{nu}^{ab}X_{b}^{m}
    
    returns V_{eff} as list for all n states(self.nn[s]).
    """
    
    from scipy.sparse.linalg import LinearOperator
    self.comega_current = comega
    veff_op = LinearOperator((self.nprod,self.nprod),
                             matvec=self.gw_vext2veffmatvec,
                             dtype=self.dtypeComplex)

    from scipy.sparse.linalg import lgmres
    resgm, info = lgmres(veff_op,
                         np.require(vext, dtype=self.dtypeComplex, requirements='C'),
                         atol=self.gw_iter_tol, maxiter=self.maxiter)
    if info != 0:
      print("LGMRES has not achieved convergence: exitCode = {}".format(info))
    return resgm 

def rsphar(r,lmax,res):
  """
    Computes (all) real spherical harmonics up to the angular momentum lmax
    Args:
      r : Cartesian coordinates defining correct theta and phi angles for spherical harmonic
      lmax : Integer, maximal angular momentum
    Result:
      1-d numpy array of float64 elements with all spherical harmonics stored in order 0,0; 1,-1; 1,0; 1,+1 ... lmax,lmax, althogether 0 : (lmax+1)**2 elements.
  """
  assert r.shape[-1]==3
  
  r_cp = np.require(r,  dtype=float, requirements='C')
  res = np.require(res,  dtype=float, requirements='CW')
  
  libnao.rsphar(r_cp.ctypes.data_as(POINTER(c_double)), c_int64(lmax), res.ctypes.data_as(POINTER(c_double)))
  return 0

#
#
# 

def rsphar_vec(rvs,lmax):
  """
    Computes (all) real spherical harmonics up to the angular momentum lmax
    Args:
      rvs : Cartesian coordinates defining correct theta and phi angles for spherical harmonic
      lmax : Integer, maximal angular momentum
    Result:
      1-d numpy array of float64 elements with all spherical harmonics stored in order 0,0; 1,-1; 1,0; 1,+1 ... lmax,lmax, althogether 0 : (lmax+1)**2 elements.
  """
  assert rvs.shape[-1]==3
  r_cp = np.require(rvs,  dtype=float, requirements='C')
  nc = len(rvs)
  res = np.require( np.zeros((nc, (lmax+1)**2)), dtype=float, requirements='CW')
  libnao.rsphar_vec(r_cp.ctypes.data_as(POINTER(c_double)), c_int64(nc), c_int64(lmax), res.ctypes.data_as(POINTER(c_double)))
  
  #for irv,rvec in enumerate(rvs): rsphar(rvec,lmax,res[irv,:])
  return res

#
#
# 

def rsphar_exp_vec(rvs,lmax):
  """
    Computes (all) real spherical harmonics up to the angular momentum lmax
    Args:
      rvs : Cartesian coordinates defining correct theta and phi angles for spherical harmonic
      lmax : Integer, maximal angular momentum
    Result:
      1-d numpy array of float64 elements with all spherical harmonics stored in order 0,0; 1,-1; 1,0; 1,+1 ... lmax,lmax, althogether 0 : (lmax+1)**2 elements.
  """  
  assert rvs.shape[0]==3
  r_cp = np.require(rvs,  dtype=np.float64, requirements='C')
  nc = rvs[0,...].size
  res = np.require( np.zeros(((lmax+1)**2,nc)), dtype=np.float64, requirements='CW')
  libnao.rsphar_exp_vec(r_cp.ctypes.data_as(POINTER(c_double)), c_int64(nc), c_int64(lmax), res.ctypes.data_as(POINTER(c_double)))
  
  #for irv,rvec in enumerate(rvs): rsphar(rvec,lmax,res[irv,:])
  return res 

def dens_libnao(crds, nspin):
  """  Compute the electronic density using library call """
  assert crds.ndim==2  
  assert crds.shape[-1]==3
  
  nc = crds.shape[0]
  crds_cp = require(crds, dtype=c_double, requirements='C')
  dens = require( zeros((nc, nspin)), dtype=c_double, requirements='CW')

  libnao.dens_libnao(
    crds_cp.ctypes.data_as(POINTER(c_double)),
    c_int64(nc),
    dens.ctypes.data_as(POINTER(c_double)),
    c_int64(nspin))

  return dens 

def get_ac_vertex_array(self, dtype=np.float64):
    """ Returns the product vertex coefficients as 3d array (dense table) """
    atom2so = self.sv.atom2s
    nfap = self.c2s[-1]
    n = self.sv.atom2s[-1]
    pab2v = np.require( np.zeros((nfap,n,n), dtype=dtype), requirements='CW')
    for atom,[sd,fd,pt,spp] in enumerate(zip(self.dpc2s,self.dpc2s[1:],self.dpc2t,self.dpc2sp)):
      if pt!=1: continue
      s,f = atom2so[atom:atom+2]
      pab2v[sd:fd,s:f,s:f] = self.prod_log.sp2vertex[spp]

    for sd,fd,pt,spp in zip(self.dpc2s,self.dpc2s[1:],self.dpc2t,self.dpc2sp):
      if pt!=2: continue
      inf= self.bp2info[spp]
      lab = einsum('dl,dab->lab', inf.cc, inf.vrtx)
      a,b = inf.atoms
      sa,fa,sb,fb = atom2so[a],atom2so[a+1],atom2so[b],atom2so[b+1]
      for c,ls,lf in zip(inf.cc2a, inf.cc2s, inf.cc2s[1:]):
        pab2v[self.c2s[c]:self.c2s[c+1],sa:fa,sb:fb] = lab[ls:lf,:,:]
        pab2v[self.c2s[c]:self.c2s[c+1],sb:fb,sa:fa] = einsum('pab->pba', lab[ls:lf,:,:])
    return pab2v 

def get_dp_vertex_array(self, dtype=np.float64):
    """ Returns the product vertex coefficients as 3d array for dominant products """
    atom2so = self.sv.atom2s
    nfdp = self.dpc2s[-1]
    n = self.sv.atom2s[-1]
    pab2v = np.require(np.zeros((nfdp,n,n), dtype=dtype), requirements='CW')
    for atom,[sd,fd,pt,spp] in enumerate(zip(self.dpc2s,self.dpc2s[1:],self.dpc2t,self.dpc2sp)):
      if pt!=1: continue
      s,f = atom2so[atom:atom+2]
      pab2v[sd:fd,s:f,s:f] = self.prod_log.sp2vertex[spp]

    for sd,fd,pt,spp in zip(self.dpc2s,self.dpc2s[1:],self.dpc2t,self.dpc2sp):
      if pt!=2: continue
      inf= self.bp2info[spp]
      a,b = inf.atoms
      sa,fa,sb,fb = atom2so[a],atom2so[a+1],atom2so[b],atom2so[b+1]
      pab2v[sd:fd,sa:fa,sb:fb] = inf.vrtx
      pab2v[sd:fd,sb:fb,sa:fa] = einsum('pab->pba', inf.vrtx)
    return pab2v 

def init_mo_from_pyscf(self, **kw):
    """ Initializing from a previous pySCF mean-field calc. """
    from pyscf.nao.m_fermi_energy import fermi_energy as comput_fermi_energy
    from pyscf.nao.m_color import color as tc
    self.telec = kw['telec'] if 'telec' in kw else 0.0000317 # 10K
    self.mf = mf = kw['mf']
    self.xc_code = mf.xc if hasattr(mf, 'xc') else 'HF'
    self.k2xyzw = np.array([[0.0,0.0,0.0,1.0]])
    
    self.mo_energy = np.asarray(mf.mo_energy)
    self.nspin = self.mo_energy.ndim
    assert self.nspin in [1,2]
    nspin,n=self.nspin,self.norbs
    self.mo_energy = require( self.mo_energy.reshape((1, nspin, n)), requirements='CW')
    self.mo_occ = require( mf.mo_occ.reshape((1,nspin,n)), requirements='CW')    
    self.mo_coeff =  require(zeros((1,nspin,n,n,1), dtype=self.dtype), requirements='CW')
    conv = conv_yzx2xyz_c(kw['gto'])
    aaux = np.asarray(mf.mo_coeff).reshape((nspin,n,n))
    for s in range(nspin):
      self.mo_coeff[0,s,:,:,0] = conv.conv_yzx2xyz_1d(aaux[s], conv.m_xyz2m_yzx).T

    self.nelec = kw['nelec'] if 'nelec' in kw else np.array([int(s2o.sum()) for s2o in self.mo_occ[0]])
    fermi = comput_fermi_energy(self.mo_energy, sum(self.nelec), self.telec)
    self.fermi_energy = kw['fermi_energy'] if 'fermi_energy' in kw else fermi 

def init_mo_coeff_fireball(self, **kw):
    """ Constructor a mean-field class from the preceeding FIREBALL calculation """
    from pyscf.nao.m_fermi_dirac import fermi_dirac_occupations
    from pyscf.nao.m_fireball_get_eigen_dat import fireball_get_eigen_dat
    from pyscf.nao.m_fireball_hsx import fireball_hsx
    self.telec = kw['telec'] if 'telec' in kw else self.telec
    self.fermi_energy = kw['fermi_energy'] if 'fermi_energy' in kw else self.fermi_energy
    self.mo_energy = require(fireball_get_eigen_dat(self.cd), dtype=self.dtype, requirements='CW')
    ksn2fd = fermi_dirac_occupations(self.telec, self.mo_energy, self.fermi_energy)
    self.mo_occ = (3-self.nspin)*ksn2fd
    if abs(self.nelectron-self.mo_occ.sum())>1e-6: raise RuntimeError("mo_occ wrong?" )
    #print(__name__, ' self.nspecies ', self.nspecies)
    #print(self.sp_mu2j)
    
    self.hsx = fireball_hsx(self, **kw)
    #print(self.telec)
    #print(self.mo_energy)
    #print(self.fermi_energy)
    #print(__name__, ' sum(self.mo_occ)', sum(self.mo_occ))
    #print(self.mo_occ) 

def get_ac_vertex_array(self, dtype=np.float64):
    """ Returns the product vertex coefficients as 3d array (dense table) """
    atom2so = self.sv.atom2s
    nfap = self.c2s[-1]
    n = self.sv.atom2s[-1]
    pab2v = np.require( zeros((nfap,n,n), dtype=dtype), requirements='CW')
    for atom,[sd,fd,pt,spp] in enumerate(zip(self.dpc2s,self.dpc2s[1:],self.dpc2t,self.dpc2sp)):
      if pt!=1: continue
      s,f = atom2so[atom:atom+2]
      pab2v[sd:fd,s:f,s:f] = self.prod_log.sp2vertex[spp]

    for sd,fd,pt,spp in zip(self.dpc2s,self.dpc2s[1:],self.dpc2t,self.dpc2sp):
      if pt!=2: continue
      inf= self.bp2info[spp]
      lab = einsum('dl,dab->lab', inf.cc, inf.vrtx)
      a,b = inf.atoms
      sa,fa,sb,fb = atom2so[a],atom2so[a+1],atom2so[b],atom2so[b+1]
      for c,ls,lf in zip(inf.cc2a, inf.cc2s, inf.cc2s[1:]):
        pab2v[self.c2s[c]:self.c2s[c+1],sa:fa,sb:fb] = lab[ls:lf,:,:]
        pab2v[self.c2s[c]:self.c2s[c+1],sb:fb,sa:fa] = einsum('pab->pba', lab[ls:lf,:,:])
    return pab2v 

def get_dp_vertex_array(self, dtype=np.float64):
    """ Returns the product vertex coefficients as 3d array for dominant products """
    atom2so = self.sv.atom2s
    nfdp = self.dpc2s[-1]
    n = self.sv.atom2s[-1]
    pab2v = np.require(zeros((nfdp,n,n), dtype=dtype), requirements='CW')
    for atom,[sd,fd,pt,spp] in enumerate(zip(self.dpc2s,self.dpc2s[1:],self.dpc2t,self.dpc2sp)):
      if pt!=1: continue
      s,f = atom2so[atom:atom+2]
      pab2v[sd:fd,s:f,s:f] = self.prod_log.sp2vertex[spp]

    for sd,fd,pt,spp in zip(self.dpc2s,self.dpc2s[1:],self.dpc2t,self.dpc2sp):
      if pt!=2: continue
      inf= self.bp2info[spp]
      a,b = inf.atoms
      sa,fa,sb,fb = atom2so[a],atom2so[a+1],atom2so[b],atom2so[b+1]
      pab2v[sd:fd,sa:fa,sb:fb] = inf.vrtx
      pab2v[sd:fd,sb:fb,sa:fa] = einsum('pab->pba', inf.vrtx)
    return pab2v 

def aos_libnao(coords, norbs):
  assert len(coords.shape) == 2
  assert coords.shape[1] == 3
  assert norbs>0

  ncoords = coords.shape[0]
  co2val = np.require( np.zeros((ncoords,norbs)), dtype=c_double, requirements='CW')
  crd_copy = np.require(coords, dtype=c_double, requirements='C')

  libnao.aos_libnao(
    c_int64(ncoords),
    crd_copy.ctypes.data_as(POINTER(c_double)),
    c_int64(norbs),
    co2val.ctypes.data_as(POINTER(c_double)),
    c_int64(norbs))

  return co2val 

def test_sumidxtab_core(self):
        nsq = 8
        ksub = 256
        cur_num = 10

        for iround in range(10):
            raw_D = np.random.random((nsq, ksub))
            raw_blk = np.random.random_integers(0, ksub-1, (cur_num, nsq))
            D = np.require(raw_D, np.float32, "C")
            blk = np.require(raw_blk, np.uint8, "C")

            self.assertLessEqual(np.abs(raw_D - D).sum(),  1e-4)
            self.assertEqual(np.abs(raw_blk - blk).sum(),  0)

            py_res = self.sumidxtab_core(D, blk)
            c_res = cext.sumidxtab_core(D, blk)
            self.assertLessEqual(np.abs(py_res - c_res).sum(),  1e-4) 

def assign(self, coords, mask=None, output=None):
        try:
            passed_coord_dtype = coords.dtype
        except AttributeError:
            coords = numpy.require(coords, dtype=coord_dtype)
        else:
            if passed_coord_dtype != coord_dtype:
                coords = numpy.require(coords, dtype=coord_dtype)

        if coords.ndim != 2:
            raise TypeError('coords must be 2-dimensional')

        if mask is None:
            mask = numpy.ones((len(coords),), dtype=numpy.bool_)
        elif len(mask) != len(coords):
            raise TypeError('mask [shape {}] has different length than coords [shape {}]'.format(mask.shape, coords.shape))

        if output is None:
            output = numpy.empty((len(coords),), dtype=index_dtype)
        elif len(output) != len(coords):
            raise TypeError('output has different length than coords')

        rectilinear_assign(coords, mask, output, self.boundaries, self._boundlens)

        return output 

def assign(self, coords, mask=None, output=None):
        if output is None:
            output = numpy.zeros((len(coords),), dtype=index_dtype)

        if mask is None:
            mask = numpy.ones((len(coords),), dtype=numpy.bool_)
        else:
            mask = numpy.require(mask, dtype=numpy.bool_)

        coord_subset = coords[mask]
        fnvals = numpy.empty((len(coord_subset), len(self.functions)), dtype=index_dtype)
        for ifn, fn in enumerate(self.functions):
            rsl = numpy.apply_along_axis(fn,0,coord_subset)
            if rsl.ndim > 1:
                # this should work like a squeeze, unless the function returned something truly
                # stupid (e.g., a 3d array with at least two dimensions greater than 1), in which
                # case a broadcast error will occur
                fnvals[:,ifn] = rsl.flat
            else:
                fnvals[:,ifn] = rsl
        amask = numpy.require(fnvals.argmax(axis=1), dtype=index_dtype)
        output[mask] = amask
        return output 

def assign(self, coords, mask=None, output=None):
        try:
            passed_coord_dtype = coords.dtype
        except AttributeError:
            coords = numpy.require(coords, dtype=coord_dtype)
        else:
            if passed_coord_dtype != coord_dtype:
                coords = numpy.require(coords, dtype=coord_dtype)

        if coords.ndim != 2:
            raise TypeError('coords must be 2-dimensional')
        if mask is None:
            mask = numpy.ones((len(coords),), dtype=numpy.bool_)
        elif len(mask) != len(coords):
            raise TypeError('mask [shape {}] has different length than coords [shape {}]'.format(mask.shape, coords.shape))

        if output is None:
            output = numpy.empty((len(coords),), dtype=index_dtype)
        elif len(output) != len(coords):
            raise TypeError('output has different length than coords')

        self.func(coords, mask, output, *self.args, **self.kwargs)

        return output 

def assign(self, coords, mask=None, output=None):
        try:
            passed_coord_dtype = coords.dtype
        except AttributeError:
            coords = numpy.require(coords, dtype=coord_dtype)
        else:
            if passed_coord_dtype != coord_dtype:
                coords = numpy.require(coords, dtype=coord_dtype)

        if coords.ndim != 2:
            raise TypeError('coords must be 2-dimensional')
        if mask is None:
            mask = numpy.ones((len(coords),), dtype=numpy.bool_)
        elif len(mask) != len(coords):
            raise TypeError('mask [shape {}] has different length than coords [shape {}]'.format(mask.shape, coords.shape))

        if output is None:
            output = numpy.empty((len(coords),), dtype=index_dtype)
        elif len(output) != len(coords):
            raise TypeError('output has different length than coords')

        apply_down(self.func, self.args, self.kwargs, coords, mask, output)

        return output 

def gw_applykernel_nspin1(self,dn):
    daux  = np.zeros(self.nprod, dtype=self.dtype)
    daux[:] = np.require(dn.real, dtype=self.dtype, requirements=["A","O"])
    vcre = self.spmv(self.nprod, 1.0, self.kernel, daux)
    
    daux[:] = np.require(dn.imag, dtype=self.dtype, requirements=["A","O"])
    vcim = self.spmv(self.nprod, 1.0, self.kernel, daux)
    return vcre,vcim 

def init_dm_libnao(sab2dm):
  from pyscf.nao.m_libnao import libnao
  from pyscf.nao.m_sv_chain_data import sv_chain_data
  from ctypes import POINTER, c_double, c_int64
  d = np.require(sab2dm, dtype=c_double, requirements='C')
  libnao.init_sab2dm_libnao.argtypes = (POINTER(c_double), POINTER(c_int64), POINTER(c_int64))
  libnao.init_sab2dm_libnao(d.ctypes.data_as(POINTER(c_double)), c_int64(d.shape[1]), c_int64(d.shape[0]) )
  return True 

def get_aoneo(self):
    """Packs the data into one array for a later transfer to the library """
    import numpy as np
    from numpy import require, float64, concatenate as conc
    nr  = self.nr
    nsp = self.nspecies
    nmt = sum(self.sp2nmult)    
    nrt = nr*nmt
    nms = nmt+nsp

    nsvn = 200 + 2*nr + 4*nsp + 2*nmt + nrt + nms
    svn = require(np.zeros(nsvn), dtype=float64, requirements='CW')
    # Simple parameters
    i = 0
    svn[i] = nsp;        i+=1;
    svn[i] = nr;         i+=1;
    svn[i] = self.rmin;  i+=1;
    svn[i] = self.rmax;  i+=1;
    svn[i] = self.kmax;  i+=1;
    svn[i] = self.jmx;   i+=1;
    svn[i] = conc(self.psi_log).sum(); i+=1;
    # Pointers to data
    i = 99
    s = 199
    svn[i] = s+1; i+=1; f=s+nr;  svn[s:f] = self.rr; s=f; # pointer to rr
    svn[i] = s+1; i+=1; f=s+nr;  svn[s:f] = self.pp; s=f; # pointer to pp
    svn[i] = s+1; i+=1; f=s+nsp; svn[s:f] = self.sp2nmult; s=f; # pointer to sp2nmult
    svn[i] = s+1; i+=1; f=s+nsp; svn[s:f] = self.sp2rcut;  s=f; # pointer to sp2rcut
    svn[i] = s+1; i+=1; f=s+nsp; svn[s:f] = self.sp2norbs; s=f; # pointer to sp2norbs
    svn[i] = s+1; i+=1; f=s+nsp; svn[s:f] = self.sp2charge; s=f; # pointer to sp2charge    
    svn[i] = s+1; i+=1; f=s+nmt; svn[s:f] = conc(self.sp_mu2j); s=f; # pointer to sp_mu2j
    svn[i] = s+1; i+=1; f=s+nmt; svn[s:f] = conc(self.sp_mu2rcut); s=f; # pointer to sp_mu2rcut
    svn[i] = s+1; i+=1; f=s+nrt; svn[s:f] = conc(self.psi_log).reshape(nrt); s=f; # pointer to psi_log
    svn[i] = s+1; i+=1; f=s+nms; svn[s:f] = conc(self.sp_mu2s); s=f; # pointer to sp_mu2s
    svn[i] = s+1; # this is a terminator to simple operation
    return svn

#
#
# 

def csc_matvec(csc, x):
    """
        Matrix vector multiplication
        using csc format
    """

    if not sparse.isspmatrix_csc(csc):
        raise Exception("Matrix must be in csc format")

    nrow, ncol = csc.shape
    nnz = csc.data.shape[0]
    if x.size != ncol:
      print(x.size, ncol)
      raise ValueError("wrong dimension!")

    xx = np.require(x, requirements="C")

    if csc.dtype == np.float32:
        y = np.zeros((nrow), dtype=np.float32)
        libsparsetools.scsc_matvec(c_int(nrow), c_int(ncol), c_int(nnz), 
                csc.indptr.ctypes.data_as(POINTER(c_int)),
                csc.indices.ctypes.data_as(POINTER(c_int)), 
                csc.data.ctypes.data_as(POINTER(c_float)),
                xx.ctypes.data_as(POINTER(c_float)), 
                y.ctypes.data_as(POINTER(c_float)))

    elif csc.dtype == np.float64:
        y = np.zeros((nrow), dtype=np.float64)
        libsparsetools.dcsc_matvec(c_int(nrow), c_int(ncol), c_int(nnz), 
                csc.indptr.ctypes.data_as(POINTER(c_int)),
                csc.indices.ctypes.data_as(POINTER(c_int)), 
                csc.data.ctypes.data_as(POINTER(c_double)),
                xx.ctypes.data_as(POINTER(c_double)), 
                y.ctypes.data_as(POINTER(c_double)))
    else:
        raise ValueError("Not implemented")

    return y 

def csc_matvecs(csc, B, transB = False, order="C"):
    """
        Matrix matrix multiplication
        using csc format
    """

    if not sparse.isspmatrix_csc(csc):
        raise Exception("Matrix must be in csc format")

    if transB:
        # Here need to be careful, since using the transpose of B
        # will change from row major to col major and vice-versa
        mat = np.require(B.T, dtype=B.dtype, requirements=["A", "O", order])
    else:
        mat = np.require(B, dtype=B.dtype, requirements=["A", "O", order])

    nrow, ncol = csc.shape
    nvecs = mat.shape[1]

    if csc.dtype == np.float32:
        C = np.zeros((nrow, nvecs), dtype=np.float32, order=order)
        libsparsetools.scsc_matvecs(c_int(nrow), c_int(ncol), c_int(nvecs), 
                csc.indptr.ctypes.data_as(POINTER(c_int)),
                csc.indices.ctypes.data_as(POINTER(c_int)), 
                csc.data.ctypes.data_as(POINTER(c_float)),
                mat.ctypes.data_as(POINTER(c_float)), 
                C.ctypes.data_as(POINTER(c_float)))

    elif csc.dtype == np.float64:
        C = np.zeros((nrow, nvecs), dtype=np.float64, order=order)
        libsparsetools.dcsc_matvecs(c_int(nrow), c_int(ncol), c_int(nvecs), 
                csc.indptr.ctypes.data_as(POINTER(c_int)),
                csc.indices.ctypes.data_as(POINTER(c_int)), 
                csc.data.ctypes.data_as(POINTER(c_double)),
                mat.ctypes.data_as(POINTER(c_double)), 
                C.ctypes.data_as(POINTER(c_double)))
    else:
        raise ValueError("Not implemented")

    return C 

def kchi0_mv(self, v, comega=None):
    """ Operator O = K chi0 """
    self.matvec_ncalls+=1
    w = self.comega_current if comega is None else comega
    chi0 = self.apply_rf0(v, w)
    
    chi0_reim = require(chi0.real, dtype=self.dtype, requirements=["A", "O"])
    matvec_real = self.spmv(self.nprod, 1.0, self.kernel, chi0_reim)
    
    chi0_reim = require(chi0.imag, dtype=self.dtype, requirements=["A", "O"])
    matvec_imag = self.spmv(self.nprod, 1.0, self.kernel, chi0_reim)

    return (matvec_real + 1.0j*matvec_imag) 

def solve_umkckc(self, vext, comega=1j*0.0, x0=None):
    """ This solves a system of linear equations 
           (1 - K chi0 K chi0 ) X = vext 
     or computes 
           X = (1 - K chi0 K chi0 )^{-1} vext 
    """
    from scipy.sparse.linalg import LinearOperator, lgmres
    assert len(vext)==len(self.moms0), "%r, %r "%(len(vext), len(self.moms0))
    self.comega_current = comega
    veff2_op = LinearOperator((self.nprod,self.nprod), matvec=self.umkckc_mv, dtype=self.dtypeComplex)

    if self.res_method == "absolute":
        tol = 0.0
        atol = self.tddft_iter_tol
    elif self.res_method == "relative":
        tol = self.tddft_iter_tol
        atol = 0.0
    elif self.res_method == "both":
        tol = self.tddft_iter_tol
        atol = self.tddft_iter_tol
    else:
        raise ValueError("Unknow res_method")

    resgm,info = lgmres(veff2_op, np.require(vext, dtype=self.dtypeComplex,
                                             requirements='C'), x0=x0, 
                        tol=tol, atol=atol, maxiter=self.maxiter)
    
    if info != 0:  print("LGMRES Warning: info = {0}".format(info))

    return resgm 

def init_dm_libnao(dm):
  from pyscf.nao.m_libnao import libnao
  from ctypes import POINTER, c_double, c_int64, byref
  
  d = np.require(dm, dtype=c_double, requirements='C')

  libnao.init_dm_libnao.argtypes = (POINTER(c_double), 
    POINTER(c_int64), # nreim 
    POINTER(c_int64), # norbs
    POINTER(c_int64), # nspin
    POINTER(c_int64), # nkpoints
    POINTER(c_int64)) # alloc_stat
  
  alloc_stat = c_int64(-999)
  
  libnao.init_dm_libnao(d.ctypes.data_as(POINTER(c_double)),
    c_int64(d.shape[-1]), 
    c_int64(d.shape[-2]),
    c_int64(d.shape[-4]),
    c_int64(d.shape[-5]),
    byref(alloc_stat))

  if alloc_stat.value!=0 : 
    raise RuntimeError('could not allocate?') 
    return None

  return dm 

def comp_apair_pp_libint(self, a1,a2):
    """ Get's the vertex coefficient and conversion coefficients for a pair of atoms given by their atom indices """
    from operator import mul
    from pyscf.nao.m_prod_biloc import prod_biloc_c
    if not hasattr(self, 'sv_pbloc_data') : raise RuntimeError('.sv_pbloc_data is absent')
    assert a1>=0
    assert a2>=0
    
    t1 = timer()
    sv = self.sv
    aos = self.sv.ao_log
    sp12 = np.require( np.array([sv.atom2sp[a] for a in (a1,a2)], dtype=c_int64), requirements='C')
    rc12 = np.require( np.array([sv.atom2coord[a,:] for a in (a1,a2)]), requirements='C')
    icc2a = np.require( np.array(self.ls_contributing(a1,a2), dtype=c_int64), requirements='C')
    npmx = aos.sp2norbs[sv.atom2sp[a1]]*aos.sp2norbs[sv.atom2sp[a2]]
    npac = sum([self.prod_log.sp2norbs[sv.atom2sp[ia]] for ia in icc2a ])
    nout = c_int64(npmx**2+npmx*npac+10)
    dout = np.require( zeros(nout.value), requirements='CW')
    
    libnao.vrtx_cc_apair( sp12.ctypes.data_as(POINTER(c_int64)), rc12.ctypes.data_as(POINTER(c_double)), icc2a.ctypes.data_as(POINTER(c_int64)), c_int64(len(icc2a)), dout.ctypes.data_as(POINTER(c_double)), nout )    
    if dout[0]<1: return None
    
    nnn = np.array(dout[0:3], dtype=int)
    nnc = np.array([dout[8],dout[7]], dtype=int)
    ncc = int(dout[9])
    if ncc!=len(icc2a): raise RuntimeError('ncc!=len(icc2a)')
    s = 10; f=s+np.prod(nnn); vrtx  = dout[s:f].reshape(nnn)
    s = f;  f=s+np.prod(nnc); ccoe  = dout[s:f].reshape(nnc)
    icc2s = np.zeros(len(icc2a)+1, dtype=np.int64)
    for icc,a in enumerate(icc2a): icc2s[icc+1] = icc2s[icc] + self.prod_log.sp2norbs[sv.atom2sp[a]]
    pbiloc = prod_biloc_c(atoms=array([a2,a1]),vrtx=vrtx,cc2a=icc2a,cc2s=icc2s,cc=ccoe)
    
    return pbiloc 

def comp_moments(self, dtype=np.float64):
    """ Computes the scalar and dipole moments for the all functions in the product basis """
    sp2mom0, sp2mom1 = self.prod_log.comp_moments()
    n = self.c2s[-1]
    mom0 = np.require(np.zeros(n, dtype=dtype), requirements='CW')
    mom1 = np.require(np.zeros((n,3), dtype=dtype), requirements='CW')
    for a,[sp,coord,s,f] in enumerate(zip(self.sv.atom2sp,self.sv.atom2coord,self.c2s,self.c2s[1:])):
      mom0[s:f],mom1[s:f,:] = sp2mom0[sp], einsum('j,k->jk', sp2mom0[sp],coord)+sp2mom1[sp]
    return mom0,mom1 

def apply_kernel4p_nspin1(self, ddm):
    reim = require(ddm.real, dtype=self.dtype, requirements=["A","O"]) # real part
    mv_real = dot(self.ss2kernel_4p[0][0], reim)
    
    reim = require(ddm.imag, dtype=self.dtype, requirements=["A","O"]) # imaginary
    mv_imag = dot(self.ss2kernel_4p[0][0], reim)
    return mv_real+1j*mv_imag 

def apply_kernel4p_nspin2(self, ddm):
    res = np.zeros((2,self.nspin,self.norbs2), dtype=self.dtype)
    aux = np.zeros((self.norbs2), dtype=self.dtype)
    s2ddm = ddm.reshape((self.nspin,self.norbs2))

    for s in range(self.nspin):
      for t in range(self.nspin):
        for ireim,sreim in enumerate(('real', 'imag')):
          aux[:] = require(getattr(s2ddm[t], sreim), dtype=self.dtype, requirements=["A","O"])
          res[ireim,s] += np.dot(self.ss2kernel_4p[s][t], aux)

    return res[0].reshape(-1)+1j*res[1].reshape(-1) 

def read_wfsx(self, fname, **kw):
    """
    An occasional reading of the SIESTA's .WFSX file
    """
    from pyscf.nao.m_siesta_wfsx import siesta_wfsx_c
    from pyscf.nao.m_siesta2blanko_denvec import _siesta2blanko_denvec
    from pyscf.nao.m_fermi_dirac import fermi_dirac_occupations

    self.wfsx = siesta_wfsx_c(fname=fname, **kw)
    
    assert self.nkpoints == self.wfsx.nkpoints
    assert self.norbs == self.wfsx.norbs 
    assert self.nspin == self.wfsx.nspin
    orb2m = self.get_orb2m()
    for k in range(self.nkpoints):
      for s in range(self.nspin):
        for n in range(self.norbs):
          _siesta2blanko_denvec(orb2m, self.wfsx.x[k,s,n,:,:])

    self.mo_coeff = require(self.wfsx.x, dtype=self.dtype, requirements='CW')
    self.mo_energy = require(self.wfsx.ksn2e, dtype=self.dtype, requirements='CW')
    self.telec = kw['telec'] if 'telec' in kw else self.hsx.telec
    self.nelec = kw['nelec'] if 'nelec' in kw else self.hsx.nelec
    self.fermi_energy = kw['fermi_energy'] if 'fermi_energy' in kw else self.fermi_energy
    ksn2fd = fermi_dirac_occupations(self.telec, self.mo_energy, self.fermi_energy)
    self.mo_occ = (3-self.nspin)*ksn2fd
    return self