jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def calc_gmm(coords, peaks, seed, iter, gmm_filename):
    X = coords.to_numpy(dtype=float)
    print(f"Calculating GMM with {peaks} means for max {iter} iterations")
    gmm = mixture.GaussianMixture(n_components=peaks,
                       covariance_type='full',
                       verbose=1,
                       n_init=1,
                       max_iter=iter,
                       random_state=seed).fit(X)
    save_gmm(gmm, gmm_filename)
//  [python]
def calc_gmm(coords, peaks, seed, iter, gmm_filename):
    X = coords.to_numpy(dtype=float)
    print(f"Calculating GMM with {peaks} means for max {iter} iterations")
    gmm = GaussianMixture(n_components=peaks,
                       covariance_type='full',
                       verbose=1,
                       n_init=1,
                       max_iter=iter,
                       random_state=seed).fit(X)
    save_gmm(gmm, gmm_filename)
//  [python]
def gmm_crit(coords, start, end, step, seed):
    X = coords.to_numpy(dtype=float)

    fig, ax = plt.subplots(figsize=(20, 10))

    models = []
    n_components = []
    for n in range(start, end, step):
        print(f"Calculating model with {n} means")
        n_components.append(n)
        models.append(mixture.GaussianMixture(n_components=n,
                                      covariance_type='full',
                                      verbose=1,
                                      random_state=seed).fit(X))

    plt.plot(n_components, [m.bic(X) for m in models], label='BIC')
    plt.plot(n_components, [m.aic(X) for m in models], label='AIC')
    plt.legend(loc='best')
    plt.xlabel('n_components')
    #plt.show()

    f = f"results/img/gmm.png"
    plt.savefig(f)
//  [python]
def gmm_crit(coords, start, end, step, seed):
    X = coords.to_numpy(dtype=float)

    fig, ax = plt.subplots(figsize=(20, 10))

    models = []
    n_components = []
    for n in range(start, end, step):
        print(f"Calculating model with {n} means")
        n_components.append(n)
        models.append(GaussianMixture(n_components=n,
                                      covariance_type='full',
                                      verbose=1,
                                      random_state=seed).fit(X))

    plt.plot(n_components, [m.bic(X) for m in models], label='BIC')
    plt.plot(n_components, [m.aic(X) for m in models], label='AIC')
    plt.legend(loc='best')
    plt.xlabel('n_components')
    #plt.show()

    f = f"results/img/gmm.png"
    plt.savefig(f)
//  [python]
def gmm(coords, peaks, seed):
    X = coords.to_numpy(dtype=float)
    xmin, xmax = -180., 180.
    ymin, ymax = -90., 90.
    print(X)

    fig, ax = plt.subplots(figsize=(20, 10))
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(ymin, ymax)
    map = Basemap()
    map.drawcoastlines(linewidth=0.5, color="black")
    map.drawcountries(linewidth=0.7, color="black")
    map.drawparallels(np.arange(ymin, ymax, 30.))
    map.drawmeridians(np.arange(xmin, xmax, 30.))
    map.drawmapboundary(fill_color='azure')
    map.fillcontinents(color='white', lake_color='azure')

    print(f"Calculating clusters of {X.shape[0]} points from GMM with {peaks} means and random seed {seed}")
    gmm = mixture.GaussianMixture(n_components=peaks, covariance_type='full', random_state=seed).fit(X)
    labels = gmm.predict(X)
    probs = gmm.predict_proba(X)
    size = probs.max(1)**2

    map.scatter(X[:, 0], X[:, 1], c=labels, s=size, cmap="turbo", zorder=5)
    #plt.scatter(X[:, 0], X[:, 1], s=4, c="black")

    w_factor = 0.2 / gmm.weights_.max()
    for pos, covar, w in zip(gmm.means_, gmm.covariances_, gmm.weights_):
        U, s, Vt = np.linalg.svd(covar)
        angle = np.degrees(np.arctan2(U[1, 0], U[0, 0]))
        width, height = 2 * np.sqrt(s)
        for nsig in range(1, 4):
            ax.add_patch(Ellipse(pos, nsig * width, nsig * height, angle,
                                 alpha=w * w_factor, color="black"))

    plt.title(f'Scatter plot of coordinated')
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    plt.show()
//  [python]
def gmm(coords, peaks, seed):
    X = coords.to_numpy(dtype=float)
    xmin, xmax = -180., 180.
    ymin, ymax = -90., 90.
    print(X)

    fig, ax = plt.subplots(figsize=(20, 10))
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(ymin, ymax)
    map = Basemap()
    map.drawcoastlines(linewidth=0.5, color="black")
    map.drawcountries(linewidth=0.7, color="black")
    map.drawparallels(np.arange(ymin, ymax, 30.))
    map.drawmeridians(np.arange(xmin, xmax, 30.))
    map.drawmapboundary(fill_color='azure')
    map.fillcontinents(color='white', lake_color='azure')

    print(f"Calculating clusters of {X.shape[0]} points from GMM with {peaks} means and random seed {seed}")
    gmm = GaussianMixture(n_components=peaks, covariance_type='full', random_state=seed).fit(X)
    labels = gmm.predict(X)
    probs = gmm.predict_proba(X)
    size = probs.max(1)**2

    map.scatter(X[:, 0], X[:, 1], c=labels, s=size, cmap="turbo", zorder=5)
    #plt.scatter(X[:, 0], X[:, 1], s=4, c="black")

    w_factor = 0.2 / gmm.weights_.max()
    for pos, covar, w in zip(gmm.means_, gmm.covariances_, gmm.weights_):
        U, s, Vt = np.linalg.svd(covar)
        angle = np.degrees(np.arctan2(U[1, 0], U[0, 0]))
        width, height = 2 * np.sqrt(s)
        for nsig in range(1, 4):
            ax.add_patch(Ellipse(pos, nsig * width, nsig * height, angle,
                                 alpha=w * w_factor, color="black"))

    plt.title(f'Scatter plot of coordinated')
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    plt.show()
//  [python]
mask = (im > im.mean()).astype(float)


img = mask + 0.3 * rng.normal(size=mask.shape)

hist, bin_edges = np.histogram(img, bins=60)
bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])

classif = mixture.GaussianMixture(n_components=2)
//  [python]
mask = (im > im.mean()).astype(float)


img = mask + 0.3 * rng.normal(size=mask.shape)

hist, bin_edges = np.histogram(img, bins=60)
bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])

classif = GaussianMixture(n_components=2)
//  [python]
def _build_word_cluster_probabilities(word_embeddings: np.ndarray, cluster_size: int, gaussian_mixture_parameters: Dict[Any, Any]) -> np.ndarray:
        gm = mixture.GaussianMixture(n_components=cluster_size, **gaussian_mixture_parameters)
        gm.fit(word_embeddings)
        return gm.predict_proba(word_embeddings)
//  [python]
def _build_word_cluster_probabilities(word_embeddings: np.ndarray, cluster_size: int, gaussian_mixture_parameters: Dict[Any, Any]) -> np.ndarray:
        gm = GaussianMixture(n_components=cluster_size, **gaussian_mixture_parameters)
        gm.fit(word_embeddings)
        return gm.predict_proba(word_embeddings)
//  [python]
def load_gmm(filename):
    data = pd.read_json(path_or_buf=filename, lines=True)
    #print(data)
    print(f"PARAM\tLOAD\tParameters for GMM with {len(data.index)} means are loaded")
    means, covs, prec, prec_chol = [], [], [], []
    for i in range(len(data.index)):
        means.append(np.array(data.at[i, "means"]))
        covs.append(np.array(data.at[i, "covariances"]))
        prec.append(np.array(data.at[i, "precisions"]))
        prec_chol.append(np.array(data.at[i, "precisions_cholesky"]))

    print(f"MODEL\tINIT\tInitialization of GMM with {len(data.index)} means")
    gmm = mixture.GaussianMixture(n_components=len(data.index), covariance_type='full', max_iter=1, verbose=0, random_state=seed)
    gmm.weights_ = data["weights"].values

    weights = gmm.weights_
    print('2 dec - Estimated number of clusters: ' + str((np.round(weights, 2) > 0).sum()))
    print('3 dec - Estimated number of clusters: ' + str((np.round(weights, 3) > 0).sum()))
    print('4 dec - Estimated number of clusters: ' + str((np.round(weights, 4) > 0).sum()))
    print('5 dec - Estimated number of clusters: ' + str((np.round(weights, 5) > 0).sum()))

    gmm.means_ = np.array(means)
    gmm.covariances_ = np.array(covs)
    gmm.precisions_ = np.array(prec)
    gmm.precisions_cholesky_ = np.array(prec_chol)

    print(f"MODEL\tSET\tParameters of GMM with {len(data.index)} means are set")
    return gmm
//  [python]
def load_gmm(filename):
    data = pd.read_json(path_or_buf=filename, lines=True)
    #print(data)
    print(f"PARAM\tLOAD\tParameters for GMM with {len(data.index)} means are loaded")
    means, covs, prec, prec_chol = [], [], [], []
    for i in range(len(data.index)):
        means.append(np.array(data.at[i, "means"]))
        covs.append(np.array(data.at[i, "covariances"]))
        prec.append(np.array(data.at[i, "precisions"]))
        prec_chol.append(np.array(data.at[i, "precisions_cholesky"]))

    print(f"MODEL\tINIT\tInitialization of GMM with {len(data.index)} means")
    gmm = GaussianMixture(n_components=len(data.index), covariance_type='full', max_iter=1, verbose=0, random_state=seed)
    gmm.weights_ = data["weights"].values

    weights = gmm.weights_
    print('2 dec - Estimated number of clusters: ' + str((np.round(weights, 2) > 0).sum()))
    print('3 dec - Estimated number of clusters: ' + str((np.round(weights, 3) > 0).sum()))
    print('4 dec - Estimated number of clusters: ' + str((np.round(weights, 4) > 0).sum()))
    print('5 dec - Estimated number of clusters: ' + str((np.round(weights, 5) > 0).sum()))

    gmm.means_ = np.array(means)
    gmm.covariances_ = np.array(covs)
    gmm.precisions_ = np.array(prec)
    gmm.precisions_cholesky_ = np.array(prec_chol)

    print(f"MODEL\tSET\tParameters of GMM with {len(data.index)} means are set")
    return gmm
//  [python]
plt.tight_layout()
plt.axis('equal')
plt.savefig('../figures/gmm_2d_data.pdf', dpi=300)
plt.show()

gm = mixture.GaussianMixture(n_components=K, n_init=10, random_state=42)
//  [python]
plt.tight_layout()
plt.axis('equal')
plt.savefig('../figures/gmm_2d_data.pdf', dpi=300)
plt.show()

gm = GaussianMixture(n_components=K, n_init=10, random_state=42)
//  [python]
def gmm_sklearn(X, k):
    from sklearn.mixture import GMM
    out = GMM(k).fit(X)
    return dict(labels=out.predict(X),
                centers=out.means_)
//  [python]
def gmm_sklearn(X, k):
    from sklearn.mixture import GaussianMixture
    out = GaussianMixture(k).fit(X)
    return dict(labels=out.predict(X),
                centers=out.means_)
//  [python]
def __init__(self, hparams, logging):
        super().__init__()
            
        """
        Initialise the Lightning Module that can scan over different GNN training regimes
        """                
        
        self.supernode_encoder = make_mlp(
            hparams["latent"],
            hparams["hidden"],
            hparams["latent"] - hparams["emb_dim"],
            hparams["nb_node_layer"],
            output_activation=hparams["hidden_activation"],
            hidden_activation=hparams["hidden_activation"],
            layer_norm=hparams["layernorm"],
        )

        # The edge network computes new edge features from connected nodes
        self.superedge_encoder = make_mlp(
            2 * hparams["latent"],
            hparams["hidden"],
            hparams["latent"],
            hparams["nb_edge_layer"],
            layer_norm=hparams["layernorm"],
            output_activation=hparams["hidden_activation"],
            hidden_activation=hparams["hidden_activation"],
        )

        # Initialize GNN blocks
        if hparams["share_weight"]:
            cell = HierarchicalGNNCell(hparams)
            hgnn_cells = [
                cell
                for _ in range(hparams["n_hierarchical_graph_iters"])
            ]
        else:
            hgnn_cells = [
                HierarchicalGNNCell(hparams)
                for _ in range(hparams["n_hierarchical_graph_iters"])
            ]
        
        self.hgnn_cells = nn.ModuleList(hgnn_cells)
        
        # output layers
        self.GMM_model = mixture.GaussianMixture(n_components = 2)
        self.super_graph_construction = DynamicGraphConstruction("sigmoid", hparams)
        self.bipartite_graph_construction = DynamicGraphConstruction("exp", hparams)
        self.register_buffer("score_cut", torch.tensor([float("inf")]))
        
        self.log = logging
        self.hparams = hparams
//  [python]
def __init__(self, hparams, logging):
        super().__init__()
            
        """
        Initialise the Lightning Module that can scan over different GNN training regimes
        """                
        
        self.supernode_encoder = make_mlp(
            hparams["latent"],
            hparams["hidden"],
            hparams["latent"] - hparams["emb_dim"],
            hparams["nb_node_layer"],
            output_activation=hparams["hidden_activation"],
            hidden_activation=hparams["hidden_activation"],
            layer_norm=hparams["layernorm"],
        )

        # The edge network computes new edge features from connected nodes
        self.superedge_encoder = make_mlp(
            2 * hparams["latent"],
            hparams["hidden"],
            hparams["latent"],
            hparams["nb_edge_layer"],
            layer_norm=hparams["layernorm"],
            output_activation=hparams["hidden_activation"],
            hidden_activation=hparams["hidden_activation"],
        )

        # Initialize GNN blocks
        if hparams["share_weight"]:
            cell = HierarchicalGNNCell(hparams)
            hgnn_cells = [
                cell
                for _ in range(hparams["n_hierarchical_graph_iters"])
            ]
        else:
            hgnn_cells = [
                HierarchicalGNNCell(hparams)
                for _ in range(hparams["n_hierarchical_graph_iters"])
            ]
        
        self.hgnn_cells = nn.ModuleList(hgnn_cells)
        
        # output layers
        self.GMM_model = GaussianMixture(n_components = 2)
        self.super_graph_construction = DynamicGraphConstruction("sigmoid", hparams)
        self.bipartite_graph_construction = DynamicGraphConstruction("exp", hparams)
        self.register_buffer("score_cut", torch.tensor([float("inf")]))
        
        self.log = logging
        self.hparams = hparams
//  [python]
def gmm(X):
    X = ((X - X.min()) / (X.max() - X.min())).reshape(-1, 1)
    gmm = mixture.GaussianMixture(n_components=2, max_iter=100, tol=1e-2, reg_covar=5e-4).fit(X)
    prob = gmm.predict_proba(X)[:, gmm.means_.argmin()]
    return prob
//  [python]
def gmm(X):
    X = ((X - X.min()) / (X.max() - X.min())).reshape(-1, 1)
    gmm = GaussianMixture(n_components=2, max_iter=100, tol=1e-2, reg_covar=5e-4).fit(X)
    prob = gmm.predict_proba(X)[:, gmm.means_.argmin()]
    return prob
//  [python]
winenorm['cluster'] = label_pred_KM.astype(str)
sns_plot = sns.pairplot(winenorm, hue = "cluster")
#sns_plot.savefig('pairplot.png')

########################################################################
##### GMM
########################################################################
aic = dict()
bic = dict()
for k in range(1,11):
	estimator = mixture.GaussianMixture(n_components = k, random_state=seed)
	res = estimator.fit(winenorm)
	ic1 = estimator.aic(winenorm)
	ic2 = estimator.bic(winenorm)
	aic[k] = ic1
	bic[k] = ic2
//  [python]
winenorm['cluster'] = label_pred_KM.astype(str)
sns_plot = sns.pairplot(winenorm, hue = "cluster")
#sns_plot.savefig('pairplot.png')

########################################################################
##### GMM
########################################################################
aic = dict()
bic = dict()
for k in range(1,11):
	estimator = GaussianMixture(n_components = k, random_state=seed)
	res = estimator.fit(winenorm)
	ic1 = estimator.aic(winenorm)
	ic2 = estimator.bic(winenorm)
	aic[k] = ic1
	bic[k] = ic2
//  [python]
def estimate_one_dim_gmm(K: int, X: np.ndarray) -> MOGMeasure:
    """
    :param K: number of mixtures
    :param X: input data
    :return: estimated Gaussian mixture model on the data X
    """
    tf.debugging.assert_shapes([(X, ("N",))])
    assert K > 0
    gm = mixture.GaussianMixture(
        n_components=K, random_state=0, covariance_type="spherical"
    ).fit(X.reshape(-1, 1))
    assert np.allclose(gm.weights_.sum(), 1.0)
    assert gm.means_.shape == (K, 1)
    assert gm.covariances_.shape == (K,)
    assert gm.weights_.shape == (K,)
    return MOGMeasure(
        weights=gm.weights_, means=gm.means_.reshape(-1), variances=gm.covariances_
    )
//  [python]
def estimate_one_dim_gmm(K: int, X: np.ndarray) -> MOGMeasure:
    """
    :param K: number of mixtures
    :param X: input data
    :return: estimated Gaussian mixture model on the data X
    """
    tf.debugging.assert_shapes([(X, ("N",))])
    assert K > 0
    gm = GaussianMixture(
        n_components=K, random_state=0, covariance_type="spherical"
    ).fit(X.reshape(-1, 1))
    assert np.allclose(gm.weights_.sum(), 1.0)
    assert gm.means_.shape == (K, 1)
    assert gm.covariances_.shape == (K,)
    assert gm.weights_.shape == (K,)
    return MOGMeasure(
        weights=gm.weights_, means=gm.means_.reshape(-1), variances=gm.covariances_
    )
//  [python]
pred = clf1.train(X)
time_end1 = time.time()
print("Runtime of GMM:", time_end1-time_start1)


time_start2 = time.time()
clf2 = mixture.GaussianMixture(n_components=4)
//  [python]
pred = clf1.train(X)
time_end1 = time.time()
print("Runtime of GMM:", time_end1-time_start1)


time_start2 = time.time()
clf2 = GaussianMixture(n_components=4)
//  [python]
def GaussianModel(embeddings):
    gmm = mixture.GaussianMixture(n_components=1, reg_covar=1e-05)
    gmm.fit(embeddings)

    log_likelihood = gmm.score_samples(embeddings)
    return log_likelihood
//  [python]
def GaussianModel(embeddings):
    gmm = GaussianMixture(n_components=1, reg_covar=1e-05)
    gmm.fit(embeddings)

    log_likelihood = gmm.score_samples(embeddings)
    return log_likelihood
//  [python]
def train_using_python(self, input_data, iters=10):
        seed = 5
        if self.components_seeded and self.clf is None:
            self.clf = mixture.GaussianMixture(n_components=self.M, covariance_type=self.cvtype, warm_start=True,\
                               max_iter=iters, init_params='random', random_state=seed, \
                               weights_init=self.components.weights, means_init=self.components.means, 
                               n_init=1)#, reg_covar=0.01)

        elif not self.components_seeded and self.clf is None:
            self.clf = mixture.GaussianMixture(n_components=self.M, covariance_type=self.cvtype, warm_start=True,\
                               max_iter=iters, init_params='random', n_init=1, random_state=seed)#, reg_covar=0.01)

        self.clf.fit(input_data)
        return self.clf.means_, self.clf.covariances_
//  [python]
def train_using_python(self, input_data, iters=10):
        seed = 5
        if self.components_seeded and self.clf is None:
            self.clf = GaussianMixture(n_components=self.M, covariance_type=self.cvtype, warm_start=True,\
                               max_iter=iters, init_params='random', random_state=seed, \
                               weights_init=self.components.weights, means_init=self.components.means, 
                               n_init=1)#, reg_covar=0.01)

        elif not self.components_seeded and self.clf is None:
            self.clf = GaussianMixture(n_components=self.M, covariance_type=self.cvtype, warm_start=True,\
                               max_iter=iters, init_params='random', n_init=1, random_state=seed)#, reg_covar=0.01)

        self.clf.fit(input_data)
        return self.clf.means_, self.clf.covariances_
//  [python]
def _process_cluster(
    X_c: np.ndarray,
    samples_c: np.ndarray,
    n_species: int,
    max_deviation: float
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Process each cluster using GMM.

    Auxiliary function to `match_features`.

    Parameters
    ----------
    X_c : array
    samples_c : array
    max_deviation : float
    n_species: int
        Number of features in the cluster, estimated with
        `estimate_features_per_cluster`.

    Returns
    -------
    label : np.ndarray
    indecisiveness : np.ndarray
    """

    # fit GMM
    gmm = mixture.GaussianMixture(n_components=n_species, covariance_type="diag")
    gmm.fit(X_c)

    # compute the deviation of the features respect to each cluster
    deviation = _get_deviation(X_c, gmm.covariances_, gmm.means_)

    # assign each feature in a sample to component in the GMM minimizing the
    # total deviation in the sample.
    label = - np.ones_like(samples_c)   # by default features are set as noise
    unique_samples = np.unique(samples_c)
    # the indecisiveness is a metric that counts the number of samples
    # where more than one feature can be potentially assigned to a species
    # that is, for each species, the number of rows in deviation with values
    # lower than max_deviation are counted as 1 and zero otherwise.
    # This is done for all samples and the indecisiveness is divided by the
    # number of samples.
    indecisiveness = np.zeros(n_species)
    for s in unique_samples:
        sample_mask = samples_c == s
        sample_deviation = deviation[sample_mask, :]
        # indecisiveness
        count = (sample_deviation < max_deviation).sum(axis=0)
        indecisiveness += (count > 1)

        # Find the best option for each feature
        best_row, best_col = linear_sum_assignment(sample_deviation)

        # features with deviation greater than max_deviation are set to noise
        valid_ft_mask = sample_deviation[best_row, best_col] <= max_deviation
        best_col = best_col[valid_ft_mask]
        best_row = best_row[valid_ft_mask]
        sample_label = - np.ones(sample_deviation.shape[0], dtype=int)
        sample_label[best_row] = best_col
        label[sample_mask] = sample_label
    indecisiveness /= unique_samples.size
    return label, indecisiveness
//  [python]
def _process_cluster(
    X_c: np.ndarray,
    samples_c: np.ndarray,
    n_species: int,
    max_deviation: float
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Process each cluster using GMM.

    Auxiliary function to `match_features`.

    Parameters
    ----------
    X_c : array
    samples_c : array
    max_deviation : float
    n_species: int
        Number of features in the cluster, estimated with
        `estimate_features_per_cluster`.

    Returns
    -------
    label : np.ndarray
    indecisiveness : np.ndarray
    """

    # fit GMM
    gmm = GaussianMixture(n_components=n_species, covariance_type="diag")
    gmm.fit(X_c)

    # compute the deviation of the features respect to each cluster
    deviation = _get_deviation(X_c, gmm.covariances_, gmm.means_)

    # assign each feature in a sample to component in the GMM minimizing the
    # total deviation in the sample.
    label = - np.ones_like(samples_c)   # by default features are set as noise
    unique_samples = np.unique(samples_c)
    # the indecisiveness is a metric that counts the number of samples
    # where more than one feature can be potentially assigned to a species
    # that is, for each species, the number of rows in deviation with values
    # lower than max_deviation are counted as 1 and zero otherwise.
    # This is done for all samples and the indecisiveness is divided by the
    # number of samples.
    indecisiveness = np.zeros(n_species)
    for s in unique_samples:
        sample_mask = samples_c == s
        sample_deviation = deviation[sample_mask, :]
        # indecisiveness
        count = (sample_deviation < max_deviation).sum(axis=0)
        indecisiveness += (count > 1)

        # Find the best option for each feature
        best_row, best_col = linear_sum_assignment(sample_deviation)

        # features with deviation greater than max_deviation are set to noise
        valid_ft_mask = sample_deviation[best_row, best_col] <= max_deviation
        best_col = best_col[valid_ft_mask]
        best_row = best_row[valid_ft_mask]
        sample_label = - np.ones(sample_deviation.shape[0], dtype=int)
        sample_label[best_row] = best_col
        label[sample_mask] = sample_label
    indecisiveness /= unique_samples.size
    return label, indecisiveness
//  [python]
def get_learned_gmm(points, n_gaussians, covariance_type='diag'):
    """
    Learn weights, means and covariances for a gmm based on input data
    Input:
         points: num_points_per_model*nummodels X 3 - xyz coordinates
         n_gaussians: scalar of number of gaussians /  3 element list of number of subdivisions for grid type
         covariance_type: Specify the type of covariance mmatrix : 'diag', 'full','tied', 'spherical' (Note that the Fisher Vector method relies on diagonal covariance matrix)
        See sklearn documentation : http://scikit-learn.org/stable/modules/generated/sklearn.mixture.GaussianMixture.html
    Return:
         gmm: gmm: instance of sklearn GaussianMixture (GMM) object Gauassian mixture model
    """
    gmm = mixture.GaussianMixture(n_components = n_gaussians, covariance_type=covariance_type)
    gmm.fit(points.astype(np.float64))
    return gmm
//  [python]
def get_learned_gmm(points, n_gaussians, covariance_type='diag'):
    """
    Learn weights, means and covariances for a gmm based on input data
    Input:
         points: num_points_per_model*nummodels X 3 - xyz coordinates
         n_gaussians: scalar of number of gaussians /  3 element list of number of subdivisions for grid type
         covariance_type: Specify the type of covariance mmatrix : 'diag', 'full','tied', 'spherical' (Note that the Fisher Vector method relies on diagonal covariance matrix)
        See sklearn documentation : http://scikit-learn.org/stable/modules/generated/sklearn.mixture.GaussianMixture.html
    Return:
         gmm: gmm: instance of sklearn GaussianMixture (GMM) object Gauassian mixture model
    """
    gmm = GaussianMixture(n_components = n_gaussians, covariance_type=covariance_type)
    gmm.fit(points.astype(np.float64))
    return gmm
//  [python]
def compute_gmm(X, n_components):
	# initialize clustering model
	model = mixture.GaussianMixture(n_components)

	# fit clustering model
	model.fit(X)

	# save clustering results
	K = n_components
	y = model.predict(X)

	# compute criterion value
	crit = model.bic(X)

	# print results
	print("%4d %4d: %8s: K = %4d, crit = %g" % (i, j, "GMM", K, crit))

	return K, y, crit
//  [python]
def compute_gmm(X, n_components):
	# initialize clustering model
	model = sklearn.mixture.GaussianMixture(n_components)

	# fit clustering model
	model.fit(X)

	# save clustering results
	K = n_components
	y = model.predict(X)

	# compute criterion value
	crit = model.bic(X)

	# print results
	print("%4d %4d: %8s: K = %4d, crit = %g" % (i, j, "GMM", K, crit))

	return K, y, crit
//  [python]
def get_3d_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_3d_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 sklearnMethods(clust_method, cl_method_pars, clust_data, n_clusters):
    """
    Find clusters in the 'clust_data' array using the selected algorithm.
    """

    if clust_method == 'KMeans':
        model = skclust.KMeans()

    elif clust_method == 'MiniBatchKMeans':
        model = skclust.MiniBatchKMeans()

    elif clust_method == 'AffinityPropagation':
        model = skclust.AffinityPropagation()

    elif clust_method == 'SpectralClustering':
        model = skclust.SpectralClustering()

    elif clust_method == 'AgglomerativeClustering':
        model = skclust.AgglomerativeClustering()

    elif clust_method == 'GaussianMixture':
        model = skmixture.GMM()

    elif clust_method == 'BayesianGaussianMixture':
        model = skmixture.BayesianGaussianMixture()

    elif clust_method == 'DBSCAN':
        model = skclust.DBSCAN()

    elif clust_method == 'OPTICS':
        model = skclust.OPTICS()

    elif clust_method == 'MeanShift':
        model = skclust.MeanShift()

    elif clust_method == 'Birch':
        model = skclust.Birch()

    elif clust_method == 'HDBSCAN':
        import hdbscan
        model = hdbscan.HDBSCAN()

    # Set parameters for the method (if any)
    if cl_method_pars:
        model.set_params(**cl_method_pars)

    # Only these methods require the number of clusters to be set
    if clust_method in (
            'KMeans', 'MiniBatchKMeans', 'SpectralClustering',
            'AgglomerativeClustering', 'GaussianMixture',
            'BayesianGaussianMixture'):

        if clust_method in ('GaussianMixture', 'BayesianGaussianMixture'):
            model.n_components = n_clusters
        else:
            model.n_clusters = n_clusters

    # Fit the model
    model.fit(clust_data)

    # Extract the labels
    if clust_method in ('GaussianMixture', 'BayesianGaussianMixture'):
        labels = model.predict(clust_data)
        # probs_gmm = model.predict_proba(clust_data).max(1)
    else:
        labels = model.labels_

    return labels
//  [python]
def sklearnMethods(clust_method, cl_method_pars, clust_data, n_clusters):
    """
    Find clusters in the 'clust_data' array using the selected algorithm.
    """

    if clust_method == 'KMeans':
        model = skclust.KMeans()

    elif clust_method == 'MiniBatchKMeans':
        model = skclust.MiniBatchKMeans()

    elif clust_method == 'AffinityPropagation':
        model = skclust.AffinityPropagation()

    elif clust_method == 'SpectralClustering':
        model = skclust.SpectralClustering()

    elif clust_method == 'AgglomerativeClustering':
        model = skclust.AgglomerativeClustering()

    elif clust_method == 'GaussianMixture':
        model = skmixture.GaussianMixture()

    elif clust_method == 'BayesianGaussianMixture':
        model = skmixture.BayesianGaussianMixture()

    elif clust_method == 'DBSCAN':
        model = skclust.DBSCAN()

    elif clust_method == 'OPTICS':
        model = skclust.OPTICS()

    elif clust_method == 'MeanShift':
        model = skclust.MeanShift()

    elif clust_method == 'Birch':
        model = skclust.Birch()

    elif clust_method == 'HDBSCAN':
        import hdbscan
        model = hdbscan.HDBSCAN()

    # Set parameters for the method (if any)
    if cl_method_pars:
        model.set_params(**cl_method_pars)

    # Only these methods require the number of clusters to be set
    if clust_method in (
            'KMeans', 'MiniBatchKMeans', 'SpectralClustering',
            'AgglomerativeClustering', 'GaussianMixture',
            'BayesianGaussianMixture'):

        if clust_method in ('GaussianMixture', 'BayesianGaussianMixture'):
            model.n_components = n_clusters
        else:
            model.n_clusters = n_clusters

    # Fit the model
    model.fit(clust_data)

    # Extract the labels
    if clust_method in ('GaussianMixture', 'BayesianGaussianMixture'):
        labels = model.predict(clust_data)
        # probs_gmm = model.predict_proba(clust_data).max(1)
    else:
        labels = model.labels_

    return labels
//  [python]
def fit_em(xs, k):
  gmm = mixture.GaussianMixture(n_components=k, covariance_type="full")
  labels = jnp.array(gmm.fit_predict(xs))
  mus = jnp.array(gmm.means_)
  covs = jnp.array(gmm.covariances_)
  weights = jnp.log(gmm.weights_)
  return labels, (mus, covs, weights)
//  [python]
def fit_em(xs, k):
  gmm = GaussianMixture(n_components=k, covariance_type="full")
  labels = jnp.array(gmm.fit_predict(xs))
  mus = jnp.array(gmm.means_)
  covs = jnp.array(gmm.covariances_)
  weights = jnp.log(gmm.weights_)
  return labels, (mus, covs, weights)
//  [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_initial_means(X, n_components,init_params, r):
    # Run a GaussianMixture with max_iter=0 to output the initialization means
    gmm = mixture.GaussianMixture(
        n_components=n_components, init_params=init_params, tol=1e-9, max_iter=0, random_state=r
    ).fit(X)
    return gmm.means_
//  [python]
def get_initial_means(X, n_components,init_params, r):
    # Run a GaussianMixture with max_iter=0 to output the initialization means
    gmm = GaussianMixture(
        n_components=n_components, init_params=init_params, tol=1e-9, max_iter=0, random_state=r
    ).fit(X)
    return gmm.means_
//  [python]
def eval_train(eval_loader,model,device,whichnet,queue):   
    CE = nn.CrossEntropyLoss(reduction='none')
    model.eval()
    num_iter = (len(eval_loader.dataset)//eval_loader.batch_size)+1
    losses = torch.zeros(len(eval_loader.dataset))    
    with torch.no_grad():
        for batch_idx, (inputs, targets, index) in enumerate(eval_loader):
            inputs, targets = inputs.to(device), targets.to(device,non_blocking=True) 
            outputs = model(inputs) 
            loss = CE(outputs, targets)  
            for b in range(inputs.size(0)):
                losses[index[b]]=loss[b]       
            sys.stdout.write('\n')
            sys.stdout.write('|%s Evaluating loss Iter[%3d/%3d]\t' %(whichnet,batch_idx,num_iter)) 
            sys.stdout.flush()    
                                    
    losses = (losses-losses.min())/(losses.max()-losses.min())    

    # fit a two-component GMM to the loss
    input_loss = losses.reshape(-1,1)
    gmm = mixture.GaussianMixture(n_components=2,max_iter=10,tol=1e-2,reg_covar=1e-3)
    gmm.fit(input_loss)
    prob = gmm.predict_proba(input_loss) 
    prob = prob[:,gmm.means_.argmin()]         
    queue.put(prob)
//  [python]
def eval_train(eval_loader,model,device,whichnet,queue):   
    CE = nn.CrossEntropyLoss(reduction='none')
    model.eval()
    num_iter = (len(eval_loader.dataset)//eval_loader.batch_size)+1
    losses = torch.zeros(len(eval_loader.dataset))    
    with torch.no_grad():
        for batch_idx, (inputs, targets, index) in enumerate(eval_loader):
            inputs, targets = inputs.to(device), targets.to(device,non_blocking=True) 
            outputs = model(inputs) 
            loss = CE(outputs, targets)  
            for b in range(inputs.size(0)):
                losses[index[b]]=loss[b]       
            sys.stdout.write('\n')
            sys.stdout.write('|%s Evaluating loss Iter[%3d/%3d]\t' %(whichnet,batch_idx,num_iter)) 
            sys.stdout.flush()    
                                    
    losses = (losses-losses.min())/(losses.max()-losses.min())    

    # fit a two-component GMM to the loss
    input_loss = losses.reshape(-1,1)
    gmm = GaussianMixture(n_components=2,max_iter=10,tol=1e-2,reg_covar=1e-3)
    gmm.fit(input_loss)
    prob = gmm.predict_proba(input_loss) 
    prob = prob[:,gmm.means_.argmin()]         
    queue.put(prob)
//  [python]
def filter_umis_gmm(guide_umis, cutoff=1e-10):
    # Init GMM
    gmm = mixture.GMM(n_components=2)
    log_counts = np.array(guide_umis['log_count'])
    # Fit and predict class probs
    gmm_fit = gmm.fit(log_counts.reshape(-1, 1))
    gmm_prob = gmm.predict_proba(log_counts.reshape(-1, 1))
    gmm_class = gmm.predict(log_counts.reshape(-1, 1))
    # Find high class
    pos_mean = np.mean(log_counts[gmm_class==0])
    neg_mean = np.mean(log_counts[gmm_class==1])
    high_group = np.argmax([pos_mean, neg_mean])
    # Assign to df
    guide_umis['is_high'] = gmm_prob[:, high_group] >= cutoff
    return guide_umis
//  [python]
def filter_umis_gmm(guide_umis, cutoff=1e-10):
    # Init GMM
    gmm = GaussianMixture(n_components=2)
    log_counts = np.array(guide_umis['log_count'])
    # Fit and predict class probs
    gmm_fit = gmm.fit(log_counts.reshape(-1, 1))
    gmm_prob = gmm.predict_proba(log_counts.reshape(-1, 1))
    gmm_class = gmm.predict(log_counts.reshape(-1, 1))
    # Find high class
    pos_mean = np.mean(log_counts[gmm_class==0])
    neg_mean = np.mean(log_counts[gmm_class==1])
    high_group = np.argmax([pos_mean, neg_mean])
    # Assign to df
    guide_umis['is_high'] = gmm_prob[:, high_group] >= cutoff
    return guide_umis
//  [python]
def test_gaussian_mixture_comp2(self):
        data = load_iris()
        X = data.data
        model = mixture.GaussianMixture(n_components=2)
        model.fit(X)
        model_onnx = convert_sklearn(
            model,
            "GM",
            [("input", DoubleTensorType([None, 4]))],
            target_opset=TARGET_OPSET,
        )
        self.assertIsNotNone(model_onnx)
        dump_data_and_model(
            X.astype(np.float64)[40:60],
            model,
            model_onnx,
            basename="GaussianMixtureC2Double",
            intermediate_steps=False,
        )
        self._test_score(model, X, TARGET_OPSET)
//  [python]
def test_gaussian_mixture_comp2(self):
        data = load_iris()
        X = data.data
        model = GaussianMixture(n_components=2)
        model.fit(X)
        model_onnx = convert_sklearn(
            model,
            "GM",
            [("input", DoubleTensorType([None, 4]))],
            target_opset=TARGET_OPSET,
        )
        self.assertIsNotNone(model_onnx)
        dump_data_and_model(
            X.astype(np.float64)[40:60],
            model,
            model_onnx,
            basename="GaussianMixtureC2Double",
            intermediate_steps=False,
        )
        self._test_score(model, X, TARGET_OPSET)
//  [python]
def test_model_gaussian_mixture_multiclass(self):
        model, X = self._fit_model_multiclass_classification(
            mixture.GaussianMixture(), load_iris()
        )
        model_onnx = convert_sklearn(
            model,
            "gaussian_mixture",
            [("input", DoubleTensorType([None, X.shape[1]]))],
            target_opset=TARGET_OPSET,
        )
        self.assertIsNotNone(model_onnx)
        dump_data_and_model(
            X, model, model_onnx, basename="SklearnMclGaussianMixtureDouble"
        )
        self._test_score(model, X, TARGET_OPSET)
//  [python]
def test_model_gaussian_mixture_multiclass(self):
        model, X = self._fit_model_multiclass_classification(
            GaussianMixture(), load_iris()
        )
        model_onnx = convert_sklearn(
            model,
            "gaussian_mixture",
            [("input", DoubleTensorType([None, X.shape[1]]))],
            target_opset=TARGET_OPSET,
        )
        self.assertIsNotNone(model_onnx)
        dump_data_and_model(
            X, model, model_onnx, basename="SklearnMclGaussianMixtureDouble"
        )
        self._test_score(model, X, TARGET_OPSET)
//  [python]
def test_model_gaussian_mixture_binary_classification(self):
        model, X = self._fit_model_binary_classification(sklearn.mixture.GaussianMixture(), load_iris())
        for tg in range(min(9, TARGET_OPSET), TARGET_OPSET + 1):
            with self.subTest(target_opset=tg):
                if tg < 11:
                    with self.assertRaises(RuntimeError):
                        model_onnx = convert_sklearn(
                            model,
                            "gaussian_mixture",
                            [("input", DoubleTensorType([None, X.shape[1]]))],
                            target_opset=tg,
                        )
                    continue
                model_onnx = convert_sklearn(
                    model,
                    "gaussian_mixture",
                    [("input", DoubleTensorType([None, X.shape[1]]))],
                    target_opset=tg,
                )
                self.assertIsNotNone(model_onnx)
                dump_data_and_model(
                    X, model, model_onnx, basename="SklearnBinGaussianMixtureDouble"
                )
                self._test_score(model, X, tg)
//  [python]
def test_model_gaussian_mixture_binary_classification(self):
        model, X = self._fit_model_binary_classification(GaussianMixture(), load_iris())
        for tg in range(min(9, TARGET_OPSET), TARGET_OPSET + 1):
            with self.subTest(target_opset=tg):
                if tg < 11:
                    with self.assertRaises(RuntimeError):
                        model_onnx = convert_sklearn(
                            model,
                            "gaussian_mixture",
                            [("input", DoubleTensorType([None, X.shape[1]]))],
                            target_opset=tg,
                        )
                    continue
                model_onnx = convert_sklearn(
                    model,
                    "gaussian_mixture",
                    [("input", DoubleTensorType([None, X.shape[1]]))],
                    target_opset=tg,
                )
                self.assertIsNotNone(model_onnx)
                dump_data_and_model(
                    X, model, model_onnx, basename="SklearnBinGaussianMixtureDouble"
                )
                self._test_score(model, X, tg)
//  [python]
def test_gaussian_mixture_full(self):
        data = load_iris()
        X = data.data
        model = mixture.GaussianMixture(n_components=2, covariance_type="full")
        model.fit(X)
        model_onnx = convert_sklearn(
            model,
            "GM",
            [("input", DoubleTensorType([None, 4]))],
            target_opset=TARGET_OPSET,
        )
        self.assertIsNotNone(model_onnx)
        dump_data_and_model(
            X.astype(np.float64)[40:60],
            model,
            model_onnx,
            basename="GaussianMixtureC2FullDouble",
            intermediate_steps=False,
        )
        self._test_score(model, X, TARGET_OPSET)
//  [python]
def test_gaussian_mixture_full(self):
        data = load_iris()
        X = data.data
        model = GaussianMixture(n_components=2, covariance_type="full")
        model.fit(X)
        model_onnx = convert_sklearn(
            model,
            "GM",
            [("input", DoubleTensorType([None, 4]))],
            target_opset=TARGET_OPSET,
        )
        self.assertIsNotNone(model_onnx)
        dump_data_and_model(
            X.astype(np.float64)[40:60],
            model,
            model_onnx,
            basename="GaussianMixtureC2FullDouble",
            intermediate_steps=False,
        )
        self._test_score(model, X, TARGET_OPSET)
//  [python]
def test_gaussian_mixture_tied(self):
        data = load_iris()
        X = data.data
        model = mixture.GaussianMixture(n_components=2, covariance_type="tied")
        model.fit(X)
        model_onnx = convert_sklearn(
            model,
            "GM",
            [("input", DoubleTensorType([None, 4]))],
            target_opset=TARGET_OPSET,
        )
        self.assertIsNotNone(model_onnx)
        dump_data_and_model(
            X.astype(np.float64)[40:60],
            model,
            model_onnx,
            basename="GaussianMixtureC2TiedDouble",
            intermediate_steps=False,
        )
        self._test_score(model, X, TARGET_OPSET)
//  [python]
def test_gaussian_mixture_tied(self):
        data = load_iris()
        X = data.data
        model = GaussianMixture(n_components=2, covariance_type="tied")
        model.fit(X)
        model_onnx = convert_sklearn(
            model,
            "GM",
            [("input", DoubleTensorType([None, 4]))],
            target_opset=TARGET_OPSET,
        )
        self.assertIsNotNone(model_onnx)
        dump_data_and_model(
            X.astype(np.float64)[40:60],
            model,
            model_onnx,
            basename="GaussianMixtureC2TiedDouble",
            intermediate_steps=False,
        )
        self._test_score(model, X, TARGET_OPSET)