transcriptomics.py

import io
import logging
import os
import re
from glob import glob
from typing import Union

import dask.dataframe as dd
import numpy as np
import pandas as pd
import validators
# from Bio.UniProt import GOA
from dask import delayed

from .database.base import Annotatable
from .io.files import get_pkg_data_filename
from .transforms.df import drop_duplicate_columns

__all__ = ['Expression', 'MessengerRNA', 'MicroRNA', 'LncRNA', ]

class Expression(object):
    """This class handles importing of any quantitative omics data that is
        in a table format (e.g. csv, tsv, excel). Pandas will load the DataFrame
        from file with the user-specified columns and genes column name, then
        tranpose it such that the rows are samples and columns are
        gene/transcript/peptides. The user will also specify the index argument,
        which specifies if the genes are ensembl genes ID or gene name, or
        transcripts id/names. The user should be careful about choosing the
        right genes index which makes it easier to annotate functional,
        sequence, and interaction data to it. The dataframe should only contain
        numeric values besides the genes_col_name and the sample barcode id
        indices.

    """
    expressions: pd.DataFrame
    def __init__(self, data, transpose, gene_index=None, usecols=None, gene_level=None, sample_level="sample_index",
                 transform_fn=None, dropna=False, npartitions=None, **kwargs):
        """This constructor will create a DataFrame
        from file with the user-specified columns and genes column name, then
        tranpose it such that the rows are samples and columns are
        gene/transcript/peptides. The user will also specify the index argument,
        which specifies if the genes are ensembl genes ID or gene name, or
        transcripts id/names. The user should be careful about choosing the
        right genes index which makes it easier to annotate functional,
        sequence, and interaction data to it. The dataframe should only contain
        numeric values besides the genes_col_name and the sample barcode id
        indices.

        Args:
            data (str, byte-like, pandas.DataFrame): Path or file stream of the
                table file to import. If a pandas DataFrame is passed, then
                import this dataframe and skip preprocessing steps.
            transpose (bool): True if given data table has samples or columns
                and variables for rows. False if the table has samples for row
                index, and gene names as columns.
            gene_index (str): The column name of gene/transcript/protein to
                index by.
            usecols: A regex string to import column names from the table.
                Columns names imported are string match, separated by "|".
            gene_level (str): {"gene", "transcript", "peptide"} Chooses the
                level of the gene/transcript/peptide of the genes list in this
                expression data. The expression DataFrame's index will be
                renamed to this.
            sample_level (str): {"sample_index", "patient_index"} Chooses the
                level of the patient/sample/aliquot indexing.
            transform_fn (bool): default False A callable function to transform
                single values.
            dropna (bool): Whether to drop rows with null values
            npartitions (int): [0-n], default 0 If 0, then uses a Pandas
                DataFrame, if >1, then creates an off-memory Dask DataFrame with
                n partitions
            **kwargs: Any arguments to pass into pd.read_table(**kwargs)
        """
        self.gene_level = gene_level
        self.sample_level = sample_level

        df = self.load_dataframe(data, transpose=transpose, usecols=usecols, gene_index=gene_index, dropna=dropna,
                                 **kwargs)

        self.expressions = self.preprocess_table(
            df,
            usecols=usecols,
            gene_index=gene_index,
            transposed=transpose,
            dropna=dropna,
        )

        # TODO load DD from file directly
        if npartitions and isinstance(self.expressions, pd.DataFrame):
            self.expressions = dd.from_pandas(self.expressions, npartitions=npartitions)

        if gene_level is not None:
            self.expressions.columns.name = gene_level

        self.expressions.index.name = self.sample_level

        if callable(transform_fn):
            self.expressions = self.expressions.applymap(transform_fn)
        elif transform_fn == "log2":
            self.expressions = self.expressions.applymap(
                lambda x: np.log2(x + 1))

    @property
    def gene_index(self):
        return self.expressions.columns.name

    def load_dataframe(self,
                       data: Union[str, pd.DataFrame, dd.DataFrame, io.StringIO],
                       transpose: bool,
                       usecols: str,
                       gene_index: str,
                       dropna: bool, **kwargs) -> pd.DataFrame:
        """Reading table data inputs to create a DataFrame.

        Args:
            data: either a file path, a glob file path (e.g. "table-*.tsv"), a
                pandas.DataFrame, or a dask DataFrame.
            transpose (bool): True if table oriented with samples columns, else
                False.
            usecols (str): A regex string to select columns. Default None.
            gene_index (str): The column name what contains the gene names or IDs.
            dropna (bool): Whether to drop rows with null values

        Returns:
            Union[pd.DataFrame, dd.DataFrame]: The loaded dataframe.
        """
        if isinstance(data, (pd.DataFrame, dd.DataFrame)):
            df = data

        elif isinstance(data, str) and "*" in data:
            # TODO implement handling for multiple file ByteIO
            df = self.load_dataframe_glob(globstring=data, usecols=usecols, gene_index=gene_index, transpose=transpose,
                                          dropna=dropna, **kwargs)

        elif isinstance(data, io.StringIO):
            # Needed since the file was previous read to extract columns information
            data.seek(0)
            df = pd.read_table(data, **kwargs)


        elif isinstance(data, str) and validators.url(data):
            dataurl, filename = os.path.split(data)
            file = get_pkg_data_filename(dataurl + "/", filename)
            df = pd.read_table(file, **kwargs)

        elif isinstance(data, str) and os.path.isfile(data):
            df = pd.read_table(data, sep=None, engine="python")


        else:
            raise FileNotFoundError(data)

        return df

    def preprocess_table(self,
        df: Union[pd.DataFrame, dd.DataFrame],
        usecols: str = None,
        gene_index: str = None,
        transposed: bool = True,
        sort_index: bool = False,
        dropna: bool = True,
    ):
        """This function preprocesses the expression table files where columns
        are samples and rows are gene/transcripts :param df: A Dask or Pandas
        DataFrame :type df: DataFrame :param usecols: A regular expression
        string for the column names to fetch. :type usecols: str :param
        gene_index: The column name containing the gene/transcript names or
        id's. :type gene_index: str :param transposed: Default True. Whether to
        transpose the dataframe so columns are genes (features) and rows are
        samples.

        Args:
            df (pd.DataFrame):
            usecols (str):
            gene_index (str):
            transposed (bool):
            sort_index (bool):
            dropna (bool):

        Returns:
            Union[pd.DataFrame, dd.DataFrame]: a processed Dask DataFrame
        """
        # Filter columns
        if usecols is not None and isinstance(usecols, str):
            if gene_index not in usecols:
                # include index column in the filter regex query
                usecols = (usecols + "|" + gene_index)

            if isinstance(df, pd.DataFrame):
                df = df.filter(regex=usecols)
            elif isinstance(df, dd.DataFrame):
                columns = list(filter(re.compile(usecols).match, df.columns))
                df = df[columns]

        elif usecols is not None and isinstance(usecols, list):
            if gene_index not in usecols:
                usecols.append(gene_index)
            df =  df[usecols]

        # Drop duplicate column names
        df = drop_duplicate_columns(df)

        # Drop NA geneID rows
        if dropna:
            df.dropna(axis=0, inplace=True)

        if gene_index is not None and df.index.name != gene_index:
            df = df.set_index(gene_index)

        # Needed for Dask Delayed
        if sort_index is True:
            df = df.sort_index(axis=0, ascending=True)

        # Select only numerical columns
        df = df.select_dtypes(include="number")

        # Transpose dataframe to sample rows and gene columns
        if transposed:
            df = df.T
            # Drop duplicate genes
            df = drop_duplicate_columns(df)

        return df

    def load_dataframe_glob(self, globstring: str, usecols: str, gene_index: str, transpose: bool, dropna: bool,
                            **kwargs):
        """
        Args:
            globstring (str):
            usecols (str):
            gene_index (str):
            transpose (bool):
        Returns:
            dd.DataFrame
        """

        def convert_numerical_to_float(df: pd.DataFrame):
            cols = df.columns[~df.dtypes.eq('object')]
            df[cols] = df[cols].astype(float)
            return df

        filenames = []

        lazy_dataframes = []
        for file_path in glob(globstring):
            filenames.append(os.path.split(file_path)[1])

            df = delayed(pd.read_table)(file_path, **kwargs)
            # df = delayed(convert_numerical_to_float)(df)
            df = delayed(self.preprocess_table)(
                df,
                usecols,
                gene_index,
                transpose,
                True, # sort_index
                dropna)
            lazy_dataframes.append(df)

        logging.info("Files matched: {}".format(filenames))

        return dd.from_delayed(lazy_dataframes, divisions=None, verify_meta=True)

    def set_genes_index(self, index: str, old_index: str):
        """
        Args:
            index (str):
            old_index (str):
        """
        assert isinstance(self, Annotatable) and isinstance(self, Expression)
        # Change gene name columns in expressions
        rename_dict = self.get_rename_dict(from_index=old_index,
                                           to_index=index)
        self.expressions.rename(columns=rename_dict, inplace=True)
        self.gene_index = index

        # Change index name in annotation
        self.set_index(index)

    def drop_genes(self, gene_ids: str):
        """Drop columns representing genes/rna/proteins in self.expressions
        dataframe.

        Args:
            gene_ids (str): list of strings that are a subset of the columns
                list
        """
        self.expressions = self.expressions.drop(gene_ids, axis=1)
        if hasattr(self, "annotations") and not self.annotations.empty:
            self.annotations = self.annotations.drop(gene_ids, axis=0)

    def drop_samples(self, sample_ids):
        """
        Args:
            sample_ids:
        """
        self.expressions = self.expressions.drop(sample_ids, axis=0)

    @classmethod
    def name(cls):
        raise NotImplementedError

    def get_genes_list(self, level: int = None):
        """
        Args:
            level (int): Default None. Only needed if gene index is a :class:`pd.MultiIndex`
        """
        index = self.expressions.columns

        if isinstance(index, pd.MultiIndex):
            return index.get_level_values(
                self.gene_index if level is None else level)
        else:
            return index

    def get_samples_list(self, level=None):
        """
        Args:
            level:
        """
        index = self.expressions.index
        if isinstance(index, pd.MultiIndex):
            return index.get_level_values(
                self.gene_index if level is None else level)
        else:
            return index

    samples = property(get_samples_list)
    features = property(get_genes_list)


class LncRNA(Expression, Annotatable):
    def __init__(
        self,
        data,
        transpose,
        gene_index=None,
        usecols=None,
        gene_level=None,
        sample_level="sample_index",
        transform_fn=None,
        dropna=False,
        npartitions=None,
        cohort_name=None,
    ):
        """
        Args:
            data:
            transpose:
            gene_index:
            usecols:
            gene_level:
            sample_level:
            transform_fn:
            dropna:
            npartitions:
            cohort_name:
        """
        super().__init__(data=data, transpose=transpose, gene_index=gene_index, usecols=usecols,
                         gene_level=gene_level, sample_level=sample_level, transform_fn=transform_fn,
                         dropna=dropna, npartitions=npartitions, cohort_name=cohort_name)

    @classmethod
    def name(cls):
        return cls.__name__


class MessengerRNA(Expression, Annotatable):
    def __init__(
        self,
        data,
        transpose,
        gene_index=None,
        usecols=None,
        gene_level=None,
        sample_level="sample_index",
        transform_fn=None,
        dropna=False,
        npartitions=None,
        cohort_name=None,
    ):
        super().__init__(data=data, transpose=transpose, gene_index=gene_index, usecols=usecols,
                         gene_level=gene_level, sample_level=sample_level, transform_fn=transform_fn,
                         dropna=dropna, npartitions=npartitions, cohort_name=cohort_name)

    @classmethod
    def name(cls):
        return cls.__name__


class MicroRNA(Expression, Annotatable):
    def __init__(
        self,
        data,
        transpose,
        gene_index=None,
        usecols=None,
        gene_level=None,
        sample_level="sample_index",
        transform_fn=None,
        dropna=False,
        npartitions=None,
        cohort_name=None,
    ):
        super().__init__(data=data, transpose=transpose, gene_index=gene_index, usecols=usecols,
                         gene_level=gene_level, sample_level=sample_level, transform_fn=transform_fn,
                         dropna=dropna, npartitions=npartitions, cohort_name=cohort_name)

    @classmethod
    def name(cls):
        return cls.__name__