mhcflow 0.2.0__py3-none-any.whl

mhcflow/__init__.py

@@ -0,0 +1,5 @@
+from mhcflow.mhcflow import run_mhcflow
+
+from ._version import version as __version__
+
+__all__ = ["run_mhcflow", "__version__"]

mhcflow/__main__.py

@@ -0,0 +1,5 @@
+import sys
+
+import mhcflow
+
+sys.exit(mhcflow.run_mhcflow())

mhcflow/_version.py

@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)

mhcflow/cli.py

@@ -0,0 +1,75 @@
+import argparse
+
+from tinyscibio import parse_path
+
+
+def parse_cmd() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--bam",
+        metavar="FILE",
+        type=parse_path,
+        required=True,
+        help="specify path to BAM file",
+    )
+    parser.add_argument(
+        "--ref",
+        metavar="FILE",
+        type=parse_path,
+        required=True,
+        help="specify path to HLA reference in Fasta format.",
+    )
+    parser.add_argument(
+        "--tag",
+        metavar="FILE",
+        type=parse_path,
+        required=True,
+        help="specify path to HLA tag sequence file.",
+    )
+    parser.add_argument(
+        "--bed",
+        metavar="FILE",
+        type=parse_path,
+        required=True,
+        help="specify path to HLA region in BED format.",
+    )
+    parser.add_argument(
+        "--freq",
+        metavar="FILE",
+        type=parse_path,
+        required=True,
+        help="specify path to HLA frequency file.",
+    )
+    parser.add_argument(
+        "--outdir",
+        metavar="DIR",
+        type=parse_path,
+        required=True,
+        help="specify path to output folder.",
+    )
+    parser.add_argument(
+        "--min-ecnt",
+        metavar="INT",
+        type=int,
+        default=999,
+        help="specify minimum # of mm events (999).",
+    )
+    parser.add_argument(
+        "--nproc",
+        metavar="INT",
+        type=int,
+        default=8,
+        help="specify # processes to use (8).",
+    )
+    parser.add_argument(
+        "--overwrite", action="store_true", help="specify to overwrite scores."
+    )
+    parser.add_argument(
+        "--no-clean",
+        action="store_true",
+        help="specify to NOT clean intermediate files.",
+    )
+    parser.add_argument(
+        "--debug", action="store_true", help="specify to enter debug mode."
+    )
+    return parser

mhcflow/dumper.py

@@ -0,0 +1,83 @@
+import multiprocessing as mp
+from functools import partial
+
+import numpy as np
+import polars as pl
+from tinyscibio import BAMetadata, _PathLike, make_dir, parse_path
+
+from .helper import (
+    FileManifest,
+    _check_rg_exists,
+    _check_single_rg,
+    _get_sm,
+    _verify_prev_run,
+)
+from .logger import logger
+from .runnable import _samtools_fastq
+
+
+def _bam2fq_from_idx(
+    idx_fspath: _PathLike,
+    bametadata: BAMetadata,
+    outdir: _PathLike,
+    nproc: int = 1,
+    overwrite: bool = False,
+) -> FileManifest:
+    logger.info(f"Start to dump reads to fastq given ids: {idx_fspath}")
+    # set up output dir
+    outdir = parse_path(outdir)
+    make_dir(outdir, parents=True, exist_ok=True)
+
+    # get sm field from rg
+    _check_rg_exists(bametadata)
+    _check_single_rg(bametadata)
+    rg = bametadata.read_groups[0]
+    sm = _get_sm(rg)
+
+    fm_json = outdir / f"{sm}.dumper.file_manifest.json"
+    if fm_json.exists():
+        logger.info(
+            f"Detected file manifest from previous dumper run: {str(fm_json)}"
+        )
+        fm = FileManifest._from_json(fm_json)
+        if _verify_prev_run(fm, overwrite):
+            return fm
+
+    fm = FileManifest()
+    logdir = outdir / "log"
+    make_dir(logdir, parents=True, exist_ok=True)
+    dumper_done = logdir / f"{sm}.dumper.done"
+    idx_df = pl.read_csv(idx_fspath, separator="\t")
+    if "qnames" not in idx_df.columns:
+        raise pl.exceptions.ColumnNotFoundError()
+
+    # split read ids into batches and output each batch to file
+    # each idx file will be input to samtools view -N
+    qnames = idx_df["qnames"].unique().to_numpy()
+    qnames_batches = np.array_split(qnames, nproc)
+    idxs = []
+    for i in range(len(qnames_batches)):
+        qname_batch_fspath = outdir / f"{sm}.fisher.idxs.{i}.txt"
+        pl.DataFrame({"qnames": qnames_batches[i]}).write_csv(
+            qname_batch_fspath, include_header=False
+        )
+        idxs.append(qname_batch_fspath)
+    # run samtools fastq to extract read and dump to fq
+    r1s, r2s = [], []
+    dones, logs = [], []
+    with mp.Pool(processes=nproc) as pool:
+        for res in pool.imap_unordered(
+            partial(_samtools_fastq, bam_fspath=bametadata.fspath), idxs
+        ):
+            r1, r2, done, log = res
+            r1s.append(r1)
+            r2s.append(r2)
+            dones.append(done)
+            logs.append(log)
+    dumper_done.touch()
+    fm._register_inputs(idx=idx_fspath, bam=bametadata.fspath)
+    fm._register_outputs(r1s=r1s, r2s=r2s, idxs=idxs)
+    fm._register_aux(done=dumper_done)
+    fm._register_intermediate_aux(dones=dones, logs=logs)
+    fm._to_json(fm_json)
+    return fm

mhcflow/finalizer.py

@@ -0,0 +1,100 @@
+from pathlib import Path
+
+import polars as pl
+from pyfaidx import Faidx
+from tinyscibio import BAMetadata, _PathLike, make_dir, parse_path
+
+from .helper import (
+    FileManifest,
+    _check_rg_exists,
+    _check_single_rg,
+    _get_sm,
+    _verify_prev_run,
+)
+from .logger import logger
+from .realigner import _run_realigner
+from .runnable import _novoindex
+
+
+def _dump_seq(allele: str, fa: Faidx, out: Path) -> None:
+    with open(out, "a") as f:
+        sequence = fa.fetch(allele, 1, fa.index[allele].rlen)
+        f.write(f">{allele}\n{str(sequence)}\n")
+
+
+def _run_finalizer(
+    bam_fspath: _PathLike,
+    ref: _PathLike,
+    fisher_fm_json: _PathLike,
+    typer_res_fspath: _PathLike,
+    outdir: _PathLike,
+    nproc: int = 1,
+    overwrite: bool = False,
+) -> FileManifest:
+    logger.info("Finalize sample HLA reference and realignemnt.")
+    bametadata = BAMetadata(str(bam_fspath))
+    _check_rg_exists(bametadata)
+    _check_single_rg(bametadata)
+    rg = bametadata.read_groups[0]
+    sm = _get_sm(rg)
+
+    outdir = parse_path(outdir)
+    make_dir(outdir, parents=True, exist_ok=True)
+
+    fm_json = outdir / f"{sm}.finalizer.file_manifest.json"
+    if fm_json.exists():
+        logger.info(
+            f"Detected file manifest from previous run: {str(fm_json)}"
+        )
+        finalizer_fm = FileManifest._from_json(fm_json)
+        if _verify_prev_run(finalizer_fm, overwrite):
+            return finalizer_fm
+
+    finalizer_fm = FileManifest()
+    logdir = outdir / "log"
+    make_dir(logdir, parents=True, exist_ok=True)
+    finalizer_done = logdir / f"{sm}.finalizer.done"
+
+    ref = parse_path(ref)
+    fai = ref.parent / parse_path(f"{ref.name}.fai")
+    if not fai.exists():
+        raise FileNotFoundError
+
+    typer_res = pl.read_csv(typer_res_fspath, separator="\t")
+    # 3 locus * 2 alleles = 6
+    if typer_res.shape[0] != 6:
+        raise ValueError
+
+    fa_out = outdir / f"{sm}.hla.fasta"
+    fa = Faidx(ref)
+    # do not forget to take unique for homozygous genotype of HLA gene.
+    # I dont want duplicated sequences in the fasta.
+    _ = list(
+        map(
+            lambda x: _dump_seq(x, fa, fa_out),
+            typer_res["allele"].unique().to_list(),
+        )
+    )
+    nix, index_log, index_done = _novoindex(fa_out)
+
+    realigner_fm = _run_realigner(
+        bam_fspath, nix, fisher_fm_json, outdir, nproc, overwrite
+    )
+    finalizer_done.touch()
+
+    # record file manifest for finalizer
+    finalizer_fm._register_inputs(hlaref=str(ref), typer_res=typer_res_fspath)
+    finalizer_fm._register_aux(done=finalizer_done, myself=fm_json)
+    finalizer_fm._register_outputs(
+        sample_hlaref=str(fa_out),
+        **realigner_fm.outputs,
+    )
+    finalizer_fm._register_intermediate(**realigner_fm.intermediates)
+    finalizer_fm._register_intermediate_aux(
+        index_log=index_log,
+        index_done=index_done,
+        **realigner_fm.intermediate_aux,
+    )
+    finalizer_fm._to_json(fm_json)
+
+    return finalizer_fm

mhcflow/fisher.py

@@ -0,0 +1,287 @@
+import multiprocessing as mp
+import time
+from functools import partial
+
+import numpy as np
+import polars as pl
+import pysam
+from tinyscibio import (
+    BAMetadata,
+    _PathLike,
+    bed_to_df,
+    make_dir,
+    parse_path,
+    walk_bam,
+)
+
+from .dumper import _bam2fq_from_idx
+from .helper import (
+    FileManifest,
+    _check_rg_exists,
+    _check_single_rg,
+    _get_sm,
+    _verify_prev_run,
+)
+from .logger import logger
+from .tag_builder import PREBUILD_TYPE, build
+
+CHR6 = ["6", "chr6", "NC00006", "CM000668"]
+
+
+def _fish_unplaced(
+    bam_fspath: _PathLike, prebuilt_tag: PREBUILD_TYPE, out: _PathLike
+) -> tuple[pl.DataFrame, _PathLike]:
+    logger.info("Fish unplaced sequence with tag pattern")
+    out = parse_path(out)
+    logdir = out.parent / "log"
+    make_dir(logdir, parents=True, exist_ok=True)
+    done = logdir / f"{out.stem}.done"
+    if done.exists():
+        logger.info(
+            "Found result of fished sequences with tag pattern "
+            f"from previous run: {done}. Skip."
+        )
+        df = pl.read_csv(out, separator="\t")
+        return (df, done)
+    qnames = []
+    start_t = time.time()
+    with pysam.AlignmentFile(str(bam_fspath), "rb") as bamf:
+        for aln in bamf.fetch(until_eof=True):
+            if not aln.is_unmapped:
+                continue
+            assert aln.query_sequence is not None
+            match = list(prebuilt_tag.iter(aln.query_sequence))
+            if match:
+                qnames += [aln.query_name]
+    logger.info(
+        f"Fish unplaced sequence with tag pattern: {time.time() - start_t} sec"
+    )
+    logger.info(f"Fished {len(qnames)} unplaced sequence with tag pattern")
+    merged_qnames = (
+        pl.DataFrame({"qnames": qnames}) if qnames else pl.DataFrame()
+    )
+    merged_qnames.write_csv(out, separator="\t")
+    done.touch()
+    return merged_qnames, done
+
+
+def _fish_one_hla(region: str, bam_fspath: str) -> pl.DataFrame:
+    return walk_bam(bam_fspath, region, exclude=0, return_qname=True)
+
+
+def _fish_multi_hla(
+    bed_fsapth: _PathLike,
+    bam_fspath: _PathLike,
+    out: _PathLike,
+    nproc: int = 4,
+) -> tuple[pl.DataFrame, _PathLike]:
+    logger.info(f"Fish sequence mapped to regions defined in {bed_fsapth}.")
+    out = parse_path(out)
+    logdir = out.parent / "log"
+    make_dir(logdir, parents=True, exist_ok=True)
+    done = logdir / f"{out.stem}.done"
+    if done.exists():
+        logger.info(
+            "Found result of fished sequences with tag pattern "
+            f"from previous run: {done}. Skip."
+        )
+        df = pl.read_csv(out, separator="\t")
+        return (df, done)
+    df = bed_to_df(bed_fsapth)
+    regions = [f"{row[0]}:{row[1]}-{row[2]}" for row in df.rows()]
+    nproc = min(len(regions), nproc)  # nproc set to minimum of these 2 values
+    qnames: list[pl.DataFrame] = []
+    start_t = time.time()
+    with mp.get_context("spawn").Pool(processes=nproc) as pool:
+        for res in pool.imap_unordered(
+            partial(_fish_one_hla, bam_fspath=str(bam_fspath)),
+            regions,
+        ):
+            if res is not None:
+                qnames += [res.select("qnames")]
+    logger.info(
+        "Fished sequence mapped to regions defined in BED file: "
+        f"{time.time() - start_t} sec."
+    )
+    logger.info(
+        f"Fished {len(qnames)} mapped to HLA regions defined in BED file"
+    )
+    merged_qnames = pl.concat(qnames) if qnames else pl.DataFrame()
+    merged_qnames.write_csv(out, separator="\t")
+    done.touch()
+    return (merged_qnames, done)
+
+
+def _fish_one_region(
+    region: tuple[str, int, int],
+    bam_fspath: _PathLike,
+    prebuilt_tag: PREBUILD_TYPE,
+) -> pl.DataFrame | None:
+    qnames = []
+    sn, start, end = region
+    with pysam.AlignmentFile(str(bam_fspath), "rb") as bamf:
+        for aln in bamf.fetch(contig=sn, start=start, stop=end):
+            assert aln.query_sequence is not None
+            match = list(prebuilt_tag.iter(aln.query_sequence))
+            if match:
+                qnames += [aln.query_name]
+    return pl.DataFrame({"qnames": qnames}) if qnames else None
+
+
+def _fish_multi_regions(
+    split_regions: list[tuple[str, int, int]],
+    bam_fspath: _PathLike,
+    prebuilt_tag: PREBUILD_TYPE,
+    out: _PathLike,
+    nproc: int = 4,
+) -> tuple[pl.DataFrame, _PathLike]:
+    logger.info("Fish sequence with tag pattern.")
+    out = parse_path(out)
+    logdir = out.parent / "log"
+    make_dir(logdir, parents=True, exist_ok=True)
+    done = logdir / f"{out.stem}.done"
+    if done.exists():
+        logger.info(
+            "Found result of fished sequences with tag pattern "
+            f"from previous run: {done}. Skip."
+        )
+        df = pl.read_csv(out, separator="\t")
+        return (df, done)
+    start_t = time.time()
+    qnames: list[pl.DataFrame] = []
+    with mp.Pool(processes=nproc) as pool:
+        for res in pool.imap_unordered(
+            partial(
+                _fish_one_region,
+                bam_fspath=bam_fspath,
+                prebuilt_tag=prebuilt_tag,
+            ),
+            split_regions,
+        ):
+            if res is not None:
+                qnames += [res]
+    logger.info(f"Fish sequence with tag pattern: {time.time() - start_t} sec")
+    logger.info(
+        f"Fished {sum([df.shape[0] for df in qnames])} readss with tag pattern."
+    )
+    # merged_qnames will be empty dataframe if there were no df returned
+    merged_qnames = pl.concat(qnames) if qnames else pl.DataFrame()
+    # empty file will be generated when merged_qnames is empty df
+    merged_qnames.write_csv(out, separator="\t")
+    done.touch()
+    return (merged_qnames, done)
+
+
+def _split_regions(
+    regions: dict[str, list[int]],
+    by: str,
+    n_splits: int = 4,
+    split_size: int = 500_000,
+) -> list[tuple[str, int, int]]:
+    if by not in ["len", "num"]:
+        raise ValueError(f"Unsupported split by method {by=}")
+    logger.info("Split regions into smaller intevals.")
+    splits = []
+    for sn, region in regions.items():
+        start, end = region
+        n_split = (
+            ((end - start + 1) // split_size) + 1 if by == "len" else n_splits
+        )
+        s = np.array_split(np.arange(end), n_split)
+        splits += [(sn, int(split[0]), int(split[-1])) for split in s]
+
+    logger.info(f"Split regions into {len(splits)} intevals.")
+    return splits
+
+
+def _run_fisher(
+    bam_fspath: _PathLike,
+    tag_fspath: _PathLike,
+    bed_fspath: _PathLike,
+    outdir: _PathLike,
+    prebuild_method: str = "ahocorasick",
+    nproc: int = 4,
+    overwrite: bool = False,
+) -> FileManifest:
+    logger.info("Start to fish HLA-relevant reads.")
+    outdir = parse_path(outdir)
+    make_dir(outdir, parents=True, exist_ok=True)
+
+    bametadata = BAMetadata(str(bam_fspath))
+    _check_rg_exists(bametadata)
+    _check_single_rg(bametadata)
+    rg = bametadata.read_groups[0]
+    sm = _get_sm(rg)
+
+    fm_json = outdir / f"{sm}.fisher.file_manifest.json"
+    if fm_json.exists():
+        logger.info(
+            f"Detected file manifest from previous fisher run: {str(fm_json)}"
+        )
+        fisher_fm = FileManifest._from_json(fm_json)
+        if _verify_prev_run(fisher_fm, overwrite):
+            return fisher_fm
+
+    fisher_fm = FileManifest()
+    logdir = outdir / "log"
+    make_dir(logdir, parents=True, exist_ok=True)
+    fisher_done = logdir / f"{sm}.fisher.done"
+
+    prebuilt_tag = build(tag_fspath, method=prebuild_method)
+
+    bam_fspath = bametadata.fspath
+    hla_qname_out = outdir / f"{sm}.fisher.hla_bed.idx"
+    hla_bed_qnames, hla_bed_done = _fish_multi_hla(
+        bed_fspath, bam_fspath, hla_qname_out
+    )
+
+    regions = {
+        sn: [1, ln] for sn, ln in bametadata.seqmap().items() if sn in CHR6
+    }
+    splits = _split_regions(regions, n_splits=nproc, by="num")
+    chr6_qname_out = outdir / f"{sm}.fisher.chr6.idx"
+    chr6_qnames, chr6_done = _fish_multi_regions(
+        splits, bam_fspath, prebuilt_tag, chr6_qname_out, nproc
+    )
+    unplaced_qname_out = outdir / f"{sm}.fisher.unplaced.idx"
+    unplaced_qnames, unplaced_done = _fish_unplaced(
+        bam_fspath, prebuilt_tag, unplaced_qname_out
+    )
+
+    merged_qnames = pl.concat(
+        [hla_bed_qnames, chr6_qnames, unplaced_qnames]
+    ).unique()
+    fisher_idx_out = outdir / f"{sm}.fisher.idx.final.tsv"
+    merged_qnames.write_csv(fisher_idx_out, separator="\t")
+
+    dumper_fm = _bam2fq_from_idx(fisher_idx_out, bametadata, outdir, nproc)
+
+    logger.info(f"Fished {merged_qnames.shape[0]} in total.")
+    fisher_done.touch()
+
+    logger.info("Register all relevant files to manifest.")
+    # register all relevant files to manifest
+    fisher_fm._register_inputs(
+        bam=bametadata.fspath, tag=tag_fspath, bed=bed_fspath
+    )
+    fisher_fm._register_aux(done=fisher_done, myself=fm_json)
+    fisher_fm._register_outputs(
+        **dumper_fm.outputs,
+        hla_fished_idx=hla_qname_out,
+        chr6_fished_idx=chr6_qname_out,
+        unplaced_fished_idx=unplaced_qname_out,
+        fished_all_idx=fisher_idx_out,
+    )
+    fisher_fm._register_intermediate(**dumper_fm.outputs)
+    fisher_fm._register_intermediate_aux(
+        **dumper_fm.intermediate_aux,
+        **dumper_fm.aux,
+        hla_bed_done=hla_bed_done,
+        chr6_done=chr6_done,
+        unplaced_done=unplaced_done,
+    )
+    # dump manifest to disk in json format.
+    logger.info(f"Dump manifest to {fm_json}")
+    fisher_fm._to_json(json_out=fm_json)
+
+    return fisher_fm