Python numpy.require方法代碼示例

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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

#
#
# 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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

#
#
# 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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) 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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) 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import require [as 別名]
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