Add benchmarking for HT_TRENDLINE and MA functions, along with unit tests for TA utility functions and MA functionality

Author: jstzwj

benchmark/bench_ht_trendline.py

@@ -0,0 +1,29 @@
+from utils import bench
+import tqdm
+import numpy as np
+
+import sys
+import os
+
+sys.path.append(
+    os.path.abspath(os.path.join(os.path.dirname(__file__), '../'))
+)
+
+from tabox import HT_TRENDLINE as this_HT_TRENDLINE
+from talib import HT_TRENDLINE as that_HT_TRENDLINE
+
+@bench
+def bench_this_ht_trendline():
+    for i in range(100, 2000):
+        close = np.random.random(i)
+        this_ret = this_HT_TRENDLINE(close)
+
+@bench
+def bench_that_ht_trendline():
+    for i in range(100, 2000):
+        close = np.random.random(i)
+        that_ret = that_HT_TRENDLINE(close)
+
+if __name__ == '__main__':
+    bench_this_ht_trendline()
+    bench_that_ht_trendline()

benchmark/bench_ma.py

@@ -16,15 +16,17 @@
 def bench_this_ma():
     for i in range(100, 2000):
         t = 14
+        matype = 0  # SMA
         close = np.random.random(i)
-        this_ret = this_MA(close, timeperiod=t)
+        this_ret = this_MA(close, timeperiod=t, matype=matype)
 
 @bench
 def bench_that_ma():
     for i in range(100, 2000):
         t = 14
+        matype = 0  # SMA
         close = np.random.random(i)
-        that_ret = that_MA(close, timeperiod=t)
+        that_ret = that_MA(close, timeperiod=t, matype=matype)
 
 if __name__ == '__main__':
     bench_this_ma()

tabox/tests/test_defs.py

@@ -0,0 +1,34 @@
+import unittest
+import numpy as np
+from tabox.ta_func.ta_defs import (
+    TA_INTEGER_MIN, TA_INTEGER_MAX, TA_REAL_MIN, TA_REAL_MAX,
+    TA_INTEGER_DEFAULT, TA_REAL_DEFAULT
+)
+
+class TestTADefs(unittest.TestCase):
+
+    def test_constants_values(self):
+        """Test that constants have expected values"""
+        self.assertEqual(TA_INTEGER_MIN, -2147483647)
+        self.assertEqual(TA_INTEGER_MAX, 2147483647)
+        self.assertEqual(TA_INTEGER_DEFAULT, -2147483648)
+        
+        self.assertEqual(TA_REAL_MIN, -3e37)
+        self.assertEqual(TA_REAL_MAX, 3e37)
+        self.assertEqual(TA_REAL_DEFAULT, -4e37)
+
+    def test_constants_relationships(self):
+        """Test relationships between constants"""
+        self.assertLess(TA_INTEGER_MIN, TA_INTEGER_MAX)
+        self.assertLess(TA_INTEGER_DEFAULT, TA_INTEGER_MIN)
+        
+        self.assertLess(TA_REAL_MIN, TA_REAL_MAX)
+        self.assertLess(TA_REAL_DEFAULT, TA_REAL_MIN)
+
+    def test_default_values_outside_range(self):
+        """Test that default values are outside valid ranges"""
+        self.assertLess(TA_INTEGER_DEFAULT, TA_INTEGER_MIN)
+        self.assertLess(TA_REAL_DEFAULT, TA_REAL_MIN)
+
+if __name__ == '__main__':
+    unittest.main()

tabox/tests/test_ma.py

@@ -0,0 +1,43 @@
+import numpy as np
+
+from tabox.ta_func.ta_MA import MA as this_MA
+from talib import MA as that_MA
+
+import unittest
+
+class TestMA(unittest.TestCase):
+
+    def test_random_vector(self):
+        for i in range(100, 300):
+            for t in [3, 5, 7, 13, 30]:
+                for matype in range(9):  # Test all 9 MA types
+                    close = np.random.random(i)
+                    this_ret = this_MA(close, timeperiod=t, matype=matype)
+                    that_ret = that_MA(close, timeperiod=t, matype=matype)
+
+                    self.assertTrue(np.allclose(this_ret, that_ret, equal_nan=True))
+
+    def test_ma_types(self):
+        """Test all supported MA types"""
+        close = np.random.random(100)
+        
+        # Test each MA type individually
+        for matype in range(9):
+            result = this_MA(close, timeperiod=14, matype=matype)
+            self.assertIsInstance(result, np.ndarray)
+            self.assertEqual(len(result), len(close))
+
+    def test_edge_cases(self):
+        """Test edge cases"""
+        close = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
+        
+        # Test minimum period
+        result = this_MA(close, timeperiod=1, matype=0)
+        np.testing.assert_array_equal(result, close)
+        
+        # Test period larger than data
+        result = this_MA(close, timeperiod=10, matype=0)
+        self.assertTrue(np.all(np.isnan(result)))
+
+if __name__ == '__main__':
+    unittest.main()

tabox/tests/test_utils.py

@@ -0,0 +1,110 @@
+import unittest
+import numpy as np
+from tabox.ta_func.ta_utils import (
+    check_array, check_timeperiod, check_begidx1, check_begidx2, 
+    check_begidx3, check_begidx4, check_length2, check_length3, 
+    check_length4, make_double_array
+)
+
+class TestTAUtils(unittest.TestCase):
+
+    def test_check_array(self):
+        """Test array validation function"""
+        # Test numpy array input
+        arr = np.array([1.0, 2.0, 3.0], dtype=np.float64)
+        result = check_array(arr)
+        np.testing.assert_array_equal(result, arr)
+        
+        # Test list input
+        lst = [1.0, 2.0, 3.0]
+        result = check_array(lst)
+        expected = np.array(lst, dtype=np.float64)
+        np.testing.assert_array_equal(result, expected)
+        
+        # Test wrong dtype
+        arr_int = np.array([1, 2, 3], dtype=np.int32)
+        with self.assertRaises(Exception):
+            check_array(arr_int)
+        
+        # Test wrong dimensions
+        arr_2d = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float64)
+        with self.assertRaises(Exception):
+            check_array(arr_2d)
+
+    def test_check_timeperiod(self):
+        """Test timeperiod validation"""
+        # Valid timeperiod
+        try:
+            check_timeperiod(5)
+        except Exception:
+            self.fail("check_timeperiod raised exception for valid timeperiod")
+        
+        # Invalid timeperiods
+        with self.assertRaises(Exception):
+            check_timeperiod(1)
+        with self.assertRaises(Exception):
+            check_timeperiod(0)
+        with self.assertRaises(Exception):
+            check_timeperiod(-1)
+
+    def test_check_begidx1(self):
+        """Test single array begin index check"""
+        # Normal array
+        arr = np.array([1.0, 2.0, 3.0], dtype=np.float64)
+        result = check_begidx1(arr)
+        self.assertEqual(result, 0)
+        
+        # Array with leading NaN
+        arr_nan = np.array([np.nan, np.nan, 1.0, 2.0], dtype=np.float64)
+        result = check_begidx1(arr_nan)
+        self.assertEqual(result, 2)
+        
+        # All NaN array
+        arr_all_nan = np.array([np.nan, np.nan, np.nan], dtype=np.float64)
+        with self.assertRaises(Exception):
+            check_begidx1(arr_all_nan)
+
+    def test_check_begidx2(self):
+        """Test dual array begin index check"""
+        arr1 = np.array([1.0, 2.0, 3.0], dtype=np.float64)
+        arr2 = np.array([4.0, 5.0, 6.0], dtype=np.float64)
+        result = check_begidx2(arr1, arr2)
+        self.assertEqual(result, 0)
+        
+        # With NaN in first array
+        arr1_nan = np.array([np.nan, 2.0, 3.0], dtype=np.float64)
+        arr2 = np.array([4.0, 5.0, 6.0], dtype=np.float64)
+        result = check_begidx2(arr1_nan, arr2)
+        self.assertEqual(result, 1)
+
+    def test_check_length2(self):
+        """Test dual array length validation"""
+        arr1 = np.array([1.0, 2.0, 3.0], dtype=np.float64)
+        arr2 = np.array([4.0, 5.0, 6.0], dtype=np.float64)
+        length = check_length2(arr1, arr2)
+        self.assertEqual(length, 3)
+        
+        # Different lengths
+        arr1 = np.array([1.0, 2.0], dtype=np.float64)
+        arr2 = np.array([4.0, 5.0, 6.0], dtype=np.float64)
+        with self.assertRaises(Exception):
+            check_length2(arr1, arr2)
+
+    def test_make_double_array(self):
+        """Test double array creation"""
+        length = 5
+        lookback = 2
+        result = make_double_array(length, lookback)
+        
+        self.assertEqual(len(result), length)
+        self.assertEqual(result.dtype, np.float64)
+        
+        # Check that first 'lookback' elements are NaN
+        for i in range(lookback):
+            self.assertTrue(np.isnan(result[i]))
+        
+        # Check that remaining elements are not NaN (they should be uninitialized but not NaN)
+        # Actually, np.empty might contain arbitrary values, so we just check length and dtype
+
+if __name__ == '__main__':
+    unittest.main()