text:
#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
domain decomposition COSMO

See also the code on github

https://github.com/filippolipparini/ddPCM

and the papers

[1] Domain decomposition for implicit solvation models.
E. Cances, Y. Maday, B. Stamm
J. Chem. Phys., 139, 054111 (2013)
http://dx.doi.org/10.1063/1.4816767

[2] Fast Domain Decomposition Algorithm for Continuum Solvation Models: Energy and First Derivatives.
F. Lipparini, B. Stamm, E. Cances, Y. Maday, B. Mennucci
J. Chem. Theory Comput., 9, 3637-3648 (2013)
http://dx.doi.org/10.1021/ct400280b

[3] Quantum, classical, and hybrid QM/MM calculations in solution: General implementation of the ddCOSMO linear scaling strategy.
F. Lipparini, G. Scalmani, L. Lagardere, B. Stamm, E. Cances, Y. Maday, J.-P.Piquemal, M. J. Frisch, B. Mennucci
J. Chem. Phys., 141, 184108 (2014)
http://dx.doi.org/10.1063/1.4901304
'''

import ctypes
import copy
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf import gto
from pyscf import df
from pyscf.dft import gen_grid, numint
from pyscf.data import radii
from pyscf.symm import sph

def ddcosmo_for_scf(mf, solvent_obj=None, dm=None):
    '''Patch ddCOSMO to SCF (HF and DFT) method.

    Kwargs:
        dm : if given, solvent does not response to the change of density
            matrix. A frozen ddCOSMO potential is added to the results.
    '''
    if getattr(mf, 'with_solvent', None):
        if solvent_obj is not None:
            mf.with_solvent = solvent_obj
        return mf

    oldMF = mf.__class__
    if solvent_obj is None:
        solvent_obj = DDCOSMO(mf.mol)

    if dm is not None:
        solvent_obj.epcm, solvent_obj.vpcm = solvent_obj.kernel(dm)
        solvent_obj.frozen = True

    class SCFWithSolvent(oldMF):
        def __init__(self, mf, solvent):
            self.__dict__.update(mf.__dict__)
            self.with_solvent = solvent
            self._keys.update(['with_solvent'])

        def dump_flags(self, verbose=None):
            oldMF.dump_flags(self)
            self.with_solvent.check_sanity()
            self.with_solvent.dump_flags()
            return self

        # Note vpcm should not be added to get_hcore for scf methods.
        # get_hcore is overloaded by many post-HF methods. Modifying
        # SCF.get_hcore may lead error.

        def get_veff(self, mol=None, dm=None, *args, **kwargs):
            vhf = oldMF.get_veff(self, mol, dm, *args, **kwargs)
            with_solvent = self.with_solvent
            if not with_solvent.frozen:
                with_solvent.epcm, with_solvent.vpcm = with_solvent.kernel(dm)
            epcm, vpcm = with_solvent.epcm, with_solvent.vpcm

            # NOTE: vpcm should not be added to vhf in this place. This is
            # because vhf is used as the reference for direct_scf in the next
            # iteration. If vpcm is added here, it may break direct SCF.
            return lib.tag_array(vhf, epcm=epcm, vpcm=vpcm)

        def get_fock(self, h1e=None, s1e=None, vhf=None, dm=None, cycle=-1,
                     diis=None, diis_start_cycle=None,
                     level_shift_factor=None, damp_factor=None):
            # DIIS was called inside oldMF.get_fock. vpcm, as a function of
            # dm, should be extrapolated as well. To enable it, vpcm has to be
            # added to the fock matrix before DIIS was called.
            if getattr(vhf, 'vpcm', None) is None:
                vhf = self.get_veff(self.mol, dm)
            return oldMF.get_fock(self, h1e, s1e, vhf+vhf.vpcm, dm, cycle, diis,
                                  diis_start_cycle, level_shift_factor, damp_factor)

        def energy_elec(self, dm=None, h1e=None, vhf=None):
            if dm is None:
                dm = self.make_rdm1()
            if getattr(vhf, 'epcm', None) is None:
                vhf = self.get_veff(self.mol, dm)
            e_tot, e_coul = oldMF.energy_elec(self, dm, h1e, vhf)
            e_tot += vhf.epcm
            logger.debug(self, '  E_diel = %.15g', vhf.epcm)
            return e_tot, e_coul

        def nuc_grad_method(self):
            from pyscf.solvent import ddcosmo_grad
            grad_method = oldMF.nuc_grad_method(self)
            return ddcosmo_grad.ddcosmo_grad(grad_method, self.with_solvent)

    mf1 = SCFWithSolvent(mf, solvent_obj)
    return mf1

def ddcosmo_for_casscf(mc, solvent_obj=None, dm=None):
    '''Patch ddCOSMO to CASSCF method.

    Kwargs:
        dm : if given, solvent does not response to the change of density
            matrix. A frozen ddCOSMO potential is added to the results.
    '''
    if getattr(mc, 'with_solvent', None):
        if solvent_obj is not None:
            mc.with_solvent = solvent_obj
        return mc

    oldCAS = mc.__class__
    if solvent_obj is None:
        if getattr(mc._scf, 'with_solvent', None):
            solvent_obj = mc._scf.with_solvent
        else:
            solvent_obj = DDCOSMO(mc.mol)

    if dm is not None:
        solvent_obj.epcm, solvent_obj.vpcm = solvent_obj.kernel(dm)
        solvent_obj.frozen = True

    class CASSCFWithSolvent(oldCAS):
        def __init__(self, mc, solvent):
            self.__dict__.update(mc.__dict__)
            self.with_solvent = solvent
            self._e_tot_without_solvent = 0
            self._keys.update(['with_solvent'])

        def dump_flags(self, verbose=None):
            oldCAS.dump_flags(self)
            self.with_solvent.check_sanity()
            self.with_solvent.dump_flags()
            if self.conv_tol < 1e-7:
                logger.warn(self, 'CASSCF+ddCOSMO may not be able to '
                            'converge to conv_tol=%g', self.conv_tol)
            return self

        def update_casdm(self, mo, u, fcivec, e_ci, eris, envs={}):
            casdm1, casdm2, gci, fcivec = \
                    oldCAS.update_casdm(self, mo, u, fcivec, e_ci, eris, envs)

# The potential is generated based on the density of current micro iteration.
# It will be added to hcore in casci function. Strictly speaking, this density
# is not the same to the CASSCF density (which was used to measure
# convergence) in the macro iterations.  When CASSCF is converged, it
# should be almost the same to the CASSCF density of the macro iterations.
            with_solvent = self.with_solvent
            if not with_solvent.frozen:
                # Code to mimic dm = self.make_rdm1(ci=fcivec)
                mocore = mo[:,:self.ncore]
                mocas = mo[:,self.ncore:self.ncore+self.ncas]
                dm = reduce(numpy.dot, (mocas, casdm1, mocas.T))
                dm += numpy.dot(mocore, mocore.T) * 2
                with_solvent.epcm, with_solvent.vpcm = with_solvent.kernel(dm)

            return casdm1, casdm2, gci, fcivec

# ddCOSMO Potential should be added to the effective potential. However, there
# is no hook to modify the effective potential in CASSCF. The workaround
# here is to modify hcore. It can affect the 1-electron operator in many CASSCF
# functions: gen_h_op, update_casdm, casci.  Note hcore is used to compute the
# energy for core density (Ecore).  The resultant total energy from casci
# function will include the contribution from ddCOSMO potential. The
# duplicated energy contribution from solvent needs to be removed.
        def get_hcore(self, mol=None):
            hcore = self._scf.get_hcore(mol)
            if self.with_solvent.vpcm is not None:
                hcore += self.with_solvent.vpcm
            return hcore

        def casci(self, mo_coeff, ci0=None, eris=None, verbose=None, envs=None):
            log = logger.new_logger(self, verbose)
            log.debug('Running CASCI with solvent. Note the total energy '
                      'has duplicated contributions from solvent.')

            # In oldCAS.casci function, dE was computed based on the total
            # energy without removing the duplicated solvent contributions.
            # However, envs['elast'] is the last total energy with correct
            # solvent effects. Hack envs['elast'] to make oldCAS.casci print
            # the correct energy difference.
            envs['elast'] = self._e_tot_without_solvent
            e_tot, e_cas, fcivec = oldCAS.casci(self, mo_coeff, ci0, eris,
                                                verbose, envs)
            self._e_tot_without_solvent = e_tot

            log.debug('Computing corrections to the total energy.')
            dm = self.make_rdm1(ci=fcivec, ao_repr=True)

            with_solvent = self.with_solvent
            if with_solvent.epcm is not None:
                edup = numpy.einsum('ij,ji->', with_solvent.vpcm, dm)
                ediel = with_solvent.epcm
                e_tot = e_tot - edup + ediel
                log.info('Removing duplication %.15g, '
                         'adding E_diel = %.15g to total energy:\n'
                         '    E(CASSCF+solvent) = %.15g', edup, ediel, e_tot)

            # Update solvent effects for next iteration if needed
            if not with_solvent.frozen:
                with_solvent.epcm, with_solvent.vpcm = with_solvent.kernel(dm)

            return e_tot, e_cas, fcivec

        def nuc_grad_method(self):
            from pyscf.solvent import ddcosmo_grad
            grad_method = oldCAS.nuc_grad_method(self)
            return ddcosmo_grad.ddcosmo_grad(grad_method, self.with_solvent)

    return CASSCFWithSolvent(mc, solvent_obj)

def ddcosmo_for_casci(mc, solvent_obj=None, dm=None):
    '''Patch ddCOSMO to CASCI method.

    Kwargs:
        dm : if given, solvent does not response to the change of density
            matrix. A frozen ddCOSMO potential is added to the results.
    '''
    if getattr(mc, 'with_solvent', None):
        if solvent_obj is not None:
            mc.with_solvent = solvent_obj
        return mc

    oldCAS = mc.__class__
    if solvent_obj is None:
        if getattr(mc._scf, 'with_solvent', None):
            solvent_obj = mc._scf.with_solvent
        else:
            solvent_obj = DDCOSMO(mc.mol)

    if dm is not None:
        solvent_obj.epcm, solvent_obj.vpcm = solvent_obj.kernel(dm)
        solvent_obj.frozen = True

    class CASCIWithSolvent(oldCAS):
        def __init__(self, mc, solvent):
            self.__dict__.update(mc.__dict__)
            self.with_solvent = solvent
            self._keys.update(['with_solvent'])

        def dump_flags(self, verbose=None):
            oldCAS.dump_flags(self, verbose)
            self.with_solvent.check_sanity()
            self.with_solvent.dump_flags()
            return self

        def get_hcore(self, mol=None):
            hcore = self._scf.get_hcore(mol)
            if self.with_solvent.vpcm is not None:
                # NOTE: get_hcore was called by CASCI to generate core
                # potential.  vpcm is added in this place to take accounts the
                # effects of solvent. Its contribution is duplicated and it
                # should be removed from the total energy.
                hcore += self.with_solvent.vpcm
            return hcore

        def kernel(self, mo_coeff=None, ci0=None, verbose=None):
            with_solvent = self.with_solvent

            log = logger.new_logger(self)
            log.info('\n** Self-consistently update the solvent effects for %s **',
                     oldCAS)
            log1 = copy.copy(log)
            log1.verbose -= 1  # Suppress a few output messages

            def casci_iter_(ci0, log):
                # self.e_tot, self.e_cas, and self.ci are updated in the call
                # to oldCAS.kernel
                e_tot, e_cas, ci0 = oldCAS.kernel(self, mo_coeff, ci0, log)[:3]

                if isinstance(self.e_cas, (float, numpy.number)):
                    dm = self.make_rdm1(ci=ci0)
                else:
                    log.debug('Computing solvent responses to DM of state %d',
                              with_solvent.state_id)
                    dm = self.make_rdm1(ci=ci0[with_solvent.state_id])

                if with_solvent.epcm is not None:
                    edup = numpy.einsum('ij,ji->', with_solvent.vpcm, dm)
                    self.e_tot += with_solvent.epcm - edup

                if not with_solvent.frozen:
                    with_solvent.epcm, with_solvent.vpcm = with_solvent.kernel(dm)
                log.debug('  E_diel = %.15g', with_solvent.epcm)
                return self.e_tot, e_cas, ci0

            if with_solvent.frozen:
                with lib.temporary_env(self, _finalize=lambda:None):
                    casci_iter_(ci0, log)
                log.note('Total energy with solvent effects')
                self._finalize()
                return self.e_tot, self.e_cas, self.ci, self.mo_coeff, self.mo_energy

            self.converged = False
            with lib.temporary_env(self, canonicalization=False):
                e_tot = e_last = 0
                for cycle in range(self.with_solvent.max_cycle):
                    log.info('\n** Solvent self-consistent cycle %d:', cycle)
                    e_tot, e_cas, ci0 = casci_iter_(ci0, log1)

                    de = e_tot - e_last
                    if isinstance(e_cas, (float, numpy.number)):
                        log.info('Sovlent cycle %d  E(CASCI+solvent) = %.15g  '
                                 'dE = %g', cycle, e_tot, de)
                    else:
                        for i, e in enumerate(e_tot):
                            log.info('Solvent cycle %d  CASCI root %d  '
                                     'E(CASCI+solvent) = %.15g  dE = %g',
                                     cycle, i, e, de[i])

                    if abs(e_tot-e_last).max() < with_solvent.conv_tol:
                        self.converged = True
                        break
                    e_last = e_tot

            # An extra cycle to canonicalize CASCI orbitals
            with lib.temporary_env(self, _finalize=lambda:None):
                casci_iter_(ci0, log)
            if self.converged:
                log.info('self-consistent CASCI+solvent converged')
            else:
                log.info('self-consistent CASCI+solvent not converged')
            log.note('Total energy with solvent effects')
            self._finalize()
            return self.e_tot, self.e_cas, self.ci, self.mo_coeff, self.mo_energy

        def nuc_grad_method(self):
            from pyscf.solvent import ddcosmo_grad
            grad_method = oldCAS.nuc_grad_method(self)
            return ddcosmo_grad.ddcosmo_grad(grad_method, self.with_solvent)

    return CASCIWithSolvent(mc, solvent_obj)

def ddcosmo_for_post_scf(method, solvent_obj=None, dm=None):
    '''Default wrapper to patch ddCOSMO to post-SCF methods (CC, CI, MP,
    TDDFT etc.)

    Kwargs:
        dm : if given, solvent does not response to the change of density
            matrix. A frozen ddCOSMO potential is added to the results.
    '''
    if getattr(method, 'with_solvent', None):
        if solvent_obj is not None:
            method.with_solvent = solvent_obj
            method._scf.with_solvent = solvent_obj
        return method

    old_method = method.__class__

    if getattr(method._scf, 'with_solvent', None):
        scf_with_solvent = method._scf
        if solvent_obj is not None:
            scf_with_solvent.with_solvent = solvent_obj
    else:
        scf_with_solvent = ddcosmo_for_scf(method._scf, solvent_obj, dm)

    # Post-HF objects access the solvent effects indirectly through the
    # underlying ._scf object.
    basic_scanner = method.as_scanner()
    basic_scanner._scf = scf_with_solvent.as_scanner()

    if dm is not None:
        solvent_obj = scf_with_solvent.with_solvent
        solvent_obj.epcm, solvent_obj.vpcm = solvent_obj.kernel(dm)
        solvent_obj.frozen = True

    class PostSCFWithSolvent(old_method):
        def __init__(self, method):
            self.__dict__.update(method.__dict__)
            self._scf = scf_with_solvent
            self.with_solvent = scf_with_solvent.with_solvent
            self._keys.update(['with_solvent'])

        def dump_flags(self, verbose=None):
            old_method.dump_flags(self)
            self.with_solvent.check_sanity()
            self.with_solvent.dump_flags()
            return self

        def kernel(self, *args, **kwargs):
            with_solvent = self.with_solvent
            # The underlying ._scf object is decorated with solvent effects.
            # The resultant Fock matrix and orbital energies both include the
            # effects from solvent. It means that solvent effects for post-HF
            # methods are automatically counted if solvent is enabled at scf
            # level.
            if with_solvent.frozen:
                return old_method.kernel(self, *args, **kwargs)

            log = logger.new_logger(self)
            log.info('\n** Self-consistently update the solvent effects for %s **',
                     old_method)
            ##TODO: Suppress a few output messages
            #log1 = copy.copy(log)
            #log1.note, log1.info = log1.info, log1.debug

            e_last = 0
            for cycle in range(self.with_solvent.max_cycle):
                log.info('\n** Solvent self-consistent cycle %d:', cycle)
                # Solvent effects are applied when accessing the
                # underlying ._scf objects. The flag frozen=True ensures that
                # the generated potential with_solvent.vpcm is passed to the
                # the post-HF object, without being updated in the implicit
                # call to the _scf iterations.
                with lib.temporary_env(with_solvent, frozen=True):
                    e_tot = basic_scanner(self.mol)
                    dm = basic_scanner.make_rdm1(ao_repr=True)

                if with_solvent.epcm is not None:
                    edup = numpy.einsum('ij,ji->', with_solvent.vpcm, dm)
                    e_tot = e_tot - edup + with_solvent.epcm
                    log.debug('  E_diel = %.15g', with_solvent.epcm)

                # To generate the solvent potential for ._scf object. Since
                # frozen is set when calling basic_scanner, the solvent
                # effects are frozen during the scf iterations.
                with_solvent.epcm, with_solvent.vpcm = with_solvent.kernel(dm)

                de = e_tot - e_last
                log.info('Sovlent cycle %d  E_tot = %.15g  dE = %g',
                         cycle, e_tot, de)

                if abs(e_tot-e_last).max() < with_solvent.conv_tol:
                    break
                e_last = e_tot

            # An extra cycle to compute the total energy
            log.info('\n** Extra cycle for solvent effects')
            res = old_method.kernel(self)
            with lib.temporary_env(with_solvent, frozen=True):
                #Update everything except the _scf object and _keys
                basic_scanner(self.mol)
                new_keys = self._keys
                self.__dict__.update(basic_scanner.__dict__)
                self._keys = new_keys
                self._scf = scf_with_solvent
            self._finalize()
            return res

        def nuc_grad_method(self):
            from pyscf.solvent import ddcosmo_grad
            grad_method = old_method.nuc_grad_method(self)
            return ddcosmo_grad.ddcosmo_grad(grad_method, self.with_solvent)

    return PostSCFWithSolvent(method)

# Inject DDCOSMO to other methods
from pyscf import scf
from pyscf import mcscf
from pyscf import mp, ci, cc
scf.hf.SCF.DDCOSMO = ddcosmo_for_scf
mcscf.casci.DDCOSMO = ddcosmo_for_casci
mcscf.mc1step.DDCOSMO = ddcosmo_for_casscf
mp.mp2.MP2.DDCOSMO = ddcosmo_for_post_scf
ci.cisd.CISD.DDCOSMO = ddcosmo_for_post_scf
cc.ccsd.CCSD.DDCOSMO = ddcosmo_for_post_scf

# TODO: Testing the value of psi (make_psi_vmat).  All intermediates except
# psi are tested against ddPCM implementation on github. Psi needs to be
# computed by the host program. It requires the numerical integration code.
def gen_ddcosmo_solver(pcmobj, verbose=None):
    '''Generate ddcosmo function to compute energy and potential matrix
    '''
    mol = pcmobj.mol
    if pcmobj.grids.coords is None:
        pcmobj.grids.build(with_non0tab=True)

    natm = mol.natm
    lmax = pcmobj.lmax

    r_vdw = pcmobj.get_atomic_radii()
    coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
    ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, lmax, True))

    fi = make_fi(pcmobj, r_vdw)
    ui = 1 - fi
    ui[ui<0] = 0
    nexposed = numpy.count_nonzero(ui==1)
    nbury = numpy.count_nonzero(ui==0)
    on_shell = numpy.count_nonzero(ui>0) - nexposed
    logger.debug(pcmobj, 'Num points exposed %d', nexposed)
    logger.debug(pcmobj, 'Num points buried %d', nbury)
    logger.debug(pcmobj, 'Num points on shell %d', on_shell)

    nlm = (lmax+1)**2
    Lmat = make_L(pcmobj, r_vdw, ylm_1sph, fi)
    Lmat = Lmat.reshape(natm*nlm,-1)

    cached_pol = cache_fake_multipoles(pcmobj.grids, r_vdw, lmax)

    def gen_vind(dm):
        pcmobj._dm = dm
        if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
            # spin-traced DM for UHF or ROHF
            dm = dm[0] + dm[1]

        phi = make_phi(pcmobj, dm, r_vdw, ui)
        L_X = numpy.linalg.solve(Lmat, phi.ravel()).reshape(natm,-1)
        psi, vmat = make_psi_vmat(pcmobj, dm, r_vdw, ui, pcmobj.grids, ylm_1sph,
                                  cached_pol, L_X, Lmat)[:2]
        dielectric = pcmobj.eps
        if dielectric > 0:
            f_epsilon = (dielectric-1.)/dielectric
        else:
            f_epsilon = 1
        pcmobj.epcm = .5 * f_epsilon * numpy.einsum('jx,jx', psi, L_X)
        pcmobj.vpcm = .5 * f_epsilon * vmat
        return pcmobj.epcm, pcmobj.vpcm
    return gen_vind

def energy(pcmobj, dm):
    '''
    ddCOSMO energy
    Es = 1/2 f(eps) \int rho(r) W(r) dr
    '''
    epcm = gen_ddcosmo_solver(pcmobj, pcmobj.verbose)(dm)[0]
    return epcm

def get_atomic_radii(pcmobj):
    mol = pcmobj.mol
    vdw_radii = pcmobj.radii_table
    atom_radii = pcmobj.atom_radii

    atom_symb = [mol.atom_symbol(i) for i in range(mol.natm)]
    r_vdw = [vdw_radii[gto.charge(x)] for x in atom_symb]
    if atom_radii is not None:
        for i in range(mol.natm):
            if atom_symb[i] in atom_radii:
                r_vdw[i] = atom_radii[atom_symb[i]]
    return numpy.asarray(r_vdw)

def regularize_xt(t, eta):
    xt = numpy.zeros_like(t)
    inner = t <= 1-eta
    on_shell = (1-eta < t) & (t < 1)
    xt[inner] = 1
    ti = t[on_shell]
# JCTC, 9, 3637
    xt[on_shell] = 1./eta**5 * (1-ti)**3 * (6*ti**2 + (15*eta-12)*ti
                                            + 10*eta**2 - 15*eta + 6)
# JCP, 139, 054111
#    xt[on_shell] = 1./eta**4 * (1-ti)**2 * (ti-1+2*eta)**2
    return xt

def make_grids_one_sphere(lebedev_order):
    ngrid_1sph = gen_grid.LEBEDEV_ORDER[lebedev_order]
    leb_grid = numpy.empty((ngrid_1sph,4))
    gen_grid.libdft.MakeAngularGrid(leb_grid.ctypes.data_as(ctypes.c_void_p),
                                    ctypes.c_int(ngrid_1sph))
    coords_1sph = leb_grid[:,:3]
    # Note the Lebedev angular grids are normalized to 1 in pyscf
    weights_1sph = 4*numpy.pi * leb_grid[:,3]
    return coords_1sph, weights_1sph

def make_L(pcmobj, r_vdw, ylm_1sph, fi):
    # See JCTC, 9, 3637, Eq (18)
    mol = pcmobj.mol
    natm = mol.natm
    lmax = pcmobj.lmax
    eta = pcmobj.eta
    nlm = (lmax+1)**2

    coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
    ngrid_1sph = weights_1sph.size
    atom_coords = mol.atom_coords()
    ylm_1sph = ylm_1sph.reshape(nlm,ngrid_1sph)

# JCP, 141, 184108 Eq (9), (12) is incorrect
# L_diag = <lm|(1/|s-s'|)|l'm'>
# Using Laplace expansion for electrostatic potential 1/r
# L_diag = 4pi/(2l+1)/|s| <lm|l'm'>
    L_diag = numpy.zeros((natm,nlm))
    p1 = 0
    for l in range(lmax+1):
        p0, p1 = p1, p1 + (l*2+1)
        L_diag[:,p0:p1] = 4*numpy.pi/(l*2+1)
    L_diag *= 1./r_vdw.reshape(-1,1)
    Lmat = numpy.diag(L_diag.ravel()).reshape(natm,nlm,natm,nlm)

    for ja in range(natm):
        # scale the weight, precontract d_nj and w_n
        # see JCTC 9, 3637, Eq (16) - (18)
        # Note all values are scaled by 1/r_vdw to make the formulas
        # consistent to Psi in JCP, 141, 184108
        part_weights = weights_1sph.copy()
        part_weights[fi[ja]>1] /= fi[ja,fi[ja]>1]
        for ka in atoms_with_vdw_overlap(ja, atom_coords, r_vdw):
            vjk = r_vdw[ja] * coords_1sph + atom_coords[ja] - atom_coords[ka]
            tjk = lib.norm(vjk, axis=1) / r_vdw[ka]
            wjk = pcmobj.regularize_xt(tjk, eta, r_vdw[ka])
            wjk *= part_weights
            pol = sph.multipoles(vjk, lmax)
            p1 = 0
            for l in range(lmax+1):
                fac = 4*numpy.pi/(l*2+1) / r_vdw[ka]**(l+1)
                p0, p1 = p1, p1 + (l*2+1)
                a = numpy.einsum('xn,n,mn->xm', ylm_1sph, wjk, pol[l])
                Lmat[ja,:,ka,p0:p1] += -fac * a
    return Lmat

def make_fi(pcmobj, r_vdw):
    coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
    mol = pcmobj.mol
    eta = pcmobj.eta
    natm = mol.natm
    atom_coords = mol.atom_coords()
    ngrid_1sph = coords_1sph.shape[0]
    fi = numpy.zeros((natm,ngrid_1sph))
    for ia in range(natm):
        for ja in atoms_with_vdw_overlap(ia, atom_coords, r_vdw):
            v = r_vdw[ia]*coords_1sph + atom_coords[ia] - atom_coords[ja]
            rv = lib.norm(v, axis=1)
            t = rv / r_vdw[ja]
            xt = pcmobj.regularize_xt(t, eta, r_vdw[ja])
            fi[ia] += xt
    fi[fi < 1e-20] = 0
    return fi

def make_phi(pcmobj, dm, r_vdw, ui):
    mol = pcmobj.mol
    natm = mol.natm
    coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
    ngrid_1sph = coords_1sph.shape[0]

    if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
        dm = dm[0] + dm[1]
    tril_dm = lib.pack_tril(dm+dm.T)
    nao = dm.shape[0]
    diagidx = numpy.arange(nao)
    diagidx = diagidx*(diagidx+1)//2 + diagidx
    tril_dm[diagidx] *= .5

    atom_coords = mol.atom_coords()
    atom_charges = mol.atom_charges()

    extern_point_idx = ui > 0
    cav_coords = (atom_coords.reshape(natm,1,3)
                  + numpy.einsum('r,gx->rgx', r_vdw, coords_1sph))

    v_phi = numpy.empty((natm,ngrid_1sph))
    for ia in range(natm):
# Note (-) sign is not applied to atom_charges, because (-) is explicitly
# included in rhs and L matrix
        d_rs = atom_coords.reshape(-1,1,3) - cav_coords[ia]
        v_phi[ia] = numpy.einsum('z,zp->p', atom_charges, 1./lib.norm(d_rs,axis=2))

    max_memory = pcmobj.max_memory - lib.current_memory()[0]
    blksize = int(max(max_memory*1e6/8/nao**2, 400))

    cav_coords = cav_coords[extern_point_idx]
    v_phi_e = numpy.empty(cav_coords.shape[0])
    int3c2e = mol._add_suffix('int3c2e')
    for i0, i1 in lib.prange(0, cav_coords.shape[0], blksize):
        fakemol = gto.fakemol_for_charges(cav_coords[i0:i1])
        v_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e, aosym='s2ij')
        v_phi_e[i0:i1] = numpy.einsum('x,xk->k', tril_dm, v_nj)
    v_phi[extern_point_idx] -= v_phi_e

    ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, pcmobj.lmax, True))
    phi = -numpy.einsum('n,xn,jn,jn->jx', weights_1sph, ylm_1sph, ui, v_phi)
    return phi

def make_psi_vmat(pcmobj, dm, r_vdw, ui, grids, ylm_1sph, cached_pol, L_X, L):
    mol = pcmobj.mol
    natm = mol.natm
    lmax = pcmobj.lmax
    nlm = (lmax+1)**2

    i1 = 0
    scaled_weights = numpy.empty(grids.weights.size)
    for ia in range(natm):
        fak_pol, leak_idx = cached_pol[mol.atom_symbol(ia)]
        i0, i1 = i1, i1 + fak_pol[0].shape[1]
        eta_nj = 0
        p1 = 0
        for l in range(lmax+1):
            fac = 4*numpy.pi/(l*2+1)
            p0, p1 = p1, p1 + (l*2+1)
            eta_nj += fac * numpy.einsum('mn,m->n', fak_pol[l], L_X[ia,p0:p1])
        scaled_weights[i0:i1] = eta_nj * grids.weights[i0:i1]

    if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
        dm = dm[0] + dm[1]
    ni = numint.NumInt()
    max_memory = pcmobj.max_memory - lib.current_memory()[0]
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dm)
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()
    den = numpy.empty(grids.weights.size)
    vmat = numpy.zeros((nao,nao))
    p1 = 0
    aow = None
    for ao, mask, weight, coords \
            in ni.block_loop(mol, grids, nao, 0, max_memory):
        p0, p1 = p1, p1 + weight.size
        den[p0:p1] = weight * make_rho(0, ao, mask, 'LDA')
        aow = numpy.ndarray(ao.shape, order='F', buffer=aow)
        aow = numpy.einsum('pi,p->pi', ao, scaled_weights[p0:p1], out=aow)
        vmat -= numint._dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc)
    ao = aow = scaled_weights = None

    nelec_leak = 0
    psi = numpy.empty((natm,nlm))
    i1 = 0
    for ia in range(natm):
        fak_pol, leak_idx = cached_pol[mol.atom_symbol(ia)]
        i0, i1 = i1, i1 + fak_pol[0].shape[1]
        nelec_leak += den[i0:i1][leak_idx].sum()
        p1 = 0
        for l in range(lmax+1):
            fac = 4*numpy.pi/(l*2+1)
            p0, p1 = p1, p1 + (l*2+1)
            psi[ia,p0:p1] = -fac * numpy.einsum('n,mn->m', den[i0:i1], fak_pol[l])
# Contribution of nuclear charge to the total density
# The factor numpy.sqrt(4*numpy.pi) is due to the product of 4*pi * Y_0^0
        psi[ia,0] += numpy.sqrt(4*numpy.pi)/r_vdw[ia] * mol.atom_charge(ia)
    logger.debug(pcmobj, 'electron leak %f', nelec_leak)

    # <Psi, L^{-1}g> -> Psi = SL the adjoint equation to LX = g
    L_S = numpy.linalg.solve(L.T.reshape(natm*nlm,-1), psi.ravel()).reshape(natm,-1)
    coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
    # JCP, 141, 184108, Eq (39)
    xi_jn = numpy.einsum('n,jn,xn,jx->jn', weights_1sph, ui, ylm_1sph, L_S)
    extern_point_idx = ui > 0
    cav_coords = (mol.atom_coords().reshape(natm,1,3)
                  + numpy.einsum('r,gx->rgx', r_vdw, coords_1sph))
    cav_coords = cav_coords[extern_point_idx]
    xi_jn = xi_jn[extern_point_idx]

    max_memory = pcmobj.max_memory - lib.current_memory()[0]
    blksize = int(max(max_memory*1e6/8/nao**2, 400))

    vmat_tril = 0
    for i0, i1 in lib.prange(0, xi_jn.size, blksize):
        fakemol = gto.fakemol_for_charges(cav_coords[i0:i1])
        v_nj = df.incore.aux_e2(mol, fakemol, intor='int3c2e', aosym='s2ij')
        vmat_tril += numpy.einsum('xn,n->x', v_nj, xi_jn[i0:i1])
    vmat += lib.unpack_tril(vmat_tril)
    return psi, vmat, L_S

def cache_fake_multipoles(grids, r_vdw, lmax):
# For each type of atoms, cache the product of last two terms in
# JCP, 141, 184108, Eq (31):
# x_{<}^{l} / x_{>}^{l+1} Y_l^m
    mol = grids.mol
    atom_grids_tab = grids.gen_atomic_grids(mol)
    r_vdw_type = {}
    for ia in range(mol.natm):
        symb = mol.atom_symbol(ia)
        if symb not in r_vdw_type:
            r_vdw_type[symb] = r_vdw[ia]

    cached_pol = {}
    for symb in atom_grids_tab:
        x_nj, w = atom_grids_tab[symb]
        r = lib.norm(x_nj, axis=1)
        # Different equations are used in JCTC, 9, 3637. r*Ys (the fake_pole)
        # is computed as r^l/r_vdw. "leak_idx" is not needed.
        # Here, the implementation is based on JCP, 141, 184108
        leak_idx = r > r_vdw_type[symb]

        pol = sph.multipoles(x_nj, lmax)
        fak_pol = []
        for l in range(lmax+1):
            # x_{<}^{l} / x_{>}^{l+1} Y_l^m  in JCP, 141, 184108, Eq (31)
            #:Ys = sph.real_sph_vec(x_nj/r.reshape(-1,1), lmax, True)
            #:rr = numpy.zeros_like(r)
            #:rr[r<=r_vdw[ia]] = r[r<=r_vdw[ia]]**l / r_vdw[ia]**(l+1)
            #:rr[r> r_vdw[ia]] = r_vdw[ia]**l / r[r>r_vdw[ia]]**(l+1)
            #:xx_ylm = numpy.einsum('n,mn->mn', rr, Ys[l])
            xx_ylm = pol[l] * (1./r_vdw_type[symb]**(l+1))
            # The line below is not needed for JCTC, 9, 3637
            xx_ylm[:,leak_idx] *= (r_vdw_type[symb]/r[leak_idx])**(2*l+1)
            fak_pol.append(xx_ylm)
        cached_pol[symb] = (fak_pol, leak_idx)
    return cached_pol

def atoms_with_vdw_overlap(atm_id, atom_coords, r_vdw):
    atm_dist = atom_coords - atom_coords[atm_id]
    atm_dist = numpy.einsum('pi,pi->p', atm_dist, atm_dist)
    atm_dist[atm_id] = 1e200
    vdw_sum = r_vdw + r_vdw[atm_id]
    atoms_nearby = numpy.where(atm_dist < vdw_sum**2)[0]
    return atoms_nearby

class DDCOSMO(lib.StreamObject):
    def __init__(self, mol):
        self.mol = mol
        self.stdout = mol.stdout
        self.verbose = mol.verbose
        self.max_memory = mol.max_memory

        #self.radii_table = radii.VDW
        self.radii_table = radii.UFF*1.1
        #self.radii_table = radii.MM3
        self.atom_radii = None
        self.lebedev_order = 17
        self.lmax = 6  # max angular momentum of spherical harmonics basis
        self.eta = .1  # regularization parameter
        self.eps = 78.3553
        self.grids = gen_grid.Grids(mol)

        # The maximum iterations and convergence tolerance to update solvent
        # effects in CASCI, CC, MP, CI, ... methods
        self.max_cycle = 20
        self.conv_tol = 1e-7
        self.state_id = 0

        # Set frozen to enable/disable the frozen ddCOSMO solvent potential.
        # If frozen is set, _dm (density matrix) needs to be specified to
        # generate the potential.
        self.frozen = False

##################################################
# don't modify the following attributes, they are not input options
        # epcm (the dielectric correction) and vpcm (the additional
        # potential) are updated during the SCF iterations
        self.epcm = None
        self.vpcm = None
        self._dm = None

        # _solver_ is a cached function returned by self.as_solver() to reduce
        # the overhead of initialization. It should be cleared whenever the
        # solvent parameters or the integration grids were changed.
        self._solver_ = None

        self._keys = set(self.__dict__.keys())

    @property
    def dm(self):
        '''Density matrix to generate the frozen ddCOSMO solvent potential.'''
        return self._dm
    @dm.setter
    def dm(self, dm):
        '''Set dm to enable/disable the frozen ddCOSMO solvent potential.
        Setting dm to None will disable the frozen potental, i.e. the
        potential will be response to the change of the density during SCF
        iterations.
        '''
        if isinstance(dm, numpy.ndarray):
            self._dm = dm
            self.epcm, self.vpcm = self.kernel(dm)
        else:
            self.epcm = self.vpcm = self._dm = None

    def __setattr__(self, key, val):
        if key in ('radii_table', 'atom_radii', 'lebedev_order', 'lmax',
                   'eta', 'eps', 'grids'):
            self._solver_ = None
        super(DDCOSMO, self).__setattr__(key, val)

    def dump_flags(self):
        logger.info(self, '******** %s ********', self.__class__)
        logger.info(self, 'lebedev_order = %s (%d grids per sphere)',
                    self.lebedev_order, gen_grid.LEBEDEV_ORDER[self.lebedev_order])
        logger.info(self, 'lmax = %s'         , self.lmax)
        logger.info(self, 'eta = %s'          , self.eta)
        logger.info(self, 'eps = %s'          , self.eps)
        logger.debug2(self, 'radii_table %s', self.radii_table)
        if self.atom_radii:
            logger.info(self, 'User specified atomic radii %s', str(self.atom_radii))
        self.grids.dump_flags()
        return self

    def kernel(self, dm):
        '''A single shot solvent effects for given density matrix.
        '''
        if (self._solver_ is None or
# If self.grids.coords is None, it is very likely caused by the updates of the
# "grids" parameters. The COSMO solver should be updated to adapt the new
# integral grids.
            self.grids.coords is None):
            self._solver_ = self.as_solver()

        epcm, vpcm = self._solver_(dm)
        return epcm, vpcm

    energy = energy
    gen_solver = as_solver = gen_ddcosmo_solver
    get_atomic_radii = get_atomic_radii

    def regularize_xt(self, t, eta, scale=1):
        # scale = eta*scale, is it correct?
        return regularize_xt(t, eta*scale)

if __name__ == '__main__':
    from pyscf import scf
    from pyscf import mcscf
    from pyscf import cc
    mol = gto.M(atom='H 0 0 0; H 0 1 1.2; H 1. .1 0; H .5 .5 1')
    natm = mol.natm
    r_vdw = [radii.VDW[gto.charge(mol.atom_symbol(i))]
             for i in range(natm)]
    r_vdw = numpy.asarray(r_vdw)
    pcmobj = DDCOSMO(mol)
    pcmobj.regularize_xt = lambda t, eta, scale: regularize_xt(t, eta)
    pcmobj.lebedev_order = 7
    pcmobj.lmax = 6
    pcmobj.eta = 0.1
    nlm = (pcmobj.lmax+1)**2
    coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
    fi = make_fi(pcmobj, r_vdw)
    ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, pcmobj.lmax, True))
    L = make_L(pcmobj, r_vdw, ylm_1sph, fi)
    print(lib.finger(L) - 6.2823493771037473)

    mol = gto.Mole()
    mol.atom = ''' O                  0.00000000    0.00000000   -0.11081188
                   H                 -0.00000000   -0.84695236    0.59109389
                   H                 -0.00000000    0.89830571    0.52404783 '''
    mol.basis = '3-21g' #cc-pvdz'
    mol.build()
    cm = DDCOSMO(mol)
    cm.verbose = 4
    mf = ddcosmo_for_scf(scf.RHF(mol), cm)#.newton()
    mf.verbose = 4
    print(mf.kernel() - -75.570364368059)
    cm.verbose = 3
    e = ddcosmo_for_casci(mcscf.CASCI(mf, 4, 4)).kernel()[0]
    print(e - -75.5743583693215)
    cc_cosmo = ddcosmo_for_post_scf(cc.CCSD(mf)).run()
    print(cc_cosmo.e_tot - -75.71096812100654)

    mol = gto.Mole()
    mol.atom = ''' Fe                 0.00000000    0.00000000   -0.11081188
                   H                 -0.00000000   -0.84695236    0.59109389
                   H                 -0.00000000    0.89830571    0.52404783 '''
    mol.basis = '3-21g' #cc-pvdz'
    mol.build()
    cm = DDCOSMO(mol)
    cm.eps = -1
    cm.verbose = 4
    mf = ddcosmo_for_scf(scf.ROHF(mol), cm).newton()
    mf.verbose=4
    mf.kernel()

max_stars_repo_name: pavanell/pyscf

max_issues_repo_name: pavanell/pyscf

max_forks_repo_name: pavanell/pyscf


text:
# -*- coding: utf-8 -*-
"""
Created on Tue Mar 20 20:03:39 2018

@author: admin
"""
import tensorflow as tf
import numpy as np
from tensorflow.python.framework import ops
from sklearn import preprocessing
import matplotlib.pyplot as plt
import os
import csv

def accuracy(predictions, labels):
  return (100.0 * np.sum(np.argmax(predictions, 1) == np.argmax(labels, 1))
          / predictions.shape[0])

if __name__ == '__main__':
    path = u'E:/DeepCID'
    #Load the mixture spectrum, its' labels and components information
    datafile1 = path+'/mixture.npy'
    Xtest = np.load(datafile1)

    datafile2 = path+'/label.npy'
    label = np.load(datafile2)

    n = label.shape[0]

    csv_reader = csv.reader(open( path+'/namedata.csv', encoding='utf-8'))
    names = [row for row in csv_reader]
    ypred = np.zeros((n*Xtest.shape[0],2))

    # Set the root directory of models and reload the models one by one
    root = path+'/model'
    list_dirs = os.walk(root)
    i=0
    for root, dirs, files in list_dirs:
        for d in dirs:
            tf.reset_default_graph()

            Y1 = label[i,:].reshape([Xtest.shape[0],1])
            Y2 = np.ones((Y1.shape)) - Y1
            Ytest = np.concatenate((Y1,Y2),axis=1)

            os.chdir(os.path.join(root, d))
            datafile =u'./X_scale.npy'
            X_scale = np.load(datafile)
            Xtest_scale = (Xtest - X_scale[0])/X_scale[1]

            with tf.Session() as sess:
                new_saver=tf.train.import_meta_graph('./compoent.ckpt.meta')
                new_saver.restore(sess,"./compoent.ckpt")
                graph = tf.get_default_graph()
                xs=graph.get_operation_by_name('xs').outputs[0]
                #ys=graph.get_operation_by_name('ys').outputs[0]
                keep_prob=graph.get_operation_by_name('keep_prob').outputs[0]

                prediction = graph.get_tensor_by_name('prediction:0')
                test_ypred = sess.run(prediction,feed_dict={xs: Xtest_scale, keep_prob : 1.0})
            ypred[i*Xtest.shape[0]:(i+1)*Xtest.shape[0],:] = test_ypred
            print('compoent', i, 'finished.','The test accuracy %.1f%%' %
                  accuracy(ypred[i*Xtest.shape[0]:(i+1)*Xtest.shape[0],:],Ytest))

            i+=1

    # Print the components' name that are exited in the miatures
    for j in range(Xtest.shape[0]):
        print('The', j ,'th sample contains: ')

        # Yreal is design for calculate the accuracy.
        y_real1 = label[:,j-Xtest.shape[0]].reshape([n,1])
        y_real2 = np.ones((y_real1.shape)) - y_real1
        y_real = np.concatenate((y_real1,y_real2),axis=1)

        ypre = np.zeros((n,2))
        for k in range(n):
            ypre[k,:] = ypred[j,:]
            j = j+Xtest.shape[0]

        for h in range(n):
            if (ypre[h,0]>=0.5):
                print(names[h])
        print('The prediction finished')

ext: py

max_stars_repo_name: XiaqiongFan/DeepCID

max_issues_repo_name: XiaqiongFan/DeepCID

max_forks_repo_name: XiaqiongFan/DeepCID

reason: import numpy


text:
import numpy as np
import Simulators
from Util import *
import scipy.linalg
import Semibandits

class Thompson(Semibandits.Semibandit):

    def __init__(self,B):
        self.B = B

    def init(self, T, params={}):
        self.T = T

        self.d = self.B.d
        self.b_vec= np.matrix(np.zeros((self.d,1)))
        self.cov = np.matrix(np.eye(self.d))
        self.lam = 0.1
        if 'lambda' in params:
            self.lam = params['lambda']
        print (self.lam)
        self.Cinv = scipy.linalg.inv(self.cov)
        self.weights = self.Cinv*self.b_vec
        self.t=1

        self.reward = []
        self.opt_reward = []

    def update(self, x, A, y_vec, r):
        self.t += 1

        features = np.matrix(x.get_ld_features())
        for i in range(len(A)):
            self.cov += features[A[i],:].T*features[A[i],:]
            self.b_vec += (y_vec[i]*features[A[i],:]).T
        if True:
            self.Cinv = scipy.linalg.inv(self.cov)
            self.weights = self.Cinv*self.b_vec

    def get_action(self,x):
        features = np.matrix(x.get_ld_features())
        K = x.get_K()
        w = np.array(self.weights.T)[0]
        mut = np.matrix(np.random.multivariate_normal(w, self.lam*self.Cinv))
        scores = features*mut.T
        scores = [scores[a,0] for a in range(K)]

        ranks = np.argsort(scores)
        return ranks[K-self.B.L:K]

reason: import numpy,import scipy


text:
from __future__ import absolute_import, division, print_function
from simdna.synthetic.core import DefaultNameMixin
from simdna import random
import numpy as np
from collections import OrderedDict

class AbstractQuantityGenerator(DefaultNameMixin):
    """Class for sampling values from a distribution.
    """

    def generateQuantity(self):
        """Sample a quantity from a distribution.

        Returns:
            The sampled value.
        """
        raise NotImplementedError()

    def getJsonableObject(self):
        """Get JSON object representation.

        Returns:
            A json-friendly object (built of dictionaries, lists and
        python primitives), which can be converted to json to
        record the exact details of what was simualted.
        """
        raise NotImplementedError()

class ChooseValueFromASet(AbstractQuantityGenerator):
    """Randomly samples a particular value from a set of values.

    Arguments:
        setOfPossibleValues: array of values that will be randomly sampled
    from.

        name: see :class:`.DefaultNameMixin`.
    """

    def __init__(self, setOfPossibleValues, name=None):
        self.setOfPossibleValues = setOfPossibleValues
        super(ChooseValueFromASet, self).__init__(name)

    def generateQuantity(self):
        """See superclass.
        """
        return self.setOfPossibleValues[int(random.random() * (len(self.setOfPossibleValues)))]

    def getJsonableObject(self):
        """See superclass.
        """
        return OrderedDict([("class", "ChooseValueFromASet"),
                            ("possibleValues", self.setOfPossibleValues)])

class UniformIntegerGenerator(AbstractQuantityGenerator):
    """Randomly samples an integer from minVal to maxVal, inclusive.

    Arguments:
        minVal: minimum integer that can be sampled

        maxVal: maximum integers that can be sampled

        name: See superclass.
    """

    def __init__(self, minVal, maxVal, name=None):
        self.minVal = minVal
        self.maxVal = maxVal
        super(UniformIntegerGenerator, self).__init__(name)

    def generateQuantity(self):
        """See superclass.
        """
        # the 1+ makes the max val inclusive
        return self.minVal + int(random.random() * (1 + self.maxVal - self.minVal))

    def getJsonableObject(self):
        """See superclass.
        """
        return OrderedDict([("class", "UniformIntegerGenerator"), ("minVal", self.minVal), ("maxVal", self.maxVal)])

class FixedQuantityGenerator(AbstractQuantityGenerator):
    """Returns a fixed number every time generateQuantity is called.

    Arguments:
        quantity: the value to return when generateQuantity is called.
    """

    def __init__(self, quantity, name=None):
        self.quantity = quantity
        super(FixedQuantityGenerator, self).__init__(name)

    def generateQuantity(self):
        """See superclass.
        """
        return self.quantity

    def getJsonableObject(self):
        """See superclass.
        """
        return "fixedQuantity-" + str(self.quantity)

class PoissonQuantityGenerator(AbstractQuantityGenerator):
    """Generates values according to a poisson distribution.

    Arguments:
        mean: the mean of the poisson distribution
    """

    def __init__(self, mean, name=None):
        self.mean = mean
        super(PoissonQuantityGenerator, self).__init__(name)

    def generateQuantity(self):
        """See superclass.
        """
        return random.poisson(self.mean)

    def getJsonableObject(self):
        """See superclass.
        """
        return "poisson-" + str(self.mean)

class BernoulliQuantityGenerator(AbstractQuantityGenerator):
    """Generates 1 or 0 according to a bernoulli distribution.

    Arguments:
        prob: probability of 1
    """

    def __init__(self, prob, name=None):
        self.prob = prob
        super(BernoulliQuantityGenerator, self).__init__(name)

    def generateQuantity(self):
        """See superclass.
        """
        return 1 if (random.random() <= self.prob) else 0

    def getJsonableObject(self):
        """See superclass.
        """
        return "bernoulli-" + str(self.prob)

class MinMaxWrapper(AbstractQuantityGenerator):
    """Compress a distribution to lie within a min and a max.

    Wrapper that restricts a distribution to only return values between
    the min and the max. If a value outside the range is returned,
    resamples until it obtains a value within the range.
    Warns every time it tries to resample 10 times without successfully
    finding a value in the correct range.

    Arguments:
        quantityGenerator: instance of :class:`.AbstractQuantityGenerator`.
    Used to draw samples from the distribution to truncate

        theMin: can be None; if so will be ignored.

        theMax: can be None; if so will be ignored.
    """

    def __init__(self, quantityGenerator, theMin=None, theMax=None, name=None):
        self.quantityGenerator = quantityGenerator
        self.theMin = theMin
        self.theMax = theMax
        assert self.quantityGenerator is not None
        super(MinMaxWrapper, self).__init__(name)

    def generateQuantity(self):
        """See superclass.
        """
        tries = 0
        while (True):
            tries += 1
            quantity = self.quantityGenerator.generateQuantity()
            if ((self.theMin is None or quantity >= self.theMin) and (self.theMax is None or quantity <= self.theMax)):
                return quantity
            if (tries % 10 == 0):
                print("warning: made " + str(tries) +
                      " tries at trying to sample from distribution with min/max limits")

    def getJsonableObject(self):
        """See superclass.
        """
        return OrderedDict([("min", self.theMin), ("max", self.theMax), ("quantityGenerator", self.quantityGenerator.getJsonableObject())])

class ZeroInflater(AbstractQuantityGenerator):
    """Inflate a particular distribution with zeros.

    Wrapper that inflates the number of zeros returned.
    Flips a coin; if positive, will return zero - otherwise will
    sample from the wrapped distribution (which may still return 0)

    Arguments:
        quantityGenerator: an instance of :class:`.AbstractQuantityGenerator`;\
    represents the distribution to sample from with probability ``1-zeroProb``

        zeroProb: the probability of just returning 0 without sampling\
    from ``quantityGenerator``

        name: see :class:`.DefaultNameMixin`.
    """

    def __init__(self, quantityGenerator, zeroProb, name=None):
        self.quantityGenerator = quantityGenerator
        self.zeroProb = zeroProb
        super(ZeroInflater, self).__init__(name)

    def generateQuantity(self):
        """See superclass.
        """
        val = random.random()
        if (val < self.zeroProb):
            return 0
        else:
            return self.quantityGenerator.generateQuantity()

    def getJsonableObject(self):
        """See superclass.
        """
        return OrderedDict([("class", "ZeroInflater"), ("zeroProb", self.zeroProb), ("quantityGenerator", self.quantityGenerator.getJsonableObject())])

max_stars_repo_name: hannahgz/simdna

max_issues_repo_name: hannahgz/simdna

max_forks_repo_name: hannahgz/simdna


text:
import torch
torch.backends.cudnn.enabled = False
import torch.optim as optim
import torch.optim.lr_scheduler as lr_sched
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.autograd import Variable
import numpy as np
import os
from torchvision import transforms
from models import RSCNN_MSN_Seg as RSCNN_MSN
from data import ShapeNetPart
import utils.pytorch_utils as pt_utils
import data.data_utils as d_utils
import argparse
import random
import yaml

torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = True

seed = 123
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

parser = argparse.ArgumentParser(description='Relation-Shape CNN Shape Part Segmentation Training')
parser.add_argument('--config', default='cfgs/config_msn_partseg.yaml', type=str)

def main():
    args = parser.parse_args()
    with open(args.config) as f:
        config = yaml.load(f)
    print("\n**************************")
    for k, v in config['common'].items():
        setattr(args, k, v)
        print('\n[%s]:'%(k), v)
    print("\n**************************\n")

    try:
        os.makedirs(args.save_path)
    except OSError:
        pass

    train_transforms = transforms.Compose([
        d_utils.PointcloudToTensor()
    ])
    test_transforms = transforms.Compose([
        d_utils.PointcloudToTensor()
    ])

    train_dataset = ShapeNetPart(root = args.data_root, num_points = args.num_points, split = 'trainval', normalize = True, transforms = train_transforms)
    train_dataloader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=int(args.workers),
        pin_memory=True
    )

    global test_dataset
    test_dataset = ShapeNetPart(root = args.data_root, num_points = args.num_points, split = 'test', normalize = True, transforms = test_transforms)
    test_dataloader = DataLoader(
        test_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=int(args.workers),
        pin_memory=True
    )

    model = RSCNN_MSN(num_classes = args.num_classes, input_channels = args.input_channels, relation_prior = args.relation_prior, use_xyz = True)
    model.cuda()
    optimizer = optim.Adam(
        model.parameters(), lr=args.base_lr, weight_decay=args.weight_decay)

    lr_lbmd = lambda e: max(args.lr_decay**(e // args.decay_step), args.lr_clip / args.base_lr)
    bnm_lmbd = lambda e: max(args.bn_momentum * args.bn_decay**(e // args.decay_step), args.bnm_clip)
    lr_scheduler = lr_sched.LambdaLR(optimizer, lr_lbmd)
    bnm_scheduler = pt_utils.BNMomentumScheduler(model, bnm_lmbd)

    if args.checkpoint is not '':
        model.load_state_dict(torch.load(args.checkpoint))
        print('Load model successfully: %s' % (args.checkpoint))

    criterion = nn.CrossEntropyLoss()
    num_batch = len(train_dataset)/args.batch_size

    # training
    train(train_dataloader, test_dataloader, model, criterion, optimizer, lr_scheduler, bnm_scheduler, args, num_batch)

def train(train_dataloader, test_dataloader, model, criterion, optimizer, lr_scheduler, bnm_scheduler, args, num_batch):
    PointcloudScaleAndTranslate = d_utils.PointcloudScaleAndTranslate()   # initialize augmentation
    global Class_mIoU, Inst_mIoU
    Class_mIoU, Inst_mIoU = 0.83, 0.85
    batch_count = 0
    model.train()
    for epoch in range(args.epochs):
        for i, data in enumerate(train_dataloader, 0):
            if lr_scheduler is not None:
                lr_scheduler.step(epoch)
            if bnm_scheduler is not None:
                bnm_scheduler.step(epoch-1)
            points, target, cls = data
            points, target = points.cuda(), target.cuda()
            points, target = Variable(points), Variable(target)
            # augmentation
            points.data = PointcloudScaleAndTranslate(points.data)

            optimizer.zero_grad()

            batch_one_hot_cls = np.zeros((len(cls), 16))   # 16 object classes
            for b in range(len(cls)):
                batch_one_hot_cls[b, int(cls[b])] = 1
            batch_one_hot_cls = torch.from_numpy(batch_one_hot_cls)
            batch_one_hot_cls = Variable(batch_one_hot_cls.float().cuda())

            pred = model(points, batch_one_hot_cls)
            pred = pred.view(-1, args.num_classes)
            target = target.view(-1,1)[:,0]
            loss = criterion(pred, target)
            loss.backward()
            optimizer.step()

            if i % args.print_freq_iter == 0:
                print('[epoch %3d: %3d/%3d] \t train loss: %0.6f \t lr: %0.5f' %(epoch+1, i, num_batch, loss.data.clone(), lr_scheduler.get_lr()[0]))
            batch_count += 1

            # validation in between an epoch
            if (epoch < 3 or epoch > 40) and args.evaluate and batch_count % int(args.val_freq_epoch * num_batch) == 0:
                validate(test_dataloader, model, criterion, args, batch_count)

def validate(test_dataloader, model, criterion, args, iter):
    global Class_mIoU, Inst_mIoU, test_dataset
    model.eval()

    seg_classes = test_dataset.seg_classes
    shape_ious = {cat:[] for cat in seg_classes.keys()}
    seg_label_to_cat = {}           # {0:Airplane, 1:Airplane, ...49:Table}
    for cat in seg_classes.keys():
        for label in seg_classes[cat]:
            seg_label_to_cat[label] = cat

    losses = []
    for _, data in enumerate(test_dataloader, 0):
        points, target, cls = data
        points, target = Variable(points, volatile=True), Variable(target, volatile=True)
        points, target = points.cuda(), target.cuda()

        batch_one_hot_cls = np.zeros((len(cls), 16))   # 16 object classes
        for b in range(len(cls)):
            batch_one_hot_cls[b, int(cls[b])] = 1
        batch_one_hot_cls = torch.from_numpy(batch_one_hot_cls)
        batch_one_hot_cls = Variable(batch_one_hot_cls.float().cuda())

        pred = model(points, batch_one_hot_cls)
        loss = criterion(pred.view(-1, args.num_classes), target.view(-1,1)[:,0])
        losses.append(loss.data.clone())
        pred = pred.data.cpu()
        target = target.data.cpu()
        pred_val = torch.zeros(len(cls), args.num_points).type(torch.LongTensor)
        # pred to the groundtruth classes (selected by seg_classes[cat])
        for b in range(len(cls)):
            cat = seg_label_to_cat[target[b, 0]]
            logits = pred[b, :, :]   # (num_points, num_classes)
            pred_val[b, :] = logits[:, seg_classes[cat]].max(1)[1] + seg_classes[cat][0]

        for b in range(len(cls)):
            segp = pred_val[b, :]
            segl = target[b, :]
            cat = seg_label_to_cat[segl[0]]
            part_ious = [0.0 for _ in range(len(seg_classes[cat]))]
            for l in seg_classes[cat]:
                if torch.sum((segl == l) | (segp == l)) == 0:
                    # part is not present in this shape
                    part_ious[l - seg_classes[cat][0]] = 1.0
                else:
                    part_ious[l - seg_classes[cat][0]] = torch.sum((segl == l) & (segp == l)) / float(torch.sum((segl == l) | (segp == l)))
            shape_ious[cat].append(np.mean(part_ious))

    instance_ious = []
    for cat in shape_ious.keys():
        for iou in shape_ious[cat]:
            instance_ious.append(iou)
        shape_ious[cat] = np.mean(shape_ious[cat])
    mean_class_ious = np.mean(list(shape_ious.values()))

    for cat in sorted(shape_ious.keys()):
        print('****** %s: %0.6f'%(cat, shape_ious[cat]))
    print('************ Test Loss: %0.6f' % (np.array(losses).mean()))
    print('************ Class_mIoU: %0.6f' % (mean_class_ious))
    print('************ Instance_mIoU: %0.6f' % (np.mean(instance_ious)))

    if mean_class_ious > Class_mIoU or np.mean(instance_ious) > Inst_mIoU:
        if mean_class_ious > Class_mIoU:
            Class_mIoU = mean_class_ious
        if np.mean(instance_ious) > Inst_mIoU:
            Inst_mIoU = np.mean(instance_ious)
        torch.save(model.state_dict(), '%s/seg_msn_iter_%d_ins_%0.6f_cls_%0.6f.pth' % (args.save_path, iter, np.mean(instance_ious), mean_class_ious))
    model.train()

if __name__ == "__main__":
    main()

max_stars_repo_name: tchaton/Relation-Shape-CNN

max_issues_repo_name: tchaton/Relation-Shape-CNN

max_forks_repo_name: tchaton/Relation-Shape-CNN


text:
from datetime import datetime

import archon.broker as broker
import archon.exchange.exchanges as exc
import archon.exchange.bitmex.bitmex as mex
import archon.facade as facade
import archon.model.models as models
from archon.brokersrv.brokerservice import BrokerService

from archon.custom_logger import setup_logger, remove_loggers

from archon.util import *

import pandas as pd
import numpy
import time
import json
import atexit
import logging
import threading
import time
import redis
from archon.metrics import calc_mid_price
from archon.model import models
#from topics import *

if __name__=='__main__':
    try:
        abroker = BrokerService(setAuto=True,initFeeder=False)
        abroker.set_keys_exchange_file(exchanges=[exc.BITMEX])
        book = abroker.orderbook(exc.BITMEX)
        book = models.conv_orderbook(book, exc.BITMEX)
        mid = calc_mid_price(book)
        print ("bitmex mid",mid)

        pos = abroker.position(exc.BITMEX)
        #print (pos)
        print (pos[0]["currentQty"])

        #oo = abroker.openorders(exc.BITMEX)
        #print (oo)

    except Exception as e:
        print ("error ",e)

max_stars_repo_name: economicnetwork/makers

max_issues_repo_name: economicnetwork/makers

max_forks_repo_name: economicnetwork/makers


text:
from unittest import TestCase

from dexpy.ccd import build_ccd
from dexpy.ccd import alpha_from_type
from dexpy.eval import det_xtxi
from dexpy.model import make_quadratic_model
import numpy as np
import patsy

class TestCentralComposite(TestCase):

    @classmethod
    def test_fcd(cls):
        """Tests a simple 2 factor face-centered central composite design."""
        ccd_data = build_ccd(2, 1.0)
        x_matrix = patsy.dmatrix(make_quadratic_model(ccd_data.columns),
                                 ccd_data,
                                 return_type="dataframe")

        answer_d = -8.55333223803211
        np.testing.assert_allclose(answer_d, det_xtxi(x_matrix), rtol=1e-5)

    @classmethod
    def test_rotatable(cls):
        """Tests a 5 factor rotatable central composite design."""
        ccd_data = build_ccd(5, "rotatable")
        x_matrix = patsy.dmatrix(make_quadratic_model(ccd_data.columns),
                                 ccd_data,
                                 return_type="dataframe")

        answer_d = -75.66313056042111
        np.testing.assert_allclose(answer_d, det_xtxi(x_matrix), rtol=1e-4)

    @classmethod
    def test_spherical(cls):
        """Tests a 3 factor spherical central composite design."""
        ccd_data = build_ccd(3, "spherical")
        x_matrix = patsy.dmatrix(make_quadratic_model(ccd_data.columns),
                                 ccd_data,
                                 return_type="dataframe")

        answer_d = -23.673909747960984
        np.testing.assert_allclose(answer_d, det_xtxi(x_matrix), rtol=1e-4)

    @classmethod
    def test_orthogonal(cls):
        """Tests a 7 factor orthogonal central composite design."""
        ccd_data = build_ccd(7, "orthogonal")
        x_matrix = patsy.dmatrix(make_quadratic_model(ccd_data.columns),
                                 ccd_data,
                                 return_type="dataframe")

        answer_d = -164.17093230784621
        np.testing.assert_allclose(answer_d, det_xtxi(x_matrix), rtol=1e-3)

    @classmethod
    def test_practical(cls):
        """Tests a 9 factor practical central composite design."""
        ccd_data = build_ccd(9, "practical")
        x_matrix = patsy.dmatrix(make_quadratic_model(ccd_data.columns),
                                 ccd_data,
                                 return_type="dataframe")

        answer_d = -314.58564545825374
        np.testing.assert_allclose(answer_d, det_xtxi(x_matrix), rtol=1e-3)

class TestAlphaParsing(TestCase):

    def test_custom(self):
        """Tests setting a custom alpha value."""
        self.assertAlmostEqual(1.31415, alpha_from_type(5, 1.31415))

    def test_rotatable(self):
        """Tests setting a rotatable alpha value."""
        self.assertAlmostEqual(1.4142135623731, alpha_from_type(2, "rotatable"))
        self.assertAlmostEqual(2.3784142300054, alpha_from_type(5, "rotatable"))
        self.assertAlmostEqual(4.7568284600109, alpha_from_type(9, "rotatable"))

    def test_spherical(self):
        """Tests setting a spherical alpha value."""
        self.assertAlmostEqual(1.4142135623731, alpha_from_type(2, "spherical"))
        self.assertAlmostEqual(2.2360679774998, alpha_from_type(5, "spherical"))
        self.assertAlmostEqual(3.0, alpha_from_type(9, "spherical"))

    def test_orthogonal(self):
        """Tests setting an orthogonal alpha value."""
        self.assertAlmostEqual(0.91017972112445, alpha_from_type(2, "orthogonal"))
        self.assertAlmostEqual(1.5265329278543, alpha_from_type(5, "orthogonal"))
        self.assertAlmostEqual(2.1121386170409, alpha_from_type(9, "orthogonal"))

    def test_practical(self):
        """Tests setting a practical alpha value."""
        self.assertAlmostEqual(1.1892071150027, alpha_from_type(2, "practical"))
        self.assertAlmostEqual(1.4953487812212, alpha_from_type(5, "practical"))
        self.assertAlmostEqual(1.7320508075689, alpha_from_type(9, "practical"))

    def test_face_centered(self):
        """Tests setting a face-centered alpha value."""
        self.assertAlmostEqual(1.0, alpha_from_type(2, "face centered"))
        self.assertAlmostEqual(1.0, alpha_from_type(5, "face centered"))
        self.assertAlmostEqual(1.0, alpha_from_type(9, "face centered"))

    def test_invalid(self):
        caught_exception = False
        try:
            alpha_from_type(45, "beeeeep")
        except ValueError:
            caught_exception = True
        self.assertTrue(caught_exception)

max_stars_repo_name: AI-Assistant/FEMAG-Python

max_issues_repo_name: AI-Assistant/FEMAG-Python

max_forks_repo_name: AI-Assistant/FEMAG-Python


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

with open("left.txt", "rb") as fo:
    left = pickle.load(fo)
with open("right.txt", "rb") as fo:
    right = pickle.load(fo)
left = [-1 for i in left]
right = [1 for i in right]
plt.hist(right,bins=1)
plt.show()


text:
import sys
import typing

import numpy as np
from scipy import signal

def mod_convolve(
  f: np.ndarray,
  g: np.ndarray,
  mod: int,
) -> np.ndarray:
  N = 15
  BASE = 1 << N
  f1, f0 = np.divmod(f, BASE)
  g1, g0 = np.divmod(g, BASE)
  a = signal.fftconvolve(f1, g1)
  c = signal.fftconvolve(f0, g0)
  a = np.rint(a).astype(np.int64) % mod
  c = np.rint(c).astype(np.int64) % mod
  b = signal.fftconvolve(f0 + f1, g0 + g1)
  b = np.rint(b).astype(np.int64) % mod
  b = (b - a - c) % mod
  h = (a << N * 2) + (b << N) + c
  return h % mod

mod = 998_244_353

def main() -> typing.NoReturn:
  n, m = map(int, input().split())
  a = np.array(
    sys.stdin.readline().split(),
    dtype=np.int64,
  )
  b = np.array(
    sys.stdin.readline().split(),
    dtype=np.int64,
  )
  print(*mod_convolve(a, b, mod))

main()

max_stars_repo_name: kagemeka/atcoder-submissions

max_issues_repo_name: kagemeka/atcoder-submissions

max_forks_repo_name: kagemeka/atcoder-submissions

reason: import numpy,from scipy


text:
import pickle
from data_extractor import get_data, get_processed_data
from viterbi_utils import generate_outfile
from viterbi_utils import evaluate_dev, read_predictions
from decoder import Trellis, TrellisMEMM
import numpy as np
import tageval
from tqdm import tqdm
# --------------------------------------------------------------------------
# Configuration
# --------------------------------------------------------------------------

# --------------------------------------------------------------------------
# Load Data
# --------------------------------------------------------------------------
# (trainX, trainY), train_tweets = get_data('data/train/train.txt', 'processed_data/train_data.pkl')
# (devX, devY), dev_tweets = get_data('data/dev/dev.txt', 'processed_data/dev_data.pkl')
# test_data = get_data('data/test/test.nolabels.txt', 'processed_data/test_data.pkl', isTest=True)

# (trainX, trainY), train_tweets = get_processed_data('processed_data/train_data.pkl', None, load=True)
# (devX, devY), dev_tweets = get_processed_data('processed_data/dev_data.pkl', None, load=True)

# --------------------------------------------------------------------------
# Train
# --------------------------------------------------------------------------
model = None

# --------------------------------------------------------------------------
# Decode
# --------------------------------------------------------------------------
dev_tweets = pickle.load(open('data/pos_pickle_file/dev_data_pos.p', 'rb'))
clf = pickle.load(open('processed_data/mlp_viterbi.pkl', 'rb'))
trel = TrellisMEMM(dev_tweets)
dev_tags = []
for tweet in tqdm(dev_tweets):
    tweet_pos_tokens = [tok[0:2] for tok in tweet]
    predicted_tags = trel.viterbi_memm(tweet_pos_tokens, clf)
    print(predicted_tags)
    dev_tags.append(predicted_tags)

f1 = tageval.evaluate_tagging('data/dev/dev.out', dev_tags)
print('F1: {}'.format(f1))

# --------------------------------------------------------------------------
# Evaluate
# --------------------------------------------------------------------------
# tag2idx = {'B': 0, 'I': 1, 'O': 2}
# dev_predictions = read_predictions('exp/2/dev_predictions.out')
# evaluate_dev(dev_tweets, dev_predictions)

max_stars_repo_name: sijunhe/twitter-ner

max_issues_repo_name: sijunhe/twitter-ner

max_forks_repo_name: sijunhe/twitter-ner


text:
from PIL import Image
import numpy as np
import pandas as pd
import os
import torch
from torch.utils.data.dataset import Dataset
import torchvision.transforms as transforms
import torchvision.transforms.functional as F
import random
import cv2
import albumentations as A
import json
import platform

from fire.dataaug_user import TrainDataAug, TestDataAug

##### Common
def getFileNames(file_dir, tail_list=['.png','.jpg','.JPG','.PNG']):
        L=[]
        for root, dirs, files in os.walk(file_dir):
            for file in files:
                if os.path.splitext(file)[1] in tail_list:
                    L.append(os.path.join(root, file))
        return L

######## dataloader

class TensorDatasetTrainClassify(Dataset):
    _print_times = 0
    def __init__(self, train_jpg, label_type, label_path, log_classname=True, transform=None):
        self.train_jpg = train_jpg
        self.label_type = label_type
        self.label_path = label_path
        self.transform = transform
        self.log_classname = log_classname

        self.label_dict = {}
        self.getLabels()
        self.cate_dirs = []

    def getLabels(self):

        if self.label_type == "DIR":
            self.cate_dirs = os.listdir(self.label_path)
            self.cate_dirs.sort()
            if TensorDatasetTrainClassify._print_times==0:
                if self.log_classname:
                    print("[INFO] Default classes names: ", self.cate_dirs)
                    TensorDatasetTrainClassify._print_times=1

            for i, img_path in enumerate(self.train_jpg):
                img_dirs = img_path.replace(self.label_path,'')
                img_dirs = img_dirs.split(os.sep)[:2]
                img_dir = img_dirs[0] if img_dirs[0] else img_dirs[1]

                y = self.cate_dirs.index(img_dir)
                self.label_dict[img_path] = y

        # User code here
        elif self.label_type == "CSV":
            df = pd.read_csv(self.label_path)
            dir_path = os.path.dirname(self.train_jpg[0])
            for index, row in df.iterrows():
                #print(row["image_id"], type(row["label"]))
                img_path = os.path.join(dir_path, row["image_id"])
                img_path = img_path.replace("\\","/")
                # print(img_path)
                # b
                self.label_dict[img_path] = row["label"]
                #b
            if TensorDatasetTrainClassify._print_times==0:
                print("[INFO] Labels count: ")
                print(df['label'].value_counts().sort_index())
                TensorDatasetTrainClassify._print_times=1

        else:
            raise Exception("[ERROR] In datatools.py getLabel() reimplement needed. ")

    def __getitem__(self, index):

        #img = Image.open(self.train_jpg[index]).convert('RGB')
        img = cv2.imread(self.train_jpg[index])

        if self.transform is not None:
            img = self.transform(img)

        y = self.label_dict[self.train_jpg[index]]

        # y_onehot = [0,0]
        # y_onehot[y] = 1

        return img, y, self.train_jpg[index]

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

class TensorDatasetTestClassify(Dataset):

    def __init__(self, train_jpg, transform=None):
        self.train_jpg = train_jpg
        if transform is not None:
            self.transform = transform
        else:
            self.transform = None

    def __getitem__(self, index):
        # img = cv2.imread(self.train_jpg[index])
        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        #img = cv2.resize(img,(180, 180))

        #img = Image.open(self.train_jpg[index]).convert('RGB')
        img = cv2.imread(self.train_jpg[index])
        #img = imgPaddingWrap(img)
        #b
        if self.transform is not None:
            img = self.transform(img)

        # path_dir = '/'.join(self.train_jpg[index].split('/')[:-1])
        # y = 0
        # if  'true' in path_dir:
        #     y = 1

        return img, self.train_jpg[index]

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

###### 3. get data loader

def getNormorlize(model_name):
    if model_name in ['mobilenetv2','mobilenetv3']:
        my_normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    elif model_name == 'xception':
        my_normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    elif "adv-eff" in model_name:
        my_normalize = transforms.Lambda(lambda img: img * 2.0 - 1.0)
    elif "resnex" in model_name or 'eff' in model_name or 'RegNet' in model_name:
        my_normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        #my_normalize = transforms.Normalize([0.4783, 0.4559, 0.4570], [0.2566, 0.2544, 0.2522])
    elif "EN-B" in model_name:
        my_normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    else:
        print("[Info] Not found normalize type! Use defalut imagenet normalize.")
        my_normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    return my_normalize

def getDataLoader(mode, input_data, cfg):

    my_normalize = getNormorlize(cfg['model_name'])

    data_aug_train = TrainDataAug(cfg['img_size'])
    data_aug_test = TestDataAug(cfg['img_size'])

    if mode=="trainval_onehot":
        my_dataloader = TensorDatasetTrainClassify

        train_loader = torch.utils.data.DataLoader(
                    my_dataloader(input_data[0],
                                transforms.Compose([
                                    data_aug_train,
                                    transforms.ToTensor(),
                                    my_normalize,
                                ])),
                        batch_size=cfg['batch_size'], shuffle=True, num_workers=cfg['num_workers'], pin_memory=cfg['pin_memory'])

        val_loader = torch.utils.data.DataLoader(
                    my_dataloader(input_data[1],
                                    transforms.Compose([
                                    data_aug_test,
                                    transforms.ToTensor(),
                                    my_normalize
                                ])),
                        batch_size=1, shuffle=False, num_workers=cfg['num_workers'], pin_memory=cfg['pin_memory'])
        return train_loader, val_loader

    elif mode=="test":
        my_dataloader = TensorDatasetTestClassify

        test_loader = torch.utils.data.DataLoader(
                my_dataloader(input_data[0],
                                transforms.Compose([
                                    data_aug_test,
                                    transforms.ToTensor(),
                                    my_normalize
                                ])
                ), batch_size=cfg['test_batch_size'], shuffle=False,
                num_workers=cfg['num_workers'], pin_memory=cfg['pin_memory']
            )

        return test_loader

    elif mode=="trainval":
        my_dataloader = TensorDatasetTrainClassify

        #auto aug

        #from .autoaugment import ImageNetPolicy
        # from libs.FastAutoAugment.data import  Augmentation
        # from libs.FastAutoAugment.archive import fa_resnet50_rimagenet
        if cfg['label_type'] == 'DIR':
            cfg['label_path'] = cfg['train_path']

        train_loader = torch.utils.data.DataLoader(
                                        my_dataloader(input_data[0],
                                            cfg['label_type'],
                                            cfg['label_path'],
                                            True,
                                            transforms.Compose([
                                                data_aug_train,
                                                #ImageNetPolicy(),  #autoaug
                                                #Augmentation(fa_resnet50_rimagenet()), #fastaa
                                                transforms.ToTensor(),
                                                my_normalize,
                                        ])),
                                batch_size=cfg['batch_size'], shuffle=True, num_workers=cfg['num_workers'], pin_memory=cfg['pin_memory'])

        if cfg['val_path']:
            cfg['label_path'] = cfg['val_path']

        val_loader = torch.utils.data.DataLoader(
                                        my_dataloader(input_data[1],
                                            cfg['label_type'],
                                            cfg['label_path'],
                                            True,
                                            transforms.Compose([
                                                data_aug_test,
                                                transforms.ToTensor(),
                                                my_normalize
                                        ])),
                                batch_size=cfg['batch_size'], shuffle=False, num_workers=cfg['num_workers'], pin_memory=cfg['pin_memory'])
        return train_loader, val_loader

    elif mode=="train":
        my_dataloader = TensorDatasetTrainClassify

        #auto aug

        #from .autoaugment import ImageNetPolicy
        # from libs.FastAutoAugment.data import  Augmentation
        # from libs.FastAutoAugment.archive import fa_resnet50_rimagenet
        if cfg['label_type'] == 'DIR':
            cfg['label_path'] = cfg['train_path']

        train_loader = torch.utils.data.DataLoader(
                                my_dataloader(input_data[0],
                                            cfg['label_type'],
                                            cfg['label_path'],
                                            False,
                                            transforms.Compose([
                                                data_aug_test,
                                                #ImageNetPolicy(),  #autoaug
                                                #Augmentation(fa_resnet50_rimagenet()), #fastaa
                                                transforms.ToTensor(),
                                                my_normalize,
                                                ])),
                                batch_size=cfg['batch_size'], shuffle=True, num_workers=cfg['num_workers'], pin_memory=cfg['pin_memory'])

        return train_loader

    elif mode=="val":
        my_dataloader = TensorDatasetTrainClassify

        #auto aug

        #from .autoaugment import ImageNetPolicy
        # from libs.FastAutoAugment.data import  Augmentation
        # from libs.FastAutoAugment.archive import fa_resnet50_rimagenet
        if cfg['label_type'] == 'DIR':
            cfg['label_path'] = cfg['val_path']

        data_loader = torch.utils.data.DataLoader(
                                my_dataloader(input_data[0],
                                            cfg['label_type'],
                                            cfg['label_path'],
                                            False,
                                            transforms.Compose([
                                                data_aug_test,
                                                #ImageNetPolicy(),  #autoaug
                                                #Augmentation(fa_resnet50_rimagenet()), #fastaa
                                                transforms.ToTensor(),
                                                my_normalize,
                                                ])),
                                batch_size=cfg['batch_size'], shuffle=True, num_workers=cfg['num_workers'], pin_memory=cfg['pin_memory'])

        return data_loader

    elif mode=="eval":
        my_dataloader = TensorDatasetTrainClassify

        #auto aug

        #from .autoaugment import ImageNetPolicy
        # from libs.FastAutoAugment.data import  Augmentation
        # from libs.FastAutoAugment.archive import fa_resnet50_rimagenet
        if cfg['label_type'] == 'DIR':
            cfg['label_path'] = cfg['eval_path']

        data_loader = torch.utils.data.DataLoader(
                                my_dataloader(input_data[0],
                                            cfg['label_type'],
                                            cfg['label_path'],
                                            False,
                                            transforms.Compose([
                                                data_aug_test,
                                                #ImageNetPolicy(),  #autoaug
                                                #Augmentation(fa_resnet50_rimagenet()), #fastaa
                                                transforms.ToTensor(),
                                                my_normalize,
                                                ])),
                                batch_size=cfg['batch_size'], shuffle=True, num_workers=cfg['num_workers'], pin_memory=cfg['pin_memory'])

        return data_loader


text:
import os
import cv2
import numpy as np
from plantcv.plantcv import print_image
from plantcv.plantcv import plot_image
from plantcv.plantcv import color_palette
from plantcv.plantcv import params
import sys
from datetime import datetime
from plantcv.plantcv import apply_mask

def cluster_contours(img, roi_objects, roi_obj_hierarchy, nrow=1, ncol=1, show_grid=False):

    """
    This function take a image with multiple contours and clusters them based on user input of rows and columns

    Inputs:
    img                     = RGB or grayscale image data for plotting
    roi_objects             = object contours in an image that are needed to be clustered.
    roi_obj_hierarchy       = object hierarchy
    nrow                    = number of rows to cluster (this should be the approximate  number of desired rows
                              in the entire image (even if there isn't a literal row of plants)
    ncol                    = number of columns to cluster (this should be the approximate number of desired columns
                              in the entire image (even if there isn't a literal row of plants)
    show_grid               = if True then the grid will get plot to show how plants are being clustered

    Returns:
    grouped_contour_indexes = contours grouped
    contours                = All inputed contours

    :param img: numpy.ndarray
    :param roi_objects: list
    :param roi_obj_hierarchy: numpy.ndarray
    :param nrow: int
    :param ncol: int
    :param show_grid: bool
    :return grouped_contour_indexes: list
    :return contours: list
    :return roi_obj_hierarchy: list
    """

    params.device += 1

    if len(np.shape(img)) == 3:
        iy, ix, iz = np.shape(img)
    else:
        iy, ix, = np.shape(img)

    # get the break groups
    if nrow == 1:
        rbreaks = [0, iy]
    else:
        rstep = np.rint(iy / nrow)
        rstep1 = np.int(rstep)
        rbreaks = range(0, iy, rstep1)
    if ncol == 1:
        cbreaks = [0, ix]
    else:
        cstep = np.rint(ix / ncol)
        cstep1 = np.int(cstep)
        cbreaks = range(0, ix, cstep1)

    # categorize what bin the center of mass of each contour
    def digitize(a, step):
        # The way cbreaks and rbreaks are calculated, step will never be an integer
        # if isinstance(step, int):
        #     i = step
        # else:
        num_bins = len(step)
        for x in range(0, num_bins):
            if x == 0:
                if a >= 0 and a < step[1]:
                    return 1
            else:
                if a >= step[x - 1] and a < step[x]:
                    return x
                elif a >= np.max(step):
                    return num_bins

    dtype = [('cx', int), ('cy', int), ('rowbin', int), ('colbin', int), ('index', int)]
    coord = []
    for i in range(0, len(roi_objects)):
        m = cv2.moments(roi_objects[i])
        if m['m00'] == 0:
            pass
        else:
            cx = int(m['m10'] / m['m00'])
            cy = int(m['m01'] / m['m00'])
            colbin = digitize(cx, cbreaks)
            rowbin = digitize(cy, rbreaks)
            a = (cx, cy, colbin, rowbin, i)
            coord.append(a)
    coord1 = np.array(coord, dtype=dtype)
    coord2 = np.sort(coord1, order=('colbin', 'rowbin'))

    # get the list of unique coordinates and group the contours with the same bin coordinates
    groups = []
    for i, y in enumerate(coord2):
        col = y[3]
        row = y[2]
        location = str(row) + ',' + str(col)
        groups.append(location)

    #unigroup = np.unique(groups)
    unigroup = ["1,1", "2,1", "3,1", "4,1", "5,1", "6,1"]
    coordgroups = []

    for i, y in enumerate(unigroup):
        col_row = y.split(',')
        col = int(col_row[0])
        row = int(col_row[1])
        for a, b in enumerate(coord2):
            if b[2] == col and b[3] == row:
                grp = i
                contour = b[4]
                coordgroups.append((grp, contour))
            else:
                pass
        if len(coordgroups) == 0:
            coordgroups.append((i, None))
        else:
            check = coordgroups[-1]
            if i != check[0]:
                coordgroups.append((i, None))

    coordlist = [[y[1] for y in coordgroups if y[0] == x] for x in range(0, (len(unigroup)))]

    contours = roi_objects
    grouped_contour_indexes = coordlist

    # Debug image is rainbow printed contours

    if params.debug is not None:
        if len(np.shape(img)) == 3:
            img_copy = np.copy(img)
        else:
            iy, ix = np.shape(img)
            img_copy = np.zeros((iy, ix, 3), dtype=np.uint8)

        rand_color = color_palette(len(coordlist))
        for i, x in enumerate(coordlist):
            for a in x:
                if a != None:
                    if roi_obj_hierarchy[0][a][3] > -1:
                        pass
                    else:
                        cv2.drawContours(img_copy, roi_objects, a, rand_color[i], -1, hierarchy=roi_obj_hierarchy)
        if show_grid:
            for y in rbreaks:
                cv2.line(img_copy, (0, y), (ix, y), (255, 0, 0), params.line_thickness)
            for x in cbreaks:
                cv2.line(img_copy, (x, 0), (x, iy), (255, 0, 0), params.line_thickness)
        if params.debug == 'print':
            print_image(img_copy, os.path.join(params.debug_outdir, str(params.device) + '_clusters.png'))
        elif params.debug == 'plot':
            plot_image(img_copy)

    return grouped_contour_indexes, contours, roi_obj_hierarchy

def cluster_contour_splitimg(rgb_img, grouped_contour_indexes, contours, hierarchy, outdir=None, file=None,
                             filenames=None):

    """
    This function takes clustered contours and splits them into multiple images, also does a check to make sure that
    the number of inputted filenames matches the number of clustered contours.

    Inputs:
    rgb_img                 = RGB image data
    grouped_contour_indexes = output of cluster_contours, indexes of clusters of contours
    contours                = contours to cluster, output of cluster_contours
    hierarchy               = hierarchy of contours, output of find_objects
    outdir                  = out directory for output images
    file                    = the name of the input image to use as a plantcv name,
                              output of filename from read_image function
    filenames               = input txt file with list of filenames in order from top to bottom left to right
                              (likely list of genotypes)

    Returns:
    output_path             = array of paths to output images

    :param rgb_img: numpy.ndarray
    :param grouped_contour_indexes: list
    :param contours: list
    :param hierarchy: numpy.ndarray
    :param outdir: str
    :param file: str
    :param filenames: str
    :return output_path: str
    """

    params.device += 1

    sys.stderr.write(
        'This function has been updated to include object hierarchy so object holes can be included\n')

    i = datetime.now()
    timenow = i.strftime('%m-%d-%Y_%H:%M:%S')

    if file is None:
        filebase = timenow
    else:
        filebase = os.path.splitext(file)[0]

    if filenames is None:
        namelist = []
        for x in range(0, len(grouped_contour_indexes)):
            namelist.append(x)
    else:
        with open(filenames, 'r') as n:
            namelist = n.read().splitlines()
        n.close()

    # make sure the number of objects matches the namelist, and if not, remove the smallest grouped countor
    # removing contours is not ideal but the lists don't match there is a warning to check output
    if len(namelist) == len(grouped_contour_indexes):
        corrected_contour_indexes = grouped_contour_indexes
    elif len(namelist) < len(grouped_contour_indexes):
        print("Warning number of names is less than number of grouped contours, attempting to fix, to double check "
              "output")
        diff = len(grouped_contour_indexes) - len(namelist)
        size = []
        for i, x in enumerate(grouped_contour_indexes):
            totallen = []
            for a in x:
                g = i
                la = len(contours[a])
                totallen.append(la)
            sumlen = np.sum(totallen)
            size.append((sumlen, g, i))

        dtype = [('len', int), ('group', list), ('index', int)]
        lencontour = np.array(size, dtype=dtype)
        lencontour = np.sort(lencontour, order='len')

        rm_contour = lencontour[diff:]
        rm_contour = np.sort(rm_contour, order='group')
        corrected_contour_indexes = []

        for x in rm_contour:
            index = x[2]
            corrected_contour_indexes.append(grouped_contour_indexes[index])

    elif len(namelist) > len(grouped_contour_indexes):
        print("Warning number of names is more than number of  grouped contours, double check output")
        diff = len(namelist) - len(grouped_contour_indexes)
        namelist = namelist[0:-diff]
        corrected_contour_indexes = grouped_contour_indexes

    # create filenames
    group_names = []
    group_names1 = []
    for i, x in enumerate(namelist):
        plantname = str(filebase) + '_' + str(x) + '_p' + str(i) + '.png'
        maskname = str(filebase) + '_' + str(x) + '_p' + str(i) + '_mask.png'
        group_names.append(plantname)
        group_names1.append(maskname)

    # split image
    output_path = []
    output_imgs = []
    output_masks = []

    for y, x in enumerate(corrected_contour_indexes):
        if outdir is not None:
            savename = os.path.join(str(outdir), group_names[y])
            savename1 = os.path.join(str(outdir), group_names1[y])
        else:
            savename = os.path.join(".", group_names[y])
            savename1 = os.path.join(".", group_names1[y])
        iy, ix, iz = np.shape(rgb_img)
        mask = np.zeros((iy, ix, 3), dtype=np.uint8)
        masked_img = np.copy(rgb_img)
        for a in x:
            if a != None:
                if hierarchy[0][a][3] > -1:
                    cv2.drawContours(mask, contours, a, (0, 0, 0), -1, lineType=8, hierarchy=hierarchy)
                else:
                    cv2.drawContours(mask, contours, a, (255, 255, 255), -1, lineType=8, hierarchy=hierarchy)

        mask_binary = mask[:, :, 0]

        if np.sum(mask_binary) == 0:
            pass
        else:
            retval, mask_binary = cv2.threshold(mask_binary, 254, 255, cv2.THRESH_BINARY)
            masked1 = apply_mask(masked_img, mask_binary, 'white')
            output_imgs.append(masked1)
            output_masks.append(mask_binary)
            if outdir is not None:
                print_image(masked1, savename)
                print_image(mask_binary, savename1)
            output_path.append(savename)

            if params.debug == 'print':
                print_image(masked1, os.path.join(params.debug_outdir, str(params.device) + '_clusters.png'))
                print_image(mask_binary, os.path.join(params.debug_outdir, str(params.device) + '_clusters_mask.png'))
            elif params.debug == 'plot':
                plot_image(masked1)
                plot_image(mask_binary, cmap='gray')

    return output_path, output_imgs, output_masks

max_stars_repo_name: massivejords/Agar-plate-leaf-area

max_issues_repo_name: massivejords/Agar-plate-leaf-area

max_forks_repo_name: massivejords/Agar-plate-leaf-area