Cythonize get_coinc_indices

Author: spxiwh

pycbc/events/coherent.py

@@ -27,8 +27,20 @@
 import logging
 import numpy as np
 
+from .eventmgr_cython import get_coinc_indexes_cython_twodet_twocoinc
+
 logger = logging.getLogger('pycbc.events.coherent')
 
+def get_coinc_indexes_twodet_twocoinc(idxarr1, idxarr2, offset1, offset2, output):
+    num_idxs = get_coinc_indexes_cython_twodet_twocoinc(
+        idxarr1,
+        idxarr2,
+        offset1,
+        offset2,
+        output
+    )
+    return num_idxs
+
 
 def get_coinc_indexes(idx_dict, time_delay_idx, min_nifos):
     """Return the indexes corresponding to coincident triggers. If only one
@@ -53,6 +65,20 @@ def get_coinc_indexes(idx_dict, time_delay_idx, min_nifos):
         List of indexes for triggers in geocent time that appear in
         multiple detectors
     """
+    if min_nifos == 2 and len(idx_dict) == 2:
+        ifos = list(idx_dict.keys())
+        # Could cache an output array if needed
+        idxarr1 = idx_dict[ifos[0]]
+        idxarr2 = idx_dict[ifos[1]]
+        outarr = np.zeros(max(len(idxarr1), len(idxarr2)), dtype=idxarr1.dtype)
+        num_idxs = get_coinc_indexes_twodet_twocoinc(
+            idxarr1,
+            idxarr2,
+            time_delay_idx[ifos[0]],
+            time_delay_idx[ifos[1]],
+            outarr
+        )
+        return outarr[:num_idxs]
     coinc_list = np.array([], dtype=int)
     for ifo in idx_dict.keys():
         # Create list of indexes above single detector threshold, in geocent

pycbc/events/eventmgr_cython.pyx

@@ -1,5 +1,6 @@
 import numpy as np
 cimport numpy as cnp
+import cython
 from cython import wraparound, boundscheck, cdivision
 from libc.math cimport M_PI, sqrt
 from libc.math cimport round as cround
@@ -306,3 +307,39 @@ def timecluster_cython(
         elif max_loc < i:
             i += 1
     return j
+
+ctypedef fused numeric_type:
+    cython.int
+    cython.long
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def get_coinc_indexes_cython_twodet_twocoinc(
+    numeric_type [::1] idxarr1,
+    numeric_type [::1] idxarr2,
+    numeric_type offset1,
+    numeric_type offset2,
+    numeric_type [::1] outputs
+):
+    cdef int arr1pos = 0
+    cdef int arr2pos = 0
+    cdef int outpos = 0
+    cdef int arr1_size = idxarr1.size
+    cdef int arr2_size = idxarr2.size
+    cdef int cpos1
+    cdef int cpos2
+
+    while arr1pos < arr1_size and arr2pos < arr2_size:
+        cpos1 = idxarr1[arr1pos] - offset1
+        cpos2 = idxarr2[arr2pos]-offset2
+        if (cpos1) == (cpos2):
+            outputs[outpos] = cpos1
+            outpos += 1
+            arr1pos += 1
+            arr2pos += 1
+        elif cpos1 < cpos2:
+            arr1pos += 1
+        else:
+            arr2pos += 1
+    return outpos
+