[routine][statistics] Update the statistics module + add returned_dtype + unify functions

mojoproject.toml

@@ -30,9 +30,10 @@ t = "clear && magic run test"
 # run individual tests to avoid overheat
 test_core = "magic run mojo test tests/core -I ./ -I ./tests/"
 test_creation = "magic run mojo test tests/routines/test_creation.mojo -I ./ -I ./tests/"
-test_random = "magic run mojo test tests/routines/test_random.mojo -I ./ -I ./tests/"
-test_manipulation = "magic run mojo test tests/routines/test_manipulation.mojo -I ./ -I ./tests/"
 test_linalg = "magic run mojo test tests/routines/test_linalg.mojo -I ./ -I ./tests/"
+test_manipulation = "magic run mojo test tests/routines/test_manipulation.mojo -I ./ -I ./tests/"
+test_random = "magic run mojo test tests/routines/test_random.mojo -I ./ -I ./tests/"
+test_statistics = "magic run mojo test tests/routines/test_statistics.mojo -I ./ -I ./tests/"
 
 # run all final checks before a commit
 final = "magic run test && magic run format && magic run package"

numojo/__init__.mojo

@@ -147,16 +147,10 @@ from numojo.routines.math import (
 from numojo.routines import statistics
 from numojo.routines.statistics import (
     mean,
-    meanall,
-    cummean,
     mode,
     median,
     variance,
     std,
-    cumpvariance,
-    cumvariance,
-    cumpstdev,
-    cumstdev,
 )
 
 from numojo.routines import bitwise

numojo/core/complex/complex_ndarray.mojo

@@ -47,7 +47,7 @@ from numojo.routines.math.extrema import maxT, minT
 from numojo.routines.math.products import prod, cumprod
 from numojo.routines.math.sums import sum, cumsum
 import numojo.routines.sorting as sorting
-from numojo.routines.statistics.averages import mean, cummean
+from numojo.routines.statistics.averages import mean
 
 
 # ===----------------------------------------------------------------------===#

numojo/core/matrix.mojo

@@ -1020,20 +1020,24 @@ struct Matrix[dtype: DType = DType.float64](
         """
         return numojo.math.extrema.max(self, axis=axis)
 
-    fn mean(self) raises -> Scalar[dtype]:
+    fn mean[
+        returned_dtype: DType = DType.float64
+    ](self) raises -> Scalar[returned_dtype]:
         """
         Calculate the arithmetic average of all items in the Matrix.
         """
-        return numojo.statistics.mean(self)
+        return numojo.statistics.mean[returned_dtype](self)
 
-    fn mean(self, axis: Int) raises -> Self:
+    fn mean[
+        returned_dtype: DType = DType.float64
+    ](self, axis: Int) raises -> Matrix[returned_dtype]:
         """
         Calculate the arithmetic average of a Matrix along the axis.
 
         Args:
             axis: 0 or 1.
         """
-        return numojo.statistics.mean(self, axis=axis)
+        return numojo.statistics.mean[returned_dtype](self, axis=axis)
 
     fn min(self) raises -> Scalar[dtype]:
         """
@@ -1103,24 +1107,28 @@ struct Matrix[dtype: DType = DType.float64](
     fn round(self, decimals: Int) raises -> Self:
         return numojo.math.rounding.round(self, decimals=decimals)
 
-    fn std(self, ddof: Int = 0) raises -> Scalar[dtype]:
+    fn std[
+        returned_dtype: DType = DType.float64
+    ](self, ddof: Int = 0) raises -> Scalar[returned_dtype]:
         """
         Compute the standard deviation.
 
         Args:
             ddof: Delta degree of freedom.
         """
-        return numojo.statistics.std(self, ddof=ddof)
+        return numojo.statistics.std[returned_dtype](self, ddof=ddof)
 
-    fn std(self, axis: Int, ddof: Int = 0) raises -> Self:
+    fn std[
+        returned_dtype: DType = DType.float64
+    ](self, axis: Int, ddof: Int = 0) raises -> Matrix[returned_dtype]:
         """
         Compute the standard deviation along axis.
 
         Args:
             axis: 0 or 1.
             ddof: Delta degree of freedom.
         """
-        return numojo.statistics.std(self, axis=axis, ddof=ddof)
+        return numojo.statistics.std[returned_dtype](self, axis=axis, ddof=ddof)
 
     fn sum(self) -> Scalar[dtype]:
         """
@@ -1167,24 +1175,30 @@ struct Matrix[dtype: DType = DType.float64](
     fn T(self) -> Self:
         return transpose(self)
 
-    fn variance(self, ddof: Int = 0) raises -> Scalar[dtype]:
+    fn variance[
+        returned_dtype: DType = DType.float64
+    ](self, ddof: Int = 0) raises -> Scalar[returned_dtype]:
         """
         Compute the variance.
 
         Args:
             ddof: Delta degree of freedom.
         """
-        return numojo.statistics.variance(self, ddof=ddof)
+        return numojo.statistics.variance[returned_dtype](self, ddof=ddof)
 
-    fn variance(self, axis: Int, ddof: Int = 0) raises -> Self:
+    fn variance[
+        returned_dtype: DType = DType.float64
+    ](self, axis: Int, ddof: Int = 0) raises -> Matrix[returned_dtype]:
         """
         Compute the variance along axis.
 
         Args:
             axis: 0 or 1.
             ddof: Delta degree of freedom.
         """
-        return numojo.statistics.variance(self, axis=axis, ddof=ddof)
+        return numojo.statistics.variance[returned_dtype](
+            self, axis=axis, ddof=ddof
+        )
 
     # ===-------------------------------------------------------------------===#
     # To other data types

numojo/core/ndarray.mojo

@@ -56,7 +56,7 @@ import numojo.routines.math.rounding as rounding
 import numojo.routines.bitwise as bitwise
 import numojo.routines.linalg as linalg
 from numojo.core.datatypes import TypeCoercion, _concise_dtype_str
-from numojo.routines.statistics.averages import mean, cummean
+from numojo.routines.statistics.averages import mean
 from numojo.routines.math.products import prod, cumprod
 from numojo.routines.math.sums import sum, cumsum
 from numojo.routines.logic.truth import any
@@ -255,11 +255,15 @@ struct NDArray[dtype: DType = DType.float64](
 
     @always_inline("nodebug")
     fn __del__(owned self):
+        """
+        Destroys all elements in the list and free its memory.
+        """
         var owndata = True
         try:
             owndata = self.flags["OWNDATA"]
         except:
             print("Invalid `OWNDATA` flag. Treat as `True`.")
+
         if owndata:
             self._buf.ptr.free()
 
@@ -1850,6 +1854,13 @@ struct NDArray[dtype: DType = DType.float64](
     fn __pow__(self, p: Int) -> Self:
         return self._elementwise_pow(p)
 
+    fn __pow__(self, rhs: Scalar[dtype]) raises -> Self:
+        """Power of items."""
+        var res = self
+        for i in range(self.size):
+            res._buf.ptr[i] = self._buf.ptr[i].__pow__(rhs)
+        return res^
+
     fn __pow__(self, p: Self) raises -> Self:
         if self.size != p.size:
             raise Error(
@@ -3243,26 +3254,29 @@ struct NDArray[dtype: DType = DType.float64](
 
         return result
 
-    fn mean(self: Self, axis: Int) raises -> Self:
+    fn mean[
+        returned_dtype: DType = DType.float64
+    ](self: Self, axis: Int) raises -> NDArray[returned_dtype]:
         """
         Mean of array elements over a given axis.
+
         Args:
-            array: NDArray.
             axis: The axis along which the mean is performed.
+
         Returns:
             An NDArray.
 
         """
-        return mean(self, axis)
+        return mean[returned_dtype](self, axis)
 
     fn mean(self) raises -> Scalar[dtype]:
         """
-        Cumulative mean of a array.
+        Mean of a array.
 
         Returns:
-            The cumulative mean of the array as a SIMD Value of `dtype`.
+            The mean of the array as a SIMD Value of `dtype`.
         """
-        return cummean[dtype](self)
+        return mean[dtype](self)
 
     # fn nonzero(self):
     #     pass
@@ -3409,6 +3423,25 @@ struct NDArray[dtype: DType = DType.float64](
         var I = NDArray[DType.index](self.shape)
         sorting._sort_inplace(self, I, axis=axis)
 
+    fn std[
+        returned_dtype: DType = DType.float64
+    ](self, ddof: Int = 0) raises -> Scalar[returned_dtype]:
+        """
+        Compute the standard deviation.
+
+        Args:
+            ddof: Delta degree of freedom.
+        """
+
+        if ddof >= self.size:
+            raise Error(
+                String("ddof {} should be smaller than size {}").format(
+                    ddof, self.size
+                )
+            )
+
+        return variance[returned_dtype](self, ddof=ddof) ** 0.5
+
     fn sum(self: Self) raises -> Scalar[dtype]:
         """
         Sum of all array elements.
@@ -3506,6 +3539,20 @@ struct NDArray[dtype: DType = DType.float64](
         """
         return self._buf.ptr
 
+    fn variance[
+        returned_dtype: DType = DType.float64
+    ](self, ddof: Int = 0) raises -> Scalar[returned_dtype]:
+        """
+        Returns the variance of array.
+
+        Parameters:
+            returned_dtype: The returned data type, defaulting to float64.
+
+        Args:
+            ddof: Delta degree of freedom.
+        """
+        return variance[returned_dtype](self, ddof=ddof)
+
 
 # ===----------------------------------------------------------------------===#
 # NDArrayIterator

numojo/routines/math/sums.mojo

@@ -39,9 +39,7 @@ fn sum[dtype: DType](A: NDArray[dtype]) -> Scalar[dtype]:
     return res
 
 
-fn sum[
-    dtype: DType
-](A: NDArray[dtype], owned axis: Int) raises -> NDArray[dtype]:
+fn sum[dtype: DType](A: NDArray[dtype], axis: Int) raises -> NDArray[dtype]:
     """
     Returns sums of array elements over a given axis.
 
@@ -52,6 +50,10 @@ fn sum[
     print(nm.sum(A, axis=0))
     ```
 
+    Raises:
+        Error: If the axis is out of bound.
+        Error: If the number of dimensions is 1.
+
     Args:
         A: NDArray.
         axis: The axis along which the sum is performed.
@@ -60,25 +62,35 @@ fn sum[
         An NDArray.
     """
 
-    if axis < 0:
-        axis += A.ndim
-    if (axis < 0) or (axis >= A.ndim):
+    var normalized_axis = axis
+    if normalized_axis < 0:
+        normalized_axis += A.ndim
+
+    if (normalized_axis < 0) or (normalized_axis >= A.ndim):
         raise Error(
-            String("Invalid index: index out of bound [0, {}).").format(A.ndim)
+            String("Axis {} out of bound [0, {}).").format(axis, A.ndim)
+        )
+    if A.ndim == 1:
+        raise Error(
+            String(
+                "`numojo.routines.math.sums.sum()`: "
+                "Cannot sum over axis for 1-d array. "
+                "Please remove the `axis` argument."
+            )
         )
 
     var result_shape: List[Int] = List[Int]()
-    var size_of_axis: Int = A.shape[axis]
+    var size_of_axis: Int = A.shape[normalized_axis]
     var slices: List[Slice] = List[Slice]()
     for i in range(A.ndim):
-        if i != axis:
+        if i != normalized_axis:
             result_shape.append(A.shape[i])
             slices.append(Slice(0, A.shape[i]))
         else:
             slices.append(Slice(0, 0))  # Temp value
     var result = zeros[dtype](NDArrayShape(result_shape))
     for i in range(size_of_axis):
-        slices[axis] = Slice(i, i + 1)
+        slices[normalized_axis] = Slice(i, i + 1)
         var arr_slice = A[slices]
         result += arr_slice
 

numojo/routines/random.mojo

@@ -323,18 +323,7 @@ fn randn[
         a normal distribution with given mean and variance.
     """
 
-    builtin_random.seed()
-
-    var result: NDArray[dtype] = NDArray[dtype](shape)
-
-    builtin_random.randn[dtype](
-        ptr=result._buf.ptr,
-        size=result.size,
-        mean=mean.cast[DType.float64](),
-        variance=variance.cast[DType.float64](),
-    )
-
-    return result^
+    return randn[dtype](shape) * mt.sqrt(variance) + mean
 
 
 fn randn[

numojo/routines/statistics/__init__.mojo

@@ -1,15 +1,9 @@
 from .averages import (
     mean,
-    meanall,
     max,
     min,
-    cummean,
     mode,
     median,
     variance,
     std,
-    cumpvariance,
-    cumvariance,
-    cumpstdev,
-    cumstdev,
 )