<commit_before># basic functions for visualization of clustering, state estimation, lineage

import matplotlib.pyplot as plt
from sklearn.decomposition import PCA

def visualize_poisson_w(w, labels, filename, method='pca', figsize=(18,10), title='', **scatter_options):
    """
    Saves a scatter plot of a visualization of W, the result from Poisson SE.
    """
    if method == 'pca':
        pca = PCA(2)
        r_dim_red = pca.fit_transform(w.T).T
    elif method == 'tsne':
        pass
    else:
        print("Method is not available. use 'pca' (default) or 'tsne'.")
        return
    visualize_dim_red(r_dim_red, labels, filename, figsize, title, **scatter_options)

def visualize_dim_red(r, labels, filename=None, figsize=(18,10), title='', legend=True, label_map=None, **scatter_options):
    """
    Saves a scatter plot of a (2,n) matrix r, where each column is a cell.

    Args:
        r (array): (2,n) matrix
        labels (array): (n,) array of ints
        filename (string): string to save the output graph. If None, then this just displays the plot.
        figsize (tuple): Default: (18, 10)
        title (string): graph title
        legend (bool): Default: True
        label_map (dict): map of labels to label names. Default: None
    """
    fig = plt.figure(figsize=figsize)
    plt.cla()
    for i in set(labels):
        label = i
        if label_map is not None:
            label = label_map[i]
        plt.scatter(r[0, labels==i], r[1, labels==i], label=label, **scatter_options)
    plt.title(title)
    if legend:
        plt.legend()
    if filename is not None:
        plt.savefig(filename, dpi=100)
        plt.close()
    return fig<commit_msg>updated vis with more options<commit_after># basic functions for visualization of clustering, state estimation, lineage

import matplotlib.pyplot as plt
from sklearn.decomposition import PCA

def visualize_poisson_w(w, labels, filename, method='pca', figsize=(18,10), title='', **scatter_options):
    """
    Saves a scatter plot of a visualization of W, the result from Poisson SE.
    """
    if method == 'pca':
        pca = PCA(2)
        r_dim_red = pca.fit_transform(w.T).T
    elif method == 'tsne':
        pass
    else:
        print("Method is not available. use 'pca' (default) or 'tsne'.")
        return
    visualize_dim_red(r_dim_red, labels, filename, figsize, title, **scatter_options)

def visualize_dim_red(r, labels, filename=None, figsize=(18,10), title='', legend=True, label_map=None, label_scale=False, **scatter_options):
    """
    Saves a scatter plot of a (2,n) matrix r, where each column is a cell.

    Args:
        r (array): (2,n) matrix
        labels (array): (n,) array of ints/strings or floats. Can be None.
        filename (string): string to save the output graph. If None, then this just displays the plot.
        figsize (tuple): Default: (18, 10)
        title (string): graph title
        legend (bool): Default: True
        label_map (dict): map of labels to label names. Default: None
        label_scale (bool): True if labels is should be treated as floats. Default: False
    """
    fig = plt.figure(figsize=figsize)
    plt.cla()
    if not label_scale:
        for i in set(labels):
            label = i
            if label_map is not None:
                label = label_map[i]
            plt.scatter(r[0, labels==i], r[1, labels==i], label=label, **scatter_options)
    else:
        if labels is None:
            plt.scatter(r[0,:], r[1,:], **scatter_options)
        else:
            plt.scatter(r[0,:], r[1,:], c=labels/labels.max(), **scatter_options)
    plt.title(title)
    if legend:
        plt.legend()
    if filename is not None:
        plt.savefig(filename, dpi=100)
        plt.close()
    return fig
{'fraction_non_alphanumeric': 0.11933174224343675, 'fraction_numerical': 0.011667992574913817, 'mean_word_length': 2.7644710578842315, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3895050', 'n_tokens_mistral': 1221, 'n_tokens_neox': 1175, 'n_words': 394}