text:
# Copyright (C) 2016 Antoine Carme <Antoine.Carme@Laposte.net>
# All rights reserved.

# This file is part of the Python Automatic Forecasting (PyAF) library and is made available under
# the terms of the 3 Clause BSD license

import pandas as pd
import numpy as np

import traceback

# for timing
import time

import threading

from . import Time as tsti
from . import Exogenous as tsexog
from . import MissingData as tsmiss
from . import Signal_Transformation as tstransf
from . import Perf as tsperf
from . import SignalDecomposition_Trend as tstr
from . import SignalDecomposition_Cycle as tscy
from . import SignalDecomposition_AR as tsar
from . import Options as tsopts
from . import TimeSeriesModel as tsmodel
from . import TimeSeries_Cutting as tscut
from . import Utils as tsutil

import copy

def sample_signal_if_needed(iInputDS, iOptions):
    logger = tsutil.get_pyaf_logger();
    lInputDS = iInputDS
    if(iOptions.mActivateSampling):
        if(iOptions.mDebugProfile):
            logger.info("PYAF_MODEL_SAMPLING_ACTIVATED " +
                        str((iOptions.mSamplingThreshold, iOptions.mSeed)));
        lInputDS = iInputDS.tail(iOptions.mSamplingThreshold);
    return lInputDS

class cSignalDecompositionOneTransform:

    def __init__(self):
        self.mSignalFrame = None
        self.mTime = None
        self.mSignal = None
        self.mTimeInfo = tsti.cTimeInfo();
        self.mTransformation = tstransf.cSignalTransform_None();

    def setParams(self , iInputDS, iTime, iSignal, iHorizon, iTransformation,
                  iDecomspositionType, iExogenousData = None):
        assert(iInputDS.shape[0] > 0)
        assert(iInputDS.shape[1] > 0)
        assert(iTime in iInputDS.columns)
        assert(iSignal in iInputDS.columns)

        # print("setParams , head", iInputDS.head());
        # print("setParams , tail", iInputDS.tail());
        # print("setParams , columns", iInputDS.columns);

        self.mTime = iTime
        self.mOriginalSignal = iSignal;

        self.mDecompositionType = iDecomspositionType

        self.mTransformation = iTransformation;
        self.mTransformation.mOriginalSignal = iSignal;
        self.mTransformation.mOptions = self.mOptions;

        self.mSignal = iTransformation.get_name(iSignal)
        self.mHorizon = iHorizon;

        self.mSplit = tscut.cCuttingInfo()
        self.mSplit.mTime = self.mTime;
        self.mSplit.mSignal = self.mSignal;
        self.mSplit.mOriginalSignal = self.mOriginalSignal;
        self.mSplit.mHorizon = self.mHorizon;
        self.mSplit.mOptions = self.mOptions;

        self.mTimeInfo = tsti.cTimeInfo();
        self.mTimeInfo.mTime = self.mTime;
        self.mTimeInfo.mSignal = self.mSignal;
        self.mTimeInfo.mOriginalSignal = self.mOriginalSignal;
        self.mTimeInfo.mHorizon = self.mHorizon;
        self.mTimeInfo.mOptions = self.mOptions;
        self.mTimeInfo.mSplit = self.mSplit;

        self.mExogenousInfo = None;
        if(iExogenousData is not None):
            self.mExogenousInfo = tsexog.cExogenousInfo();
            self.mExogenousInfo.mExogenousData = iExogenousData;
            self.mExogenousInfo.mTimeInfo = self.mTimeInfo;
            self.mExogenousInfo.mOptions = self.mOptions;

    def updatePerfsForAllModels(self , iModels):
        self.mPerfsByModel = {}
        for model in iModels.keys():
            iModels[model].updatePerfs();

        for (name, model) in iModels.items():
            # print(name, model.__dict__);
            lComplexity = model.getComplexity();
            lFitPerf = model.mFitPerf;
            lForecastPerf = model.mForecastPerf;
            lTestPerf = model.mTestPerf;
            lSplit = model.mTimeInfo.mOptions.mCustomSplit
            self.mPerfsByModel[(self.mSignal, self.mDecompositionType, lSplit, model.mOutName)] = [(self.mSignal, self.mDecompositionType, model), lComplexity, lFitPerf , lForecastPerf, lTestPerf];
        return iModels;

    def train(self , iInputDS, iSplit, iTime, iSignal,
              iHorizon, iTransformation, iDecomspositionType):
        logger = tsutil.get_pyaf_logger();

        start_time = time.time()
        lInputDS = iInputDS[[iTime, iSignal]].copy()
        lInputDS = sample_signal_if_needed(lInputDS, self.mOptions)

        self.setParams(lInputDS, iTime, iSignal, iHorizon, iTransformation, iDecomspositionType, self.mExogenousData);

        lMissingImputer = tsmiss.cMissingDataImputer()
        lMissingImputer.mOptions = self.mOptions
        self.mSignalFrame = lMissingImputer.apply(lInputDS, iTime, iSignal).copy()
        assert(self.mSignalFrame.shape[0] > 0)

        # estimate time info
        # assert(self.mTimeInfo.mSignalFrame.shape[0] == iInputDS.shape[0])
        self.mSplit.mSignalFrame = self.mSignalFrame;
        self.mSplit.estimate();
        self.mTimeInfo.mSignalFrame = self.mSignalFrame;
        self.mTimeInfo.estimate();
        self.mSignalFrame['row_number'] = np.arange(0, self.mSignalFrame.shape[0]);

        lSignal = self.mSignalFrame[self.mOriginalSignal]
        self.mTransformation.fit(lSignal);

        self.mSignalFrame[self.mSignal] = self.mTransformation.apply(lSignal);
        # self.mSignalFrame[self.mSignal] = self.mSignalFrame[self.mSignal].astype(np.float32);

        exog_start_time = time.time()
        if(self.mExogenousInfo is not None):
            self.mExogenousInfo.fit();
            if(self.mOptions.mDebugProfile):
                logger.info("EXOGENOUS_ENCODING_TIME_IN_SECONDS " + str(self.mSignal) + " " + str(time.time() - exog_start_time))

        # estimate the trend

        lTrendEstimator = tstr.cTrendEstimator()
        lTrendEstimator.mSignalFrame = self.mSignalFrame
        lTrendEstimator.mTimeInfo = self.mTimeInfo
        lTrendEstimator.mSplit = self.mSplit
        lTrendEstimator.mOptions = self.mOptions;
        lTrendEstimator.mDecompositionType = iDecomspositionType

        trend_start_time = time.time()
        lTrendEstimator.estimateTrend();
        #lTrendEstimator.plotTrend();
        if(self.mOptions.mDebugProfile):
            logger.info("TREND_TIME_IN_SECONDS "  + str(self.mSignal) + " " + str(time.time() - trend_start_time))

        # estimate cycles
        cycle_start_time = time.time()

        lCycleEstimator = tscy.cCycleEstimator();
        lCycleEstimator.mTrendFrame = lTrendEstimator.mTrendFrame;
        lCycleEstimator.mTrendList = lTrendEstimator.mTrendList;

        del lTrendEstimator;

        lCycleEstimator.mTimeInfo = self.mTimeInfo
        lCycleEstimator.mSplit = self.mSplit
        lCycleEstimator.mDecompositionType = iDecomspositionType
        lCycleEstimator.mOptions = self.mOptions;

        lCycleEstimator.estimateAllCycles();
        # if(self.mOptions.mDebugCycles):
            # lCycleEstimator.plotCycles();
        if(self.mOptions.mDebugProfile):
            logger.info("CYCLE_TIME_IN_SECONDS "  + str(self.mSignal) + " " + str( str(time.time() - cycle_start_time)))

        # autoregressive
        ar_start_time = time.time()
        lAREstimator = tsar.cAutoRegressiveEstimator();
        lAREstimator.mCycleFrame = lCycleEstimator.mCycleFrame;
        lAREstimator.mTrendList = lCycleEstimator.mTrendList;
        lAREstimator.mCycleList = lCycleEstimator.mCycleList;

        del lCycleEstimator;

        lAREstimator.mTimeInfo = self.mTimeInfo
        lAREstimator.mSplit = self.mSplit
        lAREstimator.mDecompositionType = iDecomspositionType
        lAREstimator.mExogenousInfo = self.mExogenousInfo;
        lAREstimator.mOptions = self.mOptions;
        lAREstimator.estimate();
        #lAREstimator.plotAR();
        if(self.mOptions.mDebugProfile):
            logger.info("AUTOREG_TIME_IN_SECONDS " + str(self.mSignal) + " " + str( str(time.time() - ar_start_time)))
        # forecast perfs

        perf_start_time = time.time()
        lModels = {};
        for trend in lAREstimator.mTrendList:
            for cycle in lAREstimator.mCycleList[trend]:
                cycle_residue = cycle.getCycleResidueName();
                for autoreg in lAREstimator.mARList[cycle_residue]:
                    lModel = tsmodel.cTimeSeriesModel(self.mTransformation, self.mDecompositionType,
                                                      trend, cycle, autoreg);
                    lModels[lModel.mOutName] = lModel;

        del lAREstimator;
        self.updatePerfsForAllModels(lModels);

        if(self.mOptions.mDebugProfile):
            logger.info("PERF_TIME_IN_SECONDS " + str(self.mSignal) + " " + str(len(lModels)) + " " + str( str(time.time() - perf_start_time)))

        if(self.mOptions.mDebugProfile):
            logger.info("TRAINING_TIME_IN_SECONDS "  + str(self.mSignal) + " " + str(time.time() - start_time))

class cTraining_Arg:
    def __init__(self , name):
        self.mName = name;
        self.mInputDS = None;
        self.mTime = None;
        self.mSignal = None;
        self.mHorizon = None;
        self.mTransformation = None;
        self.mSigDec = None;
        self.mSplit = None
        self.mResult = None;

class cModelSelector_OneSignal:
    def __init__(self):
        self.mOptions = None
        pass

    def collectPerformanceIndices_ModelSelection(self, iSignal, iSigDecs) :
        modelsel_start_time = time.time()
        logger = tsutil.get_pyaf_logger();

        rows_list = []
        lPerfsByModel = {}
        for (lName, sigdec) in iSigDecs.items():
            for (model , value) in sorted(sigdec.mPerfsByModel.items()):
                lPerfsByModel[model] = value
                lTranformName = sigdec.mSignal;
                lDecompType = model[1];
                lModelFormula = model
                lModelCategory = value[0][2].get_model_category()
                lSplit = value[0][2].mTimeInfo.mOptions.mCustomSplit
                #  value format : self.mPerfsByModel[lModel.mOutName] = [lModel, lComplexity, lFitPerf , lForecastPerf, lTestPerf];
                lComplexity = value[1];
                lFitPerf = value[2];
                lForecastPerf = value[3];
                lTestPerf = value[4];
                row = [lSplit, lTranformName, lDecompType, lModelFormula[3], lModelFormula, lModelCategory, lComplexity,
                       lFitPerf.getCriterionValue(self.mOptions.mModelSelection_Criterion),
                       lForecastPerf.getCriterionValue(self.mOptions.mModelSelection_Criterion),
                       lTestPerf.getCriterionValue(self.mOptions.mModelSelection_Criterion)]
                rows_list.append(row);
                if(self.mOptions.mDebugPerformance):
                    logger.info("collectPerformanceIndices : " +
                                str((self.mOptions.mModelSelection_Criterion,
                                     lModelFormula, lComplexity, round(row[7], 4))))

        self.mTrPerfDetails =  pd.DataFrame(rows_list, columns=
                                            ('Split', 'Transformation', 'DecompositionType',
                                             'Model', 'DetailedFormula', 'Category', 'Complexity',
                                             'Fit' + self.mOptions.mModelSelection_Criterion,
                                             'Forecast' + self.mOptions.mModelSelection_Criterion,
                                             'Test' + self.mOptions.mModelSelection_Criterion))
        # print(self.mTrPerfDetails.head(self.mTrPerfDetails.shape[0]));
        lIndicator = 'Forecast' + self.mOptions.mModelSelection_Criterion;
        lBestPerf = self.mTrPerfDetails[ lIndicator ].min();
        # allow a loss of one point (0.01 of MAPE) if complexity is reduced.
        assert(not np.isnan(lBestPerf))
        self.mTrPerfDetails.sort_values(by=[lIndicator, 'Complexity', 'Model'] ,
                                        ascending=[True, True, True],
                                        inplace=True);
        self.mTrPerfDetails = self.mTrPerfDetails.reset_index(drop=True);

        lInterestingModels = self.mTrPerfDetails[self.mTrPerfDetails[lIndicator] <= (lBestPerf + 0.01)].reset_index(drop=True);
        lInterestingModels.sort_values(by=['Complexity'] , ascending=True, inplace=True)
        # print(self.mTransformList);
        # print(lInterestingModels.head());
        # print(self.mPerfsByModel);
        lBestName = lInterestingModels['DetailedFormula'].iloc[0]
        lBestModel = lPerfsByModel[lBestName][0][2];
        # print("BEST_MODEL", lBestName, lBestModel)
        if(self.mOptions.mDebugProfile):
            logger.info("MODEL_SELECTION_TIME_IN_SECONDS "  + str(iSignal) + " " + str(time.time() - modelsel_start_time))
        self.mBestModel = lBestModel
        self.mPerfsByModel = lPerfsByModel
        self.mModelShortList = lInterestingModels[['Transformation', 'DecompositionType', 'Model', lIndicator, 'Complexity']]
        return (iSignal, lPerfsByModel, lBestModel, self.mModelShortList)

    def collectPerformanceIndices(self, iSignal, iSigDecs) :
        modelsel_start_time = time.time()
        logger = tsutil.get_pyaf_logger();

        rows_list = []
        lPerfsByModel = {}
        for (lName, sigdec) in iSigDecs.items():
            for (model , value) in sorted(sigdec.mPerfsByModel.items()):
                lPerfsByModel[model] = value;
                lTranformName = sigdec.mSignal;
                lDecompType = model[1];
                lModelFormula = model
                lModelCategory = value[0][2].get_model_category()
                lSplit = value[0][2].mTimeInfo.mOptions.mCustomSplit
                #  value format : self.mPerfsByModel[lModel.mOutName] = [lModel, lComplexity, lFitPerf , lForecastPerf, lTestPerf];
                lComplexity = value[1];
                lFitPerf = value[2];
                lForecastPerf = value[3];
                lTestPerf = value[4];
                row = [lSplit, lTranformName, lDecompType, lModelFormula[3], lModelFormula , lModelCategory, lComplexity,
                       lFitPerf.mCount, lFitPerf.mL1, lFitPerf.mL2, lFitPerf.mMAPE,  lFitPerf.mMASE, lFitPerf.mCRPS,
                       lForecastPerf.mCount, lForecastPerf.mL1, lForecastPerf.mL2, lForecastPerf.mMAPE, lForecastPerf.mMASE, lForecastPerf.mCRPS,
                       lTestPerf.mCount, lTestPerf.mL1, lTestPerf.mL2, lTestPerf.mMAPE, lTestPerf.mMASE, lTestPerf.mCRPS]
                rows_list.append(row);
                if(self.mOptions.mDebugPerformance):
                    lIndicatorValue = lForecastPerf.getCriterionValue(self.mOptions.mModelSelection_Criterion)
                    logger.info("collectPerformanceIndices : " +
                                str((self.mOptions.mModelSelection_Criterion, lModelFormula, lModelCategory,
                                     lComplexity, round(lIndicatorValue, 4))))

        self.mTrPerfDetails =  pd.DataFrame(rows_list, columns=
                                            ('Split', 'Transformation', 'DecompositionType', 'Model',
                                             'DetailedFormula', 'Category', 'Complexity',
                                             'FitCount', 'FitL1', 'FitL2', 'FitMAPE', 'FitMASE', 'FitCRPS',
                                             'ForecastCount', 'ForecastL1', 'ForecastL2', 'ForecastMAPE', 'ForecastMASE', 'ForecastCRPS',
                                             'TestCount', 'TestL1', 'TestL2', 'TestMAPE', 'TestMASE', 'TestCRPS'))
        # print(self.mTrPerfDetails.head(self.mTrPerfDetails.shape[0]));
        lIndicator = 'Forecast' + self.mOptions.mModelSelection_Criterion;
        lBestPerf = self.mTrPerfDetails[ lIndicator ].min();
        # allow a loss of one point (0.01 of MAPE) if complexity is reduced.
        assert(not np.isnan(lBestPerf))
        self.mTrPerfDetails.sort_values(by=[lIndicator, 'Complexity', 'Model'] ,
                                        ascending=[True, True, True],
                                        inplace=True);
        self.mTrPerfDetails = self.mTrPerfDetails.reset_index(drop=True);

        lInterestingModels = self.mTrPerfDetails[self.mTrPerfDetails[lIndicator] <= (lBestPerf + 0.01)].reset_index(drop=True);
        lInterestingModels.sort_values(by=['Complexity'] , ascending=True, inplace=True)
        # print(self.mTransformList);
        # print(lInterestingModels.head());
        lBestName = lInterestingModels['DetailedFormula'].iloc[0];
        lBestModel = lPerfsByModel[lBestName][0][2];
        if(self.mOptions.mDebugProfile):
            logger.info("MODEL_SELECTION_TIME_IN_SECONDS "  + str(self.mBestModel.mSignal) + " " + str(time.time() - modelsel_start_time))
        self.mBestModel = lBestModel
        self.mPerfsByModel = lPerfsByModel
        self.mModelShortList = lInterestingModels[['Transformation', 'DecompositionType', 'Model', lIndicator, 'Complexity']]
        return (iSignal, lPerfsByModel, lBestModel, self.mModelShortList)

    def perform_model_selection(self, lSignal, sigdecs):
        if(self.mOptions.mDebugPerformance):
            self.collectPerformanceIndices(lSignal, sigdecs);
        else:
            self.collectPerformanceIndices_ModelSelection(lSignal, sigdecs);

    def perform_model_selection_cross_validation(self):
        logger = tsutil.get_pyaf_logger();
        modelsel_start_time = time.time()
        # self.mTrPerfDetails.to_csv("perf_time_series_cross_val.csv")
        lIndicator = 'Forecast' + self.mOptions.mModelSelection_Criterion;
        lColumns = ['Category', 'Complexity', lIndicator]
        lPerfByCategory = self.mTrPerfDetails[lColumns].groupby(by=['Category'] , sort=False)[lIndicator].mean()
        lPerfByCategory_df = pd.DataFrame(lPerfByCategory).reset_index()
        lPerfByCategory_df.columns = ['Category' , lIndicator]
        print("CROSS_VAL_PERF", lPerfByCategory_df)
        # lPerfByCategory_df.to_csv("perf_time_series_cross_val_by_category.csv")
        lBestPerf = lPerfByCategory_df[ lIndicator ].min();
        lPerfByCategory_df.sort_values(by=[lIndicator, 'Category'] ,
                                ascending=[True, True],
                                inplace=True);
        lPerfByCategory_df = lPerfByCategory_df.reset_index(drop=True);

        lInterestingCategories_df = lPerfByCategory_df[lPerfByCategory_df[lIndicator] <= (lBestPerf + 0.01)].reset_index(drop=True);
        # print(lPerfByCategory_df.head());
        # print(lInterestingCategories_df.head());
        # print(self.mPerfsByModel);
        lInterestingCategories = list(lInterestingCategories_df['Category'].unique())
        self.mTrPerfDetails['IC'] = self.mTrPerfDetails['Category'].apply(lambda x :1 if x in lInterestingCategories else 0)
        lInterestingModels = self.mTrPerfDetails[self.mTrPerfDetails['IC'] == 1].copy()
        lInterestingModels.sort_values(by=['Complexity'] , ascending=True, inplace=True)
        # print(self.mTransformList);
        # print(lInterestingModels.head());
        lBestName = lInterestingModels['DetailedFormula'].iloc[0];
        lBestSplit = lInterestingModels['Split'].iloc[0];
        self.mBestModel = self.mPerfsByModel[lBestName][0][2];
        self.mModelShortList = lInterestingModels[['Model', 'Category', 'Split', lIndicator, 'IC']]
        # print("BEST_MODEL", lBestName, self.mBestModel)
        if(self.mOptions.mDebugProfile):
            logger.info("MODEL_SELECTION_TIME_IN_SECONDS "  + str(self.mBestModel.mSignal) + " " + str(time.time() - modelsel_start_time))
        pass

def run_transform_thread(arg):
    arg.mSigDec.train(arg.mInputDS, arg.mSplit, arg.mTime, arg.mSignal, arg.mHorizon, arg.mTransformation, arg.mDecompositionType);
    return arg;

def run_finalize_training(arg):
    (lSignal , sigdecs, lOptions) = arg

    lModelSelector = cModelSelector_OneSignal()
    lModelSelector.mOptions = lOptions
    lModelSelector.collectPerformanceIndices_ModelSelection(lSignal, sigdecs)
    lModelSelector.perform_model_selection(lSignal, sigdecs)
    if(lOptions.mCrossValidationOptions.mMethod is not None):
        lModelSelector.perform_model_selection_cross_validation()

    # Prediction Intervals
    lModelSelector.mBestModel.updatePerfs(compute_all_indicators = True);
    lModelSelector.mBestModel.computePredictionIntervals();
    return (lSignal, lModelSelector.mPerfsByModel, lModelSelector.mBestModel, lModelSelector.mTrPerfDetails, lModelSelector.mModelShortList)

class cSignalDecompositionTrainer:

    def __init__(self):
        self.mSigDecBySplitAndTransform = {};
        self.mOptions = tsopts.cSignalDecomposition_Options();
        self.mExogenousData = None;
        self.mTransformList = {}
        pass

    def define_splits(self):
        lSplits = [None]
        if(self.mOptions.mCrossValidationOptions.mMethod is not None):
            lFolds = self.mOptions.mCrossValidationOptions.mNbFolds
            lRatio = 1.0 / lFolds
            lSplits = [(k * lRatio , lRatio , 0.0) for k in range(lFolds // 2, lFolds)]
        return lSplits

    def train(self, iInputDS, iSplits, iTime, iSignals, iHorizon):
        self.train_all_transofrmations(iInputDS, iSplits, iTime, iSignals, iHorizon);
        self.finalize_training()
        # self.cleanup_after_model_selection()

    def finalize_training(self):
        # print([transform1.mFormula for transform1 in self.mTransformList]);
        args = [];
        for (lSignal , sigdecs) in self.mSigDecBySplitAndTransform.items():
            args = args + [(lSignal, sigdecs, self.mOptions)]

        self.mPerfsByModel = {}
        self.mTrPerfDetails = {}
        self.mModelShortList = {}
        self.mBestModels = {}
        NCores = min(len(args) , self.mOptions.mNbCores)
        if(self.mOptions.mParallelMode and NCores > 1):
            from multiprocessing import Pool
            pool = Pool(NCores)

            for res in pool.imap(run_finalize_training, args):
                # print("FINISHED_TRAINING" , res.mName);
                (lSignal, lPerfsByModel, lBestModel, lPerfDetails, lModelShortList) = res
                assert(self.mPerfsByModel.get(lSignal) is None)
                self.mPerfsByModel[lSignal] = lPerfsByModel;
                self.mBestModels[lSignal] = lBestModel
                self.mTrPerfDetails[lSignal] = lPerfDetails
                self.mModelShortList[lSignal] = lModelShortList
            pool.close()
            pool.join()
        else:
            for arg in args:
                res = run_finalize_training(arg)
                (lSignal, lPerfsByModel, lBestModel, lPerfDetails, lModelShortList) = res
                assert(self.mPerfsByModel.get(lSignal) is None)
                self.mPerfsByModel[lSignal] = lPerfsByModel;
                self.mBestModels[lSignal] = lBestModel
                self.mTrPerfDetails[lSignal] = lPerfDetails
                self.mModelShortList[lSignal] = lModelShortList

    def defineTransformations(self, iInputDS, iSplit, iTime, iSignal):
        lTransformationEstimator = tstransf.cTransformationEstimator()
        lTransformationEstimator.mOptions = self.mOptions
        lTransformationEstimator.defineTransformations(iInputDS, iTime, iSignal)
        self.mTransformList[(iSignal, iSplit)] = lTransformationEstimator.mTransformList

    def train_all_transofrmations(self , iInputDS, iSplits, iTimes, iSignals, iHorizons):
        # print([transform1.mFormula for transform1 in self.mTransformList]);
        args = [];
        for lSignal in iSignals:
            self.mSigDecBySplitAndTransform[lSignal] = {}
        lActiveDecompositionTypes = [decomp_type for (decomp_type, active_status) in self.mOptions.mActiveDecompositionTypes.items() if (active_status is True)]
        for lSplit in iSplits:
            for lSignal in iSignals:
                self.defineTransformations(iInputDS, lSplit, iTimes[lSignal], lSignal);
                for transform1 in self.mTransformList[(lSignal, lSplit)]:
                    for decomp_type in lActiveDecompositionTypes:
                        arg = cTraining_Arg(transform1.get_name(""));
                        arg.mName = (lSignal, str(lSplit) , transform1.get_name(""), decomp_type)
                        arg.mDecompositionType= decomp_type
                        arg.mSigDec = cSignalDecompositionOneTransform();
                        arg.mSigDec.mOptions = copy.copy(self.mOptions);
                        arg.mSigDec.mOptions.mCustomSplit = lSplit
                        arg.mSplit = lSplit
                        arg.mSigDec.mExogenousData = self.mExogenousData[lSignal];
                        arg.mInputDS = iInputDS;
                        arg.mTime = iTimes[lSignal];
                        arg.mSignal = lSignal;
                        arg.mHorizon = iHorizons[lSignal];
                        arg.mTransformation = transform1;
                        arg.mOptions = self.mOptions;
                        arg.mExogenousData = self.mExogenousData[lSignal];
                        arg.mResult = None;
                        args.append(arg);

        NCores = min(len(args) , self.mOptions.mNbCores)
        if(self.mOptions.mParallelMode and NCores > 1):
            from multiprocessing import Pool
            pool = Pool(NCores)
            for res in pool.imap(run_transform_thread, args):
                lSignal = res.mName[0]
                self.mSigDecBySplitAndTransform[lSignal][res.mName] = res.mSigDec;
            pool.close()
            pool.join()
        else:
            for arg in args:
                res = run_transform_thread(arg)
                lSignal = res.mName[0]
                self.mSigDecBySplitAndTransform[lSignal][res.mName] = res.mSigDec;

    def cleanup_after_model_selection(self):
        lSigDecByTransform = {}
        for (lSignal , sigdecs) in self.mSigDecBySplitAndTransform.items():
            lBestTransformationName = self.mBestModels[lSignal].mTransformation.get_name("")
            for (name, sigdec) in self.mSigDecBySplitAndTransform[lSignal].items():
                if(name == lBestTransformationName):
                    for modelname in sigdec.mPerfsByModel.keys():
                        # store only model names here.
                        sigdec.mPerfsByModel[modelname][0] = modelname
                        lSigDecByTransform[lSignal][name]  = sigdec
            # delete failing transformations
        del self.mSigDecBySplitAndTransform
        self.mSigDecBySplitAndTransform = lSigDecByTransform

def forecast_one_signal(arg):
    (lSignal, iDecomsposition, iInputDS, iHorizon) = arg
    lBestModel = iDecomsposition.mBestModels[lSignal]
    lMissingImputer = tsmiss.cMissingDataImputer()
    lMissingImputer.mOptions = iDecomsposition.mOptions
    lInputDS = iInputDS[[lBestModel.mTime, lSignal]].copy()
    lInputDS = lMissingImputer.apply(lInputDS, lBestModel.mTime, lBestModel.mOriginalSignal)
    lForecastFrame_i = lBestModel.forecast(lInputDS, iHorizon);
    return (lSignal, lBestModel.mTime, lForecastFrame_i)

class cSignalDecompositionForecaster:

    def __init__(self):
        pass

    def merge_frames(self, iFullFrame, iOneSignalFrame, iTime):
        if(iFullFrame is None):
            return iOneSignalFrame
        lTime = iFullFrame.columns[0]
        lForecastFrame = iFullFrame.merge(iOneSignalFrame, how='left', left_on=lTime, right_on=iTime);
        return lForecastFrame

    def forecast(self, iDecomsposition, iInputDS, iHorizons):
        lHorizons = {}
        for sig in iDecomsposition.mSignals:
            if(dict == type(iHorizons)):
                lHorizons[sig] = iHorizons[sig]
            else:
                lHorizons[sig] = int(iHorizons)

        lForecastFrame = None
        args = [];
        for lSignal in iDecomsposition.mSignals:
            args = args + [(lSignal, iDecomsposition, iInputDS, lHorizons[lSignal])]

        NCores = min(len(args) , iDecomsposition.mOptions.mNbCores)
        if(iDecomsposition.mOptions.mParallelMode and  NCores > 1):
            from multiprocessing import Pool
            pool = Pool(NCores)

            for res in pool.imap(forecast_one_signal, args):
                (lSignal, lTime, lForecastFrame_i) = res
                # print((lSignal, lTime, lForecastFrame_i.columns))
                lForecastFrame = self.merge_frames(lForecastFrame, lForecastFrame_i, lTime)
                del lForecastFrame_i
            pool.close()
            pool.join()
        else:
            for arg in args:
                res = forecast_one_signal(arg)
                (lSignal, lTime, lForecastFrame_i) = res
                lForecastFrame = self.merge_frames(lForecastFrame, lForecastFrame_i, lTime)
                del lForecastFrame_i

        return lForecastFrame;

class cSignalDecomposition:

    def __init__(self):
        self.mSigDecBySplitAndTransform = {};
        self.mOptions = tsopts.cSignalDecomposition_Options();
        self.mExogenousData = None;
        pass

    def checkData(self, iInputDS, iTime, iSignal, iHorizon, iExogenousData):
        if(iHorizon != int(iHorizon)):
            raise tsutil.PyAF_Error("PYAF_ERROR_NON_INTEGER_HORIZON " + str(iHorizon));
        if(iHorizon < 1):
            raise tsutil.PyAF_Error("PYAF_ERROR_NEGATIVE_OR_NULL_HORIZON " + str(iHorizon));
        if(iTime not in iInputDS.columns):
            raise tsutil.PyAF_Error("PYAF_ERROR_TIME_COLUMN_NOT_FOUND " + str(iTime));
        for lSignal in [iSignal]:
            if(lSignal not in iInputDS.columns):
                raise tsutil.PyAF_Error("PYAF_ERROR_SIGNAL_COLUMN_NOT_FOUND " + str(lSignal));
            type2 = iInputDS[lSignal].dtype
            # print(type2)
            if(type2.kind != 'i' and type2.kind != 'u' and type2.kind != 'f'):
                raise tsutil.PyAF_Error("PYAF_ERROR_SIGNAL_COLUMN_TYPE_NOT_ALLOWED '" + str(lSignal) + "' '" + str(type2) + "'");
        type1 = iInputDS[iTime].dtype
        # print(type1)
        if(type1.kind != 'M' and type1.kind != 'i' and type1.kind != 'u' and type1.kind != 'f'):
            raise tsutil.PyAF_Error("PYAF_ERROR_TIME_COLUMN_TYPE_NOT_ALLOWED '" + str(iTime) + "' '" + str(type1) + "'");
        # time in exogenous data should be the strictly same type as time in training dataset (join needed)
        if(iExogenousData is not None):
            lExogenousDataFrame = iExogenousData[0];
            lExogenousVariables = iExogenousData[1];
            if(iTime not in lExogenousDataFrame.columns):
                raise tsutil.PyAF_Error("PYAF_ERROR_TIME_COLUMN_NOT_FOUND_IN_EXOGENOUS " + str(iTime));
            for exog in lExogenousVariables:
                if(exog not in lExogenousDataFrame.columns):
                    raise tsutil.PyAF_Error("PYAF_ERROR_EXOGENOUS_VARIABLE_NOT_FOUND " + str(exog));

            type3 = lExogenousDataFrame[iTime].dtype
            if(type1 != type3):
                raise tsutil.PyAF_Error("PYAF_ERROR_INCOMPATIBLE_TIME_COLUMN_TYPE_IN_EXOGENOUS '" + str(iTime) + "' '" + str(type1)  + "' '" + str(type3) + "'");

    def reinterpret_by_signal_args(self, iTimes, iSignals, iHorizons, iExogenousData):
        # backward compatibility
        self.mSignals = iSignals
        if(str == type(iSignals)):
            self.mSignals = [iSignals]
        self.mDateColumns = {}
        for sig in self.mSignals:
            if(dict == type(iTimes)):
                self.mDateColumns[sig] = iTimes[sig]
            else:
                self.mDateColumns[sig] = iTimes
        self.mHorizons = {}
        for sig in self.mSignals:
            if(dict == type(iHorizons)):
                self.mHorizons[sig] = iHorizons[sig]
            else:
                self.mHorizons[sig] = int(iHorizons)
        self.mExogenousData = {}
        for sig in self.mSignals:
            if(dict == type(iExogenousData)):
                self.mExogenousData[sig] = iExogenousData[sig]
            else:
                self.mExogenousData[sig] = iExogenousData

    def train(self , iInputDS, iTimes, iSignals, iHorizons, iExogenousData = None):
        logger = tsutil.get_pyaf_logger();
        logger.info("START_TRAINING '" + str(iSignals) + "'")
        start_time = time.time()
        self.reinterpret_by_signal_args(iTimes, iSignals, iHorizons, iExogenousData)
        # print(iInputDS.shape, iInputDS.columns, self.mSignals, self.mDateColumns, self.mHorizons)

        for sig in self.mSignals:
            self.checkData(iInputDS, self.mDateColumns[sig], sig, self.mHorizons[sig], self.mExogenousData[sig]);

        self.mTrainingDataset = iInputDS;

        lTrainer = cSignalDecompositionTrainer()
        lSplits = lTrainer.define_splits()
        lTrainer.mOptions = self.mOptions;
        lTrainer.mExogenousData = self.mExogenousData;

        lTrainer.train(iInputDS, lSplits , self.mDateColumns, self.mSignals, self.mHorizons)
        self.mBestModels = lTrainer.mBestModels
        self.mTrPerfDetailsBySignal = lTrainer.mTrPerfDetails
        self.mModelShortListBySignal = lTrainer.mModelShortList
        # some backward compatibility
        lFirstSignal = self.mSignals[0]
        self.mTrPerfDetails = lTrainer.mTrPerfDetails[lFirstSignal]
        self.mModelShortList = lTrainer.mModelShortList[lFirstSignal]
        self.mBestModel = self.mBestModels[lFirstSignal]

        end_time = time.time()
        self.mTrainingTime = end_time - start_time;
        logger.info("END_TRAINING_TIME_IN_SECONDS '" + str(self.mSignals) + "' " + str(self.mTrainingTime))
        pass

    def forecast(self , iInputDS, iHorizon):
        start_time = time.time()
        logger = tsutil.get_pyaf_logger();
        logger.info("START_FORECASTING '" + str(self.mSignals) + "'")
        lForecaster = cSignalDecompositionForecaster()
        lInputDS = sample_signal_if_needed(iInputDS, self.mOptions)
        lForecastFrame = lForecaster.forecast(self, lInputDS, iHorizon)
        lForecastTime = time.time() - start_time;
        logger.info("END_FORECAST_TIME_IN_SECONDS  '" + str(self.mSignals) + "' " + str(lForecastTime))
        return lForecastFrame;

    def getModelFormula(self):
        lFormula = {}
        for lSignal in self.mSignals:
            lFormula[lSignal] = self.mBestModel.getFormula();
        return lFormula;

    def getModelInfo(self):
        for lSignal in self.mSignals:
            self.mBestModels[lSignal].getInfo()

    def to_dict(self, iWithOptions = False):
        dict1 = {}
        for lSignal in self.mSignals:
            dict1[lSignal] = self.mBestModels[lSignal].to_dict(iWithOptions);
        return dict1

    def standardPlots(self, name = None, format = 'png'):
        logger = tsutil.get_pyaf_logger();
        start_time = time.time()
        logger.info("START_PLOTTING")
        for lSignal in self.mSignals:
            lName = name
            if(name is not None):
                lName = str(name) + "_" + str(lSignal)
            self.mBestModels[lSignal].standardPlots(lName, format);
        lPlotTime = time.time() - start_time;
        logger.info("END_PLOTTING_TIME_IN_SECONDS " + str(lPlotTime))

    def getPlotsAsDict(self):
        lDict = {}
        for lSignal in self.mSignals:
            lDict[lSignal] = self.mBestModels[lSignal].getPlotsAsDict();
        return lDict;


text:
import argparse
import collections
import pandas
import numpy as np
import os
import gym

from keras.models import Sequential, Model
from keras.layers import Dense, Activation, Flatten, Input, Concatenate
from keras.optimizers import Adam
import tensorflow as tf

from rl.agents import NAFAgent
from rl.memory import SequentialMemory
from rl.random import OrnsteinUhlenbeckProcess
from rl.core import Processor
from noise_estimator import *

parser = argparse.ArgumentParser()
parser.add_argument('--log_dir', default='logs',
                    help='Log dir [default: logs]')
parser.add_argument('--reward', default='normal',
                    help='reward choice: normal/noisy/surrogate [default: normal]')
parser.add_argument('--weight', type=float, default=0.6,
                    help='Weight of random confusion matrix [default: 0.6]')
parser.add_argument('--noise_type', type=str, default='norm_all',
                    help='Type of noise added: norm_all/norm_one/anti_iden/max_one [default: norm_all]')
FLAGS = parser.parse_args()

REWARD = FLAGS.reward
WEIGHT = FLAGS.weight
NOISE_TYPE = FLAGS.noise_type

assert (NOISE_TYPE in ["norm_all", "norm_one", "anti_iden", "max_one"])

if REWARD == "normal":
    LOG_DIR = os.path.join(os.path.join(FLAGS.log_dir, "naf_pendulum"), "normal")
else:
    LOG_DIR = os.path.join(os.path.join(os.path.join(FLAGS.log_dir, "naf_pendulum"), NOISE_TYPE), str(WEIGHT))
ENV_NAME = 'Pendulum-v0'
# gym.undo_logger_setup()

if not os.path.exists(LOG_DIR):
    os.makedirs(LOG_DIR)
os.system('cp naf_pendulum.py %s' % (LOG_DIR)) # bkp of train procedure

def train():
    # Get the environment and extract the number of actions.
    env = gym.make(ENV_NAME)
    np.random.seed(123)
    env.seed(123)
    assert len(env.action_space.shape) == 1
    nb_actions = env.action_space.shape[0]

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    from keras import backend as K
    K.set_session(sess)

    # Build all necessary models: V, mu, and L networks.
    V_model = Sequential()
    V_model.add(Flatten(input_shape=(1,) + env.observation_space.shape))
    V_model.add(Dense(16))
    V_model.add(Activation('relu'))
    V_model.add(Dense(16))
    V_model.add(Activation('relu'))
    V_model.add(Dense(16))
    V_model.add(Activation('relu'))
    V_model.add(Dense(1))
    V_model.add(Activation('linear'))
    V_model.summary()

    mu_model = Sequential()
    mu_model.add(Flatten(input_shape=(1,) + env.observation_space.shape))
    mu_model.add(Dense(16))
    mu_model.add(Activation('relu'))
    mu_model.add(Dense(16))
    mu_model.add(Activation('relu'))
    mu_model.add(Dense(16))
    mu_model.add(Activation('relu'))
    mu_model.add(Dense(nb_actions))
    mu_model.add(Activation('linear'))
    mu_model.summary()

    action_input = Input(shape=(nb_actions,), name='action_input')
    observation_input = Input(shape=(1,) + env.observation_space.shape, name='observation_input')
    x = Concatenate()([action_input, Flatten()(observation_input)])
    x = Dense(32)(x)
    x = Activation('relu')(x)
    x = Dense(32)(x)
    x = Activation('relu')(x)
    x = Dense(32)(x)
    x = Activation('relu')(x)
    x = Dense(((nb_actions * nb_actions + nb_actions) // 2))(x)
    x = Activation('linear')(x)
    L_model = Model(inputs=[action_input, observation_input], outputs=x)
    L_model.summary()

    # Finally, we configure and compile our agent. You can use every built-in Keras optimizer and
    # even the metrics!
    memory = SequentialMemory(limit=100000, window_length=1)
    random_process = OrnsteinUhlenbeckProcess(theta=.15, mu=0., sigma=.3, size=nb_actions)

    if REWARD == "normal":
        processor = NAFPendulumProcessor()
        naf_normal = NAFAgent(nb_actions=nb_actions, V_model=V_model, L_model=L_model, mu_model=mu_model,
                         memory=memory, nb_steps_warmup=100, random_process=random_process,
                         gamma=.99, target_model_update=1e-3, processor=processor)
        naf_normal.compile(Adam(lr=.00025, clipnorm=1.), metrics=['mae'])
        history_normal = naf_normal.fit(env, nb_steps=150000, visualize=False, verbose=2, nb_max_episode_steps=200)
        naf_normal.save_weights(os.path.join(LOG_DIR, 'naf_normal_{}_weights.h5f'.format(ENV_NAME)), overwrite=True)
        naf_normal.test(env, nb_episodes=10, visualize=False, nb_max_episode_steps=200)

        pandas.DataFrame(history_normal.history).to_csv(os.path.join(LOG_DIR, "normal.csv"))

    elif REWARD == "noisy":
        # processor_noisy = PendulumSurrogateProcessor(weight=WEIGHT, surrogate=False, noise_type=NOISE_TYPE)
        processor_noisy = PendulumProcessor(weight=WEIGHT, surrogate=False, noise_type=NOISE_TYPE)
        naf_noisy = NAFAgent(nb_actions=nb_actions, V_model=V_model, L_model=L_model, mu_model=mu_model,
                             memory=memory, nb_steps_warmup=100, random_process=random_process,
                             gamma=.99, target_model_update=1e-3, processor=processor_noisy)
        naf_noisy.compile(Adam(lr=.00025, clipnorm=1.), metrics=['mae'])
        history_noisy = naf_noisy.fit(env, nb_steps=150000, visualize=False, verbose=2, nb_max_episode_steps=200)
        naf_noisy.save_weights(os.path.join(LOG_DIR, 'naf_noisy_{}_weights.h5f'.format(ENV_NAME)), overwrite=True)
        naf_noisy.test(env, nb_episodes=10, visualize=False, nb_max_episode_steps=200)

        pandas.DataFrame(history_noisy.history).to_csv(os.path.join(LOG_DIR, "noisy.csv"))

    elif REWARD == "surrogate":
        # processor_surrogate = PendulumSurrogateProcessor(weight=WEIGHT, surrogate=True, noise_type=NOISE_TYPE)
        processor_surrogate = PendulumProcessor(weight=WEIGHT, surrogate=True, noise_type=NOISE_TYPE)
        naf_surrogate = NAFAgent(nb_actions=nb_actions, V_model=V_model, L_model=L_model, mu_model=mu_model,
                                 memory=memory, nb_steps_warmup=100, random_process=random_process,
                                 gamma=.99, target_model_update=1e-3, processor=processor_surrogate)
        naf_surrogate.compile(Adam(lr=.00025, clipnorm=1.), metrics=['mae'])
        history_surrogate = naf_surrogate.fit(env, nb_steps=150000, visualize=False, verbose=2, nb_max_episode_steps=200)
        naf_surrogate.save_weights(os.path.join(LOG_DIR, 'naf_surrogate_{}_weights.h5f'.format(ENV_NAME)), overwrite=True)
        naf_surrogate.test(env, nb_episodes=10, visualize=False, nb_max_episode_steps=200)

        pandas.DataFrame(history_surrogate.history).to_csv(os.path.join(LOG_DIR, "surrogate.csv"))

    else:
        raise NotImplementedError

if __name__ == "__main__":
    train()

ext: py

max_stars_repo_name: AI-secure/rl-perturbed-reward

max_issues_repo_name: AI-secure/rl-perturbed-reward

max_forks_repo_name: AI-secure/rl-perturbed-reward

reason: import numpy


text:
"""	Plot offline data.
"""

__author__ = 'Achin Jain'
__email__ = 'achinj@seas.upenn.edu'

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec

from bayes_race.tracks import ETHZ
from bayes_race.models import Dynamic

#####################################################################
# settings

SAVE_RESULTS = False

SAMPLING_TIME = 0.02
HORIZON = 20

#####################################################################
# load track

N_SAMPLES = 390
TRACK_NAME = 'ETHZ'
LAPS_DONE = 1

track = ETHZ(reference='optimal', longer=True)

#####################################################################
# load inputs used to simulate Dynamic model

SUFFIX = 'NOCONS-'

data = np.load('../data/DYN-NMPC-{}{}.npz'.format(SUFFIX, TRACK_NAME))
time_dyn = data['time'][:N_SAMPLES+1]
states_dyn = data['states'][:,:N_SAMPLES+1]
inputs_dyn = data['inputs'][:,:N_SAMPLES]

data = np.load('../data/DYN-GPMPC-{}{}.npz'.format(SUFFIX, TRACK_NAME))
time_gp = data['time'][:N_SAMPLES+1]
states_gp = data['states'][:,:N_SAMPLES+1]
inputs_gp = data['inputs'][:,:N_SAMPLES]

data = np.load('../data/DYN-GPMPC{}-{}{}.npz'.format(LAPS_DONE, SUFFIX, TRACK_NAME))
time_gp_updated = data['time'][:N_SAMPLES+1]
states_gp_updated = data['states'][:,:N_SAMPLES+1]
inputs_gp_updated = data['inputs'][:,:N_SAMPLES]

#####################################################################
# plots

plt.figure(figsize=(6,4))
plt.axis('equal')
plt.plot(-track.y_outer, track.x_outer, 'k', lw=0.5, alpha=0.5)
plt.plot(-track.y_inner, track.x_inner, 'k', lw=0.5, alpha=0.5)
# plt.plot(-states_gp[1], states_gp[0], 'r', lw=1, label='$f_{\mathrm{corr}}$')
plt.plot(-states_gp_updated[1], states_gp_updated[0], 'r', lw=1, label='MPC (updated GP)')
plt.plot(-states_dyn[1], states_dyn[0], '--k', lw=1, label='MPC (exact)')
plt.xlabel('$x$ [m]')
plt.ylabel('$y$ [m]')
plt.legend(loc='upper center', ncol=3, bbox_to_anchor=(0.5,1.15), frameon=False)

INDEX = 0
plt.scatter(-states_gp_updated[1,INDEX], states_gp_updated[0,INDEX], color='r', marker='o', alpha=0.8, s=15)
plt.text(-states_gp[1,INDEX], states_gp[0,INDEX]+0.05, '0', color='k', fontsize=10, ha='center', va='bottom')
plt.scatter(-states_dyn[1,INDEX], states_dyn[0,INDEX], color='k', marker='o', alpha=0.8, s=15)
for INDEX in range(HORIZON+5,N_SAMPLES,HORIZON+5):
	plt.scatter(-states_gp_updated[1,INDEX], states_gp_updated[0,INDEX], color='r', marker='o', alpha=0.8, s=15)
	plt.text(-states_gp_updated[1,INDEX]+0.05, states_gp_updated[0,INDEX]+0.05, '{}'.format(INDEX*SAMPLING_TIME), color='r', fontsize=10)
	plt.scatter(-states_dyn[1,INDEX], states_dyn[0,INDEX], color='k', marker='o', alpha=0.8, s=15)
	plt.text(-states_dyn[1,INDEX]-0.05, states_dyn[0,INDEX]-0.05, '{}'.format(INDEX*SAMPLING_TIME), color='k', fontsize=10, ha='right', va='top')

filepath = 'track_mpc_lap{}.png'.format(LAPS_DONE)
if SAVE_RESULTS:
	plt.savefig(filepath, dpi=600, bbox_inches='tight')

plt.figure(figsize=(6,4.3))
gs = gridspec.GridSpec(2,1)

plt.subplot(gs[0,:])
ax = plt.gca()
plt.plot(time_gp_updated[:-1], inputs_gp_updated[1], 'r', lw=1, label='MPC (updated GP)')
plt.plot(time_dyn[:-1], inputs_dyn[1], '--k', lw=1, label='MPC (exact)')
plt.ylabel('steering $\delta$ [rad]')
plt.xlim([0, N_SAMPLES*SAMPLING_TIME])
ax.set_xticks([])
plt.legend(loc='upper center', ncol=3, bbox_to_anchor=(0.5,1.31), frameon=False)

plt.subplot(gs[1,:])
plt.plot(time_gp_updated[:-1], inputs_gp_updated[0], 'r', lw=1, label='MPC (updated GP)')
plt.plot(time_dyn[:-1], inputs_dyn[0], '--k', lw=1, label='MPC (exact)')
plt.ylabel('PWM $d$ [-]')
plt.xlabel('time [s]')
plt.xlim([0, N_SAMPLES*SAMPLING_TIME])

filepath = 'inputs_mpc_lap{}.png'.format(LAPS_DONE)
if SAVE_RESULTS:
	plt.savefig(filepath, dpi=600, bbox_inches='tight')

plt.show()

max_stars_repo_name: DaniMarts/bayesrace

max_issues_repo_name: DaniMarts/bayesrace

max_forks_repo_name: DaniMarts/bayesrace


text:
import numpy as np
import gym3
from procgen import ProcgenGym3Env

def get_procgen_venv(*, env_id, num_envs, rendering=False, **env_kwargs):
    if rendering:
        env_kwargs["render_human"] = True

    env = ProcgenGym3Env(num=num_envs, env_name=env_id, **env_kwargs)

    env = gym3.ExtractDictObWrapper(env, "rgb")

    if rendering:
        env = gym3.ViewerWrapper(env, info_key="rgb")
    return env

def get_venv(num_envs, env_name, **env_kwargs):
    venv = get_procgen_venv(num_envs=num_envs, env_id=env_name, **env_kwargs)

    return venv


text:
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import paddle

from ..backends import depth_convert
from .utils import randint, weighted_sampling

__all__ = ['depth_augment', 'spect_augment', 'random_crop1d', 'random_crop2d']

# example y = depth_augment(y,['int8','int16'],[0.8,0.1])
def depth_augment(y, choices=['int8', 'int16'], probs=[0.5, 0.5]):
    assert len(probs) == len(choices), 'number of choices {} must be equal to size of probs {}'.format(
        len(choices), len(probs))
    k = weighted_sampling(probs)
    #k = randint(len(choices))
    src_depth = y.dtype
    y1 = depth_convert(y, choices[k])
    y2 = depth_convert(y1, src_depth)
    return y2

def adaptive_spect_augment(spect, tempo_axis=0, level=0.1):

    assert spect.ndim == 2., 'only supports 2d tensor or numpy array'
    if tempo_axis == 0:
        nt, nf = spect.shape
    else:
        nf, nt = spect.shape

    time_mask_width = int(nt * level * 0.5)
    freq_mask_width = int(nf * level * 0.5)

    num_time_mask = int(10 * level)
    num_freq_mask = int(10 * level)

    # num_zeros = num_time_mask*time_mask_width*nf + num_freq_mask*freq_mask_width*nt
    # factor = (nt*nf)/(nt*nf-num_zeros)

    if tempo_axis == 0:
        for i in range(num_time_mask):
            start = randint(nt - time_mask_width)
            spect[start:start + time_mask_width, :] = 0
        for i in range(num_freq_mask):
            start = randint(nf - freq_mask_width)
            spect[:, start:start + freq_mask_width] = 0
    else:
        for i in range(num_time_mask):
            start = randint(nt - time_mask_width)
            spect[:, start:start + time_mask_width] = 0
        for i in range(num_freq_mask):
            start = randint(nf - freq_mask_width)
            spect[start:start + freq_mask_width, :] = 0

    return spect

def spect_augment(
    spect,
    tempo_axis=0,
    max_time_mask=3,
    max_freq_mask=3,
    max_time_mask_width=30,
    max_freq_mask_width=20,
):

    assert spect.ndim == 2., 'only supports 2d tensor or numpy array'
    if tempo_axis == 0:
        nt, nf = spect.shape
    else:
        nf, nt = spect.shape

    num_time_mask = randint(max_time_mask)
    num_freq_mask = randint(max_freq_mask)

    time_mask_width = randint(max_time_mask_width)
    freq_mask_width = randint(max_freq_mask_width)

    #print(num_time_mask)
    #print(num_freq_mask)

    if tempo_axis == 0:
        for i in range(num_time_mask):
            start = randint(nt - time_mask_width)
            spect[start:start + time_mask_width, :] = 0
        for i in range(num_freq_mask):
            start = randint(nf - freq_mask_width)
            spect[:, start:start + freq_mask_width] = 0
    else:
        for i in range(num_time_mask):
            start = randint(nt - time_mask_width)
            spect[:, start:start + time_mask_width] = 0
        for i in range(num_freq_mask):
            start = randint(nf - freq_mask_width)
            spect[start:start + freq_mask_width, :] = 0

    return spect

def random_crop1d(y, crop_len):
    assert y.ndim == 1, 'only accept 1d tensor or numpy array'
    n = len(y)
    idx = randint(n - crop_len)
    return y[idx:idx + crop_len]

def random_crop2d(s, crop_len, tempo_axis=0):  # random crop according to temporal direction
    assert tempo_axis < s.ndim, 'axis out of range'
    n = s.shape[tempo_axis]
    idx = randint(high=n - crop_len)
    if type(s) == np.ndarray:
        sli = [slice(None) for i in range(s.ndim)]
        sli[tempo_axis] = slice(idx, idx + crop_len)
        out = s[tuple(sli)]
    else:
        out = paddle.index_select(s, paddle.Tensor(np.array([i for i in range(idx, idx + crop_len)])), axis=tempo_axis)
    return out

max_issues_repo_name: livingbody/models

max_forks_repo_name: livingbody/models


text:
import numpy as np
from astropy.io import fits
import pandas as pd
import crossmatch
import matplotlib.pyplot as plt
plt.style.use('seaborn')
import plotly.express as px
import time

def calculate_mean(vals):
  """slow version to compute mean
  """
  return sum(vals)/len(vals)

def calc_stats(data):
  data = np.loadtxt(data, delimiter=',')
  mean = np.mean(data)
  median = np.median(data)
  return (mean, median)

def mean_csvs(datafiles):
  """get the element-wise mean of a list of csv files
  """
  n = len(datafiles)
  if n > 0:
    data = np.loadtxt(datafiles[0], delimiter=',')
    for i in range(1,n):
      data += np.loadtxt(datafiles[i], delimiter=',')

    mean = data/n
    return mean

def mean_fits(fits_files):
  """get the element-wise means of a list of fits files
  """
  n = len(fits_files)
  if n > 0:
    with fits.open(fits_files[0]) as hdulist:
      data = np.copy(hdulist[0].data)
      for i in range(1,n):
        data += fits.open(fits_files[i])[0].data

    data_mean = data/n
    return data_mean

def get_argmax_fits(fits_file):
  """get index of highest value in a fits file
  """
  with fits.open(fits_file) as hdulist:
    data = hdulist[0].data
    argmax = np.argwhere(data==data.max())[0]
  return tuple(argmax)

def mean_median(data):
  """get the mean and median of a list of values
  """
  mean = np.mean(data)
  data_sorted = np.sort(data)
  mid = int(len(data_sorted)/2)
  if len(data_sorted)%2 == 0:
    median = (data_sorted[mid]+data_sorted[mid-1])/2.0
  else:
    median = data_sorted[mid]
  median = round(median, 3)
  return (round(float(median), 2), mean)

def median_fits(fitsfiles):
  """get median of list of fits files using numpy stacking
  """
  # read in all the FITS files and store in list
  fits_list = []
  for fitsfile in fitsfiles:
    with fits.open(fitsfile) as hdulist:
      fits_list.append(hdulist[0].data)

  # stack fits arrays in 3D arrays
  fits_stack = np.dstack(fits_list)
  median = np.median(fits_stack, axis=2)

  # calculate the memory consumed by the data (kB)
  memory = fits_stack.nbytes / 1024

  return median

def welfords_method(fitsfiles):
  '''Calculates the running mean and stdev for a list of FITS files
  using Welford's method.
  https://jonisalonen.com/2013/deriving-welfords-method-for-computing-variance/
  https://www.jstor.org/stable/1266577
  '''

  for n, fitsfile in enumerate(fitsfiles, 1):
    with fits.open(fitsfile) as hdulist:
      data = hdulist[0].data
      # initialize mean and std
      if n == 1:
        mean = np.zeros_like(data)
        s = np.zeros_like(data)

      # running (x - x_bar)
      delta = data - mean
      # running x_bar
      mean += delta/n
      # running std
      s += delta*(data - mean)

  # compute total std dev
  s /= (n - 1)
  np.sqrt(s, s)

  if n < 2:
    return mean, None
  else:
    return mean, s

def median_bins(data, nbins):
  """use bin approximation method to compute median; this function
  organizes the data samples into bins
  """
  mean = np.mean(data)
  std = np.std(data)
  left_bin = 0
  bins = np.zeros(nbins)
  bin_width = 2 * std / nbins

  for val in data:
    # place data points smaller than min val to ignore bin
    if val < mean-std:
      left_bin += 1
    # stores data points within range into bins
    elif val < mean+std:
      bin = int((val - (mean-std))/bin_width)
      bins[bin] += 1
    # Ignore values above mean + std

  return mean, std, left_bin, bins

def median_approx(values, nbins):
  """use bin approximation method to compute median; this function
  calls median_bins, then approximates the median by counting
  magnitudes of each bin in sequence.
  """
  mean, std, left_bin, bins = median_bins(values, nbins)

  # position of the middle element
  mid_count = (len(values)+1)/2

  # start with the ignored bin (val < mean-std)
  count = left_bin

  # add bin count to total count
  for b, bin_count in enumerate(bins):
    count += bin_count
    if count >= mid_count:
      break

  # get the median approximation
  width = 2 * std/nbins
  median = mean - std + width*(b+0.5)
  return median

def median_bins_fits(fitsfiles, nbins):

  # calculate the mean and standard deviation using the running method
  mean, std = welfords_method(fitsfiles)

  # get dimensions of the FITS file
  dim = mean.shape

  # initialise 3D bins
  left_bin = np.zeros(dim)
  bins = np.zeros((dim[0], dim[1], nbins))
  bin_width = 2 * std / nbins

  # loop over all FITS files
  for fitsfile in fitsfiles:
    with fits.open(fitsfile) as hdulist:
      data = hdulist[0].data

      # loop over every point in the 2D array
      for i in range(dim[0]):
        for j in range(dim[1]):
          value = data[i, j]
          lower_bound = mean[i, j] - std[i, j]
          upper_bound = mean[i, j] + std[i, j]

          # adding to the ignored bin (val < x_bar-std)
          if value < lower_bound:
            left_bin[i, j] += 1
          # adding to the appropriate bin
          elif value >= lower_bound and value < upper_bound:
            bin = int( (value - lower_bound) / bin_width[i, j] )
            bins[i, j, bin] += 1
          # ignoring values higher than upper bound limit

  return mean, std, left_bin, bins

def median_approx_fits(filenames, nbins):
  mean, std, left_bin, bins = median_bins_fits(filenames, nbins)

  # get dimensions of the FITS file
  dim = mean.shape

  # position of the middle element over all files
  mid = (len(filenames) + 1) / 2

  # calculate the approximated median for each array element
  bin_width = 2 * std / nbins
  median = np.zeros(dim)
  for i in range(dim[0]):
    for j in range(dim[1]):

      count = left_bin[i, j]
      for b, bincount in enumerate(bins[i, j]):
        count += bincount
        if count >= mid:
          # stop when the cumulative count exceeds the midpoint
          break
      median[i, j] = mean[i, j] - std[i, j] + bin_width[i, j]*(b + 0.5)

  return median

def hms2dec(hours, minutes, seconds):
  """convert right ascension HMS to decimal degrees; 0 <= hours < 24;
  RA = angle from the vernal equinox to the point, going east along
  the celestial equator
  """
  return 15*(hours + minutes/60 + seconds/3600)

def dms2dec(degrees, arcminutes, arcseconds):
  """convert declination from DMS to decimal degrees; -90 <= degrees < 90;
  the anglefrom the celestial equator to the point, (+) going north
  """
  return degrees + np.sign(degrees)*(arcminutes/60 + arcseconds/3600)

def greatcirc_dist(coords1, coords2):
  """computes the great circle distance between 2 points on a sphere
  using RA and Dec; applies the haversine formula
  :param coords1: [RA, Dec] for first object (radians)
  :param coords2: [RA, Dec] for second object (radians)
  :return: angular distance (radians)
  """
  ra1 = coords1[0]
  dec1 = coords1[1]
  ra2 = coords2[0]
  dec2 = coords2[1]

  b = np.cos(dec1)*np.cos(dec2)*np.sin(np.abs(ra1 - ra2)/2)**2
  a = np.sin(np.abs(dec1 - dec2)/2)**2
  d = 2*np.arcsin(np.sqrt(a + b))

  return d

def test_exploratory():
  fluxes = [23.3, 42.1, 2.0, -3.2, 55.6]
  m = np.mean(fluxes)
  m2 = sum(fluxes)/len(fluxes)
  print(m, m2)

  h_degs = greatcirc_dist(np.deg2rad([21.07, 0.1]),
                          np.deg2rad([21.15, 8.2]))
  assert np.allclose(np.rad2deg(h_degs), 8.1003923)

if __name__ == '__main__':

  # testing feather loading
  start = time.perf_counter()
  data = pd.read_csv('outputs/data_processed.csv', index_col=0, header=0)
  print('processing time, csv (s):', time.perf_counter() - start)
  start = time.perf_counter()
  data = pd.read_feather('outputs/data_processed.feather')
  print('processing time, feather (s):', time.perf_counter() - start)

  pass
  start = time.perf_counter()
  # data = pd.read_csv('outputs/data_processed.csv', index_col=0, header=0)
  # print(data)
  # print(data.columns)

  data = data[data['gz2class'] != 'A']
  data_spiral = data[data['gz2class'] == 'spiral']
  data_ellip = data[data['gz2class'] == 'elliptical']
  # print(data_spiral)
  # print(data_ellip)

  # statistics for 4th moments
  data = data[data['mCr4_u'] > -9000]
  data = data[data['mCr4_g'] > -9000]
  data = data[data['mCr4_r'] > -9000]
  data = data[data['mCr4_i'] > -9000]
  data = data[data['mCr4_z'] > -9000]

  # box plots
  # adaptive fourth moments of objects
  mCr4s = ["mCr4_u", "mCr4_g", "mCr4_r", "mCr4_i", "mCr4_z"]
  for mCr4 in mCr4s:
    fig = px.box(data, x="gz2class", y=mCr4, notched=True, points='all')
    fig.update_layout(
      title=f"Adaptive Fourth Moments of Object, {mCr4}",
      xaxis_title="Galaxy Type",
      yaxis_title="Intensity",
      legend_title="Galaxy Type",
      font=dict(size=18))
    fig.show()

  # pflux
  cidxs = ['petro_cidx_u','petro_cidx_g','petro_cidx_r','petro_cidx_i','petro_cidx_z']
  for cidx in cidxs:
    fig = px.box(data, x="gz2class", y=cidx, notched=True, points='all')
    fig.update_layout(
      title=f"Petrosian Concentration Index, {cidx}",
      xaxis_title="Galaxy Type",
      yaxis_title="Intensity",
      legend_title="Galaxy Type",
      font=dict(size=18))
    fig.show()

  # color filter statistics plots
  bands = ["u-g", "g-r", "r-i", "i-z"]
  for band in bands:
    fig = px.box(data, x="gz2class", y=band, notched=True, points='all')
    fig.update_layout(
      title=f"Color Bands, {band}",
      xaxis_title="Galaxy Type",
      yaxis_title="Intensity",
      legend_title="Galaxy Type",
      font=dict(size=18))
    fig.show()

  # redshift statistics
  data = data[data['z1'] <= 0.25]
  data_spiral = data_spiral[data_spiral['z1'] <= 0.25]
  data_ellip = data_spiral[data_spiral['z1'] <= 0.25]
  redshift = px.box(data, x="gz2class", y="z1", notched=True, points='all')
  # redshift.show()

  # polar plot of redshift with galaxy types
  fig = px.scatter_polar(data, r="z1", theta="ra", color='gz2class',
  category_orders={'gz2class': ['elliptical', 'spiral']}, opacity=0.7)
  fig.update_traces(marker=dict(size=3,))
  fig.update_layout(legend=dict(
      yanchor="top", y=0.65, xanchor="right", x=0.85,
      font=dict(size=18, color="black"),))
  fig.show()
  print('processing time (s):', time.perf_counter() - start)

max_stars_repo_name: snpham/jacaranda

max_issues_repo_name: snpham/jacaranda

max_forks_repo_name: snpham/jacaranda

reason: import numpy,from astropy


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

import numpy as np
import os
import csv
import time
import rospy
import matplotlib.pyplot as plt
import matplotlib.colors
from scipy.interpolate import interp1d

from tinkerforge.ip_connection import IPConnection
from tinkerforge.bricklet_rs232_v2 import BrickletRS232V2
from tinkerforge.brick_master import BrickMaster

HOST = "localhost"
PORT = 4223
UID_Master = "6SwGJu"
UID_base = "Mad"

# Turns on the LEDs to record data
def turn_on_led():
    # time.sleep(0.1)
    rs232_base.write([ord(c) for c in list("ATLED1=1\r\n")])
    time.sleep(0.1)
    rs232_base.write([ord(c) for c in list("ATLED3=1\r\n")])
    time.sleep(0.1)
    rs232_base.write([ord(c) for c in list("ATLED5=1\r\n")])
    time.sleep(0.1)

# Turns off LEDs after data is collected
def turn_off_led():
    # time.sleep(0.1)
    rs232_base.write([ord(c) for c in list("ATLED1=0\r\n")])
    time.sleep(0.1)
    rs232_base.write([ord(c) for c in list("ATLED3=0\r\n")])
    time.sleep(0.1)
    rs232_base.write([ord(c) for c in list("ATLED5=0\r\n")])
    time.sleep(0.1)

# Checks if the string is float, otherwise throws an error
def is_float(s):
    try:
        float(s)
        return True
    except ValueError:
        return False

# Receives the data from the Sparkfun triad spectroscopy sensor and reformats the wavelengths into a list of floats.
def process_data(wv):
    data = (''.join(rs232_base.read(150))).split()
    data = [s.strip(',') for s in data]
    data = [float(s) for s in data if is_float(s)]
    data = [x for _, x in sorted(zip(wv, data))]
    return data

# Plots the spectra from the collected data
def plot_spectra(wv_sorted, spectre, count):
    global Y_AXIS_MAX

    # Defines a color map
    norm = matplotlib.colors.Normalize(380, 940)
    colors = [[norm(380), "black"],  # so 410nm looks better
              [norm(405), "indigo"],
              [norm(427), "midnightblue"],
              [norm(435), "darkblue"],
              [norm(460), "blue"],
              [norm(487), "cyan"],
              [norm(510), "green"],
              [norm(520), "darkgreen"],
              [norm(570), "gold"],
              [norm(585), "orange"],
              [norm(610), "orangered"],
              [norm(640), "red"],

              [norm(670), "darkred"],

              [norm(690), "maroon"],
              [norm(720), "black"],
              [norm(940), "black"]]

    cmap = matplotlib.colors.LinearSegmentedColormap.from_list("", colors)
    """ END SET COLORS """
    if max(spectre) > Y_AXIS_MAX:
        Y_AXIS_MAX = max(spectre)

    plt.figure(1)

    f = interp1d(wv_sorted, spectre, kind='cubic')
    wv_new = np.linspace(410, 940, 1000)

    plt.rc('font', size=12)

    plt.title('Spectral Response')
    plt.grid(True)
    plt.ylabel('Intensity')
    plt.xlabel('Wavelength (nm)')
    plt.ylim(0, Y_AXIS_MAX*1.1)
    plt.xlim(410, 940)
    plt.plot(wv_new, f(wv_new), label=("Trial " + str(count)))
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.subplots_adjust(left=0.15, bottom=0.15, right=0.72)
    # plt.scatter(wv_new, f(wv_new), c=wv_new, norm=norm, cmap=cmap)
    # sc = ax.scatter(wv_new, f(wv_new), c=wv_new, norm=norm, cmap=cmap)

# Calculate the vegetation indices based on the wavelength intensity data
def calculate_indices(spectre):
    # 810 nm wavelength for NIR
    nir = spectre[14]

    # 680 nm wavelength for red visible
    red = spectre[10]

    # Four different vegetation indices.
    ndvi = (nir - red) / (nir + red)
    rdvi = (nir - red) / (nir + red)**0.5
    msr = ((nir/red)-1) / ((nir/red)+1)**0.5
    savi = (1.5 * (nir - red)) / (nir + red + 0.5)
    print("[NDVI, RDVI, MSR, SAVI]: ", [ndvi, rdvi, msr, savi])
    return [ndvi, rdvi, msr, savi]

if __name__ == "__main__":
    # Connect to TinkerForge stack and initialize ROS node
    ipcon = IPConnection()
    master = BrickMaster(UID_Master, ipcon)
    rs232_base = BrickletRS232V2(UID_base, ipcon)
    rospy.init_node('spectroscopy_base', anonymous=True)
    ipcon.connect(HOST, PORT)

    # Create array of wavelengths (sorted)
    wv = [610, 680, 730, 760, 810, 860, 560, 585, 645, 705, 900, 940, 410, 435, 460, 485, 510, 535]
    wv_sorted = [410, 435, 460, 485, 510, 535, 560, 585, 610, 645, 680, 705, 730, 760, 810, 860, 900, 940]

    # Prompts user for collecting data. Pressing enter turns the LEDs on, collects data, and then turns the LEDs off.
    Y_AXIS_MAX = 0
    counter = 0
    spectre_list = []
    while True:
        user_input = input("Press enter to record spectra, enter e to exit")
        if user_input == "e":
            break
        else:
            counter += 1

            # Turn off blue LED
            print("Turning off blue LED.")
            rs232_base.write([ord(c) for c in list("ATLED0=0\r\n")])
            time.sleep(0.2)

            # Collects data after turning on LEDs
            print("Collecting data...")
            turn_on_led()
            time.sleep(0.1)
            # spectre = collect_data(wv)
            rs232_base.write([ord(c) for c in list("ATCDATA\r\n")])
            time.sleep(0.5)
            turn_off_led()
            time.sleep(0.1)

            # Process the data collected and convert to plot-able format
            print("Data collected, plotting spectra.")
            spectre = process_data(wv)
            spectre_list.append(spectre)
            # print(spectre)

            # Calculate vegetation indices
            vegetation_indices = calculate_indices(spectre)
            plot_spectra(wv_sorted, spectre, counter)

    ipcon.disconnect()

    # Place spectroscopy tool data into .csv file
    folder_path = os.path.dirname(os.path.abspath(__file__))
    print(folder_path)
    file_path = os.path.join(folder_path, 'data_collection_csv/spectroscopy_data.csv')
    print(file_path)
    with open(file_path, 'w', newline='') as my_file:
        csv_writer = csv.writer(my_file, delimiter=',')
        csv_writer.writerow(wv_sorted)

        for row in spectre_list:
            csv_writer.writerow(row)

    # Displays all the plots based on the data.
    plt.show()

reason: import numpy,from scipy


text:
#!/usr/bin/env python
"""
Abstract: This script calculates and returns LDA plots based on normalized relative
          abundances or distance matrices (for e.g. unifrac distance matrix).
"""

import sys
import argparse
from itertools import cycle
from phylotoast import util, biom_calc as bc, graph_util as gu
errors = []
try:
    import matplotlib as mpl
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
    mpl.rc("font", family="Arial")  # define font for figure text
except ImportError as ie:
    errors.append(ie)
try:
    import numpy as np
except ImportError as ie:
    errors.append(ie)
try:
    import pandas as pd
except ImportError as ie:
    errors.append(ie)
try:
    import biom
except ImportError as ie:
    errors.append(ie)
try:
    from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
except ImportError as ie:
    errors.append(ie)
try:
    from palettable.colorbrewer.qualitative import Set3_12
except ImportError as ie:
    errors.append("No module named palettable")
if len(errors) != 0:
    for item in errors:
        print("Import Error:", item)
    sys.exit()

def plot_LDA(X_lda, y_lda, class_colors, exp_var, style, fig_size, label_pad,
             font_size, sids, dim=2, zangles=None, pt_size=250, out_fp=""):
    """
    Plot transformed LDA data.
    """
    cats = class_colors.keys()
    fig = plt.figure(figsize=fig_size)
    if dim == 3:
        try:
            assert X_lda.shape[1] >= 3
        except AssertionError:
            sys.exit("\nLinear Discriminant Analysis requires at least 4 groups of "
                     "samples to create a 3D figure. Please update group information or "
                     "use the default 2D view of the results.\n")
        if sids is not None:
            print("\nPoint annotations are available only for 2D figures.\n")
        ax = fig.add_subplot(111, projection="3d")
        ax.view_init(elev=zangles[1], azim=zangles[0])
        try:
            ax.set_zlabel("LD3 (Percent Explained Variance: {:.3f}%)".
                          format(exp_var[2]*100), fontsize=font_size, labelpad=label_pad)
        except:
            ax.set_zlabel("LD3", fontsize=font_size, labelpad=label_pad)
        for i, target_name in zip(range(len(cats)), cats):
            cat_x = X_lda[:, 0][y_lda == target_name]
            cat_y = X_lda[:, 1][y_lda == target_name]
            cat_z = X_lda[:, 2][y_lda == target_name]
            ax.scatter(xs=cat_x, ys=cat_y, zs=cat_z, label=target_name,
                       c=class_colors[target_name], alpha=0.85, s=pt_size, edgecolors="k",
                       zdir="z")
    else:
        ax = fig.add_subplot(111)
        for i, target_name in zip(range(len(cats)), cats):
            cat_x = X_lda[:, 0][y_lda == target_name]
            if X_lda.shape[1] == 1:
                cat_y = np.ones((cat_x.shape[0], 1)) + i
            else:
                cat_y = X_lda[:, 1][y_lda == target_name]
            ax.scatter(x=cat_x, y=cat_y, label=target_name, alpha=0.85, s=pt_size,
                       color=class_colors[target_name], edgecolors="k")
        # Annotate data points with sample IDs
        if sids is not None:
            for sample, point, group in zip(sids, X_lda, y_lda):
                try:
                    assert len(point) >= 2
                except AssertionError:
                    point = (point[0], cats.index(group)+1)
                finally:
                    ax.annotate(s=sample, xy=point[:2], xytext=(0, -15), ha="center",
                                va="center", textcoords="offset points")
    if X_lda.shape[1] == 1:
        plt.ylim((0.5, 2.5))
    try:
        ax.set_xlabel("LD1 (Percent Explained Variance: {:.3f}%)".
                      format(exp_var[0]*100), fontsize=font_size, labelpad=label_pad)
    except:
        ax.set_xlabel("LD1", fontsize=font_size, labelpad=label_pad)
    try:
        ax.set_ylabel("LD2 (Percent Explained Variance: {:.3f}%)".
                      format(exp_var[1]*100), fontsize=font_size, labelpad=label_pad)
    except:
        ax.set_ylabel("LD2", fontsize=font_size, labelpad=label_pad)

    leg = plt.legend(loc="best", scatterpoints=3, frameon=True, framealpha=1, fontsize=15)
    leg.get_frame().set_edgecolor('k')
    if dim == 2 and style:
        gu.ggplot2_style(ax)
        fc = "0.8"
    else:
        fc = "none"

    # save or display result
    if out_fp:
        plt.savefig(out_fp, facecolor=fc, edgecolor="none", dpi=300, bbox_inches="tight",
                    pad_inches=0.1)
    else:
        plt.show()

def run_LDA(df):
    """
    Run LinearDiscriminantAnalysis on input dataframe (df) and return
    transformed data, scalings and explained variance by discriminants.
    """
    # Prep variables for sklearn LDA
    X = df.iloc[:, 1:df.shape[1]].values     # input data matrix
    y = df["Condition"].values               # data categories list

    # Calculate LDA
    sklearn_lda = LDA()
    X_lda_sklearn = sklearn_lda.fit_transform(X, y)
    try:
        exp_var = sklearn_lda.explained_variance_ratio_
    except AttributeError as ae:
        print("\n{}: explained variance cannot be computed.\nPlease check this GitHub PR:"
              " https://github.com/scikit-learn/scikit-learn/pull/6027".format(ae))
        return X_lda_sklearn, y, "NA"
    return X_lda_sklearn, y, exp_var

def handle_program_options():
    parser = argparse.ArgumentParser(description="This script calculates and returns LDA "
                                                 "plots based on normalized relative "
                                                 "abundances or distance matrices "
                                                 "(for e.g. unifrac distance matrix).")
    parser.add_argument("-m", "--map_fp", required=True,
                        help="Metadata mapping file. [REQUIRED]")
    parser.add_argument("-g", "--group_by", required=True,
                        help="A column name in the mapping file containing categorical "
                             "values that will be used to identify groups. Each sample "
                             "ID must have a group entry. Default is no categories and "
                             "all the data will be treated as a single group. [REQUIRED]")
    parser.add_argument("-c", "--colors",
                        help="A column name in the mapping file containing hexadecimal "
                             "(#FF0000) color values that will be used to color the "
                             "groups. Each sample ID must have a color entry.")
    parser.add_argument("-ot", "--otu_table", help="Input biom file format OTU table.")
    parser.add_argument("-dm", "--dist_matrix_file", help="Input distance matrix file.")
    parser.add_argument("--save_lda_input",
                        help="Save a CSV-format file of the transposed LDA-input table to"
                             " the file specifed by this option.")
    parser.add_argument("--plot_title", default=None,
                        help="Plot title. Default is no title.")
    parser.add_argument("-o", "--out_fp", default="",
                        help="The path and file name to save the LDA plot under. If "
                             "specified, the figure will be saved directly instead of "
                             "opening a window in which the plot can be viewed before "
                             "saving.")
    parser.add_argument("-d", "--dimensions", default=2, type=int, choices=[2, 3],
                        help="Choose whether to plot 2D or 3D. Default is a 2D view.")
    parser.add_argument("-s", "--point_size", default=100, type=int,
                        help="Specify the size of the circles representing each of the "
                        "samples in the plot")
    parser.add_argument("--z_angles", type=float, nargs=2, default=[45., 30.],
                        help="Specify the azimuth and elevation angles for a 3D plot.")
    parser.add_argument("--figsize", default=[14, 8], type=int, nargs=2,
                        help="Specify the 'width height' in inches for LDA plots. Default"
                        " figure size is 14x8 inches.")
    parser.add_argument("--font_size", default=12, type=int,
                        help="Sets the font size for text elements in the plot.")
    parser.add_argument("--label_padding", default=15, type=int,
                        help="Sets the spacing in points between the each axis and its "
                        "label.")
    parser.add_argument("--annotate_points", action="store_true",
                        help="If specified, each data point will be labeled with its "
                        "sample ID. Default is False.")
    parser.add_argument("--ggplot2_style", action="store_true",
                        help="Apply ggplot2 styling to the figure. Default is False.")
    return parser.parse_args()

def main():
    args = handle_program_options()

    # Parse and read mapping file
    try:
        header, imap = util.parse_map_file(args.map_fp)
        category_idx = header.index(args.group_by)
    except IOError as ioe:
        err_msg = "\nError in metadata mapping filepath (-m): {}\n"
        sys.exit(err_msg.format(ioe))

    # Obtain group colors
    try:
        assert args.colors is not None
    except AssertionError:
        categories = {v[category_idx] for k, v in imap.items()}
        color_cycle = cycle(Set3_12.hex_colors)
        class_colors = {c: color_cycle.next() for c in categories}
    else:
        class_colors = util.color_mapping(imap, header, args.group_by, args.colors)

    if args.dist_matrix_file:
        try:
            dm_data = pd.read_csv(args.dist_matrix_file, sep="\t", index_col=0)
        except IOError as ioe:
            err_msg = "\nError with unifrac distance matrix file (-d): {}\n"
            sys.exit(err_msg.format(ioe))
        dm_data.insert(0, "Condition", [imap[str(sid)][category_idx] for sid in dm_data.index])
        if args.annotate_points:
            sampleids = [str(sid) for sid in dm_data.index]
        else:
            sampleids = None
        if args.save_lda_input:
            dm_data.to_csv(args.save_lda_input, sep="\t")
        # Run LDA
        X_lda, y_lda, exp_var = run_LDA(dm_data)
    else:
        # Load biom file and calculate relative abundance
        try:
            biomf = biom.load_table(args.otu_table)
        except IOError as ioe:
            err_msg = "\nError with biom format file (-d): {}\n"
            sys.exit(err_msg.format(ioe))
        # Get normalized relative abundances
        rel_abd = bc.relative_abundance(biomf)
        rel_abd = bc.arcsine_sqrt_transform(rel_abd)
        df_rel_abd = pd.DataFrame(rel_abd).T
        df_rel_abd.insert(0, "Condition", [imap[sid][category_idx]
                                           for sid in df_rel_abd.index])
        if args.annotate_points:
            sampleids = df_rel_abd.index
        else:
            sampleids = None
        if args.save_lda_input:
            df_rel_abd.to_csv(args.save_lda_input, sep="\t")
        # Run LDA
        X_lda, y_lda, exp_var = run_LDA(df_rel_abd)

    # Plot LDA
    if args.dimensions == 3:
        plot_LDA(X_lda, y_lda, class_colors, exp_var, style=args.ggplot2_style,
                 fig_size=args.figsize, label_pad=args.label_padding,
                 font_size=args.font_size, sids=sampleids, dim=3,
                 zangles=args.z_angles, pt_size=args.point_size, out_fp=args.out_fp)
    else:
        plot_LDA(X_lda, y_lda, class_colors, exp_var, style=args.ggplot2_style,
                 fig_size=args.figsize, label_pad=args.label_padding,
                 font_size=args.font_size, sids=sampleids, pt_size=args.point_size,
                 out_fp=args.out_fp)

if __name__ == "__main__":
    sys.exit(main())


text:
import os
import shutil
import tempfile
import contextlib
import unittest
import argparse
import sys
import io
import torch
import errno
import __main__

from numbers import Number
from torch._six import string_classes
from collections import OrderedDict

import numpy as np
from PIL import Image

@contextlib.contextmanager
def get_tmp_dir(src=None, **kwargs):
    tmp_dir = tempfile.mkdtemp(**kwargs)
    if src is not None:
        os.rmdir(tmp_dir)
        shutil.copytree(src, tmp_dir)
    try:
        yield tmp_dir
    finally:
        shutil.rmtree(tmp_dir)

ACCEPT = os.getenv('EXPECTTEST_ACCEPT')
TEST_WITH_SLOW = os.getenv('PYTORCH_TEST_WITH_SLOW', '0') == '1'
# TEST_WITH_SLOW = True  # TODO: Delete this line once there is a PYTORCH_TEST_WITH_SLOW aware CI job

parser = argparse.ArgumentParser(add_help=False)
parser.add_argument('--accept', action='store_true')
args, remaining = parser.parse_known_args()
if not ACCEPT:
    ACCEPT = args.accept
for i, arg in enumerate(sys.argv):
    if arg == '--accept':
        del sys.argv[i]
        break

class MapNestedTensorObjectImpl(object):
    def __init__(self, tensor_map_fn):
        self.tensor_map_fn = tensor_map_fn

    def __call__(self, object):
        if isinstance(object, torch.Tensor):
            return self.tensor_map_fn(object)

        elif isinstance(object, dict):
            mapped_dict = {}
            for key, value in object.items():
                mapped_dict[self(key)] = self(value)
            return mapped_dict

        elif isinstance(object, (list, tuple)):
            mapped_iter = []
            for iter in object:
                mapped_iter.append(self(iter))
            return mapped_iter if not isinstance(object, tuple) else tuple(mapped_iter)

        else:
            return object

def map_nested_tensor_object(object, tensor_map_fn):
    impl = MapNestedTensorObjectImpl(tensor_map_fn)
    return impl(object)

def is_iterable(obj):
    try:
        iter(obj)
        return True
    except TypeError:
        return False

# adapted from TestCase in torch/test/common_utils to accept non-string
# inputs and set maximum binary size
class TestCase(unittest.TestCase):
    precision = 1e-5

    def assertExpected(self, output, subname=None, prec=None, strip_suffix=None):
        r"""
        Test that a python value matches the recorded contents of a file
        derived from the name of this test and subname.  The value must be
        pickable with `torch.save`. This file
        is placed in the 'expect' directory in the same directory
        as the test script. You can automatically update the recorded test
        output using --accept.

        If you call this multiple times in a single function, you must
        give a unique subname each time.

        strip_suffix allows different tests that expect similar numerics, e.g.
        "test_xyz_cuda" and "test_xyz_cpu", to use the same pickled data.
        test_xyz_cuda would pass strip_suffix="_cuda", test_xyz_cpu would pass
        strip_suffix="_cpu", and they would both use a data file name based on
        "test_xyz".
        """
        def remove_prefix_suffix(text, prefix, suffix):
            if text.startswith(prefix):
                text = text[len(prefix):]
            if suffix is not None and text.endswith(suffix):
                text = text[:len(text) - len(suffix)]
            return text
        # NB: we take __file__ from the module that defined the test
        # class, so we place the expect directory where the test script
        # lives, NOT where test/common_utils.py lives.
        module_id = self.__class__.__module__
        munged_id = remove_prefix_suffix(self.id(), module_id + ".", strip_suffix)
        test_file = os.path.realpath(sys.modules[module_id].__file__)
        expected_file = os.path.join(os.path.dirname(test_file),
                                     "expect",
                                     munged_id)

        subname_output = ""
        if subname:
            expected_file += "_" + subname
            subname_output = " ({})".format(subname)
        expected_file += "_expect.pkl"
        expected = None

        def accept_output(update_type):
            print("Accepting {} for {}{}:\n\n{}".format(update_type, munged_id, subname_output, output))
            torch.save(output, expected_file)
            MAX_PICKLE_SIZE = 50 * 1000  # 50 KB
            binary_size = os.path.getsize(expected_file)
            self.assertTrue(binary_size <= MAX_PICKLE_SIZE)

        try:
            expected = torch.load(expected_file)
        except IOError as e:
            if e.errno != errno.ENOENT:
                raise
            elif ACCEPT:
                return accept_output("output")
            else:
                raise RuntimeError(
                    ("I got this output for {}{}:\n\n{}\n\n"
                     "No expect file exists; to accept the current output, run:\n"
                     "python {} {} --accept").format(munged_id, subname_output, output, __main__.__file__, munged_id))

        if ACCEPT:
            equal = False
            try:
                equal = self.assertEqual(output, expected, prec=prec)
            except Exception:
                equal = False
            if not equal:
                return accept_output("updated output")
        else:
            self.assertEqual(output, expected, prec=prec)

    def assertEqual(self, x, y, prec=None, message='', allow_inf=False):
        """
        This is copied from pytorch/test/common_utils.py's TestCase.assertEqual
        """
        if isinstance(prec, str) and message == '':
            message = prec
            prec = None
        if prec is None:
            prec = self.precision

        if isinstance(x, torch.Tensor) and isinstance(y, Number):
            self.assertEqual(x.item(), y, prec=prec, message=message,
                             allow_inf=allow_inf)
        elif isinstance(y, torch.Tensor) and isinstance(x, Number):
            self.assertEqual(x, y.item(), prec=prec, message=message,
                             allow_inf=allow_inf)
        elif isinstance(x, torch.Tensor) and isinstance(y, torch.Tensor):
            def assertTensorsEqual(a, b):
                super(TestCase, self).assertEqual(a.size(), b.size(), message)
                if a.numel() > 0:
                    if (a.device.type == 'cpu' and (a.dtype == torch.float16 or a.dtype == torch.bfloat16)):
                        # CPU half and bfloat16 tensors don't have the methods we need below
                        a = a.to(torch.float32)
                    b = b.to(a)

                    if (a.dtype == torch.bool) != (b.dtype == torch.bool):
                        raise TypeError("Was expecting both tensors to be bool type.")
                    else:
                        if a.dtype == torch.bool and b.dtype == torch.bool:
                            # we want to respect precision but as bool doesn't support substraction,
                            # boolean tensor has to be converted to int
                            a = a.to(torch.int)
                            b = b.to(torch.int)

                        diff = a - b
                        if a.is_floating_point():
                            # check that NaNs are in the same locations
                            nan_mask = torch.isnan(a)
                            self.assertTrue(torch.equal(nan_mask, torch.isnan(b)), message)
                            diff[nan_mask] = 0
                            # inf check if allow_inf=True
                            if allow_inf:
                                inf_mask = torch.isinf(a)
                                inf_sign = inf_mask.sign()
                                self.assertTrue(torch.equal(inf_sign, torch.isinf(b).sign()), message)
                                diff[inf_mask] = 0
                        # TODO: implement abs on CharTensor (int8)
                        if diff.is_signed() and diff.dtype != torch.int8:
                            diff = diff.abs()
                        max_err = diff.max()
                        tolerance = prec + prec * abs(a.max())
                        self.assertLessEqual(max_err, tolerance, message)
            super(TestCase, self).assertEqual(x.is_sparse, y.is_sparse, message)
            super(TestCase, self).assertEqual(x.is_quantized, y.is_quantized, message)
            if x.is_sparse:
                x = self.safeCoalesce(x)
                y = self.safeCoalesce(y)
                assertTensorsEqual(x._indices(), y._indices())
                assertTensorsEqual(x._values(), y._values())
            elif x.is_quantized and y.is_quantized:
                self.assertEqual(x.qscheme(), y.qscheme(), prec=prec,
                                 message=message, allow_inf=allow_inf)
                if x.qscheme() == torch.per_tensor_affine:
                    self.assertEqual(x.q_scale(), y.q_scale(), prec=prec,
                                     message=message, allow_inf=allow_inf)
                    self.assertEqual(x.q_zero_point(), y.q_zero_point(),
                                     prec=prec, message=message,
                                     allow_inf=allow_inf)
                elif x.qscheme() == torch.per_channel_affine:
                    self.assertEqual(x.q_per_channel_scales(), y.q_per_channel_scales(), prec=prec,
                                     message=message, allow_inf=allow_inf)
                    self.assertEqual(x.q_per_channel_zero_points(), y.q_per_channel_zero_points(),
                                     prec=prec, message=message,
                                     allow_inf=allow_inf)
                    self.assertEqual(x.q_per_channel_axis(), y.q_per_channel_axis(),
                                     prec=prec, message=message)
                self.assertEqual(x.dtype, y.dtype)
                self.assertEqual(x.int_repr().to(torch.int32),
                                 y.int_repr().to(torch.int32), prec=prec,
                                 message=message, allow_inf=allow_inf)
            else:
                assertTensorsEqual(x, y)
        elif isinstance(x, string_classes) and isinstance(y, string_classes):
            super(TestCase, self).assertEqual(x, y, message)
        elif type(x) == set and type(y) == set:
            super(TestCase, self).assertEqual(x, y, message)
        elif isinstance(x, dict) and isinstance(y, dict):
            if isinstance(x, OrderedDict) and isinstance(y, OrderedDict):
                self.assertEqual(x.items(), y.items(), prec=prec,
                                 message=message, allow_inf=allow_inf)
            else:
                self.assertEqual(set(x.keys()), set(y.keys()), prec=prec,
                                 message=message, allow_inf=allow_inf)
                key_list = list(x.keys())
                self.assertEqual([x[k] for k in key_list],
                                 [y[k] for k in key_list],
                                 prec=prec, message=message,
                                 allow_inf=allow_inf)
        elif is_iterable(x) and is_iterable(y):
            super(TestCase, self).assertEqual(len(x), len(y), message)
            for x_, y_ in zip(x, y):
                self.assertEqual(x_, y_, prec=prec, message=message,
                                 allow_inf=allow_inf)
        elif isinstance(x, bool) and isinstance(y, bool):
            super(TestCase, self).assertEqual(x, y, message)
        elif isinstance(x, Number) and isinstance(y, Number):
            inf = float("inf")
            if abs(x) == inf or abs(y) == inf:
                if allow_inf:
                    super(TestCase, self).assertEqual(x, y, message)
                else:
                    self.fail("Expected finite numeric values - x={}, y={}".format(x, y))
                return
            super(TestCase, self).assertLessEqual(abs(x - y), prec, message)
        else:
            super(TestCase, self).assertEqual(x, y, message)

    def checkModule(self, nn_module, args, unwrapper=None, skip=False):
        """
        Check that a nn.Module's results in TorchScript match eager and that it
        can be exported
        """
        if not TEST_WITH_SLOW or skip:
            # TorchScript is not enabled, skip these tests
            return

        sm = torch.jit.script(nn_module)

        with freeze_rng_state():
            eager_out = nn_module(*args)

        with freeze_rng_state():
            script_out = sm(*args)
            if unwrapper:
                script_out = unwrapper(script_out)

        self.assertEqual(eager_out, script_out)
        self.assertExportImportModule(sm, args)

        return sm

    def getExportImportCopy(self, m):
        """
        Save and load a TorchScript model
        """
        buffer = io.BytesIO()
        torch.jit.save(m, buffer)
        buffer.seek(0)
        imported = torch.jit.load(buffer)
        return imported

    def assertExportImportModule(self, m, args):
        """
        Check that the results of a model are the same after saving and loading
        """
        m_import = self.getExportImportCopy(m)
        with freeze_rng_state():
            results = m(*args)
        with freeze_rng_state():
            results_from_imported = m_import(*args)
        self.assertEqual(results, results_from_imported)

@contextlib.contextmanager
def freeze_rng_state():
    rng_state = torch.get_rng_state()
    if torch.cuda.is_available():
        cuda_rng_state = torch.cuda.get_rng_state()
    yield
    if torch.cuda.is_available():
        torch.cuda.set_rng_state(cuda_rng_state)
    torch.set_rng_state(rng_state)

class TransformsTester(unittest.TestCase):

    def _create_data(self, height=3, width=3, channels=3, device="cpu"):
        tensor = torch.randint(0, 255, (channels, height, width), dtype=torch.uint8, device=device)
        pil_img = Image.fromarray(tensor.permute(1, 2, 0).contiguous().cpu().numpy())
        return tensor, pil_img

    def compareTensorToPIL(self, tensor, pil_image, msg=None):
        np_pil_image = np.array(pil_image)
        if np_pil_image.ndim == 2:
            np_pil_image = np_pil_image[:, :, None]
        pil_tensor = torch.as_tensor(np_pil_image.transpose((2, 0, 1)))
        if msg is None:
            msg = "tensor:\n{} \ndid not equal PIL tensor:\n{}".format(tensor, pil_tensor)
        self.assertTrue(tensor.cpu().equal(pil_tensor), msg)

    def approxEqualTensorToPIL(self, tensor, pil_image, tol=1e-5, msg=None, agg_method="mean"):
        np_pil_image = np.array(pil_image)
        if np_pil_image.ndim == 2:
            np_pil_image = np_pil_image[:, :, None]
        pil_tensor = torch.as_tensor(np_pil_image.transpose((2, 0, 1))).to(tensor)
        err = getattr(torch, agg_method)(tensor - pil_tensor).item()
        self.assertTrue(
            err < tol,
            msg="{}: err={}, tol={}: \n{}\nvs\n{}".format(msg, err, tol, tensor[0, :10, :10], pil_tensor[0, :10, :10])
        )

max_stars_repo_name: uiuc-arc/vision

max_issues_repo_name: uiuc-arc/vision

max_forks_repo_name: uiuc-arc/vision