Add locus-level quantification option

irescue/main.py

@@ -99,10 +99,21 @@ def parseArguments():
         metavar="STR",
         help="BAM tag containing the UMI sequence (default: %(default)s).",
     )
+    parser.add_argument(
+        "-l",
+        "--locus",
+        action="store_true",
+        help=(
+            "Perform locus-level quantification, instead of subfamily-level"
+            " (default: %(default)s)."
+        ),
+    )
     parser.add_argument(
         "--no-umi",
         action="store_true",
-        help="Ignore UMI sequence (for UMI-less datasets, such as Smart-seq).",
+        help="Ignore UMI sequence."
+            " Intended for UMI-less datasets, such as Smart-seq"
+            " (default: %(default)s).",
     )
     parser.add_argument(
         "-p",
@@ -294,8 +305,9 @@ def main():
         regions=args.regions,
         genome=args.genome,
         genomes=__genomes__,
-        tmpdir=dirs["tmp"],
-        outname="rmsk.bed",
+        outdir=dirs["out"],
+        locus=args.locus,
+        outname="rmsk.bed.gz",
     )
 
     # get list of reference names from bam
@@ -341,6 +353,7 @@ def main():
         threads=args.threads,
         outdir=dirs["mex"],
         tmpdir=dirs["tmp"],
+        locus=args.locus,
         bedtools=args.bedtools,
         verbose=args.verbose,
     )

irescue/map.py

@@ -1,17 +1,18 @@
 #!/usr/bin/env python
 
+import gzip
 import io
 import os
-from gzip import open as gzopen
+from collections import defaultdict
 
 import requests
 from pysam import AlignmentFile, idxstats, index
 
 from irescue.misc import getlen, run_shell_cmd, testGz, unGzip, writerr
 
 
-# Check if bam file is indexed
 def checkIndex(bamFile, verbose):
+    """Check if BAM file is indexed. If not, attempt to index it."""
     with AlignmentFile(bamFile) as bam:
         if not bam.has_index():
             writerr("BAM index not found. Attempting to index the BAM...")
@@ -34,32 +35,48 @@ def checkIndex(bamFile, verbose):
                 )
 
 
-# Check repeatmasker regions bed file format. Download if not provided.
-# Returns the path of the repeatmasker bed file.
-def makeRmsk(regions, genome, genomes, tmpdir, outname):
+def makeRmsk(
+    regions, genome, genomes, outdir, locus=False, outname="rmsk.bed.gz"
+):
+    """Format and/or download RepeatMasker annotation.
+
+    Check repeatmasker regions bed file format. Download if not provided.
+    Returns the path of the repeatmasker bed file.
+
+    Args:
+        regions (str): Path to repeatmasker bed file.
+            Takes priority over genome.
+        genome (str): Genome assembly name.
+        genomes (dict): Dictionary of genome assembly names and URLs.
+        outdir (str): Path to output directory.
+        locus (bool): If True, prepare for locus-level quantification.
+        outname (str): Name of the output repeatmasker bed file.
+
+    Returns:
+        str: Path to the repeatmasker bed file.
+
+    Raises:
+        SystemExit: If neither regions nor genome is provided, or if the
+                    regions file is not properly formatted.
+    """
     # if a repeatmasker bed file is provided, use that
     if regions:
-        if testGz(regions):
-            f = gzopen(regions, "rb")
+        is_gz = testGz(regions)
+        f = gzip.open(regions, "rb") if is_gz else open(regions, "r")
 
-            def rl(x):
-                return x.readline().decode()
-        else:
-            f = open(regions, "r")
-
-            def rl(x):
-                return x.readline()
+        def rl(x, decode=False):
+            return x.readline().decode() if decode else x.readline()
 
         # skip header
-        line = rl(f)
+        line = rl(f, decode=is_gz)
         while line[0] == "#":
-            line = rl(f)
+            line = rl(f, decode=is_gz)
         # check for minimum column number
         if len(line.strip().split("\t")) < 4:
             writerr(
-                "Error: please provide a tab-separated BED file with at "
-                "least 4 columns and TE feature name (e.g. subfamily) "
-                "in 4th column.",
+                "Error: please provide a tab-separated BED file with at least"
+                " 4 columns and TE feature name (e.g. locus or subfamily)"
+                " in 4th column.",
                 error=True,
             )
         f.close()
@@ -80,14 +97,13 @@ def rl(x):
                 f"Couldn't connect to host.\n\n{e}",
                 error=True,
             )
-        rmsk = gzopen(io.BytesIO(response.content), "rb")
-        out = os.path.join(tmpdir, outname)
-        with open(out, "w") as f:
+        rmsk = gzip.open(io.BytesIO(response.content), "rb")
+        out = os.path.join(outdir, outname)
+        with gzip.GzipFile(out, "wb", mtime=0) as f:
             # print header
-            h = ["#chr", "start", "end", "name", "score", "strand"]
-            h = "\t".join(h)
-            h += "\n"
-            f.write(h)
+            h = ["#chr", "start", "end", "name", "locus_index", "strand"]
+            h = "\t".join(h) + "\n"
+            f.write(h.encode())
             # skip rmsk header
             for _ in range(header_lines):
                 next(rmsk)
@@ -100,22 +116,30 @@ def rl(x):
                 "srpRNA",
                 "tRNA",
             ]
+            subfamilies = defaultdict(int)
             for line in rmsk:
                 lst = line.decode("utf-8").strip().split()
-                strand, subfamily, famclass = lst[8:11]
+                strand, repname, famclass = lst[8:11]
                 if famclass.split("/")[0] in fams_to_skip:
                     continue
                 # concatenate family and class with subfamily
-                subfamily += "#" + famclass
-                score = lst[0]
+                repname += "#" + famclass
+                subfamilies[repname] += 1
+                locus_index = subfamilies[repname]
+                if locus:
+                    # make unique locus names
+                    repname += f"~{locus_index}"
                 chr, start, end = lst[4:7]
                 # make coordinates 0-based
                 start = str(int(start) - 1)
                 if strand != "+":
                     strand = "-"
-                outl = "\t".join([chr, start, end, subfamily, score, strand])
+                outl = "\t".join(
+                    [chr, start, end, repname, str(locus_index), strand]
+                )
                 outl += "\n"
-                f.write(outl)
+                f.write(outl.encode())
+        writerr(f"Wrote RepeatMasker annotation to {out}.")
     else:
         writerr(
             "Error: it is mandatory to define either --regions OR "
@@ -125,25 +149,29 @@ def rl(x):
     return out
 
 
-# Uncompress the whitelist file if compressed.
-# Return the whitelist path, or False if not using a whitelist.
 def prepare_whitelist(whitelist, tmpdir):
+    """Uncompress the whitelist file if compressed.
+    Return the whitelist path, or False if not using a whitelist.
+    """
     if whitelist and testGz(whitelist):
         wlout = os.path.join(tmpdir, "whitelist.tsv")
         whitelist = unGzip(whitelist, wlout)
     return whitelist
 
 
-# Get list of reference names from BAM file, skipping those without reads.
 def getRefs(bamFile, bedFile):
+    """Get list of reference names from BAM file, skips those without reads
+    and checks their presence in the TE annotation bed file.
+    Returns the list of reference names to process.
+    """
     chrNames = list()
     for line in idxstats(bamFile).strip().split("\n"):
         fields = line.strip().split("\t")
         if int(fields[2]) > 0:
             chrNames.append(fields[0])
     bedChrNames = set()
     if testGz(bedFile):
-        with gzopen(bedFile, "rb") as f:
+        with gzip.open(bedFile, "rb") as f:
             for line in f:
                 bedChrNames.add(line.decode().split("\t")[0])
     else:
@@ -172,7 +200,6 @@ def getRefs(bamFile, bedFile):
         )
 
 
-# Intersect reads with repeatmasker regions. Return the intersection file path.
 def isec(
     bamFile,
     bedFile,
@@ -188,6 +215,11 @@ def isec(
     verbose,
     chrom,
 ):
+    """
+    Intersect alignments from bamFile with features from bedFile for a
+    specific chromosome (chrom). Return the path of the intersection file.
+    Intended for parallelization by chromosome.
+    """
     refdir = os.path.join(tmpdir, "refs")
     isecdir = os.path.join(tmpdir, "isec")
     os.makedirs(refdir, exist_ok=True)
@@ -234,12 +266,12 @@ def isec(
     # filter by minimum overlap between read and feature, if set
     ovfrac = f" -f {fracOverlap} " if fracOverlap else ""
     ovbp = f" $NF>={bpOverlap} " if bpOverlap else ""
-    
+
     # strand-specific intersection
     strandedness = strandedness.lower()
-    if strandedness == 'forward':
+    if strandedness == "forward":
         strand = " -s "
-    elif strandedness == 'reverse':
+    elif strandedness == "reverse":
         strand = " -S "
     else:
         strand = ""
@@ -263,8 +295,20 @@ def isec(
     return isecFile
 
 
-# Concatenate and sort data obtained from isec()
-def chrcat(filesList, threads, outdir, tmpdir, bedtools, verbose):
+def chrcat(
+    filesList,
+    threads,
+    outdir,
+    tmpdir,
+    locus=False,
+    bedtools="bedtools",
+    verbose=0,
+):
+    """
+    Concatenate and sort intersection files from isec() function.
+    Write mappings.tsv.gz, barcodes.tsv.gz and features.tsv.gz files.
+    Returns paths of the three output files.
+    """
     os.makedirs(outdir, exist_ok=True)
     mappings_file = os.path.join(tmpdir, "mappings.tsv.gz")
     barcodes_file = os.path.join(outdir, "barcodes.tsv.gz")
@@ -292,7 +336,9 @@ def chrcat(filesList, threads, outdir, tmpdir, bedtools, verbose):
     # write features.tsv.gz file
     cmd2 = f"zcat {mappings_file} "
     cmd2 += " | cut -f3 | sed 's/,/\\n/g' | gawk '!x[$1]++ { "
-    cmd2 += ' print $1"\\t"gensub(/#.+/,"",1,$1)"\\tGene Expression" }\' '
+    cmd2 += ' print $1"\\t"gensub(/#'
+    cmd2 += "[^~]" if locus else "."
+    cmd2 += '+/,"",1,$1)"\\tGene Expression" }\' '
     cmd2 += f" | LC_ALL=C sort -u | gzip > {features_file} "
 
     writerr("Concatenating mappings", level=1, send=verbose)

tests/test.yml

@@ -27,6 +27,8 @@
       md5sum: d71ee82b25107d4e104d313efb4be134
     - path: "irescue_out/tmp/mappings.tsv.gz"
       md5sum: d404e6c3123f8cfde7689b5fb7763a89
+    - path: "irescue_out/rmsk.bed.gz"
+      md5sum: 169571f538624496a00f189771be2f5e
 
 - name: multi
   tags:
@@ -123,3 +125,21 @@
       md5sum: 3f0f7ca61f7af561c2b8723e010cba37
     - path: "irescue_out/tmp/mappings.tsv.gz"
       md5sum: f5aae354f59bef7a4bdb9cf3c5c8dafe
+
+- name: locus
+  tags:
+    - locus
+  command: irescue --keeptmp --dump-ec -vv -b ./tests/data/Aligned.sortedByCoord.out.bam -g test --locus
+  files:
+    - path: "irescue_out/counts/barcodes.tsv.gz"
+      md5sum: 1a74fa12e65ac1703bbe61282854f151
+    - path: "irescue_out/counts/features.tsv.gz"
+      md5sum: 4894ae806fdafe9aad2c4989689fba31
+    - path: "irescue_out/counts/matrix.mtx.gz"
+      md5sum: 0b4c43f61ad89f330f17ccc1844ea437
+    - path: "irescue_out/ec_dump.tsv.gz"
+      md5sum: 29b3efa69a460a64af88bd37ea3794c4
+    - path: "irescue_out/tmp/mappings.tsv.gz"
+      md5sum: d7db121c92c04b36e3cc26ae0223426d
+    - path: "irescue_out/rmsk.bed.gz"
+      md5sum: 75ed0333b049a02672da8b1379cfa6bc