diff --git a/src/ann_solo/ann_solo.py b/src/ann_solo/ann_solo.py index c2f0de5..48cfec4 100644 --- a/src/ann_solo/ann_solo.py +++ b/src/ann_solo/ann_solo.py @@ -54,7 +54,7 @@ def main(args: Union[str, List[str]] = None) -> int: # Initialize logging. logging.basicConfig(format='{asctime} [{levelname}/{processName}] ' '{module}.{funcName} : {message}', - style='{', level=logging.DEBUG) + style='{', level=logging.INFO) # Load the configuration. config.parse(args) diff --git a/src/ann_solo/reader.py b/src/ann_solo/reader.py index d2d82f3..8ccf2ce 100644 --- a/src/ann_solo/reader.py +++ b/src/ann_solo/reader.py @@ -15,6 +15,7 @@ from spectrum_utils.spectrum import MsmsSpectrum from ann_solo.parsers import SplibParser +from ann_solo.sqlite_parsers import ElibParser from ann_solo.spectrum import process_spectrum @@ -23,7 +24,7 @@ class SpectralLibraryReader: Read spectra from a SpectraST spectral library .splib file. """ - _supported_extensions = ['.splib'] + _supported_extensions = ['.splib', '.elib', '.dlib'] is_recreated = False @@ -156,7 +157,11 @@ def _create_config(self) -> None: config_filename, compress=9, protocol=pickle.DEFAULT_PROTOCOL) def open(self) -> None: - self._parser = SplibParser(self._filename.encode()) + _, ext = os.path.splitext(self._filename) + if ext.lower() == ".splib": + self._parser = SplibParser(self._filename.encode()) + else: + self._parser = ElibParser(self._filename.encode()) def close(self) -> None: if self._parser is not None: @@ -213,13 +218,17 @@ def get_all_spectra(self) -> Iterator[Tuple[MsmsSpectrum, int]]: library file. """ self._parser.seek_first_spectrum() + try: - while True: - spectrum, offset = self._parser.read_spectrum() - spectrum.is_processed = False - yield spectrum, offset - except StopIteration: - return + yield from self._parser.get_all_spectra() + except AttributeError: + try: + while True: + spectrum, offset = self._parser.read_spectrum() + spectrum.is_processed = False + yield spectrum, offset + except StopIteration: + return def get_version(self) -> str: """ diff --git a/src/ann_solo/sqlite_parsers.py b/src/ann_solo/sqlite_parsers.py new file mode 100644 index 0000000..4261d92 --- /dev/null +++ b/src/ann_solo/sqlite_parsers.py @@ -0,0 +1,190 @@ +""" +This module provides support for common spectral library formats that +rely on SQLite3 databases. + +Currently these include: + - ELIB + - DLIB + +BLIB should be easy to add. +""" +import re +import sqlite3 +import zlib +from typing import Tuple, Dict, Iterator, Union, List + +import numpy as np +from spectrum_utils.spectrum import MsmsSpectrum + + +class ElibParser: + """Parse an ELIB or DLIB spectral library""" + def __init__(self, filename: str) -> None: + """ + Initialize an ELIB/DLIB spectral library parser. + + The ELIB and DLIB formats are described here: + https://bitbucket.org/searleb/encyclopedia/wiki/EncyclopeDIA%20File%20Formats + + These files should have either '.blib' or '.dlib' extensions. + + Parameters + ---------- + filename : str + The file name of the DLIB or BLIB spectral library. + """ + self._conn = sqlite3.connect(filename) + self._cursor = self._conn.cursor() + + # Get row IDs for fast look-up: + rowids = self._cursor.execute('SELECT rowid FROM entries') + self._rowids = {k: v[0] for k, v in enumerate(rowids)} + self._size = len(self._rowids) + self._pos = 0 + + # Build an decoy dict for fast look-up: + decoy_map = self._cursor.execute( + 'SELECT PeptideSeq, isDecoy FROM peptidetoprotein' + ) + self._is_decoy = dict(decoy_map) + + def _parse_spectrum(self, row: Tuple, identifier: int) -> MsmsSpectrum: + """ + Parse a single spectrum given one row of the table + + Parameters + ---------- + row : Tuple + One row of the 'entries' table. + identifier : int + The identifier for a spectrum. + + Returns + ------- + MsmsSpectrum object + """ + precursor_mz = row[0] + precursor_charge = row[1] + mods = _parse_mods(row[2]) + seq = row[3] + mz_array = _decode(row[8], dtype='>d') + int_array = _decode(row[10], dtype='>f') + + spectrum = MsmsSpectrum(identifier, precursor_mz, precursor_charge, + mz_array, int_array, + is_decoy=self._is_decoy[seq], + peptide=seq, modifications=mods) + + spectrum.annotate_peptide_fragments(10, "ppm") + return spectrum + + def _get_row(self, offset: int) -> MsmsSpectrum: + """ + Read a row at the offset + Parameters + ---------- + offset : int + The index of the row to read. + Returns + ------- + A parsed Spectrum. + """ + rowid = self._rowids[offset] + row = self._cursor.execute( + 'SELECT * FROM entries WHERE rowid=?', (str(rowid),) + ) + return self._parse_spectrum(row.fetchone(), offset) + + def seek_first_spectrum(self) -> None: + """Needed for for compatibility""" + self._pos = 0 + + def read_spectrum(self, offset: int = None) -> MsmsSpectrum: + """ + Read a spectrum from the library. + + This function is mainly to maintain a consistent API with the + splib parser. However, use this function sparingly - it is + exceedingly slow due to having lookup each spectrum from the + SQLite table individually. + + Parameters + ---------- + offset : int + The row to start reading from. + + Returns + ------- + An Spectrum object + """ + if offset is not None and offset >= 0: + self._pos = offset + + spectrum_offset = self._pos + if self._pos >= self._size: + print(self._pos, self._size) + raise StopIteration + + spectrum = self._get_row(spectrum_offset) + self._pos += 1 + return spectrum, spectrum_offset + + def get_all_spectra(self) -> Iterator[Tuple[MsmsSpectrum, int]]: + """ + Generates all spectra from the spectral library file. + For each individual spectrum a tuple consisting of the spectrum and + some additional information as a nested tuple (containing on the type + of spectral library file) are returned. + + Returns + ------- + Iterator[Tuple[Spectrum, int]] + An iterator of all spectra along with their offset in the spectral + library file. + """ + rows = self._cursor.execute('SELECT * FROM entries') + for spectrum_offset, row in enumerate(rows): + spectrum = self._parse_spectrum(row, spectrum_offset) + spectrum.is_processed = False + yield spectrum, spectrum_offset + + +def _decode(array: bytes, dtype: str) -> np.ndarray: + """ + Decode a zlib-compressed array + + Parameters + ---------- + array : bytestring + The zlib-compressed array. + dtype : str + The data type to be read by numpy. '>f' is for Big Endian + floats, '>d' is for Big Endian doubles. + """ + return np.frombuffer(zlib.decompress(array), dtype=dtype) + + +def _parse_mods(peptide: str) -> Dict[int, float]: + """ + Parse a modified peptide string. + + Parameters + ---------- + peptide : str + The peptide string with modification indicated in square + brackets. + + Returns + ------- + Dict[int, float] + The modifications for spectrum_utils. + """ + mods = re.finditer(r"\[(.+?)\]", peptide) + mod_dict = {} + offset = 0 + for mod in mods: + position = mod.start() - offset + mod_dict[position] = float(mod.groups(1)[0]) + offset += mod.end() - mod.start() + + return mod_dict