<commit_before>from __future__ import print_function, division
import numpy as np
from timeit import default_timer as timer
from pyscf.nao.tddft_iter import use_numba, tddft_iter
from pyscf.nao import gw
from scipy.linalg import blas
from pyscf.nao.m_pack2den import pack2den_u, pack2den_l
if use_numba: from pyscf.nao.m_iter_div_eigenenergy_numba import div_eigenenergy_numba

class bse_iter(gw):

  def __init__(self, **kw):
    """ 
      Iterative BSE a la PK, DF, OC JCTC 
      additionally to the fields from tddft_iter_c, we add the dipole matrix elements dab[ixyz][a,b]
      which is constructed as list of numpy arrays 
       $ d_i = \int f^a(r) r_i f^b(r) dr $
    """
    from pyscf.nao.m_fermi_dirac import fermi_dirac_occupations

    xc_code_kw = kw['xc_code'] if 'xc_code' in kw else None
    gw.__init__(self, **kw)
    self.l0_ncalls = 0
    self.dip_ab = [d.toarray() for d in self.dipole_coo()]
    self.norbs2 = self.norbs**2
    kernel_den = pack2den_l(self.kernel)
    n = self.norbs
    v_dab = self.v_dab
    cc_da = self.cc_da
    self.x = self.mo_coeff[0,0,:,:,0]
    self.kernel_4p = (((v_dab.T*(cc_da*kernel_den))*cc_da.T)*v_dab).reshape([n*n,n*n])
    #print(type(self.kernel_4p), self.kernel_4p.shape, 'this is just a reference kernel, must be removed later for sure')

    self.xc_code = self.xc_code if xc_code_kw is None else xc_code_kw 
    xc = self.xc_code.split(',')[0]
    if self.verbosity>0: 
      print(__name__, '     xc_code_mf ', self.xc_code_mf)
      print(__name__, ' xc_code_kernel ', self.xc_code_kernel)
      print(__name__, '    xc_code_scf ', self.xc_code_scf)
      print(__name__, '        xc_code ', self.xc_code)
      
    if xc=='CIS' or xc=='HF' or xc=='GW':
      self.kernel_4p -= 0.5*np.einsum('abcd->bcad', self.kernel_4p.reshape([n,n,n,n])).reshape([n*n,n*n])
    elif xc=='RPA' or xc=='LDA' or xc=='GGA':
      pass
    else :
      print(' ?? xc_code ', self.xc_code, xc)
      raise RuntimeError('??? xc_code ???')

    if xc=='GW':
      w_c_4p = (((v_dab.T*(cc_da*  self.si_c([0.0])[0].real ))*cc_da.T)*v_dab).reshape([n*n,n*n])
      self.kernel_4p -= 0.5*np.einsum('abcd->bcad', w_c_4p.reshape([n,n,n,n])).reshape([n*n,n*n])


    if hasattr(self, 'mo_energy_gw'):
      # Need to redefine a lot...
      self.mo_energy_scf = np.copy(self.mo_energy)
      self.mo_energy = self.mo_energy_gw
      self.mo_coeff_scf = np.copy(self.mo_coeff)
      self.mo_coeff = self.mo_coeff_gw
      self.x = self.mo_coeff[0,0,:,:,0]
      self.ksn2e = np.require(np.zeros((1,self.nspin,self.norbs)), dtype=self.dtype, requirements='CW')
      self.ksn2e[0,0,:] = self.mo_energy
      ksn2fd = fermi_dirac_occupations(self.telec, self.ksn2e, self.fermi_energy)
      if all(ksn2fd[0,0,:]>self.nfermi_tol):
        print(__name__, self.telec, nfermi_tol, ksn2fd[0,0,:])
        raise RuntimeError('telec is too high?')
    
      self.ksn2f = (3-self.nspin)*ksn2fd
      self.nfermi = np.argmax(ksn2fd[0,0,:]<self.nfermi_tol)
      self.vstart = np.argmax(1.0-ksn2fd[0,0,:]>=self.nfermi_tol)
      self.xocc = self.mo_coeff[0,0,0:self.nfermi,:,0]
      self.xvrt = self.mo_coeff[0,0,self.vstart:,:,0]


  def apply_l0(self, sab, comega=1j*0.0):
    """ This applies the non-interacting four point Green's function to a suitable vector (e.g. dipole matrix elements)"""
    assert sab.size==(self.norbs2), "%r,%r"%(sab.size,self.norbs2)

    sab = sab.reshape([self.norbs,self.norbs])
    self.l0_ncalls+=1

    nm2v = np.zeros((self.norbs,self.norbs), self.dtypeComplex)
    nm2v[self.vstart:, :self.nfermi] = blas.cgemm(1.0, np.dot(self.xvrt, sab), self.xocc.T)
    nm2v[:self.nfermi, self.vstart:] = blas.cgemm(1.0, np.dot(self.xocc, sab), self.xvrt.T)
    
    for n,[en,fn] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
      for m,[em,fm] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
        nm2v[n,m] = nm2v[n,m] * (fn-fm) * ( 1.0 / (comega - (em - en)))

    nb2v = blas.cgemm(1.0, nm2v, self.x)
    ab2v = blas.cgemm(1.0, self.x.T, nb2v)
    return ab2v

  def apply_l0_exp(self, sab, comega=1j*0.0):
    """ This applies the non-interacting four point Green's function to a suitable vector (e.g. dipole matrix elements)"""
    assert sab.size==(self.norbs2), "%r,%r"%(sab.size,self.norbs2)

    sab = sab.reshape([self.norbs,self.norbs])<commit_msg>easier to not redefine mo_energy<commit_after>from __future__ import print_function, division
import numpy as np
from timeit import default_timer as timer
from pyscf.nao.tddft_iter import use_numba, tddft_iter
from pyscf.nao import gw
from scipy.linalg import blas
from pyscf.nao.m_pack2den import pack2den_u, pack2den_l
if use_numba: from pyscf.nao.m_iter_div_eigenenergy_numba import div_eigenenergy_numba

class bse_iter(gw):

  def __init__(self, **kw):
    """ 
      Iterative BSE a la PK, DF, OC JCTC 
      additionally to the fields from tddft_iter_c, we add the dipole matrix elements dab[ixyz][a,b]
      which is constructed as list of numpy arrays 
       $ d_i = \int f^a(r) r_i f^b(r) dr $
    """
    from pyscf.nao.m_fermi_dirac import fermi_dirac_occupations

    xc_code_kw = kw['xc_code'] if 'xc_code' in kw else None
    gw.__init__(self, **kw)
    self.l0_ncalls = 0
    self.dip_ab = [d.toarray() for d in self.dipole_coo()]
    self.norbs2 = self.norbs**2
    kernel_den = pack2den_l(self.kernel)
    n = self.norbs
    v_dab = self.v_dab
    cc_da = self.cc_da
    self.x = self.mo_coeff[0,0,:,:,0]
    self.kernel_4p = (((v_dab.T*(cc_da*kernel_den))*cc_da.T)*v_dab).reshape([n*n,n*n])
    #print(type(self.kernel_4p), self.kernel_4p.shape, 'this is just a reference kernel, must be removed later for sure')

    self.xc_code = self.xc_code if xc_code_kw is None else xc_code_kw 
    xc = self.xc_code.split(',')[0]
    if self.verbosity>0: 
      print(__name__, '     xc_code_mf ', self.xc_code_mf)
      print(__name__, ' xc_code_kernel ', self.xc_code_kernel)
      print(__name__, '    xc_code_scf ', self.xc_code_scf)
      print(__name__, '        xc_code ', self.xc_code)
      
    if xc=='CIS' or xc=='HF' or xc=='GW':
      self.kernel_4p -= 0.5*np.einsum('abcd->bcad', self.kernel_4p.reshape([n,n,n,n])).reshape([n*n,n*n])
    elif xc=='RPA' or xc=='LDA' or xc=='GGA':
      pass
    else :
      print(' ?? xc_code ', self.xc_code, xc)
      raise RuntimeError('??? xc_code ???')

    if xc=='GW':
      w_c_4p = (((v_dab.T*(cc_da*  self.si_c([0.0])[0].real ))*cc_da.T)*v_dab).reshape([n*n,n*n])
      self.kernel_4p -= 0.5*np.einsum('abcd->bcad', w_c_4p.reshape([n,n,n,n])).reshape([n*n,n*n])


    if hasattr(self, 'mo_energy_gw') and xc=='GW':  # Need to redefine eigenvectors previously set in the class tddft_iter
      self.ksn2e = np.require(np.zeros((1,self.nspin,self.norbs)), dtype=self.dtype, requirements='CW')
      self.ksn2e[0,0,:] = self.mo_energy_gw

      ksn2fd = fermi_dirac_occupations(self.telec, self.ksn2e, self.fermi_energy)
      if all(ksn2fd[0,0,:]>self.nfermi_tol):
        print(__name__, self.telec, nfermi_tol, ksn2fd[0,0,:])
        raise RuntimeError('telec is too high?')
      self.ksn2f = (3-self.nspin)*ksn2fd

      self.x = self.mo_coeff_gw[0,0,:,:,0]
      self.nfermi = np.argmax(ksn2fd[0,0,:]<self.nfermi_tol)
      self.vstart = np.argmax(1.0-ksn2fd[0,0,:]>=self.nfermi_tol)
      self.xocc = self.mo_coeff_gw[0,0,0:self.nfermi,:,0]
      self.xvrt = self.mo_coeff_gw[0,0,self.vstart:,:,0]


  def apply_l0(self, sab, comega=1j*0.0):
    """ This applies the non-interacting four point Green's function to a suitable vector (e.g. dipole matrix elements)"""
    assert sab.size==(self.norbs2), "%r,%r"%(sab.size,self.norbs2)

    sab = sab.reshape([self.norbs,self.norbs])
    self.l0_ncalls+=1

    nm2v = np.zeros((self.norbs,self.norbs), self.dtypeComplex)
    nm2v[self.vstart:, :self.nfermi] = blas.cgemm(1.0, np.dot(self.xvrt, sab), self.xocc.T)
    nm2v[:self.nfermi, self.vstart:] = blas.cgemm(1.0, np.dot(self.xocc, sab), self.xvrt.T)
    
    for n,[en,fn] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
      for m,[em,fm] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
        nm2v[n,m] = nm2v[n,m] * (fn-fm) * ( 1.0 / (comega - (em - en)))

    nb2v = blas.cgemm(1.0, nm2v, self.x)
    ab2v = blas.cgemm(1.0, self.x.T, nb2v)
    return ab2v

  def apply_l0_exp(self, sab, comega=1j*0.0):
    """ This applies the non-interacting four point Green's function to a suitable vector (e.g. dipole matrix elements)"""
    assert sab.size==(self.norbs2), "%r,%r"%(sab.size,self.norbs2)

    sab = sab.reshape([self.norbs,self.norbs])
{'fraction_non_alphanumeric': 0.16524781341107872, 'fraction_numerical': 0.025189504373177844, 'mean_word_length': 3.559276980329612, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 13, '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': '5356913', 'n_tokens_mistral': 3871, 'n_tokens_neox': 3747, 'n_words': 654}