@pikaa-ai/pikaa 0.2.5 → 0.3.0

package/skills/pathogen-variant-surveillance/scripts/lineage_prevalence.py

@@ -0,0 +1,310 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.11"
+# dependencies = []
+# ///
+"""Weekly prevalence of one or more lineages, with intervals and a coverage flag.
+
+Answers "what is circulating, and is it growing" from live sequence counts
+rather than from memory. Every number is a count returned by the instance at
+the data version printed in the provenance block.
+
+    python3 lineage_prevalence.py XFG.1.1 XFJ.3 --instance sars-cov-2 --where country=USA
+    python3 lineage_prevalence.py XFG --sublineages --weeks 12 --growth
+    python3 lineage_prevalence.py 2.3.4.4b --instance h5n1 --where country=USA --weeks 52
+
+Proportions carry Wilson intervals because surveillance weeks are small, and
+recent weeks are flagged ``low`` when their denominator has not filled in yet
+-- see reporting_lag.py for the measured filling-in curve.
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+from datetime import date, timedelta
+
+from lapis_client import (
+    LapisError,
+    aggregated,
+    bin_weekly,
+    count,
+    data_version,
+    describe_instance,
+    emit,
+    flag_low_coverage,
+    lineage_filter,
+    logit_slope,
+    pick_date_field,
+    pick_lineage_field,
+    range_keys,
+    resolve_base_url,
+    top_values,
+    week_range,
+    wilson_interval,
+)
+
+COLUMNS = (
+    "lineage",
+    "week",
+    "n",
+    "total",
+    "proportion",
+    "ci_low",
+    "ci_high",
+    "coverage",
+)
+
+
+def parse_where(pairs: list[str], schema: dict) -> dict[str, str]:
+    """Turn ``KEY=VALUE`` arguments into filters, checking keys against the schema."""
+    filters: dict[str, str] = {}
+    types = schema.get("types", {})
+    for pair in pairs:
+        key, sep, value = pair.partition("=")
+        if not sep:
+            raise LapisError(f"--where expects KEY=VALUE, got {pair!r}")
+        base = key.split(".")[0]
+        for suffix in ("From", "To"):
+            if base.endswith(suffix) and base[: -len(suffix)] in types:
+                base = base[: -len(suffix)]
+                break
+        if base not in types:
+            near = sorted(n for n in types if base.lower() in n.lower())
+            hint = f" Did you mean: {', '.join(near[:6])}?" if near else ""
+            raise LapisError(f"{key!r} is not a field on this instance.{hint}")
+        filters[key] = value
+    return filters
+
+
+def weekly_counts(
+    base_url: str, filters: dict, date_field: str, weeks: list[str]
+) -> tuple[dict[str, int], int]:
+    """Weekly counts over the requested window, zero-filled for empty weeks."""
+    rows = aggregated(base_url, filters, [date_field])
+    binned, undated = bin_weekly(rows, date_field)
+    return {week: binned.get(week, 0) for week in weeks}, undated
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Weekly lineage prevalence from a LAPIS instance.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("lineages", nargs="*",
+                        help="lineage or clade names; 'NAME*' includes descendants. "
+                             "Omit to discover them with --top")
+    parser.add_argument("--top", type=int, metavar="N",
+                        help="discover the N most common lineages in the window instead of "
+                             "naming them (default when no names are given)")
+    parser.add_argument("--instance", default="sars-cov-2", help="registry name (default: sars-cov-2)")
+    parser.add_argument("--base-url", help="any other LAPIS deployment")
+    parser.add_argument("--lineage-field", help="override the auto-detected lineage column")
+    parser.add_argument("--date-field", help="override the auto-detected collection-date column")
+    parser.add_argument("--where", action="append", default=[], metavar="KEY=VALUE",
+                        help="extra filter, repeatable (e.g. --where country=USA)")
+    parser.add_argument("--weeks", type=int, default=26, help="window length in weeks (default: 26)")
+    parser.add_argument("--since", help="window start, YYYY-MM-DD; overrides --weeks")
+    parser.add_argument("--until", help="window end, YYYY-MM-DD (default: today)")
+    parser.add_argument("--sublineages", action="store_true", help="include descendant lineages")
+    parser.add_argument("--growth", action="store_true",
+                        help="also fit a weighted log-odds slope over the trusted weeks")
+    parser.add_argument("--include-incomplete", action="store_true",
+                        help="let low-coverage weeks into the growth fit")
+    parser.add_argument("--coverage-fraction", type=float, default=0.4,
+                        help="flag weeks below this fraction of the settled median (default: 0.4)")
+    parser.add_argument("--format", choices=("table", "tsv", "json"), default="table")
+    parser.add_argument("-o", "--output", help="write to a file instead of stdout")
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+
+    try:
+        base_url = resolve_base_url(args.instance, args.base_url)
+        schema = describe_instance(base_url)
+        lineage_field, has_index = pick_lineage_field(schema, args.lineage_field)
+        date_field = pick_date_field(schema, "collection", args.date_field)
+        where = parse_where(args.where, schema)
+    except LapisError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 2
+
+    requested_until = date.fromisoformat(args.until) if args.until else date.today()
+    requested_since = (
+        date.fromisoformat(args.since)
+        if args.since
+        else requested_until - timedelta(weeks=max(1, args.weeks))
+    )
+    if requested_since > requested_until:
+        print("error: --since is after --until", file=sys.stderr)
+        return 2
+    if args.top is not None and args.top < 1:
+        print("error: --top must be at least 1", file=sys.stderr)
+        return 2
+
+    # Snap both ends to ISO week boundaries. A window that starts mid-week gives
+    # a first row covering three days and a last row covering four, and neither
+    # is comparable to the full weeks between them -- exactly the kind of
+    # silently uneven denominator this skill exists to prevent. Widening to whole
+    # weeks keeps every row meaning the same thing; a still-unfilled final week
+    # shows up through the coverage flag rather than as a short bar.
+    since = requested_since - timedelta(days=requested_since.weekday())
+    until = requested_until + timedelta(days=6 - requested_until.weekday())
+    snapped = (since, until) != (requested_since, requested_until)
+
+    from_key, to_key = range_keys(date_field)
+    window = {from_key: since.isoformat(), to_key: until.isoformat(), **where}
+    weeks = week_range(since.isoformat(), until.isoformat())
+
+    try:
+        version = data_version(base_url)
+        totals, undated = weekly_counts(base_url, window, date_field, weeks)
+
+        rows: list[dict] = []
+        notes: list[str] = []
+        fits: dict[str, dict] = {}
+        coverage = flag_low_coverage(totals, args.coverage_fraction)
+
+        targets = list(args.lineages)
+        if args.top is not None or not targets:
+            wanted = args.top or 10
+            discovered = top_values(
+                aggregated(base_url, window, [lineage_field]), lineage_field, wanted
+            )
+            if not discovered:
+                print(
+                    f"error: no sequences in {since} to {until}"
+                    + (f" with {where}" if where else "")
+                    + ". Widen the window or relax the filters.",
+                    file=sys.stderr,
+                )
+                return 1
+            notes.append(
+                f"discovered the {len(discovered)} most common {lineage_field} "
+                f"value{'s' if len(discovered) != 1 else ''} in the window: "
+                f"{', '.join(discovered)}"
+            )
+            targets = discovered + [t for t in targets if t not in discovered]
+
+        for name in targets:
+            value = lineage_filter(name, has_index, args.sublineages)
+            counts, _ = weekly_counts(
+                base_url, {**window, lineage_field: value}, date_field, weeks
+            )
+
+            if has_index and not value.endswith("*"):
+                # Both sides must come from the same kind of query. Summing the
+                # weekly bins would drop sequences with no usable collection
+                # date, making every lineage that has some look as though it had
+                # descendants it does not.
+                exact = count(base_url, {**window, lineage_field: value})
+                inclusive = count(base_url, {**window, lineage_field: f"{value}*"})
+                if inclusive > exact:
+                    notes.append(
+                        f"{value}: {exact} sequences named exactly {value}, "
+                        f"{inclusive} including descendants. Pass --sublineages "
+                        f"or write '{value}*' for the second number."
+                    )
+
+            for week in weeks:
+                k, n = counts[week], totals[week]
+                low, high = wilson_interval(k, n)
+                rows.append(
+                    {
+                        "lineage": value,
+                        "week": week,
+                        "n": k,
+                        "total": n,
+                        "proportion": f"{(k / n):.4f}" if n else "",
+                        "ci_low": f"{low:.4f}" if n else "",
+                        "ci_high": f"{high:.4f}" if n else "",
+                        "coverage": "low" if coverage.get(week) else "ok",
+                    }
+                )
+
+            if args.growth:
+                usable = [
+                    (i, counts[w], totals[w])
+                    for i, w in enumerate(weeks)
+                    if args.include_incomplete or not coverage.get(w)
+                ]
+                fit = logit_slope([(float(t), k, n) for t, k, n in usable])
+                if fit:
+                    fits[value] = fit
+                else:
+                    notes.append(
+                        f"{value}: too few observations for a growth estimate "
+                        f"({sum(k for _, k, _ in usable)} sequences across "
+                        f"{sum(1 for _, k, _ in usable if k > 0)} non-empty weeks). "
+                        "No slope is reported rather than a slope driven by the denominator."
+                    )
+    except LapisError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 2
+
+    if args.format == "json":
+        import json
+
+        payload = {
+            "meta": {
+                "instance": schema["name"],
+                "base_url": base_url,
+                "data_version": version,
+                "lineage_field": lineage_field,
+                "lineage_index": has_index,
+                "date_field": date_field,
+                "window": [since.isoformat(), until.isoformat()],
+                "filters": where,
+                "undated_sequences": undated,
+                "notes": notes,
+            },
+            "rows": rows,
+            "growth": fits,
+        }
+        text = json.dumps(payload, indent=2)
+    else:
+        text = emit(rows, COLUMNS, args.format)
+
+    if args.output:
+        with open(args.output, "w", encoding="utf-8") as handle:
+            handle.write(text + "\n")
+    else:
+        print(text)
+
+    sys.stdout.flush()
+    if args.format != "json":
+        print(
+            f"\n# {schema['name']} via {base_url}"
+            f"\n# data version {version} | {lineage_field}"
+            f"{' (lineage-indexed)' if has_index else ' (no lineage index)'}"
+            f" | dates from {date_field}"
+            f"\n# window {since} to {until}"
+            + (f" (widened from {requested_since}..{requested_until} to whole ISO weeks)"
+               if snapped else "")
+            + (f" | filters {where}" if where else "")
+            + (f"\n# {undated} matching sequences carry no usable collection date" if undated else ""),
+            file=sys.stderr,
+        )
+        for note in notes:
+            print(f"# note: {note}", file=sys.stderr)
+        if any(coverage.values()):
+            flagged = [w for w, bad in coverage.items() if bad]
+            print(
+                f"# {len(flagged)} week(s) flagged low: denominator still filling in "
+                f"(earliest {min(flagged)}). Treat their proportions as unstable.",
+                file=sys.stderr,
+            )
+        for name, fit in fits.items():
+            print(
+                f"# growth {name}: log-odds slope {fit['slope_per_week']:+.3f}/week "
+                f"(95% CI {fit['ci_low']:+.3f} to {fit['ci_high']:+.3f}, "
+                f"{int(fit['n_weeks'])} weeks, dispersion {fit['dispersion']:.1f}). "
+                f"Descriptive only -- confounded by sampling and reporting changes.",
+                file=sys.stderr,
+            )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

package/skills/pathogen-variant-surveillance/scripts/mutation_profile.py

@@ -0,0 +1,215 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.11"
+# dependencies = []
+# ///
+"""Mutations carried by a lineage, or the difference between two lineages.
+
+Use this to answer "what distinguishes this lineage" and "does my assay target
+still match" from current sequences, instead of from a lineage's founding
+description -- which drifts as sublineages accumulate substitutions.
+
+    python3 mutation_profile.py XFG.1.1 --gene S --min-proportion 0.9
+    python3 mutation_profile.py XFG.23.1.3 --versus XFG.1.1 --gene S
+    python3 mutation_profile.py 2.3.4.4b --instance h5n1 --gene HA --min-proportion 0.8
+    python3 mutation_profile.py XFG --sublineages --nucleotide --since 2026-01-01
+
+``proportion`` is over ``coverage`` -- the sequences that actually resolved that
+site -- not over every matching sequence. A site with poor coverage can show a
+high proportion on very few reads, so read the coverage column before quoting a
+proportion.
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+
+from lapis_client import (
+    LapisError,
+    count,
+    data_version,
+    describe_instance,
+    emit,
+    lineage_filter,
+    mutations,
+    pick_date_field,
+    pick_lineage_field,
+    range_keys,
+    resolve_base_url,
+)
+
+PROFILE_COLUMNS = ("mutation", "gene", "position", "from", "to", "proportion", "count", "coverage")
+DIFF_COLUMNS = ("mutation", "gene", "position", "verdict", "prop_a", "prop_b", "n_a", "n_b")
+
+
+def parse_where(pairs: list[str], schema: dict) -> dict[str, str]:
+    """Turn ``KEY=VALUE`` arguments into filters, checking keys against the schema."""
+    filters: dict[str, str] = {}
+    types = schema.get("types", {})
+    for pair in pairs:
+        key, sep, value = pair.partition("=")
+        if not sep:
+            raise LapisError(f"--where expects KEY=VALUE, got {pair!r}")
+        if key.split(".")[0] not in types:
+            raise LapisError(f"{key!r} is not a field on this instance")
+        filters[key] = value
+    return filters
+
+
+def profile(base_url: str, filters: dict, *, amino_acid: bool, min_proportion: float,
+            gene: str | None) -> dict[str, dict]:
+    """Mutation rows keyed by mutation string, optionally restricted to one gene."""
+    rows = mutations(base_url, filters, amino_acid=amino_acid, min_proportion=min_proportion)
+    keyed: dict[str, dict] = {}
+    for row in rows:
+        if gene and str(row.get("sequenceName") or "").upper() != gene.upper():
+            continue
+        keyed[row["mutation"]] = row
+    return keyed
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Mutation profile of a lineage.")
+    parser.add_argument("lineage", help="lineage or clade name; 'NAME*' includes descendants")
+    parser.add_argument("--versus", help="second lineage to diff against")
+    parser.add_argument("--instance", default="sars-cov-2", help="registry name (default: sars-cov-2)")
+    parser.add_argument("--base-url", help="any other LAPIS deployment")
+    parser.add_argument("--lineage-field", help="override the auto-detected lineage column")
+    parser.add_argument("--date-field", help="override the auto-detected collection-date column")
+    parser.add_argument("--gene", help="restrict to one gene or segment (e.g. S, HA, seg4)")
+    parser.add_argument("--nucleotide", action="store_true",
+                        help="nucleotide instead of amino-acid mutations")
+    parser.add_argument("--min-proportion", type=float, default=0.8,
+                        help="report mutations at or above this proportion (default: 0.8)")
+    parser.add_argument("--sublineages", action="store_true", help="include descendant lineages")
+    parser.add_argument("--since", help="restrict to sequences collected on or after YYYY-MM-DD")
+    parser.add_argument("--where", action="append", default=[], metavar="KEY=VALUE",
+                        help="extra filter, repeatable")
+    parser.add_argument("--format", choices=("table", "tsv", "json"), default="table")
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+
+    try:
+        base_url = resolve_base_url(args.instance, args.base_url)
+        schema = describe_instance(base_url)
+        lineage_field, has_index = pick_lineage_field(schema, args.lineage_field)
+        where = parse_where(args.where, schema)
+        if args.since:
+            date_field = pick_date_field(schema, "collection", args.date_field)
+            where[range_keys(date_field)[0]] = args.since
+
+        # A diff must fetch below the reporting threshold on both sides,
+        # otherwise a mutation pruned out of one side is indistinguishable from
+        # one genuinely absent there and gets misreported as gained or lost.
+        # The floor tracks --min-proportion so lowering it stays correct.
+        diff_floor = min(0.05, args.min_proportion)
+
+        primary = lineage_filter(args.lineage, has_index, args.sublineages)
+        base_filters = {**where, lineage_field: primary}
+        n_primary = count(base_url, base_filters)
+        if n_primary == 0:
+            print(
+                f"error: no sequences match {lineage_field}={primary}"
+                + (f" with {where}" if where else "")
+                + ". Check the name with resolve_lineage.py.",
+                file=sys.stderr,
+            )
+            return 1
+
+        first = profile(
+            base_url,
+            base_filters,
+            amino_acid=not args.nucleotide,
+            min_proportion=args.min_proportion if not args.versus else diff_floor,
+            gene=args.gene,
+        )
+
+        if not args.versus:
+            rows = [
+                {
+                    "mutation": m["mutation"],
+                    "gene": m.get("sequenceName", ""),
+                    "position": m.get("position", ""),
+                    "from": m.get("mutationFrom", ""),
+                    "to": m.get("mutationTo", ""),
+                    "proportion": f"{m['proportion']:.3f}",
+                    "count": m.get("count", ""),
+                    "coverage": m.get("coverage", ""),
+                }
+                for m in sorted(
+                    first.values(),
+                    key=lambda r: (str(r.get("sequenceName")), int(r.get("position") or 0)),
+                )
+            ]
+            print(emit(rows, PROFILE_COLUMNS, args.format))
+            summary = (
+                f"\n# {schema['name']} via {base_url} | data version {data_version(base_url)}"
+                f"\n# {lineage_field}={primary} | {n_primary} sequences"
+                f" | {'nucleotide' if args.nucleotide else 'amino acid'} mutations"
+                f" at proportion >= {args.min_proportion}"
+                f"\n# proportion is over per-site coverage, not over all {n_primary} sequences"
+            )
+        else:
+            secondary = lineage_filter(args.versus, has_index, args.sublineages)
+            n_secondary = count(base_url, {**where, lineage_field: secondary})
+            if n_secondary == 0:
+                print(f"error: no sequences match {lineage_field}={secondary}", file=sys.stderr)
+                return 1
+            second = profile(
+                base_url,
+                {**where, lineage_field: secondary},
+                amino_acid=not args.nucleotide,
+                min_proportion=diff_floor,
+                gene=args.gene,
+            )
+
+            rows = []
+            for mutation in sorted(set(first) | set(second)):
+                a = first.get(mutation, {})
+                b = second.get(mutation, {})
+                pa = float(a.get("proportion") or 0.0)
+                pb = float(b.get("proportion") or 0.0)
+                if pa >= args.min_proportion and pb < args.min_proportion:
+                    verdict = "gained"
+                elif pb >= args.min_proportion and pa < args.min_proportion:
+                    verdict = "lost"
+                elif pa >= args.min_proportion and pb >= args.min_proportion:
+                    verdict = "shared"
+                else:
+                    continue
+                source = a or b
+                rows.append(
+                    {
+                        "mutation": mutation,
+                        "gene": source.get("sequenceName", ""),
+                        "position": source.get("position", ""),
+                        "verdict": verdict,
+                        "prop_a": f"{pa:.3f}",
+                        "prop_b": f"{pb:.3f}",
+                        "n_a": a.get("count", 0),
+                        "n_b": b.get("count", 0),
+                    }
+                )
+            rows.sort(key=lambda r: ({"gained": 0, "lost": 1, "shared": 2}[r["verdict"]],
+                                     str(r["gene"]), int(r["position"] or 0)))
+            print(emit(rows, DIFF_COLUMNS, args.format))
+            summary = (
+                f"\n# {schema['name']} via {base_url} | data version {data_version(base_url)}"
+                f"\n# A = {primary} ({n_primary} sequences), B = {secondary} ({n_secondary})"
+                f"\n# verdict is relative to a {args.min_proportion} proportion threshold"
+            )
+    except LapisError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 2
+
+    sys.stdout.flush()
+    if args.format != "json":
+        print(summary, file=sys.stderr)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())