Merge pull request #53 from MadAlex1997/mad_ndarray_compare

Add inout to flatten using Optional, fix comparison functions and min…

Author: MadAlex1997

numojo/core/array_creation_routines.mojo

@@ -44,10 +44,11 @@ fn arange[
     Returns:
         NDArray[dtype] - NDArray of datatype T with elements ranging from "start" to "stop" incremented with "step".
     """
-    if is_inttype[in_dtype]() and is_inttype[out_dtype]():
+    if is_floattype[in_dtype]() and is_inttype[out_dtype]():
         raise Error(
-            "Input and output cannot be `Int` datatype as it may lead to"
-            " precision errors"
+            """
+            If in_dtype is a float then out_dtype must also be a float
+            """
         )
 
     var num: Int = ((stop - start) / step).__int__()

numojo/core/ndarray.mojo

@@ -34,7 +34,7 @@ from ..math.statistics.cumulative_reduce import (
 )
 from ..math.check import any, all
 from ..math.arithmetic import abs
-from .ndarray_utils import _get_index, _traverse_iterative, to_numpy
+from .ndarray_utils import _get_index, _traverse_iterative, to_numpy, bool_to_numeric
 from .utility_funcs import is_inttype
 from ..math.linalg.matmul import matmul_parallelized
 
@@ -1612,7 +1612,7 @@ struct NDArray[dtype: DType = DType.float32](Stringable, CollectionElement, Size
         vectorize[vectorized_fill, simd_width](self.ndshape._size)
         return self
 
-    fn flatten(inout self, inplace: Bool = False) raises -> Self:
+    fn flatten(inout self, inplace: Bool = False) raises -> Optional[Self]:
         # inplace has some problems right now
         # if inplace:
         #     self.ndshape = NDArrayShape(self.ndshape._size, size=self.ndshape._size)
@@ -1628,15 +1628,18 @@ struct NDArray[dtype: DType = DType.float32](Stringable, CollectionElement, Size
             )
 
         vectorize[vectorized_flatten, simd_width](self.ndshape._size)
-        return res
+        if inplace:
+            self = res
+            return None
+        else:
+            return res
 
     fn item(self, *indices: Int) raises -> SIMD[dtype, 1]:  # I should add
         if indices.__len__() == 1:
             return self.data.load[width=1](indices[0])
         else:
             return self.data.load[width=1](_get_index(indices, self.stride))
 
-    #TODO:  not finished yet
     fn max(self, axis: Int = 0) raises -> Self: 
         var ndim: Int = self.ndim
         var shape: List[Int] = List[Int]()
@@ -1655,16 +1658,49 @@ struct NDArray[dtype: DType = DType.float32](Stringable, CollectionElement, Size
                 slices.append(Slice(0, 0))
         print(result_shape.__str__())
         var result: NDArray[dtype] = NDArray[dtype](NDArrayShape(result_shape))
-
-        for i in range(axis_size):
+        slices[axis] = Slice(0, 1)
+        result = self[slices]
+        for i in range(1,axis_size):
             slices[axis] = Slice(i, i + 1)
             var arr_slice = self[slices]
-            result += maxT(arr_slice)
+            var mask1 = greater(arr_slice,result)
+            var mask2 = less(arr_slice,result)
+            # Wherever result is less than the new slice it is set to zero
+            # Wherever arr_slice is greater than the old result it is added to fill those zeros
+            result = add(result * bool_to_numeric[dtype](mask2),arr_slice * bool_to_numeric[dtype](mask1))
 
         return result
 
-    fn min(self, axis: Int = 0):
-        pass
+    fn min(self, axis: Int = 0)raises-> Self:
+        var ndim: Int = self.ndim
+        var shape: List[Int] = List[Int]()
+        for i in range(ndim):
+            shape.append(self.ndshape[i])
+        if axis > ndim - 1:
+            raise Error("axis cannot be greater than the rank of the array")
+        var result_shape: List[Int] = List[Int]()
+        var axis_size: Int = shape[axis]
+        var slices: List[Slice] = List[Slice]()
+        for i in range(ndim):
+            if i != axis:
+                result_shape.append(shape[i])
+                slices.append(Slice(0, shape[i]))
+            else:
+                slices.append(Slice(0, 0))
+        
+        var result: NDArray[dtype] = NDArray[dtype](NDArrayShape(result_shape))
+        slices[axis] = Slice(0, 1)
+        result = self[slices]
+        for i in range(1,axis_size):
+            slices[axis] = Slice(i, i + 1)
+            var arr_slice = self[slices]
+            var mask1 = less(arr_slice,result)
+            var mask2 = greater(arr_slice,result)
+            # Wherever result is greater than the new slice it is set to zero
+            # Wherever arr_slice is less than the old result it is added to fill those zeros
+            result = add(result * bool_to_numeric[dtype](mask2),arr_slice * bool_to_numeric[dtype](mask1))
+
+        return result
 
     fn mean(self: Self, axis: Int) raises -> Self:
         """

numojo/core/ndarray_utils.mojo

@@ -84,6 +84,16 @@ fn _traverse_iterative[
             orig, narr, ndim, coefficients, strides, offset, index, newdepth
         )
 
+fn bool_to_numeric[dtype:DType](array: NDArray[DType.bool])raises->NDArray[dtype]:
+    # Can't use simd becuase of bit packing error
+    var res: NDArray[dtype] = NDArray[dtype](array.shape())
+    for i in range(array.size()):
+        var t = array[i]
+        if t:
+            res[i] = 1
+        else:
+            res[i] = 0
+    return res
 
 fn to_numpy[dtype: DType](array: NDArray[dtype]) raises -> PythonObject:
     try:

numojo/math/comparison.mojo

@@ -32,9 +32,17 @@ fn greater[
 
     An element of the result NDArray will be True if the corresponding element in x is greater than the corresponding element in y, and False otherwise.
     """
-    return backend()._math_func_compare_2_tensors[dtype, SIMD.__gt__](
-        array1, array2
-    )
+    # return backend()._math_func_compare_2_tensors[dtype, SIMD.__gt__](
+    #     array1, array2
+    # )
+    if array1.shape() != array2.shape():
+            raise Error(
+                "Shape Mismatch error shapes must match for this function"
+            )
+    var result_array: NDArray[DType.bool] = NDArray[DType.bool](array1.shape())
+    for i in range(result_array.size()):
+        result_array[i] = array1[i]>array2[i]
+    return result_array
 
 
 fn greater_equal[
@@ -56,10 +64,17 @@ fn greater_equal[
 
     An element of the result NDArray will be True if the corresponding element in x is greater than or equal to the corresponding element in y, and False otherwise.
     """
-    return backend()._math_func_compare_2_tensors[dtype, SIMD.__ge__](
-        array1, array2
-    )
-
+    # return backend()._math_func_compare_2_tensors[dtype, SIMD.__ge__](
+    #     array1, array2
+    # )
+    if array1.shape() != array2.shape():
+            raise Error(
+                "Shape Mismatch error shapes must match for this function"
+            )
+    var result_array: NDArray[DType.bool] = NDArray[DType.bool](array1.shape())
+    for i in range(result_array.size()):
+        result_array[i] = array1[i]>=array2[i]
+    return result_array
 
 fn less[
     dtype: DType, backend: _mf.Backend = _mf.Vectorized
@@ -80,10 +95,17 @@ fn less[
 
     An element of the result NDArray will be True if the corresponding element in x is or equal to the corresponding element in y, and False otherwise.
     """
-    return backend()._math_func_compare_2_tensors[dtype, SIMD.__lt__](
-        array1, array2
-    )
-
+    # return backend()._math_func_compare_2_tensors[dtype, SIMD.__lt__](
+    #     array1, array2
+    # )
+    if array1.shape() != array2.shape():
+            raise Error(
+                "Shape Mismatch error shapes must match for this function"
+            )
+    var result_array: NDArray[DType.bool] = NDArray[DType.bool](array1.shape())
+    for i in range(result_array.size()):
+        result_array[i] = array1[i]<array2[i]
+    return result_array
 
 fn less_equal[
     dtype: DType, backend: _mf.Backend = _mf.Vectorized
@@ -104,9 +126,17 @@ fn less_equal[
 
     An element of the result NDArray will be True if the corresponding element in x is less than or equal to the corresponding element in y, and False otherwise.
     """
-    return backend()._math_func_compare_2_tensors[dtype, SIMD.__le__](
-        array1, array2
-    )
+    # return backend()._math_func_compare_2_tensors[dtype, SIMD.__le__](
+    #     array1, array2
+    # )
+    if array1.shape() != array2.shape():
+            raise Error(
+                "Shape Mismatch error shapes must match for this function"
+            )
+    var result_array: NDArray[DType.bool] = NDArray[DType.bool](array1.shape())
+    for i in range(result_array.size()):
+        result_array[i] = array1[i]<=array2[i]
+    return result_array
 
 
 fn equal[
@@ -128,10 +158,17 @@ fn equal[
 
     An element of the result NDArray will be True if the corresponding element in x is equal to the corresponding element in y, and False otherwise.
     """
-    return backend()._math_func_compare_2_tensors[dtype, SIMD.__eq__](
-        array1, array2
-    )
-
+    # return backend()._math_func_compare_2_tensors[dtype, SIMD.__eq__](
+    #     array1, array2
+    # )
+    if array1.shape() != array2.shape():
+            raise Error(
+                "Shape Mismatch error shapes must match for this function"
+            )
+    var result_array: NDArray[DType.bool] = NDArray[DType.bool](array1.shape())
+    for i in range(result_array.size()):
+        result_array[i] = array1[i]==array2[i]
+    return result_array
 
 fn not_equal[
     dtype: DType, backend: _mf.Backend = _mf.Vectorized
@@ -152,6 +189,14 @@ fn not_equal[
 
     An element of the result NDArray will be True if the corresponding element in x is not equal to the corresponding element in y, and False otherwise.
     """
-    return backend()._math_func_compare_2_tensors[dtype, SIMD.__ne__](
-        array1, array2
-    )
+    # return backend()._math_func_compare_2_tensors[dtype, SIMD.__ne__](
+    #     array1, array2
+    # )
+    if array1.shape() != array2.shape():
+            raise Error(
+                "Shape Mismatch error shapes must match for this function"
+            )
+    var result_array: NDArray[DType.bool] = NDArray[DType.bool](array1.shape())
+    for i in range(result_array.size()):
+        result_array[i] = array1[i]!=array2[i]
+    return result_array

numojo/math/statistics/stats.mojo

@@ -6,7 +6,7 @@
 """
 # from numojo.core.NDArray import NDArray
 from ...core.ndarray import NDArray
-
+from .. import mul
 
 fn sum(array: NDArray, axis: Int = 0) raises -> NDArray[array.dtype]:
     """
@@ -96,7 +96,7 @@ fn prod(array: NDArray, axis: Int = 0) raises -> NDArray[array.dtype]:
     for i in range(1, axis_size):
         slices[axis] = Slice(i, i + 1)
         var arr_slice = array[slices]
-        result = result * arr_slice
+        result = mul[array.dtype](result, arr_slice)
 
     return result
 
@@ -137,3 +137,65 @@ fn meanall(array: NDArray) raises -> Float64:
     """
 
     return sumall(array).cast[DType.float64]() / Int32(array.ndshape._size).cast[DType.float64]()
+
+fn max[dtype:DType](array:NDArray[dtype], axis: Int = 0) raises -> NDArray[dtype]: 
+        var ndim: Int = array.ndim
+        var shape: List[Int] = List[Int]()
+        for i in range(ndim):
+            shape.append(array.ndshape[i])
+        if axis > ndim - 1:
+            raise Error("axis cannot be greater than the rank of the array")
+        var result_shape: List[Int] = List[Int]()
+        var axis_size: Int = shape[axis]
+        var slices: List[Slice] = List[Slice]()
+        for i in range(ndim):
+            if i != axis:
+                result_shape.append(shape[i])
+                slices.append(Slice(0, shape[i]))
+            else:
+                slices.append(Slice(0, 0))
+        print(result_shape.__str__())
+        var result: NDArray[dtype] = NDArray[dtype](NDArrayShape(result_shape))
+        slices[axis] = Slice(0, 1)
+        result = array[slices]
+        for i in range(1,axis_size):
+            slices[axis] = Slice(i, i + 1)
+            var arr_slice = array[slices]
+            var mask1 = greater(arr_slice,result)
+            var mask2 = less(arr_slice,result)
+            # Wherever result is less than the new slice it is set to zero
+            # Wherever arr_slice is greater than the old result it is added to fill those zeros
+            result = add(result * bool_to_numeric[dtype](mask2),arr_slice * bool_to_numeric[dtype](mask1))
+
+        return result
+
+fn min[dtype:DType](array:NDArray[dtype], axis: Int = 0)raises-> NDArray[dtype]:
+    var ndim: Int = array.ndim
+    var shape: List[Int] = List[Int]()
+    for i in range(ndim):
+        shape.append(array.ndshape[i])
+    if axis > ndim - 1:
+        raise Error("axis cannot be greater than the rank of the array")
+    var result_shape: List[Int] = List[Int]()
+    var axis_size: Int = shape[axis]
+    var slices: List[Slice] = List[Slice]()
+    for i in range(ndim):
+        if i != axis:
+            result_shape.append(shape[i])
+            slices.append(Slice(0, shape[i]))
+        else:
+            slices.append(Slice(0, 0))
+    
+    var result: NDArray[dtype] = NDArray[dtype](NDArrayShape(result_shape))
+    slices[axis] = Slice(0, 1)
+    result = array[slices]
+    for i in range(1,axis_size):
+        slices[axis] = Slice(i, i + 1)
+        var arr_slice = array[slices]
+        var mask1 = less(arr_slice,result)
+        var mask2 = greater(arr_slice,result)
+        # Wherever result is greater than the new slice it is set to zero
+        # Wherever arr_slice is less than the old result it is added to fill those zeros
+        result = add(result * bool_to_numeric[dtype](mask2),arr_slice * bool_to_numeric[dtype](mask1))
+
+    return result