text:
# -*- coding: utf-8 -*-
"""
@author: Chenglong Chen <c.chenglong@gmail.com>
@brief: extreme ensemble selection module
@note:
        - support multiprocessing (set multiprocessing = True and multiprocessing_num_cores = #cores)
        - support random weight in greedy forward model selection (set weight_opt_max_evals = 1)
"""

"""
Implement (extreme) ensemble selection with the following techniques described in [1][2]:
1. Ensemble Initialization
2. Bagged Ensemble Selection
3. Ensemble Selection with Replacement
4. Cross-Validated Ensemble Selection
5. Random Weight in Greedy Forward Model Selection

Reference
[1] Ensemble selection from libraries of models, R. Caruana, A. Niculescu-Mizil, G. Crew, and A. Ksikes.
[2] Getting the Most Out of Ensemble Selection, Rich Caruana, Art Munson, and Alexandru Niculescu-Mizil

"""

import csv
import time
from optparse import OptionParser

import scipy
import numpy as np
import pandas as pd
from joblib import Parallel, delayed
from sklearn.utils.fixes import bincount
from sklearn.cross_validation import check_random_state
from sklearn.cross_validation import BaseShuffleSplit, ShuffleSplit, StratifiedShuffleSplit
# http://stackoverflow.com/questions/2801882/generating-a-png-with-matplotlib-when-display-is-undefined
# The solution for me was to add the following code in a place
# that gets read before any other pylab/matplotlib/pyplot import:
import matplotlib
# Force matplotlib to not use any Xwindows backend.
matplotlib.use('Agg')
import matplotlib.pyplot as plt

import config
from utils import dist_utils, np_utils
from utils import logging_utils, os_utils, pkl_utils, time_utils
from get_stacking_feature_conf import get_model_list

splitter_level1 = pkl_utils._load("%s/splits_level1.pkl"%config.SPLIT_DIR)
splitter_level2 = pkl_utils._load("%s/splits_level2.pkl"%config.SPLIT_DIR)
splitter_level3 = pkl_utils._load("%s/splits_level3.pkl"%config.SPLIT_DIR)
assert len(splitter_level1) == len(splitter_level2)
assert len(splitter_level1) == len(splitter_level3)
n_iter = len(splitter_level1)

class StratifiedShuffleSplitReplacement(BaseShuffleSplit):

    def __init__(self, y, n_iter=10, test_size=0.1, train_size=None,
                 random_state=None):

        super(StratifiedShuffleSplitReplacement, self).__init__(
            len(y), n_iter, test_size, train_size, random_state)

        self.y = np.array(y)
        self.classes, self.y_indices = np.unique(y, return_inverse=True)
        n_cls = self.classes.shape[0]

    def _iter_indices(self):
        rng = np.random.RandomState(self.random_state)
        cls_count = bincount(self.y_indices)

        for n in range(self.n_iter):
            train = []
            test = []

            for i, cls in enumerate(self.classes):
                sample_size = int(cls_count[i]*(1-self.test_size))
                randint = rng.randint(cls_count[i], size=sample_size)
                aidx = np.where((self.y == cls))[0]
                iidx = aidx[randint]
                oidx = aidx[list(set(range(cls_count[i])).difference(set(randint)))]

                train.extend(iidx)
                test.extend(oidx)

            train = rng.permutation(train)
            test = rng.permutation(test)

            yield train, test

    def __repr__(self):
        return ('%s(labels=%s, n_iter=%d, test_size=%s, '
                'random_state=%s)' % (
                    self.__class__.__name__,
                    self.y,
                    self.n_iter,
                    str(self.test_size),
                    self.random_state,
                ))

    def __len__(self):
        return self.n_iter

class ExtremeEnsembleSelection:
    def __init__(self, model_folder, model_list, subm_prefix,
                weight_opt_max_evals=10, w_min=-1., w_max=1.,
                inst_subsample=0.5, inst_subsample_replacement=False,
                inst_splitter=None,
                model_subsample=1.0, model_subsample_replacement=True,
                bagging_size=10, init_top_k=5, epsilon=0.00001,
                multiprocessing=False, multiprocessing_num_cores=1,
                enable_extreme=True, random_seed=0):

        self.model_folder = model_folder
        self.model_list = model_list
        self.subm_prefix = subm_prefix
        self.weight_opt_max_evals = weight_opt_max_evals
        self.w_min = w_min
        self.w_max = w_max
        assert inst_subsample > 0 and inst_subsample <= 1.
        self.inst_subsample = inst_subsample
        self.inst_subsample_replacement = inst_subsample_replacement
        self.inst_splitter = inst_splitter
        assert model_subsample > 0
        assert (type(model_subsample) == int) or (model_subsample <= 1.)
        self.model_subsample = model_subsample
        self.model_subsample_replacement = model_subsample_replacement
        self.bagging_size = bagging_size
        self.init_top_k = init_top_k
        self.epsilon = epsilon
        self.multiprocessing = multiprocessing
        self.multiprocessing_num_cores = multiprocessing_num_cores
        self.enable_extreme = enable_extreme
        self.random_seed = random_seed
        logname = "ensemble_selection_%s.log"%time_utils._timestamp()
        self.logger = logging_utils._get_logger(config.LOG_DIR, logname)
        self.n_models = len(self.model_list)

    def _merge_pred(self, w1, p1, w2, p2):
        p = (w1 * p1 + w2 * p2) / (w1 + w2)
        # relevance is in [1,3]
        p = np.clip(p, 1., 3.)
        return p

    def _pick_random_models(self, sorted_models, seed):
        num_model = len(sorted_models)
        #### bagging for models
        rng = np.random.RandomState(seed)
        if type(self.model_subsample) == int:
            sample_size = min(num_model, self.model_subsample)
        else:
            sample_size = int(num_model * self.model_subsample)
        if self.model_subsample_replacement:
            index_base = rng.randint(num_model, size=sample_size)
        else:
            index_base = rng.permutation(num_model)[:sample_size]
        this_sorted_models = [sorted_models[i] for i in sorted(index_base)]
        return this_sorted_models

    def _ens_obj_generic(self, weight2, p1_list, weight1, p2_list,
                        true_label_list, numBSTMatrix, bst_inst_idx):
        rmse_cv = np.zeros((config.N_RUNS, config.N_FOLDS), dtype=float)
        for run in range(config.N_RUNS):
            for fold in range(config.N_FOLDS):
                numBST = numBSTMatrix[run,fold]
                bidx = bst_inst_idx[run,fold,:numBST].tolist()
                p1 = p1_list[run,fold,bidx]
                p2 = p2_list[run,fold,bidx]
                true_label = true_label_list[run,fold,bidx]
                p_ens = self._merge_pred(weight1, p1, weight2, p2)
                rmse_cv[run,fold] = dist_utils._rmse(p_ens, true_label)
        rmse_mean = np.mean(rmse_cv)
        rmse_std = np.std(rmse_cv)
        return rmse_mean, rmse_std

    def _ens_obj_scipy(self, weight2, p1_list, weight1, p2_list,
                        true_label_list, numBSTMatrix, bst_inst_idx):
        rmse_mean, rmse_std = self._ens_obj_generic(weight2, p1_list, weight1, p2_list,
            true_label_list, numBSTMatrix, bst_inst_idx)
        return rmse_mean

    ## scipy.optimize.fmin for the best weight
    def _find_optim_weight_scipy(self, p_ens_list_valid_tmp, pred_list_valid, Y_list_valid,
                                    numBSTMatrix, bst_inst_idx, w_ens, model_index_dict, model):
        this_p_list_valid = pred_list_valid[model_index_dict[model]]
        if w_ens == 0:
            this_w = 1.
            best_rmse_mean, best_rmse_std = self._ens_obj_generic(this_w, p_ens_list_valid_tmp,
                0., this_p_list_valid, Y_list_valid, numBSTMatrix, bst_inst_idx)
        else:
            obj = lambda weight2: self._ens_obj_scipy(weight2, p_ens_list_valid_tmp,
                1., this_p_list_valid, Y_list_valid, numBSTMatrix, bst_inst_idx)
            weight2_init = np.random.uniform(self.w_min, self.w_max, size=1)[0]
            xopt = scipy.optimize.fmin(obj, weight2_init, maxiter=self.weight_opt_max_evals, disp=False)
            this_w = xopt[0]
            best_rmse_mean, best_rmse_std = self._ens_obj_generic(this_w, p_ens_list_valid_tmp,
                1., this_p_list_valid, Y_list_valid, numBSTMatrix, bst_inst_idx)
            this_w *= w_ens

        return best_rmse_mean, best_rmse_std, model, this_w

    def _ens_predict(self, best_bagged_model_list, best_bagged_model_weight):
        bagging_size = len(best_bagged_model_list)
        for bagging_iter in range(bagging_size):
            w_ens = 0.
            iter = 0
            for model,w in zip(best_bagged_model_list[bagging_iter], best_bagged_model_weight[bagging_iter]):
                iter += 1
                pred_file = "%s/All/test.pred.%s.csv" % (self.model_folder, model)
                df = pd.read_csv(pred_file, dtype=float)
                this_p_valid = df["prediction"].values
                this_p_valid = np.clip(this_p_valid, 1., 3.)
                if iter == 1:
                    p_ens_valid = np.zeros((this_p_valid.shape[0]),dtype=float)
                    id_test = df["id"].values.astype(int)
                p_ens_valid = self._merge_pred(w_ens, p_ens_valid, w, this_p_valid)
                w_ens += w

            if bagging_iter == 0:
                p_ens_score_bag = np.zeros((p_ens_valid.shape[0], bagging_size), dtype=float)

            p_ens_score_bag[:,bagging_iter] = p_ens_valid
        ##
        p_ens_score_bag = np.mean(p_ens_score_bag, axis=1)

        output = pd.DataFrame({"id": id_test, "relevance": p_ens_score_bag})
        return output

    def go(self):

        ## initialization
        pred_list_valid = np.zeros((self.n_models, config.N_RUNS, config.N_FOLDS, config.VALID_SIZE_MAX), dtype=float)
        Y_list_valid = np.zeros((config.N_RUNS, config.N_FOLDS, config.VALID_SIZE_MAX), dtype=float)
        numValidMatrix = np.zeros((config.N_RUNS, config.N_FOLDS), dtype=int)
        p_ens_list_valid = np.zeros((config.N_RUNS, config.N_FOLDS, config.VALID_SIZE_MAX), dtype=float)

        bst_inst_idx = np.zeros((config.N_RUNS, config.N_FOLDS, config.VALID_SIZE_MAX), dtype=float)
        numBSTMatrix = np.zeros((config.N_RUNS, config.N_FOLDS), dtype=int)
        oob_inst_idx = np.zeros((config.N_RUNS, config.N_FOLDS, config.VALID_SIZE_MAX), dtype=float)
        numOOBMatrix = np.zeros((config.N_RUNS, config.N_FOLDS), dtype=int)

        self.logger.info("Perform Extreme Ensemble Selection...")
        ## model index
        model_index_dict = dict(zip(self.model_list, range(self.n_models)))
        model_rmse_dict = dict(zip(self.model_list, [0]*self.n_models))
        self.logger.info("="*80)
        self.logger.info("Load model...")
        for model in self.model_list:
            self.logger.info("model: %s" % model)
            model_id = model_index_dict[model]
            rmse_cv = np.zeros((config.N_RUNS, config.N_FOLDS), dtype=float)
            ## load model
            for run in range(config.N_RUNS):
                for fold in range(config.N_FOLDS):
                    path = "%s/Run%d" % (self.model_folder, run+1)
                    pred_file = "%s/valid.pred.%s.csv" % (path, model)

                    this_p_valid = pd.read_csv(pred_file, dtype=float)
                    numValidMatrix[run,fold] = this_p_valid.shape[0]
                    numValid = numValidMatrix[run,fold]
                    this_target = this_p_valid["target"].values
                    this_p_valid = this_p_valid["prediction"].values
                    pred_list_valid[model_id,run,fold,:numValid] = np.clip(this_p_valid, 1., 3.)
                    Y_list_valid[run,fold,:numValid] = this_target

                    ##
                    rmse_cv[run,fold] = dist_utils._rmse(pred_list_valid[model_id,run,fold,:numValid],
                                                            Y_list_valid[run,fold,:numValid])

            self.logger.info("rmse: %.6f (%.6f)" % (np.mean(rmse_cv), np.std(rmse_cv)))
            model_rmse_dict[model] = (np.mean(rmse_cv), np.std(rmse_cv))
        self.logger.info("%d models in total." % self.n_models)

        sorted_models = sorted(model_rmse_dict.items(), key=lambda x: x[1][0])

        # greedy ensemble
        self.logger.info("="*80)
        best_bagged_model_list = [[]]*self.bagging_size
        best_bagged_model_weight = [[]]*self.bagging_size
        score_valid_bag_mean = np.nan * np.zeros((config.N_RUNS, config.N_FOLDS, config.VALID_SIZE_MAX, self.bagging_size), dtype=float)
        rmse_cv_mean_mean_lst = [0]*self.bagging_size
        rmse_cv_mean_std_lst = [0]*self.bagging_size
        for bagging_iter in range(self.bagging_size):
            seed_model = self.random_seed + 100 * bagging_iter
            if not self.enable_extreme:
                this_sorted_models = self._pick_random_models(sorted_models, seed_model)
            #### instance level subsampling
            for run in range(config.N_RUNS):
                for fold in range(config.N_FOLDS):
                    if self.inst_splitter is None:
                        # GENERAL APPROACH
                        seed_inst = self.random_seed + 1000 * bagging_iter + 100 * run + 10 * fold
                        rng_inst = np.random.RandomState(seed_inst)
                        numValid = numValidMatrix[run,fold]
                        if self.inst_subsample_replacement:
                            sss = StratifiedShuffleSplitReplacement(Y_list_valid[run,fold,:numValid], n_iter=1,
                                test_size=1.-self.inst_subsample, random_state=seed_inst)
                            iidx, oidx = list(sss)[0]
                        else:
                            if self.inst_subsample < 1:
                                # Stratified ShuffleSplit
                                sss = ShuffleSplit(len(Y_list_valid[run,fold,:numValid]), n_iter=1,
                                    test_size=1.-self.inst_subsample, random_state=seed_inst)
                                iidx, oidx = list(sss)[0]
                            elif self.inst_subsample == 1:
                                # set iidx (trianing) the same as oidx (validation)
                                iidx = np.arange(numValid)
                                oidx = np.arange(numValid)
                    else:
                        iidx, oidx = self.inst_splitter[run]
                    numBSTMatrix[run,fold] = len(iidx)
                    bst_inst_idx[run,fold,:numBSTMatrix[run,fold]] = iidx
                    numOOBMatrix[run,fold] = len(oidx)
                    oob_inst_idx[run,fold,:numOOBMatrix[run,fold]] = oidx

            #print this_model_list
            best_model_list = []
            best_model_weight = []
            best_model_rmse = []
            best_rmse = 0
            best_rmse_std = 0
            best_model = None
            p_ens_list_valid_tmp = np.zeros((config.N_RUNS, config.N_FOLDS, config.VALID_SIZE_MAX), dtype=float)
            #### Technique: Ensemble Initialization
            iter = 0
            w_ens, this_w = 0.0, 1.0
            if self.init_top_k > 0:
                # self.logger.info("** Ensemble Initialization **")
                # init_top_k = min(init_top_k, num_model)
                rmse_cv = np.zeros((config.N_RUNS, config.N_FOLDS), dtype=float)
                for cnt in range(self.init_top_k):
                    iter += 1
                    start = time.time()
                    seed_model = self.random_seed + 100 * bagging_iter + 10 * iter
                    if self.enable_extreme:
                        this_sorted_models = self._pick_random_models(sorted_models, seed_model)
                        best_model,(rmse,rmse_std) = this_sorted_models[0]
                    else:
                        best_model,(rmse,rmse_std) = this_sorted_models[cnt]
                    this_p_list_valid = pred_list_valid[model_index_dict[best_model]]
                    for run in range(config.N_RUNS):
                        for fold in range(config.N_FOLDS):
                            numValid = numValidMatrix[run,fold]
                            numBST = numBSTMatrix[run,fold]
                            bidx = bst_inst_idx[run,fold,:numBST].tolist()
                            p_ens_list_valid_tmp[run,fold,:numValid] = self._merge_pred(
                                w_ens, p_ens_list_valid_tmp[run,fold,:numValid],
                                this_w, this_p_list_valid[run,fold,:numValid])
                            true_label = Y_list_valid[run,fold,bidx]
                            rmse_cv[run,fold] = dist_utils._rmse(p_ens_list_valid_tmp[run,fold,bidx], true_label)
                    end = time.time()
                    best_weight = this_w
                    best_rmse = np.mean(rmse_cv)
                    best_rmse_std = np.std(rmse_cv)

                    self.logger.info("Iter: %d (%.2fs)" % (iter, (end - start)))
                    self.logger.info("     model: %s" % best_model)
                    self.logger.info("     weight: %s" % best_weight)
                    self.logger.info("     rmse: %.6f (%.6f)" % (best_rmse, best_rmse_std))

                    best_model_list.append(best_model)
                    best_model_weight.append(best_weight)
                    w_ens += best_weight

            #### Technique: Ensemble Selection with Replacement
            while True:
                iter += 1
                seed_model = self.random_seed + 100 * bagging_iter + 10 * iter
                if self.enable_extreme:
                    this_sorted_models = self._pick_random_models(sorted_models, seed_model)
                if self.multiprocessing:
                    start = time.time()
                    models_tmp = [model for model,(_,_) in this_sorted_models]
                    best_trial_rmse_mean_lst, best_trial_rmse_std_lst, model_lst, this_w_lst = \
                        zip(*Parallel(n_jobs=self.multiprocessing_num_cores)(
                            delayed(self._find_optim_weight_scipy)(
                                p_ens_list_valid_tmp, pred_list_valid, Y_list_valid, numBSTMatrix,
                                bst_inst_idx, w_ens, model_index_dict, m
                                ) for m in models_tmp
                            ))
                    ##
                    ind_best = np.argmin(best_trial_rmse_mean_lst)
                    best_trial_rmse_mean = best_trial_rmse_mean_lst[ind_best]
                    best_trial_rmse_std = best_trial_rmse_std_lst[ind_best]
                    model = model_lst[ind_best]
                    this_w = this_w_lst[ind_best]
                    if best_trial_rmse_mean < best_rmse:
                        best_rmse, best_rmse_std = best_trial_rmse_mean, best_trial_rmse_std
                        best_model, best_weight = model, this_w
                    end = time.time()
                else:
                    start = time.time()
                    for model,(_,_) in this_sorted_models:
                        best_trial_rmse_mean, best_trial_rmse_std, model, this_w = \
                            self._find_optim_weight_scipy(
                            p_ens_list_valid_tmp, pred_list_valid, Y_list_valid, numBSTMatrix,
                            bst_inst_idx, w_ens, model_index_dict, model)
                        if best_trial_rmse_mean < best_rmse:
                            best_rmse, best_rmse_std = best_trial_rmse_mean, best_trial_rmse_std
                            best_model, best_weight = model, this_w
                    end = time.time()
                if best_model is None:
                    break
                if len(best_model_rmse) > 1 and (best_model_rmse[-1] - best_rmse < self.epsilon):
                    break

                ##
                self.logger.info("Iter: %d (%.2fs)" % (iter, (end - start)))
                self.logger.info("     model: %s" % best_model)
                self.logger.info("     weight: %s" % best_weight)
                self.logger.info("     rmse: %.6f (%.6f)" % (best_rmse, best_rmse_std))

                # valid
                this_p_list_valid = pred_list_valid[model_index_dict[best_model]]
                pred_raw_list = []
                true_label_list = []
                for run in range(config.N_RUNS):
                    for fold in range(config.N_FOLDS):
                        numValid = numValidMatrix[run,fold]
                        numBST = numBSTMatrix[run,fold]
                        bidx = bst_inst_idx[run,fold,:numBST].tolist()
                        p_ens_list_valid_tmp[run,fold,:numValid] = self._merge_pred(
                            w_ens, p_ens_list_valid_tmp[run,fold,:numValid],
                            best_weight, this_p_list_valid[run,fold,:numValid])

                        pred_raw_list.append( p_ens_list_valid_tmp[run,fold,bidx] )
                        true_label_list.append( Y_list_valid[run,fold,bidx] )

                best_model_list.append(best_model)
                best_model_weight.append(best_weight)
                best_model_rmse.append(best_rmse)

                best_model = None
                w_ens += best_weight

            ## compute OOB score
            rmse_cv_mean = np.zeros((config.N_RUNS, config.N_FOLDS), dtype=float)
            for run in range(config.N_RUNS):
                for fold in range(config.N_FOLDS):
                    numValid = numValidMatrix[run,fold]
                    true_label = Y_list_valid[run,fold,:numValid]
                    numOOB = numOOBMatrix[run,fold]
                    oidx = oob_inst_idx[run,fold,:numOOB].tolist()
                    pred_raw = p_ens_list_valid_tmp[run,fold,oidx]
                    ## mean
                    score_valid_bag_mean[run,fold,oidx,bagging_iter] = pred_raw
                    pred_mean = np_utils._array_mean(score_valid_bag_mean[run,fold,:numValid,:(bagging_iter+1)])
                    non_nan_idx = pred_mean != config.MISSING_VALUE_NUMERIC
                    rmse_cv_mean[run,fold] = dist_utils._rmse(pred_mean[non_nan_idx], true_label[non_nan_idx])
            self.logger.info("-"*80)
            self.logger.info( "Bag: %d"% (bagging_iter+1))
            self.logger.info( "rmse-mean: %.6f (%.6f)" % (np.mean(rmse_cv_mean), np.std(rmse_cv_mean)))
            self.logger.info("-"*80)

            best_bagged_model_list[bagging_iter] = best_model_list
            best_bagged_model_weight[bagging_iter] = best_model_weight

            ## save the current prediction
            mr = "R" + str(self.model_subsample_replacement).upper()[0]
            ir = "R" + str(self.inst_subsample_replacement).upper()[0]
            ## mean
            best_rmse_mean = np.mean(rmse_cv_mean)
            best_rmse_std = np.std(rmse_cv_mean)
            output = self._ens_predict(best_bagged_model_list[:(bagging_iter+1)],
                best_bagged_model_weight[:(bagging_iter+1)])
            sub_file = "%s_[MS%.2f_%s]_[IS%.2f_%s]_[Top%d]_[Bag%d]_[Mean%.6f]_[Std%.6f].mean.csv" % (
                self.subm_prefix, self.model_subsample, mr, self.inst_subsample, ir,
                self.init_top_k, bagging_iter+1, best_rmse_mean, best_rmse_std)
            output.to_csv(sub_file, index=False)
            rmse_cv_mean_mean_lst[bagging_iter] = best_rmse_mean
            rmse_cv_mean_std_lst[bagging_iter] = best_rmse_std

            ## plot OOB score
            x = np.arange(1,bagging_iter+2,1)
            label = "Mean (Best = %.6f, Bag = %d)"%(
                    np.min(rmse_cv_mean_mean_lst[:(bagging_iter+1)]),
                    np.argmin(rmse_cv_mean_mean_lst[:(bagging_iter+1)])+1)
            plt.errorbar(x, rmse_cv_mean_mean_lst[:(bagging_iter+1)],
                yerr=rmse_cv_mean_std_lst[:(bagging_iter+1)],
                fmt='-o', label=label)
            plt.xlim(1, self.bagging_size)
            plt.title("Extreme Ensemble Selection RMSE")
            plt.xlabel("Bag")
            plt.ylabel("CV/OOB RMSE")
            plt.legend(loc="upper right")
            fig_file = "%s/ensemble_selection_%d.pdf"%(config.FIG_DIR, bagging_iter+1)
            plt.savefig(fig_file)
            plt.clf()

#--------------------------- Main ---------------------------
def main(options):

    # create sub folder
    subm_folder = "%s/ensemble_selection"%config.SUBM_DIR
    os_utils._create_dirs( [subm_folder] )
    subm_prefix = "%s/test.pred.[%s]" % (subm_folder, options.outfile)

    # get model list
    log_folder = "%s/level%d_models"%(config.LOG_DIR, options.level-1)
    model_list = get_model_list(log_folder, options.size)

    # get instance splitter
    if options.level not in [2, 3]:
        inst_splitter = None
    elif options.level == 2:
        inst_splitter = splitter_level2
    elif options.level == 3:
        inst_splitter = splitter_level3

    ees = ExtremeEnsembleSelection(
            model_folder=config.OUTPUT_DIR,
            model_list=model_list,
            subm_prefix=subm_prefix,
            weight_opt_max_evals=options.weight_opt_max_evals,
            w_min=-1.,
            w_max=1.,
            inst_subsample=options.inst_subsample,
            inst_subsample_replacement=options.inst_subsample_replacement,
            inst_splitter=inst_splitter,
            model_subsample=options.model_subsample,
            model_subsample_replacement=options.model_subsample_replacement,
            bagging_size=options.bagging_size,
            init_top_k=options.init_top_k,
            epsilon=options.epsilon,
            multiprocessing=False,
            multiprocessing_num_cores=config.NUM_CORES,
            enable_extreme=options.enable_extreme,
            random_seed=config.RANDOM_SEED
        )
    ees.go()

def parse_args(parser):
    parser.add_option("-l", "--level", type="int", dest="level",
        default=2, help="level of base models")
    parser.add_option("-s", "--size", type="int", dest="size",
        default=10, help="size of each model")
    parser.add_option("-o", "--outfile", type="string", dest="outfile",
        default="extreme_ensemble_selection", help="output model name")
    parser.add_option("-b", "--bag", type="int", dest="bagging_size",
        default=10, help="bagging_size")
    parser.add_option("-t", "--init", type="int", dest="init_top_k",
        default=5, help="init_top_k")
    parser.add_option("-i", "--inst_subsample", type="float", dest="inst_subsample",
        default=0.5, help="inst_subsample")
    parser.add_option("-I", "--inst_subsample_replacement", action="store_true",
        dest="inst_subsample_replacement",
        default=False, help="inst_subsample_replacement")
    parser.add_option("-m", "--model_subsample", type="float", dest="model_subsample",
        default=1.0, help="model_subsample")
    parser.add_option("-M", "--model_subsample_replacement", action="store_true",
        dest="model_subsample_replacement",
        default=True, help="model_subsample_replacement")
    parser.add_option("-e", "--epsilon", type="float", dest="epsilon",
        default=0.0001, help="epsilon")
    parser.add_option("-w", "--weight_opt_max_evals", type="int", dest="weight_opt_max_evals",
        default=5, help="weight_opt_max_evals")
    parser.add_option("-x", action="store_true", dest="enable_extreme",
        default=True, help="enable_extreme")
    (options, args) = parser.parse_args()
    return options, args

if __name__ == "__main__":

    parser = OptionParser()
    options, args = parse_args(parser)
    main(options)

reason: import numpy,import scipy


text:
import unittest

import numpy as np

from unittest.mock import Mock

from bioslds.hyperopt import random_maximize

class TestRandomMaximizeObjectiveFunctionCalls(unittest.TestCase):
    def setUp(self):
        self.fct = Mock(return_value=0)
        self.n_iter = 40
        self.param_ranges = {"foo": (1, 3), "bar": (-1.0, 5.3)}
        self.param_types = {"foo": np.integer, "bar": np.inexact}
        random_maximize(self.fct, self.param_ranges, self.n_iter)

    def test_fct_called_n_iter_times(self):
        self.assertEqual(self.fct.call_count, self.n_iter)

    def test_fct_called_with_no_positional_args(self):
        for call in self.fct.call_args_list:
            self.assertEqual(len(call[0]), 0)

    def test_fct_called_with_all_necessary_kwargs(self):
        for call in self.fct.call_args_list:
            for key in self.param_ranges:
                self.assertIn(key, call[1])

    def test_fct_not_called_with_extraneous_kwargs(self):
        for call in self.fct.call_args_list:
            for key in call[1]:
                self.assertIn(key, self.param_ranges)

    def test_generated_param_values_are_in_correct_range(self):
        for call in self.fct.call_args_list:
            for key, (lo, hi) in self.param_ranges.items():
                value = call[1][key]
                self.assertGreaterEqual(value, lo)
                self.assertLess(value, hi)

    def test_generated_param_values_have_correct_type(self):
        for call in self.fct.call_args_list:
            for key, value_type in self.param_types.items():
                value = np.asarray([call[1][key]])
                self.assertTrue(np.issubdtype(value.dtype, value_type))

class TestRandomMaximizeRngCalls(unittest.TestCase):
    def setUp(self):
        self.n_iter = 40
        self.param_ranges = {"foo": (1, 3), "bar": (-1.0, 5.3)}
        self.param_types = {"foo": np.integer, "bar": np.inexact}

    def test_rng_seed_uses_default_rng(self):
        seed = 10

        fct1 = Mock(return_value=0)
        random_maximize(fct1, self.param_ranges, self.n_iter, rng=seed)

        fct2 = Mock(return_value=0)
        random_maximize(
            fct2, self.param_ranges, self.n_iter, rng=np.random.default_rng(seed)
        )

        for call1, call2 in zip(fct1.call_args_list, fct2.call_args_list):
            for key in self.param_ranges:
                value1 = call1[1][key]
                value2 = call2[1][key]

                self.assertAlmostEqual(value1, value2)

    def test_rng_as_numpy_random_state_works(self):
        fct = Mock(return_value=0)
        random_maximize(
            fct, self.param_ranges, self.n_iter, rng=np.random.RandomState(0)
        )

        fct.assert_called()

class TestRandomMaximizeProgress(unittest.TestCase):
    def test_progress_called(self):
        mock_progress = Mock(side_effect=lambda it: it)
        fct = Mock(return_value=0)

        random_maximize(fct, {"foo": (-1.2, 3.4)}, 52, progress=mock_progress)

        mock_progress.assert_called()

class TestRandomMaximizeReturnValues(unittest.TestCase):
    def setUp(self):
        self.rng = np.random.default_rng(1)
        self.n_iter = 50
        self.fct_values = self.rng.normal(size=self.n_iter)

        self.fct = Mock()
        self.fct.side_effect = [
            (fct_value, {"a": 2 * fct_value, "b": 3.0}) for fct_value in self.fct_values
        ]

        self.param_ranges = {"foo": (0, 3), "bar": (-1.3, 5.2)}
        self.out = random_maximize(self.fct, self.param_ranges, self.n_iter)

    def test_first_output_is_max(self):
        expected_max = np.max(self.fct_values)
        self.assertAlmostEqual(self.out[0], expected_max)

    def test_second_output_is_argmax(self):
        # noinspection PyTypeChecker
        i: int = np.argmax(self.fct_values)
        call_i = self.fct.call_args_list[i]

        for key in self.param_ranges:
            value = self.out[1][key]
            value_exp = call_i[1][key]

            self.assertAlmostEqual(value, value_exp)

    def test_third_output_params_contains_used_parameter_sets(self):
        for i, call in enumerate(self.fct.call_args_list):
            param_iter = self.out[2]["params"][i]
            for key in self.param_ranges:
                value = param_iter[key]
                value_exp = call[1][key]

                self.assertAlmostEqual(value, value_exp)

    def test_third_output_fct_scalar_contains_scalar_fct_values(self):
        np.testing.assert_allclose(self.out[2]["fct_scalar"], self.fct_values)

    def test_third_output_fct_full_contains_full_fct_outputs(self):
        for i, fct_exp in enumerate(self.fct.side_effect):
            fct_full = self.out[2]["fct_full"][i]
            self.assertEqual(fct_full, fct_exp)

class TestRandomMaximizeWithScalarFunction(unittest.TestCase):
    def setUp(self):
        self.rng = np.random.default_rng(1)
        self.n_iter = 50
        self.fct_values = self.rng.normal(size=self.n_iter)

        self.fct = Mock()
        self.fct.side_effect = self.fct_values

        self.param_ranges = {"foo": (0, 3), "bar": (-1.3, 5.2)}
        self.out = random_maximize(self.fct, self.param_ranges, self.n_iter)

    def test_first_output_is_max(self):
        expected_max = np.max(self.fct_values)
        self.assertAlmostEqual(self.out[0], expected_max)

    def test_second_output_is_argmax(self):
        # noinspection PyTypeChecker
        i: int = np.argmax(self.fct_values)
        call_i = self.fct.call_args_list[i]

        for key in self.param_ranges:
            value = self.out[1][key]
            value_exp = call_i[1][key]

            self.assertAlmostEqual(value, value_exp)

    def test_third_output_fct_scalar_contains_scalar_fct_values(self):
        np.testing.assert_allclose(self.out[2]["fct_scalar"], self.fct_values)

    def test_third_output_fct_full_contains_full_fct_outputs(self):
        for i, fct_exp in enumerate(self.fct.side_effect):
            fct_full = self.out[2]["fct_full"][i]
            self.assertEqual(fct_full, fct_exp)

class TestRandomMaximize(unittest.TestCase):
    def test_log_scale_rounds_integer_vars(self):
        _, _, details = random_maximize(
            lambda **kwargs: 1.0, {"i_log": (1, 5)}, 20, log_scale=["i_log"]
        )
        not_int = [
            np.abs(_["i_log"] - np.round(_["i_log"])) > 1e-10 for _ in details["params"]
        ]
        self.assertEqual(np.sum(not_int), 0)

    def test_log_scale_samples_differently_than_default_for_real_var(self):
        seed = 10
        _, _, details1 = random_maximize(
            lambda **kwargs: 1.0, {"i_log": (1.5, 3.5)}, 20, rng=seed
        )
        _, _, details2 = random_maximize(
            lambda **kwargs: 1.0,
            {"i_log": (1.5, 3.5)},
            20,
            rng=seed,
            log_scale=["i_log"],
        )

        i1 = [_["i_log"] for _ in details1["params"]]
        i2 = [_["i_log"] for _ in details2["params"]]

        diff = np.max(np.abs(np.sort(i1) - np.sort(i2)))
        self.assertGreater(diff, 1e-3)

    def test_log_scale_samples_differently_than_default_for_integer_var(self):
        seed = 10
        _, _, details1 = random_maximize(
            lambda **kwargs: 1.0, {"i_log": (1, 20)}, 100, rng=seed
        )
        _, _, details2 = random_maximize(
            lambda **kwargs: 1.0,
            {"i_log": (1, 20)},
            100,
            rng=seed,
            log_scale=["i_log"],
        )

        i1 = [_["i_log"] for _ in details1["params"]]
        i2 = [_["i_log"] for _ in details2["params"]]

        diff = np.max(np.abs(np.sort(i1) - np.sort(i2)))
        self.assertGreater(diff, 1e-3)

if __name__ == "__main__":
    unittest.main()

ext: py

max_stars_repo_name: ttesileanu/bio-time-series

max_issues_repo_name: ttesileanu/bio-time-series

max_forks_repo_name: ttesileanu/bio-time-series

reason: import numpy


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

class Bandit:
    def __init__(self, k=10, exp_rate=.3, lr=0.1, ucb=False, seed=None, c=2):
        self.k = k
        self.actions = range(self.k)
        self.exp_rate = exp_rate
        self.lr = lr
        self.total_reward = 0
        self.avg_reward = []

        self.TrueValue = []
        np.random.seed(seed)
        for i in range(self.k):
            self.TrueValue.append(np.random.randn() + 2)  # normal distribution

        self.values = np.zeros(self.k)
        self.times = 0
        self.action_times = np.zeros(self.k)

        self.ucb = ucb  # if select action using upper-confidence-bound
        self.c = c

    def chooseAction(self):
        # explore
        if np.random.uniform(0, 1) <= self.exp_rate:
            action = np.random.choice(self.actions)
        else:
            # exploit
            if self.ucb:
                if self.times == 0:
                    action = np.random.choice(self.actions)
                else:
                    confidence_bound = self.values + self.c * np.sqrt(
                        np.log(self.times) / (self.action_times + 0.1))  # c=2
                    action = np.argmax(confidence_bound)
            else:
                action = np.argmax(self.values)
        return action

    def takeAction(self, action):
        self.times += 1
        self.action_times[action] += 1
        # take action and update value estimates
        # reward = self.TrueValue[action]
        reward = np.random.randn() + self.TrueValue[action]  # add randomness to reward
        # using incremental method to propagate
        self.values[action] += self.lr * (reward - self.values[action])  # look like fixed lr converges better

        self.total_reward += reward
        self.avg_reward.append(self.total_reward / self.times)

    def play(self, n):
        for _ in range(n):
            action = self.chooseAction()
            self.takeAction(action)

if __name__ == "__main__":
    bdt = Bandit(k=5, exp_rate=0.1, seed=1234)
    bdt.play(2000)

    print("Estimated values", bdt.values)
    print("Actual values", bdt.TrueValue)
    avg_reward1 = bdt.avg_reward

    bdt = Bandit(k=5, exp_rate=0.3, seed=1234)
    bdt.play(2000)

    print("Estimated values", bdt.values)
    print("Actual values", bdt.TrueValue)

    avg_reward2 = bdt.avg_reward

    bdt = Bandit(k=5, exp_rate=0.1, seed=1234, ucb=True, c=2)
    bdt.play(2000)

    print("Estimated values", bdt.values)
    print("Actual values", bdt.TrueValue)

    avg_reward3 = bdt.avg_reward

    bdt = Bandit(k=5, exp_rate=0.1, seed=1234, ucb=True, c=5)
    bdt.play(2000)

    print("Estimated values", bdt.values)
    print("Actual values", bdt.TrueValue)

    avg_reward4 = bdt.avg_reward

    # reward plot
    plt.figure(figsize=[8, 6])
    plt.plot(avg_reward1, label="exp_rate=0.1")
    plt.plot(avg_reward2, label="exp_rate=0.3")
    plt.plot(avg_reward3, label="ucb, c=2")
    plt.plot(avg_reward4, label="ucb, c=5")

    plt.xlabel("n_iter", fontsize=14)
    plt.ylabel("avg reward", fontsize=14)
    plt.legend()


text:
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import itertools

#Import Data
mbhmms_card = pd.read_csv("../s01-card/MBhmms/201006-MBhmms_pr_analysis.txt",sep='\t',skiprows=2,usecols=[0,4,5],names=["Family","Precision", "Recall"])
mbhmms_card["Database"] = "Resfams 2.0"
oldres_card = pd.read_csv("../s01-card/OldRes/201006-res_pr_analysis.txt",sep='\t',skiprows=2,usecols=[0,4,5],names=["Family","Precision", "Recall"])
oldres_card["Database"] = "Resfams 1.0"
mbhmms_ncbi = pd.read_csv("../s02-ncbi/MBHmms/201007-MBHmms_pr_analysis.txt",sep='\t',skiprows=2,usecols=[0,4,5],names=["Family","Precision", "Recall"])
mbhmms_ncbi["Database"] = "Resfams 2.0"
oldres_ncbi = pd.read_csv("../s02-ncbi/OldRes/201007-OldRes_pr_analysis.txt",sep='\t',skiprows=2,usecols=[0,4,5],names=["Family","Precision", "Recall"])
oldres_ncbi["Database"] = "Resfams 1.0"

# #Parse Data
ncbi_frames = [mbhmms_ncbi,oldres_ncbi]
ncbi_data = pd.concat(ncbi_frames,ignore_index=True, sort=False)
ncbi_temp = pd.merge(mbhmms_ncbi,oldres_ncbi,on='Family')
mbhmms_ncbi = ncbi_temp[["Family", "Precision_x", "Recall_x","Database_x"]]
mbhmms_ncbi.columns = ["Family", "Precision", "Recall","Database"]
oldres_ncbi = ncbi_temp[["Family", "Precision_y", "Recall_y","Database_y"]]
oldres_ncbi.columns = ["Family", "Precision", "Recall","Database"]

ncbi_data = pd.concat(ncbi_frames,ignore_index=True, sort=False).fillna(value=0.0).round(4)
ncbi_data['Precision'] = ncbi_data[['Precision']] * 100
ncbi_data['Recall'] = ncbi_data[['Recall']] * 100
ncbi_pData = ncbi_data[['Family','Database','Precision']]
ncbi_pData["Analysis"] = "Precision"
ncbi_pData.columns = ['Family','Database','Value','Analysis']
ncbi_rData = ncbi_data[['Family','Database','Recall']]
ncbi_rData["Analysis"] = "Precision"
ncbi_rData.columns = ['Family','Database','Value','Analysis']

## Make Graphs
#Precision
f, ax = plt.subplots(figsize=(18,25))

ncbi_gpData = ncbi_pData.pivot('Family','Database','Value')
pHeatMap = sns.heatmap(ncbi_gpData,cmap="Spectral",ax=ax,linewidths=0.01,linecolor='black',annot=True,fmt='g')
pHeatMap.set_title('Database Precision Comparative Analysis')
pHeatMap.get_figure().savefig("../plots/201010-ncbi_precision.png",bbox_inches='tight')

#Recall
f, ax = plt.subplots(figsize=(18,25))

ncbi_grData = ncbi_rData.pivot('Family','Database','Value')
rHeatMap = sns.heatmap(ncbi_grData,cmap="Spectral",ax=ax,linewidths=0.01,linecolor='black',annot=True,fmt='g')
rHeatMap.set_title('Database Recall Comparative Analysis')
rHeatMap.get_figure().savefig("../plots/201010-ncbi_recall.png",bbox_inches='tight')


text:
from pynput.keyboard import Key
import numpy as np

EXIT          = -1
NOOP          = 0
FIRE          = 1
UP            = 2
RIGHT         = 3
LEFT          = 4
DOWN          = 5
UPRIGHT       = 6
UPLEFT        = 7
DOWNRIGHT     = 8
DOWNLEFT      = 9
UPFIRE        = 10
RIGHTFIRE     = 11
LEFTFIRE      = 12
DOWNFIRE      = 13
UPRIGHTFIRE   = 14
UPLEFTFIRE    = 15
DOWNRIGHTFIRE = 16
DOWNLEFTFIRE  = 17

class Joypad():
    def __init__(self, env):
        self.action2num = {}
        self.action_meanings = env.unwrapped.get_action_meanings()
        self.n_actions = env.action_space.n
        for idx, action in enumerate(self.action_meanings):
            self.action2num[action] = idx
        self.num2action = {v: k for k, v in self.action2num.items()}
        self.action = 0
        self.env_action = 0
        self.lastKey = None
        self.autoplay = False

    def actionString(self):
        return self.num2action[self.env_action]

    def actionOneHot(self):
        ret = np.zeros(self.n_actions, dtype=np.float64)
        np.put(ret, self.env_action, 1)
        return ret

    def actionButtons(self):
        bFIRE  = 0
        bUP    = 1
        bRIGHT = 2
        bLEFT  = 3
        bDOWN  = 4
        ret = np.zeros(5, dtype=np.float64)
        if(self.action == FIRE):
            np.put(ret, bFIRE, 1)
        elif self.action == UP:
            np.put(ret, bUP, 1)
        elif self.action == RIGHT:
            np.put(ret, bRIGHT, 1)
        elif self.action == LEFT:
            np.put(ret, bLEFT, 1)
        elif self.action == DOWN:
            np.put(ret, bDOWN, 1)
        elif self.action == UPRIGHT:
            np.put(ret, bUP, 1)
            np.put(ret, bRIGHT, 1)
        elif self.action == UPLEFT:
            np.put(ret, bUP, 1)
            np.put(ret, bLEFT, 1)
        elif self.action == DOWNRIGHT:
            np.put(ret, bDOWN, 1)
            np.put(ret, bRIGHT, 1)
        elif self.action == DOWNLEFT:
            np.put(ret, bDOWN, 1)
            np.put(ret, bLEFT, 1)
        elif self.action == UPFIRE:
            np.put(ret, bUP, 1)
            np.put(ret, bFIRE, 1)
        elif self.action == RIGHTFIRE:
            np.put(ret, bRIGHT, 1)
            np.put(ret, bFIRE, 1)
        elif self.action == LEFTFIRE:
            np.put(ret, bLEFT, 1)
            np.put(ret, bFIRE, 1)
        elif self.action == DOWNFIRE:
            np.put(ret, bDOWN, 1)
            np.put(ret, bFIRE, 1)
        elif self.action == UPRIGHTFIRE:
            np.put(ret, bUP, 1)
            np.put(ret, bRIGHT, 1)
            np.put(ret, bFIRE, 1)
        elif self.action == UPLEFTFIRE:
            np.put(ret, bUP, 1)
            np.put(ret, bLEFT, 1)
            np.put(ret, bFIRE, 1)
        elif self.action == DOWNRIGHTFIRE:
            np.put(ret, bDOWN, 1)
            np.put(ret, bRIGHT, 1)
            np.put(ret, bFIRE, 1)
        elif self.action == DOWNLEFTFIRE:
            np.put(ret, bDOWN, 1)
            np.put(ret, bLEFT, 1)
            np.put(ret, bFIRE, 1)

        return ret

    def on_press(self, key):
        self.lastKey = key
        if(key == Key.space):
            if  (self.action == NOOP):
                self.action = FIRE
            elif(self.action == UP):
                self.action = UPFIRE
            elif(self.action == RIGHT):
                self.action = RIGHTFIRE
            elif(self.action == LEFT):
                self.action = LEFTFIRE
            elif(self.action == DOWN):
                self.action = DOWNFIRE
            elif(self.action == UPRIGHT):
                self.action = UPRIGHTFIRE
            elif(self.action == UPLEFT):
                self.action = UPLEFTFIRE
            elif(self.action == DOWNRIGHT):
                self.action = DOWNRIGHTFIRE
            elif(self.action == DOWNLEFT):
                self.action = DOWNLEFTFIRE
        elif(key == Key.up):
            if  (self.action == NOOP):
                self.action = UP
            elif(self.action == FIRE):
                self.action = UPFIRE
            elif(self.action == RIGHT):
                self.action = UPRIGHT
            elif(self.action == LEFT):
                self.action = UPLEFT
            elif(self.action == RIGHTFIRE):
                self.action = UPRIGHTFIRE
            elif(self.action == LEFTFIRE):
                self.action = UPLEFTFIRE
        elif(key == Key.right):
            if  (self.action == NOOP):
                self.action = RIGHT
            elif(self.action == FIRE):
                self.action = RIGHTFIRE
            elif(self.action == UP):
                self.action = UPRIGHT
            elif(self.action == DOWN):
                self.action = DOWNRIGHT
            elif(self.action == UPFIRE):
                self.action = UPRIGHTFIRE
            elif(self.action == DOWNFIRE):
                self.action = DOWNRIGHTFIRE
        elif(key == Key.left):
            if  (self.action == NOOP):
                self.action = LEFT
            elif(self.action == FIRE):
                self.action = LEFTFIRE
            elif(self.action == UP):
                self.action = UPLEFT
            elif(self.action == DOWN):
                self.action = DOWNLEFT
            elif(self.action == UPFIRE):
                self.action = UPLEFTFIRE
            elif(self.action == DOWNFIRE):
                self.action = DOWNLEFTFIRE
        elif(key == Key.down):
            if  (self.action == NOOP):
                self.action = DOWN
            elif(self.action == FIRE):
                self.action = DOWNFIRE
            elif(self.action == RIGHT):
                self.action = DOWNRIGHT
            elif(self.action == LEFT):
                self.action = DOWNLEFT
            elif(self.action == RIGHTFIRE):
                self.action = DOWNRIGHTFIRE
            elif(self.action == LEFTFIRE):
                self.action = DOWNLEFTFIRE
        elif(key == Key.tab):
            self.autoplay = not self.autoplay
        elif(key == Key.esc):
            self.action = EXIT

    def on_release(self, key):
        if(key == Key.space):
            if  (self.action == FIRE):
                self.action = NOOP
            elif(self.action == UPFIRE):
                self.action = UP
            elif(self.action == RIGHTFIRE):
                self.action = RIGHT
            elif(self.action == LEFTFIRE):
                self.action = LEFT
            elif(self.action == DOWNFIRE):
                self.action = DOWN
            elif(self.action == UPRIGHTFIRE):
                self.action = UPRIGHT
            elif(self.action == UPLEFTFIRE):
                self.action = UPLEFT
            elif(self.action == DOWNRIGHTFIRE):
                self.action = DOWNRIGHT
            elif(self.action == DOWNLEFTFIRE):
                self.action = DOWNLEFT
        elif(key == Key.up):
            if  (self.action == UP):
                self.action = NOOP
            elif(self.action == UPFIRE):
                self.action = FIRE
            elif(self.action == UPRIGHT):
                self.action = RIGHT
            elif(self.action == UPLEFT):
                self.action = LEFT
            elif(self.action == UPRIGHTFIRE):
                self.action = RIGHTFIRE
            elif(self.action == UPLEFTFIRE):
                self.action = LEFTFIRE
        elif(key == Key.right):
            if  (self.action == RIGHT):
                self.action = NOOP
            elif(self.action == RIGHTFIRE):
                self.action = FIRE
            elif(self.action == UPRIGHT):
                self.action = UP
            elif(self.action == DOWNRIGHT):
                self.action = DOWN
            elif(self.action == UPRIGHTFIRE):
                self.action = UPFIRE
            elif(self.action == DOWNRIGHTFIRE):
                self.action = DOWNFIRE
        elif(key == Key.left):
            if  (self.action == LEFT):
                self.action = NOOP
            elif(self.action == LEFTFIRE):
                self.action = FIRE
            elif(self.action == UPLEFT):
                self.action = UP
            elif(self.action == DOWNLEFT):
                self.action = DOWN
            elif(self.action == UPLEFTFIRE):
                self.action = UPFIRE
            elif(self.action == DOWNLEFTFIRE):
                self.action = DOWNFIRE
        elif(key == Key.down):
            if  (self.action == DOWN):
                self.action = NOOP
            elif(self.action == DOWNFIRE):
                self.action = FIRE
            elif(self.action == DOWNRIGHT):
                self.action = RIGHT
            elif(self.action == DOWNLEFT):
                self.action = LEFT
            elif(self.action == DOWNRIGHTFIRE):
                self.action = RIGHTFIRE
            elif(self.action == DOWNLEFTFIRE):
                self.action = LEFTFIRE

    def set_env_action(self):
        if self.action == NOOP and "NOOP" in self.action2num:
            self.env_action = self.action2num["NOOP"]
        elif self.action == FIRE and "FIRE" in self.action2num:
            self.env_action = self.action2num["FIRE"]
        elif self.action == UP and "UP" in self.action2num:
            self.env_action = self.action2num["UP"]
        elif self.action == RIGHT and "RIGHT" in self.action2num:
            self.env_action = self.action2num["RIGHT"]
        elif self.action == LEFT and "LEFT" in self.action2num:
            self.env_action = self.action2num["LEFT"]
        elif self.action == DOWN and "DOWN" in self.action2num:
            self.env_action = self.action2num["DOWN"]
        elif self.action == UPRIGHT and "UPRIGHT" in self.action2num:
            self.env_action = self.action2num["UPRIGHT"]
        elif self.action == UPLEFT and "UPLEFT" in self.action2num:
            self.env_action = self.action2num["UPLEFT"]
        elif self.action == DOWNRIGHT and "DOWNRIGHT" in self.action2num:
            self.env_action = self.action2num["DOWNRIGHT"]
        elif self.action == DOWNLEFT and "DOWNLEFT" in self.action2num:
            self.env_action = self.action2num["DOWNLEFT"]
        elif self.action == UPFIRE and "UPFIRE" in self.action2num:
            self.env_action = self.action2num["UPFIRE"]
        elif self.action == RIGHTFIRE and "RIGHTFIRE" in self.action2num:
            self.env_action = self.action2num["RIGHTFIRE"]
        elif self.action == LEFTFIRE and "LEFTFIRE" in self.action2num:
            self.env_action = self.action2num["LEFTFIRE"]
        elif self.action == DOWNFIRE and "DOWNFIRE" in self.action2num:
            self.env_action = self.action2num["DOWNFIRE"]
        elif self.action == UPRIGHTFIRE and "UPRIGHTFIRE" in self.action2num:
            self.env_action = self.action2num["UPRIGHTFIRE"]
        elif self.action == UPLEFTFIRE and "UPLEFTFIRE" in self.action2num:
            self.env_action = self.action2num["UPLEFTFIRE"]
        elif self.action == DOWNRIGHTFIRE and "DOWNRIGHTFIRE" in self.action2num:
            self.env_action = self.action2num["DOWNRIGHTFIRE"]
        elif self.action == DOWNLEFTFIRE and "DOWNLEFTFIRE" in self.action2num:
            self.env_action = self.action2num["DOWNLEFTFIRE"]
        else:
            self.env_action = 0

max_stars_repo_name: bandreghetti/autoplay

max_issues_repo_name: bandreghetti/autoplay

max_forks_repo_name: bandreghetti/autoplay


text:
#  Routine to calibrate the image frames of Halpha for estimation of
#  absolute flux of net emission line by using line and continuum image
#  frames of standard objects and program objects.

import os.path
import numpy as np
from scipy import interpolate
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
from scipy import integrate
from scipy.interpolate import InterpolatedUnivariateSpline as interp
from astropy.io import fits
import time
from astropy.visualization import astropy_mpl_style
plt.style.use(astropy_mpl_style)
from astropy.visualization import make_lupton_rgb

def readfile(filename):
    ''' Reads a file.
    Input: Filename.
    Returns: x, y, xmin, xmax, number of points, filename.
    '''
    if os.path.isfile(filename) != True:
        print("File not in this directory!")
    else:
        file = np.loadtxt(filename, unpack=True)
        x = file[0]
        y = file[1]
        xmax = np.amax(x)
        xmin = np.amin(x)
        npts = len(file[0])
        return(x, y, xmin, xmax, npts, filename)

def obsWavelength(x, vel):
    ''' Corrects for systemic velocity.
    Input: Rest wavelength and velocity in km/s.
    Output: Observed wavelengh.
    '''
    c = 299792.458
    z = vel / c
    lambda_obs = x * (1 + z)
    return lambda_obs

def blackbody(wavelength, T):
    ''' Planck's law, which describes the electromagnetic radiation
    emitted by a black body in thermal equilibrium at a given temperature T.
    (Atomic units, and wavelength in Angstroms)
    '''
    h = 4.136e-15 # eV/s
    kb = 8.617e-5 # eV/K
    c = 2.99792458e18 # A/s
    B = 2 * h * c**2 / (wavelength**5 * (np.exp(h*c / (wavelength*kb*T)) - 1))
    return B

def heatmap(u,g,i,w,x,y):
    ''' Fits the Planck's law to each pixel of a given image.
    Input: flux in u, g and i bands, with their effective wavelengths. Dimension of the image (x,y).
    Output: array with the temperatures given by the fit.
    '''
    start = time.time()
    f = np.empty(shape=(x,y), dtype=float)
    for i in range(0,x):
        for j in range(0,y):
            fluxes = np.array([u_s[i][j], g_s[i][j], i_s[i][j]])
            if (fluxes != 0).all():
                popt, pcov = curve_fit(blackbody, wavelength, fluxes, bounds=(0,1e5))
                T = popt
            else:
                T = 0
            f[i][j] = T
    end = time.time()
    print(end - start,'s')
    return f

if __name__ == "__main__":
    # Data input
    Xl = 1.38
    print('Line filter air masses (program objects):', Xl)

    Xsl = 1.50
    print('Line filter air masses (standard objects):', Xsl)

    Xc = 1.34
    print('Air masses for the continuum filter (program):', Xc)

    Xsc = 1.54
    print('Air masses for the continuum filter (standard):', Xsc)

    kpl = 0.14
    print('mag/air mass for line:', kpl)

    kpc = 0.23
    print('mag/air mass for continuum:', kpc)

    fsl = 294.3
    print('Raw fluxes of the standard with sky subtracted (line):', fsl)

    fsc = 2565.5
    print('Raw fluxes of the standard with sky subtracted (continuum):', fsc)

    nline = 1
    print('Number of lines found in the line filter range, its rest wavelengths and fractional contribution (sum=1.0):', nline)

    rfWave = [6563,1]
    print('Rest wavelength (Angstroms) of each line as well as its fractional contibution (sum = 1.0):', rfWave)

    vsys = 45
    print('Systemic velocity of the galaxy:', vsys)

    texpl = 2400
    print('Exposure times of program frames (line):', texpl)

    texpc = 600
    print('Exposure times of program frames (continuum):', texpc)

    skyl = 67.9
    print('Sky background of the program frames in counts/pixel (line):', skyl)

    skyc = 123.8
    print('Sky background of the program frames in counts/pixel (continuum):', skyc)

    # Line filter
    path = '/Users/ziliotto/Documents/GitHub/calibrate_line'
    line = readfile(os.path.join(path, '6568.dat'))
    print('Line filter file:', line[5])
    linex = line[0]
    liney = line[1]
    xlmax = line[3]
    xlmin = line[2]

    # Continuum filter
    continuum = readfile(os.path.join(path, '6092.dat'))
    print('Continuum filter file:', continuum[5])
    xcmax = continuum[3]
    xcmin = continuum[2]
    contx = continuum[0]
    conty = continuum[1]

    # Standard
    standard = readfile(os.path.join(path, 'f34f.dat'))
    print('Standard flux file:', standard[5],'\n')
    standx = standard[0]
    standy = standard[1]
    xmin = standard[2]
    xmax = standard[3]

    # Interpolations
    lineInterp = interp(linex,liney)
    standInterp = interp(standx,standy)
    contInterp = interp(contx,conty)

    # Continuum function (for the case of constant continuum. Use blackbody function in the variable case)
    def continuum():
        ''' Returns the flux of the continuum.
        '''
        funcQ1 = lambda x: standInterp(x) * contInterp(x)
        funcQ2 = lambda x: contInterp(x)
        fc = (1 / fsc) * 10 ** ( 0.4 * kpc * ( Xc - Xsc ) ) * ( integrate.quad(funcQ1, xcmin, xcmax)[0] / integrate.quad(funcQ2, xcmin, xcmax)[0] )
        return fc

    # A
    funcA = lambda x: continuum() * lineInterp(x)
    A = integrate.quad( funcA, xlmin, xlmax, epsabs=1.49e-11 )[0]
    print('A =', A)

    # P
    funcP = lambda x: standInterp(x) * lineInterp(x) # phi_l
    P = (1 / fsl) * 10 ** ( 0.4 * kpl * ( Xl - Xsl ) ) * integrate.quad( funcP, xlmin, xlmax, epsabs=1.49e-11 )[0]
    print('P =', P)

    # Q
    funcQ1 = lambda x: standInterp(x) * contInterp(x) # phi_c
    funcQ2 = lambda x: continuum() * contInterp(x) # phi_c
    Q = (1 / fsc) * 10 ** ( 0.4 * kpc * ( Xc - Xsc ) ) * ( integrate.quad(funcQ1, xcmin, xcmax, epsabs=1.49e-11)[0] / integrate.quad(funcQ2, xcmin, xcmax, epsabs=1.49e-11)[0] )
    print('Q =', Q)

    # R
    fwhm = 1
    g = fwhm / (2 * np.sqrt(np.log(2)))
    cte = 1. / ( np.sqrt( np.pi ) * g )
    waveobs = obsWavelength(6563., vsys)
    funcR = lambda x: lineInterp(x) * np.exp(-(( x - waveobs )/g )**2)
    intR = integrate.quad(funcR, xlmin, xlmax)
    R = intR[0] * cte
    print('R =', R)

    Fl = print('F(line) =', ( 1 / R ), '* (', A, '- ( (I(line) - ', skyl,') / (I(continuum) - ', skyc, ') ) * ( ',texpc*P/(texpl*Q), ') ) ')
    print(Fl)

    u_image_file = fits.open('/Users/ziliotto/Downloads/stacks/ss_fornax_tile1_u_long_ALIGNi.003.fits')
    g_image_file = fits.open('/Users/ziliotto/Downloads/stacks/ss_fornax_tile1_g_long_ALIGNi.003.fits')
    #r_image_file = fits.open('/Users/ziliotto/Downloads/stacks/')
    i_image_file = fits.open('/Users/ziliotto/Downloads/stacks/ss_fornax_tile1_i_long.003.fits')

    u_gain = u_image_file['PRIMARY'].header['GAIN']
    g_gain = g_image_file['PRIMARY'].header['GAIN']
    i_gain = i_image_file['PRIMARY'].header['GAIN']

    u_image_data = fits.getdata('/Users/ziliotto/Downloads/stacks/ss_fornax_tile1_u_long_ALIGNi.003.fits', ext=0)
    g_image_data = fits.getdata('/Users/ziliotto/Downloads/stacks/ss_fornax_tile1_g_long_ALIGNi.003.fits', ext=0)
    i_image_data = fits.getdata('/Users/ziliotto/Downloads/stacks/ss_fornax_tile1_i_long.003.fits', ext=0)

    flux_u = u_image_data / u_gain
    flux_g = g_image_data / g_gain
    flux_i = i_image_data / i_gain

    # Effective wavelengths for u, g, i bands (Angstrom)
    # Available here: http://svo2.cab.inta-csic.es/theory/fps/index.php?mode=browse&gname=CTIO&gname2=DECam&asttype=
    u_w = 3580.78
    g_w = 4773.99
    r_w = 6444.80
    i_w = 7858.77

    wavelength = np.array([u_w,g_w,i_w])

    # Example for an individual galaxy (time of execution ~170s)
    u_s = flux_u[12220:12380,8840:9000]
    g_s = flux_g[12220:12380,8840:9000]
    i_s = flux_i[12220:12380,8840:9000]
    #print(u_s.shape)

    #image = make_lupton_rgb(i_s, g_s, u_s, Q=10, stretch=0.5)
    #plt.grid(None)
    #plt.imshow(image)
    f = heatmap(u_s,g_s,i_s,wavelength,u_s.shape[0],u_s.shape[1])
    plt.imshow(f,cmap='gray')
    plt.grid(None)
    plt.colorbar(label='Temperature (K)')
    plt.show()

reason: import numpy,from scipy,import astropy,from astropy


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

#######################################################################
# Código complementar ao Doc2VecFacil para criar os vocabs de um conjuntos de textos em pastas padronizadas
# Esse código, dicas de uso e outras informações:
#   -> https://github.com/luizanisio/Doc2VecFacil/
# Luiz Anísio
# 03/10/2021 - disponibilizado no GitHub
#######################################################################

from gensim.models import TfidfModel
from gensim.corpora import Dictionary

from util_doc2vec_facil import Documentos, TokenizadorInteligente, carregar_arquivo, map_thread
from collections import Counter
import os
import re
import pandas as pd
import numpy as np
from statistics import mean

# TFIDF: https://rockcontent.com/br/blog/tf-idf/

#tokenizador padrão sem dicionário
TOKENIZADOR = TokenizadorInteligente('',registrar_oov=False)
# tamanho máximo de token para registrar na saída
TAMANHO_MAXIMO_TOKEN = 50

# cria o posicionamento do valor no boxplot
def boxplot_str(valores, ignorar_zeros = True):
   vls = valores if not ignorar_zeros else [_ for _ in valores if _>0]
   q1, q2i, q2f, q3, irq, lw, uw = 0,0,0,0,0,0,0
   if any(vls):
        q1, q2i, q2f, q3 = np.percentile(vls,[25,49.5,50.5,75])
        irq = q3-q1
        lw = q1 - 1.5*irq
        uw = q3 + 1.5*irq
   vls = []
   #print(f'  {lw}|---[{q1}  {q3} ]---|{uw}  ')
   for v in valores:
       if ignorar_zeros and v==0:
          vls.append('')
       elif v<lw:
          vls.append('o |---[---]---|  ')
       elif v<q1:
          vls.append('  |-x-[---]---|  ')
       elif v<q2i:
          vls.append('  |---[x--]---|  ')
       elif v<q2f:
          vls.append('  |---[-x-]---|  ')
       elif v<q3:
          vls.append('  |---[--x]---|  ')
       elif v<uw:
          vls.append('  |---[---]-x-|  ')
       else:
          vls.append('  |---[---]---| o')
   return vls

def progress_bar(current_value, total, msg=''):
    increments = 25
    percentual = int((current_value / total) * 100)
    i = int(percentual // (100 / increments))
    text = "\r[{0: <{1}}] {2:.2f}%".format('=' * i, increments, percentual)
    print('{} {}           '.format(text, msg), end="\n" if percentual == 100 else "")

# gera os arquivos e retorna uma lista dos tokens (token, estava no vocab, estava no vocab quebrado)

def criar_arquivo_curadoria_termos(pasta_textos, pasta_vocab = None, complemento_arq_saida='', vocab_treino = False):
    cache_extensao = '.plan' # nome do cache do arquivo de curadoria
    # definição dos arquivos de curadoria
    arquivo_saida = os.path.join(pasta_vocab,f'curadoria_planilha_vocab{complemento_arq_saida}.xlsx')
    ####################################################################################################
    print('=============================================================')
    print(f'Análise iniciada: ', pasta_textos, ' para ',arquivo_saida)

    print('\t - carregando documentos para criar vocabulário')
    # recria o cache tokenizando tudo sem fragmentos
    docs = Documentos(pasta_vocab=pasta_vocab, pasta_textos=pasta_textos, retornar_tokens=True,
                      tokenizar_tudo=True, ignorar_cache=True, fragmentar=False, cache_extensao = cache_extensao)
    vocab_base = set(docs.tokenizer.vocab)
    vocab_removido = set(docs.tokenizer.vocab_removido)
    if vocab_treino:
        arq_vocab_treino = os.path.join(pasta_vocab,'vocab_treino.txt')
        if os.path.isfile(arq_vocab_treino):
            with open(arq_vocab_treino,'r') as f:
                vocab_base = {_.strip().replace('#','') for _ in f.read().split('\n') if _.strip()}
        else:
            print('=============================================================')
            print(f'ERRO: não foi possível encontrar o arquivo {arq_vocab_treino} para a carga do vocab')
            exit()

    # recupera a lista de tokens que entraram depois de quebrados para incluir na análise
    print(f'\t - processando documentos e compilando dicionário')
    dicionario = Dictionary(docs)
    ####################################################################################################
    print('\t - calculando pesos dos termos em cada documento processando todos os documentos')
    docs = Documentos(pasta_vocab=pasta_vocab, pasta_textos=pasta_textos, retornar_tokens=True,
                      tokenizar_tudo=True, ignorar_cache=False, fragmentar=False, cache_extensao = cache_extensao)
    print('\t - calculando modelo TFIDF')
    modelo_tfidf = TfidfModel((dicionario.doc2bow(d) for d in docs), normalize = True)
    print('\t - calculando quantidades e maior peso de cada termo nos documentos')
    docs = Documentos(pasta_vocab=pasta_vocab, pasta_textos=pasta_textos, retornar_tokens=True,
                      tokenizar_tudo=True, ignorar_cache=False, fragmentar=False, cache_extensao = cache_extensao)
    pesos_medios = dict({})
    contadores = Counter()
    contadores_docs = Counter()

    i, qtd = 1, docs.high
    for d in docs:
        if i % 100 == 0 or i == qtd:
            progress_bar(i,qtd,f'calculando pesos {i}/{qtd}')
        i += 1
        pesos = pesos_documento(modelo_tfidf, dicionario, d)
        contadores.update(Counter(d))
        contadores_docs.update(Counter(set(d)))
        for c,v in pesos.items():
            if v<=0:
               continue
            if not c in pesos_medios:
                pesos_medios[c] = []
            pesos_medios[c].append(v)

    for c,v in pesos_medios.items():
        pesos_medios[c] = mean(v)

    print('\t - incluindo termos sem contabilização dos pesos')
    termos_sem_peso = list(set(list(contadores.keys()) + list(vocab_base)))

    for t in termos_sem_peso:
        if not t in pesos_medios:
           pesos_medios[t] = 0
           contadores_docs[t] = 0
           contadores[t] = 0

    re_estranho = re.compile(r'([^rsaei])\1{2,}')
    re_estranho_consoantes = re.compile(r'[^rsaeiounl_\n]{3,}')
    re_estranho_vogais = re.compile(r'([aeiou])\1{2,}|([aeiou]{4,})')
    re_vogais = re.compile(r'[aeiouáéíóú]')
    def _estranho(termo):
        # não tem vogais ou tem letras repetidas
        if re_estranho.search(termo) or \
           re_estranho_consoantes.search(termo) or \
           re_estranho_vogais.search(termo):
            return 'S'
        if not re_vogais.search(termo):
            return 'S'
        if termo and (termo[0] == '_') or termo[-1]=='_':
            return 'S'
        return 'N'

    print('\t - preparando dataframe')
    dados_saida = []
    for c,v in pesos_medios.items():
        # tokens muito gandes são ignorados (exceto compostos)
        if len(c) > TAMANHO_MAXIMO_TOKEN and c.find('_') < 0:
            continue
        v = round(v,5)
        c = c.replace('#','')
        composto = 'S' if c.find('_') >= 0 else ''
        if composto:
            ok_vocab = 'TERMO (composto)' if c in vocab_base else 'NÃO (composto)'
        else:
            ok_vocab = 'TERMO' if c in vocab_base else ''
        ok_vocab = 'REMOVIDO' if c in vocab_removido else ok_vocab
        estranho = 'N' if ok_vocab else _estranho(c)
        quebrado = TOKENIZADOR.quebrar_token_simples(c)
        quebrado = '' if quebrado == c else quebrado
        quebrado = f'{quebrado} '.split(' ')
        prefixo =quebrado[0]
        sufixo = quebrado[1]
        # se não estiver no vocab, verifica se prefixo ou o prefixo e sufixo estão no vocab
        # - não importa ter só o sufixo
        if not ok_vocab:
            ok_vocab = 'PREFIXO' if prefixo in vocab_base else ''
            ok_vocab = f'{ok_vocab}+SUFIXO' if ok_vocab and sufixo in vocab_base else ok_vocab
            ok_vocab = 'NÃO' if not ok_vocab else ok_vocab
        linha = {'TERMO': c, 'PREFIXO':prefixo,'SUFIXO':sufixo, 'TFIDF': v, 'TAMANHO':len(c),
                 'QTD' :contadores[c],
                 'QTD_DOCS' : contadores_docs[c],
                 'VOCAB': ok_vocab, 'ESTRANHO': estranho}
        dados_saida.append(linha)
    if not any(dados_saida):
        print('===================================================================')
        print('ERRO: nenhum termo encontrado para criação da planilha de curadoria')
        print('Pasta do modelo: ', pasta_vocab)
        print('Pasta de textos: ', pasta_textos)
        print('===================================================================')
        exit()
    dados_saida = pd.DataFrame(dados_saida)
    print('\t - calculando boxplot tfidf')
    dados_saida['TFIDF Bp'] = boxplot_str(dados_saida['TFIDF'])
    print('\t - calculando boxplot tamanho')
    dados_saida['TAMANHO Bp'] = boxplot_str(dados_saida['TAMANHO'])
    print('\t - calculando boxplot qtd')
    dados_saida['QTD Bp'] = boxplot_str(dados_saida['QTD'])
    print('\t - calculando boxplot qtd_docs')
    dados_saida['QTD_DOCS Bp'] = boxplot_str(dados_saida['QTD_DOCS'])
    print('\t - finalizando arquivo excel')
    colunas = ['TERMO', 'PREFIXO', 'SUFIXO', 'TFIDF', 'TAMANHO', 'QTD', 'QTD_DOCS','VOCAB','ESTRANHO',
               'TFIDF Bp', 'TAMANHO Bp', 'QTD Bp', 'QTD_DOCS Bp',]
    dados_saida.to_excel(arquivo_saida,sheet_name=f'Curadoria de vocab para treino',
                         index = False, columns=colunas)

    print('Análise concluída: ', pasta_textos, ' para ',arquivo_saida)
    print('################################################################################')
    print('ARQUIVO DE CURADORIA CRIADO!!! ')
    print( ' - arquivo: ', arquivo_saida)
    print( ' - pasta analisada: ', pasta_textos)
    _nome_arq = os.path.split(arquivo_saida)[1]
    print(f' - você já pode abrir o arquivo {_nome_arq} no Excel')
    print('################################################################################')
    print(f'Removendo arquivos de cache da curadoria ...')
    lst_remover = docs.listar_documentos(cache_extensao)
    print(f' - removendo {len(lst_remover)} arquivos ... ', end='')
    apagar_arquivos(lst_remover)
    print(' Finalizado o/')

def apagar_arquivos(lista_de_arquivos):
    def _apagar(arquivo):
        try:
          os.remove(arquivo)
        except:
            # não tem motivo para poluir a tela com erro de exclusão
            # ou o arquivo está aberto, protegido ou o antivírus impediu tantas exclusões
            pass
    map_thread(_apagar, lista_de_arquivos, n_threads=3)

def pesos_documento(modelo_tfidf, dicionario, texto):
  tks = TOKENIZADOR.tokenizar(texto) if type(texto) is str else texto
  vec = modelo_tfidf[dicionario.doc2bow(tks)]
  return {dicionario[pair[0]]: pair[1] for pair in vec}

# faz a leitura de um arquivo de termos e quebra os tokens gerando um arquivo "nome_arquivo (QUEBRADOS).txt"
def quebrar_tokens_arquivo(arquivo):
    if not os.path.isfile(arquivo):
        print('quebrar_tokens_arquivo: aquivo não encontrado: ',arquivo)
    print('Quebrando tokens do arquivo: ', arquivo)
    arquivo_saida_quebrados = arquivo[:-4] if arquivo.lower().endswith('.txt') else arquivo.replace('.','_')
    arquivo_saida_quebrados = f'{arquivo_saida_quebrados} (QUEBRADOS).txt'
    vocab = carregar_arquivo(arquivo,juntar_linhas=True)
    vocab = TokenizadorInteligente.remover_acentos( vocab.replace('\n',' ').lower() )
    vocab = TokenizadorInteligente.retornar_vocab_texto(vocab)
    print(f'\t - gravando {len(vocab)} tokens criados')
    with open(arquivo_saida_quebrados,'w') as f:
         [f.write(f'{_}\n') for _ in sorted(vocab)]

def testar_tokenizador(texto='', pasta_vocab='./meu_modelo/doc2vecfacil/'):
    txt = texto if texto else 'Esse é um pequeno teste de tokenização de texto com ngramas 1 2 3 4 5 6 7 8 9 10: processos administrativos processo administrativo. Outra frases singularizadas administrativas politicos politicas.'
    tokenizador = TokenizadorInteligente(pasta_vocab=pasta_vocab, registrar_oov=True)
    print(f'Tokenizador carregado: {pasta_vocab}')
    print(f' - vocab carregado com {len(tokenizador.vocab)} termos')
    if tokenizador.tradutor_termos:
        print(f' - tradutor carregado com {len(tokenizador.tradutor_termos.termos)} termos')
    print(f' - texto: {txt}')
    tokens = tokenizador.tokenizar(txt)
    print(f' - tokens encontrados: {len(tokens)}')
    print(f'Tokens: |{"|".join(tokens)}|')
    print(f'OOV: |{"|".join(tokenizador.oov)}|')

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Pasta do modelo')
    parser.add_argument('-pasta', help='pasta raiz contendo as pastas do modelo e dos textos', required=False)
    parser.add_argument('-testar', help='realiza a carga do tokenizador para teste - não processa documentos', required=False, action='store_const', const=1)
    parser.add_argument('-treino', help='gera a curadoria com os textos de treino', required=False, action='store_const', const=1)
    parser.add_argument('-teste', help='gera a curadoria com os textos de teste', required=False, action='store_const', const=1)
    parser.add_argument('-vocab_treino', help='gera a curadoria com os termos do vocab de treinamento', required=False, action='store_const', const=1)
    args = parser.parse_args()

    PASTA_BASE = args.pasta if args.pasta else './meu_modelo'
    teste_tokenizar = args.testar
    pasta_treino = args.treino
    pasta_teste = args.teste and not pasta_treino
    vocab_treino = args.vocab_treino

    if not os.path.isdir(PASTA_BASE):
        msg = f'Não foi encontrada a pasta "{PASTA_BASE}" para criação dos vocabulários'
        raise Exception(msg)

    if teste_tokenizar:
        pasta_base = PASTA_BASE
        if not os.path.isdir(pasta_base):
            if os.path.isdir(os.path.join(pasta_base,'doc2vecfacil')):
                pasta_base = os.path.join(pasta_base,'doc2vecfacil')

        if not os.path.isdir(pasta_base):
            msg = f'Não foi encontrada a pasta "{pasta_base}" para teste do tokenizador'
            raise Exception(msg)

        print('###########################################################')
        print('# Carregando tokenizador para testes                      #')
        print('###########################################################')
        testar_tokenizador(pasta_vocab=pasta_base)
        exit()

    PASTA_TEXTOS_VOCAB = os.path.join(PASTA_BASE,'textos_vocab')
    PASTA_TEXTOS_TREINO = os.path.join(PASTA_BASE,'textos_treino')
    PASTA_TEXTOS_TESTE = os.path.join(PASTA_BASE,'textos_teste')
    PASTA_MODELO = os.path.join(PASTA_BASE,'doc2vecfacil')
    os.makedirs(PASTA_MODELO, exist_ok=True)

    print('################################################################################')
    print('# Criando arquivo de curadoria                            ')

    _vocab = ' - VOCAB TREINO' if vocab_treino else ''
    if pasta_treino:
        pasta = PASTA_TEXTOS_TREINO
        complemento_arq = f' (TREINO{_vocab})'
    elif pasta_teste:
        pasta = PASTA_TEXTOS_TESTE
        complemento_arq = f' (TESTE{_vocab})'
    elif not (os.path.isdir(PASTA_TEXTOS_VOCAB)):
        pasta = PASTA_TEXTOS_TREINO
        complemento_arq = f' (TREINO{_vocab})'
    else:
        pasta = PASTA_TEXTOS_VOCAB
        complemento_arq = f'{_vocab}'
    print(f'# Incluindo textos da pasta "{pasta}" para a curadoria{complemento_arq}')
    print('################################################################################')

    criar_arquivo_curadoria_termos(pasta_textos = pasta,
                                   pasta_vocab=PASTA_MODELO,
                                   complemento_arq_saida = complemento_arq,
                                   vocab_treino=vocab_treino)


text:
import sys
import math
import numpy as np
import csv

class NBClassifier(object):
    def __init__(self):
        pass
    def train(self, data):
        """
        @param: list of unprocessed data
        """
        self.raw_data = data

        self.attr_count = len(data[0]) - 2
        self.total_samples = len(data)

        self.split_data_by_class()
        self.calculate_priors()
        self.calculate_gaussian_vars()
        self.calculates_features_to_ignore()

    def split_data_by_class(self):
        self.data_by_class = {}

        for raw_data_row in self.raw_data:
            class_name = raw_data_row[-1]
            data_row = [float(x) for x in raw_data_row[1:-1]]

            if class_name not in self.data_by_class:
                self.data_by_class.setdefault(class_name, [])
            self.data_by_class[class_name].append(data_row)

    def calculates_features_to_ignore(self):
        self.features_to_ignore = []
        for class_name, gaussian_vars in self.gaussian_vars.items():
            for i, gv in enumerate(gaussian_vars):
                if gv['var'] == 0.0 and i not in self.features_to_ignore:
                    self.features_to_ignore.append(i)

    def calculate_priors(self):
        self.priors = {}
        for class_name, data_rows in self.data_by_class.items():
            prior_prob = float(len(data_rows))/float(self.total_samples)
            self.priors.setdefault(class_name, prior_prob)

    def calculate_gaussian_vars(self):
        self.gaussian_vars = {}

        for class_name, data_rows in self.data_by_class.items():
            np_data_rows = np.array(data_rows)
            #print np_data_rows
            for i in range(self.attr_count):
                attr_vals = np_data_rows[:, i]
                mean = attr_vals.mean()
                var = attr_vals.var()
                if class_name not in self.gaussian_vars:
                    self.gaussian_vars.setdefault(class_name, [])
                self.gaussian_vars[class_name].append({'mean': mean, 'var': var})

    def test(self, row, ignore=False, use_numpy_min=False):
        actual_class_name = row[-1]
        row = [float(x) for x in row[1:-1]]

        max_prob = -1.0
        predicted_class = None

        for class_name, gaussian_vars in self.gaussian_vars.items():
            calculated_prob = self.priors[class_name]

            for i, feature_val in enumerate(row):
                if ignore:
                    if i not in self.features_to_ignore:
                        if gaussian_vars[i]['var'] > 0.0:

                            mult_1 = 1.0/(math.sqrt(2*math.pi*gaussian_vars[i]['var']))
                            mult_2 = math.exp((-1.0/(2.0*gaussian_vars[i]['var']))*pow((feature_val - gaussian_vars[i]['mean']), 2))
                            calculated_prob *= (mult_1*mult_2)

                        elif use_numpy_min:
                            _var = 2.01032500e-09
                            mult_1 = 1.0/(math.sqrt(2*math.pi*_var))
                            mult_2 = math.exp((-1.0/(2.0*_var))*pow((feature_val - gaussian_vars[i]['mean']), 2))
                            calculated_prob *= (mult_1*mult_2)
                else:
                    if gaussian_vars[i]['var'] > 0.0:
                        mult_1 = 1.0/(math.sqrt(2*math.pi*gaussian_vars[i]['var']))
                        mult_2 = math.exp((-1.0/(2.0*gaussian_vars[i]['var']))*pow((feature_val - gaussian_vars[i]['mean']), 2))
                        calculated_prob *= (mult_1*mult_2)

                    elif use_numpy_min:
                        _var = 2.01032500e-09
                        mult_1 = 1.0/(math.sqrt(2*math.pi*_var))
                        mult_2 = math.exp((-1.0/(2.0*_var))*pow((feature_val - gaussian_vars[i]['mean']), 2))
                        calculated_prob *= (mult_1*mult_2)

            if calculated_prob > max_prob or predicted_class == None:
                max_prob = calculated_prob
                predicted_class = class_name

        return predicted_class#, max_prob

def main(train_file_path, test_file_path):

    train_file_csv_reader = csv.reader(open(train_file_path, 'r'))
    test_file_csv_reader = csv.reader(open(test_file_path, 'r'))

    nbClassifier = NBClassifier()
    nbClassifier.train(list(train_file_csv_reader))

    test_passed = 0
    test_failed = 0

    for row in test_file_csv_reader:
        predicted_class = nbClassifier.test(row)
        if predicted_class == row[-1]:
            test_passed +=1
        else:
            test_failed += 1
        print "actual_class: {}, predicted_class: {}".format(row[-1], predicted_class)
    print "passed: {}, failed: {}, accuracy: {}%".format(test_passed, test_failed, ((test_passed)/float(test_passed + test_failed))*100)
    return nbClassifier

if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])


text:
from numpy import random
from .settings import Settings
from .grid import Grid
from .const import VECTOR

class Game(object):

	def __init__(self):
		self.__score = 0
		self.__lastMoveScore = 0
		self.__grid = Grid()
		self.__addStartTiles()

	def __addStartTiles(self):
		for i in range(Settings.START_TILES):
			self.__addRandomTile()

	def __addRandomTile(self):
		cells = self.__grid.CellsAvailable()
		r = random.randint(len(cells))
		(x, y) = cells[r]
		if random.rand() < 0.9:
			value = 2
		else:
			value = 4
		self.__grid.set(x,y, value)

	def GetGrid(self):
		return self.__grid.getTable()

	def Move(self, direction):
		if self.__grid.Move(VECTOR(direction)):
			self.__updateScore()
			self.__addRandomTile()
		else:
			cells = self.__grid.CellsAvailable()
			if len(cells) == 0:
				self.__gameOver()
				return

	def __updateScore(self):
		score = self.__grid.getScore()
		self.__lastMoveScore = score
		self.__score += score

	def GetTotalScore(self):
		return self.__score

	def GetLastMoveScore(self):
		return self.__lastMoveScore

	def __gameOver(self):
		if self.__gameOverFunc:
			self.__gameOverFunc()

	def onGameOver(self, func):
		self.__gameOverFunc = func

	def Print(self):
		self.__grid.Print()
		print()

reason: from numpy


text:
def Chi_sq_test(df, dependant, independant):
    #takes in the names of a dependant and independant variable (column), runs a chi squared test and then outputs
    #a seaborn heatmap of the percent difference between the expected and actual values

    from scipy.stats import chi2_contingency
    from scipy.stats import chi2

    #create cotingency table
    count_table = df.groupby([dependant, independant])['id'].count().unstack()
    count_table.fillna(0, inplace=True)
    count_table = count_table.astype('int')

    #Chi Squared test is for only counts above 5, we are keeping the same ratio, but increasing min value to 5 in each column
    if count_table.isin(range(0,5)).any().any():
        for j in range(len(count_table.columns)):
            for i in range(len(count_table.index)):
                if count_table.iloc[i,j] < 1:
                    count_table.iloc[i,j] = 5
                    count_table.iloc[:,j] = count_table.iloc[:,j]*5
                elif count_table.iloc[i,j] <5:
                    count_table.iloc[:,j] = count_table.iloc[:,j]*(5/count_table.iloc[i,j])

    stat, p, dof, expected = chi2_contingency(count_table)

    #print test information
    print('P-Value = {}'.format(p))
    print('Chi Statistic = {}'.format(stat))
    print('Degrees of Freedom = {}'.format(dof))

    #caluclate and print heatmap
    plt.figure(figsize=(12,6))
    sns.heatmap(((count_table - expected) / count_table *100), annot=True, vmax=100, vmin=-100, fmt='.1f',
                annot_kws={'rotation': 90}, cmap='viridis')
    plt.title('Percent Difference of Expected vs. Actual Classes per {}'.format(str.title(independant)))
    plt.show()

reason: from scipy


text:
from sympy.liealgebras.cartan_type import CartanType
from sympy.core.compatibility import range
from sympy.matrices import Matrix
from sympy.core.backend import S

def test_type_F():
    c = CartanType("F4")
    m = Matrix(4, 4, [2, -1, 0, 0, -1, 2, -2, 0, 0, -1, 2, -1, 0, 0, -1, 2])
    assert c.cartan_matrix() == m
    assert c.dimension() == 4
    assert c.simple_root(1) == [1, -1, 0, 0]
    assert c.simple_root(2) == [0, 1, -1, 0]
    assert c.simple_root(3) == [0, 0, 0, 1]
    assert c.simple_root(4) == [-S.Half, -S.Half, -S.Half, -S.Half]
    assert c.roots() == 48
    assert c.basis() == 52
    diag = "0---0=>=0---0\n" + "   ".join(str(i) for i in range(1, 5))
    assert c.dynkin_diagram() == diag
    assert c.positive_roots() == {1: [1, -1, 0, 0], 2: [1, 1, 0, 0], 3: [1, 0, -1, 0],
            4: [1, 0, 1, 0], 5: [1, 0, 0, -1], 6: [1, 0, 0, 1], 7: [0, 1, -1, 0],
            8: [0, 1, 1, 0], 9: [0, 1, 0, -1], 10: [0, 1, 0, 1], 11: [0, 0, 1, -1],
            12: [0, 0, 1, 1], 13: [1, 0, 0, 0], 14: [0, 1, 0, 0], 15: [0, 0, 1, 0],
            16: [0, 0, 0, 1], 17: [S.Half, S.Half, S.Half, S.Half], 18: [S.Half, -S.Half, S.Half, S.Half],
            19: [S.Half, S.Half, -S.Half, S.Half], 20: [S.Half, S.Half, S.Half, -S.Half], 21: [S.Half, S.Half, -S.Half, -S.Half],
            22: [S.Half, -S.Half, S.Half, -S.Half], 23: [S.Half, -S.Half, -S.Half, S.Half], 24: [S.Half, -S.Half, -S.Half, -S.Half]}

reason: from sympy


text:
from __future__ import division
import gym
import gym_rover

import numpy as np

from gym_rover.envs.rover_env import RoverEnv

import torch
from torch.autograd import Variable

import os
import psutil
import gc

import ddpg_train
import ddpg_buffer
import ddpg_train

env = gym.make('rover-v0')

MAX_EPISODES = 1000
MAX_BUFFER = 100000
NUM_POIS = 2
NUM_AGENTS = 1
TIME_LIMIT = 20
LENGTH = 10

env.world_height = LENGTH
env.world_width = LENGTH
env.num_agents = NUM_AGENTS
env.num_pois = NUM_POIS
env.time_limit = TIME_LIMIT

env.observation_mode = 'feature'
env.actions = 'continuous'

env.set_observation_space()
env.set_action_space()
env.reset()

S_DIM = env._box_one_obs.shape[0]
A_DIM = env._agent_action.shape[0]
A_MAX = env._agent_action.high[0]

trainers = []
rams = []
for _ in range(NUM_AGENTS):
    ram = ddpg_buffer.MemoryBuffer(MAX_BUFFER)
    rams.append(ram)
    trainers.append(ddpg_train.Trainer(S_DIM, A_DIM, A_MAX, ram))

for _ep in range(MAX_EPISODES):
    observation = env.reset()
    done = False

    while not done:
        #env.render()
        actions = []
        for i in range(NUM_AGENTS):
            state = np.float32(observation[i])
            action = trainers[i].get_exploration_action(state)
            #print (action)
            actions.append(action)

        next_states, reward, done, _ = env.step(np.array(actions))

        if not done:
            for i in range(NUM_AGENTS):
                state = np.float32(observation[i])
                new_state = np.float32(next_states[i])
                rams[i].add(state, actions[i], reward, new_state)

        observation = next_states
        for t in trainers:
            t.optimize()

        if done:
            print ('EPISODE :- ' + str(_ep) + ' ' + str(reward))
            break
    gc.collect()

max_stars_repo_name: eklinkhammer/gym-rover

max_issues_repo_name: eklinkhammer/gym-rover

max_forks_repo_name: eklinkhammer/gym-rover


text:
#from mpc import mpc
#from mpc.mpc import QuadCost, LinDx, GradMethods
import torch
import numpy as np
import torch.nn as nn
import pdb
from ..scene_funcs.cnn import CNN
from ..scene_funcs.scene_funcs import scene_funcs
from .. import augmentation
import time
from .utils import *
import cv2
import trajnetbaselines

import warnings
warnings.filterwarnings("ignore")

class BicycleModel(nn.Module):
    def __init__(self, scale=None, dt=None):
        super(BicycleModel, self).__init__()

        self.scale = scale
        self.dt = dt

    def forward(self, state, action):
        squeeze = state.ndimension() == 1

        if squeeze:
            state = state.unsqueeze(0)
            action = action.unsqueeze(0)

        u = action
        # pdb.set_trace()
        # print(u.shape)
        # print(state.shape)
        a = torch.clamp(u[:, 0], -self.scale * 8, self.scale * 4)
        psi = torch.clamp(u[:, 1], -45 * 0.017, 45 * 0.017)

        beta = torch.atan(torch.tan(psi) * 1.77 / (1.77 + 1.17))
        new_x_pos = state[:, 0] + state[:, 2] * self.dt * torch.cos(state[:, 3] + beta)
        new_y_pos = state[:, 1] + state[:, 2] * self.dt * torch.sin(state[:, 3] + beta)
        new_v = state[:, 2] + a * self.dt
        new_psi = state[:, 3] + state[:, 2] / (1.77 * self.scale) * self.dt * torch.sin(beta)

        x = torch.cat((new_x_pos.unsqueeze(1), new_y_pos.unsqueeze(1), new_v.unsqueeze(1), new_psi.unsqueeze(1)), dim=1)
        if squeeze:
            x = x.squeeze(0)
        return x

def mpc_function(predicted_trajectory, obs, scale, sample_rate, n_pred):
    pred_traj = predicted_trajectory[:, :, :, :2].cpu()  # to exclude variance when predicting Gauss
    mpc_trajc = pred_traj.view(-1, n_pred, 2)

    dx = BicycleModel(scale=scale[0, 0, 0].item(), dt=sample_rate[0, 0, 0].item())

    n_batch, n_state, n_ctrl, T = mpc_trajc.shape[0], 4, 2, n_pred + 1
    ctrl_penalty = 0.1
    n_sc = n_state + n_ctrl

    start_point = torch.zeros(mpc_trajc.shape[0], 2)
    compteur = 0
    for j in range(pred_traj.shape[0]):
        for k in range(pred_traj.shape[1]):
            start_point[compteur, :] = obs[j, -1, 0, :]
            compteur += 1

    v0 = torch.tensor(
        torch.sqrt((mpc_trajc[:, 0, 0] - start_point[:, 0]) ** 2 + (mpc_trajc[:, 0, 1] - start_point[:, 1]) ** 2) /
        sample_rate[0, 0, 0].item()).unsqueeze(1)
    psi0 = torch.tensor((torch.atan2((mpc_trajc[:, 0, 1] - start_point[:, 1]),
                                     (mpc_trajc[:, 0, 0] - start_point[:, 0])))).unsqueeze(1)

    x0 = start_point[:, 0].unsqueeze(1)
    y0 = start_point[:, 1].unsqueeze(1)
    x_init = torch.cat((x0, y0, v0, psi0), 1)
    output_network = torch.cat((mpc_trajc, torch.zeros(n_batch, 10, 2)), dim=2)
    goal_state = torch.cat((x_init.unsqueeze(1), output_network), dim=1).permute(1, 0, 2).unsqueeze(2)
    goal_weights = torch.Tensor((5., 5., 0., 0.)).repeat(T, n_batch, 1, 1)
    px = -torch.sqrt(goal_weights) * goal_state
    q = torch.cat((goal_weights, ctrl_penalty * torch.ones(n_ctrl).repeat(T, n_batch, 1, 1)), 3)
    p = torch.cat((px, torch.zeros(n_ctrl).repeat(T, n_batch, 1, 1)), 3).squeeze(2).to(device=obs.device)
    Q = (q * torch.eye((n_sc))).to(device=obs.device)

    u_lower = -torch.cat(
        (torch.ones(T, n_batch, 1) * 8 * scale[0, 0, 0].item(), torch.ones(T, n_batch, 1) * 45 * 0.017), dim=2).squeeze(
        2).to(device=obs.device)
    u_upper = torch.cat((torch.ones(T, n_batch, 1) * 4 * scale[0, 0, 0].item(), torch.ones(T, n_batch, 1) * 45 * 0.017),
                        dim=2).squeeze(2).to(device=obs.device)
    # pdb.set_trace()

    # .to(device=obs.device)
    x_pos, _, _ = mpc.MPC(
        n_state=n_state,
        n_ctrl=n_ctrl,
        T=T,
        u_lower=u_lower,
        u_upper=u_upper,
        lqr_iter=30,
        verbose=-1,
        n_batch=n_batch,
        grad_method=GradMethods.AUTO_DIFF,
        # grad_method=GradMethods.FINITE_DIFF,
        exit_unconverged=False,
        detach_unconverged=False,
        backprop=True,
        slew_rate_penalty=0.1,
    )(x_init.to(device=obs.device), QuadCost(Q, p), dx)

    del u_upper, u_lower, Q, p
    mpc_position = x_pos.permute(1, 0, 2)[:, 1:, :2].view(pred_traj.shape[0], pred_traj.shape[1], pred_traj.shape[2], 2)
    # mpc_position = x_pos.permute(1, 0, 2)[:,1:,:2].contiguous().view(pred_traj.shape[0],pred_traj.shape[1],pred_traj.shape[2],2)
    # pdb.set_trace()
    if (pred_traj.shape[3] == 4):  # if it is gaussian
        mpc_position = torch.cat([mpc_position, predicted_trajectory[:, :, :, 2:]], dim=3)

    return mpc_position

def min_distance_2_points(beg_x, beg_y, end_x, end_y, line):
    # min_dist_beg, indice_beg = torch.min(((beg_x - line[0]) ** 2 + (beg_y - line[1]) ** 2) ** 0.5, 0)
    min_dist_beg, indice_beg = torch.min((abs(beg_x - line[0]) + abs(beg_y - line[1])), 0)
    # min_dist_end, indice_end = torch.min(((end_x - line[0][indice_beg.item():]) ** 2 + (end_y - line[1][indice_beg:]) ** 2) ** 0.5, 0)
    min_dist_end, indice_end = torch.min((abs(end_x - line[0][indice_beg.item():]) + abs(end_y - line[1][indice_beg:])),
                                         0)
    if (indice_end + 20 > len(line[0])):
        dist = len(line[0]) - (indice_end + 20 + 1)
    else:
        dist = 0

    return min_dist_beg + min_dist_end, indice_beg - 2, (indice_end + indice_beg) + 20 + dist

def best_road(obs, pred, center_lines, device, k, iteration, epochs=None, num_roads=1, rotated_scene=None, margin=None):
    dist_beg_obs = torch.sum(abs(center_lines - obs[0:2].unsqueeze(0).unsqueeze(2).type(torch.float64)), dim=1)
    dist_end_obs = torch.sum(abs(center_lines - obs[2:].unsqueeze(0).unsqueeze(2).type(torch.float64)), dim=1)
    dist_beg_pred = torch.sum(abs(center_lines - pred[0:2].unsqueeze(0).unsqueeze(2).type(torch.float64)), dim=1)
    dist_end_pred = torch.sum(abs(center_lines - pred[2:].unsqueeze(0).unsqueeze(2).type(torch.float64)), dim=1)
    min_dist_beg_obs, min_dist_end_obs = torch.min(dist_beg_obs, dim=1), torch.min(dist_end_obs, dim=1)
    min_dist_beg_pred, min_dist_end_pred = torch.min(dist_beg_pred, dim=1), torch.min(dist_end_pred, dim=1)

    boolean_obs = torch.ones((min_dist_beg_obs.values.shape[0]), dtype=torch.float64).cuda()
    boolean_pred = torch.ones((min_dist_beg_pred.values.shape[0]), dtype=torch.float64).cuda()
    # l2 = time.time()

    boolean_obs = (min_dist_beg_obs.indices > min_dist_end_obs.indices).type(torch.float64) * 10000 + 1
    boolean_pred = (min_dist_beg_pred.indices > min_dist_end_pred.indices).type(torch.float64) * 10000 + 1
    # l3 = time.time()
    if (num_roads == 1):
        #
        min_dist_obs = (min_dist_beg_obs.values + min_dist_end_obs.values) * boolean_obs
        min_dist_pred = (min_dist_beg_pred.values + min_dist_end_pred.values) * boolean_pred
        sorted_obs, best_centers_obs = torch.sort(min_dist_obs,
                                                  dim=0)  # the starting point is more mportant than the end because we want the vehicle to continue from it's current position
        best_center = torch.min(min_dist_obs * 2 + min_dist_pred, dim=0).indices.item()
        # pdb.set_trace()
    else:
        ###This part should be fixed like the previous bock
        distances_fixMePleaseLikePreviousBlock = (min_dist_beg.values + min_dist_end.values) * boolean
        min_args = torch.argsort(distances,
                                 dim=0)  # the starting point is more mportant than the end because we want the vehicle to continue from it's current position
        best_center = min_args[0].item()
        second_center = min_args[0].item()
        max_of_min_dist_end = min_dist_end.values[best_center]
        for i in min_args:
            if (distances[i] < distances[best_center] + 20 and min_dist_end.values[i] > max_of_min_dist_end):
                second_center = i.item()
                max_of_min_dist_end = min_dist_end.values[i]

    indice_beg_final = min_dist_beg_pred.indices[best_center]

    if (num_roads == 1):
        return (center_lines[best_center][0][indice_beg_final:]), (center_lines[best_center][1][indice_beg_final:])
    return (center_lines[best_center][0][indice_beg_final:]), (center_lines[best_center][1][indice_beg_final:]), (
    center_lines[second_center][0][indice_beg_final:]), (center_lines[second_center][1][indice_beg_final:])

def extract_center_line_of_interset(obsrvation, prediction_p, center_line, device, iteration, epochs, n_pred,
                                    num_roads=1):
    if (num_roads == 1):
        center_lines = torch.zeros((prediction_p.shape[0], prediction_p.shape[1], 600, 2), device=device)
    else:
        center_lines = torch.zeros((prediction_p.shape[0], num_roads, 600, 2), device=device)
    for i in range(int(prediction_p.shape[0])):  # batch size
        for j in range(int(prediction_p.shape[1])):  # pred modes
            # lap1 = time.time()
            pred = torch.tensor([], device=device)
            obs = torch.tensor([], device=device)
            beg_pred_x, beg_pred_y, end_pred_x, end_pred_y = prediction_p[i, j, 0, 0], prediction_p[i, j, 0, 1], \
                                                             prediction_p[i, j, n_pred - 1, 0], prediction_p[
                                                                 i, j, n_pred - 1, 1]
            beg_obs_x, beg_obs_y, end_obs_x, end_obs_y = obsrvation[i, 0, 0, 0], obsrvation[i, 0, 0, 1], obsrvation[
                i, -1, 0, 0], obsrvation[i, -1, 0, 1]

            pred = torch.cat((pred, beg_pred_x.unsqueeze(-1), beg_pred_y.unsqueeze(-1), end_pred_x.unsqueeze(-1),
                              end_pred_y.unsqueeze(-1)))
            obs = torch.cat((obs, beg_obs_x.unsqueeze(-1), beg_obs_y.unsqueeze(-1), end_obs_x.unsqueeze(-1),
                             end_obs_y.unsqueeze(-1)))
            result = best_road(obs, pred, center_line[i], device, i, iteration, epochs, num_roads)
            ##lap2 = time.time()
            if (len(result[0]) > 600):
                center_lines[i, j, :, 0] = result[0][:600]
                center_lines[i, j, :, 1] = result[1][:600]
            else:
                center_lines[i, j, :len(result[0]), 0] = result[0]
                center_lines[i, j, :len(result[0]), 1] = result[1]
            if (num_roads is not 1):
                if (len(result[2]) > 600):
                    center_lines[i, 1, :, 0] = result[2][:600]
                    center_lines[i, 1, :, 1] = result[3][:600]
                else:
                    center_lines[i, 1, :len(result[2]), 0] = result[2]
                    center_lines[i, 1, :len(result[2]), 1] = result[3]
    return center_lines

max_stars_repo_name: vita-epfl/RRB