[core][ComplexNDArray] Improve ComplexNDArray methods (#275)

## Pull Request Overview (From Copilot)

This PR enhances ComplexNDArray functionality by adding comparison
operators, trait methods, statistical/reduction methods, and array
manipulation capabilities. It also introduces temporary Int conversions
for strides/shape operations and implements SIMD load/store methods for
vectorized calculations.

### Key Changes
- Added trait implementations (ImplicitlyCopyable, Movable) and
conversion methods (__bool__, __int__, __float__) for ComplexNDArray
- Implemented magnitude-based comparison operators (__lt__, __le__,
__gt__, __ge__) for complex arrays
- Added statistical methods (all, any, sum, prod, mean, max, min,
argmax, argmin, cumsum, cumprod) and array manipulation methods
(flatten, fill, row, col, clip, round, T, diagonal, trace, tolist,
resize)
- Changed internal buffer types from `UnsafePointer[Int]` to
`UnsafePointer[Scalar[DType.int]]` in NDArrayShape, NDArrayStrides, and
Item structs
- Added SIMD load/store methods (load, store, unsafe_load, unsafe_store)
for Item, Shape, and Strides

<details>
<summary>Show a summary per file</summary>

| File | Description |
| ---- | ----------- |
| numojo/routines/indexing.mojo | Added Int conversions for stride
operations in compress function |
| numojo/routines/creation.mojo | Removed duplicate import statements |
| numojo/core/ndstrides.mojo | Changed buffer type to Scalar[DType.int],
updated __setitem__ validation, added SIMD load/store methods |
| numojo/core/ndshape.mojo | Changed buffer type to Scalar[DType.int],
updated __setitem__ validation, added SIMD load/store methods, modified
size_of_array calculation |
| numojo/core/ndarray.mojo | Added Int conversions for stride/shape
buffer accesses throughout |
| numojo/core/item.mojo | Changed buffer type to Scalar[DType.int],
removed Item.__init__(idx, shape) constructor and offset() method, added
SIMD load/store methods |
| numojo/core/complex/complex_simd.mojo | Added ImplicitlyCopyable and
Movable traits to ComplexSIMD |
| numojo/core/complex/complex_ndarray.mojo | Added comparison operators,
conversion methods, power operations, statistical methods, and array
manipulation methods; added Int conversions for stride operations |
</details>

---------

Co-authored-by: ZHU Yuhao 朱宇浩 <[email protected]>

Author: shivasankarka

numojo/core/complex/complex_ndarray.mojo

@@ -27,7 +27,9 @@ alias CScalar[cdtype: ComplexDType] = ComplexSIMD[cdtype, width=1]
 
 
 @register_passable("trivial")
-struct ComplexSIMD[cdtype: ComplexDType, width: Int = 1](Stringable, Writable):
+struct ComplexSIMD[cdtype: ComplexDType, width: Int = 1](
+    ImplicitlyCopyable, Movable, Stringable, Writable
+):
     """
     A SIMD-enabled complex number type that supports vectorized operations.
 

numojo/core/complex/complex_simd.mojo

@@ -29,7 +29,11 @@ struct Item(
     Specifies the indices of an item of an array.
     """
 
-    var _buf: UnsafePointer[Int]
+    # Aliases
+    alias _type: DType = DType.int
+
+    # Fields
+    var _buf: UnsafePointer[Scalar[Self._type]]
     var ndim: Int
 
     @always_inline("nodebug")
@@ -42,7 +46,7 @@ struct Item(
         Args:
             args: Initial values.
         """
-        self._buf = UnsafePointer[Int]().alloc(args.__len__())
+        self._buf = UnsafePointer[Scalar[Self._type]]().alloc(args.__len__())
         self.ndim = args.__len__()
         for i in range(args.__len__()):
             self._buf[i] = index(args[i])
@@ -58,7 +62,7 @@ struct Item(
             args: Initial values.
         """
         self.ndim = len(args)
-        self._buf = UnsafePointer[Int]().alloc(self.ndim)
+        self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim)
         for i in range(self.ndim):
             (self._buf + i).init_pointee_copy(index(args[i]))
 
@@ -70,7 +74,7 @@ struct Item(
             args: Initial values.
         """
         self.ndim = len(args)
-        self._buf = UnsafePointer[Int]().alloc(self.ndim)
+        self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim)
         for i in range(self.ndim):
             (self._buf + i).init_pointee_copy(Int(args[i]))
 
@@ -110,61 +114,14 @@ struct Item(
 
         if ndim == 0:
             self.ndim = 0
-            self._buf = UnsafePointer[Int]()
+            self._buf = UnsafePointer[Scalar[Self._type]]()
         else:
             self.ndim = ndim
-            self._buf = UnsafePointer[Int]().alloc(ndim)
+            self._buf = UnsafePointer[Scalar[Self._type]]().alloc(ndim)
             if initialized:
                 for i in range(ndim):
                     (self._buf + i).init_pointee_copy(0)
 
-    fn __init__(out self, idx: Int, shape: NDArrayShape) raises:
-        """
-        Get indices of the i-th item of the array of the given shape.
-        The item traverse the array in C-order.
-
-        Args:
-            idx: The i-th item of the array.
-            shape: The strides of the array.
-
-        Examples:
-
-        The following example demonstrates how to get the indices (coordinates)
-        of the 123-th item of a 3D array with shape (20, 30, 40).
-
-        ```console
-        >>> from numojo.prelude import *
-        >>> var item = Item(123, Shape(20, 30, 40))
-        >>> print(item)
-        Item at index: (0,3,3)  Length: 3
-        ```
-        """
-
-        if (idx < 0) or (idx >= shape.size_of_array()):
-            raise Error(
-                IndexError(
-                    message=String(
-                        "Linear index {} out of range [0, {})."
-                    ).format(idx, shape.size_of_array()),
-                    suggestion=String(
-                        "Ensure 0 <= idx < total size ({})."
-                    ).format(shape.size_of_array()),
-                    location=String(
-                        "Item.__init__(idx: Int, shape: NDArrayShape)"
-                    ),
-                )
-            )
-
-        self.ndim = shape.ndim
-        self._buf = UnsafePointer[Int]().alloc(self.ndim)
-
-        var strides = NDArrayStrides(shape, order="C")
-        var remainder = idx
-
-        for i in range(self.ndim):
-            (self._buf + i).init_pointee_copy(remainder // strides._buf[i])
-            remainder %= strides._buf[i]
-
     @always_inline("nodebug")
     fn __copyinit__(out self, other: Self):
         """Copy construct the tuple.
@@ -173,7 +130,7 @@ struct Item(
             other: The tuple to copy.
         """
         self.ndim = other.ndim
-        self._buf = UnsafePointer[Int]().alloc(self.ndim)
+        self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim)
         memcpy(self._buf, other._buf, self.ndim)
 
     @always_inline("nodebug")
@@ -232,7 +189,7 @@ struct Item(
                 )
             )
         var normalized_idx: Int = self.normalize_index(index_int(idx))
-        return self._buf[normalized_idx]
+        return Int(self._buf[normalized_idx])
 
     @always_inline("nodebug")
     fn __getitem__(self, slice_index: Slice) raises -> Self:
@@ -458,33 +415,6 @@ struct Item(
 
         return new_item^
 
-    fn offset(self, strides: NDArrayStrides) -> Int:
-        """
-        Calculates the offset of the item according to strides.
-
-        Args:
-            strides: The strides of the array.
-
-        Returns:
-            The offset of the item.
-
-        Examples:
-
-        ```mojo
-        from numojo.prelude import *
-        var item = Item(1, 2, 3)
-        var strides = nm.Strides(4, 3, 2)
-        print(item.offset(strides))
-        # This prints `16`.
-        ```
-        .
-        """
-
-        var offset: Int = 0
-        for i in range(self.ndim):
-            offset += self._buf[i] * strides._buf[i]
-        return offset
-
     # ===-------------------------------------------------------------------===#
     # Other private methods
     # ===-------------------------------------------------------------------===#
@@ -539,7 +469,7 @@ struct Item(
         if axis == self.ndim - 1:
             return result^
 
-        var value: Int = result._buf[axis]
+        var value: Scalar[Self._type] = result._buf[axis]
         for i in range(axis, result.ndim - 1):
             result._buf[i] = result._buf[i + 1]
         result._buf[result.ndim - 1] = value
@@ -670,6 +600,98 @@ struct Item(
 
         return (start, step, length)
 
+    fn load[width: Int = 1](self, idx: Int) raises -> SIMD[Self._type, width]:
+        """
+        Load a SIMD vector from the Item at the specified index.
+
+        Parameters:
+            width: The width of the SIMD vector.
+
+        Args:
+            idx: The starting index to load from.
+
+        Returns:
+            A SIMD vector containing the loaded values.
+
+        Raises:
+            Error: If the load exceeds the bounds of the Item.
+        """
+        if idx < 0 or idx + width > self.ndim:
+            raise Error(
+                IndexError(
+                    message=String(
+                        "Load operation out of bounds: idx={} width={} ndim={}"
+                    ).format(idx, width, self.ndim),
+                    suggestion=(
+                        "Ensure that idx and width are within valid range."
+                    ),
+                    location="Item.load",
+                )
+            )
+
+        return self._buf.load[width=width](idx)
+
+    fn store[
+        width: Int = 1
+    ](self, idx: Int, value: SIMD[Self._type, width]) raises:
+        """
+        Store a SIMD vector into the Item at the specified index.
+
+        Parameters:
+            width: The width of the SIMD vector.
+
+        Args:
+            idx: The starting index to store to.
+            value: The SIMD vector to store.
+
+        Raises:
+            Error: If the store exceeds the bounds of the Item.
+        """
+        if idx < 0 or idx + width > self.ndim:
+            raise Error(
+                IndexError(
+                    message=String(
+                        "Store operation out of bounds: idx={} width={} ndim={}"
+                    ).format(idx, width, self.ndim),
+                    suggestion=(
+                        "Ensure that idx and width are within valid range."
+                    ),
+                    location="Item.store",
+                )
+            )
+
+        self._buf.store[width=width](idx, value)
+
+    fn unsafe_load[width: Int = 1](self, idx: Int) -> SIMD[Self._type, width]:
+        """
+        Unsafely load a SIMD vector from the Item at the specified index.
+
+        Parameters:
+            width: The width of the SIMD vector.
+
+        Args:
+            idx: The starting index to load from.
+
+        Returns:
+            A SIMD vector containing the loaded values.
+        """
+        return self._buf.load[width=width](idx)
+
+    fn unsafe_store[
+        width: Int = 1
+    ](self, idx: Int, value: SIMD[Self._type, width]):
+        """
+        Unsafely store a SIMD vector into the Item at the specified index.
+
+        Parameters:
+            width: The width of the SIMD vector.
+
+        Args:
+            idx: The starting index to store to.
+            value: The SIMD vector to store.
+        """
+        self._buf.store[width=width](idx, value)
+
 
 struct _ItemIter[
     forward: Bool = True,

numojo/core/item.mojo

@@ -409,7 +409,7 @@ struct NDArray[dtype: DType = DType.float64](
         """
         var index_of_buffer: Int = 0
         for i in range(self.ndim):
-            index_of_buffer += indices[i] * self.strides._buf[i]
+            index_of_buffer += indices[i] * Int(self.strides._buf[i])
         return self._buf.ptr[index_of_buffer]
 
     fn _getitem(self, indices: List[Int]) -> Scalar[dtype]:
@@ -437,7 +437,7 @@ struct NDArray[dtype: DType = DType.float64](
         """
         var index_of_buffer: Int = 0
         for i in range(self.ndim):
-            index_of_buffer += indices[i] * self.strides._buf[i]
+            index_of_buffer += indices[i] * Int(self.strides._buf[i])
         return self._buf.ptr[index_of_buffer]
 
     fn __getitem__(self) raises -> SIMD[dtype, 1]:
@@ -636,15 +636,15 @@ struct NDArray[dtype: DType = DType.float64](
         for lin in range(total):
             var rem = lin
             for d in range(out_ndim - 1, -1, -1):
-                var dim = dst.shape._buf[d]
+                var dim = Int(dst.shape._buf[d])
                 coords[d] = rem % dim
                 rem //= dim
             var off = base
             for d in range(out_ndim):
                 off += coords[d] * src.strides._buf[d + 1]
             var dst_off = 0
             for d in range(out_ndim):
-                dst_off += coords[d] * dst.strides._buf[d]
+                dst_off += coords[d] * Int(dst.strides._buf[d])
             dst._buf.ptr[dst_off] = src._buf.ptr[off]
 
     fn __getitem__(self, var *slices: Slice) raises -> Self:
@@ -1839,7 +1839,7 @@ struct NDArray[dtype: DType = DType.float64](
         """
         var index_of_buffer: Int = 0
         for i in range(self.ndim):
-            index_of_buffer += indices[i] * self.strides._buf[i]
+            index_of_buffer += indices[i] * Int(self.strides._buf[i])
         self._buf.ptr[index_of_buffer] = val
 
     fn __setitem__(self, idx: Int, val: Self) raises:
@@ -1964,14 +1964,14 @@ struct NDArray[dtype: DType = DType.float64](
         for lin in range(total):
             var rem = lin
             for d in range(out_ndim - 1, -1, -1):
-                var dim = src.shape._buf[d]
+                var dim = Int(src.shape._buf[d])
                 coords[d] = rem % dim
                 rem //= dim
             var dst_off = base
             var src_off = 0
             for d in range(out_ndim):
-                var stride_src = src.strides._buf[d]
-                var stride_dst = dst.strides._buf[d + 1]
+                var stride_src = Int(src.strides._buf[d])
+                var stride_dst = Int(dst.strides._buf[d + 1])
                 var c = coords[d]
                 dst_off += c * stride_dst
                 src_off += c * stride_src
@@ -3815,7 +3815,7 @@ struct NDArray[dtype: DType = DType.float64](
         """
         Returns length of 0-th dimension.
         """
-        return self.shape._buf[0]
+        return Int(self.shape._buf[0])
 
     fn __iter__(
         self,
@@ -5851,7 +5851,7 @@ struct _NDIter[is_mutable: Bool, //, origin: Origin[is_mutable], dtype: DType](
                     (indices._buf + i).init_pointee_copy(
                         remainder // self.strides_compatible._buf[i]
                     )
-                    remainder %= self.strides_compatible._buf[i]
+                    remainder %= Int(self.strides_compatible._buf[i])
             (indices._buf + self.axis).init_pointee_copy(remainder)
 
         else:
@@ -5860,7 +5860,7 @@ struct _NDIter[is_mutable: Bool, //, origin: Origin[is_mutable], dtype: DType](
                     (indices._buf + i).init_pointee_copy(
                         remainder // self.strides_compatible._buf[i]
                     )
-                    remainder %= self.strides_compatible._buf[i]
+                    remainder %= Int(self.strides_compatible._buf[i])
             (indices._buf + self.axis).init_pointee_copy(remainder)
 
         return self.ptr[_get_offset(indices, self.strides)]
@@ -5893,15 +5893,15 @@ struct _NDIter[is_mutable: Bool, //, origin: Origin[is_mutable], dtype: DType](
                     (indices._buf + i).init_pointee_copy(
                         remainder // self.strides_compatible._buf[i]
                     )
-                    remainder %= self.strides_compatible._buf[i]
+                    remainder %= Int(self.strides_compatible._buf[i])
             (indices._buf + self.axis).init_pointee_copy(remainder)
         else:
             for i in range(self.ndim - 1, -1, -1):
                 if i != self.axis:
                     (indices._buf + i).init_pointee_copy(
                         remainder // self.strides_compatible._buf[i]
                     )
-                    remainder %= self.strides_compatible._buf[i]
+                    remainder %= Int(self.strides_compatible._buf[i])
             (indices._buf + self.axis).init_pointee_copy(remainder)
 
         return self.ptr[_get_offset(indices, self.strides)]