fixed linitng issue from unit tests.

Mojo-Numerics-and-Algorithms-group · Sep 7, 2024 · 702f522 · 702f522
1 parent 2e89312
commit 702f522
Show file tree

Hide file tree

Showing 4 changed files with 209 additions and 53 deletions.
diff --git a/tests/test_array_manipulation.mojo b/tests/test_array_manipulation.mojo
@@ -4,9 +4,10 @@ from testing.testing import assert_true, assert_almost_equal, assert_equal
 from utils_for_test import check, check_is_close
 from python import Python
 
+
 def test_arr_manipulation():
     var np = Python.import_module("numpy")
-    
+
     # Test arange
     var A = nm.arange[nm.i16](1, 7, 1)
     var np_A = np.arange(1, 7, 1, dtype=np.int16)

diff --git a/tests/test_bool_masks.mojo b/tests/test_bool_masks.mojo
@@ -6,43 +6,43 @@ from python import Python
 
 # def test_bool_masks():
 #     var np = Python.import_module("numpy")
-    
+
 #     # Create NumPy and NuMojo arrays using arange and reshape
 #     var np_A = np.arange(0, 24, dtype=np.int16).reshape((3, 2, 4))
 #     var A = nm.arange[nm.i16](0, 24)
 #     A.reshape(3, 2, 4)
-    
+
 #     # Test greater than
 #     var np_gt = np_A > 10
 #     var gt = A > Scalar[nm.i16](10)
 #     check(gt, np_gt, "Greater than mask")
-    
+
 #     # Test greater than or equal
 #     var np_ge = np_A >= 10
 #     var ge = A >= Scalar[nm.i16](10)
 #     check(ge, np_ge, "Greater than or equal mask")
-    
+
 #     # Test less than
 #     var np_lt = np_A < 10
 #     var lt = A < Scalar[nm.i16](10)
 #     check(lt, np_lt, "Less than mask")
-    
+
 #     # Test less than or equal
 #     var np_le = np_A <= 10
 #     var le = A <= Scalar[nm.i16](10)
 #     check(le, np_le, "Less than or equal mask")
-    
+
 #     # Test equal
 #     var np_eq = np_A == 10
 #     var eq = A == Scalar[nm.i16](10)
 #     check(eq, np_eq, "Equal mask")
-    
+
 #     # Test not equal
 #     var np_ne = np_A != 10
 #     var ne = A != Scalar[nm.i16](10)
 #     check(ne, np_ne, "Not equal mask")
-    
+
 #     # Test masked array
 #     var np_mask = np_A[np_A > 10]
 #     var mask = A[A > Scalar[nm.i16](10)]
-#     check(mask, np_mask, "Masked array")
+#     check(mask, np_mask, "Masked array")
diff --git a/tests/test_constructors.mojo b/tests/test_constructors.mojo
@@ -3,9 +3,10 @@ from numojo import *
 from testing.testing import assert_true, assert_almost_equal, assert_equal
 from utils_for_test import check, check_is_close
 
+
 def test_constructors():
     # Test NDArray constructor with different input types
-    var arr1 = NDArray[f32](3,4,5)
+    var arr1 = NDArray[f32](3, 4, 5)
     assert_true(arr1.ndim == 3, "NDArray constructor: ndim")
     assert_true(arr1.ndshape[0] == 3, "NDArray constructor: shape element 0")
     assert_true(arr1.ndshape[1] == 4, "NDArray constructor: shape element 1")
@@ -14,25 +15,64 @@ def test_constructors():
     assert_true(arr1.dtype == DType.float32, "NDArray constructor: dtype")
 
     var arr2 = NDArray[f32](VariadicList[Int](3, 4, 5))
-    assert_true(arr2.ndshape[0] == 3, "NDArray constructor with VariadicList: shape element 0")
-    assert_true(arr2.ndshape[1] == 4, "NDArray constructor with VariadicList: shape element 1")
-    assert_true(arr2.ndshape[2] == 5, "NDArray constructor with VariadicList: shape element 2")
+    assert_true(
+        arr2.ndshape[0] == 3,
+        "NDArray constructor with VariadicList: shape element 0",
+    )
+    assert_true(
+        arr2.ndshape[1] == 4,
+        "NDArray constructor with VariadicList: shape element 1",
+    )
+    assert_true(
+        arr2.ndshape[2] == 5,
+        "NDArray constructor with VariadicList: shape element 2",
+    )
 
     var arr3 = NDArray[f32](VariadicList[Int](3, 4, 5), fill=Scalar[f32](10.0))
     # maybe it's better to return a scalar for arr[0, 0, 0]
-    assert_equal(arr3[0,0,0].get_scalar(0), 10.0, "NDArray constructor with fill value")
+    assert_equal(
+        arr3[0, 0, 0].get_scalar(0), 10.0, "NDArray constructor with fill value"
+    )
 
     var arr4 = NDArray[f32](List[Int](3, 4, 5))
-    assert_true(arr4.ndshape[0] == 3, "NDArray constructor with List: shape element 0")
-    assert_true(arr4.ndshape[1] == 4, "NDArray constructor with List: shape element 1")
-    assert_true(arr4.ndshape[2] == 5, "NDArray constructor with List: shape element 2")
+    assert_true(
+        arr4.ndshape[0] == 3, "NDArray constructor with List: shape element 0"
+    )
+    assert_true(
+        arr4.ndshape[1] == 4, "NDArray constructor with List: shape element 1"
+    )
+    assert_true(
+        arr4.ndshape[2] == 5, "NDArray constructor with List: shape element 2"
+    )
 
-    var arr5 = NDArray[f32](NDArrayShape(3,4,5))
-    assert_true(arr5.ndshape[0] == 3, "NDArray constructor with NDArrayShape: shape element 0")
-    assert_true(arr5.ndshape[1] == 4, "NDArray constructor with NDArrayShape: shape element 1")
-    assert_true(arr5.ndshape[2] == 5, "NDArray constructor with NDArrayShape: shape element 2")
+    var arr5 = NDArray[f32](NDArrayShape(3, 4, 5))
+    assert_true(
+        arr5.ndshape[0] == 3,
+        "NDArray constructor with NDArrayShape: shape element 0",
+    )
+    assert_true(
+        arr5.ndshape[1] == 4,
+        "NDArray constructor with NDArrayShape: shape element 1",
+    )
+    assert_true(
+        arr5.ndshape[2] == 5,
+        "NDArray constructor with NDArrayShape: shape element 2",
+    )
 
-    var arr6 = NDArray[f32](data=List[SIMD[f32, 1]](1,2,3,4,5,6,7,8,9,10), shape=List[Int](2,5))
-    assert_true(arr6.ndshape[0] == 2, "NDArray constructor with data and shape: shape element 0")
-    assert_true(arr6.ndshape[1] == 5, "NDArray constructor with data and shape: shape element 1")
-    assert_equal(arr6[1,4].get_scalar(0), 10.0, "NDArray constructor with data: value check")
+    var arr6 = NDArray[f32](
+        data=List[SIMD[f32, 1]](1, 2, 3, 4, 5, 6, 7, 8, 9, 10),
+        shape=List[Int](2, 5),
+    )
+    assert_true(
+        arr6.ndshape[0] == 2,
+        "NDArray constructor with data and shape: shape element 0",
+    )
+    assert_true(
+        arr6.ndshape[1] == 5,
+        "NDArray constructor with data and shape: shape element 1",
+    )
+    assert_equal(
+        arr6[1, 4].get_scalar(0),
+        10.0,
+        "NDArray constructor with data: value check",
+    )
diff --git a/tests/test_random.mojo b/tests/test_random.mojo
@@ -5,27 +5,46 @@ from python import Python, PythonObject
 from utils_for_test import check, check_is_close
 from testing.testing import assert_true, assert_almost_equal
 
+
 def test_rand():
     """Test random array generation with specified shape."""
     var arr = random.rand[nm.f64](3, 5, 2)
     assert_true(arr.ndshape[0] == 3, "Shape of random array")
     assert_true(arr.ndshape[1] == 5, "Shape of random array")
     assert_true(arr.ndshape[2] == 2, "Shape of random array")
 
+
 def test_randminmax():
     """Test random array generation with min and max values."""
     var arr_variadic = random.rand[nm.f64](10, 10, 10, min=1, max=2)
     var arr_list = random.rand[nm.f64](List[Int](10, 10, 10), min=3, max=4)
     var arr_variadic_mean = nm.cummean(arr_variadic)
     var arr_list_mean = nm.cummean(arr_list)
-    assert_almost_equal(arr_variadic_mean, 1.5, msg="Mean of random array within min and max", atol=0.1)
-    assert_almost_equal(arr_list_mean, 3.5, msg="Mean of random array within min and max", atol=0.1)
+    assert_almost_equal(
+        arr_variadic_mean,
+        1.5,
+        msg="Mean of random array within min and max",
+        atol=0.1,
+    )
+    assert_almost_equal(
+        arr_list_mean,
+        3.5,
+        msg="Mean of random array within min and max",
+        atol=0.1,
+    )
+
 
 def test_randn():
     """Test random array generation with normal distribution."""
-    var arr_variadic_01 = random.randn[nm.f64](20, 20, 20, mean=0.0, variance=1.0)
-    var arr_variadic_31 = random.randn[nm.f64](20, 20, 20, mean=3.0, variance=1.0)
-    var arr_variadic_12 = random.randn[nm.f64](20, 20, 20, mean=1.0, variance=3.0)
+    var arr_variadic_01 = random.randn[nm.f64](
+        20, 20, 20, mean=0.0, variance=1.0
+    )
+    var arr_variadic_31 = random.randn[nm.f64](
+        20, 20, 20, mean=3.0, variance=1.0
+    )
+    var arr_variadic_12 = random.randn[nm.f64](
+        20, 20, 20, mean=1.0, variance=3.0
+    )
 
     var arr_variadic_mean01 = nm.cummean(arr_variadic_01)
     var arr_variadic_mean31 = nm.cummean(arr_variadic_31)
@@ -34,19 +53,56 @@ def test_randn():
     var arr_variadic_var31 = nm.cumvariance(arr_variadic_31)
     var arr_variadic_var12 = nm.cumvariance(arr_variadic_12)
 
-    assert_almost_equal(arr_variadic_mean01, 0, msg="Mean of random array with mean 0 and variance 1", atol=0.1)
-    assert_almost_equal(arr_variadic_mean31, 3, msg="Mean of random array with mean 3 and variance 1", atol=0.1)
-    assert_almost_equal(arr_variadic_mean12, 1, msg="Mean of random array with mean 1 and variance 2", atol=0.1)
+    assert_almost_equal(
+        arr_variadic_mean01,
+        0,
+        msg="Mean of random array with mean 0 and variance 1",
+        atol=0.1,
+    )
+    assert_almost_equal(
+        arr_variadic_mean31,
+        3,
+        msg="Mean of random array with mean 3 and variance 1",
+        atol=0.1,
+    )
+    assert_almost_equal(
+        arr_variadic_mean12,
+        1,
+        msg="Mean of random array with mean 1 and variance 2",
+        atol=0.1,
+    )
+
+    assert_almost_equal(
+        arr_variadic_var01,
+        1,
+        msg="Variance of random array with mean 0 and variance 1",
+        atol=0.1,
+    )
+    assert_almost_equal(
+        arr_variadic_var31,
+        1,
+        msg="Variance of random array with mean 3 and variance 1",
+        atol=0.1,
+    )
+    assert_almost_equal(
+        arr_variadic_var12,
+        3,
+        msg="Variance of random array with mean 1 and variance 2",
+        atol=0.1,
+    )
 
-    assert_almost_equal(arr_variadic_var01, 1, msg="Variance of random array with mean 0 and variance 1", atol=0.1)
-    assert_almost_equal(arr_variadic_var31, 1, msg="Variance of random array with mean 3 and variance 1", atol=0.1)
-    assert_almost_equal(arr_variadic_var12, 3, msg="Variance of random array with mean 1 and variance 2", atol=0.1)
 
 def test_randn_list():
     """Test random array generation with normal distribution."""
-    var arr_list_01 = random.randn[nm.f64](List[Int](20, 20, 20), mean=0.0, variance=1.0)
-    var arr_list_31 = random.randn[nm.f64](List[Int](20, 20, 20), mean=3.0, variance=1.0)
-    var arr_list_12 = random.randn[nm.f64](List[Int](20, 20, 20), mean=1.0, variance=3.0)
+    var arr_list_01 = random.randn[nm.f64](
+        List[Int](20, 20, 20), mean=0.0, variance=1.0
+    )
+    var arr_list_31 = random.randn[nm.f64](
+        List[Int](20, 20, 20), mean=3.0, variance=1.0
+    )
+    var arr_list_12 = random.randn[nm.f64](
+        List[Int](20, 20, 20), mean=1.0, variance=3.0
+    )
 
     var arr_list_mean01 = nm.cummean(arr_list_01)
     var arr_list_mean31 = nm.cummean(arr_list_31)
@@ -55,36 +111,95 @@ def test_randn_list():
     var arr_list_var31 = nm.cumvariance(arr_list_31)
     var arr_list_var12 = nm.cumvariance(arr_list_12)
 
-    assert_almost_equal(arr_list_mean01, 0, msg="Mean of random array with mean 0 and variance 1", atol=0.1)
-    assert_almost_equal(arr_list_mean31, 3, msg="Mean of random array with mean 3 and variance 1", atol=0.1)
-    assert_almost_equal(arr_list_mean12, 1, msg="Mean of random array with mean 1 and variance 2", atol=0.1)
+    assert_almost_equal(
+        arr_list_mean01,
+        0,
+        msg="Mean of random array with mean 0 and variance 1",
+        atol=0.1,
+    )
+    assert_almost_equal(
+        arr_list_mean31,
+        3,
+        msg="Mean of random array with mean 3 and variance 1",
+        atol=0.1,
+    )
+    assert_almost_equal(
+        arr_list_mean12,
+        1,
+        msg="Mean of random array with mean 1 and variance 2",
+        atol=0.1,
+    )
+
+    assert_almost_equal(
+        arr_list_var01,
+        1,
+        msg="Variance of random array with mean 0 and variance 1",
+        atol=0.1,
+    )
+    assert_almost_equal(
+        arr_list_var31,
+        1,
+        msg="Variance of random array with mean 3 and variance 1",
+        atol=0.1,
+    )
+    assert_almost_equal(
+        arr_list_var12,
+        3,
+        msg="Variance of random array with mean 1 and variance 2",
+        atol=0.1,
+    )
 
-    assert_almost_equal(arr_list_var01, 1, msg="Variance of random array with mean 0 and variance 1", atol=0.1)
-    assert_almost_equal(arr_list_var31, 1, msg="Variance of random array with mean 3 and variance 1", atol=0.1)
-    assert_almost_equal(arr_list_var12, 3, msg="Variance of random array with mean 1 and variance 2", atol=0.1)
 
 def test_rand_exponential():
     """Test random array generation with exponential distribution."""
     var arr_variadic = random.rand_exponential[nm.f64](20, 20, 20, rate=2.0)
-    var arr_list = random.rand_exponential[nm.f64](List[Int](20, 20, 20), rate=0.5)
+    var arr_list = random.rand_exponential[nm.f64](
+        List[Int](20, 20, 20), rate=0.5
+    )
 
     var arr_variadic_mean = nm.cummean(arr_variadic)
     var arr_list_mean = nm.cummean(arr_list)
 
     # For exponential distribution, mean = 1 / rate
-    assert_almost_equal(arr_variadic_mean, 1/2, msg="Mean of exponential distribution with rate 2.0", atol=0.1)
-    assert_almost_equal(arr_list_mean, 1/0.5, msg="Mean of exponential distribution with rate 0.5", atol=0.2)
+    assert_almost_equal(
+        arr_variadic_mean,
+        1 / 2,
+        msg="Mean of exponential distribution with rate 2.0",
+        atol=0.1,
+    )
+    assert_almost_equal(
+        arr_list_mean,
+        1 / 0.5,
+        msg="Mean of exponential distribution with rate 0.5",
+        atol=0.2,
+    )
 
     # For exponential distribution, variance = 1 / (rate^2)
     var arr_variadic_var = nm.cumvariance(arr_variadic)
     var arr_list_var = nm.cumvariance(arr_list)
 
-    assert_almost_equal(arr_variadic_var, 1/(2), msg="Variance of exponential distribution with rate 2.0", atol=0.1)
-    assert_almost_equal(arr_list_var, 1/(0.5), msg="Variance of exponential distribution with rate 0.5", atol=0.5)
+    assert_almost_equal(
+        arr_variadic_var,
+        1 / (2),
+        msg="Variance of exponential distribution with rate 2.0",
+        atol=0.1,
+    )
+    assert_almost_equal(
+        arr_list_var,
+        1 / (0.5),
+        msg="Variance of exponential distribution with rate 0.5",
+        atol=0.5,
+    )
 
     # Test that all values are non-negative
     for i in range(arr_variadic.num_elements()):
-        assert_true(arr_variadic.data[i] >= 0, "Exponential distribution should only produce non-negative values")
+        assert_true(
+            arr_variadic.data[i] >= 0,
+            "Exponential distribution should only produce non-negative values",
+        )
 
     for i in range(arr_list.num_elements()):
-        assert_true(arr_list.data[i] >= 0, "Exponential distribution should only produce non-negative values")
+        assert_true(
+            arr_list.data[i] >= 0,
+            "Exponential distribution should only produce non-negative values",
+        )