Skills › Scientific Agent Skills › pathogen-variant-surveillance › lineage_prevalence.py
lineage_prevalence.py
333 lines · skills/pathogen-variant-surveillance/scripts/lineage_prevalence.py at 5f2bc03
1#!/usr/bin/env python3
2# /// script
3# requires-python = ">=3.11"
4# dependencies = []
5# ///
6"""Weekly prevalence of one or more lineages, with intervals and a coverage flag.
8Answers "what is circulating, and is it growing" from live sequence counts
9rather than from memory. Every number is a count returned by the instance at
10the data version printed in the provenance block.
12 python3 lineage_prevalence.py XFG.1.1 XFJ.3 --instance sars-cov-2 --where country=USA
13 python3 lineage_prevalence.py XFG --sublineages --weeks 12 --growth
14 python3 lineage_prevalence.py 2.3.4.4b --instance h5n1 --where country=USA --weeks 52
16Proportions carry Wilson intervals. Low-count/age flags are heuristics, not
17measured completeness; reporting_lag.py separately describes observed delays.
18"""
19from __future__ import annotations
21import argparse
22import sys
23from datetime import date, timedelta
25from lapis_client import (
26 LapisError,
27 aggregated,
28 bin_weekly,
29 count,
30 data_version,
31 describe_instance,
32 emit,
33 flag_low_coverage,
34 lineage_filter,
35 logit_slope,
36 pick_date_field,
37 pick_lineage_field,
38 range_keys,
39 resolve_base_url,
40 surveillance_filters,
41 top_values,
42 week_range,
43 wilson_interval,
44)
46COLUMNS = (
47 "lineage",
48 "week",
49 "n",
50 "total",
51 "proportion",
52 "ci_low",
53 "ci_high",
54 "coverage",
55)
58def parse_where(pairs: list[str], schema: dict) -> dict[str, str]:
59 """Turn ``KEY=VALUE`` arguments into filters, checking keys against the schema."""
60 filters: dict[str, str] = {}
61 types = schema.get("types", {})
62 for pair in pairs:
63 key, sep, value = pair.partition("=")
64 if not sep:
65 raise LapisError(f"--where expects KEY=VALUE, got {pair!r}")
66 base = key.split(".")[0]
67 for suffix in ("From", "To"):
68 if base.endswith(suffix) and base[: -len(suffix)] in types:
69 base = base[: -len(suffix)]
70 break
71 if base not in types:
72 near = sorted(n for n in types if base.lower() in n.lower())
73 hint = f" Did you mean: {', '.join(near[:6])}?" if near else ""
74 raise LapisError(f"{key!r} is not a field on this instance.{hint}")
75 filters[key] = value
76 return filters
79def weekly_counts(
80 base_url: str, filters: dict, date_field: str, weeks: list[str], upper_field: str | None = None
81) -> tuple[dict[str, int], int]:
82 """Weekly counts over the requested window, zero-filled for empty weeks."""
83 rows = aggregated(base_url, filters, [date_field] + ([upper_field] if upper_field else []))
84 excluded = sum(int(r.get("count") or 0) for r in rows if upper_field and r.get(upper_field) != r.get(date_field))
85 rows = [r for r in rows if not upper_field or r.get(upper_field) == r.get(date_field)]
86 binned, undated = bin_weekly(rows, date_field)
87 return {week: binned.get(week, 0) for week in weeks}, undated + excluded
90def build_parser() -> argparse.ArgumentParser:
91 parser = argparse.ArgumentParser(
92 description="Weekly lineage prevalence from a LAPIS instance.",
93 formatter_class=argparse.RawDescriptionHelpFormatter,
94 )
95 parser.add_argument("lineages", nargs="*",
96 help="lineage or clade names; 'NAME*' includes descendants. "
97 "Omit to discover them with --top")
98 parser.add_argument("--top", type=int, metavar="N",
99 help="discover the N most common lineages in the window instead of "
100 "naming them (default when no names are given)")
101 parser.add_argument("--instance", default="sars-cov-2", help="registry name (default: sars-cov-2)")
102 parser.add_argument("--base-url", help="any other LAPIS deployment")
103 parser.add_argument("--lineage-field", help="override the auto-detected lineage column")
104 parser.add_argument("--date-field", help="override the auto-detected collection-date column")
105 parser.add_argument("--where", action="append", default=[], metavar="KEY=VALUE",
106 help="extra filter, repeatable (e.g. --where country=USA)")
107 parser.add_argument("--weeks", type=int, default=26, help="window length in weeks (default: 26)")
108 parser.add_argument("--since", help="window start, YYYY-MM-DD; overrides --weeks")
109 parser.add_argument("--until", help="window end, YYYY-MM-DD (default: today)")
110 parser.add_argument("--sublineages", action="store_true", help="include descendant lineages")
111 parser.add_argument("--growth", action="store_true",
112 help="also fit a descriptive weighted log-odds slope")
113 parser.add_argument("--lag-days", type=int, default=0,
114 help="flag weeks ending within this many days of today (set from lag review)")
115 parser.add_argument("--include-incomplete", action="store_true",
116 help="let low-coverage weeks into the growth fit")
117 parser.add_argument("--coverage-fraction", type=float, default=0.4,
118 help="flag weeks below this fraction of older-half median (default: 0.4)")
119 parser.add_argument("--format", choices=("table", "tsv", "json"), default="table")
120 parser.add_argument("-o", "--output", help="write to a file instead of stdout")
121 return parser
124def main(argv: list[str] | None = None) -> int:
125 args = build_parser().parse_args(argv)
127 try:
128 if args.weeks < 1 or args.lag_days < 0 or not 0 <= args.coverage_fraction <= 1:
129 raise LapisError("--weeks must be positive, --lag-days nonnegative, --coverage-fraction in [0, 1]")
130 base_url = resolve_base_url(args.instance, args.base_url)
131 schema = describe_instance(base_url)
132 lineage_field, has_index = pick_lineage_field(schema, args.lineage_field)
133 date_field = pick_date_field(schema, "collection", args.date_field)
134 where = surveillance_filters(schema, parse_where(args.where, schema))
135 reserved = {date_field, *range_keys(date_field), lineage_field}
136 if any(k.split('.')[0] in reserved for k in where):
137 raise LapisError("use --since/--until and positional lineages instead of --where date/lineage overrides")
138 upper_field = date_field.replace("RangeLower", "RangeUpper") if date_field.endswith("RangeLower") else None
139 if upper_field and upper_field not in schema["types"]:
140 raise LapisError("collection range lacks an upper bound; cannot identify exact dates")
141 requested_until = date.fromisoformat(args.until) if args.until else date.today()
142 requested_since = date.fromisoformat(args.since) if args.since else requested_until - timedelta(weeks=args.weeks)
143 except (LapisError, ValueError) as exc:
144 print(f"error: {exc}", file=sys.stderr)
145 return 2
147 if requested_since > requested_until:
148 print("error: --since is after --until", file=sys.stderr)
149 return 2
150 if args.top is not None and args.top < 1:
151 print("error: --top must be at least 1", file=sys.stderr)
152 return 2
154 # Snap both ends to ISO week boundaries. A window that starts mid-week gives
155 # a first row covering three days and a last row covering four, and neither
156 # is comparable to the full weeks between them -- exactly the kind of
157 # silently uneven denominator this skill exists to prevent. Widening to whole
158 # weeks keeps every row meaning the same thing; a still-unfilled final week
159 # shows up through the coverage flag rather than as a short bar.
160 since = requested_since - timedelta(days=requested_since.weekday())
161 until = requested_until + timedelta(days=6 - requested_until.weekday())
162 snapped = (since, until) != (requested_since, requested_until)
164 from_key, to_key = range_keys(date_field)
165 window = {from_key: since.isoformat(), to_key: until.isoformat(), **where}
166 weeks = week_range(since.isoformat(), until.isoformat())
168 try:
169 version = data_version(base_url)
170 totals, undated = weekly_counts(base_url, window, date_field, weeks, upper_field)
172 rows: list[dict] = []
173 notes: list[str] = []
174 fits: dict[str, dict] = {}
175 coverage = flag_low_coverage(totals, args.coverage_fraction)
176 recent_boundary = date.today() - timedelta(days=args.lag_days)
177 for week in weeks:
178 if date.fromisoformat(week) + timedelta(days=6) >= recent_boundary:
179 coverage[week] = True
180 notes.append("coverage is a count/age heuristic, not measured completeness; choose --lag-days after reviewing reporting lag")
182 targets = list(args.lineages)
183 if args.top is not None or not targets:
184 wanted = args.top or 10
185 discovery_fields = [lineage_field] + ([date_field, upper_field] if upper_field else [])
186 grouped = aggregated(base_url, window, discovery_fields)
187 if upper_field:
188 # Rank the same exact-date subset used by weekly denominators.
189 by_label: dict[str, int] = {}
190 for row in grouped:
191 if row.get(date_field) and row.get(date_field) == row.get(upper_field):
192 label = row.get(lineage_field)
193 by_label[label] = by_label.get(label, 0) + int(row.get("count") or 0)
194 grouped = [{lineage_field: label, "count": n} for label, n in by_label.items()]
195 discovered = top_values(grouped, lineage_field, wanted)
196 if not discovered:
197 print(
198 f"error: no sequences in {since} to {until}"
199 + (f" with {where}" if where else "")
200 + ". Widen the window or relax the filters.",
201 file=sys.stderr,
202 )
203 return 1
204 notes.append(
205 f"discovered the {len(discovered)} most common {lineage_field} "
206 f"value{'s' if len(discovered) != 1 else ''} in the window: "
207 f"{', '.join(discovered)}"
208 )
209 targets = discovered + [t for t in targets if t not in discovered]
211 for name in targets:
212 value = lineage_filter(name, has_index, args.sublineages)
213 counts, _ = weekly_counts(
214 base_url, {**window, lineage_field: value}, date_field, weeks, upper_field
215 )
217 if has_index and not value.endswith("*"):
218 # Both sides must come from the same kind of query. Summing the
219 # weekly bins would drop sequences with no usable collection
220 # date, making every lineage that has some look as though it had
221 # descendants it does not.
222 exact = count(base_url, {**window, lineage_field: value})
223 inclusive = count(base_url, {**window, lineage_field: f"{value}*"})
224 if inclusive > exact:
225 notes.append(
226 f"{value}: {exact} sequences named exactly {value}, "
227 f"{inclusive} including descendants. Pass --sublineages "
228 f"or write '{value}*' for the second number."
229 )
231 for week in weeks:
232 k, n = counts[week], totals[week]
233 if k > n:
234 raise LapisError("lineage count exceeds denominator; inconsistent filters or snapshot")
235 low, high = wilson_interval(k, n)
236 rows.append(
237 {
238 "lineage": value,
239 "week": week,
240 "n": k,
241 "total": n,
242 "proportion": f"{(k / n):.4f}" if n else "",
243 "ci_low": f"{low:.4f}" if n else "",
244 "ci_high": f"{high:.4f}" if n else "",
245 "coverage": "low" if coverage.get(week) else "ok",
246 }
247 )
249 if args.growth:
250 usable = [
251 (i, counts[w], totals[w])
252 for i, w in enumerate(weeks)
253 if args.include_incomplete or not coverage.get(w)
254 ]
255 fit = logit_slope([(float(t), k, n) for t, k, n in usable])
256 if fit:
257 fits[value] = fit
258 else:
259 notes.append(
260 f"{value}: too few observations for a growth estimate "
261 f"({sum(k for _, k, _ in usable)} sequences across "
262 f"{sum(1 for _, k, _ in usable if k > 0)} non-empty weeks). "
263 "No slope is reported rather than a slope driven by the denominator."
264 )
265 except LapisError as exc:
266 print(f"error: {exc}", file=sys.stderr)
267 return 2
269 if args.format == "json":
270 import json
272 payload = {
273 "meta": {
274 "instance": schema["name"],
275 "base_url": base_url,
276 "data_version": version,
277 "lineage_field": lineage_field,
278 "lineage_index": has_index,
279 "date_field": date_field,
280 "window": [since.isoformat(), until.isoformat()],
281 "filters": where,
282 "undated_sequences": undated,
283 "lag_days": args.lag_days,
284 "notes": notes,
285 },
286 "rows": rows,
287 "growth": fits,
288 }
289 text = json.dumps(payload, indent=2)
290 else:
291 text = emit(rows, COLUMNS, args.format)
293 if args.output:
294 with open(args.output, "w", encoding="utf-8") as handle:
295 handle.write(text + "\n")
296 else:
297 print(text)
299 sys.stdout.flush()
300 print(
301 f"\n# {schema['name']} via {base_url}"
302 f"\n# data version {version} | {lineage_field}"
303 f"{' (lineage-indexed)' if has_index else ' (no lineage index)'}"
304 f" | dates from {date_field}"
305 f"\n# window {since} to {until}"
306 + (f" (widened from {requested_since}..{requested_until} to whole ISO weeks)"
307 if snapped else "")
308 + (f" | filters {where}" if where else "")
309 + (f"\n# {undated} returned records excluded for missing/imprecise collection date" if undated else ""),
310 file=sys.stderr,
311 )
312 for note in notes:
313 print(f"# note: {note}", file=sys.stderr)
314 if any(coverage.values()):
315 flagged = [w for w, bad in coverage.items() if bad]
316 print(
317 f"# {len(flagged)} week(s) flagged low by denominator/age heuristic "
318 f"(earliest {min(flagged)}). Other weeks are not certified complete.",
319 file=sys.stderr,
320 )
321 for name, fit in fits.items():
322 print(
323 f"# growth {name}: log-odds slope {fit['slope_per_week']:+.3f}/week "
324 f"(95% CI {fit['ci_low']:+.3f} to {fit['ci_high']:+.3f}, "
325 f"{int(fit['n_weeks'])} weeks, dispersion {fit['dispersion']:.1f}). "
326 f"Descriptive only -- confounded by sampling and reporting changes.",
327 file=sys.stderr,
328 )
329 return 0
332if __name__ == "__main__":
333 sys.exit(main())