Refactored the recompile process

Author: NimaSarajpoor

stumpy/cache.py

@@ -9,6 +9,8 @@
 import site
 import warnings
 
+import numba
+
 CACHE_WARNING = "Caching `numba` functions is purely for experimental purposes "
 CACHE_WARNING += "and should never be used or depended upon as it is not supported! "
 CACHE_WARNING += "All caching capabilities are not tested and may be removed/changed "
@@ -107,3 +109,43 @@ def _get_cache():
     site_pkg_dir = site.getsitepackages()[0]
     numba_cache_dir = site_pkg_dir + "/stumpy/__pycache__"
     return [f.name for f in pathlib.Path(numba_cache_dir).glob("*nb*") if f.is_file()]
+
+
+def _recompile(func=None, fastmath=None):
+    """
+    Recompile a jit/njit decorated function. If `func` is None, then it wiil
+    recompile all njit functions of STUMPY.
+
+    Parameters
+    ----------
+    func : a njit function, default None
+        The numba function to recompile. If None, then all njit functions
+        of STUMPY will be recompiled.
+
+    fastmath : bool or set, default None
+        The fastmath flags to use. If None, then the func's fastmath flags
+        will not be changed. This is only used when `func` is provided.
+
+    Returns
+    -------
+    None
+    """
+    warnings.warn(CACHE_WARNING)
+    if func is None:
+        njit_funcs = get_njit_funcs()
+        for module_name, func_name in njit_funcs:
+            module = importlib.import_module(f".{module_name}", package="stumpy")
+            func = getattr(module, func_name)
+            func.recompile()
+
+    else:
+        if not numba.extending.is_jitted(func):
+            msg = "The function `func` must be a (n)jit function."
+            raise ValueError(msg)
+
+        if fastmath is not None:
+            func.targetoptions["fastmath"] = fastmath
+
+        func.recompile()
+
+    return

tests/test_precision.py

@@ -1,5 +1,4 @@
 import functools
-import importlib
 from unittest.mock import patch
 
 import naive
@@ -149,75 +148,48 @@ def test_snippets():
         cmp_regimes,
     ) = stumpy.snippets(T, m, k, s=s, mpdist_T_subseq_isconstant=isconstant_custom_func)
 
-    if np.allclose(ref_snippets, cmp_snippets) or numba.config.DISABLE_JIT:
-        npt.assert_almost_equal(
-            ref_snippets, cmp_snippets, decimal=config.STUMPY_TEST_PRECISION
-        )
-        npt.assert_almost_equal(
-            ref_indices, cmp_indices, decimal=config.STUMPY_TEST_PRECISION
-        )
-        npt.assert_almost_equal(
-            ref_profiles, cmp_profiles, decimal=config.STUMPY_TEST_PRECISION
-        )
-        npt.assert_almost_equal(
-            ref_fractions, cmp_fractions, decimal=config.STUMPY_TEST_PRECISION
-        )
-        npt.assert_almost_equal(
-            ref_areas, cmp_areas, decimal=config.STUMPY_TEST_PRECISION
-        )
-        npt.assert_almost_equal(ref_regimes, cmp_regimes)
-    else:
-        # Revise fastmath flag, recompile, and re-calculate snippets,
-        # and then revert the changes
-
+    if not np.allclose(ref_snippets, cmp_snippets) and not numba.config.DISABLE_JIT:
+        # Revise fastmath flags by removing reassoc (to improve precision),
+        # recompile njit functions, and re-compute snippets.
         config.STUMPY_FASTMATH_FLAGS = {"nsz", "arcp", "contract", "afn"}
-        core._calculate_squared_distance.targetoptions["fastmath"] = (
-            config.STUMPY_FASTMATH_FLAGS
+        cache._recompile(
+            core._calculate_squared_distance, fastmath=config.STUMPY_FASTMATH_FLAGS
         )
-
-        njit_funcs = cache.get_njit_funcs()
-        for module_name, func_name in njit_funcs:
-            module = importlib.import_module(f".{module_name}", package="stumpy")
-            func = getattr(module, func_name)
-            func.recompile()
+        cache._recompile()
 
         (
-            cmp_snippets_NOreassoc,
-            cmp_indices_NOreassoc,
-            cmp_profiles_NOreassoc,
-            cmp_fractions_NOreassoc,
-            cmp_areas_NOreassoc,
-            cmp_regimes_NOreassoc,
+            cmp_snippets,
+            cmp_indices,
+            cmp_profiles,
+            cmp_fractions,
+            cmp_areas,
+            cmp_regimes,
         ) = stumpy.snippets(
             T, m, k, s=s, mpdist_T_subseq_isconstant=isconstant_custom_func
         )
 
-        config._reset("STUMPY_FASTMATH_FLAGS")
-
-        core._calculate_squared_distance.targetoptions["fastmath"] = (
-            config.STUMPY_FASTMATH_FLAGS
-        )
-        for module_name, func_name in njit_funcs:
-            module = importlib.import_module(f".{module_name}", package="stumpy")
-            func = getattr(module, func_name)
-            func.recompile()
+    npt.assert_almost_equal(
+        ref_snippets, cmp_snippets, decimal=config.STUMPY_TEST_PRECISION
+    )
+    npt.assert_almost_equal(
+        ref_indices, cmp_indices, decimal=config.STUMPY_TEST_PRECISION
+    )
+    npt.assert_almost_equal(
+        ref_profiles, cmp_profiles, decimal=config.STUMPY_TEST_PRECISION
+    )
+    npt.assert_almost_equal(
+        ref_fractions, cmp_fractions, decimal=config.STUMPY_TEST_PRECISION
+    )
+    npt.assert_almost_equal(ref_areas, cmp_areas, decimal=config.STUMPY_TEST_PRECISION)
+    npt.assert_almost_equal(ref_regimes, cmp_regimes)
 
-        npt.assert_almost_equal(
-            ref_snippets, cmp_snippets_NOreassoc, decimal=config.STUMPY_TEST_PRECISION
-        )
-        npt.assert_almost_equal(
-            ref_indices, cmp_indices_NOreassoc, decimal=config.STUMPY_TEST_PRECISION
-        )
-        npt.assert_almost_equal(
-            ref_profiles, cmp_profiles_NOreassoc, decimal=config.STUMPY_TEST_PRECISION
-        )
-        npt.assert_almost_equal(
-            ref_fractions, cmp_fractions_NOreassoc, decimal=config.STUMPY_TEST_PRECISION
-        )
-        npt.assert_almost_equal(
-            ref_areas, cmp_areas_NOreassoc, decimal=config.STUMPY_TEST_PRECISION
+    if not numba.config.DISABLE_JIT:
+        # Revert fastmath flag back to their default values
+        config._reset("STUMPY_FASTMATH_FLAGS")
+        cache._recompile(
+            core._calculate_squared_distance, fastmath=config.STUMPY_FASTMATH_FLAGS
         )
-        npt.assert_almost_equal(ref_regimes, cmp_regimes_NOreassoc)
+        cache._recompile()
 
 
 @pytest.mark.filterwarnings("ignore", category=NumbaPerformanceWarning)