From 7a0e9d7155c1267477db0b18461bb07a13964914 Mon Sep 17 00:00:00 2001
From: Benedetto Polimeni <34317613+bepoli@users.noreply.github.com>
Date: Tue, 21 Oct 2025 12:04:26 +0200
Subject: [PATCH 1/6] update docstrings

---
 irescue/map.py | 42 ++++++++++++++++++++++++++++++++++--------
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/irescue/map.py b/irescue/map.py
index b7a4347..86134bc 100644
--- a/irescue/map.py
+++ b/irescue/map.py
@@ -10,8 +10,8 @@
 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,9 +34,27 @@ 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, tmpdir, outname="rmsk.bed"):
+    """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.
+        tmpdir (str): Path to temporary directory.
+        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):
@@ -125,17 +143,21 @@ 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")
@@ -172,7 +194,6 @@ def getRefs(bamFile, bedFile):
         )
 
 
-# Intersect reads with repeatmasker regions. Return the intersection file path.
 def isec(
     bamFile,
     bedFile,
@@ -188,6 +209,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)

From 88bbb47d9b72bffbde5056e02af0b36b5e5ede4d Mon Sep 17 00:00:00 2001
From: Benedetto Polimeni <34317613+bepoli@users.noreply.github.com>
Date: Tue, 21 Oct 2025 12:04:57 +0200
Subject: [PATCH 2/6] fix formatting

---
 irescue/map.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/irescue/map.py b/irescue/map.py
index 86134bc..c71fe72 100644
--- a/irescue/map.py
+++ b/irescue/map.py
@@ -50,7 +50,7 @@ def makeRmsk(regions, genome, genomes, tmpdir, outname="rmsk.bed"):
 
     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.
@@ -260,12 +260,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 = ""

From 0dc132ef657c80ab79de9299bb754eab9545d8c4 Mon Sep 17 00:00:00 2001
From: Benedetto Polimeni <34317613+bepoli@users.noreply.github.com>
Date: Tue, 21 Oct 2025 16:18:46 +0200
Subject: [PATCH 3/6] Add locus-level quantification option

---
 irescue/main.py | 15 ++++++++++-
 irescue/map.py  | 66 +++++++++++++++++++++++++++++++------------------
 2 files changed, 56 insertions(+), 25 deletions(-)

diff --git a/irescue/main.py b/irescue/main.py
index e171afd..525b706 100644
--- a/irescue/main.py
+++ b/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",
@@ -295,6 +306,7 @@ def main():
         genome=args.genome,
         genomes=__genomes__,
         tmpdir=dirs["tmp"],
+        locus=args.locus,
         outname="rmsk.bed",
     )
 
@@ -341,6 +353,7 @@ def main():
         threads=args.threads,
         outdir=dirs["mex"],
         tmpdir=dirs["tmp"],
+        locus=args.locus,
         bedtools=args.bedtools,
         verbose=args.verbose,
     )
diff --git a/irescue/map.py b/irescue/map.py
index c71fe72..14f4904 100644
--- a/irescue/map.py
+++ b/irescue/map.py
@@ -1,8 +1,9 @@
 #!/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
@@ -34,7 +35,9 @@ def checkIndex(bamFile, verbose):
                 )
 
 
-def makeRmsk(regions, genome, genomes, tmpdir, outname="rmsk.bed"):
+def makeRmsk(
+    regions, genome, genomes, tmpdir, locus=False, outname="rmsk.bed"
+):
     """Format and/or download RepeatMasker annotation.
 
     Check repeatmasker regions bed file format. Download if not provided.
@@ -46,6 +49,7 @@ def makeRmsk(regions, genome, genomes, tmpdir, outname="rmsk.bed"):
         genome (str): Genome assembly name.
         genomes (dict): Dictionary of genome assembly names and URLs.
         tmpdir (str): Path to temporary directory.
+        locus (bool): If True, prepare for locus-level quantification.
         outname (str): Name of the output repeatmasker bed file.
 
     Returns:
@@ -57,27 +61,22 @@ def makeRmsk(regions, genome, genomes, tmpdir, outname="rmsk.bed"):
     """
     # if a repeatmasker bed file is provided, use that
     if regions:
-        if testGz(regions):
-            f = gzopen(regions, "rb")
-
-            def rl(x):
-                return x.readline().decode()
-        else:
-            f = open(regions, "r")
+        is_gz = testGz(regions)
+        f = gzip.open(regions, "rb") if is_gz else 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, is_gz)
         while line[0] == "#":
-            line = rl(f)
+            line = rl(f, 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()
@@ -98,7 +97,7 @@ def rl(x):
                 f"Couldn't connect to host.\n\n{e}",
                 error=True,
             )
-        rmsk = gzopen(io.BytesIO(response.content), "rb")
+        rmsk = gzip.open(io.BytesIO(response.content), "rb")
         out = os.path.join(tmpdir, outname)
         with open(out, "w") as f:
             # print header
@@ -118,20 +117,25 @@ 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
+                repname += "#" + famclass
+                if locus:
+                    # make unique locus names
+                    subfamilies[repname] += 1
+                    repname += f"~{subfamilies[repname]}"
                 score = lst[0]
                 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, score, strand])
                 outl += "\n"
                 f.write(outl)
     else:
@@ -165,7 +169,7 @@ def getRefs(bamFile, bedFile):
             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:
@@ -289,8 +293,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")
@@ -318,7 +334,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)

From 7d06c5a6008a44b41ea6b8e77f3fc44526289fdf Mon Sep 17 00:00:00 2001
From: Benedetto Polimeni <34317613+bepoli@users.noreply.github.com>
Date: Tue, 21 Oct 2025 17:04:28 +0200
Subject: [PATCH 4/6] Fix argument name in rl function call

---
 irescue/map.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/irescue/map.py b/irescue/map.py
index 14f4904..881525f 100644
--- a/irescue/map.py
+++ b/irescue/map.py
@@ -68,9 +68,9 @@ def rl(x, decode=False):
             return x.readline().decode() if decode else x.readline()
 
         # skip header
-        line = rl(f, is_gz)
+        line = rl(f, decode=is_gz)
         while line[0] == "#":
-            line = rl(f, is_gz)
+            line = rl(f, decode=is_gz)
         # check for minimum column number
         if len(line.strip().split("\t")) < 4:
             writerr(

From ee84cfc4ca2eca3c1cef25fb43e5f5c9dce08758 Mon Sep 17 00:00:00 2001
From: Benedetto Polimeni <34317613+bepoli@users.noreply.github.com>
Date: Tue, 21 Oct 2025 17:09:05 +0200
Subject: [PATCH 5/6] add test for locus-level

---
 tests/test.yml | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tests/test.yml b/tests/test.yml
index 704f2b2..525f04b 100644
--- a/tests/test.yml
+++ b/tests/test.yml
@@ -123,3 +123,19 @@
       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

From 792db379ea58f6dba70955ac19638cf724513778 Mon Sep 17 00:00:00 2001
From: Benedetto Polimeni <34317613+bepoli@users.noreply.github.com>
Date: Wed, 22 Oct 2025 09:53:44 +0200
Subject: [PATCH 6/6] write on-the-fly-generated TE annotation to output
 directory

---
 irescue/main.py |  4 ++--
 irescue/map.py  | 30 ++++++++++++++++--------------
 tests/test.yml  |  4 ++++
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/irescue/main.py b/irescue/main.py
index 525b706..dc7d5e5 100644
--- a/irescue/main.py
+++ b/irescue/main.py
@@ -305,9 +305,9 @@ def main():
         regions=args.regions,
         genome=args.genome,
         genomes=__genomes__,
-        tmpdir=dirs["tmp"],
+        outdir=dirs["out"],
         locus=args.locus,
-        outname="rmsk.bed",
+        outname="rmsk.bed.gz",
     )
 
     # get list of reference names from bam
diff --git a/irescue/map.py b/irescue/map.py
index 881525f..06ddecf 100644
--- a/irescue/map.py
+++ b/irescue/map.py
@@ -36,7 +36,7 @@ def checkIndex(bamFile, verbose):
 
 
 def makeRmsk(
-    regions, genome, genomes, tmpdir, locus=False, outname="rmsk.bed"
+    regions, genome, genomes, outdir, locus=False, outname="rmsk.bed.gz"
 ):
     """Format and/or download RepeatMasker annotation.
 
@@ -48,7 +48,7 @@ def makeRmsk(
             Takes priority over genome.
         genome (str): Genome assembly name.
         genomes (dict): Dictionary of genome assembly names and URLs.
-        tmpdir (str): Path to temporary directory.
+        outdir (str): Path to output directory.
         locus (bool): If True, prepare for locus-level quantification.
         outname (str): Name of the output repeatmasker bed file.
 
@@ -98,13 +98,12 @@ def rl(x, decode=False):
                 error=True,
             )
         rmsk = gzip.open(io.BytesIO(response.content), "rb")
-        out = os.path.join(tmpdir, outname)
-        with open(out, "w") as f:
+        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)
@@ -125,19 +124,22 @@ def rl(x, decode=False):
                     continue
                 # concatenate family and class with subfamily
                 repname += "#" + famclass
+                subfamilies[repname] += 1
+                locus_index = subfamilies[repname]
                 if locus:
                     # make unique locus names
-                    subfamilies[repname] += 1
-                    repname += f"~{subfamilies[repname]}"
-                score = lst[0]
+                    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, repname, 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 "
@@ -335,7 +337,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 else "."
     cmd2 += '+/,"",1,$1)"\\tGene Expression" }\' '
     cmd2 += f" | LC_ALL=C sort -u | gzip > {features_file} "
 
diff --git a/tests/test.yml b/tests/test.yml
index 525f04b..a743008 100644
--- a/tests/test.yml
+++ b/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:
@@ -139,3 +141,5 @@
       md5sum: 29b3efa69a460a64af88bd37ea3794c4
     - path: "irescue_out/tmp/mappings.tsv.gz"
       md5sum: d7db121c92c04b36e3cc26ae0223426d
+    - path: "irescue_out/rmsk.bed.gz"
+      md5sum: 75ed0333b049a02672da8b1379cfa6bc