[translate 2/3] Split the shared export machinery out of translate.py

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