text:
# see a2/save_movies_overlay.py
from __future__ import division
from __future__ import print_function
import numpy
from numpy import nan, pi
import tables as PT
import tables.flavor

tables.flavor.restrict_flavors(keep=["numpy"])  # ensure pytables 2.x
import pytz  # from http://pytz.sourceforge.net/
import datetime
import sys
from optparse import OptionParser
import pylab
import flydra_core.reconstruct
import flydra_analysis.analysis.result_utils as result_utils
import matplotlib.colors

def auto_subplot(fig, n, n_rows=2, n_cols=3):
    # 2 rows and n_cols

    rrow = n // n_cols  # reverse row
    row = n_rows - rrow - 1  # row number
    col = n % n_cols

    x_space = 0.02 / n_cols
    # y_space = 0.0125
    y_space = 0.03
    y_size = (1.0 / n_rows) - (2 * y_space)

    left = col * (1.0 / n_cols) + x_space
    bottom = row * y_size + y_space
    w = (1.0 / n_cols) - x_space
    h = y_size - 2 * y_space
    return fig.add_axes([left, bottom, w, h])

def show_it(
    fig, filename, frame_start=None, frame_stop=None,
):

    results = result_utils.get_results(filename, mode="r")
    reconstructor = flydra_core.reconstruct.Reconstructor(results)
    camn2cam_id, cam_id2camns = result_utils.get_caminfo_dicts(results)
    data2d = results.root.data2d_distorted  # make sure we have 2d data table

    print("reading frames...")
    frames = data2d.read(field="frame")
    print("OK")

    if frame_start is not None:
        print("selecting frames after start")
        after_start = numpy.nonzero(frames >= frame_start)[0]
    else:
        after_start = None

    if frame_stop is not None:
        print("selecting frames before stop")
        before_stop = numpy.nonzero(frames <= frame_stop)[0]
    else:
        before_stop = None

    print("finding all frames")
    if after_start is not None and before_stop is not None:
        use_idxs = numpy.intersect1d(after_start, before_stop)
    elif after_start is not None:
        use_idxs = after_start
    elif before_stop is not None:
        use_idxs = before_stop
    else:
        use_idxs = numpy.arange(data2d.nrows)

    # OK, we have data coords, plot
    print("reading cameras")
    frames = frames[use_idxs]
    camns = data2d.read(field="camn")
    camns = camns[use_idxs]
    unique_camns = numpy.unique1d(camns)
    unique_cam_ids = list(set([camn2cam_id[camn] for camn in unique_camns]))
    unique_cam_ids.sort()
    print("%d cameras with data" % (len(unique_cam_ids),))

    if len(unique_cam_ids) == 1:
        n_rows = 1
        n_cols = 1
    else:
        n_rows = 2
        n_cols = 3

    subplot_by_cam_id = {}
    for i, cam_id in enumerate(unique_cam_ids):
        ax = auto_subplot(fig, i, n_rows=n_rows, n_cols=n_cols)
        ax.text(
            0.5,
            0.95,
            "%s: %s" % (cam_id, str(cam_id2camns[cam_id])),
            horizontalalignment="center",
            verticalalignment="top",
            transform=ax.transAxes,
        )
        ##        ax.set_xticks([])
        ##        ax.set_yticks([])
        subplot_by_cam_id[cam_id] = ax

    cdict = {
        "red": ((0, 0, 0), (1, 1, 1)),
        "green": ((0, 0, 0), (0, 0, 0)),
        "blue": ((0, 0, 0), (1, 1, 1)),
    }
    #    black2magenta = matplotlib.colors.LinearSegmentedColormap('black2magenta',cdict,256)

    for camn in unique_camns:
        cam_id = camn2cam_id[camn]
        ax = subplot_by_cam_id[cam_id]
        this_camn_idxs = use_idxs[camns == camn]

        w, h = reconstructor.get_resolution(cam_id)
        xbins = numpy.linspace(0, w, 40)
        ybins = numpy.linspace(0, h, 30)

        xs = data2d.read_coordinates(this_camn_idxs, field="x")
        ys = data2d.read_coordinates(this_camn_idxs, field="y")

        hist, xedges, yedges = numpy.histogram2d(xs, ys, bins=(xbins, ybins))
        hist = numpy.ma.masked_where(hist == 0, hist)
        im = getattr(results.root.images, cam_id)

        ax.imshow(im, origin="lower", extent=[0, w - 1, 0, h - 1], cmap=pylab.cm.gray)
        pcolor_im = ax.pcolor(
            xbins, ybins, hist.T, alpha=0.5,
        )  # cmap = black2magenta )
        # fig.colorbar( pcolor_im, cax )

        res = reconstructor.get_resolution(cam_id)
        ax.set_xlim([0, res[0]])
        # ax.set_ylim([0,res[1]])
        ax.set_ylim([res[1], 0])

    if 1:
        return

    # Old code below:
    # Do same as above for Kalman-filtered data

    if kalman_filename is None:
        return

    kresults = PT.open_file(kalman_filename, mode="r")
    kobs = kresults.root.ML_estimates
    kframes = kobs.read(field="frame")
    if frame_start is not None:
        k_after_start = numpy.nonzero(kframes >= frame_start)[0]
        # k_after_start = kobs.read_coordinates(idxs)
        # k_after_start = kobs.get_where_list(
        #    'frame>=frame_start')
    else:
        k_after_start = None
    if frame_stop is not None:
        k_before_stop = numpy.nonzero(kframes <= frame_stop)[0]
        # k_before_stop = kobs.read_coordinates(idxs)
        # k_before_stop = kobs.get_where_list(
        #    'frame<=frame_stop')
    else:
        k_before_stop = None

    if k_after_start is not None and k_before_stop is not None:
        k_use_idxs = numpy.intersect1d(k_after_start, k_before_stop)
    elif k_after_start is not None:
        k_use_idxs = k_after_start
    elif k_before_stop is not None:
        k_use_idxs = k_before_stop
    else:
        k_use_idxs = numpy.arange(kobs.nrows)

    obj_ids = kobs.read(field="obj_id")[k_use_idxs]
    # obj_ids = kobs.read_coordinates( k_use_idxs,
    #                                field='obj_id')
    obs_2d_idxs = kobs.read(field="obs_2d_idx")[k_use_idxs]
    # obs_2d_idxs = kobs.read_coordinates( k_use_idxs,
    #                                    field='obs_2d_idx')
    kframes = kframes[k_use_idxs]  # kobs.read_coordinates( k_use_idxs,
    # field='frame')

    kobs_2d = kresults.root.ML_estimates_2d_idxs
    xys_by_obj_id = {}
    for obj_id, kframe, obs_2d_idx in zip(obj_ids, kframes, obs_2d_idxs):
        obs_2d_idx_find = int(obs_2d_idx)  # XXX grr, why can't pytables do this?
        obj_id_save = int(obj_id)  # convert from possible numpy scalar
        xys_by_cam_id = xys_by_obj_id.setdefault(obj_id_save, {})
        kobs_2d_data = kobs_2d.read(start=obs_2d_idx_find, stop=obs_2d_idx_find + 1)
        assert len(kobs_2d_data) == 1
        kobs_2d_data = kobs_2d_data[0]
        this_camns = kobs_2d_data[0::2]
        this_camn_idxs = kobs_2d_data[1::2]

        this_use_idxs = use_idxs[frames == kframe]
        if debugADS:
            print()
            print(kframe, "===============")
            print("this_use_idxs", this_use_idxs)

        d2d = data2d.read_coordinates(this_use_idxs)
        if debugADS:
            print("d2d ---------------")
            for row in d2d:
                print(row)
        for this_camn, this_camn_idx in zip(this_camns, this_camn_idxs):
            this_idxs_tmp = numpy.nonzero(d2d["camn"] == this_camn)[0]
            this_camn_d2d = d2d[d2d["camn"] == this_camn]
            found = False
            for this_row in this_camn_d2d:  # XXX could be sped up
                if this_row["frame_pt_idx"] == this_camn_idx:
                    found = True
                    break
            if not found:
                if 0:
                    print("WARNING:point not found in data!?")
                    continue
                else:
                    raise RuntimeError("point not found in data!?")
            # this_row = this_camn_d2d[this_camn_idx]
            this_cam_id = camn2cam_id[this_camn]
            xys = xys_by_cam_id.setdefault(this_cam_id, ([], []))
            xys[0].append(this_row["x"])
            xys[1].append(this_row["y"])

    for obj_id in xys_by_obj_id:
        xys_by_cam_id = xys_by_obj_id[obj_id]
        for cam_id, (xs, ys) in xys_by_cam_id.items():
            ax = subplot_by_cam_id[cam_id]
            if 0:
                ax.plot(xs, ys, label="obs: %d" % obj_id)
            else:
                ax.plot(xs, ys, "x-", label="obs: %d" % obj_id)
            ax.text(xs[0], ys[0], "%d:" % (obj_id,))
            ax.text(xs[-1], ys[-1], ":%d" % (obj_id,))

    for cam_id in subplot_by_cam_id.keys():
        ax = subplot_by_cam_id[cam_id]
        ax.legend()
    print("note: could/should also plot re-projection of Kalman filtered/smoothed data")

def main():
    usage = "%prog FILE [options]"

    parser = OptionParser(usage)

    parser.add_option(
        "-f",
        "--file",
        dest="filename",
        type="string",
        help="hdf5 file with data to display FILE",
        metavar="FILE",
    )

    ## parser.add_option('-k', "--kalman-file", dest="kalman_filename", type='string',
    ##                   help="hdf5 file with kalman data to display KALMANFILE",
    ##                   metavar="KALMANFILE")

    parser.add_option(
        "--start", type="int", help="first frame to plot", metavar="START"
    )

    parser.add_option("--stop", type="int", help="last frame to plot", metavar="STOP")

    (options, args) = parser.parse_args()

    if options.filename is not None:
        args.append(options.filename)

    if len(args) > 1:
        print("arguments interpreted as FILE supplied more than once", file=sys.stderr)
        parser.print_help()
        return

    if len(args) < 1:
        parser.print_help()
        return

    h5_filename = args[0]

    fig = pylab.figure()
    show_it(
        fig,
        h5_filename,
        ## kalman_filename = options.kalman_filename,
        frame_start=options.start,
        frame_stop=options.stop,
    )
    pylab.show()

if __name__ == "__main__":
    main()

max_stars_repo_name: elhananby/flydra

max_issues_repo_name: elhananby/flydra

max_forks_repo_name: elhananby/flydra

reason: import numpy,from numpy


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

'''
CASSCF analytical nuclear gradients

Ref.
J. Comput. Chem., 5, 589
'''

import time
from functools import reduce
import numpy
from pyscf import lib
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.grad import rhf as rhf_grad
from pyscf.grad.mp2 import _shell_prange

def kernel(mc, mo_coeff=None, ci=None, atmlst=None, mf_grad=None,
           verbose=None):
    if mo_coeff is None: mo_coeff = mc.mo_coeff
    if ci is None: ci = mc.ci
    if mf_grad is None: mf_grad = mc._scf.nuc_grad_method()
    if mc.frozen is not None:
        raise NotImplementedError

    mol = mc.mol
    ncore = mc.ncore
    ncas = mc.ncas
    nocc = ncore + ncas
    nelecas = mc.nelecas
    nao, nmo = mo_coeff.shape
    nao_pair = nao * (nao+1) // 2

    mo_occ = mo_coeff[:,:nocc]
    mo_core = mo_coeff[:,:ncore]
    mo_cas = mo_coeff[:,ncore:nocc]

    casdm1, casdm2 = mc.fcisolver.make_rdm12(ci, ncas, nelecas)

# gfock = Generalized Fock, Adv. Chem. Phys., 69, 63
    dm_core = numpy.dot(mo_core, mo_core.T) * 2
    dm_cas = reduce(numpy.dot, (mo_cas, casdm1, mo_cas.T))
    aapa = ao2mo.kernel(mol, (mo_cas, mo_cas, mo_occ, mo_cas), compact=False)
    aapa = aapa.reshape(ncas,ncas,nocc,ncas)
    vj, vk = mc._scf.get_jk(mol, (dm_core, dm_cas))
    h1 = mc.get_hcore()
    vhf_c = vj[0] - vk[0] * .5
    vhf_a = vj[1] - vk[1] * .5
    gfock = reduce(numpy.dot, (mo_occ.T, h1 + vhf_c + vhf_a, mo_occ)) * 2
    gfock[:,ncore:nocc] = reduce(numpy.dot, (mo_occ.T, h1 + vhf_c, mo_cas, casdm1))
    gfock[:,ncore:nocc] += numpy.einsum('uviw,vuwt->it', aapa, casdm2)
    dme0 = reduce(numpy.dot, (mo_occ, (gfock+gfock.T)*.5, mo_occ.T))
    aapa = vj = vk = vhf_c = vhf_a = h1 = gfock = None

    dm1 = dm_core + dm_cas
    vhf1c, vhf1a = mf_grad.get_veff(mol, (dm_core, dm_cas))
    hcore_deriv = mf_grad.hcore_generator(mol)
    s1 = mf_grad.get_ovlp(mol)

    diag_idx = numpy.arange(nao)
    diag_idx = diag_idx * (diag_idx+1) // 2 + diag_idx
    casdm2_cc = casdm2 + casdm2.transpose(0,1,3,2)
    dm2buf = ao2mo._ao2mo.nr_e2(casdm2_cc.reshape(ncas**2,ncas**2), mo_cas.T,
                                (0, nao, 0, nao)).reshape(ncas**2,nao,nao)
    dm2buf = lib.pack_tril(dm2buf)
    dm2buf[:,diag_idx] *= .5
    dm2buf = dm2buf.reshape(ncas,ncas,nao_pair)
    casdm2 = casdm2_cc = None

    if atmlst is None:
        atmlst = range(mol.natm)
    aoslices = mol.aoslice_by_atom()
    de = numpy.zeros((len(atmlst),3))

    max_memory = mc.max_memory - lib.current_memory()[0]
    blksize = int(max_memory*.9e6/8 / ((aoslices[:,3]-aoslices[:,2]).max()*nao_pair))
    blksize = min(nao, max(2, blksize))

    for k, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = aoslices[ia]
        h1ao = hcore_deriv(ia)
        de[k] += numpy.einsum('xij,ij->x', h1ao, dm1)
        de[k] -= numpy.einsum('xij,ij->x', s1[:,p0:p1], dme0[p0:p1]) * 2

        q1 = 0
        for b0, b1, nf in _shell_prange(mol, 0, mol.nbas, blksize):
            q0, q1 = q1, q1 + nf
            dm2_ao = lib.einsum('ijw,pi,qj->pqw', dm2buf, mo_cas[p0:p1], mo_cas[q0:q1])
            shls_slice = (shl0,shl1,b0,b1,0,mol.nbas,0,mol.nbas)
            eri1 = mol.intor('int2e_ip1', comp=3, aosym='s2kl',
                             shls_slice=shls_slice).reshape(3,p1-p0,nf,nao_pair)
            de[k] -= numpy.einsum('xijw,ijw->x', eri1, dm2_ao) * 2
            eri1 = None
        de[k] += numpy.einsum('xij,ij->x', vhf1c[:,p0:p1], dm1[p0:p1]) * 2
        de[k] += numpy.einsum('xij,ij->x', vhf1a[:,p0:p1], dm_core[p0:p1]) * 2

    de += mf_grad.grad_nuc(mol, atmlst)
    return de

def as_scanner(mcscf_grad, envs=None):
    '''Generating a nuclear gradients scanner/solver (for geometry optimizer).

    The returned solver is a function. This function requires one argument
    "mol" as input and returns energy and first order nuclear derivatives.

    The solver will automatically use the results of last calculation as the
    initial guess of the new calculation.  All parameters assigned in the
    nuc-grad object and SCF object (DIIS, conv_tol, max_memory etc) are
    automatically applied in the solver.

    Note scanner has side effects.  It may change many underlying objects
    (_scf, with_df, with_x2c, ...) during calculation.

    Examples:

    >>> from pyscf import gto, scf, mcscf
    >>> mol = gto.M(atom='N 0 0 0; N 0 0 1.1', verbose=0)
    >>> mc_grad_scanner = mcscf.CASSCF(scf.RHF(mol), 4, 4).nuc_grad_method().as_scanner()
    >>> etot, grad = mc_grad_scanner(gto.M(atom='N 0 0 0; N 0 0 1.1'))
    >>> etot, grad = mc_grad_scanner(gto.M(atom='N 0 0 0; N 0 0 1.5'))
    '''
    from pyscf import gto
    if isinstance(mcscf_grad, lib.GradScanner):
        return mcscf_grad

    logger.info(mcscf_grad, 'Create scanner for %s', mcscf_grad.__class__)

    class CASSCF_GradScanner(mcscf_grad.__class__, lib.GradScanner):
        def __init__(self, g):
            lib.GradScanner.__init__(self, g)
        def __call__(self, mol_or_geom, **kwargs):
            if isinstance(mol_or_geom, gto.Mole):
                mol = mol_or_geom
            else:
                mol = self.mol.set_geom_(mol_or_geom, inplace=False)

            #print "we are here 111", envs['cycle']
            mc_scanner = self.base
            e_tot = mc_scanner(mol)
            #print "aaa", mc_scanner
            self.mol = mol
            de = self.kernel(**kwargs)
            return e_tot, de
    return CASSCF_GradScanner(mcscf_grad)

class Gradients(lib.StreamObject):
    '''Non-relativistic restricted Hartree-Fock gradients'''
    def __init__(self, mc):
        self.base = mc
        self.mol = mc.mol
        self.stdout = mc.stdout
        self.verbose = mc.verbose
        self.max_memory = mc.max_memory
        self.atmlst = None
        self.de = None
        self._keys = set(self.__dict__.keys())

    def dump_flags(self):
        log = logger.Logger(self.stdout, self.verbose)
        log.info('\n')
        if not self.base.converged:
            log.warn('Ground state CASSCF not converged')
        log.info('******** %s for %s ********',
                 self.__class__, self.base.__class__)
        log.info('max_memory %d MB (current use %d MB)',
                 self.max_memory, lib.current_memory()[0])
        return self

    def kernel(self, mo_coeff=None, ci=None, atmlst=None, mf_grad=None,
               verbose=None, envs=None):
        cput0 = (time.clock(), time.time())
        log = logger.new_logger(self, verbose)
        if atmlst is None:
            atmlst = self.atmlst
        else:
            self.atmlst = atmlst

        if self.verbose >= logger.WARN:
            self.check_sanity()
        if self.verbose >= logger.INFO:
            self.dump_flags()

        if self.verbose >= logger.WARN:
            self.check_sanity()
        if self.verbose >= logger.INFO:
            self.dump_flags()

        self.de = kernel(self.base, mo_coeff, ci, atmlst, mf_grad, log)
        log.timer('CASSCF gradients', *cput0)
        self._finalize()
        return self.de

    def _finalize(self):
        if self.verbose >= logger.NOTE:
            logger.note(self, '--------------- %s gradients ---------------',
                        self.base.__class__.__name__)
            rhf_grad._write(self, self.mol, self.de, self.atmlst)
            logger.note(self, '----------------------------------------------')

    as_scanner = as_scanner

Grad = Gradients

if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import mcscf

    mol = gto.Mole()
    mol.atom = 'N 0 0 0; N 0 0 1.2; H 1 1 0; H 1 1 1.2'
    mol.basis = '631g'
    mol.build()
    mf = scf.RHF(mol).run()
    mc = mcscf.CASSCF(mf, 4, 4).run()
    de = Grad(mc).kernel()
    print(lib.finger(de) - 0.019602220578635747)

    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = 'N 0 0 0; N 0 0 1.2'
    mol.basis = 'sto3g'
    mol.build()
    mf = scf.RHF(mol).run()
    mc = mcscf.CASSCF(mf, 4, 4).run()
    de = kernel(mc)

    mcs = mc.as_scanner()
    mol.set_geom_('N 0 0 0; N 0 0 1.201')
    e1 = mcs(mol)
    mol.set_geom_('N 0 0 0; N 0 0 1.199')
    e2 = mcs(mol)
    print(de[1,2], (e1-e2)/0.002*lib.param.BOHR)

ext: py

max_stars_repo_name: LeonOtis/pyscf

max_issues_repo_name: LeonOtis/pyscf

max_forks_repo_name: LeonOtis/pyscf

reason: import numpy


text:
import jax.numpy as jnp

from nmax import Parameter, Module
from cvax.utils import config

class BatchNorm2d(Module):

    gamma: Parameter
    beta: Parameter

    def __init__(self, n_channels):
        self.gamma = jnp.ones((1, n_channels, 1, 1))
        self.beta = jnp.zeros((1, n_channels, 1, 1))

        self.momentum = 0.1

        self.mean = jnp.zeros((1, n_channels, 1, 1))
        self.var = jnp.ones((1, n_channels, 1, 1))

    def forward(self, x):
        if self._mode == 'eval':
            x_mean = self.mean
            x_var = self.var
        elif self._mode == 'train':
            x_mean = jnp.mean(x, axis=(0, 2, 3), keepdims=True)
            x_var = jnp.var(x, axis=(0, 2, 3), keepdims=True)

            # Update population statistics
            self.mean = self.momentum * x_mean + (1 - self.momentum) * self.mean
            self.var = self.momentum * x_var + (1 - self.momentum) * self.var

        x_normalised = (x - x_mean) / jnp.sqrt(x_var + config.EPSILON)

        return self.gamma * x_normalised + self.beta

reason: import jax


text:
import xarray
import geo
import defopt
import numpy
import re
from horizgrid import HorizGrid
import mint
# to read the target points from file
from latlonreader import LatLonReader
import timeobj

EARTH_RADIUS = 6371000.0

class Field(object):

    def __init__(self, tFile, uFile, vFile, lonLatZPoints, sverdrup=False):

        self.sverdrup = sverdrup

        # read the cell bounds
        with xarray.open_dataset(tFile) as nc:
            bounds_lat = nc['bounds_lat'][:]
            bounds_lon = nc['bounds_lon'][:]
            self.bounds_depth = nc['deptht_bounds'][:]

        self.lonmin = bounds_lon.min()
        self.lonmax = bounds_lon.max()
        self.latmin = bounds_lat.min()
        self.latmax = bounds_lat.max()
        print(f'lon-lat box: {self.lonmin}, {self.latmin} -> {self.lonmax}, {self.latmax}')

        self.timeIndex = 0
        self.ncU = xarray.open_dataset(uFile)
        self.ncV = xarray.open_dataset(vFile)

        # read the time values
        self.timeObj = timeobj.TimeObj(self.ncU)

        self.nt, self.nz, self.ny, self.nx = self.getSizes()

        self.gr = HorizGrid(tFile)
        self.plis = []
        pli = mint.PolylineIntegral()
        pli.setGrid(self.gr.getMintGrid())
        pli.buildLocator(numCellsPerBucket=128, periodX=360., enableFolding=False)
        for lonlatpts in lonLatZPoints:
            pli.computeWeights(numpy.array(lonlatpts), counterclock=False)
            self.plis.append(pli)

        self.thickness = self.bounds_depth[:, 1] - self.bounds_depth[:, 0]

        numCells = self.ny * self.nx
        self.dx = min((self.lonmax - self.lonmin)/float(self.nx), (self.latmax - self.latmin)/float(self.ny))
        self.arcLengths = numpy.zeros((numCells, 4), numpy.float64)
        self.computeArcLengths()

        # compute the edge fluxes from the vector fields
        self.edgeFluxesUArray = numpy.zeros((numCells,), numpy.float64)
        self.edgeFluxesVArray = numpy.zeros((numCells,), numpy.float64)
        self.integratedVelocity = numpy.zeros((numCells, 4), numpy.float64)
        self.maxAbsFlux = 0.

        # read/get the staggered field integrated over the depth
        uVerticallyIntegrated, vVerticallyIntegrated = self.getUV()
        self.computeIntegratedFlux(uVerticallyIntegrated, vVerticallyIntegrated)
        print(f'max vertically integrated edge |flux|: {self.maxAbsFlux}')

        self.buildEdgeUVGrids(bounds_lon, bounds_lat)

        # create a polyline grid along the target points and place vectors on them
        vectorPoints = []
        self.uVectors = []
        for lonlatpts in lonLatZPoints:
            for i in range(len(lonlatpts) - 1):
                begPoint = numpy.array(lonlatpts[i])
                endPoint = numpy.array(lonlatpts[i + 1])
                u = endPoint - begPoint
                distance = numpy.sqrt(u.dot(u))
                # normalize
                u /= distance
                nvpts = max(2, int(distance / self.dx))
                vdx = distance / float(nvpts - 1)
                for j in range(nvpts):
                    vectorPoints.append(begPoint + u*j*vdx)
                    self.uVectors.append(u)
            self.vectorPoints = numpy.array(vectorPoints)

        # compute the vector at the target line
        self.vinterp = mint.VectorInterp()
        self.vinterp.setGrid(self.gr.getMintGrid())
        self.vinterp.buildLocator(numCellsPerBucket=128, periodX=360.)
        self.vinterp.findPoints(self.vectorPoints, tol2=1.e-12)
        self.vectorValues = self.vinterp.getFaceVectors(
               self.integratedVelocity, placement=0)

    def getFluxText(self):

        txt = ""
        for pli in self.plis:
            txt += f"{pli.getIntegral(self.integratedVelocity):4.3g}, "
        if self.sverdrup:
            txt += "(Sv) "
        else:
            txt += "(A m^2/s) "
        txt = re.sub(r',\s*\(', ' (', txt)
        return txt

    def update(self):

        # update the data
        uVerticallyIntegrated, vVerticallyIntegrated = self.getUV()
        # this will update self.integratedVelocity
        self.computeIntegratedFlux(uVerticallyIntegrated, vVerticallyIntegrated)

        self.vectorValues[:] = self.vinterp.getFaceVectors(
               self.integratedVelocity, placement=0)

    def getSizes(self):
        # get the sizes
        nt, nz, ny, nx = 1, 1, 0, 0
        shapeU = self.ncU['uo'].shape
        try:
            nt, nz, ny, nx = shapeU
        except:
            try:
                nz, ny, nx = shapeU
            except:
                try:
                    ny, nx = shapeU
                except:
                    raise RuntimeError("ERROR: uo's shape does not match (t, z, y, x), (z, y, x) or (y, x)")
        return nt, nz, ny, nx

    def buildEdgeUVGrids(self, bounds_lon, bounds_lat):

        # points, 4 points per cell, 3D
        self.lonlat = numpy.zeros((self.ny, self.nx, 4, 3), numpy.float64)
        self.lonlat[..., 0] = bounds_lon
        self.lonlat[..., 1] = bounds_lat

    def readField(self, nc, fieldName):

        # read fieldName
        try:
            field = nc[fieldName][self.timeIndex, :, :, :]
        except:
            try:
                field = nc[fieldName][...]
            except:
                raise RuntimeError(f'ERROR: could not read {fieldName} field')

        # set to zero where missing
        field = field.fillna(0.0)

        # integrate vertically, multiplying by the thickness of the layers,
        # this is the sum of multiplying axis 0 of thickness with axis 0 of field...
        fieldVerticallyIntegrated = numpy.tensordot(self.thickness, field, axes=(0, 0))

        return fieldVerticallyIntegrated

    def getUV(self):
        uVerticallyIntegrated = self.readField(self.ncU, 'uo')
        vVerticallyIntegrated = self.readField(self.ncV, 'vo')
        return uVerticallyIntegrated, vVerticallyIntegrated

    def computeArcLengths(self):

        # lon, lat, elev coords
        points = self.gr.getPoints() # array of size (numCells, 4, 3)

        # convert to Cartesian
        xyz = geo.lonLat2XYZArray(points, radius=geo.EARTH_RADIUS) # array of size (numCells, 4, 3)

        # compute the arc lengths
        for i0 in range(4):
            i1 = (i0 + 1) % 4
            self.arcLengths[:, i0] = geo.getArcLengthArray(xyz[:, i0, :], xyz[:, i1, :], radius=geo.EARTH_RADIUS)

    def computeIntegratedFlux(self, uVerticallyIntegrated, vVerticallyIntegrated):

        numCells = self.ny * self.nx
        #        ^
        #        |
        #  3-----V-----2
        #  |           |
        #  |     T     U-->
        #  |           |
        #  0-----------1

        # cross product means that the flux in V points to the negative y
        self.edgeFluxesUArray[:] = + uVerticallyIntegrated.reshape((numCells,)) * self.arcLengths[:, 1]
        self.edgeFluxesVArray[:] = - vVerticallyIntegrated.reshape((numCells,)) * self.arcLengths[:, 2]

        #        2
        #        ^
        #        |
        #  3-----V-----2
        #  |           |
        #  |->3  0     U->1
        #  |     ^     |
        #  |     |     |
        #  0-----------1

        # east
        self.integratedVelocity[:, 1] = self.edgeFluxesUArray
        # north
        self.integratedVelocity[:, 2] = self.edgeFluxesVArray

        # these should provide a different view to the array, NOT A COPY
        eU = self.edgeFluxesUArray.reshape((self.ny, self.nx))
        eV = self.edgeFluxesVArray.reshape((self.ny, self.nx))
        iV = self.integratedVelocity.reshape((self.ny, self.nx, 4))

        # south, (ny, nx, 4)
        iV[1:, :, 0] = eV[:-1, :] # will set self.integratedVelocity on the south side
        # west
        iV[:, 1:, 3] = eU[:, :-1]
        # periodic BCs
        iV[:, 0, 3] = eU[:, -1]

        if self.sverdrup:
            eU *= EARTH_RADIUS / 1.e6
            eV *= EARTH_RADIUS / 1.e6
            iV *= EARTH_RADIUS / 1.e6

        # from now on, edge fluxes are abs values!
        self.edgeFluxesUArray[:] = numpy.fabs(self.edgeFluxesUArray)
        self.edgeFluxesVArray[:] = numpy.fabs(self.edgeFluxesVArray)

        self.maxAbsFlux = max(self.maxAbsFlux, self.edgeFluxesUArray.max(), self.edgeFluxesVArray.max())

def __del__(self):
    self.ncU.close()
    self.ncV.close()

def main(*, tFile: str, uFile: str, vFile: str, lonLatPoints: str='', iFile: str='', sverdrup: bool=False):
    """Visualize fluxes
    :param tFile: netcdf file holding the T-grid
    :param uFile: netcdf file holding u data
    :param vFile: netcdf file holding v data
    :param lonLatPoints: target points "(lon0, lat0), (lon1, lat1),..."
    :param iFile: alternatively read target points from text file
    """
    if lonLatPoints:
        xyVals = numpy.array(eval(lonLatPoints))
    elif iFile:
        llreader = LatLonReader(iFile)
        xyVals = llreader.getLonLats()
    else:
        raise RuntimeError('ERROR must provide either iFile or lonLatPoints!')
    print(f'target points:\n {xyVals}')
    numTargetPoints = xyVals.shape[0]
    lonLatZPoints = numpy.zeros((numTargetPoints, 3), numpy.float64)
    lonLatZPoints[:, :2] = xyVals
    field = Field(tFile, uFile, vFile, lonLatZPoints, sverdrup)

if __name__ == '__main__':
    defopt.run(main)

max_stars_repo_name: pletzer/nemoflux

max_issues_repo_name: pletzer/nemoflux

max_forks_repo_name: pletzer/nemoflux


text:
"""
A few auxiliary methods that work with supported collection formats.
Used widely within AgentNet.
"""

from collections import OrderedDict
from warnings import warn

import lasagne
import numpy as np

def is_layer(var):
    """checks if var is lasagne layer"""
    return isinstance(var, lasagne.layers.Layer)

def is_theano_object(var):
    """checks if var is a theano input, transformation, constant or shared variable"""
    return type(var).__module__.startswith("theano")

def is_numpy_object(var):
    """checks if var is a theano input, transformation, constant or shared variable"""
    return type(var).__module__.startswith("numpy")

supported_sequences = (tuple, list)

def check_sequence(variables):
    """
    Ensure that variables is one of supported_sequences or converts to one.
    If naive conversion fails, throws an error.
    """
    if any(isinstance(variables, seq) for seq in supported_sequences):
        return variables
    else:
        # If it is a numpy or theano array, excluding numpy array of objects, return a list with single element
        # Yes, i know it's messy. Better options are welcome for pull requests :)
        if (is_theano_object(variables) or is_numpy_object(variables)) and variables.dtype != np.object:
            return [variables]
        elif hasattr(variables, '__iter__'):
            # Elif it is a different kind of sequence try casting to tuple. If cannot, treat that it will be treated
            # as an atomic object.
            try:
                tupled_variables = tuple(variables)
                message = """{variables} of type {var_type} will be treated as a sequence of {len_casted} elements,
                not a single element.
                If you want otherwise, please pass it as a single-element list/tuple.
                """
                warn(message.format(variables=variables, var_type=type(variables), len_casted=len(tupled_variables)))
                return tupled_variables
            except:
                message = """
                {variables} of type {var_type} will be treated as a single input/output tensor,
                and not a collection of such.
                If you want otherwise, please cast it to list/tuple.
                """
                warn(message.format(variables=variables, var_type=type(variables)))
                return [variables]
        else:
            # otherwise it's a one-element list
            return [variables]

def check_list(variables):
    """Ensure that variables is a list or converts to one.
    If naive conversion fails, throws an error
    :param variables: sequence expected
    """
    return list(check_sequence(variables))

def check_tuple(variables):
    """Ensure that variables is a list or converts to one.
    If naive conversion fails, throws an error
    :param variables: sequence expected
    """
    return tuple(check_sequence(variables))

def check_ordered_dict(variables):
    """Ensure that variables is an OrderedDict
    :param variables: dictionary expected
    """
    assert isinstance(variables, dict)
    try:
        return OrderedDict(list(variables.items()))
    except:
        raise ValueError("Could not convert {variables} to an ordered dictionary".format(variables=variables))

def unpack_list(array, parts_lengths):
    """
    Returns slices of the input list a.
    unpack_list(a, [2,3,5]) -> a[:2], a[2:2+3], a[2+3:2+3+5]

    :param array: array-like or tensor variable
    :param parts_lengths: lengths of subparts

    """
    borders = np.concatenate([[0], np.cumsum(parts_lengths)])

    groups = []
    for low, high in zip(borders[:-1], borders[1:]):
        groups.append(array[low:high])

    return groups


text:
import pytest
import numpy as np
from impyute.imputation.cs import mean
from impyute.ops import error
from impyute.ops import wrapper

# pylint:disable=redefined-builtin
try:
    raise ModuleNotFoundError
except NameError:
    class ModuleNotFoundError(Exception):
        "placeholder required for python2.7"
        pass
except ModuleNotFoundError:
    pass

@wrapper.wrappers
def wrappers_mul(arr):
    """Some function that performs an inplace operation on the input. Accepts kwargs"""
    arr *= 25
    return arr

def test_wrappers_inplace_false():
    """Input should be unchanged if inplace set to false"""
    A = np.ones((5, 5))
    A_copy = A.copy()
    wrappers_mul(A, inplace=False)
    assert A[0, 0] == A_copy[0, 0]

def test_wrappers_inplace_true():
    """Input may be changed if inplace set to true and operation is inplace"""
    A = np.ones((5, 5))
    A_copy = A.copy()
    wrappers_mul(A, inplace=True)
    assert A[0, 0] != A_copy[0, 0]

def test_wrappers_pandas_input():
    """ Input: DataFrame, Output: DataFrame """
    # Skip this test if you don't have pandas
    pytest.importorskip('pandas')
    import pandas as pd
    # Create a DataFrame with a NaN
    A = np.arange(25).reshape((5, 5)).astype(np.float)
    A[0, 0] = np.nan
    A = pd.DataFrame(A)
    # Assert that the output is a DataFrame
    assert isinstance(mean(A), pd.DataFrame)

@wrapper.checks
def some_fn(data):
    """Dummy fn that has form of np.array -> np.array"""
    return data

def test_correct_input():
    """ Test that an array that should satisfy all checks, no BadInputError should be raised"""
    # Integer np.ndarray (check: `_is_ndarray`, `_shape_2d`, `_nan_exists`)
    arr = np.array([[np.nan, 2], [3, 4]])
    # Cast integer array to float (check: `_dtype_float`)
    arr.dtype = np.float
    try:
        some_fn(arr)
    except error.BadInputError:
        assert False

def test_1d():
    """ Check 1d array, BadInputError raised"""
    arr = np.array([np.nan, 2])
    with pytest.raises(error.BadInputError) as excinfo:
        some_fn(arr)
    assert str(excinfo.value) == "No support for arrays that aren't 2D yet."

def test_not_nparray():
    """ If not an np.array, BadInputError raised"""
    with pytest.raises(error.BadInputError) as excinfo:
        some_fn([[np.nan, 2.], [3, 4]])
    assert str(excinfo.value) == "Not a np.ndarray."

def test_nan_exists():
    """ If no NaN, BadInputError raised"""
    with pytest.raises(error.BadInputError) as excinfo:
        some_fn(np.array([[1.]]))
    assert str(excinfo.value) == "No NaN's in given data"

@wrapper.conform_output
def conform_output_dummy(x):
    return x

def is_between_0_1(x):
    return 0 <= x <= 1

def coerce_between_0_1(arr, x_i, y_i):
    val = arr[x_i, y_i]
    if val < 0:
        return 0
    elif val > 1:
        return 1
    else:
        return x

def test_conform_output_not_used():
    """ If neither args passed, don't do anything"""
    assert "some input" == conform_output_dummy("some input")

def test_conform_output_valid_coerce():
    """ Check value valid and coerce invalid values"""
    arr = np.array([[1.1, 0.5], [0.2, -1]])
    actual = conform_output_dummy(
            arr,
            valid_fn=is_between_0_1,
            coerce_fn=coerce_between_0_1,
            )
    expected = np.array([[1.0, 0.5], [0.2, 0.0]])
    assert np.array_equal(expected, actual)

def test_conform_output_coerce():
    """ Coerce function doesn't run if no valid_fn passed  """
    arr = np.array([[1.1, 0.5], [0.2, -1]])
    actual = conform_output_dummy(
            arr,
            coerce_fn=coerce_between_0_1,
            )
    expected = np.array([[1.1, 0.5], [0.2, -1]])
    assert np.array_equal(expected, actual)

def test_conform_output_valid():
    """ No coerce_fn with valid_fn will raise BadOutputeError if invalid values
    encountered. First invalid value is 1.1
    """
    arr = np.array([[1.1, 0.5], [0.2, -1]])
    with pytest.raises(error.BadOutputError) as excinfo:
        conform_output_dummy(arr, valid_fn=is_between_0_1)
    assert str(excinfo.value) == "1.1 does not conform"

max_stars_repo_name: eltonlaw/impy

max_issues_repo_name: eltonlaw/impy

max_forks_repo_name: eltonlaw/impy


text:
""" Module for the model instance.
Creates the general road model in which the car agents reside.
"""
import numpy as np
from mesa import Model
from mesa.time import RandomActivation
from mesa.datacollection import DataCollector
import cargrid as car
from lane_grid import LaneSpace
from data_collection import avg_speed, cars_in_lane, track_params
from data_collection import track_run, avg_slowdown
np.warnings.filterwarnings('ignore')

class RoadSim(Model):

    """
    Hosts the road model and the mesa grid.
    Contains methods to generate new car agents and collect data.

    Attributes:
        uid (int): Hash of specific model parameters, used to
            conveniently group duplicate runs for data analysis.
        current_id (int): Initial value of the agent id generator,
            advances by 1 each time an id is generated through the
            next_id() method from the mesa scheduler.
        length (int): Length in meters of the highway lanes.
            Floored to an integer if a float is provided.
        lanes (int): Number of lanes in the highway model
        spawn_chance (float): Probability to generate a car each
            second, scales with the relative time_step to maintain
            a consistent spawn rate.
        time_step (float): How many seconds each iteration advances,
            e.g. 0.1 -> 10 steps for a second in 'real time'.
        agression (float): Agression of the car agents.
        min_gap (float): Minimal gap the car agents maintain.
        grid (obj): Instance of the LaneSpace class, contains
            the grid with all cars' locations.
        schedule (obj): Instance of the mesa scheduler.
        speed (float): Speed of the cars in m/s
        cars (list): List of all car agent objects.
        datacollector (obj): Mesa datacollector object to
            report data during a simulation.
    """
    # pylint: disable=too-many-instance-attributes
    # pylint: disable=too-many-arguments

    def __init__(self, lanes=3, length=5000, spawn=0.4, agression=0.1,
                 speed=100, time_step=0.1, min_gap=1.6):
        """
        Args:
            lanes (int): number of lanes
            length (float): length of each lane in meters, floored to an
                integer value
            spawn (float): spawn probability of a car on each lane every step
            agression (float): agression of the car agents
            speed (float): general speed of the car agents in km/h
            time_step (float): time in seconds to be advanced each step
            min_gap (float): minimum gap cars maintain in meters
        """
        super().__init__()
        self.uid = hash((spawn, agression, lanes))
        self.current_id = 0
        self.length = int(length)
        self.lanes = lanes
        self.spawn_chance = spawn
        self.time_step = time_step
        self.agression = agression
        self.min_gap = min_gap

        self.grid = LaneSpace(self.length, self.lanes, self.time_step,
                              scale=0.5)

        self.schedule = RandomActivation(self)
        self.speed = speed/3.6
        self.cars = []
        self.datacollector = DataCollector(
            model_reporters={
                "Avg_speed": avg_speed,
                'Cars_in_lane': cars_in_lane,
                'Avg_slowdown': avg_slowdown,
                'Model Params': track_params,
                'Run': track_run}
            )
        # Initialize 2 cars
        self.new_car()
        self.new_car(start_lane=1)

    def get_free_lanes(self):
        """
        Determines if each lane has enough space for a car to spawn: should
        be at least the speed of the car.

        Returns:
            A boolean array with the value of each lane.
        """
        return np.count_nonzero(self.grid.positions <
                                self.speed*self.time_step, axis=1) == 0

    def init_cars(self):
        """
        Loops over all lanes and randomly creates a new car object if
        sufficient space is available.
        """
        free_lanes = self.get_free_lanes()
        for i in range(self.lanes):
            if free_lanes[i] and np.random.rand()\
             < (self.spawn_chance*self.time_step):
                self.new_car(i)

    def new_car(self, start_lane=0):
        """
        Generates a new car object and adds it to the model scheduler.
        If the index of the agent is already occupied in the lane grid
        a new one is generated.
        """
        new_car = car.Car(self.next_id(), self, start_lane, self.speed,
                          self.agression, self.min_gap)

        while not self.grid.place_agent(new_car):
            new_car = car.Car(self.next_id(), self, start_lane, self.speed,
                              self.agression, self.min_gap)
        self.cars.append(new_car)
        self.schedule.add(new_car)

    def move(self, agent, new_lane):
        """
        Wrapper method over the grid move_agent method. Determines if
        the move remained inbounds and removes the car from the model
        otherwise.
        """
        has_moved = self.grid.move_agent(agent, new_lane)
        if not has_moved:
            self.grid.remove_agent(agent)
            self.schedule.remove(agent)
            self.cars.remove(agent)

    def step(self):
        """
        Step method for the mesa scheduler. Moves the car agents, initializes
        new cars, and collects model data in order.
        """
        self.schedule.step()
        self.init_cars()
        self.datacollector.collect(self)

max_stars_repo_name: Potgront/ABM

max_issues_repo_name: Potgront/ABM

max_forks_repo_name: Potgront/ABM


text:
import mdtraj as md
from mdtraj.utils import ensure_type
from mdtraj.geometry import _geometry
from rdkit import Chem
import glob
import pandas as pd
from rdkit.Chem import Descriptors
from rdkit.Chem import Lipinski
import subprocess
import glob
from rdkit.Chem import AllChem
import numpy as np
from pymol import cmd
import os
from scipy.spatial import distance

class ComposerGMX:

    def __init__(self):
        pass

    @classmethod
    def gro2pdb(self, gro_file, pdb_output = None):
        """
        Converts a gro file into a pdb file.

        Parameters
        ----------
        gro_file: str
            coordinates filename (.gro)
        pdb_output: str, optional
            name of the output pdb file. If not specified, then the pdb output file name be the same as the one of the gro_file.

        Returns
        ----------
        output.pdb: file
            pdb file
        """

        if pdb_output == None:
            pdb_output = os.path.splitext(gro_file)[0] + '.pdb'
        else:
            if str.split(pdb_output, '.')[-1] != "pdb":
                pdb_output = pdb_output + '.pdb'

        if os.path.isfile(pdb_output):
            return

        if os.path.isfile(gro_file):
            cmd.reinitialize()
            cmd.load(gro_file, 'test')
            cmd.save(pdb_output, 'test')

        if os.path.isfile(pdb_output):
            return
        else:
            try:
                args_subproc = "gmx editconf -f " + gro_file + " -o " + pdb_output + " -pbc no"
                subprocess.call(args_subproc, shell=True)
            except:
                print("Error: {} could not be converted to pdb".format(gro_file))

    @classmethod
    def solute_solvent_split(self, topology, solute_residue_name = None):
        """
        Obtain atom indexes of solute and solvent atoms from mdtraj topology file

        Parameters
        ----------
        topology: topology
            mdtraj topology file. It can be obtained from loaded structure or trajectory as object.topology
        solute_residue_name: str or list, optional
            name or list of names of the residue(s) to consider as solute (Default = 'LIG')

        Returns
        ----------
        solute_atoms: list
            List of indexes of the solute atoms
        solvent_atoms: list
            List of indexes of the solvent atoms
        """

        if solute_residue_name == None:
            solute_residue_name = 'LIG'

        if isinstance(solute_residue_name, list):
            solute_atoms = []
            for res_name in solute_residue_name:
                solute_atoms = solute_atoms + [atom.index for atom in topology.atoms if atom.residue.name == res_name]
        else:
            solute_atoms = [atom.index for atom in topology.atoms if atom.residue.name == solute_residue_name]
            if len(solute_atoms) == 0:
                print("The topology file does not containg any residue named '{}'. No solute atom extracted.".format(solute_residue_name))
        solvent_atoms = [atom.index for atom in topology.atoms if atom.residue.name != 'LIG']
        return solute_atoms, solvent_atoms

    @classmethod
    def get_stats(self, df):
        """Calcute mean, standard deviation, and median from list or array of numbers.

        Parameters
        ----------
        values: list or array
            list or array of numbers

        Returns
        ----------
        stats: list
            mean, standard deviation, and median of the provided input numbers
        """
        stats = [np.mean(df), np.std(df), np.median(df)]
        return stats

    @classmethod
    def shrake_rupley(self, traj, probe_radius=0.14, n_sphere_points=960, mode='residue', change_radii=None):
        """Compute the solvent accessible surface area of each atom or residue in each simulation frame.

        Modified from MDTraj

        Parameters
        ----------
        traj : Trajectory
            An mtraj trajectory.
        probe_radius : float, optional
            The radius of the probe, in nm.
        n_sphere_points : int, optional
            The number of points representing the surface of each atom, higher
            values leads to more accuracy.
        mode : {'atom', 'residue'}
            In mode == 'atom', the extracted areas are resolved per-atom
            In mode == 'residue', this is consolidated down to the
            per-residue SASA by summing over the atoms in each residue.
        change_radii : dict, optional
            A partial or complete dict containing the radii to change from the
            defaults. Should take the form {"Symbol" : radii_in_nm }, e.g.
            {"Cl" : 0.181 } to change the radii of Chlorine to 181 pm for the ionic Cl-.
        Returns
        -------
        areas : np.array, shape=(n_frames, n_features)
            The accessible surface area of each atom or residue in every frame.
            If mode == 'atom', the second dimension will index the atoms in
            the trajectory, whereas if mode == 'residue', the second
            dimension will index the residues.
        """

        _ATOMIC_RADII = {'H'   : 0.120, 'He'  : 0.140, 'Li'  : 0.076, 'Be' : 0.059,
                         'B'   : 0.192, 'C'   : 0.170, 'N'   : 0.155, 'O'  : 0.152,
                         'F'   : 0.147, 'Ne'  : 0.154, 'Na'  : 0.102, 'Mg' : 0.086,
                         'Al'  : 0.184, 'Si'  : 0.210, 'P'   : 0.180, 'S'  : 0.180,
                         'Cl'  : 0.175, 'Ar'  : 0.188, 'K'   : 0.138, 'Ca' : 0.114,
                         'Sc'  : 0.211, 'Ti'  : 0.200, 'V'   : 0.200, 'Cr' : 0.200,
                         'Mn'  : 0.200, 'Fe'  : 0.200, 'Co'  : 0.200, 'Ni' : 0.163,
                         'Cu'  : 0.140, 'Zn'  : 0.139, 'Ga'  : 0.187, 'Ge' : 0.211,
                         'As'  : 0.185, 'Se'  : 0.190, 'Br'  : 0.185, 'Kr' : 0.202,
                         'Rb'  : 0.303, 'Sr'  : 0.249, 'Y'   : 0.200, 'Zr' : 0.200,
                         'Nb'  : 0.200, 'Mo'  : 0.200, 'Tc'  : 0.200, 'Ru' : 0.200,
                         'Rh'  : 0.200, 'Pd'  : 0.163, 'Ag'  : 0.172, 'Cd' : 0.158,
                         'In'  : 0.193, 'Sn'  : 0.217, 'Sb'  : 0.206, 'Te' : 0.206,
                         'I'   : 0.198, 'Xe'  : 0.216, 'Cs'  : 0.167, 'Ba' : 0.149,
                         'La'  : 0.200, 'Ce'  : 0.200, 'Pr'  : 0.200, 'Nd' : 0.200,
                         'Pm'  : 0.200, 'Sm'  : 0.200, 'Eu'  : 0.200, 'Gd' : 0.200,
                         'Tb'  : 0.200, 'Dy'  : 0.200, 'Ho'  : 0.200, 'Er' : 0.200,
                         'Tm'  : 0.200, 'Yb'  : 0.200, 'Lu'  : 0.200, 'Hf' : 0.200,
                         'Ta'  : 0.200, 'W'   : 0.200, 'Re'  : 0.200, 'Os' : 0.200,
                         'Ir'  : 0.200, 'Pt'  : 0.175, 'Au'  : 0.166, 'Hg' : 0.155,
                         'Tl'  : 0.196, 'Pb'  : 0.202, 'Bi'  : 0.207, 'Po' : 0.197,
                         'At'  : 0.202, 'Rn'  : 0.220, 'Fr'  : 0.348, 'Ra' : 0.283,
                         'Ac'  : 0.200, 'Th'  : 0.200, 'Pa'  : 0.200, 'U'  : 0.186,
                         'Np'  : 0.200, 'Pu'  : 0.200, 'Am'  : 0.200, 'Cm' : 0.200,
                         'Bk'  : 0.200, 'Cf'  : 0.200, 'Es'  : 0.200, 'Fm' : 0.200,
                         'Md'  : 0.200, 'No'  : 0.200, 'Lr'  : 0.200, 'Rf' : 0.200,
                         'Db'  : 0.200, 'Sg'  : 0.200, 'Bh'  : 0.200, 'Hs' : 0.200,
                         'Mt'  : 0.200, 'Ds'  : 0.200, 'Rg'  : 0.200, 'Cn' : 0.200,
                         'Uut' : 0.200, 'Fl'  : 0.200, 'Uup' : 0.200, 'Lv' : 0.200,
                         'Uus' : 0.200, 'Uuo' : 0.200}

        xyz = ensure_type(traj.xyz, dtype=np.float32, ndim=3, name='traj.xyz', shape=(None, None, 3), warn_on_cast=False)

        if mode == 'atom':
            dim1 = xyz.shape[1]
            atom_mapping = np.arange(dim1, dtype=np.int32)

        elif mode == 'residue':
            dim1 = traj.n_residues
            if dim1 == 1:
                atom_mapping = np.array([0] * xyz.shape[1], dtype=np.int32)
            else:
                atom_mapping = np.array([a.residue.index for a in traj.top.atoms], dtype=np.int32)
                if not np.all(np.unique(atom_mapping) == np.arange(1 + np.max(atom_mapping))):
                    raise ValueError('residues must have contiguous integer indices starting from zero')
        else:
            raise ValueError('mode must be one of "residue", "atom". "{}" supplied'.format(mode))

        modified_radii = {}
        if change_radii is not None:
            # in case _ATOMIC_RADII is in use elsehwere...
            modified_radii = deepcopy(_ATOMIC_RADII)
            # Now, modify the values specified in 'change_radii'
            for k, v in change_radii.items(): modified_radii[k] = v

        out = np.zeros((xyz.shape[0], dim1), dtype=np.float32)
        atom_radii = []
        for atom in traj.topology.atoms:
            atom_name = "{}".format(atom).split('-')[1]
            element = ''.join(i for i in atom_name if not i.isdigit())
            if bool(modified_radii):
                atom_radii.append(modified_radii[element])
            else:
                atom_radii.append(_ATOMIC_RADII[element])

        radii = np.array(atom_radii, np.float32) + probe_radius

        _geometry._sasa(xyz, radii, int(n_sphere_points), atom_mapping, out)

        return out

    @classmethod
    def calc_shape_from_SMILES(self, smiles):
        """ Calculate 2D-shape parameter of compounds from SMILES.
        The 2D-shape parameter describes the elongation of the molecule and is calculated as
        the ratio between the eigenvalues of the covariance matrix of the 2D coordinates.
        (Help function for Properties2D_From_SMILES)

        Parameters
        ----------
        SMILES: SMILES
            SMILES of the molecule

        Returns
        ----------
        2D-shape: float
           2D-shape parameter of the input molecule
        """
        mol = Chem.MolFromSmiles(smiles)
        AllChem.Compute2DCoords(mol)
        conf = mol.GetConformer( 0 )
        xyz_array = np.array(conf.GetPositions())
        xy_array = xyz_array[:,:2]
        xy_mean = np.mean(xy_array, axis=0)
        xy_array_rescale = np.subtract(xy_array, xy_mean)
        mean_vec = np.mean(xy_array_rescale, axis=0)
        cov_mat = (xy_array_rescale - mean_vec).T.dot((xy_array_rescale - mean_vec)) / (xy_array_rescale.shape[0]-1)
        eig_vals, eig_vecs = np.linalg.eig(cov_mat)
        ratio_eig_vals = eig_vals[0]/eig_vals[1]
        return(ratio_eig_vals)

    @classmethod
    def calc_shape_from_Mol(self, molecule):
        """ Calculate 2D-shape parameter of compounds from RDKit molecule.
        The 2D-shape parameter describes the elongation of the molecule and is calculated as
        the ratio between the eigenvalues of the covariance matrix of the 2D coordinates.
        (Help function for Properties2D_From_Mol)

        Parameters
        ----------
        Molecule: RDKit molecule
            Molecule read in using one of the RDKit functions. See RDKit documentation.

        Returns
        ----------
        2D-shape: float
           2D-shape parameter of the input molecule
        """
        mol = molecule
        AllChem.Compute2DCoords(mol)
        conf = mol.GetConformer( 0 )
        xyz_array = np.array(conf.GetPositions())
        xy_array = xyz_array[:,:2]
        xy_mean = np.mean(xy_array, axis=0)
        xy_array_rescale = np.subtract(xy_array, xy_mean)
        mean_vec = np.mean(xy_array_rescale, axis=0)
        cov_mat = (xy_array_rescale - mean_vec).T.dot((xy_array_rescale - mean_vec)) / (xy_array_rescale.shape[0]-1)
        eig_vals, eig_vecs = np.linalg.eig(cov_mat)
        ratio_eig_vals = eig_vals[0]/eig_vals[1]
        return(ratio_eig_vals)

    @classmethod
    def extract_Zwitterionic_Label_From_Mol(self, mol, dist_cutoff = None, cmpd_name = None):
        """
        It returns a dictionary containing a binary label for zwitterionic compounds.
        The label is "0" if the compound is not zwitterionic, while it is "1" otherwise.
        A compound is considered zwitterinic if it contains both a positively and a negatively charged groups and
        if the groups are at a distance from each other smaller than the cutoff distance (dist_cutoff).
        If the cutoff distance is not specified, then every compound containing both a positively and a negatively charged groups
        is considered zwitterionic.
        If cmpd_name is specified, it is returned in the output dictionary.

        Parameters
        ----------
        Molecule: RDKit molecule
            Molecule read in using one of the RDKit functions. See RDKit documentation.
        dist_cutoff: float, optional
            Cutoff distance for the definition of a zwitterionic compound. See description above. (Default = None)
        cmpd_name: str, optional
            Name of the compound. If specified, it is returned in the output dictionary. (Default = None)

        Returns
        ----------
        dict_zwit: dict
            Dictionary key is "is_zwit". It contains a flag for zwitterionic compounds. See the description above.
            If cmpd_name is specified, it is returned in the dictionary.
        """

        # SMARTS of charged groups
        NR4 = Chem.MolFromSmarts('[NX4+]')
        NX3_no_amide = Chem.MolFromSmarts('[NX3;!$([NX3][CX3](=[OX1])[#6])]')
        COO = Chem.MolFromSmarts('[CX3](=O)[O-]')
        SO3 = Chem.MolFromSmarts('[#16X4](=[OX1])(=[OX1])[O-]')
        nR3 = Chem.MolFromSmarts('[$([nX4+]),$([nX3]);!$(*=*)&!$(*:*)]')
        OR4 = Chem.MolFromSmarts('[O-]')
        SR4 = Chem.MolFromSmarts('[S-]')

        if mol:
            #check for carboxylate
            idx_COO = mol.GetSubstructMatches(COO)  #((19, 21, 20),)
            idx_COO_O = [id1[2] for id1 in idx_COO]
            # check for sulfonic acids
            idx_SO3 = mol.GetSubstructMatches(SO3)
            idx_SO3_O = [id1[3] for id1 in idx_SO3]
            # Thiols
            idx_SR4 = mol.GetSubstructMatches(SR4)
            idx_SR4_S = [id1[0] for id1 in idx_SR4]
            # Alchols
            idx_OR4 = mol.GetSubstructMatches(OR4)
            idx_OR4_O = [id1[0] for id1 in idx_OR4]
            negatively_charged_atoms = list(set(idx_COO_O + idx_SO3_O + idx_SR4_S + idx_OR4_O))
            if negatively_charged_atoms:
                #check for charged nitrogens
                idx_NR4 = mol.GetSubstructMatches(NR4)
                idx_NR4_n = [id1[0] for id1 in idx_NR4]
                idx_nR3 = mol.GetSubstructMatches(nR3)
                idx_nR3_n = [id1[0] for id1 in idx_nR3]
                idx_NR3 = mol.GetSubstructMatches(NX3_no_amide)
                idx_NR3_n = [id1[0] for id1 in idx_NR3]
                positively_charged_atoms = list(set(idx_NR4_n + idx_nR3_n + idx_NR3_n))
                #ionizable_atoms_tmp = idx_NR4_n + idx_nR3_n + idx_COO_O + idx_SO3_O + idx_SR4_S + idx_OR4_O + idx_NR3_n
                #ionizable_atoms = list(set(ionizable_atoms_tmp))
                if positively_charged_atoms:
                    conf = mol.GetConformer( -1 )
                    dist = []
                    pts_pos =  np.array([ list( conf.GetAtomPosition(atmidx) ) for atmidx in positively_charged_atoms ])
                    pts_neg = np.array([ list( conf.GetAtomPosition(atmidx) ) for atmidx in negatively_charged_atoms ])
                    for pt1, pt2 in zip(pts_pos, pts_neg):
                        dist.append(distance.euclidean(pt1, pt2))

                    dist_min = min(dist)
                    if dist_cutoff != None:
                        if dist_min <= dist_cutoff:
                            is_zwit = 1
                        else:
                            is_zwit = 0
                    else:
                        is_zwit = 1

                else:
                    is_zwit = 0
            else:
                is_zwit = 0

        dict_zwit = {"is_zwit": is_zwit}

        if cmpd_name == None:
            return(dict_zwit)
        else:
            dict_zwit.update({"cmpd_name": cmpd_name})
            return(dict_zwit)

    @classmethod
    def Properties2D_From_Mol(self, molecule, cmpd_name = None, includeSandP = None, dist_cutoff = None):
        """
        It returns a dictionary of 2D-counts and properties for the input RDKit molecule.
        These are: number of heavy atoms ("HA_count"), number of rotatable bonds ("RB_count"), number of N, O, F, P, S, Cl, Br, and I atoms ("ATOM_count" where ATOM is the element), molecular weight ("MW"), number of hydrogen bond donor ("HBD_count") and acceptor atoms ("HBA_count"), binary zwitterionic flag ("is_zwit"), 2D-shape parameter ("2d_shape"), and the topological polar surface area ("2d_psa").
        The 2D-shape parameter descibes the elongation of the molecule (see the function calc_shape_from_Mol for details).

        Parameters:
        ----------
        Molecule: RDKit molecule
            Molecule read in using one of the RDKit functions. See RDKit documentation.
        includeSandP: bool, optional
            Set to False to exclude the S and P atoms from the calculation of the TPSA. (Default = True)
        dist_cutoff: float, optional
            Cutoff distance for the definition of a zwitterionic compound. (Default = None)
            If the compound contains both a positively and a negatively charged chemical group and these are at a distance from each other smaller than the cutoff distance (dist_cutoff),
            than the zwitterionic flag ("is_zwit") is set to 1. Otherwise the compound is considered not zwitterionic ("is_zwit" = 0).
            If the cutoff distance is not specified, then every compound containing both a positively and a negatively charged groups is considered zwitterionic.
        cmpd_name: str, optional
            Name of the compound. If specified, it is returned in the output dictionary. (Default = None)

        Returns
        ----------
        dict_2dcounts: dict
            Dictionary of 2D-counts and properties (see description above).
            If cmpd_name is specified, it is returned in the dictionary.
        """

        if includeSandP == None:
            includeSandP = True

        # define SMARTS
        RotatableBond = Chem.MolFromSmarts('[!$(*#*)&!D1]-&!@[!$(*#*)&!D1]')
        H_acceptor = Chem.MolFromSmarts('[!$([#6,F,Cl,Br,I,o,s,nX3,#7v5,#15v5,#16v4,#16v6,*+1,*+2,*+3])]')
        H_donor = Chem.MolFromSmarts('[!H0;#7,#8,#9,#16]')

        # Counts
        HA_count = molecule.GetNumHeavyAtoms()
        N_count  = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#7]'))) # nitrogens
        O_count  = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#8]'))) # oxygens
        F_count  = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#9]'))) # fluorines
        P_count  = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#15]'))) # phosphorous
        S_count  = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#16]'))) # sulfurs
        Cl_count = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#17]'))) # chlorines
        Br_count = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#35]'))) # bromines
        I_count  = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#53]'))) # iodines
        PSA2D = (Descriptors.TPSA(molecule, includeSandP=includeSandP))/100 #/100 to have nm instead of Angstrom
        N_donors = len(molecule.GetSubstructMatches(H_donor))
        N_acceptors = len(molecule.GetSubstructMatches(H_acceptor))
        RB_count  = len(molecule.GetSubstructMatches(RotatableBond))
        MW = Chem.Descriptors.ExactMolWt(molecule)
        shape_2d = self.calc_shape_from_Mol(molecule)
        is_zwit = self.extract_Zwitterionic_Label_From_Mol(molecule, dist_cutoff=dist_cutoff)

        dict_2dcounts = {"MW": MW, "HA_count": HA_count, "RB_count": RB_count, "N_count": N_count, "O_count": O_count, "F_count": F_count, "P_count": P_count, "S_count": S_count, "Cl_count": Cl_count, "Br_count": Br_count, "I_count": I_count, "HBD_count": N_donors, "HBA_count": N_acceptors, "2d_shape": shape_2d, "2d_psa": PSA2D}

        dict_2dcounts.update(is_zwit)

        if cmpd_name == None:
            return(dict_2dcounts)
        else:
            dict_2dcounts.update({"cmpd_name": cmpd_name})
            return(dict_2dcounts)

    @classmethod
    def Properties2D_From_SMILES(self, smiles, cmpd_name = None, includeSandP = None):
        """
        It returns a dictionary of 2D-counts and properties for the input molecule provided as SMILES.
        These are: number of heavy atoms ("HA_count"), number of rotatable bonds ("RB_count"), number of N, O, F, P, S, Cl, Br, and I atoms ("ATOM_count" where ATOM is the element), molecular weight ("MW"), number of hydrogen bond donor ("HBD_count") and acceptor atoms ("HBA_count"), binary zwitterionic flag ("is_zwit"), 2D-shape parameter ("2d_shape"), and the topological polar surface area ("2d_psa").
        The 2D-shape parameter descibes the elongation of the molecule (see the function calc_shape_from_Mol for details).

        Parameters:
        ----------
        SMILES: SMILES
            SMILES of the molecule
        includeSandP: bool, optional
            Set to False to exclude the S and P atoms from the calculation of the TPSA. (Default = True)
        dist_cutoff: float, optional
            Cutoff distance for the definition of a zwitterionic compound. (Default = None)
            If the compound contains both a positively and a negatively charged chemical group and these are at a distance from each other smaller than the cutoff distance (dist_cutoff),
            than the zwitterionic flag ("is_zwit") is set to 1. Otherwise the compound is considered not zwitterionic ("is_zwit" = 0).
            If the cutoff distance is not specified, then every compound containing both a positively and a negatively charged groups is considered zwitterionic.
        cmpd_name: str, optional
            Name of the compound. If specified, it is returned in the output dictionary. (Default = None)

        Returns
        ----------
        dict_2dcounts: dict
            Dictionary of 2D-counts and properties (see description above).
            If cmpd_name is specified, it is returned in the dictionary.
        """

        if includeSandP == None:
            includeSandP = True

        molecule = Chem.MolFromSmiles(smiles)

        # define SMARTS
        RotatableBond = Chem.MolFromSmarts('[!$(*#*)&!D1]-&!@[!$(*#*)&!D1]')
        H_acceptor = Chem.MolFromSmarts('[!$([#6,F,Cl,Br,I,o,s,nX3,#7v5,#15v5,#16v4,#16v6,*+1,*+2,*+3])]')
        H_donor = Chem.MolFromSmarts('[!H0;#7,#8,#9,#16]')

        # Counts
        HA_count = molecule.GetNumHeavyAtoms()
        N_count  = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#7]'))) # nitrogens
        O_count  = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#8]'))) # oxygens
        F_count  = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#9]'))) # fluorines
        P_count  = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#15]'))) # phosphorous
        S_count  = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#16]'))) # sulfurs
        Cl_count = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#17]'))) # chlorines
        Br_count = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#35]'))) # bromines
        I_count  = len(molecule.GetSubstructMatches(Chem.MolFromSmarts('[#53]'))) # iodines
        PSA2D = (Descriptors.TPSA(molecule, includeSandP=includeSandP))/100 #/100 to have nm instead of Angstrom
        N_donors = len(molecule.GetSubstructMatches(H_donor))
        N_acceptors = len(molecule.GetSubstructMatches(H_acceptor))
        RB_count  = len(molecule.GetSubstructMatches(RotatableBond))
        MW = Chem.Descriptors.ExactMolWt(molecule)
        shape_2d = self.calc_shape_from_Mol(molecule)
        is_zwit = self.extract_Zwitterionic_Label_From_Mol(molecule)

        dict_2dcounts = {"MW": MW, "HA_count": HA_count, "RB_count": RB_count, "N_count": N_count, "O_count": O_count, "F_count": F_count, "P_count": P_count, "S_count": S_count, "Cl_count": Cl_count, "Br_count": Br_count, "I_count": I_count, "HBD_count": N_donors, "HBA_count": N_acceptors, "2d_shape": shape_2d, "2d_psa": PSA2D}

        dict_2dcounts.update(is_zwit)

        if cmpd_name == None:
            return(dict_2dcounts)
        else:
            dict_2dcounts.update({"cmpd_name": cmpd_name})
            return(dict_2dcounts)

    @classmethod
    def extract_dipole_moments(self, solute_traj, prm_obj, solute_atoms, cmpd_name = None):
        """
        It returns Dipole Moment Terms.
        Especially, it returns a dictionary with mean, median and standard deviation of the dipole moment (mu)
        and of the dipole moment components along x-, y-, and z-axis (mu_x, mu_y, mu_z).
        It uses the dipole_moments function of MDTraj.
        It requires as inputs the trajectory of the solute molecule, partial charges charges of the solute atoms, which are contained in the topology file, and the indices of the solute atoms.
        If cmpd_name is specified, it is returned in the dictionary.
        See also Tutorial2.ipynb.

        Parameters:
        ----------
        solute_traj: MDTraj trajectory object
            trajectory of the solute. The trajectory should be read in using the MDTraj functions.
            Example: solute_traj = md.load(traj_file, top=pdb_file, atom_indices = solute_atoms)
        prm_obj: object containing structure's topology information
            This object can be obtained from the topology file (.top for GROMACS) using the parmed library.
            For the GROMCS topology file, the command is: prm_obj = pmd.gromacs.GromacsTopologyFile(top_file).
        solute_atoms: list
            indeces of solute atoms. Needed to extract the partial charges of the solute atoms from the topology object.
        cmpd_name: str, optional
            Name of the compound. If specified, it is returned in the output dictionary. (Default = None)

        Returns
        ----------
        dict_dip_mom: dict
            Dictionary containing mean, standard deviation, and median of the the dipole moment magnitude and components calculated over the simulation trajectory.
            If cmpd_name is specified, it is returned in the dictionary.

        """

        solute_charges = [i.charge for idx,i in enumerate(prm_obj.atoms) if idx in solute_atoms ]
        dip_moments = md.dipole_moments(solute_traj, solute_charges)
        dipole_magnitude = np.sqrt(np.square(dip_moments).sum(axis=1))
        # <mu_x>, <mu_y>, <mu_z>, std_mu_x, std_mu_y, std_mu_z, med_mu_x, med_mu_y, med_mu_z, <mu>, std_mu, med_mu
        dip_components_av = list(dip_moments.mean(axis=0))  # <mu_x>, <mu_y>, <mu_z>
        dip_components_std = list(dip_moments.std(axis=0))  # std_mu_x, std_mu_y, std_mu_z
        dip_components_med = list(np.median(dip_moments, axis=0))  #med_mu_x, med_mu_y, med_mu_z
        dip_mom_stats = self.get_stats(dipole_magnitude)  #<mu>, std_mu, med_mu

        dict_dip_mom = {'av_mu_x': dip_components_av[0], 'av_mu_y': dip_components_av[1], 'av_mu_z': dip_components_av[2], 'std_mu_x': dip_components_std[0], 'std_mu_y': dip_components_std[1], 'std_mu_z': dip_components_std[2], 'med_mu_x': dip_components_med[0], 'med_mu_y': dip_components_med[1], 'med_mu_z': dip_components_med[2], 'av_mu': dip_mom_stats[0], 'std_mu': dip_mom_stats[1], 'med_mu': dip_mom_stats[2]}

        if cmpd_name == None:
            return(dict_dip_mom)
        else:
            dict_dip_mom.update({"cmpd_name": cmpd_name})
            return(dict_dip_mom)

    @classmethod
    def extract_rgyr(self, mdtraj_obj, cmpd_name = None):
        """
        It returns a dictionary containing Mean, Median and Standard Deviation of the Radius of Gyration (Rgyr).
        It uses the compute_rg function of MDTraj.
        If cmpd_name is specified, it is returned in the dictionary.

        Parameters:
        ----------
        solute_traj: MDTraj trajectory object
            trajectory of the solute. The trajectory should be read in using the MDTraj functions.
            Example: solute_traj = md.load(traj_file, top=pdb_file, atom_indices = solute_atoms)
        cmpd_name: str, optional
            Name of the compound. If specified, it is returned in the output dictionary. (Default = None)

        Returns
        ----------
        dict_rgyr: dict
            Dictionary containing mean, standard deviation, and median of the Rgyr calculated over the simulation trajectory.
            If cmpd_name is specified, it is returned in the dictionary.
        """

        df = list(md.compute_rg(mdtraj_obj, masses = np.array([a.element.mass for a in mdtraj_obj.topology.atoms])))
        stats = list(self.get_stats(df))

        dict_rgyr = {'wat_rgyr_av': stats[0], 'wat_rgyr_std': stats[1], 'wat_rgyr_med': stats[2]}

        if cmpd_name == None:
            return(dict_rgyr)
        else:
            dict_rgyr.update({"cmpd_name": cmpd_name})
            return(dict_rgyr)

    @classmethod
    def extract_sasa(self, mdtraj_obj, cmpd_name = None):
        """
        It returns a dictionary containing Mean, Median and Standard Deviation of the Solvent Accessible Surface Area (SASA).
        It uses the shrake_rupley function of MDTraj.
        If cmpd_name is specified, it is returned in the dictionary.

        Parameters:
        ----------
        solute_traj: MDTraj trajectory object
            trajectory of the solute. The trajectory should be read in using the MDTraj functions.
            Example: solute_traj = md.load(traj_file, top=pdb_file, atom_indices = solute_atoms)
        cmpd_name: str, optional
            Name of the compound. If specified, it is returned in the output dictionary. (Default = None)

        Returns
        ----------
        dict_rgyr: dict
            Dictionary containing mean, standard deviation, and median of the Rgyr calculated over the simulation trajectory.
            If cmpd_name is specified, it is returned in the dictionary.
        """

        df = list(self.shrake_rupley(mdtraj_obj, mode='residue'))
        stats = list(self.get_stats(df))

        dict_sasa = {'wat_sasa_av': stats[0], 'wat_sasa_std': stats[1], 'wat_sasa_med': stats[2]}

        if cmpd_name == None:
            return(dict_sasa)
        else:
            dict_sasa.update({"cmpd_name": cmpd_name})
            return(dict_sasa)

    @classmethod
    def read_xvg(self, energy_file_xvg):
        """
        It reads the GROMACS xvg file into a pandas dataframe.
        (Help function for EnergyTermsGmx).

        Parameters:
        ----------
        energy_file_xvg: str
            Name of the xvg file to read in

        Returns
        ----------
        energy_df: df
            Pandas dataframe of the energy terms and simulation parameters contained in the xvg file
        """
        if os.path.isfile(energy_file_xvg):
            data = pd.read_csv(energy_file_xvg, sep="\t", header = None)
            N_lines_header = data[data[0].str.contains('@')].shape[0] + data[data[0].str.contains('#')].shape[0]
            energy_df = pd.read_csv(energy_file_xvg, header = None, skiprows = N_lines_header, delim_whitespace=True)
            N_col = energy_df.shape[1]
            colnames = ["time"]
            for i in range(N_col-1):
                colnames.append(str.split(str(data[data[0].str.contains('s{}'.format(i))]), " ")[-1].replace('"',''))
            energy_df.columns = colnames
            return energy_df
        else:
            print("Error: Energy file {} does not exist".format(energy_file_xvg))
            return

    @classmethod
    def EnergyTermsGmx(self, energy_file_xvg, cmpd_name = None):
        """
        Energy Terms are extracted from the Gromacs xvg files.
        It returns Energy Terms (as described in the original publication).
        Espacially, it returns a dictionary that contains as keys the mean, standard deviation and median of the energy terms.
        If cmpd_name is specified, it is returned in the dictionary.

        Parameters:
        --------------------
        energy_file_xvg: file
            Gromacs xvg file containing energy terms
        cmpd_name: str, optional
            Name of the compound. If specified, it is returned in the output dictionary.

        Returns
        ----------
        dict_ene: dict
            Keys are mean, standard deviation, and median of the energy terms.
            If cmpd_name is specified, it is returned in the dictionary.
        """

        if os.path.isfile(energy_file_xvg):
            energy_df = self.read_xvg(energy_file_xvg)
        else:
            print("Error: Energy file {} does not exist".format(energy_file_xvg))
            return

        energy_array = np.array(energy_df)
        idx_coul_sr_lig = [i for i, s in enumerate(energy_df) if 'Coul-SR:LIG' in s][0] #5
        idx_coul_14_lig = [i for i, s in enumerate(energy_df) if 'Coul-14:LIG' in s][0] #7
        idx_lj_sr_lig = [i for i, s in enumerate(energy_df) if 'LJ-SR:LIG' in s][0] #6
        idx_lj_14_lig = [i for i, s in enumerate(energy_df) if 'LJ-14:LIG' in s][0] #8
        idx_coul_sr_wat_lig = [i for i, s in enumerate(energy_df) if 'Coul-SR:Water' in s][0] #1
        idx_coul_14_wat_lig = [i for i, s in enumerate(energy_df) if 'Coul-14:Water' in s][0] #3
        idx_lj_sr_wat_lig = [i for i, s in enumerate(energy_df) if 'LJ-SR:Water' in s][0] #2
        idx_lj_14_wat_lig = [i for i, s in enumerate(energy_df) if 'LJ-14:Water' in s][0] #4
        intra_crf_stats = self.get_stats(energy_array[:,idx_coul_sr_lig] + energy_array[:,idx_coul_14_lig])
        intra_lj_stats = self.get_stats(energy_array[:,idx_lj_sr_lig] + energy_array[:,idx_lj_14_lig])
        tot_crf_stats = self.get_stats(energy_array[:,idx_coul_sr_lig] + energy_array[:,idx_coul_14_lig] + energy_array[:,idx_coul_sr_wat_lig] + energy_array[:,idx_coul_14_wat_lig])
        tot_lj_stats = self.get_stats(energy_array[:,idx_lj_sr_lig] + energy_array[:,idx_lj_14_lig] + energy_array[:,idx_lj_sr_wat_lig] + energy_array[:,idx_lj_14_wat_lig])
        intra_ene_stats = self.get_stats(np.sum(energy_array[:,[idx_coul_sr_lig, idx_coul_14_lig, idx_lj_sr_lig, idx_lj_14_lig]], axis=1))
        tot_ene_stats = self.get_stats(np.sum(energy_array[:,[idx_coul_sr_lig, idx_coul_14_lig, idx_lj_sr_lig, idx_lj_14_lig, idx_coul_sr_wat_lig, idx_coul_14_wat_lig, idx_lj_sr_wat_lig, idx_lj_14_wat_lig]], axis=1))
        intra_crf_av = intra_crf_stats[0]
        intra_crf_std = intra_crf_stats[1]
        intra_crf_med = intra_crf_stats[2]
        intra_lj_av = intra_lj_stats[0]
        intra_lj_std = intra_lj_stats[1]
        intra_lj_med = intra_lj_stats[2]
        tot_crf_av = tot_crf_stats[0]
        tot_crf_std = tot_crf_stats[1]
        tot_crf_med = tot_crf_stats[2]
        tot_lj_av = tot_lj_stats[0]
        tot_lj_std = tot_lj_stats[1]
        tot_lj_med = tot_lj_stats[2]
        intra_ene_av = intra_ene_stats[0]
        intra_ene_std = intra_ene_stats[1]
        intra_ene_med = intra_ene_stats[2]
        tot_ene_av = tot_ene_stats[0]
        tot_ene_std = tot_ene_stats[1]
        tot_ene_med = tot_ene_stats[2]

        dict_ene = {'intra_crf_av_wat': intra_crf_av, 'intra_crf_std_wat': intra_crf_std, 'intra_crf_med_wat': intra_crf_med, 'intra_lj_av_wat': intra_lj_av, 'intra_lj_std_wat': intra_lj_std, 'intra_lj_med_wat': intra_lj_med, 'total_crf_av_wat': tot_crf_av, 'total_crf_std_wat': tot_crf_std, 'total_crf_med_wat': tot_crf_med, 'total_lj_av_wat': tot_lj_av, 'total_lj_std_wat': tot_lj_std, 'total_lj_med_wat': tot_lj_med, 'intra_ene_av_wat': intra_ene_av, 'intra_ene_std_wat': intra_ene_std, 'intra_ene_med_wat': intra_ene_med, 'total_ene_av_wat': tot_ene_av, 'total_ene_std_wat': tot_ene_std, 'total_ene_med_wat': tot_ene_med}

        if cmpd_name != None:
            dict_ene.update({"cmpd_name": cmpd_name})

        return(dict_ene)

    @classmethod
    def calc_psa3d(self, obj_list=None, include_SandP: bool = True, atom_to_remove = None):
        """
        Help function to calculate the 3d polar surface area (3D-PSA) of molecules in Interface_Pymol for all the snapshots in a MD trajectory.
        (Contribution by Benjamin Schroeder)

        Parameters
        ----------
        obj_list: list, optional
            list of pymol objects (Default = "cmpd1")
        include_SandP: bool, optional
            Set to False to exclude the S and P atoms from the calculation of the 3D-PSA. (Default = True)
        atom_to_remove: str, optional
            Single atom name of the atom to remove from the selection (Default = None).
            Useful if you want to include only S or only P in the calculation of the 3D-PSA.

        Returns
        ----------
        obj_psa_dict: list
            Values correspond to mean, standard deviation, and median of the 3D-PSA calculated over the simulation time
        """

        # IO
        if (obj_list is None):
            obj_list = cmd.get_names("objects")
        # Loop over objects
        obj_psa_dict = {}
        for obj in obj_list:
            cmd.frame(0)
            states = range(1, cmd.count_states(obj) + 1)  # get all states of the object
            ##Loop over all states
            psa = []
            for state in states:
                ###select all needed atoms by partialCSelection or element or H next to (O ,N)
                if atom_to_remove != None and isinstance(atom_to_remove, str):
                    if include_SandP:
                        select_string = "resn LIG and (elem N or elem O or elem S or elem P or (elem H and (neighbor elem N+O+S+P))) and " + obj + " and not name {}".format(atom_to_remove)  #@carmen add: "or elem S"
                    else:
                        select_string = "resn LIG and (elem N or elem O or (elem H and (neighbor elem N+O))) and " + obj  + " and not name {}".format(atom_to_remove)  #@carmen add: "or elem S"
                else:
                    if include_SandP:
                        select_string = "resn LIG and (elem N or elem O or elem S or elem P or (elem H and (neighbor elem N+O+S+P))) and " + obj   #@carmen add: "or elem S"
                    else:
                        select_string = "resn LIG and (elem N or elem O or (elem H and (neighbor elem N+O))) and " + obj  #@carmen add: "or elem S"
                cmd.select("noh", select_string)
                ###calc surface area
                psa.append(float(cmd.get_area("noh", state=state)))
                ###gather data
                #obj_psa_dict.update({obj: psa})
        obj_psa_dict = [np.mean(psa)/100, np.std(psa)/100, np.median(psa)/100]  #/100 to have nm instead of Angstrom

        return obj_psa_dict

    @classmethod
    def extract_psa3d(self, traj_file, gro_file, obj_list=None, include_SandP = None, cmpd_name = None):
        """
        Calculates the 3d polar surface area (3D-PSA) of molecules in Interface_Pymol for all the snapshots in a MD trajectory.
        (Contribution by Benjamin Schroeder)

        Parameters:
        -------------
        traj_file: str
            trajectory filename
        gro_file: str
            coordinates filename (.gro or .pdb)
        cmpd_name: str, optional
            Name of the compound used as object name (default is "cmpd1")
        include_SandP: bool, optional
            Set to False to exclude the S and P atoms from the calculation of the 3D-PSA. (Default = True)
        atom_to_remove: str, optional
            Single atom name of the atom to remove from the selection (Default = None).
            Useful if you want to include only S or only P in the calculation of the 3D-PSA.

        Returns
        ----------
        dict_psa3d: dict
            Keys are mean (3d_psa_av), standard deviation (3d_psa_sd), and median (3d_psa_med) of the 3D-PSA calculated over the simulation time.
            If cmpd_name is specified, it is returned in the dictionary.
        """

        if cmpd_name == None:
            cmpd_name = "cmpd1"

        # Load trajectory and remove solvent and salts
        obj1 = cmpd_name
        cmd.reinitialize()
        cmd.load(gro_file, object=obj1)
        cmd.load_traj(traj_file, object=obj1)
        cmd.remove("solvent")
        cmd.remove("resn Cl-")
        cmd.remove("resn Na+")

        atom_names = []
        cmd.iterate_state(-1, selection=obj1 + " and not elem H", expression="atom_names.append(name)", space=locals())
        total_psa = self.calc_psa3d(include_SandP = include_SandP)

        dict_psa3d = {'3d_psa_av': total_psa[0], '3d_psa_sd': total_psa[1], '3d_psa_med': total_psa[2]}

        if cmpd_name == "cmpd1":
            return(dict_psa3d)
        else:
            dict_psa3d.update({"cmpd_name": cmpd_name})
            return(dict_psa3d)

max_stars_repo_name: greglandrum/mdfptools

max_issues_repo_name: greglandrum/mdfptools

max_forks_repo_name: greglandrum/mdfptools

reason: import numpy,from scipy


text:
import numpy as np
import tensorflow as tf
import cv2
import os
import random
import time
import sys

# number
number = ['1', '2', '3', 'b', 'c', 'm', 'n', 'v', 'x', 'z']

# 图像大小
IMAGE_HEIGHT = 22  # 80
IMAGE_WIDTH = 62  # 160
MAX_CAPTCHA = 4

char_set = number
CHAR_SET_LEN = len(char_set)  #10
valid_path = "TestImg/"
model_path = "model/"

X = tf.placeholder(tf.float32, [None, IMAGE_HEIGHT * IMAGE_WIDTH])
Y = tf.placeholder(tf.float32, [None, MAX_CAPTCHA * CHAR_SET_LEN])
keep_prob = tf.placeholder(tf.float32)  # dropout

# 定义CNN
def crack_captcha_cnn(w_alpha=0.01, b_alpha=0.1):
    x = tf.reshape(X, shape=[-1, IMAGE_HEIGHT, IMAGE_WIDTH, 1])

    # 3 conv layer
    w_c1 = tf.Variable(w_alpha * tf.random_normal([3, 3, 1, 32]))
    b_c1 = tf.Variable(b_alpha * tf.random_normal([32]))
    conv1 = tf.nn.relu(tf.nn.bias_add(tf.nn.conv2d(x, w_c1, strides=[1, 1, 1, 1], padding='SAME'), b_c1))
    conv1 = tf.nn.max_pool(conv1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
    conv1 = tf.nn.dropout(conv1, keep_prob)

    w_c2 = tf.Variable(w_alpha * tf.random_normal([3, 3, 32, 64]))
    b_c2 = tf.Variable(b_alpha * tf.random_normal([64]))
    conv2 = tf.nn.relu(tf.nn.bias_add(tf.nn.conv2d(conv1, w_c2, strides=[1, 1, 1, 1], padding='SAME'), b_c2))
    conv2 = tf.nn.max_pool(conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
    conv2 = tf.nn.dropout(conv2, keep_prob)

    w_c3 = tf.Variable(w_alpha * tf.random_normal([3, 3, 64, 64]))
    b_c3 = tf.Variable(b_alpha * tf.random_normal([64]))
    conv3 = tf.nn.relu(tf.nn.bias_add(tf.nn.conv2d(conv2, w_c3, strides=[1, 1, 1, 1], padding='SAME'), b_c3))
    conv3 = tf.nn.max_pool(conv3, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
    conv3 = tf.nn.dropout(conv3, keep_prob)

    # Fully connected layer
    w_d = tf.Variable(w_alpha * tf.random_normal([3 * 8 * 64, 1024]))
    b_d = tf.Variable(b_alpha * tf.random_normal([1024]))
    dense = tf.reshape(conv3, [-1, w_d.get_shape().as_list()[0]])
    dense = tf.nn.relu(tf.add(tf.matmul(dense, w_d), b_d))
    dense = tf.nn.dropout(dense, keep_prob)

    w_out = tf.Variable(w_alpha * tf.random_normal([1024, MAX_CAPTCHA * CHAR_SET_LEN]))
    b_out = tf.Variable(b_alpha * tf.random_normal([MAX_CAPTCHA * CHAR_SET_LEN]))
    out = tf.add(tf.matmul(dense, w_out), b_out)
    # out = tf.nn.softmax(out)
    return out

# 向量转回文本
def vec2text(vec):
    char_pos = vec.nonzero()[0]
    text = []
    for i, c in enumerate(char_pos):
        text.append(char_set[c % CHAR_SET_LEN])
    return "".join(text)

def predict_captcha(captcha_image):
    output = crack_captcha_cnn()

    saver = tf.train.Saver()
    with tf.Session() as sess:
        saver.restore(sess, tf.train.latest_checkpoint(model_path))

        predict = tf.argmax(tf.reshape(output, [-1, MAX_CAPTCHA, CHAR_SET_LEN]), 2)
        text_list = sess.run(predict, feed_dict={X: [captcha_image], keep_prob: 1})

        text = text_list[0].tolist()
        vector = np.zeros(MAX_CAPTCHA * CHAR_SET_LEN)
        i = 0
        for n in text:
            vector[i * CHAR_SET_LEN + n] = 1
            i += 1
        return vec2text(vector)

if not os.path.exists(valid_path):
    print('Image does not exist, please check!, path:"{}"'.format(os.path.abspath(valid_pathb)))
    sys.exit()
image_list = os.listdir(valid_path)

output = crack_captcha_cnn()

saver = tf.train.Saver()
with tf.Session() as sess:
    saver.restore(sess, tf.train.latest_checkpoint(model_path))

    predict = tf.argmax(tf.reshape(output, [-1, MAX_CAPTCHA, CHAR_SET_LEN]), 2)

    count = 0
    acceptCount = 0
    for image_ in image_list:
        count += 1
        text_ = image_.split('.')[0]
        image_p = os.path.join(valid_path, image_)
        # 单张图片预测
        image = np.float32(cv2.imread(image_p, 0))
        image = image.flatten() / 255

        text_list = sess.run(predict, feed_dict={X: [image], keep_prob: 1})

        text = text_list[0].tolist()
        vector = np.zeros(MAX_CAPTCHA * CHAR_SET_LEN)
        i = 0
        for n in text:
            vector[i * CHAR_SET_LEN + n] = 1
            i += 1
        predict_text= vec2text(vector)

        print("真实值: {0}  预测值: {1}".format(text_, predict_text),text_ == predict_text)
        if text_ == predict_text: acceptCount += 1;
    print("测试集准确率",acceptCount,count,acceptCount/count)


text:
import os
import csv
import sys
import time
import json
import h5py
import logging
import pickle as pkl
import argparse
import random
from collections import OrderedDict

import torch
import numpy as np
from tqdm import tqdm

from nglib.common import utils
from nglib.narrative.dataset import sample_mcnc
from nglib.narrative.narrative_graph import get_pp_ridx2distr
from nglib.narrative.narrative_graph import get_pp_ridx2distr_coref

def get_arguments(argv):
    parser = argparse.ArgumentParser(description='prepare coref chains for MCNC')
    parser.add_argument('config_file', metavar='CONFIG_FILE',
                        help='ng config file')
    parser.add_argument('input_dir', metavar='INPUT_DIR',
                        help='input dir')
    parser.add_argument('prefix', metavar='PREFIX',
                        help='output file prefix')
    parser.add_argument('output_dir', metavar='OUTPUT_DIR',
                        help='output dir')

    parser.add_argument('--is_training_set', action='store_true', default=False,
                        help='is training set')

    parser.add_argument("--chain_len", type=int, default=9,
                        help="fixed chain length (default: 9).")
    parser.add_argument("--n_choices", type=int, default=5,
                        help="number of choices (default: 5).")

    parser.add_argument("--n_per_file", type=int, default=5000,
                        help="number of graphs per file.")
    parser.add_argument("--sample_n_graphs", type=int, default=-1,
                        help="sample n graphs.")

    parser.add_argument("--seed", type=int, default=135,
                        help="random seed for initialization")
    parser.add_argument('-v', '--verbose', action='store_true', default=False,
                        help='show info messages')
    parser.add_argument('-d', '--debug', action='store_true', default=False,
                        help='show debug messages')
    args = parser.parse_args(argv)
    return args

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

def load_gid2docid(fpath):
    logger.info('loading {}...'.format(fpath))
    gid2docid = OrderedDict()
    with open(fpath, 'r') as fr:
        for line in fr:
            line = line.rstrip('\n')
            sp = line.split('\t')
            gid2docid[int(sp[0])] = sp[1]
    logger.info('{} documents loaded.'.format(len(gid2docid)))
    return gid2docid

def load_chain_info(fpath):
    logger.info('loading {}...'.format(fpath))
    out = OrderedDict()
    with open(fpath, 'r') as fr:
        for line in fr:
            d = json.loads(line)
            out[d['doc_id']] = d
    logger.info('{} documents loaded.'.format(len(out)))
    return out

def main():
    config = json.load(open(args.config_file))
    assert config["config_target"] == "narrative_graph"

    logger.warning('!!!!!!!!Note that NO_ENTITY={}. Please make sure it is correct'.format(config['no_entity']))

    rtype2idx = config['rtype2idx']
    coref_ridx = rtype2idx['cnext']

    new_gid = 0
    new_gid2docid = OrderedDict()

    t2 = time.time()
    fs = sorted([f for f in os.listdir(args.input_dir) if f.endswith('.h5')])
    count_fw = 0
    questions = {}
    fw = os.path.join(args.output_dir, '{}_{}.h5'.format(args.prefix, count_fw))
    logger.info('writing {}...'.format(fw))
    fw = h5py.File(fw, 'w')
    for f in fs:
        fname = f.split('.')[0]
        fpath = os.path.join(args.input_dir, '{}.txt'.format(fname))
        gid2docid = load_gid2docid(fpath)

        fpath = os.path.join(args.input_dir, '{}.json'.format(fname))
        doc_chain_info = load_chain_info(fpath)

        fpath = os.path.join(args.input_dir, f)
        logger.info('processing {}...'.format(fpath))

        fr = h5py.File(fpath, 'r')
        for gn in tqdm(fr.keys()):
            gid = int(gn.split('_')[-1])
            docid = gid2docid[gid]
            if docid not in doc_chain_info:
                logger.debug('{} does not have coref chains'.format(docid))
                continue

            # read
            ng_edges = fr[gn]['ng_edges'][:]
            bert_inputs = fr[gn]['bert_inputs'][:]
            bert_target_idxs = fr[gn]['bert_target_idxs'][:]
            bert_nid2rows = fr[gn]['bert_nid2rows'][:]
            coref_chains = fr[gn]['coref_chains'][:]

            # fixed chain length
            chain_info = doc_chain_info[docid]['chains'][0][:args.chain_len]
            coref_nids = [m['nid'] for m in chain_info]

            if not args.is_training_set:
                # use only one chain for each graph
                n_nodes = bert_nid2rows.shape[0]

                # random sample
                correct, choices, target_edge_idx = sample_mcnc(
                    ng_edges, n_nodes, coref_nids, coref_ridx, args.n_choices)

                neg_nids = [choices[i][2] for i in range(args.n_choices) if i != correct]
                fw['graph_{}/negative_coref_nids'.format(new_gid)] = np.array(neg_nids, dtype=np.int64)

                questions['graph_{}'.format(new_gid)] = (correct, choices, target_edge_idx, chain_info)

            # dump output
            new_gid2docid[new_gid] = docid

            fw['graph_{}/bert_inputs'.format(new_gid)] = bert_inputs
            fw['graph_{}/bert_target_idxs'.format(new_gid)] = bert_target_idxs
            fw['graph_{}/bert_nid2rows'.format(new_gid)] = bert_nid2rows
            fw['graph_{}/ng_edges'.format(new_gid)] = ng_edges
            fw['graph_{}/coref_nids'.format(new_gid)] = np.array(coref_nids, dtype=np.int64)

            new_gid += 1
            if new_gid % args.n_per_file == 0:
                logger.info('sampling {} files takes {} s'.format(args.n_per_file, time.time()-t2))

                fw.close()
                if not args.is_training_set:
                    fpath = os.path.join(args.output_dir, 'q_{}_{}.pkl'.format(args.prefix, count_fw))
                    logger.info('writing {}...'.format(fpath))
                    pkl.dump(questions, open(fpath, 'wb'))

                # next file
                count_fw += 1
                fpath = os.path.join(args.output_dir, '{}_{}.h5'.format(args.prefix, count_fw))
                logger.info('writing {}...'.format(fpath))
                fw = h5py.File(fpath, 'w')

                questions = {}

                # restart timer
                t2 = time.time()
            if args.sample_n_graphs !=-1 and new_gid >= args.sample_n_graphs:
                break
        if args.sample_n_graphs != -1 and new_gid >= args.sample_n_graphs:
            break

        fr.close()
    fw.close()
    if not args.is_training_set:
        fpath = os.path.join(args.output_dir, '{}_{}.pkl'.format(args.prefix, count_fw))
        logger.info('writing {}...'.format(fpath))
        pkl.dump(questions, open(fpath, 'wb'))

    fpath = os.path.join(args.output_dir, 'gid2docid.txt')
    with open(fpath, 'w') as fw:
        for gid, did in new_gid2docid.items():
            fw.write('{}\t{}\n'.format(gid, did))

    logger.info('#graphs = {}'.format(len(new_gid2docid)))

if __name__ == "__main__":
    args = utils.bin_config(get_arguments)
    logger = utils.get_root_logger(args)
    main()


text:
from nose.tools import eq_, ok_
from distutils.version import LooseVersion

import nose
import sys
import pandas as pd
import numpy as np
import scipy as sp
import matplotlib
import torch

def setup_module():
    pass

def test_library_versions():

    # We use Python's distutils to verify versions. It seems sufficient for our purpose.
    # Alternatively, we could use pkg_resources.parse_version for a more robust parse.

    min_python = '3.6'
    min_numpy = '1.13'
    min_scipy = '1.0'
    min_torch = '0.3'
    min_pandas = '0.21'
    min_matplotlib = '2.0'

    ok_(LooseVersion(sys.version) > LooseVersion(min_python))
    ok_(LooseVersion(np.__version__) > LooseVersion(min_numpy))
    ok_(LooseVersion(sp.__version__) > LooseVersion(min_scipy))
    ok_(LooseVersion(torch.__version__) > LooseVersion(min_torch))
    ok_(LooseVersion(pd.__version__) > LooseVersion(min_pandas))
    ok_(LooseVersion(matplotlib.__version__) > LooseVersion(min_matplotlib))

reason: import numpy,import scipy


text:
import sys
import glob
import numpy as np
from scipy.io import wavfile
import spectrogram as sp
import os
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

# read in file names
# change audio_dir to directory where the audio files are
#audio_dir = '/media/sdcard/calls/' #/media/omacaodh/DATA/Data/AIStats_data/calls/'
#audio_files = glob.glob(audio_dir + '*.wav')
#print 'total number of audio files', len(audio_files)

# select a single file
# change the number of the file_id to see different bat calls
#file_id = 0
#file_name = os.path.basename(audio_files[file_id])

#print audio_files[file_id]

# read sys arguments
file_name = sys.argv[1]
#print file_name
print 'filename', file_name

## read audio file
##sample_rate, x = wavfile.read(audio_files[file_id])
sample_rate, x = wavfile.read(file_name)
#sample_rate = 44100
file_length = x.shape[0] / float(sample_rate)
#file_length = 1.0
print 'sampling rate', sample_rate, ', length (secs)', file_length

## create spectrogram
spec = sp.gen_mag_spectrogram_mod(x, sample_rate, 0.001)
print 'number of frequency bins', spec.shape[0]
print 'number of time bins', spec.shape[1]

## plot so we can look at it
plt.figure(1)
plt.imshow(spec)
plt.title("Recording")  #(str(file_id) + ' : ' + file_name)
plt.gca().invert_yaxis()  # flip the frequency axis
plt.xlabel('time (secs)')
plt.ylabel('frequency (kHz)')
#plt.show()

## you can save the spectrogram to as an image
## again we flip it so that the lowest frequencies are at the bottom
plt.imsave(file_name[:-4] + '.png', np.flipud(spec))

## as the calls are too high fequency to hear we slow them down
## by changing the sampling rate (here factor of 10) and save to disk
#op_file_name = file_name[:-4] + '_slow.wav'
#wavfile.write(op_file_name, sample_rate/10.0, x)
#

max_stars_repo_licenses: Intel Unlicense

max_issues_repo_licenses: Intel Unlicense

max_forks_repo_licenses: Intel Unlicense


text:
import pandas as pd
import numpy as np

data = pd.read_csv("datasets/raw.csv")

temp = data[["tag_name", "song_id"]]
groups = temp.groupby("tag_name").count()
groups = groups.reset_index()
groups.columns = ["tag_name", "count"]

greaterfive = groups[groups["count"] > 5]
tofive = groups[groups["count"] <= 5]

ones = tofive[tofive["count"] == 1]
twotofive = tofive[tofive["count"] > 1]

'''
Total tags: 529086
Tags appearing once: 61%
Tags appearing two to five times: 22%
Tags appearing more than five times: 17%
'''

ones = ones.reset_index().drop("index", axis=1)
rnd = np.random.randint(len(ones), size=300)

chris = ones.loc[rnd[0:100], :]
torsten = ones.loc[rnd[100:200], :]
mohsen = ones.loc[rnd[200:300], :]

chris.to_csv("chris.csv", index=False)
torsten.to_csv("torsten.txt", index=False)
mohsen.to_csv("mohsen.txt", index=False)