obikmer/benchmark/verify_merge_count.py

#!/usr/bin/env python3
"""Verify the merged count index against all per-specimen reference sets.

Streams `obikmer dump` once on the merged index, accumulates per-specimen
kmer+count pairs from each column, then compares each against its reference .npz.

Output to stdout: one CSV row per specimen (same columns as verify_count.py)
  species,strain,ref_kmers,idx_kmers,false_neg,false_pos,count_mismatch,
  fn_pct,fp_pct,cm_pct
"""
import argparse
import subprocess
import sys
from pathlib import Path

import numpy as np


# ── encoding ──────────────────────────────────────────────────────────────────

_ENCODE = {'A': 0, 'C': 1, 'G': 2, 'T': 3,
           'a': 0, 'c': 1, 'g': 2, 't': 3}

_DECODE = ['A', 'C', 'G', 'T']


def encode_kmer(s: str) -> int:
    kmer = 0
    for c in s:
        kmer = (kmer << 2) | _ENCODE[c]
    return kmer


def decode_kmer(val: int, k: int) -> str:
    bases = []
    for _ in range(k):
        bases.append(_DECODE[val & 3])
        val >>= 2
    return ''.join(reversed(bases))


# ── single-pass dump ──────────────────────────────────────────────────────────

def stream_merged_dump(obikmer_bin: str, index_dir: str,
                       ) -> tuple[list[str], dict[str, tuple[list[int], list[int]]]]:
    """Stream the merged dump once.

    Returns:
        specimen_names : column labels in dump order
        per_specimen   : mapping label → (kmer_ints, counts) for entries > 0
    """
    cmd = [obikmer_bin, 'dump', index_dir]
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
                            text=True)

    header_line = proc.stdout.readline().rstrip('\n')
    cols = header_line.split(',')
    specimen_names = cols[1:]
    per_specimen: dict[str, tuple[list[int], list[int]]] = {
        name: ([], []) for name in specimen_names}

    for line in proc.stdout:
        parts = line.rstrip('\n').split(',')
        kmer_int = encode_kmer(parts[0])
        for i, name in enumerate(specimen_names):
            count = int(parts[i + 1])
            if count > 0:
                per_specimen[name][0].append(kmer_int)
                per_specimen[name][1].append(count)

    proc.wait()
    if proc.returncode != 0:
        print(f'ERROR: obikmer dump exited {proc.returncode}', file=sys.stderr)
        sys.exit(1)

    return specimen_names, per_specimen


# ── per-specimen comparison ───────────────────────────────────────────────────

def compare_specimen(name: str,
                     kmer_list: list[int],
                     count_list: list[int],
                     ref_dir: Path,
                     k: int,
                     save_fn: Path | None,
                     save_fp: Path | None,
                     save_cm: Path | None,
                     ) -> str:
    ref_path = ref_dir / f'{name}.npz'
    if not ref_path.exists():
        print(f'  SKIP {name}: no reference at {ref_path}', file=sys.stderr)
        return ''

    species = name.split('--')[0]
    strain  = name[len(species) + 2:]

    npz        = np.load(ref_path)
    ref_kmers  = npz['kmers']    # sorted uint64
    ref_counts = npz['counts']   # uint32

    order      = np.argsort(np.array(kmer_list, dtype=np.uint64), kind='stable')
    idx_kmers  = np.array(kmer_list,  dtype=np.uint64)[order]
    idx_counts = np.array(count_list, dtype=np.uint32)[order]

    false_neg = np.setdiff1d(ref_kmers, idx_kmers, assume_unique=True)
    false_pos = np.setdiff1d(idx_kmers, ref_kmers, assume_unique=True)

    # Count mismatches among shared kmers
    pos_in_idx     = np.searchsorted(idx_kmers, ref_kmers)
    pos_in_idx     = np.clip(pos_in_idx, 0, len(idx_kmers) - 1)
    shared_mask    = idx_kmers[pos_in_idx] == ref_kmers
    mismatch_mask  = ref_counts[shared_mask] != idx_counts[pos_in_idx[shared_mask]]
    cm_kmers       = ref_kmers[shared_mask][mismatch_mask]
    cm_ref         = ref_counts[shared_mask][mismatch_mask]
    cm_idx         = idx_counts[pos_in_idx[shared_mask]][mismatch_mask]

    n_shared = int(shared_mask.sum())
    fn_pct   = 100.0 * len(false_neg) / len(ref_kmers) if len(ref_kmers) else 0.0
    fp_pct   = 100.0 * len(false_pos) / len(idx_kmers) if len(idx_kmers) else 0.0
    cm_pct   = 100.0 * len(cm_kmers)  / n_shared       if n_shared        else 0.0

    print(f'  {name}: ref={len(ref_kmers):,}  idx={len(idx_kmers):,}  '
          f'fn={len(false_neg):,} ({fn_pct:.4f}%)  '
          f'fp={len(false_pos):,} ({fp_pct:.4f}%)  '
          f'cm={len(cm_kmers):,} ({cm_pct:.4f}%)',
          file=sys.stderr)

    if save_fn and len(false_neg):
        fn_file = save_fn / f'{name}_fn.txt'
        fn_file.write_text('\n'.join(decode_kmer(int(v), k) for v in false_neg) + '\n')

    if save_fp and len(false_pos):
        fp_file = save_fp / f'{name}_fp.txt'
        fp_file.write_text('\n'.join(decode_kmer(int(v), k) for v in false_pos) + '\n')

    if save_cm and len(cm_kmers):
        cm_file = save_cm / f'{name}_cm.csv'
        lines = ['kmer,ref_count,idx_count']
        for v, rc, ic in zip(cm_kmers, cm_ref, cm_idx):
            lines.append(f'{decode_kmer(int(v), k)},{rc},{ic}')
        cm_file.write_text('\n'.join(lines) + '\n')

    return (f'{species},{strain},'
            f'{len(ref_kmers)},{len(idx_kmers)},'
            f'{len(false_neg)},{len(false_pos)},{len(cm_kmers)},'
            f'{fn_pct:.4f},{fp_pct:.4f},{cm_pct:.4f}')


# ── main ─────────────────────────────────────────────────────────────────────

def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('index',     metavar='INDEX_DIR', nargs='?',
                    help='Merged count index directory')
    ap.add_argument('ref_dir',   metavar='REF_DIR',   nargs='?',
                    help='Directory containing per-specimen .npz reference files')
    ap.add_argument('--obikmer', default='obikmer')
    ap.add_argument('--header',  action='store_true',
                    help='Print CSV header and exit')
    ap.add_argument('--save-fn', metavar='DIR',
                    help='Directory for false-negative kmer lists')
    ap.add_argument('--save-fp', metavar='DIR',
                    help='Directory for false-positive kmer lists')
    ap.add_argument('--save-cm', metavar='DIR',
                    help='Directory for count-mismatch CSV files')
    args = ap.parse_args()

    if args.header:
        print('species,strain,ref_kmers,idx_kmers,'
              'false_neg,false_pos,count_mismatch,'
              'fn_pct,fp_pct,cm_pct')
        return

    ref_dir = Path(args.ref_dir)
    save_fn = Path(args.save_fn) if args.save_fn else None
    save_fp = Path(args.save_fp) if args.save_fp else None
    save_cm = Path(args.save_cm) if args.save_cm else None
    for d in (save_fn, save_fp, save_cm):
        if d: d.mkdir(parents=True, exist_ok=True)

    out1 = subprocess.check_output(
        [args.obikmer, 'dump', '--head', '1', args.index],
        stderr=subprocess.DEVNULL, text=True)
    k = len(out1.splitlines()[1].split(',')[0])

    print(f'k={k}  streaming merged dump: {args.index}', file=sys.stderr)
    specimen_names, per_specimen = stream_merged_dump(args.obikmer, args.index)
    print(f'{len(specimen_names)} specimen columns loaded', file=sys.stderr)

    for name in specimen_names:
        kmers, counts = per_specimen[name]
        row = compare_specimen(name, kmers, counts, ref_dir, k,
                               save_fn, save_fp, save_cm)
        if row:
            print(row)


if __name__ == '__main__':
    main()