Skills › Scientific Agent Skills › onekgpd › onekgpd_meta.py
onekgpd_meta.py
485 lines · skills/onekgpd/scripts/onekgpd_meta.py at 5f2bc03
1# /// script
2# requires-python = ">=3.11"
3# dependencies = []
4# ///
5"""OneKGPd — sample & population metadata (offline) over the 1000 Genomes Project.
7Six commands answering population/pedigree questions from a data file bundled in
8the skill (``onekgpd/assets/kgpe.json``): no network, no credentials, and no
9third-party dependencies. The sample identifier is the same name used by the
10variant/kinship commands (e.g. ``NA19240``), so the two layers compose.
12Every command writes full JSON to a file (``--output``, or a temp file by
13default) and prints a concise summary to stdout.
15Examples
16--------
17 uv run scripts/onekgpd_meta.py list-superpopulations
18 uv run scripts/onekgpd_meta.py sample-metadata --samples NA19240,HG00096
19 uv run scripts/onekgpd_meta.py population-stats --populations YRI --populations CHS
20 uv run scripts/onekgpd_meta.py select-samples-by-population --population YRI --limit 20
22MIT License. Author: Dnaerys.
23"""
25from __future__ import annotations
27import argparse
28import gzip
29import json
30import os
31import sys
32import tempfile
33from pathlib import Path
34from typing import Any, NoReturn
36# ---------------------------------------------------------------------------
37# Constants
38# ---------------------------------------------------------------------------
40_ASSETS = Path(__file__).resolve().parent.parent / "assets"
41# Prefer the plain-text asset (inspectable, scanner-clean); fall back to a
42# gzipped copy if that is the only form present. Both hold identical JSON.
43DATA_PATH = _ASSETS / "kgpe.json"
44if not DATA_PATH.exists():
45 DATA_PATH = _ASSETS / "kgpe.json.gz"
46DEFAULT_LIMIT = 50 # MetaClient.java:48
47MAX_LIMIT = 3202 # MetaClient.java:49
48PREVIEW_ROWS = 10 # rows shown in a stdout summary preview
50_ABSENT = (None, "", "0") # pid/mid "absent" sentinel values
53# ---------------------------------------------------------------------------
54# Small generic helpers (mirrors onekgpd_api.py; kept local to stay dnaerys-free)
55# ---------------------------------------------------------------------------
58def _fail(msg: str) -> NoReturn:
59 """Print a clean one-line message to stderr and exit non-zero."""
60 print(msg, file=sys.stderr)
61 sys.exit(1)
64def _split_csv(s: str) -> list[str]:
65 """Split a comma-separated value into a list of non-empty trimmed tokens."""
66 return [part.strip() for part in s.split(",") if part.strip()]
69def _save_json(data: Any, prefix: str, output_path: str | None = None) -> str:
70 """Write *data* as indented JSON to *output_path* or a temp file; return path."""
71 if output_path:
72 path = output_path
73 with open(path, "w") as fh:
74 json.dump(data, fh, indent=2)
75 else:
76 fd, path = tempfile.mkstemp(prefix=f"onekgpd_{prefix}_", suffix=".json", text=True)
77 with os.fdopen(fd, "w") as fh:
78 json.dump(data, fh, indent=2)
79 return path
82def _emit(data: Any, prefix: str, summary_lines: list[str], output: str | None) -> None:
83 """Save full JSON to a file and print the summary + saved-path line to stdout."""
84 path = _save_json(data, prefix, output)
85 for line in summary_lines:
86 print(line)
87 print(f"[*] Full JSON saved to {path}")
90# ---------------------------------------------------------------------------
91# Data loading / indexing
92# ---------------------------------------------------------------------------
94_RECORDS: list[dict] | None = None
97def _load_records() -> list[dict]:
98 """Load and cache the bundled pedigree records (once per process)."""
99 global _RECORDS
100 if _RECORDS is None:
101 opener = gzip.open if DATA_PATH.suffix == ".gz" else open
102 with opener(DATA_PATH, "rt", encoding="utf-8") as fh:
103 _RECORDS = json.load(fh)
104 return _RECORDS
107def _present(x: str | None) -> bool:
108 """True if a pid/mid value names a real parent (not absent/'0'/empty)."""
109 return x not in _ABSENT
112def _none_if_empty(x: str | None) -> str | None:
113 """Java nullIfEmpty: None for empty/None (MetaClient.java:537-539)."""
114 return None if x in (None, "") else x
117def _none_if_absent(x: str | None) -> str | None:
118 """Java nullIfAbsent: None for None/empty/'0' (MetaClient.java:533-535)."""
119 return None if x in _ABSENT else x
122def _children_index(records: list[dict]) -> dict[str, list[str]]:
123 """Map parent externalID -> child externalIDs.
125 Replicates the LEFT JOIN ``(c.pid = s OR c.mid = s) AND c.pid != '0' AND
126 c.mid != '0'`` (MetaClient.java:155-156): a child counts only when BOTH of
127 its parents are recorded.
128 """
129 idx: dict[str, list[str]] = {}
130 for c in records:
131 if _present(c["pid"]) and _present(c["mid"]):
132 cid = c["externalIDs"]
133 idx.setdefault(c["pid"], []).append(cid)
134 idx.setdefault(c["mid"], []).append(cid)
135 return idx
138def _valid_pop_lower(records: list[dict]) -> set[str]:
139 """Lowercased set of all valid population codes and full names."""
140 s: set[str] = set()
141 for r in records:
142 s.add(r["pop"].lower())
143 s.add(r["Population"].lower())
144 return s
147def _valid_reg_lower(records: list[dict]) -> set[str]:
148 """Lowercased set of all valid superpopulation codes and full names."""
149 s: set[str] = set()
150 for r in records:
151 s.add(r["reg"].lower())
152 s.add(r["region"].lower())
153 return s
156# ---------------------------------------------------------------------------
157# Stats aggregation (shared by population-stats and superpopulation-summary)
158# ---------------------------------------------------------------------------
161def _population_stats(records_subset: list[dict]) -> dict[tuple, dict]:
162 """Group a subset by (pop, Population, reg, region) -> count aggregates.
164 Mirrors the COUNT(CASE WHEN …) columns in MetaClient.java:292-298.
165 """
166 groups: dict[tuple, dict] = {}
167 for r in records_subset:
168 key = (r["pop"], r["Population"], r["reg"], r["region"])
169 g = groups.setdefault(key, {"n": 0, "m": 0, "f": 0, "p3": 0, "trio": 0})
170 g["n"] += 1
171 if r["gender"] == "male":
172 g["m"] += 1
173 elif r["gender"] == "female":
174 g["f"] += 1
175 if r["phase3"] == "TRUE":
176 g["p3"] += 1
177 if _present(r["pid"]) and _present(r["mid"]):
178 g["trio"] += 1
179 return groups
182def _stats_obj(key: tuple, g: dict) -> dict:
183 return {
184 "population_code": key[0],
185 "population": key[1],
186 "superpopulation_code": key[2],
187 "superpopulation": key[3],
188 "sample_count": g["n"],
189 "male_count": g["m"],
190 "female_count": g["f"],
191 "phase3_count": g["p3"],
192 "trio_count": g["trio"],
193 }
196# ---------------------------------------------------------------------------
197# Command handlers
198# ---------------------------------------------------------------------------
201def cmd_sample_metadata(args) -> None:
202 ids = _split_csv(args.samples)
203 if not ids:
204 _fail("Error: Parameter 'sampleIds' must not be null or empty")
205 records = _load_records()
206 by_id = {r["externalIDs"]: r for r in records}
207 unknown = [i for i in ids if i not in by_id] # case-sensitive sample IDs
208 if unknown:
209 _fail(f"Error: Unknown sample IDs: [{', '.join(unknown)}]")
211 children = _children_index(records)
212 samples = []
213 for sid in sorted(set(ids)): # ORDER BY externalIDs, distinct
214 r = by_id[sid]
215 kids = sorted(set(children.get(sid, [])))
216 samples.append({
217 "sample_id": r["externalIDs"],
218 "family_id": _none_if_empty(r["familyId"]),
219 "gender": r["gender"],
220 "paternal_id": _none_if_absent(r["pid"]),
221 "maternal_id": _none_if_absent(r["mid"]),
222 "relationship": _none_if_empty(r["Relationship"]),
223 "children": kids,
224 "population_code": r["pop"],
225 "population": r["Population"],
226 "superpopulation_code": r["reg"],
227 "superpopulation": r["region"],
228 "phase3": r["phase3"],
229 })
231 data = {"command": "sample-metadata", "samples": samples}
232 summary = [f"Metadata for {len(samples)} sample(s)"]
233 for s in samples[:PREVIEW_ROWS]:
234 kids = ", ".join(s["children"]) if s["children"] else "-"
235 summary.append(
236 f" {s['sample_id']} {s['population_code']}/{s['superpopulation_code']} "
237 f"{s['gender']} {s['relationship'] or '-'} "
238 f"family {s['family_id'] or '-'} children: {kids}"
239 )
240 _emit(data, "sample_metadata", summary, args.output)
243def cmd_list_populations(args) -> None:
244 records = _load_records()
245 counts: dict[tuple, int] = {}
246 for r in records:
247 key = (r["pop"], r["Population"], r["reg"], r["region"])
248 counts[key] = counts.get(key, 0) + 1
249 # ORDER BY reg, pop (MetaClient.java:220)
250 rows = sorted(counts.items(), key=lambda kv: (kv[0][2], kv[0][0]))
251 populations = [{
252 "population_code": k[0],
253 "population": k[1],
254 "superpopulation_code": k[2],
255 "superpopulation": k[3],
256 "sample_count": cnt,
257 } for k, cnt in rows]
259 data = {"command": "list-populations", "populations": populations}
260 n_super = len({k[2] for k in counts})
261 summary = [f"{len(populations)} populations across {n_super} superpopulations"]
262 for p in populations[:PREVIEW_ROWS]:
263 summary.append(
264 f" {p['population_code']} {p['population']} "
265 f"{p['superpopulation_code']} {p['sample_count']}"
266 )
267 if len(populations) > PREVIEW_ROWS:
268 summary.append(f" ... {len(populations) - PREVIEW_ROWS} more (see file)")
269 _emit(data, "list_populations", summary, args.output)
272def cmd_list_superpopulations(args) -> None:
273 records = _load_records()
274 groups: dict[str, dict] = {}
275 for r in records:
276 reg = r["reg"]
277 g = groups.setdefault(reg, {"region": r["region"], "count": 0, "pops": set()})
278 g["count"] += 1
279 g["pops"].add(r["pop"])
280 # ORDER BY reg (MetaClient.java:250)
281 rows = sorted(groups.items(), key=lambda kv: kv[0])
282 superpopulations = [{
283 "superpopulation_code": reg,
284 "superpopulation": g["region"],
285 "sample_count": g["count"],
286 "populations": sorted(g["pops"]),
287 } for reg, g in rows]
289 data = {"command": "list-superpopulations", "superpopulations": superpopulations}
290 summary = [f"{len(superpopulations)} superpopulations"]
291 for sp in superpopulations:
292 summary.append(
293 f" {sp['superpopulation_code']} {sp['superpopulation']} "
294 f"n={sp['sample_count']} pops={len(sp['populations'])}"
295 )
296 _emit(data, "list_superpopulations", summary, args.output)
299def cmd_population_stats(args) -> None:
300 vals = [v.strip() for v in args.populations if v.strip()]
301 if not vals:
302 _fail("Error: Parameter 'populations' must not be null or empty")
303 records = _load_records()
304 valid = _valid_pop_lower(records)
305 unknown = [v for v in vals if v.lower() not in valid]
306 if unknown:
307 _fail(f"Error: Unrecognised population values: [{', '.join(unknown)}]")
309 wanted = {v.lower() for v in vals}
310 subset = [r for r in records
311 if r["pop"].lower() in wanted or r["Population"].lower() in wanted]
312 groups = _population_stats(subset)
313 # ORDER BY pop (MetaClient.java:302)
314 rows = sorted(groups.items(), key=lambda kv: kv[0][0])
315 populations = [_stats_obj(k, g) for k, g in rows]
317 data = {"command": "population-stats", "populations": populations}
318 summary = [f"Stats for {len(populations)} population(s)"]
319 for p in populations:
320 summary.append(
321 f" {p['population_code']} {p['population']} n={p['sample_count']} "
322 f"M={p['male_count']} F={p['female_count']} "
323 f"phase3={p['phase3_count']} trios={p['trio_count']}"
324 )
325 _emit(data, "population_stats", summary, args.output)
328def cmd_superpopulation_summary(args) -> None:
329 vals = [v.strip() for v in args.superpopulations if v.strip()]
330 if not vals:
331 _fail("Error: Parameter 'superpopulations' must not be null or empty")
332 records = _load_records()
333 valid = _valid_reg_lower(records)
334 unknown = [v for v in vals if v.lower() not in valid]
335 if unknown:
336 _fail(f"Error: Unrecognised superpopulation values: [{', '.join(unknown)}]")
338 wanted = {v.lower() for v in vals}
339 subset = [r for r in records
340 if r["reg"].lower() in wanted or r["region"].lower() in wanted]
341 groups = _population_stats(subset)
342 # Per-population rows ordered by (reg, pop); group by reg preserving order.
343 pop_rows = sorted(groups.items(), key=lambda kv: (kv[0][2], kv[0][0]))
344 by_super: dict[str, dict] = {}
345 for key, g in pop_rows:
346 reg, region = key[2], key[3]
347 sg = by_super.setdefault(reg, {"region": region, "pops": []})
348 sg["pops"].append(_stats_obj(key, g))
350 superpopulations = []
351 for reg, sg in by_super.items():
352 pops = sg["pops"]
353 superpopulations.append({
354 "superpopulation_code": reg,
355 "superpopulation": sg["region"],
356 "sample_count": sum(p["sample_count"] for p in pops),
357 "male_count": sum(p["male_count"] for p in pops),
358 "female_count": sum(p["female_count"] for p in pops),
359 "phase3_count": sum(p["phase3_count"] for p in pops),
360 "trio_count": sum(p["trio_count"] for p in pops),
361 "populations": pops,
362 })
364 data = {"command": "superpopulation-summary", "superpopulations": superpopulations}
365 summary = [f"Summary for {len(superpopulations)} superpopulation(s)"]
366 for sp in superpopulations:
367 summary.append(
368 f" {sp['superpopulation_code']} {sp['superpopulation']} "
369 f"n={sp['sample_count']} M={sp['male_count']} F={sp['female_count']} "
370 f"phase3={sp['phase3_count']} trios={sp['trio_count']} "
371 f"({len(sp['populations'])} populations)"
372 )
373 _emit(data, "superpopulation_summary", summary, args.output)
376def cmd_select_samples_by_population(args) -> None:
377 pop = args.population.strip() if args.population and args.population.strip() else None
378 sup = args.superpopulation.strip() if args.superpopulation and args.superpopulation.strip() else None
379 if pop is None and sup is None:
380 _fail("Error: At least one parameter ('population' or 'superpopulation') must be provided")
382 skip = args.skip if args.skip is not None else 0
383 limit = args.limit if args.limit is not None else DEFAULT_LIMIT
384 if skip < 0:
385 _fail(f"Error: Invalid parameter: 'skip' must be >= 0, actual: {skip}")
386 if limit < 1 or limit > MAX_LIMIT:
387 _fail(f"Error: Invalid parameter: 'limit' must be between 1 and {MAX_LIMIT}, actual: {limit}")
389 records = _load_records()
390 if pop is not None and pop.lower() not in _valid_pop_lower(records):
391 _fail(f"Error: Unrecognised population: '{pop}'")
392 if sup is not None and sup.lower() not in _valid_reg_lower(records):
393 _fail(f"Error: Unrecognised superpopulation: '{sup}'")
395 pop_l = pop.lower() if pop is not None else None
396 sup_l = sup.lower() if sup is not None else None
398 def match(r: dict) -> bool:
399 if pop_l is not None and pop_l not in (r["pop"].lower(), r["Population"].lower()):
400 return False
401 if sup_l is not None and sup_l not in (r["reg"].lower(), r["region"].lower()):
402 return False
403 return True
405 matched = sorted(r["externalIDs"] for r in records if match(r)) # ORDER BY externalIDs
406 page = matched[skip:skip + limit]
408 data = {
409 "command": "select-samples-by-population",
410 "count": len(page),
411 "samples": page,
412 "request": {"population": pop, "superpopulation": sup, "skip": skip, "limit": limit},
413 }
414 parts = []
415 if pop is not None:
416 parts.append(f"population {pop}")
417 if sup is not None:
418 parts.append(f"superpopulation {sup}")
419 label = " & ".join(parts)
420 summary = [f"{len(page)} samples in {label} (rows {skip}..{skip + len(page)} of {len(matched)} total)"]
421 for s in page[:PREVIEW_ROWS]:
422 summary.append(f" {s}")
423 _emit(data, "select_samples_by_population", summary, args.output)
426# ---------------------------------------------------------------------------
427# Argument parsing
428# ---------------------------------------------------------------------------
431def build_parser() -> argparse.ArgumentParser:
432 conn = argparse.ArgumentParser(add_help=False)
433 conn.add_argument(
434 "--output",
435 help="Write full JSON to this path (default: a temp file in the system temp dir).",
436 )
438 parser = argparse.ArgumentParser(
439 prog="onekgpd_meta.py",
440 description="Sample & population metadata for the 1000 Genomes Project "
441 "(offline; reads a bundled data file, no network).",
442 )
443 sub = parser.add_subparsers(dest="command", required=True)
445 p = sub.add_parser("sample-metadata", parents=[conn], help="Pedigree/population metadata for given sample IDs.")
446 p.add_argument("--samples", required=True, metavar="CSV", help="Comma-separated sample IDs (case-sensitive).")
447 p.set_defaults(func=cmd_sample_metadata)
449 p = sub.add_parser("list-populations", parents=[conn], help="List all populations with superpopulation and sample count.")
450 p.set_defaults(func=cmd_list_populations)
452 p = sub.add_parser("list-superpopulations", parents=[conn], help="List all superpopulations with sample count and constituent populations.")
453 p.set_defaults(func=cmd_list_superpopulations)
455 p = sub.add_parser("population-stats", parents=[conn], help="Per-population stats: sex split, phase3, trio membership.")
456 p.add_argument("--populations", required=True, action="append", metavar="VALUE",
457 help="Population code or full name; repeat for multiple (case-insensitive). "
458 "Repeated rather than comma-separated because full names contain commas.")
459 p.set_defaults(func=cmd_population_stats)
461 p = sub.add_parser("superpopulation-summary", parents=[conn], help="Per-superpopulation summary with per-population breakdown.")
462 p.add_argument("--superpopulations", required=True, action="append", metavar="VALUE",
463 help="Superpopulation code or full name; repeat for multiple (case-insensitive).")
464 p.set_defaults(func=cmd_superpopulation_summary)
466 p = sub.add_parser("select-samples-by-population", parents=[conn], help="Select sample IDs by population and/or superpopulation.")
467 p.add_argument("--population", metavar="P", help="Population code or full name (case-insensitive).")
468 p.add_argument("--superpopulation", metavar="R", help="Superpopulation code or full name (case-insensitive).")
469 p.add_argument("--skip", type=int, help="Number of results to skip (default 0).")
470 p.add_argument("--limit", type=int, help=f"Max results to return (default {DEFAULT_LIMIT}, max {MAX_LIMIT}).")
471 p.set_defaults(func=cmd_select_samples_by_population)
473 return parser
476def main(argv: list[str] | None = None) -> None:
477 args = build_parser().parse_args(argv)
478 try:
479 args.func(args)
480 except (ValueError, OSError) as e:
481 _fail(f"Error: {e}")
484if __name__ == "__main__":
485 main()