training/audit_variant_sequence_encoding.py

#!/usr/bin/env python3
"""Audit whether sequence strings contain REF or ALT alleles near the center.

The prepared sequence files were fetched as roughly +/-512 bp around each
variant, so the variant should begin around index 512 in the sequence string.
This script checks that assumption for SNVs and small indels.
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path

import pandas as pd


ORIGINAL_SPLIT_FILES = {
    "train": "train_with_sequences.csv",
    "val": "val_with_sequences.csv",
    "test": "test_with_sequences.csv",
}

ALT_SPLIT_FILES = {
    "train": "train_with_alt_sequences.csv",
    "val": "val_with_alt_sequences.csv",
    "test": "test_with_alt_sequences.csv",
}

REQUIRED_COLUMNS = {"sequence", "REF", "ALT"}
OPTIONAL_EXAMPLE_COLUMNS = [
    "variant_id",
    "CHROM",
    "POS",
    "CLNHGVS",
    "gene_symbol",
    "GENEINFO",
]

CENTER_INDEX = 512
SNIPPET_FLANK = 20
MAX_INDEL_SIZE = 50
VALID_BASES = set("ACGTN")


@dataclass
class AuditCounts:
    total_rows_seen: int = 0
    total_rows_checked: int = 0
    snv_rows_checked: int = 0
    indel_rows_checked: int = 0
    reference_matches: int = 0
    alternate_matches: int = 0
    mismatches: int = 0
    skipped_missing_values: int = 0
    skipped_short_sequence: int = 0
    skipped_multiple_alt: int = 0
    skipped_symbolic_alt: int = 0
    skipped_non_acgtn_allele: int = 0
    skipped_not_snv_or_small_indel: int = 0

    def add(self, other: "AuditCounts") -> None:
        for field_name in self.__dataclass_fields__:
            setattr(self, field_name, getattr(self, field_name) + getattr(other, field_name))


@dataclass
class DatasetChoice:
    data_dir: Path
    split_files: dict[str, str]
    is_alternate_dataset: bool


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--data-dir",
        type=Path,
        default=None,
        help=(
            "Directory containing train/val/test CSV files. Defaults to "
            "training/csv_files_20k_alt, training/csv_files_10k_alt, "
            "training/csv_files_large_alt, training/csv_files_alt, "
            "training/csv_files_20k, training/csv_files_10k, "
            "training/csv_files_large, training/csv_files, then data/processed."
        ),
    )
    parser.add_argument(
        "--center-index",
        type=int,
        default=CENTER_INDEX,
        help="0-based index where the variant is expected to start. Default: 512.",
    )
    return parser.parse_args()


def project_root() -> Path:
    return Path(__file__).resolve().parents[1]


def has_all_files(directory: Path, split_files: dict[str, str]) -> bool:
    return all((directory / filename).exists() for filename in split_files.values())


def choose_dataset(root: Path, requested_dir: Path | None) -> DatasetChoice:
    if requested_dir is not None:
        data_dir = requested_dir.expanduser().resolve()
        if has_all_files(data_dir, ALT_SPLIT_FILES):
            return DatasetChoice(data_dir, ALT_SPLIT_FILES, True)
        if has_all_files(data_dir, ORIGINAL_SPLIT_FILES):
            return DatasetChoice(data_dir, ORIGINAL_SPLIT_FILES, False)
        raise FileNotFoundError(
            "The requested directory does not contain a complete alternate or original dataset:\n"
            f"{data_dir}"
        )

    candidates: list[tuple[Path, dict[str, str], bool]] = [
        (root / "training" / "csv_files_20k_alt", ALT_SPLIT_FILES, True),
        (root / "training" / "csv_files_10k_alt", ALT_SPLIT_FILES, True),
        (root / "training" / "csv_files_large_alt", ALT_SPLIT_FILES, True),
        (root / "training" / "csv_files_alt", ALT_SPLIT_FILES, True),
        (root / "training" / "csv_files_20k", ORIGINAL_SPLIT_FILES, False),
        (root / "training" / "csv_files_10k", ORIGINAL_SPLIT_FILES, False),
        (root / "training" / "csv_files_large", ORIGINAL_SPLIT_FILES, False),
        (root / "training" / "csv_files", ORIGINAL_SPLIT_FILES, False),
        (root / "data" / "processed", ORIGINAL_SPLIT_FILES, False),
    ]

    for directory, split_files, is_alternate_dataset in candidates:
        if has_all_files(directory, split_files):
            return DatasetChoice(directory, split_files, is_alternate_dataset)

    searched = "\n".join(str(path) for path, _split_files, _is_alt in candidates)
    raise FileNotFoundError(
        "Could not find a complete alternate or original sequence dataset.\n"
        f"Searched:\n{searched}"
    )


def normalize_sequence(value: object) -> str:
    return str(value).strip().upper()


def normalize_allele(value: object) -> str:
    return str(value).strip().upper()


def is_symbolic_alt(alt: str) -> bool:
    return alt.startswith("<") or alt.endswith(">") or "[" in alt or "]" in alt


def is_acgtn(value: str) -> bool:
    return bool(value) and set(value).issubset(VALID_BASES)


def is_snv(ref: str, alt: str) -> bool:
    return len(ref) == 1 and len(alt) == 1


def is_small_indel(ref: str, alt: str) -> bool:
    return ref != alt and abs(len(ref) - len(alt)) <= MAX_INDEL_SIZE


def center_snippet(sequence: str, center_index: int) -> str:
    start = max(0, center_index - SNIPPET_FLANK)
    end = min(len(sequence), center_index + SNIPPET_FLANK + 1)
    snippet = sequence[start:end]
    marker_position = center_index - start
    if 0 <= marker_position < len(snippet):
        return snippet[:marker_position] + "[" + snippet[marker_position] + "]" + snippet[marker_position + 1 :]
    return snippet


def classify_row(row: pd.Series, center_index: int, is_alternate_dataset: bool) -> tuple[str, str, str]:
    sequence = normalize_sequence(row["sequence"])
    ref = normalize_allele(row["REF"])
    alt = normalize_allele(row["ALT"])

    if not sequence or not ref or not alt or sequence == "NAN" or ref == "NAN" or alt == "NAN":
        return "skipped_missing_values", "", ""

    if "," in alt:
        return "skipped_multiple_alt", "", ""

    if is_symbolic_alt(alt):
        return "skipped_symbolic_alt", "", ""

    if not is_acgtn(ref) or not is_acgtn(alt):
        return "skipped_non_acgtn_allele", "", ""

    if len(sequence) <= center_index:
        return "skipped_short_sequence", "", ""

    if is_snv(ref, alt):
        observed = sequence[center_index]
        if is_alternate_dataset:
            if observed == alt:
                return "alternate_match_snv", observed, center_snippet(sequence, center_index)
            if observed == ref:
                return "reference_match_snv", observed, center_snippet(sequence, center_index)
        else:
            if observed == ref:
                return "reference_match_snv", observed, center_snippet(sequence, center_index)
            if observed == alt:
                return "alternate_match_snv", observed, center_snippet(sequence, center_index)
        return "mismatch_snv", observed, center_snippet(sequence, center_index)

    if is_small_indel(ref, alt):
        expected_allele = alt if is_alternate_dataset else ref
        if len(sequence) < center_index + len(expected_allele):
            return "skipped_short_sequence", "", ""

        ref_window = sequence[center_index : center_index + len(ref)] if len(sequence) >= center_index + len(ref) else ""
        alt_window = sequence[center_index : center_index + len(alt)] if len(sequence) >= center_index + len(alt) else ""

        if is_alternate_dataset:
            if alt_window == alt:
                return "alternate_match_indel", alt_window, center_snippet(sequence, center_index)
            if ref_window == ref:
                return "reference_match_indel", ref_window, center_snippet(sequence, center_index)
            return "mismatch_indel", alt_window, center_snippet(sequence, center_index)

        if ref_window == ref:
            return "reference_match_indel", ref_window, center_snippet(sequence, center_index)
        if alt_window == alt:
            return "alternate_match_indel", alt_window, center_snippet(sequence, center_index)
        return "mismatch_indel", ref_window, center_snippet(sequence, center_index)

    return "skipped_not_snv_or_small_indel", "", ""


def update_counts(counts: AuditCounts, status: str) -> None:
    if status.startswith("skipped_"):
        current_value = getattr(counts, status)
        setattr(counts, status, current_value + 1)
        return

    counts.total_rows_checked += 1

    if status.endswith("_snv"):
        counts.snv_rows_checked += 1
    elif status.endswith("_indel"):
        counts.indel_rows_checked += 1

    if status.startswith("reference_match"):
        counts.reference_matches += 1
    elif status.startswith("alternate_match"):
        counts.alternate_matches += 1
    elif status.startswith("mismatch"):
        counts.mismatches += 1


def build_example(row: pd.Series, split_name: str, status: str, observed: str, snippet: str) -> dict[str, object]:
    example = {
        "split": split_name,
        "status": status,
        "REF": normalize_allele(row["REF"]),
        "ALT": normalize_allele(row["ALT"]),
        "observed_at_center": observed,
        "center_sequence_snippet": snippet,
    }

    for column in OPTIONAL_EXAMPLE_COLUMNS:
        if column in row.index:
            example[column] = row[column]

    return example


def audit_split(
    split_name: str,
    csv_path: Path,
    center_index: int,
    is_alternate_dataset: bool,
) -> tuple[AuditCounts, list[dict[str, object]]]:
    print("=" * 80)
    print(f"{split_name.upper()} SPLIT")
    print("=" * 80)
    print(f"Path: {csv_path}")

    df = pd.read_csv(csv_path)
    counts = AuditCounts(total_rows_seen=len(df))
    examples: list[dict[str, object]] = []

    missing_columns = sorted(REQUIRED_COLUMNS - set(df.columns))
    if missing_columns:
        print(f"ERROR: missing required columns: {missing_columns}")
        print()
        return counts, examples

    for _, row in df.iterrows():
        status, observed, snippet = classify_row(row, center_index, is_alternate_dataset)
        update_counts(counts, status)

        if len(examples) < 10 and not status.startswith("skipped_"):
            examples.append(build_example(row, split_name, status, observed, snippet))

    print_counts(counts)
    print()
    return counts, examples


def print_counts(counts: AuditCounts) -> None:
    print(f"Rows seen: {counts.total_rows_seen:,}")
    print(f"Total rows checked: {counts.total_rows_checked:,}")
    print(f"SNV rows checked: {counts.snv_rows_checked:,}")
    print(f"Indel rows checked: {counts.indel_rows_checked:,}")
    print(f"Reference allele matches: {counts.reference_matches:,}")
    print(f"Alternate allele matches: {counts.alternate_matches:,}")
    print(f"Mismatches: {counts.mismatches:,}")
    print("Skipped rows:")
    print(f"  missing sequence/REF/ALT: {counts.skipped_missing_values:,}")
    print(f"  sequence too short for center check: {counts.skipped_short_sequence:,}")
    print(f"  multiple ALT alleles: {counts.skipped_multiple_alt:,}")
    print(f"  symbolic ALT allele: {counts.skipped_symbolic_alt:,}")
    print(f"  non-ACGTN allele: {counts.skipped_non_acgtn_allele:,}")
    print(f"  not SNV or small indel: {counts.skipped_not_snv_or_small_indel:,}")


def print_examples(examples: list[dict[str, object]]) -> None:
    print("=" * 80)
    print("FIRST 10 CHECKED EXAMPLES")
    print("=" * 80)

    if not examples:
        print("No checked examples available.")
        print()
        return

    examples_df = pd.DataFrame(examples)
    display_columns = [
        column
        for column in [
            "split",
            "variant_id",
            "CHROM",
            "POS",
            "gene_symbol",
            "REF",
            "ALT",
            "status",
            "observed_at_center",
            "center_sequence_snippet",
            "CLNHGVS",
            "GENEINFO",
        ]
        if column in examples_df.columns
    ]

    with pd.option_context("display.max_columns", None, "display.width", 220, "display.max_colwidth", 90):
        print(examples_df[display_columns].to_string(index=False))
    print()


def print_conclusion(counts: AuditCounts) -> None:
    print("=" * 80)
    print("OVERALL SUMMARY")
    print("=" * 80)
    print_counts(counts)

    checked = counts.total_rows_checked
    if checked == 0:
        conclusion = "Could not determine clearly."
    else:
        ref_fraction = counts.reference_matches / checked
        alt_fraction = counts.alternate_matches / checked
        print(f"Reference match fraction: {ref_fraction:.2%}")
        print(f"Alternate match fraction: {alt_fraction:.2%}")

        if ref_fraction >= 0.80 and counts.reference_matches > counts.alternate_matches:
            conclusion = "The sequence appears to be reference sequence."
        elif alt_fraction >= 0.80 and counts.alternate_matches > counts.reference_matches:
            conclusion = "The sequence appears to contain alternate alleles."
        else:
            conclusion = "Could not determine clearly."

    print()
    print(f"Conclusion: {conclusion}")


def main() -> None:
    args = parse_args()
    root = project_root()
    dataset = choose_dataset(root, args.data_dir)

    print("Variant sequence encoding audit")
    print(f"Selected data directory: {dataset.data_dir}")
    print(f"Auditing alternate dataset: {dataset.is_alternate_dataset}")
    print(f"Expected variant start index: {args.center_index}")
    print()

    total_counts = AuditCounts()
    all_examples: list[dict[str, object]] = []

    for split_name, filename in dataset.split_files.items():
        csv_path = dataset.data_dir / filename
        if not csv_path.exists():
            raise FileNotFoundError(f"Missing required CSV file: {csv_path}")

        split_counts, split_examples = audit_split(
            split_name,
            csv_path,
            args.center_index,
            dataset.is_alternate_dataset,
        )
        total_counts.add(split_counts)
        remaining_example_slots = 10 - len(all_examples)
        if remaining_example_slots > 0:
            all_examples.extend(split_examples[:remaining_example_slots])

    print_examples(all_examples)
    print_conclusion(total_counts)


if __name__ == "__main__":
    main()