custom pairwise threaded

Author: chrisholder

aeon/distances/_distance.py

@@ -1,10 +1,11 @@
 __maintainer__ = []
 
-import warnings
 from enum import Enum
 from typing import Any, Callable, Optional, TypedDict, Union
 
 import numpy as np
+from joblib import Parallel, delayed
+from numba import set_num_threads
 from typing_extensions import Unpack
 
 from aeon.distances._mpdist import mp_distance, mp_pairwise_distance
@@ -84,6 +85,7 @@
     squared_pairwise_distance,
 )
 from aeon.utils.conversion._convert_collection import _convert_collection_to_numba_list
+from aeon.utils.validation import check_n_jobs
 from aeon.utils.validation.collection import _is_numpy_list_multivariate
 
 
@@ -264,40 +266,54 @@ def _custom_func_pairwise(
     n_jobs: int = 1,
     **kwargs: Unpack[DistanceKwargs],
 ) -> np.ndarray:
+    n_jobs = check_n_jobs(n_jobs)
+    set_num_threads(n_jobs)
     if dist_func is None:
         raise ValueError("dist_func must be a callable")
 
-    if n_jobs != 1:
-        warnings.warn(
-            "You are using a custom distance function with n_jobs > 1. "
-            "Aeon does not support parallelization for custom distance "
-            "functions. If it is an existing aeon distance try using the "
-            "string name instead.",
-            UserWarning,
-            stacklevel=2,
-        )
-
     multivariate_conversion = _is_numpy_list_multivariate(X, y)
     X, _ = _convert_collection_to_numba_list(X, "X", multivariate_conversion)
+
+    if n_jobs > 1:
+        X = np.array(X)
+
     if y is None:
         # To self
-        return _custom_pairwise_distance(X, dist_func, **kwargs)
+        return _custom_pairwise_distance(X, dist_func, n_jobs=n_jobs, **kwargs)
     y, _ = _convert_collection_to_numba_list(y, "y", multivariate_conversion)
-    return _custom_from_multiple_to_multiple_distance(X, y, dist_func, **kwargs)
+    if n_jobs > 1:
+        y = np.array(y)
+    return _custom_from_multiple_to_multiple_distance(
+        X, y, dist_func, n_jobs=n_jobs, **kwargs
+    )
 
 
 def _custom_pairwise_distance(
     X: Union[np.ndarray, list[np.ndarray]],
     dist_func: DistanceFunction,
+    n_jobs: int = 1,
     **kwargs: Unpack[DistanceKwargs],
 ) -> np.ndarray:
     n_cases = len(X)
     distances = np.zeros((n_cases, n_cases))
 
-    for i in range(n_cases):
-        for j in range(i + 1, n_cases):
+    def compute_single_distance(i, j):
+        return i, j, dist_func(X[i], X[j], **kwargs)
+
+    indices = [(i, j) for i in range(n_cases) for j in range(i + 1, n_cases)]
+
+    if n_jobs == 1:
+        for i, j in indices:
             distances[i, j] = dist_func(X[i], X[j], **kwargs)
-            distances[j, i] = distances[i, j]
+            distances[j, i] = distances[i, j]  # Mirror for symmetry
+    else:
+        results = Parallel(n_jobs=n_jobs)(
+            delayed(compute_single_distance)(i, j) for i, j in indices
+        )
+
+        for i, j, dist in results:
+            distances[i, j] = dist
+            distances[j, i] = dist  # Mirror for symmetry
 
     return distances
 
@@ -306,15 +322,29 @@ def _custom_from_multiple_to_multiple_distance(
     x: Union[np.ndarray, list[np.ndarray]],
     y: Union[np.ndarray, list[np.ndarray]],
     dist_func: DistanceFunction,
+    n_jobs: int = 1,
     **kwargs: Unpack[DistanceKwargs],
 ) -> np.ndarray:
     n_cases = len(x)
     m_cases = len(y)
     distances = np.zeros((n_cases, m_cases))
 
-    for i in range(n_cases):
-        for j in range(m_cases):
+    def compute_single_distance(i, j):
+        return i, j, dist_func(x[i], y[j], **kwargs)
+
+    indices = [(i, j) for i in range(n_cases) for j in range(m_cases)]
+
+    if n_jobs == 1:
+        for i, j in indices:
             distances[i, j] = dist_func(x[i], y[j], **kwargs)
+    else:
+        results = Parallel(n_jobs=n_jobs)(
+            delayed(compute_single_distance)(i, j) for i, j in indices
+        )
+
+        for i, j, dist in results:
+            distances[i, j] = dist
+
     return distances
 
 

aeon/utils/numba/general.py

@@ -772,3 +772,29 @@ def get_all_subsequences(X: np.ndarray, length: int, dilation: int) -> np.ndarra
     out_shape = (n_timestamps - (length - 1) * dilation, n_features, np.int64(length))
     strides = (s1, s0, s1 * dilation)
     return np.lib.stride_tricks.as_strided(X, shape=out_shape, strides=strides)
+
+
+def is_numba_function(func) -> bool:
+    """Determine if a function is compiled with Numba.
+
+    Parameters
+    ----------
+    func : callable
+        The function to check.
+
+    Returns
+    -------
+    bool
+        True if the function is compiled with Numba.
+    """
+    if hasattr(func, "nopython_signatures") or hasattr(func, "__numba__"):
+        return True
+
+    if hasattr(func, "_numba_type_"):
+        return True
+
+    module_name = getattr(func, "__module__", "")
+    if module_name and module_name.startswith("numba."):
+        return True
+
+    return False