jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def gmm_clustering(cY, components, config):
    """Uses GMM models to cluster text lines based on their y values."""
    from sklearn.mixture import GMM
    Ms = list(range(1, 30))

    lower_bounds = []
    all_means = []
    for m in Ms:
        gmm = GMM(n_components=m, random_state=0).fit(np.expand_dims(cY, 1))
        lower_bounds.append(gmm.lower_bound_)        
        means = gmm.means_.squeeze()

        # Sort if multiple means, or turn into an array is just one
        try:
            means.sort()
        except:
            means = np.array([means])

        all_means.append(means)


    # Different methods for selecting the number of components
    n = determine_components_n(Ms, lower_bounds, all_means, config)

    # Perform analysis with determined number of components n
    gmm = GMM(n_components=n, random_state=0).fit(np.expand_dims(cY, 1))
    cluster_means = (gmm.means_.squeeze()).astype(np.int32)
    cluster_means.sort()
    return cluster_means
//  [python]
def gmm_clustering(cY, components, config):
    """Uses GMM models to cluster text lines based on their y values."""
    from sklearn.mixture import GaussianMixture
    Ms = list(range(1, 30))

    lower_bounds = []
    all_means = []
    for m in Ms:
        gmm = GaussianMixture(n_components=m, random_state=0).fit(np.expand_dims(cY, 1))
        lower_bounds.append(gmm.lower_bound_)        
        means = gmm.means_.squeeze()

        # Sort if multiple means, or turn into an array is just one
        try:
            means.sort()
        except:
            means = np.array([means])

        all_means.append(means)


    # Different methods for selecting the number of components
    n = determine_components_n(Ms, lower_bounds, all_means, config)

    # Perform analysis with determined number of components n
    gmm = GaussianMixture(n_components=n, random_state=0).fit(np.expand_dims(cY, 1))
    cluster_means = (gmm.means_.squeeze()).astype(np.int32)
    cluster_means.sort()
    return cluster_means
//  [python]
def load_gmm(train_data, gmmdir='gmm/fv0', nmix=12):
    if not os.path.exists(gmmdir):
        os.mkdir(gmmdir)
    weights_pth = os.path.join(gmmdir, "weights.npy")
    means_pth = os.path.join(gmmdir, "means.npy")
    covs_pth = os.path.join(gmmdir, "covs.npy")
    if os.path.exists(weights_pth):
        print("----load pretrained gmm from %s-----" % gmmdir)
        weights = np.load(weights_pth)
        means = np.load(means_pth)
        covs = np.load(covs_pth)
        return weights, means, covs

    descriptors = []
    for item in train_data:
        gene, nimg, gene_label, timestep = item
        descriptors.append(nimg)

    descriptors = np.concatenate(descriptors)
    gmm = mixture.GMM(n_components=nmix, covariance_type='diag')
    print("--------fit gmm-------------", weights_pth)
    gmm.fit(descriptors)

    w = gmm.weights_
    m = gmm.means_
    c = gmm.covariances_

    thr = 0.5 / nmix
    weights = np.float32([v for k, v in zip(range(nmix), w) if w[k] > thr])
    means = np.float32([v for k, v in zip(range(nmix), m) if w[k] > thr])
    covs = np.float32([v for k, v in zip(range(nmix), c) if w[k] > thr])

    np.save(weights_pth, weights)
    np.save(means_pth, means)
    np.save(covs_pth, covs)

    return weights, means, covs
//  [python]
def load_gmm(train_data, gmmdir='gmm/fv0', nmix=12):
    if not os.path.exists(gmmdir):
        os.mkdir(gmmdir)
    weights_pth = os.path.join(gmmdir, "weights.npy")
    means_pth = os.path.join(gmmdir, "means.npy")
    covs_pth = os.path.join(gmmdir, "covs.npy")
    if os.path.exists(weights_pth):
        print("----load pretrained gmm from %s-----" % gmmdir)
        weights = np.load(weights_pth)
        means = np.load(means_pth)
        covs = np.load(covs_pth)
        return weights, means, covs

    descriptors = []
    for item in train_data:
        gene, nimg, gene_label, timestep = item
        descriptors.append(nimg)

    descriptors = np.concatenate(descriptors)
    # gmm = GaussianMixture(n_components=nmix, covariance_type='diag')
    gmm = GaussianMixture(n_components=nmix)
    print("--------fit gmm-------------", weights_pth)
    gmm.fit(descriptors)

    w = gmm.weights_
    m = gmm.means_
    c = gmm.covariances_

    thr = 0.5 / nmix
    weights = np.float32([v for k, v in zip(range(nmix), w) if w[k] > thr])
    means = np.float32([v for k, v in zip(range(nmix), m) if w[k] > thr])
    covs = np.float32([v for k, v in zip(range(nmix), c) if w[k] > thr])

    np.save(weights_pth, weights)
    np.save(means_pth, means)
    np.save(covs_pth, covs)

    return weights, means, covs
//  [python]
def cluster_embeddings(
        self,
        cluster_method,
        num_clusters,
        use_decomposed,
        additional_params={"random_state": 0},
    ):
        """
        @param cluster_method clustering method, one of {KMeans, Spectral, GaussianMix, BayesGaussMix}
        @param num_clusters number of clusters
        @param use_decomposed boolean, whether to use decomposed or raw representations
        @param additional_params {param_name: value} of parameters accepted by the sklearn clustering function.
            Cannot include n_components.

        Predicts cluster assignments for data points on raw or decomposed embedding representations.
        """
        self.cluster_method = cluster_method
        self.num_clusters = num_clusters

        if use_decomposed:
            vec_df = self.decomposed_embedding_representation
        else:
            vec_df = self.embedding_representation

        if self.cluster_method == "KMeans":
            self.predicted_labels = KMeans(
                n_clusters=self.num_clusters, **additional_params
            ).fit_predict(vec_df)

        if self.cluster_method == "Spectral":
            self.predicted_labels = SpectralClustering(
                n_clusters=self.num_clusters,
                affinity="cosine",
                assign_labels="discretize",
                **additional_params,
            ).fit_predict(vec_df)

        if self.cluster_method == "GaussianMix":
            gm_model = mixture.GaussianMixture(
                n_components=self.num_clusters, **additional_params
            ).fit(vec_df)

            self.predicted_labels = gm_model.predict(vec_df)

        if self.cluster_method == "BayesGaussMix":
            bgm_model = BayesianGaussianMixture(
                n_components=self.num_clusters, **additional_params
            ).fit(vec_df)

            self.predicted_labels = bgm_model.predict(vec_df)
            self.num_clusters = len(
                set(self.predicted_labels)
            )
//  [python]
def cluster_embeddings(
        self,
        cluster_method,
        num_clusters,
        use_decomposed,
        additional_params={"random_state": 0},
    ):
        """
        @param cluster_method clustering method, one of {KMeans, Spectral, GaussianMix, BayesGaussMix}
        @param num_clusters number of clusters
        @param use_decomposed boolean, whether to use decomposed or raw representations
        @param additional_params {param_name: value} of parameters accepted by the sklearn clustering function.
            Cannot include n_components.

        Predicts cluster assignments for data points on raw or decomposed embedding representations.
        """
        self.cluster_method = cluster_method
        self.num_clusters = num_clusters

        if use_decomposed:
            vec_df = self.decomposed_embedding_representation
        else:
            vec_df = self.embedding_representation

        if self.cluster_method == "KMeans":
            self.predicted_labels = KMeans(
                n_clusters=self.num_clusters, **additional_params
            ).fit_predict(vec_df)

        if self.cluster_method == "Spectral":
            self.predicted_labels = SpectralClustering(
                n_clusters=self.num_clusters,
                affinity="cosine",
                assign_labels="discretize",
                **additional_params,
            ).fit_predict(vec_df)

        if self.cluster_method == "GaussianMix":
            gm_model = GaussianMixture(
                n_components=self.num_clusters, **additional_params
            ).fit(vec_df)

            self.predicted_labels = gm_model.predict(vec_df)

        if self.cluster_method == "BayesGaussMix":
            bgm_model = BayesianGaussianMixture(
                n_components=self.num_clusters, **additional_params
            ).fit(vec_df)

            self.predicted_labels = bgm_model.predict(vec_df)
            self.num_clusters = len(
                set(self.predicted_labels)
            )
//  [python]
def gmm(data, n_components, max_iter=150, random_state=0, **kwds):
    """
    gaussian mixture model by sklearn
    """
    model = s.GaussianMixture(
        n_components, max_iter=max_iter, random_state=random_state, **kwds
    )
    model.fit(data)
    return model
//  [python]
def gmm(data, n_components, max_iter=150, random_state=0, **kwds):
    """
    gaussian mixture model by sklearn
    """
    model = mix.GaussianMixture(
        n_components, max_iter=max_iter, random_state=random_state, **kwds
    )
    model.fit(data)
    return model
//  [python]
def execute(self, namespace):
        from sklearn.mixture import GMM
        from PYME.IO import MetaDataHandler

        points = namespace[self.input_points]
        X = np.stack([points['x'], points['y'], points['z']], axis=1)

        if self.mode == 'n':
            gmm = GMM(n_components=self.n,
                                  covariance_type=self.covariance,
                                  max_iter=self.max_iter,
                                  init_params=self.init_params)
            predictions = gmm.fit_predict(X) + 1  # PYME labeling scheme
            log_prob = gmm.score(X)
            if not gmm.converged_:
                logger.error('GMM fitting did not converge')
                predictions = np.zeros(len(points), int)
                log_prob = - np.inf * np.ones(len(points))
        
        elif self.mode == 'bic':
            n_components = range(1, self.n + 1)
            bic = np.zeros(len(n_components))
            for ind in range(len(n_components)):
                gmm = GMM(n_components=n_components[ind],
                                      covariance_type=self.covariance,
                                      max_iter=self.max_iter,
                                      init_params=self.init_params)
                gmm.fit(X)
                bic[ind] = gmm.bic(X)
                logger.debug('%d BIC: %f' % (n_components[ind], bic[ind]))

            best = n_components[np.argmin(bic)]
            if best == self.n or (self.n > 10 and best > 0.9 * self.n):
                logger.warning('BIC optimization selected n components near n max')
            
            gmm = GMM(n_components=best,
                                  covariance_type=self.covariance,
                                  max_iter=self.max_iter,
                                  init_params=self.init_params)
            predictions = gmm.fit_predict(X) + 1  # PYME labeling scheme
            log_prob = gmm.score(X)
            if not gmm.converged_:
                logger.error('GMM fitting did not converge')
                predictions = np.zeros(len(points), int)
                log_prob = - np.inf * np.ones(len(points))
        
        elif self.mode == 'bayesian':
            # BayesianGaussianMixture has no direct equivalent in older sklearn versions.  This section needs a rewrite for compatibility.
            raise NotImplementedError("BayesianGaussianMixture is not supported in the deprecated sklearn API.")
        
        out = tabular.MappingFilter(points)
        try:
            out.mdh = MetaDataHandler.DictMDHandler(points.mdh)
        except AttributeError:
            pass

        out.addColumn(self.label_key, predictions)
        out.addColumn(self.label_key + '_log_prob', log_prob)
        avg_log_prob = np.empty_like(log_prob)
        for label in np.unique(predictions):
            mask = label == predictions
            avg_log_prob[mask] = np.mean(log_prob[mask])
        out.addColumn(self.label_key + '_avg_log_prob', avg_log_prob)
        namespace[self.output_labeled] = out
//  [python]
scan_dir = 'datasets/referit3d/scan_data'
output_dir = os.path.join(scan_dir, 'instance_id_to_gmm_color')
os.makedirs(output_dir, exist_ok=True)

for scan_file in glob.glob(os.path.join(scan_dir, 'pcd_with_global_alignment', '*')):
    scan_id = os.path.basename(scan_file).split('.')[0]
    print(scan_file)
    
    data = torch.load(scan_file) # xyz, rgb, semantic_labels, instance_labels
    colors = data[1]
    instance_labels = data[3]

    if instance_labels is None:
        continue

    # normalize
    colors = colors / 127.5 - 1 
    
    clustered_colors = []
    for i in range(instance_labels.max() + 1):
        mask = instance_labels == i     # time consuming
        obj_colors = colors[mask]
        
        gm = mixture.GMM(n_components=3, covariance_type='full', random_state=0).fit(obj_colors)
        clustered_colors.append({
            'weights': gm.weights_.tolist(),
            'means': gm.means_.tolist(),
        })
        
    json.dump(
        clustered_colors,
        open(os.path.join(output_dir, '%s.json'%scan_id), 'w')
    )
//  [python]
scan_dir = 'datasets/referit3d/scan_data'
output_dir = os.path.join(scan_dir, 'instance_id_to_gmm_color')
os.makedirs(output_dir, exist_ok=True)

for scan_file in glob.glob(os.path.join(scan_dir, 'pcd_with_global_alignment', '*')):
    scan_id = os.path.basename(scan_file).split('.')[0]
    print(scan_file)
    
    data = torch.load(scan_file) # xyz, rgb, semantic_labels, instance_labels
    colors = data[1]
    instance_labels = data[3]

    if instance_labels is None:
        continue

    # normalize
    colors = colors / 127.5 - 1 
    
    clustered_colors = []
    for i in range(instance_labels.max() + 1):
        mask = instance_labels == i     # time consuming
        obj_colors = colors[mask]
        
        gm = GaussianMixture(n_components=3, covariance_type='full', random_state=0).fit(obj_colors)
        clustered_colors.append({
            'weights': gm.weights_.tolist(),
            'means': gm.means_.tolist(),
        })
        
    json.dump(
        clustered_colors,
        open(os.path.join(output_dir, '%s.json'%scan_id), 'w')
    )
//  [python]
def fit(self, imgs):
        means = self._compute_means(imgs)
        self.mixture = mixture.GaussianMixture(self.nb_clusters)
        self.mixture.fit(means)
//  [python]
def fit(self, imgs):
        means = self._compute_means(imgs)
        self.mixture = GaussianMixture(self.nb_clusters)
        self.mixture.fit(means)
//  [python]
def init_params_scikit(self, data, cov_type='full'):
        from sklearn.mixture import BayesianGaussianMixture, GaussianMixture
        gmm_init = GaussianMixture(self.nb_states, covariance_type=cov_type, n_init=5, init_params='random')
        gmm_init.fit(data)

        self.mu = gmm_init.means_
        if cov_type == 'diag':
            self.sigma = np.array([np.diag(gmm_init.covariances_[i]) for i in range(self.nb_states)])
        else:
            self.sigma = gmm_init.covariances_

        self.priors = gmm_init.weights_

        self.Trans = np.ones((self.nb_states, self.nb_states)) * 0.01

        self.init_priors = np.ones(self.nb_states) * 1. / self.nb_states
//  [python]
def init_params_scikit(self, data, cov_type='full'):
        from sklearn.mixture import BayesianGaussianMixture, GaussianMixture
        gmm_init = GaussianMixture(self.nb_states, cov_type, n_init=5, init_params='random')
        gmm_init.fit(data)

        self.mu = gmm_init.means_
        if cov_type == 'diag':
            self.sigma = np.array([np.diag(gmm_init.covariances_[i]) for i in range(self.nb_states)])
        else:
            self.sigma = gmm_init.covariances_

        self.priors = gmm_init.weights_

        self.Trans = np.ones((self.nb_states, self.nb_states)) * 0.01

        self.init_priors = np.ones(self.nb_states) * 1. / self.nb_states
//  [python]
def pre_train(self,dataloader,pre_epoch=10):

        if  not os.path.exists('./pretrain_model.pk'):

            Loss=nn.MSELoss()
            opti=Adam(itertools.chain(self.encoder.parameters(),self.decoder.parameters()))

            print('Pretraining......')
            epoch_bar=tqdm(range(pre_epoch))
            for _ in epoch_bar:
                L=0
                for x,y in dataloader:
                    if self.args.cuda:
                        x=x.cuda()

                    z,_=self.encoder(x)
                    x_=self.decoder(z)
                    loss=Loss(x,x_)

                    L+=loss.detach().cpu().numpy()

                    opti.zero_grad()
                    loss.backward()
                    opti.step()

                epoch_bar.write('L2={:.4f}'.format(L/len(dataloader)))

            self.encoder.log_sigma2_l.load_state_dict(self.encoder.mu_l.state_dict())

            Z = []
            Y = []
            with torch.no_grad():
                for x, y in dataloader:
                    if self.args.cuda:
                        x = x.cuda()

                    z1, z2 = self.encoder(x)
                    assert F.mse_loss(z1, z2) == 0
                    Z.append(z1)
                    Y.append(y)

            Z = torch.cat(Z, 0).detach().cpu().numpy()
            Y = torch.cat(Y, 0).detach().numpy()

            gmm = mixture.GaussianMixture(n_components=self.args.nClusters, covariance_type='diag')

            pre = gmm.fit_predict(Z)
            print('Acc={:.4f}%'.format(cluster_acc(pre, Y)[0] * 100))

            self.pi_.data = torch.from_numpy(gmm.weights_).cuda().float()
            self.mu_c.data = torch.from_numpy(gmm.means_).cuda().float()
            self.log_sigma2_c.data = torch.log(torch.from_numpy(gmm.covariances_).cuda().float())

            torch.save(self.state_dict(), './pretrain_model.pk')

        else:


            self.load_state_dict(torch.load('./pretrain_model.pk'))
//  [python]
def pre_train(self,dataloader,pre_epoch=10):

        if  not os.path.exists('./pretrain_model.pk'):

            Loss=nn.MSELoss()
            opti=Adam(itertools.chain(self.encoder.parameters(),self.decoder.parameters()))

            print('Pretraining......')
            epoch_bar=tqdm(range(pre_epoch))
            for _ in epoch_bar:
                L=0
                for x,y in dataloader:
                    if self.args.cuda:
                        x=x.cuda()

                    z,_=self.encoder(x)
                    x_=self.decoder(z)
                    loss=Loss(x,x_)

                    L+=loss.detach().cpu().numpy()

                    opti.zero_grad()
                    loss.backward()
                    opti.step()

                epoch_bar.write('L2={:.4f}'.format(L/len(dataloader)))

            self.encoder.log_sigma2_l.load_state_dict(self.encoder.mu_l.state_dict())

            Z = []
            Y = []
            with torch.no_grad():
                for x, y in dataloader:
                    if self.args.cuda:
                        x = x.cuda()

                    z1, z2 = self.encoder(x)
                    assert F.mse_loss(z1, z2) == 0
                    Z.append(z1)
                    Y.append(y)

            Z = torch.cat(Z, 0).detach().cpu().numpy()
            Y = torch.cat(Y, 0).detach().numpy()

            gmm = GaussianMixture(n_components=self.args.nClusters, covariance_type='diag')

            pre = gmm.fit_predict(Z)
            print('Acc={:.4f}%'.format(cluster_acc(pre, Y)[0] * 100))

            self.pi_.data = torch.from_numpy(gmm.weights_).cuda().float()
            self.mu_c.data = torch.from_numpy(gmm.means_).cuda().float()
            self.log_sigma2_c.data = torch.log(torch.from_numpy(gmm.covariances_).cuda().float())

            torch.save(self.state_dict(), './pretrain_model.pk')

        else:


            self.load_state_dict(torch.load('./pretrain_model.pk'))
//  [python]
def make_gif(max_iter=MAX_ITER, seed=0):
    """
    this doesn't work currently because the random initial guess diverges
    and the kmeans initialization is too good
    """
    colors = ["blue", "green", "red"]
    labels = ["g", "e", "f"]

    # grab data
    (g_samples, e_samples, f_samples) = iq_data = grab_data()
    train_stack = np.vstack([g_samples, e_samples, f_samples])
    samples_sep = np.array([g_samples, e_samples, f_samples])
    g_mean = np.mean(g_samples, axis=0)
    e_mean = np.mean(e_samples, axis=0)
    f_mean = np.mean(f_samples, axis=0)
    g_ind = 0
    e_ind = 1
    f_ind = 2
    precisions_init = np.array([
        [[1e-6, 1e-7],
         [1e-7, 1e-6],],
        [[1e-6, 1e-7],
         [1e-7, 1e-6],],
        np.linalg.inv([[2.32531177e-6, 3.05499850e-7],
                       [3.05499850e-7, 2.83755594e-6],]),

    ])
    
    for i in range(1, max_iter + 1):
        random.seed(0)
        np.random.seed(0)

        plot_file_path = os.path.join(SAVE_PATH, "{:05d}_iqp.png".format(i))
        gmm = mixture.GaussianMixture(COMPONENT_COUNT, max_iter=i, init_params="random",
                              means_init=np.array([g_mean, e_mean, f_mean]),
                              weights_init=np.ones(COMPONENT_COUNT) / COMPONENT_COUNT,
                              precisions_init=precisions_init,
                              random_state=seed,
        )
        gmm.fit(train_stack)

        # g_ind = find_nearest_index_2d(g_mean, gmm.means_)
        # e_ind = find_nearest_index_2d(e_mean, gmm.means_)
        # f_ind = find_nearest_index_2d(f_mean, gmm.means_)
        gmm_colors = [colors[g_ind], colors[e_ind], colors[f_ind]]
        
        plot_gmm(gmm, samples=samples_sep, colors=colors, gmm_colors=gmm_colors,
                 labels=labels, ylim=(-0.025, -0.015), xlim=(-0.01, 0.005),
                 plot_file_path=plot_file_path, dpi=DPI)
        print("pfp: {}".format(plot_file_path))
//  [python]
def make_gif(max_iter=MAX_ITER, seed=0):
    """
    this doesn't work currently because the random initial guess diverges
    and the kmeans initialization is too good
    """
    colors = ["blue", "green", "red"]
    labels = ["g", "e", "f"]

    # grab data
    (g_samples, e_samples, f_samples) = iq_data = grab_data()
    train_stack = np.vstack([g_samples, e_samples, f_samples])
    samples_sep = np.array([g_samples, e_samples, f_samples])
    g_mean = np.mean(g_samples, axis=0)
    e_mean = np.mean(e_samples, axis=0)
    f_mean = np.mean(f_samples, axis=0)
    g_ind = 0
    e_ind = 1
    f_ind = 2
    precisions_init = np.array([
        [[1e-6, 1e-7],
         [1e-7, 1e-6],],
        [[1e-6, 1e-7],
         [1e-7, 1e-6],],
        np.linalg.inv([[2.32531177e-6, 3.05499850e-7],
                       [3.05499850e-7, 2.83755594e-6],]),

    ])
    
    for i in range(1, max_iter + 1):
        random.seed(0)
        np.random.seed(0)

        plot_file_path = os.path.join(SAVE_PATH, "{:05d}_iqp.png".format(i))
        gmm = GaussianMixture(COMPONENT_COUNT, max_iter=i, init_params="random",
                              means_init=np.array([g_mean, e_mean, f_mean]),
                              weights_init=np.ones(COMPONENT_COUNT) / COMPONENT_COUNT,
                              precisions_init=precisions_init,
                              random_state=seed,
        )
        gmm.fit(train_stack)

        # g_ind = find_nearest_index_2d(g_mean, gmm.means_)
        # e_ind = find_nearest_index_2d(e_mean, gmm.means_)
        # f_ind = find_nearest_index_2d(f_mean, gmm.means_)
        gmm_colors = [colors[g_ind], colors[e_ind], colors[f_ind]]
        
        plot_gmm(gmm, samples=samples_sep, colors=colors, gmm_colors=gmm_colors,
                 labels=labels, ylim=(-0.025, -0.015), xlim=(-0.01, 0.005),
                 plot_file_path=plot_file_path, dpi=DPI)
        print("pfp: {}".format(plot_file_path))
//  [python]
def validate_plot():
    """
    Test the efficiency of the gmm and plot it.

    Arguments: None
    Returns: None
    """
    # for reproducability
    random.seed(0)
    np.random.seed(0)
    colors = ["blue", "green", "red"]
    labels = ["g", "e", "f"]

    # grab data
    (g_samples, e_samples, f_samples) = grab_data()

    # if you are using this for an experiment you should
    # use all of your data in training like so:
    # (g_samples, e_samples, f_samples) = grab_data()
    # train_stack = np.vstack([g_samples, e_samples, f_samples])
    # gmm = GaussianMixture(COMPONENT_COUNT)
    # gmm.fit(train_stack)

    # do k-slice validation for illustrative purposes
    g_accuracy = e_accuracy = f_accuracy = 0
    for i in range(VALIDATION_COUNT):
        (g_train, g_validate) = slice_for_validation(g_samples, i)
        (e_train, e_validate) = slice_for_validation(e_samples, i)
        (f_train, f_validate) = slice_for_validation(f_samples, i)
        train_stack = np.vstack([g_train, e_train, f_train])

        # fit the gmm to data
        gmm = mixture.GaussianMixture(COMPONENT_COUNT)
        gmm.fit(train_stack)

        # determine the labels of each component by proximity to mean
        g_mean = np.mean(g_train, axis=0)
        e_mean = np.mean(e_train, axis=0)
        f_mean = np.mean(f_train, axis=0)
        g_ind = find_nearest_index_2d(g_mean, gmm.means_)
        e_ind = find_nearest_index_2d(e_mean, gmm.means_)
        f_ind = find_nearest_index_2d(f_mean, gmm.means_)
        # print("means:\n{}\ngmm.means_:\n{}\ninds:\n{}"
        #       "".format(np.array([g_mean, e_mean, f_mean]), gmm.means_,
        #                 np.array([g_ind, e_ind, f_ind])))

        # validate
        g_predicted = gmm.predict(g_validate)
        e_predicted = gmm.predict(e_validate)
        f_predicted = gmm.predict(f_validate)
        g_accuracy = g_accuracy + np.count_nonzero(g_predicted == g_ind) / g_validate.shape[0]
        e_accuracy = e_accuracy + np.count_nonzero(e_predicted == e_ind) / e_validate.shape[0]
        f_accuracy = f_accuracy + np.count_nonzero(f_predicted == f_ind) / f_validate.shape[0]
    #ENDFOR
    # average the accuracies
    g_accuracy = g_accuracy / VALIDATION_COUNT
    e_accuracy = e_accuracy / VALIDATION_COUNT
    f_accuracy = f_accuracy / VALIDATION_COUNT

    # plot the gmm and data
    gmm_colors = [colors[g_ind], colors[e_ind], colors[f_ind]]
    samples_sep = np.array([g_samples, e_samples, f_samples])
    plot_gmm(gmm, samples=samples_sep,
             colors=colors, labels=labels, gmm_colors=gmm_colors,
             ylim=(-0.025, -0.015), xlim=(-0.01, 0.005))
    print("g_acc: {}, e_acc: {}, f_acc: {}"
          "".format(g_accuracy, e_accuracy, f_accuracy))
//  [python]
def validate_plot():
    """
    Test the efficiency of the gmm and plot it.

    Arguments: None
    Returns: None
    """
    # for reproducability
    random.seed(0)
    np.random.seed(0)
    colors = ["blue", "green", "red"]
    labels = ["g", "e", "f"]

    # grab data
    (g_samples, e_samples, f_samples) = grab_data()

    # if you are using this for an experiment you should
    # use all of your data in training like so:
    # (g_samples, e_samples, f_samples) = grab_data()
    # train_stack = np.vstack([g_samples, e_samples, f_samples])
    # gmm = GaussianMixture(COMPONENT_COUNT)
    # gmm.fit(train_stack)

    # do k-slice validation for illustrative purposes
    g_accuracy = e_accuracy = f_accuracy = 0
    for i in range(VALIDATION_COUNT):
        (g_train, g_validate) = slice_for_validation(g_samples, i)
        (e_train, e_validate) = slice_for_validation(e_samples, i)
        (f_train, f_validate) = slice_for_validation(f_samples, i)
        train_stack = np.vstack([g_train, e_train, f_train])

        # fit the gmm to data
        gmm = GaussianMixture(COMPONENT_COUNT)
        gmm.fit(train_stack)

        # determine the labels of each component by proximity to mean
        g_mean = np.mean(g_train, axis=0)
        e_mean = np.mean(e_train, axis=0)
        f_mean = np.mean(f_train, axis=0)
        g_ind = find_nearest_index_2d(g_mean, gmm.means_)
        e_ind = find_nearest_index_2d(e_mean, gmm.means_)
        f_ind = find_nearest_index_2d(f_mean, gmm.means_)
        # print("means:\n{}\ngmm.means_:\n{}\ninds:\n{}"
        #       "".format(np.array([g_mean, e_mean, f_mean]), gmm.means_,
        #                 np.array([g_ind, e_ind, f_ind])))

        # validate
        g_predicted = gmm.predict(g_validate)
        e_predicted = gmm.predict(e_validate)
        f_predicted = gmm.predict(f_validate)
        g_accuracy = g_accuracy + np.count_nonzero(g_predicted == g_ind) / g_validate.shape[0]
        e_accuracy = e_accuracy + np.count_nonzero(e_predicted == e_ind) / e_validate.shape[0]
        f_accuracy = f_accuracy + np.count_nonzero(f_predicted == f_ind) / f_validate.shape[0]
    #ENDFOR
    # average the accuracies
    g_accuracy = g_accuracy / VALIDATION_COUNT
    e_accuracy = e_accuracy / VALIDATION_COUNT
    f_accuracy = f_accuracy / VALIDATION_COUNT

    # plot the gmm and data
    gmm_colors = [colors[g_ind], colors[e_ind], colors[f_ind]]
    samples_sep = np.array([g_samples, e_samples, f_samples])
    plot_gmm(gmm, samples=samples_sep,
             colors=colors, labels=labels, gmm_colors=gmm_colors,
             ylim=(-0.025, -0.015), xlim=(-0.01, 0.005))
    print("g_acc: {}, e_acc: {}, f_acc: {}"
          "".format(g_accuracy, e_accuracy, f_accuracy))
//  [python]
y_third = pd.Series(2 * X_third + 0.5*np.random.randn(100,))

X = pd.concat([X_first, X_second, X_third], axis=0).to_frame()
y = pd.concat([y_first, y_second, y_third], axis=0).to_frame()

scaler = MinMaxScaler()
gaussian_mixture = mixture.GaussianMixture(n_components=3)
//  [python]
y_third = pd.Series(2 * X_third + 0.5*np.random.randn(100,))

X = pd.concat([X_first, X_second, X_third], axis=0).to_frame()
y = pd.concat([y_first, y_second, y_third], axis=0).to_frame()

scaler = MinMaxScaler()
gaussian_mixture = GaussianMixture(n_components=3)
//  [python]
def initialize_means_gmm(self,k=1):
        gm = mixture.GaussianMixture(n_components=k).fit(self.data)
        means = gm.means_
        return means
//  [python]
def initialize_means_gmm(self,k=1):
        gm = GaussianMixture(n_components=k).fit(self.data)
        means = gm.means_
        return means
//  [python]
skf = StratifiedKFold(n_splits=5) # 
skf.get_n_splits(X, y)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.4, random_state=0)

# Extract the number of classes
num_classes = len(np.unique(y_train))

# Build GMM
classifier = mixture.GMM(n_components=num_classes, covariance_type='full', 
        init_params='kmeans', max_iter=20)
//  [python]
skf = StratifiedKFold(n_splits=5) # 
skf.get_n_splits(X, y)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.4, random_state=0)

# Extract the number of classes
num_classes = len(np.unique(y_train))

# Build GMM
classifier = GaussianMixture(n_components=num_classes, covariance_type='full', 
        init_params='kmeans', max_iter=20)
//  [python]
ax.plot(*mean, color='black', marker='o')


plt.tight_layout()
plt.show()

### Mixtures of Normals

df = sp500_px.loc[sp500_px.index >= '2011-01-01', ['XOM', 'CVX']]
mclust = mixture.GaussianMixture(n_components=2).fit(df)
//  [python]
ax.plot(*mean, color='black', marker='o')


plt.tight_layout()
plt.show()

### Mixtures of Normals

df = sp500_px.loc[sp500_px.index >= '2011-01-01', ['XOM', 'CVX']]
mclust = GaussianMixture(n_components=2).fit(df)
//  [python]
def get_2d_grid_gmm(subdivisions=[5,5], variance=0.04):
    """
    Compute the weight, mean and covariance of a 2D gmm placed on a 2D grid
    Input:
        subdivisions: 2 element list of number of subdivisions of the 2D space in each axes to form the grid
        variance: scalar for spherical gmm.p
    Return:
         gmm: gmm: instance of sklearn GaussianMixture (GMM) object Gauassian mixture model
    """

    # n_gaussians = reduce(lambda x, y: x*y,subdivisions)
    n_gaussians = np.prod(np.array(subdivisions))

    step = [ 1.0/(subdivisions[0]),  1.0/(subdivisions[1])]

    means = np.mgrid[ step[0]-1 : 1.0-step[0]: complex(0,subdivisions[0]),
                      step[1]-1 : 1.0-step[1]: complex(0,subdivisions[1])]
    means = np.reshape(means, [2,-1]).T
    covariances = variance*np.ones_like(means)
    weights = (1.0/n_gaussians)*np.ones(n_gaussians)
    gmm = mixture.GaussianMixture(n_components=n_gaussians, covariance_type='diag')
    gmm.weights_ = weights
    gmm.covariances_ = covariances
    gmm.means_ = means
    from sklearn.mixture.gaussian_mixture import _compute_precision_cholesky
    gmm.precisions_cholesky_ = _compute_precision_cholesky(covariances, 'diag')
    return gmm
//  [python]
def get_2d_grid_gmm(subdivisions=[5,5], variance=0.04):
    """
    Compute the weight, mean and covariance of a 2D gmm placed on a 2D grid
    Input:
        subdivisions: 2 element list of number of subdivisions of the 2D space in each axes to form the grid
        variance: scalar for spherical gmm.p
    Return:
         gmm: gmm: instance of sklearn GaussianMixture (GMM) object Gauassian mixture model
    """

    # n_gaussians = reduce(lambda x, y: x*y,subdivisions)
    n_gaussians = np.prod(np.array(subdivisions))

    step = [ 1.0/(subdivisions[0]),  1.0/(subdivisions[1])]

    means = np.mgrid[ step[0]-1 : 1.0-step[0]: complex(0,subdivisions[0]),
                      step[1]-1 : 1.0-step[1]: complex(0,subdivisions[1])]
    means = np.reshape(means, [2,-1]).T
    covariances = variance*np.ones_like(means)
    weights = (1.0/n_gaussians)*np.ones(n_gaussians)
    gmm = GaussianMixture(n_components=n_gaussians, covariance_type='diag')
    gmm.weights_ = weights
    gmm.covariances_ = covariances
    gmm.means_ = means
    from sklearn.mixture.gaussian_mixture import _compute_precision_cholesky
    gmm.precisions_cholesky_ = _compute_precision_cholesky(covariances, 'diag')
    return gmm
//  [python]
def get_grid_gmm(subdivisions=[5,5,5], variance=0.04):
    """
    Compute the weight, mean and covariance of a gmm placed on a 3D grid
    Input:
        subdivisions: 2 element list of number of subdivisions of the 3D space in each axes to form the grid
        variance: scalar for spherical gmm.p
    Return:
         gmm: gmm: instance of sklearn GaussianMixture (GMM) object Gauassian mixture model
    """
    # n_gaussians = reduce(lambda x, y: x*y,subdivisions)
    n_gaussians = np.prod(np.array(subdivisions))
    step = [ 1.0/(subdivisions[0]),  1.0/(subdivisions[1]),  1.0/(subdivisions[2])]

    means = np.mgrid[ step[0]-1 : 1.0-step[0]: complex(0,subdivisions[0]),
                      step[1]-1 : 1.0-step[1]: complex(0,subdivisions[1]),
                      step[2]-1 : 1.0-step[2]: complex(0,subdivisions[2])]
    means = np.reshape(means, [3,-1]).T
    covariances = variance*np.ones_like(means)
    weights = (1.0/n_gaussians)*np.ones(n_gaussians)
    gmm = mixture.GaussianMixture(n_components=n_gaussians, covariance_type='diag')
    gmm.weights_ = weights
    gmm.covariances_ = covariances
    gmm.means_ = means
    from sklearn.mixture.gaussian_mixture import _compute_precision_cholesky
    gmm.precisions_cholesky_ = _compute_precision_cholesky(covariances, 'diag')
    return gmm
//  [python]
def get_grid_gmm(subdivisions=[5,5,5], variance=0.04):
    """
    Compute the weight, mean and covariance of a gmm placed on a 3D grid
    Input:
        subdivisions: 2 element list of number of subdivisions of the 3D space in each axes to form the grid
        variance: scalar for spherical gmm.p
    Return:
         gmm: gmm: instance of sklearn GaussianMixture (GMM) object Gauassian mixture model
    """
    # n_gaussians = reduce(lambda x, y: x*y,subdivisions)
    n_gaussians = np.prod(np.array(subdivisions))
    step = [ 1.0/(subdivisions[0]),  1.0/(subdivisions[1]),  1.0/(subdivisions[2])]

    means = np.mgrid[ step[0]-1 : 1.0-step[0]: complex(0,subdivisions[0]),
                      step[1]-1 : 1.0-step[1]: complex(0,subdivisions[1]),
                      step[2]-1 : 1.0-step[2]: complex(0,subdivisions[2])]
    means = np.reshape(means, [3,-1]).T
    covariances = variance*np.ones_like(means)
    weights = (1.0/n_gaussians)*np.ones(n_gaussians)
    gmm = GaussianMixture(n_components=n_gaussians, covariance_type='diag')
    gmm.weights_ = weights
    gmm.covariances_ = covariances
    gmm.means_ = means
    from sklearn.mixture.gaussian_mixture import _compute_precision_cholesky
    gmm.precisions_cholesky_ = _compute_precision_cholesky(covariances, 'diag')
    return gmm
//  [python]
def create(docs, word_vec, n_components=10):
    """Create scdv vectors

    Args:
        docs: np.array()
        word_vec: Loaded word2vectors
        n_components (int, optional): Number of components

    Returns:
        swem: Created scdv vectors
    """
    n_wv_embed = word_vec.vector_size

    # Create vocab set of w2v model and corpus
    vocab_model = set(k for k in word_vec.vocab.keys())
    vocab_docs = set([w for doc in docs for w in doc])
    out_of_vocabs = len(vocab_docs) - len(vocab_docs & vocab_model)
    print('out of vocabs: {out_of_vocabs}'.format(**locals()))
    use_words = list(vocab_docs & vocab_model)

    df_use = pd.DataFrame()
    df_use['word'] = use_words
    df_idf = create_idf_dataframe(docs)
    df_use = pd.merge(df_use, df_idf, on='word', how='left')
    idf = df_use['idf'].values

    use_word_vectors = np.array([word_vec[w] for w in use_words])

    clf = mixture.GaussianMixture(n_components=n_components, covariance_type='tied', verbose=2)
    clf.fit(use_word_vectors)

    word_probs = clf.predict_proba(use_word_vectors)
    # (n_vocabs, n_components,)
    word_cluster_vector = use_word_vectors[:, None, :] * word_probs[:, :, None]
    # (n_vocabs, n_components, n_wv_embed)

    topic_vector = word_cluster_vector.reshape(-1, n_components * n_wv_embed) * idf[:, None]

    topic_vector[np.isnan(topic_vector)] = 0
    word_to_topic = dict(zip(use_words, topic_vector))
    n_embedding = topic_vector.shape[1]

    cdv_vector = create_document_vector(docs, word_to_topic, n_embedding)
    compressed = compress_document_vector(cdv_vector)

    return compressed
//  [python]
def create(docs, word_vec, n_components=10):
    """Create scdv vectors

    Args:
        docs: np.array()
        word_vec: Loaded word2vectors
        n_components (int, optional): Number of components

    Returns:
        swem: Created scdv vectors
    """
    n_wv_embed = word_vec.vector_size

    # Create vocab set of w2v model and corpus
    vocab_model = set(k for k in word_vec.vocab.keys())
    vocab_docs = set([w for doc in docs for w in doc])
    out_of_vocabs = len(vocab_docs) - len(vocab_docs & vocab_model)
    print('out of vocabs: {out_of_vocabs}'.format(**locals()))
    use_words = list(vocab_docs & vocab_model)

    df_use = pd.DataFrame()
    df_use['word'] = use_words
    df_idf = create_idf_dataframe(docs)
    df_use = pd.merge(df_use, df_idf, on='word', how='left')
    idf = df_use['idf'].values

    use_word_vectors = np.array([word_vec[w] for w in use_words])

    clf = GaussianMixture(n_components=n_components, covariance_type='tied', verbose=2)
    clf.fit(use_word_vectors)

    word_probs = clf.predict_proba(use_word_vectors)
    # (n_vocabs, n_components,)
    word_cluster_vector = use_word_vectors[:, None, :] * word_probs[:, :, None]
    # (n_vocabs, n_components, n_wv_embed)

    topic_vector = word_cluster_vector.reshape(-1, n_components * n_wv_embed) * idf[:, None]

    topic_vector[np.isnan(topic_vector)] = 0
    word_to_topic = dict(zip(use_words, topic_vector))
    n_embedding = topic_vector.shape[1]

    cdv_vector = create_document_vector(docs, word_to_topic, n_embedding)
    compressed = compress_document_vector(cdv_vector)

    return compressed
//  [python]
def make_gmm(n_components=N_COMPONENTS, max_iter=MAX_ITER, 
             n_init=N_INIT, random_state=RANDOM_STATE):
    """fn: create gmm object"""
    model_kwds = dict(n_components=n_components, 
                      max_iter=max_iter,
                      n_init=n_init,
                      init_params='random',
                      random_state=random_state)

    gmm = s.GMM(**model_kwds)
    return gmm
//  [python]
def make_gmm(n_components=N_COMPONENTS, max_iter=MAX_ITER, 
             n_init=N_INIT, random_state=RANDOM_STATE):
    """fn: create gmm object"""
    model_kwds = dict(n_components=n_components, 
                      max_iter=max_iter,
                      n_init=n_init,
                      init_params='random',
                      random_state=random_state)

    gmm = mix.GaussianMixture(**model_kwds)
    return gmm
//  [python]
def calc_gmm_acc(params: dict, data: ndarray) -> float:
    """
    Estimates the performance of GaussianMixture model with hyper-parameters values.

    Parameters
    ----------
    params : dict
        Hyper-parameters values.

    data : array-like of shape (n_samples, n_features)
        Training sample.

    Returns
    -------
    score : float
        Performance score.
    """
    skf = KFold(n_splits=3, random_state=42, shuffle=True)
    res_metric = []
    for index in skf.split(data):
        x_train, x_test = data[index[0], :], data[index[1], :]
        model = mixture.GaussianMixture(n_components=int(params["n_components"]), covariance_type=params['covariance_type'],
                                random_state=42)
        model.fit(x_train)
        res_metric.append(model.score(x_test))
    score = -np.mean(res_metric)
    return score
//  [python]
def calc_gmm_acc(params: dict, data: ndarray) -> float:
    """
    Estimates the performance of GaussianMixture model with hyper-parameters values.

    Parameters
    ----------
    params : dict
        Hyper-parameters values.

    data : array-like of shape (n_samples, n_features)
        Training sample.

    Returns
    -------
    score : float
        Performance score.
    """
    skf = KFold(n_splits=3, random_state=42, shuffle=True)
    res_metric = []
    for index in skf.split(data):
        x_train, x_test = data[index[0], :], data[index[1], :]
        model = GaussianMixture(n_components=int(params["n_components"]), covariance_type=params['covariance_type'],
                                random_state=42)
        model.fit(x_train)
        res_metric.append(model.score(x_test))
    score = -np.mean(res_metric)
    return score
//  [python]
def compute_metrics_mog_gmm(means, covs, samples):
    """Compute of bunch of metrics (KL between the modes, mean MSE, cov MSE, mode weight) between a
    the real mixture and a GMM fitted approximaion

    Args:
        means (torch.Tensor (n_modes, dim)): Means of the target
        covs (torch.Tensor (n_modes, dim, dim)): Covariances of the target
        samples (torch.Tensor (n_samples, dim)): Approximate samples

    Returns:
        metrics (dict): Metrics
    """

    # Disable gradients
    with torch.no_grad():
        # Make a output dictionnary
        ret = {}
        # Fit the GMM
        try:
            gm = mixture.GaussianMixture(n_components=means.shape[0],
                                 weights_init=np.array([2. / 3., 1. / 3.]),
                                 means_init=means.cpu().numpy(),
                                 precisions_init=torch.linalg.inv(covs).cpu().numpy())
            gm = gm.fit(samples.cpu().numpy())
        except BaseException:
            return {}
        # Get the components
        mask, weight = compute_relative_weights(means[0], means[1], samples, return_mask=True)
        ret['weight_gmm'] = float(np.max(gm.weights_))
        # Compute the means and covs
        means_est = torch.from_numpy(gm.means_).to(samples.device)
        covs_est = torch.from_numpy(gm.covariances_).to(samples.device)
        try:
            ret['kl_mode1_gmm'] = compute_kl_gaussians(means[0], covs[0], means_est[0], covs_est[0])
        except BaseException:
            ret['kl_mode1_gmm'] = torch.nan
        try:
            ret['kl_mode2_gmm'] = compute_kl_gaussians(means[1], covs[1], means_est[1], covs_est[1])
        except BaseException:
            ret['kl_mode2_gmm'] = torch.nan
        ret['mean_err_mode1_gmm'] = float(torch.linalg.norm(means_est[0] - means[0]).cpu())
        ret['mean_err_mode2_gmm'] = float(torch.linalg.norm(means_est[1] - means[1]).cpu())
        ret['cov_err_mode1_gmm'] = float(torch.linalg.norm(covs[0] - covs_est[0], ord='fro').cpu())
        ret['cov_err_mode2_gmm'] = float(torch.linalg.norm(covs[1] - covs_est[1], ord='fro').cpu())
    return ret
//  [python]
def compute_metrics_mog_gmm(means, covs, samples):
    """Compute of bunch of metrics (KL between the modes, mean MSE, cov MSE, mode weight) between a
    the real mixture and a GMM fitted approximaion

    Args:
        means (torch.Tensor (n_modes, dim)): Means of the target
        covs (torch.Tensor (n_modes, dim, dim)): Covariances of the target
        samples (torch.Tensor (n_samples, dim)): Approximate samples

    Returns:
        metrics (dict): Metrics
    """

    # Disable gradients
    with torch.no_grad():
        # Make a output dictionnary
        ret = {}
        # Fit the GMM
        try:
            gm = GaussianMixture(n_components=means.shape[0],
                                 weights_init=np.array([2. / 3., 1. / 3.]),
                                 means_init=means.cpu().numpy(),
                                 precisions_init=torch.linalg.inv(covs).cpu().numpy())
            gm = gm.fit(samples.cpu().numpy())
        except BaseException:
            return {}
        # Get the components
        mask, weight = compute_relative_weights(means[0], means[1], samples, return_mask=True)
        ret['weight_gmm'] = float(np.max(gm.weights_))
        # Compute the means and covs
        means_est = torch.from_numpy(gm.means_).to(samples.device)
        covs_est = torch.from_numpy(gm.covariances_).to(samples.device)
        try:
            ret['kl_mode1_gmm'] = compute_kl_gaussians(means[0], covs[0], means_est[0], covs_est[0])
        except BaseException:
            ret['kl_mode1_gmm'] = torch.nan
        try:
            ret['kl_mode2_gmm'] = compute_kl_gaussians(means[1], covs[1], means_est[1], covs_est[1])
        except BaseException:
            ret['kl_mode2_gmm'] = torch.nan
        ret['mean_err_mode1_gmm'] = float(torch.linalg.norm(means_est[0] - means[0]).cpu())
        ret['mean_err_mode2_gmm'] = float(torch.linalg.norm(means_est[1] - means[1]).cpu())
        ret['cov_err_mode1_gmm'] = float(torch.linalg.norm(covs[0] - covs_est[0], ord='fro').cpu())
        ret['cov_err_mode2_gmm'] = float(torch.linalg.norm(covs[1] - covs_est[1], ord='fro').cpu())
    return ret
//  [python]
def clustering(vectors, mode, n_clusters, cluster_method, out_path, ind=False):
    index = []
    if mode == 'feature':
        vectors = vectors.T
        for i in range(len(vectors)):
            index.append('F%d' % (i + 1))
    else:
        for i in range(len(vectors)):
            index.append('S%d' % (i + 1))

    labels = np.zeros(len(vectors))
    if cluster_method == 'AP':
        ap = AffinityPropagation().fit(vectors)
        cluster_centers_indices = ap.cluster_centers_indices_
        labels = ap.labels_
        # plot_ap(vectors, cluster_centers_indices, labels, out_path, ind)
    elif cluster_method == 'DBSCAN':
        data = StandardScaler().fit_transform(vectors)
        db = DBSCAN().fit(data)
        labels = db.labels_
    elif cluster_method == 'GMM':
        gm = mixture.GaussianMixture(n_components=n_clusters).fit(vectors)
        labels = gm.predict(vectors)
    elif cluster_method == 'AGNES':
        # plot_hc(vectors, index, out_path, ind)
        connectivity = kneighbors_graph(vectors, n_neighbors=10, include_self=False)
        ward = AgglomerativeClustering(n_clusters=n_clusters, connectivity=connectivity,
                                       linkage='ward').fit(vectors)
        labels = ward.labels_
    elif cluster_method == 'Kmeans':
        labels = KMeans(n_clusters=n_clusters).fit_predict(vectors)

    res = []
    for i in range(len(vectors)):
        res.append([index[i], labels[i]])
    return res