Document assignment artifact generator

Author: joewandy

vimms/AssignmentChemicalArtifact.py

@@ -1,3 +1,5 @@
+"""Materialize assignment-noise peak tables as VIMMS chemicals and mzML sidecars."""
+
 from __future__ import annotations
 
 import json
@@ -54,6 +56,8 @@ class AssignmentChemicalArtifactConfig:
 def _chromatogram_for_peak(
     apex_seconds: float, width_seconds: float
 ) -> tuple[EmpiricalChromatogram, float]:
+    """Build a simple empirical chromatogram centered on one picked feature."""
+
     half_width = max(float(width_seconds), 1.0)
     start = max(0.0, float(apex_seconds) - half_width)
     apex = max(float(apex_seconds), start + 0.25)
@@ -69,6 +73,8 @@ def _make_ms2_children(
     row: pd.Series,
     rng: np.random.Generator,
 ) -> list[MSN]:
+    """Create fragment children for chemicals with available/plausible MS2 support."""
+
     score = float(row.get("best_ms2_score", 0.0))
     available = float(row.get("ms2_available", 0.0)) > 0.0
     if not available and score < 0.25:
@@ -97,6 +103,8 @@ def _materialize_assignment_chemicals(
     config: AssignmentChemicalArtifactConfig,
     rng: np.random.Generator,
 ) -> tuple[list[UnknownChemical], pd.DataFrame]:
+    """Convert every truth/artifact peak row into a VIMMS UnknownChemical."""
+
     chemicals: list[UnknownChemical] = []
     truth_rows: list[dict[str, Any]] = []
 
@@ -109,6 +117,9 @@ def _materialize_assignment_chemicals(
             float(row.get("peak_width", 0.08)) * 60.0 * 1.5,
         )
         chromatogram, start_rt = _chromatogram_for_peak(target_rt_seconds, width_seconds)
+        # UnknownChemical stores a neutral mass, but the picked-table contract is
+        # already m/z based. Subtracting a proton gives a chemical that emits at
+        # the target m/z without changing model-facing values.
         neutral_for_unknown = max(1.0, target_mz - PROTON_MASS)
         max_intensity = max(
             float(row.get("intensity", np.exp(row.get("log_intensity", 12.0)))), 1.0
@@ -149,6 +160,8 @@ def _materialize_assignment_chemicals(
 
 
 def _output_path(config: AssignmentChemicalArtifactConfig) -> Path | None:
+    """Return an output directory only when mzML writing is requested."""
+
     if not config.write_mzml:
         return None
     if config.output_dir is None:
@@ -162,6 +175,8 @@ def _run_scan_simulation(
     chemicals: list[UnknownChemical],
     config: AssignmentChemicalArtifactConfig,
 ) -> tuple[Environment, Path | None]:
+    """Run VIMMS TopN scan simulation for optional audit/export mzML output."""
+
     out_dir = _output_path(config)
     out_file = f"{config.prefix}.mzML" if out_dir is not None else None
     mass_spec = IndependentMassSpectrometer(
@@ -195,6 +210,8 @@ def _run_scan_simulation(
 
 
 def _scan_summary(scans: Mapping[int, list[Any]]) -> pd.DataFrame:
+    """Flatten VIMMS controller scans into a lightweight audit table."""
+
     rows: list[dict[str, Any]] = []
     for ms_level, level_scans in sorted(scans.items()):
         for scan in level_scans:
@@ -211,12 +228,16 @@ def _scan_summary(scans: Mapping[int, list[Any]]) -> pd.DataFrame:
 
 
 def _empty_scan_summary() -> pd.DataFrame:
+    """Return the scan-summary schema used when mzML writing is disabled."""
+
     return pd.DataFrame(
         columns=["scan_id", "ms_level", "rt_seconds", "n_peaks", "tic"]
     )
 
 
 def _truth_table_from_peaks(peak_table: pd.DataFrame) -> pd.DataFrame:
+    """Extract peak-level truth columns from the model-facing peak table."""
+
     columns = [
         "peak_id",
         "true_label",
@@ -238,6 +259,8 @@ def _truth_table_from_peaks(peak_table: pd.DataFrame) -> pd.DataFrame:
 
 
 def _jsonable(value: Any) -> Any:
+    """Recursively convert dataclass metadata values to JSON-safe objects."""
+
     if isinstance(value, Path):
         return str(value)
     if isinstance(value, tuple):
@@ -276,6 +299,8 @@ def generate_assignment_chemical_artifact(
         rng=rng,
     )
     if config.write_mzml:
+        # Scans are sidecars only. The returned peak table stays bit-identical
+        # to the direct picked-table generator used by IonClassifier.
         env, mzml_path = _run_scan_simulation(chemicals, config)
         scan_summary = _scan_summary(env.controller.scans)
     else:

vimms/AssignmentNoise.py

@@ -1,3 +1,5 @@
+"""Generate labelled picked-peak assignment scenarios for IonClassifier."""
+
 from __future__ import annotations
 
 import json
@@ -161,6 +163,8 @@ def role_properties(role: str, carbon_count: float) -> dict[str, float]:
 
 
 def _validate_config(config: AssignmentScenarioConfig) -> None:
+    """Validate dimensions and explicit candidate-presence overrides."""
+
     profile = config.profile
     if profile.n_candidates < 1:
         raise ValueError("n_candidates must be at least 1")
@@ -180,12 +184,16 @@ def _presence_pattern(
     profile: AssignmentNoiseProfile,
     present_pattern: tuple[int, ...] | None,
 ) -> np.ndarray:
+    """Sample which candidates are truly present in a local assignment problem."""
+
     if present_pattern is not None:
         return np.asarray(present_pattern, dtype=np.int8)
     present = (rng.random(profile.n_candidates) < profile.p_present).astype(np.int8)
     if not present.any():
         present[0] = 1
     if profile.n_candidates > 1 and present.all() and profile.p_matched_decoy > 0.0:
+        # Keep one plausible absent candidate available so matched-decoy
+        # clusters can test false positive rejection.
         present[-1] = 0
     return present
 
@@ -195,6 +203,8 @@ def _make_candidate_table(
     profile: AssignmentNoiseProfile,
     present: np.ndarray,
 ) -> pd.DataFrame:
+    """Create candidate metadata around one target-like local mass/RT region."""
+
     base_mass = float(rng.uniform(profile.mass_min, profile.mass_max))
     base_rt = float(rng.uniform(profile.rt_min, profile.rt_max))
     base_c = int(rng.integers(5, 45))
@@ -236,6 +246,8 @@ def _make_candidate_table(
 
 
 def _make_ion_table(candidates: pd.DataFrame, roles: tuple[str, ...]) -> pd.DataFrame:
+    """Expand candidates into theoretical role/adduct/isotope ions."""
+
     rows: list[dict[str, Any]] = []
     for _, candidate in candidates.iterrows():
         candidate_index = int(candidate["candidate_index"])
@@ -264,6 +276,8 @@ def _make_ion_truth_table(
     candidate_table: pd.DataFrame,
     ion_table: pd.DataFrame,
 ) -> pd.DataFrame:
+    """Record observed and missing true companion ions for every candidate role."""
+
     observed = (
         peak_table.loc[peak_table["true_label"] > 0, ["true_label", "peak_id", "source_type"]]
         .drop_duplicates("true_label")
@@ -321,13 +335,17 @@ def generate_assignment_peak_table(
     latent_profiles: dict[int, np.ndarray] = {}
     for _, candidate in candidate_table.iterrows():
         candidate_index = int(candidate["candidate_index"])
+        # Candidate ions share a latent abundance profile so isotope/adduct
+        # companions look correlated, while artifacts can partially mimic it.
         latent_profiles[candidate_index] = rng.normal(
             float(candidate["expected_log_intensity"]),
             rng.uniform(0.35, 0.95),
             size=profile.n_bio_samples,
         )
 
     peaks: list[dict[str, Any]] = []
+    # Add true evidence first, then layer increasingly adversarial background
+    # profiles around the same theoretical ions.
     _add_candidate_ions(peaks, rng, candidate_table, ion_table, profile, roles, latent_profiles)
     _add_structured_interferents(peaks, rng, candidate_table, ion_table, profile)
     _add_coeluting_isobars(peaks, rng, candidate_table, ion_table, profile)

vimms/AssignmentNoisePeaks.py

@@ -1,3 +1,5 @@
+"""Peak-row generation helpers for assignment-noise scenarios."""
+
 from __future__ import annotations
 
 from typing import Any, Mapping
@@ -9,6 +11,8 @@
 
 
 def _skew(values: np.ndarray) -> float:
+    """Return sample skewness with a stable value for near-constant profiles."""
+
     arr = np.asarray(values, dtype=float)
     sd = arr.std()
     if sd < 1e-8:
@@ -17,6 +21,8 @@ def _skew(values: np.ndarray) -> float:
 
 
 def _beta_score(rng: np.random.Generator, true: bool, fp_prob: float) -> float:
+    """Sample an MS2-like support score for true and false candidate evidence."""
+
     if true:
         return float(rng.beta(9.0, 2.2))
     if rng.random() < fp_prob:
@@ -49,6 +55,8 @@ def _peak_row(
     ms2_true_score: bool = False,
     metadata: Mapping[str, Any] | None = None,
 ) -> dict[str, Any]:
+    """Create one normalized picked-peak row with feature and truth metadata."""
+
     row: dict[str, Any] = {
         "mz": float(mz),
         "rt": float(rt),
@@ -93,6 +101,8 @@ def _peak_row(
     available_prob = 0.70 if ms2_true_score else min(0.85, 0.25 + 0.55 * fp_prob)
     available = bool(rng.random() < available_prob)
     if not available:
+        # Keep weak residual scores so the feature exists, but make it clear
+        # that no useful spectrum should support this row.
         for candidate_index in range(n_candidates):
             row[f"ms2_score_cand_{candidate_index}"] *= float(rng.uniform(0.0, 0.25))
         ms2_scores = [
@@ -107,6 +117,8 @@ def _peak_row(
 
 
 def _finalize_peak_table(peaks: list[dict[str, Any]], rng: np.random.Generator) -> pd.DataFrame:
+    """Compute local context features and assign stable peak identifiers."""
+
     df = pd.DataFrame(peaks)
     if df.empty:
         return df
@@ -151,6 +163,8 @@ def _add_diffuse_background(
     profile: AssignmentNoiseProfile,
     n: int,
 ) -> None:
+    """Append unstructured clutter and blank-contaminant background peaks."""
+
     mz_min = float(candidates["neutral_mass"].min() - 25.0)
     mz_max = float(candidates["neutral_mass"].max() + 45.0)
     rt_center = float(candidates["pred_rt"].mean())
@@ -193,6 +207,8 @@ def _add_candidate_ions(
     roles: tuple[str, ...],
     latent_profiles: dict[int, np.ndarray],
 ) -> None:
+    """Append true candidate-role peaks plus split/merged artifacts."""
+
     for _, candidate in candidates.iterrows():
         candidate_index = int(candidate["candidate_index"])
         if int(candidate["present"]) <= 0:
@@ -272,6 +288,8 @@ def _add_candidate_ions(
             emitted_any = True
 
             if rng.random() < profile.p_split:
+                # Split/shoulder artifacts resemble the parent ion in m/z/RT and
+                # abundance profile but remain background labels.
                 split_profile = peak_profile + rng.normal(0.0, 0.15, size=profile.n_bio_samples)
                 peaks.append(
                     _peak_row(
@@ -300,6 +318,9 @@ def _add_candidate_ions(
                 )
 
             if rng.random() < profile.p_merge:
+                # Merged peaks intentionally have good-looking evidence but are
+                # labelled background because the picked feature is not the clean
+                # theoretical ion.
                 merged_profile = peak_profile + rng.normal(0.0, 0.18, size=profile.n_bio_samples)
                 peaks.append(
                     _peak_row(
@@ -335,6 +356,8 @@ def _add_structured_interferents(
     ions: pd.DataFrame,
     profile: AssignmentNoiseProfile,
 ) -> None:
+    """Append wrong-source peaks near theoretical ions and candidate RTs."""
+
     for _, ion in ions.iterrows():
         if rng.random() >= profile.p_interferent:
             continue
@@ -408,6 +431,8 @@ def _add_coeluting_isobars(
     ions: pd.DataFrame,
     profile: AssignmentNoiseProfile,
 ) -> None:
+    """Append high-quality background peaks that coelute with theoretical ions."""
+
     for _, ion in ions.iterrows():
         if rng.random() >= profile.p_coeluting_isobar:
             continue
@@ -460,6 +485,8 @@ def _add_matched_decoy_clusters(
     profile: AssignmentNoiseProfile,
     latent_profiles: dict[int, np.ndarray],
 ) -> None:
+    """Append plausible isotope/adduct clusters for absent or wrong candidates."""
+
     candidate_indices = list(candidates["candidate_index"].astype(int))
     absent = list(
         candidates.loc[candidates["present"].astype(int) == 0, "candidate_index"].astype(int)
@@ -479,6 +506,8 @@ def _add_matched_decoy_clusters(
         cand_ions = ions[ions["candidate_index"] == candidate_index].sort_values("role_index")
         decoy_log_a = float(candidate["expected_log_intensity"] + rng.normal(-0.05, 0.55))
         if latent_profiles and rng.random() < 0.65:
+            # Reusing a true latent profile makes a decoy cluster look like a
+            # coherent biological signal instead of independent random clutter.
             base_profile = latent_profiles[int(rng.choice(list(latent_profiles.keys())))].copy()
             base_profile = base_profile - base_profile.mean() + decoy_log_a
         else:
@@ -541,6 +570,8 @@ def _add_matched_decoy_clusters(
                 },
             )
             if rng.random() < profile.p_decoy_false_ms2:
+                # Some decoy clusters get candidate-specific false MS2 support,
+                # matching the failure mode the classifier must reject.
                 row[f"ms2_score_cand_{candidate_index}"] = float(
                     max(row.get(f"ms2_score_cand_{candidate_index}", 0.0), rng.beta(5.0, 2.0))
                 )
@@ -590,6 +621,8 @@ def _trim_peaks(
     rng: np.random.Generator,
     max_peaks: int,
 ) -> list[dict[str, Any]]:
+    """Limit peak count while preserving true ions and hard negative artifacts."""
+
     if len(peaks) <= max_peaks:
         return peaks
     true_rows = [p for p in peaks if int(p["true_label"]) > 0]
@@ -605,6 +638,8 @@ def _trim_peaks(
     easy_rows = [
         p for p in peaks if int(p["true_label"]) == 0 and p["source_type"] not in hard_types
     ]
+    # True ions are always retained. Hard negatives are retained before easy
+    # diffuse clutter so stress profiles remain adversarial after trimming.
     remaining = max(0, max_peaks - len(true_rows))
     if len(hard_rows) > remaining:
         idx = rng.choice(len(hard_rows), size=remaining, replace=False)

vimms/AssignmentScanSimulation.py

@@ -1,3 +1,5 @@
+"""Backward-compatible wrappers for the old assignment scan-simulation API."""
+
 from __future__ import annotations
 
 import warnings
@@ -13,6 +15,8 @@
 
 
 def generate_assignment_scan_artifact(*args, **kwargs):
+    """Deprecated alias for ``generate_assignment_chemical_artifact``."""
+
     warnings.warn(
         "generate_assignment_scan_artifact is deprecated; use "
         "generate_assignment_chemical_artifact instead.",
@@ -23,6 +27,8 @@ def generate_assignment_scan_artifact(*args, **kwargs):
 
 
 def write_assignment_scan_artifact(*args, **kwargs):
+    """Deprecated alias for ``write_assignment_chemical_artifact``."""
+
     warnings.warn(
         "write_assignment_scan_artifact is deprecated; use "
         "write_assignment_chemical_artifact instead.",