diff --git a/pyproject.toml b/pyproject.toml index 8f866cb..ab4fb89 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ line-length = 79 [tool.flake8] max-line-length = 79 -extend-ignore = ["E231", "E221"] +extend-ignore = ["E203", "E231", "E221"] exclude = [ ".github/scripts/license_message.py", "sentieon_cli/scripts/combine_cnv.py", diff --git a/sentieon_cli/command_strings.py b/sentieon_cli/command_strings.py index 4bdebde..1748d98 100644 --- a/sentieon_cli/command_strings.py +++ b/sentieon_cli/command_strings.py @@ -248,45 +248,20 @@ def cmd_pyexec_hybrid_select( threads: int, slop_size: int = 1000, ) -> Pipeline: - select_cmd = Command( - sys.executable, - str(hybrid_select), - "-v", - str(vcf), - "-t", - str(threads), - "-", - ) - view_cmd = Command( - "bcftools", - "view", - "-f", - "PASS,.", - "-", - ) - query_cmd = Command( - "bcftools", - "query", - "-f", - "%CHROM\\t%POS0\\t%END\\n", - "-", - ) - slop_cmd = Command( - "bedtools", - "slop", - "-b", - str(slop_size), - "-g", - str(ref_fai), - "-i", - "-", - ) return Pipeline( - select_cmd, - view_cmd, - query_cmd, - slop_cmd, - file_output=out_bed, + Command( + sys.executable, + str(hybrid_select), + "-v", + str(vcf), + "-t", + str(threads), + "--reference-fai", + str(ref_fai), + "--slop-size", + str(slop_size), + str(out_bed), + ) ) @@ -311,6 +286,39 @@ def cmd_pyexec_hybrid_anno( return Pipeline(Command(*cmd)) +def cmd_pyexec_hybrid_transfer( + out_vcf: pathlib.Path, + raw_vcf: pathlib.Path, + population_vcf: pathlib.Path, + reference_fai: pathlib.Path, + temp_dir: pathlib.Path, + hybrid_transfer: pathlib.Path, + threads: int, + workers: int, +) -> Pipeline: + """Transfer population annotations in one bounded ordered job.""" + + return Pipeline( + Command( + sys.executable, + str(hybrid_transfer), + "--raw-vcf", + str(raw_vcf), + "--population-vcf", + str(population_vcf), + "--reference-fai", + str(reference_fai), + "--temp-dir", + str(temp_dir), + "--threads", + str(threads), + "--workers", + str(workers), + str(out_vcf), + ) + ) + + def hybrid_stage1( out_aln: pathlib.Path, reference: pathlib.Path, diff --git a/sentieon_cli/dnascope_hybrid.py b/sentieon_cli/dnascope_hybrid.py index e1a9e28..a882bab 100644 --- a/sentieon_cli/dnascope_hybrid.py +++ b/sentieon_cli/dnascope_hybrid.py @@ -45,7 +45,6 @@ parse_fai, vcf_contigs, ) -from .transfer import build_transfer_jobs logger = get_logger(__name__) @@ -61,6 +60,9 @@ "ONT": packaging.version.Version("1.2"), } +HYBRID_TRANSFER_WORKERS = 32 +FINAL_NORM_PROCESSES = 3 + class RgInfo: """A container class for short and long-read readgroups""" @@ -651,8 +653,7 @@ def build_dag(self) -> DAG: concat_job, rm_job5, anno_job, - transfer_jobs, - transfer_concat, + transfer_job, apply_job, norm_job, ) = self.call_variants(sr_aln, lr_aln, rg_info) @@ -670,11 +671,9 @@ def build_dag(self) -> DAG: dag.add_job(anno_job, {concat_job}) apply_dependencies = {anno_job} - if transfer_jobs and transfer_concat: - for job in transfer_jobs: - dag.add_job(job, {anno_job}) - dag.add_job(transfer_concat, set(transfer_jobs)) - apply_dependencies = {transfer_concat} + if transfer_job: + dag.add_job(transfer_job, {anno_job}) + apply_dependencies = {transfer_job} if apply_job: dag.add_job(apply_job, apply_dependencies) @@ -715,7 +714,6 @@ def call_variants( Job, Job, Job, - Optional[List[Job]], Optional[Job], Optional[Job], Optional[Job], @@ -781,7 +779,7 @@ def call_variants( threads=self.cores, ), "hybrid-select", - 0, + self.cores, ) mapq0_bed = self.tmp_dir.joinpath("hybrid_mapq0.bed") @@ -1010,11 +1008,10 @@ def call_variants( self.cores, ), "anno-calls", - 0, + self.cores, ) - transfer_jobs: Optional[List[Job]] = None - transfer_concat_job: Optional[Job] = None + transfer_job: Optional[Job] = None input_to_apply = anno_target if self.pop_vcf: @@ -1024,15 +1021,29 @@ def call_variants( if self.skip_model_apply: transfer_target = self.output_vcf - transfer_jobs, transfer_concat_job = build_transfer_jobs( - transfer_target, - self.pop_vcf, - anno_target, - self.tmp_dir, - self.shards, - self.pop_vcf_contigs, - self.fai_data, - self.dry_run, + hybrid_transfer = pathlib.Path( + str( + files("sentieon_cli.scripts").joinpath( + "hybrid_transfer.py" + ) + ) + ).resolve() + transfer_workers = min( + HYBRID_TRANSFER_WORKERS, + max(1, self.cores - 1), + ) + transfer_job = Job( + cmds.cmd_pyexec_hybrid_transfer( + out_vcf=transfer_target, + raw_vcf=anno_target, + population_vcf=self.pop_vcf, + reference_fai=ref_fai, + temp_dir=self.tmp_dir, + hybrid_transfer=hybrid_transfer, + threads=self.cores, + workers=transfer_workers, + ), + "population-transfer", self.cores, ) input_to_apply = transfer_target @@ -1056,8 +1067,7 @@ def call_variants( concat_job, rm_job5, anno_job, - transfer_jobs, - transfer_concat_job, + transfer_job, None, None, ) @@ -1089,7 +1099,7 @@ def call_variants( exclude_homref=not self.gvcf, ), "final-norm", - 0, + min(FINAL_NORM_PROCESSES, self.cores), ) return ( call_job, @@ -1109,8 +1119,7 @@ def call_variants( concat_job, rm_job5, anno_job, - transfer_jobs, - transfer_concat_job, + transfer_job, apply_job, norm_job, ) diff --git a/sentieon_cli/scripts/hybrid_anno.py b/sentieon_cli/scripts/hybrid_anno.py index 7a4ba81..ec8f518 100644 --- a/sentieon_cli/scripts/hybrid_anno.py +++ b/sentieon_cli/scripts/hybrid_anno.py @@ -1,109 +1,529 @@ #!/usr/bin/env python -from __future__ import print_function +from __future__ import annotations + import argparse -import math +import bisect +import collections import multiprocessing as mp +import os +import pathlib +import re +import shutil +import subprocess import sys +import tempfile +from typing import BinaryIO, Iterable, Iterator, Sequence + import vcflib -import os -import bisect -from vcflib.compat import * +from vcflib import bgzf, tabix -extra_headers = ( +EXTRA_HEADERS = ( '##INFO=', ) -remove_headers = ( -) +DEFAULT_STEP_SIZE = 10_000_000 +MAX_ANNOTATION_WORKERS = 96 +MAX_COMPRESSION_THREADS = 32 + +_worker_vcf_path: str | None = None +_worker_bed_intervals: dict[str, list[tuple[int, int, int]]] | None = None +_worker_temp_dir: str | None = None +_worker_vcf: BinaryIO | None = None +_worker_index: tabix.Tabix | None = None + + +class HybridAnnoError(RuntimeError): + pass + + +def load_bed(path: pathlib.Path) -> dict[str, list[tuple[int, int, int]]]: + if not path.exists(): + raise HybridAnnoError(f"input bed file {path} does not exist") + + intervals: dict[str, list[tuple[int, int, int]]] = {} + with path.open() as bed_file: + for line_number, line in enumerate(bed_file, 1): + columns = line.rstrip().split("\t") + if len(columns) < 4: + raise HybridAnnoError( + f"wrong format in {path} line {line_number}: {line.rstrip()}" + ) + try: + start = int(columns[1]) + end = int(columns[2]) + count = int(columns[3]) + except ValueError as error: + raise HybridAnnoError( + f"non-integer BED value in {path} line {line_number}" + ) from error + if start < 0 or end < start: + raise HybridAnnoError( + f"invalid BED interval in {path} line {line_number}" + ) + interval = (start, end, count) + contig_intervals = intervals.setdefault(columns[0], []) + if contig_intervals and interval < contig_intervals[-1]: + raise HybridAnnoError( + f"BED intervals are not sorted in {path} line {line_number}" + ) + contig_intervals.append(interval) + return intervals + + +def copy_header_lines(input_vcf: vcflib.VCF) -> list[str]: + headers: collections.OrderedDict[ + str, collections.OrderedDict[str | None, str] + ] = collections.OrderedDict() + pattern = re.compile(r"^##([^=]+)=()?") + + for line in input_vcf.headers: + match = pattern.match(line) + if match is None: + field, identifier = line, None + else: + field, identifier = match.group(1), match.group(3) + headers.setdefault(field, collections.OrderedDict())[identifier] = line + + for line in EXTRA_HEADERS: + match = pattern.match(line) + if match is None: + continue + field, identifier = match.group(1), match.group(3) + headers.setdefault(field, collections.OrderedDict())[identifier] = line + + output = [ + line + for field_headers in headers.values() + for line in field_headers.values() + if not line.startswith("#CHROM") + ] + columns = ["#CHROM", "POS", "ID", "REF", "ALT", "QUAL", "FILTER", "INFO"] + if input_vcf.samples: + columns.append("FORMAT") + columns.extend(input_vcf.samples) + output.append("\t".join(columns)) + return output + + +def load_vcf_metadata( + path: pathlib.Path, +) -> tuple[bytes, list[tuple[str, int]]]: + if not path.exists(): + raise HybridAnnoError(f"input file {path} does not exist") + if not str(path).endswith(".gz"): + raise HybridAnnoError("input VCF must be BGZF-compressed") + if not pathlib.Path(f"{path}.tbi").exists() and not pathlib.Path( + f"{path}.csi" + ).exists(): + raise HybridAnnoError(f"input VCF index is missing for {path}") + + input_vcf = vcflib.VCF(str(path), "r") + try: + contigs: list[tuple[str, int]] = [] + for contig, description in input_vcf.contigs.items(): + if "length" not in description: + raise HybridAnnoError( + f"contig {contig} has no length in the VCF header" + ) + try: + length = int(description["length"]) + except ValueError as error: + raise HybridAnnoError( + f"contig {contig} has an invalid length" + ) from error + if length <= 0: + raise HybridAnnoError( + f"contig {contig} has a non-positive length" + ) + contigs.append((contig, length)) + if not contigs: + raise HybridAnnoError("input VCF declares no contigs") + header = ("\n".join(copy_header_lines(input_vcf)) + "\n").encode() + finally: + input_vcf.close() + return header, contigs + + +def cut_shards( + contigs: Sequence[tuple[str, int]], step_size: int +) -> list[tuple[int, str, int, int]]: + shards: list[tuple[int, str, int, int]] = [] + filled = 0 + for contig, length in contigs: + start = 0 + while start < length: + chunk_size = min(length - start, step_size - filled) + end = start + chunk_size + shards.append((len(shards), contig, start, end)) + start = end + filled += chunk_size + if filled == step_size: + filled = 0 + return shards + + +def init_worker( + vcf_path: str, + bed_intervals: dict[str, list[tuple[int, int, int]]], + temp_dir: str, +) -> None: + global _worker_vcf_path + global _worker_bed_intervals + global _worker_temp_dir + global _worker_vcf + global _worker_index + + _worker_vcf_path = vcf_path + _worker_bed_intervals = bed_intervals + _worker_temp_dir = temp_dir + _worker_vcf = bgzf.open(vcf_path, "rb") + _worker_index = tabix.Tabix(vcf_path, "r") + + +def close_worker() -> None: + global _worker_vcf + if _worker_vcf is not None: + _worker_vcf.close() + _worker_vcf = None -def annotate_lhc(invcf, outvcf, ctg2s_e_c): - start = getattr(invcf, 'start', -1) - bed_ctg = '' - bed_idx = 0 - for v in invcf: - ctg = v.chrom - pos = v.pos - hap_cnt = -1 - if (bed_ctg != ctg) : - if ctg in ctg2s_e_c: - bed_ctg = ctg - bed_idx = bisect.bisect_left(ctg2s_e_c[ctg], (pos, pos+1, 0)) - if (bed_idx > 0): - bed_idx -= 1 - if (bed_ctg == ctg) : - while (bed_idx < len(ctg2s_e_c[ctg])): - if (ctg2s_e_c[ctg][bed_idx][1] <= pos): - bed_idx += 1 - else : + +def parse_record(line: bytes) -> tuple[list[bytes], str, int, int]: + stripped = line.rstrip() + columns = stripped.split(b"\t", 8) + if len(columns) < 8: + raise HybridAnnoError( + f"VCF record has fewer than eight columns: {stripped!r}" + ) + try: + contig = columns[0].decode() + position = int(columns[1]) - 1 + except (UnicodeDecodeError, ValueError) as error: + raise HybridAnnoError(f"invalid VCF record: {stripped!r}") from error + if position < 0: + raise HybridAnnoError(f"invalid VCF position: {stripped!r}") + + record_end = position + len(columns[3]) + for item in columns[7].split(b";"): + if item.startswith(b"END="): + try: + record_end = int(item[4:]) + except ValueError as error: + raise HybridAnnoError( + f"invalid INFO/END in VCF record: {stripped!r}" + ) from error + return columns, contig, position, record_end + + +def iter_shard_records( + contig: str, start: int, end: int +) -> Iterator[tuple[list[bytes], int]]: + if _worker_vcf is None or _worker_index is None: + raise HybridAnnoError("annotation worker was not initialized") + + ranges = list(_worker_index.query(contig, start, end)) + sequential = False + for range_start, range_end in ranges: + _worker_vcf.seek(range_start) + while True: + line = _worker_vcf.readline() + if not line: + return + if line.startswith(b"#"): + continue + columns, record_contig, position, record_end = parse_record(line) + if record_contig != contig or position >= end: + return + if not sequential and record_end <= start: + if _worker_vcf.tell() >= range_end: break - if (bed_idx < len(ctg2s_e_c[ctg])): - if (ctg2s_e_c[ctg][bed_idx][0] <= pos): - hap_cnt = ctg2s_e_c[ctg][bed_idx][2] - - if hap_cnt != -1: - cols = v.line.split('\t') - cols[7] += ';LHC='+str(hap_cnt) - v.line='\t'.join(cols) - if v.pos >= start: - outvcf.emit(v) - -def main(args): - if not os.path.exists(args.bed): - print('Error: input bed file %s does not exist' % args.bed) - return -1 - - ctg2s_e_c=dict() - - with open(args.bed, 'r') as bedf: - for line in bedf: - cols = line.rstrip().split('\t') - if len(cols) < 4: - print('Error: wrong format in line %s' % line) - return -1 - ctg = cols[0] - if ctg not in ctg2s_e_c: - ctg2s_e_c[ctg] = list() - ctg2s_e_c[ctg].append((int(cols[1]), int(cols[2]), int(cols[3]))) - - if not os.path.exists(args.vcf): - print('Error: input file %s does not exist' % args.vcf) - return -1 - - invcf = vcflib.VCF(args.vcf, 'r') - - outvcf = vcflib.VCF(args.output, 'w') - outvcf.copy_header(invcf, extra_headers, remove_headers) - outvcf.emit_header() - - if args.threads < 2: - annotate_lhc(invcf, outvcf, ctg2s_e_c) + continue + sequential = True + yield columns, position + if sequential: + return + + +def annotate_shard( + shard: tuple[int, str, int, int], +) -> tuple[int, str, int, int]: + shard_number, contig, start, end = shard + if _worker_bed_intervals is None or _worker_temp_dir is None: + raise HybridAnnoError("annotation worker was not initialized") + + output_path = os.path.join( + _worker_temp_dir, f"fragment.{shard_number:08d}.vcf" + ) + contig_intervals = _worker_bed_intervals.get(contig) + bed_index = 0 + bed_initialized = False + record_count = 0 + annotated_count = 0 + + with open(output_path, "wb") as output: + for columns, position in iter_shard_records(contig, start, end): + hap_count = -1 + if contig_intervals: + if not bed_initialized: + bed_index = bisect.bisect_left( + contig_intervals, (position, position + 1, 0) + ) + if bed_index > 0: + bed_index -= 1 + bed_initialized = True + while ( + bed_index < len(contig_intervals) + and contig_intervals[bed_index][1] <= position + ): + bed_index += 1 + if ( + bed_index < len(contig_intervals) + and contig_intervals[bed_index][0] <= position + ): + hap_count = contig_intervals[bed_index][2] + + if position < start: + continue + if hap_count != -1: + columns[7] += f";LHC={hap_count}".encode() + annotated_count += 1 + output.write(b"\t".join(columns)) + output.write(b"\n") + record_count += 1 + + return shard_number, output_path, record_count, annotated_count + + +def generate_fragments( + shards: Sequence[tuple[int, str, int, int]], + workers: int, + vcf_path: pathlib.Path, + bed_intervals: dict[str, list[tuple[int, int, int]]], + temp_dir: pathlib.Path, +) -> tuple[list[pathlib.Path], int, int]: + results: Iterable[tuple[int, str, int, int]] + initializer_arguments = (str(vcf_path), bed_intervals, str(temp_dir)) + + if workers == 1: + init_worker(*initializer_arguments) + try: + results = [annotate_shard(shard) for shard in shards] + finally: + close_worker() else: - sharder = vcflib.Sharder(args.threads) + context = mp.get_context() + with context.Pool( + processes=workers, + initializer=init_worker, + initargs=initializer_arguments, + ) as pool: + results = pool.imap(annotate_shard, shards, chunksize=1) + results = list(results) + + ordered = sorted(results) + expected_numbers = list(range(len(shards))) + observed_numbers = [result[0] for result in ordered] + if observed_numbers != expected_numbers: + raise HybridAnnoError("annotation shards were not returned exactly once") + paths = [pathlib.Path(result[1]) for result in ordered] + record_count = sum(result[2] for result in ordered) + annotated_count = sum(result[3] for result in ordered) + return paths, record_count, annotated_count + + +def compress_fragments( + bgzip_path: str, + header: bytes, + fragments: Sequence[pathlib.Path], + threads: int, + output_path: pathlib.Path, +) -> None: + with output_path.open("wb") as compressed_output: + process = subprocess.Popen( + [bgzip_path, "-@", str(threads), "-c"], + stdin=subprocess.PIPE, + stdout=compressed_output, + stderr=subprocess.PIPE, + ) + if process.stdin is None or process.stderr is None: + process.kill() + raise HybridAnnoError("failed to open bgzip pipes") try: - contig_lengths = [ - (contig, 0, int(d["length"])) - for contig, d in invcf.contigs.items() - ] - except KeyError: - return ( - "This script requires a VCF with contig lengths in" - " the header when using multiple threads" + process.stdin.write(header) + for fragment in fragments: + with fragment.open("rb") as fragment_file: + shutil.copyfileobj( + fragment_file, process.stdin, length=16 * 1024 * 1024 + ) + process.stdin.close() + except BaseException: + process.kill() + process.wait() + raise + error_output = process.stderr.read().decode(errors="replace") + return_code = process.wait() + if return_code != 0: + raise HybridAnnoError( + f"bgzip failed with exit code {return_code}: {error_output.strip()}" + ) + + +def publish_output( + output_path: pathlib.Path, + header: bytes, + fragments: Sequence[pathlib.Path], + compression_threads: int, + bgzip_path: str, + tabix_path: str, +) -> None: + output_parent = output_path.parent + if not output_parent.is_dir(): + raise HybridAnnoError( + f"output directory {output_parent} does not exist" + ) + + descriptor, partial_name = tempfile.mkstemp( + prefix=f".{output_path.name}.partial.", + suffix=".vcf.gz", + dir=output_parent, + ) + os.close(descriptor) + partial_path = pathlib.Path(partial_name) + partial_index = pathlib.Path(f"{partial_path}.tbi") + try: + compress_fragments( + bgzip_path, + header, + fragments, + compression_threads, + partial_path, + ) + result = subprocess.run( + [tabix_path, "-f", "-p", "vcf", str(partial_path)], + capture_output=True, + text=True, + ) + if result.returncode != 0: + raise HybridAnnoError( + "tabix failed with exit code " + f"{result.returncode}: {result.stderr.strip()}" ) - shards = sharder.cut(contig_lengths, args.step_size) - sharder.run(shards, annotate_lhc, [], invcf, outvcf, ctg2s_e_c) + os.replace(partial_path, output_path) + os.replace(partial_index, pathlib.Path(f"{output_path}.tbi")) + stale_csi = pathlib.Path(f"{output_path}.csi") + if stale_csi.exists(): + stale_csi.unlink() + finally: + for temporary_path in (partial_path, partial_index): + if temporary_path.exists(): + temporary_path.unlink() - outvcf.close() - invcf.close() + +def main(args: argparse.Namespace) -> int: + if args.threads < 1: + raise HybridAnnoError("--threads must be at least 1") + if args.step_size is not None and args.step_size < 1: + raise HybridAnnoError("--step-size must be at least 1") + + input_path = pathlib.Path(args.vcf).resolve() + output_path = pathlib.Path(args.output).resolve() + bed_path = pathlib.Path(args.bed).resolve() + if input_path == output_path: + raise HybridAnnoError("input and output VCF paths must differ") + if not str(output_path).endswith(".gz"): + raise HybridAnnoError("output VCF must end in .gz") + + bgzip_path = shutil.which("bgzip") + tabix_path = shutil.which("tabix") + if bgzip_path is None: + raise HybridAnnoError("bgzip is not available on PATH") + if tabix_path is None: + raise HybridAnnoError("tabix is not available on PATH") + + bed_intervals = load_bed(bed_path) + header, contigs = load_vcf_metadata(input_path) + workers = min(args.threads, MAX_ANNOTATION_WORKERS) + compression_threads = min(args.threads, MAX_COMPRESSION_THREADS) + step_size = args.step_size or DEFAULT_STEP_SIZE + shards = cut_shards(contigs, step_size) + + temp_dir = pathlib.Path( + tempfile.mkdtemp( + prefix="hybrid-anno.", + dir=os.getenv("SENTIEON_TMPDIR"), + ) + ) + success = False + try: + print( + "hybrid_anno: " + f"workers={workers} compression_threads={compression_threads} " + f"step_size={step_size} shards={len(shards)} " + f"tmpdir={temp_dir}", + file=sys.stderr, + ) + fragments, record_count, annotated_count = generate_fragments( + shards, + workers, + input_path, + bed_intervals, + temp_dir, + ) + publish_output( + output_path, + header, + fragments, + compression_threads, + bgzip_path, + tabix_path, + ) + print( + "hybrid_anno: " + f"records={record_count} annotated={annotated_count}", + file=sys.stderr, + ) + success = True + finally: + if success: + shutil.rmtree(temp_dir) + else: + print( + f"hybrid_anno: preserved failed work directory {temp_dir}", + file=sys.stderr, + ) return 0 -if __name__ == '__main__': - parser = argparse.ArgumentParser(prog='sentieon pyexec hybrid_anno.py', usage='%(prog)s [options] -v VCF -b BED output') - parser.add_argument('output', help='Output vcf file name') - parser.add_argument('-v','--vcf',required=True, help='Input vcf file name') - parser.add_argument('-b','--bed',required=True, help='region haplotype count from long reads') - parser.add_argument('-t','--threads',type=int,default=mp.cpu_count(),help='number of threads') - parser.add_argument('--step-size',type=int,default=10*1000*1000,help=argparse.SUPPRESS) - sys.exit(main(parser.parse_args())) -# vim: ts=4 sw=4 expandtab +if __name__ == "__main__": + parser = argparse.ArgumentParser( + prog="sentieon pyexec hybrid_anno.py", + usage="%(prog)s [options] -v VCF -b BED output", + ) + parser.add_argument("output", help="Output vcf file name") + parser.add_argument( + "-v", "--vcf", required=True, help="Input vcf file name" + ) + parser.add_argument( + "-b", + "--bed", + required=True, + help="region haplotype count from long reads", + ) + parser.add_argument( + "-t", + "--threads", + type=int, + default=mp.cpu_count(), + help="number of threads", + ) + parser.add_argument( + "--step-size", + type=int, + default=None, + help=argparse.SUPPRESS, + ) + try: + sys.exit(main(parser.parse_args())) + except HybridAnnoError as error: + print(f"Error: {error}", file=sys.stderr) + sys.exit(1) diff --git a/sentieon_cli/scripts/hybrid_select.py b/sentieon_cli/scripts/hybrid_select.py index 9d69ed2..c3f24a8 100644 --- a/sentieon_cli/scripts/hybrid_select.py +++ b/sentieon_cli/scripts/hybrid_select.py @@ -1,157 +1,480 @@ #!/usr/bin/env python -from __future__ import print_function +from __future__ import annotations + import argparse -import math import multiprocessing as mp +import os +import pathlib +import shutil import sys +import tempfile +from typing import BinaryIO, Iterable, Iterator, Sequence + import vcflib -import os -from vcflib.compat import * - -class HybridFilter: - - params = { - 'min_conf_longread' : [ 30., "Minimum call confidence in long read", ""], - 'min_depth_longread' : [ 2, "Minimum depth in long read", ""], - 'min_conf_shortread' : [ 30., "Minimum call confidence in short read", ""], - } - - @classmethod - def add_arguments(cls, parser): - for k,v in cls.params.items(): - parser.add_argument('--'+k, default=v[0], type=type(v[0]), help=v[1] + ' (default: '+str(v[0])+')', metavar=v[2]) - - def __init__(self, args): - self.args = args - - def applyFilters(self, in_vcf, out_vcf): - start = getattr(in_vcf, 'start', -1) - for v in in_vcf: - filters = [] - lad = v.samples[0].get('LAD') - sad = v.samples[0].get('SAD') - lpl = v.samples[0].get('LPL') - spl = v.samples[0].get('SPL') - - if not lad: - lad=[0,0,0] - if not lpl: - lpl=[0,0,0] - if not spl: - spl=[0,0,0] - - lminpos = lpl.index(0) #same as GT - sminpos = spl.index(0) - lpl_ref = lpl[0] - lpl.remove(0) - lpl_conf = min(lpl) #2nd min - spl.remove(0) - spl_conf = min(spl) #2nd min - - if not filters: - if sum(lad) < self.args.min_depth_longread: - filters.append('longread_lowdepth') - - if not filters: - if ((v.info.get('STR') is None) and (lminpos != 0)): - lpl_conf = lpl_ref - if ( lpl_conf < self.args.min_conf_longread): - filters.append('longread_lowconf') - - if not filters: - if ( spl_conf >= self.args.min_conf_shortread): - if (lminpos == sminpos): - filters.append('same_gt') - elif (lminpos == 0) and (v.info.get('STR') is not None): - filters.append('longread_str') - - flds = v.line.split('\t') - flds[6] = filters and ';'.join(sorted(set(filters))) or 'PASS' - v.line = '\t'.join(flds) - if v.pos >= start: - out_vcf.emit(v) - return - -extra_headers = ( - '##FILTER=', - '##FILTER=', - '##FILTER=', - '##FILTER=', - '##FILTER=', - '##FILTER=', -) - -remove_headers = ( - '##source=*', -) - -expect_types = { - 'GT': {'Number': '1', 'Type': 'String' }, -} - -def check_header_types(vcf): - vcf_types = {} - for k, v in iteritems(vcf.infos): - if k in expect_types: - dd = {} - for kk, vv in iteritems(v): - if kk in ('Number', 'Type'): - dd[kk] = vv - vcf_types[k] = dd - for k, v in iteritems(vcf.formats): - if k in expect_types: - dd = {} - for kk, vv in iteritems(v): - if kk in ('Number', 'Type'): - dd[kk] = vv - vcf_types[k] = dd - return vcf_types == expect_types - -def main(args): - if not os.path.exists(args.vcf): - print('Error: input file %s does not exist' % args.vcf) - return -1 - - invcf = vcflib.VCF(args.vcf, 'r') - - if not check_header_types(invcf): - print('Error: vcf format is not expected') - return -1 - - filter = HybridFilter(args) - - outvcf = vcflib.VCF(args.output, 'w') - outvcf.copy_header(invcf, extra_headers, remove_headers) - outvcf.emit_header() - - if args.threads < 2: - filter.applyFilters(invcf, outvcf) - else: - sharder = vcflib.Sharder(args.threads) +from vcflib import bgzf, tabix + +DEFAULT_STEP_SIZE = 10_000_000 +DEFAULT_SLOP_SIZE = 1_000 +MAX_SELECT_WORKERS = 96 + +_worker_vcf: BinaryIO | None = None +_worker_index: tabix.Tabix | None = None +_worker_fai_lengths: dict[str, int] | None = None +_worker_temp_dir: str | None = None +_worker_slop_size: int | None = None + + +class HybridSelectError(RuntimeError): + pass + + +def load_fai(path: pathlib.Path) -> dict[str, int]: + if not path.is_file(): + raise HybridSelectError(f"reference index {path} does not exist") + + lengths: dict[str, int] = {} + with path.open() as reference_index: + for line_number, line in enumerate(reference_index, 1): + columns = line.rstrip().split("\t") + if len(columns) < 2: + raise HybridSelectError( + f"reference index line {line_number} has fewer than two fields" + ) + try: + length = int(columns[1]) + except ValueError as error: + raise HybridSelectError( + f"reference index line {line_number} has an invalid length" + ) from error + if length <= 0: + raise HybridSelectError( + f"reference index line {line_number} has a non-positive length" + ) + if columns[0] in lengths: + raise HybridSelectError( + f"reference index contains duplicate contig {columns[0]}" + ) + lengths[columns[0]] = length + if not lengths: + raise HybridSelectError("reference index declares no contigs") + return lengths + + +def load_vcf_contigs(path: pathlib.Path) -> list[tuple[str, int]]: + if not path.is_file(): + raise HybridSelectError(f"input file {path} does not exist") + if not str(path).endswith(".gz"): + raise HybridSelectError("input VCF must be BGZF-compressed") + if ( + not pathlib.Path(f"{path}.tbi").is_file() + and not pathlib.Path(f"{path}.csi").is_file() + ): + raise HybridSelectError(f"input VCF index is missing for {path}") + + input_vcf = vcflib.VCF(str(path), "r") + try: + genotype = input_vcf.formats.get("GT") + observed_type = None + if genotype is not None: + observed_type = { + key: value + for key, value in genotype.items() + if key in ("Number", "Type") + } + expected_type = {"Number": "1", "Type": "String"} + if observed_type != expected_type: + raise HybridSelectError( + "VCF FORMAT/GT is not Number=1,Type=String" + ) + if not input_vcf.samples: + raise HybridSelectError("input VCF contains no samples") + + contigs: list[tuple[str, int]] = [] + for contig, description in input_vcf.contigs.items(): + if "length" not in description: + raise HybridSelectError( + f"contig {contig} has no length in the VCF header" + ) + try: + length = int(description["length"]) + except ValueError as error: + raise HybridSelectError( + f"contig {contig} has an invalid length" + ) from error + if length <= 0: + raise HybridSelectError( + f"contig {contig} has a non-positive length" + ) + contigs.append((contig, length)) + if not contigs: + raise HybridSelectError("input VCF declares no contigs") + return contigs + finally: + input_vcf.close() + + +def cut_shards( + contigs: Sequence[tuple[str, int]], step_size: int +) -> list[tuple[int, str, int, int]]: + shards: list[tuple[int, str, int, int]] = [] + filled = 0 + for contig, length in contigs: + start = 0 + while start < length: + chunk_size = min(length - start, step_size - filled) + end = start + chunk_size + shards.append((len(shards), contig, start, end)) + start = end + filled += chunk_size + if filled == step_size: + filled = 0 + return shards + + +def parse_integer_list(value: bytes | None, field: str) -> list[int] | None: + if value is None or value in (b"", b"."): + return None + try: + return [int(item) for item in value.split(b",")] + except ValueError as error: + raise HybridSelectError( + f"invalid FORMAT/{field} value {value!r}" + ) from error + + +def sample_fields(columns: list[bytes]) -> dict[bytes, bytes]: + if len(columns) < 10: + raise HybridSelectError("VCF record contains no sample column") + keys = columns[8].split(b":") + values = columns[9].split(b":") + return dict(zip(keys, values)) + + +def has_info_field(info: bytes, name: bytes) -> bool: + return any(item.split(b"=", 1)[0] == name for item in info.split(b";")) + + +def record_passes(columns: list[bytes]) -> bool: + sample = sample_fields(columns) + lad = parse_integer_list(sample.get(b"LAD"), "LAD") or [0, 0, 0] + lpl = parse_integer_list(sample.get(b"LPL"), "LPL") or [0, 0, 0] + spl = parse_integer_list(sample.get(b"SPL"), "SPL") or [0, 0, 0] + + try: + lminpos = lpl.index(0) + sminpos = spl.index(0) + lpl_ref = lpl[0] + lpl = lpl.copy() + spl = spl.copy() + lpl.remove(0) + spl.remove(0) + lpl_conf = min(lpl) + spl_conf = min(spl) + except (ValueError, IndexError) as error: + raise HybridSelectError( + "invalid LPL/SPL genotype likelihood vector" + ) from error + + if sum(lad) < 2: + return False + + is_str = has_info_field(columns[7], b"STR") + if not is_str and lminpos != 0: + lpl_conf = lpl_ref + if lpl_conf < 30.0: + return False + + if spl_conf >= 30.0: + if lminpos == sminpos: + return False + if lminpos == 0 and is_str: + return False + return True + + +def parse_record(line: bytes) -> tuple[list[bytes], str, int, int]: + stripped = line.rstrip() + columns = stripped.split(b"\t") + if len(columns) < 10: + raise HybridSelectError( + f"VCF record has fewer than ten columns: {stripped!r}" + ) + try: + contig = columns[0].decode() + position = int(columns[1]) - 1 + except (UnicodeDecodeError, ValueError) as error: + raise HybridSelectError(f"invalid VCF record: {stripped!r}") from error + if position < 0: + raise HybridSelectError(f"invalid VCF position: {stripped!r}") + + record_end = position + len(columns[3]) + for item in columns[7].split(b";"): + if item.startswith(b"END="): + try: + record_end = int(item[4:]) + except ValueError as error: + raise HybridSelectError( + f"invalid INFO/END in VCF record: {stripped!r}" + ) from error + if record_end <= position: + raise HybridSelectError(f"invalid VCF record end: {stripped!r}") + return columns, contig, position, record_end + + +def init_worker( + vcf_path: str, + fai_lengths: dict[str, int], + temp_dir: str, + slop_size: int, +) -> None: + global _worker_vcf + global _worker_index + global _worker_fai_lengths + global _worker_temp_dir + global _worker_slop_size + + _worker_vcf = bgzf.open(vcf_path, "rb") + _worker_index = tabix.Tabix(vcf_path, "r") + _worker_fai_lengths = fai_lengths + _worker_temp_dir = temp_dir + _worker_slop_size = slop_size + + +def close_worker() -> None: + global _worker_vcf + if _worker_vcf is not None: + _worker_vcf.close() + _worker_vcf = None + + +def iter_shard_records( + contig: str, start: int, end: int +) -> Iterator[tuple[list[bytes], int, int]]: + if _worker_vcf is None or _worker_index is None: + raise HybridSelectError("selection worker was not initialized") + + ranges = list(_worker_index.query(contig, start, end)) + sequential = False + for range_start, range_end in ranges: + _worker_vcf.seek(range_start) + while True: + line = _worker_vcf.readline() + if not line: + return + if line.startswith(b"#"): + continue + columns, record_contig, position, record_end = parse_record(line) + if record_contig != contig or position >= end: + return + if not sequential and record_end <= start: + if _worker_vcf.tell() >= range_end: + break + continue + sequential = True + if position >= start: + yield columns, position, record_end + if sequential: + return + + +def select_shard( + shard: tuple[int, str, int, int], +) -> tuple[int, str, int, int]: + shard_number, contig, start, end = shard + if ( + _worker_fai_lengths is None + or _worker_temp_dir is None + or _worker_slop_size is None + ): + raise HybridSelectError("selection worker was not initialized") + if contig not in _worker_fai_lengths: + raise HybridSelectError( + f"contig {contig} is absent from the reference index" + ) + + output_path = os.path.join( + _worker_temp_dir, f"fragment.{shard_number:08d}.bed" + ) + record_count = 0 + selected_count = 0 + contig_length = _worker_fai_lengths[contig] + with open(output_path, "wb") as output: + for columns, position, record_end in iter_shard_records( + contig, start, end + ): + record_count += 1 + if not record_passes(columns): + continue + bed_start = max(0, position - _worker_slop_size) + bed_end = min(contig_length, record_end + _worker_slop_size) + output.write(f"{contig}\t{bed_start}\t{bed_end}\n".encode()) + selected_count += 1 + return shard_number, output_path, record_count, selected_count + + +def generate_fragments( + shards: Sequence[tuple[int, str, int, int]], + workers: int, + vcf_path: pathlib.Path, + fai_lengths: dict[str, int], + temp_dir: pathlib.Path, + slop_size: int, +) -> tuple[list[pathlib.Path], int, int]: + initializer_arguments = ( + str(vcf_path), + fai_lengths, + str(temp_dir), + slop_size, + ) + results: Iterable[tuple[int, str, int, int]] + if workers == 1: + init_worker(*initializer_arguments) try: - contig_lengths = [ - (contig, 0, int(d["length"])) - for contig, d in invcf.contigs.items() - ] - except KeyError: - return ( - "This script requires a VCF with contig lengths in" - " the header when using multiple threads" + results = [select_shard(shard) for shard in shards] + finally: + close_worker() + else: + context = mp.get_context() + with context.Pool( + processes=workers, + initializer=init_worker, + initargs=initializer_arguments, + ) as pool: + results = list(pool.imap(select_shard, shards, chunksize=1)) + + ordered = sorted(results) + if [item[0] for item in ordered] != list(range(len(shards))): + raise HybridSelectError( + "selection shards were not returned exactly once" + ) + return ( + [pathlib.Path(item[1]) for item in ordered], + sum(item[2] for item in ordered), + sum(item[3] for item in ordered), + ) + + +def publish_output( + output_path: pathlib.Path, fragments: Sequence[pathlib.Path] +) -> None: + if not output_path.parent.is_dir(): + raise HybridSelectError( + f"output directory {output_path.parent} does not exist" + ) + descriptor, partial_name = tempfile.mkstemp( + prefix=f".{output_path.name}.partial.", + suffix=".bed", + dir=output_path.parent, + ) + os.close(descriptor) + partial_path = pathlib.Path(partial_name) + try: + with partial_path.open("wb") as output: + for fragment in fragments: + with fragment.open("rb") as input_fragment: + shutil.copyfileobj( + input_fragment, output, length=16 * 1024 * 1024 + ) + os.replace(partial_path, output_path) + finally: + if partial_path.exists(): + partial_path.unlink() + + +def main(args: argparse.Namespace) -> int: + if args.threads < 1: + raise HybridSelectError("--threads must be at least 1") + if args.step_size < 1: + raise HybridSelectError("--step-size must be at least 1") + if args.slop_size < 0: + raise HybridSelectError("--slop-size must not be negative") + + input_path = pathlib.Path(args.vcf).resolve() + output_path = pathlib.Path(args.output).resolve() + reference_index = pathlib.Path(args.reference_fai).resolve() + contigs = load_vcf_contigs(input_path) + fai_lengths = load_fai(reference_index) + for contig, _length in contigs: + if contig not in fai_lengths: + raise HybridSelectError( + f"contig {contig} is absent from the reference index" ) - shards = sharder.cut(contig_lengths, args.step_size) - sharder.run(shards, filter.applyFilters, [], invcf, outvcf) - outvcf.close() - invcf.close() + workers = min(args.threads, MAX_SELECT_WORKERS) + shards = cut_shards(contigs, args.step_size) + temp_dir = pathlib.Path( + tempfile.mkdtemp( + prefix="hybrid-select.", + dir=os.getenv("SENTIEON_TMPDIR"), + ) + ) + success = False + try: + print( + "hybrid_select: " + f"workers={workers} step_size={args.step_size} " + f"shards={len(shards)} slop_size={args.slop_size} " + f"tmpdir={temp_dir}", + file=sys.stderr, + ) + fragments, record_count, selected_count = generate_fragments( + shards, + workers, + input_path, + fai_lengths, + temp_dir, + args.slop_size, + ) + publish_output(output_path, fragments) + print( + f"hybrid_select: records={record_count} selected={selected_count}", + file=sys.stderr, + ) + success = True + finally: + if success: + shutil.rmtree(temp_dir) + else: + print( + f"hybrid_select: preserved failed work directory {temp_dir}", + file=sys.stderr, + ) return 0 -if __name__ == '__main__': - parser = argparse.ArgumentParser(prog='sentieon pyexec hybrid_select.py', usage='%(prog)s [options] -v VCF output.vcf.gz') - parser.add_argument('output', help='Output vcf file name') - parser.add_argument('-v','--vcf',required=True, help='Input vcf file name') - parser.add_argument('-t','--threads',type=int,default=mp.cpu_count(),help='number of threads') - parser.add_argument('--step-size',type=int,default=10*1000*1000,help=argparse.SUPPRESS) - HybridFilter.add_arguments(parser) - sys.exit(main(parser.parse_args())) -# vim: ts=4 sw=4 expandtab +if __name__ == "__main__": + parser = argparse.ArgumentParser( + prog="sentieon pyexec hybrid_select.py", + usage="%(prog)s [options] -v VCF --reference-fai REF.fai output.bed", + ) + parser.add_argument("output", help="output BED file name") + parser.add_argument( + "-v", "--vcf", required=True, help="input VCF file name" + ) + parser.add_argument( + "--reference-fai", required=True, help="reference FASTA index" + ) + parser.add_argument( + "-t", + "--threads", + type=int, + default=mp.cpu_count(), + help="number of worker processes", + ) + parser.add_argument( + "--step-size", + type=int, + default=DEFAULT_STEP_SIZE, + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--slop-size", + type=int, + default=DEFAULT_SLOP_SIZE, + help=argparse.SUPPRESS, + ) + try: + sys.exit(main(parser.parse_args())) + except HybridSelectError as error: + print(f"Error: {error}", file=sys.stderr) + sys.exit(1) diff --git a/sentieon_cli/scripts/hybrid_transfer.py b/sentieon_cli/scripts/hybrid_transfer.py new file mode 100644 index 0000000..9e2ce7a --- /dev/null +++ b/sentieon_cli/scripts/hybrid_transfer.py @@ -0,0 +1,880 @@ +#!/usr/bin/env python +"""Bounded, ordered population annotation transfer for Hybrid VCFs.""" + +from __future__ import annotations + +import argparse +import multiprocessing as mp +import os +import pathlib +import re +import shutil +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from typing import BinaryIO, Iterator, Sequence, cast + +DEFAULT_STEP_SIZE = 10_000_000 +DEFAULT_WORKERS = 32 +MAX_WORKERS = 64 + +_HEADER_FIELD_PATTERN = re.compile(r'(.*?)=(".*?"|.*?)(?:,|$)') + +_worker_raw_vcf: str | None = None +_worker_population_vcf: str | None = None +_worker_merge_rules: str | None = None +_worker_info_numbers: dict[bytes, bytes] | None = None +_worker_format_numbers: dict[bytes, bytes] | None = None +_worker_temp_dir: str | None = None + + +class HybridTransferError(RuntimeError): + """A population-transfer contract or subprocess failed.""" + + +@dataclass(frozen=True) +class WorkItem: + """One legacy-compatible reference shard or raw-only contig.""" + + number: int + contig: str + start: int + stop: int + merge_population: bool + + +@dataclass(frozen=True) +class WorkResult: + """Header/body fragment produced by one ordered work item.""" + + number: int + contig: str + header_path: str + body_path: str + input_records: int + output_records: int + + +def require_index(path: pathlib.Path, label: str) -> None: + """Require a tabix or CSI index without attempting discovery/fallback.""" + + if not path.is_file(): + raise HybridTransferError(f"{label} {path} does not exist") + tbi = pathlib.Path(f"{path}.tbi") + csi = pathlib.Path(f"{path}.csi") + if not tbi.is_file() and not csi.is_file(): + raise HybridTransferError(f"{label} index is missing for {path}") + + +def run_checked(command: Sequence[str], label: str) -> bytes: + """Run a bounded metadata command and return stdout.""" + + try: + result = subprocess.run( + command, + check=False, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + except FileNotFoundError as error: + raise HybridTransferError( + f"required executable {command[0]} is missing" + ) from error + if result.returncode != 0: + message = result.stderr.decode(errors="replace").strip() + raise HybridTransferError( + f"{label} failed with exit code {result.returncode}: {message}" + ) + return result.stdout + + +def parse_structured_header(line: str) -> dict[str, str]: + """Parse the simple structured-header grammar used by the legacy code.""" + + try: + start = line.index("<") + end = line.index(">") + except ValueError as error: + raise HybridTransferError( + f"malformed structured VCF header: {line}" + ) from error + return dict(_HEADER_FIELD_PATTERN.findall(line[start + 1 : end])) + + +def population_contract( + raw_vcf: pathlib.Path, population_vcf: pathlib.Path +) -> tuple[str, set[str], bytes]: + """Return legacy merge rules, population contigs, and merged header.""" + + population_header = run_checked( + ["bcftools", "view", "--no-version", "-h", str(population_vcf)], + "population header read", + ) + info_fields: list[str] = [] + population_contigs: set[str] = set() + for raw_line in population_header.decode(errors="strict").splitlines(): + if raw_line.startswith("##INFO") and ",Number=A" in raw_line: + fields = parse_structured_header(raw_line) + if "ID" not in fields: + raise HybridTransferError( + f"population INFO header has no ID: {raw_line}" + ) + info_fields.append(fields["ID"]) + elif raw_line.startswith("##contig"): + fields = parse_structured_header(raw_line) + if "ID" not in fields: + raise HybridTransferError( + f"population contig header has no ID: {raw_line}" + ) + population_contigs.add(fields["ID"]) + + merge_rules = ",".join(f"{field}:sum" for field in info_fields) + merged_header = run_checked( + [ + "bcftools", + "merge", + "--print-header", + "--no-version", + "--regions-overlap", + "pos", + "-m", + "all", + "-i", + merge_rules, + str(raw_vcf), + str(population_vcf), + ], + "merged header construction", + ) + if not merged_header.endswith(b"\n"): + raise HybridTransferError( + "merged VCF header is not newline terminated" + ) + return merge_rules, population_contigs, merged_header + + +def parse_number_schemas( + merged_header: bytes, +) -> tuple[dict[bytes, bytes], dict[bytes, bytes]]: + """Collect Number=A/R/G INFO and FORMAT declarations for trimming.""" + + infos: dict[bytes, bytes] = {} + formats: dict[bytes, bytes] = {} + for line in merged_header.decode(errors="strict").splitlines(): + target: dict[bytes, bytes] | None = None + if line.startswith("##INFO"): + target = infos + elif line.startswith("##FORMAT"): + target = formats + if target is None: + continue + fields = parse_structured_header(line) + number = fields.get("Number") + identifier = fields.get("ID") + if identifier is not None and number in ("A", "R", "G"): + target[identifier.encode()] = number.encode() + return infos, formats + + +def load_fai(path: pathlib.Path) -> list[tuple[str, int]]: + """Load reference contigs in their declared order.""" + + if not path.is_file(): + raise HybridTransferError(f"reference index {path} does not exist") + contigs: list[tuple[str, int]] = [] + seen: set[str] = set() + with path.open() as reference_index: + for line_number, line in enumerate(reference_index, 1): + columns = line.rstrip().split("\t") + if len(columns) < 2: + raise HybridTransferError( + f"reference index line {line_number} has fewer than " + "two fields" + ) + contig = columns[0] + try: + length = int(columns[1]) + except ValueError as error: + raise HybridTransferError( + f"reference index line {line_number} has an invalid length" + ) from error + if contig in seen: + raise HybridTransferError( + f"reference index contains duplicate contig {contig}" + ) + if length <= 0: + raise HybridTransferError( + f"reference index line {line_number} has a " + "non-positive length" + ) + seen.add(contig) + contigs.append((contig, length)) + if not contigs: + raise HybridTransferError("reference index declares no contigs") + return contigs + + +def build_work_items( + contigs: Sequence[tuple[str, int]], + population_contigs: set[str], + step_size: int, +) -> list[WorkItem]: + """Reproduce the legacy 10 Mb shard and unusual-contig ownership.""" + + items: list[WorkItem] = [] + for contig, length in contigs: + if contig not in population_contigs: + items.append( + WorkItem( + number=len(items), + contig=contig, + start=0, + stop=length, + merge_population=False, + ) + ) + continue + + start = 1 + while start <= length: + stop = min(start + step_size - 1, length) + items.append( + WorkItem( + number=len(items), + contig=contig, + start=start, + stop=stop, + merge_population=True, + ) + ) + start = stop + 1 + return items + + +def parse_key_value(value: bytes) -> tuple[bytes, bytes | None]: + """Parse one INFO key/value, representing flags with ``None``.""" + + fields = value.split(b"=", 1) + if len(fields) == 2: + return fields[0], fields[1] + return fields[0], None + + +def filter_values(values: bytes, keep: Sequence[bool]) -> bytes: + """Filter a comma-delimited VCF vector using legacy zip truncation.""" + + return b",".join( + value for value, retained in zip(values.split(b","), keep) if retained + ) + + +def trim_record( + line: bytes, + info_numbers: dict[bytes, bytes], + format_numbers: dict[bytes, bytes], +) -> bytes | None: + """Port ``trimalt.py`` record semantics without object construction.""" + + values = line.rstrip().split(b"\t") + if len(values) < 9 or values[8] == b".": + return None + + alts = values[4].split(b",") + keep_alt: list[bool] | None = None + info: list[tuple[bytes, bytes | None]] | None = None + + if b"" in alts: + nonref_index = alts.index(b"") + keep_alt = [True] * (nonref_index + 1) + keep_alt.extend([False] * (len(alts) - nonref_index - 1)) + elif len(values) >= 8 and values[7] != b".": + info = [parse_key_value(item) for item in values[7].split(b";")] + for key, value in info: + if key == b"AF": + if value is None: + raise HybridTransferError("INFO/AF is declared as a flag") + keep_alt = [item != b"." for item in value.split(b",")] + break + + if keep_alt is None: + return None + if all(keep_alt): + return line + + keep_ref = [True, *keep_alt] + if info is None and len(values) >= 8 and values[7] != b".": + info = [parse_key_value(item) for item in values[7].split(b";")] + + if values[4] != b".": + retained_alts = [ + alt for alt, retained in zip(alts, keep_alt) if retained + ] + reference = values[3] + try: + suffix_length = ( + min( + len(reference), + min( + len(alt) + for alt in retained_alts + if alt != b"" + ), + ) + - 1 + ) + except ValueError: + suffix_length = 0 + if suffix_length > 0: + reference = reference[:-suffix_length] + retained_alts = [ + (alt if alt == b"" else alt[:-suffix_length]) + for alt in retained_alts + ] + values[3] = reference + values[4] = b",".join(retained_alts) + + if info is not None: + rewritten_info: list[bytes] = [] + for key, value in info: + number = info_numbers.get(key) if value != b"." else None + if number == b"A": + if value is None: + raise HybridTransferError( + "Number=A INFO/" + f"{key.decode(errors='replace')} is a flag" + ) + rewritten_info.append( + key + b"=" + filter_values(value, keep_alt) + ) + elif number == b"R": + if value is None: + raise HybridTransferError( + "Number=R INFO/" + f"{key.decode(errors='replace')} is a flag" + ) + rewritten_info.append( + key + b"=" + filter_values(value, keep_ref) + ) + elif value is None: + rewritten_info.append(key) + else: + rewritten_info.append(key + b"=" + value) + values[7] = b";".join(rewritten_info) + + if len(values) >= 10 and values[9] != b".": + format_keys = values[8].split(b":") + diploid_genotypes = [ + first and second + for index, first in enumerate(keep_ref) + for second in keep_ref[: index + 1] + ] + for sample_index, sample in enumerate(values[9:], 9): + sample_values = sample.split(b":") + for field_index, value in enumerate(sample_values): + if field_index >= len(format_keys): + break + number = ( + format_numbers.get(format_keys[field_index]) + if value != b"." + else None + ) + if number == b"A": + sample_values[field_index] = filter_values(value, keep_alt) + elif number == b"R": + sample_values[field_index] = filter_values(value, keep_ref) + elif number == b"G": + vector_length = len(value.split(b",")) + if vector_length == len(keep_ref): + keep_genotype = keep_ref + elif vector_length == len(diploid_genotypes): + keep_genotype = diploid_genotypes + else: + name = format_keys[field_index].decode( + errors="replace" + ) + raise HybridTransferError( + f"FORMAT/{name} Number=G vector has " + "unexpected length" + ) + sample_values[field_index] = filter_values( + value, keep_genotype + ) + values[sample_index] = b":".join(sample_values) + + return b"\t".join(values) + b"\n" + + +def init_worker( + raw_vcf: str, + population_vcf: str, + merge_rules: str, + info_numbers: dict[bytes, bytes], + format_numbers: dict[bytes, bytes], + temp_dir: str, +) -> None: + """Initialize immutable process-local transfer state.""" + + global _worker_raw_vcf + global _worker_population_vcf + global _worker_merge_rules + global _worker_info_numbers + global _worker_format_numbers + global _worker_temp_dir + + _worker_raw_vcf = raw_vcf + _worker_population_vcf = population_vcf + _worker_merge_rules = merge_rules + _worker_info_numbers = info_numbers + _worker_format_numbers = format_numbers + _worker_temp_dir = temp_dir + + +def worker_command(item: WorkItem, region_path: pathlib.Path) -> list[str]: + """Build the exact legacy merge or unusual-contig view primitive.""" + + if _worker_raw_vcf is None: + raise HybridTransferError("transfer worker was not initialized") + if item.merge_population: + if _worker_population_vcf is None or _worker_merge_rules is None: + raise HybridTransferError("transfer worker was not initialized") + return [ + "bcftools", + "merge", + "--regions-file", + str(region_path), + "--no-version", + "--regions-overlap", + "pos", + "-m", + "all", + "-i", + _worker_merge_rules, + _worker_raw_vcf, + _worker_population_vcf, + ] + return [ + "bcftools", + "view", + "--no-version", + "--regions-file", + str(region_path), + _worker_raw_vcf, + ] + + +def consume_records( + stream: BinaryIO, + header: BinaryIO, + body: BinaryIO, + trim: bool, +) -> tuple[int, int]: + """Split one bcftools stream and optionally apply exact trimming.""" + + if _worker_info_numbers is None or _worker_format_numbers is None: + raise HybridTransferError("transfer worker was not initialized") + input_records = 0 + output_records = 0 + saw_column_header = False + for line in stream: + if line.startswith(b"#"): + if input_records: + raise HybridTransferError( + "VCF header appeared after record data" + ) + header.write(line) + if line.startswith(b"#CHROM\t"): + saw_column_header = True + continue + if not saw_column_header: + raise HybridTransferError("bcftools output has no #CHROM header") + input_records += 1 + output_line = ( + trim_record(line, _worker_info_numbers, _worker_format_numbers) + if trim + else line + ) + if output_line is not None: + body.write(output_line) + output_records += 1 + if not saw_column_header: + raise HybridTransferError("bcftools output has no #CHROM header") + return input_records, output_records + + +def process_work_item(item: WorkItem) -> WorkResult: + """Run one merge/view shard and emit an uncompressed ordered fragment.""" + + if _worker_temp_dir is None: + raise HybridTransferError("transfer worker was not initialized") + base = pathlib.Path(_worker_temp_dir) + prefix = f"fragment.{item.number:08d}" + region_path = base / f"{prefix}.bed" + header_path = base / f"{prefix}.header.vcf" + body_path = base / f"{prefix}.body.vcf" + error_path = base / f"{prefix}.stderr" + region_path.write_text(f"{item.contig}\t{item.start}\t{item.stop}\n") + + command = worker_command(item, region_path) + try: + with error_path.open("wb") as error_file: + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=error_file, + ) + if process.stdout is None: + process.terminate() + raise HybridTransferError( + "bcftools stdout pipe was not created" + ) + process_stdout = cast(BinaryIO, process.stdout) + try: + with ( + header_path.open("wb") as header, + body_path.open("wb") as body, + ): + input_records, output_records = consume_records( + process_stdout, + header, + body, + trim=item.merge_population, + ) + except BaseException: + process.terminate() + process.wait() + raise + finally: + process_stdout.close() + return_code = process.wait() + except FileNotFoundError as error: + raise HybridTransferError( + "required executable bcftools is missing" + ) from error + + if return_code != 0: + message = error_path.read_text(errors="replace").strip() + raise HybridTransferError( + f"work item {item.number} ({item.contig}) failed with exit code " + f"{return_code}: {message}" + ) + return WorkResult( + number=item.number, + contig=item.contig, + header_path=str(header_path), + body_path=str(body_path), + input_records=input_records, + output_records=output_records, + ) + + +def column_header(path: pathlib.Path) -> bytes: + """Return the #CHROM line for a fragment header.""" + + with path.open("rb") as header: + for line in header: + if line.startswith(b"#CHROM\t"): + return line + raise HybridTransferError(f"fragment header {path} has no #CHROM line") + + +def copy_path(path: pathlib.Path, output: BinaryIO) -> None: + """Copy a fragment with a large bounded userspace buffer.""" + + with path.open("rb") as fragment: + shutil.copyfileobj(fragment, output, length=16 * 1024 * 1024) + + +def partial_paths( + output_path: pathlib.Path, +) -> tuple[pathlib.Path, pathlib.Path]: + """Allocate a unique sibling path accepted by bcftools write-index.""" + + descriptor, partial_name = tempfile.mkstemp( + prefix=f".{output_path.name}.partial.", + suffix=".vcf.gz", + dir=output_path.parent, + ) + os.close(descriptor) + partial_path = pathlib.Path(partial_name) + partial_path.unlink() + return partial_path, pathlib.Path(f"{partial_path}.tbi") + + +def publish_pair( + partial_path: pathlib.Path, + partial_index: pathlib.Path, + output_path: pathlib.Path, +) -> None: + """Replace the completed VCF/index pair only after both validate.""" + + if not partial_path.is_file() or not partial_index.is_file(): + raise HybridTransferError("bcftools did not create the VCF/index pair") + run_checked( + ["bcftools", "index", "--nrecords", str(partial_path)], + "output index validation", + ) + output_index = pathlib.Path(f"{output_path}.tbi") + backup_vcf = pathlib.Path(f"{partial_path}.previous-vcf") + backup_index = pathlib.Path(f"{partial_path}.previous-index") + had_vcf = output_path.exists() + had_index = output_index.exists() + try: + if had_vcf: + os.replace(output_path, backup_vcf) + if had_index: + os.replace(output_index, backup_index) + os.replace(partial_path, output_path) + os.replace(partial_index, output_index) + except BaseException: + if output_path.exists(): + output_path.unlink() + if output_index.exists(): + output_index.unlink() + if had_vcf and backup_vcf.exists(): + os.replace(backup_vcf, output_path) + if had_index and backup_index.exists(): + os.replace(backup_index, output_index) + raise + finally: + if backup_vcf.exists(): + backup_vcf.unlink() + if backup_index.exists(): + backup_index.unlink() + + +def stream_results( + results: Iterator[WorkResult], + output_path: pathlib.Path, + compression_threads: int, + temp_dir: pathlib.Path, +) -> tuple[int, int]: + """Publish ordered fragment bodies through one bcftools BGZF writer.""" + + partial_path, partial_index = partial_paths(output_path) + publisher_error_path = temp_dir / "publisher.stderr" + publisher: subprocess.Popen[bytes] | None = None + expected_number = 0 + expected_column_header: bytes | None = None + input_records = 0 + output_records = 0 + try: + with publisher_error_path.open("wb") as publisher_error: + publisher = subprocess.Popen( + [ + "bcftools", + "view", + "--no-version", + "-O", + "z", + "--threads", + str(compression_threads), + "-o", + str(partial_path), + "-W=tbi", + "-", + ], + stdin=subprocess.PIPE, + stderr=publisher_error, + ) + if publisher.stdin is None: + publisher.terminate() + raise HybridTransferError( + "bcftools stdin pipe was not created" + ) + publisher_stdin = cast(BinaryIO, publisher.stdin) + try: + for result in results: + if result.number != expected_number: + raise HybridTransferError( + "transfer work results are not complete and " + "ordered" + ) + header_path = pathlib.Path(result.header_path) + body_path = pathlib.Path(result.body_path) + observed_column_header = column_header(header_path) + if expected_column_header is None: + expected_column_header = observed_column_header + copy_path(header_path, publisher_stdin) + elif observed_column_header != expected_column_header: + raise HybridTransferError( + f"fragment {result.number} sample header differs" + ) + copy_path(body_path, publisher_stdin) + input_records += result.input_records + output_records += result.output_records + header_path.unlink() + body_path.unlink() + expected_number += 1 + except BaseException: + publisher_stdin.close() + publisher.terminate() + publisher.wait() + raise + publisher_stdin.close() + return_code = publisher.wait() + + if expected_number == 0: + raise HybridTransferError("transfer constructed no work items") + if return_code != 0: + message = publisher_error_path.read_text(errors="replace").strip() + raise HybridTransferError( + f"final VCF publication failed with exit code {return_code}: " + f"{message}" + ) + publish_pair(partial_path, partial_index, output_path) + return input_records, output_records + except FileNotFoundError as error: + raise HybridTransferError( + "required executable bcftools is missing" + ) from error + finally: + if publisher is not None and publisher.poll() is None: + publisher.terminate() + publisher.wait() + if partial_path.exists(): + partial_path.unlink() + if partial_index.exists(): + partial_index.unlink() + + +def run_transfer(args: argparse.Namespace) -> int: + """Validate, execute, and atomically publish the fused transfer.""" + + if args.threads < 1: + raise HybridTransferError("--threads must be at least 1") + if args.workers < 1 or args.workers > MAX_WORKERS: + raise HybridTransferError( + f"--workers must be between 1 and {MAX_WORKERS}" + ) + if args.workers > args.threads: + raise HybridTransferError("--workers must not exceed --threads") + if args.step_size < 1: + raise HybridTransferError("--step-size must be at least 1") + + raw_vcf = pathlib.Path(args.raw_vcf).resolve() + population_vcf = pathlib.Path(args.population_vcf).resolve() + reference_fai = pathlib.Path(args.reference_fai).resolve() + output_path = pathlib.Path(args.output).resolve() + base_temp_dir = pathlib.Path(args.temp_dir).resolve() + require_index(raw_vcf, "raw VCF") + require_index(population_vcf, "population VCF") + if not output_path.parent.is_dir(): + raise HybridTransferError( + f"output directory {output_path.parent} does not exist" + ) + if not base_temp_dir.is_dir(): + raise HybridTransferError( + f"temporary directory {base_temp_dir} does not exist" + ) + + merge_rules, population_contigs, merged_header = population_contract( + raw_vcf, population_vcf + ) + info_numbers, format_numbers = parse_number_schemas(merged_header) + contigs = load_fai(reference_fai) + items = build_work_items(contigs, population_contigs, args.step_size) + worker_budget = max(1, (args.threads - 1) // 2) + workers = min(args.workers, len(items), worker_budget) + compression_threads = max(0, args.threads - (2 * workers) - 1) + temp_dir = pathlib.Path( + tempfile.mkdtemp(prefix="hybrid-transfer.", dir=base_temp_dir) + ) + success = False + pool: mp.pool.Pool | None = None + try: + print( + "hybrid_transfer: " + f"workers={workers} compression_threads={compression_threads} " + f"step_size={args.step_size} shards={len(items)} " + f"tmpdir={temp_dir}", + file=sys.stderr, + ) + context = mp.get_context() + pool = context.Pool( + processes=workers, + initializer=init_worker, + initargs=( + str(raw_vcf), + str(population_vcf), + merge_rules, + info_numbers, + format_numbers, + str(temp_dir), + ), + ) + results = pool.imap(process_work_item, items, chunksize=1) + input_records, output_records = stream_results( + results, + output_path, + compression_threads, + temp_dir, + ) + pool.close() + pool.join() + pool = None + print( + "hybrid_transfer: " + f"input_records={input_records} output_records={output_records}", + file=sys.stderr, + ) + success = True + finally: + if pool is not None: + pool.terminate() + pool.join() + if success: + shutil.rmtree(temp_dir) + else: + print( + f"hybrid_transfer: preserved failed work directory {temp_dir}", + file=sys.stderr, + ) + return 0 + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + prog="sentieon pyexec hybrid_transfer.py", + usage=( + "%(prog)s --raw-vcf RAW --population-vcf POP " + "--reference-fai REF.fai --temp-dir DIR [options] output.vcf.gz" + ), + ) + parser.add_argument("output", help="output BGZF VCF") + parser.add_argument( + "--raw-vcf", required=True, help="annotated Hybrid VCF" + ) + parser.add_argument( + "--population-vcf", required=True, help="population annotation VCF" + ) + parser.add_argument( + "--reference-fai", required=True, help="reference FASTA index" + ) + parser.add_argument( + "--temp-dir", required=True, help="existing scratch directory" + ) + parser.add_argument( + "-t", + "--threads", + type=int, + default=mp.cpu_count(), + help="total thread budget", + ) + parser.add_argument( + "--workers", + type=int, + default=DEFAULT_WORKERS, + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--step-size", + type=int, + default=DEFAULT_STEP_SIZE, + help=argparse.SUPPRESS, + ) + try: + sys.exit(run_transfer(parser.parse_args())) + except HybridTransferError as error: + print(f"Error: {error}", file=sys.stderr) + sys.exit(1) diff --git a/sentieon_cli/scripts/trimalt.py b/sentieon_cli/scripts/trimalt.py index a216200..e6a3b37 100644 --- a/sentieon_cli/scripts/trimalt.py +++ b/sentieon_cli/scripts/trimalt.py @@ -76,11 +76,11 @@ def parse_kv(kv): if info: for i,(k,v) in enumerate(info): - n = infos.get(k) if v != '.' else None - if n == 'A': + number = infos.get(k) if v != '.' else None + if number == 'A': v = v.split(',') info[i] = k + '=' + ','.join(x for x,y in zip(v, ka) if y) - elif n == 'R': + elif number == 'R': v = v.split(',') info[i] = k + '=' + ','.join(x for x,y in zip(v, kr) if y) elif v is True: @@ -95,14 +95,14 @@ def parse_kv(kv): for j,val in enumerate(vals[9:]): val = val.split(':') for i,v in enumerate(val): - n = fmts.get(fmt[i]) if v != '.' else None - if n == 'A': + number = fmts.get(fmt[i]) if v != '.' else None + if number == 'A': v = v.split(',') val[i] = ','.join(x for x,y in zip(v, ka) if y) - elif n == 'R': + elif number == 'R': v = v.split(',') val[i] = ','.join(x for x,y in zip(v, kr) if y) - elif n == 'G': + elif number == 'G': v = v.split(',') kx = kr if len(v) == len(kr) else kg if len(v) == len(kg) else None assert kx diff --git a/tests/unit/test_dag_construction.py b/tests/unit/test_dag_construction.py index c42159f..21a598b 100644 --- a/tests/unit/test_dag_construction.py +++ b/tests/unit/test_dag_construction.py @@ -18,6 +18,7 @@ from sentieon_cli.dnascope_longread import DNAscopeLRPipeline from sentieon_cli.dag import DAG from sentieon_cli.job import Job +from sentieon_cli.scheduler import ThreadScheduler from sentieon_cli.shell_pipeline import Pipeline, Command @@ -246,6 +247,22 @@ def test_job_failure_tolerance(self): ) assert job.shell.nodes[0].fail_ok is False + def test_full_budget_jobs_do_not_oversubscribe(self): + """Two full-budget worker pools cannot be scheduled together.""" + + dag = DAG() + first = Job(Pipeline(Command("first")), "first", 8) + second = Job(Pipeline(Command("second")), "second", 8) + trivial = Job(Pipeline(Command("trivial")), "trivial", 0) + dag.add_job(first) + dag.add_job(second) + dag.add_job(trivial) + + scheduler = ThreadScheduler(dag, threads=8).schedule() + scheduled = scheduler.send(None) + assert trivial in scheduled + assert len({first, second} & scheduled) == 1 + class TestLongReadDAGConstruction: """Test DAG construction for long-read pipeline""" diff --git a/tests/unit/test_dnascope_hybrid_pop_vcf.py b/tests/unit/test_dnascope_hybrid_pop_vcf.py index 7add9f1..b18712b 100644 --- a/tests/unit/test_dnascope_hybrid_pop_vcf.py +++ b/tests/unit/test_dnascope_hybrid_pop_vcf.py @@ -11,7 +11,9 @@ from unittest.mock import patch, MagicMock # Add the parent directory to the path to import sentieon_cli -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +) from sentieon_cli.dnascope_hybrid import DNAscopeHybridPipeline from sentieon_cli.dag import DAG @@ -50,10 +52,10 @@ def create_pipeline(self): # We need to mock sys.exit to avoid exiting during initialization if validation fails with patch("sys.exit"): pipeline = DNAscopeHybridPipeline() - + # Setup mocks pipeline.logger = MagicMock() - + # Configure arguments pipeline.output_vcf = self.mock_vcf pipeline.reference = self.mock_ref @@ -65,14 +67,14 @@ def create_pipeline(self): pipeline.tmp_dir = self.mock_dir pipeline.pop_vcf = None pipeline.skip_pop_vcf_id_check = False - + # Mock fai related pipeline.fai_data = {"chr1": {"length": 1000}} pipeline.shards = [MagicMock()] pipeline.shards[0].contig = "chr1" pipeline.shards[0].start = 1 pipeline.shards[0].stop = 1000 - + # Mock pop vcf contigs pipeline.pop_vcf_contigs = {"chr1": 1000} @@ -85,7 +87,7 @@ def create_pipeline(self): pipeline.longread_tech = "ONT" pipeline.sr_aln = self.mock_sr_aln pipeline.lr_aln = self.mock_lr_aln - + return pipeline @patch("sentieon_cli.dnascope_hybrid.ar_load") @@ -93,76 +95,125 @@ def create_pipeline(self): def test_validation_requires_pop_vcf(self, mock_vcf_id, mock_ar_load): """Test that validation fails if model bundle requires pop_vcf but none provided""" pipeline = self.create_pipeline() - + # Mock bundle info requiring SentieonVcfID mock_bundle_info = { "longReadPlatform": "ONT", "shortReadPlatform": "Illumina", "SentieonVcfID": "some_id", - "minScriptVersion": "2.0" + "minScriptVersion": "2.0", } mock_ar_load.return_value = json.dumps(mock_bundle_info).encode() pipeline.dry_run = False - + with patch("sys.exit") as mock_exit: - pipeline.validate_bundle() - mock_exit.assert_called_with(2) + pipeline.validate_bundle() + mock_exit.assert_called_with(2) @patch("sentieon_cli.dnascope_hybrid.ar_load") def test_validation_pop_vcf_mismatch(self, mock_ar_load): """Test validation fails if pop_vcf ID mismatches""" pipeline = self.create_pipeline() pipeline.pop_vcf = self.mock_pop_vcf - + mock_bundle_info = { "longReadPlatform": "ONT", "shortReadPlatform": "Illumina", "SentieonVcfID": "expected_id", - "minScriptVersion": "2.0" + "minScriptVersion": "2.0", } mock_ar_load.return_value = json.dumps(mock_bundle_info).encode() pipeline.dry_run = False - with patch("sentieon_cli.dnascope_hybrid.vcf_id", return_value="wrong_id"), \ - patch("sys.exit") as mock_exit: - + with ( + patch( + "sentieon_cli.dnascope_hybrid.vcf_id", return_value="wrong_id" + ), + patch("sys.exit") as mock_exit, + ): + pipeline.validate_bundle() mock_exit.assert_called_with(2) - def test_transfer_jobs_creation(self): - """Test that transfer jobs are added to the DAG when pop_vcf is present""" + def test_fused_transfer_job_creation(self): + """The Hybrid DAG has one bounded population-transfer job.""" pipeline = self.create_pipeline() pipeline.pop_vcf = self.mock_pop_vcf - + # Build DAG - with patch("sentieon_cli.dnascope_hybrid.check_version", return_value=True): - dag = pipeline.build_dag() - - # Verify transfer jobs exist - job_names = [job.name for job in dag.waiting_jobs] - - # Expect merge-trim jobs - assert any("merge-trim" in name for name in job_names) - assert "merge-trim-concat" in job_names + with patch( + "sentieon_cli.dnascope_hybrid.check_version", return_value=True + ): + dag = pipeline.build_dag() + + transfer_jobs = [ + job + for job in dag.waiting_jobs + if job.name == "population-transfer" + ] + assert len(transfer_jobs) == 1 + transfer_job = transfer_jobs[0] + assert transfer_job.threads == pipeline.cores + assert [job.name for job in dag.waiting_jobs[transfer_job]] == [ + "anno-calls" + ] + assert not any( + job.name.startswith("merge-trim") for job in dag.waiting_jobs + ) def test_model_apply_input_with_pop_vcf(self): """Test that model apply uses the transfer output when pop_vcf is present""" pipeline = self.create_pipeline() pipeline.pop_vcf = self.mock_pop_vcf - - with patch("sentieon_cli.dnascope_hybrid.check_version", return_value=True): - dag = pipeline.build_dag() - + + with patch( + "sentieon_cli.dnascope_hybrid.check_version", return_value=True + ): + dag = pipeline.build_dag() + # Find model-apply job apply_job = None for job in dag.waiting_jobs: if job.name == "model-apply": apply_job = job break - + assert apply_job is not None - - # Check dependency: apply_job should depend on merge-trim-concat + + # ModelApply consumes the single atomically published transfer VCF. deps = dag.waiting_jobs[apply_job] dep_names = [j.name for j in deps] - assert "merge-trim-concat" in dep_names + assert dep_names == ["population-transfer"] + + def test_multithreaded_python_jobs_claim_their_budget(self): + """Python worker pools cannot bypass scheduler accounting.""" + + pipeline = self.create_pipeline() + pipeline.pop_vcf = self.mock_pop_vcf + with patch( + "sentieon_cli.dnascope_hybrid.check_version", return_value=True + ): + dag = pipeline.build_dag() + + all_jobs = [*dag.ready_jobs, *dag.waiting_jobs] + by_name = {job.name: job for job in all_jobs} + for name in ( + "hybrid-select", + "anno-calls", + "population-transfer", + ): + assert by_name[name].threads == pipeline.cores + + final_norm = by_name["final-norm"] + assert final_norm.threads == min(3, pipeline.cores) + assert [node.executable for node in final_norm.shell.nodes] == [ + "bcftools", + "bcftools", + "sentieon", + ] + + zero_thread_jobs = {job.name for job in all_jobs if job.threads == 0} + assert "hybrid-select" not in zero_thread_jobs + assert "anno-calls" not in zero_thread_jobs + assert "population-transfer" not in zero_thread_jobs + assert "final-norm" not in zero_thread_jobs diff --git a/tests/unit/test_hybrid_anno.py b/tests/unit/test_hybrid_anno.py new file mode 100644 index 0000000..e629932 --- /dev/null +++ b/tests/unit/test_hybrid_anno.py @@ -0,0 +1,266 @@ +import gzip +import os +import pathlib +import shutil +import subprocess +import sys + +import pytest + +from sentieon_cli.command_strings import cmd_pyexec_hybrid_anno +from sentieon_cli.scripts.hybrid_anno import cut_shards + + +REPOSITORY_ROOT = pathlib.Path(__file__).parents[2] +HYBRID_ANNO = REPOSITORY_ROOT / "sentieon_cli" / "scripts" / "hybrid_anno.py" + +INPUT_VCF = """\ +##fileformat=VCFv4.2 +##contig= +##contig= +##FILTER= +##INFO= +##INFO= +##FORMAT= +#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tSAMPLE +chr1\t1\t.\tA\t\t.\tPASS\tEND=5\tGT\t0/0 +chr1\t6\t.\tC\t\t.\tPASS\tEND=15\tGT\t0/0 +chr1\t11\t.\tA\tG\t50\tPASS\t.\tGT\t0/1 +chr1\t20\t.\tA\tT\t40\tPASS\tDP=3\tGT\t0/1 +chr1\t21\t.\tA\tT\t40\tPASS\tLHC=9\tGT\t0/1 +chr1\t30\t.\tG\tC\t30\tPASS\t.\tGT\t1/1 +chr1\t31\t.\tT\tC\t30\tPASS\t.\tGT\t1/1 +chr2\t1\t.\tA\tG\t20\tPASS\t.\tGT\t0/1 +chr2\t10\t.\tC\tT\t20\tPASS\t.\tGT\t0/1 +""" + +BED = """\ +chr1\t0\t10\t1 +chr1\t5\t20\t2 +chr1\t20\t30\t3 +chr2\t0\t5\t4 +""" + +EXPECTED_ANNOTATED_VCF = """\ +##fileformat=VCFv4.2 +##contig= +##contig= +##FILTER= +##INFO= +##INFO= +##FORMAT= +#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tSAMPLE +chr1\t1\t.\tA\t\t.\tPASS\tEND=5;LHC=1\tGT\t0/0 +chr1\t6\t.\tC\t\t.\tPASS\tEND=15;LHC=1\tGT\t0/0 +chr1\t11\t.\tA\tG\t50\tPASS\t.;LHC=2\tGT\t0/1 +chr1\t20\t.\tA\tT\t40\tPASS\tDP=3;LHC=2\tGT\t0/1 +chr1\t21\t.\tA\tT\t40\tPASS\tLHC=9;LHC=3\tGT\t0/1 +chr1\t30\t.\tG\tC\t30\tPASS\t.;LHC=3\tGT\t1/1 +chr1\t31\t.\tT\tC\t30\tPASS\t.\tGT\t1/1 +chr2\t1\t.\tA\tG\t20\tPASS\t.;LHC=4\tGT\t0/1 +chr2\t10\t.\tC\tT\t20\tPASS\t.\tGT\t0/1 +""" + +EXPECTED_EMPTY_BED_VCF = EXPECTED_ANNOTATED_VCF.replace( + "END=5;LHC=1", "END=5" +).replace( + "END=15;LHC=1", "END=15" +).replace( + ".;LHC=2", "." +).replace( + "DP=3;LHC=2", "DP=3" +).replace( + "LHC=9;LHC=3", "LHC=9" +).replace( + ".;LHC=3", "." +).replace( + ".;LHC=4", "." +) + + +def require_htslib_tools() -> tuple[str, str]: + bgzip = shutil.which("bgzip") + tabix = shutil.which("tabix") + assert bgzip is not None, "bgzip is required for hybrid_anno tests" + assert tabix is not None, "tabix is required for hybrid_anno tests" + return bgzip, tabix + + +def make_indexed_vcf(tmp_path: pathlib.Path) -> pathlib.Path: + bgzip, tabix = require_htslib_tools() + source = tmp_path / "input.vcf" + source.write_text(INPUT_VCF) + compressed = tmp_path / "input.vcf.gz" + with compressed.open("wb") as output: + subprocess.run([bgzip, "-c", str(source)], check=True, stdout=output) + subprocess.run([tabix, "-p", "vcf", str(compressed)], check=True) + return compressed + + +def run_annotator( + tmp_path: pathlib.Path, + input_vcf: pathlib.Path, + bed_text: str, + threads: int, + step_size: int | None = None, +) -> tuple[pathlib.Path, subprocess.CompletedProcess[str]]: + bed = tmp_path / f"stage1.{threads}.bed" + bed.write_text(bed_text) + output = tmp_path / f"annotated.{threads}.vcf.gz" + scratch = tmp_path / f"scratch.{threads}" + scratch.mkdir() + command = [ + sys.executable, + str(HYBRID_ANNO), + "-v", + str(input_vcf), + "-b", + str(bed), + "-t", + str(threads), + ] + if step_size is not None: + command.extend(["--step-size", str(step_size)]) + command.append(str(output)) + environment = os.environ.copy() + environment["SENTIEON_TMPDIR"] = str(scratch) + result = subprocess.run( + command, + capture_output=True, + text=True, + env=environment, + ) + return output, result + + +@pytest.mark.parametrize( + ("threads", "step_size"), + ((1, None), (4, 13)), +) +def test_decompressed_output_matches_v170_oracle( + tmp_path: pathlib.Path, + threads: int, + step_size: int | None, +) -> None: + input_vcf = make_indexed_vcf(tmp_path) + output, result = run_annotator( + tmp_path, input_vcf, BED, threads, step_size + ) + assert result.returncode == 0, result.stderr + expected_step_size = step_size or 10_000_000 + assert f"step_size={expected_step_size}" in result.stderr + assert pathlib.Path(f"{output}.tbi").is_file() + with gzip.open(output, "rt") as annotated: + assert annotated.read() == EXPECTED_ANNOTATED_VCF + + _, tabix = require_htslib_tools() + observed = subprocess.check_output( + [tabix, str(output), "chr1:6-21"], + text=True, + ) + expected = "\n".join( + line + for line in EXPECTED_ANNOTATED_VCF.splitlines() + if line.startswith("chr1\t") + and 6 <= int(line.split("\t", 2)[1]) <= 21 + ) + assert observed == f"{expected}\n" + + +def test_empty_bed_preserves_records_and_adds_header( + tmp_path: pathlib.Path, +) -> None: + input_vcf = make_indexed_vcf(tmp_path) + output, result = run_annotator(tmp_path, input_vcf, "", 2, 10) + assert result.returncode == 0, result.stderr + with gzip.open(output, "rt") as annotated: + assert annotated.read() == EXPECTED_EMPTY_BED_VCF + + +def test_shard_boundaries_carry_across_contigs_like_vcflib() -> None: + assert cut_shards([("chr1", 15), ("chr2", 10)], 10) == [ + (0, "chr1", 0, 10), + (1, "chr1", 10, 15), + (2, "chr2", 0, 5), + (3, "chr2", 5, 10), + ] + + +def test_missing_index_fails_without_touching_input( + tmp_path: pathlib.Path, +) -> None: + input_vcf = make_indexed_vcf(tmp_path) + pathlib.Path(f"{input_vcf}.tbi").unlink() + output, result = run_annotator(tmp_path, input_vcf, BED, 1) + assert result.returncode == 1 + assert "input VCF index is missing" in result.stderr + assert input_vcf.is_file() + assert not output.exists() + + +def test_invalid_thread_count_fails(tmp_path: pathlib.Path) -> None: + input_vcf = make_indexed_vcf(tmp_path) + output, result = run_annotator(tmp_path, input_vcf, BED, 0) + assert result.returncode == 1 + assert "--threads must be at least 1" in result.stderr + assert not output.exists() + + +def test_missing_contig_length_fails(tmp_path: pathlib.Path) -> None: + bgzip, tabix = require_htslib_tools() + source = tmp_path / "missing-length.vcf" + source.write_text( + INPUT_VCF.replace( + "##contig=", + "##contig=", + ) + ) + input_vcf = tmp_path / "missing-length.vcf.gz" + with input_vcf.open("wb") as output_file: + subprocess.run([bgzip, "-c", str(source)], check=True, stdout=output_file) + subprocess.run([tabix, "-p", "vcf", str(input_vcf)], check=True) + output, result = run_annotator(tmp_path, input_vcf, BED, 1) + assert result.returncode == 1 + assert "contig chr1 has no length" in result.stderr + assert not output.exists() + + +@pytest.mark.parametrize( + "bed_text", + ( + "chr1\t0\t10\n", + "chr1\tbad\t10\t1\n", + "chr1\t10\t5\t1\n", + "chr1\t10\t20\t1\nchr1\t0\t5\t2\n", + ), +) +def test_malformed_bed_fails( + tmp_path: pathlib.Path, + bed_text: str, +) -> None: + input_vcf = make_indexed_vcf(tmp_path) + output, result = run_annotator(tmp_path, input_vcf, bed_text, 1) + assert result.returncode == 1 + assert "Error:" in result.stderr + assert not output.exists() + + +def test_command_builder_interface_is_unchanged(tmp_path: pathlib.Path) -> None: + output = tmp_path / "output.vcf.gz" + input_vcf = tmp_path / "input.vcf.gz" + bed = tmp_path / "stage1_hap.bed" + script = tmp_path / "hybrid_anno.py" + pipeline = cmd_pyexec_hybrid_anno(output, input_vcf, bed, script, 128) + assert len(pipeline.nodes) == 1 + command = pipeline.nodes[0] + assert command.executable == sys.executable + assert command.args == [ + str(script), + "-v", + str(input_vcf), + "-b", + str(bed), + "-t", + "128", + str(output), + ] diff --git a/tests/unit/test_hybrid_select.py b/tests/unit/test_hybrid_select.py new file mode 100644 index 0000000..2f260fb --- /dev/null +++ b/tests/unit/test_hybrid_select.py @@ -0,0 +1,270 @@ +import os +import pathlib +import shutil +import subprocess +import sys + +import pytest +import vcflib + +from sentieon_cli.command_strings import cmd_pyexec_hybrid_select +from sentieon_cli.scripts.hybrid_select import cut_shards + +REPOSITORY_ROOT = pathlib.Path(__file__).parents[2] +HYBRID_SELECT = ( + REPOSITORY_ROOT / "sentieon_cli" / "scripts" / "hybrid_select.py" +) + +INPUT_VCF = """\ +##fileformat=VCFv4.2 +##contig= +##contig= +##INFO= +##INFO= +##FORMAT= +##FORMAT= +##FORMAT= +##FORMAT= +#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tSAMPLE +chr1\t1\t.\tA\t\t.\t.\tEND=5\tGT:LAD:LPL:SPL\t0/0:1,1,0:0,40,50:0,5,10 +chr1\t6\t.\tC\tT\t50\tPASS\t.\tGT:LAD:LPL:SPL\t0/1:1,0,0:0,40,50:0,5,10 +chr1\t10\t.\tA\tG\t50\tPASS\t.\tGT:LAD:LPL:SPL\t0/1:2,2,0:0,10,20:0,5,10 +chr1\t15\t.\tA\tG\t50\tPASS\t.\tGT:LAD:LPL:SPL\t0/1:2,2,0:0,40,50:0,40,50 +chr1\t20\t.\tA\tG\t50\tPASS\tSTR\tGT:LAD:LPL:SPL\t0/1:2,2,0:0,40,50:40,0,50 +chr1\t25\t.\tAA\tA\t50\tPASS\t.\tGT:LAD:LPL:SPL\t0/1:2,2,0:40,0,50:0,5,10 +chr2\t1\t.\tG\tC\t50\tPASS\t.\tGT:LAD:LPL:SPL\t0/1:2,2,0:40,0,50:0,5,10 +""" + +EXPECTED_BED = """\ +chr1\t0\t8 +chr1\t21\t29 +chr2\t0\t4 +""" + +LEGACY_EXTRA_HEADERS = ( + '##FILTER=', + '##FILTER=', + '##FILTER=', + '##FILTER=', + '##FILTER=', + '##FILTER=', +) + + +def run_v170_object_pipeline( + tmp_path: pathlib.Path, + input_vcf: pathlib.Path, + reference_index: pathlib.Path, +) -> bytes: + """Run the deployed vcflib -> bcftools -> bedtools selector oracle.""" + + legacy_vcf = tmp_path / "legacy.filtered.vcf.gz" + input_handle = vcflib.VCF(str(input_vcf), "r") + output_handle = vcflib.VCF(str(legacy_vcf), "w") + try: + output_handle.copy_header( + input_handle, + LEGACY_EXTRA_HEADERS, + ("##source=*",), + ) + output_handle.emit_header() + for variant in input_handle: + filters = [] + lad = variant.samples[0].get("LAD") or [0, 0, 0] + lpl = variant.samples[0].get("LPL") or [0, 0, 0] + spl = variant.samples[0].get("SPL") or [0, 0, 0] + long_minimum = lpl.index(0) + short_minimum = spl.index(0) + long_reference = lpl[0] + lpl.remove(0) + spl.remove(0) + long_confidence = min(lpl) + short_confidence = min(spl) + + if sum(lad) < 2: + filters.append("longread_lowdepth") + if not filters: + if variant.info.get("STR") is None and long_minimum != 0: + long_confidence = long_reference + if long_confidence < 30.0: + filters.append("longread_lowconf") + if not filters and short_confidence >= 30.0: + if long_minimum == short_minimum: + filters.append("same_gt") + elif long_minimum == 0 and variant.info.get("STR") is not None: + filters.append("longread_str") + + columns = variant.line.split("\t") + columns[6] = ";".join(sorted(set(filters))) if filters else "PASS" + variant.line = "\t".join(columns) + output_handle.emit(variant) + finally: + output_handle.close() + input_handle.close() + + view = subprocess.run( + ["bcftools", "view", "-f", "PASS,.", str(legacy_vcf)], + check=True, + stdout=subprocess.PIPE, + ) + query = subprocess.run( + ["bcftools", "query", "-f", "%CHROM\\t%POS0\\t%END\\n", "-"], + check=True, + input=view.stdout, + stdout=subprocess.PIPE, + ) + slop = subprocess.run( + [ + "bedtools", + "slop", + "-b", + "3", + "-g", + str(reference_index), + "-i", + "-", + ], + check=True, + input=query.stdout, + stdout=subprocess.PIPE, + ) + return slop.stdout + + +def require_htslib_tools() -> tuple[str, str]: + bgzip = shutil.which("bgzip") + tabix = shutil.which("tabix") + assert bgzip is not None, "bgzip is required for hybrid_select tests" + assert tabix is not None, "tabix is required for hybrid_select tests" + return bgzip, tabix + + +def make_indexed_vcf(tmp_path: pathlib.Path) -> pathlib.Path: + bgzip, tabix = require_htslib_tools() + source = tmp_path / "input.vcf" + source.write_text(INPUT_VCF) + compressed = tmp_path / "input.vcf.gz" + with compressed.open("wb") as output: + subprocess.run([bgzip, "-c", str(source)], check=True, stdout=output) + subprocess.run([tabix, "-p", "vcf", str(compressed)], check=True) + return compressed + + +def run_selector( + tmp_path: pathlib.Path, + input_vcf: pathlib.Path, + threads: int, + step_size: int, +) -> tuple[pathlib.Path, subprocess.CompletedProcess[str]]: + reference_index = tmp_path / "reference.fa.fai" + reference_index.write_text("chr1\t50\t0\t50\t51\nchr2\t25\t51\t25\t26\n") + output = tmp_path / f"selected.{threads}.bed" + scratch = tmp_path / f"scratch.{threads}" + scratch.mkdir() + environment = os.environ.copy() + environment["SENTIEON_TMPDIR"] = str(scratch) + result = subprocess.run( + [ + sys.executable, + str(HYBRID_SELECT), + "-v", + str(input_vcf), + "--reference-fai", + str(reference_index), + "--slop-size", + "3", + "--step-size", + str(step_size), + "-t", + str(threads), + str(output), + ], + capture_output=True, + text=True, + env=environment, + ) + return output, result + + +@pytest.mark.parametrize("threads", (1, 4)) +def test_direct_bed_matches_v170_pipeline_semantics( + tmp_path: pathlib.Path, threads: int +) -> None: + input_vcf = make_indexed_vcf(tmp_path) + output, result = run_selector(tmp_path, input_vcf, threads, 13) + assert result.returncode == 0, result.stderr + assert output.read_text() == EXPECTED_BED + assert "records=7 selected=3" in result.stderr + + +def test_cross_contig_shard_carry_matches_vcflib() -> None: + assert cut_shards([("chr1", 15), ("chr2", 10)], 10) == [ + (0, "chr1", 0, 10), + (1, "chr1", 10, 15), + (2, "chr2", 0, 5), + (3, "chr2", 5, 10), + ] + + +@pytest.mark.skipif( + shutil.which("bedtools") is None, + reason="exact deployed selector oracle requires legacy bedtools", +) +def test_direct_bed_matches_exact_v170_object_pipeline( + tmp_path: pathlib.Path, +) -> None: + input_vcf = make_indexed_vcf(tmp_path) + output, result = run_selector(tmp_path, input_vcf, 4, 13) + assert result.returncode == 0, result.stderr + reference_index = tmp_path / "reference.fa.fai" + assert output.read_bytes() == run_v170_object_pipeline( + tmp_path, + input_vcf, + reference_index, + ) + + +def test_missing_index_fails_atomically(tmp_path: pathlib.Path) -> None: + input_vcf = make_indexed_vcf(tmp_path) + pathlib.Path(f"{input_vcf}.tbi").unlink() + output, result = run_selector(tmp_path, input_vcf, 1, 10) + assert result.returncode == 1 + assert "input VCF index is missing" in result.stderr + assert not output.exists() + + +def test_invalid_thread_count_fails_atomically(tmp_path: pathlib.Path) -> None: + input_vcf = make_indexed_vcf(tmp_path) + output, result = run_selector(tmp_path, input_vcf, 0, 10) + assert result.returncode == 1 + assert "--threads must be at least 1" in result.stderr + assert not output.exists() + + +def test_command_builder_is_one_direct_process(tmp_path: pathlib.Path) -> None: + output = tmp_path / "selected.bed" + input_vcf = tmp_path / "input.vcf.gz" + reference_index = tmp_path / "reference.fa.fai" + script = tmp_path / "hybrid_select.py" + pipeline = cmd_pyexec_hybrid_select( + output, + input_vcf, + reference_index, + script, + 128, + ) + assert len(pipeline.nodes) == 1 + command = pipeline.nodes[0] + assert command.executable == sys.executable + assert command.args == [ + str(script), + "-v", + str(input_vcf), + "-t", + "128", + "--reference-fai", + str(reference_index), + "--slop-size", + "1000", + str(output), + ] diff --git a/tests/unit/test_hybrid_transfer.py b/tests/unit/test_hybrid_transfer.py new file mode 100644 index 0000000..0c62204 --- /dev/null +++ b/tests/unit/test_hybrid_transfer.py @@ -0,0 +1,433 @@ +import gzip +import pathlib +import shutil +import subprocess +import sys + +import pytest + +from sentieon_cli.command_strings import cmd_pyexec_hybrid_transfer +from sentieon_cli.scripts.hybrid_transfer import ( + build_work_items, + trim_record, +) + +REPOSITORY_ROOT = pathlib.Path(__file__).parents[2] +HYBRID_TRANSFER = ( + REPOSITORY_ROOT / "sentieon_cli" / "scripts" / "hybrid_transfer.py" +) +TRIMALT = REPOSITORY_ROOT / "sentieon_cli" / "scripts" / "trimalt.py" + +RAW_VCF = """\ +##fileformat=VCFv4.2 +##FILTER= +##contig= +##contig= +##INFO= +##INFO= +##INFO= +##INFO= +##FORMAT= +##FORMAT= +##FORMAT= +##FORMAT= +#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tSAMPLE +chr1\t1\t.\tA\tC\t50\tPASS\tAF=0.5;IA=11;IR=1,2\tGT:FA:FR:FG\t0/1:11:1,2:0,10,20 +chr1\t2\t.\tA\tC,G\t50\tPASS\tAF=0.1,0.2;IA=12,13;IR=1,2,3\tGT:FA:FR:FG\t1/2:12,13:1,2,3:0,10,20,30,40,50 +chr1\t3\t.\tG\t\t.\tPASS\tEND=5\tGT:FA:FR:FG\t0/0:.:1,.:0,10,20 +chr1\t9\t.\tT\tC\t40\tPASS\tAF=0.3;IA=19;IR=3,4\tGT:FA:FR:FG\t0/1:19:3,4:0,11,22 +chr1\t10\t.\tA\tG\t40\tPASS\tAF=0.4;IA=20;IR=4,5\tGT:FA:FR:FG\t0/1:20:4,5:0,12,24 +chr1\t11\t.\tA\tT\t40\tPASS\tAF=0.5;IA=21;IR=5,6\tGT:FA:FR:FG\t0/1:21:5,6:0,13,26 +chr1\t12\ta\tC\tG\t40\tPASS\tAF=0.6;IA=22;IR=6,7\tGT:FA:FR:FG\t0/1:22:6,7:0,14,28 +chr1\t12\tb\tC\tT\t40\tPASS\tAF=0.7;IA=23;IR=7,8\tGT:FA:FR:FG\t1:23:7,8:0,15 +chr1\t15\t.\tAT\tCT\t40\tPASS\tAF=0.8;IA=24;IR=8,9\tGT:FA:FR:FG\t0/1:24:8,9:0,16,32 +chr1\t20\t.\tG\tA\t40\tPASS\tAF=.;IA=25;IR=9,10\tGT:FA:FR:FG\t./.:25:9,10:0,17,34 +chrUn\t1\t.\tA\tT\t30\tPASS\tAF=0.9;IA=31;IR=10,11\tGT:FA:FR:FG\t0/1:31:10,11:0,18,36 +chrUn\t7\t.\tC\tG\t30\tPASS\tAF=0.4;IA=32;IR=11,12\tGT:FA:FR:FG\t0/1:32:11,12:0,19,38 +""" + +POPULATION_VCF = """\ +##fileformat=VCFv4.2 +##FILTER= +##contig= +##INFO= +##INFO= +#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO +chr1\t2\t.\tA\tC,T\t.\tPASS\tPOPA=5,7;POPR=20,5,7 +chr1\t3\t.\tG\tA\t.\tPASS\tPOPA=8;POPR=20,8 +chr1\t4\t.\tC\tT\t.\tPASS\tPOPA=9;POPR=20,9 +chr1\t10\t.\tA\tG,C\t.\tPASS\tPOPA=10,11;POPR=20,10,11 +chr1\t11\t.\tA\tC\t.\tPASS\tPOPA=12;POPR=20,12 +chr1\t12\t.\tC\tA\t.\tPASS\tPOPA=13;POPR=20,13 +chr1\t15\t.\tAT\tGT\t.\tPASS\tPOPA=14;POPR=20,14 +chr1\t21\t.\tT\tA\t.\tPASS\tPOPA=15;POPR=20,15 +""" + + +def require_tools() -> tuple[str, str, str]: + bcftools = shutil.which("bcftools") + bgzip = shutil.which("bgzip") + tabix = shutil.which("tabix") + assert bcftools is not None, "bcftools is required" + assert bgzip is not None, "bgzip is required" + assert tabix is not None, "tabix is required" + return bcftools, bgzip, tabix + + +def make_indexed_vcf( + tmp_path: pathlib.Path, name: str, contents: str +) -> pathlib.Path: + _, bgzip, tabix = require_tools() + source = tmp_path / f"{name}.vcf" + source.write_text(contents) + compressed = tmp_path / f"{name}.vcf.gz" + with compressed.open("wb") as output: + subprocess.run([bgzip, "-c", str(source)], check=True, stdout=output) + subprocess.run([tabix, "-p", "vcf", str(compressed)], check=True) + return compressed + + +def make_inputs( + tmp_path: pathlib.Path, +) -> tuple[pathlib.Path, pathlib.Path, pathlib.Path]: + raw = make_indexed_vcf(tmp_path, "raw", RAW_VCF) + population = make_indexed_vcf(tmp_path, "population", POPULATION_VCF) + reference_index = tmp_path / "reference.fa.fai" + reference_index.write_text("chr1\t25\t0\t25\t26\nchrUn\t7\t26\t7\t8\n") + return raw, population, reference_index + + +def run_pipeline( + commands: list[list[str]], +) -> tuple[bytes, bytes]: + stdin: bytes | None = None + stderr = bytearray() + for command in commands: + result = subprocess.run( + command, + input=stdin, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stderr.extend(result.stderr) + assert result.returncode == 0, stderr.decode(errors="replace") + stdin = result.stdout + return stdin or b"", bytes(stderr) + + +def run_legacy_transfer( + tmp_path: pathlib.Path, + raw: pathlib.Path, + population: pathlib.Path, + reference_index: pathlib.Path, + step_size: int, +) -> pathlib.Path: + bcftools, _, _ = require_tools() + population_header = subprocess.check_output( + [bcftools, "view", "-h", str(population)], text=True + ) + info_fields = [] + for line in population_header.splitlines(): + if line.startswith("##INFO") and ",Number=A" in line: + info_fields.append(line.split("ID=", 1)[1].split(",", 1)[0]) + merge_rules = ",".join(f"{field}:sum" for field in info_fields) + items = build_work_items([("chr1", 25), ("chrUn", 7)], {"chr1"}, step_size) + shards: list[pathlib.Path] = [] + for item in items: + bed = tmp_path / f"legacy.{item.number}.bed" + bed.write_text(f"{item.contig}\t{item.start}\t{item.stop}\n") + shard = tmp_path / f"legacy.{item.number}.vcf.gz" + if item.merge_population: + merged, _ = run_pipeline( + [ + [ + bcftools, + "merge", + "--regions-file", + str(bed), + "--no-version", + "--regions-overlap", + "pos", + "-m", + "all", + "-i", + merge_rules, + str(raw), + str(population), + ], + [sys.executable, str(TRIMALT)], + ] + ) + result = subprocess.run( + [ + bcftools, + "view", + "--no-version", + "-W=tbi", + "-o", + str(shard), + ], + input=merged, + capture_output=True, + ) + else: + result = subprocess.run( + [ + bcftools, + "view", + "--no-version", + "-W=tbi", + "-O", + "z", + "-o", + str(shard), + "--regions-file", + str(bed), + str(raw), + ], + capture_output=True, + ) + assert result.returncode == 0, result.stderr.decode(errors="replace") + shards.append(shard) + + output = tmp_path / "legacy.vcf.gz" + result = subprocess.run( + [ + bcftools, + "concat", + "-W=tbi", + "--output", + str(output), + "--no-version", + "--threads", + "4", + *(str(shard) for shard in shards), + ], + capture_output=True, + ) + assert result.returncode == 0, result.stderr.decode(errors="replace") + return output + + +def run_fused_transfer( + tmp_path: pathlib.Path, + raw: pathlib.Path, + population: pathlib.Path, + reference_index: pathlib.Path, +) -> tuple[pathlib.Path, subprocess.CompletedProcess[str]]: + output = tmp_path / "fused.vcf.gz" + scratch = tmp_path / "scratch" + scratch.mkdir() + result = subprocess.run( + [ + sys.executable, + str(HYBRID_TRANSFER), + "--raw-vcf", + str(raw), + "--population-vcf", + str(population), + "--reference-fai", + str(reference_index), + "--temp-dir", + str(scratch), + "--step-size", + "10", + "--threads", + "4", + "--workers", + "2", + str(output), + ], + capture_output=True, + text=True, + ) + return output, result + + +def decompressed(path: pathlib.Path) -> str: + with gzip.open(path, "rt") as input_vcf: + return input_vcf.read() + + +def record_body(path: pathlib.Path) -> list[str]: + return [ + line + for line in decompressed(path).splitlines() + if not line.startswith("#") + ] + + +def test_fused_transfer_matches_legacy_pipeline( + tmp_path: pathlib.Path, +) -> None: + raw, population, reference_index = make_inputs(tmp_path) + legacy = run_legacy_transfer( + tmp_path, raw, population, reference_index, step_size=10 + ) + fused, result = run_fused_transfer( + tmp_path, raw, population, reference_index + ) + assert result.returncode == 0, result.stderr + assert pathlib.Path(f"{fused}.tbi").is_file() + assert record_body(fused) == record_body(legacy) + assert "workers=1 compression_threads=1" in result.stderr + + _, _, tabix = require_tools() + for region in ("chr1:1-25", "chr1:10-12", "chrUn:1-7"): + expected = subprocess.check_output([tabix, str(legacy), region]) + observed = subprocess.check_output([tabix, str(fused), region]) + assert observed == expected + + +def test_work_items_preserve_legacy_position_ownership() -> None: + items = build_work_items([("chr1", 25), ("chrUn", 7)], {"chr1"}, 10) + assert [ + ( + item.number, + item.contig, + item.start, + item.stop, + item.merge_population, + ) + for item in items + ] == [ + (0, "chr1", 1, 10, True), + (1, "chr1", 11, 20, True), + (2, "chr1", 21, 25, True), + (3, "chrUn", 0, 7, False), + ] + + +def test_trim_record_matches_number_arg_semantics() -> None: + line = ( + b"chr1\t2\t.\tAT\tCT,GT\t50\tPASS\t" + b"AF=0.2,.;IA=7,8;IR=1,2,3;FLAG\t" + b"GT:FA:FR:FG\t0/1:7,8:1,2,3:0,10,20,30,40,50\n" + ) + assert trim_record( + line, + {b"AF": b"A", b"IA": b"A", b"IR": b"R"}, + {b"FA": b"A", b"FR": b"R", b"FG": b"G"}, + ) == ( + b"chr1\t2\t.\tA\tC\t50\tPASS\t" + b"AF=0.2;IA=7;IR=1,2;FLAG\t" + b"GT:FA:FR:FG\t0/1:7:1,2:0,10,20\n" + ) + + +def test_missing_index_fails_without_replacing_output( + tmp_path: pathlib.Path, +) -> None: + raw, population, reference_index = make_inputs(tmp_path) + pathlib.Path(f"{population}.tbi").unlink() + output = tmp_path / "fused.vcf.gz" + output.write_bytes(b"existing output") + scratch = tmp_path / "scratch" + scratch.mkdir() + result = subprocess.run( + [ + sys.executable, + str(HYBRID_TRANSFER), + "--raw-vcf", + str(raw), + "--population-vcf", + str(population), + "--reference-fai", + str(reference_index), + "--temp-dir", + str(scratch), + "--threads", + "2", + "--workers", + "1", + str(output), + ], + capture_output=True, + text=True, + ) + assert result.returncode == 1 + assert "population VCF index is missing" in result.stderr + assert output.read_bytes() == b"existing output" + assert not pathlib.Path(f"{output}.tbi").exists() + + +@pytest.mark.parametrize( + ("threads", "workers", "message"), + ( + (0, 1, "--threads must be at least 1"), + (2, 0, "--workers must be between 1 and 64"), + (2, 3, "--workers must not exceed --threads"), + ), +) +def test_invalid_budget_fails_atomically( + tmp_path: pathlib.Path, + threads: int, + workers: int, + message: str, +) -> None: + raw, population, reference_index = make_inputs(tmp_path) + output = tmp_path / "fused.vcf.gz" + scratch = tmp_path / "scratch" + scratch.mkdir() + result = subprocess.run( + [ + sys.executable, + str(HYBRID_TRANSFER), + "--raw-vcf", + str(raw), + "--population-vcf", + str(population), + "--reference-fai", + str(reference_index), + "--temp-dir", + str(scratch), + "--threads", + str(threads), + "--workers", + str(workers), + str(output), + ], + capture_output=True, + text=True, + ) + assert result.returncode == 1 + assert message in result.stderr + assert not output.exists() + + +def test_command_builder_uses_one_process(tmp_path: pathlib.Path) -> None: + output = tmp_path / "output.vcf.gz" + raw = tmp_path / "raw.vcf.gz" + population = tmp_path / "population.vcf.gz" + reference_index = tmp_path / "reference.fa.fai" + scratch = tmp_path / "scratch" + script = tmp_path / "hybrid_transfer.py" + pipeline = cmd_pyexec_hybrid_transfer( + output, + raw, + population, + reference_index, + scratch, + script, + 128, + 32, + ) + assert len(pipeline.nodes) == 1 + command = pipeline.nodes[0] + assert command.executable == sys.executable + assert command.args == [ + str(script), + "--raw-vcf", + str(raw), + "--population-vcf", + str(population), + "--reference-fai", + str(reference_index), + "--temp-dir", + str(scratch), + "--threads", + "128", + "--workers", + "32", + str(output), + ]