Datasets
Models
Diff of
/openomics/database/sequence.py
[000000]
..
[548210]
Switch to side-by-side view
--- a +++ b/openomics/database/sequence.py @@ -0,0 +1,1025 @@ +import os +import re +import traceback +from abc import abstractmethod +from collections import defaultdict, OrderedDict +from typing import Union, List, Callable, Dict, Tuple, Optional, Iterable + +import numpy as np +import pandas as pd +import tqdm +from Bio import SeqIO +from Bio.SeqFeature import ExactPosition +from dask import dataframe as dd +from logzero import logger +from pyfaidx import Fasta +from six.moves import intern + +import openomics +from openomics.io.read_gtf import read_gtf +from .base import Database +from ..io.files import select_files_with_ext +from ..transforms.agg import get_agg_func +from ..transforms.df import drop_duplicate_columns + +__all__ = ['GENCODE', 'UniProt', 'MirBase', 'RNAcentral'] + +SEQUENCE_COL = 'sequence' + + +class SequenceDatabase(Database): + """Provides a series of methods to extract sequence data from + SequenceDataset. + """ + + def __init__(self, **kwargs): + """ + Args: + **kwargs: + """ + super().__init__(**kwargs) + self.close() + + @abstractmethod + def load_sequences(self, fasta_file: str, index=None, keys: Union[pd.Index, List[str]] = None, blocksize=None) \ + -> pd.DataFrame: + """Returns a pandas DataFrame containing the fasta sequence entries. + With a column named 'sequence'. + + Args: + index (): + fasta_file (str): path to the fasta file, usually as + self.file_resources[<file_name>] + keys (pd.Index): list of keys to + blocksize: + """ + raise NotImplementedError + + @abstractmethod + def get_sequences(self, index: str, omic: str, agg: str, **kwargs) -> Union[pd.Series, Dict]: + """Returns a dictionary where keys are 'index' and values are + sequence(s). + + Args: + index (str): {"gene_id", "gene_name", "transcript_id", + "transcript_name"} + omic (str): {"lncRNA", "microRNA", "messengerRNA"} + agg (str): {"all", "shortest", "longest"} + **kwargs: any additional argument to pass to + SequenceDataset.get_sequences() + """ + raise NotImplementedError + + @staticmethod + def aggregator_fn(agg: Union[str, Callable] = None) -> Callable: + """Returns a function used aggregate a list of sequences from a groupby + on a given key. + + Args: + agg: One of ("all", "shortest", "longest", "first"), default "all". If "all", + then return a list of sequences. + """ + if agg == "all": + agg_func = lambda x: list(x) if not isinstance(x, str) else x + elif agg == "shortest": + agg_func = lambda x: min(x, key=len) if isinstance(x, list) else x + elif agg == "longest": + agg_func = lambda x: max(x, key=len) if isinstance(x, list) else x + elif agg == 'first': + agg_func = lambda x: x[0] if isinstance(x, list) else x + elif callable(agg): + return agg + else: + raise Exception( + "agg_sequences argument must be one of {'all', 'shortest', 'longest'}" + ) + return agg_func + +class GENCODE(SequenceDatabase): + """Loads the GENCODE database from https://www.gencodegenes.org/ . + + Default path: ftp://ftp.ebi.ac.uk/pub/databases/gencode/Gencode_human/release_32/ . + Default file_resources: { + "basic.annotation.gtf": "gencode.v32.basic.annotation.gtf.gz", + "long_noncoding_RNAs.gtf": "gencode.v32.long_noncoding_RNAs.gtf.gz", + "lncRNA_transcripts.fa": "gencode.v32.lncRNA_transcripts.fa.gz", + "transcripts.fa": "gencode.v32.transcripts.fa.gz", + } + """ + def __init__( + self, + path="ftp://ftp.ebi.ac.uk/pub/databases/gencode/Gencode_human/release_32/", + file_resources=None, + col_rename=None, + blocksize=0, + remove_version_num=False, + **kwargs + ): + """ + Args: + path: + file_resources: + col_rename: + blocksize: + remove_version_num (bool): Whether to drop the version number on the + ensembl ID. + """ + if file_resources is None: + file_resources = { + "basic.annotation.gtf.gz": "gencode.v32.basic.annotation.gtf.gz", + "long_noncoding_RNAs.gtf.gz": "gencode.v32.long_noncoding_RNAs.gtf.gz", + "lncRNA_transcripts.fa.gz": "gencode.v32.lncRNA_transcripts.fa.gz", + "transcripts.fa.gz": "gencode.v32.transcripts.fa.gz", + } + + self.remove_version_num = remove_version_num + + super().__init__(path=path, file_resources=file_resources, col_rename=col_rename, blocksize=blocksize, **kwargs) + + def load_dataframe(self, file_resources, blocksize=None): + """ + Loads the GENCODE annotation file into a pandas or dask DataFrame. + Args: + file_resources: + blocksize: + + Returns: + + """ + gtf_files = select_files_with_ext(file_resources, ".gtf") + if not gtf_files: + gtf_files = select_files_with_ext(file_resources, ".gtf.gz") + + ddfs = [] + for filename, filepath in gtf_files.items(): + if blocksize and not isinstance(filepath, str): continue + ddf = read_gtf(filepath, blocksize=blocksize, + compression="gzip" if filename.endswith(".gz") else None) + + ddfs.append(ddf) + + annotation_df = dd.concat(ddfs, interleave_partitions=True, + ignore_unknown_divisions=True) if blocksize else pd.concat(ddfs) + + if self.remove_version_num: + annotation_df = annotation_df.assign( + gene_id=annotation_df["gene_id"].str.replace("[.]\d*", "", regex=True), + transcript_id=annotation_df["transcript_id"].str.replace("[.]\d*", "", regex=True)) + + return annotation_df + + def load_sequences(self, fasta_file: str, index=None, keys: pd.Index = None, blocksize=None): + """ + Args: + index (): + keys (): + fasta_file: + blocksize: + """ + if hasattr(self, '_seq_df_dict') and fasta_file in self._seq_df_dict: + return self._seq_df_dict[fasta_file] + + def get_transcript_id(x): + key = x.split('|')[0] # transcript_id + if self.remove_version_num: + return re.sub("[.]\d*", "", key) + else: + return key + + fa = Fasta(fasta_file, key_function=get_transcript_id, as_raw=True) + + entries = [] + for key, record in tqdm.tqdm(fa.items(), desc=str(fasta_file)): + if keys is not None and key not in keys: continue + + attrs = record.long_name.split("|") + record_dict = { + "transcript_id": attrs[0], + "gene_id": attrs[1], + "gene_name": attrs[5], + "transcript_name": attrs[4], + "transcript_length": attrs[6], + "transcript_biotype": intern(attrs[7]), + SEQUENCE_COL: str(record), + } + + entries.append(record_dict) + + seq_df = pd.DataFrame(entries) + if blocksize: + seq_df = dd.from_pandas(seq_df, chunksize=blocksize) + + if self.remove_version_num: + seq_df["gene_id"] = seq_df["gene_id"].str.replace("[.]\d*", "", regex=True) + seq_df["transcript_id"] = seq_df["transcript_id"].str.replace("[.]\d*", "", regex=True) + + # Cache the seq_df + if not hasattr(self, '_seq_df_dict'): + self._seq_df_dict = {} + if keys is not None: + self._seq_df_dict[fasta_file] = seq_df + + return seq_df + + def get_sequences(self, index: Union[str, Tuple[str]], omic: str, agg: str = 'all', biotypes: List[str] = None): + """ + Args: + index (str): + omic (str): + agg (str): + biotypes (List[str]): + """ + agg_func = self.aggregator_fn(agg) + + # Parse lncRNA & mRNA fasta + if omic == openomics.MessengerRNA.name(): + fasta_file = self.file_resources["transcripts.fa"] + elif omic == openomics.LncRNA.name(): + fasta_file = self.file_resources["lncRNA_transcripts.fa"] + else: + raise Exception("omic argument must be one of {'MessengerRNA', 'LncRNA'}") + + assert isinstance(fasta_file, str), \ + f"Fasta file provided in `file_resources` must be an uncompressed .fa file. Given {fasta_file}." + + seq_df = self.load_sequences(fasta_file) + + if "gene" in index: + if biotypes: + seq_df = seq_df[seq_df["transcript_biotype"].isin(biotypes)] + else: + print("INFO: You can pass in a list of transcript biotypes to filter using the argument 'biotypes'.") + + return seq_df.groupby(index)[SEQUENCE_COL].agg(agg_func) + + else: + return seq_df.groupby(index)[SEQUENCE_COL].first() + + def get_rename_dict(self, from_index="gene_id", to_index="gene_name"): + """ + Args: + from_index: + to_index: + """ + ensembl_id_to_gene_name = pd.Series( + self.data[to_index].values, index=self.data[from_index]).to_dict() + return ensembl_id_to_gene_name + + +class UniProt(SequenceDatabase): + COLUMNS_RENAME_DICT = { + # idmapping_selected.tab + "UniProtKB-AC": 'protein_id', + "UniProtKB-ID": 'protein_name', + "Ensembl": "gene_id", + "Ensembl_TRS": "transcript_id", + "Ensembl_PRO": "protein_embl_id", + "NCBI-taxon": "species_id", + "GeneID (EntrezGene)": "entrezgene_id", + "GO": "go_id", + # FASTA headers + "OS": 'species', "OX": 'species_id', 'GN': 'gene_name', 'PE': 'ProteinExistence', 'SV': "version", + # UniProt XML headers + "accession": "UniProtKB-AC", "name": "protein_name", "gene": "gene_name", "keyword": "keywords", + "geneLocation": "subcellular_location", + + } + + SPECIES_ID_NAME = { + '10090': 'MOUSE', '10116': 'RAT', '226900': 'BACCR', '243273': 'MYCGE', '284812': 'SCHPO', '287': 'PSEAI', + '3702': 'ARATH', '99287': 'SALTY', '44689': 'DICDI', '4577': 'MAIZE', '559292': 'YEAST', '6239': 'CAEEL', + '7227': 'DROME', '7955': 'DANRE', '83333': 'ECOLI', '9606': 'HUMAN', '9823': 'PIG', } + + SPECIES_ID_TAXONOMIC = { + 'HUMAN': 'human', 'MOUSE': 'rodents', 'RAT': 'rodents', 'BACCR': 'bacteria', 'MYCGE': 'bacteria', + 'SCHPO': 'fungi', 'PSEAI': 'bacteria', 'ARATH': 'plants', 'SALTY': 'bacteria', 'DICDI': 'bacteria', + 'MAIZE': 'plants', 'YEAST': 'fungi', 'CAEEL': 'vertebrates', 'DROME': 'invertebrates', 'DANRE': 'vertebrates', + 'ECOLI': 'bacteria', 'PIG': 'mammals', + } + + def __init__(self, path="https://ftp.uniprot.org/pub/databases/uniprot/current_release/", + file_resources: Dict[str, str] = None, + species_id: str = "9606", remove_version_num=True, + index_col='UniProtKB-AC', keys=None, + col_rename=COLUMNS_RENAME_DICT, + **kwargs): + """ + Loads the UniProt database from https://uniprot.org/ . + + Default path: 'https://ftp.uniprot.org/pub/databases/uniprot/current_release/' + Default file_resources: { + file_resources['uniprot_sprot.xml.gz'] = "knowledgebase/complete/uniprot_sprot.xml.gz + file_resources['uniprot_trembl.xml.gz'] = "knowledgebase/complete/uniprot_trembl.xml.gz + file_resources["idmapping_selected.tab.gz"] = "knowledgebase/idmapping/idmapping_selected.tab.gz' + file_resources["proteomes.tsv"] = "https://rest.uniprot.org/proteomes/stream?compressed=true& + fields=upid%2Corganism%2Corganism_id&format=tsv&query=%28%2A%29%20AND%20%28proteome_type%3A1%29" + file_resources['speclist.txt'] = 'https://ftp.uniprot.org/pub/databases/uniprot/current_release/ + knowledgebase/complete/docs/speclist' + } + + Args: + path: + file_resources: + col_rename: + verbose: + blocksize: + """ + self.species_id = species_id + self.species = UniProt.SPECIES_ID_NAME[species_id] if species_id in UniProt.SPECIES_ID_NAME else None + self.taxonomic_id = UniProt.SPECIES_ID_TAXONOMIC[ + self.species] if self.species in UniProt.SPECIES_ID_TAXONOMIC else None + self.remove_version_num = remove_version_num + + if file_resources is None: + file_resources = {} + + file_resources['uniprot_sprot.xml.gz'] = os.path.join(path, "knowledgebase/complete/uniprot_sprot.xml.gz") + file_resources['uniprot_trembl.xml.gz'] = os.path.join(path, "knowledgebase/complete/uniprot_trembl.xml.gz") + file_resources["idmapping_selected.tab.gz"] = os.path.join( + path, "knowledgebase/idmapping/idmapping_selected.tab.gz") + + if self.species: + file_resources['uniprot_sprot.xml.gz'] = os.path.join( + path, "knowledgebase/taxonomic_divisions/", f'uniprot_sprot_{self.taxonomic_id}.xml.gz') + file_resources['uniprot_trembl.xml.gz'] = os.path.join( + path, "knowledgebase/taxonomic_divisions/", f'uniprot_trembl_{self.taxonomic_id}.xml.gz') + file_resources["idmapping_selected.tab.gz"] = os.path.join( + path, "knowledgebase/idmapping/by_organism/", + f'{self.species}_{self.species_id}_idmapping_selected.tab.gz') + + file_resources["proteomes.tsv"] = \ + "https://rest.uniprot.org/proteomes/stream?compressed=true&fields=upid%2Corganism%2Corganism_id&format=tsv&query=%28%2A%29%20AND%20%28proteome_type%3A1%29" + + super().__init__(path=path, file_resources=file_resources, index_col=index_col, keys=keys, + col_rename=col_rename, + **kwargs) + + def load_dataframe(self, file_resources, blocksize=None): + """ + Args: + file_resources: + blocksize: + """ + # Load idmapping_selected.tab + args = dict( + names=['UniProtKB-AC', 'UniProtKB-ID', 'GeneID (EntrezGene)', 'RefSeq', 'GI', 'PDB', 'GO', 'UniRef100', + 'UniRef90', 'UniRef50', 'UniParc', 'PIR', 'NCBI-taxon', 'MIM', 'UniGene', 'PubMed', 'EMBL', + 'EMBL-CDS', 'Ensembl', 'Ensembl_TRS', 'Ensembl_PRO', 'Additional PubMed'], + usecols=['UniProtKB-AC', 'UniProtKB-ID', 'GeneID (EntrezGene)', 'RefSeq', 'GI', 'PDB', 'GO', + 'NCBI-taxon', 'Ensembl', 'Ensembl_TRS', 'Ensembl_PRO'], + dtype='str') + + if blocksize: + if "idmapping_selected.parquet" in file_resources and \ + isinstance(file_resources["idmapping_selected.parquet"], str): + idmapping = dd.read_parquet(file_resources["idmapping_selected.parquet"]) + + elif "idmapping_selected.tab" in file_resources and \ + isinstance(file_resources["idmapping_selected.tab"], str): + idmapping = dd.read_table(file_resources["idmapping_selected.tab"], blocksize=blocksize, **args) + else: + idmapping = dd.read_table(file_resources["idmapping_selected.tab.gz"], compression="gzip", **args, ) + + idmapping: dd.DataFrame + else: + if "idmapping_selected.parquet" in file_resources and \ + isinstance(file_resources["idmapping_selected.parquet"], str): + idmapping = pd.read_parquet(file_resources["idmapping_selected.parquet"]) + else: + idmapping = pd.read_table(file_resources["idmapping_selected.tab"], index_col=self.index_col, **args) + + # Filter UniProt accession keys + if self.keys is not None and idmapping.index.name == self.index_col: + idmapping = idmapping.loc[idmapping.index.isin(self.keys)] + elif self.keys is not None and idmapping.index.name != self.index_col: + idmapping = idmapping.loc[idmapping[self.index_col].isin(self.keys)] + + if idmapping.index.name != self.index_col: + idmapping = idmapping.set_index(self.index_col, sorted=False) + if not idmapping.known_divisions: + idmapping.divisions = idmapping.compute_current_divisions() + + # Transform list columns + if isinstance(idmapping, dd.DataFrame): + idmapping = idmapping.assign(**self.assign_transforms(idmapping)) + else: + idmapping = idmapping.assign(**self.assign_transforms(idmapping)) + + # Join metadata from uniprot_sprot.parquet + if any(fn.startswith('uniprot') and fn.endswith('.parquet') for fn in file_resources): + uniprot_anns = self.load_uniprot_parquet(file_resources, blocksize=blocksize) + uniprot_anns = uniprot_anns[uniprot_anns.columns.difference(idmapping.columns)] + uniprot_anns = drop_duplicate_columns(uniprot_anns) + assert idmapping.index.name == uniprot_anns.index.name, f"{idmapping.index.name} != {uniprot_anns.index.name}" + idmapping = idmapping.join(uniprot_anns, on=idmapping.index.name, how='left') + + # Load proteome.tsv + if "proteomes.tsv" in file_resources: + proteomes = pd.read_table(file_resources["proteomes.tsv"], + usecols=['Organism Id', 'Proteome Id'], + dtype={'Organism Id': 'str', 'Proteome Id': 'str'}) \ + .rename(columns={'Organism Id': 'NCBI-taxon', 'Proteome Id': 'proteome_id'}) \ + .dropna().set_index('NCBI-taxon') + idmapping = idmapping.join(proteomes, on='NCBI-taxon') + + # Load species info from speclist.txt + if 'speclist.txt' in file_resources: + speclist = self.get_species_list(file_resources['speclist.txt']) + idmapping = idmapping.join(speclist, on='NCBI-taxon') + + return idmapping + + def assign_transforms(self, idmapping: pd.DataFrame) -> Dict[str, Union[dd.Series, pd.Series]]: + # Convert string of list elements to a np.array + list2array = lambda x: np.array(x) if isinstance(x, Iterable) else x + assign_fn = {} + for col in {'PDB', 'GI', 'GO', 'RefSeq'}.intersection(idmapping.columns): + try: + # Split string to list + assign_fn[col] = idmapping[col].str.split("; ").map(list2array) + except: + continue + + for col in {'Ensembl', 'Ensembl_TRS', 'Ensembl_PRO'}.intersection(idmapping.columns): + # Removing .# ENGS gene version number at the end + try: + if self.remove_version_num: + series = idmapping[col].str.replace("[.]\d*", "", regex=True) + else: + series = idmapping[col] + + assign_fn[col] = series.str.split("; ").map(list2array) + + if col == 'Ensembl_PRO': + # Prepend species_id to ensembl protein ids to match with STRING PPI + concat = dd.concat([idmapping["NCBI-taxon"], assign_fn[col]]) \ + if isinstance(idmapping, dd.DataFrame) else \ + pd.concat([idmapping["NCBI-taxon"], assign_fn[col]]) + + assign_fn['protein_external_id'] = concat.apply( + lambda row: np.char.add(row['NCBI-taxon'] + ".", row['Ensembl_PRO']) \ + if isinstance(row['Ensembl_PRO'], Iterable) else None, + axis=1) + except: + continue + + return assign_fn + + def load_uniprot_parquet(self, file_resources: Dict[str, str], blocksize=None) -> Union[dd.DataFrame, pd.DataFrame]: + dfs = [] + for filename, file_path in file_resources.items(): + if not ('uniprot' in filename and filename.endswith('.parquet')): continue + if blocksize: + df: dd.DataFrame = dd.read_parquet(file_path) \ + .rename(columns=UniProt.COLUMNS_RENAME_DICT) + if df.index.name in UniProt.COLUMNS_RENAME_DICT: + df.index = df.index.rename(UniProt.COLUMNS_RENAME_DICT[df.index.name]) + + if self.keys is not None: + if self.index_col in df.columns: + df = df.loc[df[self.index_col].isin(self.keys)] + elif df.index.name == self.index_col: + df = df.loc[df.index.isin(self.keys)] + + if df.index.size.compute() == 0: continue + + if df.index.name != self.index_col: + try: + df = df.set_index(self.index_col, sorted=True) + except Exception as e: + print(file_path, e) + df = df.set_index(self.index_col, sorted=False) + + if not df.known_divisions: + df.divisions = df.compute_current_divisions() + + else: + df = pd.read_parquet(file_path).rename(columns=UniProt.COLUMNS_RENAME_DICT).set_index(self.index_col) + + if self.keys is not None: + df_keys = df.index if df.index.name == self.index_col else df[self.index_col] + df = df.loc[df_keys.isin(self.keys)] + if df.index.size == 0: continue + + dfs.append(df) + + if dfs: + dfs = dd.concat(dfs, interleave_partitions=True) if blocksize else pd.concat(dfs) + + return dfs + else: + return None + + def load_uniprot_xml(self, file_path: str, keys=None, blocksize=None) -> pd.DataFrame: + records = [] + seqfeats = [] + if isinstance(keys, str): + index = keys + keys_set = self.data.index if keys == self.data.index.name else self.data[keys] + elif isinstance(keys, (dd.Index, dd.Series)): + index = keys.name + keys_set = keys.compute() + else: + index = keys_set = None + + for record in tqdm.tqdm(SeqIO.parse(file_path, "uniprot-xml"), desc=str(file_path)): + # Sequence features + annotations = defaultdict(None, record.annotations) + record_dict = { + 'protein_id': record.id, + "protein_name": record.name, + 'gene_name': annotations['gene_name_primary'], + 'description': record.description, + 'molecule_type': annotations['molecule_type'], + 'created': annotations['created'], + 'ec_id': annotations['type'], + 'subcellular_location': annotations['comment_subcellularlocation_location'], + 'taxonomy': annotations['taxonomy'], + 'keywords': annotations['keywords'], + 'sequence_mass': annotations['sequence_mass'], + SEQUENCE_COL: str(record.seq), + } + if index is not None: + if record_dict[keys] not in keys_set: continue + + records.append(record_dict) + + # Sequence interval features + _parse_interval = lambda sf: pd.Interval(left=sf.location.start, right=sf.location.end, ) + feature_type_intervals = defaultdict(lambda: []) + for sf in record.features: + if isinstance(sf.location.start, ExactPosition) and isinstance(sf.location.end, ExactPosition): + feature_type_intervals[sf.type].append(_parse_interval(sf)) + + features_dict = {type: pd.IntervalIndex(intervals, name=type) \ + for type, intervals in feature_type_intervals.items()} + seqfeats.append({"protein_id": record.id, **features_dict}) + + records_df = pd.DataFrame(records) if not blocksize else dd.from_pandas(records, chunksize=blocksize) + records_df = records_df.set_index(['protein_id']) + + seqfeats_df = pd.DataFrame(seqfeats) if not blocksize else dd.from_pandas(seqfeats, chunksize=blocksize) + seqfeats_df = seqfeats_df.set_index(['protein_id']) + seqfeats_df.columns = [f"seq/{col}" for col in seqfeats_df.columns] + + # Join new metadata to self.data + if SEQUENCE_COL not in self.data.columns: + exclude_cols = records_df.columns.intersection(self.data.columns) + self.data = self.data.join(records_df.drop(columns=exclude_cols, errors="ignore"), + on='protein_id', how="left") + else: + self.data.update(records_df, overwrite=False) + + # Add new seq features + if len(seqfeats_df.columns.difference(self.data.columns)): + self.data = self.data.join(seqfeats_df.drop(columns=seqfeats_df.columns.intersection(self.data.columns)), + on='protein_id', how="left") + # Fillna seq features + if len(seqfeats_df.columns.intersection(self.data.columns)): + self.data.update(seqfeats_df.filter(seqfeats_df.columns.intersection(self.data.columns), axis='columns'), + overwrite=False) + + return records_df + + @classmethod + def get_species_list(cls, file_path): + speclist = pd.read_fwf(file_path, + names=['species_code', 'Taxon', 'species_id', 'attr'], + comment="==", skipinitialspace=True, skiprows=59, skipfooter=4) + speclist = speclist.drop(index=speclist.index[~speclist['attr'].str.contains("=")]) + speclist['species_id'] = speclist['species_id'].str.rstrip(":") + speclist = speclist.fillna(method='ffill') + speclist = speclist.groupby(speclist.columns[:3].tolist())['attr'] \ + .apply('|'.join) \ + .apply(lambda s: dict(map(str.strip, sub.split('=', 1)) for sub in s.split("|") if '=' in sub)) \ + .apply(pd.Series) + speclist = speclist.rename(columns={'N': 'Official (scientific) name', 'C': 'Common name', 'S': 'Synonym'}) \ + .reset_index() \ + .set_index('species_id') + speclist['Taxon'] = speclist['Taxon'].replace( + {'A': 'archaea', 'B': 'bacteria', 'E': 'eukaryota', 'V': 'viruses', 'O': 'others'}) + speclist.index.name = 'NCBI-taxon' + return speclist + + def load_sequences(self, fasta_file: str, index=None, keys: Union[pd.Index, List[str]] = None, blocksize=None) \ + -> OrderedDict: + def get_id(s: str): + if index == 'protein_id': + return s.split('|')[1] + elif index == 'protein_name': + return s.split('|')[2] + else: + return s.split('|')[1] + + fa = Fasta(fasta_file, key_function=get_id, as_raw=True, ) + + return fa.records + + def get_sequences(self, index: str, omic: str = None, agg: str = "first", **kwargs): + assert index, '`index` must be either "protein_id" or "protein_name"' + + # Parse lncRNA & mRNA fasta + seq_df = self.load_sequences(self.file_resources["uniprot_sprot.fasta"], index=index, blocksize=self.blocksize) + if "uniprot_trembl.fasta" in self.file_resources: + trembl_seq_df = self.load_sequences(self.file_resources["uniprot_trembl.fasta"], index=index, + blocksize=self.blocksize) + seq_df.update(trembl_seq_df) + + return seq_df + +class MirBase(SequenceDatabase): + """Loads the MirBase database from https://mirbase.org . + + Default path: "ftp://mirbase.org/pub/mirbase/CURRENT/" . + Default file_resources: { + "aliases.txt": "aliases.txt.gz", + "mature.fa": "mature.fa.gz", + "hairpin.fa": "hairpin.fa.gz", + "rnacentral.mirbase.tsv": "ftp://ftp.ebi.ac.uk/pub/databases/RNAcentral/current_release/id_mapping/database_mappings/mirbase.tsv", + } + """ + + def __init__( + self, + path="http://mirbase.org/ftp/CURRENT/", + file_resources=None, + species_id: Optional[str] = '9606', + index_col: str = "mirbase_id", + col_rename=None, + **kwargs, + ): + """ + Args: + path: + file_resources: + sequence (str): + species_id (str): Species code, e.g., 9606 for human + col_rename: + blocksize: + """ + if file_resources is None: + file_resources = {} + file_resources["aliases.txt.gz"] = "aliases.txt.gz" + file_resources["mature.fa.gz"] = "mature.fa.gz" + file_resources["hairpin.fa.gz"] = "hairpin.fa.gz" + + if 'rnacentral.mirbase.tsv' not in file_resources: + file_resources["rnacentral.mirbase.tsv"] = "ftp://ftp.ebi.ac.uk/pub/databases/RNAcentral/current_release/" \ + "id_mapping/database_mappings/mirbase.tsv" + + self.species_id = species_id + super().__init__(path=path, file_resources=file_resources, index_col=index_col, col_rename=col_rename, **kwargs) + + def load_dataframe(self, file_resources, blocksize=None): + """ + Args: + file_resources: dict of file name and path + blocksize: + """ + rnacentral_mirbase = pd.read_table( + file_resources["rnacentral.mirbase.tsv"], low_memory=True, header=None, + names=["RNAcentral id", "database", "mirbase_id", "species_id", "RNA type", "NA"], + usecols=["RNAcentral id", "database", "mirbase_id", "species_id", "RNA type"], + index_col="RNAcentral id", + dtype={'mirbase_id': 'str', "species_id": "category", 'database': 'category', 'RNA type': 'category'}) + + if isinstance(self.species_id, str): + rnacentral_mirbase = rnacentral_mirbase[rnacentral_mirbase["species_id"] == self.species_id] + elif isinstance(self.species_id, Iterable): + rnacentral_mirbase = rnacentral_mirbase[rnacentral_mirbase["species_id"].isin(set(self.species_id))] + + mirbase_df = pd.read_table(file_resources["aliases.txt"], low_memory=True, header=None, + names=["mirbase_id", "mirbase_name"], index_col=self.index_col, + dtype='str', ) + if mirbase_df.index.name == 'mirbase id': + mirbase_df = mirbase_df.join(rnacentral_mirbase, on=self.index_col, how="left", rsuffix='_rnacentral') + else: + mirbase_df = mirbase_df.merge(rnacentral_mirbase, on=self.index_col, how="left") + + # Expanding miRNA names in each MirBase Ascension ID + mirbase_df['mirbase_name'] = mirbase_df['mirbase_name'].str.rstrip(";").str.split(";") + + seq_dfs = [] + for filename in file_resources: + if filename.endswith('.fa') or filename.endswith('.fasta'): + assert isinstance(file_resources[filename], str), f"Must provide a path to an uncompressed .fa file. " \ + f"Given {file_resources[filename]}" + df = self.load_sequences(file_resources[filename], index=self.index_col, keys=self.keys) + seq_dfs.append(df) + + if len(seq_dfs): + seq_dfs = pd.concat(seq_dfs, axis=0) + mirbase_df = mirbase_df.join(seq_dfs, how='left', on=self.index_col) + else: + logger.info('Missing sequence data because no "hairpin.fa" or "mature.fa" file were given.') + + # mirbase_df = mirbase_df.explode(column='gene_name') + # mirbase_name["miRNA name"] = mirbase_name["miRNA name"].str.lower() + # mirbase_name["miRNA name"] = mirbase_name["miRNA name"].str.replace("-3p.*|-5p.*", "") + + return mirbase_df + + def load_sequences(self, fasta_file, index=None, keys=None, blocksize=None): + """ + Args: + fasta_file: + index (): + keys (): + blocksize: + """ + if hasattr(self, '_seq_df_dict') and fasta_file in self._seq_df_dict: + return self._seq_df_dict[fasta_file] + + fa = Fasta(fasta_file, read_long_names=True, as_raw=True) + mirna_types = {'stem-loop', 'stem', 'type', 'loop'} + + entries = [] + for key, record in tqdm.tqdm(fa.items(), desc=str(fasta_file)): + attrs: List[str] = record.long_name.split(" ") + + if attrs[-1] in mirna_types: + if attrs[-2] in mirna_types: + mirna_type = intern(' '.join(attrs[-2:])) + gene_name_idx = -3 + else: + mirna_type = intern(attrs[-1]) + gene_name_idx = -2 + else: + mirna_type = None + gene_name_idx = -1 + + record_dict = { + "gene_id": attrs[0], + "mirbase_id": attrs[1], + "species": intern(" ".join(attrs[2:gene_name_idx])), + "gene_name": attrs[gene_name_idx], + "mirna_type": mirna_type, + SEQUENCE_COL: str(record), + } + if keys is not None and index: + if record_dict[index] not in keys: + del record_dict + continue + + entries.append(record_dict) + + df = pd.DataFrame(entries) + if index: + df = df.set_index(index) + # if blocksize: + # df = dd.from_pandas(df, chunksize=blocksize) + + if not hasattr(self, '_seq_df_dict'): + self._seq_df_dict = {} + self._seq_df_dict[fasta_file] = df + + return df + + def get_sequences(self, + index="gene_name", + omic=None, + agg="all", + **kwargs): + """ + Args: + index: + omic: + agg: + **kwargs: + """ + dfs = [] + for filename in self.file_resources: + if filename.endswith('.fa'): + seq_df = self.load_sequences(self.file_resources[filename]) + dfs.append(seq_df) + seq_df = pd.concat(dfs, axis=0) + + seq_df = seq_df.groupby(index)[SEQUENCE_COL].agg(self.aggregator_fn(agg)) + + return seq_df + + +class RNAcentral(SequenceDatabase): + """ + Loads the RNAcentral database from https://rnacentral.org/ and provides a series of methods to extract sequence data from it. + + Default path: https://ftp.ebi.ac.uk/pub/databases/RNAcentral/current_release/ . + Default file_resources: { + "rnacentral_rfam_annotations.tsv": "go_annotations/rnacentral_rfam_annotations.tsv.gz", + "database_mappings/gencode.tsv": "id_mapping/database_mappings/gencode.tsv", + "gencode.fasta": "sequences/by-database/gencode.fasta", + ... + } + """ + COLUMNS_RENAME_DICT = { + 'ensembl_gene_id': 'gene_id', + 'external id': 'transcript_id', + 'GO terms': 'go_id', + 'gene symbol': 'gene_id', + } + + def __init__(self, path="https://ftp.ebi.ac.uk/pub/databases/RNAcentral/current_release/", file_resources=None, + col_rename=COLUMNS_RENAME_DICT, species_id: Union[List[str], str, None] = None, + index_col="RNAcentral id", keys=None, + remove_version_num=True, remove_species_suffix=True, **kwargs): + """ + Provide + + Args: + path (): + file_resources (): + col_rename (): + species_id (): + index_col (): + keys (): + remove_version_num (): + remove_species_suffix (): + **kwargs (): + """ + self.species_id = species_id + self.remove_version_num = remove_version_num + self.remove_species_suffix = remove_species_suffix + + if file_resources is None: + file_resources = {} + file_resources["rnacentral_rfam_annotations.tsv.gz"] = "go_annotations/rnacentral_rfam_annotations.tsv.gz" + file_resources["database_mappings/ensembl_gencode.tsv"] = "id_mapping/database_mappings/ensembl_gencode.tsv" + file_resources["database_mappings/mirbase.tsv"] = "id_mapping/database_mappings/mirbase.tsv" + + super().__init__(path=path, file_resources=file_resources, col_rename=col_rename, index_col=index_col, + keys=keys, **kwargs) + + def load_dataframe(self, file_resources, blocksize=None): + """ + Args: + file_resources: + blocksize: + """ + transcripts_df = [] + # Concatenate transcripts ids by combining `database_mappings/` files from multiple RNAcentral databases + for filename in (fname for fname in file_resources if "database_mappings" in fname): + args = dict(low_memory=True, header=None, + names=["RNAcentral id", "database", "external id", "species_id", "RNA type", "gene symbol"], + dtype={'gene symbol': 'str', + 'database': 'category', 'species_id': 'category', 'RNA type': 'category', }) + + if blocksize: + if filename.endswith('.tsv'): + id_mapping: dd.DataFrame = dd.read_table( + file_resources[filename], blocksize=None if isinstance(blocksize, bool) else blocksize, **args) + elif filename.endswith('.parquet'): + id_mapping: dd.DataFrame = dd.read_parquet( + file_resources[filename], blocksize=None if isinstance(blocksize, bool) else blocksize, ) + else: + id_mapping = None + else: + if filename.endswith('.tsv'): + id_mapping = pd.read_table(file_resources[filename], **args) + elif filename.endswith('.parquet'): + id_mapping = pd.read_parquet(file_resources[filename]) + else: + id_mapping = None + + if id_mapping is None: + raise Exception("Must provide a file with 'database_mappings/(*).tsv' in file_resources") + + # Filter by species + if isinstance(self.species_id, str): + id_mapping = id_mapping.where(id_mapping["species_id"] == self.species_id) + elif isinstance(self.species_id, Iterable): + id_mapping = id_mapping.where(id_mapping["species_id"].isin(self.species_id)) + + # Filter by index + if self.keys and id_mapping.index.name == self.index_col: + id_mapping = id_mapping.loc[id_mapping.index.isin(self.keys)] + elif self.keys and id_mapping.index.name != self.index_col: + id_mapping = id_mapping.loc[id_mapping[self.index_col].isin(self.keys)] + + # Add species_id prefix to index values to match the sequence ids + if "RNAcentral id" in id_mapping.columns: + id_mapping["RNAcentral id"] = id_mapping["RNAcentral id"] + "_" + id_mapping["species_id"].astype(str) + elif "RNAcentral id" == id_mapping.index.name: + id_mapping.index = id_mapping.index + "_" + id_mapping["species_id"].astype(str) + + # Add sequence column if a FASTA file provided for the database + fasta_filename = f"{filename.split('/')[-1].split('.')[0]}.fasta" + if fasta_filename in file_resources: + seq_df = self.load_sequences(file_resources[fasta_filename]) + id_mapping = id_mapping.merge(seq_df, how='left', + left_on="RNAcentral id", + left_index=True if id_mapping.index.name == "RNAcentral id" else False, + right_index=True) + else: + logger.info(f"{fasta_filename} not provided for `{filename}` so missing sequencing data") + + if self.remove_version_num and 'gene symbol' in id_mapping.columns: + id_mapping["gene symbol"] = id_mapping["gene symbol"].str.replace("[.].\d*", "", regex=True) + if self.remove_species_suffix: + id_mapping["RNAcentral id"] = id_mapping["RNAcentral id"].str.replace("_(\d*)", '', regex=True) + + # Set index + args = dict(sorted=True) if blocksize else {} + id_mapping = id_mapping.set_index(self.index_col, **args) + if isinstance(id_mapping, dd.DataFrame) and not id_mapping.known_divisions: + id_mapping.divisions = id_mapping.compute_current_divisions() + + transcripts_df.append(id_mapping) + + # Concatenate multiple `database_mappings` files from different databases + if blocksize: + transcripts_df = dd.concat(transcripts_df, axis=0, interleave_partitions=True, join='outer') + else: + transcripts_df = pd.concat(transcripts_df, axis=0, join='outer') + + # Join go_id and Rfams annotations to each "RNAcentral id" from 'rnacentral_rfam_annotations.tsv' + try: + transcripts_df = self.add_rfam_annotation(transcripts_df, file_resources, blocksize) + except Exception as e: + logger.warning(f"Failed to add Rfam annotations to transcripts_df:") + traceback.print_exc() + + return transcripts_df + + def add_rfam_annotation(self, transcripts_df: Union[pd.DataFrame, dd.DataFrame], + file_resources, blocksize=None) -> Union[pd.DataFrame, dd.DataFrame]: + args = dict(low_memory=True, names=["RNAcentral id", "GO terms", "Rfams"]) + + if blocksize: + if 'rnacentral_rfam_annotations.tsv' in file_resources and isinstance( + file_resources['rnacentral_rfam_annotations.tsv'], str): + anns = dd.read_table(file_resources["rnacentral_rfam_annotations.tsv"], **args) + else: + anns = dd.read_table(file_resources["rnacentral_rfam_annotations.tsv.gz"], compression="gzip", **args) + anns = anns.set_index("RNAcentral id", sorted=True) + + # Filter annotations by "RNAcentral id" in `transcripts_df` + anns = anns.loc[anns.index.isin(transcripts_df.index.compute())] + + if not anns.known_divisions: + anns.divisions = anns.compute_current_divisions() + + # Groupby on index + anns_groupby: dd.DataFrame = anns \ + .groupby(by=lambda idx: idx) \ + .agg({col: get_agg_func('unique', use_dask=True) for col in ["GO terms", 'Rfams']}) + + else: + anns = pd.read_table(file_resources["rnacentral_rfam_annotations.tsv"], index_col='RNAcentral id', **args) + idx = transcripts_df.index.compute() if isinstance(transcripts_df, dd.DataFrame) else transcripts_df.index + anns = anns.loc[anns.index.isin(set(idx))] + anns_groupby = anns.groupby("RNAcentral id").agg({col: 'unique' for col in ["GO terms", 'Rfams']}) + + transcripts_df = transcripts_df.merge(anns_groupby, how='left', left_index=True, right_index=True) + return transcripts_df + + def load_sequences(self, fasta_file: str, index=None, keys=None, blocksize=None): + """ + Args: + index (): + fasta_file: + keys (): + blocksize: + """ + fa = Fasta(fasta_file, as_raw=True) + + entries = [] + for key, record in tqdm.tqdm(fa.items(), desc=str(fasta_file)): + id = re.sub("_(\d*)", '', key) if self.remove_species_suffix else key + if keys is not None and self.index_col == 'RNAcentral id' and id not in keys: + continue + desc = record.long_name.split(" ", maxsplit=1)[-1] + + record_dict = { + 'RNAcentral id': key, + 'description': desc, + SEQUENCE_COL: str(record), + } + + entries.append(record_dict) + + df = pd.DataFrame(entries).set_index("RNAcentral id") + + return df + + def get_sequences(self, + index="RNAcentral id", + omic=None, + agg="all", + **kwargs): + """ + Args: + index: + omic: + agg: + **kwargs: + """ + dfs = [] + for filename in self.file_resources: + if filename.endswith('.fa') or filename.endswith('.fasta'): + seq_df = self.load_sequences(self.file_resources[filename]) + dfs.append(seq_df) + seq_df = pd.concat(dfs, axis=0) + + seq_df = seq_df.groupby(index)[SEQUENCE_COL].agg(self.aggregator_fn(agg)) + + return seq_df