<commit_before>from __future__ import print_function, division
import numpy as np
from scipy.sparse import coo_matrix
from pyscf.nao.lsofcsr import lsofcsr_c
from numpy import array, einsum, zeros, int64, sqrt
from ctypes import POINTER, c_double, c_int64, byref
from pyscf.nao.m_libnao import libnao
from timeit import default_timer as timer

libnao.init_vrtx_cc_apair.argtypes = (POINTER(c_double), POINTER(c_int64))

libnao.init_vrtx_cc_batch.argtypes = (POINTER(c_double), POINTER(c_int64))

libnao.vrtx_cc_apair.argtypes = (
  POINTER(c_int64),   # sp12(1:2)     ! chemical species indices
  POINTER(c_double),  # rc12(1:3,1:2) ! positions of species
  POINTER(c_int64),   # lscc(ncc)     ! list of contributing centers
  POINTER(c_int64),   # ncc           ! number of contributing centers 
  POINTER(c_double),  # dout(nout)    ! vertex & converting coefficients
  POINTER(c_int64))   # nout          ! size of the buffer for vertex & converting coefficients

libnao.vrtx_cc_batch.argtypes = (
  POINTER(c_int64),   # npairs        ! chemical species indices
  POINTER(c_double),  # p2srncc       ! species indices, positions of species, number of cc, cc
  POINTER(c_int64),   # ncc           ! leading dimension of p2srncc
  POINTER(c_int64))   # p2ndp         ! pair -> number of dominant products in this pair

libnao.get_vrtx_cc_batch.argtypes = (
  POINTER(c_int64),   # ps            ! start pair 
  POINTER(c_int64),   # pf            ! finish pair
  POINTER(c_double),  # data          ! output data buffer
  POINTER(c_int64))   # ndat          ! size of data buffer

#
#
#
class prod_basis():
  '''
  Holder of local and bilocal product functions and vertices.
  Args:
    system_vars, i.e. holder of the geometry, and orbitals description
    tol : tolerance to keep the linear combinations
  Returns:
    For each specie returns a set of radial functions defining a product basis
    These functions are sufficient to represent the products of original atomic orbitals
    via a product vertex coefficients and conversion coefficients.
    
  Examples:
  '''
  def __init__(self, nao=None, **kw):
    """ Variable belonging to the class prod_basis_c:
        From input:
            self.sv: copy of sv (system variable), probably not necessary??
            self.tol_loc: tolerance for local basis
            self.tol_biloc: tolerance for bilocal basis
            self.ac_rcut_ratio: ac rcut ratio??
            self.ac_npc_max: maximal number of participating centers
        Output:
            self.prod_log: Holder of (local) product functions and vertices
            self.hkernel_csr: hartree kernel: local part of Coulomb interaction
            self.c2s: global product Center (atom) -> start in case of atom-centered basis
            self.bp2info: some information including indices of atoms, list of contributing centres, conversion coefficients
            self.dpc2s, self.dpc2t, self.dpc2sp: product Center -> list of the size of the basis set in this center,of center's types,of product species
    """
    
    self.nao = nao
    self.tol_loc = kw['tol_loc'] if 'tol_loc' in kw else 1e-5
    self.tol_biloc = kw['tol_biloc'] if 'tol_biloc' in kw else 1e-6
    self.ac_rcut_ratio = kw['ac_rcut_ratio'] if 'ac_rcut_ratio' in kw else 1.0
    self.ac_npc_max = kw['ac_npc_max'] if 'ac_npc_max' in kw else 8
    self.algorithm = kw['algorithm'].lower() if 'algorithm' in kw else 'pp'
    self.jcutoff = kw['jcutoff'] if 'jcutoff' in kw else 14
    self.metric_type = kw['metric_type'] if 'metric_type' in kw else 2
    self.optimize_centers = kw['optimize_centers'] if 'optimize_centers' in kw else 0
    self.ngl = kw['ngl'] if 'ngl' in kw else 96
    
    if nao is None: return
    
    if self.algorithm=='pp':
      self.init_prod_basis_pp_batch(nao, **kw)
    elif self.algorithm=='fp':
      from m_pb_ae import pb_ae
      pb_ae(self, nao, self.tol_loc, self.tol_biloc, self.ac_rcut_ratio)
    else:
      print( 'self.algorithm', self.algorithm )
      raise RuntimeError('unknown algorithm')
      
    return  

  def init_prod_basis_pp(self, sv, **kvargs):
    """ Talman's procedure should be working well with Pseudo-Potential starting point."""
    from pyscf.nao.m_prod_biloc import prod_biloc_c

    #t1 = timer()    
    self.init_inp_param_prod_log_dp(sv, **kvargs)
    data = self.chain_data()
    libnao.init_vrtx_cc_apair(data.ctypes.data_as(POINTER(c_double)), c_int64(len(data)))
    self.sv_pbloc_data = True
    
    #t2 = timer(); print(t2-t1); t1=timer();
    self.bp2info = [] # going to be some information including indices of atoms, list of contributing centres, conversion coefficients
    for ia1 in range(sv.natoms):
      rc1 = sv.ao_log.sp2rcut[sv.atom2sp[ia1]]
      for ia2 in range(ia1+1,sv.natoms):
        rc2,dist = sv.ao_log.sp2rcut[sv.atom2sp[ia2]], sqrt(((sv.atom2coord[ia1]-sv.atom2coord[ia2])**2).sum())
        if dist>rc1+rc2 : continue
        pbiloc = self.comp_apair_pp_libint(ia1,ia2)
        if pbiloc is not None : self.bp2info.append(pbiloc)
    
    self.dpc2s,self.dpc2t,self.dpc2sp = self.init_c2s_domiprod() # dominant product's counting
    self.npdp = self.dpc2s[-1]
    self.norbs = self.sv.norbs
    return self

  def init_prod_basis_pp_batch(self, nao, **kw):
    """ Talman's procedure should be working well with Pseudo-Potential starting point."""
    from pyscf.nao import prod_log_c
    from pyscf.nao.m_prod_biloc import prod_biloc_c

    sv = nao
    t1 = timer()
    self.norbs = sv.norbs
    self.init_inp_param_prod_log_dp(sv, **kw)
    #t2 = timer(); print(' after init_inp_param_prod_log_dp ', t2-t1); t1=timer()
    data = self.chain_data()
    libnao.init_vrtx_cc_batch(data.ctypes.data_as(POINTER(c_double)), c_int64(len(data)))
    self.sv_pbloc_data = True

    aos = sv.ao_log
    p2srncc,p2npac,p2atoms = [],[],[]
    for a1,[sp1,ra1] in enumerate(zip(sv.atom2sp, sv.atom2coord)):
      rc1 = aos.sp2rcut[sp1]
      for a2,[sp2,ra2] in enumerate(zip(sv.atom2sp[a1+1:], sv.atom2coord[a1+1:])):
        a2+=a1+1
        rc2,dist = aos.sp2rcut[sp2], sqrt(((ra1-ra2)**2).sum())
        if dist>rc1+rc2 : continue
        cc2atom = self.ls_contributing(a1,a2)
        p2atoms.append([a1,a2])
        p2srncc.append([sp1,sp2]+list(ra1)+list(ra2)+[len(cc2atom)]+list(cc2atom))
        p2npac.append( sum([self.prod_log.sp2norbs[sv.atom2sp[ia]] for ia in cc2atom ]))

    #print(np.asarray(p2srncc))
    p2ndp = np.require( zeros(len(p2srncc), dtype=np.int64), requirements='CW')
    p2srncc_cp = np.require(  np.asarray(p2srncc), requirements='C')
    npairs = p2srncc_cp.shape[0]
    self.npairs = npairs
    self.bp2info = [] # going to be indices of atoms, list of contributing centres, conversion coefficients
    if npairs>0 : # Conditional fill of the self.bp2info if there are bilocal pairs (natoms>1)
      ld = p2srncc_cp.shape[1]
      #print('npairs  p2srncc_cp.shape', npairs, p2srncc_cp.shape)
      #t2 = timer(); print('call vrtx_cc_batch ', t2-t1, 'npairs ', npairs); t1=timer()
      libnao.vrtx_cc_batch( c_int64(npairs), p2srncc_cp.ctypes.data_as(POINTER(c_double)), 
        c_int64(ld), p2ndp.ctypes.data_as(POINTER(c_int64)))
      #t2 = timer(); print('after vrtx_cc_batch ', t2-t1); t1=timer()
      nout = 0
      sp2norbs = sv.ao_log.sp2norbs
      for srncc,ndp,npac in zip(p2srncc,p2ndp,p2npac):
        sp1,sp2 = srncc[0],srncc[1]
        nout = nout + ndp*sp2norbs[sp1]*sp2norbs[sp2]+npac*ndp
      
      dout = np.require( zeros(nout), requirements='CW')<commit_msg>messages before calling library in prod_basis.py<commit_after>from __future__ import print_function, division
import numpy as np
from scipy.sparse import coo_matrix
from pyscf.nao.lsofcsr import lsofcsr_c
from numpy import array, einsum, zeros, int64, sqrt
from ctypes import POINTER, c_double, c_int64, byref
from pyscf.nao.m_libnao import libnao
from timeit import default_timer as timer

libnao.init_vrtx_cc_apair.argtypes = (POINTER(c_double), POINTER(c_int64))

libnao.init_vrtx_cc_batch.argtypes = (POINTER(c_double), POINTER(c_int64))

libnao.vrtx_cc_apair.argtypes = (
  POINTER(c_int64),   # sp12(1:2)     ! chemical species indices
  POINTER(c_double),  # rc12(1:3,1:2) ! positions of species
  POINTER(c_int64),   # lscc(ncc)     ! list of contributing centers
  POINTER(c_int64),   # ncc           ! number of contributing centers 
  POINTER(c_double),  # dout(nout)    ! vertex & converting coefficients
  POINTER(c_int64))   # nout          ! size of the buffer for vertex & converting coefficients

libnao.vrtx_cc_batch.argtypes = (
  POINTER(c_int64),   # npairs        ! chemical species indices
  POINTER(c_double),  # p2srncc       ! species indices, positions of species, number of cc, cc
  POINTER(c_int64),   # ncc           ! leading dimension of p2srncc
  POINTER(c_int64))   # p2ndp         ! pair -> number of dominant products in this pair

libnao.get_vrtx_cc_batch.argtypes = (
  POINTER(c_int64),   # ps            ! start pair 
  POINTER(c_int64),   # pf            ! finish pair
  POINTER(c_double),  # data          ! output data buffer
  POINTER(c_int64))   # ndat          ! size of data buffer

#
#
#
class prod_basis():
  '''
  Holder of local and bilocal product functions and vertices.
  Args:
    system_vars, i.e. holder of the geometry, and orbitals description
    tol : tolerance to keep the linear combinations
  Returns:
    For each specie returns a set of radial functions defining a product basis
    These functions are sufficient to represent the products of original atomic orbitals
    via a product vertex coefficients and conversion coefficients.
    
  Examples:
  '''
  def __init__(self, nao=None, **kw):
    """ Variable belonging to the class prod_basis_c:
        From input:
            self.sv: copy of sv (system variable), probably not necessary??
            self.tol_loc: tolerance for local basis
            self.tol_biloc: tolerance for bilocal basis
            self.ac_rcut_ratio: ac rcut ratio??
            self.ac_npc_max: maximal number of participating centers
        Output:
            self.prod_log: Holder of (local) product functions and vertices
            self.hkernel_csr: hartree kernel: local part of Coulomb interaction
            self.c2s: global product Center (atom) -> start in case of atom-centered basis
            self.bp2info: some information including indices of atoms, list of contributing centres, conversion coefficients
            self.dpc2s, self.dpc2t, self.dpc2sp: product Center -> list of the size of the basis set in this center,of center's types,of product species
    """
    
    self.nao = nao
    self.tol_loc = kw['tol_loc'] if 'tol_loc' in kw else 1e-5
    self.tol_biloc = kw['tol_biloc'] if 'tol_biloc' in kw else 1e-6
    self.ac_rcut_ratio = kw['ac_rcut_ratio'] if 'ac_rcut_ratio' in kw else 1.0
    self.ac_npc_max = kw['ac_npc_max'] if 'ac_npc_max' in kw else 8
    self.algorithm = kw['algorithm'].lower() if 'algorithm' in kw else 'pp'
    self.jcutoff = kw['jcutoff'] if 'jcutoff' in kw else 14
    self.metric_type = kw['metric_type'] if 'metric_type' in kw else 2
    self.optimize_centers = kw['optimize_centers'] if 'optimize_centers' in kw else 0
    self.ngl = kw['ngl'] if 'ngl' in kw else 96
    
    if nao is None: return
    
    if self.algorithm=='pp':
      self.init_prod_basis_pp_batch(nao, **kw)
    elif self.algorithm=='fp':
      from m_pb_ae import pb_ae
      pb_ae(self, nao, self.tol_loc, self.tol_biloc, self.ac_rcut_ratio)
    else:
      print( 'self.algorithm', self.algorithm )
      raise RuntimeError('unknown algorithm')
      
    return  

  def init_prod_basis_pp(self, sv, **kvargs):
    """ Talman's procedure should be working well with Pseudo-Potential starting point."""
    from pyscf.nao.m_prod_biloc import prod_biloc_c

    #t1 = timer()    
    self.init_inp_param_prod_log_dp(sv, **kvargs)
    data = self.chain_data()
    libnao.init_vrtx_cc_apair(data.ctypes.data_as(POINTER(c_double)), c_int64(len(data)))
    self.sv_pbloc_data = True
    
    #t2 = timer(); print(t2-t1); t1=timer();
    self.bp2info = [] # going to be some information including indices of atoms, list of contributing centres, conversion coefficients
    for ia1 in range(sv.natoms):
      rc1 = sv.ao_log.sp2rcut[sv.atom2sp[ia1]]
      for ia2 in range(ia1+1,sv.natoms):
        rc2,dist = sv.ao_log.sp2rcut[sv.atom2sp[ia2]], sqrt(((sv.atom2coord[ia1]-sv.atom2coord[ia2])**2).sum())
        if dist>rc1+rc2 : continue
        pbiloc = self.comp_apair_pp_libint(ia1,ia2)
        if pbiloc is not None : self.bp2info.append(pbiloc)
    
    self.dpc2s,self.dpc2t,self.dpc2sp = self.init_c2s_domiprod() # dominant product's counting
    self.npdp = self.dpc2s[-1]
    self.norbs = self.sv.norbs
    return self

  def init_prod_basis_pp_batch(self, nao, **kw):
    """ Talman's procedure should be working well with Pseudo-Potential starting point."""
    from pyscf.nao import prod_log_c
    from pyscf.nao.m_prod_biloc import prod_biloc_c

    sv = nao
    t1 = timer()
    self.norbs = sv.norbs
    self.init_inp_param_prod_log_dp(sv, **kw)
    #t2 = timer(); print(' after init_inp_param_prod_log_dp ', t2-t1); t1=timer()
    data = self.chain_data()
    libnao.init_vrtx_cc_batch(data.ctypes.data_as(POINTER(c_double)), c_int64(len(data)))
    self.sv_pbloc_data = True

    aos = sv.ao_log
    p2srncc,p2npac,p2atoms = [],[],[]
    for a1,[sp1,ra1] in enumerate(zip(sv.atom2sp, sv.atom2coord)):
      rc1 = aos.sp2rcut[sp1]
      for a2,[sp2,ra2] in enumerate(zip(sv.atom2sp[a1+1:], sv.atom2coord[a1+1:])):
        a2+=a1+1
        rc2,dist = aos.sp2rcut[sp2], sqrt(((ra1-ra2)**2).sum())
        if dist>rc1+rc2 : continue
        cc2atom = self.ls_contributing(a1,a2)
        p2atoms.append([a1,a2])
        p2srncc.append([sp1,sp2]+list(ra1)+list(ra2)+[len(cc2atom)]+list(cc2atom))
        p2npac.append( sum([self.prod_log.sp2norbs[sv.atom2sp[ia]] for ia in cc2atom ]))

    #print(np.asarray(p2srncc))
    p2ndp = np.require( zeros(len(p2srncc), dtype=np.int64), requirements='CW')
    p2srncc_cp = np.require(  np.asarray(p2srncc), requirements='C')
    npairs = p2srncc_cp.shape[0]
    self.npairs = npairs
    self.bp2info = [] # going to be indices of atoms, list of contributing centres, conversion coefficients
    if npairs>0 : # Conditional fill of the self.bp2info if there are bilocal pairs (natoms>1)
      ld = p2srncc_cp.shape[1]
      #print('npairs  p2srncc_cp.shape', npairs, p2srncc_cp.shape)
      if nao.verbosity>0:
        t2 = timer(); print('call vrtx_cc_batch ', t2-t1, 'npairs ', npairs); t1=timer()
      libnao.vrtx_cc_batch( c_int64(npairs), p2srncc_cp.ctypes.data_as(POINTER(c_double)), 
        c_int64(ld), p2ndp.ctypes.data_as(POINTER(c_int64)))
      if nao.verbosity>0:
        t2 = timer(); print('after vrtx_cc_batch ', t2-t1); t1=timer()
      nout = 0
      sp2norbs = sv.ao_log.sp2norbs
      for srncc,ndp,npac in zip(p2srncc,p2ndp,p2npac):
        sp1,sp2 = srncc[0],srncc[1]
        nout = nout + ndp*sp2norbs[sp1]*sp2norbs[sp2]+npac*ndp
      
      dout = np.require( zeros(nout), requirements='CW')
{'fraction_non_alphanumeric': 0.1110081419209638, 'fraction_numerical': 0.02621301383464619, 'mean_word_length': 3.337639965546942, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6779448', 'n_tokens_mistral': 5681, 'n_tokens_neox': 5515, 'n_words': 1300}