text:
#!/usr/bin/env python
# coding: utf-8

import os
import argparse
import logging

import numpy as np
from astropy.table import Table, vstack
from matplotlib import pyplot as plt

from pycorr import TwoPointCorrelationFunction, TwoPointEstimator, KMeansSubsampler, utils, setup_logging
from LSS.tabulated_cosmo import TabulatedDESI

logger = logging.getLogger('xirunpc')

def get_zlims(tracer, tracer2=None, option=None):

    if tracer2 is not None:
        zlims1 = get_zlims(tracer, option=option)
        zlims2 = get_zlims(tracer2, option=option)
        return [zlim for zlim in zlims1 if zlim in zlims2]

    if tracer.startswith('LRG'):
        zlims = [0.4, 0.6, 0.8, 1.1]

    if tracer.startswith('ELG'):# or type == 'ELG_HIP':
        zlims = [0.8, 1.1, 1.6]
        if option:
            if option == 'safez':
                zlims = [0.9, 1.48]
            if 'extended' in option:
                print('extended is no longer a meaningful option')
                #zlims = [0.8, 1.1, 1.6]
            if 'smallshells' in option:
                zlims = [0.8, 0.9,1.0,1.1,1.2,1.3,1.4,1.5,1.6]

    if tracer.startswith('QSO'):
        zlims = [0.8, 1.1, 1.6, 2.1,3.5]
        if option == 'highz':
            zlims = [2.1, 3.5]
        if option == 'lowz':
            zlims = [0.8, 2.1]

    if tracer.startswith('BGS'):
        zlims = [0.1, 0.3, 0.5]
        if option == 'lowz':
            zlims = [0.1, 0.3]
        if option == 'highz':
            zlims = [0.3, 0.5]

    if option == 'fullonly':
        zlims = [zlims[0], zlims[-1]]

    return zlims

def get_regions(survey, rec=False):
    regions = ['N', 'S']#, '']
    #if survey in ['main', 'DA02']:
    #    regions = ['DN', 'DS', 'N', 'S']
    #    if rec: regions = ['DN', 'N']
    return regions

def select_region(ra, dec, region):
    mask_ra = (ra > 100 - dec)
    mask_ra &= (ra < 280 + dec)
    if region == 'DN':
        mask = dec < 32.375
        mask &= mask_ra
    elif region == 'DS':
        mask = dec > -25
        mask &= ~mask_ra
    else:
        raise ValueError('Input region must be one of ["DN", "DS"].')
    return mask

def catalog_dir(survey='main', verspec='guadalupe', version='test', base_dir='/global/cfs/cdirs/desi/survey/catalogs'):
    return os.path.join(base_dir, survey, 'LSS', verspec, 'LSScats', version)

def catalog_fn(tracer='ELG', region='', ctype='clustering', name='data', rec_type=False, nrandoms=4, cat_dir=None, survey='main', **kwargs):
    if cat_dir is None:
        cat_dir = catalog_dir(survey=survey, **kwargs)
    if survey in ['main', 'DA02']:
        tracer += 'zdone'
    if ctype == 'full':
        region = ''
    dat_or_ran = name[:3]
    if name == 'randoms' and tracer == 'LRG_main' and ctype == 'full':
        tracer = 'LRG'
    if region: region = '_' + region
    if rec_type:
        dat_or_ran = '{}.{}'.format(rec_type, dat_or_ran)
    if name == 'data':
        return os.path.join(cat_dir, '{}{}_{}.{}.fits'.format(tracer, region, ctype, dat_or_ran))
    return [os.path.join(cat_dir, '{}{}_{:d}_{}.{}.fits'.format(tracer, region, iran, ctype, dat_or_ran)) for iran in range(nrandoms)]

def get_clustering_positions_weights(catalog, distance, zlim=(0., np.inf), weight_type='default', name='data', return_mask=False, option=None):

    mask = (catalog['Z'] >= zlim[0]) & (catalog['Z'] < zlim[1])
    if option:
        if 'elgzmask' in option:
            zmask = ((catalog['Z'] >= 1.49) & (catalog['Z'] < 1.52))
            mask &= ~zmask
    logger.info('Using {:d} rows for {}.'.format(mask.sum(), name))
    positions = [catalog['RA'][mask], catalog['DEC'][mask], distance(catalog['Z'][mask])]
    weights = np.ones_like(positions[0])

    if 'completeness_only' in weight_type and 'bitwise' in weight_type:
        raise ValueError('inconsistent choices were put into weight_type')

    if name == 'data':
        if 'zfail' in weight_type:
            weights *= catalog['WEIGHT_ZFAIL'][mask]
        if 'default' in weight_type and 'bitwise' not in weight_type:
            weights *= catalog['WEIGHT'][mask]
        if 'RF' in weight_type:
            weights *= catalog['WEIGHT_RF'][mask]*catalog['WEIGHT_COMP'][mask]
        if 'completeness_only' in weight_type:
            weights = catalog['WEIGHT_COMP'][mask]
        if 'FKP' in weight_type:
            weights *= catalog['WEIGHT_FKP'][mask]
        if 'bitwise' in weight_type:
            if catalog['BITWEIGHTS'].ndim == 2: weights = list(catalog['BITWEIGHTS'][mask].T) + [weights]
            else: weights = [catalog['BITWEIGHTS'][mask]] + [weights]

    if name == 'randoms':
        if 'default' in weight_type:
            weights *= catalog['WEIGHT'][mask]
        if 'RF' in weight_type:
            weights *= catalog['WEIGHT_RF'][mask]*catalog['WEIGHT_COMP'][mask]
        if 'zfail' in weight_type:
            weights *= catalog['WEIGHT_ZFAIL'][mask]
        if 'completeness_only' in weight_type:
            weights = catalog['WEIGHT_COMP'][mask]
        if 'FKP' in weight_type:
            weights *= catalog['WEIGHT_FKP'][mask]

    if return_mask:
        return positions, weights, mask
    return positions, weights

def read_clustering_positions_weights(distance, zlim=(0., np.inf), weight_type='default', name='data', concatenate=False, option=None, **kwargs):

    cat_fns = catalog_fn(ctype='clustering', name=name, **kwargs)
    logger.info('Loading {}.'.format(cat_fns))
    isscalar = not isinstance(cat_fns, (tuple, list))
    if isscalar:
        cat_fns = [cat_fns]

    toret = [get_clustering_positions_weights(Table.read(cat_fn), distance, zlim=zlim, weight_type=weight_type, name=name, option=option) for cat_fn in cat_fns]
    if isscalar:
        return toret[0]
    positions_weights = [[tmp[0] for tmp in toret], [tmp[1] for tmp in toret]]
    if concatenate:
        for iarray, arrays in enumerate(positions_weights):
            if isinstance(arrays[0], (tuple, list)):  # e.g., list of bitwise weights
                array = [np.concatenate([arr[iarr] for arr in arrays], axis=0) for iarr in range(len(arrays[0]))]
            else:
                array = np.concatenate(arrays, axis=0)
            positions_weights[iarray] = array
    return positions_weights

def get_full_positions_weights(catalog, name='data', weight_type='default', fibered=False, region='', return_mask=False):

    mask = np.ones(len(catalog), dtype='?')
    if region in ['DS', 'DN']:
        mask &= select_region(catalog['RA'], catalog['DEC'], region)
    elif region:
        mask &= catalog['PHOTSYS'] == region.strip('_')

    if fibered: mask &= catalog['LOCATION_ASSIGNED']
    positions = [catalog['RA'][mask], catalog['DEC'][mask], catalog['DEC'][mask]]
    if fibered and 'bitwise' in weight_type:
        if catalog['BITWEIGHTS'].ndim == 2: weights = list(catalog['BITWEIGHTS'][mask].T)
        else: weights = [catalog['BITWEIGHTS'][mask]]
    else: weights = np.ones_like(positions[0])
    if return_mask:
        return positions, weights, mask
    return positions, weights

def read_full_positions_weights(name='data', weight_type='default', fibered=False, region='', **kwargs):

    cat_fn = catalog_fn(ctype='full', name=name, **kwargs)
    logger.info('Loading {}.'.format(cat_fn))
    if isinstance(cat_fn, (tuple, list)):
        catalog = vstack([Table.read(fn) for fn in cat_fn])
    else:
        catalog = Table.read(cat_fn)
    return get_full_positions_weights(catalog, name=name, weight_type=weight_type, fibered=fibered, region=region)

def compute_angular_weights(nthreads=8, dtype='f8', tracer='ELG', tracer2=None, mpicomm=None, mpiroot=None, **kwargs):

    autocorr = tracer2 is None
    catalog_kwargs = kwargs

    fibered_data_positions1, fibered_data_weights1, fibered_data_positions2, fibered_data_weights2 = None, None, None, None
    parent_data_positions1, parent_data_weights1, parent_data_positions2, parent_data_weights2 = None, None, None, None
    parent_randoms_positions1, parent_randoms_weights1, parent_randoms_positions2, parent_randoms_weights2 = None, None, None, None

    if mpicomm is None or mpicomm.rank == mpiroot:

        fibered_data_positions1, fibered_data_weights1 = read_full_positions_weights(name='data', fibered=True, tracer=tracer, **catalog_kwargs)
        parent_data_positions1, parent_data_weights1 = read_full_positions_weights(name='data', fibered=False, tracer=tracer, **catalog_kwargs)
        parent_randoms_positions1, parent_randoms_weights1 = read_full_positions_weights(name='randoms', tracer=tracer, **catalog_kwargs)
        if not autocorr:
            fibered_data_positions2, fibered_data_weights2 = read_full_positions_weights(name='data', fibered=True, tracer=tracer2, **catalog_kwargs)
            parent_data_positions2, parent_data_weights2 = read_full_positions_weights(name='data', fibered=False, tracer=tracer2, **catalog_kwargs)
            parent_randoms_positions2, parent_randoms_weights2 = read_full_positions_weights(name='randoms', tracer=tracer2, **catalog_kwargs)

    tedges = np.logspace(-4., 0.5, 41)
    # First D1D2_parent/D1D2_PIP angular weight
    wangD1D2 = TwoPointCorrelationFunction('theta', tedges, data_positions1=fibered_data_positions1, data_weights1=fibered_data_weights1,
                                            data_positions2=fibered_data_positions2, data_weights2=fibered_data_weights2,
                                            randoms_positions1=parent_data_positions1, randoms_weights1=parent_data_weights1,
                                            randoms_positions2=parent_data_positions2, randoms_weights2=parent_data_weights2,
                                            estimator='weight', engine='corrfunc', position_type='rdd', nthreads=nthreads,
                                            dtype=dtype, mpicomm=mpicomm, mpiroot=mpiroot)

    # First D1R2_parent/D1R2_IIP angular weight
    # Input bitwise weights are automatically turned into IIP
    if autocorr:
         parent_randoms_positions2, parent_randoms_weights2 = parent_randoms_positions1, parent_randoms_weights1
    wangD1R2 = TwoPointCorrelationFunction('theta', tedges, data_positions1=fibered_data_positions1, data_weights1=fibered_data_weights1,
                                            data_positions2=parent_randoms_positions2, data_weights2=parent_randoms_weights2,
                                            randoms_positions1=parent_data_positions1, randoms_weights1=parent_data_weights1,
                                            randoms_positions2=parent_randoms_positions2, randoms_weights2=parent_randoms_weights2,
                                            estimator='weight', engine='corrfunc', position_type='rdd', nthreads=nthreads,
                                            dtype=dtype, mpicomm=mpicomm, mpiroot=mpiroot)
    wangR1D2 = None
    if not autocorr:
        wangR1D2 = TwoPointCorrelationFunction('theta', tedges, data_positions1=parent_randoms_positions1, data_weights1=parent_randoms_weights1,
                                               data_positions2=fibered_data_positions2, data_weights2=fibered_data_weights2,
                                               randoms_positions1=parent_randoms_positions1, randoms_weights1=parent_randoms_weights1,
                                               randoms_positions2=parent_data_positions2, randoms_weights2=parent_data_weights2,
                                               estimator='weight', engine='corrfunc', position_type='rdd', nthreads=nthreads,
                                               dtype=dtype, mpicomm=mpicomm, mpiroot=mpiroot)

    wang = {}
    wang['D1D2_twopoint_weights'] = wangD1D2
    wang['D1R2_twopoint_weights'] = wangD1R2
    wang['R1D2_twopoint_weights'] = wangR1D2

    return wang

def compute_correlation_function(corr_type, edges, distance, nthreads=8, dtype='f8', wang=None, split_randoms_above=30., weight_type='default', tracer='ELG', tracer2=None, rec_type=None, njack=120,option=None, mpicomm=None, mpiroot=None, **kwargs):

    autocorr = tracer2 is None
    catalog_kwargs = kwargs.copy()
    catalog_kwargs['weight_type'] = weight_type
    with_shifted = rec_type is not None

    if 'angular' in weight_type and wang is None:

        wang = compute_angular_weights(nthreads=nthreads, dtype=dtype, weight_type=weight_type, tracer=tracer, tracer2=tracer2, mpicomm=mpicomm, mpiroot=mpiroot, **kwargs)

    data_positions1, data_weights1, data_samples1, data_positions2, data_weights2, data_samples2 = None, None, None, None, None, None
    randoms_positions1, randoms_weights1, randoms_samples1, randoms_positions2, randoms_weights2, randoms_samples2 = None, None, None, None, None, None
    shifted_positions1, shifted_weights1, shifted_samples1, shifted_positions2, shifted_weights2, shifted_samples2 = None, None, None, None, None, None
    jack_positions = None

    if mpicomm is None or mpicomm.rank == mpiroot:

        if with_shifted:
            data_positions1, data_weights1 = read_clustering_positions_weights(distance, name='data', rec_type=rec_type, tracer=tracer,option=option, **catalog_kwargs)
            shifted_positions1, shifted_weights1 = read_clustering_positions_weights(distance, name='randoms', rec_type=rec_type, tracer=tracer,option=option, **catalog_kwargs)
        else:
            data_positions1, data_weights1 = read_clustering_positions_weights(distance, name='data', rec_type=rec_type, tracer=tracer,option=option, **catalog_kwargs)
        randoms_positions1, randoms_weights1 = read_clustering_positions_weights(distance, name='randoms', rec_type=rec_type, tracer=tracer,option=option, **catalog_kwargs)
        jack_positions = data_positions1

        if not autocorr:
            if with_shifted:
                data_positions2, data_weights2 = read_clustering_positions_weights(distance, name='data', rec_type=rec_type, tracer=tracer2,option=option, **catalog_kwargs)
                shifted_positions2, shifted_weights2 = read_clustering_positions_weights(distance, name='randoms', rec_type=rec_type, tracer=tracer2,option=option, **catalog_kwargs)
            else:
                data_positions2, data_weights2 = read_clustering_positions_weights(distance, name='data', rec_type=rec_type, tracer=tracer2,option=option, **catalog_kwargs)
            randoms_positions2, randoms_weights2 = read_clustering_positions_weights(distance, name='randoms', rec_type=rec_type, tracer=tracer2,option=option, **catalog_kwargs)
            jack_positions = [np.concatenate([p1, p2], axis=0) for p1, p2 in zip(jack_positions, data_positions2)]

    if njack >= 2:
        subsampler = KMeansSubsampler('angular', positions=jack_positions, nsamples=njack, nside=512, random_state=42, position_type='rdd',
                                      dtype=dtype, mpicomm=mpicomm, mpiroot=mpiroot)

        if mpicomm is None or mpicomm.rank == mpiroot:
            data_samples1 = subsampler.label(data_positions1)
            randoms_samples1 = [subsampler.label(p) for p in randoms_positions1]
            if with_shifted:
                shifted_samples1 = [subsampler.label(p) for p in shifted_positions1]
            if not autocorr:
                data_samples2 = subsampler.label(data_positions2)
                randoms_samples2 = [subsampler.label(p) for p in randoms_positions2]
                if with_shifted:
                    shifted_samples2 = [subsampler.label(p) for p in shifted_positions2]

    kwargs = {}
    kwargs.update(wang or {})
    randoms_kwargs = dict(randoms_positions1=randoms_positions1, randoms_weights1=randoms_weights1, randoms_samples1=randoms_samples1,
                          randoms_positions2=randoms_positions2, randoms_weights2=randoms_weights2, randoms_samples2=randoms_samples2,
                          shifted_positions1=shifted_positions1, shifted_weights1=shifted_weights1, shifted_samples1=shifted_samples1,
                          shifted_positions2=shifted_positions2, shifted_weights2=shifted_weights2, shifted_samples2=shifted_samples2)

    zedges = np.array(list(zip(edges[0][:-1], edges[0][1:])))
    mask = zedges[:,0] >= split_randoms_above
    zedges = [zedges[~mask], zedges[mask]]
    split_edges, split_randoms = [], []
    for ii, zedge in enumerate(zedges):
        if zedge.size:
            split_edges.append([np.append(zedge[:,0], zedge[-1,-1])] + list(edges[1:]))
            split_randoms.append(ii > 0)

    results = []
    if mpicomm is None:
        nran = len(randoms_positions1)
    else:
        nran = mpicomm.bcast(len(randoms_positions1) if mpicomm.rank == mpiroot else None, root=mpiroot)
    for i_split_randoms, edges in zip(split_randoms, split_edges):
        result = 0
        D1D2 = None
        for iran in range(nran if i_split_randoms else 1):
            tmp_randoms_kwargs = {}
            if i_split_randoms:
                # On scales above split_randoms_above, sum correlation function over multiple randoms
                for name, arrays in randoms_kwargs.items():
                    if arrays is None:
                        continue
                    else:
                        tmp_randoms_kwargs[name] = arrays[iran]
            else:
                # On scales below split_randoms_above, concatenate randoms
                for name, arrays in randoms_kwargs.items():
                    if arrays is None:
                        continue
                    elif isinstance(arrays[0], (tuple, list)):  # e.g., list of bitwise weights
                        array = [np.concatenate([arr[iarr] for arr in arrays], axis=0) for iarr in range(len(arrays[0]))]
                    else:
                        array = np.concatenate(arrays, axis=0)
                    tmp_randoms_kwargs[name] = array
            tmp = TwoPointCorrelationFunction(corr_type, edges, data_positions1=data_positions1, data_weights1=data_weights1, data_samples1=data_samples1,
                                              data_positions2=data_positions2, data_weights2=data_weights2, data_samples2=data_samples2,
                                              engine='corrfunc', position_type='rdd', nthreads=nthreads, dtype=dtype, **tmp_randoms_kwargs, **kwargs,
                                              D1D2=D1D2, mpicomm=mpicomm, mpiroot=mpiroot)
            D1D2 = tmp.D1D2
            result += tmp
        results.append(result)
    return results[0].concatenate_x(*results), wang

def get_edges(corr_type='smu', bin_type='lin'):

    if bin_type == 'log':
        sedges = np.geomspace(0.01, 100., 49)
    elif bin_type == 'lin':
        sedges = np.linspace(0., 200, 201)
    else:
        raise ValueError('bin_type must be one of ["log", "lin"]')
    if corr_type == 'smu':
        edges = (sedges, np.linspace(-1., 1., 201)) #s is input edges and mu evenly spaced between -1 and 1
    elif corr_type == 'rppi':
        if bin_type == 'lin':
            edges = (sedges, sedges) #transverse and radial separations are coded to be the same here
        else:
            edges = (sedges, np.linspace(0., 40., 41))
    elif corr_type == 'theta':
        edges = np.linspace(0., 4., 101)
    else:
        raise ValueError('corr_type must be one of ["smu", "rppi", "theta"]')
    return edges

def corr_fn(file_type='npy', region='', tracer='ELG', tracer2=None, zmin=0, zmax=np.inf, rec_type=False, weight_type='default', bin_type='lin', njack=0, out_dir='.',option=None):
    if tracer2: tracer += '_' + tracer2
    if rec_type: tracer += '_' + rec_type
    if region: tracer += '_' + region
    if option:
        zmax = str(zmax) + option
    root = '{}_{}_{}_{}_{}_njack{:d}'.format(tracer, zmin, zmax, weight_type, bin_type, njack)
    if file_type == 'npy':
        return os.path.join(out_dir, 'allcounts_{}.npy'.format(root))
    return os.path.join(out_dir, '{}_{}.txt'.format(file_type, root))

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--tracer', help='tracer(s) to be selected - 2 for cross-correlation', type=str, nargs='+', default=['ELG'])
    parser.add_argument('--basedir', help='where to find catalogs', type=str, default='/global/cfs/cdirs/desi/survey/catalogs')
    parser.add_argument('--survey', help='e.g., SV3 or main', type=str, choices=['SV3', 'DA02', 'main'], default='SV3')
    parser.add_argument('--verspec', help='version for redshifts', type=str, default='guadalupe')
    parser.add_argument('--version', help='catalog version', type=str, default='test')
    parser.add_argument('--region', help='regions; by default, run on all regions', type=str, nargs='*', choices=['N', 'S'], default=None)
    parser.add_argument('--zlim', help='z-limits, or options for z-limits, e.g. "highz", "lowz", "fullonly"', type=str, nargs='*', default=None)
    parser.add_argument('--corr_type', help='correlation type', type=str, nargs='*', choices=['smu', 'rppi', 'theta'], default=['smu', 'rppi'])
    parser.add_argument('--weight_type', help='types of weights to use; use "default_angular_bitwise" for PIP with angular upweighting; "default" just uses WEIGHT column', type=str, default='default')
    parser.add_argument('--bin_type', help='binning type', type=str, choices=['log', 'lin'], default='lin')
    parser.add_argument('--nran', help='number of random files to combine together (1-18 available)', type=int, default=4)
    parser.add_argument('--split_ran_above', help='separation scale above which RR are summed over each random file;\
                                                   typically, most efficient for xi < 1, i.e. sep > 10 Mpc/h;\
                                                   see https://arxiv.org/pdf/1905.01133.pdf',
                        type=float, default=np.inf)
    parser.add_argument('--njack', help='number of jack-knife subsamples; 0 for no jack-knife error estimates', type=int, default=60)
    parser.add_argument('--nthreads', help='number of threads', type=int, default=64)
    parser.add_argument('--outdir', help='base directory for output', type=str, default=None)
    #parser.add_argument('--mpi', help='whether to use MPI', action='store_true', default=False)
    parser.add_argument('--vis', help='show plot of each xi?', action='store_true', default=False)

    #only relevant for reconstruction
    parser.add_argument('--rec_type', help='reconstruction algorithm + reconstruction convention', choices=['IFTrecsym', 'IFTreciso', 'MGrecsym', 'MGreciso'], type=str, default=None)

    setup_logging()
    args = parser.parse_args()

    mpicomm, mpiroot = None, None
    if True:#args.mpi:
        from pycorr import mpi
        mpicomm = mpi.COMM_WORLD
        mpiroot = 0

    if args.basedir == '/global/project/projectdirs/desi/users/acarnero/mtl_mock000_univ1/':
        cat_dir = args.basedir
        args.region = ['']
    else:
        cat_dir = catalog_dir(base_dir=args.basedir, survey=args.survey, verspec=args.verspec, version=args.version)
    out_dir = os.path.join(os.environ['CSCRATCH'], args.survey)
    if args.outdir is not None: out_dir = args.outdir
    tracer, tracer2 = args.tracer[0], None
    if len(args.tracer) > 1:
        tracer2 = args.tracer[1]
        if len(args.tracer) > 2:
            raise ValueError('Provide <= 2 tracers!')
    if tracer2 == tracer:
        tracer2 = None # otherwise counting of self-pairs
    catalog_kwargs = dict(tracer=tracer, tracer2=tracer2, survey=args.survey, cat_dir=cat_dir, rec_type=args.rec_type) # survey required for zdone
    distance = TabulatedDESI().comoving_radial_distance

    regions = args.region
    if regions is None:
        regions = get_regions(args.survey, rec=bool(args.rec_type))

    option = None
    if args.zlim is None:
        zlims = get_zlims(tracer, tracer2=tracer2)
    elif not args.zlim[0].replace('.', '').isdigit():
        option=args.zlim[0]
        zlims = get_zlims(tracer, tracer2=tracer2,option=option )
    else:
        zlims = [float(zlim) for zlim in args.zlim]
    zlims = list(zip(zlims[:-1], zlims[1:])) + ([(zlims[0], zlims[-1])] if len(zlims) > 2 else []) # len(zlims) == 2 == single redshift range

    rebinning_factors = [1, 4, 5, 10] if 'lin' in args.bin_type else [1,2,4]
    pi_rebinning_factors = [1, 4, 5, 10] if 'log' in args.bin_type else [1]
    if mpicomm is None or mpicomm.rank == mpiroot:
        logger.info('Computing correlation functions {} in regions {} in redshift ranges {}.'.format(args.corr_type, regions, zlims))

    for zmin, zmax in zlims:
        base_file_kwargs = dict(tracer=tracer, tracer2=tracer2, zmin=zmin, zmax=zmax, rec_type=args.rec_type, weight_type=args.weight_type, bin_type=args.bin_type, njack=args.njack, option=option)
        for region in regions:
            wang = None
            for corr_type in args.corr_type:
                if mpicomm is None or mpicomm.rank == mpiroot:
                    logger.info('Computing correlation function {} in region {} in redshift range {}.'.format(corr_type, region, (zmin, zmax)))
                edges = get_edges(corr_type=corr_type, bin_type=args.bin_type)
                result, wang = compute_correlation_function(corr_type, edges=edges, distance=distance, nrandoms=args.nran, split_randoms_above=args.split_ran_above, nthreads=args.nthreads, region=region, zlim=(zmin, zmax), weight_type=args.weight_type, njack=args.njack, wang=wang, mpicomm=mpicomm, mpiroot=mpiroot,option=option, **catalog_kwargs)
                #save pair counts
                result.save(corr_fn(file_type='npy', region=region, out_dir=os.path.join(out_dir, corr_type), **base_file_kwargs))

        # Save combination and .txt files
        for corr_type in args.corr_type:
            all_regions = regions.copy()
            if 'N' in regions and 'S' in regions:  # let's combine
                corr = sum([TwoPointCorrelationFunction.load(
                            corr_fn(file_type='npy', region=region, out_dir=os.path.join(out_dir, corr_type), **base_file_kwargs)).normalize() for region in 'NS'])
                corr.save(corr_fn(file_type='npy', region='NScomb', out_dir=os.path.join(out_dir, corr_type), **base_file_kwargs))
                all_regions.append('NScomb')
            for region in all_regions:
                txt_kwargs = base_file_kwargs.copy()
                txt_kwargs.update(region=region, out_dir=os.path.join(out_dir, corr_type))
                result = TwoPointCorrelationFunction.load(corr_fn(file_type='npy', **txt_kwargs))
                for factor in rebinning_factors:
                    #result = TwoPointEstimator.load(fn)
                    rebinned = result[:(result.shape[0]//factor)*factor:factor]
                    txt_kwargs.update(bin_type=args.bin_type+str(factor))
                    if corr_type == 'smu':
                        fn_txt = corr_fn(file_type='xismu', **txt_kwargs)
                        rebinned.save_txt(fn_txt)
                        fn_txt = corr_fn(file_type='xipoles', **txt_kwargs)
                        rebinned.save_txt(fn_txt, ells=(0, 2, 4))
                        fn_txt = corr_fn(file_type='xiwedges', **txt_kwargs)
                        rebinned.save_txt(fn_txt, wedges=(-1., -2./3, -1./3, 0., 1./3, 2./3, 1.))
                    elif corr_type == 'rppi':
                        fn_txt = corr_fn(file_type='wp', **txt_kwargs)
                        rebinned.save_txt(fn_txt, pimax=40.)
                        for pifac in pi_rebinning_factors:
                            rebinned = result[:(result.shape[0]//factor)*factor:factor,:(result.shape[1]//pifac)*pifac:pifac]
                            txt_kwargs.update(bin_type=args.bin_type+str(factor)+'_'+str(pifac))
                            fn_txt = corr_fn(file_type='xirppi', **txt_kwargs)
                            rebinned.save_txt(fn_txt)
                    elif corr_type == 'theta':
                        fn_txt = corr_fn(file_type='theta', **txt_kwargs)
                        rebinned.save_txt(fn_txt)

                    if args.vis and (mpicomm is None or mpicomm.rank == mpiroot):
                        if corr_type == 'smu':
                            sep, xis = rebinned(ells=(0, 2, 4), return_sep=True, return_std=False)
                        elif corr_type == 'rppi':
                            sep, xis = rebinned(pimax=40, return_sep=True, return_std=False)
                        else:
                            sep, xis = rebinned(return_sep=True, return_std=False)
                        if args.bin_type == 'log':
                            for xi in xis: plt.loglog(sep, xi)
                        if args.bin_type == 'lin':
                            for xi in xis: plt.plot(sep, sep**2 * xi)
                        tracers = tracer
                        if tracer2 is not None: tracers += ' x ' + tracer2
                        plt.title('{} {:.2f} < z {:.2f} in {}'.format(tracers, zmin, zmax, region))
                        plt.show()

max_stars_repo_name: echaussidon/LSS

max_issues_repo_name: echaussidon/LSS

max_forks_repo_name: echaussidon/LSS

reason: import numpy,from astropy


text:
import os
import tempfile
import numpy as np
import pandas as pd

import matplotlib.pyplot as plt
from urllib import request

ROOT_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..')
DATA_DIR = os.path.join(ROOT_DIR, 'data')
RESULTS_DIR = os.path.join(ROOT_DIR, 'results')
FIGURES_DIR = os.path.join(ROOT_DIR, 'figures')

CMU_URL = 'http://www.cs.cmu.edu/~keystroke/'
CMU_FNAME = 'DSL-StrongPasswordData.csv'

BIOSIG_URL = 'http://www.epaymentbiometrics.ensicaen.fr/wp-content/uploads/2015/04/'
BIOSIG_FNAME = 'greyc-nislab-keystroke-benchmark-dataset.xls'

def maybe_mkdir(dirpath=None):
    if dirpath is None:
        dirpath = tempfile.mkdtemp()
    elif not os.path.exists(dirpath):
        os.mkdir(dirpath)
    return dirpath

def maybe_download(source_url, filename, destdir):
    """Download the data from Yann's website, unless it's already here."""
    if not os.path.exists(destdir):
        os.makedirs(destdir, mode=0o755)
    filepath = os.path.join(destdir, filename)
    if not os.path.exists(filepath):
        filepath, _ = request.urlretrieve(source_url + filename, filepath)
        with open(filepath) as f:
            size = f.seek(0, 2)
        print('Successfully downloaded', filename, size, 'bytes.')
    return filepath

def load_cmu():
    """
    Load the CMU benchmark keystroke dataset as a dataframe with index: subject, repetition
    """
    fname = maybe_download(CMU_URL, CMU_FNAME, DATA_DIR)
    df = pd.read_csv(fname, index_col=[0, 1, 2])
    repetition = (df.index.get_level_values(1) - 1) * 50 + df.index.get_level_values(2)
    df = df.reset_index(level=[1, 2], drop=True)
    df = df.set_index(repetition, append=True)
    df.index.names = ['subject', 'repetition']
    return df

def load_biosig(i, hands=2):
    """
    Load the Biosig keystroke dataset (given the number of typing hands) as a dataframe with index: subject, repetition
    """
    fname = maybe_download(BIOSIG_URL, BIOSIG_FNAME, DATA_DIR)
    df = pd.read_excel(fname, sheetname=i)

    df = df[df['Class'] == hands]

    df['repetition'] = 0
    df['subject'] = df['User_ID']

    def make_repetitions(x):
        x['repetition'] = np.arange(1, len(x) + 1)
        return x

    df = df.groupby(['subject', 'repetition']).apply(make_repetitions)

    def make_features(x):
        return pd.Series(x['Keystroke Template Vector'].split()).astype(int) / 1e7

    features = df.apply(make_features, axis=1)

    features['subject'] = df['subject']
    features['repetition'] = df['repetition']

    features = features.set_index(['subject', 'repetition'])

    return features

def save_fig(name, ext='pdf'):
    plt.savefig(os.path.join(FIGURES_DIR, name + '.%s' % ext), bbox_inches='tight')
    plt.close()
    return

def save_results(df, name):
    df.to_csv(os.path.join(RESULTS_DIR, name + '.csv'))
    return

def load_results(name, **args):
    return pd.read_csv(os.path.join(RESULTS_DIR, name + '.csv'), **args)

ext: py

max_stars_repo_name: vmonaco/transfer-encoder

max_issues_repo_name: vmonaco/transfer-encoder

max_forks_repo_name: vmonaco/transfer-encoder

reason: import numpy


text:
import pandas as pd
import numpy as np
import dask.array as da
import dask.dataframe as dd
import dask.delayed as delayed
import time
import math
import itertools
from netCDF4 import Dataset
import os,datetime,sys,fnmatch
import matplotlib.pyplot as plt
import dask
from dask_jobqueue import SLURMCluster
from dask.distributed import as_completed
from dask.distributed import Client
from dask.distributed import wait

def read_MODIS_level2_data(MOD06_file,MOD03_file):
    #print('Reading The Cloud Mask From MOD06_L2 Product:')
    myd06 = Dataset(MOD06_file, "r")
    CM = myd06.variables["Cloud_Mask_1km"][:,:,:] # Reading Specific Variable 'Cloud_Mask_1km'.
    CM   = (np.array(CM[:,:,0],dtype='byte') & 0b00000110) >>1
    CM = np.array(CM).byteswap().newbyteorder()
   # print('The Level-2 Cloud Mask Array Shape',CM.shape)
    print(' ')

    myd03 = Dataset(MOD03_file, "r")
    #print('Reading The Latitude-Longitude From MOD03 Product:')
    latitude = myd03.variables["Latitude"][:,:] # Reading Specific Variable 'Latitude'.
    latitude = np.array(latitude).byteswap().newbyteorder() # Addressing Byteswap For Big Endian Error.
    longitude = myd03.variables["Longitude"][:,:] # Reading Specific Variable 'Longitude'.
    longitude = np.array(longitude).byteswap().newbyteorder() # Addressing Byteswap For Big Endian Error
    #print('The Level2 Latitude-Longitude Array Shape',latitude.shape)
    print(' ')

    return latitude,longitude,CM

def countzero(x, axis=1):
    #print(x)
    count0 = 0
    count1 = 0
    for i in x:
        if i <= 1:
            count0 +=1
    #print(count0/len(x))
    return count0/len(x)

MOD03_path ='/home/dprakas1/jianwu_common/MODIS_Aggregation/deepak-code/test/' #'input-data/MYD3'
MOD06_path ='/home/dprakas1/jianwu_common/MODIS_Aggregation/deepak-code/test/' #'input-data/MYD6'

MOD03_filepath = 'MYD03.A*.hdf'
MOD06_filepath = 'MYD06_L2.A*.hdf'
MOD03_filename, MOD06_filename =[],[]
MOD03_filename2, MOD06_filename2 =[],[]

for MOD06_filelist in  os.listdir(MOD06_path):
    if fnmatch.fnmatch(MOD06_filelist, MOD06_filepath):
        MOD06_filename = MOD06_filelist
        MOD06_filename2.append(MOD06_filelist)

for MOD03_filelist in  os.listdir(MOD03_path):
    if fnmatch.fnmatch(MOD03_filelist, MOD03_filepath):
        MOD03_filename = MOD03_filelist
        MOD03_filename2.append(MOD03_filelist)

#if MOD03_filename and MOD06_filename:
#    print('Reading Level 2 GeoLocation & Cloud Data')
#    Lat,Lon,CM = read_MODIS_level2_data(MOD06_path+MOD06_filename,MOD03_path+MOD03_filename)

print('The Number Of Files In The MODO3 List: ')
print(len(MOD03_filename2))
print(' ')
print('The Number Of Files In The MODO6_L2 List: ')
print(len(MOD06_filename2))
print(' ')

if __name__ == '__main__':

    cluster = SLURMCluster(cores=1, memory='100GB', project='pi_jianwu', queue='high_mem', walltime='02:00:00', job_extra=['--exclusive', '--qos=medium+'])

    cluster.scale(3)

    client = Client(cluster)

    dask.config.set(temporary_directory='/home/dprakas1/jianwu_common/MODIS_Aggregation/deepak-code/test-tmp')

    t0 = time.time()

    cm = np.zeros((2030,1354), dtype=np.float32)

    for MOD06_file in MOD06_filename2:
        MOD06_file2 = MOD06_path + MOD06_file
        myd06 = Dataset(MOD06_file2, "r")
        CM = myd06.variables["Cloud_Mask_1km"][:,:,0]# Reading Specific Variable 'Cloud_Mask_1km'.
        CM = (np.array(CM,dtype='byte') & 0b00000110) >>1
        CM = np.array(CM).byteswap().newbyteorder()
        cm = np.concatenate((cm,CM))

    print('The Cloud Mask Array Shape Is: ',cm.shape)

    lat = np.zeros((2030,1354), dtype=np.float32)
    lon = np.zeros((2030,1354), dtype=np.float32)
    for MOD03_file in MOD03_filename2:
        MOD03_file2 = MOD03_path + MOD03_file
        myd03 = Dataset(MOD03_file2, "r")
        latitude = myd03.variables["Latitude"][:,:] # Reading Specific Variable 'Latitude'.
        lat = np.concatenate((lat,latitude))
        longitude = myd03.variables["Longitude"][:,:] # Reading Specific Variable 'Longitude'.
        lon = np.concatenate((lon,longitude))

    print('Longitude Shape Is: ',lon.shape)
    print('Latitude Shape Is: ',lat.shape)

    cm=np.ravel(cm)
    lat=np.ravel(lat)
    lon=np.ravel(lon)

    lon=lon.astype(int)
    lat=lat.astype(int)
    cm=cm.astype(int)

    d = {'Latitude' :  pd.Series(lat), 'Longitude' : pd.Series(lon),'CM':pd.Series(cm)} #Reading the values in a pandas dataframe
    df = pd.DataFrame(d,columns=['Latitude', 'Longitude','CM']) #Creating a pandas dataframe

    df2=df.groupby(['Longitude','Latitude']).CM.apply(countzero)
    df3=df2.reset_index()

    combs=[]
    for x in range(-89,91):
        for y in range(-179,181):
            combs.append((x, y))

    df_1=pd.DataFrame(combs)
    df_1.columns=['Latitude','Longitude']

    #df4=pd.merge(df_1, df3,on=('Longitude','Latitude'), how='left')

    #df5=df4['CM'].values

    #b=df5.reshape(180,360)

    client.close()

    t1 = time.time()
    total = t1-t0
    print("total execution time (Seconds):" + str(total))

    #plt.figure(figsize=(14,7))
    #plt.contourf(range(-180,180), range(-90,90), b, 100, cmap = "jet")
    #plt.xlabel("Longitude", fontsize = 14)
    #plt.ylabel("Latitude", fontsize = 14)
    #plt.title("Level 3 Cloud Fraction Aggregation for January 2008", fontsize = 16)
    #plt.colorbar()
    #plt.savefig("pandas_oneday.png")

max_stars_repo_name: supriyascode/MODIS-Aggregation

max_issues_repo_name: supriyascode/MODIS-Aggregation

max_forks_repo_name: supriyascode/MODIS-Aggregation


text:
# Copyright (c) 2016, the Cap authors.
#
# This file is subject to the Modified BSD License and may not be distributed
# without copyright and license information. Please refer to the file LICENSE
# for the text and further information on this license.

from numpy import trapz, count_nonzero, array, append, power, argsort
from matplotlib import pyplot
from copy import copy
from .PyCap import PropertyTree
from .charge_discharge import Charge, Discharge
from .data_helpers import initialize_data, save_data
from .observer_pattern import Experiment, Observer

__all__ = ['plot_ragone', 'retrieve_performance_data',
           'RagonePlot', 'RagoneAnalysis']

def plot_ragone(data, figure=None, ls='r-o'):
    power = data['power']
    energy = data['energy']
    plot_linewidth = 3
    label_fontsize = 30
    tick_fontsize = 20
    if figure:
        pyplot.figure(figure.number)
    else:
        pyplot.figure(figsize=(16, 12))
    pyplot.plot(power, energy, ls, lw=plot_linewidth)
    pyplot.xscale('log')
    pyplot.yscale('log')
    pyplot.xlabel(r'$\mathrm{Power\  [W]}$', fontsize=label_fontsize)
    pyplot.ylabel(r'$\mathrm{Energy\ [J]}$', fontsize=label_fontsize)
    pyplot.gca().get_xaxis().set_tick_params(labelsize=tick_fontsize)
    pyplot.gca().get_yaxis().set_tick_params(labelsize=tick_fontsize)

def run_discharge(device, ptree):
    data = initialize_data()

    # (re)charge the device
    initial_voltage = ptree.get_double('initial_voltage')

    charge_database = PropertyTree()
    charge_database.put_string('charge_mode', 'constant_current')
    charge_database.put_double('charge_current', 0.5)
    charge_database.put_string('charge_stop_at_1', 'voltage_greater_than')
    charge_database.put_double('charge_voltage_limit', initial_voltage)
    charge_database.put_bool('charge_voltage_finish', True)
    charge_database.put_double('charge_voltage_finish_current_limit', 1e-6)
    charge_database.put_double('charge_voltage_finish_max_time', 600)
    charge_database.put_double('charge_rest_time', 0)
    charge_database.put_double('time_step', 0.1)

    charge = Charge(charge_database)
    charge.run(device, data)

    data['time'] -= data['time'][-1]

    # discharge at constant power
    discharge_power = ptree.get_double('discharge_power')
    final_voltage = ptree.get_double('final_voltage')
    time_step = ptree.get_double('time_step')

    discharge_database = PropertyTree()
    discharge_database.put_string('discharge_mode', 'constant_power')
    discharge_database.put_double('discharge_power', discharge_power)
    discharge_database.put_string('discharge_stop_at_1', 'voltage_less_than')
    discharge_database.put_double('discharge_voltage_limit', final_voltage)
    discharge_database.put_double('discharge_rest_time', 10 * time_step)
    discharge_database.put_double('time_step', time_step)

    discharge = Discharge(discharge_database)
    discharge.run(device, data)

    return data

def examine_discharge(data):
    time = data['time']
    current = data['current']
    voltage = data['voltage']
    power = current[:] * voltage[:]
    mask = time[:] <= 0
    energy_in = trapz(power[mask], time[mask])
    mask = time[:] >= 0
    energy_out = trapz(power[mask], time[mask])

    return [energy_in, energy_out]

def retrieve_performance_data(fin):
    path = 'ragone_chart_data'
    performance_data = {'power': array([], dtype=float),
                        'energy': array([], dtype=float)}
    for key in fin[path].keys():
        start = key.find('=') + 1
        end = key.find('W')
        power = float(key[start:end])
        for measurement in ['second', 'first']:
            if measurement in fin[path][key].keys():
                discharge_data = fin[path][key][measurement]
                (energy_in, energy_out) = examine_discharge(discharge_data)
                performance_data['power'] = append(performance_data['power'],
                                                   power)
                performance_data['energy'] = append(performance_data['energy'],
                                                    -energy_out)
                break
    sort = argsort(performance_data['power'])
    performance_data['power'] = performance_data['power'][sort]
    performance_data['energy'] = performance_data['energy'][sort]

    return performance_data

class RagonePlot(Observer):
    '''Ragone plot.

    Plots the values of specific energy versus specific power. Both axes are
    logarithmic.
    '''
    def __new__(cls, *args, **kwargs):
        return object.__new__(RagonePlot)

    def __init__(self, filename=None):
        self._figure = pyplot.figure(figsize=(14, 14))
        self._filename = filename

    def update(self, subject, *args, **kwargs):
        plot_ragone(subject._data, figure=self._figure)
        if self._filename is not None:
            pyplot.savefig(self._filename, bbox_inches='tight')
Observer._builders['RagonePlot'] = RagonePlot

# TODO: Probably want to make that class more general to do discharges at
# constant current as well.
# For batteries plot voltage vs capacity

class RagoneAnalysis(Experiment):
    '''Record a Ragone plot

    Performs a series of discharge at various rates and neasures the energy.

    Attributes
    ----------
    _discharge_power_lower_limit : float
    _discharge_upper_lower_limit : float
    _steps_per_decade : int
    _ptree : PropertyTree
    _data : dict
        Stores power and energy as numpy.array(s) of floating point numbers.

    See Also
    --------
    RagonePlot
    '''
    def __new__(cls, *args, **kwargs):
        return object.__new__(RagoneAnalysis)

    def __init__(self, ptree):
        Experiment.__init__(self)
        self._discharge_power_lower_limit = ptree.get_double(
            'discharge_power_lower_limit')
        self._discharge_power_upper_limit = ptree.get_double(
            'discharge_power_upper_limit')
        self._steps_per_decade = ptree.get_int('steps_per_decade')
        self._min_steps_per_discharge = ptree.get_int(
            'min_steps_per_discharge')
        self._max_steps_per_discharge = ptree.get_int(
            'max_steps_per_discharge')
        self._time_step_initial_guess = ptree.get_double('time_step')
        self._ptree = copy(ptree)
        self.reset()

    def reset(self):
        self._ptree.put_double('time_step', self._time_step_initial_guess)
        self._data = {
            'energy': array([], dtype=float),
            'power': array([], dtype=float)
        }

    def run(self, device, fout=None):
        discharge_power = self._discharge_power_lower_limit
        while discharge_power <= self._discharge_power_upper_limit:
            # print discharge_power
            self._ptree.put_double('discharge_power', discharge_power)
            try:
                # this loop control the number of time steps in the discharge
                for measurement in ['first', 'second']:
                    data = run_discharge(device, self._ptree)
                    if fout:
                        path = 'ragone_chart_data'
                        path += '/power=' + str(discharge_power) + 'W'
                        path += '/' + measurement
                        save_data(data, path, fout)
                    energy_in, energy_out = examine_discharge(data)
                    steps = count_nonzero(data['time'] > 0)
                    if steps >= self._min_steps_per_discharge:
                        break
                    else:
                        time_step = data['time'][-1]
                        time_step /= self._max_steps_per_discharge
                        self._ptree.put_double('time_step', time_step)

            except RuntimeError:
                print('Failed to discharge at {0} watt'.format(
                    discharge_power))
                break
            self._data['energy'] = append(self._data['energy'],
                                          -energy_out)
            self._data['power'] = append(self._data['power'],
                                         discharge_power)
            discharge_power *= power(10.0, 1.0 / self._steps_per_decade)
            self.notify()
Experiment._builders['RagoneAnalysis'] = RagoneAnalysis

max_stars_repo_name: iiscsahoo/EnergyData

max_issues_repo_name: iiscsahoo/EnergyData

max_forks_repo_name: iiscsahoo/EnergyData

reason: from numpy


text:
from capreolus.extractor import Extractor
from capreolus.tokenizer import Tokenizer
from capreolus.utils.common import padlist
from capreolus.utils.loginit import get_logger
from tqdm import tqdm
import numpy as np
from collections import Counter, defaultdict

logger = get_logger(__name__)  # pylint: disable=invalid-name

class BagOfWords(Extractor):
    """ Bag of Words (or bag of trigrams when `datamode=trigram`) extractor. Used with the DSSM reranker. """

    pad = 0
    pad_tok = "<pad>"
    tokenizer_name = "anserini"

    def __init__(self, *args, **kwargs):
        super(BagOfWords, self).__init__(*args, **kwargs)
        self.itos = {self.pad: self.pad_tok}
        self.stoi = {self.pad_tok: self.pad}
        self.embeddings = self.stoi
        self.idf = defaultdict(lambda: 0)

    @staticmethod
    def config():
        datamode = "unigram"  # DSSM requires trigram vector as input
        keepstops = False
        return locals().copy()  # ignored by sacred

    def build_unigram_stoi(self, toks_list, keepstops, calculate_idf):
        nexti = len(self.stoi)
        for toks in tqdm(toks_list, desc="text to embed", unit_scale=True):
            for tok in toks:
                if tok not in self.stoi:
                    self.stoi[tok] = nexti
                    nexti += 1

                    if calculate_idf:
                        # TODO this will break if self.tokenizer is more restrictive (eg eliminates stops but index tokenizer does not)
                        # index_toks = self.index.tokenizer(self.index).tokenize(tok)
                        # index_tok = "" if len(index_toks) == 0 else index_toks[0]
                        # self.idf[tok] = self.index.getidf(index_tok)

                        # TODO: This is a temp hack. Refactor so that we query from an index where stop words are kept
                        self.idf[tok] = self.index.getidf(tok)

    def get_trigrams_for_toks(self, toks_list):
        return [("#%s#" % tok)[i : i + 3] for tok in toks_list for i in range(len(tok))]

    def build_trigram_stoi(self, toks_list, keepstops, calculate_idf):
        nexti = len(self.stoi)
        trigrams_list = []
        for toks in toks_list:
            trigrams_list.extend(self.get_trigrams_for_toks(toks))

        for trigram in trigrams_list:
            if trigram not in self.stoi:
                self.stoi[trigram] = nexti
                nexti += 1

    def build_stoi(self, toks_list, keepstops, calculate_idf):
        """
        Builds an stoi dict that stores a unique number for each token
        toks_list - An array of arrays. Each element (i.e array) in toks_list represent tokens of an entire document
        """
        if self.index is None:
            raise ValueError("Index cannot be None")

        if self.pipeline_config["datamode"] == "unigram":
            self.build_unigram_stoi(toks_list, keepstops, calculate_idf)
        elif self.pipeline_config["datamode"] == "trigram":
            self.build_trigram_stoi(toks_list, keepstops, calculate_idf)
        else:
            raise Exception("Unknown datamode")

    def build_from_benchmark(self, keepstops, *args, **kwargs):
        if not all([self.collection, self.benchmark, self.index]):
            raise ValueError("The Feature class was not initialized with a collection, benchmark and index.")

        tokenizer = Tokenizer.ALL[self.tokenizer_name].get_tokenizer_instance(self.index, keepstops=keepstops)
        tokenizer.create()
        self.index.open()

        qids = set(self.benchmark.pred_pairs.keys()).union(self.benchmark.train_pairs.keys())
        doc_ids = set()
        for qdocs in self.benchmark.pred_pairs.values():
            doc_ids.update(qdocs)
        for qdocs in self.benchmark.train_pairs.values():
            doc_ids.update(qdocs)

        queries = self.collection.topics[self.benchmark.query_type]
        query_text_list = [queries[qid] for qid in qids]
        query_tocs_list = [tokenizer.tokenize(query) for query in query_text_list]

        logger.info("loading %s queries", len(queries))
        self.build_stoi(query_tocs_list, keepstops, True)

        self.doc_id_to_doc_toks = tokenizer.tokenizedocs(doc_ids)
        self.build_stoi([toks for doc_id, toks in self.doc_id_to_doc_toks.items()], keepstops, False)
        self.embeddings = self.stoi
        self.itos = {i: s for s, i in self.stoi.items()}
        logger.debug("The vocabulary size: %s", len(self.stoi))

    def transform_qid_posdocid_negdocid(self, q_id, posdoc_id, negdoc_id=None):
        if not all([self.collection, self.benchmark, self.index]):
            raise ValueError("The Feature class was not initialized with a collection, benchmark and index")

        tokenizer = Tokenizer.ALL[self.tokenizer_name].get_tokenizer_instance(
            self.index, keepstops=self.pipeline_config["keepstops"]
        )
        queries = self.collection.topics[self.benchmark.query_type]

        posdoc_toks = self.doc_id_to_doc_toks.get(posdoc_id)
        query_toks = tokenizer.tokenize(queries[q_id])

        if not posdoc_toks:
            logger.debug("missing docid %s", posdoc_id)
            return None
        transformed_query = self.transform_txt(query_toks, self.pipeline_config["maxqlen"])
        idfs = [self.idf[self.itos[tok]] for tok, count in enumerate(transformed_query)]
        transformed = {
            "qid": q_id,
            "posdocid": posdoc_id,
            "negdocid": negdoc_id,
            "query": transformed_query,
            "posdoc": self.transform_txt(posdoc_toks, self.pipeline_config["maxdoclen"]),
            "query_idf": np.array(idfs, dtype=np.float32),
        }

        if negdoc_id is not None:
            negdoc_toks = self.doc_id_to_doc_toks.get(negdoc_id)
            if not negdoc_toks:
                logger.debug("missing docid %s", negdoc_id)
                return None

            transformed["negdoc"] = self.transform_txt(negdoc_toks, self.pipeline_config["maxdoclen"])

        return transformed

    def transform_txt(self, term_list, maxlen):
        nvocab = len(self.stoi)
        bog_txt = np.zeros(nvocab, dtype=np.float32)

        if self.pipeline_config["datamode"] == "unigram":
            toks = [self.stoi.get(term, 0) for term in term_list]
            tok_counts = Counter(toks)
        elif self.pipeline_config["datamode"] == "trigram":
            trigrams = self.get_trigrams_for_toks(term_list)
            toks = [self.stoi.get(trigram, 0) for trigram in trigrams]
            tok_counts = Counter(toks)
        else:
            raise Exception("Unknown datamode")

        for tok, count in tok_counts.items():
            bog_txt[tok] = count

        return bog_txt

max_stars_repo_name: bpiwowar/capreolus-xpm

max_issues_repo_name: bpiwowar/capreolus-xpm

max_forks_repo_name: bpiwowar/capreolus-xpm


text:
"""
Tests for metrics.py

"""
import pandas as pd
from tests.datasets_to_test import *
from predictions_analyzer.models import *
from predictions_analyzer.metrics import *

import numpy as np
import pandas as pd
from typing import List

def test_classification_models():
    all_models = AllClassificationModels()

    for model in all_models._get_models_list():
        pass

def test_ScikitModels():

    for dataset in ClassificationDatasets():

        print("\n")
        print(dataset.name)
        print(dataset.path)

        for model in AllClassificationModels():

            print("\n")
            print(model.name)
            model.fit(dataset.X_val, dataset.y_val)
            preds = model.predict(dataset.X_test)
            print(preds.shape)
            print(dataset.y_test.shape)

            model.report_speeds()

            classification_metrics = AllClassificationMetrics()

            for metric in classification_metrics:
                print(metric.name)
                the_score = metric.score(dataset.y_test, preds)
                print(the_score)

def test_AllRegressionModels():
    allreg = AllRegressionModels()
    print(allreg)

max_stars_repo_name: ThomasMeli/MeliML

max_issues_repo_name: ThomasMeli/MeliML

max_forks_repo_name: ThomasMeli/MeliML


text:
import numpy as np
import pandas as pd
import scipy as sp
import scipy.ndimage
from scipy import ndimage
np.seterr(divide='ignore', invalid='ignore')

def RF_patterns(isveg, rvl_params):
    """
    inputs:
      isveg: binary array of vegetation
      rvl_params: dictionary of parameters specifying how features are computed:
         edge : "padding" added to all distances to boundaries, to reduce edge effects.
         saturate : maximum size of elements in the feature arrays
         gsigma : smoothing lengthscales (in the Gaussian filter)

    output:
      pattern_dict: a dictionary of feature maps, each with shape (ncol x nrow)

     does not include: d2divide (ncol x nrow), the distance to divide
    """
    isvegc = np.array(isveg, dtype = float)
    ncol = isvegc.shape[0]
    nrow = isvegc.shape[1]

    edge = int(rvl_params['edge'])
    saturate = int(rvl_params['saturate'])
    gsigma = rvl_params['gsigma']

    if gsigma == 0:
        gsigma = []
    elif type(gsigma) == int:
        gsigma = [gsigma]

    d2uB = func_d2uB(isvegc, edge, saturate)
    d2dB = func_d2dB(isvegc, edge, saturate)

    d2lB = func_d2lB(isvegc, edge, saturate)
    d2rB = func_d2rB(isvegc, edge, saturate)

    d2xB = np.minimum(d2lB, d2rB)
    d2xB[d2xB > saturate] = saturate

    d2uV = func_d2uV(isvegc, edge, saturate)
    d2dV = func_d2dV(isvegc, edge, saturate)

    d2lV = func_d2lV(isvegc, edge, saturate)
    d2rV = func_d2rV(isvegc, edge, saturate)
    d2xV = np.minimum(d2lV, d2rV)

    patchL,patchLB = get_patchL(isvegc, saturate)
    bareL, bareLV = get_bareL(isvegc, saturate)

    # assemble base pattern_dict
    pattern_dict = {'isvegc' : isvegc,
                    'd2uB' : d2uB,
                    'd2dB' : d2dB,
                    'd2xB' : d2xB,
                    'd2uV' : d2uV,
                    'd2dV' : d2dV,
                    'd2xV' : d2xV,
                    'patchLB' : patchLB,
                    'bareLV' : bareLV,
                  }

    pos_sigma = [gs for gs in gsigma if gs > 0]

    for gs in pos_sigma:
        for key in ['d2uV','d2dV','d2xV','bareLV']:
            # smooth bare features
            pattern_dict[key + '_s{0}'.format(gs)] =   smoothB(pattern_dict[key], isvegc, gs)

        for key in ['d2uB', 'd2dB', 'd2xB','patchLB']:
            # smooth vegetion features
            pattern_dict[key + '_s{0}'.format(gs)] =   smoothV(pattern_dict[key], isvegc, gs)

    # delete unsmoothed features if 0 not in gsigma
    if 0 not in gsigma:

        for key in ['d2uV','d2dV','d2xV','bareLV',  \
                    'd2uB', 'd2dB', 'd2xB','patchLB']:
            del pattern_dict[key]

    return pattern_dict

def get_feature_matrix(sim, rvl_params, target = None):
    """
    Compute feature matrix for an input vegetation field (sim['isvegc'])

    input:
         sim: dictionary containing vegetation map (isvegc) and grid resolution (dx)

     veg features, bare features

    tasks:
      get pattern_dict from RF_patterns  (in ravel_fxns_RF.py)

      for all the keys in features:
        if in pattern_dict, np.ravel and add to pattern_ravel
        if in sim dict (non-array), add as list with equiv size
        skip 'local_path' key
        add 'd2divide' manually (units : grid cell)

    """
    isvegc = sim['isvegc'].astype(float)
    dx = sim['dx']

    ncol = isvegc.shape[0]
    nrow = isvegc.shape[1]

    xc = np.arange(0, ncol*dx, dx)  + dx/2
    yc = np.arange(0, nrow*dx, dx)  + dx/2
    xc, yc = np.meshgrid(xc, yc)

    xc = xc.T
    yc = yc.T

    pattern_dict = RF_patterns(isvegc, rvl_params)
    features =  pattern_dict.keys() + ["d2divide"]

    pattern_ravel = {}
    eqv_size =  int(np.size(isvegc.ravel()))

    for key in features:
        if key in pattern_dict:
            pattern_ravel[key] = np.ravel(pattern_dict[key])

        elif key in sim.keys() and type(sim[key]) != np.ndarray:
            pattern_ravel[key] = np.array([sim[key]]*eqv_size)

        elif key =='d2divide':
            pattern_ravel[key] = (nrow - yc/dx).ravel() # ( nrow - yc/sim.dx).ravel()

        else:
            print key

    pattern_ravel['isvegc'] = np.ravel(isvegc)
    pattern_ravel['fV'] = np.array([np.mean(isvegc)]*eqv_size)

    if target:
        pattern_ravel[target] = np.ravel(sim[target])

    pattern_ravel = pd.DataFrame(pattern_ravel)

    return pattern_ravel

def smoothB(U, isvegc, gsigma):
    """
    Smooths array U over bare soil cells with a Gaussian filter,
      ignoring vegetated cells

    Inputs:
        U : array to smooth over
        isvegc : vegetated cells
        gsigma : standard deviation of the Gaussian kernel, determining length scale of smoothing
    Outputs:
        Z : array containing U smoothed over the bare soil cells (ignoring vegetated)
    """
    U = U.astype(float)
    U[isvegc == 1] = np.nan
    V=U.copy()
    V[U!=U]=0
    VV=sp.ndimage.gaussian_filter(V,gsigma)

    W=0*U.copy()+1
    W[U!=U]=0
    WW=sp.ndimage.gaussian_filter(W,gsigma)

    Z=VV/WW
    Z = Z.astype(int)
    Z[isvegc ==1] = 0
    return Z

def smoothV(U, isvegc, gsigma):
    """
    Smooths array U over vegetated cells with a Gaussian filter,
      ignoring bare soil cells

    Inputs:
        U : array to smooth over
        isvegc : vegetated cells
        gsigma : standard deviation of the Gaussian kernel, determining length scale of smoothing
    Outputs:
        Z : array containing U smoothed over the vegetated cells
    """
    U = U.astype(float)
    U[isvegc == 0] = np.nan
    V=U.copy()
    V[U!=U]=0
    VV=sp.ndimage.gaussian_filter(V,gsigma)

    W=0*U.copy()+1
    W[U!=U]=0
    WW=sp.ndimage.gaussian_filter(W,gsigma)

    Z=VV/WW
    Z = Z.astype(int)
    Z[isvegc ==0] = 0
    return Z

def func_d2uB(isvegc, edge, saturate):
    """
    Computes the distance to nearest upslope bare soil cell
    Input:
        isvegc: binary array of vegetation locations
    Returns:
        df1 : array containing the distance to the nearest upslope bare cell
             =  0 for bare cells
             =  1 for veg cells with a neighboring bare cell upslope
             >  1 for veg cells with a bare cell further upslope
    """
    ncol = isvegc.shape[0]
    nrow = isvegc.shape[1]
    arr = np.ones([ncol+ 2*edge, nrow + 2*edge]).T
    arr[edge:-edge, edge:-edge] = np.flipud(isvegc.T)
    a = pd.DataFrame(arr) != 0

    df1 = a.cumsum()-a.cumsum().where(~a).ffill().fillna(0).astype(int)
    df1 = np.array(df1)[edge:-edge, edge:-edge]
    df1 = np.flipud(df1).T
    df1[isvegc == 0] = 0

    df1[df1>saturate] = saturate

    return df1

def func_d2uV(isvegc, edge, saturate):
    """
    Computes the distance to nearest upslope vegetated cell
    Input:
        isvegc: binary array of vegetation locations
    Returns:
        df1 : array containing the distance to the nearest upslope vegetated cell
             =  0 for vegetated cells
             =  1 for bare cells with a neighboring veg cell upslope
             >  1 for bare cells with a veg cell further upslope
    """

    ncol = isvegc.shape[0]
    nrow = isvegc.shape[1]
    arr = np.ones([ncol+ 2*edge, nrow + 2*edge]).T
    arr[edge:-edge, edge:-edge] = 1 - np.flipud(isvegc.T)
    a = pd.DataFrame(arr) != 0

    df1 = a.cumsum()-a.cumsum().where(~a).ffill().fillna(0).astype(int)
    df1 = np.array(df1)[edge:-edge, edge:-edge]
    df1 = np.flipud(df1).T
    df1[isvegc == 1] = 0

    df1[df1>saturate] = saturate
    return df1

def func_d2dB(isvegc, edge, saturate):
    """
    Computes the distance to nearest downslope bare cell
    Input:
        isvegc: binary array of vegetation locations
    Returns:
        df1 : array containing the distance to the nearest downslope bare cell
    """
    ncol = isvegc.shape[0]
    nrow = isvegc.shape[1]

    arr = np.ones([ncol+ 2*edge, nrow + 2*edge]).T
    arr[edge:-edge, edge:-edge] = isvegc.T
    a = pd.DataFrame(arr) != 0

    df1 = a.cumsum()-a.cumsum().where(~a).ffill().fillna(0).astype(int)
    df1 = np.array(df1)[edge:-edge, edge:-edge]
    df1 = df1.T
    df1[isvegc == 0] = 0

    df1[df1>saturate] = saturate

    return df1

def func_d2dV(isvegc, edge, saturate):
    """
    Computes the distance to nearest downslope vegetated cell
    Input:
        isvegc: binary array of vegetation locations
    Returns:
        df1 : array containing the distance to the nearest downslope vegetated cell
    """

    ncol = isvegc.shape[0]
    nrow = isvegc.shape[1]
    arr = np.ones([ncol+ 2*edge, nrow + 2*edge]).T
    arr[edge:-edge, edge:-edge] = 1 - isvegc.T
    a = pd.DataFrame(arr) != 0

    df1 = a.cumsum()-a.cumsum().where(~a).ffill().fillna(0).astype(int)
    df1 = np.array(df1)[edge:-edge, edge:-edge]
    df1 =  df1.T
    df1[isvegc == 1] = 0

    df1[df1>saturate] = saturate

    return df1

def func_d2lB(isvegc, edge, saturate):
    """
    input:
      isvegc : [ncol x nrow] array of vegetation field

    output :
      d2lB : [ncol x nrow] array of distance to nearest left bare
        =  0 for bare cells
        =  1 for veg cells with a bare cell immediately left
    """

    ncol = isvegc.shape[0]
    nrow = isvegc.shape[1]
    arr = np.ones([ncol+ 2*edge, nrow + 2*edge])
    arr[edge:-edge, edge:-edge] = isvegc
    a = pd.DataFrame(arr) != 0

    df1 = a.cumsum()-a.cumsum().where(~a).ffill().fillna(0).astype(int)
    df1 = np.array(df1)[edge:-edge, edge:-edge]
    df1[isvegc == 0] = 0

    df1[df1>saturate] = saturate

    return df1

def func_d2lV(isvegc, edge, saturate):
    """
    input:
      isvegc : [ncol x nrow] array of vegetation field

    output :
      d2lV : [ncol x nrow] array, distance to nearest veg cell to the left
        =  0 for veg cells
        =  1 for bare cells with a veg cell immediately left
    """

    ncol = isvegc.shape[0]
    nrow = isvegc.shape[1]
    arr = np.ones([ncol+ 2*edge, nrow + 2*edge])
    arr[edge:-edge, edge:-edge] = 1 - isvegc
    a = pd.DataFrame(arr) != 0

    df1 = a.cumsum()-a.cumsum().where(~a).ffill().fillna(0).astype(int)
    df1 = np.array(df1)[edge:-edge, edge:-edge]
    df1[isvegc == 1] = 0

    df1[df1>saturate] = saturate

    return df1

def func_d2rB(isvegc, edge, saturate):
    """
    input:
      isvegc : [ncol x nrow] array of vegetation field

    output :
      d2rB : [ncol x nrow] array;  distance to nearest bare cell to right
        =  0 for bare cells
        =  1 for veg cells with a veg cell immediately to right
    """

    ncol = isvegc.shape[0]
    nrow = isvegc.shape[1]
    arr = np.ones([ncol+ 2*edge, nrow + 2*edge])
    arr[edge:-edge, edge:-edge] = np.flipud(isvegc)
    a = pd.DataFrame(arr) != 0

    df1 = a.cumsum()-a.cumsum().where(~a).ffill().fillna(0).astype(int)
    df1 = np.array(df1)[edge:-edge, edge:-edge]
    df1 = np.flipud(df1)
    df1[isvegc == 0] = 0

    df1[df1>saturate] = saturate

    return df1

def func_d2rV(isvegc, edge, saturate):
    """
    input:
      isvegc : [ncol x nrow] array of vegetation field

    output :
      d2rV : [ncol x nrow] array;  distance to nearest veg cell to right
        =  0 for veg cells
        =  1 for bare cells with a veg cell immediately to right
    """
    ncol = isvegc.shape[0]
    nrow = isvegc.shape[1]
    arr = np.ones([ncol+ 2*edge, nrow + 2*edge])
    arr[edge:-edge, edge:-edge] = 1- np.flipud(isvegc)
    a = pd.DataFrame(arr) != 0

    df1 = a.cumsum()-a.cumsum().where(~a).ffill().fillna(0).astype(int)
    df1 = np.array(df1)[edge:-edge, edge:-edge]
    df1 = np.flipud(df1)
    df1[isvegc == 1] = 0

    df1[df1>saturate] = saturate

    return df1

def get_patchL(isvegc, saturate):
    """
    input : isvegc from get_source(df)

    output :

      patchLv:  vegetated patch length
      patchLB:  upslope interspace patch length (paired to veg patch)
      patchLc:  charcteristic length  Lv/(Lv + Lb)

      Ldict:  dictionary of veg patch lengths.
         Ldict key :  downslope patch coordinate
      Bdict:  dictionary of paired upslope interspace lengths.

    usage:
      patchLv,patchLB,patchLc,Ldict,Bdict = get_patchL(isvegc)
    """
    ncol = isvegc.shape[0]
    nrow = isvegc.shape[1]
    patchLv = np.zeros(isvegc.shape, dtype = float)  # veg patch length
    patchLB = np.zeros(isvegc.shape, dtype = float)  # upslope interspace patch length (paired to veg patch)

    for i in range(ncol):  # loop over across-slope direction first
        count = 0
        for j in range(nrow):
            if isvegc[i, j] == 1:    #  if veg patch, add 1
                if j >= (nrow -1):  # if we're at the top of the hill
                  patchLv[i, j-count:] = count  # record veg patch length
                count += 1

            # if [i,j] is bare and the slope cell is vegetated, record.
            # each patch starts at [i,j-count] and ends at [i,j-1]
            elif isvegc[i,j] == 0 and isvegc[i, j-1] == 1:
                if j > 0:
                  # veg patch starts at j-count and ends at j
                  patchLv[i, j-count:j] = count
                  try:
                      # find the nearest upslope veg cell
                      Lb = np.where(isvegc[i,j:] == 1)[0][0]
                      patchLB[i,j-count:j] = Lb
                  except IndexError:  # bare patch extends to top of hill
                      patchLB[i,j-count:j] = nrow - j
                  count = 0
    patchLv[patchLv > saturate] = saturate
    patchLB[patchLB > saturate] = saturate

    return  patchLv, patchLB

def get_bareL(isvegc, saturate, skipflag = 0):
    """
    input : isvegc from get_source(df)

    output :

      bareL:  bare patch length
      bareLV:  upslope vegeted patch length (paired to bare patch)

    """
    ncol = isvegc.shape[0]
    nrow = isvegc.shape[1]
    bareLV = np.zeros(isvegc.shape, dtype = float)  # veg patch length
    bareL = np.zeros(isvegc.shape, dtype = float)  # upslope interspace patch length (paired to veg patch)

    for i in range(ncol):  # loop over across-slope direction first
        count = 0
        for j in range(nrow):
            if isvegc[i, j] == 0:    #  if bare, add 1
                if j >= (nrow -1):  # if we're at the top of the hill
                  bareL[i, j-count:] = count  # record bare length
                count += 1

            # if [i,j] is veg and the slope cell is bare, record.
            # each patch starts at [i,j-count] and ends at [i,j-1]
            elif isvegc[i,j] == 1 and isvegc[i, j-1] == 0:
                if j > 0:
                  # veg patch starts at j-count and ends at j
                  bareL[i, j-count:j] = count
                  if skipflag == 1 and j-count ==0:
                      bareL[i, j-count:j] = 0
                  try:
                      # find the nearest upslope bare cell
                      Lb = np.where(isvegc[i,j:] == 0)[0][0]
                      bareLV[i,j-count:j] = Lb
                  except IndexError:  # bare patch extends to top of hill
                      bareLV[i,j-count:j] = nrow - j
                  count = 0
    bareLV[bareLV > saturate] = saturate
    bareL[bareL > saturate] = saturate

    return  bareL, bareLV

max_stars_repo_name: octavia-crompton/SVE-R

max_issues_repo_name: octavia-crompton/SVE-R

max_forks_repo_name: octavia-crompton/SVE-R

reason: import numpy,import scipy,from scipy


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon May 28 11:46:20 2018

@author: sungkyun

- Live Data Loader for LJ-Speech Dataset.
"""
import pandas as pd
import numpy as np
import torch
import string, glob
import librosa
from torch.utils.data.dataset import Dataset
from nnmnkwii.datasets import FileDataSource, FileSourceDataset, MemoryCacheDataset

#DATA_ROOT = '/mnt/ssd2/data/LJSpeech-1.1'
N_TRAIN = 13000 # N_TEST = 100 (13000~13099)
X_SPEC_MAX = 193.99077 # Required for feature normalization
X_MELSPEC_MAX = 0.04071
OMIT_DATA_ROWS =[945, 1119, 1149, 1371, 1853, 1929, 1939, 2965, 3405, 4246, 4670,
                  4684, 4699, 4706, 4713, 5163, 5164, 5165, 5166, 5167, 5169, 5172,
                   5174, 5182, 5185, 5532, 5935, 5964, 5974, 6039] # Reduce sentences with foreign characters.
MAX_LEN_TEXT           = 250+2 # Max-lengths are required for z-padding
MAX_LEN_MELSPEC        = 217+2
MAX_LEN_SPEC           = 868+8
#MAX_LEN_PAIRED_TEXT    = 360
#MAX_LEN_PAIRED_MELSPEC = 291# Max-lengths are required for z-padding
#MAX_LEN_PAIRED_SPEC    = 1164

class SpecSource(FileDataSource):
    '''
    USAGE:
        Mel-spec-only : mode = ['melspec']
        Spec, Mel-spec: mode = ['SSRN']
    '''
    def __init__(self, wav_data_root=None, file_sel_range=None, target_sr = 22050, output_mode=['melspec'], norm_factors=[X_SPEC_MAX,X_MELSPEC_MAX], max_num_files=-1):

        file_list = sorted(glob.glob(wav_data_root + "/*.wav")) # All *.wav files into list
        for i in sorted(OMIT_DATA_ROWS, reverse=True):
            del file_list[i] # Omitting foreign language files
        self.input_wav_files = file_list

        self.target_sr = target_sr
        self.output_mode = output_mode
        self.norm_factors = norm_factors # [max_X_spec, max_X_melspec]
        if max_num_files is not -1 :self.input_wav_files = self.input_wav_files[:max_num_files]
        if file_sel_range is not None: self.input_wav_files = self.input_wav_files[file_sel_range[0]:file_sel_range[1]]

        self.mel_filter = librosa.filters.mel(self.target_sr, 1024, n_mels=80)

    def collect_files(self):
        # This class method is required..
        return self.input_wav_files

    def collect_features(self, path):
        x, fs = librosa.load(path, sr=self.target_sr, mono=True, dtype=np.float32)
        x_spec = np.abs(librosa.core.stft(y=x, n_fft=1024, hop_length=256, window='hann')) #STFT size : [513,4T]
        x_spec = np.power(x_spec / self.norm_factors[0], 0.6)  # normalize by max(abs(X_spec))
        x_spec_fixed_t = int(x_spec.shape[1]/4) * 4
        x_spec = x_spec[:, :x_spec_fixed_t]

        x_melspec = np.matmul(self.mel_filter, x_spec)
        x_melspec = x_melspec[:, np.arange(0, x_spec_fixed_t, 4)]
        #x_melspec = np.power(x_melspec / self.norm_factors[1], 0.6) # normalize by max(abs(X_melspec))
        x_melspec = np.power((x_melspec / np.max(x_melspec)) * np.random.uniform(0.9,1.0), 0.6)

        if 'SSRN' in self.output_mode:
            return x_spec, x_melspec
        else:
            return x_melspec

#%%
class LJSpeechDataset(Dataset):
    '''
    Live Data Loader for LJ-Speech Dataset.
    OUTPUT MODE:
        mode = ['melspec'] : return index, text, melspec
        mode = ['SSRN']    : return index, spec, melspec
    '''
    def __init__(self, data_root_dir=None, train_mode=False , output_mode='melspec', transform=None, data_sel=None):

        self.wav_root_dir = data_root_dir + '/wavs/'
        self.train_mode  = train_mode
        self.output_mode = output_mode
        self.transform   = transform
        self.data_sel = data_sel

        self.max_len_text    = MAX_LEN_TEXT
        self.max_len_melspec = MAX_LEN_MELSPEC
        self.max_len_spec    = MAX_LEN_SPEC

#        self.max_len_paired_text = MAX_LEN_PAIRED_TEXT # Max-lengths are required for z-padding
#        self.max_len_paired_spec = MAX_LEN_PAIRED_SPEC
#        self.max_len_paired_melspec = MAX_LEN_PAIRED_MELSPEC

        # Preparing Text:
        self.text_csv_path = data_root_dir + '/metadata.csv'
        self.reduce_punc_table = str.maketrans(string.ascii_uppercase, string.ascii_lowercase,
                                         '0123456789!#"$%&\()*+/:;<=>?@[\\]^_`{|}~')
        self.chr2int_table = dict(zip(" ',-." + string.ascii_lowercase, np.arange(0, 31)))

        df = pd.read_csv(self.text_csv_path, index_col=False, sep='|', header=None, memory_map=True) # memory_map: speed-up reading.
        nan_rows = df[df[2].isnull()].index.values
        df.iloc[nan_rows,2] = df.iloc[nan_rows,1]  # fixing dataset NaN value bugs...
        df = df.drop(1, axis=1);  df.columns=['file_id', 'text']
        df = df.drop(OMIT_DATA_ROWS, axis=0).reset_index(drop=True) # Omitting foreign language..

        if self.train_mode is True:
            self.file_ids = df.iloc[0:N_TRAIN, 0] # file_ids: LJ**-**** (13,000)
            self.texts    = df.iloc[0:N_TRAIN, 1]
            if self.data_sel is not None:
                self.file_ids = self.file_ids[self.data_sel].reset_index(drop=True)
                self.texts    = self.texts[self.data_sel].reset_index(drop=True)
        else:
            self.file_ids = df.iloc[N_TRAIN:, 0].reset_index(drop=True) # (100)
            self.texts    = df.iloc[N_TRAIN:, 1].reset_index(drop=True)
            if self.data_sel is not None:
                self.file_ids = self.file_ids[self.data_sel].reset_index(drop=True)
                self.texts    = self.texts[self.data_sel].reset_index(drop=True)

        # Prepraing Audio:
        if self.train_mode is True:
            self.spec_features = MemoryCacheDataset(
                    FileSourceDataset(
                            SpecSource(
                                    wav_data_root=self.wav_root_dir,
                                    file_sel_range=[0, N_TRAIN],
                                    output_mode=self.output_mode)),
                            cache_size=len(self.file_ids))
        else:
            self.spec_features = MemoryCacheDataset(
                    FileSourceDataset(
                            SpecSource(
                                    wav_data_root=self.wav_root_dir,
                                    file_sel_range=[N_TRAIN, None],
                                    output_mode=self.output_mode)),
                            cache_size=len(self.file_ids))
        assert(len(self.file_ids) == len(self.spec_features))

#        # Pairing: Sort and divide by feature lengths, then concat small + large
#        lengths       = np.load('mspec_length_train_13000.npy')
#        sorted_by_len = np.argsort(lengths)
#        n_org       = len(sorted_by_len)
#        n_pairs     = int(n_org / 2)
#        self.paired_items = list()
#        for i in range(n_pairs):
#            self.paired_items.append([sorted_by_len[i], sorted_by_len[n_org - 1 - i]])

        return None

    def __getitem__(self, index): # = Index of self.paired_items

#        # Decouple paired index to idx1, idx2
#        _choice = np.random.choice([0,1])
#        idx1, idx2 = paired_items[index][_choice], paired_items[index][1-_choice]
        if self.output_mode is 'melspec':
            text = self.chr2int(self.texts[index])
            melspec = self.spec_features[index]

            text, nz_text = self.zeropad(text, self.max_len_text)
            melspec, nz_melspec = self.zeropad(melspec, self.max_len_melspec)
            return index, torch.LongTensor(text), melspec, np.asarray([nz_text, nz_melspec])

        else: # if self.output_mode is 'SSRN':
            spec, melspec = self.spec_features[index]

            spec, nz_spec = self.zeropad(spec, self.max_len_spec)
            melspec, nz_melspec = self.zeropad(melspec, self.max_len_melspec)

            return index, spec, melspec, np.asarray([nz_spec, nz_melspec])

        #np.flip(paired_items[30], axis=0)

    def __len__(self):
        #return len(self.paired_items)
        return len(self.file_ids)

    def chr2int(self, text):
        # 'City$,' ==> ['c','i','t','y',','] ==> [7,13,24,29,2]
        text = list(text.translate(self.reduce_punc_table))
        return np.asarray([self.chr2int_table[c] for c in text])

    def zeropad(self, x, target_length):

        if len(x.shape) is 1:
            # 1D input:
            n_zeros = target_length - len(x)
            #x = np.pad(x, (n_zeros,0), 'constant', constant_values=(0,0))
            x = np.pad(x, (1,n_zeros-1), 'constant', constant_values=(0,0))  # letf 1 + right (all-1)-zpading!
        else:
            # 2D input: D x T
            n_zeros = target_length - x.shape[1]
            xz = np.zeros((x.shape[0], x.shape[1]+n_zeros))
            #xz[:, n_zeros:] = x
            xz[:, 1:(xz.shape[1] - n_zeros + 1)] = x # right-zpading!!
            x = xz
        return x, n_zeros