Revised test functions

Author: NimaSarajpoor

tests/test_fastmath.py

@@ -1,37 +1,36 @@
 import numpy as np
+from numba import njit
 
 from stumpy import fastmath
 
 
-def test_fastmath():
+def test_set():
     # 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
+    # the fastmath._add_assoc function, taken from the following link:
+    # https://numba.pydata.org/numba-doc/dev/user/performance-tips.html#fastmath
+    py_func = fastmath._add_assoc.py_func
 
-    x, y = 0.0, np.inf
+    x = 0.0
+    y = np.inf
+    fastmath_flags = [False, {"reassoc", "nsz"}, {"reassoc"}, {"nsz"}]
+    for flag in fastmath_flags:
+        ref = njit(fastmath=flag)(py_func)(x, y)
 
-    # 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._set("fastmath", "_add_assoc", flag)
+        comp = fastmath._add_assoc(x, y)
 
-    # fastmath={'reassoc'}
-    fastmath._set("fastmath", "_add_assoc", {"reassoc"})
-    out = fastmath._add_assoc(x, y)
-    assert np.isnan(out)
+        if np.isnan(ref) and np.isnan(comp):
+            assert True
+        else:
+            assert ref == comp
 
-    # 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)
+def test_reset():
+    # 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, taken from the following link:
+    # https://numba.pydata.org/numba-doc/dev/user/performance-tips.html#fastmath
+    fastmath._set("fastmath", "_add_assoc", False)
     fastmath._reset("fastmath", "_add_assoc")
-    out = fastmath._add_assoc(x, y)
-    assert out == 0.0
+    assert fastmath._add_assoc(0.0, np.inf) == 0.0