From ece53c64706b6a84fab0beeabc19e07de4656f70 Mon Sep 17 00:00:00 2001
From: Jayaram Kancherla <jayaram.kancherla@gmail.com>
Date: Wed, 24 Jan 2024 09:36:03 -0800
Subject: [PATCH] Compute pair-wise distances (#70)

---
 src/genomicranges/GenomicRanges.py | 22 ++++++++++++++++++++++
 tests/test_gr_methods_search.py    | 11 +++++++++++
 2 files changed, 33 insertions(+)

diff --git a/src/genomicranges/GenomicRanges.py b/src/genomicranges/GenomicRanges.py
index fc0496e..6627da0 100644
--- a/src/genomicranges/GenomicRanges.py
+++ b/src/genomicranges/GenomicRanges.py
@@ -2237,6 +2237,28 @@ def follow(
 
         return rev_map
 
+    def distance(self, query: Union["GenomicRanges", IRanges]) -> np.ndarray:
+        """Compute the pair-wise distance with intervals in query.
+
+        Args:
+            query:
+                Query `GenomicRanges` or `IRanges`.
+
+        Returns:
+            Numpy vector containing distances for each interval in query.
+        """
+        if not isinstance(query, (IRanges, GenomicRanges)):
+            raise TypeError("'query' is not a `GenomicRanges` or `IRanges` object.")
+
+        if len(self) != len(query):
+            raise ValueError("'query' does not contain the same number of intervals.")
+
+        _qranges = query
+        if isinstance(query, GenomicRanges):
+            _qranges = query.get_ranges()
+
+        return self._ranges.distance(_qranges)
+
     def match(self, query: "GenomicRanges") -> List[List[int]]:
         """Element wise comparison to find exact match ranges.
 
diff --git a/tests/test_gr_methods_search.py b/tests/test_gr_methods_search.py
index 0e1452f..8cc33f8 100644
--- a/tests/test_gr_methods_search.py
+++ b/tests/test_gr_methods_search.py
@@ -69,3 +69,14 @@ def test_follow():
 
     assert query_hits is not None
     assert query_hits == [[4], [], []]
+
+
+def test_distance():
+    query = GenomicRanges(["A", "A", "A"], ranges=IRanges([1, 2, 10], [5, 7, 2]))
+    subject = GenomicRanges(["A", "A", "A"], ranges=IRanges([6, 5, 13], [5, 6, 3]))
+
+    distance = subject.distance(query)
+
+    assert distance is not None
+    assert isinstance(distance, np.ndarray)
+    assert distance.tolist() == [0, 0, 1]