<commit_before>import palettable
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.gridspec import GridSpec
from matplotlib import cm
import scipy.spatial.distance as distance
import scipy.cluster.hierarchy as sch
import numpy as np
import matplotlib as mpl
import pylab
import itertools

from corrplots import scatterfit

__all__ = ['plotHCluster',
            'plotHColCluster',
            'plotCorrHeatmap',
            'mapColors2Labels',
            'computeDMat',
            'computeHCluster']

def clean_axis(ax):
    """Remove ticks, tick labels, and frame from axis"""
    ax.get_xaxis().set_ticks([])
    ax.get_yaxis().set_ticks([])
    for sp in ax.spines.values():
        sp.set_visible(False)

def mapColors2Labels(labels, setStr = 'Set3', cmap = None):
    """Return pd.Series of colors based on labels"""
    if cmap is None:
        N = max(3,min(12,len(np.unique(labels))))
        cmap = palettable.colorbrewer.get_map(setStr,'Qualitative',N).mpl_colors
    cmapLookup = {k:col for k,col in zip(sorted(np.unique(labels)),itertools.cycle(cmap))}
    return labels.map(cmapLookup.get)

def computeDMat(df, metric = None, minN = 1, dfunc = None):
    if dfunc is None:
        if metric in ['spearman', 'pearson']:
            """Anti-correlations are also considered as high similarity and will cluster together"""
            """dmat = 1 - df.corr(method = metric, min_periods = minN).values
            dmat[np.isnan(dmat)] = 1
            """
            dmat = 1 - df.corr(method = metric, min_periods = minN).values**2
            dmat[np.isnan(dmat)] = 1
        elif metric in ['spearman-signed', 'pearson-signed']:
            """Anti-correlations are considered as dissimilar and will NOT cluster together"""
            dmat = (1 - df.corr(method = metric.replace('-signed',''), min_periods = minN).values) / 2
            dmat[np.isnan(dmat)] = 1
        else:
            dmat = distance.squareform(distance.pdist(df, metric = metric))
    else:
        ncols = df.shape[1]
        dmat = np.zeros((ncols,ncols))
        for i in range(ncols):
            for j in range(ncols):
                """Assume its symetrical"""
                if i<=j:
                    tmpdf = df.iloc[:,[i,j]]
                    tmpdf = tmpdf.dropna()
                    if tmpdf.shape[0] >= minN:
                        d = dfunc(df.iloc[:,i],df.iloc[:,j])
                    else:
                        d = np.nan
                    dmat[i,j] = d
                    dmat[j,i] = d
    return dmat

def computeHCluster(df, method = 'complete', metric = 'euclidean', dmat = None, minN = 1):
    """Compute dmat, clusters and dendrogram of df using
    the linkage method and distance metric given"""
    if dmat is None:
        dmat = computeDMat(df, metric, minN = minN)
    clusters = sch.linkage(dmat, method = method)
    den = sch.dendrogram(clusters, color_threshold = np.inf, no_plot=True)
    return dmat, clusters, den

def testData(rows=50,columns=20):
    data=np.random.multivariate_normal(rand(columns),rand(columns,columns),rows)
    df=pd.DataFrame(data,columns=[''.join([lett]*9) for lett in 'ABCDEFGHIJKLMNOPQRST'])
    rowLabels=pd.Series(rand(rows).round(),index=df.index)
    columnLabels=pd.Series(rand(columns).round(),index=df.columns)
    return {'df':df,'row_labels':rowLabels,'col_labels':columnLabels}

def addColorbar(fig,cb_ax,data_ax,label='Correlation'):
    """Colorbar"""
    cb = fig.colorbar(data_ax,cb_ax) # note that we could pass the norm explicitly with norm=my_norm
    cb.set_label(label)
    """Make colorbar labels smaller"""
    for t in cb.ax.yaxis.get_ticklabels():
        t.set_fontsize('small')<commit_msg>Incremental.<commit_after>import palettable
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.gridspec import GridSpec
from matplotlib import cm
import scipy.spatial.distance as distance
import scipy.cluster.hierarchy as sch
import numpy as np
import matplotlib as mpl
import pylab
import itertools

from corrplots import scatterfit

__all__ = ['plotHCluster',
            'plotHColCluster',
            'plotCorrHeatmap',
            'mapColors2Labels',
            'computeDMat',
            'computeHCluster']

def clean_axis(ax):
    """Remove ticks, tick labels, and frame from axis"""
    ax.get_xaxis().set_ticks([])
    ax.get_yaxis().set_ticks([])
    for sp in ax.spines.values():
        sp.set_visible(False)

def mapColors2Labels(labels, setStr = 'Set3', cmap = None):
    """Return pd.Series of colors based on labels"""
    if cmap is None:
        N = max(3,min(12,len(np.unique(labels))))
        cmap = palettable.colorbrewer.get_map(setStr,'Qualitative',N).mpl_colors
    cmapLookup = {k:col for k,col in zip(sorted(np.unique(labels)),itertools.cycle(cmap))}
    return labels.map(cmapLookup.get)

def computeDMat(df, metric = None, minN = 1, dfunc = None):
    if dfunc is None:
        if metric in ['spearman', 'pearson']:
            """Anti-correlations are also considered as high similarity and will cluster together"""
            """dmat = 1 - df.corr(method = metric, min_periods = minN).values
            dmat[np.isnan(dmat)] = 1
            """
            dmat = 1 - df.corr(method = metric, min_periods = minN).values**2
            dmat[np.isnan(dmat)] = 1
        elif metric in ['spearman-signed', 'pearson-signed']:
            """Anti-correlations are considered as dissimilar and will NOT cluster together"""
            dmat = (1 - df.corr(method = metric.replace('-signed',''), min_periods = minN).values) / 2
            dmat[np.isnan(dmat)] = 1
        else:
            dmat = distance.squareform(distance.pdist(df.T, metric = metric))
    else:
        ncols = df.shape[1]
        dmat = np.zeros((ncols,ncols))
        for i in range(ncols):
            for j in range(ncols):
                """Assume its symetrical"""
                if i<=j:
                    tmpdf = df.iloc[:,[i,j]]
                    tmpdf = tmpdf.dropna()
                    if tmpdf.shape[0] >= minN:
                        d = dfunc(df.iloc[:,i],df.iloc[:,j])
                    else:
                        d = np.nan
                    dmat[i,j] = d
                    dmat[j,i] = d
    assert dmat.shape[0] == dmat.shape[1]
    assert dmat.shape[0] == df.shape[1]
    return dmat

def computeHCluster(df, method = 'complete', metric = 'euclidean', dmat = None, minN = 1):
    """Compute dmat, clusters and dendrogram of df using
    the linkage method and distance metric given"""
    if dmat is None:
        dmat = computeDMat(df, metric, minN = minN)
    clusters = sch.linkage(dmat, method = method)
    den = sch.dendrogram(clusters, color_threshold = np.inf, no_plot=True)
    return dmat, clusters, den

def testData(rows=50,columns=20):
    data=np.random.multivariate_normal(rand(columns),rand(columns,columns),rows)
    df=pd.DataFrame(data,columns=[''.join([lett]*9) for lett in 'ABCDEFGHIJKLMNOPQRST'])
    rowLabels=pd.Series(rand(rows).round(),index=df.index)
    columnLabels=pd.Series(rand(columns).round(),index=df.columns)
    return {'df':df,'row_labels':rowLabels,'col_labels':columnLabels}

def addColorbar(fig,cb_ax,data_ax,label='Correlation'):
    """Colorbar"""
    cb = fig.colorbar(data_ax,cb_ax) # note that we could pass the norm explicitly with norm=my_norm
    cb.set_label(label)
    """Make colorbar labels smaller"""
    for t in cb.ax.yaxis.get_ticklabels():
        t.set_fontsize('small')
{'fraction_non_alphanumeric': 0.12763140473983847, 'fraction_numerical': 0.006619886137958427, 'mean_word_length': 2.70839469808542, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, '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': '6562786', 'n_tokens_mistral': 2502, 'n_tokens_neox': 2418, 'n_words': 589}