Fix fasthmath precision issue

fastmath.py

@@ -13,7 +13,7 @@ def get_njit_funcs(pkg_dir):
     Parameters
     ----------
     pkg_dir : str
-       The path to the directory containing some .py files
+        The path to the directory containing some .py files
 
     Returns
     -------

stumpy/__init__.py

@@ -1,9 +1,12 @@
+import importlib
 import os.path
 from importlib.metadata import distribution
 from site import getsitepackages
 
+import numba
 from numba import cuda
 
+from . import cache, config
 from .aamp import aamp  # noqa: F401
 from .aamp_mmotifs import aamp_mmotifs  # noqa: F401
 from .aamp_motifs import aamp_match, aamp_motifs  # noqa: F401
@@ -32,6 +35,18 @@
 from .stumped import stumped  # noqa: F401
 from .stumpi import stumpi  # noqa: F401
 
+# Get the default fastmath flags for all njit functions
+# and update the _STUMPY_DEFAULTS dictionary
+
+if not numba.config.DISABLE_JIT:  # pragma: no cover
+    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)
+        key = module_name + "." + func_name  # e.g., core._mass
+        key = "STUMPY_FASTMATH_" + key.upper()  # e.g., STUMPY_FASTHMATH_CORE._MASS
+        config._STUMPY_DEFAULTS[key] = func.targetoptions["fastmath"]
+
 if cuda.is_available():
     from .gpu_aamp import gpu_aamp  # noqa: F401
     from .gpu_aamp_ostinato import gpu_aamp_ostinato  # noqa: F401

stumpy/aamp.py

@@ -13,7 +13,7 @@
 @njit(
     # "(f8[:], f8[:], i8, b1[:], b1[:], f8, i8[:], i8, i8, i8, f8[:, :, :],"
     # "f8[:, :], f8[:, :], i8[:, :, :], i8[:, :], i8[:, :], b1)",
-    fastmath=True,
+    fastmath=config.STUMPY_FASTMATH_TRUE,
 )
 def _compute_diagonal(
     T_A,
@@ -186,7 +186,7 @@ def _compute_diagonal(
 @njit(
     # "(f8[:], f8[:], i8, b1[:], b1[:], i8[:], b1, i8)",
     parallel=True,
-    fastmath=True,
+    fastmath=config.STUMPY_FASTMATH_TRUE,
 )
 def _aamp(
     T_A,

stumpy/cache.py

@@ -4,10 +4,13 @@
 
 import ast
 import importlib
+import inspect
 import pathlib
 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 "
@@ -74,7 +77,15 @@ def _enable():
     -------
     None
     """
-    warnings.warn(CACHE_WARNING)
+    frame = inspect.currentframe()
+    caller_name = inspect.getouterframes(frame)[1].function
+    if caller_name != "_save":
+        msg = (
+            "The 'cache._enable()' function is deprecated and no longer supported. "
+            + "Please use 'cache.save()' instead"
+        )
+        warnings.warn(msg, DeprecationWarning, stacklevel=2)
+
     njit_funcs = get_njit_funcs()
     for module_name, func_name in njit_funcs:
         module = importlib.import_module(f".{module_name}", package="stumpy")
@@ -94,12 +105,29 @@ def _clear():
     -------
     None
     """
-    warnings.warn(CACHE_WARNING)
     site_pkg_dir = site.getsitepackages()[0]
     numba_cache_dir = site_pkg_dir + "/stumpy/__pycache__"
     [f.unlink() for f in pathlib.Path(numba_cache_dir).glob("*nb*") if f.is_file()]
 
 
+def clear():
+    """
+    Clear numba cache directory
+
+    Parameters
+    ----------
+    None
+
+    Returns
+    -------
+    None
+    """
+    warnings.warn(CACHE_WARNING)
+    _clear()
+
+    return
+
+
 def _get_cache():
     """
     Retrieve a list of cached numba functions
@@ -117,3 +145,69 @@ 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():
+    """
+    Recompile all njit functions
+
+    Parameters
+    ----------
+    None
+
+    Returns
+    -------
+    None
+
+    Notes
+    -----
+    If the `numba` cache is enabled, this results in saving (and/or overwriting)
+    the cached numba functions to disk.
+    """
+    for module_name, func_name in get_njit_funcs():
+        module = importlib.import_module(f".{module_name}", package="stumpy")
+        func = getattr(module, func_name)
+        func.recompile()
+
+    return
+
+
+def _save():
+    """
+    Save all njit functions
+
+    Parameters
+    ----------
+    None
+
+    Returns
+    -------
+    None
+    """
+    _enable()
+    _recompile()
+
+    return
+
+
+def save():
+    """
+    Save/overwrite all the cache data files of
+    all-so-far compiled njit functions.
+
+    Parameters
+    ----------
+    None
+
+    Returns
+    -------
+    None
+    """
+    if numba.config.DISABLE_JIT:
+        msg = "Could not save/cache function because NUMBA JIT is disabled"
+        warnings.warn(msg)
+    else:
+        warnings.warn(CACHE_WARNING)
+        _save()
+
+    return

stumpy/config.py

@@ -2,15 +2,72 @@
 # Copyright 2019 TD Ameritrade. Released under the terms of the 3-Clause BSD license.
 # STUMPY is a trademark of TD Ameritrade IP Company, Inc. All rights reserved.
 
+import warnings
+
 import numpy as np
 
-STUMPY_THREADS_PER_BLOCK = 512
-STUMPY_MEAN_STD_NUM_CHUNKS = 1
-STUMPY_MEAN_STD_MAX_ITER = 10
-STUMPY_DENOM_THRESHOLD = 1e-14
-STUMPY_STDDEV_THRESHOLD = 1e-7
-STUMPY_P_NORM_THRESHOLD = 1e-14
-STUMPY_TEST_PRECISION = 5
-STUMPY_MAX_P_NORM_DISTANCE = np.finfo(np.float64).max
-STUMPY_MAX_DISTANCE = np.sqrt(STUMPY_MAX_P_NORM_DISTANCE)
-STUMPY_EXCL_ZONE_DENOM = 4
+_STUMPY_DEFAULTS = {
+    "STUMPY_THREADS_PER_BLOCK": 512,
+    "STUMPY_MEAN_STD_NUM_CHUNKS": 1,
+    "STUMPY_MEAN_STD_MAX_ITER": 10,
+    "STUMPY_DENOM_THRESHOLD": 1e-14,
+    "STUMPY_STDDEV_THRESHOLD": 1e-7,
+    "STUMPY_P_NORM_THRESHOLD": 1e-14,
+    "STUMPY_TEST_PRECISION": 5,
+    "STUMPY_MAX_P_NORM_DISTANCE": np.finfo(np.float64).max,
+    "STUMPY_MAX_DISTANCE": np.sqrt(np.finfo(np.float64).max),
+    "STUMPY_EXCL_ZONE_DENOM": 4,
+    "STUMPY_FASTMATH_TRUE": True,
+    "STUMPY_FASTMATH_FLAGS": {"nsz", "arcp", "contract", "afn", "reassoc"},
+}
+
+# In addition to these configuration variables, there exist config variables
+# that have the default value of the fastmath flag of the njit functions. The
+# name of this config variable has the following format:
+# STUMPY_FASTMATH_<module_name>.<function_name>
+# See __init__.py for more details
+
+STUMPY_THREADS_PER_BLOCK = _STUMPY_DEFAULTS["STUMPY_THREADS_PER_BLOCK"]
+STUMPY_MEAN_STD_NUM_CHUNKS = _STUMPY_DEFAULTS["STUMPY_MEAN_STD_NUM_CHUNKS"]
+STUMPY_MEAN_STD_MAX_ITER = _STUMPY_DEFAULTS["STUMPY_MEAN_STD_MAX_ITER"]
+STUMPY_DENOM_THRESHOLD = _STUMPY_DEFAULTS["STUMPY_DENOM_THRESHOLD"]
+STUMPY_STDDEV_THRESHOLD = _STUMPY_DEFAULTS["STUMPY_STDDEV_THRESHOLD"]
+STUMPY_P_NORM_THRESHOLD = _STUMPY_DEFAULTS["STUMPY_P_NORM_THRESHOLD"]
+STUMPY_TEST_PRECISION = _STUMPY_DEFAULTS["STUMPY_TEST_PRECISION"]
+STUMPY_MAX_P_NORM_DISTANCE = _STUMPY_DEFAULTS["STUMPY_MAX_P_NORM_DISTANCE"]
+STUMPY_MAX_DISTANCE = _STUMPY_DEFAULTS["STUMPY_MAX_DISTANCE"]
+STUMPY_EXCL_ZONE_DENOM = _STUMPY_DEFAULTS["STUMPY_EXCL_ZONE_DENOM"]
+STUMPY_FASTMATH_TRUE = _STUMPY_DEFAULTS["STUMPY_FASTMATH_TRUE"]
+STUMPY_FASTMATH_FLAGS = _STUMPY_DEFAULTS["STUMPY_FASTMATH_FLAGS"]
+
+
+def _reset(var=None):
+    """
+    Reset the value of a configuration variable(s) to their default value(s)
+
+    Parameters
+    ----------
+    var : str, default None
+        The name of the configuration variable. If None, then all
+        configuration variables are reset to their default values.
+
+    Returns
+    -------
+    None
+    """
+    config_vars = [
+        k for k, _ in globals().items() if k.isupper() and k.startswith("STUMPY")
+    ]
+
+    if var is None:
+        for config_var in config_vars:
+            globals()[config_var] = _STUMPY_DEFAULTS[config_var]
+    elif var in config_vars:
+        globals()[var] = _STUMPY_DEFAULTS[var]
+    else:  # pragma: no cover
+        msg = (
+            f"Configuration reset was skipped for unrecognized '_STUMPY_DEFAULT[{var}]'"
+        )
+        warnings.warn(msg)
+
+    return