SampleSpecificIsoforms.py

#!/usr/bin/env python3

'''
This is a script to identify sample-specific isoforms based on
an isoform read count matrix generated by ESPRESSO (rows are 
detected isoforms and columns are samples)
'''

# Load required libraries
import argparse
import numpy as np
import pandas as pd
import concurrent.futures as cf
from scipy.stats import binom, chi2, chi2_contingency
from statsmodels.stats.multitest import multipletests

def ParseMatrix(infile):
    # Read infile as a pandas dataframe
    matrix = pd.read_csv(infile, sep='\t', header=0)

    # Drop all entries with NA gene ID (corresponds to unknown genes)
    matrix = matrix[matrix['gene_ID'].notna()]

    # Remove any entries where assignment of gene-isoform assignment is not one-to-one
    # The gene_ID column will harbor at least one comma
    matrix = matrix[~matrix['gene_ID'].str.contains(',')]

    # Sort rows by gene ID (and reset row indices); also round all counts to the nearest integer
    matrix = matrix.sort_values(by = ['gene_ID', 'transcript_ID']).reset_index(drop=True).round()

    # Extract set of gene IDs
    genelist = matrix['gene_ID'].unique().tolist()

    return matrix, genelist

def FirstPass(chunk, matrix, cutoff):
    # Initialize an output dataframe for FirstPass
    colNames = ['gene_ID', 'p_value']
    outputDF = pd.DataFrame(columns = colNames)

    # Iterate over gene IDs in given chunk
    for geneID in chunk:
        # Extract dataframe associated with geneID
        rcMatrix = matrix[matrix['gene_ID'] == geneID]

        # Drop isoforms where total count is 0.0
        rcMatrix = rcMatrix[rcMatrix.iloc[:,3:].sum(axis=1) > 0.0]
        nrow = rcMatrix.shape[0]

        # Only work with genes with more than 1 isoform (nrow > 1)
        if nrow > 1:
            # Remove columns where sum of total read counts is 0
            geneCounts = rcMatrix.iloc[:,3:].sum(axis=0)

            # Identify tissues where geneCounts == 0 and drop them from rcMatrix
            rcMatrix = rcMatrix.drop(geneCounts[geneCounts == 0.0].index.tolist(), axis=1)
            ncol = rcMatrix.shape[1]

            # If only one tissue remains, ncol = 4 (only proceed if ncol > 4)
            if ncol > 4:
                # Determine total read count threshold using Cohen's w formula
                # Use significance level of cutoff% (and assume effect size of 0.5)
                threshold = 4*chi2.ppf(1-cutoff,df=(nrow-1)*(ncol-4))
                totalRC = rcMatrix.iloc[:,3:].to_numpy().sum()
                if totalRC > threshold:
                    # Run a chi-square test of homogeneity
                    pval = chi2_contingency(rcMatrix.iloc[:,3:])[1]

                    # Update outputDF
                    outputDF = outputDF.append({'gene_ID': geneID, 'p_value': pval}, ignore_index = True)

    return outputDF

def SecondPass(chunk, matrix):
    # Initialize an output dataframe for SecondPass
    colNames = ['gene_ID', 'transcript_id', 'sample', 'raw_read_count', 'isoform_proportion_global', 'isoform_proportion_sample', 'p_value']
    outputDF = pd.DataFrame(columns = colNames)

    # Iterate over gene IDs in given chunk
    for geneID in chunk:
        # Extract dataframe associated with geneID
        # All of these genes should have at least one isoform, at least two tissues with nonzero gene-level read counts,
        # and a sufficiently large number of total reads (do not need to run filters again)
        rcMatrix = matrix[matrix['gene_ID'] == geneID]
        
        # Drop isoforms where total count is 0.0
        rcMatrix = rcMatrix[rcMatrix.iloc[:,3:].sum(axis=1) > 0.0]

        # Remove columns where sum of total read counts is 0
        geneCounts = rcMatrix.iloc[:,3:].sum(axis=0)
        
        # Identify tissues where geneCounts == 0 and drop them from rcMatrix
        rcMatrix = rcMatrix.drop(geneCounts[geneCounts == 0.0].index.tolist(), axis=1)
        nrow, ncol = rcMatrix.shape

        # Retrieve global isoform proportions from rcMatrix
        rowCounts = rcMatrix.iloc[:,3:].sum(axis=1)
        expected = (rowCounts/rowCounts.sum()).to_list()
        
        # Re-compute geneCounts (after dropping out zero-count columns)
        geneCounts = rcMatrix.iloc[:,3:].sum(axis=0)

        # One-tailed binomial test p-values
        pvalMatrix = np.array([binom.sf(rcMatrix.iloc[i,3:].to_list(), geneCounts.to_list(), expected[i]) +
            binom.pmf(rcMatrix.iloc[i,3:].to_list(), geneCounts.to_list(), expected[i]) for i in range(nrow)])

        ## Build output matrix
        # Flatten pvalMatrix and rcMatrix row-wise
        pvalArr, rcArr = pvalMatrix.flatten(), rcMatrix.iloc[:,3:].to_numpy().flatten()

        # Generate a matrix of isoform proportions (flattened)
        propArr = (rcMatrix.iloc[:,3:]/geneCounts).to_numpy().flatten()
        
        # Repeat gene ID, transcript ID, tissue, and expected isoform proportions
        geneIDArr, transcriptIDArr = np.repeat(geneID, nrow*(ncol-3)), np.repeat(rcMatrix['transcript_ID'].to_numpy(),ncol-3)
        tissueArr, expectedArr = np.tile(rcMatrix.columns[3:],nrow), np.repeat(expected, ncol-3)

        resultDF = pd.DataFrame(list(zip(geneIDArr, transcriptIDArr, tissueArr, rcArr, expectedArr, propArr, pvalArr)), columns = colNames)
        outputDF = outputDF.append(resultDF, ignore_index = True)

    return outputDF

def SampleSpecificIsoforms(infile, threads, cutoff, outfile):
    # Parse isoform read count matrix and extract list of genes
    print('Parsing isoform read count matrix...', flush=True)
    matrix, genelist = ParseMatrix(infile)

    # First pass test: Run a chi-square test of homogeneity on the isoform read count matrix
    # and return a dataframe with each processed gene and corresponding p-value
    # This test will identify genes in which there exists at least one tissue whose isoform 
    # proportions deviate from the global expected proportion (averaged across all tissues)

    print('Running first-pass test...', flush=True)

    # Split genelist into chunks for multiprocessing
    genechunks = [x.tolist() for x in np.array_split(genelist, threads)]

    with cf.ThreadPoolExecutor(max_workers=threads) as executor:
        processes = [executor.submit(FirstPass, chunk, matrix, cutoff) for chunk in genechunks]
    
    # Merge together individual dataframes generated from running FirstPass on chunks of gene IDs
    # df1 is a dataframe with two columns: (i) gene_ID, (ii) p-value
    results1 = [p.result() for p in processes]
    df1 = pd.concat(results1, ignore_index=True)

    # Perform FDR-adjustment of p-values in df1 and keep genes with FDR < cutoff%
    pAdj1 = multipletests(df1['p_value'].tolist(), method='fdr_bh', is_sorted=False, returnsorted=False)[1].tolist()
    df1['p_adj'] = pAdj1

    # Keep gene IDs for which FDR < cutoff% from FirstPass
    filteredGenes = df1[df1['p_adj'] < cutoff]['gene_ID'].to_list()

    # Second pass test: Run a one-tailed binomial test on each entry of the isoform read
    # count matrix. This test will identify isoform-tissue pairs in which the isoform
    # proportion in the given tissue is significantly higher than expected
    # Expected proportions are taken as the global isoform proportions across all samples

    print('Running second-pass test...', flush=True)

    # Split filteredGenes into chunks for multiprocessing
    filterChunks = [x.tolist() for x in np.array_split(filteredGenes, threads)]

    with cf.ThreadPoolExecutor(max_workers=threads) as executor:
        processes = [executor.submit(SecondPass, chunk, matrix) for chunk in filterChunks]
    
    # Merge together individual dataframes generated from running SecondPass on chunks of gene IDs
    results2 = [p.result() for p in processes]
    outDF = pd.concat(results2, ignore_index=True)

    # Perform FDR-adjustment of p-values in outDF and keep isoform-tissue pairs with FDR < cutoff%
    pAdj = multipletests(outDF['p_value'].tolist(), method='fdr_bh', is_sorted=False, returnsorted=False)[1].tolist()
    outDF['p_adj'] = pAdj

    # Drop entries where pAdj < cutoff
    outDF = outDF[outDF['p_adj'] < cutoff]

    # Sort dataframe by adjusted p-value
    outDF = outDF.sort_values(by = 'p_adj')

    # Print outDF to outfile
    outDF.to_csv(outfile, sep='\t', index=False)

def main():
    moduleSummary = 'This is a script to call sample-specific isoforms from an isoform read count matrix generated by ESPRESSO'
    parser = argparse.ArgumentParser(description=moduleSummary)

    # Add arguments
    parser.add_argument('-i', metavar='/path/to/read/count/matrix', required=True,
        help='path to isoform read count matrix generated by ESPRESSO')
    parser.add_argument('-t', metavar='###', required=True,
        help='number of worker threads')
    parser.add_argument('-c', metavar='###', required=True,
        help='FDR threshold (between 0 and 1)') 
    parser.add_argument('-o', metavar='/path/to/output/file', required=True,
        help='path to output file')
    
    # Parse command-line arguments
    args = parser.parse_args()
    infile, threads, cutoff, outfile = args.i, int(args.t), float(args.c), args.o

    print('Isoform read count matrix: ' + infile, flush=True)
    print('Number of threads: ' + str(threads), flush=True)
    print('FDR cutoff: ' + str(cutoff), flush=True)
    print('Output file: ' + outfile, flush=True)

    # Run SampleSpecificIsoforms
    SampleSpecificIsoforms(infile, threads, cutoff, outfile)

if __name__ == '__main__':
    main()