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