Implement fragments aggregation into instances: modified main.py and map.py

irescue/main.py

@@ -100,12 +100,15 @@ def parseArguments():
         help="BAM tag containing the UMI sequence (default: %(default)s).",
     )
     parser.add_argument(
-        "-l",
-        "--locus",
-        action="store_true",
+        "--locus-level",
+        type=str,
+        metavar="STR",
+        choices=["disabled", "fragment", "instance"],
+        default="disabled", 
         help=(
-            "Perform locus-level quantification, instead of subfamily-level"
-            " (default: %(default)s)."
+            "Perform locus-level quantification. "
+            "One of: disabled, fragment, instance (default: %(default)s). "
+            "Disabled means subfamily-level quantification."
         ),
     )
     parser.add_argument(
@@ -306,7 +309,7 @@ def main():
         genome=args.genome,
         genomes=__genomes__,
         outdir=dirs["out"],
-        locus=args.locus,
+        locus=args.locus_level,
         outname="rmsk.bed.gz",
     )
 
@@ -353,7 +356,7 @@ def main():
         threads=args.threads,
         outdir=dirs["mex"],
         tmpdir=dirs["tmp"],
-        locus=args.locus,
+        locus=args.locus_level,
         bedtools=args.bedtools,
         verbose=args.verbose,
     )

irescue/map.py

@@ -35,9 +35,7 @@ def checkIndex(bamFile, verbose):
                 )
 
 
-def makeRmsk(
-    regions, genome, genomes, outdir, locus=False, outname="rmsk.bed.gz"
-):
+def makeRmsk(regions, genome, genomes, outdir, locus="disabled", outname="rmsk.bed.gz"):
     """Format and/or download RepeatMasker annotation.
 
     Check repeatmasker regions bed file format. Download if not provided.
@@ -49,7 +47,7 @@ def makeRmsk(
         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.
+        locus (str): If not disabled, prepare for locus-level quantification.
         outname (str): Name of the output repeatmasker bed file.
 
     Returns:
@@ -97,18 +95,20 @@ def rl(x, decode=False):
                 f"Couldn't connect to host.\n\n{e}",
                 error=True,
             )
-        rmsk = gzip.open(io.BytesIO(response.content), "rb")
         out = os.path.join(outdir, outname)
-        with gzip.GzipFile(out, "wb", mtime=0) as f:
+        with (
+            gzip.open(io.BytesIO(response.content), "rb") as rmsk_in,
+            gzip.GzipFile(out, "wb", mtime=0) as rmsk_out,
+        ):
             # print header
-            h = ["#chr", "start", "end", "name", "locus_index", "strand"]
+            h = ["#chr", "start", "end", "name", "repInstance", "strand"]
             h = "\t".join(h) + "\n"
-            f.write(h.encode())
+            rmsk_out.write(h.encode())
             # skip rmsk header
             for _ in range(header_lines):
-                next(rmsk)
+                next(rmsk_in)
             # parse rmsk
-            fams_to_skip = [
+            repeats_to_skip = [
                 "Low_complexity",
                 "Simple_repeat",
                 "rRNA",
@@ -117,36 +117,87 @@ def rl(x, decode=False):
                 "tRNA",
             ]
             subfamilies = defaultdict(int)
-            for line in rmsk:
-                lst = line.decode("utf-8").strip().split()
-                strand, repname, famclass = lst[8:11]
-                if famclass.split("/")[0] in fams_to_skip:
+            rmsk_dict = dict()
+            for line in rmsk_in:
+                fields = line.decode("utf-8").strip().split()
+                chr, start, end = fields[4:7]
+                strand, subfamily, clasfam = fields[8:11]
+                instance = fields[14]
+                clas, family = clasfam.split("/") if "/" in clasfam else (clasfam, "")
+                # skip problematic repeats
+                if clas in repeats_to_skip:
                     continue
                 # concatenate family and class with subfamily
-                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)
+                name = f"{subfamily}#{clasfam}"
+
+                # repeat names for locus-level quantifications
+                if locus != "disabled":
+                    subfamilies[name] += 1
+                    locus_index = subfamilies[name]
+                    name += f"~{locus_index}"
+                    if locus == "instance":
+                        if instance in rmsk_dict.keys():
+                            rmsk_dict[instance]["repSubfam"].add(subfamily)
+                            rmsk_dict[instance]["repFam"].add(family)
+                            rmsk_dict[instance]["repClass"].add(clas)
+                        else:
+                            rmsk_dict[instance] = {
+                                "repSubfam": {subfamily},
+                                "repFam": {family},
+                                "repClass": {clas},
+                            }
+
+                start = str(int(start) - 1)  # make coordinates 0-based
                 if strand != "+":
                     strand = "-"
-                outl = "\t".join(
-                    [chr, start, end, repname, str(locus_index), strand]
-                )
+                outl = "\t".join([chr, start, end, name, instance, strand])
                 outl += "\n"
-                f.write(outl.encode())
-        writerr(f"Wrote RepeatMasker annotation to {out}.")
+                rmsk_out.write(outl.encode())
+
+        # if locus-level quantification by instance is requested,
+        # prepare a separate file with instance-level annotation
+        if locus == "instance":
+            outname = "rmsk_instance.bed.gz"
+            out_instance = os.path.join(outdir, outname)
+            with (
+                gzip.open(out, "rb") as rmsk_in,
+                gzip.GzipFile(out_instance, "wb", mtime=0) as rmsk_out,
+            ):
+                next(rmsk_in)  # skip header row
+                # print header
+                h = ["#chr", "start", "end", "name", "repInstance", "strand"]
+                h = "\t".join(h) + "\n"
+                rmsk_out.write(h.encode())
+                for line in rmsk_in:
+                    chr, start, end, name, instance, strand = (
+                        line.decode("utf-8").strip().split("\t")
+                    )
+                    repSubfam = "|".join(sorted(rmsk_dict[instance]["repSubfam"]))
+                    repFam = "|".join(sorted(rmsk_dict[instance]["repFam"]))
+                    repFam = (
+                        "" if not repFam else "/" + repFam
+                    )  # to take into account cases in which repFamily is not annotated
+                    repClass = "|".join(sorted(rmsk_dict[instance]["repClass"]))
+
+                    repname = f"{repSubfam}#{repClass}{repFam}~{instance}"
+                    outl = "\t".join([chr, start, end, repname, instance, strand])
+                    outl += "\n"
+                    rmsk_out.write(outl.encode())
+
+        level = "subfamily" if locus == "disabled" else locus
+        outfile = out if locus != "instance" else out_instance
+        writerr(f"Wrote RepeatMasker {level}-level annotation to {outfile}.")
+
     else:
         writerr(
-            "Error: it is mandatory to define either --regions OR "
-            "--genome parameter.",
+            "Error: it is mandatory to define either --regions OR --genome parameter.",
             error=True,
         )
-    return out
+
+    if locus == "instance":
+        return out_instance
+    else:
+        return out
 
 
 def prepare_whitelist(whitelist, tmpdir):
@@ -242,15 +293,11 @@ def isec(
     else:
         stream = f" <({samtools} view -h {bamFile} {chrom} | "
     stream += ' gawk \'!($1~/^@/) { split("", tags); '
-    stream += (
-        ' for (i=12;i<=NF;i++) {split($i,tag,":"); tags[tag[1]]=tag[3]}; '
-    )
+    stream += ' for (i=12;i<=NF;i++) {split($i,tag,":"); tags[tag[1]]=tag[3]}; '
     # Discard records without CB tag, unvalid STARSolo CBs, missing UMI tag,
     # UMIs with Ns and homopolymer UMIs
     if UMItag:
-        stream += (
-            f' if(tags["{CBtag}"]~/^(|-)$/ || tags["{UMItag}"]~/.*N.*/ || '
-        )
+        stream += f' if(tags["{CBtag}"]~/^(|-)$/ || tags["{UMItag}"]~/.*N.*/ || '
         stream += f' tags["{UMItag}"]~/^$|^(A+|G+|T+|C+)$/) {{next}}; '
         # Append CB and UMI to read name
         stream += f' $1=$1"/"tags["{CBtag}"]"/"tags["{UMItag}"]; '
@@ -300,7 +347,7 @@ def chrcat(
     threads,
     outdir,
     tmpdir,
-    locus=False,
+    locus="disabled",
     bedtools="bedtools",
     verbose=0,
 ):
@@ -337,7 +384,7 @@ def chrcat(
     cmd2 = f"zcat {mappings_file} "
     cmd2 += " | cut -f3 | sed 's/,/\\n/g' | gawk '!x[$1]++ { "
     cmd2 += ' print $1"\\t"gensub(/#'
-    cmd2 += "[^~]" if locus else "."
+    cmd2 += "[^~]" if locus != "disabled" else "."
     cmd2 += '+/,"",1,$1)"\\tGene Expression" }\' '
     cmd2 += f" | LC_ALL=C sort -u | gzip > {features_file} "
 

tests/test.yml

@@ -28,7 +28,7 @@
     - path: "irescue_out/tmp/mappings.tsv.gz"
       md5sum: d404e6c3123f8cfde7689b5fb7763a89
     - path: "irescue_out/rmsk.bed.gz"
-      md5sum: 169571f538624496a00f189771be2f5e
+      md5sum: 663134e9bb5519d136201d2bd5d35d63
 
 - name: multi
   tags:
@@ -126,10 +126,10 @@
     - path: "irescue_out/tmp/mappings.tsv.gz"
       md5sum: f5aae354f59bef7a4bdb9cf3c5c8dafe
 
-- name: locus
+- name: fragment
   tags:
     - locus
-  command: irescue --keeptmp --dump-ec -vv -b ./tests/data/Aligned.sortedByCoord.out.bam -g test --locus
+  command: irescue --keeptmp --dump-ec -vv -b ./tests/data/Aligned.sortedByCoord.out.bam -g test --locus-level fragment
   files:
     - path: "irescue_out/counts/barcodes.tsv.gz"
       md5sum: 1a74fa12e65ac1703bbe61282854f151
@@ -142,4 +142,24 @@
     - path: "irescue_out/tmp/mappings.tsv.gz"
       md5sum: d7db121c92c04b36e3cc26ae0223426d
     - path: "irescue_out/rmsk.bed.gz"
-      md5sum: 75ed0333b049a02672da8b1379cfa6bc
+      md5sum: 050b0a484a04e5d466afd54c794e35fd
+
+- name: instance
+  tags:
+    - locus
+  command: irescue --keeptmp --dump-ec -vv -b ./tests/data/Aligned.sortedByCoord.out.bam -g test --locus-level instance
+  files:
+    - path: "irescue_out/counts/barcodes.tsv.gz"
+      md5sum: 1a74fa12e65ac1703bbe61282854f151
+    - path: "irescue_out/counts/features.tsv.gz"
+      md5sum: 045273e35f85522ec70a7be127fd5d7f
+    - path: "irescue_out/counts/matrix.mtx.gz"
+      md5sum: a63c59d45b79bd4bfae1d7ffb26b6c6e
+    - path: "irescue_out/ec_dump.tsv.gz"
+      md5sum: dc7a39ce64959f0b83dd53c158041746
+    - path: "irescue_out/tmp/mappings.tsv.gz"
+      md5sum: e0986db4bea3929a8f14eb61f2596a41
+    - path: "irescue_out/rmsk.bed.gz"
+      md5sum: 050b0a484a04e5d466afd54c794e35fd
+    - path: "irescue_out/rmsk_instance.bed.gz"
+      md5sum: 7ba1ba91b19292a6d49a34e818470630