Add test function for fastmath

Author: NimaSarajpoor

stumpy/fastmath.py

@@ -1,8 +1,36 @@
 import importlib
 
+from numba import njit
+
 from stumpy import config
 
 
+@njit(fastmath=True)
+def _add_assoc(x, y):  # pragma: no cover
+    """
+    A dummy function to test the fastmath module
+
+    Parameters
+    ----------
+    x : float
+        A float value
+
+    y : floatf
+        A float value
+
+    Returns
+    -------
+    out : float
+        The ouput valus
+
+    Notes
+    -----
+    This is taken from the following link:
+    https://numba.pydata.org/numba-doc/dev/user/performance-tips.html#fastmath
+    """
+    return (x - y) + y
+
+
 def _set(module_name, func_name, flag):
     """
     Set fastmath flag for a given function

tests/test_fastmath.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+from stumpy import fastmath
+
+
+def test_fastmath():
+    # Test the _set and _reset function in fastmath.py
+    # The test is done by changing the value of fastmath flag for
+    # the fastmath._add_assoc function
+    # See: https://numba.pydata.org/numba-doc/dev/user/performance-tips.html#fastmath
+
+    x, y = 0.0, np.inf
+
+    # fastmath=False
+    fastmath._set("fastmath", "_add_assoc", False)
+    out = fastmath._add_assoc(x, y)
+    assert np.isnan(out)
+
+    # fastmath={'reassoc', 'nsz'}
+    fastmath._set("fastmath", "_add_assoc", {"reassoc", "nsz"})
+    out = fastmath._add_assoc(x, y)
+    assert out == 0.0
+
+    # fastmath={'reassoc'}
+    fastmath._set("fastmath", "_add_assoc", {"reassoc"})
+    out = fastmath._add_assoc(x, y)
+    assert np.isnan(out)
+
+    # fastmath={'nsz'}
+    fastmath._set("fastmath", "_add_assoc", {"nsz"})
+    out = fastmath._add_assoc(x, y)
+    assert np.isnan(out)
+
+    # reset value of fastmath (default is True)
+    fastmath._reset("fastmath", "_add_assoc")
+    out = fastmath._add_assoc(x, y)
+    assert out == 0.0