gmm_2d_data_maker.py

import numpy as np
from itertools import chain

def gmm_2d_data_maker(pi_list_of_mixing_weights, 
                      mu_list_of_vectors_means, 
                      R_list_of_arrays_covs, 
                      N_number_of_samples):  
    """
    Creates a two dimensional list of points about N_number_of_samples long
    with labels.
    """
    
    assert np.isclose(sum(pi_list_of_mixing_weights), 1.0), "Your mixing probabilities {} don't add up to 1! They currently sum to {}.".format(pi_list_of_mixing_weights, sum(pi_list_of_mixing_weights))
    
    numb_to_generate = [N_number_of_samples*pi_k for pi_k in pi_list_of_mixing_weights]
    numb_to_generate = list(map(int, numb_to_generate))
    X = []
    y = []
    for k, num_samples_k in enumerate(numb_to_generate):
        X.append(np.random.multivariate_normal(mean=mu_list_of_vectors_means[k], cov=R_list_of_arrays_covs[k], size=num_samples_k))
        y.append([k]*num_samples_k)
    
    X = np.vstack(X)
    y = np.array(list(chain.from_iterable(y)))
    Xy = np.column_stack((X,y))
    np.random.shuffle(Xy)
    
    X = Xy[:,0:2]
    y = Xy[:,-1]
    
    return X, y