code: from __future__ import print_function, division
  from numpy import array, argmax
  from pyscf.nao import tddft_iter
  
  
  class tddft_iter_x_zip(tddft_iter):
    """ Iterative TDDFT with a high-energy part of the KS eigenvectors compressed """
  
    def __init__(self, **kw):
      from pyscf.nao.m_fermi_dirac import fermi_dirac_occupations
      
      tddft_iter.__init__(self, **kw)
      self.x_zip = kw['x_zip'] if 'x_zip' in kw else False
      self.x_zip_eps = kw['x_zip_eps'] if 'x_zip_eps' in kw else 0.05
      self.x_zip_emax = kw['x_zip_emax'] if 'x_zip_emax' in kw else 0.25
  
      if self.x_zip: # redefine the eigenvectors
        sm2e,sma2x = self.build_x_zip()
        if self.verbosity>0: 
          print(__name__, 'self.mo_energy.shape =', self.mo_energy.shape)
          print(__name__, 'sm2e.shape =', sm2e.shape)
        self.ksn2e = array([sm2e])
        ksn2fd = fermi_dirac_occupations(self.telec, self.ksn2e, self.fermi_energy)
        for s,n2fd in enumerate(ksn2fd[0]):
          if not all(n2fd>self.nfermi_tol): continue
          print(self.telec, s, nfermi_tol, n2fd)
          raise RuntimeError(__name__, 'telec is too high?')
          
        self.ksn2f = (3-self.nspin)*ksn2fd
        self.nfermi = array([argmax(ksn2fd[0,s,:]<self.nfermi_tol) for s in range(self.nspin)], dtype=int)
        self.vstart = array([argmax(1.0-ksn2fd[0,s,:]>=self.nfermi_tol) for s in range(self.nspin)], dtype=int)
        self.xocc = [ma2x[:nfermi,:] for ma2x,nfermi in zip(sma2x,self.nfermi)]
        self.xvrt = [ma2x[vstart:,:] for ma2x,vstart in zip(sma2x,self.vstart)]
  
    def build_x_zip(self):
      """ define compressed eigenvectors """
      from pyscf.nao.m_x_zip import x_zip
      sm2e = []
      sma2x = []
      for n2e,na2x in zip(self.mo_energy[0], self.mo_coeff[0,:,:,:,0]):
        vst, i2w,i2dos, m2e, ma2x = x_zip(n2e, na2x, eps=self.x_zip_eps, emax=self.x_zip_emax)
        sm2e.append(m2e)
        sma2x.append(ma2x)
      sm2e = array(sm2e)    
      return sm2e, sma2x
   # Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  from __future__ import print_function, division
  from pyscf.nao.m_overlap_ni import overlap_ni
  
  def overlap_lil(sv, ao_log=None, funct=overlap_ni,**kvargs):
    """
    Computes the overlap matrix and returns it in List of Lists format (easy to index)
    Args:
    sv : (System Variables), this must have arrays of coordinates and species, etc
    Returns:
    overlap (real-space overlap) for the whole system
    """
    from pyscf.nao.m_ao_matelem import ao_matelem_c
    from scipy.sparse import lil_matrix
    from numpy import array, int64, zeros
  
    aome = ao_matelem_c(sv.ao_log.rr, sv.ao_log.pp)
    me = aome.init_one_set(sv.ao_log) if ao_log is None else aome.init_one_set(ao_log)
    atom2s = zeros((sv.natm+1), dtype=int64)
    for atom,sp in enumerate(sv.atom2sp): atom2s[atom+1]=atom2s[atom]+me.ao1.sp2norbs[sp]
    sp2rcut = array([max(mu2rcut) for mu2rcut in me.ao1.sp_mu2rcut])
  
    n = atom2s[-1]
    lil_mat = lil_matrix((n,n))
  
    for atom1,[sp1,rv1,s1,f1] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
      for atom2,[sp2,rv2,s2,f2] in enumerate(zip(sv.atom2sp,sv.atom2coord,atom2s,atom2s[1:])):
        if (sp2rcut[sp1]+sp2rcut[sp2])**2<=sum((rv1-rv2)**2) : continue
        lil_mat[s1:f1,s2:f2] = funct(me,sp1,rv1,sp2,rv2,**kvargs)
  
    return lil_mat
  
apis: numpy.array numpy.argmax
apis: scipy.sparse.lil_matrix numpy.zeros