Repository: chrovis/cljam
Pull request #329 (2024-12-23 13:28:37)

Add a function to read mate alignments
#### Summary
This PR adds a function `cljam.io.sam/read-mate-alignments` for BAM file which reads R1/R2 counterpart alignments of given alignments.

#### Added functions
This function can be used for collecting alignments around some regions keeping pairs like `samtools view --fetch-pairs`.

```clj
(let [xs (sam/read-alignments reader {:chr "chr1", :start 1, :end 100000})]
  (concat xs (sam/read-mate-alignments reader xs))
```

In the example above, it is necessary to recheck which alignment pairs with which. To make things easier, I added another function `cljam.io.sam/make-pairs` that returns a sequence with paired alignments grouped together.

```clj
(let [xs (sam/read-alignments reader {:chr "chr1", :start 1, :end 100000})]
  (sam/make-pairs reader xs))
```

To achieve these functionalities, I implemented a function `cljam.io.bam-index/get-spans-for-regions` that queries chunks corresponding to multiple regions at once against the BAI index.

#### Implementation details
`cljam.io.sam/read-mate-alignments` for SAM and CRAM is not supported yet.

When searching for mates, a large number of alignment blocks that do not meet the criteria must be discarded. Therefore, the current implementation decodes only the minimal set of fields necessary for condition evaluation and immediately rejects blocks that do not satisfy the criteria.

As a result, a significant part of the execution time is spent decompressing BGZF blocks, so reducing the number of chunks accessed would make a big difference.
However, the current implementation is limited to what can be achieved by combining existing functionalities, leaving room for further optimization in the future.

#### Tests
Since we don't have a good BAM file with paired reads of appropriate size, I added auto-generated test cases for `cljam.io.sam/read-mate-alignments`. It writes a temporary BAM file containing paired alignments on the fly, and then perform reading and checking.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/chrovis/cljam/pull/329?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=chrovis) Report
Attention: Patch coverage is `97.77778%` with `2 lines` in your changes missing coverage. Please review.
> Project coverage is 90.10%. Comparing base [(`18f8cca`)](https://app.codecov.io/gh/chrovis/cljam/commit/18f8cca90f5f210f0239dfce2134fbc02e8c044c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=chrovis) to head [(`6337013`)](https://app.codecov.io/gh/chrovis/cljam/commit/6337013d38eef9bb78cb4c4445a7d8554056640c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=chrovis).

| [Files with missing lines](https://app.codecov.io/gh/chrovis/cljam/pull/329?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=chrovis) | Patch % | Lines |
|---|---|---|
| [src/cljam/io/bam/reader.clj](https://app.codecov.io/gh/chrovis/cljam/pull/329?src=pr&el=tree&filepath=src%2Fcljam%2Fio%2Fbam%2Freader.clj&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=chrovis#diff-c3JjL2NsamFtL2lvL2JhbS9yZWFkZXIuY2xq) | 97.91% | [1 Missing :warning: ](https://app.codecov.io/gh/chrovis/cljam/pull/329?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=chrovis) |
| [src/cljam/io/sam.clj](https://app.codecov.io/gh/chrovis/cljam/pull/329?src=pr&el=tree&filepath=src%2Fcljam%2Fio%2Fsam.clj&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=chrovis#diff-c3JjL2NsamFtL2lvL3NhbS5jbGo=) | 95.45% | [1 Missing :warning: ](https://app.codecov.io/gh/chrovis/cljam/pull/329?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=chrovis) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master     #329      +/-   ##
==========================================
+ Coverage   89.97%   90.10%   +0.12%     
==========================================
  Files         104      104              
  Lines        9341     9435      +94     
  Branches      490      491       +1     
==========================================
+ Hits         8405     8501      +96     
+ Misses        446      443       -3     
- Partials      490      491       +1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/chrovis/cljam/pull/329?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=chrovis).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=chrovis).


Patch:
diff --git a/project.clj b/project.clj
index 09da317..b306279 100644
--- a/project.clj
+++ b/project.clj
@@ -5,14 +5,14 @@
             :url "https://www.apache.org/licenses/LICENSE-2.0"}
   :dependencies [[org.clojure/core.memoize "1.1.266"]
                  [org.clojure/tools.logging "1.3.0"]
-                 [org.apache.commons/commons-compress "1.26.1"]
-                 [org.tukaani/xz "1.9"] ; necessary for CRAM LZMA compression
+                 [org.apache.commons/commons-compress "1.27.1"]
+                 [org.tukaani/xz "1.10"] ; necessary for CRAM LZMA compression
                  [digest "1.4.10"]
                  [bgzf4j "0.1.2"]
                  [com.climate/claypoole "1.1.4"]
                  [camel-snake-kebab "0.4.3"]
                  [proton "0.2.3"]]
-  :profiles {:dev {:dependencies [[org.clojure/clojure "1.11.3"]
+  :profiles {:dev {:dependencies [[org.clojure/clojure "1.12.0"]
                                   [cavia "0.5.1"]
                                   [criterium "0.4.6"]
                                   [net.totakke/libra "0.1.1"]
@@ -31,7 +31,8 @@
              :1.7 {:dependencies [[org.clojure/clojure "1.7.0"]]}
              :1.8 {:dependencies [[org.clojure/clojure "1.8.0"]]}
              :1.9 {:dependencies [[org.clojure/clojure "1.9.0"]]}
-             :1.10 {:dependencies [[org.clojure/clojure "1.10.3"]]}}
+             :1.10 {:dependencies [[org.clojure/clojure "1.10.3"]]}
+             :1.11 {:dependencies [[org.clojure/clojure "1.11.4"]]}}
   :deploy-repositories [["snapshots" {:url "https://clojars.org/repo/"
                                       :username [:env/clojars_username :gpg]
                                       :password [:env/clojars_password :gpg]}]]
diff --git a/src/cljam/io/bam/reader.clj b/src/cljam/io/bam/reader.clj
index be33758..f7686b0 100644
--- a/src/cljam/io/bam/reader.clj
+++ b/src/cljam/io/bam/reader.clj
@@ -3,15 +3,18 @@
   (:require [cljam.io.protocols :as protocols]
             [cljam.io.sam.util.refs :as refs]
             [cljam.io.sam.util.header :as header]
-            [cljam.io.bam-index.core :as bai]
+            [cljam.io.bam-index :as bai]
             [cljam.io.bam.decoder :as decoder]
-            [cljam.io.util.lsb.data-io :as lsb])
+            [cljam.io.util.lsb.data-io :as lsb]
+            [cljam.io.sam.util.flag :as flag]
+            [cljam.io.util.byte-buffer :as bb])
   (:import [java.io Closeable FileNotFoundException]
            [cljam.io.bam.decoder BAMRawBlock]
            [bgzf4j BGZFInputStream]))
 
 (declare read-blocks-sequentially*
-         read-blocks-randomly*)
+         read-blocks-randomly*
+         read-mate-alignments*)
 
 ;; BAMReader
 ;; ---------
@@ -47,6 +50,10 @@
       (if (nil? chr)
         (read-blocks-sequentially* this 64 decoder)
         (read-blocks-randomly* this chr start end 64 decoder))))
+  (read-mate-alignments [this alignments]
+    (protocols/read-mate-alignments this {} alignments))
+  (read-mate-alignments [this options alignments]
+    (read-mate-alignments* this options alignments))
   (read-blocks [this]
     (protocols/read-blocks this {} {}))
   (read-blocks [this region]
@@ -163,3 +170,75 @@
                                      :len  l-ref})))))]
     {:header header
      :refs refs}))
+
+(defn- mate?-fn
+  "Returns a predicate that checks if an alignment `block` is a mate of an
+  alignment in the given `alignments`."
+  [alignments]
+  (let [query (transduce
+               (map (fn [{:keys [pnext qname flag]}]
+                      [[(flag/r1? flag) pnext (count qname)] qname]))
+               (completing (fn [r [ks v]] (update-in r ks (fnil conj #{}) v)))
+               {}
+               alignments)]
+    (fn mate?
+      [^BAMRawBlock block]
+      (let [b (bb/make-lsb-byte-buffer (.data block))
+            flag (bit-and 0xFFFF (.getShort b 14))]
+        (when (flag/primary? flag)
+          (when-let [q' (query (not (flag/r1? flag)))] ; check R1/R2
+            (when-let [q'' (q' (inc (.getInt b 4)))]   ; check POS
+              (let [l-read-name (dec (bit-and 0xFF (.get b 8)))]
+                (when-let [q''' (q'' l-read-name)]     ; check length of QNAME
+                  (.position b 32)                     ; offset of qname
+                  (q''' (String. ^bytes (bb/read-bytes b l-read-name))))))))))))
+
+(defn- read-mate-blocks-in-chr
+  "Reads mate alignment blocks in a specific `chr`."
+  [^BAMReader rdr {:keys [chunk-size] :or {chunk-size 64}} [chr xs]]
+  (let [refs (.refs rdr)
+        rid (refs/ref-id refs chr)]
+    (if (= chr "*")
+      ;; unmapped
+      (do (.seek ^BGZFInputStream (.reader rdr)
+                 (ffirst (bai/get-unplaced-spans @(.index-delay rdr))))
+          (eduction (filter (mate?-fn xs)) (read-to-finish rdr chunk-size)))
+      ;; mapped
+      (->> xs
+           (sequence
+            (comp
+             (map :pnext)
+             (distinct)
+             (map (fn [p] {:start p, :end p}))))
+           (map (juxt :start :end))
+           (bai/get-spans-for-regions @(.index-delay rdr) rid)
+           (eduction
+            (comp
+             (mapcat
+              (fn [[begin finish]]
+                (read-to-finish rdr begin finish chunk-size)))
+             (filter (mate?-fn xs))))))))
+
+(defn- ref-order
+  "Returns a map of RNAME to its index in the header of BAM file."
+  [refs]
+  (into {}
+        (map-indexed (fn [i m] [(:name m) i]))
+        (concat refs [{:name "*"}])))
+
+(defn- read-mate-alignments*
+  "Implementation for `cljam.io.sam/read-mate-alignments`.
+
+  The following options are availabe:
+
+  - :chunk-size  A size of chunks used for the internal block sequence"
+  ([^BAMReader rdr alignments]
+   (read-mate-alignments* rdr {} alignments))
+  ([^BAMReader rdr options alignments]
+   (->> alignments
+        (group-by (fn [{:keys [rname rnext]}] (case rnext "=" rname rnext)))
+        (sort-by (comp (ref-order (.refs rdr)) key))
+        (eduction
+         (comp
+          (mapcat (partial read-mate-blocks-in-chr rdr options))
+          (map (partial decoder/decode-alignment (.refs rdr))))))))
diff --git a/src/cljam/io/bam_index.clj b/src/cljam/io/bam_index.clj
index 443eaa6..8f2c758 100644
--- a/src/cljam/io/bam_index.clj
+++ b/src/cljam/io/bam_index.clj
@@ -19,6 +19,19 @@
   [bai ref-idx beg end]
   (util-bin/get-spans bai ref-idx beg end))
 
+(defn get-spans-for-regions
+  "Returns a sequence of [start end) pairs of virtual file offsets of a BAM file
+  that may contain alignments overlapping one of the given genomic `regions`
+  which is a sequence of pairs of integers, begin and end."
+  [bai ^long ref-idx regions]
+  (util-bin/get-spans-for-regions bai ref-idx regions))
+
+(defn get-unplaced-spans
+  "Returns a sequence of [start end) pairs of virtual file offsets that may
+  contain alignments that don't have RNAME."
+  [bai]
+  (bai-core/get-unplaced-spans bai))
+
 (defn bam-index
   "Returns a cljam.bam-index.core.BAMIndex."
   [f]
diff --git a/src/cljam/io/protocols.clj b/src/cljam/io/protocols.clj
index eef8724..cbf8d02 100644
--- a/src/cljam/io/protocols.clj
+++ b/src/cljam/io/protocols.clj
@@ -31,6 +31,8 @@
     "Returns references of the SAM/BAM file.")
   (read-alignments [this] [this region]
     "Reads alignments of the SAM/BAM file, returning the alignments as an eduction.")
+  (read-mate-alignments [this alignments] [this option alignments]
+    "Reads mate alignments of given `alignments`.")
   (read-blocks [this] [this region] [this region option]
     "Reads alignment blocks of the SAM/BAM file, returning the blocks as an eduction."))
 
diff --git a/src/cljam/io/sam.clj b/src/cljam/io/sam.clj
index 5517daa..7ed8c1e 100644
--- a/src/cljam/io/sam.clj
+++ b/src/cljam/io/sam.clj
@@ -7,7 +7,8 @@
             [cljam.io.sam.writer :as sam-writer]
             [cljam.io.bam.core :as bam-core]
             [cljam.io.protocols :as protocols]
-            [cljam.io.util :as io-util])
+            [cljam.io.util :as io-util]
+            [cljam.io.sam.util.flag :as flag])
   (:import java.io.Closeable
            cljam.io.sam.reader.SAMReader
            cljam.io.sam.writer.SAMWriter
@@ -69,6 +70,46 @@
   ([rdr] (protocols/read-alignments rdr))
   ([rdr region] (protocols/read-alignments rdr region)))
 
+(defn read-mate-alignments
+  "Reads the mate alignments corresponding to the given `alignments` from the
+  `rdr`. A mate alignment here refers to a primary alignment that has the same
+  QNAME as the given alignment and is the opposite R1/R2 pair. The `rdr` must
+  be indexed. Returns an eduction."
+  ([rdr alignments]
+   (protocols/read-mate-alignments rdr alignments))
+  ([rdr options alignments]
+   (protocols/read-mate-alignments rdr options alignments)))
+
+(defn make-pairs
+  "Pairs mate alignments from the given `alignments` and `rdr`.
+
+  This is a convenient wrapper for `read-mate-alignments`. This function
+  processes the alignments provided in the `alignments` sequence, pairing mate
+  alignments based on their QNAMEs. If a mate alignment is missing from
+  `alignments`, it attempts to read the mate alignment from the reader `rdr`.
+  Non-primary alignments are ignored. The order of the paired alignments in the
+  output is not guaranteed."
+  ([rdr alignments]
+   (make-pairs rdr {} alignments))
+  ([rdr options alignments]
+   (letfn [(has-mate? [{:keys [flag]}]
+             (and (flag/multiple? flag)
+                  (flag/primary? flag)))]
+     (let [{solos true,
+            pairs false} (->> alignments
+                              (group-by :qname)
+                              (keep
+                               (fn [[_ vs]] (not-empty (filterv has-mate? vs))))
+                              (group-by #(= 1 (count %))))
+           solos (mapv first solos)
+           mates (->> solos
+                      (read-mate-alignments rdr options)
+                      (into {} (map (juxt :qname identity))))]
+       (concat pairs
+               (sequence
+                (comp (map (juxt identity (comp mates :qname)))
+                      (filter second)) solos))))))
+
 (defn read-blocks
   "Reads alignment blocks of the SAM/BAM file, returning the blocks as an eduction."
   ([rdr] (protocols/read-blocks rdr))
diff --git a/src/cljam/io/util/bin.clj b/src/cljam/io/util/bin.clj
index 8111465..23363be 100644
--- a/src/cljam/io/util/bin.clj
+++ b/src/cljam/io/util/bin.clj
@@ -1,5 +1,6 @@
 (ns cljam.io.util.bin
-  (:require [cljam.io.util.chunk :as util-chunk]))
+  (:require [cljam.io.util.chunk :as util-chunk]
+            [cljam.util.region :as reg]))
 
 (defn max-pos
   "Returns a maximum position of a binning index. The value is identical to the
@@ -101,3 +102,23 @@
         min-offset (get-min-offset index-data ref-idx beg)]
     (->> (util-chunk/optimize-chunks chunks min-offset)
          (map vals))))
+
+(defn get-spans-for-regions
+  "Calculates span information for multiple regions."
+  [index-data ^long ref-idx regions]
+  (when (seq regions)
+    (let [min-shift (get-min-shift index-data)
+          depth (get-depth index-data)]
+      ;; Merge query regions before computing bins.
+      (->> regions
+           sort
+           (map (fn [[b e]] {:start b :end e}))
+           (reg/merge-regions 1000)
+           (mapcat
+            (fn [{beg :start end :end}]
+              (util-chunk/optimize-chunks
+               (get-chunks index-data ref-idx
+                           (reg->bins beg end min-shift depth))
+               (get-min-offset index-data ref-idx beg))))
+           (#(util-chunk/optimize-chunks % 0))
+           (map vals)))))


Test patch:
diff --git a/test/cljam/io/bam/reader_test.clj b/test/cljam/io/bam/reader_test.clj
new file mode 100644
index 0000000..a218e8c
--- /dev/null
+++ b/test/cljam/io/bam/reader_test.clj
@@ -0,0 +1,101 @@
+(ns cljam.io.bam.reader-test
+  (:require [clojure.java.io :as io]
+            [clojure.test :refer [deftest is testing]]
+            [cljam.util :as util]
+            [cljam.util.chromosome :as chr]
+            [cljam.io.protocols :as prot]
+            [cljam.io.bam.core :as bam.core]
+            [cljam.io.sam.util.flag :as flag]))
+
+(defn- conj-flag [flag & xs]
+  (flag/encode (apply conj (flag/decode flag) xs)))
+
+(defn- unmap [xs]
+  (let [[unmapped mapped] (if (< (double (rand)) 0.5)
+                            [0 1]
+                            [1 0])]
+    (-> xs
+        (update unmapped assoc
+                :rname "*" :pos 0 :end 0
+                :mapq 0 :cigar "" :tlen 0)
+        (update-in [unmapped :flag] conj-flag :unmapped)
+        (update-in [unmapped :rnext] #(if (= "=" %)
+                                        (:rname (first xs))
+                                        %))
+        (update mapped assoc
+                :rnext "*" :pnext 0 :tlen 0)
+        (update-in [mapped :flag] conj-flag :next-unmapped))))
+
+(defn- generate-paired-bam [n]
+  (let [refs (->>
+              (range 24)
+              (map (fn [i] {:SN (str "chr" (inc (long i))),
+                            :LN (+ 100000000 (long (rand-int 50000000)))}))
+              (sort-by (comp chr/chromosome-order-key :SN)))]
+    (->>
+     (range n)
+     (map
+      (fn [i]
+        (let [{:keys [SN LN] :as r} (rand-nth refs)
+              pos (int (rand-int (- (long LN) 300)))
+              chimeric? (< (double (rand)) 0.3)
+              unmapped? (< (double (rand)) 0.2)
+              rnext (if chimeric?
+                      (rand-nth (filterv (complement #{r}) refs))
+                      r)
+              pnext (int (if chimeric?
+                           (rand-int (- (long (:LN rnext)) 300))
+                           (+ pos (int 70) (int (rand-int 200)))))
+              qname (str "read" (inc (long i)))]
+          (cond-> [(prot/->SAMAlignment
+                    qname (flag/encode #{:multiple :first})
+                    SN pos (int (+ pos 150)) 60 "151M"
+                    (if chimeric? (:SN rnext) "=") pnext 0 "*" "*" [])
+                   (prot/->SAMAlignment
+                    qname (flag/encode #{:multiple :last})
+                    (:SN rnext) pnext (int (+ pnext 150)) 60 "151M"
+                    (if chimeric? SN "=") pos 0 "*" "*" [])]
+            unmapped? unmap))))
+     (apply concat)
+     (sort-by (juxt (comp chr/chromosome-order-key :rname) :pos))
+     (vector {:HD {:VN "1.4", :SO "coordinate"}, :SQ refs}))))
+
+(defn- generate-test-case [n alignments]
+  (let [in (->> alignments
+                shuffle
+                (take n)
+                (group-by :qname)
+                vals
+                (map first)) ; take either R1 or R2
+        mate? (into #{} (map (juxt :qname (comp not flag/r1? :flag)) in))]
+    [in (filter (comp mate? (juxt :qname (comp flag/r1? :flag))) alignments)]))
+
+(deftest read-mate-alignments-test
+  ;; 1. Generates a sequence of paired alignments with random positions.
+  ;; 2. Write the alignments to a BAM file with an index.
+  ;; 3. Randomly pick some subsets of the alignments as queries.
+  ;; 4. Compute the mate alignments for each query using (1).
+  ;; 5. Read the mate alignments from (2) and compare with (4).
+  (util/with-temp-dir [d "read-pairs-test"]
+    (let [temp-bam (io/file d "tmp.bam")
+          n-total-alignments 10000
+          n-query-alignments 50
+          n-tests 5
+          [header alignments] (generate-paired-bam n-total-alignments)]
+
+      ;; write a BAM with a BAI
+      (with-open [w (bam.core/writer temp-bam true)]
+        (doto w
+          (prot/write-header header)
+          (prot/write-refs header)
+          (prot/write-alignments alignments header)))
+
+      ;; read and compare
+      (with-open [r (bam.core/reader temp-bam)]
+        (doseq [[i [input expected]] (->> #(generate-test-case
+                                            n-query-alignments alignments)
+                                          (repeatedly n-tests)
+                                          (map vector (range)))]
+          (testing (inc (long i))
+            (is (= expected
+                   (prot/read-mate-alignments r input)))))))))
diff --git a/test/cljam/io/sam_test.clj b/test/cljam/io/sam_test.clj
index cca8ccf..29f8419 100644
--- a/test/cljam/io/sam_test.clj
+++ b/test/cljam/io/sam_test.clj
@@ -468,3 +468,14 @@
       temp-bam-file
       (cio/file temp-bam-file)
       (cio/as-url (cio/file temp-bam-file)))))
+
+(deftest make-pairs-test
+  (testing "BAM"
+    (with-open [r (sam/reader test-sorted-bam-file)]
+      (let [[r001 r002 r003 r004 r003' r001'] (vec (sam/read-alignments r))]
+        (is (= [[r001 r001']]
+               (sam/make-pairs r [r001])))
+        (is (= [[r001 r001']]
+               (sam/make-pairs r [r001 r001'])))
+        (is (= [[r001 r001']] ;; r003 is not flagged as paired
+               (sam/make-pairs r [r001  r002 r003 r003' r004])))))))
diff --git a/test/cljam/io/util/bin_test.clj b/test/cljam/io/util/bin_test.clj
index dff2ff4..8783aad 100644
--- a/test/cljam/io/util/bin_test.clj
+++ b/test/cljam/io/util/bin_test.clj
@@ -197,3 +197,16 @@
                       (number? (first %))
                       (number? (second %)))
                 (util-bin/get-spans csi* 0 1 100000)))))
+
+(deftest get-spans-for-regions-test
+  (testing "tabix"
+    (let [tabix (tabix/read-index test-tabix-file)]
+      (is (= [[0 50872]]
+             (util-bin/get-spans-for-regions tabix 0 [[1 100]])))
+      (is (= [[0 50872]]
+             (util-bin/get-spans-for-regions
+              tabix 0 [[1 10] [10 20] [30 40] [40 50] [50 100]])))))
+  (testing "csi"
+    (let [csi (csi/read-index test-csi-file)]
+      (is (= [[3904 3973]]
+             (util-bin/get-spans-for-regions csi 0 [[1 100000]]))))))