optimze compute_diagonal for k=1

TDAmeritrade · Nov 6, 2022 · c9a2f8d · c9a2f8d
1 parent ea29eaf
commit c9a2f8d
Showing 1 changed file with 32 additions and 17 deletions.
diff --git a/stumpy/stump.py b/stumpy/stump.py
@@ -205,24 +205,32 @@ def _compute_diagonal(
                 # 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, :k], pearson)
-                    core._shift_insert_at_index(
-                        ρ[thread_idx, i, :k], idx, pearson, shift="left"
-                    )
-                    core._shift_insert_at_index(
-                        I[thread_idx, i, :k], idx, i + g, shift="left"
-                    )
-
-                if ignore_trivial:  # self-joins only
-                    if pearson > ρ[thread_idx, i + g, 0]:
-                        idx = np.searchsorted(ρ[thread_idx, i + g, :k], pearson)
+                    if k == 1:
+                        ρ[thread_idx, i, 0] = pearson
+                        I[thread_idx, i, 0] = i + g
+                    else:
+                        idx = np.searchsorted(ρ[thread_idx, i, :k], pearson)
                         core._shift_insert_at_index(
-                            ρ[thread_idx, i + g, :k], idx, pearson, shift="left"
+                            ρ[thread_idx, i, :k], idx, pearson, shift="left"
                         )
                         core._shift_insert_at_index(
-                            I[thread_idx, i + g, :k], idx, i, shift="left"
+                            I[thread_idx, i, :k], idx, i + g, shift="left"
                         )
 
+                if ignore_trivial:  # self-joins only
+                    if pearson > ρ[thread_idx, i + g, 0]:
+                        if k == 1:
+                            ρ[thread_idx, i + g, 0] = pearson
+                            I[thread_idx, i + g, 0] = i
+                        else:
+                            idx = np.searchsorted(ρ[thread_idx, i + g, :k], pearson)
+                            core._shift_insert_at_index(
+                                ρ[thread_idx, i + g, :k], idx, pearson, shift="left"
+                            )
+                            core._shift_insert_at_index(
+                                I[thread_idx, i + g, :k], idx, i, shift="left"
+                            )
+
                     if i < i + g:
                         # left pearson correlation and left matrix profile index
                         if pearson > ρ[thread_idx, i + g, k]:
@@ -610,12 +618,19 @@ def _stump(
     _I = I[0, :, :k][:, ::-1]
 
     PL = core._transform_pearson_to_distance(ρ[0, :, k], m)
-    IL = I[0, :, k]
-
     PR = core._transform_pearson_to_distance(ρ[0, :, k + 1], m)
-    IR = I[0, :, k + 1]
 
-    return P, PL, PR, _I, IL, IR
+    return P, PL, PR, _I, I[0, :, k], I[0, :, k + 1]
+    # P = core._transform_pearson_to_distance(ρ[0], m)
+
+    # return (
+    #    P[:, :k][:, ::-1],  # main
+    #    P[:, k],  # left
+    #    P[:, k + 1],  # right
+    #    I[0, :, :k][:, ::-1],  # main
+    #    I[0, :, k],  # left
+    #    I[0, :, k + 1],  # right
+    # )
 
 
 @core.non_normalized(aamp)