Python scipy.sparse.linalg.lgmres() 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 si_c_check (self, tol = 1e-5):
    """
    This compares np.solve and LinearOpt-lgmres methods for solving linear equation (1-v\chi_{0}) * W_c = v\chi_{0}v
    """
    import time
    import numpy as np
    ww = 1j*self.ww_ia
    t = time.time()
    si0_1 = self.si_c(ww)      #method 1:  numpy.linalg.solve
    t1 = time.time() - t
    print('numpy: {} sec'.format(t1))
    t2 = time.time()
    si0_2 = self.si_c2(ww)     #method 2:  scipy.sparse.linalg.lgmres
    t3 = time.time() - t2
    print('lgmres: {} sec'.format(t3))
    summ = abs(si0_1 + si0_2).sum()
    diff = abs(si0_1 - si0_2).sum() 
    if diff/summ < tol and diff/si0_1.size < tol:
       print('OK! scipy.lgmres methods and np.linalg.solve have identical results')
    else:
       print('Results (W_c) are NOT similar!')     
    return [[diff/summ] , [np.amax(abs(diff))] ,[tol]]

  #@profile 

def seff(self, sext, comega=1j*0.0):
    """ This computes an effective two point field (scalar non-local potential) given an external two point field.
        L = L0 (1 - K L0)^-1
        We want therefore an effective X_eff for a given X_ext
        X_eff = (1 - K L0)^-1 X_ext   or   we need to solve linear equation
        (1 - K L0) X_eff = X_ext  

        The operator (1 - K L0) is named self.sext2seff_matvec """
    
    from scipy.sparse.linalg import LinearOperator
    from scipy.sparse.linalg import lgmres as gmres_alias
    #from spipy.sparse.linalg import gmres as gmres_alias
    nsnn = self.nspin*self.norbs2
    assert sext.size==nsnn
    
    self.comega_current = comega
    op = LinearOperator((nsnn,nsnn), matvec=self.sext2seff_matvec, dtype=self.dtypeComplex)
    sext_shape = np.require(sext.reshape(nsnn), dtype=self.dtypeComplex, requirements='C')
    resgm,info = gmres_alias(op, sext_shape, tol=self.tddft_iter_tol)
    return (resgm.reshape(-1),info) 

def solve_linear(model):
    logger.info('solving problem with %d DOFs...'%model.DOF)
    K_,f_=model.K_,model.f_
#    M_x = lambda x: sl.spsolve(P, x)
#    M = sl.LinearOperator((n, n), M_x)
    #print(sl.spsolve(K_,f_))
    delta,info=sl.lgmres(K_,f_.toarray())
    model.is_solved=True
    logger.info('Done!')
    model.d_=delta.reshape((model.node_count*6,1))
    model.r_=model.K*model.d_ 

def si_c2(self,ww):
    """
    This computes the correlation part of the screened interaction using LinearOpt and lgmres
    lgmres method is much slower than np.linalg.solve !!
    """
    import numpy as np
    from scipy.sparse.linalg import lgmres
    from scipy.sparse.linalg import LinearOperator
    rf0 = si0 = self.rf0(ww)    
    for iw,w in enumerate(ww):                                
      k_c = np.dot(self.kernel_sq, rf0[iw,:,:])                                         
      b = np.dot(k_c, self.kernel_sq)               
      self.comega_current = w
      k_c_opt = LinearOperator((self.nprod,self.nprod), matvec=self.gw_vext2veffmatvec, dtype=self.dtypeComplex)  
      for m in range(self.nprod): 
         si0[iw,m,:],exitCode = lgmres(k_c_opt, b[m,:], atol=self.gw_iter_tol, maxiter=self.maxiter)   
      if exitCode != 0: print("LGMRES has not achieved convergence: exitCode = {}".format(exitCode))
      #np.allclose(np.dot(k_c, si0), b, atol=1e-05) == True  #Test   
    return si0 

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 comp_veff(self, vext, comega=1j*0.0, x0=None):
    """ This computes an effective field (scalar potential) given the external scalar potential """
    from scipy.sparse.linalg import LinearOperator, lgmres
    nsp = self.nspin*self.nprod

    assert len(vext)==nsp, "{} {}".format(len(vext), nsp)
    self.comega_current = comega
    veff_op = LinearOperator((nsp,nsp), matvec=self.vext2veff_matvec, 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(veff_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 test_scipy_gmres_den(self):
    """ This is a test on gmres method with dense matrix in scipy """
    x_itr,info = linalg.lgmres(A, b)
    derr = abs(x_ref-x_itr).sum()/x_ref.size
    self.assertLess(derr, 1e-6) 

def test_scipy_gmres_linop(self):
    """ This is a test on gmres method with linear operators in scipy """
    linop = linalg.LinearOperator((n,n), matvec=mvop, dtype=np.complex64)
    x_itr,info = linalg.lgmres(linop, b)
    derr = abs(x_ref-x_itr).sum()/x_ref.size
    self.assertLess(derr, 1e-6) 

def WhichLinearSolvers(self):
        return {"direct":["superlu", "umfpack", "mumps", "pardiso"],
                "iterative":["cg", "bicg", "cgstab", "bicgstab", "gmres", "lgmres"],
                "amg":["cg", "bicg", "cgstab", "bicgstab", "gmres", "lgmres"],
                "petsc":["cg", "bicgstab", "gmres"]} 

def gij(m,i=0,delta=0.01,e=0.0):
  """Calculate a single row of the Green function"""
  v0 = np.zeros(m.shape[0])
  v0[i] = 1.
  iden = eye(v0.shape[0]) # identity matrix
  g = iden*(e+1j*delta) - csc_matrix(m) # matrix to invert
#  print(type(g)) ; exit()
  (b,info) = slg.lgmres(g,v0) # solve the equation  
  go = (b*np.conjugate(b)).real
  return go 

def get_snmw2sf_iter(self, optimize="greedy"):
    """ 
    This computes a matrix elements of W_c: <\Psi(r)\Psi(r) | W_c(r,r',\omega) |\Psi(r')\Psi(r')>.
    sf[spin,n,m,w] = X^n V_mu X^m W_mu_nu X^n V_nu X^m,
    where n runs from s...f, m runs from 0...norbs, w runs from 0...nff_ia, spin=0...1 or 2.
    1- XVX is calculated using dominant product in COO format: gw_xvx('dp_coo')
    2- I_nm = W XVX = (1-v\chi_0)^{-1}v\chi_0v
    3- S_nm = XVX W XVX = XVX * I_nm
    """

    from scipy.sparse.linalg import LinearOperator,lgmres
    
    ww = 1j*self.ww_ia
    xvx= self.gw_xvx('blas')
    snm2i = []
    #convert k_c as full matrix into Operator
    k_c_opt = LinearOperator((self.nprod,self.nprod),
                             matvec=self.gw_vext2veffmatvec,
                             dtype=self.dtypeComplex)

    for s in range(self.nspin):
        sf_aux = np.zeros((len(self.nn[s]), self.norbs, self.nprod), dtype=self.dtypeComplex)
        inm = np.zeros((len(self.nn[s]), self.norbs, len(ww)), dtype=self.dtypeComplex)
        
        # w is complex plane
        for iw,w in enumerate(ww):
            self.comega_current = w                            
            #print('k_c_opt',k_c_opt.shape)
            for n in range(len(self.nn[s])):    
                for m in range(self.norbs):
                    # v XVX
                    a = np.dot(self.kernel_sq, xvx[s][n,m,:])
                    # \chi_{0}v XVX by using matrix vector
                    b = self.gw_chi0_mv(a, self.comega_current)
                    # v\chi_{0}v XVX, this should be equals to bxvx in last approach
                    a = np.dot(self.kernel_sq, b)
                    sf_aux[n,m,:],exitCode = lgmres(k_c_opt, a,
                                                     atol=self.gw_iter_tol,
                                                     maxiter=self.maxiter)
                    if exitCode != 0:
                      print("LGMRES has not achieved convergence: exitCode = {}".format(exitCode))
            # I= XVX I_aux
            inm[:,:,iw]=np.einsum('nmp,nmp->nm',xvx[s], sf_aux, optimize=optimize)
        snm2i.append(np.real(inm))

    if (self.write_w==True):
        from pyscf.nao.m_restart import write_rst_h5py
        print(write_rst_h5py(data = snm2i, filename= 'SCREENED_COULOMB.hdf5'))

    return snm2i 

def check_veff(self, optimize="greedy"):
    """
    This checks the equality of effective field (scalar potential) given the external
    scalar potential obtained from lgmres(linearopt, v_ext) and np.solve(dense matrix, vext). 
    """

    from numpy.linalg import solve

    ww = 1j*self.ww_ia
    rf0 = self.rf0(ww)
    #V_{\mu}^{ab}
    v_pab = self.pb.get_ac_vertex_array()
    for s in range(self.nspin):
      v_eff = np.zeros((len(self.nn[s]), self.nprod), dtype=self.dtype)
      v_eff_ref = np.zeros((len(self.nn[s]), self.nprod), dtype=self.dtype)
      # X_{a}^{n}
      xna = self.mo_coeff[0,s,self.nn[s],:,0]
      # X_{b}^{m}
      xmb = self.mo_coeff[0,s,:,:,0]
      # X_{a}^{n}V_{\mu}^{ab}X_{b}^{m}
      xvx = np.einsum('na,pab,mb->nmp', xna, v_pab, xmb, optimize=optimize)
      for iw,w in enumerate(ww):     
          # v\chi_{0} 
          k_c = np.dot(self.kernel_sq, rf0[iw,:,:])
          # v\chi_{0}v 
          b = np.dot(k_c, self.kernel_sq)
          #(1-v\chi_{0})
          k_c = np.eye(self.nprod)-k_c
          
          #v\chi_{0}v * X_{a}^{n}V_{\nu}^{ab}X_{b}^{m}
          bxvx = np.einsum('pq,nmq->nmp', b, xvx, optimize=optimize)
          #V_{ext}=X_{a}^{n}V_{\mu}^{ab}X_{b}^{m} * v\chi_{0}v * X_{a}^{n}V_{\nu}^{ab}X_{b}^{m}
          xvxbxvx = np.einsum ('nmp,nlp->np', xvx, bxvx, optimize=optimize)
          
          for n in range (len(self.nn[s])):
              # compute v_eff in tddft_iter class as referance
              v_eff_ref[n,:] = self.gw_comp_veff(xvxbxvx[n,:])
              # linear eq. for finding V_{eff} --> (1-v\chi_{0})V_{eff}=V_{ext}
              v_eff[n,:]=solve(k_c, xvxbxvx[n,:])

    # compares both V_{eff}
    if np.allclose(v_eff,v_eff_ref,atol=1e-4)== True:
      return v_eff