Helper function to easily extract parts of a symmetrical N by N matrix

parse_triangle.py

import numpy as np
import pandas as pd
def parse_triangle(df, condition='upper'):
    '''
    This function grabs the upper triangle of a correlation matrix
    by masking out the bottom triangle (tril) and returns the values. 
    You can use scipy.spatial.distance.squareform to recreate matrix from upper triangle
    
    Args:
      df: pandas or numpy correlation matrix
      condition: 'upper': grabs the upper triangle
                 'pairs': grabs pair of subjects, skipping each diagonal
                 'nonpairs': grabs nonpairs, upper - pairs 
    Returns 
      df: masked df
    '''
    try:
        assert(type(df)==np.ndarray)
    except:
        if type(df)==pd.DataFrame:
            df = df.as_matrix()
        else:
            raise TypeError('Must be np.ndarray or pd.DataFrame')
    if condition =='upper':
        mask = np.triu_indices(df.shape[0], k=1)
        return df[mask] 
    else:
        noDyads = int(df.shape[0]/2)
        if condition =='pairs':
            return np.diag(df,k=1)[range(0,noDyads*2,2)]
        elif condition =='nonpairs':
            mask = np.triu(np.ones(df.shape),k=1).astype(np.bool)
            ix, iy = np.arange(0,noDyads*2,2), np.arange(1,noDyads*2,2)
            for i in np.arange(0,len(ix)):
                mask[ix[i],iy[i]] = False
            return df[mask]
        else:
            raise ValueError('Condition,'+ str(condition) +' not recognized')