From 8d1ff642f008a6da86a37ef621221756f43231fb Mon Sep 17 00:00:00 2001 From: dachengx Date: Tue, 3 Dec 2024 08:17:30 -0600 Subject: [PATCH] Accelerate Euclidean distance by numba --- .../peaklets/peaklet_classification_som.py | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/straxen/plugins/peaklets/peaklet_classification_som.py b/straxen/plugins/peaklets/peaklet_classification_som.py index 8f761c83e..d4d0ef08a 100644 --- a/straxen/plugins/peaklets/peaklet_classification_som.py +++ b/straxen/plugins/peaklets/peaklet_classification_som.py @@ -1,6 +1,5 @@ import numpy as np import numpy.lib.recfunctions as rfn -from scipy.spatial.distance import cdist from straxen.plugins.peaklets.peaklet_classification_vanilla import PeakletClassificationVanilla import numba @@ -146,12 +145,27 @@ def generate_color_ref_map(color_image, unique_colors, xdim, ydim): return ref_map +@numba.njit +def euclidean_dist(XA, XB): + # mimicking scipy.spatial.distance.cdist when metric='euclidean' + nA, dA = XA.shape + nB, dB = XB.shape + assert dA == dB, "Dimensions of points in XA and XB must match." + distances = np.empty((nA, nB)) + for i in range(nA): + for j in range(nB): + dist = 0.0 + for k in range(dA): + diff = XA[i, k] - XB[j, k] + dist += diff * diff + distances[i, j] = np.sqrt(dist) + return distances + + def som_cls_recall(array_to_fill, data_in_som_fmt, weight_cube, reference_map): som_xdim, som_ydim, _ = weight_cube.shape # for data_point in data_in_SOM_fmt: - distances = cdist( - weight_cube.reshape(-1, weight_cube.shape[-1]), data_in_som_fmt, metric="euclidean" - ) + distances = euclidean_dist(weight_cube.reshape(-1, weight_cube.shape[-1]), data_in_som_fmt) w_neuron = np.argmin(distances, axis=0) x_idx, y_idx = np.unravel_index(w_neuron, (som_xdim, som_ydim)) array_to_fill["som_sub_type"] = reference_map[x_idx, y_idx]