[WIP] Add Top-K Nearest Neighbors for Normalized Matrix Profile #592

stumpy/stump.py

@@ -185,10 +185,14 @@ def _compute_diagonal(
             iter_range = range(-g, min(n_A - m + 1, n_B - m + 1 - g))
 
         for i in iter_range:
-            if i == 0 or (g < 0 and i == -g):
+            uint64_i = np.uint64(i)
+            uint64_j = np.uint64(i + g)
+
+            if uint64_i == 0 or uint64_j == 0:
                 cov = (
                     np.dot(
-                        (T_B[i + g : i + g + m] - M_T[i + g]), (T_A[i : i + m] - μ_Q[i])
+                        (T_B[uint64_j : uint64_j + uint64_m] - M_T[uint64_j]),
+                        (T_A[uint64_i : uint64_i + uint64_m] - μ_Q[uint64_i]),
                     )
                     * m_inverse
                 )
@@ -200,52 +204,53 @@ def _compute_diagonal(
                 #     - (T_B[i + k - 1] - M_T_m_1[i + k]) * (T_A[i - 1] - μ_Q_m_1[i])
                 # )
                 cov = cov + constant * (
-                    cov_a[i + g] * cov_b[i] - cov_c[i + g] * cov_d[i]
+                    cov_a[uint64_j] * cov_b[uint64_i]
+                    - cov_c[uint64_j] * cov_d[uint64_i]
                 )
 
-            if T_B_subseq_isfinite[i + g] and T_A_subseq_isfinite[i]:
+            if T_B_subseq_isfinite[uint64_j] and T_A_subseq_isfinite[uint64_i]:
                 # Neither subsequence contains NaNs
-                if T_B_subseq_isconstant[i + g] or T_A_subseq_isconstant[i]:
+                if T_B_subseq_isconstant[uint64_j] or T_A_subseq_isconstant[uint64_i]:
                     pearson = 0.5
                 else:
-                    pearson = cov * Σ_T_inverse[i + g] * σ_Q_inverse[i]
+                    pearson = cov * Σ_T_inverse[uint64_j] * σ_Q_inverse[uint64_i]
 
-                if T_B_subseq_isconstant[i + g] and T_A_subseq_isconstant[i]:
+                if T_B_subseq_isconstant[uint64_j] and T_A_subseq_isconstant[uint64_i]:
                     pearson = 1.0
 
                 # `ρ[thread_idx, i, :]` is sorted ascendingly and MUST be updated
                 # when the newly-calculated `pearson` value becomes greater than the
                 # first (i.e. smallest) element in this array. Note that a higher
                 # pearson value corresponds to a lower distance.
                 if pearson > ρ[thread_idx, i, 0]:
-                    idx = np.searchsorted(ρ[thread_idx, i], pearson)
+                    idx = np.searchsorted(ρ[thread_idx, uint64_i], pearson)
                     core._shift_insert_at_index(
-                        ρ[thread_idx, i], idx, pearson, shift="left"
+                        ρ[thread_idx, uint64_i], idx, pearson, shift="left"
                     )
                     core._shift_insert_at_index(
-                        I[thread_idx, i], idx, i + g, shift="left"
+                        I[thread_idx, uint64_i], idx, uint64_j, shift="left"
                     )
 
                 if ignore_trivial:  # self-joins only
-                    if pearson > ρ[thread_idx, i + g, 0]:
-                        idx = np.searchsorted(ρ[thread_idx, i + g], pearson)
+                    if pearson > ρ[thread_idx, uint64_j, 0]:
+                        idx = np.searchsorted(ρ[thread_idx, uint64_j], pearson)
                         core._shift_insert_at_index(
-                            ρ[thread_idx, i + g], idx, pearson, shift="left"
+                            ρ[thread_idx, uint64_j], idx, pearson, shift="left"
                         )
                         core._shift_insert_at_index(
-                            I[thread_idx, i + g], idx, i, shift="left"
+                            I[thread_idx, uint64_j], idx, uint64_i, shift="left"
                         )
 
-                    if i < i + g:
+                    if uint64_i < uint64_j:
                         # left pearson correlation and left matrix profile index
-                        if pearson > ρL[thread_idx, i + g]:
-                            ρL[thread_idx, i + g] = pearson
-                            IL[thread_idx, i + g] = i
+                        if pearson > ρL[thread_idx, uint64_j]:
+                            ρL[thread_idx, uint64_j] = pearson
+                            IL[thread_idx, uint64_j] = uint64_i
 
                         # right pearson correlation and right matrix profile index
-                        if pearson > ρR[thread_idx, i]:
-                            ρR[thread_idx, i] = pearson
-                            IR[thread_idx, i] = i + g
+                        if pearson > ρR[thread_idx, uint64_i]:
+                            ρR[thread_idx, uint64_i] = pearson
+                            IR[thread_idx, uint64_i] = uint64_j
 
     return
 

tests/naive.py

@@ -1845,25 +1845,30 @@ def merge_topk_PI(PA, PB, IA, IB):
                 IA[i] = IB[i]
         return
 
-    k = PA.shape[1]
-    for i in range(PA.shape[0]):
-        _, _, overlap_idx_B = np.intersect1d(IA[i], IB[i], return_indices=True)
-        PB[i, overlap_idx_B] = np.inf
-        IB[i, overlap_idx_B] = -1
+    else:
+        k = PA.shape[1]
+        for i in range(PA.shape[0]):
+            _, _, overlap_idx_B = np.intersect1d(IA[i], IB[i], return_indices=True)
+            PB[i, overlap_idx_B] = np.inf
+            IB[i, overlap_idx_B] = -1
 
-    profile = np.column_stack((PA, PB))
-    indices = np.column_stack((IA, IB))
-    IDX = np.argsort(profile, axis=1, kind="mergesort")
-    profile[:, :] = np.take_along_axis(profile, IDX, axis=1)
-    indices[:, :] = np.take_along_axis(indices, IDX, axis=1)
+        profile = np.column_stack((PA, PB))
+        indices = np.column_stack((IA, IB))
+        IDX = np.argsort(profile, axis=1, kind="mergesort")
+        profile[:, :] = np.take_along_axis(profile, IDX, axis=1)
+        indices[:, :] = np.take_along_axis(indices, IDX, axis=1)
 
-    PA[:, :] = profile[:, :k]
-    IA[:, :] = indices[:, :k]
+        PA[:, :] = profile[:, :k]
+        IA[:, :] = indices[:, :k]
+
+        return
 
 
 def merge_topk_ρI(ρA, ρB, IA, IB):
-    # this is to merge two pearson profiles `ρA` and `ρB`, where each is a 2D array
-    # and each row is sorted ascendingly. we want to keep top-k largest values in
+    # This function merges two pearson profiles `ρA` and `ρB`, and updates `ρA`
+    # and `IA` accordingly. When the inputs are 1D, `ρA[i]` is updated if
+    #  `ρA[i] < ρB[i]` and IA[i] != IB[i]. When the inputs are 2D, each row in
+    #  `ρA` and `ρB` is sorted ascendingly. we want to keep top-k largest values in
     # merging row `ρA[i]` and `ρB[i]`.
 
     # In case of ties between `ρA` and `ρB`, the priority is with `ρA`. In case
@@ -1879,31 +1884,34 @@ def merge_topk_ρI(ρA, ρB, IA, IB):
 
     # For the same example:
     # merging `ρB` and `ρA` ascendingly while choosing `ρB` over `ρA` in case of
-    # ties: [0_B, 0_A, 0'_A, 1_B, 1'_B, 1_A], and we just need to keep the second
-    # half of this array, and discard the first half.
+    # ties: [0_B, 0_A, 0'_A, 1_B, 1'_B, 1_A], and the second half of this array
+    # is the desribale outcome.
     if ρA.ndim == 1:
         for i in range(ρA.shape[0]):
             if ρB[i] > ρA[i] and IB[i] != IA[i]:
                 ρA[i] = ρB[i]
                 IA[i] = IB[i]
         return
 
-    k = ρA.shape[1]
-    for i in range(ρA.shape[0]):
-        _, _, overlap_idx_B = np.intersect1d(IA[i], IB[i], return_indices=True)
-        ρB[i, overlap_idx_B] = np.NINF
-        IB[i, overlap_idx_B] = -1
+    else:
+        k = ρA.shape[1]
+        for i in range(ρA.shape[0]):
+            _, _, overlap_idx_B = np.intersect1d(IA[i], IB[i], return_indices=True)
+            ρB[i, overlap_idx_B] = np.NINF
+            IB[i, overlap_idx_B] = -1
 
-    profile = np.column_stack((ρB, ρA))
-    indices = np.column_stack((IB, IA))
+        profile = np.column_stack((ρB, ρA))
+        indices = np.column_stack((IB, IA))
 
-    idx = np.argsort(profile, axis=1, kind="mergesort")
-    profile[:, :] = np.take_along_axis(profile, idx, axis=1)
-    indices[:, :] = np.take_along_axis(indices, idx, axis=1)
+        idx = np.argsort(profile, axis=1, kind="mergesort")
+        profile[:, :] = np.take_along_axis(profile, idx, axis=1)
+        indices[:, :] = np.take_along_axis(indices, idx, axis=1)
 
-    # keep the last k elements (top-k largest values)
-    ρA[:, :] = profile[:, k:]
-    IA[:, :] = indices[:, k:]
+        # keep the last k elements (top-k largest values)
+        ρA[:, :] = profile[:, k:]
+        IA[:, :] = indices[:, k:]
+
+        return
 
 
 def find_matches(D, excl_zone, max_distance, max_matches=None):