diff --git a/alphabase/spectral_library/translate.py b/alphabase/spectral_library/translate.py index 948c0944..57d693a0 100644 --- a/alphabase/spectral_library/translate.py +++ b/alphabase/spectral_library/translate.py @@ -1,235 +1,136 @@ +"""Translate AlphaBase spectral libraries to a SWATH/Spectronaut transition list. + +The shared export machinery lives in :mod:`alphabase.spectral_library.translate_core`; +this module holds the SWATH column names, the precursor mapping and the tsv writer. The +names it used to define are re-exported below, so importing them from here keeps +working. +""" + import multiprocessing as mp -import numpy as np import pandas as pd import tqdm -from alphabase.constants.modification import MOD_DF, ModificationKeys -from alphabase.numba_wrapper import numba_njit from alphabase.spectral_library.base import SpecLibBase -from alphabase.utils import explode_multiple_columns - - -# @numba.njit #(cannot use numba for pd.Series) -def create_modified_sequence( - seq_mods_sites: tuple, # must be ('sequence','mods','mod_sites') - translate_mod_dict: dict = None, - mod_sep="[]", - nterm="_", - cterm="_", -): - """ - Translate `(sequence, mods, mod_sites)` into a modified sequence. Used by `df.apply()`. - For example, `('ABCDEFG','Mod1@A;Mod2@E','1;5')`->`_A[Mod1@A]BCDE[Mod2@E]FG_`. - - Parameters - ---------- - seq_mods_sites : List - must be `(sequence, mods, mod_sites)` - - translate_mod_dict : dict - A dict to map AlphaX modification names to other software, - use unimod name if None. - Defaults to None. - - mod_sep : str - '[]' or '()', default '[]' - - """ - mod_seq, mods, mod_sites = seq_mods_sites - if mods: - mods = mods.split(ModificationKeys.SEPARATOR) - mod_sites = [int(i) for i in mod_sites.split(ModificationKeys.SEPARATOR)] - rev_order = np.argsort(mod_sites)[::-1] - mod_sites = [mod_sites[rev_order[i]] for i in range(len(mod_sites))] - mods = [mods[rev_order[i]] for i in range(len(mods))] - if translate_mod_dict is None: - mods = [mod[: mod.find(ModificationKeys.SITE_SEPARATOR)] for mod in mods] - else: - mods = [translate_mod_dict[mod] for mod in mods] - for _site, mod in zip(mod_sites, mods): - if _site > 0: - mod_seq = ( - mod_seq[:_site] + mod_sep[0] + mod + mod_sep[1] + mod_seq[_site:] - ) - elif _site == -1: - cterm += mod_sep[0] + mod + mod_sep[1] - elif _site == 0: - nterm += mod_sep[0] + mod + mod_sep[1] - else: - mod_seq = ( - mod_seq[:_site] + mod_sep[0] + mod + mod_sep[1] + mod_seq[_site:] - ) - return nterm + mod_seq + cterm - - -@numba_njit -def _get_frag_info_from_column_name(column: str): - """ - Only used when converting alphabase libraries into other libraries - """ - idx = column.rfind("_") - frag_type = column[:idx] - charge = column[idx + 2 :] - if len(frag_type) == 1: - loss_type = "noloss" - else: - idx = frag_type.find("_") - loss_type = frag_type[idx + 1 :] - frag_type = frag_type[0] - return frag_type, loss_type, charge - - -def _get_frag_num(columns, rows, frag_len): - frag_nums = [] - for r, c in zip(rows, columns): - if is_nterm_frag(c): - frag_nums.append(r + 1) - else: - frag_nums.append(frag_len - r) - return frag_nums - - -def merge_precursor_fragment_df( +from alphabase.spectral_library.translate_core import ( + CCS_COLUMNS, + MOBILITY_COLUMNS, + RT_COLUMNS, + FragmentTableCols, + create_modified_sequence, + first_present_column, + get_fragment_table, + get_precursor_mz, + is_nterm_frag, + join_fragments, + mod_to_unimod_dict, +) + +# the SWATH names for the canonical fragment columns, in output order +SWATH_FRAGMENT_COLUMNS = { + FragmentTableCols.FRAG_TYPE: "FragmentType", + FragmentTableCols.MZ: "FragmentMz", + FragmentTableCols.INTENSITY: "RelativeIntensity", + FragmentTableCols.CHARGE: "FragmentCharge", + FragmentTableCols.SERIES_NUMBER: "FragmentNumber", + FragmentTableCols.LOSS_TYPE: "FragmentLossType", +} + +__all__ = [ + # re-exported from translate_core, where these now live + "create_modified_sequence", + "is_nterm_frag", + "mod_to_unimod_dict", + "get_precursor_mz", + # this module's own + "WritingProcess", + "speclib_to_single_df", + "speclib_to_swath_df", + "translate_to_tsv", +] + +# pandas renamed `to_csv`'s newline argument in 1.5; alphabase does not pin a minimum +_CSV_NEWLINE = ( + {"lineterminator": "\n"} + if tuple(int(part) for part in pd.__version__.split(".")[:2]) >= (1, 5) + else {"line_terminator": "\n"} +) + + +def _precursors_to_swath_df( # noqa: PLR0913 precursor_df: pd.DataFrame, fragment_mz_df: pd.DataFrame, - fragment_inten_df: pd.DataFrame, - top_n_inten: int, - frag_type_head: str = "FragmentType", - frag_mass_head: str = "FragmentMz", - frag_inten_head: str = "RelativeIntensity", - frag_charge_head: str = "FragmentCharge", - frag_series_head: str = "FragmentNumber", - frag_loss_head: str = "FragmentLossType", + fragment_intensity_df: pd.DataFrame, + *, + translate_mod_dict: dict = None, + keep_k_highest_fragments: int = 12, + min_frag_mz=200, + max_frag_mz=2000, + min_frag_intensity=0.01, + min_frag_nAA=0, + modloss: str = "H3PO4", verbose=True, -): - """ - Convert alphabase library into a single dataframe. - This method is not important, as it will be only - used by DiaNN, or spectronaut, or others - """ - df = precursor_df.copy() - frag_columns = fragment_mz_df.columns.values.astype("U") - frag_type_list = [] - frag_loss_list = [] - frag_charge_list = [] - frag_mass_list = [] - frag_inten_list = [] - frag_num_list = [] - iters = enumerate(df[["frag_start_idx", "frag_stop_idx"]].values) - if verbose: - iters = tqdm.tqdm(iters) - for _i, (start, end) in iters: - intens = fragment_inten_df.iloc[start:end, :].to_numpy( - copy=True - ) # is loc[start:end-1,:] faster? - max_inten = np.amax(intens) - if max_inten > 0: - intens /= max_inten - masses = fragment_mz_df.iloc[start:end, :].values - sorted_idx = np.argsort(intens.reshape(-1))[-top_n_inten:][::-1] - idx_in_df = np.unravel_index(sorted_idx, masses.shape) - - frag_len = end - start - rows = np.arange(frag_len, dtype=np.int32)[idx_in_df[0]] - columns = frag_columns[idx_in_df[1]] - - frag_types, loss_types, charges = zip( - *[_get_frag_info_from_column_name(_) for _ in columns] - ) +) -> pd.DataFrame: + """Convert precursor and fragment frames to a SWATH transition list. - frag_nums = _get_frag_num(columns, rows, frag_len) - - frag_type_list.append(frag_types) - frag_loss_list.append(loss_types) - frag_charge_list.append(charges) - frag_mass_list.append(masses[idx_in_df]) - frag_inten_list.append(intens[idx_in_df]) - frag_num_list.append(frag_nums) - - df[frag_type_head] = frag_type_list - df[frag_mass_head] = frag_mass_list - df[frag_inten_head] = frag_inten_list - df[frag_charge_head] = frag_charge_list - df[frag_series_head] = frag_num_list - df[frag_loss_head] = frag_loss_list - - return explode_multiple_columns( - df, - [ - frag_type_head, - frag_mass_head, - frag_inten_head, - frag_charge_head, - frag_series_head, - frag_loss_head, - ], + The dataframe-in form of :func:`speclib_to_single_df`, so that one batch of + precursors can be converted without standing up a `SpecLibBase` around it. + """ + df = pd.DataFrame() + df["ModifiedPeptide"] = precursor_df[["sequence", "mods", "mod_sites"]].apply( + create_modified_sequence, + axis=1, + translate_mod_dict=translate_mod_dict, + mod_sep="[]", ) - # try: - # return df.explode([ - # frag_type_head, - # frag_mass_head, - # frag_inten_head, - # frag_charge_head, - # frag_loss_head, - # frag_num_head - # ]) - # except ValueError: - # # df.explode does not allow mulitple columns before pandas version 1.x.x. - # df = df.explode(frag_type_head) + df["PrecursorCharge"] = precursor_df["charge"] - # df[frag_mass_head] = _flatten(frag_mass_list) - # df[frag_inten_head] = _flatten(frag_inten_list) - # df[frag_charge_head] = _flatten(frag_charge_list) - # df[frag_loss_head] = _flatten(frag_loss_list) - # df[frag_num_head] = _flatten(frag_num_list) - # return df + rt = first_present_column(precursor_df, RT_COLUMNS) + if rt is None: + raise ValueError("precursor_df must contain the RT columns") + df["RT"] = rt + # ion mobility and CCS are optional: the column is omitted, not defaulted + mobility = first_present_column(precursor_df, MOBILITY_COLUMNS) + if mobility is not None: + df["IonMobility"] = mobility -mod_to_unimod_dict = {} -for mod_name, unimod_id in MOD_DF[["mod_name", "unimod_id"]].values: - if unimod_id == -1 or unimod_id == "-1": - continue - mod_to_unimod_dict[mod_name] = f"UniMod:{unimod_id}" + ccs = first_present_column(precursor_df, CCS_COLUMNS) + if ccs is not None: + df["CCS"] = ccs + df["StrippedPeptide"] = precursor_df["sequence"] -def is_nterm_frag(frag_type: str): - return frag_type[0] in "abc" + df["PrecursorMz"] = get_precursor_mz(precursor_df) + # this format prefers uniprot_ids; the DIA-NN one splits them over two columns + proteins = first_present_column(precursor_df, ["uniprot_ids", "proteins"]) + if proteins is not None: + df["ProteinID"] = proteins -def mask_fragment_intensity_by_mz_( - fragment_mz_df: pd.DataFrame, - fragment_intensity_df: pd.DataFrame, - min_frag_mz, - max_frag_mz, -): - fragment_intensity_df.mask( - (fragment_mz_df > max_frag_mz) | (fragment_mz_df < min_frag_mz), 0, inplace=True - ) + if "genes" in precursor_df.columns: + df["Genes"] = precursor_df["genes"] + if "decoy" in precursor_df.columns: + df["Decoy"] = precursor_df["decoy"] -def mask_fragment_intensity_by_frag_nAA( - fragment_intensity_df: pd.DataFrame, precursor_df: pd.DataFrame, max_mask_frag_nAA -): - if max_mask_frag_nAA <= 0: - return - b_mask = np.zeros(len(fragment_intensity_df), dtype=np.bool_) - y_mask = b_mask.copy() - for i_frag in range(max_mask_frag_nAA): - b_mask[precursor_df.frag_start_idx.values + i_frag] = True - y_mask[precursor_df.frag_stop_idx.values - i_frag - 1] = True - - masks = np.zeros( - (len(fragment_intensity_df), len(fragment_intensity_df.columns)), dtype=np.bool_ + fragments = get_fragment_table( + precursor_df["frag_start_idx"].to_numpy(), + precursor_df["frag_stop_idx"].to_numpy(), + fragment_mz_df, + fragment_intensity_df, + keep_k_highest=keep_k_highest_fragments, + min_frag_mz=min_frag_mz, + max_frag_mz=max_frag_mz, + min_frag_nAA=min_frag_nAA, + verbose=verbose, ) - for i, col in enumerate(fragment_intensity_df.columns.values): - if is_nterm_frag(col): - masks[:, i] = b_mask - else: - masks[:, i] = y_mask + df = join_fragments(df, fragments, SWATH_FRAGMENT_COLUMNS) + df = df[df["RelativeIntensity"] > min_frag_intensity] + df.loc[df["FragmentLossType"] == "modloss", "FragmentLossType"] = modloss - fragment_intensity_df.mask(masks, 0, inplace=True) + return df def speclib_to_single_df( @@ -242,18 +143,12 @@ def speclib_to_single_df( min_frag_intensity=0.01, min_frag_nAA=0, modloss: str = "H3PO4", - frag_type_head: str = "FragmentType", - frag_mass_head: str = "FragmentMz", - frag_inten_head: str = "RelativeIntensity", - frag_charge_head: str = "FragmentCharge", - frag_loss_head: str = "FragmentLossType", - frag_series_head: str = "FragmentNumber", verbose=True, ) -> pd.DataFrame: - """ - Convert alphabase library to diann (or Spectronaut) library dataframe - This method is not important, as it will be only - used by DiaNN, or spectronaut, or others + """Convert an alphabase library into a SWATH/Spectronaut transition list. + + The in-memory counterpart of :func:`translate_to_tsv`, which writes the same + table to a file. Parameters ---------- @@ -262,8 +157,12 @@ def speclib_to_single_df( use unimod name if None. Defaults to None. - keep_k_highest_peaks : int - only keep highest fragments for each precursor. Default: 12 + keep_k_highest_fragments : int + Keep only the k most intense fragments per precursor. Default: 12 + + min_frag_mz, max_frag_mz : float + Fragment m/z range; fragments outside it are dropped. Pass 0 for no lower bound + and `np.inf` for no upper bound. Returns ------- @@ -271,91 +170,19 @@ def speclib_to_single_df( a single dataframe in the SWATH-like format """ - df = pd.DataFrame() - df["ModifiedPeptide"] = speclib._precursor_df[ - ["sequence", "mods", "mod_sites"] - ].apply( - create_modified_sequence, - axis=1, - translate_mod_dict=translate_mod_dict, - mod_sep="[]", - ) - - df["frag_start_idx"] = speclib._precursor_df["frag_start_idx"] - df["frag_stop_idx"] = speclib._precursor_df["frag_stop_idx"] - - df["PrecursorCharge"] = speclib._precursor_df["charge"] - - for rt_col in ["irt_pred", "rt_pred", "rt", "irt", "rt_norm"]: - if rt_col in speclib.precursor_df.columns: - df["RT"] = speclib.precursor_df[rt_col] - break - if "RT" not in df.columns: - raise ValueError("precursor_df must contain the RT columns") - - for im_col in ["mobility_pred", "mobility"]: - if im_col in speclib.precursor_df.columns: - df["IonMobility"] = speclib.precursor_df[im_col] - break - - for ccs_col in ["ccs_pred", "ccs"]: - if ccs_col in speclib.precursor_df.columns: - df["CCS"] = speclib.precursor_df[ccs_col] - break - - # df['LabelModifiedSequence'] = df['ModifiedPeptide'] - df["StrippedPeptide"] = speclib.precursor_df["sequence"] - - if "precursor_mz" not in speclib._precursor_df.columns: - speclib.calc_precursor_mz() - df["PrecursorMz"] = speclib._precursor_df["precursor_mz"] - - for prot_col in ["uniprot_ids", "proteins"]: - if prot_col in speclib.precursor_df.columns: - df["ProteinID"] = speclib.precursor_df[prot_col] - break - - if "genes" in speclib._precursor_df.columns: - df["Genes"] = speclib._precursor_df["genes"] - - if "decoy" in speclib._precursor_df.columns: - df["Decoy"] = speclib._precursor_df["decoy"] - - # if 'protein_group' in speclib._precursor_df.columns: - # df['ProteinGroups'] = speclib._precursor_df['protein_group'] - - if min_frag_mz > 0 or max_frag_mz > 0: - mask_fragment_intensity_by_mz_( - speclib._fragment_mz_df, - speclib._fragment_intensity_df, - min_frag_mz, - max_frag_mz, - ) - - if min_frag_nAA > 0: - mask_fragment_intensity_by_frag_nAA( - speclib._fragment_intensity_df, - speclib._precursor_df, - max_mask_frag_nAA=min_frag_nAA - 1, - ) - - df = merge_precursor_fragment_df( - df, + return _precursors_to_swath_df( + speclib._precursor_df, speclib._fragment_mz_df, speclib._fragment_intensity_df, - top_n_inten=keep_k_highest_fragments, - frag_type_head=frag_type_head, - frag_mass_head=frag_mass_head, - frag_inten_head=frag_inten_head, - frag_charge_head=frag_charge_head, - frag_loss_head=frag_loss_head, - frag_series_head=frag_series_head, + translate_mod_dict=translate_mod_dict, + keep_k_highest_fragments=keep_k_highest_fragments, + min_frag_mz=min_frag_mz, + max_frag_mz=max_frag_mz, + min_frag_intensity=min_frag_intensity, + min_frag_nAA=min_frag_nAA, + modloss=modloss, verbose=verbose, ) - df = df[df["RelativeIntensity"] > min_frag_intensity] - df.loc[df[frag_loss_head] == "modloss", frag_loss_head] = modloss - - return df.drop(["frag_start_idx", "frag_stop_idx"], axis=1) def speclib_to_swath_df( @@ -387,17 +214,13 @@ def run(self): df, batch = self.task_queue.get() if df is None: break - if tuple([int(i) for i in pd.__version__.split(".")[:2]]) >= (1, 5): - newline = dict(lineterminator="\n") - else: - newline = dict(line_terminator="\n") df.to_csv( self.tsv, header=(batch == 0), sep="\t", mode="a", index=False, - **newline, + **_CSV_NEWLINE, ) @@ -414,6 +237,19 @@ def translate_to_tsv( translate_mod_dict: dict = None, multiprocessing: bool = True, ): + """Translate an alphabase library into a SWATH/Spectronaut transition-list tsv. + + Precursors are converted in batches and appended to the tsv, so large libraries do + not need to be held in memory at once in the flat, one-row-per-fragment format. + + Parameters + ---------- + min_frag_mz, max_frag_mz : float + Fragment m/z range; fragments outside it are dropped. Pass 0 for no lower bound + and `np.inf` for no upper bound. + + See :func:`speclib_to_single_df`, whose parameters this shares, for the rest. + """ if multiprocessing: queue_size = 1000000 // batch_size if queue_size < 2: @@ -423,47 +259,36 @@ def translate_to_tsv( df_head_queue = mp.Queue(maxsize=queue_size) writing_process = WritingProcess(df_head_queue, tsv) writing_process.start() - mask_fragment_intensity_by_mz_( - speclib._fragment_mz_df, - speclib._fragment_intensity_df, - min_frag_mz, - max_frag_mz, - ) - if min_frag_nAA > 0: - mask_fragment_intensity_by_frag_nAA( - speclib._fragment_intensity_df, - speclib._precursor_df, - max_mask_frag_nAA=min_frag_nAA - 1, - ) if isinstance(tsv, str): with open(tsv, "w"): pass - # process precursors in batches: the flat (one row per fragment) format is much larger - # than the compact library, so batching keeps peak memory bounded for large libraries - batch_speclib = SpecLibBase() - batch_speclib._fragment_intensity_df = speclib._fragment_intensity_df - batch_speclib._fragment_mz_df = speclib._fragment_mz_df + + # only the precursors are batched -- the fragment indices are absolute precursor_df = speclib._precursor_df - for i in tqdm.tqdm(range(0, len(precursor_df), batch_size)): - batch_speclib._precursor_df = precursor_df.iloc[i : i + batch_size] - df = speclib_to_single_df( - batch_speclib, + for first_row in tqdm.tqdm(range(0, len(precursor_df), batch_size)): + df = _precursors_to_swath_df( + precursor_df.iloc[first_row : first_row + batch_size], + speclib._fragment_mz_df, + speclib._fragment_intensity_df, translate_mod_dict=translate_mod_dict, keep_k_highest_fragments=keep_k_highest_fragments, - min_frag_mz=0, - max_frag_mz=0, + min_frag_mz=min_frag_mz, + max_frag_mz=max_frag_mz, min_frag_intensity=min_frag_intensity, - min_frag_nAA=0, + min_frag_nAA=min_frag_nAA, verbose=False, ) if multiprocessing: - df_head_queue.put((df, i)) + df_head_queue.put((df, first_row)) else: - if tuple([int(i) for i in pd.__version__.split(".")[:2]]) >= (1, 5): - newline = dict(lineterminator="\n") - else: - newline = dict(line_terminator="\n") - df.to_csv(tsv, header=(i == 0), sep="\t", mode="a", index=False, **newline) + df.to_csv( + tsv, + header=(first_row == 0), + sep="\t", + mode="a", + index=False, + **_CSV_NEWLINE, + ) if multiprocessing: df_head_queue.put((None, None)) print( diff --git a/alphabase/spectral_library/translate_core.py b/alphabase/spectral_library/translate_core.py new file mode 100644 index 00000000..360c9369 --- /dev/null +++ b/alphabase/spectral_library/translate_core.py @@ -0,0 +1,375 @@ +"""Code shared by the two spectral library export formats. + +`translate` writes a SWATH/Spectronaut transition list. `translate_diann` writes a +DIA-NN 1.9.1+ parquet library. Both turn the same alphabase library into one row per +precursor and fragment, and only the output format differs. + +The parts they share live here: rendering modified sequences, picking which fragments +to keep, and finding the precursor columns to export. +""" + +import warnings +from typing import Optional, Union + +import numpy as np +import pandas as pd +import tqdm + +from alphabase.constants.modification import MOD_DF, ModificationKeys +from alphabase.numba_wrapper import numba_njit +from alphabase.peptide.precursor import update_precursor_mz +from alphabase.psm_reader.keys import ConstantsClass, PsmDfCols +from alphabase.utils import explode_multiple_columns + +# Candidate precursor columns in order of precedence, for libraries that carry more than +# one. The `*_pred` names are peptdeep's prediction outputs, which take priority over a +# measured value; `irt_pred` outranks `rt_pred` because an indexed RT is what a +# third-party library wants. +RT_COLUMNS = ["irt_pred", "rt_pred", PsmDfCols.RT, "irt", PsmDfCols.RT_NORM] +MOBILITY_COLUMNS = ["mobility_pred", PsmDfCols.MOBILITY] +CCS_COLUMNS = ["ccs_pred", PsmDfCols.CCS] + +# AlphaBase modification name -> UniMod id, for the formats that name mods by id. +# Modifications without a UniMod id are absent, so looking one up raises rather than +# writing a name the target software cannot parse. +mod_to_unimod_dict = { + mod_name: f"UniMod:{unimod_id}" + for mod_name, unimod_id in MOD_DF[["mod_name", "unimod_id"]].to_numpy() + if unimod_id not in (-1, "-1") +} + + +class FragmentTableCols(metaclass=ConstantsClass): + """Canonical columns of the flattened fragment table. + + Each export renames these to its own dialect. ``PRECURSOR_ROW`` is the positional + row of the precursor a fragment belongs to; it is what joins the table back to the + precursors and is not written out. + """ + + PRECURSOR_ROW = "precursor_row" + FRAG_TYPE = "frag_type" + MZ = "mz" + INTENSITY = "intensity" + CHARGE = "charge" + SERIES_NUMBER = "series_number" + LOSS_TYPE = "loss_type" + + +# the per-fragment columns, in the order the exports emit them +FRAGMENT_VALUE_COLUMNS = [ + FragmentTableCols.FRAG_TYPE, + FragmentTableCols.MZ, + FragmentTableCols.INTENSITY, + FragmentTableCols.CHARGE, + FragmentTableCols.SERIES_NUMBER, + FragmentTableCols.LOSS_TYPE, +] + + +def get_precursor_mz(precursor_df: pd.DataFrame) -> pd.Series: + """Return the precursors' m/z, leaving `precursor_df` alone. + + The read-only counterpart of + :func:`alphabase.peptide.precursor.update_precursor_mz`, which writes its result + into the frame it is handed. + """ + if PsmDfCols.PRECURSOR_MZ in precursor_df.columns: + return precursor_df[PsmDfCols.PRECURSOR_MZ] + return update_precursor_mz(precursor_df.copy())[PsmDfCols.PRECURSOR_MZ] + + +def first_present_column( + precursor_df: pd.DataFrame, + candidates: list[str], + default: Union[str, float, None] = None, +) -> Union[pd.Series, str, float, None]: + """Return the first present candidate column of `precursor_df`, else `default`. + + Parameters + ---------- + precursor_df : pd.DataFrame + The precursor frame to look in. + + candidates : list of str + Column names in order of precedence. + + default : str or float or None + Returned when the frame carries none of the candidates. Defaults to None, which + lets a caller tell "absent" from a legitimate value and omit the output column. + + Returns + ------- + pd.Series or str or float or None + The first candidate column present, else `default`. + + """ + for column in candidates: + if column in precursor_df.columns: + return precursor_df[column] + return default + + +def create_modified_sequence( + seq_mods_sites: tuple, # must be ('sequence','mods','mod_sites') + translate_mod_dict: Optional[dict] = None, + mod_sep: str = "[]", + nterm: str = "_", + cterm: str = "_", +) -> str: + """Translate `(sequence, mods, mod_sites)` into a modified sequence. + + Used by `df.apply()`. Sites are 1-based, 0 is the N-terminus and -1 the + C-terminus:: + + ('ABCDEFG', 'Mod1@A;Mod2@E', '1;5') -> _A[Mod1]BCDE[Mod2]FG_ + ('PEPTIDE', 'Acetyl@Protein_N-term', '0') -> _[Acetyl]PEPTIDE_ + ('PEPTIDE', 'Amidated@Any_C-term', '-1') -> _PEPTIDE_[Amidated] + + Mods are inserted from the C-terminal end inwards, so an earlier insertion + cannot shift a later site. + + Parameters + ---------- + seq_mods_sites : tuple + Must be `(sequence, mods, mod_sites)`. + + translate_mod_dict : dict + A dict to map AlphaX modification names to other software; the bare AlphaBase + name (everything before the `@`) is used if None. Defaults to None. + + mod_sep : str + '[]' or '()', default '[]'. + + nterm : str + Rendered before the sequence, and carries a site-0 modification. + + cterm : str + Rendered after the sequence, and carries a site--1 modification. + + Returns + ------- + str + The modified sequence. + + """ + mod_seq, mods, mod_sites = seq_mods_sites + if mods: + mods = mods.split(ModificationKeys.SEPARATOR) + mod_sites = [int(i) for i in mod_sites.split(ModificationKeys.SEPARATOR)] + rev_order = np.argsort(mod_sites)[::-1] + mod_sites = [mod_sites[rev_order[i]] for i in range(len(mod_sites))] + mods = [mods[rev_order[i]] for i in range(len(mods))] + if translate_mod_dict is None: + mods = [mod[: mod.find(ModificationKeys.SITE_SEPARATOR)] for mod in mods] + else: + mods = [translate_mod_dict[mod] for mod in mods] + for _site, mod in zip(mod_sites, mods): + if _site == -1: + cterm += mod_sep[0] + mod + mod_sep[1] + elif _site == 0: + nterm += mod_sep[0] + mod + mod_sep[1] + else: + mod_seq = ( + mod_seq[:_site] + mod_sep[0] + mod + mod_sep[1] + mod_seq[_site:] + ) + return nterm + mod_seq + cterm + + +def is_nterm_frag(frag_type: str) -> bool: + """Whether a fragment column name is an N-terminal (a/b/c) series.""" + return frag_type[0] in "abc" + + +@numba_njit +def _get_frag_info_from_column_name(column: str) -> tuple: + """Split a fragment column name into `(frag_type, loss_type, charge)`. + + For example `y_modloss_z2` -> `('y', 'modloss', '2')` and `b_z1` -> `('b', + 'noloss', '1')`. The charge is left as a string, as it is only written out. + """ + idx = column.rfind("_") + frag_type = column[:idx] + charge = column[idx + 2 :] + if len(frag_type) == 1: + loss_type = "noloss" + else: + idx = frag_type.find("_") + loss_type = frag_type[idx + 1 :] + frag_type = frag_type[0] + return frag_type, loss_type, charge + + +def _get_frag_num(columns: np.ndarray, rows: np.ndarray, frag_len: int) -> list: + """Number each fragment within its series. + + N-terminal series are numbered from the start of the peptide and C-terminal ones + from the end, so row `r` of a precursor with `frag_len` fragment rows is `r + 1` + for a b-ion and `frag_len - r` for a y-ion. + """ + return [ + row + 1 if is_nterm_frag(column) else frag_len - row + for row, column in zip(rows, columns) + ] + + +def get_fragment_table( # noqa: PLR0913 + frag_start_idx: np.ndarray, + frag_stop_idx: np.ndarray, + fragment_mz_df: pd.DataFrame, + fragment_intensity_df: pd.DataFrame, + *, + keep_k_highest: int, + min_frag_mz: float = 0, + max_frag_mz: float = np.inf, + min_frag_nAA: int = 0, # noqa: N803 + verbose: bool = True, +) -> pd.DataFrame: + """Flatten each precursor's most intense fragments into one row per fragment. + + Works on a per-precursor copy, so the library's fragment frames are untouched. + Fragments outside the m/z window are dropped, as are empty slots. Intensities are + normalized to each precursor's most intense kept fragment, and the `keep_k_highest` + highest are kept in descending order. The default bounds `0` and `np.inf` accept + every fragment, so an unbounded window needs no special handling. + + Parameters + ---------- + frag_start_idx, frag_stop_idx : np.ndarray + Per precursor, the half-open row range into the fragment frames. These are + absolute offsets, so batching the precursors leaves the fragment frames whole. + + fragment_mz_df : pd.DataFrame + The library's fragment m/z frame. + + fragment_intensity_df : pd.DataFrame + The library's fragment intensity frame. + + keep_k_highest : int + Keep this many fragments per precursor. + + min_frag_mz : float + Drop fragments below this m/z. 0 for no lower bound, as m/z is positive. + + max_frag_mz : float + Drop fragments above this m/z. `np.inf` for no upper bound. + + min_frag_nAA : int + Drop the smallest `min_frag_nAA - 1` fragments of each series; 0 disables. The + off-by-one is the existing meaning of the export parameter of the same name. + + verbose : bool + Show a progress bar over the precursors. + + Returns + ------- + pd.DataFrame + One row per kept fragment, in :class:`FragmentTableCols` columns. + + """ + frag_columns = fragment_mz_df.columns.to_numpy().astype("U") + is_nterm = np.array([is_nterm_frag(column) for column in frag_columns]) + n_masked_per_terminus = max(min_frag_nAA - 1, 0) + + if min_frag_mz == 0 and max_frag_mz == 0: + warnings.warn( + "Disabling the fragment m/z window with min_frag_mz=0, max_frag_mz=0 is " + "deprecated; pass max_frag_mz=np.inf instead. min_frag_mz=0 already means " + "no lower bound, as m/z is positive.", + FutureWarning, + ) + max_frag_mz = np.inf + + frag_types = [] + frag_losses = [] + frag_charges = [] + frag_masses = [] + frag_intensities = [] + frag_numbers = [] + frag_idx_ranges = zip(frag_start_idx, frag_stop_idx) + if verbose: + frag_idx_ranges = tqdm.tqdm(frag_idx_ranges) + for start, end in frag_idx_ranges: + masses = fragment_mz_df.iloc[start:end, :].to_numpy() + keep_mask = (masses > 0) & (masses >= min_frag_mz) & (masses <= max_frag_mz) + if n_masked_per_terminus: + # b numbers count from the first row, y numbers from the last, so the + # smallest of each series sit at opposite ends of the block. `max(..., 0)` + # because a negative slice start wraps rather than clamping. + cterm_start = max(len(keep_mask) - n_masked_per_terminus, 0) + keep_mask[:n_masked_per_terminus, is_nterm] = False + keep_mask[cterm_start:, ~is_nterm] = False + + # `copy=True`, so normalizing and zeroing below cannot reach the library + intens = fragment_intensity_df.iloc[start:end, :].to_numpy(copy=True) + intens[~keep_mask] = 0 + max_inten = np.amax(intens) + if max_inten > 0: + intens /= max_inten + + sorted_idx = np.argsort(intens.reshape(-1))[-keep_k_highest:][::-1] + # a filtered-out slot can still be selected when a precursor has fewer than + # `keep_k_highest` fragments left, so drop those rather than export them + sorted_idx = sorted_idx[keep_mask.reshape(-1)[sorted_idx]] + idx_in_df = np.unravel_index(sorted_idx, masses.shape) + + frag_len = end - start + rows = np.arange(frag_len, dtype=np.int32)[idx_in_df[0]] + columns = frag_columns[idx_in_df[1]] + + infos = [_get_frag_info_from_column_name(column) for column in columns] + types, losses, charges = zip(*infos) if infos else ((), (), ()) + frag_types.append(types) + frag_losses.append(losses) + frag_charges.append(charges) + frag_masses.append(masses[idx_in_df]) + frag_intensities.append(intens[idx_in_df]) + frag_numbers.append(_get_frag_num(columns, rows, frag_len)) + + fragments_df = pd.DataFrame( + { + FragmentTableCols.PRECURSOR_ROW: np.arange(len(frag_start_idx)), + FragmentTableCols.FRAG_TYPE: frag_types, + FragmentTableCols.MZ: frag_masses, + FragmentTableCols.INTENSITY: frag_intensities, + FragmentTableCols.CHARGE: frag_charges, + FragmentTableCols.SERIES_NUMBER: frag_numbers, + FragmentTableCols.LOSS_TYPE: frag_losses, + } + ) + fragments_df = explode_multiple_columns(fragments_df, FRAGMENT_VALUE_COLUMNS) + # a precursor that kept nothing explodes to one all-NaN row; drop those + return fragments_df.dropna(subset=[FragmentTableCols.MZ]) + + +def join_fragments( + precursor_df: pd.DataFrame, + fragment_df: pd.DataFrame, + columns: dict, +) -> pd.DataFrame: + """Repeat each precursor row across its fragments, renamed to `columns`. + + Parameters + ---------- + precursor_df : pd.DataFrame + The export's precursor rows, in the order `fragment_df`'s `precursor_row` + indexes them. + + fragment_df : pd.DataFrame + A :func:`get_fragment_table` result. + + columns : dict + Maps :class:`FragmentTableCols` names to this format's output names. Its order + is the order the fragment columns are appended in. + + Returns + ------- + pd.DataFrame + One row per precursor/fragment pair, keeping `precursor_df`'s index. + + """ + rows = fragment_df[FragmentTableCols.PRECURSOR_ROW].to_numpy() + joined = precursor_df.iloc[rows].copy() + for canonical, name in columns.items(): + joined[name] = fragment_df[canonical].to_numpy() + return joined diff --git a/alphabase/spectral_library/translate_diann.py b/alphabase/spectral_library/translate_diann.py index 756e825f..c2b5d34c 100644 --- a/alphabase/spectral_library/translate_diann.py +++ b/alphabase/spectral_library/translate_diann.py @@ -1,22 +1,25 @@ """Translate AlphaBase spectral libraries to DIA-NN 1.9.1+ parquet format. -This reuses shared export helpers from -:mod:`alphabase.spectral_library.translate`; it is kept as a separate module in -preparation for a larger refactor of the library-export code. +The shared export machinery lives in :mod:`alphabase.spectral_library.translate_core`; +this module holds the DIA-NN schema, the precursor mapping and the parquet writer. """ -from typing import Optional, Union +from typing import Optional import pandas as pd import tqdm from alphabase.psm_reader.keys import ConstantsClass, LibPsmDfCols, PsmDfCols from alphabase.spectral_library.base import SpecLibBase -from alphabase.spectral_library.translate import ( +from alphabase.spectral_library.translate_core import ( + MOBILITY_COLUMNS, + RT_COLUMNS, + FragmentTableCols, create_modified_sequence, - mask_fragment_intensity_by_frag_nAA, - mask_fragment_intensity_by_mz_, - merge_precursor_fragment_df, + first_present_column, + get_fragment_table, + get_precursor_mz, + join_fragments, mod_to_unimod_dict, ) @@ -24,8 +27,8 @@ class DiannParquetCols(metaclass=ConstantsClass): """Column names of a DIA-NN 1.9.1+ `.parquet` spectral library. - DIA-NN uses its report-style dot notation here. ``SIGNATURE`` is listed for - completeness but is not written: DIA-NN requires third-party libraries to omit it. + DIA-NN uses its report-style dot notation here. ``Signature`` is deliberately + absent: DIA-NN requires third-party libraries to omit it. """ PRECURSOR_ID = "Precursor.Id" @@ -57,17 +60,16 @@ class DiannParquetCols(metaclass=ConstantsClass): GENES = "Genes" FLAGS = "Flags" SOURCE_ID = "Source.Id" - SIGNATURE = "Signature" -# fragment column names passed to `merge_precursor_fragment_df` -DIANN_PARQUET_FRAG_HEADS = { - "frag_type_head": DiannParquetCols.FRAGMENT_TYPE, - "frag_mass_head": DiannParquetCols.PRODUCT_MZ, - "frag_inten_head": DiannParquetCols.RELATIVE_INTENSITY, - "frag_charge_head": DiannParquetCols.FRAGMENT_CHARGE, - "frag_series_head": DiannParquetCols.FRAGMENT_SERIES_NUMBER, - "frag_loss_head": DiannParquetCols.FRAGMENT_LOSS_TYPE, +# the DIA-NN names for the canonical fragment columns, in output order +DIANN_FRAGMENT_COLUMNS = { + FragmentTableCols.FRAG_TYPE: DiannParquetCols.FRAGMENT_TYPE, + FragmentTableCols.MZ: DiannParquetCols.PRODUCT_MZ, + FragmentTableCols.INTENSITY: DiannParquetCols.RELATIVE_INTENSITY, + FragmentTableCols.CHARGE: DiannParquetCols.FRAGMENT_CHARGE, + FragmentTableCols.SERIES_NUMBER: DiannParquetCols.FRAGMENT_SERIES_NUMBER, + FragmentTableCols.LOSS_TYPE: DiannParquetCols.FRAGMENT_LOSS_TYPE, } # dtype tokens for DIANN_PARQUET_SCHEMA (INT64 / FLOAT=float32 / str) @@ -121,22 +123,10 @@ class DiannParquetCols(metaclass=ConstantsClass): _DIANN_FLAG_FIRST_FRAGMENT = 1 << 4 -def _get_first_present_column( +def _precursors_to_diann_df( # noqa: PLR0913 precursor_df: pd.DataFrame, - candidates: list[str], - default: Union[str, float, None] = None, -) -> Union[pd.Series, str, float, None]: - """Return the first present candidate column of `precursor_df`, else `default`.""" - for col in candidates: - if col in precursor_df.columns: - return precursor_df[col] - return default - - -# TODO: go for an OOP approach: a writer class holding the export settings as state, with -# the precursor mapping / fragment explosion / dtype casting as methods. -def speclib_to_diann_df( # noqa: PLR0913, PLR0915 - speclib: SpecLibBase, + fragment_mz_df: pd.DataFrame, + fragment_intensity_df: pd.DataFrame, *, translate_mod_dict: Optional[dict] = None, keep_k_highest_fragments: int = 12, @@ -147,57 +137,14 @@ def speclib_to_diann_df( # noqa: PLR0913, PLR0915 modloss: str = "H3PO4", verbose: bool = True, ) -> pd.DataFrame: - """Convert an alphabase library to a DIA-NN 1.9.1+ parquet-format dataframe. - - Emits DIA-NN's report-style dot-notation columns (see ``DIANN_PARQUET_SCHEMA``) and - ``(UniMod:N)`` modified sequences, importable by DIA-NN 1.9.1+ and readable back with - :class:`alphabase.spectral_library.reader.LibraryReaderBase`. Columns a ``SpecLibBase`` - has no value for are filled with defaults matching a DIA-NN predicted library (q-values - and scores 0, ``PTM.Site.Confidence`` 1, ``Source.Id`` empty). - - ``N.Term``/``C.Term`` are protein-terminus flags taken from ``is_prot_nterm``/ - ``is_prot_cterm`` (alphabase FASTA digestion) if present, else 0. ``Signature`` is not - written, as DIA-NN requires for third-party libraries. - - Parameters - ---------- - speclib : SpecLibBase - The alphabase spectral library to convert. - - translate_mod_dict : dict - Maps AlphaBase modification names to other software; defaults to UniMod ids. - - keep_k_highest_fragments : int - Keep only the k most intense fragments per precursor. Default: 12 - - min_frag_mz, max_frag_mz : float - Fragment m/z range; fragments outside it are dropped. Set both to 0 to disable. - - min_frag_intensity : float - Drop fragments whose relative intensity is at or below this value. - - min_frag_nAA : int - Mask the smallest ``min_frag_nAA - 1`` b/y fragments per precursor; 0 disables. - - modloss : str - Loss label written for modification-loss fragments. Default: "H3PO4" - - verbose : bool - Show a progress bar while exploding fragments. - - Returns - ------- - pd.DataFrame - A long-format dataframe in the DIA-NN parquet library schema. + """Convert precursor and fragment frames to a DIA-NN parquet-format dataframe. + The dataframe-in form of :func:`speclib_to_diann_df`, so that one batch of + precursors can be converted without standing up a `SpecLibBase` around it. """ if translate_mod_dict is None: translate_mod_dict = mod_to_unimod_dict - if PsmDfCols.PRECURSOR_MZ not in speclib.precursor_df.columns: - speclib.calc_precursor_mz() - precursor_df = speclib.precursor_df - df = pd.DataFrame(index=precursor_df.index) df[DiannParquetCols.MODIFIED_SEQUENCE] = precursor_df[ @@ -215,39 +162,35 @@ def speclib_to_diann_df( # noqa: PLR0913, PLR0915 df[DiannParquetCols.PRECURSOR_ID] = df[DiannParquetCols.MODIFIED_SEQUENCE] + df[ DiannParquetCols.PRECURSOR_CHARGE ].astype(str) - df[DiannParquetCols.PRECURSOR_MZ] = precursor_df[PsmDfCols.PRECURSOR_MZ] + df[DiannParquetCols.PRECURSOR_MZ] = get_precursor_mz(precursor_df) - rt = _get_first_present_column( - precursor_df, ["irt_pred", "rt_pred", PsmDfCols.RT, "irt", PsmDfCols.RT_NORM] - ) + rt = first_present_column(precursor_df, RT_COLUMNS) if rt is None: raise ValueError("precursor_df must contain a retention time column") df[DiannParquetCols.RT] = rt - df[DiannParquetCols.IM] = _get_first_present_column( - precursor_df, ["mobility_pred", PsmDfCols.MOBILITY], 0.0 - ) + df[DiannParquetCols.IM] = first_present_column(precursor_df, MOBILITY_COLUMNS, 0.0) - df[DiannParquetCols.PROTEIN_GROUP] = _get_first_present_column( + df[DiannParquetCols.PROTEIN_GROUP] = first_present_column( precursor_df, [PsmDfCols.PROTEINS, PsmDfCols.UNIPROT_IDS], "" ) - df[DiannParquetCols.PROTEIN_IDS] = _get_first_present_column( + df[DiannParquetCols.PROTEIN_IDS] = first_present_column( precursor_df, [PsmDfCols.UNIPROT_IDS, PsmDfCols.PROTEINS], "" ) - df[DiannParquetCols.PROTEIN_NAMES] = _get_first_present_column( + df[DiannParquetCols.PROTEIN_NAMES] = first_present_column( precursor_df, ["protein_names"], "" ) - df[DiannParquetCols.GENES] = _get_first_present_column( + df[DiannParquetCols.GENES] = first_present_column( precursor_df, [PsmDfCols.GENES], "" ) - df[DiannParquetCols.DECOY] = _get_first_present_column( + df[DiannParquetCols.DECOY] = first_present_column( precursor_df, [PsmDfCols.DECOY], 0 ) # N.Term/C.Term mark peptides at the protein N-/C-terminus (from FASTA digestion) - df[DiannParquetCols.N_TERM] = _get_first_present_column( + df[DiannParquetCols.N_TERM] = first_present_column( precursor_df, ["is_prot_nterm"], 0 ) - df[DiannParquetCols.C_TERM] = _get_first_present_column( + df[DiannParquetCols.C_TERM] = first_present_column( precursor_df, ["is_prot_cterm"], 0 ) @@ -265,31 +208,18 @@ def speclib_to_diann_df( # noqa: PLR0913, PLR0915 df[DiannParquetCols.EXCLUDE_FROM_QUANT] = 0 df[DiannParquetCols.SOURCE_ID] = "" - df[LibPsmDfCols.FRAG_START_IDX] = precursor_df[LibPsmDfCols.FRAG_START_IDX] - df[LibPsmDfCols.FRAG_STOP_IDX] = precursor_df[LibPsmDfCols.FRAG_STOP_IDX] - - if min_frag_mz > 0 or max_frag_mz > 0: - mask_fragment_intensity_by_mz_( - speclib.fragment_mz_df, - speclib.fragment_intensity_df, - min_frag_mz, - max_frag_mz, - ) - if min_frag_nAA > 0: - mask_fragment_intensity_by_frag_nAA( - speclib.fragment_intensity_df, - speclib.precursor_df, - max_mask_frag_nAA=min_frag_nAA - 1, - ) - - df = merge_precursor_fragment_df( - df, - speclib.fragment_mz_df, - speclib.fragment_intensity_df, - top_n_inten=keep_k_highest_fragments, + fragments = get_fragment_table( + precursor_df[LibPsmDfCols.FRAG_START_IDX].to_numpy(), + precursor_df[LibPsmDfCols.FRAG_STOP_IDX].to_numpy(), + fragment_mz_df, + fragment_intensity_df, + keep_k_highest=keep_k_highest_fragments, + min_frag_mz=min_frag_mz, + max_frag_mz=max_frag_mz, + min_frag_nAA=min_frag_nAA, verbose=verbose, - **DIANN_PARQUET_FRAG_HEADS, ) + df = join_fragments(df, fragments, DIANN_FRAGMENT_COLUMNS) df = df[df[DiannParquetCols.RELATIVE_INTENSITY] > min_frag_intensity] df.loc[ df[DiannParquetCols.FRAGMENT_LOSS_TYPE] == "modloss", @@ -305,8 +235,6 @@ def speclib_to_diann_df( # noqa: PLR0913, PLR0915 ].idxmax() df.loc[base_peak_idx, DiannParquetCols.FLAGS] |= _DIANN_FLAG_FIRST_FRAGMENT - df = df.drop([LibPsmDfCols.FRAG_START_IDX, LibPsmDfCols.FRAG_STOP_IDX], axis=1) - for name, dtype in DIANN_PARQUET_SCHEMA: if dtype == "str": df[name] = df[name].fillna("").astype(str) @@ -315,6 +243,77 @@ def speclib_to_diann_df( # noqa: PLR0913, PLR0915 return df[DIANN_PARQUET_COLUMN_ORDER] +def speclib_to_diann_df( # noqa: PLR0913 + speclib: SpecLibBase, + *, + translate_mod_dict: Optional[dict] = None, + keep_k_highest_fragments: int = 12, + min_frag_mz: float = 200, + max_frag_mz: float = 2000, + min_frag_intensity: float = 0.01, + min_frag_nAA: int = 0, # noqa: N803 + modloss: str = "H3PO4", + verbose: bool = True, +) -> pd.DataFrame: + """Convert an alphabase library to a DIA-NN 1.9.1+ parquet-format dataframe. + + Emits DIA-NN's report-style dot-notation columns (see ``DIANN_PARQUET_SCHEMA``) and + ``(UniMod:N)`` modified sequences, importable by DIA-NN 1.9.1+ and readable back with + :class:`alphabase.spectral_library.reader.LibraryReaderBase`. Columns a ``SpecLibBase`` + has no value for are filled with defaults matching a DIA-NN predicted library (q-values + and scores 0, ``PTM.Site.Confidence`` 1, ``Source.Id`` empty). + + ``N.Term``/``C.Term`` are protein-terminus flags taken from ``is_prot_nterm``/ + ``is_prot_cterm`` (alphabase FASTA digestion) if present, else 0. + + Parameters + ---------- + speclib : SpecLibBase + The alphabase spectral library to convert. + + translate_mod_dict : dict + Maps AlphaBase modification names to other software; defaults to UniMod ids. + + keep_k_highest_fragments : int + Keep only the k most intense fragments per precursor. Default: 12 + + min_frag_mz, max_frag_mz : float + Fragment m/z range; fragments outside it are dropped. Pass 0 for no lower bound + and `np.inf` for no upper bound. + + min_frag_intensity : float + Drop fragments whose relative intensity is at or below this value. + + min_frag_nAA : int + Mask the smallest ``min_frag_nAA - 1`` b/y fragments per precursor; 0 disables. + + modloss : str + Loss label written for modification-loss fragments. Default: "H3PO4" + + verbose : bool + Show a progress bar while exploding fragments. + + Returns + ------- + pd.DataFrame + A long-format dataframe in the DIA-NN parquet library schema. + + """ + return _precursors_to_diann_df( + speclib.precursor_df, + speclib.fragment_mz_df, + speclib.fragment_intensity_df, + translate_mod_dict=translate_mod_dict, + keep_k_highest_fragments=keep_k_highest_fragments, + min_frag_mz=min_frag_mz, + max_frag_mz=max_frag_mz, + min_frag_intensity=min_frag_intensity, + min_frag_nAA=min_frag_nAA, + modloss=modloss, + verbose=verbose, + ) + + def translate_to_parquet( # noqa: PLR0913 speclib: SpecLibBase, parquet_path: str, @@ -348,7 +347,8 @@ def translate_to_parquet( # noqa: PLR0913 Keep only the k most intense fragments per precursor. Default: 12 min_frag_mz, max_frag_mz : float - Fragment m/z range; fragments outside it are dropped. Set both to 0 to disable. + Fragment m/z range; fragments outside it are dropped. Pass 0 for no lower bound + and `np.inf` for no upper bound. min_frag_intensity : float Drop fragments whose relative intensity is at or below this value. @@ -376,46 +376,25 @@ def translate_to_parquet( # noqa: PLR0913 [(name, arrow_type[dtype]()) for name, dtype in DIANN_PARQUET_SCHEMA] ) - if min_frag_mz > 0 or max_frag_mz > 0: - mask_fragment_intensity_by_mz_( - speclib.fragment_mz_df, - speclib.fragment_intensity_df, - min_frag_mz, - max_frag_mz, - ) - if min_frag_nAA > 0: - mask_fragment_intensity_by_frag_nAA( - speclib.fragment_intensity_df, - speclib.precursor_df, - max_mask_frag_nAA=min_frag_nAA - 1, - ) - - # process precursors in batches: the flat (one row per fragment) format is much larger - # than the compact library, so batching keeps peak memory bounded for large libraries. - # SpecLibBase has no public setters for the fragment frames, and its precursor_df setter - # would refine/reorder the batch, so the private frames are assigned directly here. - batch_speclib = SpecLibBase() - batch_speclib._fragment_intensity_df = speclib.fragment_intensity_df # noqa: SLF001 - batch_speclib._fragment_mz_df = speclib.fragment_mz_df # noqa: SLF001 - precursor_df = speclib.precursor_df - writer = pq.ParquetWriter(parquet_path, schema) try: - for i in tqdm.tqdm(range(0, len(precursor_df), batch_size)): - # Only the precursors are batched: frag_start_idx/frag_stop_idx are absolute offsets into - # the full fragment frames, so those stay whole for the lookup to stay in sync. - batch_speclib._precursor_df = precursor_df.iloc[i : i + batch_size] # noqa: SLF001 - df = speclib_to_diann_df( - batch_speclib, + # only the precursors are batched -- the fragment indices are absolute + precursor_df = speclib.precursor_df + for first_row in tqdm.tqdm(range(0, len(precursor_df), batch_size)): + df = _precursors_to_diann_df( + precursor_df.iloc[first_row : first_row + batch_size], + speclib.fragment_mz_df, + speclib.fragment_intensity_df, translate_mod_dict=translate_mod_dict, keep_k_highest_fragments=keep_k_highest_fragments, - min_frag_mz=0, - max_frag_mz=0, + min_frag_mz=min_frag_mz, + max_frag_mz=max_frag_mz, min_frag_intensity=min_frag_intensity, - min_frag_nAA=0, + min_frag_nAA=min_frag_nAA, verbose=False, ) - table = pa.Table.from_pandas(df, schema=schema, preserve_index=False) - writer.write_table(table) + writer.write_table( + pa.Table.from_pandas(df, schema=schema, preserve_index=False) + ) finally: writer.close() diff --git a/docs/modules_spectral_library.rst b/docs/modules_spectral_library.rst index 14f774ef..46256a54 100644 --- a/docs/modules_spectral_library.rst +++ b/docs/modules_spectral_library.rst @@ -10,4 +10,5 @@ alphabase.spectral_library spectral_library/reader spectral_library/peaks_reader spectral_library/translate + spectral_library/translate_core spectral_library/translate_diann diff --git a/docs/spectral_library/translate.rst b/docs/spectral_library/translate.rst index 3c07c29d..76bd29e8 100644 --- a/docs/spectral_library/translate.rst +++ b/docs/spectral_library/translate.rst @@ -2,7 +2,8 @@ alphabase.spectral_library.translate ===================================== Export a spectral library to tsv (:func:`translate_to_tsv`). For DIA-NN ``.parquet`` -export see :doc:`translate_diann `. +export see :doc:`translate_diann `; for the machinery both formats +share, :doc:`translate_core `. .. automodule:: alphabase.spectral_library.translate :members: diff --git a/docs/spectral_library/translate_core.rst b/docs/spectral_library/translate_core.rst new file mode 100644 index 00000000..511a7166 --- /dev/null +++ b/docs/spectral_library/translate_core.rst @@ -0,0 +1,11 @@ +alphabase.spectral_library.translate_core +========================================== + +Machinery shared by the spectral library export formats: the modified-sequence +rendering, the fragment selection and the candidate precursor columns used by both +:doc:`translate ` and :doc:`translate_diann `. + +.. automodule:: alphabase.spectral_library.translate_core + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/spectral_library/translate_diann.rst b/docs/spectral_library/translate_diann.rst index 22be29cc..98288efb 100644 --- a/docs/spectral_library/translate_diann.rst +++ b/docs/spectral_library/translate_diann.rst @@ -3,7 +3,8 @@ alphabase.spectral_library.translate_diann Write an alphabase library to a DIA-NN 1.9.1+ ``.parquet`` spectral library (:func:`translate_to_parquet`); see the -:doc:`DIA-NN parquet notebook <../nbs/diann_parquet_library>`. +:doc:`DIA-NN parquet notebook <../nbs/diann_parquet_library>`. For the machinery this +shares with the tsv export, see :doc:`translate_core `. .. automodule:: alphabase.spectral_library.translate_diann :members: diff --git a/tests/unit/spectral_library/test_translate.py b/tests/unit/spectral_library/test_translate.py index 850d80dd..3ac089ce 100644 --- a/tests/unit/spectral_library/test_translate.py +++ b/tests/unit/spectral_library/test_translate.py @@ -3,17 +3,11 @@ These tests pin the behaviour of `alphabase.spectral_library.translate` as it is today, so that the upcoming restructuring can be shown to change nothing. -Six tests deliberately pin *buggy* behaviour that a later commit fixes. Each -carries a `CHARACTERIZATION (bug)` note in its docstring and is expected to be -rewritten there: - -1. the export zeroes intensities in the caller's library and adds `precursor_mz` -2. so a second export of the same library silently differs from a fresh one -3. with the m/z window disabled, empty fragment slots leak into the output -4. `translate_to_tsv` masks by m/z unconditionally, so disabling the window - yields a file whose every fragment is at m/z 0 -5. `rt_norm_pred` is not accepted as a retention time column -6. the exploded fragment columns are object dtype, with a string `FragmentCharge` +Two tests still pin *buggy* behaviour that a later commit fixes. Each carries a +`CHARACTERIZATION (bug)` note in its docstring: + +1. `rt_norm_pred` is not accepted as a retention time column +2. the exploded fragment columns are object dtype, with a string `FragmentCharge` One more, `test_speclib_to_swath_df_returns_none`, pins a function that a later commit removes rather than fixes. @@ -112,7 +106,7 @@ def _export(speclib: SpecLibBase, **kwargs) -> pd.DataFrame: def _unfiltered(speclib: SpecLibBase, **kwargs) -> pd.DataFrame: """Export with every fragment filter disabled, to see the raw selection.""" return _export( - speclib, min_frag_mz=0, max_frag_mz=0, min_frag_intensity=0.0, **kwargs + speclib, min_frag_mz=0, max_frag_mz=np.inf, min_frag_intensity=0.0, **kwargs ) @@ -281,6 +275,21 @@ def test_mz_window_selects_fragments_inside_it() -> None: assert df["FragmentMz"].between(600, 900).all() +def test_unbounded_mz_window_spellings_agree() -> None: + """0 and `np.inf` mean no bound; the old 0/0 sentinel warns but still works.""" + unbounded = _export(_build_speclib(), min_frag_mz=0, max_frag_mz=np.inf) + pd.testing.assert_frame_equal( + unbounded, _export(_build_speclib(), min_frag_mz=-np.inf, max_frag_mz=np.inf) + ) + + with pytest.warns(FutureWarning, match="max_frag_mz=np.inf"): + deprecated = _export(_build_speclib(), min_frag_mz=0, max_frag_mz=0) + pd.testing.assert_frame_equal(unbounded, deprecated) + + # a bound of 0 on its own now means what it says, rather than "no bound" + assert len(_export(_build_speclib(), max_frag_mz=0)) == 0 + + def test_min_frag_nAA_masks_the_smallest_fragments() -> None: """`min_frag_nAA=n` removes the n-1 smallest b and y fragments per terminus.""" for min_frag_nAA in (2, 3, 4): @@ -290,6 +299,17 @@ def test_min_frag_nAA_masks_the_smallest_fragments() -> None: assert numbers.min() == min_frag_nAA, (series, min_frag_nAA) +def test_min_frag_nAA_wider_than_any_precursor() -> None: + """Regression guard: a mask wider than the block covers all of it, not part. + + A `min_frag_nAA` larger than any precursor's fragment count is not a request + the export defines an answer to -- it is pinned only because the masking + works on row offsets, where an unclamped negative start silently masks the + wrong end. This is what main does too. + """ + assert len(_export(_build_speclib(), min_frag_nAA=12)) == 0 + + def test_modloss_label() -> None: """The `modloss` loss label is replaced by the `modloss` argument's value.""" default = _export(_build_speclib()) @@ -300,73 +320,54 @@ def test_modloss_label() -> None: assert set(custom["FragmentLossType"]) <= {"noloss", "H2O"} -def test_export_modifies_the_source_library() -> None: - """CHARACTERIZATION (bug): the export edits the library it is handed. - - The m/z window is applied by zeroing intensities *in the library*, and - `precursor_mz` is computed onto the caller's precursor frame. A later commit - filters on a copy instead. - """ +def test_export_leaves_the_source_library_untouched() -> None: + """The export reads the library and writes nothing back to it.""" speclib = _build_speclib() intensities_before = speclib.fragment_intensity_df.to_numpy(copy=True) + mz_before = speclib.fragment_mz_df.to_numpy(copy=True) columns_before = set(speclib.precursor_df.columns) _export(speclib) - zeroed = (intensities_before != 0) & (speclib.fragment_intensity_df.to_numpy() == 0) - assert zeroed.sum() > 0 - assert set(speclib.precursor_df.columns) - columns_before == {"precursor_mz"} + np.testing.assert_array_equal( + speclib.fragment_intensity_df.to_numpy(), intensities_before + ) + np.testing.assert_array_equal(speclib.fragment_mz_df.to_numpy(), mz_before) + # `precursor_mz` in particular is computed on a copy, not onto the library + assert set(speclib.precursor_df.columns) == columns_before -def test_second_export_of_the_same_library_differs_from_a_fresh_one() -> None: - """CHARACTERIZATION (bug): the in-place edit corrupts later exports. +def test_second_export_of_the_same_library_matches_a_fresh_one() -> None: + """Exporting twice at different m/z windows is not order-dependent. - Exporting at 200-2000 and then at 100-3000 cannot recover the fragments the - first window zeroed, so the second export silently loses fragments a fresh - library would have kept, and substitutes others in their top-k slots. + A narrow window followed by a wider one used to be unable to recover the + fragments the first export had zeroed in the library. """ - - def fragment_keys(df: pd.DataFrame) -> set: - return set( - zip( - df["ModifiedPeptide"], - df["FragmentType"], - df["FragmentNumber"], - df["FragmentCharge"], - df["FragmentLossType"], - ) - ) - reused = _build_speclib() _export(reused, min_frag_mz=DEFAULT_MIN_FRAG_MZ, max_frag_mz=DEFAULT_MAX_FRAG_MZ) second = _export(reused, min_frag_mz=100, max_frag_mz=3000) fresh = _export(_build_speclib(), min_frag_mz=100, max_frag_mz=3000) - assert fragment_keys(second) != fragment_keys(fresh) - assert len(fragment_keys(fresh) - fragment_keys(second)) > 0 - + pd.testing.assert_frame_equal(second, fresh) -def test_disabled_mz_window_keeps_empty_fragment_slots() -> None: - """CHARACTERIZATION (bug): with the window off, m/z 0 padding is exported. - A precursor's `*_modloss` slots are 0 unless it carries a loss-bearing mod. - The m/z window is what removes them, so disabling it emits them as if they - were real fragments -- mostly labelled with the `modloss` loss label. A later - commit skips empty slots regardless of the window. +def test_disabled_mz_window_skips_empty_fragment_slots() -> None: + """Empty fragment slots are never exported, m/z window or not. - The fixture's seeded intensities give those empty slots a positive intensity, - which is what makes them win a top-k slot here; a predictor would leave them - near 0. The selection has no notion of an empty slot either way. + A precursor's `*_modloss` slots carry m/z 0 unless it has a loss-bearing mod. + The m/z window used to be the only thing removing them, so disabling it + emitted them as if they were real fragments -- and they took top-k slots away + from fragments that do exist. """ df = _unfiltered(_build_speclib()) - padding = df["FragmentMz"] == 0 - assert padding.sum() > 0 - # and the padding rows carry the modification-loss label - assert (df.loc[padding, "FragmentLossType"] == "H3PO4").all() - # only the three phospho-bearing precursors have any real loss fragment - real_loss = df[(df["FragmentLossType"] == "H3PO4") & (df["FragmentMz"] > 0)] - assert real_loss["ModifiedPeptide"].nunique() < padding.sum() + assert (df["FragmentMz"] > 0).all() + # the freed slots go to real fragments, so disabling the window cannot yield + # fewer of them than the default window does + assert len(df) >= len(_export(_build_speclib())) + # only the loss-bearing precursors carry a loss fragment + with_loss = set(df.loc[df["FragmentLossType"] == "H3PO4", "StrippedPeptide"]) + assert with_loss <= {"SVIVSPYSTGAK", "LHDSTPPPYK"} def test_exploded_fragment_columns_are_object_dtype() -> None: @@ -455,34 +456,38 @@ def test_translate_to_tsv_multiprocessing_matches_single_process(tmp_path) -> No assert len(set(digests)) == 1 -def test_translate_to_tsv_disabled_mz_window_writes_only_empty_slots( +def test_translate_to_tsv_disabled_mz_window_matches_the_in_memory_export( tmp_path, ) -> None: - """CHARACTERIZATION (bug): `min_frag_mz=0, max_frag_mz=0` writes m/z 0 rows. + """`min_frag_mz=0, max_frag_mz=np.inf` disables the filter in both entry points. - Unlike `speclib_to_single_df`, `translate_to_tsv` masks by m/z without - checking whether the window is disabled, so 0/0 -- the documented way to turn - the filter off -- zeroes every real fragment's intensity and leaves only the - empty slots to be selected. The file is a valid tsv of unusable fragments, - written without a warning. + `translate_to_tsv` used to mask by m/z without checking whether the window + was disabled, so 0/0 -- the documented way to turn the filter off -- zeroed + every real fragment and wrote a file of nothing but empty slots, while + `speclib_to_single_df` given the same arguments kept the real ones. """ tsv = str(tmp_path / "lib.tsv") translate_to_tsv( _build_speclib(), tsv, min_frag_mz=0, - max_frag_mz=0, + max_frag_mz=np.inf, min_frag_intensity=0.0, multiprocessing=False, ) written = pd.read_csv(tsv, sep="\t") assert len(written) > 0 - assert (written["FragmentMz"] == 0).all() + assert (written["FragmentMz"] > 0).all() - # the in-memory export disagrees: it guards the mask, so real fragments survive - in_memory = _unfiltered(_build_speclib()) - assert (in_memory["FragmentMz"] > 0).any() + expected = _unfiltered(_build_speclib()) + numeric = ["FragmentMz", "RelativeIntensity", "FragmentCharge", "FragmentNumber"] + expected[numeric] = expected[numeric].apply(pd.to_numeric) + pd.testing.assert_frame_equal( + written.reset_index(drop=True), + expected.reset_index(drop=True), + check_dtype=False, + ) def test_translate_to_tsv_writes_a_readable_library(tmp_path) -> None: diff --git a/tests/unit/spectral_library/test_translate_diann.py b/tests/unit/spectral_library/test_translate_diann.py index 03a3aa99..18b9e35c 100644 --- a/tests/unit/spectral_library/test_translate_diann.py +++ b/tests/unit/spectral_library/test_translate_diann.py @@ -55,7 +55,11 @@ def test_speclib_to_diann_df_columns_and_mod_format() -> None: speclib = _build_speclib() df = speclib_to_diann_df( - speclib, min_frag_mz=0, max_frag_mz=0, min_frag_intensity=0.0, verbose=False + speclib, + min_frag_mz=0, + max_frag_mz=np.inf, + min_frag_intensity=0.0, + verbose=False, ) # exact DIA-NN column set/order, and `Signature` must NOT be present @@ -96,7 +100,11 @@ def test_speclib_to_diann_df_flags() -> None: speclib = _build_speclib() df = speclib_to_diann_df( - speclib, min_frag_mz=0, max_frag_mz=0, min_frag_intensity=0.0, verbose=False + speclib, + min_frag_mz=0, + max_frag_mz=np.inf, + min_frag_intensity=0.0, + verbose=False, ) # every row has the base bit (1 << 0) @@ -119,7 +127,7 @@ def test_translate_to_parquet_roundtrip(tmp_path) -> None: out_path = str(tmp_path / "lib.parquet") translate_to_parquet( - speclib, out_path, min_frag_mz=0, max_frag_mz=0, min_frag_intensity=0.0 + speclib, out_path, min_frag_mz=0, max_frag_mz=np.inf, min_frag_intensity=0.0 ) exported = pd.read_parquet(out_path) @@ -157,7 +165,11 @@ def test_speclib_to_diann_df_rt_column_precedence(present: list, expected: str) speclib._precursor_df = speclib._precursor_df.drop(columns=["rt"]).assign(**values) df = speclib_to_diann_df( - speclib, min_frag_mz=0, max_frag_mz=0, min_frag_intensity=0.0, verbose=False + speclib, + min_frag_mz=0, + max_frag_mz=np.inf, + min_frag_intensity=0.0, + verbose=False, ) assert df["RT"].unique().tolist() == [values[expected]] @@ -173,7 +185,11 @@ def test_speclib_to_diann_df_rejects_rt_norm_pred() -> None: with pytest.raises(ValueError, match="must contain a retention time column"): speclib_to_diann_df( - speclib, min_frag_mz=0, max_frag_mz=0, min_frag_intensity=0.0, verbose=False + speclib, + min_frag_mz=0, + max_frag_mz=np.inf, + min_frag_intensity=0.0, + verbose=False, ) @@ -190,7 +206,11 @@ def test_speclib_to_diann_df_flags_group_by_precursor_id() -> None: n_precursors = len(speclib.precursor_df) df = speclib_to_diann_df( - speclib, min_frag_mz=0, max_frag_mz=0, min_frag_intensity=0.0, verbose=False + speclib, + min_frag_mz=0, + max_frag_mz=np.inf, + min_frag_intensity=0.0, + verbose=False, ) # every precursor row is exported, but only the distinct ids get a base peak @@ -198,33 +218,34 @@ def test_speclib_to_diann_df_flags_group_by_precursor_id() -> None: assert int((df["Flags"] & (1 << 4) > 0).sum()) == n_precursors // 2 -def test_speclib_to_diann_df_modifies_the_source_library() -> None: - """CHARACTERIZATION (bug): the export edits the library it is handed. - - The m/z window is applied by zeroing intensities *in the library*, and - `precursor_mz` is computed onto the caller's precursor frame. A later commit - filters on a copy instead. - """ +def test_speclib_to_diann_df_leaves_the_source_library_untouched() -> None: + """The export reads the library and writes nothing back to it.""" speclib = _build_speclib() intensities_before = speclib.fragment_intensity_df.to_numpy(copy=True) + mz_before = speclib.fragment_mz_df.to_numpy(copy=True) columns_before = set(speclib.precursor_df.columns) speclib_to_diann_df(speclib, verbose=False) - zeroed = (intensities_before != 0) & (speclib.fragment_intensity_df.to_numpy() == 0) - assert zeroed.sum() > 0 - assert set(speclib.precursor_df.columns) - columns_before == {"precursor_mz"} + np.testing.assert_array_equal( + speclib.fragment_intensity_df.to_numpy(), intensities_before + ) + np.testing.assert_array_equal(speclib.fragment_mz_df.to_numpy(), mz_before) + assert set(speclib.precursor_df.columns) == columns_before -def test_translate_to_parquet_modifies_the_source_library(tmp_path) -> None: - """CHARACTERIZATION (bug): the streaming export edits the library too.""" +def test_translate_to_parquet_leaves_the_source_library_untouched(tmp_path) -> None: + """The streaming export writes nothing back to the library either.""" speclib = _build_speclib() intensities_before = speclib.fragment_intensity_df.to_numpy(copy=True) + columns_before = set(speclib.precursor_df.columns) translate_to_parquet(speclib, str(tmp_path / "lib.parquet")) - zeroed = (intensities_before != 0) & (speclib.fragment_intensity_df.to_numpy() == 0) - assert zeroed.sum() > 0 + np.testing.assert_array_equal( + speclib.fragment_intensity_df.to_numpy(), intensities_before + ) + assert set(speclib.precursor_df.columns) == columns_before def test_translate_to_parquet_batching_does_not_change_the_output(tmp_path) -> None: @@ -236,7 +257,7 @@ def test_translate_to_parquet_batching_does_not_change_the_output(tmp_path) -> N _build_speclib(), path, min_frag_mz=0, - max_frag_mz=0, + max_frag_mz=np.inf, min_frag_intensity=0.0, batch_size=batch_size, )