Skills › Scientific Agent Skills › matchms › library_search.py
library_search.py
624 lines · skills/matchms/scripts/library_search.py at 5f2bc03
1#!/usr/bin/env python3
2"""Search query MS/MS spectra against a reference library with matchms.
4Examples:
5 uv run python library_search.py queries.mgf library.msp hits.csv
6 uv run python library_search.py queries.mgf library.msp hits.csv \
7 --metric modified --tolerance 0.02 --top-k 10 \
8 --min-score 0.6 --min-matches 5
9 uv run python library_search.py queries.mgf library.msp hits.csv \
10 --metric flash-entropy --max-pairs 20000000
12This script targets matchms 0.33.1. It refuses pickle inputs, bounds the
13query-reference Cartesian product, applies identical peak processing to both
14collections, and handles both scalar and structured matchms score records.
15"""
17from __future__ import annotations
19import argparse
20import csv
21import math
22import sys
23from importlib.metadata import PackageNotFoundError, version
24from pathlib import Path
25from typing import Any
27try:
28 from matchms import SpectrumProcessor, calculate_scores
29 from matchms.filtering import (
30 add_compound_name,
31 add_precursor_mz,
32 clean_compound_name,
33 correct_charge,
34 derive_adduct_from_name,
35 derive_formula_from_name,
36 derive_ionmode,
37 interpret_pepmass,
38 make_charge_int,
39 normalize_intensities,
40 reduce_to_number_of_peaks,
41 remove_peaks_around_precursor_mz,
42 require_minimum_number_of_peaks,
43 require_precursor_mz,
44 select_by_mz,
45 select_by_relative_intensity,
46 )
47 from matchms.importing import load_spectra
48 from matchms.similarity import (
49 BlinkCosine,
50 CosineGreedy,
51 CosineHungarian,
52 CosineLinear,
53 FlashSimilarity,
54 ModifiedCosineGreedy,
55 ModifiedCosineHungarian,
56 NeutralLossesCosine,
57 )
58except ImportError as exc:
59 raise SystemExit(
60 'matchms is required. Install the verified release with: '
61 'uv pip install "matchms==0.33.1"'
62 ) from exc
65TARGET_VERSION = "0.33.1"
66UNSAFE_PICKLE_SUFFIXES = {".pickle", ".pkl"}
67SUPPORTED_SUFFIXES = {".json", ".mgf", ".msp", ".mzml", ".mzxml"}
68STRUCTURED_METRICS = {
69 "cosine",
70 "cosine-exact",
71 "cosine-linear",
72 "modified",
73 "modified-exact",
74 "neutral-loss",
75 "blink",
76}
77PRECURSOR_METRICS = {
78 "modified",
79 "modified-exact",
80 "neutral-loss",
81 "flash-modified",
82}
83DEFAULT_METADATA_FILTERS = (
84 make_charge_int,
85 add_compound_name,
86 derive_adduct_from_name,
87 derive_formula_from_name,
88 clean_compound_name,
89 interpret_pepmass,
90 add_precursor_mz,
91 derive_ionmode,
92 correct_charge,
93)
96def build_parser() -> argparse.ArgumentParser:
97 parser = argparse.ArgumentParser(
98 description=__doc__,
99 formatter_class=argparse.RawDescriptionHelpFormatter,
100 )
101 parser.add_argument("queries", type=Path, help="Query spectra: MGF/MSP/mzML/mzXML/JSON")
102 parser.add_argument("references", type=Path, help="Reference library: MGF/MSP/mzML/mzXML/JSON")
103 parser.add_argument("output", type=Path, help="Output CSV path")
104 parser.add_argument(
105 "--metric",
106 choices=[
107 "cosine",
108 "cosine-exact",
109 "cosine-linear",
110 "modified",
111 "modified-exact",
112 "neutral-loss",
113 "blink",
114 "flash-entropy",
115 "flash-cosine",
116 "flash-modified",
117 ],
118 default="modified",
119 help="Similarity method (default: modified)",
120 )
121 parser.add_argument(
122 "--tolerance",
123 type=float,
124 default=0.02,
125 help="Fragment matching tolerance in daltons (default: 0.02)",
126 )
127 parser.add_argument("--top-k", type=int, default=10, help="Hits per query (default: 10)")
128 parser.add_argument(
129 "--min-score",
130 type=float,
131 default=0.0,
132 help="Minimum similarity in [0, 1] (default: 0)",
133 )
134 parser.add_argument(
135 "--min-matches",
136 type=int,
137 default=0,
138 help="Minimum matched peaks for metrics that report matches (default: 0)",
139 )
140 parser.add_argument(
141 "--array-type",
142 choices=["numpy", "sparse"],
143 default="numpy",
144 help="matchms calculation array type (default: numpy)",
145 )
146 parser.add_argument(
147 "--max-pairs",
148 type=int,
149 default=5_000_000,
150 help="Refuse comparisons above this reference×query count (default: 5,000,000)",
151 )
152 parser.add_argument(
153 "--relative-intensity",
154 type=float,
155 default=0.01,
156 help="Drop peaks below this fraction of maximum; 0 disables (default: 0.01)",
157 )
158 parser.add_argument(
159 "--min-peaks",
160 type=int,
161 default=5,
162 help="Reject spectra with fewer peaks; 0 disables (default: 5)",
163 )
164 parser.add_argument(
165 "--max-peaks",
166 type=int,
167 default=None,
168 help="Keep at most this many highest-intensity peaks",
169 )
170 parser.add_argument("--mz-min", type=float, default=None, help="Minimum fragment m/z")
171 parser.add_argument("--mz-max", type=float, default=None, help="Maximum fragment m/z")
172 parser.add_argument(
173 "--remove-precursor-window",
174 type=float,
175 default=None,
176 metavar="DA",
177 help="Remove peaks within this precursor-centered window",
178 )
179 parser.add_argument(
180 "--no-default-filters",
181 action="store_true",
182 help="Skip matchms metadata default_filters",
183 )
184 parser.add_argument(
185 "--no-normalize",
186 action="store_true",
187 help="Skip intensity normalization",
188 )
189 parser.add_argument(
190 "--query-id-field",
191 default=None,
192 help="Preferred metadata key for query identifiers",
193 )
194 parser.add_argument(
195 "--reference-id-field",
196 default=None,
197 help="Preferred metadata key for reference identifiers",
198 )
199 parser.add_argument(
200 "--bin-width",
201 type=float,
202 default=0.001,
203 help="BLINK bin width in daltons (default: 0.001)",
204 )
205 parser.add_argument(
206 "--blink-top-k",
207 type=int,
208 default=None,
209 help="Optional number of most intense peaks retained by BLINK",
210 )
211 parser.add_argument("--force", action="store_true", help="Overwrite an existing CSV")
212 parser.add_argument("--quiet", action="store_true", help="Hide processing progress bars")
213 return parser
216def validate_args(args: argparse.Namespace) -> None:
217 for path, label in ((args.queries, "query"), (args.references, "reference")):
218 if not path.is_file():
219 raise ValueError(f"{label} file does not exist: {path}")
220 suffix = path.suffix.lower()
221 if suffix in UNSAFE_PICKLE_SUFFIXES:
222 raise ValueError(
223 f"pickle input is intentionally refused because unpickling can execute code: {path}"
224 )
225 if suffix not in SUPPORTED_SUFFIXES:
226 raise ValueError(
227 f"unsupported {label} suffix {suffix!r}; choose MGF, MSP, mzML, mzXML, or JSON"
228 )
230 if args.output.exists() and not args.force:
231 raise ValueError(f"output already exists (pass --force to replace it): {args.output}")
232 if not args.output.parent.is_dir():
233 raise ValueError(f"output directory does not exist: {args.output.parent}")
234 if args.output.resolve() in {args.queries.resolve(), args.references.resolve()}:
235 raise ValueError("output must not replace an input spectrum file")
236 for name in ("tolerance", "min_score", "relative_intensity", "mz_min", "mz_max",
237 "remove_precursor_window", "bin_width"):
238 value = getattr(args, name)
239 if value is not None and not math.isfinite(value):
240 raise ValueError(f"--{name.replace('_', '-')} must be finite")
241 if args.tolerance <= 0:
242 raise ValueError("--tolerance must be positive")
243 if args.top_k <= 0:
244 raise ValueError("--top-k must be positive")
245 if not 0.0 <= args.min_score <= 1.0:
246 raise ValueError("--min-score must be between 0 and 1")
247 if args.min_matches < 0:
248 raise ValueError("--min-matches cannot be negative")
249 if args.max_pairs <= 0:
250 raise ValueError("--max-pairs must be positive")
251 if not 0.0 <= args.relative_intensity <= 1.0:
252 raise ValueError("--relative-intensity must be between 0 and 1")
253 if args.min_peaks < 0:
254 raise ValueError("--min-peaks cannot be negative")
255 if args.max_peaks is not None and args.max_peaks <= 0:
256 raise ValueError("--max-peaks must be positive")
257 if args.mz_min is not None and args.mz_max is not None and args.mz_min >= args.mz_max:
258 raise ValueError("--mz-min must be smaller than --mz-max")
259 if args.remove_precursor_window is not None and args.remove_precursor_window <= 0:
260 raise ValueError("--remove-precursor-window must be positive")
261 if args.bin_width <= 0:
262 raise ValueError("--bin-width must be positive")
263 if args.blink_top_k is not None and args.blink_top_k <= 0:
264 raise ValueError("--blink-top-k must be positive")
265 if args.metric not in STRUCTURED_METRICS and args.min_matches > 0:
266 raise ValueError(f"--metric {args.metric!r} does not report matched-peak counts")
269def installed_matchms_version() -> str:
270 try:
271 return version("matchms")
272 except PackageNotFoundError:
273 return "unknown"
276def require_finite_positive_precursor(spectrum_in: Any, clone: bool = True) -> Any:
277 """Reject NaN, infinity and zero that upstream range checks can accept."""
278 if spectrum_in is None:
279 return None
280 precursor = spectrum_in.get("precursor_mz")
281 if (isinstance(precursor, bool) or not isinstance(precursor, (int, float))
282 or not math.isfinite(precursor) or precursor <= 0):
283 return None
284 return spectrum_in.clone() if clone else spectrum_in
287def create_processor(args: argparse.Namespace) -> SpectrumProcessor:
288 filters: list[Any] = []
289 if not args.no_default_filters:
290 # Expand default_filters so SpectrumProcessor can preserve the registered
291 # metadata-before-peaks order. The aggregate callable is otherwise
292 # treated as a custom filter and moved to the end.
293 filters.extend(DEFAULT_METADATA_FILTERS)
295 needs_precursor = (
296 args.metric in PRECURSOR_METRICS
297 or args.remove_precursor_window is not None
298 )
299 if needs_precursor:
300 filters.append((require_precursor_mz, {"minimum_accepted_mz": 0.0}))
301 if not args.no_normalize:
302 filters.append(normalize_intensities)
303 if args.mz_min is not None or args.mz_max is not None:
304 filters.append(
305 (
306 select_by_mz,
307 {
308 "mz_from": 0.0 if args.mz_min is None else args.mz_min,
309 "mz_to": float("inf") if args.mz_max is None else args.mz_max,
310 },
311 )
312 )
313 if args.remove_precursor_window is not None:
314 filters.append(
315 (
316 remove_peaks_around_precursor_mz,
317 {"mz_tolerance": args.remove_precursor_window},
318 )
319 )
320 if args.relative_intensity > 0:
321 filters.append(
322 (
323 select_by_relative_intensity,
324 {"intensity_from": args.relative_intensity},
325 )
326 )
327 if args.max_peaks is not None:
328 filters.append((reduce_to_number_of_peaks, {"n_max": args.max_peaks}))
329 if args.min_peaks > 0:
330 filters.append((require_minimum_number_of_peaks, {"n_required": args.min_peaks}))
331 processor = SpectrumProcessor(filters)
332 if needs_precursor:
333 # Place after metadata harmonization and before any precursor-dependent
334 # peak operation; custom filters otherwise default to the end.
335 names = [step[0] if isinstance(step, tuple) else step
336 for step in processor.processing_steps]
337 processor.parse_and_add_filter(
338 require_finite_positive_precursor,
339 filter_position=names.index("require_precursor_mz"),
340 )
341 return processor
344def load_and_process(
345 path: Path,
346 processor: SpectrumProcessor,
347 *,
348 show_progress: bool,
349) -> tuple[list[Any], int]:
350 raw = list(load_spectra(str(path)))
351 cleaned, _ = processor.process_spectra(
352 raw,
353 progress_bar=show_progress,
354 create_report=False,
355 )
356 return cleaned, len(raw)
359def create_metric(args: argparse.Namespace) -> Any:
360 common = {"tolerance": args.tolerance}
361 if args.metric == "cosine":
362 return CosineGreedy(**common)
363 if args.metric == "cosine-exact":
364 return CosineHungarian(**common)
365 if args.metric == "cosine-linear":
366 return CosineLinear(**common)
367 if args.metric == "modified":
368 return ModifiedCosineGreedy(**common)
369 if args.metric == "modified-exact":
370 return ModifiedCosineHungarian(**common)
371 if args.metric == "neutral-loss":
372 return NeutralLossesCosine(**common)
373 if args.metric == "blink":
374 return BlinkCosine(
375 tolerance=args.tolerance,
376 bin_width=args.bin_width,
377 min_relative_intensity=args.relative_intensity,
378 top_k=args.blink_top_k,
379 sparse_score_min=args.min_score,
380 )
381 if args.metric == "flash-entropy":
382 return FlashSimilarity(
383 score_type="spectral_entropy",
384 matching_mode="fragment",
385 tolerance=args.tolerance,
386 noise_cutoff=args.relative_intensity,
387 )
388 if args.metric == "flash-cosine":
389 return FlashSimilarity(
390 score_type="cosine",
391 matching_mode="fragment",
392 tolerance=args.tolerance,
393 noise_cutoff=args.relative_intensity,
394 )
395 if args.metric == "flash-modified":
396 return FlashSimilarity(
397 score_type="cosine",
398 matching_mode="hybrid",
399 tolerance=args.tolerance,
400 noise_cutoff=args.relative_intensity,
401 )
402 raise ValueError(f"unsupported metric: {args.metric}")
405def choose_score_fields(score_names: tuple[str, ...]) -> tuple[str, str | None]:
406 score_name = next(
407 (name for name in score_names if name.endswith("_score")),
408 score_names[0],
409 )
410 matches_name = next(
411 (name for name in score_names if name.endswith("_matches")),
412 None,
413 )
414 return score_name, matches_name
417def numeric_field(value: Any, name: str) -> float:
418 dtype = getattr(value, "dtype", None)
419 names = getattr(dtype, "names", None)
420 if names and name in names:
421 return float(value[name])
422 return float(value)
425def matched_peaks(value: Any, name: str | None) -> int | None:
426 if name is None:
427 return None
428 dtype = getattr(value, "dtype", None)
429 names = getattr(dtype, "names", None)
430 if names and name in names:
431 return int(value[name])
432 return None
435def spectrum_id(
436 spectrum: Any,
437 *,
438 preferred: str | None,
439 index: int,
440 prefix: str,
441) -> str:
442 keys = [
443 preferred,
444 "spectrum_id",
445 "id",
446 "feature_id",
447 "scan_number",
448 "compound_name",
449 "title",
450 ]
451 for key in keys:
452 if key:
453 value = spectrum.get(key)
454 if value not in (None, ""):
455 return str(value)
456 return f"{prefix}-{index}"
459def metadata_value(spectrum: Any, key: str) -> str:
460 value = spectrum.get(key)
461 return "" if value is None else str(value)
464def write_hits(
465 output: Path,
466 *,
467 scores: Any,
468 queries: list[Any],
469 references: list[Any],
470 args: argparse.Namespace,
471) -> int:
472 score_name, matches_name = choose_score_fields(scores.score_names)
473 reference_indices = {id(spectrum): index for index, spectrum in enumerate(references)}
474 fieldnames = [
475 "query_index",
476 "query_id",
477 "query_precursor_mz",
478 "reference_index",
479 "reference_id",
480 "reference_compound_name",
481 "reference_inchikey",
482 "reference_precursor_mz",
483 "rank",
484 "metric",
485 "score_field",
486 "score",
487 "matched_peaks",
488 ]
490 rows_written = 0
491 mode = "w"
492 with output.open(mode, newline="", encoding="utf-8") as handle:
493 writer = csv.DictWriter(handle, fieldnames=fieldnames)
494 writer.writeheader()
496 for query_index, query in enumerate(queries):
497 ranked = scores.scores_by_query(query, name=score_name, sort=True)
498 accepted = 0
499 for reference, value in ranked:
500 score = numeric_field(value, score_name)
501 matches = matched_peaks(value, matches_name)
502 if not math.isfinite(score) or score < args.min_score:
503 continue
504 if matches is not None and matches < args.min_matches:
505 continue
507 reference_index = reference_indices.get(id(reference))
508 if reference_index is None:
509 raise RuntimeError("a score referenced a spectrum outside the input library")
511 accepted += 1
512 writer.writerow(
513 {
514 "query_index": query_index,
515 "query_id": spectrum_id(
516 query,
517 preferred=args.query_id_field,
518 index=query_index,
519 prefix="query",
520 ),
521 "query_precursor_mz": metadata_value(query, "precursor_mz"),
522 "reference_index": reference_index,
523 "reference_id": spectrum_id(
524 reference,
525 preferred=args.reference_id_field,
526 index=reference_index,
527 prefix="reference",
528 ),
529 "reference_compound_name": metadata_value(
530 reference,
531 "compound_name",
532 ),
533 "reference_inchikey": metadata_value(reference, "inchikey"),
534 "reference_precursor_mz": metadata_value(
535 reference,
536 "precursor_mz",
537 ),
538 "rank": accepted,
539 "metric": args.metric,
540 "score_field": score_name,
541 "score": format(score, ".12g"),
542 "matched_peaks": "" if matches is None else matches,
543 }
544 )
545 rows_written += 1
546 if accepted >= args.top_k:
547 break
548 return rows_written
551def run(args: argparse.Namespace) -> int:
552 validate_args(args)
553 matchms_version = installed_matchms_version()
554 if matchms_version != TARGET_VERSION:
555 print(
556 f"warning: script targets matchms {TARGET_VERSION}; found {matchms_version}",
557 file=sys.stderr,
558 )
560 processor = create_processor(args)
561 queries, raw_query_count = load_and_process(
562 args.queries,
563 processor,
564 show_progress=not args.quiet,
565 )
566 references, raw_reference_count = load_and_process(
567 args.references,
568 processor,
569 show_progress=not args.quiet,
570 )
571 if not queries:
572 raise ValueError("no query spectra remain after processing")
573 if not references:
574 raise ValueError("no reference spectra remain after processing")
576 pair_count = len(queries) * len(references)
577 if pair_count > args.max_pairs:
578 raise ValueError(
579 f"comparison would evaluate {pair_count:,} pairs, above --max-pairs "
580 f"{args.max_pairs:,}; raise the limit only after reviewing memory/runtime"
581 )
583 print(
584 f"matchms {matchms_version}; queries {raw_query_count}->{len(queries)}; "
585 f"references {raw_reference_count}->{len(references)}; pairs {pair_count:,}",
586 file=sys.stderr,
587 )
588 print(f"processing steps: {processor.processing_steps}", file=sys.stderr)
590 metric = create_metric(args)
591 scores = calculate_scores(
592 references=references,
593 queries=queries,
594 similarity_function=metric,
595 array_type=args.array_type,
596 is_symmetric=False,
597 )
598 rows_written = write_hits(
599 args.output,
600 scores=scores,
601 queries=queries,
602 references=references,
603 args=args,
604 )
605 print(
606 f"wrote {rows_written} hit rows to {args.output} "
607 f"(score fields: {', '.join(scores.score_names)})",
608 file=sys.stderr,
609 )
610 return 0
613def main() -> int:
614 parser = build_parser()
615 args = parser.parse_args()
616 try:
617 return run(args)
618 except (AssertionError, OSError, RuntimeError, TypeError, ValueError) as exc:
619 print(f"error: {exc}", file=sys.stderr)
620 return 2
623if __name__ == "__main__":
624 raise SystemExit(main())