Sync with iranges's nearest and parallelization methods (#158)

Author: jkanche

CHANGELOG.md

@@ -1,9 +1,9 @@
 # Changelog
 
-## Version 0.7.0 - 0.7.2
+## Version 0.7.0 - 0.7.3
 
 - Changes to switch to LTLA/nclist-cpp in the iranges package for overlap and search operations.
-- Improve performance of search operations, bump version of iranges to 0.5.2.
+- Improve performance of search operations, bump version of iranges to 0.5.4.
 - Optimize group by operations that `GenomicRanges` uses internally for the inter-range operations.
 - Reset cached indexes when ranges were modified.
 

README.md

@@ -272,10 +272,10 @@ Performance comparison between Python and R GenomicRanges implementations. The q
 
 | Operation | Python/GenomicRanges | Python/GenomicRanges (5 threads) | R/GenomicRanges |
 |-----------|---------------------|-----------------------------------|-----------------|
-| Overlap | 3.02s | 2.13s | 4.40s |
-| Overlap (single chromosome) | 6.98s | 5.36s | 10.06s |
-| Nearest | 50.1s | 32.3s | 42.16s |
-| Nearest (single chromosome) | 15.5s | 11.4s | 11.01s |
+| Overlap | 2.80s | 2.06s | 4.40s |
+| Overlap (single chromosome) | 6.73s | 5.19s | 10.06s |
+| Nearest | 2.27s | 1.5s | 42.16s |
+| Nearest (single chromosome) | 4.7s | 4.67s | 11.01s |
 
 > [!NOTE]
 > The single chromosome benchmark ignores chromosome/sequence information and performs overlap operations solely on intervals.

setup.cfg

@@ -50,7 +50,7 @@ python_requires = >=3.9
 install_requires =
     importlib-metadata; python_version<"3.8"
     biocframe>=0.6.2
-    iranges>=0.5.3
+    iranges>=0.5.4
     biocutils>=0.2.1
     numpy
 

src/genomicranges/GenomicRanges.py

@@ -15,7 +15,6 @@
     group_by_indices,
     sanitize_strand_vector,
     wrapper_follow_precede,
-    wrapper_nearest,
 )
 
 __author__ = "jkanche"
@@ -2279,10 +2278,10 @@ def find_overlaps(
 
             all_s_hits, all_q_hits = find_overlaps_groups(
                 self._ranges.get_start().astype(np.int32),
-                self._ranges.get_end().astype(np.int32) + 1,
+                self._ranges.get_end_exclusive().astype(np.int32),
                 [s.astype(np.int32) for s in self_groups],
                 query._ranges.get_start().astype(np.int32),
-                query._ranges.get_end().astype(np.int32) + 1,
+                query._ranges.get_end_exclusive().astype(np.int32),
                 [q.astype(np.int32) for q in query_groups],
                 query_type,
                 select,
@@ -2309,10 +2308,10 @@ def find_overlaps(
 
             all_q_hits, all_s_hits = find_overlaps_groups(
                 query._ranges.get_start().astype(np.int32),
-                query._ranges.get_end().astype(np.int32) + 1,
+                query._ranges.get_end_exclusive().astype(np.int32),
                 [q.astype(np.int32) for q in query_groups],
                 self._ranges.get_start().astype(np.int32),
-                self._ranges.get_end().astype(np.int32) + 1,
+                self._ranges.get_end_exclusive().astype(np.int32),
                 [s.astype(np.int32) for s in self_groups],
                 query_type,
                 select,
@@ -2461,6 +2460,7 @@ def nearest(
         select: Literal["all", "arbitrary"] = "arbitrary",
         ignore_strand: bool = False,
         num_threads: int = 1,
+        adjacent_equals_overlap: bool = True,
     ) -> Union[np.ndarray, BiocFrame]:
         """Search nearest positions both upstream and downstream that overlap with each range in ``query``.
 
@@ -2479,6 +2479,16 @@ def nearest(
                 Number of threads to use.
                 Defaults to 1.
 
+            adjacent_equals_overlap:
+                Whether to consider immediately-adjacent subject intervals to be
+                equally "nearest" to the query as an overlapping subject interval.
+
+                If true, both overlapping and immediately-adjacent subject intervals
+                (i.e., a gap of zero) will be reported in matches.
+
+                Otherwise, immediately-adjacent subjects will only be reported if
+                overlapping subjects are not present.
+
         Returns:
             If select="arbitrary":
                 A numpy array of integers with length matching query, containing indices
@@ -2493,36 +2503,37 @@ def nearest(
         if not isinstance(query, GenomicRanges):
             raise TypeError("'query' is not a `GenomicRanges` object.")
 
-        effective_threads = min(num_threads, cpu_count()) if num_threads > 0 else cpu_count()
+        from iranges.lib_iranges import nearest_groups
 
         self_groups, query_groups = self._get_query_common_groups(query)
 
-        tasks = [
-            (
-                s_group,
-                q_group,
-                self._ranges,
-                query._ranges,
-                self._strand,
-                query._strand,
-                ignore_strand,
-            )
-            for s_group, q_group in zip(self_groups, query_groups)
-        ]
+        all_q_hits, all_s_hits = nearest_groups(
+            self._ranges.get_start().astype(np.int32),
+            self._ranges.get_end_exclusive().astype(np.int32),
+            [s.astype(np.int32) for s in self_groups],
+            query._ranges.get_start().astype(np.int32),
+            query._ranges.get_end_exclusive().astype(np.int32),
+            [q.astype(np.int32) for q in query_groups],
+            select,
+            num_threads,
+            adjacent_equals_overlap,
+        )
 
-        if effective_threads == 1 or len(self_groups) <= 1:
-            results = [wrapper_nearest(task) for task in tasks]
-        else:
-            with Pool(processes=effective_threads) as pool:
-                results = pool.map(wrapper_nearest, tasks)
+        if ignore_strand is False:
+            s_strands = self._strand[all_s_hits]
+            q_strands = query._strand[all_q_hits]
 
-        if results:
-            all_qhits_list, all_shits_list = zip(*results)
-        else:
-            all_qhits_list, all_shits_list = [], []
+            mask = s_strands == q_strands
+            # to allow '*' with any strand from query
+            mask[s_strands == 0] = True
+            mask[q_strands == 0] = True
+            all_q_hits = all_q_hits[mask]
+            all_s_hits = all_s_hits[mask]
 
-        final_qhits = np.concatenate(all_qhits_list) if all_qhits_list else np.array([], dtype=np.int32)
-        final_shits = np.concatenate(all_shits_list) if all_shits_list else np.array([], dtype=np.int32)
+        order = np.argsort(all_q_hits, stable=True)
+
+        final_qhits = all_q_hits[order]
+        final_shits = all_s_hits[order]
 
         if select == "arbitrary":
             ret_result = np.full(len(query), None)

src/genomicranges/utils.py

@@ -228,37 +228,6 @@ def extract_groups_from_granges(x, ignore_strand=False):
         return groups
 
 
-def wrapper_nearest(args):
-    """Processes a single pair of self and query groups to find the nearest ranges.
-    This function is designed to be called by a multiprocessing pool.
-    """
-    (
-        s_group,
-        q_group,
-        self_ranges,
-        query_ranges,
-        self_strand,
-        query_strand,
-        ignore_strand,
-    ) = args
-
-    res_idx = self_ranges[s_group].nearest(query=query_ranges[q_group], select="all")
-
-    _q_hits = np.asarray([q_group[j] for j in res_idx.get_column("query_hits")])
-    _s_hits = np.asarray([s_group[x] for x in res_idx.get_column("self_hits")])
-
-    if not ignore_strand:
-        s_strands = self_strand[s_group][res_idx.get_column("self_hits")]
-        q_strands = query_strand[q_group][res_idx.get_column("query_hits")]
-
-        mask = (s_strands == q_strands) | (s_strands == 0) | (q_strands == 0)
-
-        _q_hits = _q_hits[mask]
-        _s_hits = _s_hits[mask]
-
-    return _q_hits, _s_hits
-
-
 def wrapper_follow_precede(args):
     """Processes a single group for precede/follow operations.
     This function is designed to be called by a multiprocessing pool.

tests/test_gr_search.py

@@ -35,7 +35,7 @@ def test_nearest():
     query_hits = gr.nearest(q_gr)
 
     assert query_hits is not None
-    assert np.all(query_hits == [3, 3])
+    assert np.all(query_hits == [1, 1])  # R returns [3,3], select is arbitrary so its ok
 
     query_hits = q_gr.nearest(gr)
     assert np.all(
@@ -47,7 +47,7 @@ def test_nearest():
     assert np.all(query_hits.get_column("self_hits") == [1, 2, 3, 1, 2, 3])
 
     query_hits = gr.nearest(q_gr, ignore_strand=True)
-    assert np.all(query_hits == [3, 3])
+    assert np.all(query_hits == [1, 1])  # R returns [3,3], select is arbitrary so its ok
 
     query_hits = gr.nearest(q_gr, select="all", ignore_strand=True)
     assert np.all(query_hits.get_column("query_hits") == [0, 0, 0, 1, 1, 1])