text:
"""
@author: cvincentcuaz

Torch implementation of semi-relaxed Gromov-Wasserstein dictionary learning,
detailed in section 4 of the main paper. Also contains the srgw completion framework
detailed in this latter section.

It Supports all kinds of regularization for solving the unmixing problem
detailed in the section 3 of the main paper.

"""

from data_handler import dataloader
import numpy as np
from tqdm import tqdm
import os
import pandas as pd
import pickle
from sklearn.model_selection import train_test_split
import pylab as pl
from sklearn.cluster import KMeans
from  scipy.sparse.csgraph import shortest_path
import torch as th
import srGW_algorithms.srGW as srGW

#%%

str2rpzfunctions ={'ADJ': (lambda x: x),
                   'SP':shortest_path}

class srFGW_DL():

    def __init__(self,
                 graphs:list=None,
                 features:list=None,
                 masses:list=None,
                 y:list=None,
                 dataset_name:str=None,
                 mode:str='ADJ',
                 Ntarget:int=None,
                 alpha:float=0.5,
                 experiment_repo:str=None,
                 experiment_name:str=None,
                 degrees:bool=False,
                 completion_parameters:dict={},
                 data_path:str='../data/',
                 dtype:type=th.float64,
                 device:str='cpu'):
        """
        Parameters
        ----------
        graphs: list of torch arrays (N_k, N_k).
                If set to "None", graphs will be downloaded from the specified "dataset_name"
        features: list of torch arrays (N_k, d).
                If set to "None", features will be downloaded from the specified "dataset_name"
        masses: list of numpy arrays matching respectively graphs sizes.
                If set to "None", computed based on downloaded graphs from "dataset_name"
        y: numpy array (N_k,)
            If set to "None", labels will be downloaded from the specified "dataset_name". Used for analysis of the unsupervised learning process.
        dataset_name :
            name of the dataset to experiment on. To match our data loaders it is restricted to ['imdb-b','imdb-m','balanced_clustertoy','clustertoy2C']
        mode :
            representations for input graphs. (e.g) 'ADJ':adjacency / 'SP': shortest path
        Ntarget:
            size of target graph which summarizes the dataset
        alpha:
            Trade-off parameter involved in the Fused Gromov-Wasserstein loss function
        experiment_repo :
            subrepository to save results of the experiment
        experiment_name :
            subrepository to save results of the experiment under the 'experiment repo' repository
        degrees:
            either to use uniform distribution (False) for each graph, else use degree distribution (True)
        completion_parameters:
            dict used to handle split into D_train and D_test
        data_path :
            path where data is. The default is '../data/'.
        """
        """
        Parameters
        ----------
        dataset_name : name of the dataset to experiment on. To match our data loaders it is restricted to ['imdb-b','imdb-m','balanced_clustertoy','clustertoy2C']
        mode : representations for input graphs. (e.g) 'ADJ':adjacency / 'SP': shortest path
        Ntarget: size of target graph which summarizes the dataset
        experiment_repo : subrepository to save results of the experiment
        experiment_name : subrepository to save results of the experiment under the 'experiment repo' repository
        degrees: either to use uniform distribution (False) for each graph, else use degree distribution (True)
        denoising_parameters:   If learning dictionary on noisy observations
                                must contain keys ['n_samples', 'n_trials', 'list_probas']
        data_path : path where data is. The default is '../data/'.
        """

        self.experiment_repo= experiment_repo
        self.experiment_name = experiment_name
        print('dataset_name:', dataset_name)
        str_to_method = {'ADJ': 'adjacency', 'SP':'shortest_path','LAP':'laplacian'}
        self.degrees=degrees
        self.Ntarget = Ntarget
        self.alpha=alpha
        self.dataset_name= dataset_name
        self.mode = mode
        self.completion_parameters = completion_parameters
        self.dtype = dtype
        self.device = device

        if graphs is None:
            if dataset_name in ['mutag','ptc']: #onehot representation for features
                one_hot = True
                X, self.y = dataloader.load_local_data(data_path, dataset_name, one_hot=one_hot)
            else:
                X, self.y = dataloader.load_local_data(data_path, dataset_name)

            if completion_parameters =={}:
                if self.mode in str_to_method.keys():
                    self.graphs= [th.tensor(X[t].distance_matrix(method=str_to_method[mode]), dtype=self.dtype, device=self.device) for t in range(X.shape[0])]
                    self.features = [th.tensor(X[t].values(), dtype=self.dtype, device=self.device) for t in range(X.shape[0])]
                    if not self.degrees:
                        #uniform distributions
                        self.masses= [th.ones(Cs.shape[0], dtype=self.dtype, device=self.device)/Cs.shape[0] for Cs in self.graphs]
                    else:
                        print('computing degree distributions')
                        self.masses =[]
                        for C in self.graphs:
                            h = C.sum(axis=0)
                            self.masses.append( h / h.sum())
                else:
                    raise 'unknown mode /graph representation'
            else:
                assert self.mode =='ADJ'
                y = self.y.copy()
                list_X = [th.tensor(X[t].distance_matrix(method=str_to_method['ADJ']), dtype=self.dtype, device=self.device) for t in range(X.shape[0])]
                list_F = [th.tensor(X[t].values(), dtype=self.dtype, device=self.device) for t in range(X.shape[0])]
                N = len(list_X)
                train_idx,test_idx = train_test_split(np.arange(N), test_size=self.completion_parameters['split_rate'], stratify=np.array(y), random_state = self.completion_parameters['split_seed'])
                self.raw_train_graphs, self.raw_test_graphs= [list_X[i] for i in train_idx], [list_X[i] for i in test_idx]
                self.raw_train_features, self.raw_test_features= [list_F[i] for i in train_idx], [list_F[i] for i in test_idx]
                self.train_y, self.test_y= [y[i] for i in train_idx], [y[i] for i in test_idx]
                self.graphs = self.raw_train_graphs
                self.features = [list_F[idx] for idx in train_idx]

                if not self.degrees:
                    #uniform distributions
                    self.masses= [th.ones(C.shape[0], dtype=self.dtype, device=self.device)/C.shape[0] for C in self.raw_train_graphs]

                else:
                    print('computing degree distributions')
                    self.masses =[]
                    for C in self.raw_train_graphs:
                        h = C.sum(axis=0)
                        self.masses.append( h / h.sum())
                print('number of graphs in the train dataset for completion experiments:', len(self.graphs))
        else:# The graphs to learn on are already given
            assert len(graphs) == len(features) == len(masses)
            self.mode = mode
            self.dataset_name = dataset_name
            self.graphs = graphs
            self.features = features
            self.masses = masses
            self.y = y

        self.d= self.features[0].shape[1] # dimension of features

        self.dataset_size = len(self.graphs)
        if self.dataset_name in ['mutag', 'ptc', 'bzr', 'cox2', 'enzymes', 'protein']:
            self.undirected = True
        else:
            self.undirected = th.all([th.all(self.graphs[i] == self.graphs[i].T) for i in range(self.dataset_size)])
            print('All graphs in the dataset are undirected?', self.undirected)

    def init_dictionary(self, seed:int=0, use_checkpoint:bool=True, verbose:bool=False):
        """
            Initialize the graph structure atom following strategies:
                - graph atom ~ N(0.5,0.01)

            use_checkpoint instantiates the placeholder used for tracking with the early stopping scheme.

        """
        np.random.seed(seed)
        x = th.tensor(np.random.normal(loc=0.5, scale=0.01, size=(self.Ntarget,self.Ntarget)), dtype=self.dtype, device=self.device)

        if self.proj in ['nsym','sym']:
            self.Ctarget = (x + x.T) / 2.
        if self.proj == 'nsym':
            self.Ctarget = self.Ctarget.clamp(min=0.)

        #Kmeans on the features to initialize the features of the dictionary
        stacked_features = th.cat(self.features)
        stacked_features = stacked_features.detach().cpu().numpy()

        km = KMeans(n_clusters = self.Ntarget, n_init=10, random_state = seed).fit(stacked_features)
        self.Ftarget = th.tensor(km.cluster_centers_, dtype=self.dtype, device=self.device)

        if use_checkpoint:
            self.checkpoint_Ctarget = []
            self.checkpoint_Ftarget = []

    def init_optimizer(self):
        #Initialization for our numpy implementation of adam optimizer
        self.C_adam_moment1 = th.zeros((self.Ntarget, self.Ntarget), dtype=self.dtype, device=self.device)#Initialize first  moment vector
        self.C_adam_moment2 = th.zeros((self.Ntarget, self.Ntarget), dtype=self.dtype, device=self.device)#Initialize second moment vector
        self.F_adam_moment1 = th.zeros((self.Ntarget, self.d), dtype=self.dtype, device=self.device)#Initialize first  moment vector
        self.F_adam_moment2 = th.zeros((self.Ntarget, self.d), dtype=self.dtype, device=self.device)#Initialize second moment vector
        self.adam_count = 1

    def create_srFGW_operator(self,init_mode:str='product',
                             eps_inner:float=10**(-6),
                             max_iter_inner:int=1000,
                             gamma_entropy:float=0,
                             lambda_reg:float=None,
                             eps_inner_MM:float=10**(-6),
                             max_iter_MM:int=50,
                             use_warmstart_MM:bool=True,
                             seed:int=0):
        """
        Parameters
        ----------
        init_mode : str, optional
            Initialization mode for srGW's conditional gradient solver among ['product','random']. The default is 'product' (i.e h_1.h_2^T) .
        eps_inner : float, optional
            convergence precision used in the CG solver. The default is 10**(-6).
        max_iter_inner : int, optional
            maximum number of iterations for the CG solver if it has not converged yet. The default is 1000.
        gamma_entropy : float, optional
            Entropic parameter used in our Mirror Descent algorithm.
            The default is 0. If set to 0 we use the CG solver otherwise the MD solver.
        lambda_reg : float, optional
            Regularization parameter for promoting sparsity of our embeddings using the MM solver.
            The default is None. If set in [None, 0] we call the CG or MD solvers depending on gamma_entropy.
        eps_inner_MM : float, optional
            convergence precision used in the MM solver (outer loop containing CG or MD solver iteration).
            The default is 10**(-6).
        max_iter_MM : int, optional
            maximum number of iterations for the MM solver if it has not converged yet. The default is 50.
        use_warmstart_MM : bool, optional
            Specify either to reuse the previously computed OT plan in the MM solver outer loop. The default is True.
        seed : int, optional
            Random seed if random initialization is required (init_mode=='random'). The default is 0.

        Returns
        -------
        Instantiate self.srGW_operator :
            the srGW solver function used for the dictionary learning,taking as inputs,

            C1: np.array of size (N,N), input graph.
            h1: np.array of size (N,), input node distribution.
            C2: np.array of size (Nbar,Nbar), graph atom.
            T_init: np.array of size (N,Nbar), optional.
                    Transport plan to initialize the srGW, if set to None, related init parameters of the function will be used.
        """
        # Just create an operator for unmixing step involved in each iteration of the dictionary learning
        if lambda_reg == 0: # This condition means We do not use concave sparsity promoting regularization with MM solver.
            if gamma_entropy ==0:
                self.srFGW_operator = (lambda C1, F1, h1, C2, F2, T_init: srGW.cg_semirelaxed_fused_gromov_wasserstein(C1, F1, h1, C2, F2, self.alpha, self.undirected, init_mode, T_init, False,
                                                                                eps_inner, max_iter_inner, seed, False, self.device, self.dtype))
            else:
                self.srFGW_operator = (lambda C1, F1, h1, C2, F2, T_init: srGW.md_semirelaxed_fused_gromov_wasserstein(C1, F1, h1, C2, F2, gamma_entropy, self.alpha, self.undirected, init_mode, T_init, False,
                                                                                eps_inner, max_iter_inner, seed, False, self.device, self.dtype))
        else:  # We use Majorization-Minimization solver.

            self.srFGW_operator = (lambda C1, F1, h1, C2, F2, T_init: srGW.mm_lpl1_semirelaxed_fused_gromov_wasserstein(C1, F1, h1, C2, F2, self.alpha, gamma_entropy, T_init, init_mode, self.undirected,
                                                                                0.5, lambda_reg, False, use_warmstart_MM, eps_inner, eps_inner_MM, max_iter_inner,
                                                                                max_iter_MM, seed, False, False, self.dtype, self.device))

    def Learn_dictionary(self,
                         lambda_reg:float,
                         max_iter_inner:int,
                         eps_inner:float,
                         lrC:float, lrF:float, batch_size:int, epochs:int, algo_seed:int,
                         max_iter_MM:int=0,
                         eps_inner_MM:float = 0.,
                         use_warmstart_MM:bool=True,
                         gamma_entropy:float=0.,
                         beta_1:float=0.9,
                         beta_2:float=0.99,
                         use_optimizer:bool=True,
                         checkpoint_freq:int=100,
                         earlystopping_patience:int = 2,
                         use_checkpoint:bool = True,
                         proj:str= 'nsym',
                         init_GW:str='product',
                         draw_loss:bool=False):
        """
        Stochastic Algorithm to learn srGW dictionaries
        described in Section 4 of the main paper and Algorithm 2,
        further details in the supplementary material.

        Parameters
        ----------
        lambda_reg : sparse regularization coefficient
        max_iter_inner : maximum number of iterations for the Conditional Gradient algorithm on {wk}
        eps_inner: precision to stop "srGW solver" based on relative variation of the loss
        max_iter_MM : maximum number of iterations for the Majorization minimization algorithm on {wk}
                    > only used if lambda_reg >0.
        eps_inner_MM: precision to stop "srGW MM- solver" based on relative variation of the loss
                    > only used if lambda_reg >0.
        gamma_entropy: regularization coefficient of the negative quadratic regularization on unmixings
        eps : precision to stop our learning process based on relative variation of the loss
        lr : Initial learning rate of Adam optimizer
        batch_size : batch size
        algo_seed : initialization random seed
        OT_loss : GW discrepency ground cost. The default is 'square_loss'.
        beta_1 : Adam parameter on gradient. The default is 0.9.
        beta_2 : Adam parameter on gradient**2. The default is 0.99.
        use_checkpoint : To save dictionary state and corresponding unmixing at different time steps. The default is False.
        verbose : Check the good evolution of the loss. The default is False.
        """
        np.random.seed(algo_seed)
        th.manual_seed(algo_seed)
        if lambda_reg == 0:
            self.settings = {'Ntarget':self.Ntarget, 'max_iter_inner':max_iter_inner,
                             'eps_inner':eps_inner, 'epochs':epochs,
                             'lrC':lrC, 'lrF':lrF, 'batch_size':batch_size,
                             'algo_seed':algo_seed, 'beta1':beta_1, 'beta2':beta_2,'l2_reg':0,'lambda_reg':0,
                             'use_optimizer':use_optimizer,'init_GW':init_GW, 'proj':proj}
        else:
            self.settings = {'Ntarget':self.Ntarget, 'max_iter_FW':max_iter_inner,'eps_inner_FW':eps_inner,'max_iter_MM':max_iter_MM,'eps_inner_MM':eps_inner_MM,
                             'lrC':lrC, 'lrF':lrF, 'batch_size':batch_size,'epochs':epochs, 'algo_seed':algo_seed, 'beta1':beta_1, 'beta2':beta_2,'lambda_reg':lambda_reg,
                             'use_optimizer':use_optimizer,'init_GW':init_GW, 'proj':proj,'use_warmstart_MM':use_warmstart_MM}

        if gamma_entropy !=0:
            self.settings['gamma_entropy'] = gamma_entropy
        self.proj = proj

        self.init_dictionary(algo_seed, use_checkpoint)
        algo_seed=None
        if use_optimizer:
            self.init_optimizer()
        self.create_srFGW_operator(init_mode=init_GW,eps_inner=eps_inner, max_iter_inner=max_iter_inner,
                                    eps_inner_MM=eps_inner_MM, max_iter_MM=max_iter_MM,lambda_reg=lambda_reg,
                                    gamma_entropy=gamma_entropy,use_warmstart_MM=use_warmstart_MM,seed=algo_seed)

        T = len(self.graphs)
        self.log ={'batch_loss':[], 'epoch_loss':[]}

        best_epoch_global_rec = np.inf
        consecutive_global_rec_drops =0 # Used to decide on when to stop learning.

        for epoch in tqdm(range(epochs), desc='epochs'):
            seen_graphs_count = 0
            epoch_global_rec = 0
            while seen_graphs_count < self.dataset_size:
                #batch sampling
                seen_graphs_count+=batch_size
                batch_t = np.random.choice(range(T), size=batch_size, replace=False)
                #print('batch idx:', batch_t)
                best_T = []
                batch_loss = 0
                for k,t in enumerate(batch_t):
                    local_T, local_loss= self.srFGW_operator(self.graphs[t], self.features[t], self.masses[t], self.Ctarget, self.Ftarget, None)
                    best_T.append(local_T)
                    batch_loss += local_loss
                self.log['batch_loss'].append(batch_loss.item())
                epoch_global_rec += batch_loss

                #Stochastic update
                gradC= th.zeros_like(self.Ctarget)
                gradF= th.zeros_like(self.Ftarget)
                for k,t in enumerate(batch_t):
                    hk = best_T[k].sum(0)
                    gradC += self.Ctarget * (hk[:, None] @ hk[None, :]) - (best_T[k].T) @ self.graphs[t] @ best_T[k]
                    gradF += th.diag(hk) @ self.Ftarget - (best_T[k].T) @ self.features[t]
                gradC *= 2. / batch_size
                gradF *= 2. / batch_size
                if not use_optimizer:
                    self.Ctarget -= lrC * gradC
                    self.Ftarget -= lrF * gradF
                else:

                    C_m1_t = beta_1 * self.C_adam_moment1 + (1 - beta_1) * gradC
                    C_m2_t = beta_2 * self.C_adam_moment2 + (1 - beta_2) * (gradC**2)
                    C_m1_t_unbiased = C_m1_t / (1 - beta_1**self.adam_count)
                    C_m2_t_unbiased = C_m2_t / (1 - beta_2**self.adam_count)
                    self.Ctarget -= lrC * C_m1_t_unbiased / (th.sqrt(C_m2_t_unbiased) + 1e-15)
                    F_m1_t = beta_1 * self.F_adam_moment1 + (1-beta_1) * gradF
                    F_m2_t = beta_2 * self.F_adam_moment2 + (1-beta_2) * (gradF**2)
                    F_m1_t_unbiased = F_m1_t / (1-beta_1**self.adam_count)
                    F_m2_t_unbiased = F_m2_t / (1-beta_2**self.adam_count)
                    self.Ftarget -= lrF * F_m1_t_unbiased / (th.sqrt(F_m2_t_unbiased) + 1e-15)

                    self.C_adam_moment1 = C_m1_t
                    self.C_adam_moment2 = C_m2_t
                    self.F_adam_moment1 = F_m1_t
                    self.F_adam_moment2 = F_m2_t
                    self.adam_count +=1
                #projection on nonnegative matrices
                #projection on nonnegative matrices
                if proj == 'nsym':
                    self.Ctarget = th.clamp(self.Ctarget, min=0.)
            self.log['epoch_loss'].append(epoch_global_rec.item())

            if epoch == 0:
                print('saved settings:', self.settings)
                self.save_elements(save_settings=True, use_checkpoint = use_checkpoint)
            elif epoch > 0 and (epoch % checkpoint_freq == 0):
                print('checkpoint_step to evaluate embeddings and decide on early stopping')
                self.save_elements(save_settings=False, use_checkpoint = use_checkpoint)
                if draw_loss:
                    pl.figure(1, (10,5))
                    pl.clf()
                    pl.subplot(121)
                    pl.plot(self.log['batch_loss'])
                    pl.title('loss evolution by batches')
                    pl.xlabel('iterations');pl.ylabel('reconstruction loss')
                    pl.subplot(122)
                    pl.plot(self.log['epoch_loss'])
                    pl.title('loss evolution by epochs')
                    pl.xlabel('iterations');pl.ylabel('reconstruction loss')
                    pl.tight_layout()
                    pl.show()
                _, list_losses = self.compute_unmixing(use_checkpoint = False)
                mean_rec = np.mean(list_losses)
                if mean_rec < best_epoch_global_rec:
                    best_epoch_global_rec = mean_rec
                    consecutive_global_rec_drops = 0
                    print('[unmixings check] epoch:%s / new best epoch global rec :%s'%(epoch, best_epoch_global_rec))
                else:
                    consecutive_global_rec_drops += 1
                    print('[not improved- unmixings check] epoch :%s / current epoch loss :%s / fails:%s '%(epoch, mean_rec, consecutive_global_rec_drops))
                    if consecutive_global_rec_drops > earlystopping_patience:
                        break

    def compute_unmixing(self, use_checkpoint:bool = False):
        """
        Parameters
        ----------
        use_checkpoint : bool, optional. The default is False.
            If set to False, compute unmixings on self.Ctarget (current state).
            Else if set to True, compute unmixings on all dictionary states saved in self.checkpoint_Ctarget.
        Returns
        -------
        If use_checkpoint ==False:
            best_T, best_losses: (list of np.array, list) corresponding to OT providing the unmixings and the corresponding srGW divergences.
        If use_checkpoint ==True:
            list_best_T, list_best_losses: (list of lists of np.array, list of lists) OT and losses for each dictionary state stored in self.checkpoint_Ctarget.
        """
        if not use_checkpoint:
            print('computing srGW unmixings on current dictionary state')
        else:
            print('computing srGW unmixings on all saved dictionary states')

        T= len(self.graphs)

        if not use_checkpoint :
            best_T = []
            best_losses = []
            for t in range(T):
                local_T, local_loss = self.srFGW_operator(self.graphs[t], self.features[t], self.masses[t], self.Ctarget, self.Ftarget, T_init=None)
                best_T.append(local_T)
                best_losses.append(local_loss.item())
            return best_T, best_losses
        else: #ran over all saved dictionary graph state
            list_best_T =[]
            list_best_losses=[]
            for i in range(len(self.checkpoint_Ctarget)):
                local_list_bestT = []
                local_list_bestlosses = []
                for t in range(T):
                    local_T, local_loss = self.srFGW_operator(self.graphs[t], self.features[t], self.masses[t], self.checkpoint_Ctarget[i], self.checkpoint_Ftarget[i], T_init=None)
                    local_list_bestT.append(local_T)
                    local_list_bestlosses.append(local_loss.item())

                list_best_T.append(local_list_bestT)
                list_best_losses.append(local_list_bestlosses)
            return list_best_T, list_best_losses

    def complete_patch(self,
                       patch_graph:np.array,
                       patch_features:np.array,
                       Nfullpatch:int,
                       lrC:float=0.01,
                       lrF:float=0.01,
                       max_iter:int=100,
                       eps:float=10**(-6),
                       proj:str='nsym',
                       algo_seed:int = 0,
                       use_optimizer:bool=True,
                       beta_1:float=0.9,
                       beta_2:float=0.99,
                       use_warmstart:float=False,
                       use_log:bool=False,
                       init_patch_graph:str='random',
                       init_patch_features:str='random_range'):
        np.random.seed(0)
        th.manual_seed(0)
        if not (proj in ['nsym','sym']):
            raise "only proj in ['nsym','sym'] is supported for now"
        if use_log:
            local_log = {'loss':[]}
        else:
            local_log = None

        Npatch = patch_graph.shape[0]
        learnable_mask_graph = th.ones((Nfullpatch, Nfullpatch), dtype=self.dtype, device=self.device)
        learnable_mask_features = th.ones((Nfullpatch, self.d), dtype=self.dtype, device=self.device)
        learnable_mask_graph[:Npatch, :Npatch] = 0.
        learnable_mask_features[:Npatch, :]= 0.

        assert Nfullpatch >= Npatch
        # Initialization of the patch (graph, features)
        ### 1. Dealing with init of the graph part
        if init_patch_graph == 'random':
            np.random.seed(algo_seed)
            #x = np.random.uniform(low=0.1, high=0.9, size=(self.Ntarget,self.Ntarget))
            x = th.tensor( np.random.normal(loc=0.5, scale=0.01, size=(Nfullpatch,Nfullpatch)), dtype=self.dtype, device=self.device)

            completed_patch_graph = (x + x.T) / 2.
            th.fill_diagonal(completed_patch_graph, 0.)# no diagonal as we do not seek for super nodes
        else:
            raise 'unsupported init_patch_graph : %s'%init_patch_graph
        ### 2. Then dealing with init of the features part
        if init_patch_features == 'random_range':
            #get the range of features and generate features randomly in this range
            min_ = patch_features.min(0).detach().numpy()
            max_ = patch_features.max(0).detach().numpy()
            completed_patch_features = th.tensor(np.random.uniform(low=min_, high=max_, size=(Nfullpatch, self.d)), dtype=self.dtype, device=self.device)
        else:
            raise 'unsupported init_patch_features : %s'%init_patch_features

        #Init of Adam Optimizer is use_optimizer is True.
        if use_optimizer:
            graph_adam_moment1 = th.zeros((Nfullpatch, Nfullpatch), dtype=self.dtype, device=self.device)#Initialize first  moment vector
            graph_adam_moment2 = th.zeros((Nfullpatch, Nfullpatch), dtype=self.dtype, device=self.device)#Initialize second moment vector
            features_adam_moment1 = th.zeros((Nfullpatch, self.d), dtype=self.dtype, device=self.device)#Initialize first  moment vector
            features_adam_moment2 = th.zeros((Nfullpatch, self.d), dtype=self.dtype, device=self.device)#Initialize second moment vector
            adam_count = 1

        completed_patch_masses = th.ones(Nfullpatch, dtype=self.dtype, device=self.device)/Nfullpatch
        weight_mask_graph = completed_patch_masses[:, None] @ completed_patch_masses[None,:]
        weight_mask_features = th.diag(completed_patch_masses)
        completed_patch_graph[:Npatch, :Npatch] = patch_graph
        completed_patch_features[:Npatch, :] = patch_features
        init_completed_patch_graph = completed_patch_graph.clone()
        init_completed_patch_features = completed_patch_features.clone()

        curr_loss = th.tensor(1e15, dtype=self.dtype, device=self.device)
        best_loss = np.inf
        best_completed_patch_graph = completed_patch_graph.clone()
        best_completed_patch_features = completed_patch_features.clone()
        convergence_criterion = np.inf
        count = 0
        T_init = None

        while (convergence_criterion >= eps) and (count < max_iter):
            prev_loss = curr_loss.clone()
            #print('count :%s /curr_loss : %s'%(count,curr_loss))
            # compute transport between completed patch and dictionary
            local_OT, curr_loss = self.srFGW_operator(completed_patch_graph, completed_patch_features, completed_patch_masses,
                                               self.Ctarget, self.Ftarget, T_init)
            if use_warmstart:
                T_init = local_OT.clone()

            if curr_loss < best_loss:
                best_loss  = curr_loss.item()
                best_completed_patch_graph = completed_patch_graph.clone()
                best_completed_patch_features = completed_patch_features.clone()

            # compute gradient to update the completed_patch
            if not use_optimizer:
                completed_patch_graph -= 2 *self.alpha * lrC * learnable_mask_graph * (completed_patch_graph *  weight_mask_graph - local_OT @ self.Ctarget @ local_OT.T)
                completed_patch_features -= 2 * (1 - self.alpha) * lrF * learnable_mask_features * (weight_mask_features @ completed_patch_features - local_OT @ self.Ftarget)

            else:
                graph_grad = 2 * self.alpha * learnable_mask_graph *(completed_patch_graph * weight_mask_graph - local_OT @ self.Ctarget @ local_OT.T)
                graph_m1_t = beta_1 * graph_adam_moment1 + (1 - beta_1) * graph_grad
                graph_m2_t = beta_2 * graph_adam_moment2 + (1 - beta_2) * (graph_grad**2)
                graph_m1_t_unbiased = graph_m1_t / (1 - beta_1**adam_count)
                graph_m2_t_unbiased = graph_m2_t /(1 - beta_2**adam_count)
                completed_patch_graph -= lrC * graph_m1_t_unbiased / (th.sqrt(graph_m2_t_unbiased)+1e-15)
                graph_adam_moment1 = graph_m1_t
                graph_adam_moment2 = graph_m2_t

                features_grad = 2 * (1 - self.alpha) * learnable_mask_features * (weight_mask_features @ completed_patch_features - local_OT @ self.Ftarget)
                features_m1_t = beta_1 * features_adam_moment1 + (1 - beta_1) * features_grad
                features_m2_t = beta_2 * features_adam_moment2 + (1 - beta_2) * (features_grad**2)
                features_m1_t_unbiased = features_m1_t / (1 - beta_1**adam_count)
                features_m2_t_unbiased = features_m2_t / (1 - beta_2**adam_count)
                completed_patch_features -= lrF * features_m1_t_unbiased / (th.sqrt(features_m2_t_unbiased)+1e-15)
                features_adam_moment1 = features_m1_t
                features_adam_moment2 = features_m2_t
                adam_count += 1

            if proj == 'nsym':
                completed_patch_graph[completed_patch_graph < 0.] = 0.
            if prev_loss.item() != 0:
                convergence_criterion = abs(prev_loss.item() -curr_loss.item())/ abs(prev_loss.item())
            else:
                convergence_criterion = abs(prev_loss.item() -curr_loss.item())/ abs(prev_loss.item() + 1e-15)
            count+=1

            if use_log:
                local_log['loss'].append(curr_loss.item())

        return best_completed_patch_graph, best_completed_patch_features, best_loss, local_log, init_completed_patch_graph, init_completed_patch_features

    def save_elements(self, save_settings=False, use_checkpoint = False):
        """
            DL saver used while learning the srGW graph atom. [ TO DO WITH PICKLE]
        """
        path = os.path.abspath('../')+self.experiment_repo
        print('path',path)
        if not os.path.exists(path+self.experiment_name):
            os.makedirs(path+self.experiment_name)
            print('made dir', path+self.experiment_name)
        if not use_checkpoint:
            pickle.dump(self.Ctarget, open(path+'%s/Ctarget.pkl'%self.experiment_name,'wb'))
            pickle.dump(self.Ftarget, open(path+'%s/Ftarget.pkl'%self.experiment_name,'wb'))

        else:
            self.checkpoint_Ctarget.append(self.Ctarget.clone())
            self.checkpoint_Ftarget.append(self.Ftarget.clone())
            print('#checkpoints:', len(self.checkpoint_Ctarget))

            pickle.dump(self.checkpoint_Ctarget, open(path+'%s/checkpoint_Ctarget.pkl'%self.experiment_name,'wb'))
            pickle.dump(self.checkpoint_Ftarget, open(path+'%s/checkpoint_Ftarget.pkl'%self.experiment_name,'wb'))

        for key in self.log.keys():
            np.save(path+'%s/%s.npy'%(self.experiment_name,key), np.array(self.log[key]))

        if save_settings:
            pd.DataFrame(self.settings, index=self.settings.keys()).to_csv(path+'%s/settings'%self.experiment_name)

    def load_elements(self, use_checkpoint=False):
        """
            DL loader for analysing learned dictionaries.
        """
        path = os.path.abspath('../')+self.experiment_repo
        if not use_checkpoint:
            self.Ctarget = pickle.load(open(path+'%s/Ctarget.pkl'%self.experiment_name, 'rb'))
            self.Ftarget = pickle.load(open(path+'%s/Ftarget.pkl'%self.experiment_name, 'rb'))
            self.Ntarget = self.Ctarget.shape[0]
        else:
            self.checkpoint_Ctarget = pickle.load(open(path+'%s/checkpoint_Ctarget.pkl'%self.experiment_name, 'rb'))
            self.checkpoint_Ftarget = pickle.load(open(path+'%s/checkpoint_Ftarget.pkl'%self.experiment_name, 'rb'))
            self.Ntarget = self.checkpoint_Ctarget[0].shape[-1]

ext: py

max_stars_repo_name: cedricvincentcuaz/srGW

max_issues_repo_name: cedricvincentcuaz/srGW

max_forks_repo_name: cedricvincentcuaz/srGW

reason: import numpy


text:
import argparse

import numpy as np
import sklearn
from parallelm.mlops import mlops as mlops
# use below import if user wants to user RegressionMetrics predefined metrics names.
from parallelm.mlops.metrics_constants import RegressionMetrics
from parallelm.mlops.stats.bar_graph import BarGraph
from sklearn.datasets import make_regression
from sklearn.svm import SVR

def parse_args():
    """
    Parse Arguments from component
    :return:
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("--num_samples", help="# samples")
    parser.add_argument("--num_features", help="# features")

    parser.add_argument("--kernel", help="Kernel")
    parser.add_argument("--degree", help="Degree")
    parser.add_argument("--gamma", help="Gamma")
    parser.add_argument("--tol", help="Tol")
    parser.add_argument("--max_iter", dest="max_iter", type=int, required=False, default=100,
                        help='Maximum number of iterations')
    parser.add_argument("--output-model", help="Data file to save model")
    options = parser.parse_args()
    return options

def main():
    pm_options = parse_args()

    print("PM: Configuration:")
    print("PM: # Sample:                    [{}]".format(pm_options.num_samples))
    print("PM: # Features:                  [{}]".format(pm_options.num_features))

    print("PM: Kernel:                      [{}]".format(pm_options.kernel))
    print("PM: Degree:                      [{}]".format(pm_options.degree))
    print("PM: Gamma:                       [{}]".format(pm_options.gamma))
    print("PM: Tolerance:                   [{}]".format(pm_options.tol))
    print("PM: Maximum iterations:          [{}]".format(pm_options.max_iter))

    print("PM: Output model:                [{}]".format(pm_options.output_model))

    # Initialize MLOps Library
    mlops.init()

    num_samples = int(pm_options.num_samples)
    num_features = int(pm_options.num_features)

    # Create synthetic data using scikit learn
    X, y = make_regression(n_samples=num_samples,
                           n_features=num_features,
                           n_informative=2,
                           random_state=42)

    # for making labels all positive
    y = y + -1 * np.min(y)

    # Separate into features and labels
    features = X
    labels = y

    # Add noise to the data
    noisy_features = np.random.uniform(0, 10) * \
                     np.random.normal(0, 1,
                                      (num_samples, num_features))
    features = features + noisy_features

    # Output Health Statistics to MCenter
    # MLOps API to report the distribution statistics of each feature in the data
    mlops.set_data_distribution_stat(features)

    hist, bin_edges = np.histogram(labels)

    # Output label distribution as a BarGraph using MCenter
    bar = BarGraph().name("User Defined: Label Distribution") \
        .cols((bin_edges).astype(str).tolist()) \
        .data((hist).tolist()) \
        .as_continuous()

    mlops.set_stat(bar)

    # Create a model that should be deployed into production
    final_model = SVR(kernel=pm_options.kernel,
                      degree=int(pm_options.degree),
                      gamma=str(pm_options.gamma),
                      tol=float(pm_options.tol),
                      max_iter=int(pm_options.max_iter))

    final_model.fit(features, labels)

    labels_pred = final_model.predict(features)
    hist_pred, bin_edges_pred = np.histogram(labels_pred)

    # Output prediction label distribution as a BarGraph using MCenter
    pred_label_bar = BarGraph().name("User Defined: Prediction Label Distribution") \
        .cols((bin_edges_pred).astype(str).tolist()) \
        .data((hist_pred).tolist()) \
        .as_continuous()

    mlops.set_stat(pred_label_bar)

    y_true = labels
    y_pred = labels_pred

    # Regression Metrics
    ##########################################################################
    #################### Start: Output Explained Variance ####################
    ##########################################################################

    evs = sklearn.metrics.explained_variance_score(y_true, y_pred)

    #################### OLD WAY ####################
    # First Way
    # mlops.set_stat("User Defined: Explained Variance ", evs)
    #################### DONE OLD WAY ####################

    #################### NEW WAY ####################
    # Second Way
    mlops.set_stat(RegressionMetrics.EXPLAINED_VARIANCE_SCORE, evs)

    # OR

    # Third Way
    mlops.metrics.explained_variance_score(y_true=labels, y_pred=labels_pred)
    #################### DONE NEW WAY ####################

    ########################################################################
    #################### End: Output Explained Variance ####################
    ########################################################################

    ######################################################################
    #################### Start: Output Mean Abs Error ####################
    ######################################################################

    mae = sklearn.metrics.mean_absolute_error(y_true, y_pred)

    #################### OLD WAY ####################
    # First Way
    # mlops.set_stat("User Defined: Mean Abs Error", mae)
    #################### DONE OLD WAY ####################

    #################### NEW WAY ####################
    # Second Way
    mlops.set_stat(RegressionMetrics.MEAN_ABSOLUTE_ERROR, mae)

    # OR

    # Third Way
    mlops.metrics.mean_absolute_error(y_true=labels, y_pred=labels_pred)
    #################### DONE NEW WAY ####################

    ####################################################################
    #################### End: Output Mean Abs Error ####################
    ####################################################################

    ##########################################################################
    #################### Start: Output Mean Squared Error ####################
    ##########################################################################

    mse = sklearn.metrics.mean_squared_error(y_true, y_pred)

    #################### OLD WAY ####################
    # First Way
    # mlops.set_stat("User Defined: Mean Squared Error", mse)
    #################### DONE OLD WAY ####################

    #################### NEW WAY ####################
    # Second Way
    mlops.set_stat(RegressionMetrics.MEAN_SQUARED_ERROR, mse)

    # OR

    # Third Way
    mlops.metrics.mean_squared_error(y_true=labels, y_pred=labels_pred)
    #################### DONE NEW WAY ####################

    ########################################################################
    #################### End: Output Mean Squared Error ####################
    ########################################################################

    ##############################################################################
    #################### Start: Output Mean Squared Log Error ####################
    ##############################################################################

    msle = sklearn.metrics.mean_squared_log_error(y_true, y_pred)

    #################### OLD WAY ####################
    # First Way
    # mlops.set_stat("User Defined: Mean Squared Log Error", msle)
    #################### DONE OLD WAY ####################

    #################### NEW WAY ####################
    # Second Way
    mlops.set_stat(RegressionMetrics.MEAN_SQUARED_LOG_ERROR, msle)

    # OR

    # Third Way
    mlops.metrics.mean_squared_log_error(y_true=labels, y_pred=labels_pred)

    #################### DONE NEW WAY ####################

    ############################################################################
    #################### End: Output Mean Squared Log Error ####################
    ############################################################################

    ########################################################################
    #################### Start: Output Median Abs Error ####################
    ########################################################################

    median_ae = sklearn.metrics.median_absolute_error(y_true, y_pred)

    #################### OLD WAY ####################
    # First Way
    # mlops.set_stat("User Defined: Median Abs Error", median_ae)
    #################### DONE OLD WAY ####################

    #################### NEW WAY ####################
    # Second Way
    mlops.set_stat(RegressionMetrics.MEDIAN_ABSOLUTE_ERROR, median_ae)

    # OR

    # Third Way
    mlops.metrics.median_absolute_error(y_true=labels, y_pred=labels_pred)
    #################### DONE NEW WAY ####################

    ######################################################################
    #################### End: Output Median Abs Error ####################
    ######################################################################

    ################################################################
    #################### Start: Output R2 Score ####################
    ################################################################

    r2_s = sklearn.metrics.r2_score(y_true, y_pred)

    #################### OLD WAY ####################
    # First Way
    # mlops.set_stat("User Defined: R2 Score", r2_s)
    #################### DONE OLD WAY ####################

    #################### NEW WAY ####################
    # Second Way
    mlops.set_stat(RegressionMetrics.R2_SCORE, r2_s)

    # OR

    # Third Way
    mlops.metrics.r2_score(y_true=labels, y_pred=labels_pred)
    #################### DONE NEW WAY ####################

    ##############################################################
    #################### End: Output R2 Score ####################
    ##############################################################

    # Save the model
    import pickle
    model_file = open(pm_options.output_model, 'wb')
    pickle.dump(final_model, model_file)
    model_file.close()
    # Terminate MLOPs
    mlops.done()

if __name__ == "__main__":
    main()

max_stars_repo_name: mlpiper/mlhub

max_issues_repo_name: mlpiper/mlhub

max_forks_repo_name: mlpiper/mlhub


text:
#!/usr/bin/python
# -*- coding: utf-8 -*-
from eospac import eV2K_cst
from eospac.rh import RankineHugoniot
import numpy as np
from numpy.testing import assert_allclose
import os
from nose.plugins.skip import Skip, SkipTest

def setup():
    global rh_options
    rh_options = dict(backend='eospac',
            eos_pars={'.*': {'create_tzero': True, 'linear': False}},
            root_opts={'method': 'hybr'},
                  material=3720)

def test_consitency():
    """Checking that computing RH by different methods gives the
    same result """
    state0 = {"temp": 300, "rho": 2.7}
    #state1 = {"rho": np.linspace(1, 4, 100)*state0['rho']}
    state1 = {"pres": np.linspace(1e2, 2.5e3, 100)} # GPa # inital conditions so it doesn't fail...

    res1 = RankineHugoniot.solve(state0, state1, **rh_options)
    state1 = {"rho": res1['rho1']}
    res2 = RankineHugoniot.solve(state0, state1, **rh_options)

    for key in res1:
        yield assert_allclose, res1[key], res2[key]

def test_compare_with_original():
    """Compare with original SESAME implementation by Kerley """
    raise SkipTest
    basedir = os.path.dirname(os.path.realpath(__file__))
    f = np.loadtxt(os.path.join(basedir, 'hugoniot_data/hugoniot_3720.txt'))
    state0 = {'rho': 2.7, 'temp': 300, 'pres': 9.05001730E-03}
    res_ref = {}
    for idx, key in enumerate(['rho1', 'pres1', 'eint1', 'temp1', 'u_s', 'u_p']):
        res_ref[key] = f[:,idx]
    res_ref['temp1'] *= eV2K_cst
    res_ref['pres1'] = res_ref['pres1']*100 + state0['pres'] # to Mbar -> GPa

    state1 = {'pres': res_ref['pres1']}
    res = RankineHugoniot.solve(state0, state1, **rh_options)
    for key in res_ref:
        yield assert_allclose, res_ref[key], np.array(res[key]), 1.2e-2
        # data is constistant to up to 12%
        # the difference might come from convergence criteria or some
        # numeric. But that's good enough.

max_stars_repo_name: ojschumann/pyeospac

max_issues_repo_name: ojschumann/pyeospac

max_forks_repo_name: ojschumann/pyeospac

reason: import numpy,from numpy


text:
"""
t1vsflux.py

Authors: Daniel Weiss, Helin Zhang
"""

import mpmath
import itertools
import os

import scqubits as qubit
import matplotlib
matplotlib.rcParams['text.usetex'] = False
matplotlib.rcParams.update({'font.size': 22})
import matplotlib.pyplot as plt
import numpy as np
import scipy.linalg as la
import scipy.constants as const

from qutip import *
from scqubits.utils.spectrum_utils import get_matrixelement_table, matrix_element
from scqubits.utils.sweep_plotting import bare_spectrum, dressed_spectrum, difference_spectrum
from scqubits import HilbertSpace, InteractionTerm, ParameterSweep

EXPERIMENT_META = "spin"
WDIR = os.environ["ROBUST_QOC_PATH"]
OUT_PATH = os.path.join(WDIR, "out", EXPERIMENT_META)
SAVE_FILE_PATH = os.path.join(OUT_PATH, "t1vsflux.png")

def main():
    # (1/(2pi hbar 10^9))
    # J_to_GHz = 1. / (2. * np.pi * const.hbar * (10**9))
    # GHz_to_J = 1./J_to_GHz

    # resonator = qubit.Oscillator(
    #     E_osc = 5.7286,
    #     truncated_dim=5
    # )

    fluxonium = qubit.Fluxonium(
        EJ = 3.395,
        EL = 0.132,
        EC = 0.479,
        flux = 0.5,
        cutoff = 110,
        truncated_dim = 60
    )

    # hilbertspc = HilbertSpace([fluxonium, resonator])

    # adag = resonator.creation_operator()
    # a = resonator.annihilation_operator()

    # int_term = InteractionTerm(
    #     g_strength = 0.076,
    #     subsys1 = fluxonium,
    #     op1 = fluxonium.n_operator(),
    #     subsys2 = resonator,
    #     op2 = a+adag
    # )

    # interaction_list = [int_term]
    # hilbertspc.interaction_list = interaction_list

    # dressed_hamiltonian = hilbertspc.get_hamiltonian()

    # evals, evecs = dressed_hamiltonian.eigenstates(eigvals=10)

    # # name of varying external parameter
    # param_name = 'flux'
    # # parameter values
    # param_vals = np.linspace(0.25, 0.5, 100)
    # # list of HilbertSpace subsystems which are affected by parameter changes
    # subsys_update_list = [fluxonium]

    # # function that shows how Hilbert space components are updated
    # def update_hilbertspace(param_val):
    #     fluxonium.flux = param_val

    # sweep = ParameterSweep(
    #     param_name=param_name,
    #     param_vals=param_vals,
    #     evals_count=50,
    #     hilbertspace=hilbertspc,
    #     subsys_update_list=subsys_update_list,
    #     update_hilbertspace=update_hilbertspace
    # )

    # dressed_spectrum(sweep, xlim=(0.3,0.4),ylim=(0,7))

    # temp=0.05

    # def nthermal(omega):
    #     return(1./(np.exp(omega*GHz_to_J/(const.k*temp))-1.))

    # def Pthermalprob(n, resonator):
    #     return((1-np.exp(-resonator.E_osc*GHz_to_J/(const.k*temp)))
    #            *np.exp(-n*resonator.E_osc*GHz_to_J/(const.k*temp)))

    # def gamma_purcell_sweep_old(i, j, sweep, kappa, k, flux_val):
    #     #i -> j
    #     evals_dressed, evecs_dressed = (sweep.dressed_specdata.energy_table[k],
    #                                     sweep.dressed_specdata.state_table[k])
    #     fluxonium.flux = flux_val
    #     evals_fluxonium, evecs_fluxonium = fluxonium.eigensys(evals_count=20)
    #     gammap = 0.0
    #     if (evals_fluxonium[i] > evals_fluxonium[j]):
    #         #n1 ranges over initial states of the resonator
    #         for n1 in range(2):
    #             #n2 over final
    #             for n2 in range(2):
    #                 indexin1 = sweep.lookup.dressed_index((i,n1), param_index=k)
    #                 indexjn2 = sweep.lookup.dressed_index((j,n2), param_index=k)
    #                 gammap += (Pthermalprob(n1, resonator)
    #                            * np.abs(matrix_element(
    #                                evecs_dressed[indexjn2],
    #                                sweep.hilbertspace.identity_wrap(
    #                                    resonator.annihilation_operator(),
    #                                    resonator), evecs_dressed[indexin1]))**2)
    #         omegaij = evals_fluxonium[i]-evals_fluxonium[j]
    #         return(10**9 * kappa * (nthermal(omegaij) + 1) * gammap)
    #     else:
    #         #n1 ranges over initial states of the resonator
    #         for n1 in range(2):
    #             #n2 over final
    #             for n2 in range(2):
    #                 indexin1 = sweep.lookup.dressed_index((i,n1), param_index=k)
    #                 indexjn2 = sweep.lookup.dressed_index((j,n2), param_index=k)
    #                 gammap += (Pthermalprob(n1, resonator)
    #                             *np.abs(matrix_element(
    #                                 evecs_dressed[indexjn2],
    #                                 sweep.hilbertspace.identity_wrap(
    #                                     resonator.creation_operator(),
    #                                     resonator),evecs_dressed[indexin1]))**2)
    #         omegaij = evals_fluxonium[j]-evals_fluxonium[i]
    #         return(10**9 * kappa * (nthermal(omegaij)) * gammap)
    #     #ENDDEF

    # def gamma_purcell_sweep(i, j, sweep, kappa, k, flux_val):
    #     #i -> j
    #     evals_dressed, evecs_dressed = (sweep.dressed_specdata.energy_table[k],
    #                                     sweep.dressed_specdata.state_table[k])
    #     fluxonium.flux = flux_val
    #     evals_fluxonium, evecs_fluxonium = fluxonium.eigensys(evals_count=20)
    #     gammap = 0.0
    #     if (evals_fluxonium[i] > evals_fluxonium[j]):
    #         for n1 in range(2):
    #             for n2 in range(2):
    #                 indexin1 = sweep.lookup.dressed_index((i,n1), param_index=k)
    #                 indexjn2 = sweep.lookup.dressed_index((j,n2), param_index=k)
    #                 omegaij = np.abs(evals_dressed[indexin1]-evals_dressed[indexjn2])
    #                 gammap += (Pthermalprob(n1, resonator) * (nthermal(omegaij) + 1)
    #                            *np.abs(matrix_element(
    #                                evecs_dressed[indexjn2],
    #                                hilbertspc.identity_wrap(resonator.annihilation_operator(),
    #                                                         resonator),evecs_dressed[indexin1]))**2)
    #         return(10**9 * kappa * gammap)
    #     else:
    #         for n1 in range(2):
    #             for n2 in range(2):
    #                 indexin1 = sweep.lookup.dressed_index((i,n1), param_index=k)
    #                 indexjn2 = sweep.lookup.dressed_index((j,n2), param_index=k)
    #                 omegaij = np.abs(evals_dressed[indexin1]-evals_dressed[indexjn2])
    #                 gammap += (Pthermalprob(n1, resonator) * (nthermal(omegaij))
    #                             *np.abs(matrix_element(
    #                                 evecs_dressed[indexjn2],
    #                                 hilbertspc.identity_wrap(resonator.creation_operator(),
    #                                                          resonator),evecs_dressed[indexin1]))**2)
    #     return(10**9 * kappa * gammap)
    # #ENDDEF

    # gamma_purcell_list = np.zeros(len(param_vals))
    # gamma_purcell_0up = np.zeros_like(gamma_purcell_list)
    # gamma_purcell_1up = np.zeros_like(gamma_purcell_list)
    # for k, flux_val in enumerate(param_vals):
    #     gamma_purcell_list[k] = (gamma_purcell_sweep(0,1,sweep,0.00165, k, flux_val)
    #                              + gamma_purcell_sweep(1,0,sweep,0.00165, k, flux_val))
    #     #j ranges over higher fluxonium levels
    #     gamma_purcell_0up[k] = np.sum([gamma_purcell_sweep(0,j,sweep,0.00165,k,flux_val) for j in range(1,15)])
    #     gamma_purcell_1up[k] = np.sum([gamma_purcell_sweep(1,j,sweep,0.00165,k,flux_val) for j in range(2,15)])
    # #ENDFOR

    # blist = np.linspace(0.25, 0.5, 100)
    blist = [0.0, 0.1, 0.5, 1.0]
    spect = fluxonium.get_spectrum_vs_paramvals('flux', blist, evals_count=2, subtract_ground=True)
    spect_table = spect.energy_table.T[1]*2*np.pi*1e9
    wq = spect_table

    levls = 8
    philist = []
    nlist = []
    sinphilist = []
    for ii in blist:
        fluxonium.flux = ii
        phi = abs((fluxonium.matrixelement_table('phi_operator', evals_count=levls)[0, 1]))**2
        n =  abs((fluxonium.matrixelement_table('n_operator', evals_count=levls)[0, 1]))**2
    #     sinphi =  abs((fluxonium.matrixelement_table('sinphi_operator', evals_count=levls)[0, 1]))**2
        philist.append(phi)
        nlist.append(n)
    #     sinphilist.append(sinphi)
    philist = np.array(philist)
    nlist = np.array(nlist)
    # sinphilist = np.array(sinphilist)

    T_cap = 0.042
    EC = 0.479e9*const.h
    Gamma_cap = const.hbar * wq * wq / np.tanh(const.hbar*wq/const.k/T_cap/2) * philist / (8 * EC/(9e-6))

    gamma = Gamma_cap
    t1 = 1 / gamma
    print("t1:\n{}"
          "".format(t1))
    exit(0)

    # # fig, ax = plt.subplots(1,1, figsize=(10,5))
    # # ax.set_yscale('log')
    # # ax.plot(blist,1/Gamma_cap*1e6)
    # print(1/Gamma_cap[-1]*1e6)

    # T_ind = 0.042
    # EL = 0.132e9*const.h
    # Gamma_ind = 2*EL/const.hbar/np.tanh(const.hbar*wq/const.k/T_ind/2)*philist/(6e8)

    # # fig, ax = plt.subplots(1,1, figsize=(12,5))
    # # ax.set_yscale('log')
    # # ax.plot(blist,1/Gamma_ind*1e6)
    # print(1/Gamma_ind[-1]*1e6)

    # t1_cap = 1/Gamma_cap
    # t1_ind = 1/Gamma_ind
    # T1total = 1/(Gamma_cap + Gamma_ind + gamma_purcell_list + gamma_purcell_0up + gamma_purcell_1up)
    # T1purcell = 1./(gamma_purcell_list + gamma_purcell_0up + gamma_purcell_1up)
    # T1purcell01 = 1./gamma_purcell_list
    # T1purcell0up = 1./gamma_purcell_0up
    # T1purcell1up = 1./gamma_purcell_1up
    # #T1purcellperturb = 1./gamma_purcell_perturb_list
    # fig=plt.figure(figsize=(12,5))
    # plt.yscale("log")
    # plt.plot(param_vals, T1total, 'k-')
    # plt.plot(param_vals, T1purcell, 'b:')
    # # plt.plot(param_vals, T1purcell01, color="yellow")
    # # plt.plot(param_vals, T1purcell0up, color="red")
    # # plt.plot(param_vals, T1purcell1up, color="green")

    # # plt.plot(param_vals, t1_charge, color="orange")
    # plt.plot(param_vals, t1_cap, "r--")
    # plt.plot(param_vals, t1_ind, "g-.")
    # #plt.plot(param_vals, T1purcellperturb, color="purple")
    # errlist = 10**(-6)*np.array([78.034, 70.57, 20.71, 20.93, 12.92, 66.93, 137.06, 319.19,
    #                  194.84, 146.87, 170.584, 374.582, 322.533, 125.10, 105.987, 58.612, 22.295, 11.29])
    # t1_data = 10**(-6)*np.array([1597.923, 1627.93, 301.86, 269.03, 476.33, 1783.19, 2131.76, 2634.50,
    #                    4364.68, 2587.82, 1661.915, 1794.468, 2173.88, 1188.83,
    #                    1576.493, 965.183, 560.251, 310.88])
    # t1_flux = np.array([0.26, 0.28, 0.32, 0.34, 0.36, 0.38, 0.4, 0.42, 0.44, 0.46, 0.465,
    #                    0.47, 0.475, 0.48, 0.484, 0.488, 0.492, 0.5])
    # plt.errorbar(t1_flux,t1_data,errlist,fmt=' ',ecolor='k')
    # plt.plot(t1_flux,t1_data,'ko')
    # # plt.legend(("T1purcell","T1purcell01","T1purcell0up","T1purcell1up",
    # #             "T1total","t1_charge","t1_cap","t1_ind","data"), fontsize=12)
    # plt.legend(("T1total","T1purcell","t1_cap","t1_ind","data"), fontsize=12)
    # plt.ylim(10**(-4),2.2*10**(-2))
    # plt.xlim(0.245,0.505)
    # # plt.xlabel("flux")
    # # plt.ylabel("s")
    # # plt.show()
    # fig.savefig(SAVE_FILE_PATH)
#ENDDEF

if __name__ == "__main__":
    main()

max_stars_repo_name: SchusterLab/rbqoc

max_issues_repo_name: SchusterLab/rbqoc

max_forks_repo_name: SchusterLab/rbqoc

reason: import numpy,import scipy,import mpmath


text:
import pandas as pd
import numpy as np
from random import choices, sample
import yaml
from FRETboard.helper_functions import numeric_timestamp

# FRETBOARD MODEL TEMPLATE
# Using this template you can create your own FRETboard model. For a working model, you need to implement at most 4 steps:
# 1: a model generation/training procedure
# 2: a predict function
# 3: Optionally, parameter return functions
# 4: Optinally, model saving and loading functions
#
# You are completely free to implement any model, however as FRETboard is meant to rapidly estimate and re-estimate models
# to react to the user's input we advise to keep your model implementations light, e.g. fully training a neural network
# from the ground up might work but will probably slow down your analysis a lot.
#
# If you have written a model that you think might be useful for others, please consider submitting it as a pull request
# on the project github page:
#
# https://github.com/cvdelannoy/FRETboard.git)
#
# or get in touch with me:
#
# carlos.delannoy@wur.nl

class Classifier(object):
    """ A template for your own FRETboard-model!
    """
    def __init__(self, nb_states, data, **kwargs):
        """
        :param nb_states: number of states to detect
        :param data: object of class MainTable
        :param gui: Gui object [optional]
        :param buffer: int, size of buffer area around regular classes [required if gui not given]
        """
        self.trained = None
        self.framerate = None
        self.timestamp = numeric_timestamp()
        self.nb_threads = kwargs.get('nb_threads', 8)
        self.feature_list = kwargs['features']
        self.supervision_influence = kwargs['supervision_influence']

        self.nb_states = nb_states
        self.data = data

    # --- training ---
    def train(self, data_dict, supervision_influence=1.0):
        """
        Generate trained hmm and predict examples
        """
        self.supervision_influence = supervision_influence
        self.trained = self.get_trained_model(data_dict)
        self.timestamp = numeric_timestamp()

    def get_trained_model(self, data_dict, bootstrap=False):
        """
        STEP 1:
        Write a function that returns your trained model. The model should be trained on the traces provided
        in data_dict. If bootstrap is provided your model should be trained on a bootstrapped sample of data_dict.
        Don' worry about that though, a general implementation of bootstrapping that modifies the data_dict variable
        is already there waiting for you!

        :param data_dict: dict of traces. Keys are trace ids, values are pandas dataframes of shape
                          [nb_timepoints, nb_features + time]
        :param bootstrap: bool, if true, train on bootstrapped sample of data_dict instead
        :return: trained model
        """
        if self.framerate is None:
            self.framerate = 1 / np.concatenate([data_dict[tr].time.iloc[1:].to_numpy() -
                                                 data_dict[tr].time.iloc[:-1].to_numpy() for tr in data_dict]).mean()

        idx_list = list(data_dict)
        nb_labeled = self.data.manual_table.is_labeled.sum()
        nb_unlabeled = len(idx_list) - nb_labeled
        unlabeled_idx = [idx for idx in idx_list if idx not in self.data.manual_table.query('is_labeled').index]

        # Take bootstrapped /subsampled sample
        if bootstrap:
            labeled_seqs = [si for si in self.data.manual_table.query('is_labeled').index if si in data_dict]
            labeled_seqs = choices(labeled_seqs, k=nb_labeled)
            if nb_unlabeled <= 100:
                # bootstrap size m == n
                unlabeled_seqs = choices(unlabeled_idx, k=nb_unlabeled)
            else:
                # subsampling, m = 100
                unlabeled_seqs = sample(unlabeled_idx, k=100)
            seq_idx = labeled_seqs + unlabeled_seqs
            data_dict = {si: data_dict[si] for si in seq_idx}
        elif nb_unlabeled > 100:
            labeled_seqs = self.data.manual_table.query('is_labeled').index.to_list()
            unlabeled_seqs = sample(unlabeled_idx, k=100)
            seq_idx = labeled_seqs + unlabeled_seqs
            data_dict = {si: data_dict[si] for si in seq_idx}

        # Your implementation goes here

        return my_awesome_model

    def predict(self, trace_df, model=None):
        """
        STEP 2:
        implement your predict function. Given a pandas data frame containing the trace (trace_df), it should return
        a numpy array of integers indicating the class for each time point, and the log-probability of the classification
        being correct as a single float (log_prob). The latter will be reported in the FRETboard GUI in a histogram,
        however if you're just experimenting or your classifier can't generate such a measure for some reason, just
        provide a 0. Predictions should be made with the model in self.trained, but if a model is provided as a parameter
        (see below), please allow for this model to be used instead, to ensure proper functioning of bootstrapping
        confidence intervals.

        :param trace_df: a pandas dataframe containing a single trace of shape [nb_timepoints, nb_features + time]
        :param model: a model that is structured and works exactly as self.trained. Added for bootstrapping functionality.
        :returns: a 1D numpy array of states of length nb_timepoints, a single float
        """
        if model is None: model = self.trained

        # Your implementation goes here

        return state_list, logprob

    # --- parameters and performance measures ---
    def get_data_tm(self, trace_dict, nb_bootstrap_iters):
        """
        STEP 3.1 (OPTIONAL):
        To provide a confidence interval for transition rates FRETboard can train nb_bootstrap_iters models,
        classify data each time and derive transition rates from the new classification, as implemented here. For some
        types of models it makes more sense to derive transition matrices some other way; for example hidden markov
        models already estimates a transition matrix to work with. If this is also the case for you, you can
        change the implementation in this method to make the bootstrapping a lot faster or better.

        :param trace_dict:  dict of traces. Keys are trace ids, values are pandas dataframes of shape
                          [nb_timepoints, nb_features + time]
        :param nb_bootstrap_iters: integer number of bootstrap iterations to perform
        :returns tm: a 2D numpy array of size [nb_classes, nb_classes] containing transition rates estimated on final model
        :returns ci_mat: a 3D numpy array of size [nb_classes, nb_classes, 2]. [nb_classes, nb_classes 0] contains
                         lower bounds for CI for transition rates, [nb_classes, nb_classes, 1] contains upper bounds.
        """

        # Ensure junk traces didn't end up in trace_dict
        invalid_indices = [idx for idx, tup in self.data.manual_table.iterrows() if tup.is_labeled or tup.is_junk]
        idx_list = [idx for idx in self.data.index_table.index if idx not in invalid_indices]
        idx_list = np.concatenate((idx_list, self.data.manual_table.query('is_labeled').index), axis=0)
        trace_dict = {tr: trace_dict[tr] for tr in trace_dict if tr in idx_list}

        tm = self.tm_from_seq([self.predict(trace_dict[idx])[0] for idx in idx_list])

        # get bootstrapped CIs
        tm_array = []
        for _ in range(nb_bootstrap_iters):
            model = self.get_trained_model(trace_dict, bootstrap=True)
            seqs = [self.predict(trace_dict[idx], model)[0] for idx in idx_list]
            tm_array.append(self.tm_from_seq(seqs))
        tm_mat = np.stack(tm_array, axis=-1)
        sd_mat = np.std(tm_mat, axis=-1)
        mu_mat = np.mean(tm_mat, axis=-1)
        ci_mat = np.tile(np.expand_dims(mu_mat, -1), (1, 1, 2))
        ci_mat[:, :, 0] -= sd_mat * 2
        ci_mat[:, :, 1] += sd_mat * 2
        return tm, ci_mat

    def tm_from_seq(self, seq_list):
        """
        Used for the standard bootstrapping implementation, do not remove if that is still in use.
        """
        tm_out = np.zeros([self.nb_states, self.nb_states], dtype=int)
        # inventory of transitions
        for seq in seq_list:
            for tr in zip(seq[:-1], seq[1:]):
                tm_out[tr[0], tr[1]] += 1
        return tm_out / np.expand_dims(tm_out.sum(axis=1), -1)

    def get_mus(self, feature):
        """
        STEP 3.2 (OPTIONAL):
        Given a feature name, return a list of means of feature values, one for each state. This may be derived from your
        model or from your data. Values are reported in the GUI and the generated report. If not implemented,
        just use an array of zeros
        :param feature: string, name of feature for which mean values per class should be returned
        :return: list of floats of length nb_classes
        """
        mu_list = [0.0 for _ in self.nb_states]

        return mu_list

    def get_sds(self, feature):
        """
        STEP 3.3 (OPTIONAL):
        Given a feature name, return a list of standard deviation feature values, one for each state.
        This may be derived from your model or from your data. Values are reported in the GUI and the generated report.
        If not implemented, just use an array of zeros.
        :param feature: string, name of feature for which mean values per class should be returned
        :return: list of floats of length nb_classes
        """
        sd_list = [0.0 for _ in self.nb_states]
        return sd_list

    # --- saving/loading models ---
    def get_params(self):
        """
        STEP 4.1 (OPTIONAL):
        Return a string representation of your model that can be loaded into some other instance using the load_params
        function below. You are completely free on whether or how to implement this. If you do implement it, your
        implementation is vastly more useful as trained models can now be re-used.

        Tip: If you save your parameters in dicts you can easily store them in strings using yaml.dump, which can
             be easily decoded into dicts again using yaml.load

        :return: string
        """

        return 'MODEL SAVING NOT IMPLEMENTED'

    def load_params(self, file_contents):
        """
        STEP 4.2 (OPTIONAL):
        Load a previously trained model. Naturally, this only makes sense if the function save_params is also
        implemented. The result should be that the model in self.trained functions exactly the same way as it did
        when the model was first trained.
        """

        # Your implementation goes here

        self.timestamp = numeric_timestamp()  # time stamp for model updating methods, please leave as is

max_stars_repo_name: cvdelannoy/FRETboard

max_issues_repo_name: cvdelannoy/FRETboard

max_forks_repo_name: cvdelannoy/FRETboard


text:
# -*- coding: utf-8 -*-
# @Time    : 2019/11/10 10:26
# @Author  : Esbiya
# @Email   : 18829040039@163.com
# @File    : pic_locate.py
# @Software: PyCharm

import cv2
import numpy as np
from PIL import Image
import os

def _pic_download(session, url, type):
    """
    图片下载
    :param url:
    :param type:
    :return:
    """
    save_path = os.path.abspath('...') + '\\' + 'images'
    if not os.path.exists(save_path):
        os.mkdir(save_path)

    img_path = save_path + '\\' + '{}.jpg'.format(type)
    # 请求多次保证图片下载成功, 有时第一次图片下载失败
    for _ in range(10):
        session.get(url)
    img_data = session.get(url).content
    with open(img_path, 'wb') as f:
        f.write(img_data)
    return img_path

def get_distance(session, slider_url, captcha_url):
    """
    获取缺口距离
    :param session:
    :param slider_url:
    :param captcha_url:
    :return:
    """
    save_path = os.path.abspath('...') + '\\' + 'images'
    if not os.path.exists(save_path):
        os.mkdir(save_path)

    # 引用上面的图片下载
    slider_path = _pic_download(session, slider_url, 'slider')
    w, h = Image.open(slider_path).size

    # 图片还原
    captcha_path = _pic_download(session, captcha_url, 'captcha')

    # # 计算拼图还原距离
    target = cv2.imread(slider_path, 0)
    template = cv2.imread(captcha_path, 0)
    temp = save_path + '\\' + 'temp.jpg'
    targ = save_path + '\\' + 'targ.jpg'
    cv2.imwrite(targ, target)
    cv2.imwrite(temp, template)

    target = cv2.imread(targ)
    target = cv2.cvtColor(target, cv2.COLOR_BGR2GRAY)
    target = abs(255 - target)
    cv2.imwrite(targ, target)

    target = cv2.imread(targ)
    template = cv2.imread(temp)
    result = cv2.matchTemplate(target, template, cv2.TM_CCOEFF_NORMED)
    x, y = np.unravel_index(result.argmax(), result.shape)

    # 调用PIL Image 做测试
    image = Image.open(captcha_path)
    xy = (y + 23, x, y + w, x + h)
    # 切割
    imagecrop = image.crop(xy)
    # 保存切割的缺口
    imagecrop.convert('RGB').save(save_path + '\\' + "new_captcha.jpg")
    # imagecrop.show()
    return int(y + 23)


text:
"""
MIT License

Copyright (c) 2020 Daniil Shchepetilnikov

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import networkx as nx
import matplotlib.pyplot as plt
import EoN as eon

class Graph:
    """Class, represting a directed graph"""

    def __init__(self, adjlist, root=int()):
        self.__root = root
        self.__graph = nx.parse_adjlist(adjlist,
                                        create_using=nx.DiGraph,
                                        nodetype=int)

    @property
    def root(self) -> int:
        """Get root of graph"""

        return self.__root

    def childrens_of(self, node) -> list:
        """Get list of child nodes of specified node"""

        return list(self.__graph.successors(node))

    def parent_of(self, node):
        """Get parent node of specified node"""

        try:
            return list(self.__graph.predecessors(node))[0]
        except IndexError:
            return None

    def visualize(self):
        """Visualize graph like tree"""

        tree = nx.bfs_tree(self.__graph, self.__root)
        plt.figure(figsize=(8, 8))
        pos = eon.hierarchy_pos(tree, self.__root)
        nx.draw(self.__graph, pos=pos, with_labels=True, arrows=True)
        plt.show()

reason: import networkx


text:
import random
from collections import defaultdict

import numpy as np
from torch.utils.data import Dataset

class SampleDataset(Dataset):
    """
    SampleDataset to be used by TaskGenerator
    """

    def __init__(self, data, labels, sampled_task):
        self.data = data
        self.label = labels
        self.sampled_task = sampled_task

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.label[idx]

class MetaDataset(Dataset):
    """

    **Descritpion**

    It wraps a torch dataset by creating a map of target to indices.
    This comes in handy when we want to sample elements randomly for a particular label.

    Notes:
        For l2l to work its important that the dataset returns a (data, target) tuple.
        If your dataset doesn't return that, it should be trivial to wrap your dataset
        with another class to do that.
        #TODO : Add example for wrapping a non standard l2l dataset

    **Arguments**

    * **dataset** (Dataset) -  A torch dataset.
    * **labels_to_indices** (Dict) -  A dictionary mapping label to their indices.
                                     If not specified then we loop through all the datapoints to understand the mapping. (default: None)

    **Example**
    ~~~python
    mnist = torchvision.datasets.MNIST(root="/tmp/mnist", train=True)
    mnist = l2l.data.MetaDataset(mnist)
    ~~~
    """

    def __init__(self, dataset, labels_to_indices=None):

        if not isinstance(dataset, Dataset):
            raise TypeError("MetaDataset only accepts a torch dataset as input")

        self.dataset = dataset

        # TODO add tests for `labels_to_indices`
        # TODO add assertion to ensure `labels_to_indices` is valid
        self.labels_to_indices = labels_to_indices or self.get_dict_of_labels_to_indices()

        if not isinstance(self.labels_to_indices, dict):
            raise TypeError(
                "labels_to_indices should be only a dict mapping labels to keys")

        self.labels = list(self.labels_to_indices.keys())

    def __getitem__(self, item):
        return self.dataset[item]

    def __len__(self):
        return len(self.dataset)

    def get_dict_of_labels_to_indices(self):
        """ Iterates over the entire dataset and creates a map of target to indices.

        Returns: A dict with key as the label and value as list of indices.

        """

        classes_to_indices = defaultdict(list)
        for i in range(len(self.dataset)):
            try:
                # if label is a Tensor, then take get the scala value
                label = self.dataset[i][1].item()
            except AttributeError:
                # if label is a scalar then use as is
                label = self.dataset[i][1]
            except ValueError as e:
                raise ValueError("Currently l2l only supports scalar labels. \n" + str(e))

            classes_to_indices[label].append(i)
        return classes_to_indices

class LabelEncoder:
    def __init__(self, classes):
        """ Encodes a list of classes into indices, starting from 0.

        Args:
            classes: List of classes
        """
        assert len(set(classes)) == len(classes), "Classes contains duplicate values"
        self.class_to_idx = dict()
        self.idx_to_class = dict()
        for idx, old_class in enumerate(classes):
            self.class_to_idx.update({old_class: idx})
            self.idx_to_class.update({idx: old_class})

class TaskGenerator:
    """

    [[Source]](https://github.com/learnables/learn2learn/blob/master/learn2learn/data/task_generator.py)

    **Description**

    A wrapper to generate few-shot classification tasks.

    `tasks` can both indicate predefined tasks, or just the number of tasks to sample.
    If specified as an int, a list of size `task` would be generated from which we'll sample.
    If specified as a list, then that list of tasks would be used to sample always.

    The acceptable shape of list would be `n * w`, with n the number of tasks to sample and w the number of ways.

    Each of the task should have w distinct elements all of which are required to be a subset of ways.

    **Arguments**

    * **ways** (int, *optional*, default=2) - Number of labels to sample from.
    * **shots** (int, *optional*, default=1) - Number of data points per task to sample.
    * **dataset** (MetaDataset or Dataset) - The (meta-) dataset to wrap.
    * **classes** (list, *optional*, default=None) - List of classes to sample from,
        if none then sample from all available classes in dataset. (default: None)
    * **tasks** (int or list, *optional*, default=None) - Tasks to be generated.
                If none, then all possible permutations are chosen from n classes and w ways.
    """

    def __init__(self, dataset, ways=2, shots=1, classes=None, tasks=None):
        self.dataset = dataset
        self.ways = ways
        self.classes = classes
        self.shots = shots

        if not isinstance(dataset, MetaDataset):
            self.dataset = MetaDataset(dataset)

        if classes is None:
            self.classes = self.dataset.labels

        assert len(self.classes) >= ways, ValueError("Ways are more than the number of classes available")
        self._check_classes(self.classes)

        if tasks is None:
            self.tasks = None
        elif isinstance(tasks, int):
            self.tasks = self.generate_n_tasks(tasks)
        elif isinstance(tasks, list):
            self.tasks = tasks
        else:
            # TODO : allow numpy array as an input
            raise TypeError("tasks is none of None/int/list but rather {}".format(type(tasks)))

        # used for next(taskgenerator)
        self.tasks_idx = 0

        self._check_tasks(self.tasks)

        # TODO : assert that shots are always less than equal to min_samples for each class

    def _get_sample_task(self):
        return random.sample(self.classes, k=self.ways)

    def generate_n_tasks(self, n):
        # Args:
        #     n: Number of tasks to generate

        # Returns: A list of shape `n * w` where n is the number of tasks to generate and w is the ways.

        return [self._get_sample_task() for _ in range(n)]

    def __iter__(self):
        self.tasks_idx = 0
        return self

    def __len__(self):
        if self.tasks is not None:
            return len(self.tasks)
        else:
            # return 0 when tasks aren't specified
            return 0

    def __next__(self):
        if self.tasks is not None:
            try:
                task = self.sample(task=self.tasks[self.tasks_idx])
                self.tasks_idx += 1
                return task
            except IndexError:
                raise StopIteration()
        else:
            task = self.sample()
            return task

    def sample(self, shots=None, task=None):
        """

        **Description**

        Returns a dataset and the labels that we have sampled.

        The dataset is of length `shots * ways`.
        The length of labels we have sampled is the same as `shots`.

        **Arguments**

        **shots** (int, *optional*, default=None) - Number of data points to return per class, if None gets self.shots.
        **task** (list, *optional*, default=None) - List of labels you want to sample from.

        **Returns**

        * Dataset - Containing the sampled task.

        """

        # 1) Get how many number of data points to sample for every class

        # If shots isn't defined, then try to inherit from object
        if shots is None:
            if self.shots is None:
                raise ValueError(
                    "Shots is undefined in object definition neither while calling the sample method.")
            shots = self.shots

        # 2) Get a list of classes (aka tasks) that we want to sample from

        # If tasks isn't specified while calling the function, then we can
        # sample from all the tasks mentioned during the initialization of the TaskGenerator
        # except when tasks are specified None during initialization
        if task is None:
            # select few classes that will be selected for this task (for eg, 6,4,7 from 0-9 in MNIST when ways are 3)
            if self.tasks is not None:
                rand_idx = random.randint(0, len(self.tasks) - 1)
                task_to_sample = self.tasks[rand_idx]
            else:
                # Then generate a sample task on the fly
                task_to_sample = self._get_sample_task()
                assert self._check_task(task_to_sample), ValueError("Randomly generated task is malformed.")

        else:
            task_to_sample = task
            assert self._check_task(task_to_sample), ValueError("Task is malformed.")

        # 3) Map the classes into 0 index int

        # encode labels (map 6,4,7 to 0,1,2 so that we can do a BCELoss)
        label_encoder = LabelEncoder(task_to_sample)

        # 4) Get required number of data points (step 1) from the classes to sample from (step 2)
        data_indices = []
        data_labels = []
        for _class in task_to_sample:
            # select subset of indices from each of the classes and add it to data_indices
            data_indices.extend(np.random.choice(self.dataset.labels_to_indices[_class], shots, replace=False))
            # add those labels to data_labels (6 mapped to 0, so add 0's initially then 1's (for 4) and so on)
            data_labels.extend(np.full(shots, fill_value=label_encoder.class_to_idx[_class]))

        # 5) Return data from step 4  wrapped in Sample Dataset

        # map data indices to actual data
        data = [self.dataset[idx][0] for idx in data_indices]
        return SampleDataset(data, data_labels, task_to_sample)

    def _check_classes(self, classes):
        """ ensure that classes are a subset of dataset.labels """
        assert len(set(classes) - set(self.dataset.labels)) == 0, "classes contains a label that isn't in dataset"

    def _check_task(self, task) -> bool:
        """ check if each individual task is a subset of self.classes and has no duplicates """
        return (len(set(task) - set(self.classes)) == 0) and (len(set(task)) - len(task) == 0)

    def _check_tasks(self, tasks):
        """ ensure that all tasks are correctly defined. """

        # This can be possible when our `tasks=None` during TG initialization
        if tasks is not None:
            invalid_tasks = list(filter(lambda task: not self._check_task(task), tasks))
            assert len(invalid_tasks) == 0, f"Following task in mentioned tasks are unacceptable. \n {invalid_tasks}"