Skills › Scientific Agent Skills › waypoint-bio › profiler_to_waypoint.py
profiler_to_waypoint.py
520 lines · skills/waypoint-bio/scripts/profiler_to_waypoint.py at 5f2bc03
1#!/usr/bin/env python3
2"""Convert microbiome profiler output into waypoint format.
4`waypoint prepare-dataset` reads a plain abundance matrix whose row or column
5labels are already ``;``-separated lineage strings. Real profilers rarely emit
6that: MetaPhlAn separates ranks with ``|``, Kraken2 encodes the hierarchy as
7indentation with no lineage string at all, and QIIME 2 / SILVA prefixes the
8domain with ``d__``, which the Waypoint tokenizer does not recognise and
9silently skips.
11This script normalises those layouts into a waypoint-format table with the
12``Taxa`` / ``Relative Abundances`` list-columns the CLI expects.
14Examples
15--------
16 python profiler_to_waypoint.py --input merged_metaphlan.tsv \
17 --format metaphlan --output dataset.parquet
19 python profiler_to_waypoint.py --input reports/*.kreport \
20 --format kraken --rank species --output dataset.parquet
22 python profiler_to_waypoint.py --input feature-table.tsv \
23 --format qiime2 --taxonomy-column taxonomy --output dataset.parquet
24"""
26from __future__ import annotations
28import argparse
29import math
30import re
31import sys
32from pathlib import Path
34from typing import TYPE_CHECKING
36if TYPE_CHECKING:
37 import pandas as pd
39# Ranks the Waypoint tokenizer understands, most specific first.
40RANK_PREFIXES: dict[str, str] = {
41 "species": "s__",
42 "genus": "g__",
43 "family": "f__",
44 "order": "o__",
45 "class": "c__",
46 "phylum": "p__",
47 "kingdom": "k__",
48}
50# Kraken2 single-letter rank codes -> tokenizer prefixes. Domain is folded into
51# kingdom because the tokenizer has no domain rank.
52KRAKEN_RANK_CODES: dict[str, str] = {
53 "D": "k__",
54 "K": "k__",
55 "P": "p__",
56 "C": "c__",
57 "O": "o__",
58 "F": "f__",
59 "G": "g__",
60 "S": "s__",
61}
63# Prefixes emitted by other databases that mean the same rank as a tokenizer
64# prefix. SILVA and Greengenes2 write the domain as ``d__``.
65PREFIX_ALIASES: dict[str, str] = {"d__": "k__", "sk__": "k__"}
67_VALID_PREFIXES = set(RANK_PREFIXES.values())
68_RANK_CODE_RE = re.compile(r"^[A-Z-]\d*$")
69_UNASSIGNED = {"unassigned", "unclassified", "", "na", "nan", "none"}
72# ---------------------------------------------------------------------------
73# Lineage normalisation
74# ---------------------------------------------------------------------------
77def normalise_lineage(lineage: str) -> str:
78 """Rewrite a lineage string into the form the Waypoint tokenizer reads.
80 Accepts ``|`` or ``;`` separators, rewrites aliased rank prefixes
81 (``d__`` -> ``k__``), drops strain (``t__``) and empty segments, and joins
82 with ``"; "``.
83 """
84 text = str(lineage).strip().strip('"')
85 separator = "|" if "|" in text else ";"
86 segments: list[str] = []
87 for raw in text.split(separator):
88 segment = raw.strip()
89 if not segment:
90 continue
91 for alias, replacement in PREFIX_ALIASES.items():
92 if segment.startswith(alias):
93 segment = replacement + segment[len(alias) :]
94 break
95 if segment.startswith("t__"):
96 continue # strain-level; below the tokenizer's deepest rank
97 # Anything without a recognised rank prefix -- "Unassigned", "root",
98 # "cellular organisms" -- is invisible to the tokenizer, so drop it
99 # rather than let it become a bogus taxon label. Bare prefixes such as
100 # "g__" carry no name and go too.
101 if segment[:3] not in _VALID_PREFIXES or len(segment) <= 3:
102 continue
103 segments.append(segment)
104 return "; ".join(segments)
107def deepest_rank(lineage: str) -> str | None:
108 """Return the most specific rank name present in a normalised lineage."""
109 present = {
110 rank
111 for rank, prefix in RANK_PREFIXES.items()
112 if any(seg.strip().startswith(prefix) for seg in lineage.split(";"))
113 }
114 for rank in RANK_PREFIXES: # dict order is most specific first
115 if rank in present:
116 return rank
117 return None
120# ---------------------------------------------------------------------------
121# Parsers: each returns a samples x lineage abundance matrix
122# ---------------------------------------------------------------------------
125def parse_metaphlan(path: Path, rank: str) -> pd.DataFrame:
126 """Parse a merged MetaPhlAn table into a samples x lineage matrix."""
127 import pandas as pd
129 header: list[str] | None = None
130 rows: list[list[str]] = []
131 with open(path, encoding="utf-8") as handle:
132 for line in handle:
133 line = line.rstrip("\n")
134 if not line.strip():
135 continue
136 fields = line.split("\t")
137 if line.startswith("#"):
138 # The real header is the comment line naming the clade column.
139 if header is None and any(
140 f.lstrip("#").strip().lower() in {"clade_name", "taxonomy"}
141 for f in fields
142 ):
143 header = [f.lstrip("#").strip() for f in fields]
144 continue
145 if header is None:
146 header = [f.strip() for f in fields]
147 continue
148 rows.append(fields)
150 if header is None or not rows:
151 raise ValueError(f"{path}: no MetaPhlAn table found")
153 frame = pd.DataFrame(rows, columns=header)
154 clade_col = header[0]
155 for dropped in ("NCBI_tax_id", "clade_taxid", "taxid"):
156 if dropped in frame.columns:
157 frame = frame.drop(columns=[dropped])
159 # Select the requested cumulative rank before removing unsupported strain
160 # segments: otherwise a species row and its t__ descendants count twice.
161 strain_rows = frame[clade_col].map(
162 lambda value: any(part.strip().startswith("t__") for part in re.split(r"[|;]", str(value)))
163 )
164 frame = frame.loc[~strain_rows].copy()
165 frame[clade_col] = frame[clade_col].map(normalise_lineage)
166 frame = frame[frame[clade_col].map(deepest_rank) == rank]
167 if frame.empty:
168 raise ValueError(f"{path}: no rows at rank {rank!r}")
170 frame = frame.set_index(clade_col)
171 matrix = frame.apply(pd.to_numeric, errors="raise").fillna(0.0).T
172 matrix.index.name = "sample_id"
173 return matrix
176def parse_kraken_report(path: Path, rank: str) -> pd.Series:
177 """Parse one Kraken2 report into a lineage -> clade-read-count Series."""
178 import pandas as pd
180 prefix = RANK_PREFIXES[rank]
181 stack: list[tuple[int, str]] = [] # (indent depth, prefixed name)
182 counts: dict[str, float] = {}
184 with open(path, encoding="utf-8") as handle:
185 for line in handle:
186 if not line.strip() or line.startswith("#"):
187 continue
188 fields = line.rstrip("\n").split("\t")
189 if len(fields) < 6:
190 continue
192 # Locate the rank-code column: layouts differ between plain reports
193 # (6 columns) and --report-minimizer-data (8 columns).
194 code_idx = next(
195 (
196 i
197 for i, f in enumerate(fields[:-2])
198 if _RANK_CODE_RE.match(f.strip())
199 ),
200 None,
201 )
202 if code_idx is None:
203 continue
205 code = fields[code_idx].strip()
206 name_field = "\t".join(fields[code_idx + 2 :])
207 name = name_field.strip()
208 if not name or name.lower() in _UNASSIGNED:
209 continue
211 depth = (len(name_field) - len(name_field.lstrip(" "))) // 2
212 try:
213 clade_reads = float(fields[1])
214 except ValueError:
215 continue
217 stack = [entry for entry in stack if entry[0] < depth]
219 base_code = code[0]
220 if len(code) > 1 or base_code not in KRAKEN_RANK_CODES:
221 continue # sub-rank (D1, S1, ...) or U/R: keeps depth, no token
222 stack.append((depth, KRAKEN_RANK_CODES[base_code] + name))
224 if KRAKEN_RANK_CODES[base_code] == prefix:
225 lineage = "; ".join(entry[1] for entry in stack)
226 counts[lineage] = counts.get(lineage, 0.0) + clade_reads
228 return pd.Series(counts, dtype=float)
231def parse_kraken(paths: list[Path], rank: str) -> pd.DataFrame:
232 """Parse many Kraken2 reports into a samples x lineage matrix."""
233 import pandas as pd
235 per_sample = {}
236 for path in paths:
237 if path.stem in per_sample:
238 raise ValueError(f"duplicate sample ID from report filenames: {path.stem}")
239 series = parse_kraken_report(path, rank)
240 if series.empty:
241 print(f"warning: {path} has no rows at rank {rank!r}", file=sys.stderr)
242 per_sample[path.stem] = series
243 matrix = pd.DataFrame(per_sample).T.fillna(0.0)
244 matrix.index.name = "sample_id"
245 return matrix
248def parse_table(
249 path: Path,
250 *,
251 taxonomy_column: str | None,
252 orientation: str,
253) -> pd.DataFrame:
254 """Parse a QIIME 2 / biom / generic delimited table into samples x lineage."""
255 import pandas as pd
257 sep = "," if path.suffix.lower() == ".csv" else "\t"
258 with open(path, encoding="utf-8") as handle:
259 first = handle.readline()
260 skiprows = 1 if first.startswith("# Constructed from biom file") else 0
262 frame = pd.read_csv(path, sep=sep, skiprows=skiprows)
263 frame.columns = [str(c).lstrip("#").strip() for c in frame.columns]
265 if taxonomy_column and taxonomy_column not in frame.columns:
266 raise ValueError(f"taxonomy column {taxonomy_column!r} not found in {path}")
267 if taxonomy_column:
268 lineages = frame[taxonomy_column].astype(str).map(normalise_lineage)
269 feature_col = frame.columns[0]
270 drop = {taxonomy_column, feature_col}
271 values = frame.drop(columns=[c for c in frame.columns if c in drop])
272 values = values.apply(pd.to_numeric, errors="raise").fillna(0.0)
273 values.index = lineages
274 matrix = values.T
275 else:
276 label_col = frame.columns[0]
277 indexed = frame.set_index(label_col)
278 numeric = indexed.apply(pd.to_numeric, errors="raise").fillna(0.0)
279 taxa_as_rows = orientation == "taxa_as_rows" or (
280 orientation == "auto"
281 and str(label_col).lower()
282 in {"taxonomy", "lineage", "taxon", "otu", "otu id", "feature id"}
283 )
284 if taxa_as_rows:
285 numeric.index = [normalise_lineage(i) for i in numeric.index]
286 matrix = numeric.T
287 else:
288 numeric.columns = [normalise_lineage(c) for c in numeric.columns]
289 matrix = numeric
291 # A boolean mask preserves each duplicate lineage once for later summation.
292 keep = [bool(str(c).strip()) for c in matrix.columns]
293 matrix = matrix.loc[:, keep]
294 matrix.index = matrix.index.astype(str)
295 matrix.index.name = "sample_id"
296 return matrix
299# ---------------------------------------------------------------------------
300# Matrix -> waypoint format
301# ---------------------------------------------------------------------------
304def matrix_to_waypoint(
305 matrix: pd.DataFrame,
306 *,
307 normalize: bool = True,
308 drop_zeros: bool = True,
309 min_abundance: float = 0.0,
310) -> pd.DataFrame:
311 """Convert a samples x lineage matrix into waypoint format.
313 Duplicate lineage columns are summed first: MetaPhlAn and Kraken can both
314 produce the same normalised lineage from different rows.
315 """
316 import pandas as pd
318 if matrix.empty:
319 raise ValueError("abundance matrix is empty")
320 if not matrix.index.is_unique:
321 raise ValueError("sample IDs must be unique")
322 if not math.isfinite(min_abundance) or min_abundance < 0:
323 raise ValueError("min_abundance must be finite and nonnegative")
325 values = matrix.astype(float)
326 if any(not math.isfinite(v) or v < 0 for v in values.to_numpy().flat):
327 raise ValueError("abundances must be finite and nonnegative")
328 if values.columns.duplicated().any():
329 values = values.T.groupby(level=0).sum().T
331 if normalize:
332 totals = values.sum(axis=1)
333 empty = totals == 0
334 if empty.any():
335 names = ", ".join(map(str, values.index[empty][:5]))
336 raise ValueError(f"samples with zero total abundance: {names}")
337 values = values.div(totals, axis=0)
339 taxa: list[list[str]] = []
340 abundances: list[list[float]] = []
341 columns = list(values.columns)
342 for _, row in values.iterrows():
343 pairs = [
344 (str(col), float(val))
345 for col, val in zip(columns, row.to_numpy())
346 if (not drop_zeros or val > 0) and val >= min_abundance
347 ]
348 taxa.append([t for t, _ in pairs])
349 abundances.append([a for _, a in pairs])
351 out = pd.DataFrame(
352 {"Taxa": taxa, "Relative Abundances": abundances},
353 index=values.index,
354 )
355 out.index.name = values.index.name or "sample_id"
356 return out
359def attach_metadata(frame: pd.DataFrame, metadata_path: Path) -> pd.DataFrame:
360 """Join per-sample metadata, indexed by sample ID, onto a waypoint frame."""
361 import pandas as pd
363 suffix = metadata_path.suffix.lower()
364 if suffix == ".parquet":
365 meta = pd.read_parquet(metadata_path)
366 if isinstance(meta.index, pd.RangeIndex) and meta.index.name is None:
367 raise ValueError("parquet metadata must have a sample-ID index")
368 else:
369 sep = "\t" if suffix in {".tsv", ".tab"} else ","
370 id_column = pd.read_csv(metadata_path, sep=sep, nrows=0).columns[0]
371 meta = pd.read_csv(metadata_path, sep=sep, index_col=0, dtype={id_column: str})
372 meta.index = meta.index.astype(str)
373 if not meta.index.is_unique:
374 raise ValueError("metadata sample IDs must be unique")
376 joined = frame.join(meta, how="left")
377 missing = int((~frame.index.isin(meta.index)).sum())
378 if missing:
379 print(
380 f"warning: {missing}/{len(joined)} samples had no metadata match; "
381 "check that sample IDs agree",
382 file=sys.stderr,
383 )
384 return joined
387# ---------------------------------------------------------------------------
388# CLI
389# ---------------------------------------------------------------------------
392def build_matrix(args: argparse.Namespace, paths: list[Path]) -> pd.DataFrame:
393 import pandas as pd
395 if args.format == "metaphlan":
396 frames = [parse_metaphlan(p, args.rank) for p in paths]
397 return pd.concat(frames) if len(frames) > 1 else frames[0]
398 if args.format == "kraken":
399 return parse_kraken(paths, args.rank)
400 # qiime2 and generic share the delimited-table reader.
401 frames = [
402 parse_table(
403 p,
404 taxonomy_column=args.taxonomy_column,
405 orientation=args.orientation,
406 )
407 for p in paths
408 ]
409 return pd.concat(frames) if len(frames) > 1 else frames[0]
412def main(argv: list[str] | None = None) -> int:
413 parser = argparse.ArgumentParser(
414 description="Convert profiler output into waypoint format.",
415 formatter_class=argparse.ArgumentDefaultsHelpFormatter,
416 )
417 parser.add_argument(
418 "--input",
419 nargs="+",
420 required=True,
421 help="Input file(s). Kraken takes one report per sample; the others take one table.",
422 )
423 parser.add_argument(
424 "--format",
425 required=True,
426 choices=["metaphlan", "kraken", "qiime2", "generic"],
427 help="Input layout.",
428 )
429 parser.add_argument(
430 "--output",
431 required=True,
432 help="Output path (.parquet recommended; .csv and .tsv also supported).",
433 )
434 parser.add_argument(
435 "--rank",
436 default="species",
437 choices=list(RANK_PREFIXES),
438 help="Deepest rank to keep. Used by metaphlan and kraken to pick leaf rows.",
439 )
440 parser.add_argument(
441 "--taxonomy-column",
442 default=None,
443 help="For qiime2/generic: column holding the lineage string (e.g. 'taxonomy').",
444 )
445 parser.add_argument(
446 "--orientation",
447 default="auto",
448 choices=["auto", "samples_as_rows", "taxa_as_rows"],
449 help="For qiime2/generic tables without a taxonomy column.",
450 )
451 parser.add_argument(
452 "--metadata",
453 default=None,
454 help="CSV/TSV/parquet of per-sample metadata, indexed by sample ID, merged as extra columns.",
455 )
456 parser.add_argument(
457 "--min-abundance",
458 type=float,
459 default=0.0,
460 help="Drop taxa below this relative abundance (Atlas used 1e-4).",
461 )
462 parser.add_argument(
463 "--no-normalize",
464 action="store_true",
465 help="Skip row-normalisation (use when values are already relative abundances summing to 1).",
466 )
467 parser.add_argument(
468 "--keep-zeros",
469 action="store_true",
470 help="Keep zero-abundance entries in each sample's lists.",
471 )
472 args = parser.parse_args(argv)
474 paths = [Path(p) for p in args.input]
475 for path in paths:
476 if not path.exists():
477 parser.error(f"input not found: {path}")
478 if args.format != "kraken" and len(paths) > 1:
479 print(
480 f"note: concatenating {len(paths)} {args.format} tables by sample",
481 file=sys.stderr,
482 )
484 matrix = build_matrix(args, paths)
485 frame = matrix_to_waypoint(
486 matrix,
487 normalize=not args.no_normalize,
488 drop_zeros=not args.keep_zeros,
489 min_abundance=args.min_abundance,
490 )
492 if args.metadata:
493 frame = attach_metadata(frame, Path(args.metadata))
495 out_path = Path(args.output)
496 out_path.parent.mkdir(parents=True, exist_ok=True)
497 suffix = out_path.suffix.lower()
498 if suffix == ".parquet":
499 frame.to_parquet(out_path)
500 elif suffix in {".csv", ".tsv", ".tab"}:
501 frame.to_csv(out_path, sep="\t" if suffix in {".tsv", ".tab"} else ",")
502 else:
503 parser.error(f"unsupported output format: {suffix!r}")
505 n_taxa = [len(t) for t in frame["Taxa"]]
506 print(f"Wrote {len(frame)} samples to {out_path}")
507 print(
508 f"Taxa per sample: min {min(n_taxa)}, median {sorted(n_taxa)[len(n_taxa) // 2]}, "
509 f"max {max(n_taxa)}"
510 )
511 if min(n_taxa) < 10:
512 print(
513 "warning: some samples have fewer than 10 taxa; Atlas filtered these out",
514 file=sys.stderr,
515 )
516 return 0
519if __name__ == "__main__":
520 raise SystemExit(main())