Merge pull request #40 from baraline/numba_cache_install

Adding numba compilation options and an example of how to set them in examples folder. Should only be useful for some specific use case such as HPC clusters.
baraline · Mar 11, 2023 · 4913df1 · 4913df1
2 parents 74379c3 + 07c7119
commit 4913df1
Show file tree

Hide file tree

Showing 12 changed files with 123 additions and 43 deletions.
diff --git a/convst/__init__.py b/convst/__init__.py
@@ -1,5 +1,10 @@
 
 __author__ = 'Antoine Guillaume antoine.guillaume45@gmail.com'
-__version__ = "0.2.5"
+__version__ = "0.2.6"
 
 __all__ = ['transformers', 'classifiers', 'utils', 'interpreters']
+
+__USE_NUMBA_CACHE__ = True
+__USE_NUMBA_FASTMATH__ = True
+__USE_NUMBA_NOGIL__ = True
+__USE_NUMBA_PARALLEL__ = True
diff --git a/convst/transformers/_commons.py b/convst/transformers/_commons.py
@@ -9,8 +9,15 @@
 #                                                                             #
 ###############################################################################
 
+from convst import (
+    __USE_NUMBA_CACHE__, __USE_NUMBA_FASTMATH__,
+    __USE_NUMBA_NOGIL__
+)
+
+
+
 @njit(
-  fastmath=True, cache=True, nogil=True
+  fastmath=__USE_NUMBA_FASTMATH__, cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 def euclidean(x, y):
     s = 0
@@ -19,7 +26,7 @@ def euclidean(x, y):
     return sqrt(s)
 
 @njit(
-  fastmath=True, cache=True, nogil=True
+  fastmath=__USE_NUMBA_FASTMATH__, cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 def squared_euclidean(x, y):
     s = 0
@@ -28,7 +35,7 @@ def squared_euclidean(x, y):
     return s
 
 @njit(
-  fastmath=True, cache=True, nogil=True
+  fastmath=__USE_NUMBA_FASTMATH__, cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 def manhattan(x, y):
     s = 0
@@ -43,7 +50,7 @@ def manhattan(x, y):
 ###############################################################################
 
 @njit(
-  cache=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 def generate_strides_1D(X, window_size, dilation, use_phase):
     if use_phase:
@@ -52,7 +59,7 @@ def generate_strides_1D(X, window_size, dilation, use_phase):
         return _generate_strides_1D(X, window_size, dilation)
 
 @njit(
-  cache=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 
 def generate_strides_2D(X, window_size, dilation, use_phase):
@@ -63,7 +70,7 @@ def generate_strides_2D(X, window_size, dilation, use_phase):
 
 
 @njit(
-  cache=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 def _generate_strides_1D(X, window_size, dilation):
     """
@@ -93,7 +100,7 @@ def _generate_strides_1D(X, window_size, dilation):
     return X_new
 
 @njit(
-  cache=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 
 def _generate_strides_2D(X, window_size, dilation):
@@ -127,7 +134,7 @@ def _generate_strides_2D(X, window_size, dilation):
 
 
 @njit(
-  cache=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 def _generate_strides_1D_phase(X, window_size, dilation):
     """
@@ -157,7 +164,7 @@ def _generate_strides_1D_phase(X, window_size, dilation):
 
 
 @njit(
-  cache=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 def _generate_strides_2D_phase(X, window_size, dilation):
     """
@@ -189,7 +196,7 @@ def _generate_strides_2D_phase(X, window_size, dilation):
 
 
 @njit(
-  cache=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 def get_subsequence(X, index, length, d, normalize, use_phase):
     if use_phase:
@@ -203,7 +210,7 @@ def get_subsequence(X, index, length, d, normalize, use_phase):
 
 
 @njit(
-  cache=True, fastmath=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__
 )
 def _get_subsequence(X, i_start, length, d, normalize):
     """
@@ -242,7 +249,7 @@ def _get_subsequence(X, i_start, length, d, normalize):
 
 
 @njit(
-  cache=True, fastmath=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__
 )
 def _get_subsequence_phase(X, i_start, length, d, normalize):
     """
@@ -289,7 +296,7 @@ def _get_subsequence_phase(X, i_start, length, d, normalize):
 
 
 @njit(
-  cache=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 def compute_shapelet_dist_vector(
     x, values, length, dilation, dist_func, normalize, use_phase
@@ -315,7 +322,7 @@ def compute_shapelet_dist_vector(
 
 
 @njit(
-  cache=True, fastmath=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__
 )
 def _compute_shapelet_dist_vector(x, values, length, dilation, dist_func):
     """
@@ -351,7 +358,7 @@ def _compute_shapelet_dist_vector(x, values, length, dilation, dist_func):
 
 
 @njit(
-  cache=True, fastmath=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__
 )
 def _compute_shapelet_dist_vector_norm(x, values, length, dilation, dist_func):
     """
@@ -388,7 +395,7 @@ def _compute_shapelet_dist_vector_norm(x, values, length, dilation, dist_func):
 
 
 @njit(
-  cache=True, fastmath=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__
 )
 def _compute_shapelet_dist_vector_phase(x, values, length, dilation, dist_func):
     """
@@ -424,7 +431,7 @@ def _compute_shapelet_dist_vector_phase(x, values, length, dilation, dist_func):
 
 
 @njit(
-  cache=True, fastmath=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__
 )
 def _compute_shapelet_dist_vector_norm_phase(x, values, length, dilation, dist_func):
     """
@@ -462,7 +469,7 @@ def _compute_shapelet_dist_vector_norm_phase(x, values, length, dilation, dist_f
 
 
 @njit(
-  cache=True, fastmath=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__
 )
 def apply_one_shapelet_one_sample_univariate(x, values, threshold, dist_func):
     """
@@ -514,7 +521,7 @@ def apply_one_shapelet_one_sample_univariate(x, values, threshold, dist_func):
 
 
 @njit(
-  cache=True, fastmath=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__
 )
 def apply_one_shapelet_one_sample_multivariate(x, values, threshold, dist_func):
     """
@@ -569,7 +576,7 @@ def apply_one_shapelet_one_sample_multivariate(x, values, threshold, dist_func):
 
 
 @njit(
-  cache=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 def _combinations_1d(x,y):
     """
@@ -608,7 +615,7 @@ def _combinations_1d(x,y):
     return combinations
 
 @njit(
-  cache=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 def prime_up_to(n):
     is_p = zeros(n+1, dtype=bool_)
@@ -618,7 +625,7 @@ def prime_up_to(n):
 
 
 @njit(
-  cache=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__
 )
 def is_prime(n):
     if (n % 2 == 0 and n > 2) or n == 0: 
@@ -629,7 +636,7 @@ def is_prime(n):
     return True
 
 @njit(
-  cache=True, fastmath=True, nogil=True
+  cache=__USE_NUMBA_CACHE__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__
 )
 def choice_log(n_choice, n_sample):
     if n_choice > 1:

diff --git a/convst/transformers/_input_transformers.py b/convst/transformers/_input_transformers.py
@@ -20,6 +20,12 @@
 from sklearn.preprocessing import StandardScaler, MinMaxScaler
 from sklearn.base import BaseEstimator, TransformerMixin
 
+
+from convst import (
+    __USE_NUMBA_CACHE__, __USE_NUMBA_PARALLEL__
+)
+
+
 class c_StandardScaler(StandardScaler):
     def fit(self, X, y=None):
         self.usefull_atts = np.where(np.std(X, axis=0) != 0)[0]
@@ -36,15 +42,15 @@ def fit(self, X, y=None):
     def transform(self, X):
         return super().transform(X[:, self.usefull_atts])
 
-@njit(cache=True)
+@njit(cache=__USE_NUMBA_CACHE__)
 def z_norm_one_sample(x):
     n_features, n_timestamps = x.shape
     x_new = np.empty((n_features, n_timestamps))
     for i in prange(n_features):
         x_new[i] = (x[i] - x[i].mean()) / (x[i].std() + 1e-8)
     return x_new
 
-@njit(cache=True, parallel=True)
+@njit(cache=__USE_NUMBA_CACHE__, parallel=__USE_NUMBA_PARALLEL__)
 def z_norm_all_samples(X):
     n_samples, n_features, n_timestamps = X.shape
     X_new = np.empty((n_samples, n_features, n_timestamps))

diff --git a/convst/transformers/_multivariate_same_length.py b/convst/transformers/_multivariate_same_length.py
@@ -17,7 +17,13 @@
 
 from numba import njit, prange
 
-@njit(cache=True, nogil=True)
+from convst import (
+    __USE_NUMBA_CACHE__, __USE_NUMBA_FASTMATH__,
+    __USE_NUMBA_NOGIL__, __USE_NUMBA_PARALLEL__
+)
+
+
+@njit(cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__)
 def M_SL_init_random_shapelet_params(
     n_shapelets, shapelet_sizes, n_timestamps, p_norm, max_channels, prime_scheme
 ):
@@ -96,7 +102,7 @@ def M_SL_init_random_shapelet_params(
 
     return values, lengths, dilations, threshold, normalize, n_channels, channel_ids
 
-@njit(cache=True, parallel=True, nogil=True)
+@njit(cache=__USE_NUMBA_CACHE__, parallel=__USE_NUMBA_PARALLEL__, nogil=__USE_NUMBA_NOGIL__)
 def M_SL_generate_shapelet(
     X, y, n_shapelets, shapelet_sizes, r_seed, p_norm, p_min, p_max, alpha,
     dist_func, use_phase, max_channels, prime_scheme
@@ -273,7 +279,7 @@ def M_SL_generate_shapelet(
         channel_ids[:a2]
     )
 
-@njit(cache=True, parallel=True, fastmath=True, nogil=True)
+@njit(cache=__USE_NUMBA_CACHE__, parallel=__USE_NUMBA_PARALLEL__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__)
 def M_SL_apply_all_shapelets(
     X, shapelets, dist_func, use_phase
 ):

diff --git a/convst/transformers/_multivariate_variable_length.py b/convst/transformers/_multivariate_variable_length.py
@@ -17,7 +17,13 @@
 
 from numba import njit, prange
 
-@njit(cache=True, nogil=True)
+from convst import (
+    __USE_NUMBA_CACHE__, __USE_NUMBA_FASTMATH__,
+    __USE_NUMBA_NOGIL__, __USE_NUMBA_PARALLEL__
+)
+
+
+@njit(cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__)
 def M_VL_init_random_shapelet_params(
     n_shapelets, shapelet_sizes, n_timestamps, p_norm, max_channels, prime_scheme
 ):
@@ -96,7 +102,7 @@ def M_VL_init_random_shapelet_params(
 
     return values, lengths, dilations, threshold, normalize, n_channels, channel_ids
 
-@njit(cache=True, parallel=True, nogil=True)
+@njit(cache=__USE_NUMBA_CACHE__, parallel=__USE_NUMBA_PARALLEL__, nogil=__USE_NUMBA_NOGIL__)
 def M_VL_generate_shapelet(
     X, y, n_shapelets, shapelet_sizes, r_seed, p_norm, p_min, p_max, alpha,
     dist_func, use_phase, max_channels, min_len, X_len, prime_scheme
@@ -302,7 +308,7 @@ def M_VL_generate_shapelet(
         channel_ids[:a2]
     )
 
-@njit(cache=True, parallel=True, fastmath=True, nogil=True)
+@njit(cache=__USE_NUMBA_CACHE__, parallel=__USE_NUMBA_PARALLEL__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__)
 def M_VL_apply_all_shapelets(
     X, shapelets, dist_func, use_phase, X_len
 ):

diff --git a/convst/transformers/_univariate_same_length.py b/convst/transformers/_univariate_same_length.py
@@ -16,7 +16,12 @@
 
 from numba import njit, prange
 
-@njit(cache=True, nogil=True)
+from convst import (
+    __USE_NUMBA_CACHE__, __USE_NUMBA_FASTMATH__,
+    __USE_NUMBA_NOGIL__, __USE_NUMBA_PARALLEL__
+)
+
+@njit(cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__)
 def U_SL_init_random_shapelet_params(
     n_shapelets, shapelet_sizes, n_timestamps, p_norm, prime_scheme
 ):
@@ -81,7 +86,7 @@ def U_SL_init_random_shapelet_params(
 
     return values, lengths, dilations, threshold, normalize
 
-@njit(cache=True, parallel=True, nogil=True)
+@njit(cache=__USE_NUMBA_CACHE__, parallel=__USE_NUMBA_PARALLEL__, nogil=__USE_NUMBA_NOGIL__)
 def U_SL_generate_shapelet(
     X, y, n_shapelets, shapelet_sizes, r_seed, p_norm, p_min, p_max, alpha,
     dist_func, use_phase, prime_scheme
@@ -227,7 +232,7 @@ def U_SL_generate_shapelet(
     )
 
 
-@njit(cache=True, parallel=True, fastmath=True, nogil=True)
+@njit(cache=__USE_NUMBA_CACHE__, parallel=__USE_NUMBA_PARALLEL__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__)
 def U_SL_apply_all_shapelets(
     X, shapelets, dist_func, use_phase
 ):

diff --git a/convst/transformers/_univariate_variable_length.py b/convst/transformers/_univariate_variable_length.py
@@ -17,9 +17,12 @@
 
 from numba import njit, prange
 
-# TODO : check if numba could support Tuple of variable length numpy arrays as input
+from convst import (
+    __USE_NUMBA_CACHE__, __USE_NUMBA_FASTMATH__,
+    __USE_NUMBA_NOGIL__, __USE_NUMBA_PARALLEL__
+)
 
-@njit(cache=True, nogil=True)
+@njit(cache=__USE_NUMBA_CACHE__, nogil=__USE_NUMBA_NOGIL__)
 def U_VL_init_random_shapelet_params(
     n_shapelets, shapelet_sizes, n_timestamps, p_norm, prime_scheme
 ):
@@ -87,7 +90,7 @@ def U_VL_init_random_shapelet_params(
     return values, lengths, dilations, threshold, normalize
 
 
-@njit(cache=True, parallel=True, nogil=True)
+@njit(cache=__USE_NUMBA_CACHE__, parallel=__USE_NUMBA_PARALLEL__, nogil=__USE_NUMBA_NOGIL__)
 def U_VL_generate_shapelet(
     X, y, n_shapelets, shapelet_sizes, r_seed, p_norm, p_min, p_max, alpha,
     dist_func, use_phase, min_len, X_len, prime_scheme
@@ -259,7 +262,7 @@ def U_VL_generate_shapelet(
     )
 
 
-@njit(cache=True, parallel=True, fastmath=True, nogil=True)
+@njit(cache=__USE_NUMBA_CACHE__, parallel=__USE_NUMBA_PARALLEL__, fastmath=__USE_NUMBA_FASTMATH__, nogil=__USE_NUMBA_NOGIL__)
 def U_VL_apply_all_shapelets(
     X, shapelets, dist_func, use_phase, X_len
 ):