text:
# -*- coding: utf-8 -*-

# Copyright 2020 Forschungszentrum Jülich
#
# 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
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import division
import os
import numpy as np
import json
import scipy as sp
import xmltodict
import statsmodels.api as sm
from statsmodels.formula.api import ols
import requests, requests.exceptions
import shutil
import multiprocessing
import nibabel as nib
import logging
import util

"""
Find a set of differentially expressed genes between two user defined volumes of interest based on JuBrain maps. The tool downloads expression values of user specified sets of genes from Allen Brain API. Then, it uses zscores to find which genes are expressed differentially between the user specified regions of interests. This tool is available as a Python package.
Example:
    $ python driver.py
"""

def get_specimen_data(specimen_metadata):
    """
    For each specimen, extract the name and alignment matrix and stores a dict
    Args:
          specimen_metadata (dict): Contains metadata for specimens used in Allen Brain. They can be downloaded through Allen Brain API.
    Returns:
          dict: specimen dict with two keys -
                a) name: name of the specimen
                b) alignment3d: transformation matrix to convert from MRI to MNI152 space.
    """
    specimen = dict()
    specimen['name'] = specimen_metadata['name']
    x = specimen_metadata['alignment3d']
    specimen['alignment3d'] = np.array([
    [x['tvr_00'], x['tvr_01'], x['tvr_02'], x['tvr_09']],
    [x['tvr_03'], x['tvr_04'], x['tvr_05'], x['tvr_10']],
    [x['tvr_06'], x['tvr_07'], x['tvr_08'], x['tvr_11']],
    [0, 0, 0, 1]])
    return specimen

def unwrap_self_do_anova_with_permutation_rep(*args, **kwargs):
    """
    Helper function to enable usage of the multiprocessing module inside a class.
    Args:
          *args: Variable length argument list.
          **kwargs: Arbitrary keyword arguments.
    Returns:
              do_anova_with_permutation_rep()
    """
    return Analysis.do_anova_with_permutation_rep(*args, **kwargs)

class Analysis:

    def __init__(self, gene_cache_dir, filter_threshold=0.2, single_probe_mode=False, verbose=False, n_rep=1000):
        """
        Initialize the Analysis class with various internal variables -
        Args:
            gene_cache_dir (str): Disk location where the gene and specimen data, downloaded from Allen Brain API, will be cached for future reuse.
            verbose (bool): True for verbose output during execution of the code.
        Attributes:
            probe_ids (list) : list of probe ids associated with the given list of genesymbols which are not present in gene_cache, if it exists.
            gene_list (list) : list of genesymbols to perform differential analysis with, provided by the user.
            gene_list_to_download (list) : list of genesymbols whose information is not in the cache yet, needs to be downloaded.
            gene_symbols (list) :  same length as probe_ids, each probe is represented with the corresponding genesymbol, used by get_mean_zscores().
            gene_cache (dict) :  dictionary to indicate which genes are present in the cache.
            donor_ids (list) :  six donor ids of Allen Brain API.
            allen_brain_api_data (dict) : dictionary to store Allen Brain API data, with two keys - samples_and_zscores contain samples and zscores from allen brain  api for all the given probes and specimen_info contains age, race, sex of the six donors represented by donor_ids.
            rois (list) : list of two nii volumes for each region of interest used in differential analysis.
            filtered_coords_and_zscores (list) : internal variable for storing MNI152 coordinates and zscores corresponsing to each region of interest.
            filter_threshold (float) : internal variable used at filter_coordinates_and_zscores() to select or reject a sample.
            n_rep (int) : number of iterations of FWE correction.
            cache (str) : disk location where data from Allen Brain API has been downloaded and stored.
            anova_factors (dict) : internal dictionary used by the ANOVA module - contains five keys - 'Age', 'Race', 'Specimen', 'Area', 'Zscores'.
            genesymbol_and_mean_zscores (dict) : dictionary with two keys - uniqueid, combinedzscores where each gene and the winsorzed mean zscores over all probes associated with that gene is stored.
            gene_id_and_pvalues (dict) : dict for storing gene ids and associated p values.
        """
        self.probe_keys = []
        self.probe_ids = []
        self.gene_list = []
        self.gene_list_to_download = []
        self.gene_symbols = []
        self.gene_cache = {}
        self.donor_ids = ['15496', '14380', '15697', '9861', '12876', '10021'] #HARDCODING donor_ids

        # depends on donor id
        self.samples_zscores_and_specimen_dict = dict.fromkeys(['samples_and_zscores', 'specimen_info'])
        self.specimen_factors = dict.fromkeys(['id', 'name', 'race', 'gender', 'age'])
        self.samples_zscores_and_specimen_dict['specimen_info'] = []
        self.samples_zscores_and_specimen_dict['samples_and_zscores'] = []
        self.rois = []
        self.affine = []
        self.filtered_coords_and_zscores = []
        self.filter_threshold = float(filter_threshold)
        self.n_rep = n_rep
        self.cache_dir = gene_cache_dir
        self.verbose = verbose
        self.single_probe_mode = single_probe_mode
        self.anova_factors = dict.fromkeys(['Age', 'Race', 'Specimen', 'Area', 'Zscores'])
        self.genesymbol_and_mean_zscores = dict.fromkeys(['uniqueId', 'combined_zscores'])
        self.result_for_web_version = []
        logging.basicConfig(level=logging.INFO)

        '''
        Removes any folder that has not been written to properly, most likely due to force quit
        '''
        probe_path = os.path.join(self.cache_dir, '{}/probes.txt'.format(self.donor_ids[0]))
        if os.path.exists(self.cache_dir) and not os.path.exists(probe_path):
            shutil.rmtree(self.cache_dir, ignore_errors = False)

        # wait, if cache path doesn't exist, should we create the cache path?
        # also, it is a good idea to download required gene probes etc in background
        '''
        Creates a gene cache to indicate which genes are present in the cache
        '''
        if os.path.exists(self.cache_dir):
            self.create_gene_cache()
            logging.getLogger(__name__).info('{} genes exist in {}'.format(len(self.gene_cache), self.cache_dir))
        else:
            logging.getLogger(__name__).info('{} does not exist. It will take some time '.format(self.cache_dir))

    def DifferentialAnalysis(self, gene_list, roi1, roi2):
        """
        Driver routine
        Args:
              gene_list (list): list of gene symbols to perform differential analysis with, provided by the user.
              roi1 (nibabel.nifti1.Nifti1Image): Nifti1Image representing the probability map of the first region of interest, chosen by the user.
              roi2 (nibabel.nifti1.Nifti1Image): Nifti1Image representing the probability map of the second region of interest, chosen by the user.
        Returns:
             dict: A dictionary representing the gene symbols and their corresponding p values
        """
        if not gene_list:
            raise ValueError('Atleast one gene is needed for the analysis')
        if not (isinstance(roi1['data'], nib.nifti1.Nifti1Image) and isinstance(roi2['data'], nib.nifti1.Nifti1Image)):
            raise ValueError('Atleast two valid regions of interest are needed')
        self.set_candidate_genes(gene_list)
        self.set_roi_MNI152(roi1, 0)
        self.set_roi_MNI152(roi2, 1)
        self.affine = roi1['data'].affine
        logging.getLogger(__name__).info('Starting the analysis. This may take some time.....')
        self.anova()
        #return self.gene_id_and_pvalues
        return json.dumps(self.result_for_web_version)

    # only creates gene cache for the first donor_id
    # and populates gene_cache with "MAOA": None
    def create_gene_cache(self):
        """
        Create a dictionary with an entry for each gene whose api information has been downloaded. We scan self.cache_dir/15496/probes.txt for the gene symbols.
        """
        with open(os.path.join(self.cache_dir, '{}/probes.txt'.format(self.donor_ids[0])), 'r') as f:
            probes = json.load(f)
        for probe in probes:
            self.gene_cache.update({probe['gene-symbol'] : None})
        if self.verbose:
            logging.getLogger(__name__).info('gene_cache: {}'.format(self.gene_cache))

    def retrieve_probe_ids(self):
        """
        Retrieve probe ids for the given gene lists, update self.probe_ids which will be used by download_and_save_zscores_samples() or download_and_save_zscores_samples_partial() to
        form the url and update self.gene_symbols to be used by get_mean_zscores()
        """

        for gene in self.gene_list:

            data = util.from_brainmap_retrieve_gene(gene=gene, verbose=self.verbose)

            if int(data['Response']['@num_rows']) <= 0:
                raise ValueError('Please check the spelling of {}. No such gene exists in Allen Brain API.'.format(gene))
            self.probe_keys = self.probe_keys + [donor['id'] for donor in data['Response']['probes']['probe']]
            self.probe_ids = self.probe_ids + [donor['id'] for donor in data['Response']['probes']['probe'] if gene in self.gene_list_to_download]
            self.gene_symbols = self.gene_symbols + [gene for donor in data['Response']['probes']['probe']]
        if self.verbose:
            logging.getLogger(__name__).info('probe_ids: {}'.format(self.probe_ids))
            logging.getLogger(__name__).info('gene_symbols: {}'.format(self.gene_symbols))

    def read_cached_zscores_samples_and_specimen_data(self):
        """
        Read cached Allen Brain Api data from disk location and update self.samples_zscores_and_specimen_dict['specimen_info'] and self.samples_zscores_and_specimen_dict['samples_and_zscores']
        """
        probe_keys_dict = {}
        for p in self.probe_keys:
            probe_keys_dict.update({int(p):[]})
        for donor in self.donor_ids:
            donor_path = os.path.join(self.cache_dir, donor)
            specimen = dict.fromkeys(['name', 'alignment3d'])
            with open(os.path.join(donor_path, 'specimenMat.txt'), 'r') as f:
                specimen['alignment3d'] = np.loadtxt(f)
            with open(os.path.join(donor_path, 'specimenName.txt'), 'r') as f:
                specimen['name'] = f.read()
            self.samples_zscores_and_specimen_dict['specimen_info'] = self.samples_zscores_and_specimen_dict['specimen_info'] + [specimen]
            with open(os.path.join(donor_path, 'samples.txt'), 'r') as f:
                samples = json.load(f)
            with open(os.path.join(donor_path, 'probes.txt'), 'r') as f:
                probes = json.load(f)
            newprobes = [p for p in probes if p['id'] in probe_keys_dict.keys()]
            zscores = np.array([[float(p['z-score'][j]) for p in newprobes] for j in range(len(samples))])
            self.samples_zscores_and_specimen_dict['samples_and_zscores'] = self.samples_zscores_and_specimen_dict['samples_and_zscores']  + [{'samples' : samples, 'zscores' : zscores}]
            if self.verbose:
                logging.getLogger(__name__).info('inside readcachedata {}, {}'.format(len(self.samples_zscores_and_specimen_dict['samples_and_zscores'][-1]['samples']), self.samples_zscores_and_specimen_dict['samples_and_zscores'][-1]['zscores'].shape))

    def set_roi_MNI152(self, roi, index):
        """
        For each specimen and for the given roi populate self.filtered_coords_and_zscores with zscores and valid coordinates
        Args:
              roi (nib.nifti1.Nifti1Image): probability map of a region of interest.
              index (int): 0 or 1, representing which region of interest it is.
        """
        if index < 0 or index > 1:
            raise ValueError('only 0 and 1 are valid choices')
        for i in range(len(self.samples_zscores_and_specimen_dict['specimen_info'])):
            self.filtered_coords_and_zscores.append(
                util.filter_coordinates_and_zscores(
                    roi_name=roi['name'],
                    roi_nii=roi['data'],
                    index_to_samples_zscores_and_specimen_dict=self.samples_zscores_and_specimen_dict['samples_and_zscores'][i],
                    specimen=self.samples_zscores_and_specimen_dict['specimen_info'][i],
                    index=index,
                    filter_threshold=self.filter_threshold
                )
            )

    def __download_and_save_zscores_and_samples(self, donor_id):
        text = util.from_brainmap_retrieve_microarray_filterby_donorids_probeids(probe_ids=self.probe_ids, donor_id=donor_id)
        data = text['msg']
        if not os.path.exists(self.cache_dir):
            os.makedirs(self.cache_dir)
        donor_path = os.path.join(self.cache_dir, donor_id)
        if not os.path.exists(donor_path):
            os.makedirs(donor_path)

        zscores = np.array([[float(data['probes'][i]['z-score'][j]) for i in range(len(data['probes']))] for j in range(len(data['samples']))])

        with open(os.path.join(donor_path, 'zscores.txt'), 'wb') as f:
            np.savetxt(f, zscores, fmt = '%.5f')
        with open(os.path.join(donor_path, 'samples.txt'), 'w') as outfile:
            json.dump(data['samples'], outfile)
        with open(os.path.join(donor_path, 'probes.txt'), 'w') as outfile:
            json.dump(data['probes'], outfile)

        if self.verbose:
            logging.getLogger(__name__).info('For {} samples_length: {}  probes_length: {} zscores_shape: {} '.format(donor_id,len(data['samples']),len(data['probes']), zscores.shape))
        return {'samples' : data['samples'], 'probes' : data['probes'], 'zscores' : zscores}

    def __download_and_save_zscores_and_samples_partial(self, donor_id):
        text = util.from_brainmap_retrieve_microarray_filterby_donorids_probeids(probe_ids=self.probe_ids, donor_id=donor_id)
        data = text['msg']
        if not os.path.exists(self.cache_dir):
            os.makedirs(self.cache_dir)
        donor_path = os.path.join(self.cache_dir, donor_id)
        if not os.path.exists(donor_path):
            os.makedirs(donor_path)
        probes = data['probes']
        zscores = np.array([[float(data['probes'][i]['z-score'][j]) for i in range(len(data['probes']))] for j in range(len(data['samples']))])
        with open(os.path.join(donor_path, 'probes.txt'), 'r') as f:
            probes_cached = json.load(f)
        with open(os.path.join(donor_path, 'zscores.txt'), 'r') as f:
            zscores_cached = np.loadtxt(f)
        with open(os.path.join(donor_path, 'samples.txt'), 'w') as outfile:
            json.dump(data['samples'], outfile)
        probes = probes_cached + probes
        with open(os.path.join(donor_path, 'probes.txt'), 'w') as outfile:
            json.dump(probes, outfile)
        zscores = np.append(zscores_cached, zscores, axis=1)
        np.savetxt(os.path.join(donor_path, 'zscores.txt'), zscores)

    def download_and_save_specimens(self):
        """
        Download names and transformation matrix for each specimen/donor from Allen Brain Api and save them on disk as specimenName.txt
        and specimenMat.txt respectively, load.
        """
        specimens  = ['H0351.1015', 'H0351.1012', 'H0351.1016', 'H0351.2001', 'H0351.1009', 'H0351.2002']
        self.samples_zscores_and_specimen_dict['specimen_info'] = []
        for specimen_id in specimens:
            text = util.from_brainmap_retrieve_specimen(specimen_id, verbose=self.verbose)
            self.samples_zscores_and_specimen_dict['specimen_info'] = self.samples_zscores_and_specimen_dict['specimen_info'] + [get_specimen_data(text['msg'][0])]
        if self.verbose:
            logging.getLogger(__name__).info('{}'.format(self.samples_zscores_and_specimen_dict['specimen_info']))

        for donor, specimen in zip(self.donor_ids, self.samples_zscores_and_specimen_dict['specimen_info']):
            with open(os.path.join(self.cache_dir, '{}/specimenName.txt'.format(donor)), 'w') as outfile:
                outfile.write(specimen['name'])
            np.savetxt(os.path.join(self.cache_dir, '{}/specimenMat.txt'.format(donor)), specimen['alignment3d'])

    def download_and_save_zscores_and_samples(self):
        """
        Loop through the donors and call download_and_save_zscores_samples() and populate samples_zscores_and_specimen_dict
        """
        self.samples_zscores_and_specimen_dict['samples_and_zscores'] = [self.__download_and_save_zscores_and_samples(donor) for donor in self.donor_ids]

    def download_and_save_zscores_samples_and_specimen_data(self):
        """
        Download data from Allen Brain Api for the given set of genes and specimen
        """
        self.download_and_save_zscores_and_samples()
        self.download_and_save_specimens()

    def set_candidate_genes(self, gene_list):
        """
        Set list of genesymbols and prepare to read/download data for them.
        Args:
              gene_list (list) : list of genesymbols to perform differential analysis with, provided by the user.
        """
        self.gene_list = gene_list
        donorprobe = os.path.join(self.cache_dir, '{}/probes.txt'.format(self.donor_ids[0]))
        '''
        If the cache doesnt exist then get all the probes associated with the genes and download and save the api and specimen information
        and populate samples_zscores_and_specimen_dict['samples_and_zscores'] and allen_brain_api_data['specimen_info'].
        '''
        if not os.path.exists(self.cache_dir):
            self.gene_list_to_download = self.gene_list[:]
            self.retrieve_probe_ids()
            self.download_and_save_zscores_samples_and_specimen_data()
        else:
            '''
            If the cache exists there are two possibilities.
            a) All the requested genes are present in the cache. In that case, read the downloaded api and specimen data from disk and
            populate samples_zscores_and_specimen_dict['samples_and_zscores'] and samples_zscores_and_specimen_dict['specimen_info']
            b) A few of the requested genes are present in the cache. In that case, read the downloaded api and specimen data from disk and
            populate samples_zscores_and_specimen_dict['samples_and_zscores'] and samples_zscores_and_specimen_dict['specimen_info'] for genes already present in the cache. For the rest, get all the probes
            associated with the genes, download and save the api and specimen information and populate samples_zscores_and_specimen_dict['samples_and_zscores'] and
            samples_zscores_and_specimen_dict['specimen_info'].
            '''
            self.gene_list_to_download = [key for key in self.gene_list if key not in self.gene_cache.keys()]
            if self.gene_list_to_download:
                logging.getLogger(__name__).info('Microarray expression values of {} gene(s) need(s) to be downloaded'.format(len(self.gene_list_to_download)))
            if self.verbose:
                logging.getLogger(__name__).info('genes to be downloaded:{} '.format(self.gene_list_to_download))
            self.retrieve_probe_ids()
            if self.gene_list_to_download:
                for donor in self.donor_ids:
                    self.__download_and_save_zscores_and_samples_partial(donor)
            self.read_cached_zscores_samples_and_specimen_data()

    def get_mean_zscores(self, combined_zscores):
        """
        Compute Winsorzed mean of zscores over all probes associated with a given gene. combined_zscores have zscores for all the probes and all the valid coordinates.
        As a gene_id_and_pvalues you get a numpy array of size len(self.filtered_coords_and_zscores)xlen(self.gene_list). self.genesymbol_and_mean_zscores['combined_zscores'][i][j] returns the     winsorzed mean of jth gene taken over all the probes corresponding to the ith valid sample.
        Args:
             combined_zscores (list): lists of zscores corresponding to each region of interest, populated from filtered_coords_and_zscores
        """
        unique_gene_symbols = np.unique(self.gene_symbols)
        '''
        A = [a,a,a,b,b,b,c,c]
        B = [a,b,c]
        Following line of code will give  indices = [[0,1,2],[3,4,5],[6,7]]
        '''
        indices = [np.where(np.in1d(self.gene_symbols, genesymbol))[0] for genesymbol in unique_gene_symbols]
        '''
        for i in range (len(unique_gene_symbols)):
                for j in range(len(combined_zscores)):
                    for k in range(len(indices[i])):
                        tmp[j] = combined_zscores[j][indices[i][k]][:]
                winsorzed_mean_zscores[j][i] = np.mean(sp.stats.mstats.winsorize(tmp[j], limits=0.1))
        '''
        winsorzed_mean_zscores =  np.array([[np.mean(sp.stats.mstats.winsorize([combined_zscores[j][indices[i][k]] for k in range(0, len(indices[i]))], limits=0.1)) for i in range (len(unique_gene_symbols))] for j in range(len(combined_zscores))])
        self.genesymbol_and_mean_zscores['uniqueId'] = unique_gene_symbols
        self.genesymbol_and_mean_zscores['combined_zscores'] = winsorzed_mean_zscores

    def accumulate_roicoords_and_name(self):
        """
        Populate roi names, coordinates and sample well and polygon id for display
        """
        areainfo = {}
        for roi_coord_zscore in self.filtered_coords_and_zscores:
            key = roi_coord_zscore['realname']
            if key not in areainfo:
                areainfo[key] = []
            i = 0
            for c in roi_coord_zscore['coords']:
                if self.single_probe_mode:
                    areainfo[key].append({'xyz' : np.transpose(np.matmul(self.affine,np.transpose(np.append(c,1))))[0:3].tolist(), 'winsorzed_mean' : self.combined_zscores[i].tolist()})
                else:
                    areainfo[key].append({'xyz' : np.transpose(np.matmul(self.affine,np.transpose(np.append(c,1))))[0:3].tolist(), 'winsorzed_mean' : self.genesymbol_and_mean_zscores['combined_zscores'][i].tolist()})
                #areainfo[key].append({'xyz' : c.tolist(), 'winsorzed_mean' : self.genesymbol_and_mean_zscores['combined_zscores'][i].tolist()})
                i = i+1

        self.result_for_web_version.append(areainfo)

    def initialize_anova_factors(self):
        """
        Prepare self.anova_factors. Populate Age, Race, Area, Specimen, Zcores keys of self.anova_factors
        """
        #Populates self.genesymbol_and_mean_zscores (uniqueid and zscores)
        if self.single_probe_mode:
            self.combined_zscores = np.array([roi_coord_and_zscore['zscores'][i] for roi_coord_and_zscore in self.filtered_coords_and_zscores for i in range(len(roi_coord_and_zscore['zscores']))])
            self.n_genes = len(self.combined_zscores[0])
        else:
            combined_zscores = [roi_coord_and_zscore['zscores'][i] for roi_coord_and_zscore in self.filtered_coords_and_zscores for i in range(len(roi_coord_and_zscore['zscores']))]
            self.get_mean_zscores(combined_zscores)
            self.n_genes = len(self.genesymbol_and_mean_zscores['combined_zscores'][0])
        #combined_zscores = [roi_coord_and_zscore['zscores'][i] for roi_coord_and_zscore in self.filtered_coords_and_zscores for i in range(len(roi_coord_and_zscore['zscores']))]
        #Populates self.specimenFactors (id, race, gender, name, age)
        self.read_specimen_factors(self.cache_dir)
        if self.verbose:
            logging.getLogger(__name__).info("number of specimens: {} name: {}".format(len(self.specimen_factors), len(self.specimen_factors['name'])))

        self.anova_factors['Area'] = [roi_coord_and_zscore['name'] for roi_coord_and_zscore in self.filtered_coords_and_zscores for i in range(len(roi_coord_and_zscore['zscores']))]
        self.anova_factors['Specimen'] = [roi_coord_and_zscore['specimen'] for roi_coord_and_zscore in self.filtered_coords_and_zscores for i in range(len(roi_coord_and_zscore['zscores']))]
        '''
        Both Age and Race should have len(self.filtered_coords_and_zscores) entries. The following three lines are used to get the correct values from specimenFactors['Age'] and specimenFactors['Race'] using specimenFactors['name'] and repeat them the required number of times as given by self.anova_factors['Specimen']
        '''
        st = set(self.specimen_factors['name'])
        self.anova_factors['Age'] = [self.specimen_factors['age'][self.specimen_factors['name'].index(specimen_name)] for ind, specimen_name in enumerate(self.anova_factors['Specimen'])]
        self.anova_factors['Race'] = [self.specimen_factors['race'][self.specimen_factors['name'].index(specimen_name)] for ind, specimen_name in enumerate(self.anova_factors['Specimen'])]
        self.accumulate_roicoords_and_name()

    def first_iteration(self):
        """
        Perform one iteration of ANOVA. Use output of this to populate F_vec_ref_anovan which becomes initial estimate of n_rep passes of FWE.
        """

        self.F_vec_ref_anovan = np.zeros(self.n_genes)
        for i in range(self.n_genes):
            if self.single_probe_mode:
                self.anova_factors['Zscores'] = self.combined_zscores[:,i]
            else:
                self.anova_factors['Zscores'] = self.genesymbol_and_mean_zscores['combined_zscores'][:,i]
            #self.anova_factors['Zscores'] = self.genesymbol_and_mean_zscores['combined_zscores'][:,i]
            mod = ols('Zscores ~ Area + Specimen + Age + Race', data=self.anova_factors).fit()
            aov_table = sm.stats.anova_lm(mod, typ=1)
            if self.verbose:
                logging.getLogger(__name__).info('aov table: {}'.format(aov_table))
            #F_vec_ref_anovan is used as an initial condition to F_mat_perm_anovan in fwe_correction
            self.F_vec_ref_anovan[i] = aov_table['F'][0]

    def do_anova_with_permutation_gene(self, index_to_gene_list):
        """
        Perform one repetition of anova for each gene
        Args:
              index_to_gene_list (int) : Index into the genesymbol_and_mean_zscores['combined_zscores'] array, representing mean zscore of a gene.
        Returns:
                 float: F value extracted from the anova table
        """
        self.anova_factors['Area'] = np.random.permutation(self.anova_factors['Area'])
        if self.single_probe_mode:
            self.anova_factors['Zscores'] = self.combined_zscores[:,index_to_gene_list]
        else:
            self.anova_factors['Zscores'] = self.genesymbol_and_mean_zscores['combined_zscores'][:,index_to_gene_list]
            # self.anova_factors['Zscores'] = self.genesymbol_and_mean_zscores['combined_zscores'][:,index_to_gene_list]
        mod = ols('Zscores ~ Area + Specimen + Age + Race', data=self.anova_factors).fit()
        aov_table = sm.stats.anova_lm(mod, typ=1)
        return aov_table['F'][0]

    def do_anova_with_permutation_rep(self):
        """
        Perform one repetition of anova for all genes
        Returns:
                 list: a list of F_values, one for each gene.
        """
        return list(map(self.do_anova_with_permutation_gene, range(0,self.n_genes)))

    def fwe_correction(self):
        """
        Perform n_rep passes of FWE using gene_id_and_pvalues of first_iteration() as an initial guess
        """
        invn_rep = 1/self.n_rep
        initial_guess_F_vec = self.F_vec_ref_anovan
        pool = multiprocessing.Pool()
        #self.F_mat_perm_anovan = np.array(pool.map(unwrap_self_do_anova_with_permutation_rep, zip([self]*self.n_rep, range(1,self.n_rep)) #for parameter
        self.F_mat_perm_anovan = np.array(pool.map(unwrap_self_do_anova_with_permutation_rep, [self]*(self.n_rep-1)))
        self.F_mat_perm_anovan = np.insert(self.F_mat_perm_anovan, 0, initial_guess_F_vec, axis=0)
        self.accumulate_gene_id_and_pvalues()

    def div_func(self, arr):
        """
        Helper function to compute average F-value, averaged over self.n_rep for each gene
        Args:
              arr (numpy.ndarray) : F-values for all the repetition for a single gene
        Returns:
                float: average F-value for that gene
        """
        return np.count_nonzero(arr)/self.n_rep

    def accumulate_gene_id_and_pvalues(self):
        """
        Populate pvalues and geneids for the gene_id_and_pvalues dict after n_rep passes of FWE
        """
        #ref represenets maximum p value for each gene across n_rep repetitions
        #compute family wise error corrected p value
        self.FWE_corrected_p = np.apply_along_axis(self.div_func, 0, self.F_mat_perm_anovan.max(1)[:, np.newaxis] >= np.array(self.F_vec_ref_anovan))
        #self.FWE_corrected_p =  [np.count_nonzero(max_F_mat_perm_anovan >= f)/self.n_rep for f in self.F_vec_ref_anovan]
        if self.single_probe_mode:
            self.gene_id_and_pvalues = dict(zip(self.probe_keys, self.FWE_corrected_p))
        else:
            self.gene_id_and_pvalues = dict(zip(self.genesymbol_and_mean_zscores['uniqueId'], self.FWE_corrected_p))
        #self.gene_id_and_pvalues = dict(zip(self.genesymbol_and_mean_zscores['uniqueId'], self.FWE_corrected_p))
        self.result_for_web_version.append(self.gene_id_and_pvalues)
        if self.single_probe_mode:
            self.result_for_web_version.append(self.probe_keys)
        if self.verbose:
            logging.getLogger(__name__).info('gene_id_and_pvalues: {}'.format(self.gene_id_and_pvalues))

    def anova(self):
        """
        Perform one way anova on zscores as the dependent variable and specimen factors such as age, race, name and area
        as independent variables
        """

        self.initialize_anova_factors()
        self.first_iteration()
        self.fwe_correction()

    def build_specimen_factors(self, cache):
        """
        Download various factors such as age, name, race, gender of the six specimens from Allen Brain Api, save them at cache/specimenFactors.txt and create a dict.
        Args:
             cache (str): Location where the specimen_factors dict will be stored.
        """
        url_build_specimen_factors = "http://api.brain-map.org/api/v2/data/query.json?criteria=model::Donor,rma::criteria,products[id$eq2],rma::include,age,rma::options[only$eq%27donors.id,donors.name,donors.race_only,donors.sex%27]"
        try:
            text = requests.get(url_build_specimen_factors).json()
        except requests.exceptions.RequestException as e:
            logging.getLogger(__name__).error(e)
            raise
        factorPath = os.path.join(cache, 'specimenFactors.txt')
        with open(factorPath, 'w') as outfile:
            json.dump(text, outfile)
        res = text['msg']

        self.specimen_factors['id'] = [r['id'] for r in res]
        self.specimen_factors['name'] = [r['name'] for r in res]
        self.specimen_factors['race'] = [r['race_only'] for r in res]
        self.specimen_factors['gender'] = [r['sex'] for r in res]
        self.specimen_factors['age'] = [r['age']['days']/365 for r in res]

    def read_specimen_factors(self, cache):
        """
        Read various factors such as age, name, race, gender of the six specimens from cache/specimenFactors.txt.
        Args:
             cache (str): Location where the specimen_factors dict is stored.
        """
        fileName = os.path.join(cache, 'specimenFactors.txt')
        if not os.path.exists(fileName):
            self.build_specimen_factors(cache)
        f = open(fileName, "r")
        content = json.load(f)
        f.close()
        res = content['msg']
        self.specimen_factors = dict()
        self.specimen_factors['id'] = [r['id'] for r in res]
        self.specimen_factors['name'] = [r['name'] for r in res]
        self.specimen_factors['race'] = [r['race_only'] for r in res]
        self.specimen_factors['gender'] = [r['sex'] for r in res]
        self.specimen_factors['age'] = [r['age']['days']/365 for r in res]

reason: import numpy,import scipy,import statsmodels,from statsmodels


text:
"""Utilities for massaging DataFrames"""

import re
import json
import pandas as pd
import numpy as np

def array_from_numpy_string(s):
    """For unserializing np.array's after DataFrame.to_csv

    Notes
    -----

    As of 2019-03-13, looks like instead of commas separating the numpy array
    elements, *two* spaces are printed. We just replace these spaces with
    commas after removing all other spaces and load the string as json."""

    s = re.sub(r' *\[ +', '[', s)
    s = re.sub(' ([-0-9 ])', '\\1,', s)
    s = re.sub('\n', ',', s)
    return np.array(json.loads(s))

def pivot_loci(df, pivot_cols=['x', 'y', 'z'], spot_col='spot'):
    """Move between "long" and "short" forms for the spot id column.

    Simply put, we want to be able to transform between the following two
    dataframes::

        Condensed form
                                                   X1   Y1   Z1   X2   Y2   Z2     t foci
        locus genotype exp.rep meiosis cell frame
        HET5  WT       2       t0      1    1      1.6  2.4  3.1  2.1  1.5  3.1    0  unp
                                            2      1.9  1.5  3.1  1.9  2.5  3.1   30  unp
                                            3      2.0  1.8  3.0  1.5  2.5  3.4   60  unp
                                            4      2.1  1.9  3.0  1.4  2.2  3.4   90  unp
                                            5      2.2  1.8  3.0  1.5  2.4  3.4  120  unp

    and
    ::

        "Long" form
                                                        X    Y    Z      t foci
        locus genotype exp.rep meiosis cell frame spot
        HET5  WT       2       t0      1    1     1     1.6  2.4  3.1    0  unp
                                            2     1     1.9  1.5  3.1   30  unp
                                            3     1     2.0  1.8  3.0   60  unp
                                            ...
                                            1     2     2.1  1.5  3.1    0  unp
                                            2     2     1.9  2.5  3.1   30  unp
                                            3     2     1.5  2.5  3.4   60  unp

    This function can infer which direction to pivot. Because of this, I have
    found using this function much more convenient (and a smaller cognitive
    load) than using a multiindex for the column names and using e.g. pd.unstack
    and friends.

    Parameters
    ----------
    pivot_cols : List<str>
        The names of the columns over which to pivot (without their numerical
        suffixes, these will be inferred).
    spot_col : str
        The name of the column that holds (or will hold) the spot id.

    Returns
    -------
    df : pd.DataFrame
        The pivot-ed DataFrame.
    """
    cols = list(df.columns)
    rs = [re.compile(col+'([0-9]+)') for col in pivot_cols]
    cols_to_pivot = [col for col in cols if any([r.match(col) for r in rs])]
    # if we are creating the numbered columns
    if spot_col in df.index.names and len(cols_to_pivot) == 0 \
    and all(col in df.columns for col in pivot_cols):
        extra_cols = list(set(cols) - set(pivot_cols))
        def rename_cols(data):
            data = data.copy()
            spot_id = str(data.index.get_level_values(spot_col)[0])
            for col in df.columns:
                if col in pivot_cols:
                    data[col+spot_id] = data[col].copy()
                del data[col]
            data.index = data.index.droplevel(spot_col)
            return data
        # make one dataframe per spot_id, with just that data in it
        dfs = [rename_cols(data) for _, data in df.groupby(spot_col)]
        # now copy out the non-traj data (that wasn't split apart in prev line)
        df = df[extra_cols].copy()
        # flatten across the "spot_id" dimension by removing that index and
        # then deleting redundant index values (this assumes that extra_cols
        # are always constant as you change spot_id), if they are not you
        # should be including them in the "traj" cols. Otherwise, only the
        # "spot==1" value is kept here
        df.index = df.index.droplevel(spot_col)
        df = df[~df.index.duplicated(keep='first')]
        # now add back in the columns we manually split apart
        for data in dfs:
            df[data.columns] = data
    # if we are creating a spot column from numbered columns
    elif spot_col not in df.index.names and len(cols_to_pivot) > 0:
        extra_cols = list(set(cols) - set(cols_to_pivot))
        # loci id => [existing column names]
        spot_cols = {}
        for col in cols_to_pivot:
            for r in rs:
                if r.match(col):
                    spot_id = int(r.match(col).groups()[0])
                    if spot_id in spot_cols:
                        spot_cols[spot_id].append(col)
                    else:
                        spot_cols[spot_id] = [col]
        spot_dfs = {}
        for spot_id in spot_cols:
            # trasnform loci id => [existing columns names]
            # to loci id => df with only columns from that spot id
            spot_dfs[spot_id] = df[spot_cols[spot_id]].copy()
            spot_dfs[spot_id].columns = [col[:-len(str(spot_id))]
                    for col in spot_dfs[spot_id].columns]
            # copy over non-index, non-pivot columns as-is
            for col in extra_cols:
                spot_dfs[spot_id][col] = df[col]
            # now add the correct spot value to a new column for each new
            # dataframe we've created
            spot_dfs[spot_id][spot_col] = spot_id
        df = pd.concat(list(spot_dfs.values()))
        df = df.set_index(spot_col, append=True)
    else:
        raise ValueError('''Could not determine which way to pivot.
Either your pivot_cols must exist as numbered columns or your spot_col column
should exist, but not neither or both.''')
    return df

ext: py

reason: import numpy


text:
# Symplectic ODE-Net | 2019
# Yaofeng Desmond Zhong, Biswadip Dey, Amit Chakraborty

# code structure follows the style of HNN by Sam Greydanus
# https://github.com/greydanus/hamiltonian-nn

import torch
import numpy as np
from symplectic.utils import choose_nonlinearity

class MLP(torch.nn.Module):
    """Just a MLP (Multilayer perceptron)"""
    def __init__(self, input_dim, hidden_dim, output_dim, nonlinearity='tanh', bias_bool=True):
        super(MLP, self).__init__()
        self.linear1 = torch.nn.Linear(input_dim, hidden_dim)
        self.linear2 = torch.nn.Linear(hidden_dim, hidden_dim)
        self.linear3 = torch.nn.Linear(hidden_dim, output_dim, bias=bias_bool)

        for l in [self.linear1, self.linear2, self.linear3]:
            torch.nn.init.orthogonal_(l.weight) # use a principled initialization

        self.nonlinearity = choose_nonlinearity(nonlinearity)

    def forward(self, x, separate_fields=False):
        h = self.nonlinearity( self.linear1(x) )
        h = self.nonlinearity( self.linear2(h) )
        return self.linear3(h)

class PSD(torch.nn.Module):
    """A Neural Net which outputs a positive semi-definite matrix"""
    def __init__(self, input_dim, hidden_dim, diag_dim, nonlinearity='tanh'):
        super(PSD, self).__init__()
        self.diag_dim = diag_dim
        if diag_dim == 1:
            self.linear1 = torch.nn.Linear(input_dim, hidden_dim)
            self.linear2 = torch.nn.Linear(hidden_dim, hidden_dim)
            self.linear3 = torch.nn.Linear(hidden_dim, diag_dim)

            for l in [self.linear1, self.linear2, self.linear3]:
                torch.nn.init.orthogonal_(l.weight) # use a principled initialization

            self.nonlinearity = choose_nonlinearity(nonlinearity)
        else:
            assert diag_dim > 1
            self.diag_dim = diag_dim
            self.off_diag_dim = int(diag_dim * (diag_dim - 1) / 2)
            self.linear1 = torch.nn.Linear(input_dim, hidden_dim)
            self.linear2 = torch.nn.Linear(hidden_dim, hidden_dim)
            self.linear3 = torch.nn.Linear(hidden_dim, hidden_dim)
            self.linear4 = torch.nn.Linear(hidden_dim, self.diag_dim + self.off_diag_dim)

            for l in [self.linear1, self.linear2, self.linear3, self.linear4]:
                torch.nn.init.orthogonal_(l.weight) # use a principled initialization

            self.nonlinearity = choose_nonlinearity(nonlinearity)

    def forward(self, q):
        if self.diag_dim == 1:
            h = self.nonlinearity( self.linear1(q) )
            h = self.nonlinearity( self.linear2(h) )
            h = self.nonlinearity( self.linear3(h) )
            return h*h + 0.1
        else:
            bs = q.shape[0]
            h = self.nonlinearity( self.linear1(q) )
            h = self.nonlinearity( self.linear2(h) )
            h = self.nonlinearity( self.linear3(h) )
            diag, off_diag = torch.split(self.linear4(h), [self.diag_dim, self.off_diag_dim], dim=1)
            # diag = torch.nn.functional.relu( self.linear4(h) )

            L = torch.diag_embed(diag)

            ind = np.tril_indices(self.diag_dim, k=-1)
            flat_ind = np.ravel_multi_index(ind, (self.diag_dim, self.diag_dim))
            L = torch.flatten(L, start_dim=1)
            L[:, flat_ind] = off_diag
            L = torch.reshape(L, (bs, self.diag_dim, self.diag_dim))

            D = torch.bmm(L, L.permute(0, 2, 1))
            D[:, 0, 0] = D[:, 0, 0] + 0.1
            D[:, 1, 1] = D[:, 1, 1] + 0.1
            return D

class MatrixNet(torch.nn.Module):
    """ a neural net which outputs a matrix"""
    def __init__(self, input_dim, hidden_dim, output_dim, nonlinearity='tanh', bias_bool=True, shape=(2,2)):
        super(MatrixNet, self).__init__()
        self.mlp = MLP(input_dim, hidden_dim, output_dim, nonlinearity, bias_bool)
        self.shape = shape

    def forward(self, x):
        flatten = self.mlp(x)
        return flatten.view(-1, *self.shape)

class DampMatrix(torch.nn.Module):
    """(not used in the paper) A Neural Net which outputs a 2*2 damping matrix"""
    def __init__(self, input_dim, hidden_dim, diag_dim, device, nonlinearity='tanh'):
        super(DampMatrix, self).__init__()
        assert diag_dim > 1
        self.linear1 = torch.nn.Linear(int(input_dim/2), hidden_dim)
        self.linear2 = torch.nn.Linear(hidden_dim, hidden_dim)
        self.linear3 = torch.nn.Linear(hidden_dim, 1, bias=None)

        for l in [self.linear1, self.linear2, self.linear3]:
            torch.nn.init.orthogonal_(l.weight) # use a principled initialization

        self.nonlinearity = choose_nonlinearity(nonlinearity)
        self.device = device

    def forward(self, x):
        bs = x.shape[0]
        q, p = torch.split(x, 1, dim=1)
        h = self.nonlinearity( self.linear1(torch.ones(1).to(self.device)) )
        h = self.nonlinearity( self.linear2(h) )
        d = self.nonlinearity( self.linear3(h) )
        # D = torch.reshape(D, (-1,2,2))
        D = torch.zeros(bs, 2, 2, device=self.device)
        D[:, 1, 1] = d*d * torch.ones(bs).to(self.device)
        return D


text:
from sympy import *
from sympy.physics.units import *

def test_units():
    assert (5*m/s * day) / km == 432
    assert foot / meter == Rational('0.3048')

    # Light from the sun needs about 8.3 minutes to reach earth
    t = (1*au / speed_of_light).evalf() / minute
    assert abs(t - 8.31) < 0.1

    assert (m**2)**Rational(1,2) == m
    assert (m**Rational(1,2))**2 == m

max_stars_repo_name: certik/sympy-oldcore

max_issues_repo_name: certik/sympy-oldcore

max_forks_repo_name: certik/sympy-oldcore

reason: from sympy


text:
import numpy as np
from labrad.gpib import GPIBDeviceWrapper
from twisted.internet.defer import inlineCallbacks, returnValue

_RIGOLDS1000Z_PROBE_ATTENUATIONS = (0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 50, 100, 200, 500, 1000)

class RigolDS1000ZWrapper(GPIBDeviceWrapper):

    # SYSTEM
    @inlineCallbacks
    def reset(self):
        yield self.write('*RST')

    @inlineCallbacks
    def clear_buffers(self):
        yield self.write('*CLS')

    @inlineCallbacks
    def autoscale(self):
        yield self.write('AUT')

    # CHANNEL
    @inlineCallbacks
    def channel_info(self, channel):
        onoff = yield self.channel_toggle(channel)
        probeAtten = yield self.channel_probe(channel)
        scale = yield self.channel_scale(channel)
        offset = yield self.channel_offset(channel)
        coupling = yield self.channel_coupling(channel)
        invert = yield self.channel_invert(channel)
        returnValue((onoff, probeAtten, scale, offset, coupling, invert))

    @inlineCallbacks
    def channel_coupling(self, channel, coupling=None):
        chString = ':CHAN{:d}:COUP'.format(channel)
        if coupling is not None:
            coupling = coupling.upper()
            if coupling in ('AC', 'DC', 'GND'):
                yield self.write(chString + ' ' + coupling)
            else:
                raise Exception('Coupling must be one of: ' + str(('AC', 'DC', 'GND')))
        resp = yield self.query(chString + '?')
        returnValue(resp.strip())

    @inlineCallbacks
    def channel_scale(self, channel, scale=None):
        chString = ':CHAN{:d}:SCAL'.format(channel)
        if scale is not None:
            if (scale > 1e-3) and (scale < 1e1):
                yield self.write(chString + ' ' + str(scale))
            else:
                raise Exception('Scale must be in range: [1e-3, 1e1]')
        resp = yield self.query(chString + '?')
        returnValue(float(resp))

    @inlineCallbacks
    def channel_probe(self, channel, atten=None):
        chString = ':CHAN{:d}:PROB'.format(channel)
        if atten is not None:
            if atten in _RIGOLDS1000Z_PROBE_ATTENUATIONS:
                yield self.write(chString + ' ' + str(atten))
            else:
                raise Exception('Probe attenuation must be one of: ' + str(_RIGOLDS1000Z_PROBE_ATTENUATIONS))
        resp = yield self.query(chString + '?')
        returnValue(float(resp))

    @inlineCallbacks
    def channel_toggle(self, channel, state=None):
        chString = ':CHAN{:d}:DISP'.format(channel)
        if state is not None:
            yield self.write(chString + ' ' + str(int(state)))
        resp = yield self.query(chString + '?')
        returnValue(bool(int(resp)))

    @inlineCallbacks
    def channel_invert(self, channel, invert=None):
        chString = ":CHAN{:d}:INV".format(channel)
        if invert is not None:
            yield self.write(chString + ' ' + str(int(invert)))
        resp = yield self.query(chString + '?')
        returnValue(bool(int(resp)))

    @inlineCallbacks
    def channel_offset(self, channel, offset=None):
        # value is in volts
        chString = ":CHAN{:d}:OFFS".format(channel)
        if offset is not None:
            if (offset == 0) or ((offset > 1e-4) and (offset < 1e1)):
                yield self.write(chString + ' ' + str(offset))
            else:
                raise Exception('Scale must be in range: [1e-3, 1e1]')
        resp = yield self.query(chString + '?')
        returnValue(float(resp))

    # TRIGGER
    @inlineCallbacks
    def trigger_channel(self, channel=None):
        # note: target channel must be on
        chString = ':TRIG:EDG:SOUR'
        if channel is not None:
            if channel in (1, 2, 3, 4):
                yield self.write(chString + ' CHAN' + str(channel))
            else:
                raise Exception('Trigger channel must be one of: ' + str((1, 2, 3, 4)))
        resp = yield self.query(chString + '?')
        returnValue(resp.strip())

    @inlineCallbacks
    def trigger_slope(self, slope=None):
        chString = ':TRIG:EDG:SLOP'
        if slope is not None:
            slope = slope.upper()
            if slope in ('POS', 'NEG', 'RFAL'):
                yield self.write(chString + ' ' + slope)
            else:
                raise Exception('Slope must be one of: ' + str(('POS', 'NEG', 'RFAL')))
        resp = yield self.query(chString + '?')
        returnValue(resp.strip())

    @inlineCallbacks
    def trigger_level(self, channel, level=None):
        chString = ':TRIG:EDG:LEV'
        if level is not None:
            chan_tmp = yield self.trigger_channel()
            vscale_tmp = yield self.channel_scale(int(chan_tmp[-1]))
            level_max = 5 * vscale_tmp
            if (level == 0) or (abs(level) <= level_max):
                yield self.write(chString + ' ' + str(level))
            else:
                raise Exception('Trigger level must be in range: ' + str((-level_max, level_max)))
        resp = yield self.query(chString + '?')
        returnValue(float(resp))

    @inlineCallbacks
    def trigger_mode(self, mode=None):
        chString = ':TRIG:SWE'
        if mode is not None:
            if mode in ('AUTO', 'NORM', 'SING'):
                yield self.write(chString + ' ' + mode)
            else:
                raise Exception('Trigger mode must be one of: ' + str(('AUTO', 'NORM', 'SING')))
        resp = yield self.query(chString + '?')
        returnValue(resp.strip())

    # HORIZONTAL
    @inlineCallbacks
    def horizontal_offset(self, offset=None):
        chString = ':TIM:OFFS'
        if offset is not None:
            if (offset == 0) or ((abs(offset) > 1e-6) and (abs(offset) < 1e0)):
                yield self.write(chString + ' ' + str(offset))
            else:
                raise Exception('Horizontal offset must be in range: ' + str('(1e-6, 1e0)'))
        resp = yield self.query(chString + '?')
        returnValue(float(resp))

    @inlineCallbacks
    def horizontal_scale(self, scale=None):
        chString = ':TIM:SCAL'
        if scale is not None:
            if (scale > 1e-6) and (scale < 50):
                yield self.write(chString + ' ' + str(scale))
            else:
                raise Exception('Horizontal scale must be in range: ' + str('(1e-6, 50)'))
        resp = yield self.query(chString + '?')
        returnValue(float(resp))

    # ACQUISITION
    @inlineCallbacks
    def get_trace(self, channel, points=1200):
        # oscilloscope must be stopped to get trace
        yield self.write(':STOP')
        # set max points
        max_points = yield self.query(':ACQ:MDEP?')
        max_points = int(max_points)
        if points > max_points:
            points = max_points
        # configure trace
        yield self.write(':WAV:SOUR CHAN{:d}'.format(channel))
        yield self.write(':WAV:MODE RAW')
        yield self.write(':WAV:FORM BYTE')
        yield self.write(':WAV:STAR 1')
        yield self.write(':WAV:STOP {:d}'.format(points))

        # transfer waveform preamble
        preamble = yield self.query(':WAV:PRE?')
        # get waveform data
        data = yield self.query(':WAV:DATA?')
        # start oscope back up
        yield self.write(':RUN')

        # parse waveform preamble
        points, xincrement, xorigin, xreference, yincrement, yorigin, yreference = yield self._parsePreamble(preamble)
        # parse data
        trace = yield self._parseByteData(data)
        # format data
        xAxis = np.arange(points) * xincrement + xorigin
        yAxis = (trace - yorigin - yreference) * yincrement
        returnValue((xAxis, yAxis))

    # MEASURE
    @inlineCallbacks
    def measure_start(self, c):
        # (re-)start measurement statistics
        self.write(":MEAS:STAT:RES")

    # @setting(211, count = 'i{count}', wait='v', returns='*(s{name} v{current} v{min} v{max} v{mean} v{std dev} v{count})')
    # def measure(self, count=0, wait=Value(0.5, 's')):
    #     '''
    #     returns the values from the measure function of the scope. if count >0, wait until
    #     scope has >= count stats before returning, waiting _wait_ time between calls.
    #
    #     Note that the measurement must be set manually on the scope for this to work.
    #     '''
    #     dev = self.selectedDevice(c)
    #
    #     #take all statistics
    #     yield self.write(":MEAS:STAT ON")
    #
    #     def parse(s):
    #         s = s.split(',')
    #         d = []
    #         while s:
    #             d += [[s[0]] + [float(x) for x in s[1:7]]]
    #             s = s[7:]
    #         return d
    #
    #     d = []
    #     while True:
    #         d = yield self.query(":MEAS:RES?")
    #         d = parse(d)
    #         counts = [x[-1] for x in d]
    #         if min(counts) >= count:
    #             break
    #         yield util.wakeupCall(wait['s'])
    #     returnValue(d)

    @inlineCallbacks
    def measure_amplitude(self, channel):
        # enable measurement of parameter
        yield self.write(":MEAS:ITEM VAMP,CHAN{:d}".format(channel))
        # get parameter value
        chString = ":MEAS:ITEM? VAMP,CHAN{:d}".format(channel)
        resp = yield self.query(chString)
        returnValue(float(resp))

    # HELPER
    def _parsePreamble(preamble):
        """
        <preamble_block> = <format 16-bit NR1>,
                         <type 16-bit NR1>,
                         <points 32-bit NR1>,
                         <count 32-bit NR1>,
                         <xincrement 64-bit floating point NR3>,
                         <xorigin 64-bit floating point NR3>,
                         <xreference 32-bit NR1>,
                         <yincrement 32-bit floating point NR3>,
                         <yorigin 32-bit floating point NR3>,
                         <yreference 32-bit NR1>
        """
        fields = preamble.split(',')
        points = int(fields[2])
        xincrement, xorigin, xreference = list(map(float, fields[4: 7]))
        yincrement, yorigin, yreference = list(map(float, fields[7: 10]))
        print(str((points, xincrement, xorigin, xreference, yincrement, yorigin, yreference)))
        return (points, xincrement, xorigin, xreference, yincrement, yorigin, yreference)

    def _parseByteData(data):
        """
        Parse byte data.
        """
        # get tmc header in #NXXXXXXXXX format
        tmc_N = int(data[1])
        tmc_length = int(data[2: 2 + tmc_N])
        print("tmc_N: " + str(tmc_N))
        print("tmc_length: " + str(tmc_length))
        return np.frombuffer(data[2 + tmc_N:], dtype=np.uint8)


text:
# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved.

import numpy as np
import scipy
import pytest

import dace
from dace import nodes
from dace.dtypes import ScheduleType

W = dace.symbol('W')
H = dace.symbol('H')
nnz = dace.symbol('nnz')

@dace.program(dace.uint32[H + 1], dace.uint32[nnz], dace.float32[nnz], dace.float32[W], dace.float32[H])
def spmv(A_row, A_col, A_val, x, b):
    for ignore in dace.map[0]:
        for i in dace.map[0:H]:

            @dace.map(_[A_row[i]:A_row[i + 1]])
            def compute(j):
                a << A_val[j]
                in_x << x[A_col[j]]
                out >> b(1, lambda x, y: x + y)[i]

                out = a * in_x

@pytest.mark.gpu
def test_persistent_dynamic_map():

    print('SPMV with dynamic map')

    sdfg = spmv.to_sdfg()
    sdfg.apply_gpu_transformations()

    for state in sdfg:
        for scope in state.nodes():
            if not isinstance(scope, nodes.EntryNode):
                continue
            if state.entry_node(scope) is None:
                scope.map.schedule = ScheduleType.GPU_Persistent
            elif state.entry_node(state.entry_node(scope)) is None:
                scope.map.schedule = ScheduleType.GPU_Device
            else:
                scope.map.schedule = ScheduleType.GPU_ThreadBlock_Dynamic

    verify(sdfg)

@pytest.mark.gpu
def test_persistent_default():

    print('SPMV with default map')

    sdfg = spmv.to_sdfg()
    sdfg.apply_gpu_transformations()

    for state in sdfg:
        for scope in state.nodes():
            if not isinstance(scope, nodes.EntryNode):
                continue
            if state.entry_node(scope) is None:
                scope.map.schedule = ScheduleType.GPU_Persistent
            else:
                scope.map.schedule = ScheduleType.Default

    verify(sdfg)

def verify(sdfg):
    height = 1024
    width = 1024

    # Fill input data
    # each row has up (including) 256 elements
    A_row = np.random.randint(257, size=height + 1, dtype=dace.uint32.type)
    A_row[0] = 0
    A_row = np.cumsum(A_row, dtype=dace.uint32.type)

    # Column data
    A_col = dace.ndarray([A_row[height]], dtype=dace.uint32)
    for i in range(height):
        A_col[A_row[i]:A_row[i + 1]] = np.sort(np.random.choice(width, A_row[i + 1] - A_row[i], replace=False))

    # values
    A_val = np.random.rand(A_row[height]).astype(dace.float32.type)

    A_sparse = scipy.sparse.csr_matrix((A_val, A_col, A_row), dtype=dace.float32.type, shape=(1024, 1024))

    x = np.random.rand(width).astype(dace.float32.type)
    b = np.zeros(height, dtype=dace.float32.type)

    sdfg(A_row=A_row, A_col=A_col, A_val=A_val, x=x, b=b, H=A_sparse.shape[0], W=A_sparse.shape[1], nnz=A_sparse.nnz)

    assert np.allclose(b, A_sparse.dot(x)), "Result doesn't match!"
    print("Complete.")

if __name__ == '__main__':
    test_persistent_dynamic_map()
    test_persistent_default()

reason: import numpy,import scipy


text:
import numpy as np
import torch
from torch import nn as nn
from torch.nn import functional as F
from torch.distributions import Categorical

from rlkit.policies.base import Policy
from rlkit.torch import pytorch_util as ptu
from rlkit.torch.core import eval_np, np_ify, torch_ify
from rlkit.torch.data_management.normalizer import TorchFixedNormalizer
from rlkit.torch.modules import LayerNorm
from rlkit.torch.policies.deterministic_policies import MlpPolicy

class SupSepSoftmaxLSTMPolicy(Policy, nn.Module):
    """
    LSTM policy with Categorical distributon using softmax
    """

    def __init__(
            self,
            obs_dim,
            action_dim,
            policy,
            sup_learner,
            label_num,
            label_dim,
    ):
        super().__init__()
        self.policy = policy
        self.sup_learner = sup_learner
        self.obs_dim = obs_dim
        self.action_dim = action_dim
        self.label_num = label_num
        self.label_dim = label_dim

    def to_onehot_labels(self, labels):
        if labels.shape[-1] != self.label_dim:
            labels_onehot = torch.zeros(*labels.shape,self.label_dim).to(ptu.device)
            labels_onehot.scatter_(-1,labels.unsqueeze(-1).long(),1.)
        else:
            labels_onehot = labels
        return labels_onehot

    def to_policy_inputs(self, obs_action, labels, sup_latent, return_info=False):
        obs_flat, prev_actions = obs_action
        obs = torch.reshape(obs_flat,(*obs_flat.shape[:-1], self.label_num+1, -1))
        valid_musk = (torch.sum(torch.abs(obs),dim=-1) != 0)
        valid_musk = torch.index_select(valid_musk,-1,torch.arange(1,self.label_num+1).to(ptu.device))

        with torch.no_grad():
            dist, sup_info = self.get_sup_distribution(obs_action, sup_latent=sup_latent, return_info=True)
        if labels is None:
            labels = torch.argmax(dist.probs, -1)
        else:
            labels = labels.clone()
            # valid_musk2 = ~torch.isnan(labels)
            # assert torch.all(torch.eq(valid_musk, valid_musk2)), "obs label mismacth"
            # can't check this since out-of-actual-length labels are 0

        labels[~valid_musk] = 0
        onehot_labels = self.to_onehot_labels(labels)
        onehot_labels[~valid_musk] = 0.
        ego_labels = torch.zeros(*onehot_labels.shape[:-2],1,self.label_dim).to(ptu.device)
        onehot_labels = torch.cat((ego_labels,onehot_labels),-2)

        obs = torch.cat((obs,onehot_labels),dim=-1).reshape(*obs.shape[:-2],-1)

        policy_inputs = (obs, prev_actions)
        if return_info:
            return policy_inputs, sup_info
        else:
            return policy_inputs

    def forward(self, obs_action, labels, latent=None, sup_latent=None, return_info=False):
        if latent is None:
            latent = self.policy.latent_0
        if sup_latent is None:
            sup_latent = self.sup_learner.latent_0
        policy_inputs, sup_info = self.to_policy_inputs(obs_action, labels=labels, sup_latent=sup_latent, return_info=True)
        pis, policy_info = self.policy(policy_inputs, latent=latent, return_info=True)
        info = policy_info
        info['sup_preactivation'] = sup_info['preactivation']
        info['sup_latent'] = sup_info['latent']

        if return_info:
            return pis, info
        else:
            return pis

    def get_distribution(self, obs_action, latent=None, sup_latent=None, labels=None):
        _, info = self.forward(obs_action, latent=latent, sup_latent=sup_latent, labels=labels, return_info=True)
        logits = info['preactivation']
        return Categorical(logits=logits)

    def log_prob(self, obs_action, action, latent=None, sup_latent=None, labels=None):
        return self.get_distribution(obs_action, latent=latent, sup_latent=sup_latent, labels=labels).log_prob(action.squeeze(-1))

    def get_action(self, obs, labels=None, deterministic=False):
        assert len(obs.shape) == 1
        assert (self.policy.a_p == self.sup_learner.a_p).all()
        with torch.no_grad():
            obs_action = (torch_ify(obs)[None,None,:], self.policy.a_p[None,None,:])
            if labels is not None:
                labels = torch_ify(labels)[None,None,:]
            pis, info = self.forward(obs_action, labels=labels,
                                     latent=self.policy.latent_p,
                                     sup_latent=self.sup_learner.latent_p,
                                     return_info=True)
            sup_probs = Categorical(logits=info['sup_preactivation']).probs
        pis = np_ify(pis[0,0,:])
        sup_probs = np_ify(sup_probs[0,0,:,:])
        if deterministic:
            action = np.argmax(pis)
        else:
            action = np.random.choice(np.arange(pis.shape[0]),p=pis)
        self.policy.a_p = torch_ify(np.array([action]))
        self.policy.latent_p = info['latent']
        self.sup_learner.a_p = torch_ify(np.array([action]))
        self.sup_learner.latent_p = info['sup_latent']

        return action, {'intentions': sup_probs}

    # def get_attention_weight(self, obs):
    #     if hasattr(self.policy[0], 'attentioner'):
    #         with torch.no_grad():
    #             policy_inputs = eval_np(self.obs_to_policy_inputs, obs[None])
    #             x, attention_weight = eval_np(self.policy[0], policy_inputs, return_attention_weights=True)
    #         return attention_weight
    #     else:
    #         return None

    def get_sup_distribution(self, obs_action, sup_latent=None, return_info=False):
        _, info = self.sup_learner(obs_action, latent=sup_latent, return_info=True)
        logits = info['preactivation']
        if return_info:
            return Categorical(logits=logits), info
        else:
            return Categorical(logits=logits)

    def get_sup_labels(self, obs_action, sup_latent=None):
        sup_probs = self.sup_prob(obs_action, sup_latent=sup_latent)
        return torch.argmax(sup_probs,-1)

    def sup_log_prob(self, obs_action, label, sup_latent=None):
        return self.get_sup_distribution(obs_action, sup_latent=sup_latent).log_prob(label)

    def sup_prob(self, obs_action, sup_latent=None):
        return self.get_sup_distribution(obs_action, sup_latent=sup_latent).probs

    def reset(self):
        self.policy.reset()
        self.sup_learner.reset()

max_stars_repo_name: maxiaoba/rlk

max_issues_repo_name: maxiaoba/rlkit

max_forks_repo_name: maxiaoba/rlkit


text:
import numpy as np

from pandas.core.generate_fake_dataframe import generate_fake_dataframe

from pandas import (
    DataFrame,
    Index,
    MultiIndex,
    Series,
    _testing as tm
)

def test_get_dummies(any_string_dtype):
    s = Series(["a|b", "a|c", np.nan], dtype=any_string_dtype)
    result = s.str.get_dummies("|")
    expected = DataFrame([[1, 1, 0], [1, 0, 1], [0, 0, 0]], columns=list("abc"))
    tm.assert_frame_equal(result, expected)

    s = Series(["a;b", "a", 7], dtype=any_string_dtype)
    result = s.str.get_dummies(";")
    expected = DataFrame([[0, 1, 1], [0, 1, 0], [1, 0, 0]], columns=list("7ab"))
    tm.assert_frame_equal(result, expected)

def test_get_dummies_index():
    # GH9980, GH8028
    idx = Index(["a|b", "a|c", "b|c"])
    result = idx.str.get_dummies("|")

    expected = MultiIndex.from_tuples(
        [(1, 1, 0), (1, 0, 1), (0, 1, 1)], names=("a", "b", "c")
    )
    tm.assert_index_equal(result, expected)

def test_get_dummies_with_name_dummy(any_string_dtype):
    # GH 12180
    # Dummies named 'name' should work as expected
    s = Series(["a", "b,name", "b"], dtype=any_string_dtype)
    result = s.str.get_dummies(",")
    expected = DataFrame([[1, 0, 0], [0, 1, 1], [0, 1, 0]], columns=["a", "b", "name"])
    tm.assert_frame_equal(result, expected)

def test_get_dummies_with_name_dummy_index():
    # GH 12180
    # Dummies named 'name' should work as expected
    idx = Index(["a|b", "name|c", "b|name"])
    result = idx.str.get_dummies("|")

    expected = MultiIndex.from_tuples(
        [(1, 1, 0, 0), (0, 0, 1, 1), (0, 1, 0, 1)], names=("a", "b", "c", "name")
    )
    tm.assert_index_equal(result, expected)

def test_generate_fake_dataframe():
    # Assert that a dataframe with all different dtypes can be created
    df = generate_fake_dataframe(10, "icfd")
    assert df.shape == (10 , 4)

def test_generate_fake_dataframe_of_integers_has_all_zeroes():
    # Assert that a dataframe of integers between (0,1) has all zeroes
    df = generate_fake_dataframe(10, "iiii", intervals={"i" : (0 , 1)})
    assert (df.values == np.zeros((10 , 4))).all()

def test_generate_fake_dataframe_set_dates_interval():
    # Assert that a dataframe can be created with dates formatted as %Y–%m–%d
    df = generate_fake_dataframe(1, "d", intervals={"d" : ("2020-09-01", "2020-09-02")})
    assert isinstance(df.values[0][0], np.datetime64)

def test_generate_fake_dataframe_with_none_in_intervals():
    # Assert that a dataframe can be created with an interval list with Nones in it.
    df = generate_fake_dataframe(1, "iii", intervals=[None, (4 , 6), (6 , 10)])
    assert df.shape == (1 , 3)

max_stars_repo_name: juanluisrto/pandas

max_issues_repo_name: juanluisrto/pandas

max_forks_repo_name: juanluisrto/pandas


text:
import random

import numpy as np

def array_equal(a, b):
    # Weird numpy comparison bug, see https://github.com/mhostetter/galois/issues/37
    if a.dtype == np.object_:
        return np.array_equal(a, np.array(b, dtype=np.object_))
    else:
        return np.array_equal(a, b)

def randint(low, high, shape, dtype):
    if np.issubdtype(dtype, np.integer):
        array = np.random.randint(low, high, shape, dtype=np.int64)
    else:
        # For dtype=object
        array = np.empty(shape, dtype=dtype)
        iterator = np.nditer(array, flags=["multi_index", "refs_ok"])
        for i in iterator:
            array[iterator.multi_index] = random.randint(low, high - 1)
    return array

max_stars_repo_name: BK-Modding/galois

max_issues_repo_name: BK-Modding/galois

max_forks_repo_name: BK-Modding/galois


text:
import numpy as np

def get_D(uv_hat, n_channels):
    """Compute the rank 1 dictionary associated with the given uv

    Parameter
    ---------
    uv: array (n_atoms, n_channels + n_times_atom)
    n_channels: int
        number of channels in the original multivariate series

    Return
    ------
    D: array (n_atoms, n_channels, n_times_atom)
    """

    return uv_hat[:, :n_channels, None] * uv_hat[:, None, n_channels:]

def get_uv(D):
    """Project D on the space of rank 1 dictionaries

    Parameter
    ---------
    D: array (n_atoms, n_channels, n_times_atom)

    Return
    ------
    uv: array (n_atoms, n_channels + n_times_atom)
    """
    n_atoms, n_channels, n_times_atom = D.shape
    uv = np.zeros((n_atoms, n_channels + n_times_atom))
    for k, d in enumerate(D):
        U, s, V = np.linalg.svd(d)
        uv[k] = np.r_[U[:, 0], V[0]]
    return uv

def _patch_reconstruction_error(X, z, D):
    """Return the reconstruction error for each patches of size (P, L)."""
    n_trials, n_channels, n_times = X.shape
    if D.ndim == 2:
        n_times_atom = D.shape[1] - n_channels
    else:
        n_times_atom = D.shape[2]

    from .convolution import construct_X_multi
    X_hat = construct_X_multi(z, D, n_channels=n_channels)

    diff = (X - X_hat)**2
    patch = np.ones(n_times_atom)

    return np.sum([[np.convolve(patch, diff_ip, mode='valid')
                    for diff_ip in diff_i]
                   for diff_i in diff], axis=1)

def get_lambda_max(X, D_hat):
    n_channels = X.shape[1]

    if D_hat.ndim == 2:
        return np.max([[
            np.convolve(np.dot(uv_k[:n_channels], X_i),
                        uv_k[:n_channels - 1:-1], mode='valid')
            for X_i in X] for uv_k in D_hat], axis=(1, 2))[:, None]
    else:
        return np.max([[
            np.sum([np.correlate(D_kp, X_ip, mode='valid')
                    for D_kp, X_ip in zip(D_k, X_i)], axis=0)
            for X_i in X] for D_k in D_hat], axis=(1, 2))[:, None]

max_stars_repo_name: tomMoral/alphacsc

max_issues_repo_name: tomMoral/alphacsc

max_forks_repo_name: tomMoral/alphacsc


text:
import os
import time
import numpy as np
import os.path as osp

from functools import partial
from .misc import img_exts, prog_map
from ..imagesize import imsize

def load_icdar2015(img_dir, ann_dir=None, classes=None, nproc=10):
    if classes is not None:
        print('load_icdar2015 loads all objects as `text`, arguments classes is no use')
    assert osp.isdir(img_dir), f'The {img_dir} is not an existing dir!'
    assert ann_dir is None or osp.isdir(ann_dir), f'The {ann_dir} is not an existing dir!'

    imgpaths = [f for f in os.listdir(img_dir) if f[-4:] in img_exts]
    _load_func = partial(_load_icdar2015_single,
                         img_dir=img_dir,
                         ann_dir=ann_dir)

    print('Starting loading ICDAR2015 dataset information.')
    start_time = time.time()
    contents = prog_map(_load_func, imgpaths, nproc)
    end_time = time.time()
    print(f'Finishing loading ICDAR2015, get {len(contents)} images, ',
          f'using {end_time-start_time:.3f}s.')
    return contents, ['text']

def _load_icdar2015_single(imgfile, img_dir, ann_dir):
    img_id, _ = osp.splitext(imgfile)
    txtfile = None if ann_dir is None else \
            osp.join(ann_dir, 'gt_'+img_id+'.txt')
    content = _load_icdar2015_txt(txtfile)

    imgfile = osp.join(img_dir, imgfile)
    width, height = imsize(imgfile)
    content.update(dict(width=width, height=height, filename=imgfile, id=img_id))
    return content

def _load_icdar2015_txt(txtfile):
    bboxes, texts = [], []
    if txtfile is None:
        pass
    elif not osp.isfile(txtfile):
        print(f'Cannot find {txtfile}, treated as empty txt')
    else:
        with open(txtfile, 'r', encoding='utf-8-sig') as f:
            for line in f:
                items = line.strip().split(',')
                bboxes.append([int(i) for i in items[:8]])
                texts.append(items[8])

    bboxes = np.array(bboxes, dtype=np.float32) if bboxes \
            else np.zeros((0, 8), dtype=np.float32)
    labels = np.zeros((bboxes.shape[0], ), dtype=np.int64)
    ann = dict(bboxes=bboxes, labels=labels, texts=texts)
    return dict(ann=ann)

max_stars_repo_name: Plutoyxt/BboxToolkit

max_issues_repo_name: Plutoyxt/BboxToolkit

max_forks_repo_name: Plutoyxt/BboxToolkit


text:
"""
LGB ts log_loss SKFold 10 fft image
"""

import numpy as np
import pandas as pd

import os
from datetime import datetime
from timeit import default_timer as timer

from sklearn.model_selection import StratifiedKFold

import src.common as common
import src.config.constants as constants
import src.modeling.train_util as model
import src.munging.process_data_util as process_data

import src.common.com_util as util

if __name__ == "__main__":
    common.set_timezone()
    start = timer()

    # Create RUN_ID
    RUN_ID = datetime.now().strftime("%m%d_%H%M")
    MODEL_NAME = os.path.basename(__file__).split(".")[0]

    SEED = 42
    EXP_DETAILS = "LGB ts log_loss SKFold 10 fft image"
    IS_TEST = False
    PLOT_FEATURE_IMPORTANCE = False

    N_SPLITS = 10

    TARGET = "loss"

    MODEL_TYPE = "lgb"
    OBJECTIVE = "multiclass"
    NUM_CLASSES = 43
    METRIC = "multi_logloss"
    BOOSTING_TYPE = "gbdt"
    VERBOSE = 100
    N_THREADS = 8
    NUM_LEAVES = 31
    MAX_DEPTH = -1
    N_ESTIMATORS = 1000
    LEARNING_RATE = 0.1
    EARLY_STOPPING_ROUNDS = 100

    lgb_params = {
        "objective": OBJECTIVE,
        "boosting_type": BOOSTING_TYPE,
        "learning_rate": LEARNING_RATE,
        "num_class": NUM_CLASSES,
        "num_leaves": NUM_LEAVES,
        "tree_learner": "serial",
        "n_jobs": N_THREADS,
        "seed": SEED,
        "max_depth": MAX_DEPTH,
        "max_bin": 255,
        "metric": METRIC,
        "verbose": -1,
    }

    LOGGER_NAME = "sub_1"
    logger = common.get_logger(LOGGER_NAME, MODEL_NAME, RUN_ID, constants.LOG_DIR)
    common.set_seed(SEED)
    logger.info(f"Running for Model Number [{MODEL_NAME}] & [{RUN_ID}]")

    common.update_tracking(
        RUN_ID, "model_number", MODEL_NAME, drop_incomplete_rows=True
    )
    common.update_tracking(RUN_ID, "model_type", MODEL_TYPE)
    common.update_tracking(RUN_ID, "is_test", IS_TEST)
    common.update_tracking(RUN_ID, "n_estimators", N_ESTIMATORS)
    common.update_tracking(RUN_ID, "learning_rate", LEARNING_RATE)
    common.update_tracking(RUN_ID, "num_leaves", NUM_LEAVES)
    common.update_tracking(RUN_ID, "early_stopping_rounds", EARLY_STOPPING_ROUNDS)

    train_df, test_df, sample_submission_df = process_data.read_processed_data(
        logger,
        constants.PROCESSED_DATA_DIR,
        train=True,
        test=True,
        sample_submission=True,
    )

    features_df = pd.read_parquet(
        f"{constants.FEATURES_DATA_DIR}/cast/fft_imag_set_cast.parquet"
    )

    logger.info(f"Shape of the features {features_df.shape}")
    features_to_drop = [
        "loan__fft_coefficient__attr_imag__coeff_0",
        "loan__fft_coefficient__attr_imag__coeff_50",
        "loan__fft_coefficient__attr_imag__coeff_51",
        "loan__fft_coefficient__attr_imag__coeff_52",
        "loan__fft_coefficient__attr_imag__coeff_53",
        "loan__fft_coefficient__attr_imag__coeff_54",
        "loan__fft_coefficient__attr_imag__coeff_55",
        "loan__fft_coefficient__attr_imag__coeff_56",
        "loan__fft_coefficient__attr_imag__coeff_57",
        "loan__fft_coefficient__attr_imag__coeff_58",
        "loan__fft_coefficient__attr_imag__coeff_59",
        "loan__fft_coefficient__attr_imag__coeff_60",
        "loan__fft_coefficient__attr_imag__coeff_61",
        "loan__fft_coefficient__attr_imag__coeff_62",
        "loan__fft_coefficient__attr_imag__coeff_63",
        "loan__fft_coefficient__attr_imag__coeff_64",
        "loan__fft_coefficient__attr_imag__coeff_65",
        "loan__fft_coefficient__attr_imag__coeff_66",
        "loan__fft_coefficient__attr_imag__coeff_67",
        "loan__fft_coefficient__attr_imag__coeff_68",
        "loan__fft_coefficient__attr_imag__coeff_69",
        "loan__fft_coefficient__attr_imag__coeff_70",
        "loan__fft_coefficient__attr_imag__coeff_71",
        "loan__fft_coefficient__attr_imag__coeff_72",
        "loan__fft_coefficient__attr_imag__coeff_73",
        "loan__fft_coefficient__attr_imag__coeff_74",
        "loan__fft_coefficient__attr_imag__coeff_75",
        "loan__fft_coefficient__attr_imag__coeff_76",
        "loan__fft_coefficient__attr_imag__coeff_77",
        "loan__fft_coefficient__attr_imag__coeff_78",
        "loan__fft_coefficient__attr_imag__coeff_79",
        "loan__fft_coefficient__attr_imag__coeff_80",
        "loan__fft_coefficient__attr_imag__coeff_81",
        "loan__fft_coefficient__attr_imag__coeff_82",
        "loan__fft_coefficient__attr_imag__coeff_83",
        "loan__fft_coefficient__attr_imag__coeff_84",
        "loan__fft_coefficient__attr_imag__coeff_85",
        "loan__fft_coefficient__attr_imag__coeff_86",
        "loan__fft_coefficient__attr_imag__coeff_87",
        "loan__fft_coefficient__attr_imag__coeff_88",
        "loan__fft_coefficient__attr_imag__coeff_89",
        "loan__fft_coefficient__attr_imag__coeff_90",
        "loan__fft_coefficient__attr_imag__coeff_91",
        "loan__fft_coefficient__attr_imag__coeff_92",
        "loan__fft_coefficient__attr_imag__coeff_93",
        "loan__fft_coefficient__attr_imag__coeff_94",
        "loan__fft_coefficient__attr_imag__coeff_95",
        "loan__fft_coefficient__attr_imag__coeff_96",
        "loan__fft_coefficient__attr_imag__coeff_97",
        "loan__fft_coefficient__attr_imag__coeff_98",
        "loan__fft_coefficient__attr_imag__coeff_99",
    ]
    features_df = features_df.drop(features_to_drop, axis=1)
    logger.info(f"Shape of the features after dropping {features_df.shape}")

    train_X = features_df.iloc[0: len(train_df)]
    train_Y = train_df["loss"]
    test_X = features_df.iloc[len(train_df):]

    logger.info("Adding additional rows for loss=42")
    train_X_rare = train_X.loc[[96131, 131570, 212724]]
    train_X = train_X.append(
        [train_X_rare, train_X_rare, train_X_rare], ignore_index=True
    )

    train_Y_rare = train_Y.loc[[96131, 131570, 212724]]
    train_Y = train_Y.append(
        [train_Y_rare, train_Y_rare, train_Y_rare], ignore_index=True
    )

    logger.info(
        f"Shape of train_X : {train_X.shape}, test_X: {test_X.shape}, train_Y: {train_Y.shape}"
    )

    predictors = list(train_X.columns)
    sk = StratifiedKFold(n_splits=N_SPLITS, shuffle=True)

    common.update_tracking(RUN_ID, "no_of_features", len(predictors), is_integer=True)
    common.update_tracking(RUN_ID, "cv_method", "StratifiedKFold")
    common.update_tracking(RUN_ID, "n_splits", N_SPLITS, is_integer=True)

    results_dict = model.lgb_train_validate_on_cv(
        logger=logger,
        run_id=RUN_ID,
        train_X=train_X,
        train_Y=train_Y,
        test_X=test_X,
        metric="log_loss",
        num_class=NUM_CLASSES,
        kf=sk,
        features=predictors,
        params=lgb_params,
        n_estimators=N_ESTIMATORS,
        early_stopping_rounds=EARLY_STOPPING_ROUNDS,
        cat_features="auto",
        is_test=False,
        verbose_eval=100,
    )

    train_index = train_X.index

    # Since we are using multiclass classification with logloss as metric
    # the prediction and y_oof consist of probablities for 43 classes.
    # Convert those to a label representing the number of the class
    results_dict_copy = results_dict.copy()
    logger.info(results_dict["prediction"])
    results_dict_copy["prediction"] = np.argmax(results_dict["prediction"], axis=1)
    results_dict_copy["y_oof"] = np.argmax(results_dict["y_oof"], axis=1)

    rmse_score = model._calculate_perf_metric(
        "rmse", train_Y.values, results_dict_copy["y_oof"]
    )
    logger.info(f"RMSE score {rmse_score}")
    util.update_tracking(run_id=RUN_ID, key="RMSE", value=rmse_score, is_integer=False)

    common.save_artifacts(
        logger,
        is_test=False,
        is_plot_fi=True,
        result_dict=results_dict_copy,
        submission_df=sample_submission_df,
        train_index=train_index,
        model_number=MODEL_NAME,
        run_id=RUN_ID,
        sub_dir=constants.SUBMISSION_DIR,
        oof_dir=constants.OOF_DIR,
        fi_dir=constants.FI_DIR,
        fi_fig_dir=constants.FI_FIG_DIR,
    )

    end = timer()
    common.update_tracking(RUN_ID, "training_time", end - start, is_integer=True)
    common.update_tracking(RUN_ID, "comments", EXP_DETAILS)
    logger.info("Execution Complete")