corems.mass_spectra.calc.feature_grouping

   1"""
   2Natural-abundance isotope + singly-charged adduct grouping for consensus features.
   3
   4Unit isotope edges (``Atoms``, ``|z| = 1``) and pairwise ``ion_type_dict``
   5adduct links → Pearson apex-height gate → roll-up / merge labels.
   6Not for tracer labeling or multi-charge envelopes.
   7"""
   8
   9from __future__ import annotations
  10
  11from dataclasses import dataclass, replace
  12import re
  13import time
  14from typing import Dict, MutableMapping, Optional, Sequence, Tuple, Union
  15
  16import numpy as np
  17import pandas as pd
  18from scipy import sparse
  19from scipy.spatial import KDTree
  20from scipy.stats import pearsonr
  21
  22from corems.encapsulation.constant import Atoms, ION_TYPE_DICT
  23
  24_ION_TYPE_CHARGE_RE = re.compile(r"(\d*)([+-])\s*$")
  25
  26GROUP_COLUMNS = (
  27    "feature_group_id",
  28    "ion_role",
  29    "ion_type",
  30    "possible_ion_types",
  31    "isotope_state",
  32    "mono_cluster_id",
  33)
  34
  35# Separator for multi-candidate ion type strings (alphabetical, no preference).
  36POSSIBLE_ION_TYPES_SEP = ";"
  37
  38# Fixed quant-gate policy. Not user-selectable switches.
  39CORR_METHOD = "pearson"
  40HEIGHT_COL = "intensity"  # apex peak height; not integrated area
  41
  42# Explicit allow-list for feature-group adduct search (membership matters).
  43# Must be keys in ION_TYPE_DICT with |z| = 1. Residual ties keep all matches in
  44# possible_ion_types; ion_type is set only when a single interpretation fits.
  45DEFAULT_ION_TYPES: Tuple[str, ...] = (
  46    "[M+H]+",
  47    "[M+H-H2O]+",
  48    "[M-H]-",
  49    "[M+Na]+",
  50    "[M+H-NH3]+",
  51    "[M+NH4]+",
  52    "[M-H-H2O]-",
  53    "[M-H+2Na]+",
  54    "[M-H+H2O]-",
  55    "[M+NH4-H2O]+",
  56    "[M+H+H2O]+",
  57    "[M+K]+",
  58    "[M+H-2H2O]+",
  59    "[M]+",
  60)
  61
  62PolarityLike = Union[str, int, None]
  63
  64
  65def normalize_ms_polarity(polarity: PolarityLike) -> Optional[str]:
  66    """Normalize sample polarity to ``'positive'``, ``'negative'``, or ``None``.
  67
  68    Accepts common LCMS encodings: ``'positive'`` / ``'negative'``,
  69    ``1`` / ``-1``, ``'+'`` / ``'-'``, and short forms ``'pos'`` / ``'neg'``.
  70    Empty or unrecognized values return ``None``.
  71    """
  72    if polarity is None:
  73        return None
  74    if isinstance(polarity, (int, np.integer)):
  75        if int(polarity) > 0:
  76            return "positive"
  77        if int(polarity) < 0:
  78            return "negative"
  79        return None
  80    s = str(polarity).strip().lower()
  81    if not s:
  82        return None
  83    if s in ("positive", "pos", "+", "1"):
  84        return "positive"
  85    if s in ("negative", "neg", "-", "-1"):
  86        return "negative"
  87    return None
  88
  89
  90def ion_type_charge(ion_type: str) -> int:
  91    """Absolute charge from an ion-type key (``[M+H]+`` → 1, ``[M+2H]2+`` → 2)."""
  92    m = _ION_TYPE_CHARGE_RE.search(str(ion_type).strip())
  93    if m is None:
  94        raise ValueError(f"ion_type {ion_type!r} has no trailing charge marker")
  95    digits = m.group(1)
  96    return int(digits) if digits else 1
  97
  98
  99def is_allowed_adduct_type_pair(type_a: str, type_b: str) -> bool:
 100    """True if both keys are distinct singly-charged ion types."""
 101    if type_a == type_b:
 102        return False
 103    return ion_type_charge(type_a) == 1 and ion_type_charge(type_b) == 1
 104
 105
 106def filter_ion_types_for_polarity(
 107    ion_types: Sequence[str],
 108    polarity: PolarityLike,
 109) -> Tuple[str, ...]:
 110    """Keep only ion types whose charge sign matches sample polarity.
 111
 112    Prevents negative adducts such as ``[M+HCOO]-`` / ``[M+CH3COO]-`` from
 113    being considered on positive-mode data (and the reverse).
 114
 115    Parameters
 116    ----------
 117    ion_types :
 118        Candidate ``ion_type_dict`` keys (order preserved).
 119    polarity :
 120        Sample/collection polarity. If unknown (``None`` / unrecognized),
 121        ``ion_types`` is returned unchanged so unit tests and offline array
 122        callers are not forced to set polarity.
 123
 124    Returns
 125    -------
 126    tuple of str
 127        Filtered ion types matching ``ion_type_dict`` polarity.
 128    """
 129    pol = normalize_ms_polarity(polarity)
 130    if pol is None:
 131        return tuple(ion_types)
 132    return tuple(it for it in ion_types if ION_TYPE_DICT[it]["polarity"] == pol)
 133
 134
 135def params_with_polarity_filtered_ion_types(
 136    params: "FeatureGroupParams",
 137    polarity: PolarityLike,
 138) -> "FeatureGroupParams":
 139    """Return params with ``ion_types`` filtered to ``polarity`` (no-op if unknown)."""
 140    filtered = filter_ion_types_for_polarity(params.ion_types, polarity)
 141    if filtered == params.ion_types:
 142        return params
 143    return replace(params, ion_types=filtered)
 144
 145
 146@dataclass(frozen=True)
 147class FeatureGroupParams:
 148    """Parameters for consensus feature grouping (isotopes + adducts).
 149
 150    **Isotopes:** natural-abundance rare forms in ``Atoms`` at or above
 151    ``min_isotope_abundance``. Not for tracer / labeled experiments.
 152
 153    **Adducts:** ``ion_types`` lists ``ION_TYPE_DICT`` keys (``|z| = 1``).
 154    Pairwise neutral-mass edges; residual ties keep all matches in
 155    ``possible_ion_types`` with no preferred ``ion_type``.
 156
 157    ``rt_tol`` / ``mz_tol_ppm`` come from collection alignment settings via
 158    :meth:`from_lcms_collection_settings`. Correlation is Pearson on apex
 159    ``intensity``.
 160    """
 161
 162    rt_tol: float = 0.4
 163    mz_tol_ppm: float = 5.0
 164    # Mono elements for natural-abundance isotope edge search in feature grouping
 165    # (maps from LCMSCollectionSettings.feature_group_isotope_atoms).
 166    isotope_atoms: Tuple[str, ...] = ("C",)
 167    min_isotope_abundance: float = 0.01
 168    max_isotope_offset: int = 4
 169    corr_threshold: float = 0.80
 170    min_shared_sample_fraction: float = 0.15
 171    # Allow-list of ION_TYPE_DICT keys (|z|=1); membership matters, not order
 172    ion_types: Tuple[str, ...] = DEFAULT_ION_TYPES
 173    partition_size: int = 5000
 174    cores: int = 1
 175
 176    def min_shared_count(self, n_samples: int) -> int:
 177        """Minimum shared non-zero sample count for the correlation gate."""
 178        if n_samples < 1:
 179            return 1
 180        return max(1, int(np.ceil(self.min_shared_sample_fraction * n_samples)))
 181
 182    @classmethod
 183    def from_lcms_collection_settings(
 184        cls,
 185        settings,
 186        polarity: PolarityLike = None,
 187    ) -> "FeatureGroupParams":
 188        """Build params from LCMSCollectionSettings (or duck-typed object).
 189
 190        RT and m/z tolerances reuse ``alignment_rt_tol`` and
 191        ``alignment_mz_tol_ppm`` (not separate feature-group settings).
 192        Isotope spacing is the singly-charged Atoms mass difference.
 193        Mono elements for natural-abundance isotope spacing use
 194        ``feature_group_isotope_atoms``.
 195
 196        Parameters
 197        ----------
 198        settings :
 199            ``LCMSCollectionSettings`` (or duck-typed equivalent).
 200        polarity :
 201            Optional sample/collection polarity. When provided, ``ion_types``
 202            are filtered with :func:`filter_ion_types_for_polarity` so that
 203            e.g. formate/acetate negative adducts are not used on positive
 204            data. ``group_consensus_features()`` supplies collection polarity.
 205        """
 206        ion_types = filter_ion_types_for_polarity(
 207            tuple(settings.feature_group_ion_types), polarity
 208        )
 209        return cls(
 210            rt_tol=float(settings.alignment_rt_tol),
 211            mz_tol_ppm=float(settings.alignment_mz_tol_ppm),
 212            isotope_atoms=tuple(settings.feature_group_isotope_atoms),
 213            min_isotope_abundance=float(settings.feature_group_min_isotope_abundance),
 214            max_isotope_offset=int(settings.feature_group_max_isotope_offset),
 215            corr_threshold=float(settings.feature_group_corr_threshold),
 216            min_shared_sample_fraction=float(
 217                settings.feature_group_min_shared_sample_fraction
 218            ),
 219            ion_types=tuple(ion_types),
 220            partition_size=int(settings.feature_group_partition_size),
 221            cores=int(settings.cores),
 222        )
 223
 224
 225def validate_feature_group_params(params: FeatureGroupParams) -> None:
 226    """Raise ValueError if parameters are invalid."""
 227    if params.rt_tol <= 0:
 228        raise ValueError("alignment_rt_tol (feature grouping RT window) must be > 0")
 229    if params.mz_tol_ppm <= 0:
 230        raise ValueError("alignment_mz_tol_ppm (feature grouping m/z tol) must be > 0")
 231    if not params.isotope_atoms:
 232        raise ValueError(
 233            "feature_group_isotope_atoms must be non-empty "
 234            "(mono elements for natural-abundance feature-group isotope edges)"
 235        )
 236    if not (0.0 <= params.min_isotope_abundance <= 1.0):
 237        raise ValueError("feature_group_min_isotope_abundance must be in [0, 1]")
 238    if params.max_isotope_offset < 1:
 239        raise ValueError("feature_group_max_isotope_offset must be >= 1")
 240    if not (0.0 < params.min_shared_sample_fraction <= 1.0):
 241        raise ValueError("feature_group_min_shared_sample_fraction must be in (0, 1]")
 242    if params.corr_threshold < -1.0 or params.corr_threshold > 1.0:
 243        raise ValueError("feature_group_corr_threshold must be in [-1, 1]")
 244    if params.partition_size < 1:
 245        raise ValueError("feature_group_partition_size must be >= 1")
 246    # Validate ion types resolve in ion_type_dict (empty = isotopes only)
 247    seen_it = set()
 248    for it in params.ion_types:
 249        if it in seen_it:
 250            raise ValueError(f"Duplicate ion_type {it!r} in feature_group_ion_types")
 251        seen_it.add(it)
 252        _ion_type_mass_offset(it)
 253        if ion_type_charge(it) != 1:
 254            raise ValueError(
 255                f"feature_group_ion_types entry {it!r} is not singly charged; "
 256                "feature grouping currently supports |z| = 1 only"
 257            )
 258    for atom in params.isotope_atoms:
 259        rare_isotope_entries(
 260            atom, min_abundance=params.min_isotope_abundance
 261        )  # raises if unknown / none above floor
 262
 263
 264def _validate_mono_element(mono_symbol: str) -> None:
 265    """Raise ValueError if mono_symbol is not a usable Atoms mono element."""
 266    if mono_symbol not in Atoms.isotopes:
 267        raise ValueError(
 268            f"Unknown mono element '{mono_symbol}' in feature_group_isotope_atoms; "
 269            "must be a key in Atoms.isotopes"
 270        )
 271    if mono_symbol not in Atoms.atomic_masses:
 272        raise ValueError(
 273            f"Mono element '{mono_symbol}' missing from Atoms.atomic_masses"
 274        )
 275
 276
 277def rare_isotope_entries(
 278    mono_symbol: str,
 279    min_abundance: float = 0.01,
 280) -> Tuple[Tuple[str, float, float], ...]:
 281    """
 282    Natural-abundance rare isotopes for a mono element with signed mass deltas.
 283
 284    Includes **every** rare form listed for the element in ``Atoms.isotopes``
 285    whose **natural** (terrestrial) abundance is at least ``min_abundance`` (from
 286    ``Atoms.isotopic_abundance``). Multi-isotope elements (e.g. Se) therefore
 287    contribute multiple Δm targets, not only the first listed rare form.
 288
 289    These are natural-abundance isotopologues only—not enriched/tracer labels.
 290
 291    Parameters
 292    ----------
 293    mono_symbol : str
 294        Most-abundant (natural) isotope symbol (e.g. ``"C"``, ``"Fe"``, ``"Se"``).
 295    min_abundance : float
 296        Minimum natural abundance fraction (0–1). Default 0.01. Isotopes missing
 297        from ``Atoms.isotopic_abundance`` are skipped.
 298
 299    Returns
 300    -------
 301    tuple of (rare_label, signed_delta, abundance)
 302        ``signed_delta = m(rare) - m(mono)``. Positive for heavier rare forms
 303        (¹³C), negative when the listed rare isotope is lighter (⁵⁴Fe).
 304        Ordered by decreasing natural abundance (then by |signed_delta|).
 305    """
 306    _validate_mono_element(mono_symbol)
 307    heavies = Atoms.isotopes[mono_symbol][1]
 308    out = []
 309    for rare in heavies:
 310        if rare is None:
 311            continue
 312        if rare not in Atoms.atomic_masses:
 313            raise ValueError(
 314                f"Rare isotope '{rare}' for '{mono_symbol}' missing from "
 315                "Atoms.atomic_masses"
 316            )
 317        abun = Atoms.isotopic_abundance.get(rare)
 318        if abun is None:
 319            continue
 320        if float(abun) < float(min_abundance):
 321            continue
 322        signed = Atoms.atomic_masses[rare] - Atoms.atomic_masses[mono_symbol]
 323        if signed == 0:
 324            continue
 325        out.append((rare, float(signed), float(abun)))
 326    if not out:
 327        raise ValueError(
 328            f"Element '{mono_symbol}' has no rare isotope in Atoms with "
 329            f"natural abundance >= {min_abundance}"
 330        )
 331    out.sort(key=lambda t: (-t[2], abs(t[1])))
 332    return tuple(out)
 333
 334
 335def _heavy_isotope_label(mono_symbol: str, min_abundance: float = 0.01) -> str:
 336    """Primary natural-abundance rare isotope label (highest abundance above the floor)."""
 337    return rare_isotope_entries(mono_symbol, min_abundance=min_abundance)[0][0]
 338
 339
 340def isotope_mass_delta(
 341    mono_symbol: str, charge: int = 1, min_abundance: float = 0.01
 342) -> float:
 343    """
 344    Signed mass difference (primary natural-abundance rare − mono) / |charge|.
 345
 346    Primary rare = highest natural-abundance rare form meeting ``min_abundance``
 347    in ``Atoms``. May be negative (e.g. ⁵⁴Fe − ⁵⁶Fe). Never hard-coded.
 348    """
 349    _rare, signed, _ab = rare_isotope_entries(mono_symbol, min_abundance=min_abundance)[
 350        0
 351    ]
 352    return signed / abs(int(charge))
 353
 354
 355def isotope_state_label(rare_or_mono_symbol: str, n: int) -> str:
 356    """Build isotope_state string, e.g. ``13C1`` or ``54Fe1``."""
 357    if n <= 0:
 358        return "M+0"
 359    return f"{rare_or_mono_symbol}{n}"
 360
 361
 362def empty_group_labels(cluster_ids: Sequence) -> pd.DataFrame:
 363    """Return unlabeled group columns for the given cluster ids."""
 364    idx = pd.Index(cluster_ids, name="cluster")
 365    return pd.DataFrame(
 366        {
 367            "feature_group_id": pd.Series(pd.NA, index=idx, dtype="Int64"),
 368            "ion_role": pd.Series(None, index=idx, dtype=object),
 369            "ion_type": pd.Series(None, index=idx, dtype=object),
 370            "possible_ion_types": pd.Series(None, index=idx, dtype=object),
 371            "isotope_state": pd.Series(None, index=idx, dtype=object),
 372            "mono_cluster_id": pd.Series(pd.NA, index=idx, dtype="Int64"),
 373        }
 374    )
 375
 376
 377def format_possible_ion_types(types: Sequence[str]) -> Optional[str]:
 378    """Join unique ion-type strings alphabetically. Empty → None."""
 379    seen = set()
 380    ordered = []
 381    for t in types:
 382        if not t or t in seen:
 383            continue
 384        seen.add(t)
 385        ordered.append(t)
 386    ordered.sort()
 387    return POSSIBLE_ION_TYPES_SEP.join(ordered) if ordered else None
 388
 389
 390def parse_possible_ion_types(val) -> list:
 391    """Split a ``possible_ion_types`` cell; null/empty → []."""
 392    if val is None:
 393        return []
 394    if isinstance(val, (list, tuple, np.ndarray, pd.Series, pd.Index)):
 395        out = []
 396        for x in val:
 397            out.extend(parse_possible_ion_types(x))
 398        return out
 399    if _isna(val):
 400        return []
 401    return [p for p in str(val).split(POSSIBLE_ION_TYPES_SEP) if p]
 402
 403
 404def merge_possible_ion_types(*vals) -> Optional[str]:
 405    """Union possible-type strings; first occurrence wins."""
 406    ordered = []
 407    for v in vals:
 408        ordered.extend(parse_possible_ion_types(v))
 409    return format_possible_ion_types(ordered)
 410
 411
 412def _isna(val) -> bool:
 413    """True for None / scalar NA. Sequences of identifications are not NA."""
 414    if val is None:
 415        return True
 416    if isinstance(val, (list, tuple, set, dict, np.ndarray, pd.Series, pd.Index)):
 417        return False
 418    try:
 419        result = pd.isna(val)
 420    except (ValueError, TypeError):
 421        return False
 422    if hasattr(result, "ndim") and getattr(result, "ndim", 0) > 0:
 423        return False
 424    return bool(result)
 425
 426
 427def is_adduct_endpoint_eligible(ion_role) -> bool:
 428    """Adduct endpoints: mono or unlabeled; not ``isotope``."""
 429    return _isna(ion_role) or ion_role != "isotope"
 430
 431
 432def constrain_annotation_active(labels, constrain_flag: bool) -> bool:
 433    """True when constrain flag is on and a non-empty labels table exists."""
 434    return bool(constrain_flag) and labels is not None and len(labels) > 0
 435
 436
 437def should_skip_isotope_for_annotation(ion_role, constrain_active: bool) -> bool:
 438    """Skip consensus isotopes when annotation constrain is active."""
 439    return constrain_active and ion_role == "isotope"
 440
 441
 442def allowed_ion_types_for_row(ion_type, possible_ion_types):
 443    """Allow-list for ID ion types, or None if no filter for this row."""
 444    parsed = parse_possible_ion_types(possible_ion_types)
 445    if parsed:
 446        return set(parsed)
 447    parsed_it = parse_possible_ion_types(ion_type)
 448    if parsed_it:
 449        return set(parsed_it)
 450    return None
 451
 452
 453def ion_type_allowed(id_ion_type, allowed) -> bool:
 454    """True if identification ion type is allowed (``allowed is None`` → keep).
 455
 456    ``id_ion_type`` may be a scalar or a sequence (MS2 hits store parallel
 457    lists). A sequence is allowed if any token is in ``allowed``.
 458    """
 459    if allowed is None:
 460        return True
 461    tokens = parse_possible_ion_types(id_ion_type)
 462    if not tokens:
 463        return False
 464    allowed_l = {str(a).lower() for a in allowed}
 465    return any(t.lower() in allowed_l for t in tokens)
 466
 467
 468def subset_hits_by_allowed_ion_types(result, allowed, ion_attr: str = "ref_ion_type"):
 469    """Keep parallel hit arrays on a search-result object whose ion type is allowed.
 470
 471    Returns ``result`` unchanged when ``allowed is None``. Returns ``None`` if
 472    no hits remain. Sequences (FlashEntropy ``ref_ion_type`` lists) are subset
 473    in place on list/tuple/ndarray attributes of matching length.
 474    """
 475    if allowed is None or result is None:
 476        return result
 477    ion_types = getattr(result, ion_attr, None)
 478    if ion_types is None and hasattr(result, "get"):
 479        ion_types = result.get(ion_attr)
 480    if ion_types is None:
 481        return result
 482    if isinstance(ion_types, str) or not hasattr(ion_types, "__len__"):
 483        return result if ion_type_allowed(ion_types, allowed) else None
 484    n = len(ion_types)
 485    if n == 0:
 486        return result
 487    keep = [i for i, it in enumerate(ion_types) if ion_type_allowed(it, allowed)]
 488    if not keep:
 489        return None
 490    if len(keep) == n:
 491        return result
 492    keep_idx = np.asarray(keep, dtype=int)
 493    items = (
 494        result.items()
 495        if hasattr(result, "items") and not hasattr(result, "__dict__")
 496        else vars(result).items()
 497    )
 498    for key, val in list(items):
 499        if key.startswith("_"):
 500            continue
 501        if isinstance(val, np.ndarray) and val.shape[:1] == (n,):
 502            new_val = val[keep_idx]
 503        elif isinstance(val, (list, tuple)) and len(val) == n:
 504            seq = [val[i] for i in keep]
 505            new_val = type(val)(seq) if isinstance(val, tuple) else seq
 506        elif isinstance(val, pd.Series) and len(val) == n:
 507            new_val = val.iloc[keep]
 508        else:
 509            continue
 510        if hasattr(result, "__dict__"):
 511            setattr(result, key, new_val)
 512        else:
 513            result[key] = new_val
 514    return result
 515
 516
 517def merge_feature_group_labels_into_frame(
 518    labels: pd.DataFrame,
 519    df: Optional[pd.DataFrame],
 520) -> Optional[pd.DataFrame]:
 521    """Left-join GROUP_COLUMNS from cluster-indexed ``labels`` onto a feature table."""
 522    if df is None or len(df) == 0 or labels is None or len(labels) == 0:
 523        return df
 524    if "cluster" not in df.columns:
 525        raise ValueError("mass feature table must have a 'cluster' column")
 526
 527    out = df.copy()
 528    index_name = out.index.name or "coll_mf_id"
 529    out = out.reset_index(drop=False)
 530    index_col = index_name if index_name in out.columns else out.columns[0]
 531
 532    drop_cols = [c for c in GROUP_COLUMNS if c in out.columns]
 533    if drop_cols:
 534        out = out.drop(columns=drop_cols)
 535
 536    lab = labels.reset_index()
 537    if lab.columns[0] != "cluster":
 538        lab = lab.rename(columns={lab.columns[0]: "cluster"})
 539    out = out.merge(lab, on="cluster", how="left")
 540    out = out.set_index(index_col)
 541    out.index.name = index_name
 542    return out
 543
 544
 545def merge_feature_group_labels_into_frames(
 546    labels: pd.DataFrame,
 547    mass_features_df: Optional[pd.DataFrame],
 548    induced_df: Optional[pd.DataFrame] = None,
 549):
 550    """Join labels onto mass-feature and induced tables."""
 551    return (
 552        merge_feature_group_labels_into_frame(labels, mass_features_df),
 553        merge_feature_group_labels_into_frame(labels, induced_df),
 554    )
 555
 556
 557def annotation_meta_for_sample_mf(
 558    mass_features_df: pd.DataFrame,
 559    labels: Optional[pd.DataFrame],
 560    sample_id,
 561    mf_id,
 562) -> Optional[dict]:
 563    """Group-label fields for ``(sample_id, mf_id)``, or None if ungrouped / no labels."""
 564    if labels is None or len(labels) == 0:
 565        return None
 566    hit = mass_features_df[
 567        (mass_features_df["sample_id"] == sample_id)
 568        & (mass_features_df["mf_id"] == mf_id)
 569    ]
 570    if hit.empty:
 571        return None
 572    cluster = hit.iloc[0]["cluster"]
 573    if _isna(cluster):
 574        return None
 575    cid = int(cluster)
 576    row = labels.loc[cid]
 577    if _isna(row["feature_group_id"]):
 578        return None
 579    return {
 580        "cluster": cid,
 581        "feature_group_id": row["feature_group_id"],
 582        "ion_role": row["ion_role"],
 583        "ion_type": row["ion_type"],
 584        "possible_ion_types": row["possible_ion_types"],
 585        "isotope_state": row["isotope_state"],
 586        "mono_cluster_id": row["mono_cluster_id"],
 587    }
 588
 589
 590def _mz_tol_abs(mz_a: float, mz_b: float, mz_tol_ppm: float) -> float:
 591    return max(mz_a, mz_b) * mz_tol_ppm * 1e-6
 592
 593
 594def _get_ion_type_dict() -> Dict:
 595    """``ION_TYPE_DICT`` from encapsulation.constant."""
 596    return ION_TYPE_DICT
 597
 598
 599def _atom_count_mass(atom_counts: Dict[str, int]) -> float:
 600    """Exact mass for an atom-count dict using ``Atoms.atomic_masses``."""
 601    total = 0.0
 602    for symbol, n in atom_counts.items():
 603        if n == 0:
 604            continue
 605        if symbol not in Atoms.atomic_masses:
 606            raise ValueError(
 607                f"Unknown atom {symbol!r} in ion_type_dict entry; "
 608                "must be a key in Atoms.atomic_masses"
 609            )
 610        total += float(Atoms.atomic_masses[symbol]) * int(n)
 611    return total
 612
 613
 614def _ion_type_mass_offset(ion_type: str) -> float:
 615    """
 616    Neutral-formula mass offset for an ion_type_dict key.
 617
 618    offset = mass(atoms to add) − mass(atoms to subtract).
 619    Same convention as ``LCMSMetabolomicsExport.get_ion_formula``.
 620    """
 621    ion_type_dict = _get_ion_type_dict()
 622    if ion_type not in ion_type_dict:
 623        raise ValueError(
 624            f"Unknown ion_type {ion_type!r} for feature grouping; "
 625            f"must be a key in ion_type_dict "
 626            f"(e.g. '[M+H]+', '[M+NH4]+'). "
 627            f"Known: {sorted(ion_type_dict.keys())}"
 628        )
 629    entry = ion_type_dict[ion_type]
 630    return _atom_count_mass(entry["add"]) - _atom_count_mass(entry["sub"])
 631
 632
 633def ion_type_mass_delta(
 634    ion_type_a: str,
 635    ion_type_b: str,
 636    charge: int = 1,
 637) -> float:
 638    """
 639    m/z spacing for **same** absolute charge: ion_type_b − ion_type_a.
 640
 641    Example: ``[M+NH4]+`` − ``[M+H]+`` = m(N) + 3·m(H) at |z|=1.
 642    Sign indicates which form is heavier at that charge; no preferred “base”.
 643
 644    For pairs with **different** charges (e.g. ``[M+H]+`` vs ``[M+2H]2+``),
 645    m/z spacing depends on neutral mass ``M`` and is **not** a constant;
 646    use :func:`neutral_mass_from_mz` / adduct edge search instead.
 647    """
 648    signed = _ion_type_mass_offset(ion_type_b) - _ion_type_mass_offset(ion_type_a)
 649    return float(signed) / abs(int(charge))
 650
 651
 652# Backward-compatible alias
 653adduct_mass_delta = ion_type_mass_delta
 654
 655
 656def neutral_mass_from_mz(mz: float, ion_type: str) -> float:
 657    """Neutral mass implied by observed m/z and an ion-type assignment.
 658
 659    ``M = |z| * m/z − offset(ion_type)`` with ``offset`` from
 660    ``ion_type_dict`` atom add/subtract and ``|z|`` from the key suffix
 661    (e.g. ``2+`` → 2).
 662    """
 663    z = ion_type_charge(ion_type)
 664    offset = _ion_type_mass_offset(ion_type)
 665    return float(z) * float(mz) - float(offset)
 666
 667
 668def find_isotope_edges(
 669    cluster_ids: np.ndarray,
 670    mz: np.ndarray,
 671    rt: np.ndarray,
 672    params: FeatureGroupParams,
 673) -> pd.DataFrame:
 674    """
 675    Find unit-step mono→natural-abundance isotope edges via KDTree pair matrices.
 676
 677    Stage 1 only: natural-abundance rare isotopes from ``Atoms`` (not tracer
 678    enrichment). Pattern matches ``LCMSBase.find_c13_mass_features``:
 679
 680    1. Sort features by ascending m/z so ``triu`` means light → heavy.
 681    2. Sparse RT pairs within ``params.rt_tol``.
 682    3. Sparse m/z pairs within the largest unit isotope spacing (+ ppm tol).
 683    4. Keep the intersection; retain only pairs whose Δm matches a **unit**
 684       natural-abundance Atoms spacing (one rare substitution, not 2×, 3×, …).
 685
 686    Higher-order natural isotopologues (¹³C₂, …) are **not** linked as mono→M+n
 687    here. They appear later by roll-up along successive unit edges
 688    (mono→¹³C₁→¹³C₂), so you never get M+n without M+(n−1).
 689
 690    Parent is the chemical mono (most-abundant natural form) side of the unit
 691    step (lighter for ¹³C, heavier for ⁵⁴Fe). Unit spacing is the raw Atoms
 692    mass difference (singly charged). Multi-charge envelopes are out of scope.
 693
 694    Returns
 695    -------
 696    DataFrame
 697        Unit edges only. Columns: parent_idx, child_idx, parent_cluster,
 698        child_cluster, n (=1), atom, rare_label, charge, abs_dm.
 699    """
 700    edge_columns = [
 701        "parent_idx",
 702        "child_idx",
 703        "parent_cluster",
 704        "child_cluster",
 705        "n",
 706        "atom",
 707        "rare_label",
 708        "charge",
 709        "abs_dm",
 710    ]
 711    n_feat = len(cluster_ids)
 712    if n_feat < 2:
 713        return pd.DataFrame(columns=edge_columns)
 714
 715    # ------------------------------------------------------------------
 716    # 1) Unit isotope spacings only (signed). Higher n comes from roll-up.
 717    #    signed_unit = m_rare - m_mono  (|z| = 1 only)
 718    # ------------------------------------------------------------------
 719    spacings = []
 720    for atom in params.isotope_atoms:
 721        for rare_label, signed_delta, _abun in rare_isotope_entries(
 722            atom, min_abundance=params.min_isotope_abundance
 723        ):
 724            signed_unit = float(signed_delta)
 725            if signed_unit == 0:
 726                continue
 727            spacings.append((atom, rare_label, 1, signed_unit))
 728
 729    if not spacings:
 730        return pd.DataFrame(columns=edge_columns)
 731
 732    max_unit = max(abs(s[3]) for s in spacings)
 733
 734    # ------------------------------------------------------------------
 735    # 2) Sort by m/z so sparse triu is always light → heavy (same as C13).
 736    # ------------------------------------------------------------------
 737    order = np.argsort(mz, kind="mergesort")
 738    mz_s = np.asarray(mz, dtype=float)[order]
 739    rt_s = np.asarray(rt, dtype=float)[order]
 740    # Map sorted row/col back to original array indices
 741    # order[k] = original index of sorted position k
 742
 743    # Absolute ppm window used to bound the m/z KD query (generous).
 744    mz_tol_pad = float(np.max(mz_s)) * params.mz_tol_ppm * 1e-6
 745    max_mz_gap = max_unit + mz_tol_pad
 746
 747    # ------------------------------------------------------------------
 748    # 3) Sparse pair matrices (KDTree.sparse_distance_matrix + triu).
 749    # ------------------------------------------------------------------
 750    # RT: pairs with |ΔRT| <= rt_tol
 751    tree_rt = KDTree(rt_s.reshape(-1, 1))
 752    sdm_rt = tree_rt.sparse_distance_matrix(
 753        tree_rt, params.rt_tol, output_type="coo_matrix"
 754    )
 755    sdm_rt = sparse.triu(sdm_rt, k=1)
 756    sdm_rt.data = np.ones_like(sdm_rt.data)
 757
 758    # m/z: pairs with 0 < Δm/z <= max unit spacing (+ pad)
 759    tree_mz = KDTree(mz_s.reshape(-1, 1))
 760    sdm_mz = tree_mz.sparse_distance_matrix(
 761        tree_mz, max_mz_gap, output_type="coo_matrix"
 762    )
 763    sdm_mz = sparse.triu(sdm_mz, k=1)
 764
 765    # Intersection: keep m/z distances only where RT also coelutes
 766    cand = sdm_mz.multiply(sdm_rt).tocoo()
 767    if cand.nnz == 0:
 768        return pd.DataFrame(columns=edge_columns)
 769
 770    # ------------------------------------------------------------------
 771    # 4) Keep pairs whose Δm matches a unit isotope spacing (Atoms).
 772    # ------------------------------------------------------------------
 773    rows = []
 774    # Track (parent_orig, child_orig, atom, rare) to avoid duplicate edges
 775    seen = set()
 776
 777    for r, c, dm in zip(cand.row, cand.col, cand.data):
 778        # r, c are indices into sorted arrays; r is lighter, c is heavier
 779        i_light = int(order[r])
 780        i_heavy = int(order[c])
 781        dm = float(dm)
 782        tol = _mz_tol_abs(mz_s[r], mz_s[c], params.mz_tol_ppm)
 783
 784        # Collect all unit spacings that fit; pick smallest residual so near-isobar
 785        # shifts (e.g. |Δ⁷⁸Se| ≈ |Δ⁸²Se|) disambiguate by exact Atoms mass.
 786        best = None  # (residual, atom, rare_label, z, signed_unit, parent, child)
 787        for atom, rare_label, z, signed_unit in spacings:
 788            unit = abs(signed_unit)
 789            residual = abs(dm - unit)
 790            if residual > tol:
 791                continue
 792            if signed_unit > 0:
 793                parent_idx, child_idx = i_light, i_heavy
 794            else:
 795                parent_idx, child_idx = i_heavy, i_light
 796            cand_t = (residual, atom, rare_label, z, signed_unit, parent_idx, child_idx)
 797            if best is None or residual < best[0]:
 798                best = cand_t
 799
 800        if best is None:
 801            continue
 802
 803        residual, atom, rare_label, z, signed_unit, parent_idx, child_idx = best
 804        key = (parent_idx, child_idx, atom, rare_label)
 805        if key in seen:
 806            continue
 807        seen.add(key)
 808
 809        rows.append(
 810            {
 811                "parent_idx": parent_idx,
 812                "child_idx": child_idx,
 813                "parent_cluster": cluster_ids[parent_idx],
 814                "child_cluster": cluster_ids[child_idx],
 815                "n": 1,  # unit step; depth assigned at roll-up
 816                "atom": atom,
 817                "rare_label": rare_label,
 818                "charge": z,
 819                "abs_dm": dm,
 820            }
 821        )
 822
 823    if not rows:
 824        return pd.DataFrame(columns=edge_columns)
 825    return pd.DataFrame(rows)
 826
 827
 828_ADDUCT_EDGE_COLUMNS = [
 829    "parent_idx",
 830    "child_idx",
 831    "parent_cluster",
 832    "child_cluster",
 833    "parent_ion_type",
 834    "child_ion_type",
 835    "parent_possible_ion_types",
 836    "child_possible_ion_types",
 837    "charge",
 838    "abs_dm",
 839    "residual",
 840    "n_interpretations",
 841]
 842
 843# Residual equality for “tied” ion-type interpretations (Da of neutral mass).
 844_RESIDUAL_TIE_EPS = 1e-6
 845
 846
 847def find_adduct_edges(
 848    cluster_ids: np.ndarray,
 849    mz: np.ndarray,
 850    rt: np.ndarray,
 851    params: FeatureGroupParams,
 852) -> pd.DataFrame:
 853    """
 854    Find ion-type edges using pairwise neutral-mass consistency.
 855
 856    ``M = |z| · m/z − offset(ion_type)`` must agree within ppm for the two
 857    assigned types on a coeluting peak pair. Configured types must be
 858    singly charged (``|z| = 1``).
 859
 860    When several type assignments fit with the same residual (within
 861    ``_RESIDUAL_TIE_EPS``), all are kept in ``possible_ion_types``.
 862    ``parent_ion_type`` / ``child_ion_type`` on the edge row are a
 863    deterministic handle for shared-M geometry only; labels do not treat
 864    that pick as a preferred assignment when ``n_interpretations`` > 1.
 865
 866    Returns
 867    -------
 868    DataFrame
 869        One row per peak pair. Parent = lower m/z of the pair.
 870    """
 871    edge_columns = list(_ADDUCT_EDGE_COLUMNS)
 872    n_feat = len(cluster_ids)
 873    ion_types = tuple(params.ion_types)
 874    if n_feat < 2 or len(ion_types) < 2:
 875        return pd.DataFrame(columns=edge_columns)
 876
 877    typed: list[Tuple[str, int, float]] = []
 878    for t in ion_types:
 879        typed.append((t, ion_type_charge(t), float(_ion_type_mass_offset(t))))
 880
 881    mz = np.asarray(mz, dtype=float)
 882    rt = np.asarray(rt, dtype=float)
 883    order = np.argsort(mz, kind="mergesort")
 884    rt_s = rt[order]
 885
 886    tree_rt = KDTree(rt_s.reshape(-1, 1))
 887    sdm_rt = tree_rt.sparse_distance_matrix(
 888        tree_rt, params.rt_tol, output_type="coo_matrix"
 889    )
 890    sdm_rt = sparse.triu(sdm_rt, k=1)
 891    if sdm_rt.nnz == 0:
 892        return pd.DataFrame(columns=edge_columns)
 893
 894    rows = []
 895    seen_pairs = set()
 896    ppm = float(params.mz_tol_ppm)
 897
 898    for r, c in zip(sdm_rt.row, sdm_rt.col):
 899        i_lo = int(order[r])
 900        i_hi = int(order[c])
 901        if (i_lo, i_hi) in seen_pairs:
 902            continue
 903        mz_lo = float(mz[i_lo])
 904        mz_hi = float(mz[i_hi])
 905        dm = abs(mz_hi - mz_lo)
 906
 907        cands = []
 908        for t_lo, z_lo, off_lo in typed:
 909            M_lo = z_lo * mz_lo - off_lo
 910            if M_lo <= 0:
 911                continue
 912            for t_hi, z_hi, off_hi in typed:
 913                if not is_allowed_adduct_type_pair(t_lo, t_hi):
 914                    continue
 915                M_hi = z_hi * mz_hi - off_hi
 916                if M_hi <= 0:
 917                    continue
 918                residual = abs(M_lo - M_hi)
 919                tol = max(M_lo, M_hi) * ppm * 1e-6
 920                if residual > tol:
 921                    continue
 922                cands.append((residual, t_lo, t_hi, z_lo, z_hi))
 923
 924        if not cands:
 925            continue
 926
 927        cands.sort()  # residual, then alphabetical t_lo, t_hi
 928        best_res = cands[0][0]
 929        ties = [c for c in cands if abs(c[0] - best_res) <= _RESIDUAL_TIE_EPS]
 930        residual, t_lo, t_hi, z_lo, z_hi = ties[0]
 931        parent_poss = format_possible_ion_types([c[1] for c in ties])
 932        child_poss = format_possible_ion_types([c[2] for c in ties])
 933
 934        seen_pairs.add((i_lo, i_hi))
 935        rows.append(
 936            {
 937                "parent_idx": i_lo,
 938                "child_idx": i_hi,
 939                "parent_cluster": cluster_ids[i_lo],
 940                "child_cluster": cluster_ids[i_hi],
 941                "parent_ion_type": t_lo,
 942                "child_ion_type": t_hi,
 943                "parent_possible_ion_types": parent_poss,
 944                "child_possible_ion_types": child_poss,
 945                "charge": int(max(z_lo, z_hi)),
 946                "abs_dm": dm,
 947                "residual": float(residual),
 948                "n_interpretations": int(len(ties)),
 949            }
 950        )
 951
 952    if not rows:
 953        return pd.DataFrame(columns=edge_columns)
 954    return pd.DataFrame(rows)
 955
 956
 957def sort_adduct_edges(edges: pd.DataFrame) -> pd.DataFrame:
 958    """Order edges for merge: residual ↑, then cluster ids."""
 959    if edges is None or edges.empty:
 960        return (
 961            edges if edges is not None else pd.DataFrame(columns=_ADDUCT_EDGE_COLUMNS)
 962        )
 963    out = edges.copy()
 964    if "residual" not in out.columns:
 965        out["residual"] = 0.0
 966    return out.sort_values(
 967        by=["residual", "parent_cluster", "child_cluster"],
 968        kind="mergesort",
 969    ).reset_index(drop=True)
 970
 971
 972def filter_edges_by_height_correlation(
 973    edges: pd.DataFrame,
 974    heights: np.ndarray,
 975    params: FeatureGroupParams,
 976) -> pd.DataFrame:
 977    """
 978    Keep edges with enough shared non-zero samples and Pearson r >= threshold.
 979
 980    Correlation method is fixed to Pearson (``CORR_METHOD``). Heights must be
 981    apex peak intensities (``HEIGHT_COL``), not integrated areas.
 982
 983    Pearson is computed only on samples where **both** heights are > 0
 984    (pairwise-complete; zeros / missing do not enter the correlation).
 985
 986    Parameters
 987    ----------
 988    heights : ndarray, shape (n_features, n_samples)
 989        Apex intensity matrix (cluster × sample).
 990    """
 991    if edges.empty:
 992        return edges
 993
 994    n_samples = int(heights.shape[1]) if heights.ndim == 2 else 1
 995    min_shared = params.min_shared_count(n_samples)
 996
 997    keep = []
 998    for row in edges.itertuples(index=False):
 999        hi = heights[row.parent_idx]
1000        hj = heights[row.child_idx]
1001        # Correlate only where both features were detected (drop zeros / missing)
1002        shared = (hi > 0) & (hj > 0)
1003        n_shared = int(np.count_nonzero(shared))
1004        if n_shared < min_shared:
1005            keep.append(False)
1006            continue
1007        hi_s = hi[shared]
1008        hj_s = hj[shared]
1009        if np.std(hi_s) == 0 or np.std(hj_s) == 0:
1010            keep.append(False)
1011            continue
1012        r, _ = pearsonr(hi_s, hj_s)
1013        if np.isnan(r) or r < params.corr_threshold:
1014            keep.append(False)
1015            continue
1016        keep.append(True)
1017
1018    return edges.loc[np.asarray(keep)].reset_index(drop=True)
1019
1020
1021def assign_isotope_labels(
1022    cluster_ids: np.ndarray,
1023    mz: np.ndarray,
1024    rt: np.ndarray,
1025    heights: np.ndarray,
1026    edges: pd.DataFrame,
1027    params: FeatureGroupParams,
1028) -> pd.DataFrame:
1029    """
1030    Assign feature groups by roll-up along unit natural-abundance isotope edges.
1031
1032    Same idea as ``find_c13_mass_features`` for natural-abundance envelopes:
1033
1034    - Roots = features that appear as edge parents but never as children
1035      (chemical monoisotopes of a family).
1036    - Wave 1: unit children of roots → rare¹ (e.g. natural ¹³C₁).
1037    - Wave 2: unit children of wave 1 → rare² (e.g. ¹³C₂), etc.
1038
1039    Higher-order labels only appear if intermediate steps exist, so there is
1040    no ¹³C₂ without ¹³C₁. Depth is capped by ``max_isotope_offset``.
1041
1042    Chemical mono (``ion_role="mono"`` / ``M+0``) is the roll-up root from
1043    geometry (Atoms side of each unit step), not necessarily the tallest peak
1044    in the envelope. No mono-vs-family height prior is applied; geometry +
1045    correlation are the gates. Not designed for labeled/enriched isotope series.
1046
1047    ``heights`` is accepted for API compatibility with the correlation stage but
1048    is not used during roll-up labeling.
1049    """
1050    labels = empty_group_labels(cluster_ids)
1051    if edges.empty:
1052        return labels
1053
1054    # Adjacency: parent_idx -> list of (child_idx, rare_label, atom, z)
1055    children_of = {}
1056    all_parents = set()
1057    all_children = set()
1058    for row in edges.itertuples(index=False):
1059        p, c = int(row.parent_idx), int(row.child_idx)
1060        children_of.setdefault(p, []).append(
1061            (c, row.rare_label, row.atom, int(row.charge))
1062        )
1063        all_parents.add(p)
1064        all_children.add(c)
1065
1066    # Roots: monoisotopic candidates (never a unit-step child)
1067    roots = sorted(all_parents - all_children)
1068    if not roots:
1069        return labels
1070
1071    assigned = set()  # feature indices already labeled into some family
1072    next_group_id = 0
1073
1074    for mono_idx in roots:
1075        if mono_idx in assigned:
1076            continue
1077
1078        # Roll-up waves from this root (depth = isotope count for that rare chain)
1079        # child_idx -> (depth, rare_label, mono_atom, z)
1080        family_children = {}
1081        # frontier: list of (node_idx, depth, rare_label used to reach node)
1082        frontier = [(mono_idx, 0, None)]
1083        family_nodes = {mono_idx}
1084
1085        while frontier:
1086            next_frontier = []
1087            for node, depth, _rare_in in frontier:
1088                if depth >= params.max_isotope_offset:
1089                    continue
1090                for c, rare_label, atom, z in children_of.get(node, []):
1091                    if c in family_nodes or c in assigned:
1092                        continue
1093                    new_depth = depth + 1
1094                    if new_depth > params.max_isotope_offset:
1095                        continue
1096                    # Prefer closer RT if c already staged with worse RT (rare)
1097                    if c in family_children:
1098                        continue
1099                    family_children[c] = (new_depth, rare_label, atom, z)
1100                    family_nodes.add(c)
1101                    next_frontier.append((c, new_depth, rare_label))
1102            frontier = next_frontier
1103
1104        if not family_children:
1105            continue
1106
1107        mono_cluster = int(cluster_ids[mono_idx])
1108        gid = next_group_id
1109        next_group_id += 1
1110
1111        # ion_type filled later from pairwise adduct edges (or single configured type)
1112        single_ion_type = params.ion_types[0] if len(params.ion_types) == 1 else None
1113        labels.loc[cluster_ids[mono_idx], "feature_group_id"] = gid
1114        labels.loc[cluster_ids[mono_idx], "ion_role"] = "mono"
1115        labels.loc[cluster_ids[mono_idx], "ion_type"] = single_ion_type
1116        labels.loc[cluster_ids[mono_idx], "isotope_state"] = "M+0"
1117        labels.loc[cluster_ids[mono_idx], "mono_cluster_id"] = mono_cluster
1118        assigned.add(mono_idx)
1119
1120        for c, (depth, rare_label, _atom, _z) in family_children.items():
1121            cid = cluster_ids[c]
1122            labels.loc[cid, "feature_group_id"] = gid
1123            labels.loc[cid, "ion_role"] = "isotope"
1124            labels.loc[cid, "ion_type"] = single_ion_type
1125            labels.loc[cid, "isotope_state"] = isotope_state_label(rare_label, depth)
1126            labels.loc[cid, "mono_cluster_id"] = mono_cluster
1127            assigned.add(c)
1128
1129    return labels
1130
1131
1132def _remap_group_id(labels: pd.DataFrame, old_gid, new_gid) -> None:
1133    """In-place remap feature_group_id old → new."""
1134    if pd.isna(old_gid) or pd.isna(new_gid) or old_gid == new_gid:
1135        return
1136    mask = labels["feature_group_id"] == old_gid
1137    labels.loc[mask, "feature_group_id"] = new_gid
1138
1139
1140def _form_members_for_endpoint(
1141    labels: pd.DataFrame,
1142    endpoint: int,
1143    target_ion_type: str,
1144) -> list:
1145    """
1146    Isotope form subtree for an adduct-edge endpoint.
1147
1148    Members share the endpoint's chemical ``mono_cluster_id`` when an isotope
1149    form exists, otherwise the endpoint itself. ``ion_type`` is null or equal
1150    to ``target_ion_type``. Does **not** include other ion forms already in
1151    the same feature_group_id.
1152    """
1153    if endpoint not in labels.index:
1154        return [endpoint]
1155
1156    mono_id = labels.loc[endpoint, "mono_cluster_id"]
1157    if pd.isna(mono_id):
1158        mono_id = endpoint
1159    else:
1160        mono_id = int(mono_id)
1161
1162    members = []
1163    for cid in labels.index:
1164        cid_i = int(cid)
1165        mid = labels.loc[cid, "mono_cluster_id"]
1166        if pd.isna(mid):
1167            same_mono = cid_i == endpoint or cid_i == mono_id
1168        else:
1169            same_mono = int(mid) == mono_id or cid_i == mono_id
1170        if not same_mono:
1171            continue
1172        it = labels.loc[cid, "ion_type"]
1173        if _isna(it) or it == target_ion_type:
1174            members.append(cid_i)
1175    if endpoint not in members:
1176        members.append(int(endpoint))
1177    return members
1178
1179
1180def _form_has_isotope(labels: pd.DataFrame, form_members: Sequence[int]) -> bool:
1181    """True if any form member is labeled ``ion_role=isotope``."""
1182    for cid in form_members:
1183        if cid in labels.index and labels.loc[cid, "ion_role"] == "isotope":
1184            return True
1185    return False
1186
1187
1188def _form_mono_id(
1189    labels: pd.DataFrame, form_members: Sequence[int], endpoint: int
1190) -> int:
1191    monos = [
1192        c
1193        for c in form_members
1194        if c in labels.index and labels.loc[c, "ion_role"] == "mono"
1195    ]
1196    if monos:
1197        return int(monos[0])
1198    if endpoint in form_members:
1199        return int(endpoint)
1200    return int(form_members[0])
1201
1202
1203def _type_compatible(existing, preferred: str, possible_str) -> bool:
1204    """True if cluster can accept preferred type (null, same, or in possible)."""
1205    if _isna(existing):
1206        return True
1207    if existing == preferred:
1208        return True
1209    return str(existing) in parse_possible_ion_types(possible_str)
1210
1211
1212def merge_adduct_edges_into_labels(
1213    labels: pd.DataFrame,
1214    cluster_ids: np.ndarray,
1215    adduct_edges: pd.DataFrame,
1216    params: FeatureGroupParams,
1217    mz_by_cluster: Optional[Dict[int, float]] = None,
1218) -> pd.DataFrame:
1219    """
1220    Merge forms across ion-type edges under correctness constraints.
1221
1222    - Paint **form subtrees only** (isotope family of that form, or the
1223      unlabeled endpoint when no isotope was observed).
1224    - ``mono_cluster_id`` / ``ion_role=mono`` only when the form has an isotope.
1225    - Require shared neutral mass between form reference clusters
1226      (if ``mz_by_cluster``).
1227    - Reject if an endpoint already has a conflicting non-null ``ion_type``
1228      not listed among that edge's possible types.
1229    - Store residual-tied alternate types in ``possible_ion_types``.
1230    - Edges should be pre-sorted (see :func:`sort_adduct_edges`).
1231    """
1232    if adduct_edges is None or adduct_edges.empty:
1233        return labels
1234
1235    labels = labels.copy()
1236    if "possible_ion_types" not in labels.columns:
1237        labels["possible_ion_types"] = pd.Series(None, index=labels.index, dtype=object)
1238
1239    next_gid = 0
1240    if labels["feature_group_id"].notna().any():
1241        next_gid = int(labels["feature_group_id"].max()) + 1
1242
1243    ppm = float(params.mz_tol_ppm)
1244    if mz_by_cluster is None:
1245        mz_by_cluster = {}
1246
1247    def _M(cid: int, ion_type: str) -> Optional[float]:
1248        if cid not in mz_by_cluster:
1249            return None
1250        try:
1251            return neutral_mass_from_mz(mz_by_cluster[cid], ion_type)
1252        except Exception:
1253            return None
1254
1255    for row in adduct_edges.itertuples(index=False):
1256        parent_c = int(row.parent_cluster)
1257        child_c = int(row.child_cluster)
1258        type_light = row.parent_ion_type
1259        type_heavy = row.child_ion_type
1260        poss_light = getattr(row, "parent_possible_ion_types", None) or type_light
1261        poss_heavy = getattr(row, "child_possible_ion_types", None) or type_heavy
1262
1263        if parent_c not in labels.index or child_c not in labels.index:
1264            continue
1265
1266        it_p = labels.loc[parent_c, "ion_type"]
1267        it_c = labels.loc[child_c, "ion_type"]
1268        if not _type_compatible(it_p, type_light, poss_light):
1269            continue
1270        if not _type_compatible(it_c, type_heavy, poss_heavy):
1271            continue
1272
1273        # Prefer already-assigned primary type when compatible with possibles
1274        paint_light = type_light if _isna(it_p) else it_p
1275        paint_heavy = type_heavy if _isna(it_c) else it_c
1276
1277        light_form = _form_members_for_endpoint(labels, parent_c, paint_light)
1278        heavy_form = _form_members_for_endpoint(labels, child_c, paint_heavy)
1279
1280        light_mono = _form_mono_id(labels, light_form, parent_c)
1281        heavy_mono = _form_mono_id(labels, heavy_form, child_c)
1282
1283        # Shared-M check *before* any label writes (avoids orphan mono rows)
1284        M_l = _M(light_mono, paint_light)
1285        M_h = _M(heavy_mono, paint_heavy)
1286        if M_l is not None and M_h is not None:
1287            if M_l <= 0 or M_h <= 0:
1288                continue
1289            tol = max(M_l, M_h) * ppm * 1e-6
1290            if abs(M_l - M_h) > tol:
1291                continue
1292
1293        g_p = labels.loc[parent_c, "feature_group_id"]
1294        g_c = labels.loc[child_c, "feature_group_id"]
1295
1296        if pd.isna(g_p) and pd.isna(g_c):
1297            keep_gid = next_gid
1298            next_gid += 1
1299        elif pd.isna(g_p):
1300            keep_gid = int(g_c)
1301        elif pd.isna(g_c):
1302            keep_gid = int(g_p)
1303        else:
1304            keep_gid = int(g_p)
1305            if int(g_c) != keep_gid:
1306                _remap_group_id(labels, int(g_c), keep_gid)
1307
1308        # Do not claim ion_role=mono / mono_cluster_id without an isotope.
1309        # Adduct-only members still get feature_group_id and ion_type.
1310
1311        def _paint_form(form_members, preferred, poss_str, mono_id):
1312            set_mono_id = _form_has_isotope(labels, form_members)
1313            poss_n = len(set(parse_possible_ion_types(poss_str) or [preferred]))
1314            unique_type = poss_n <= 1
1315            for cid in form_members:
1316                if cid not in labels.index:
1317                    continue
1318                it = labels.loc[cid, "ion_type"]
1319                if not _isna(it) and it != preferred:
1320                    if str(it) not in parse_possible_ion_types(poss_str):
1321                        continue
1322                    labels.loc[cid, "feature_group_id"] = keep_gid
1323                    labels.loc[cid, "possible_ion_types"] = merge_possible_ion_types(
1324                        labels.loc[cid, "possible_ion_types"],
1325                        preferred,
1326                        poss_str,
1327                        it,
1328                    )
1329                    if set_mono_id:
1330                        labels.loc[cid, "mono_cluster_id"] = mono_id
1331                    continue
1332                labels.loc[cid, "feature_group_id"] = keep_gid
1333                labels.loc[cid, "possible_ion_types"] = merge_possible_ion_types(
1334                    labels.loc[cid, "possible_ion_types"],
1335                    preferred,
1336                    poss_str,
1337                    it,
1338                )
1339                if unique_type:
1340                    labels.loc[cid, "ion_type"] = preferred
1341                if set_mono_id:
1342                    labels.loc[cid, "mono_cluster_id"] = mono_id
1343
1344        _paint_form(light_form, paint_light, poss_light, light_mono)
1345        _paint_form(heavy_form, paint_heavy, poss_heavy, heavy_mono)
1346
1347    return labels
1348
1349
1350def validate_feature_group_labels(
1351    labels: pd.DataFrame,
1352    mz_by_cluster: Dict[int, float],
1353    params: FeatureGroupParams,
1354) -> pd.DataFrame:
1355    """
1356    Enforce group invariants; unlabel violator clusters.
1357
1358    - At most one mono per ion_type within a group.
1359    - Form monos share neutral mass within ppm.
1360    - Unique (ion_type, isotope_state) pairs within a group.
1361    """
1362    if labels is None or labels.empty:
1363        return labels
1364    labels = labels.copy()
1365    ppm = float(params.mz_tol_ppm)
1366
1367    grouped = labels.dropna(subset=["feature_group_id"])
1368    for gid, sub in grouped.groupby("feature_group_id"):
1369        to_clear = set()
1370
1371        # Unique (ion_type, isotope_state)
1372        if sub["ion_type"].notna().any():
1373            key = sub.apply(lambda r: (r["ion_type"], r["isotope_state"]), axis=1)
1374            for k, cnt in key.value_counts().items():
1375                if (
1376                    cnt > 1
1377                    and k[0] is not None
1378                    and not (isinstance(k[0], float) and np.isnan(k[0]))
1379                ):
1380                    dups = sub.index[key == k].tolist()
1381                    # keep lowest cluster id
1382                    for cid in sorted(dups)[1:]:
1383                        to_clear.add(int(cid))
1384
1385        # One mono per ion_type
1386        for ion_type, form_sub in sub.groupby(sub["ion_type"], dropna=True):
1387            monos = form_sub.index[form_sub["ion_role"] == "mono"].tolist()
1388            if len(monos) > 1:
1389                # keep mono with highest intensity proxy: lowest m/z as mono of form
1390                # deterministic: keep smallest cluster id
1391                keep = min(int(c) for c in monos)
1392                for cid in monos:
1393                    if int(cid) != keep:
1394                        to_clear.add(int(cid))
1395                monos = [keep]
1396
1397            if len(monos) != 1:
1398                continue
1399            mono_c = int(monos[0])
1400            # isotopes must point at this mono
1401            for cid, row in form_sub.iterrows():
1402                if row["ion_role"] == "isotope":
1403                    mid = row["mono_cluster_id"]
1404                    if pd.isna(mid) or int(mid) != mono_c:
1405                        to_clear.add(int(cid))
1406
1407        # Shared M among form monos still in group
1408        form_Ms = []
1409        for ion_type, form_sub in sub.groupby(sub["ion_type"], dropna=True):
1410            monos = [
1411                int(c)
1412                for c in form_sub.index[form_sub["ion_role"] == "mono"].tolist()
1413                if int(c) not in to_clear
1414            ]
1415            if not monos:
1416                continue
1417            mono_c = monos[0]
1418            if mono_c not in mz_by_cluster:
1419                continue
1420            try:
1421                M = neutral_mass_from_mz(mz_by_cluster[mono_c], str(ion_type))
1422            except Exception:
1423                continue
1424            form_Ms.append((str(ion_type), mono_c, M))
1425
1426        if len(form_Ms) >= 2:
1427            ref_M = form_Ms[0][2]
1428            for ion_type, mono_c, M in form_Ms[1:]:
1429                tol = max(abs(ref_M), abs(M), 1.0) * ppm * 1e-6
1430                if abs(M - ref_M) > tol:
1431                    form_mask = sub["ion_type"] == ion_type
1432                    for cid in sub.index[form_mask]:
1433                        to_clear.add(int(cid))
1434
1435        for cid in to_clear:
1436            if cid not in labels.index:
1437                continue
1438            labels.loc[cid, "feature_group_id"] = pd.NA
1439            labels.loc[cid, "ion_role"] = None
1440            labels.loc[cid, "ion_type"] = None
1441            if "possible_ion_types" in labels.columns:
1442                labels.loc[cid, "possible_ion_types"] = None
1443            labels.loc[cid, "isotope_state"] = None
1444            labels.loc[cid, "mono_cluster_id"] = pd.NA
1445
1446    # Clear orphan role/state rows (e.g. historical seed-before-reject bugs)
1447    orphan = labels["ion_role"].notna() & labels["feature_group_id"].isna()
1448    for cid in labels.index[orphan]:
1449        labels.loc[cid, "ion_role"] = None
1450        labels.loc[cid, "ion_type"] = None
1451        if "possible_ion_types" in labels.columns:
1452            labels.loc[cid, "possible_ion_types"] = None
1453        labels.loc[cid, "isotope_state"] = None
1454        labels.loc[cid, "mono_cluster_id"] = pd.NA
1455
1456    # Ensure possible_ion_types always lists at least preferred ion_type
1457    if "possible_ion_types" in labels.columns and "ion_type" in labels.columns:
1458        for cid in labels.index:
1459            it = labels.loc[cid, "ion_type"]
1460            if _isna(it):
1461                continue
1462            poss = labels.loc[cid, "possible_ion_types"]
1463            labels.loc[cid, "possible_ion_types"] = merge_possible_ion_types(it, poss)
1464
1465    return labels
1466
1467
1468def group_features_arrays(
1469    cluster_ids: np.ndarray,
1470    mz: np.ndarray,
1471    rt: np.ndarray,
1472    heights: np.ndarray,
1473    params: FeatureGroupParams,
1474    timings_out: Optional[MutableMapping[str, float]] = None,
1475) -> pd.DataFrame:
1476    """
1477    Run natural-abundance isotope + adduct feature grouping on array inputs.
1478
1479    1. Unit isotope edges (``Atoms``) + Pearson gate + roll-up.
1480    2. Base↔adduct edges (``ion_type_dict`` mass offsets) + Pearson gate.
1481    3. Merge groups across adduct edges; set ``ion_type`` / roles.
1482
1483    Does not model tracer/labeled enrichment. No mono-vs-family height prior.
1484
1485    Parameters
1486    ----------
1487    cluster_ids : array-like, shape (N,)
1488    mz, rt : array-like, shape (N,)
1489    heights : array-like, shape (N, S)
1490        Apex intensity matrix (not area).
1491    params : FeatureGroupParams
1492    timings_out : mutable mapping, optional
1493        If provided, filled with stage wall times in seconds
1494        (``isotope_edges``, ``isotope_corr``, ``isotope_labels``,
1495        ``adduct_edges``, ``adduct_corr``, ``adduct_merge``, ``total``,
1496        plus edge/type counts).
1497
1498    Returns
1499    -------
1500    DataFrame
1501        Index = cluster_ids; columns GROUP_COLUMNS.
1502    """
1503    t_all = time.perf_counter()
1504    validate_feature_group_params(params)
1505
1506    cluster_ids = np.asarray(cluster_ids)
1507    mz = np.asarray(mz, dtype=float)
1508    rt = np.asarray(rt, dtype=float)
1509    heights = np.asarray(heights, dtype=float)
1510    if heights.ndim == 1:
1511        heights = heights.reshape(-1, 1)
1512
1513    if len(cluster_ids) != len(mz) or len(mz) != len(rt):
1514        raise ValueError("cluster_ids, mz, and rt must have the same length")
1515    if heights.shape[0] != len(cluster_ids):
1516        raise ValueError("heights must have shape (N, S) matching cluster_ids")
1517
1518    if timings_out is not None:
1519        timings_out.clear()
1520        timings_out["n_clusters"] = float(len(cluster_ids))
1521        timings_out["n_ion_types"] = float(len(params.ion_types))
1522        timings_out["n_samples"] = float(heights.shape[1])
1523
1524    if len(cluster_ids) < 2:
1525        if timings_out is not None:
1526            timings_out["total"] = time.perf_counter() - t_all
1527        return empty_group_labels(cluster_ids)
1528
1529    # Isotope stage
1530    t0 = time.perf_counter()
1531    iso_edges = find_isotope_edges(cluster_ids, mz, rt, params)
1532    if timings_out is not None:
1533        timings_out["isotope_edges"] = time.perf_counter() - t0
1534        timings_out["n_isotope_edges_geom"] = float(len(iso_edges))
1535
1536    t0 = time.perf_counter()
1537    iso_edges = filter_edges_by_height_correlation(iso_edges, heights, params)
1538    if timings_out is not None:
1539        timings_out["isotope_corr"] = time.perf_counter() - t0
1540        timings_out["n_isotope_edges"] = float(len(iso_edges))
1541
1542    t0 = time.perf_counter()
1543    labels = assign_isotope_labels(cluster_ids, mz, rt, heights, iso_edges, params)
1544    if timings_out is not None:
1545        timings_out["isotope_labels"] = time.perf_counter() - t0
1546
1547    mz_by_cluster = {int(cid): float(m) for cid, m in zip(cluster_ids, mz)}
1548    # Adduct stage: only mono or unlabeled endpoints (never ion_role=isotope)
1549    if len(params.ion_types) >= 2:
1550        eligible_mask = np.array(
1551            [
1552                is_adduct_endpoint_eligible(
1553                    labels.loc[cid, "ion_role"] if cid in labels.index else None
1554                )
1555                for cid in cluster_ids
1556            ],
1557            dtype=bool,
1558        )
1559        elig_ids = cluster_ids[eligible_mask]
1560        elig_mz = mz[eligible_mask]
1561        elig_rt = rt[eligible_mask]
1562        elig_heights = heights[eligible_mask]
1563
1564        t0 = time.perf_counter()
1565        if len(elig_ids) >= 2:
1566            add_edges = find_adduct_edges(elig_ids, elig_mz, elig_rt, params)
1567        else:
1568            add_edges = pd.DataFrame(columns=list(_ADDUCT_EDGE_COLUMNS))
1569        if timings_out is not None:
1570            timings_out["adduct_edges"] = time.perf_counter() - t0
1571            timings_out["n_adduct_edges_geom"] = float(len(add_edges))
1572            timings_out["n_adduct_endpoints"] = float(len(elig_ids))
1573
1574        t0 = time.perf_counter()
1575        add_edges = filter_edges_by_height_correlation(add_edges, elig_heights, params)
1576        if timings_out is not None:
1577            timings_out["adduct_corr"] = time.perf_counter() - t0
1578            timings_out["n_adduct_edges"] = float(len(add_edges))
1579
1580        t0 = time.perf_counter()
1581        add_edges = sort_adduct_edges(add_edges)
1582        if timings_out is not None:
1583            timings_out["adduct_sort"] = time.perf_counter() - t0
1584
1585        t0 = time.perf_counter()
1586        labels = merge_adduct_edges_into_labels(
1587            labels,
1588            cluster_ids,
1589            add_edges,
1590            params,
1591            mz_by_cluster=mz_by_cluster,
1592        )
1593        if timings_out is not None:
1594            timings_out["adduct_merge"] = time.perf_counter() - t0
1595    elif timings_out is not None:
1596        timings_out["adduct_edges"] = 0.0
1597        timings_out["adduct_corr"] = 0.0
1598        timings_out["adduct_sort"] = 0.0
1599        timings_out["adduct_merge"] = 0.0
1600        timings_out["n_adduct_edges"] = 0.0
1601
1602    t0 = time.perf_counter()
1603    labels = validate_feature_group_labels(labels, mz_by_cluster, params)
1604    if timings_out is not None:
1605        timings_out["validate"] = time.perf_counter() - t0
1606        timings_out["total"] = time.perf_counter() - t_all
1607
1608    return labels
1609
1610
1611def detected_cluster_mz_median(
1612    mass_features_df: pd.DataFrame,
1613    cluster_ids: Sequence,
1614    cluster_col: str = "cluster",
1615    mz_col: str = "mz",
1616) -> np.ndarray:
1617    """Median m/z of detected features, in ``cluster_ids`` order.
1618
1619    ``mass_features_df`` is the non-gap-filled table. Isotope and adduct
1620    geometry uses this median, so a gap-filled apex m/z cannot move
1621    ``ion_role``.
1622    """
1623    if mass_features_df is None or cluster_col not in mass_features_df.columns:
1624        raise ValueError(f"mass features missing '{cluster_col}' column")
1625    if mz_col not in mass_features_df.columns:
1626        raise ValueError(f"mass features missing '{mz_col}' column")
1627
1628    detected = mass_features_df.dropna(subset=[cluster_col, mz_col]).copy()
1629    detected[cluster_col] = detected[cluster_col].astype(int)
1630    median_by_cluster = detected.groupby(cluster_col, sort=False)[mz_col].median()
1631
1632    ids = pd.Index(np.asarray(cluster_ids).astype(int))
1633    missing = ids.difference(median_by_cluster.index)
1634    if len(missing) > 0:
1635        shown = ", ".join(str(cluster_id) for cluster_id in list(missing[:8]))
1636        raise ValueError(
1637            "ion-role m/z median needs a detected feature for every cluster; "
1638            f"missing {shown}"
1639        )
1640    return median_by_cluster.reindex(ids).to_numpy(dtype=float)
1641
1642
1643def build_height_matrix_from_features(
1644    mass_features_df: pd.DataFrame,
1645    cluster_ids: Sequence,
1646    sample_ids: Sequence,
1647    induced_df: Optional[pd.DataFrame] = None,
1648    intensity_col: str = HEIGHT_COL,
1649    cluster_col: str = "cluster",
1650    sample_col: str = "sample_id",
1651) -> np.ndarray:
1652    """
1653    Build (N, S) height matrix for ordered cluster_ids × sample_ids.
1654
1655    Library grouping always passes apex ``intensity`` (``HEIGHT_COL``). There is
1656    no integrated-area path in ``group_consensus_features``. The ``intensity_col``
1657    argument is only for offline diagnostics outside the public API.
1658
1659    Missing cluster/sample pairs are 0.0. If multiple rows exist for the same
1660    (cluster, sample), the maximum intensity is used.
1661    """
1662    frames = [mass_features_df]
1663    if induced_df is not None and len(induced_df) > 0:
1664        frames.append(induced_df)
1665    mf = pd.concat(frames, axis=0, ignore_index=False)
1666    if cluster_col not in mf.columns:
1667        raise ValueError(f"mass features missing '{cluster_col}' column")
1668    if intensity_col not in mf.columns:
1669        raise ValueError(f"mass features missing '{intensity_col}' column")
1670    if sample_col not in mf.columns:
1671        raise ValueError(f"mass features missing '{sample_col}' column")
1672
1673    mf = mf.dropna(subset=[cluster_col])
1674    mf = mf.copy()
1675    mf[cluster_col] = mf[cluster_col].astype(int)
1676
1677    pivot = (
1678        mf.groupby([cluster_col, sample_col], as_index=False)[intensity_col]
1679        .max()
1680        .pivot(index=cluster_col, columns=sample_col, values=intensity_col)
1681    )
1682    # Reindex to full cluster × sample grid
1683    pivot = pivot.reindex(index=list(cluster_ids), columns=list(sample_ids))
1684    pivot = pivot.fillna(0.0)
1685    return pivot.to_numpy(dtype=float)