text:
#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Dirac Hartree-Fock
'''

import ctypes
import time
from functools import reduce
import numpy
import scipy.linalg
from pyscf import lib
from pyscf.lib import logger
from pyscf.scf import hf
from pyscf.scf import _vhf
import pyscf.scf.chkfile

def kernel(mf, conv_tol=1e-9, conv_tol_grad=None,
           dump_chk=True, dm0=None, callback=None):
    '''the modified SCF kernel for Dirac-Hartree-Fock.  In this kernel, the
    SCF is carried out in three steps.  First the 2-electron part is
    approximated by large component integrals (LL|LL); Next, (SS|LL) the
    interaction between large and small components are added; Finally,
    converge the SCF with the small component contributions (SS|SS)
    '''
    if conv_tol_grad is None:
        conv_tol_grad = numpy.sqrt(conv_tol)
        logger.info(mf, 'Set gradient conv threshold to %g', conv_tol_grad)
    if dm0 is None:
        dm = mf.get_init_guess()
    else:
        dm = dm0

    mf._coulomb_now = 'LLLL'
    if dm0 is None and mf._coulomb_now.upper() == 'LLLL':
        scf_conv, e_tot, mo_energy, mo_coeff, mo_occ \
                = hf.kernel(mf, 1e-2, 1e-1,
                            dump_chk, dm0=dm, callback=callback)
        dm = mf.make_rdm1(mo_coeff, mo_occ)
        mf._coulomb_now = 'SSLL'

    if dm0 is None and (mf._coulomb_now.upper() == 'SSLL' or
                        mf._coulomb_now.upper() == 'LLSS'):
        scf_conv, e_tot, mo_energy, mo_coeff, mo_occ \
                = hf.kernel(mf, 1e-3, 1e-1,
                            dump_chk, dm0=dm, callback=callback)
        dm = mf.make_rdm1(mo_coeff, mo_occ)
        mf._coulomb_now = 'SSSS'

    if mf.with_ssss:
        mf._coulomb_now = 'SSSS'
    else:
        mf._coulomb_now = 'SSLL'

    return hf.kernel(mf, conv_tol, conv_tol_grad, dump_chk, dm0=dm,
                     callback=callback)

def get_jk_coulomb(mol, dm, hermi=1, coulomb_allow='SSSS',
                   opt_llll=None, opt_ssll=None, opt_ssss=None):
    if coulomb_allow.upper() == 'LLLL':
        logger.info(mol, 'Coulomb integral: (LL|LL)')
        j1, k1 = _call_veff_llll(mol, dm, hermi, opt_llll)
        n2c = j1.shape[1]
        vj = numpy.zeros_like(dm)
        vk = numpy.zeros_like(dm)
        vj[...,:n2c,:n2c] = j1
        vk[...,:n2c,:n2c] = k1
    elif coulomb_allow.upper() == 'SSLL' \
      or coulomb_allow.upper() == 'LLSS':
        logger.info(mol, 'Coulomb integral: (LL|LL) + (SS|LL)')
        vj, vk = _call_veff_ssll(mol, dm, hermi, opt_ssll)
        j1, k1 = _call_veff_llll(mol, dm, hermi, opt_llll)
        n2c = j1.shape[1]
        vj[...,:n2c,:n2c] += j1
        vk[...,:n2c,:n2c] += k1
    else: # coulomb_allow == 'SSSS'
        logger.info(mol, 'Coulomb integral: (LL|LL) + (SS|LL) + (SS|SS)')
        vj, vk = _call_veff_ssll(mol, dm, hermi, opt_ssll)
        j1, k1 = _call_veff_llll(mol, dm, hermi, opt_llll)
        n2c = j1.shape[1]
        vj[...,:n2c,:n2c] += j1
        vk[...,:n2c,:n2c] += k1
        j1, k1 = _call_veff_ssss(mol, dm, hermi, opt_ssss)
        vj[...,n2c:,n2c:] += j1
        vk[...,n2c:,n2c:] += k1
    return vj, vk

def get_jk(mol, dm, hermi=1, coulomb_allow='SSSS'):
    return get_jk_coulomb(mol, dm, hermi=hermi, coulomb_allow=coulomb_allow)

def get_hcore(mol):
    n2c = mol.nao_2c()
    n4c = n2c * 2
    c = lib.param.LIGHT_SPEED

    t  = mol.intor_symmetric('cint1e_spsp') * .5
    vn = mol.intor_symmetric('cint1e_nuc')
    wn = mol.intor_symmetric('cint1e_spnucsp')
    h1e = numpy.empty((n4c, n4c), numpy.complex)
    h1e[:n2c,:n2c] = vn
    h1e[n2c:,:n2c] = t
    h1e[:n2c,n2c:] = t
    h1e[n2c:,n2c:] = wn * (.25/c**2) - t
    return h1e

def get_ovlp(mol):
    n2c = mol.nao_2c()
    n4c = n2c * 2
    c = lib.param.LIGHT_SPEED

    s = mol.intor_symmetric('cint1e_ovlp')
    t = mol.intor_symmetric('cint1e_spsp')
    s1e = numpy.zeros((n4c, n4c), numpy.complex)
    s1e[:n2c,:n2c] = s
    s1e[n2c:,n2c:] = t * (.5/c)**2
    return s1e

def make_rdm1(mo_coeff, mo_occ):
    return numpy.dot(mo_coeff*mo_occ, mo_coeff.T.conj())

def init_guess_by_minao(mol):
    '''Initial guess in terms of the overlap to minimal basis.'''
    dm = hf.init_guess_by_minao(mol)
    return _proj_dmll(mol, dm, mol)

def init_guess_by_1e(mol):
    '''Initial guess from one electron system.'''
    mf = UHF(mol)
    return mf.init_guess_by_1e(mol)

def init_guess_by_atom(mol):
    '''Initial guess from atom calculation.'''
    dm = hf.init_guess_by_atom(mol)
    return _proj_dmll(mol, dm, mol)

def init_guess_by_chkfile(mol, chkfile_name, project=True):
    from pyscf.scf import addons
    chk_mol, scf_rec = pyscf.scf.chkfile.load_scf(chkfile_name)

    if numpy.iscomplexobj(scf_rec['mo_coeff']):
        mo = scf_rec['mo_coeff']
        mo_occ = scf_rec['mo_occ']
        if project:
            dm = make_rdm1(addons.project_mo_r2r(chk_mol, mo, mol), mo_occ)
        else:
            dm = make_rdm1(mo, mo_occ)
    else:
        if scf_rec['mo_coeff'].ndim == 2: # nr-RHF
            mo = scf_rec['mo_coeff']
            mo_occ = scf_rec['mo_occ']
            dm = reduce(numpy.dot, (mo*mo_occ, mo.T))
        else: # nr-UHF
            mo = scf_rec['mo_coeff']
            mo_occ = scf_rec['mo_occ']
            dm = reduce(numpy.dot, (mo[0]*mo_occ[0], mo[0].T)) \
               + reduce(numpy.dot, (mo[1]*mo_occ[1], mo[1].T))
        dm = _proj_dmll(chk_mol, dm, mol)
    return dm

def get_init_guess(mol, key='minao'):
    if callable(key):
        return key(mol)
    elif key.lower() == '1e':
        return init_guess_by_1e(mol)
    elif key.lower() == 'atom':
        return init_guess_by_atom(mol)
    elif key.lower() == 'chkfile':
        raise RuntimeError('Call pyscf.scf.hf.init_guess_by_chkfile instead')
    else:
        return init_guess_by_minao(mol)

def time_reversal_matrix(mol, mat):
    ''' T(A_ij) = A[T(i),T(j)]^*
    '''
    n2c = mol.nao_2c()
    tao = numpy.asarray(mol.time_reversal_map())
    # tao(i) = -j  means  T(f_i) = -f_j
    # tao(i) =  j  means  T(f_i) =  f_j
    idx = abs(tao)-1 # -1 for C indexing convention
    #:signL = [(1 if x>0 else -1) for x in tao]
    #:sign = numpy.hstack((signL, signL))

    #:tmat = numpy.empty_like(mat)
    #:for j in range(mat.__len__()):
    #:    for i in range(mat.__len__()):
    #:        tmat[idx[i],idx[j]] = mat[i,j] * sign[i]*sign[j]
    #:return tmat.conjugate()
    sign_mask = tao<0
    if mat.shape[0] == n2c*2:
        idx = numpy.hstack((idx, idx+n2c))
        sign_mask = numpy.hstack((sign_mask, sign_mask))

    tmat = mat.take(idx,axis=0).take(idx,axis=1)
    tmat[sign_mask,:] *= -1
    tmat[:,sign_mask] *= -1
    return tmat.T

def analyze(mf, verbose=logger.DEBUG, **kwargs):
    #from pyscf.tools import dump_mat
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(mf.stdout, verbose)
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ
    #mo_coeff = mf.mo_coeff

    log.info('**** MO energy ****')
    for i in range(len(mo_energy)):
        if mo_occ[i] > 0:
            log.info('occupied MO #%d energy= %.15g occ= %g', \
                     i+1, mo_energy[i], mo_occ[i])
        else:
            log.info('virtual MO #%d energy= %.15g occ= %g', \
                     i+1, mo_energy[i], mo_occ[i])
#TODO    if mf.verbose >= logger.DEBUG:
#TODO        log.debug(' ** MO coefficients **')
#TODO        label = mf.mol.spinor_labels(True)
#TODO        dump_mat.dump_rec(mf.stdout, mo_coeff, label, start=1)
#TODO    dm = mf.make_rdm1(mo_coeff, mo_occ)
#TODO    return mf.mulliken_pop(mf.mol, dm, mf.get_ovlp(), log)

def get_grad(mo_coeff, mo_occ, fock_ao):
    '''DHF Gradients'''
    occidx = numpy.where(mo_occ> 0)[0]
    viridx = numpy.where(mo_occ==0)[0]
    nocc = len(occidx)
    nvir = len(viridx)

    fock = reduce(numpy.dot, (mo_coeff.T.conj(), fock_ao, mo_coeff))
    g = fock[occidx[:,None],viridx].T
    return g.reshape(-1)

class UHF(hf.SCF):
    __doc__ = hf.SCF.__doc__ + '''
    Attributes for Dirac-Hartree-Fock
        with_ssss : bool, for Dirac-Hartree-Fock only
            If False, ignore small component integrals (SS|SS).  Default is True.
        with_gaunt : bool, for Dirac-Hartree-Fock only
            Default is False.
        with_breit : bool, for Dirac-Hartree-Fock only
            Gaunt + gauge term.  Default is False.

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1', basis='ccpvdz', verbose=0)
    >>> mf = scf.RHF(mol)
    >>> e0 = mf.scf()
    >>> mf = scf.DHF(mol)
    >>> e1 = mf.scf()
    >>> print('Relativistic effects = %.12f' % (e1-e0))
    Relativistic effects = -0.000008854205
    '''
    def __init__(self, mol):
        hf.SCF.__init__(self, mol)
        self.conv_tol = 1e-8
        self.with_ssss = True
        self._coulomb_now = 'SSSS' # 'SSSS' ~ LLLL+LLSS+SSSS
        self.with_gaunt = False
        self.with_breit = False

        self.opt = (None, None, None, None) # (opt_llll, opt_ssll, opt_ssss, opt_gaunt)
        self._keys = set(self.__dict__.keys())

    def dump_flags(self):
        hf.SCF.dump_flags(self)
        logger.info(self, 'with_ssss %s, with_gaunt %s, with_breit %s',
                    self.with_ssss, self.with_gaunt, self.with_breit)
        logger.info(self, 'light speed = %s', lib.param.LIGHT_SPEED)
        return self

    def get_hcore(self, mol=None):
        if mol is None:
            mol = self.mol
        return get_hcore(mol)

    def get_ovlp(self, mol=None):
        if mol is None:
            mol = self.mol
        return get_ovlp(mol)

    def get_grad(self, mo_coeff, mo_occ, fock=None):
        if fock is None:
            dm1 = self.make_rdm1(mo_coeff, mo_occ)
            fock = self.get_hcore(self.mol) + self.get_veff(self.mol, dm1)
        return get_grad(mo_coeff, mo_occ, fock)

    def init_guess_by_minao(self, mol=None):
        '''Initial guess in terms of the overlap to minimal basis.'''
        if mol is None: mol = self.mol
        return init_guess_by_minao(mol)

    def init_guess_by_atom(self, mol=None):
        if mol is None: mol = self.mol
        return init_guess_by_atom(mol)

    def init_guess_by_chkfile(self, chkfile=None, project=True):
        if chkfile is None: chkfile = self.chkfile
        return init_guess_by_chkfile(self.mol, chkfile, project=project)

    def build(self, mol=None):
        if self.verbose >= logger.WARN:
            self.check_sanity()
        if self.direct_scf:
            self.opt = self.init_direct_scf(self.mol)

    def get_occ(self, mo_energy=None, mo_coeff=None):
        if mo_energy is None: mo_energy = self.mo_energy
        mol = self.mol
        c = lib.param.LIGHT_SPEED
        n4c = len(mo_energy)
        n2c = n4c // 2
        mo_occ = numpy.zeros(n2c * 2)
        if mo_energy[n2c] > -1.999 * c**2:
            mo_occ[n2c:n2c+mol.nelectron] = 1
        else:
            n = 0
            for i, e in enumerate(mo_energy):
                if e > -1.999 * c**2 and n < mol.nelectron:
                    mo_occ[i] = 1
                    n += 1
        if self.verbose >= logger.INFO:
            logger.info(self, 'HOMO %d = %.12g  LUMO %d = %.12g',
                        n2c+mol.nelectron, mo_energy[n2c+mol.nelectron-1],
                        n2c+mol.nelectron+1, mo_energy[n2c+mol.nelectron])
            logger.debug(self, 'NES  mo_energy = %s', mo_energy[:n2c])
            logger.debug(self, 'PES  mo_energy = %s', mo_energy[n2c:])
        return mo_occ

    # full density matrix for UHF
    def make_rdm1(self, mo_coeff=None, mo_occ=None):
        if mo_coeff is None: mo_coeff = self.mo_coeff
        if mo_occ is None: mo_occ = self.mo_occ
        return make_rdm1(mo_coeff, mo_occ)

    def init_direct_scf(self, mol=None):
        if mol is None: mol = self.mol
        def set_vkscreen(opt, name):
            opt._this.contents.r_vkscreen = _vhf._fpointer(name)
        opt_llll = _vhf.VHFOpt(mol, 'cint2e', 'CVHFrkbllll_prescreen',
                               'CVHFrkbllll_direct_scf',
                               'CVHFrkbllll_direct_scf_dm')
        opt_llll.direct_scf_tol = self.direct_scf_tol
        set_vkscreen(opt_llll, 'CVHFrkbllll_vkscreen')
        opt_ssss = _vhf.VHFOpt(mol, 'cint2e_spsp1spsp2',
                               'CVHFrkbllll_prescreen',
                               'CVHFrkbssss_direct_scf',
                               'CVHFrkbssss_direct_scf_dm')
        opt_ssss.direct_scf_tol = self.direct_scf_tol
        set_vkscreen(opt_ssss, 'CVHFrkbllll_vkscreen')
        opt_ssll = _vhf.VHFOpt(mol, 'cint2e_spsp1',
                               'CVHFrkbssll_prescreen',
                               'CVHFrkbssll_direct_scf',
                               'CVHFrkbssll_direct_scf_dm')
        opt_ssll.direct_scf_tol = self.direct_scf_tol
        set_vkscreen(opt_ssll, 'CVHFrkbssll_vkscreen')
#TODO: prescreen for gaunt
        opt_gaunt = None
        return opt_llll, opt_ssll, opt_ssss, opt_gaunt

    def get_jk(self, mol=None, dm=None, hermi=1):
        if mol is None: mol = self.mol
        if dm is None: dm = self.make_rdm1()
        t0 = (time.clock(), time.time())
        verbose_bak, mol.verbose = mol.verbose, self.verbose
        stdout_bak,  mol.stdout  = mol.stdout , self.stdout
        if self.direct_scf and self.opt[0] is None:
            self.opt = self.init_direct_scf(mol)
        opt_llll, opt_ssll, opt_ssss, opt_gaunt = self.opt

        vj, vk = get_jk_coulomb(mol, dm, hermi, self._coulomb_now,
                                opt_llll, opt_ssll, opt_ssss)

        if self.with_breit:
            if 'SSSS' in self._coulomb_now.upper() or not self.with_ssss:
                vj1, vk1 = _call_veff_gaunt_breit(mol, dm, hermi, opt_gaunt, True)
                logger.info(self, 'Add Breit term')
                vj += vj1
                vk += vk1
        elif self.with_gaunt and 'SS' in self._coulomb_now.upper():
            logger.info(self, 'Add Gaunt term')
            vj1, vk1 = _call_veff_gaunt_breit(mol, dm, hermi, opt_gaunt, False)
            vj += vj1
            vk += vk1

        mol.verbose = verbose_bak
        mol.stdout  = stdout_bak
        logger.timer(self, 'vj and vk', *t0)
        return vj, vk

    def get_veff(self, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
        '''Dirac-Coulomb'''
        if mol is None: mol = self.mol
        if dm is None: dm = self.make_rdm1()
        if self.direct_scf:
            ddm = numpy.array(dm, copy=False) - numpy.array(dm_last, copy=False)
            vj, vk = self.get_jk(mol, ddm, hermi=hermi)
            return numpy.array(vhf_last, copy=False) + vj - vk
        else:
            vj, vk = self.get_jk(mol, dm, hermi=hermi)
            return vj - vk

    def scf(self, dm0=None):
        cput0 = (time.clock(), time.time())

        self.build()
        self.dump_flags()
        self.converged, self.e_tot, \
                self.mo_energy, self.mo_coeff, self.mo_occ \
                = kernel(self, self.conv_tol, self.conv_tol_grad,
                         dm0=dm0, callback=self.callback)

        logger.timer(self, 'SCF', *cput0)
        self._finalize()
        return self.e_tot

    def analyze(self, verbose=None):
        if verbose is None: verbose = self.verbose
        return analyze(self, verbose)

    def x2c(self):
        import pyscf.scf.x2c
        x2chf = pyscf.scf.x2c.UHF(self.mol)
        x2chf.__dict__.update(self.__dict__)
        return x2chf

class HF1e(UHF):
    def scf(self, *args):
        logger.info(self, '\n')
        logger.info(self, '******** 1 electron system ********')
        self.converged = True
        h1e = self.get_hcore(self.mol)
        s1e = self.get_ovlp(self.mol)
        self.mo_energy, self.mo_coeff = self.eig(h1e, s1e)
        self.mo_occ = numpy.zeros_like(self.mo_energy)
        n2c = len(self.mo_occ) // 2
        self.mo_occ[n2c] = 1
        self.e_tot = self.mo_energy[n2c] + self.mol.energy_nuc()
        return self.e_tot

class RHF(UHF):
    '''Dirac-RHF'''
    def __init__(self, mol):
        if mol.nelectron.__mod__(2) != 0:
            raise ValueError('Invalid electron number %i.' % mol.nelectron)
        UHF.__init__(self, mol)

    # full density matrix for RHF
    def make_rdm1(self, mo_coeff=None, mo_occ=None):
        r'''D/2 = \psi_i^\dag\psi_i = \psi_{Ti}^\dag\psi_{Ti}
        D(UHF) = \psi_i^\dag\psi_i + \psi_{Ti}^\dag\psi_{Ti}
        RHF average the density of spin up and spin down:
        D(RHF) = (D(UHF) + T[D(UHF)])/2
        '''
        if mo_coeff is None: mo_coeff = self.mo_coeff
        if mo_occ is None: mo_occ = self.mo_occ
        dm = make_rdm1(mo_coeff, mo_occ)
        return (dm + time_reversal_matrix(self.mol, dm)) * .5

    def get_occ(self, mo_energy=None, mo_coeff=None):
        if mo_energy is None: mo_energy = self.mo_energy
        mol = self.mol
        c = lib.param.LIGHT_SPEED
        n4c = len(mo_energy)
        n2c = n4c // 2
        mo_occ = numpy.zeros(n2c * 2)
        if mo_energy[n2c] > -1.999 * c**2:
            mo_occ[n2c:n2c+mol.nelectron] = 1
        else:
            n = 0
            for i, e in enumerate(mo_energy):
                if e > -1.999 * c**2 and n < mol.nelectron:
                    mo_occ[i] = 1
                    n += 1
        if self.verbose >= logger.INFO:
            logger.info(self, 'HOMO %d = %.12g, LUMO %d = %.12g,',
                        (n2c+mol.nelectron)//2, mo_energy[n2c+mol.nelectron-1],
                        (n2c+mol.nelectron)//2+1, mo_energy[n2c+mol.nelectron])
            logger.debug(self, 'NES  mo_energy = %s', mo_energy[:n2c])
            logger.debug(self, 'PES  mo_energy = %s', mo_energy[n2c:])
        return mo_occ

def _jk_triu_(vj, vk, hermi):
    if hermi == 0:
        if vj.ndim == 2:
            vj = lib.hermi_triu(vj, 1)
        else:
            for i in range(vj.shape[0]):
                vj[i] = lib.hermi_triu(vj[i], 1)
    else:
        if vj.ndim == 2:
            vj = lib.hermi_triu(vj, hermi)
            vk = lib.hermi_triu(vk, hermi)
        else:
            for i in range(vj.shape[0]):
                vj[i] = lib.hermi_triu(vj[i], hermi)
                vk[i] = lib.hermi_triu(vk[i], hermi)
    return vj, vk

def _call_veff_llll(mol, dm, hermi=1, mf_opt=None):
    if isinstance(dm, numpy.ndarray) and dm.ndim == 2:
        n2c = dm.shape[0] // 2
        dms = dm[:n2c,:n2c].copy()
    else:
        n2c = dm[0].shape[0] // 2
        dms = []
        for dmi in dm:
            dms.append(dmi[:n2c,:n2c].copy())
    vj, vk = _vhf.rdirect_mapdm('cint2e', 's8',
                                ('ji->s2kl', 'jk->s1il'), dms, 1,
                                mol._atm, mol._bas, mol._env, mf_opt)
    return _jk_triu_(vj, vk, hermi)

def _call_veff_ssll(mol, dm, hermi=1, mf_opt=None):
    if isinstance(dm, numpy.ndarray) and dm.ndim == 2:
        n_dm = 1
        n2c = dm.shape[0] // 2
        dmll = dm[:n2c,:n2c].copy()
        dmsl = dm[n2c:,:n2c].copy()
        dmss = dm[n2c:,n2c:].copy()
        dms = (dmll, dmss, dmsl)
    else:
        n_dm = len(dm)
        n2c = dm[0].shape[0] // 2
        dms = [dmi[:n2c,:n2c].copy() for dmi in dm] \
            + [dmi[n2c:,n2c:].copy() for dmi in dm] \
            + [dmi[n2c:,:n2c].copy() for dmi in dm]
    jks = ('lk->s2ij',) * n_dm \
        + ('ji->s2kl',) * n_dm \
        + ('jk->s1il',) * n_dm
    c1 = .5 / lib.param.LIGHT_SPEED
    vx = _vhf.rdirect_bindm('cint2e_spsp1', 's4', jks, dms, 1,
                            mol._atm, mol._bas, mol._env, mf_opt) * c1**2
    vj = numpy.zeros((n_dm,n2c*2,n2c*2), dtype=numpy.complex)
    vk = numpy.zeros((n_dm,n2c*2,n2c*2), dtype=numpy.complex)
    vj[:,n2c:,n2c:] = vx[      :n_dm  ,:,:]
    vj[:,:n2c,:n2c] = vx[n_dm  :n_dm*2,:,:]
    vk[:,n2c:,:n2c] = vx[n_dm*2:      ,:,:]
    if n_dm == 1:
        vj = vj.reshape(vj.shape[1:])
        vk = vk.reshape(vk.shape[1:])
    return _jk_triu_(vj, vk, hermi)

def _call_veff_ssss(mol, dm, hermi=1, mf_opt=None):
    c1 = .5 / lib.param.LIGHT_SPEED
    if isinstance(dm, numpy.ndarray) and dm.ndim == 2:
        n2c = dm.shape[0] // 2
        dms = dm[n2c:,n2c:].copy()
    else:
        n2c = dm[0].shape[0] // 2
        dms = []
        for dmi in dm:
            dms.append(dmi[n2c:,n2c:].copy())
    vj, vk = _vhf.rdirect_mapdm('cint2e_spsp1spsp2', 's8',
                                ('ji->s2kl', 'jk->s1il'), dms, 1,
                                mol._atm, mol._bas, mol._env, mf_opt) * c1**4
    return _jk_triu_(vj, vk, hermi)

def _call_veff_gaunt_breit(mol, dm, hermi=1, mf_opt=None, with_breit=False):
    if with_breit:
        intor_prefix = 'cint2e_breit_'
    else:
        intor_prefix = 'cint2e_'
    if isinstance(dm, numpy.ndarray) and dm.ndim == 2:
        n_dm = 1
        n2c = dm.shape[0] // 2
        dmls = dm[:n2c,n2c:].copy()
        dmsl = dm[n2c:,:n2c].copy()
        dmll = dm[:n2c,:n2c].copy()
        dmss = dm[n2c:,n2c:].copy()
        dms = [dmsl, dmsl, dmls, dmll, dmss]
    else:
        n_dm = len(dm)
        n2c = dm[0].shape[0] // 2
        dmll = [dmi[:n2c,:n2c].copy() for dmi in dm]
        dmls = [dmi[:n2c,n2c:].copy() for dmi in dm]
        dmsl = [dmi[n2c:,:n2c].copy() for dmi in dm]
        dmss = [dmi[n2c:,n2c:].copy() for dmi in dm]
        dms = dmsl + dmsl + dmls + dmll + dmss
    vj = numpy.zeros((n_dm,n2c*2,n2c*2), dtype=numpy.complex)
    vk = numpy.zeros((n_dm,n2c*2,n2c*2), dtype=numpy.complex)

    jks = ('lk->s1ij',) * n_dm \
        + ('jk->s1il',) * n_dm
    vx = _vhf.rdirect_bindm(intor_prefix+'ssp1ssp2', 's1', jks, dms[:n_dm*2], 1,
                            mol._atm, mol._bas, mol._env, mf_opt)
    vj[:,:n2c,n2c:] = vx[:n_dm,:,:]
    vk[:,:n2c,n2c:] = vx[n_dm:,:,:]

    jks = ('lk->s1ij',) * n_dm \
        + ('li->s1kj',) * n_dm \
        + ('jk->s1il',) * n_dm
    vx = _vhf.rdirect_bindm(intor_prefix+'ssp1sps2', 's1', jks, dms[n_dm*2:], 1,
                            mol._atm, mol._bas, mol._env, mf_opt)
    vj[:,:n2c,n2c:]+= vx[      :n_dm  ,:,:]
    vk[:,n2c:,n2c:] = vx[n_dm  :n_dm*2,:,:]
    vk[:,:n2c,:n2c] = vx[n_dm*2:      ,:,:]

    if hermi == 1:
        vj[:,n2c:,:n2c] = vj[:,:n2c,n2c:].transpose(0,2,1).conj()
        vk[:,n2c:,:n2c] = vk[:,:n2c,n2c:].transpose(0,2,1).conj()
    elif hermi == 2:
        vj[:,n2c:,:n2c] = -vj[:,:n2c,n2c:].transpose(0,2,1).conj()
        vk[:,n2c:,:n2c] = -vk[:,:n2c,n2c:].transpose(0,2,1).conj()
    else:
        raise NotImplementedError
    if n_dm == 1:
        vj = vj.reshape(n2c*2,n2c*2)
        vk = vk.reshape(n2c*2,n2c*2)
    c1 = .5 / lib.param.LIGHT_SPEED
    if with_breit:
        return vj*c1**2, vk*c1**2
    else:
        return -vj*c1**2, -vk*c1**2

def _proj_dmll(mol_nr, dm_nr, mol):
    from pyscf.scf import addons
    proj = addons.project_mo_nr2r(mol_nr, 1, mol)

    n2c = proj.shape[0]
    n4c = n2c * 2
    dm = numpy.zeros((n4c,n4c), dtype=complex)
    # *.5 because alpha and beta are summed in project_mo_nr2r
    dm_ll = reduce(numpy.dot, (proj, dm_nr*.5, proj.T.conj()))
    dm[:n2c,:n2c] = (dm_ll + time_reversal_matrix(mol, dm_ll)) * .5
    return dm

if __name__ == '__main__':
    import pyscf.gto
    mol = pyscf.gto.Mole()
    mol.verbose = 5
    mol.output = 'out_dhf'

    mol.atom.extend([['He', (0.,0.,0.)], ])
    mol.basis = {
        'He': [(0, 0, (1, 1)),
               (0, 0, (3, 1)),
               (1, 0, (1, 1)), ]}
    mol.build()

##############
# SCF result
    method = UHF(mol)
    energy = method.scf() #-2.38146942868
    print(energy)
    method.with_gaunt = True
    print(method.scf()) # -2.38138339005
    method.with_breit = True
    print(method.scf()) # -2.38138339005


text:
#! /usr/bin/env python3
"""与加载了NBOW+MLP Classifier导出的Servable的TensorFlow Serving进行通信
"""

import numpy as np
import jieba
import tensorlayer as tl
from grpc.beta import implementations

import predict_pb2
import prediction_service_pb2
from packages import text_regularization as tr

def text_tensor(text, wv):
    """获取文本向量
    Args:
        text: 待检测文本
        wv: 词向量模型
    Returns:
        [[ 3.80905056   1.94315064  -0.20703495  -1.31589055   1.9627794
           ...
           2.16935492   2.95426321  -4.71534014  -3.25034237 -11.28901672]]
    """
    text = tr.extractWords(text)
    words = jieba.cut(text.strip())
    text_embedding = np.zeros(200, dtype=np.float32)
    for word in words:
        try:
            text_embedding += wv[word]
        except KeyError:
            text_embedding += wv['UNK']

    sample = text_embedding.reshape(1, 200)
    return sample

print(" ".join(jieba.cut('分词初始化')))
wv = tl.files.load_npy_to_any(name='../word2vec/output/model_word2vec_200.npy')

host, port = ('localhost', '9000')
channel = implementations.insecure_channel(host, int(port))
stub = prediction_service_pb2.beta_create_PredictionService_stub(channel)
request = predict_pb2.PredictRequest()
request.model_spec.name = 'antispam'

ext: py

max_stars_repo_name: cunshen/text-antispam

max_issues_repo_name: cunshen/text-antispam

max_forks_repo_name: cunshen/text-antispam

reason: import numpy


text:
from copy import deepcopy
from typing import List
from typing import Tuple

import numpy as np
import pandas as pd
import pytest
from pandas.testing import assert_frame_equal

from etna.datasets import generate_ar_df
from etna.datasets.tsdataset import TSDataset
from etna.transforms import AddConstTransform
from etna.transforms import FilterFeaturesTransform
from etna.transforms import LagTransform
from etna.transforms import MaxAbsScalerTransform
from etna.transforms import OneHotEncoderTransform
from etna.transforms import SegmentEncoderTransform
from etna.transforms import TimeSeriesImputerTransform

@pytest.fixture()
def tsdf_with_exog(random_seed) -> TSDataset:
    df_1 = pd.DataFrame.from_dict({"timestamp": pd.date_range("2021-02-01", "2021-07-01", freq="1d")})
    df_2 = pd.DataFrame.from_dict({"timestamp": pd.date_range("2021-02-01", "2021-07-01", freq="1d")})
    df_1["segment"] = "Moscow"
    df_1["target"] = [x ** 2 + np.random.uniform(-2, 2) for x in list(range(len(df_1)))]
    df_2["segment"] = "Omsk"
    df_2["target"] = [x ** 0.5 + np.random.uniform(-2, 2) for x in list(range(len(df_2)))]
    classic_df = pd.concat([df_1, df_2], ignore_index=True)

    df = TSDataset.to_dataset(classic_df)

    classic_df_exog = generate_ar_df(start_time="2021-01-01", periods=600, n_segments=2)
    classic_df_exog.rename(columns={"target": "exog"}, inplace=True)
    df_exog = TSDataset.to_dataset(classic_df_exog)

    ts = TSDataset(df=df, df_exog=df_exog, freq="1D")
    return ts

@pytest.fixture()
def df_and_regressors() -> Tuple[pd.DataFrame, pd.DataFrame, List[str]]:
    timestamp = pd.date_range("2021-01-01", "2021-02-01")
    df_1 = pd.DataFrame({"timestamp": timestamp, "target": 11, "segment": "1"})
    df_2 = pd.DataFrame({"timestamp": timestamp[5:], "target": 12, "segment": "2"})
    df = pd.concat([df_1, df_2], ignore_index=True)
    df = TSDataset.to_dataset(df)

    timestamp = pd.date_range("2020-12-01", "2021-02-11")
    df_1 = pd.DataFrame({"timestamp": timestamp, "regressor_1": 1, "regressor_2": 2, "segment": "1"})
    df_2 = pd.DataFrame({"timestamp": timestamp[5:], "regressor_1": 3, "regressor_2": 4, "segment": "2"})
    df_exog = pd.concat([df_1, df_2], ignore_index=True)
    df_exog = TSDataset.to_dataset(df_exog)

    return df, df_exog, ["regressor_1", "regressor_2"]

@pytest.fixture()
def df_and_regressors_flat() -> Tuple[pd.DataFrame, pd.DataFrame]:
    """Return flat versions of df and df_exog."""
    timestamp = pd.date_range("2021-01-01", "2021-02-01")
    df_1 = pd.DataFrame({"timestamp": timestamp, "target": 11, "segment": "1"})
    df_2 = pd.DataFrame({"timestamp": timestamp[5:], "target": 12, "segment": "2"})
    df = pd.concat([df_1, df_2], ignore_index=True)

    timestamp = pd.date_range("2020-12-01", "2021-02-11")
    df_1 = pd.DataFrame(
        {"timestamp": timestamp, "regressor_1": 1, "regressor_2": "3", "regressor_3": 5, "segment": "1"}
    )
    df_2 = pd.DataFrame(
        {"timestamp": timestamp[5:], "regressor_1": 2, "regressor_2": "4", "regressor_3": 6, "segment": "2"}
    )
    df_exog = pd.concat([df_1, df_2], ignore_index=True)
    df_exog["regressor_2"] = df_exog["regressor_2"].astype("category")
    df_exog["regressor_3"] = df_exog["regressor_3"].astype("category")

    return df, df_exog

@pytest.fixture
def ts_with_categoricals():
    timestamp = pd.date_range("2021-01-01", "2021-01-05")
    df_1 = pd.DataFrame({"timestamp": timestamp, "target": 11, "segment": "1"})
    df_2 = pd.DataFrame({"timestamp": timestamp, "target": 12, "segment": "2"})
    df = pd.concat([df_1, df_2], ignore_index=True)
    df = TSDataset.to_dataset(df)

    timestamp = pd.date_range("2021-01-01", "2021-01-06")
    categorical_values = ["1", "2", "1", "2", "1", "2"]
    df_1 = pd.DataFrame(
        {"timestamp": timestamp, "regressor": categorical_values, "not_regressor": categorical_values, "segment": "1"}
    )
    df_2 = pd.DataFrame(
        {"timestamp": timestamp, "regressor": categorical_values, "not_regressor": categorical_values, "segment": "2"}
    )
    df_exog = pd.concat([df_1, df_2], ignore_index=True)
    df_exog = TSDataset.to_dataset(df_exog)

    ts = TSDataset(df=df, freq="D", df_exog=df_exog, known_future=["regressor"])
    return ts

@pytest.fixture()
def ts_future(example_reg_tsds):
    future = example_reg_tsds.make_future(10)
    return future

@pytest.fixture
def df_segments_int():
    """DataFrame with integer segments."""
    timestamp = pd.date_range("2021-01-01", "2021-02-01")
    df1 = pd.DataFrame({"timestamp": timestamp, "target": 3, "segment": 1})
    df2 = pd.DataFrame({"timestamp": timestamp, "target": 4, "segment": 2})
    df = pd.concat([df1, df2], ignore_index=True)
    return df

def test_check_endings_error():
    """Check that _check_endings method raises exception if some segments end with nan."""
    timestamp = pd.date_range("2021-01-01", "2021-02-01")
    df1 = pd.DataFrame({"timestamp": timestamp, "target": 11, "segment": "1"})
    df2 = pd.DataFrame({"timestamp": timestamp[:-5], "target": 12, "segment": "2"})
    df = pd.concat([df1, df2], ignore_index=True)
    df = TSDataset.to_dataset(df)
    ts = TSDataset(df=df, freq="D")

    with pytest.raises(ValueError):
        ts._check_endings()

def test_check_endings_pass():
    """Check that _check_endings method passes if there is no nans at the end of all segments."""
    timestamp = pd.date_range("2021-01-01", "2021-02-01")
    df1 = pd.DataFrame({"timestamp": timestamp, "target": 11, "segment": "1"})
    df2 = pd.DataFrame({"timestamp": timestamp, "target": 12, "segment": "2"})
    df = pd.concat([df1, df2], ignore_index=True)
    df = TSDataset.to_dataset(df)
    ts = TSDataset(df=df, freq="D")
    ts._check_endings()

def test_check_known_future_wrong_literal():
    """Check that _check_known_future raises exception if wrong literal is given."""
    with pytest.raises(ValueError, match="The only possible literal is 'all'"):
        _ = TSDataset._check_known_future("wrong-literal", None)

def test_check_known_future_error_no_df_exog():
    """Check that _check_known_future raises exception if there are no df_exog, but known_future isn't empty."""
    with pytest.raises(ValueError, match="Some features in known_future are not present in df_exog"):
        _ = TSDataset._check_known_future(["regressor_1"], None)

def test_check_known_future_error_not_matching(df_and_regressors):
    """Check that _check_known_future raises exception if df_exog doesn't contain some features in known_future."""
    _, df_exog, known_future = df_and_regressors
    known_future.append("regressor_new")
    with pytest.raises(ValueError, match="Some features in known_future are not present in df_exog"):
        _ = TSDataset._check_known_future(known_future, df_exog)

def test_check_known_future_pass_all_empty():
    """Check that _check_known_future passes if known_future and df_exog are empty."""
    regressors = TSDataset._check_known_future([], None)
    assert len(regressors) == 0

@pytest.mark.parametrize(
    "known_future, expected_columns",
    [
        ([], []),
        (["regressor_1"], ["regressor_1"]),
        (["regressor_1", "regressor_2"], ["regressor_1", "regressor_2"]),
        (["regressor_1", "regressor_1"], ["regressor_1"]),
        ("all", ["regressor_1", "regressor_2"]),
    ],
)
def test_check_known_future_pass_non_empty(df_and_regressors, known_future, expected_columns):
    _, df_exog, _ = df_and_regressors
    """Check that _check_known_future passes if df_exog is not empty."""
    regressors = TSDataset._check_known_future(known_future, df_exog)
    assert regressors == expected_columns

def test_categorical_after_call_to_pandas():
    classic_df = generate_ar_df(periods=30, start_time="2021-06-01", n_segments=2)
    classic_df["categorical_column"] = [0] * 30 + [1] * 30
    classic_df["categorical_column"] = classic_df["categorical_column"].astype("category")
    df = TSDataset.to_dataset(classic_df[["timestamp", "segment", "target"]])
    exog = TSDataset.to_dataset(classic_df[["timestamp", "segment", "categorical_column"]])
    ts = TSDataset(df, "D", exog)
    flatten_df = ts.to_pandas(flatten=True)
    assert flatten_df["categorical_column"].dtype == "category"

@pytest.mark.parametrize(
    "borders, true_borders",
    (
        (
            ("2021-02-01", "2021-06-20", "2021-06-21", "2021-07-01"),
            ("2021-02-01", "2021-06-20", "2021-06-21", "2021-07-01"),
        ),
        (
            ("2021-02-03", "2021-06-20", "2021-06-22", "2021-07-01"),
            ("2021-02-03", "2021-06-20", "2021-06-22", "2021-07-01"),
        ),
        (
            ("2021-02-01", "2021-06-20", "2021-06-21", "2021-06-28"),
            ("2021-02-01", "2021-06-20", "2021-06-21", "2021-06-28"),
        ),
        (
            ("2021-02-01", "2021-06-20", "2021-06-23", "2021-07-01"),
            ("2021-02-01", "2021-06-20", "2021-06-23", "2021-07-01"),
        ),
        ((None, "2021-06-20", "2021-06-23", "2021-06-28"), ("2021-02-01", "2021-06-20", "2021-06-23", "2021-06-28")),
        (("2021-02-03", "2021-06-20", "2021-06-23", None), ("2021-02-03", "2021-06-20", "2021-06-23", "2021-07-01")),
        ((None, "2021-06-20", "2021-06-23", None), ("2021-02-01", "2021-06-20", "2021-06-23", "2021-07-01")),
        ((None, "2021-06-20", None, None), ("2021-02-01", "2021-06-20", "2021-06-21", "2021-07-01")),
        ((None, None, "2021-06-21", None), ("2021-02-01", "2021-06-20", "2021-06-21", "2021-07-01")),
    ),
)
def test_train_test_split(borders, true_borders, tsdf_with_exog):
    train_start, train_end, test_start, test_end = borders
    train_start_true, train_end_true, test_start_true, test_end_true = true_borders
    train, test = tsdf_with_exog.train_test_split(
        train_start=train_start, train_end=train_end, test_start=test_start, test_end=test_end
    )
    assert isinstance(train, TSDataset)
    assert isinstance(test, TSDataset)
    assert (train.df == tsdf_with_exog.df[train_start_true:train_end_true]).all().all()
    assert (train.df_exog == tsdf_with_exog.df_exog).all().all()
    assert (test.df == tsdf_with_exog.df[test_start_true:test_end_true]).all().all()
    assert (test.df_exog == tsdf_with_exog.df_exog).all().all()

@pytest.mark.parametrize(
    "test_size, true_borders",
    (
        (11, ("2021-02-01", "2021-06-20", "2021-06-21", "2021-07-01")),
        (9, ("2021-02-01", "2021-06-22", "2021-06-23", "2021-07-01")),
        (1, ("2021-02-01", "2021-06-30", "2021-07-01", "2021-07-01")),
    ),
)
def test_train_test_split_with_test_size(test_size, true_borders, tsdf_with_exog):
    train_start_true, train_end_true, test_start_true, test_end_true = true_borders
    train, test = tsdf_with_exog.train_test_split(test_size=test_size)
    assert isinstance(train, TSDataset)
    assert isinstance(test, TSDataset)
    assert (train.df == tsdf_with_exog.df[train_start_true:train_end_true]).all().all()
    assert (train.df_exog == tsdf_with_exog.df_exog).all().all()
    assert (test.df == tsdf_with_exog.df[test_start_true:test_end_true]).all().all()
    assert (test.df_exog == tsdf_with_exog.df_exog).all().all()

@pytest.mark.parametrize(
    "test_size, borders, true_borders",
    (
        (
            10,
            ("2021-02-01", "2021-06-20", "2021-06-21", "2021-07-01"),
            ("2021-02-01", "2021-06-20", "2021-06-21", "2021-07-01"),
        ),
        (
            15,
            ("2021-02-03", "2021-06-20", "2021-06-22", "2021-07-01"),
            ("2021-02-03", "2021-06-20", "2021-06-22", "2021-07-01"),
        ),
        (11, ("2021-02-02", None, None, "2021-06-28"), ("2021-02-02", "2021-06-17", "2021-06-18", "2021-06-28")),
        (
            4,
            ("2021-02-03", "2021-06-20", None, "2021-07-01"),
            ("2021-02-03", "2021-06-20", "2021-06-28", "2021-07-01"),
        ),
        (
            4,
            ("2021-02-03", "2021-06-20", None, None),
            ("2021-02-03", "2021-06-20", "2021-06-21", "2021-06-24"),
        ),
    ),
)
def test_train_test_split_both(test_size, borders, true_borders, tsdf_with_exog):
    train_start, train_end, test_start, test_end = borders
    train_start_true, train_end_true, test_start_true, test_end_true = true_borders
    train, test = tsdf_with_exog.train_test_split(
        train_start=train_start, train_end=train_end, test_start=test_start, test_end=test_end, test_size=test_size
    )
    assert isinstance(train, TSDataset)
    assert isinstance(test, TSDataset)
    assert (train.df == tsdf_with_exog.df[train_start_true:train_end_true]).all().all()
    assert (train.df_exog == tsdf_with_exog.df_exog).all().all()
    assert (test.df == tsdf_with_exog.df[test_start_true:test_end_true]).all().all()
    assert (test.df_exog == tsdf_with_exog.df_exog).all().all()

@pytest.mark.parametrize(
    "borders, match",
    (
        (("2021-01-01", "2021-06-20", "2021-06-21", "2021-07-01"), "Min timestamp in df is"),
        (("2021-02-01", "2021-06-20", "2021-06-21", "2021-08-01"), "Max timestamp in df is"),
    ),
)
def test_train_test_split_warning(borders, match, tsdf_with_exog):
    train_start, train_end, test_start, test_end = borders
    with pytest.warns(UserWarning, match=match):
        tsdf_with_exog.train_test_split(
            train_start=train_start, train_end=train_end, test_start=test_start, test_end=test_end
        )

@pytest.mark.parametrize(
    "test_size, borders, match",
    (
        (
            10,
            ("2021-02-01", None, "2021-06-21", "2021-07-01"),
            "test_size, test_start and test_end cannot be applied at the same time. test_size will be ignored",
        ),
    ),
)
def test_train_test_split_warning2(test_size, borders, match, tsdf_with_exog):
    train_start, train_end, test_start, test_end = borders
    with pytest.warns(UserWarning, match=match):
        tsdf_with_exog.train_test_split(
            train_start=train_start, train_end=train_end, test_start=test_start, test_end=test_end, test_size=test_size
        )

@pytest.mark.parametrize(
    "test_size, borders, match",
    (
        (
            None,
            ("2021-02-03", None, None, "2021-07-01"),
            "At least one of train_end, test_start or test_size should be defined",
        ),
        (
            17,
            ("2021-02-01", "2021-06-20", None, "2021-07-01"),
            "The beginning of the test goes before the end of the train",
        ),
        (
            17,
            ("2021-02-01", "2021-06-20", "2021-06-26", None),
            "test_size is 17, but only 6 available with your test_start",
        ),
    ),
)
def test_train_test_split_failed(test_size, borders, match, tsdf_with_exog):
    train_start, train_end, test_start, test_end = borders
    with pytest.raises(ValueError, match=match):
        tsdf_with_exog.train_test_split(
            train_start=train_start, train_end=train_end, test_start=test_start, test_end=test_end, test_size=test_size
        )

def test_train_test_split_pass_regressors_to_output(df_and_regressors):
    df, df_exog, known_future = df_and_regressors
    ts = TSDataset(df=df, df_exog=df_exog, freq="D", known_future=known_future)
    train, test = ts.train_test_split(test_size=5)
    assert train.regressors == ts.regressors
    assert test.regressors == ts.regressors

def test_dataset_datetime_conversion():
    classic_df = generate_ar_df(periods=30, start_time="2021-06-01", n_segments=2)
    classic_df["timestamp"] = classic_df["timestamp"].astype(str)
    df = TSDataset.to_dataset(classic_df[["timestamp", "segment", "target"]])
    # todo: deal with pandas datetime format
    assert df.index.dtype == "datetime64[ns]"

def test_dataset_datetime_conversion_during_init():
    classic_df = generate_ar_df(periods=30, start_time="2021-06-01", n_segments=2)
    classic_df["categorical_column"] = [0] * 30 + [1] * 30
    classic_df["categorical_column"] = classic_df["categorical_column"].astype("category")
    df = TSDataset.to_dataset(classic_df[["timestamp", "segment", "target"]])
    exog = TSDataset.to_dataset(classic_df[["timestamp", "segment", "categorical_column"]])
    df.index = df.index.astype(str)
    exog.index = df.index.astype(str)
    ts = TSDataset(df, "D", exog)
    assert ts.df.index.dtype == "datetime64[ns]"

def test_to_dataset_segment_conversion(df_segments_int):
    """Test that `TSDataset.to_dataset` makes casting of segment to string."""
    df = TSDataset.to_dataset(df_segments_int)
    assert np.all(df.columns.get_level_values("segment") == ["1", "2"])

def test_dataset_segment_conversion_during_init(df_segments_int):
    """Test that `TSDataset.__init__` makes casting of segment to string."""
    df = TSDataset.to_dataset(df_segments_int)
    # make conversion back to integers
    columns_frame = df.columns.to_frame()
    columns_frame["segment"] = columns_frame["segment"].astype(int)
    df.columns = pd.MultiIndex.from_frame(columns_frame)
    ts = TSDataset(df=df, freq="D")
    assert np.all(ts.columns.get_level_values("segment") == ["1", "2"])

@pytest.mark.xfail
def test_make_future_raise_error_on_diff_endings(ts_diff_endings):
    with pytest.raises(ValueError, match="All segments should end at the same timestamp"):
        ts_diff_endings.make_future(10)

def test_make_future_with_imputer(ts_diff_endings, ts_future):
    imputer = TimeSeriesImputerTransform(in_column="target")
    ts_diff_endings.fit_transform([imputer])
    future = ts_diff_endings.make_future(10)
    assert_frame_equal(future.df, ts_future.df)

def test_make_future():
    timestamp = pd.date_range("2020-01-01", periods=100, freq="D")
    df1 = pd.DataFrame({"timestamp": timestamp, "target": 1, "segment": "segment_1"})
    df2 = pd.DataFrame({"timestamp": timestamp, "target": 2, "segment": "segment_2"})
    df = pd.concat([df1, df2], ignore_index=False)
    ts = TSDataset(TSDataset.to_dataset(df), freq="D")
    ts_future = ts.make_future(10)
    assert np.all(ts_future.index == pd.date_range(ts.index.max() + pd.Timedelta("1D"), periods=10, freq="D"))
    assert set(ts_future.columns.get_level_values("feature")) == {"target"}

def test_make_future_small_horizon():
    timestamp = np.arange(np.datetime64("2021-01-01"), np.datetime64("2021-02-01"))
    target1 = [np.sin(i) for i in range(len(timestamp))]
    target2 = [np.cos(i) for i in range(len(timestamp))]
    df1 = pd.DataFrame({"timestamp": timestamp, "target": target1, "segment": "1"})
    df2 = pd.DataFrame({"timestamp": timestamp, "target": target2, "segment": "2"})
    df = pd.concat([df1, df2], ignore_index=True)
    df = TSDataset.to_dataset(df)
    ts = TSDataset(df, freq="D")
    train = TSDataset(ts[: ts.index[10], :, :], freq="D")
    with pytest.warns(UserWarning, match="TSDataset freq can't be inferred"):
        assert len(train.make_future(1).df) == 1

def test_make_future_with_exog():
    timestamp = pd.date_range("2020-01-01", periods=100, freq="D")
    df1 = pd.DataFrame({"timestamp": timestamp, "target": 1, "segment": "segment_1"})
    df2 = pd.DataFrame({"timestamp": timestamp, "target": 2, "segment": "segment_2"})
    df = pd.concat([df1, df2], ignore_index=False)
    exog = df.copy()
    exog.columns = ["timestamp", "exog", "segment"]
    ts = TSDataset(df=TSDataset.to_dataset(df), df_exog=TSDataset.to_dataset(exog), freq="D")
    ts_future = ts.make_future(10)
    assert np.all(ts_future.index == pd.date_range(ts.index.max() + pd.Timedelta("1D"), periods=10, freq="D"))
    assert set(ts_future.columns.get_level_values("feature")) == {"target", "exog"}

def test_make_future_with_regressors(df_and_regressors):
    df, df_exog, known_future = df_and_regressors
    ts = TSDataset(df=df, df_exog=df_exog, freq="D", known_future=known_future)
    ts_future = ts.make_future(10)
    assert np.all(ts_future.index == pd.date_range(ts.index.max() + pd.Timedelta("1D"), periods=10, freq="D"))
    assert set(ts_future.columns.get_level_values("feature")) == {"target", "regressor_1", "regressor_2"}

def test_make_future_inherits_regressors(df_and_regressors):
    df, df_exog, known_future = df_and_regressors
    ts = TSDataset(df=df, df_exog=df_exog, freq="D", known_future=known_future)
    ts_future = ts.make_future(10)
    assert ts_future.regressors == ts.regressors

def test_make_future_warn_not_enough_regressors(df_and_regressors):
    """Check that warning is thrown if regressors don't have enough values for the future."""
    df, df_exog, known_future = df_and_regressors
    ts = TSDataset(df=df, df_exog=df_exog, freq="D", known_future=known_future)
    with pytest.warns(UserWarning, match="Some regressors don't have enough values"):
        ts.make_future(ts.df_exog.shape[0] + 100)

@pytest.mark.parametrize("exog_starts_later,exog_ends_earlier", ((True, False), (False, True), (True, True)))
def test_check_regressors_error(exog_starts_later: bool, exog_ends_earlier: bool):
    """Check that error is raised if regressors don't have enough values for the train data."""
    start_time_main = "2021-01-01"
    end_time_main = "2021-02-01"
    start_time_regressors = "2021-01-10" if exog_starts_later else start_time_main
    end_time_regressors = "2021-01-20" if exog_ends_earlier else end_time_main

    timestamp = pd.date_range("2021-01-01", "2021-02-01")
    df1 = pd.DataFrame({"timestamp": timestamp, "target": 11, "segment": "1"})
    df2 = pd.DataFrame({"timestamp": timestamp[5:], "target": 12, "segment": "2"})
    df = pd.concat([df1, df2], ignore_index=True)
    df = TSDataset.to_dataset(df)

    timestamp = pd.date_range(start_time_regressors, end_time_regressors)
    df1 = pd.DataFrame({"timestamp": timestamp, "regressor_aaa": 1, "segment": "1"})
    df2 = pd.DataFrame({"timestamp": timestamp[5:], "regressor_aaa": 2, "segment": "2"})
    df_regressors = pd.concat([df1, df2], ignore_index=True)
    df_regressors = TSDataset.to_dataset(df_regressors)

    with pytest.raises(ValueError):
        TSDataset._check_regressors(df=df, df_regressors=df_regressors)

def test_check_regressors_pass(df_and_regressors):
    """Check that regressors check on creation passes with correct regressors."""
    df, df_exog, _ = df_and_regressors
    _ = TSDataset._check_regressors(df=df, df_regressors=df_exog)

def test_check_regressors_pass_empty(df_and_regressors):
    """Check that regressors check on creation passes with no regressors."""
    df, _, _ = df_and_regressors
    _ = TSDataset._check_regressors(df=df, df_regressors=pd.DataFrame())

def test_getitem_only_date(tsdf_with_exog):
    df_date_only = tsdf_with_exog["2021-02-01"]
    assert df_date_only.name == pd.Timestamp("2021-02-01")
    pd.testing.assert_series_equal(tsdf_with_exog.df.loc["2021-02-01"], df_date_only)

def test_getitem_slice_date(tsdf_with_exog):
    df_slice = tsdf_with_exog["2021-02-01":"2021-02-03"]
    expected_index = pd.DatetimeIndex(pd.date_range("2021-02-01", "2021-02-03"), name="timestamp")
    pd.testing.assert_index_equal(df_slice.index, expected_index)
    pd.testing.assert_frame_equal(tsdf_with_exog.df.loc["2021-02-01":"2021-02-03"], df_slice)

def test_getitem_second_ellipsis(tsdf_with_exog):
    df_slice = tsdf_with_exog["2021-02-01":"2021-02-03", ...]
    expected_index = pd.DatetimeIndex(pd.date_range("2021-02-01", "2021-02-03"), name="timestamp")
    pd.testing.assert_index_equal(df_slice.index, expected_index)
    pd.testing.assert_frame_equal(tsdf_with_exog.df.loc["2021-02-01":"2021-02-03"], df_slice)

def test_getitem_first_ellipsis(tsdf_with_exog):
    df_slice = tsdf_with_exog[..., "target"]
    df_expected = tsdf_with_exog.df.loc[:, [["Moscow", "target"], ["Omsk", "target"]]]
    pd.testing.assert_frame_equal(df_expected, df_slice)

def test_getitem_all_indexes(tsdf_with_exog):
    df_slice = tsdf_with_exog[:, :, :]
    df_expected = tsdf_with_exog.df
    pd.testing.assert_frame_equal(df_expected, df_slice)

def test_finding_regressors_marked(df_and_regressors):
    """Check that ts.regressors property works correctly when regressors set."""
    df, df_exog, known_future = df_and_regressors
    ts = TSDataset(df=df, df_exog=df_exog, freq="D", known_future=["regressor_1", "regressor_2"])
    assert sorted(ts.regressors) == ["regressor_1", "regressor_2"]

def test_finding_regressors_unmarked(df_and_regressors):
    """Check that ts.regressors property works correctly when regressors don't set."""
    df, df_exog, _ = df_and_regressors
    ts = TSDataset(df=df, df_exog=df_exog, freq="D")
    assert sorted(ts.regressors) == []

def test_head_default(tsdf_with_exog):
    assert np.all(tsdf_with_exog.head() == tsdf_with_exog.df.head())

def test_tail_default(tsdf_with_exog):
    np.all(tsdf_with_exog.tail() == tsdf_with_exog.df.tail())

def test_right_format_sorting():
    """Need to check if to_dataset method does not mess up with data and column names,
    sorting it with no respect to each other
    """
    df = pd.DataFrame({"timestamp": pd.date_range("2020-01-01", periods=100)})
    df["segment"] = "segment_1"
    # need names and values in inverse fashion
    df["reg_2"] = 1
    df["reg_1"] = 2
    tsd = TSDataset(TSDataset.to_dataset(df), freq="D")
    inv_df = tsd.to_pandas(flatten=True)
    pd.testing.assert_series_equal(df["reg_1"], inv_df["reg_1"])
    pd.testing.assert_series_equal(df["reg_2"], inv_df["reg_2"])

def test_to_flatten_simple(example_df):
    """Check that TSDataset.to_flatten works correctly in simple case."""
    flat_df = example_df
    sorted_columns = sorted(flat_df.columns)
    expected_df = flat_df[sorted_columns]
    obtained_df = TSDataset.to_flatten(TSDataset.to_dataset(flat_df))[sorted_columns]
    assert np.all(expected_df.columns == obtained_df.columns)
    assert np.all(expected_df.dtypes == obtained_df.dtypes)
    assert np.all(expected_df.values == obtained_df.values)

def test_to_flatten_with_exog(df_and_regressors_flat):
    """Check that TSDataset.to_flatten works correctly with exogenous features."""
    df, df_exog = df_and_regressors_flat
    # add a category type
    flat_df = pd.merge(left=df, right=df_exog, left_on=["timestamp", "segment"], right_on=["timestamp", "segment"])
    sorted_columns = sorted(flat_df.columns)
    expected_df = flat_df[sorted_columns]
    # add values to absent timestamps at one segment
    to_append = pd.DataFrame({"timestamp": df["timestamp"][:5], "segment": ["2"] * 5})
    expected_df = expected_df.append(to_append).sort_values(by=["segment", "timestamp"]).reset_index(drop=True)
    # rebuild category type according to new values

    obtained_df = TSDataset.to_flatten(TSDataset.to_dataset(flat_df))[sorted_columns].sort_values(
        by=["segment", "timestamp"]
    )
    assert np.all(sorted_columns == obtained_df.columns)
    assert np.all(expected_df.dtypes == obtained_df.dtypes)
    assert expected_df.equals(obtained_df)

def test_transform_raise_warning_on_diff_endings(ts_diff_endings):
    with pytest.warns(Warning, match="Segments contains NaNs in the last timestamps."):
        ts_diff_endings.transform([])

def test_fit_transform_raise_warning_on_diff_endings(ts_diff_endings):
    with pytest.warns(Warning, match="Segments contains NaNs in the last timestamps."):
        ts_diff_endings.fit_transform([])

def test_gather_common_data(df_and_regressors):
    """Check that TSDataset._gather_common_data correctly finds common data for info/describe methods."""
    df, df_exog, known_future = df_and_regressors
    ts = TSDataset(df=df, df_exog=df_exog, freq="D", known_future=known_future)
    common_data = ts._gather_common_data()
    assert common_data["num_segments"] == 2
    assert common_data["num_exogs"] == 2
    assert common_data["num_regressors"] == 2
    assert common_data["num_known_future"] == 2
    assert common_data["freq"] == "D"

def test_gather_segments_data(df_and_regressors):
    """Check that TSDataset._gather_segments_data correctly finds segment data for info/describe methods."""
    df, df_exog, known_future = df_and_regressors
    # add NaN in the middle
    df.iloc[-5, 0] = np.NaN
    # add NaNs at the end
    df.iloc[-3:, 1] = np.NaN
    ts = TSDataset(df=df, df_exog=df_exog, freq="D", known_future=known_future)
    segments = ts.segments
    segments_dict = ts._gather_segments_data(segments)
    segment_df = pd.DataFrame(segments_dict, index=segments)

    assert np.all(segment_df.index == ts.segments)
    assert segment_df.loc["1", "start_timestamp"] == pd.Timestamp("2021-01-01")
    assert segment_df.loc["2", "start_timestamp"] == pd.Timestamp("2021-01-06")
    assert segment_df.loc["1", "end_timestamp"] == pd.Timestamp("2021-02-01")
    assert segment_df.loc["2", "end_timestamp"] == pd.Timestamp("2021-01-29")
    assert segment_df.loc["1", "length"] == 32
    assert segment_df.loc["2", "length"] == 24
    assert segment_df.loc["1", "num_missing"] == 1
    assert segment_df.loc["2", "num_missing"] == 0

def test_describe(df_and_regressors):
    """Check that TSDataset.describe works correctly."""
    df, df_exog, known_future = df_and_regressors
    # add NaN in the middle
    df.iloc[-5, 0] = np.NaN
    # add NaNs at the end
    df.iloc[-3:, 1] = np.NaN
    ts = TSDataset(df=df, df_exog=df_exog, freq="D", known_future=known_future)
    description = ts.describe()

    assert np.all(description.index == ts.segments)
    assert description.loc["1", "start_timestamp"] == pd.Timestamp("2021-01-01")
    assert description.loc["2", "start_timestamp"] == pd.Timestamp("2021-01-06")
    assert description.loc["1", "end_timestamp"] == pd.Timestamp("2021-02-01")
    assert description.loc["2", "end_timestamp"] == pd.Timestamp("2021-01-29")
    assert description.loc["1", "length"] == 32
    assert description.loc["2", "length"] == 24
    assert description.loc["1", "num_missing"] == 1
    assert description.loc["2", "num_missing"] == 0
    assert np.all(description["num_segments"] == 2)
    assert np.all(description["num_exogs"] == 2)
    assert np.all(description["num_regressors"] == 2)
    assert np.all(description["num_known_future"] == 2)
    assert np.all(description["freq"] == "D")

@pytest.fixture()
def ts_with_regressors(df_and_regressors):
    df, df_exog, regressors = df_and_regressors
    ts = TSDataset(df=df, freq="D", df_exog=df_exog, known_future="all")
    return ts

def _test_update_regressors_transform(ts, transforms, expected_regressors):
    fitted_transforms = [transform.fit(ts.df) for transform in transforms]
    ts.transform(fitted_transforms)
    regressors = ts.regressors
    assert sorted(regressors) == sorted(expected_regressors)

def _test_update_regressors_fit_transform(ts, transforms, expected_regressors):
    ts.fit_transform(transforms)
    regressors = ts.regressors
    assert sorted(regressors) == sorted(expected_regressors)

@pytest.mark.parametrize(
    "transforms, expected_regressors",
    (
        ([SegmentEncoderTransform()], ["regressor_1", "regressor_2", "segment_code"]),
        (
            [LagTransform(in_column="target", lags=[1, 2], out_column="regressor_lag")],
            ["regressor_1", "regressor_2", "regressor_lag_1", "regressor_lag_2"],
        ),
    ),
)
def test_update_regressors_with_futuremixin_transform(ts_with_regressors, transforms, expected_regressors):
    _test_update_regressors_transform(deepcopy(ts_with_regressors), deepcopy(transforms), expected_regressors)
    _test_update_regressors_fit_transform(deepcopy(ts_with_regressors), deepcopy(transforms), expected_regressors)

@pytest.mark.parametrize(
    "transforms, expected_regressors",
    (
        (
            [OneHotEncoderTransform(in_column="regressor", out_column="regressor_ohe")],
            ["regressor_ohe_0", "regressor_ohe_1", "regressor"],
        ),
        ([OneHotEncoderTransform(in_column="not_regressor")], ["regressor"]),
    ),
)
def test_update_regressors_with_onehotencoder_transform(ts_with_categoricals, transforms, expected_regressors):
    _test_update_regressors_transform(deepcopy(ts_with_categoricals), deepcopy(transforms), expected_regressors)
    _test_update_regressors_fit_transform(deepcopy(ts_with_categoricals), deepcopy(transforms), expected_regressors)

@pytest.mark.parametrize(
    "transforms, expected_regressors",
    (
        (
            [MaxAbsScalerTransform(in_column="regressor_1", inplace=False, out_column="scaled")],
            ["regressor_1", "regressor_2", "scaled_regressor_1"],
        ),
        (
            [MaxAbsScalerTransform(in_column=["regressor_1", "regressor_2"], inplace=False, out_column=None)],
            [
                "regressor_1",
                "regressor_2",
                MaxAbsScalerTransform(in_column=["regressor_1"], inplace=False, out_column=None).__repr__(),
                MaxAbsScalerTransform(in_column=["regressor_2"], inplace=False, out_column=None).__repr__(),
            ],
        ),
        (
            [
                AddConstTransform(
                    in_column="regressor_1", value=2, inplace=False, out_column="regressor_add_constant_regressor_1"
                )
            ],
            ["regressor_1", "regressor_2", "regressor_add_constant_regressor_1"],
        ),
    ),
)
def test_update_regressors_with_regressor_in_column(ts_with_regressors, transforms, expected_regressors):
    _test_update_regressors_transform(deepcopy(ts_with_regressors), deepcopy(transforms), expected_regressors)
    _test_update_regressors_fit_transform(deepcopy(ts_with_regressors), deepcopy(transforms), expected_regressors)

@pytest.mark.parametrize(
    "transforms, expected_regressors",
    (
        (
            [MaxAbsScalerTransform(in_column="target", inplace=False, out_column="scaled_target")],
            ["regressor_1", "regressor_2"],
        ),
        (
            [AddConstTransform(in_column="target", value=2, inplace=False, out_column="add_constant_target")],
            ["regressor_1", "regressor_2"],
        ),
    ),
)
def test_update_regressors_not_add_not_regressors(ts_with_regressors, transforms, expected_regressors):
    _test_update_regressors_transform(deepcopy(ts_with_regressors), deepcopy(transforms), expected_regressors)
    _test_update_regressors_fit_transform(deepcopy(ts_with_regressors), deepcopy(transforms), expected_regressors)

@pytest.mark.parametrize(
    "transforms, expected_regressors",
    (
        (
            [FilterFeaturesTransform(exclude=["regressor_1"])],
            ["regressor_2"],
        ),
        (
            [FilterFeaturesTransform(exclude=["regressor_1", "regressor_2"])],
            [],
        ),
    ),
)
def test_update_regressors_after_filter(ts_with_regressors, transforms, expected_regressors):
    _test_update_regressors_transform(deepcopy(ts_with_regressors), deepcopy(transforms), expected_regressors)
    _test_update_regressors_fit_transform(deepcopy(ts_with_regressors), deepcopy(transforms), expected_regressors)


text:
import unittest
import cv2
import numpy as np
from connectdots import connect_dots
from os import remove

class TestDots(unittest.TestCase):

    def test_simple_triangle(self):
        # Set up Canvas
        input_image = np.zeros((100, 100, 3))
        # Color pixels at the corners
        input_image[0, 0, 0] = 255
        input_image[input_image.shape[0] - 1, 0, 1] = 255
        input_image[0, input_image.shape[0] - 1, 2] = 255
        input_image.astype('uint8')

        color = (255, 255, 255)
        thickness = 1

        # Construct expected output
        output_expected = cv2.line(input_image.copy(), (0, 0), (0, input_image.shape[0]),
                                   color, thickness)
        output_expected = cv2.line(output_expected, (0, 0), (input_image.shape[0], 0),
                                   color, thickness)
        output_expected = cv2.line(output_expected, (input_image.shape[0] - 1, 0), (0, input_image.shape[0] - 1),
                                   color, thickness)

        # Save test input image and generate output
        cv2.imwrite('test_img.png', input_image)
        output = connect_dots('test_img.png', return_array=True)

        # Remove test image
        remove('test_img.png')

        self.assertTrue((output == output_expected).all())

if __name__ == '__main__':
    unittest.main()

max_stars_repo_name: espenthaem/connectdots

max_issues_repo_name: espenthaem/connectdots

max_forks_repo_name: espenthaem/connectdots


text:
import streamlit as st
from PIL import Image
import classification
import numpy as np
import base64
from tempfile import NamedTemporaryFile

sign_names = {
        0: 'Speed limit (20km/h)',
        1: 'Speed limit (30km/h)',
        2: 'Speed limit (50km/h)',
        3: 'Speed limit (60km/h)',
        4: 'Speed limit (70km/h)',
        5: 'Speed limit (80km/h)',
        6: 'End of speed limit (80km/h)',
        7: 'Speed limit (100km/h)',
        8: 'Speed limit (120km/h)',
        9: 'No passing',
        10: 'No passing for vehicles over 3.5 metric tons',
        11: 'Right-of-way at the next intersection',
        12: 'Priority road',
        13: 'Yield',
        14: 'Stop',
        15: 'No vehicles',
        16: 'Vehicles over 3.5 metric tons prohibited',
        17: 'No entry',
        18: 'General caution',
        19: 'Dangerous curve to the left',
        20: 'Dangerous curve to the right',
        21: 'Double curve',
        22: 'Bumpy road',
        23: 'Slippery road',
        24: 'Road narrows on the right',
        25: 'Road work',
        26: 'Traffic signals',
        27: 'Pedestrians',
        28: 'Children crossing',
        29: 'Bicycles crossing',
        30: 'Beware of ice/snow',
        31: 'Wild animals crossing',
        32: 'End of all speed and passing limits',
        33: 'Turn right ahead',
        34: 'Turn left ahead',
        35: 'Ahead only',
        36: 'Go straight or right',
        37: 'Go straight or left',
        38: 'Keep right',
        39: 'Keep left',
        40: 'Roundabout mandatory',
        41: 'End of no passing',
        42: 'End of no passing by vehicles over 3.5 metric tons'}

st.markdown('',unsafe_allow_html=True)

main_bg = "route-66.jpg"
main_bg_ext = "jpg"

st.markdown(
    f"""

    """,
    unsafe_allow_html=True
)

st.markdown("
Traffic Sign Classification
", unsafe_allow_html=True)

button = st.markdown("""
""", unsafe_allow_html=True)

uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg"])

if uploaded_file is not None:

        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)

        st.write("")

        if st.button('predict'):
                label = classification.predict_img(image)
                label = label.argmax()
                res = sign_names.get(label)
                st.markdown(f"
{res}
", unsafe_allow_html=True)


text:
import numpy as np
import math
class Optimizer:
    def __init__(self,
            model,
            Gradient,
            loss_function,
            batch_size=1):

        self.batch_size = batch_size
        self.gradient = Gradient(model)
        self.model = model
        self.loss_func = loss_function

        self.use_SGD()

    #call one of this method to tell optimizer class
    #what optimizer function to use
    def use_SGD(self, learning_rate=0.001):
        self.learning_rate = learning_rate
        self.optim_func = self.SGD

    def use_Momentum(self,learning_rate=0.001, momentum=0.1):
        self.learning_rate = learning_rate
        self.momentum = momentum
        self.optim_func = self.Momentum

        self.old_grad_weight = np.array([], dtype=object)
        self.old_grad_bias = np.array([], dtype=object)

    def use_Adam(self, learning_rate=0.001, B1=0.9, B2=0.999, epsilon=1e-8):
        self.learning_rate = learning_rate
        self.B1 = B1
        self.B2 = B2
        self.epsilon = epsilon

        self.old_m_weight = np.array([], dtype=object)
        self.old_m_bias = np.array([], dtype=object)

        self.old_v_weight = np.array([], dtype=object)
        self.old_v_bias = np.array([], dtype=object)

    def SGD(self, grad_weight, grad_bias): #Socasitic gradient decent
        self.model.weight -= self.learning_rate * grad_weight
        self.model.bias -= self.learning_rate * grad_bias

    def Momentum(self, grad_weight, grad_bias): #momentum optimizer
        if not self.old_grad_weight.shape[0]:
            self.old_grad_weight = self.learning_rate * grad_weight
            self.old_grad_bias = self.learning_rate * grad_bias
        else:
            self.old_grad_weight = self.learning_rate * grad_weight +\
                self.momentum * self.old_grad_weight
            self.old_grad_bias = self.learning_rate * grad_bias +\
                self.momentum * self.old_grad_bias

        self.model.weight -= self.old_grad_weight
        self.model.bias -= self.old_grad_bias

    def Adam(self, grad_weight, grad_bias): #Adam optimizer.
        if not self.old_m_weight.shape[0]:
            self.old_m_weight = (1 - self.B1) * grad_weight
            self.old_m_bias = (1 - self.B1) * grad_bias

            self.old_v_weight = (1 - self.B2) * grad_weight**2
            self.old_v_bias = (1 - self.B2) * grad_bias**2
        else:
            self.old_m_weight = self.B1 * self.old_m_weight +\
                    (1 - self.B1) * grad_weight
            self.old_m_bias = self.B1 * self.old_m_bias +\
                    (1 - self.B1) * grad_bias

            self.old_v_weight = self.B2 * self.old_v_weight +\
                    (1 - self.B2) * grad_weight**2
            self.old_v_bias = self.B2 * self.old_v_bias +\
                    (1 - self.B2) * grad_bias**2

        m_bar_weight = self.old_m_weight / (1 - self.B1)
        m_bar_bias = self.old_m_bias / (1 - self.B1)

        v_bar_weight = self.old_v_weight / (1 - self.B2)
        v_bar_bias = self.old_v_bias / (1 - self.B2)

        self.model.weight -= self.learning_rate * \
                m_bar_weight / \
                (np.square(v_bar_weight) + self.epsilon)
        self.model.bias -= self.learning_rate * \
                m_bar_bias / \
                (np.square(v_bar_bias) + self.epsilon)

    def fit(self, input, target, cal_loss=False): #calculate gradient with one input and one target
        predicted = self.model.forward(input)
        target = np.array(target, dtype=np.longdouble)
        G_Loss = self.loss_func(predicted, target, Derivative=True)
        if cal_loss:
            self.Loss = self.loss_func(predicted, target, Derivative=False)

        grad = self.gradient.grad(G_Loss)
        self.step(grad[0], grad[1])

    def batch_fit(self,  #calculate gradient with batche datas
            input,
            target,
            batch_size,
            flatten=False,
            cal_loss=False):

        grad = np.array([], dtype=object)
        target = np.array(target, dtype=np.longdouble)
        self.Loss = 0
        for i in range(batch_size):
            if flatten:
                predicted = self.model.forward(input[i].flatten())
            else: predicted = self.model.forward(input[i])
            G_Loss = self.loss_func(predicted, target[i], Derivative=True)
            if cal_loss:
                self.Loss += self.loss_func(predicted,
                        target[i],
                        Derivative=False)

            if not i:
                grad = self.gradient.grad(G_Loss)
            else: grad += self.gradient.grad(G_Loss)

        self.step(grad[0], grad[1])

    def batch_loss(self,input, target, batch_size, flatten=False): #calculate loss with batches data
        self.Loss = 0
        for i in range(batch_size):
            if flatten:
                predicted = self.model.forward(input[i].flatten())
            else: predicted = self.model.forward(input[i])
            self.Loss += self.loss_func(predicted,
                    target[i],
                    Derivative=False)
        return self.Loss

    def step(self, weight_grad, bias_grad): #update weight and bias
        self.optim_func(weight_grad, bias_grad)

max_stars_repo_name: sokmontrey/neural-tinet

max_issues_repo_name: sokmontrey/neural-tinet

max_forks_repo_name: sokmontrey/neural-tinet


text:
import cv2
import numpy as np

# img1 = cv2.imread("../images/968_1477.png")
# img2 = cv2.imread("../images/968_1478.png")
# img3 = cv2.imread("../images/968_1479.png")
# print(img1.shape)

# # horizontal stack
# # hor = np.hstack((img1,img2,img3))
# hor = np.concatenate((img1,img2),axis=1)
# hor2 = np.concatenate((hor,img3),axis=1)
# cv2.imshow('Hor',hor2)

# cv2.waitKey(0)

n_start = 36
n_end =  0
e_start = 64
e_end = 116

# for i in range(n_start,(n_end-1),-1):
#     n = i
#     horimgs = None
#     for j in range(e_start, e_end + 1):
#         e = str(j).zfill(3)
#         coordinates = "N{0}E{1}".format(str(n),e)
#         file_name = coordinates + ".jpg"
#         if j == e_start:
#             print(file_name)
#             # horimgs = cv2.imread(path + "\\" + file_name)
#         else:
#             print(file_name)
#             # newimg = cv2.imread(path + "\\" + file_name)
#             # addimg = np.concatenate((horimgs, newimg),axis=1)
#             # horimgs = addimg
#         break

# cv2.imwrite(h + "\\" + "stitch.jpg", horImgs)
# print("Image written")

for i in range(n_start,(n_end-1),-1):
    n = i
    for j in range(e_start,e_end + 1):
        print(i,j)
        # break
    break

max_stars_repo_name: deepakxyz/get-earth

max_issues_repo_name: deepakxyz/get-earth

max_forks_repo_name: deepakxyz/get-earth


text:
from evolution import Phenotype
from evolution.generator import Generator
import torch
import evolution
from evolution import Genome, Linear, Conv2d
import logging
import torch.autograd as autograd
from torch import Tensor
from util import config
from util import tools
from sklearn.metrics import roc_auc_score
import numpy as np
from torch.nn.functional import binary_cross_entropy_with_logits

logger = logging.getLogger(__name__)

class Discriminator(Phenotype):

    labels = None
    fake_labels = None
    selected_loss = None

    def __init__(self, output_size=1, genome=None, input_shape=None, optimizer_conf=None):
        super().__init__(output_size=output_size, genome=genome, input_shape=input_shape)
        self.output_size = output_size
        self.optimizer_conf = optimizer_conf or config.gan.discriminator.optimizer

        if genome is None:
            if config.gan.discriminator.fixed:
                self.genome = Genome(random=False, add_layer_prob=0, rm_layer_prob=0, gene_mutation_prob=0,
                                     mutate_gan_type_prob=0)
                self.genome.add(Conv2d(256, stride=2, activation_type="LeakyReLU", activation_params={"negative_slope": 0.2}))
                self.genome.add(Conv2d(128, stride=2, activation_type="LeakyReLU", activation_params={"negative_slope": 0.2}))
                self.genome.add(Conv2d(64, stride=2, activation_type="LeakyReLU", activation_params={"negative_slope": 0.2}))
                # self.genome.add(SelfAttention())
            else:
                self.genome = Genome(random=not config.evolution.sequential_layers)
                self.genome.possible_genes = [(getattr(evolution, l), {}) for l in config.gan.discriminator.possible_layers]

            if not config.gan.discriminator.fixed:
                self.genome.input_genes = [Conv2d(stride=1, normalize="spectral")]
            else:
                self.genome.input_genes = []
            if not config.gan.discriminator.fixed:
                self.genome.output_genes = [Linear(1, activation_type=None, normalize="spectral", bias=False)]
            else:
                self.genome.output_genes = [Linear(1, activation_type=None, normalize="none", bias=False)]

    def forward(self, x):
        out = super().forward(x)
        out = out.view(out.size(0), -1)
        return out

    def train_step(self, G, images):
        """Train the discriminator on real+fake"""
        self.zero_grad()

        if self.labels is None:
            self.labels = tools.cuda(Tensor(torch.ones(images.size(0))))
            self.labels = self.labels * 0.9 if config.gan.label_smoothing else self.labels

        if self.fake_labels is None:
            self.fake_labels = tools.cuda(Tensor(torch.zeros(images.size(0))))
            self.fake_labels = self.fake_labels + 0.1 if config.gan.label_smoothing else self.fake_labels

        error, real_decision, fake_decision, fake_data = self.loss(G, images)

        if self.use_gradient_penalty():
            gradient_penalty = self.gradient_penalty(images.data, fake_data.data)
            error += gradient_penalty

        error.backward()
        self.optimizer.step()

        # clip weights for WGAN
        if self.genome.gan_type == "wgan" and not self.use_gradient_penalty():
            clip_value = 0.01
            for p in self.parameters():
                p.data.clamp_(-clip_value, clip_value)

        self.calc_metrics(G, error.item(), real_decision, fake_decision)
        return error.item()

    def loss(self, G, images):
        real_decision = self(images)
        fake_data = G(G.generate_noise(images.size()[0])).detach()  # detach to avoid training G on these labels
        fake_decision = self(fake_data)
        loss_function = getattr(self, f"loss_{self.genome.gan_type}")
        error = loss_function(real_decision, fake_decision)
        return error, real_decision, fake_decision, fake_data

    def loss_wgan(self, real_decision, fake_decision):
        return -real_decision.mean() + fake_decision.mean()

    def loss_wgan_gp(self, real_decision, fake_decision):
        return self.loss_wgan(real_decision, fake_decision)

    def loss_rsgan(self, real_decision, fake_decision):
        return binary_cross_entropy_with_logits(real_decision.view(-1) - fake_decision.view(-1), self.labels)

    def loss_rasgan(self, real_decision, fake_decision):
        error = (binary_cross_entropy_with_logits(real_decision.view(-1) - torch.mean(fake_decision.view(-1)),
                                                  self.labels) +
                 binary_cross_entropy_with_logits(fake_decision.view(-1) - torch.mean(real_decision.view(-1)),
                                                  self.fake_labels)) / 2
        return error

    def loss_lsgan(self, real_decision, fake_decision):
        return (torch.mean((real_decision - 1) ** 2) + torch.mean(fake_decision ** 2)) / 2

    def loss_gan(self, real_decision, fake_decision):
        real_error = binary_cross_entropy_with_logits(real_decision.view(-1), self.labels)
        fake_error = binary_cross_entropy_with_logits(fake_decision.view(-1), self.fake_labels)
        return (real_error + fake_error) / 2

    def loss_hinge(self, real_decision, fake_decision):
        real_error = torch.mean(torch.nn.ReLU(inplace=True)(1 - real_decision))
        fake_error = torch.mean(torch.nn.ReLU(inplace=True)(1 + fake_decision))
        return real_error + fake_error

    def eval_step(self, G, images):
        error, real_decision, fake_decision, _ = self.loss(G, images)
        self.calc_metrics(G, error.item(), real_decision, fake_decision)
        return error.item()

    def use_gradient_penalty(self):
        return config.gan.discriminator.use_gradient_penalty or self.genome.gan_type == "wgan_gp"

    def calc_global_metrics(self, best_generators, images):
        if Generator.noise_images is None:
            Generator.noise_images = tools.cuda(torch.FloatTensor(len(images), *images[0].shape).uniform_(-1, 1))
        G = tools.cuda(best_generators[0])
        labels = tools.cuda(Tensor(torch.ones(images.size(0))))
        fake_labels = tools.cuda(Tensor(torch.zeros(images.size(0))))
        if config.evolution.fitness.discriminator.startswith("validation_loss_"):
            fake_data = G(G.generate_noise(images.size(0)))
            loss_function = getattr(self, config.evolution.fitness.discriminator[11:])
            error = loss_function(self(images), self(fake_data)).item()
            self.fitness_values = [error]
        elif config.evolution.fitness.discriminator == "rel_avg":
            with torch.no_grad():
                real_decision = self(images)
                fake_data = G(G.generate_noise(images.size(0)))
                fake_decision = self(fake_data)

                noise_score = torch.mean(torch.sigmoid(self(Generator.noise_images)))
                train_score = torch.mean(torch.sigmoid(real_decision))
                gen_score = torch.mean(torch.sigmoid(fake_decision))

                d_conf = (1 + train_score - noise_score) / 2
                value = -d_conf * (1 - gen_score)
                self.fitness_values = [value.item()]
        elif config.evolution.fitness.discriminator == "rel_avg2":
            with torch.no_grad():
                real_decision = self(images)
                fake_data = G(G.generate_noise(images.size(0)))
                fake_decision = self(fake_data)
                noise_decision = self(Generator.noise_images)
                error = (binary_cross_entropy_with_logits(real_decision.view(-1) - torch.mean(fake_decision.view(-1)), labels) +
                         binary_cross_entropy_with_logits(fake_decision.view(-1) - torch.mean(real_decision.view(-1)), fake_labels) +
                         binary_cross_entropy_with_logits(noise_decision.view(-1) - torch.mean(real_decision.view(-1)), fake_labels)) / 3
                self.fitness_values = [error.item()]
        if config.evolution.fitness.discriminator == "rel_avg3":
            with torch.no_grad():
                real_decision = self(images)
                fake_data = G(G.generate_noise(images.size(0)))
                fake_decision = self(fake_data)
                noise_decision = self(Generator.noise_images)
                mean_noise = torch.mean(noise_decision)
                error = (binary_cross_entropy_with_logits(real_decision.view(-1) - (torch.mean(fake_decision.view(-1)) + mean_noise)/2, labels) +
                         binary_cross_entropy_with_logits(fake_decision.view(-1) - (torch.mean(real_decision.view(-1)) + mean_noise)/2, fake_labels)) / 2
                self.fitness_values = [error.item()]
        G.cpu()

    def calc_metrics(self, G, error, real_decision, fake_decision):
        self.calc_win_rate(torch.sigmoid(real_decision), torch.sigmoid(fake_decision), G)
        if config.evolution.fitness.discriminator == "rel_avg":
            pass
            # with torch.no_grad():
            #     c, w, h = self.input_shape[-3], self.input_shape[-2], self.input_shape[-1]
            #     if Generator.noise_images is None:
            #         Generator.noise_images = tools.cuda(torch.FloatTensor(real_decision.shape[0], c, w, h).uniform_(-1, 1))
            #     noise_score = torch.mean(torch.sigmoid(self(Generator.noise_images)))
            #     train_score = torch.mean(torch.sigmoid(real_decision))
            #     gen_score = torch.mean(torch.sigmoid(fake_decision))
            #     value = -train_score*(1-noise_score)*(1-gen_score)
            #     self.fitness_values.append(value.item())
        elif config.evolution.fitness.discriminator == "loss":
            self.fitness_values.append(error)
        elif config.evolution.fitness.discriminator.startswith("loss_"):
            loss_function = getattr(self, config.evolution.fitness.discriminator)
            error = loss_function(real_decision, fake_decision).item()
            self.fitness_values.append(error)
        elif config.evolution.fitness.discriminator == "AUC":
            full_decision = np.concatenate((torch.sigmoid(real_decision).detach().cpu().numpy().flatten(),
                                            torch.sigmoid(fake_decision).detach().cpu().numpy().flatten()))
            full_labels = np.concatenate((np.ones(real_decision.size()[0]), np.zeros(fake_decision.size()[0])))
            self.fitness_values.append(1 - roc_auc_score(full_labels, full_decision))
            # self.fitness_value -= average_precision_score(full_labels, full_decision)
            # self.fitness_value += 1 - np.mean(
            #     [accuracy_score(np.zeros(fake_decision.size()[0]), fake_decision.cpu().data.numpy().flatten() > 0.5),
            #      accuracy_score(np.ones(real_decision.size()[0]), real_decision.cpu().data.numpy().flatten() > 0.5)])
        elif config.evolution.fitness.discriminator == "BCE":
            self.fitness_values.append(self.loss_gan(real_decision, fake_decision).item())
        elif config.evolution.fitness.discriminator == "random_loss":
            if Discriminator.selected_loss is None:
                Discriminator.selected_loss = np.random.choice(config.gan.possible_gan_types)
                logger.info(f"using random loss function as fitness: {Discriminator.selected_loss}")
            loss_function = getattr(self, f"loss_{Discriminator.selected_loss}")
            self.fitness_values.append(loss_function(real_decision, fake_decision).item())

    def calc_win_rate(self, real_decision, fake_decision, G):
        if self.skill_rating_enabled():
            self.win_rate += (sum((real_decision > 0.5).float()) + sum((fake_decision < 0.5).float())).item()/(len(real_decision) + len(fake_decision))
            # self.win_rate += (torch.mean(real_decision).item() + (1-torch.mean(fake_decision)).item())/2
            # self.win_rate += (1-torch.mean(fake_decision)).item()

            # for match in [torch.mean((real_decision > 0.5).float()).item(), torch.mean((real_decision < 0.5).float()).item()]:
            # for match in [torch.mean(real_decision) > 0.5, torch.mean(fake_decision).item() < 0.5]:
            #     match = int(match)
            #     self.skill_rating_games.append((match, G.skill_rating))
            #     G.skill_rating_games.append((1 - match, self.skill_rating))

            # matches = [(p > 0.5) for p in real_decision] + [(p < 0.5) for p in fake_decision]
            # for match in matches:
            #     match = match.float().item()
            #     self.skill_rating_games.append((match, G.skill_rating))
            #     G.skill_rating_games.append((1 - match, self.skill_rating))

    def gradient_penalty(self, real_data, fake_data, epsilon=1e-12):
        batch_size = real_data.size()[0]
        alpha = torch.rand(batch_size, 1, 1, 1)
        alpha = tools.cuda(alpha.expand_as(real_data))

        interpolates = alpha * real_data + ((1 - alpha) * fake_data)
        interpolates = interpolates.clone().detach().requires_grad_(True)

        disc_interpolates = tools.cuda(self(interpolates))
        gradients = autograd.grad(outputs=disc_interpolates, inputs=interpolates,
                                  grad_outputs=tools.cuda(torch.ones(disc_interpolates.size())),
                                  create_graph=True, retain_graph=True, only_inputs=True)[0]
        gradients_norm = torch.sqrt(torch.sum(gradients ** 2, dim=1) + epsilon)
        return ((gradients_norm - 1) ** 2).mean() * config.gan.discriminator.gradient_penalty_lambda

    def finish_generation(self, **kwargs):
        if config.evolution.fitness.discriminator == "skill_rating":
            self.fitness_values.append(-self.skill_rating.mu) #+ 2*self.skill_rating.phi
        elif config.evolution.fitness.discriminator == "random":
            self.fitness_values = [np.random.rand()]

max_stars_repo_name: vfcosta/coegan-trained

max_issues_repo_name: vfcosta/coegan-trained

max_forks_repo_name: vfcosta/coegan-trained


text:
import numpy as np
import cv2
import matplotlib.pyplot as plt

__all__ = ['opencv_fitEllipse', 'draw_ellipse']

def opencv_fitEllipse(img, method="Direct"):
    #assert binary_mask.min() >= 0.0 and binary_mask.max() <= 1.0
    # points = np.argwhere(binary_mask == 1)  # TODO: tune threshold

    contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    # c_img = cv2.drawContours(img, contours, -1, (0, 255, 0), 1)
    # plt.figure()
    # plt.imshow(c_img)
    # plt.show()

    for c in contours:
        ellipse = cv2.fitEllipse(c)

    if method == "AMS":
        for c in contours:
            ellipse = cv2.fitEllipseAMS(c)
    elif method == "Direct":
        for c in contours:
            ellipse = cv2.fitEllipseDirect(c)
    elif method == "Simple":
        for c in contours:
            ellipse = cv2.fitEllipse(c)
    else:
        raise ValueError("Wrong method")

    return ellipse

def draw_ellipse(img, rgb_img):
    ellipse = opencv_fitEllipse(img)

    # (xx, yy), (MA, ma), angle = opencv_fitEllipse(img)
    #
    # print(ellipse)

    cv2.ellipse(
        rgb_img,
        ellipse,
        color=(255, 0, 0),
        thickness=3,
    )

    # filename = "ellipse_on_crop_mask.jpg"
    # cv2.imwrite(filename, rgb_img)
    # plt.imshow(rgb_img)
    # plt.show()

    # ####
    # cv2.ellipse(
    #     img,
    #     ellipse,
    #     color=(255, 0, 0),
    #     thickness=3,
    # )
    # contours, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    #
    #
    # perimeter = cv2.arcLength(contours, True)
    # print("Perimeter: ", perimeter)

    # c_img = cv2.drawContours(img, contours, -1, (0, 255, 0), 1)
    # plt.figure()
    # plt.imshow(c_img)
    # plt.show()

    return ellipse


text:
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

n = 20
xr = np.linspace(0, 10, n)
yr = np.linspace(0, 10, n)
xs, ys = np.meshgrid(xr, yr)
xs = np.reshape(xs, n*n)
ys = np.reshape(ys, n*n)

fs = 2*xs + 1.3*ys + 4
rs = np.random.randn(n*n)
rfs = fs + 3*rs

D = np.column_stack((xs, ys, rfs))
print(D.shape)
# np.savetxt('linreg-multi-synthetic-2.csv', D, delimiter=',')
# print(D)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(xs, ys, rfs)
plt.show()


text:
"""
Author: Rodrigo Coelho
Shan-Chen model (multiphase)
Segregation of two phases

Copywrite notice: All rights reserved. Do not share this code without permission.
"""
import numpy as np
import matplotlib.pyplot as plt

tmax = 10000
tau = 0.8           # visc = (tau - 0.5)/3
L = 50              #system size
Q = 9               # number of velocity vectors
f=np.zeros((L,L, Q))
Ux=np.zeros((L,L))
Uy=np.zeros((L,L))
Dens=np.ones((L,L))
Psi=np.ones((L,L))

#weights and velocityvectors
w = [4.0/9.0, 1.0/9.0, 1.0/9.0, 1.0/9.0, 1.0/9.0, 1.0/36.0,
             1.0/36.0, 1.0/36.0, 1.0/36.0]
ex = [0, 1, 0, -1, 0, 1, -1, -1, 1]
ey = [0, 0, 1, 0, -1, 1, 1, -1, -1]

#equilibrium distribution function
def feq(k, dens, ux, uy):
    return w[k]*dens*(1.0+3.0*(ex[k]*ux+ey[k]*uy)
            + 4.5*(ex[k]*ux+ey[k]*uy)**2 - 1.5*(ux*ux+uy*uy))

#initial conditions
def init():
    for j in np.arange(L):
        for i in np.arange(L):
            Ux[i][j] = 0.0
            Uy[i][j] = 0.0
            Dens[i][j] = 0.99 + 0.02*np.random.random_sample()
            for k in np.arange(Q):
                f[i][j][k] = feq(k, Dens[i][j], Ux[i][j], Uy[i][j])

#calculation of density and velocity
def macro():
    for j in np.arange(L):
        for i in np.arange(L):
            dens = 0.0
            ux = 0.0
            uy = 0.0
            for k1 in np.arange(Q):
                dens = dens + f[i][j][k1]
            for k2 in np.arange(Q):
                ux = ux + f[i][j][k2]*ex[k2]
            for k3 in np.arange(Q):
                uy = uy + f[i][j][k3]*ey[k3]
            Dens[i][j] = dens
            Ux[i][j] = ux/dens
            Uy[i][j] = uy/dens

#multiphase model
def ShanChen():
    for j in np.arange(L):
        for i in np.arange(L):
            G = -5.5            #temperature related parameter
            forceX = 0.0
            forceY = 0.0
            for k in np.arange(Q):
                inext = (i + ex[k]) % L
                jnext = (j + ey[k]) % L
                psinext = (1.0 - np.exp(-Dens[inext][jnext]))
                forceX = forceX + psinext*ex[k]*w[k]
                forceY = forceY + psinext*ey[k]*w[k]
            psi = 1.0 - np.exp(-Dens[i][j])
            forceX = -forceX*G*psi
            forceY = -forceY*G*psi
            Ux[i][j] = Ux[i][j] + forceX*tau/Dens[i][j]
            Uy[i][j] = Uy[i][j] + forceY*tau/Dens[i][j]

#Boltzmann equation: collision + streamming
def collision():
    for j in np.arange(L):
        for i in np.arange(L):
            for k in np.arange(Q):
                eq = feq(k, Dens[i][j], Ux[i][j], Uy[i][j])
                f[i][j][k] = f[i][j][k] + (eq - f[i][j][k])/tau

def stream():
    fbuffer = np.copy(f)
    for j in np.arange(L):
        for i in np.arange(L):
            for k in np.arange(Q):
                iprev = (i - ex[k]) % L
                jprev = (j - ey[k]) % L
                f[i][j][k] = fbuffer[iprev][jprev][k]

##########################################
#main
init()

for t in np.arange(tmax):
    collision()
    stream()
    macro()
    ShanChen()
    if((t%10)==0):
        plt.imshow(Dens, interpolation="bilinear", cmap='jet')
        #plt.savefig("mp-%d.png"%t, dpi=200)
        plt.colorbar()
        plt.pause(0.05)
        plt.clf()

max_stars_repo_name: rcvcoelho/lbm-python

max_issues_repo_name: rcvcoelho/lbm-python

max_forks_repo_name: rcvcoelho/lbm-python


text:
import os
import warnings

import numpy as np
from torch.utils.data import Dataset

warnings.filterwarnings('ignore')

def pc_normalize(pc):
    centroid = np.mean(pc, axis=0)
    pc = pc - centroid
    m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
    pc = pc / m
    return pc, centroid, m

def farthest_point_sample(point, npoint):
    """
    Input:
        xyz: pointcloud data, [N, D]
        npoint: number of samples
    Return:
        centroids: sampled pointcloud index, [npoint, D]
    """
    N, D = point.shape
    xyz = point[:, :3]
    centroids = np.zeros((npoint,))
    distance = np.ones((N,)) * 1e10
    farthest = np.random.randint(0, N)
    for i in range(npoint):
        centroids[i] = farthest
        centroid = xyz[farthest, :]
        dist = np.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = np.argmax(distance, -1)
    point = point[centroids.astype(np.int32)]
    return point

class TranslationDataLoader(Dataset):
    def __init__(self, root, npoint=1024, split='train', uniform=False, normal_channel=True, cache_size=15000):
        self.root = root
        self.npoints = npoint
        self.uniform = uniform
        self.split = split

        self.normal_channel = normal_channel

        shape_ids = {}
        shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'train', 'pointdata_train.txt'))]
        shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 'test', 'pointdata_test.txt'))]

        assert (split == 'train' or split == 'test')
        # list of (shape_name, shape_txt_file_path) tuple
        self.datapath = [(os.path.join(self.root, self.split, shape_ids[split][i])) for i in range(len(shape_ids[split]))]
        print('The size of %s data is %d' % (split, len(self.datapath)))

        self.cache_size = cache_size  # how many data points to cache in memory
        self.cache = {}  # from index to (point_set, cls) tuple

    def __len__(self):
        return len(self.datapath)

    def _get_item(self, index):
        if index in self.cache:
            point_set = self.cache[index]
        else:
            fn = self.datapath[index]  # file name
            point_set = np.loadtxt(fn).astype(np.float32)
            if len(self.cache) < self.cache_size:
                self.cache[index] = point_set

        point_set[:, 0:3], centroid, m = pc_normalize(point_set[:, 0:3])

        if not self.normal_channel:
            point_set = point_set[:, 0:3]

        return point_set, centroid, m

    def __getitem__(self, index):
        return self._get_item(index)

if __name__ == '__main__':
    import torch

    data = TranslationDataLoader('./data/point_data/', split='train', uniform=False, normal_channel=True)
    DataLoader = torch.utils.data.DataLoader(data, batch_size=12, shuffle=True)
    for point, label in DataLoader:
        import ipdb;

        ipdb.set_trace()
        print(point.shape)
        print(label.shape)


text:
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

plt.rcParams['pdf.fonttype'] = 42

def plot_one(ax, data):

    ax.axhline(0, color='.8', lw=2)

    # palette = {'Neumann': 'C2', 'Neumann_res': 'C8', 'Neumann_analytic': 'C4'}
    palette = {'NeuMiss': 'C2', 'NeuMiss_res': 'C8', 'NeuMiss_analytic': 'C4'}

    # Plot train
    sns.lineplot(
        data=data.query('train_test == "train"'), x='depth',
        y='r2', hue='method', ax=ax, ci=None, palette=palette,
        legend=False, style="train_test", dashes=6 * [(1, 2)])

    # Plot test
    sns.lineplot(
        data=data.query('train_test == "test"'), x='depth', palette=palette,
        y='r2', hue='method', ax=ax, estimator=np.median)

    ll = plt.legend()
    ll = plt.legend(ll.legendHandles[1:],
                    [h.get_label() for h in ll.legendHandles[1:]], fontsize=9,
                    handlelength=1, borderaxespad=.2)
    # Add the legend, so that we can add another
    ax.add_artist(ll)

    # Set axes
    ax.set_xlabel('depth')
    ax.set_xscale('log')
    ax.set_xlim(1, data['depth'].max())

    ax.set_ylabel('{} score - Bayes rate'.format('R2'))
    ax.set_ylim(-0.18, 0)

    ax.grid(True)

if __name__ == '__main__':

    scores = pd.read_csv('../results/simu_archi.csv', index_col=0)

    # Separate Bayes rate from other methods performances
    br = scores.query('method == "Bayes_rate"')
    scores = scores.query('method != "Bayes_rate"')
    scores.train_test.fillna('test', inplace=True)

    # Adjust for the Bayes rate
    for it in scores.iter.unique():
        for split in scores.train_test.unique():
            mask_br = (br.iter == it) & (br.train_test == split)
            br_val = br.loc[mask_br, 'r2']
            br_val = float(br_val)
            mask_data = (scores.iter == it) & (scores.train_test == split)
            scores.loc[mask_data, 'r2'] = scores.loc[mask_data, 'r2'] - br_val

    scores['method'] = scores['method'].replace({
            'Neumann_res': 'NeuMiss_res',
            'Neumann_analytic': 'NeuMiss_analytic',
            'Neumann': 'NeuMiss'
            })

    fig, ax = plt.subplots(1, 1, figsize=(3, 2.5))
    plot_one(ax, scores)

    plt.savefig('../figures/simu_architecture.pdf', bbox_inches='tight',
                edgecolor='none', facecolor='none', dpi=100)

    plt.close()

max_stars_repo_name: alexprz/NeuMiss

max_issues_repo_name: alexprz/NeuMiss

max_forks_repo_name: alexprz/NeuMiss


text:
"""
Peter Sonneveld,
CGS: A fast Lanczos-Type Solver for Nonsymmetric Linear Systems,
SIAM J. Sci. Stat. Comput.,
10(1):36-52, 1989,
<https://doi.org/10.1137/0910004>.

https://www.netlib.org/templates/templates.pdf
"""
from typing import Callable, Optional

import numpy as np
from numpy.typing import ArrayLike

from ._helpers import (
    Identity,
    Info,
    LinearOperator,
    aslinearoperator,
    get_default_inner,
)

def cgs(
    A: LinearOperator,
    b: ArrayLike,
    M: Optional[LinearOperator] = None,
    x0: Optional[ArrayLike] = None,
    inner: Optional[Callable] = None,
    tol: float = 1e-5,
    atol: float = 1.0e-15,
    maxiter: Optional[int] = None,
    callback: Optional[Callable] = None,
):
    b = np.asarray(b)

    assert len(A.shape) == 2
    assert A.shape[0] == A.shape[1]
    assert A.shape[1] == b.shape[0]

    A = aslinearoperator(A)

    M = Identity() if M is None else aslinearoperator(M)

    inner = get_default_inner(b.shape) if inner is None else inner

    def _norm(x):
        xx = inner(x, M @ x)
        if np.any(xx.imag != 0.0):
            raise ValueError("inner product <x, x> gave nonzero imaginary part")
        return np.sqrt(xx.real)

    if x0 is None:
        x = np.zeros_like(b)
        r0 = b.copy()
    else:
        x = np.array(x0)
        r0 = b - A @ x

    # common but arbitrary choice:
    rp = r0

    r = r0.copy()

    if callback:
        callback(x, r)

    resnorms = [_norm(r)]

    rho = 1.0
    alpha = None

    p = np.zeros_like(b)
    q = np.zeros_like(b)

    k = 0
    success = False
    criterion = np.maximum(tol * resnorms[0], atol)
    while True:
        if np.all(resnorms[-1] <= criterion):
            resnorms[-1] = _norm(b - A @ x)
            if np.all(resnorms[-1] <= criterion):
                success = True
                break

        if k == maxiter:
            break

        rho_old = rho
        rho = inner(rp, r)

        # TODO break-down for rho==0?

        beta = rho / np.where(rho_old != 0.0, rho_old, 1.0)
        u = r + beta * q
        p = u + beta * (q + beta * p)

        v = A @ (M @ p)

        s = inner(rp, v)
        alpha = rho / np.where(s != 0.0, s, 1.0)

        q = u - alpha * v

        u_ = M @ (u + q)

        x += alpha * u_
        r -= alpha * (A @ u_)

        if callback:
            callback(x, r)

        resnorms.append(_norm(r))

        k += 1

    return x if success else None, Info(success, x, k, resnorms)

max_stars_repo_name: ju-liu/krylov

max_issues_repo_name: ju-liu/krylov

max_forks_repo_name: ju-liu/krylov

reason: import numpy,from numpy