diff --git a/README.MD b/README.MD index b71af871..71ac1454 100644 --- a/README.MD +++ b/README.MD @@ -205,18 +205,18 @@ backend = {name = "pixi-build-mojo", version = "0.*"} name = "your_package_name" [package.host-dependencies] -modular = ">=25.5.0,<26" +modular = ">=25.7.0,<26" [package.build-dependencies] -modular = ">=25.5.0,<26" +modular = ">=25.7.0,<26" numojo = { git = "https://github.com/Mojo-Numerics-and-Algorithms-group/NuMojo", branch = "main"} [package.run-dependencies] -modular = ">=25.5.0,<26" +modular = ">=25.7.0,<26" numojo = { git = "https://github.com/Mojo-Numerics-and-Algorithms-group/NuMojo", branch = "main"} [dependencies] -max = "=25.5.0" +modular = ">=25.7.0,<26" numojo = { git = "https://github.com/Mojo-Numerics-and-Algorithms-group/NuMojo", branch = "main"} ``` @@ -227,7 +227,7 @@ pixi install **Branch Selection:** - **`main` branch**: Provides stable release. Currently supports NuMojo v0.7.0, compatible with Mojo 25.3.0. For earlier NuMojo versions, use Method 2. -- **`pre-x.y` branches**: Active development branch supporting the latest Mojo version (currently 25.5.0). Note that this branch receives frequent updates and may have breaking changes in features and syntax. +- **`pre-x.y` branches**: Active development branch supporting the latest Mojo version (currently 25.7.0). Note that this branch receives frequent updates and may have breaking changes in features and syntax. The package will be automatically available in your Pixi environment, and VSCode LSP will provide intelligent code hints. diff --git a/numojo/core/__init__.mojo b/numojo/core/__init__.mojo index b8cf4ec5..dcf4fc46 100644 --- a/numojo/core/__init__.mojo +++ b/numojo/core/__init__.mojo @@ -12,6 +12,7 @@ from .complex import ( ComplexSIMD, ComplexScalar, CScalar, + `1j`, ComplexNDArray, ComplexDType, ci8, diff --git a/numojo/core/complex/__init__.mojo b/numojo/core/complex/__init__.mojo index 243acad3..76dbbfff 100644 --- a/numojo/core/complex/__init__.mojo +++ b/numojo/core/complex/__init__.mojo @@ -1,4 +1,4 @@ -from .complex_simd import ComplexSIMD, ComplexScalar, CScalar +from .complex_simd import ComplexSIMD, ComplexScalar, CScalar, `1j` from .complex_ndarray import ComplexNDArray from .complex_dtype import ( ComplexDType, diff --git a/numojo/core/complex/complex_dtype.mojo b/numojo/core/complex/complex_dtype.mojo index 38c88c5a..aa2eca6c 100644 --- a/numojo/core/complex/complex_dtype.mojo +++ b/numojo/core/complex/complex_dtype.mojo @@ -86,16 +86,6 @@ struct ComplexDType( `ComplexDType` behaves like an enum rather than a typical object. You don't instantiate it, but instead use its compile-time constants (aliases) to declare data types for complex SIMD vectors, tensors, and other data structures. - - Example: - - ```mojo - import numojo as nm - var A = nm.CScalar[nm.cf32](re=1.0, im=2.0) - print("A:", A) # A: (1.0 + 2.0i) - var A1 = nm.ComplexSIMD[nm.cf32, 2](SIMD[nm.f32, 2](1.0, 1.0), SIMD[nm.f32, 2](2.0, 2.0)) - print("A1:", A1) # A1: ([1.0, 1.0], [2.0, 2.0] j) - ``` """ # ===-------------------------------------------------------------------===# @@ -219,6 +209,10 @@ struct ComplexDType( return Self._from_str(str.removeprefix("ComplexDType.")) elif str == "int8": return ComplexDType.int8 + elif str == "int16": + return ComplexDType.int16 + elif str == "int32": + return ComplexDType.int32 elif str == "int64": return ComplexDType.int64 elif str == "int128": @@ -476,7 +470,10 @@ struct ComplexDType( Returns: Returns True if the input type parameter is an integer. """ - return self in (DType.int, DType.uint) or self._is_non_index_integral() + return ( + self in (ComplexDType.int, ComplexDType.uint) + or self._is_non_index_integral() + ) @always_inline("nodebug") fn is_floating_point(self) -> Bool: @@ -596,6 +593,14 @@ struct ComplexDType( 2 * 8 * self.size_of() ) # 2 * because complex number has real and imaginary parts + fn component_bitwidth(self) -> Int: + """Returns the size in bits of the component type of the current ComplexDType. + + Returns: + Returns the size in bits of the component type of the current ComplexDType. + """ + return self.bitwidth() // 2 + # ===-------------------------------------------------------------------===# # __mlir_type # ===-------------------------------------------------------------------===# @@ -661,11 +666,18 @@ struct ComplexDType( return abort[__mlir_type.`!kgen.deferred`]("invalid dtype") + fn component_dtype(self) -> DType: + return self._dtype + fn _concise_dtype_str(cdtype: ComplexDType) -> String: """Returns a concise string representation of the complex data type.""" if cdtype == ci8: return "ci8" + elif cdtype == ci16: + return "ci16" + elif cdtype == ci32: + return "ci32" elif cdtype == ci64: return "ci64" elif cdtype == ci128: diff --git a/numojo/core/complex/complex_ndarray.mojo b/numojo/core/complex/complex_ndarray.mojo index 5d509303..c0518ceb 100644 --- a/numojo/core/complex/complex_ndarray.mojo +++ b/numojo/core/complex/complex_ndarray.mojo @@ -41,7 +41,7 @@ from builtin.type_aliases import Origin from collections.optional import Optional from math import log10, sqrt from memory import memset_zero, memcpy -from memory import LegacyUnsafePointer as UnsafePointer +from memory import LegacyUnsafePointer from python import PythonObject from sys import simd_width_of from utils import Variant @@ -102,28 +102,59 @@ import numojo.routines.searching as searching # Implements N-Dimensional Complex Array # ===----------------------------------------------------------------------=== # struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( - Copyable, Movable, Representable, Sized, Stringable, Writable + Copyable, + FloatableRaising, + IntableRaising, + Movable, + Representable, + Sized, + Stringable, + Writable, ): """ - Represents a Complex N-Dimensional Array. + N-dimensional Complex array. + + ComplexNDArray represents an N-dimensional array whose elements are complex numbers, supporting efficient storage, indexing, and mathematical operations. Each element consists of a real and imaginary part, stored in separate buffers. Parameters: - cdtype: Complex data type. + cdtype: The complex data type of the array elements (default: ComplexDType.float64). + + Attributes: + - _re: NDArray[Self.dtype] + Buffer for real parts. + - _im: NDArray[Self.dtype] + Buffer for imaginary parts. + - ndim: Int + Number of dimensions. + - shape: NDArrayShape + Shape of the array. + - size: Int + Total number of elements. + - strides: NDArrayStrides + Stride information for each dimension. + - flags: Flags + Memory layout information. + - print_options: PrintOptions + Formatting options for display. + + Notes: + - The array is uniquely defined by its data buffers, shape, strides, and element datatype. + - Supports both row-major (C) and column-major (F) memory order. + - Provides rich indexing, slicing, and broadcasting semantics. + - ComplexNDArray should be created using factory functions in `nomojo.routines.creation` module for convenience. """ - # ===----------------------------------------------------------------------===# - # Aliases - # ===----------------------------------------------------------------------===# - - alias dtype: DType = cdtype._dtype # corresponding real data type - - # ===----------------------------------------------------------------------===# - # FIELDS - # ===----------------------------------------------------------------------===# + # --- Aliases --- + alias dtype: DType = cdtype._dtype + """corresponding real data type""" + # --- FIELDS --- var _re: NDArray[Self.dtype] + """Buffer for real parts.""" var _im: NDArray[Self.dtype] - # It's redundant, but better to have the following as fields. + """Buffer for imaginary parts.""" + + # TODO: add methods to for users to access the following properties directly from _re, _im and remove them from here. var ndim: Int """Number of Dimensions.""" var shape: NDArrayShape @@ -135,11 +166,9 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( var flags: Flags "Information about the memory layout of the array." var print_options: PrintOptions + """Per-instance print options (formerly global).""" - # ===-------------------------------------------------------------------===# - # Life cycle methods - # ===-------------------------------------------------------------------===# - + # --- Life cycle methods --- @always_inline("nodebug") fn __init__( out self, var re: NDArray[Self.dtype], var im: NDArray[Self.dtype] @@ -183,19 +212,20 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( order: String = "C", ) raises: """ - Initialize a ComplexNDArray with given shape. - - The memory is not filled with values. + Initialize a ComplexNDArray with given shape. The memory is not filled with values. Args: shape: Variadic shape. order: Memory order C or F. Example: - ```mojo - from numojo.prelude import * - var A = nm.ComplexNDArray[cf32](Shape(2,3,4)) - ``` + ```mojo + from numojo.prelude import * + var A = nm.ComplexNDArray[cf32](Shape(2,3,4)) + ``` + + Notes: + This constructor should not be used by users directly. Use factory functions in `numojo.routines.creation` module instead. """ self._re = NDArray[Self.dtype](shape, order) self._im = NDArray[Self.dtype](shape, order) @@ -220,6 +250,15 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( Args: shape: List of shape. order: Memory order C or F. + + Example: + ```mojo + from numojo.prelude import * + var A = nm.ComplexNDArray[cf32](List[Int](2,3,4)) + ``` + + Notes: + This constructor should not be used by users directly. Use factory functions in `numojo.routines.creation` module instead. """ self._re = NDArray[Self.dtype](shape, order) self._im = NDArray[Self.dtype](shape, order) @@ -244,6 +283,15 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( Args: shape: Variadic List of shape. order: Memory order C or F. + + Example: + ```mojo + from numojo.prelude import * + var A = nm.ComplexNDArray[cf32](VariadicList(2,3,4)) + ``` + + Notes: + This constructor should not be used by users directly. Use factory functions in `numojo.routines.creation` module instead. """ self._re = NDArray[Self.dtype](shape, order) self._im = NDArray[Self.dtype](shape, order) @@ -263,7 +311,26 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( strides: List[Int], ) raises: """ - Extremely specific ComplexNDArray initializer. + Initialize a ComplexNDArray with a specific shape, offset, and strides. + + Args: + shape: List of integers specifying the shape of the array. + offset: Integer offset into the underlying buffer. + strides: List of integers specifying the stride for each dimension. + + Example: + ```mojo + from numojo.prelude import * + var shape = List[Int](2, 3) + var offset = 0 + var strides = List[Int](3, 1) + var arr = ComplexNDArray[cf32](shape, offset, strides) + ``` + + Notes: + - This constructor is intended for advanced use cases requiring precise control over memory layout. + - The resulting array is uninitialized and should be filled before use. + - Both real and imaginary buffers are created with the same shape, offset, and strides. """ self._re = NDArray[Self.dtype](shape, offset, strides) self._im = NDArray[Self.dtype](shape, offset, strides) @@ -285,16 +352,19 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( flags: Flags, ): """ - Constructs an extremely specific ComplexNDArray, with value uninitialized. - The properties do not need to be compatible and are not checked. - For example, it can construct a 0-D array (numojo scalar). + Initialize a ComplexNDArray with explicit shape, strides, number of dimensions, size, and flags. This constructor creates an uninitialized ComplexNDArray with the provided properties. No compatibility checks are performed between shape, strides, ndim, size, or flags. This allows construction of arrays with arbitrary metadata, including 0-D arrays (scalars). Args: - shape: Shape of array. - strides: Strides of array. + shape: Shape of the array. + strides: Strides for each dimension. ndim: Number of dimensions. - size: Size of array. - flags: Flags of array. + size: Total number of elements. + flags: Memory layout flags. + + Notes: + - This constructor is intended for advanced or internal use cases requiring manual control. + - The resulting array is uninitialized; values must be set before use. + - No validation is performed on the consistency of the provided arguments. """ self.shape = shape @@ -311,21 +381,32 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( fn __init__( out self, shape: NDArrayShape, - ref buffer_re: UnsafePointer[Scalar[Self.dtype]], - ref buffer_im: UnsafePointer[Scalar[Self.dtype]], + ref buffer_re: LegacyUnsafePointer[Scalar[Self.dtype]], + ref buffer_im: LegacyUnsafePointer[Scalar[Self.dtype]], offset: Int, strides: NDArrayStrides, ) raises: """ - Initialize an ComplexNDArray view with given shape, buffer, offset, and strides. - ***Unsafe!*** This function is currently unsafe. Only for internal use. + Initialize a ComplexNDArray view with explicit shape, raw buffers, offset, and strides. + + This constructor creates a view over existing memory buffers for the real and imaginary parts, + using the provided shape, offset, and stride information. It is intended for advanced or internal + use cases where direct control over memory layout is required. + + ***Unsafe!*** This function is unsafe and should only be used internally. The caller is responsible + for ensuring that the buffers are valid and that the shape, offset, and strides are consistent. Args: - shape: Shape of the array. - buffer_re: Unsafe pointer to the real part of the buffer. - buffer_im: Unsafe pointer to the imaginary part of the buffer. - offset: Offset value. - strides: Strides of the array. + shape: NDArrayShape specifying the dimensions of the array. + buffer_re: Unsafe pointer to the buffer containing the real part data. + buffer_im: Unsafe pointer to the buffer containing the imaginary part data. + offset: Integer offset into the buffers. + strides: NDArrayStrides specifying the stride for each dimension. + + Notes: + - No validation is performed on the buffers or metadata. + - The resulting ComplexNDArray shares memory with the provided buffers. + - Incorrect usage may lead to undefined behavior. """ self._re = NDArray(shape, buffer_re, offset, strides) self._im = NDArray(shape, buffer_im, offset, strides) @@ -366,18 +447,10 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( self.flags = existing.flags self.print_options = existing.print_options - # Explicit deallocation - # @always_inline("nodebug") - # fn __del__(var self): - # """ - # Deallocate memory. - # """ - # self._re.__del__() - # self._im.__del__() - # ===-------------------------------------------------------------------===# # Indexing and slicing # Getter dunders and other getter methods + # FIXME: currently most of the getitem and setitem methods don't match exactly between NDArray and ComplexNDArray in it's implementation, docstring, argument mutability etc. Fix this. # 1. Basic Indexing Operations # fn _getitem(self, *indices: Int) -> ComplexSIMD[cdtype] # Direct unsafe getter @@ -407,6 +480,32 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( # fn load[width: Int](self, *indices: Int) raises -> ComplexSIMD[Self.dtype, width] # Load SIMD at coordinates # ===-------------------------------------------------------------------===# + @always_inline + fn normalize(self, idx: Int, dim: Int) -> Int: + """ + Normalize a potentially negative index to its positive equivalent + within the bounds of the given dimension. + + Args: + idx: The index to normalize. Can be negative to indicate indexing + from the end (e.g., -1 refers to the last element). + dim: The size of the dimension to normalize against. + + Returns: + The normalized index as a non-negative integer. + + Example: + ```mojo + from numojo.prelude import * + var mat = Matrix[f32](shape=(3, 4)) + var norm_idx = mat.normalize(-1, mat.shape[0]) # Normalize -1 to 2 + ``` + """ + var idx_norm = idx + if idx_norm < 0: + idx_norm = dim + idx_norm + return idx_norm + fn _getitem(self, *indices: Int) -> ComplexSIMD[cdtype]: """ Get item at indices and bypass all boundary checks. @@ -418,23 +517,22 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( Returns: The element of the array at the indices. - Notes: - This function is unsafe and should be used only on internal use. - Examples: + ```mojo + import numojo as nm + var A = nm.ones[nm.cf32](nm.Shape(2,3,4)) + print(A._getitem(1,2,3)) + ``` - ```mojo - import numojo as nm - var A = nm.ones[nm.cf32](nm.Shape(2,3,4)) - print(A._getitem(1,2,3)) - ``` + Notes: + This function is unsafe and should be used only on internal use. """ var index_of_buffer: Int = 0 for i in range(self.ndim): index_of_buffer += indices[i] * Int(self.strides._buf[i]) return ComplexSIMD[cdtype]( - re=self._re._buf.ptr.load[width=1](index_of_buffer), - im=self._im._buf.ptr.load[width=1](index_of_buffer), + re=self._re._buf.ptr[index_of_buffer], + im=self._im._buf.ptr[index_of_buffer], ) fn _getitem(self, indices: List[Int]) -> ComplexScalar[cdtype]: @@ -448,26 +546,25 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( Returns: The element of the array at the indices. - Notes: - This function is unsafe and should be used only on internal use. - Examples: + ```mojo + import numojo as nm + var A = nm.ones[nm.cf32](numojo.Shape(2,3,4)) + print(A._getitem(List[Int](1,2,3))) + ``` - ```mojo - import numojo as nm - var A = nm.ones[nm.cf32](numojo.Shape(2,3,4)) - print(A._getitem(List[Int](1,2,3))) - ``` + Notes: + This function is unsafe and should be used only on internal use. """ var index_of_buffer: Int = 0 for i in range(self.ndim): index_of_buffer += indices[i] * Int(self.strides._buf[i]) return ComplexSIMD[cdtype]( - re=self._re._buf.ptr.load[width=1](index_of_buffer), - im=self._im._buf.ptr.load[width=1](index_of_buffer), + re=self._re._buf.ptr[index_of_buffer], + im=self._im._buf.ptr[index_of_buffer], ) - fn __getitem__(self) raises -> ComplexSIMD[cdtype]: + fn __getitem__(self) raises -> ComplexSIMD[cdtype, 1]: """ Gets the value of the 0-D Complex array. @@ -479,10 +576,10 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( Examples: - ```console - >>> import numojo as nm - >>> var A = nm.ones[nm.f32](nm.Shape(2,3,4)) - >>> print(A[]) # gets values of the 0-D array. + ```mojo + import numojo as nm + var a = nm.arange[nm.cf32](3)[0] + print(a[]) # gets values of the 0-D complex array. ```. """ if self.ndim != 0: @@ -504,7 +601,7 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( im=self._im._buf.ptr[], ) - fn __getitem__(self, index: Item) raises -> ComplexSIMD[cdtype]: + fn __getitem__(self, index: Item) raises -> ComplexSIMD[cdtype, 1]: """ Get the value at the index list. @@ -587,15 +684,15 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( stride-based copier is used for both components. (Future: return a non-owning view). - Examples: + Example: ```mojo import numojo as nm from numojo.prelude import * - var a = nm.arange[cf32](CScalar[cf32](0, 0), CScalar[cf32](12, 12), CScalar[cf32](1, 1)).reshape(nm.Shape(3, 4)) + var a = nm.arange[cf32](CScalar[cf32](0), CScalar[cf32](12), CScalar[cf32](1)).reshape(Shape(3, 4)) print(a.shape) # (3,4) print(a[1].shape) # (4,) -- 1-D slice print(a[-1].shape) # (4,) -- negative index - var b = nm.arange[cf32](CScalar[cf32](6, 6)).reshape(nm.Shape(6)) + var b = nm.arange[cf32](CScalar[cf32](6)).reshape(nm.Shape(6)) print(b[2]) # 0-D array (scalar wrapper) ``` """ @@ -638,15 +735,17 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( ) ) - var out_shape = self.shape[1:] - var alloc_order = String("C") + var out_shape: NDArrayShape = self.shape[1:] + var alloc_order: String = String("C") if self.flags.F_CONTIGUOUS: alloc_order = String("F") - var result = ComplexNDArray[cdtype](shape=out_shape, order=alloc_order) + var result: ComplexNDArray[cdtype] = ComplexNDArray[cdtype]( + shape=out_shape, order=alloc_order + ) # Fast path for C-contiguous if self.flags.C_CONTIGUOUS: - var block = self.size // self.shape[0] + var block: Int = self.size // self.shape[0] memcpy( dest=result._re._buf.ptr, src=self._re._buf.ptr + norm * block, @@ -658,11 +757,15 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( count=block, ) return result^ - - # F layout - self[Self.dtype]._re._copy_first_axis_slice(self._re, norm, result._re) - self[Self.dtype]._im._copy_first_axis_slice(self._im, norm, result._im) - return result^ + else: + # F layout + self[Self.dtype]._re._copy_first_axis_slice( + self._re, norm, result._re + ) + self[Self.dtype]._im._copy_first_axis_slice( + self._im, norm, result._im + ) + return result^ fn __getitem__(self, var *slices: Slice) raises -> Self: """ @@ -720,7 +823,7 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( var narr: Self = self[slice_list^] return narr^ - fn _calculate_strides_efficient(self, shape: List[Int]) -> List[Int]: + fn _calculate_strides(self, shape: List[Int]) -> List[Int]: var strides = List[Int](capacity=len(shape)) if self.flags.C_CONTIGUOUS: # C_CONTIGUOUS @@ -761,7 +864,7 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( - This method supports advanced slicing similar to NumPy's multi-dimensional slicing. - The returned array shares data with the original array if possible. - Examples: + Example: ```mojo import numojo as nm from numojo.prelude import * @@ -821,12 +924,13 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( ncoefficients.append(1) # only C & F order are supported - var nstrides: List[Int] = self._calculate_strides_efficient( + var nstrides: List[Int] = self._calculate_strides( nshape, ) var narr = ComplexNDArray[cdtype]( offset=noffset, shape=nshape, strides=nstrides ) + # TODO: combine the two traverses into one. var index_re: List[Int] = List[Int](length=ndims, fill=0) _traverse_iterative[Self.dtype]( self._re, @@ -956,7 +1060,7 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( var shape = indices.shape.join(self.shape._pop(0)) var result: ComplexNDArray[cdtype] = ComplexNDArray[cdtype](shape) - var size_per_item = self.size // self.shape[0] + var size_per_item: Int = self.size // self.shape[0] # Fill in the values for i in range(indices.size): @@ -1193,8 +1297,7 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( ) ) - if index < 0: - index += self.size + index = self.normalize(index, self.size) if (index < 0) or (index >= self.size): raise Error( @@ -1321,8 +1424,7 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( ```. """ - if index < 0: - index += self.size + index = self.normalize(index, self.size) if (index >= self.size) or (index < 0): raise Error( @@ -1343,7 +1445,9 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( im=self._im._buf.ptr[index], ) - fn load[width: Int = 1](self, index: Int) raises -> ComplexSIMD[cdtype]: + fn load[ + width: Int = 1 + ](self, index: Int) raises -> ComplexSIMD[cdtype, width]: """ Safely loads a ComplexSIMD element of size `width` at `index` from the underlying buffer. @@ -1373,7 +1477,7 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( ) ) - return ComplexSIMD[cdtype]( + return ComplexSIMD[cdtype, width]( re=self._re._buf.ptr.load[width=1](index), im=self._im._buf.ptr.load[width=1](index), ) @@ -1422,23 +1526,31 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( ) ) + # NOTE: if we take in an owned instances of indices, we can modify it in place. + var indices_list: List[Int] = List[Int](capacity=self.ndim) for i in range(self.ndim): - if (indices[i] < 0) or (indices[i] >= self.shape[i]): + var idx_i = indices[i] + if idx_i < 0 or idx_i >= self.shape[i]: raise Error( IndexError( message=String( - "Index {} out of range for dim {} (size {})." - ).format(indices[i], i, self.shape[i]), + "Index out of range at dim {}: got {}; valid range" + " is [0, {})." + ).format(i, idx_i, self.shape[i]), suggestion=String( - "Valid range for dim {} is [0, {})." - ).format(i, self.shape[i]), + "Clamp or validate indices against the dimension" + " size ({})." + ).format(self.shape[i]), location=String( - "ComplexNDArray.load[width](*indices: Int)" + "NDArray.load[width: Int = 1](*indices: Int) ->" + " SIMD[dtype, width]" ), ) ) + idx_i = self.normalize(idx_i, self.shape[i]) + indices_list.append(idx_i) - var idx: Int = _get_offset(indices, self.strides) + var idx: Int = _get_offset(indices_list, self.strides) return ComplexSIMD[cdtype, width=width]( re=self._re._buf.ptr.load[width=width](idx), im=self._im._buf.ptr.load[width=width](idx), @@ -1516,7 +1628,7 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( end = slice_list[i].end.value() if end < 0: end += dim_size - # Clamp to valid bounds once + # NOTE: Clamp to valid bounds once. This is an implicit behavior right now instead of raising errors. not sure if this should be kept. if step > 0: end = 0 if end < 0 else ( dim_size if end > dim_size else end @@ -1539,7 +1651,23 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( fn _setitem(self, *indices: Int, val: ComplexSIMD[cdtype]): """ (UNSAFE! for internal use only.) - Get item at indices and bypass all boundary checks. + Set item at indices and bypass all boundary checks. + + Args: + indices: Indices to set the value. + val: Value to set. + + Notes: + This function is unsafe and for internal use only. + + Examples: + + ```mojo + import numojo as nm + from numojo.prelude import * + var A = nm.full[cf32](Shape(2, 2), CScalar[cf32](1.0, 1.0)) + A._setitem(0, 1, val=CScalar[cf32](3.0, 4.0)) + ``` """ var index_of_buffer: Int = 0 for i in range(self.ndim): @@ -1579,8 +1707,7 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( ) var norm = idx - if norm < 0: - norm += self.shape[0] + norm = self.normalize(norm, self.shape[0]) if (norm < 0) or (norm >= self.shape[0]): raise Error( IndexError( @@ -1617,21 +1744,6 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( self._im._buf.ptr.store(norm, val._im._buf.ptr.load[width=1](0)) return - if val.ndim != self.ndim - 1: - raise Error( - ShapeError( - message=String( - "Shape mismatch: expected {} dims in source but got {}." - ).format(self.ndim - 1, val.ndim), - suggestion=String("Ensure RHS has shape {}.").format( - self.shape[1:] - ), - location=String( - "ComplexNDArray.__setitem__(idx: Int, val: Self)" - ), - ) - ) - if val.shape != self.shape[1:]: raise Error( ShapeError( @@ -1651,21 +1763,6 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( if self.flags.C_CONTIGUOUS & val.flags.C_CONTIGUOUS: var block = self.size // self.shape[0] - if val.size != block: - raise Error( - ShapeError( - message=String( - "Internal mismatch: computed block {} but" - " val.size {}." - ).format(block, val.size), - suggestion=String( - "Report this issue; unexpected size mismatch." - ), - location=String( - "ComplexNDArray.__setitem__(idx: Int, val: Self)" - ), - ) - ) memcpy( dest=self._re._buf.ptr + norm * block, src=val._re._buf.ptr, @@ -1682,27 +1779,62 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( self[Self.dtype]._re._write_first_axis_slice(self._re, norm, val._re) self[Self.dtype]._im._write_first_axis_slice(self._im, norm, val._im) - fn __setitem__(mut self, index: Item, val: ComplexSIMD[cdtype]) raises: + fn __setitem__(mut self, var index: Item, val: ComplexSIMD[cdtype]) raises: """ - Set the value at the index list. + Sets the value at the index list. + + Args: + index: Index list. + val: Value to set. + + Raises: + Error: If the length of index does not match the number of dimensions. + Error: If any of the indices is out of bound. + + Examples: + + ```mojo + import numojo as nm + from numojo.prelude import * + var A = nm.full[cf32](Shape(2, 2), CScalar[cf32](1.0)) + A[Item(0, 1)] = CScalar[cf32](3.0, 4.0) + ``` """ if index.__len__() != self.ndim: - var message = String( - "Error: Length of `index` does not match the number of" - " dimensions!\n" - "Length of indices is {}.\n" - "The array dimension is {}." - ).format(index.__len__(), self.ndim) - raise Error(message) + raise Error( + IndexError( + message=String( + "Invalid index length: expected {} but got {}." + ).format(self.ndim, index.__len__()), + suggestion=String( + "Pass exactly {} indices (one per dimension)." + ).format(self.ndim), + location=String( + "ComplexNDArray.__setitem__(index: Item, val:" + " Scalar[dtype])" + ), + ) + ) for i in range(index.__len__()): if index[i] >= self.shape[i]: - var message = String( - "Error: `index` exceeds the size!\n" - "For {}-th dimension:\n" - "The index value is {}.\n" - "The size of the corresponding dimension is {}" - ).format(i, index[i], self.shape[i]) - raise Error(message) + raise Error( + IndexError( + message=String( + "Index out of range at dim {}: got {}; valid range" + " is [0, {})." + ).format(i, index[i], self.shape[i]), + suggestion=String( + "Clamp or validate indices against the dimension" + " size ({})." + ).format(self.shape[i]), + location=String( + "NDArray.__setitem__(index: Item, val:" + " Scalar[dtype])" + ), + ) + ) + index[i] = self.normalize(index[i], self.shape[i]) + var idx: Int = _get_offset(index, self.strides) self._re._buf.ptr.store(idx, val.re) self._im._buf.ptr.store(idx, val.im) @@ -1726,7 +1858,9 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( if mask._im._buf.ptr.load[width=1](i): self._im._buf.ptr.store(i, value.im) - fn __setitem__(mut self, var *slices: Slice, val: Self) raises: + fn __setitem__( + mut self, var *slices: Slice, val: ComplexNDArray[cdtype] + ) raises: """ Retreive slices of an ComplexNDArray from variadic slices. @@ -1739,7 +1873,9 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( # self.__setitem__(slices=slice_list, val=val) self[slice_list^] = val - fn __setitem__(mut self, var slices: List[Slice], val: Self) raises: + fn __setitem__( + mut self, slices: List[Slice], val: ComplexNDArray[cdtype] + ) raises: """ Sets the slices of an ComplexNDArray from list of slices and ComplexNDArray. @@ -1752,6 +1888,7 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( var spec: List[Int] = List[Int]() var slice_list: List[Slice] = self._adjust_slice(slices) for i in range(n_slices): + # TODO: these conditions can be removed since _adjust_slice takes care of them. But verify it once before removing. if ( slice_list[i].start.value() >= self.shape[i] or slice_list[i].end.value() > self.shape[i] @@ -1845,31 +1982,48 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( val._im, self._im, nshape, ncoefficients, nstrides, noffset, index ) - ### compiler doesn't accept this. - # fn __setitem__(self, var *slices: Variant[Slice, Int], val: NDArray[Self.dtype]) raises: - # """ - # Get items by a series of either slices or integers. - # """ - # var n_slices: Int = slices.__len__() - # if n_slices > self.ndim: - # raise Error("Error: No of slices greater than rank of array") - # var slice_list: List[Slice] = List[Slice]() - - # var count_int = 0 - # for i in range(len(slices)): - # if slices[i].isa[Slice](): - # slice_list.append(slices[i]._get_ptr[Slice]()[0]) - # elif slices[i].isa[Int](): - # count_int += 1 - # var int: Int = slices[i]._get_ptr[Int]()[0] - # slice_list.append(Slice(int, int + 1)) - - # if n_slices < self.ndim: - # for i in range(n_slices, self.ndim): - # var size_at_dim: Int = self.shape[i] - # slice_list.append(Slice(0, size_at_dim)) - - # self.__setitem__(slices=slice_list, val=val) + ## compiler doesn't accept this. + fn __setitem__( + self, var *slices: Variant[Slice, Int], val: ComplexNDArray[cdtype] + ) raises: + """ + Get items by a series of either slices or integers. + """ + var n_slices: Int = slices.__len__() + if n_slices > self.ndim: + raise Error( + IndexError( + message=String( + "Too many indices or slices: received {} but array has" + " only {} dimensions." + ).format(n_slices, self.ndim), + suggestion=String( + "Pass at most {} indices/slices (one per dimension)." + ).format(self.ndim), + location=String( + "NDArray.__setitem__(*slices: Variant[Slice, Int], val:" + " Self)" + ), + ) + ) + var slice_list: List[Slice] = List[Slice]() + + var count_int = 0 + for i in range(len(slices)): + if slices[i].isa[Slice](): + slice_list.append(slices[i]._get_ptr[Slice]()[0]) + elif slices[i].isa[Int](): + count_int += 1 + var int: Int = slices[i]._get_ptr[Int]()[0] + slice_list.append(Slice(int, int + 1, 1)) + + if n_slices < self.ndim: + for i in range(n_slices, self.ndim): + var size_at_dim: Int = self.shape[i] + slice_list.append(Slice(0, size_at_dim, 1)) + + # self.__setitem__(slices=slice_list, val=val) + self[slice_list^] = val fn __setitem__(self, index: NDArray[DType.int], val: Self) raises: """ @@ -1886,6 +2040,7 @@ struct ComplexNDArray[cdtype: ComplexDType = ComplexDType.float64]( Int(index.load(i)), rebind[Scalar[Self.dtype]](val._im.load(i)) ) + # TODO: implement itemset(). fn __setitem__( mut self, mask: ComplexNDArray[cdtype], @@ -4070,8 +4225,8 @@ struct _ComplexNDArrayIter[ # FIELDS var index: Int - var re_ptr: UnsafePointer[Scalar[Self.dtype]] - var im_ptr: UnsafePointer[Scalar[Self.dtype]] + var re_ptr: LegacyUnsafePointer[Scalar[Self.dtype]] + var im_ptr: LegacyUnsafePointer[Scalar[Self.dtype]] var dimension: Int var length: Int var shape: NDArrayShape diff --git a/numojo/core/complex/complex_simd.mojo b/numojo/core/complex/complex_simd.mojo index 9c23f775..96e67c7b 100644 --- a/numojo/core/complex/complex_simd.mojo +++ b/numojo/core/complex/complex_simd.mojo @@ -11,62 +11,77 @@ This module provides a ComplexSIMD type that represents complex numbers using SI operations for efficient computation. It supports basic arithmetic operations like addition, subtraction, multiplication, and division, as well as other complex number operations like conjugation and absolute value. - -The implementation allows for vectorized operations on complex numbers which can -significantly improve performance for numerical computations. """ -from math import sqrt - +from math import sqrt, sin, cos from numojo.core.complex.complex_dtype import ComplexDType -# ComplexScalar alias is for internal purposes -alias ComplexScalar[cdtype: ComplexDType] = ComplexSIMD[cdtype, width=1] -# CScalar is short alias for ComplexScalar for user convenience -alias CScalar[cdtype: ComplexDType] = ComplexSIMD[cdtype, width=1] +alias ComplexScalar = ComplexSIMD[_, width=1] +"""ComplexScalar alias is for internal purposes (width=1 specialization).""" + +alias CScalar = ComplexSIMD[_, width=1] +"""User-friendly alias for scalar complex numbers.""" + +alias `1j` = ImaginaryUnit() +"""Constant representing the imaginary unit complex number 0 + 1j. +Enables Python like syntax for complex numbers, e.g., (3 + 4 * `1j`).""" +# TODO: add overloads for arithmetic functions to accept Scalar[dtype]. @register_passable("trivial") -struct ComplexSIMD[cdtype: ComplexDType, width: Int = 1]( +struct ComplexSIMD[cdtype: ComplexDType = ComplexDType.float64, width: Int = 1]( ImplicitlyCopyable, Movable, Stringable, Writable ): """ - A SIMD-enabled complex number type that supports vectorized operations. - - Parameters: - cdtype: The complex data type (like cf32 or cf64) that determines precision. - width: The SIMD vector width, defaulting to 1 for scalar operations. - - The struct contains two SIMD vectors - one for the real part and one for the - imaginary part. This allows complex arithmetic to be performed efficiently using - SIMD operations. When width=1 it acts as a regular complex scalar type. - - Example: - ```mojo - import numojo as nm - var A = nm.ComplexSIMD[nm.cf32](1.0, 2.0) - var B = nm.ComplexSIMD[nm.cf32](3.0, 4.0) - var C = A + B - print(C) # Output: (4.0 + 6.0 j) - - var A1 = nm.ComplexSIMD[nm.cf32, 2](SIMD[nm.f32](1.0, 1.0), SIMD[nm.f32](2.0, 2.0)) - print(A1) # Output: ([1.0, 1.0] + [2.0, 2.0] j) - ``` + A SIMD-enabled complex number container (SoA layout). + + Fields: + re: SIMD vector of real parts. + im: SIMD vector of imaginary parts. + + The parameter `cdtype` determines the component precision (e.g. cf32, cf64). + The parameter `width` is the SIMD lane count; when `width == 1` this acts like a scalar complex number. + + Examples: + ```mojo + from numojo.prelude import * + var a = ComplexSIMD[cf32](1.0, 2.0) + var b = ComplexSIMD[cf32](3.0, 4.0) + print(a + b) # (4.0 + 6.0 j) + + # SIMD width=2: + var a2 = ComplexSIMD[cf32, 2]( + SIMD[cf32._dtype, 2](1.0, 1.5), + SIMD[cf32._dtype, 2](2.0, -0.5) + ) + print(a2) # ( [1.0 2.0] + [1.5 -0.5]j ) + ``` + Convenience factories: + ComplexSIMD[cf64].zero() + ComplexSIMD[cf64].one() + ComplexSIMD[cf64].i() + ComplexSIMD[cf64].from_polar(2.0, 0.5) """ - # FIELDS - alias dtype: DType = cdtype._dtype # the corresponding DType - # The underlying data real and imaginary parts of the complex number. + alias dtype: DType = cdtype._dtype + """Component dtype alias (underlying real/imag dtype).""" + var re: SIMD[Self.dtype, width] var im: SIMD[Self.dtype, width] + # --- Internal helper for broadcasting scalar to SIMD lanes --- + @staticmethod + @always_inline + fn _broadcast(val: Scalar[Self.dtype]) -> SIMD[Self.dtype, Self.width]: + return SIMD[Self.dtype, Self.width](val) + + # --- Constructors --- @always_inline fn __init__(out self, other: Self): """ - Initializes a ComplexSIMD instance by copying another instance. + Copy constructor for ComplexSIMD. - Arguments: - other: Another ComplexSIMD instance to copy from. + Initializes a new ComplexSIMD instance by copying the values from another instance. """ self = other @@ -77,20 +92,11 @@ struct ComplexSIMD[cdtype: ComplexDType, width: Int = 1]( im: SIMD[Self.dtype, Self.width], ): """ - Initializes a ComplexSIMD instance with specified real and imaginary parts. - - Arguments: - re: The real part of the complex number. - im: The imaginary part of the complex number. + Constructs a ComplexSIMD from SIMD vectors of real and imaginary parts. - Example: - ```mojo - import numojo as nm - var A = nm.ComplexSIMD[nm.cf32](1.0, 2.0) - var B = nm.ComplexSIMD[nm.cf32](3.0, 4.0) - var C = A + B - print(C) - ``` + Args: + re: SIMD vector containing the real components. + im: SIMD vector containing the imaginary components. """ self.re = re self.im = im @@ -98,94 +104,582 @@ struct ComplexSIMD[cdtype: ComplexDType, width: Int = 1]( @always_inline fn __init__(out self, val: SIMD[Self.dtype, Self.width]): """ - Initializes a ComplexSIMD instance with specified real and imaginary parts. + Constructs a ComplexSIMD where both real and imaginary parts are set to the same SIMD value. - Arguments: - re: The real part of the complex number. - im: The imaginary part of the complex number. + Args: + val: SIMD vector to broadcast to both real and imaginary components. """ self.re = val self.im = val + # Factory constructors. + @staticmethod + fn zero() -> Self: + """ + Returns a ComplexSIMD instance with all real and imaginary components set to zero. + + Example: + ```mojo + from numojo.prelude import * + var comp = ComplexSIMD[cf64].zero() # (0 + 0j) + ``` + """ + return Self(Self._broadcast(0), Self._broadcast(0)) + + @staticmethod + fn one() -> Self: + """ + Returns a ComplexSIMD instance representing the complex number 1 + 0j. + + Example: + ```mojo + from numojo.prelude import * + var comp = ComplexSIMD[cf64].one() # (1 + 0j) + ``` + """ + return Self(Self._broadcast(1), Self._broadcast(0)) + + @staticmethod + fn i() -> Self: + """ + Returns a ComplexSIMD instance representing the imaginary unit 0 + 1j. + + Returns: + ComplexSIMD instance with real part 0 and imaginary part 1 for all lanes. + + Examples: + ```mojo + from numojo.prelude import * + + # Create imaginary unit for different types + var i_f64 = ComplexSIMD[cf64].i() # (0 + 1j) + var i_f32 = ComplexSIMD[cf32].i() # (0 + 1j) + print(i_f64) # (0 + 1j) + + # Use in complex arithmetic + var z = 3.0 + 4.0 * ComplexSIMD[cf64].i() # 3 + 4j + ``` + """ + return Self(Self._broadcast(0), Self._broadcast(1)) + + @staticmethod + fn from_real_imag(re: Scalar[Self.dtype], im: Scalar[Self.dtype]) -> Self: + """ + Constructs a ComplexSIMD instance from scalar real and imaginary values. + + Args: + re: Scalar value for the real component. + im: Scalar value for the imaginary component. + + Example: + ```mojo + from numojo.prelude import * + var comp = ComplexSIMD[cf64].from_real_imag(2.0, 3.0) # (2.0 + 3.0j) + ``` + """ + return Self(re, im) + + @staticmethod + fn from_polar(r: Scalar[Self.dtype], theta: Scalar[Self.dtype]) -> Self: + """ + Constructs a ComplexSIMD instance from polar coordinates. + + Args: + r: Magnitude (radius). + theta: Angle (in radians). + + Returns: + ComplexSIMD instance with real part r * cos(theta) and imaginary part r * sin(theta). + + Example: + ```mojo + from numojo.prelude import * + var comp = ComplexSIMD[cf64].from_polar(2.0, 0.5) + ``` + """ + return Self( + Self._broadcast(r * cos(theta)), + Self._broadcast(r * sin(theta)), + ) + + # --- Arithmetic operators --- fn __add__(self, other: Self) -> Self: """ - Adds two ComplexSIMD instances. + Returns the element-wise sum of two ComplexSIMD instances. - Arguments: - other: The ComplexSIMD instance to add. + Args: + other: Another ComplexSIMD instance. Returns: - Self: A new ComplexSIMD instance representing the sum. + ComplexSIMD instance where each lane is the sum of corresponding lanes. """ return Self(self.re + other.re, self.im + other.im) + fn __add__(self, other: Scalar[Self.dtype]) -> Self: + """ + Returns the sum of this ComplexSIMD instance and a scalar added to the real part. + + Args: + other: Scalar value to add to the real component. + + Returns: + ComplexSIMD instance where each lane's real part is increased by the scalar. + """ + return Self(self.re + Self._broadcast(other), self.im) + + # FIXME: currently mojo doesn't allow overloading with both SIMD[Self.dtype, Self.width] and SIMD[*_, size=Self.width]. So keep SIMD[*_, size=Self.width] only for now. We need this method to create complex numbers with syntax like (1 + 2 * `1j`). + fn __add__(self, other: SIMD[*_, size = Self.width]) -> Self: + """ + Returns the sum of this ComplexSIMD instance and a SIMD vector added to the real part. + + Args: + other: SIMD vector to add to the real component. + + Returns: + ComplexSIMD instance where each lane's real part is increased by the corresponding lane in the SIMD vector. + """ + return Self(self.re + other.cast[Self.dtype](), self.im) + + fn __add__(self, other: ImaginaryUnit) -> Self: + """ + Returns the sum of this ComplexSIMD instance and the imaginary unit 1j. + + Args: + other: Imaginary unit (1j) to add to this complex number. + + Returns: + ComplexSIMD instance where each lane's imaginary part is increased by 1. + If self = a + bj, then result = a + (b+1)j. + + Examples: + ```mojo + from numojo.prelude import * + + var z = ComplexSIMD[cf64, 1](3.0, 2.0) # 3 + 2j + var result = z + `1j` # 3 + 3j + print(result) # (3 + 3j) + ``` + """ + return Self(self.re, self.im + Self._broadcast(1)) + fn __iadd__(mut self, other: Self): """ - Performs in-place addition of another ComplexSIMD instance. + In-place addition of another ComplexSIMD instance. - Arguments: - other: The ComplexSIMD instance to add. + Args: + other: Another ComplexSIMD instance. """ self.re += other.re self.im += other.im + fn __iadd__(mut self, other: Scalar[Self.dtype]): + """ + In-place addition of a scalar to the real part of this ComplexSIMD instance. + + Args: + other: Scalar value to add to the real component. + """ + self.re += Self._broadcast(other) + + fn __iadd__(mut self, other: SIMD[*_, size = Self.width]): + """ + In-place addition of a SIMD vector to the real part of this ComplexSIMD instance. + + Args: + other: SIMD vector to add to the real component. + """ + self.re += other.cast[Self.dtype]() + + fn __iadd__(mut self, other: ImaginaryUnit): + """ + In-place addition of the imaginary unit 1j to this ComplexSIMD instance. + + Args: + other: Imaginary unit (1j) to add to this complex number. + + Examples: + ```mojo + from numojo.prelude import * + + var z = ComplexSIMD[cf64, 1](3.0, 2.0) # 3 + 2j + z += `1j` # Now z = 3 + 3j + print(z) # (3 + 3j) + ``` + """ + self.im += Self._broadcast(1) + + fn __radd__(self, other: Scalar[Self.dtype]) -> Self: + """ + Returns the sum of a scalar and this ComplexSIMD instance, adding to the real part. + + Args: + other: Scalar value to add to the real component. + + Returns: + ComplexSIMD instance where each lane's real part is increased by the scalar. + """ + return Self(Self._broadcast(other) + self.re, self.im) + + fn __radd__(self, other: SIMD[*_, size = Self.width]) -> Self: + """ + Returns the sum of a SIMD vector and this ComplexSIMD instance, adding to the real part. + + Args: + other: SIMD vector to add to the real component. + + Returns: + ComplexSIMD instance where each lane's real part is increased by the corresponding lane in the SIMD vector. + """ + return Self(other.cast[Self.dtype]() + self.re, self.im) + + fn __radd__(self, other: ImaginaryUnit) -> Self: + """ + Returns the sum of the imaginary unit 1j and this ComplexSIMD instance. + + Args: + other: Imaginary unit (1j) to add to this complex number. + + Returns: + ComplexSIMD instance where each lane's imaginary part is increased by 1. + If self = a + bj, then result = a + (b+1)j. + + Examples: + ```mojo + from numojo.prelude import * + + var z = ComplexSIMD[cf64, 1](3.0, 2.0) # 3 + 2j + var result = `1j` + z # 3 + 3j + print(result) # (3 + 3j) + ``` + """ + return Self(self.re, self.im + Self._broadcast(1)) + fn __sub__(self, other: Self) -> Self: """ - Subtracts another ComplexSIMD instance from this one. + Returns the element-wise difference of two ComplexSIMD instances. - Arguments: - other: The ComplexSIMD instance to subtract. + Args: + other: Another ComplexSIMD instance. Returns: - Self: A new ComplexSIMD instance representing the difference. + ComplexSIMD instance where each lane is the difference of corresponding lanes. """ return Self(self.re - other.re, self.im - other.im) + fn __sub__(self, other: Scalar[Self.dtype]) -> Self: + """ + Returns the difference of this ComplexSIMD instance and a scalar subtracted from the real part. + + Args: + other: Scalar value to subtract from the real component. + + Returns: + ComplexSIMD instance where each lane's real part is decreased by the scalar. + """ + return Self(self.re - Self._broadcast(other), self.im) + + fn __sub__(self, other: SIMD[*_, size = Self.width]) -> Self: + """ + Returns the difference of this ComplexSIMD instance and a SIMD vector subtracted from the real part. + + Args: + other: SIMD vector to subtract from the real component. + + Returns: + ComplexSIMD instance where each lane's real part is decreased by the corresponding lane in the SIMD vector. + """ + return Self(self.re - other.cast[Self.dtype](), self.im) + + fn __sub__(self, other: ImaginaryUnit) -> Self: + """ + Subtracts the imaginary unit 1j from this ComplexSIMD instance. + + Args: + other: Imaginary unit (1j) to subtract from this complex number. + + Returns: + ComplexSIMD instance where each lane's imaginary part is decreased by 1. + If self = a + bj, then result = a + (b-1)j. + + Examples: + ```mojo + from numojo.prelude import * + + var z = ComplexSIMD[cf64, 1](3.0, 2.0) # 3 + 2j + var result = z - `1j` # 3 + 1j + print(result) # (3 + 1j) + ``` + """ + return Self(self.re, self.im - Self._broadcast(1)) + fn __isub__(mut self, other: Self): """ - Performs in-place subtraction of another ComplexSIMD instance. + In-place subtraction of another ComplexSIMD instance. - Arguments: - other: The ComplexSIMD instance to subtract. + Args: + other: Another ComplexSIMD instance. """ self.re -= other.re self.im -= other.im + fn __isub__(mut self, other: Scalar[Self.dtype]): + """ + In-place subtraction of a scalar from the real part of this ComplexSIMD instance. + + Args: + other: Scalar value to subtract from the real component. + """ + self.re -= Self._broadcast(other) + + fn __isub__(mut self, other: SIMD[*_, size = Self.width]): + """ + In-place subtraction of a SIMD vector from the real part of this ComplexSIMD instance. + + Args: + other: SIMD vector to subtract from the real component. + """ + self.re -= other.cast[Self.dtype]() + + fn __isub__(mut self, other: ImaginaryUnit): + """ + In-place subtraction of the imaginary unit 1j from this ComplexSIMD instance. + + Args: + other: Imaginary unit (1j) to subtract from this complex number. + + Examples: + ```mojo + from numojo.prelude import * + + var z = ComplexSIMD[cf64, 1](3.0, 2.0) # 3 + 2j + z -= `1j` # Now z = 3 + 1j + print(z) # (3 + 1j) + ``` + """ + self.im -= Self._broadcast(1) + + fn __rsub__(self, other: Scalar[Self.dtype]) -> Self: + """ + Returns the difference of a scalar and this ComplexSIMD instance, subtracting from the real part. + + Args: + other: Scalar value to subtract from the real component. + + Returns: + ComplexSIMD instance where each lane's real part is (scalar - self.re). + """ + return Self(Self._broadcast(other) - self.re, -self.im) + + fn __rsub__(self, other: SIMD[*_, size = Self.width]) -> Self: + """ + Returns the difference of a SIMD vector and this ComplexSIMD instance, subtracting from the real part. + + Args: + other: SIMD vector to subtract from the real component. + + Returns: + ComplexSIMD instance where each lane's real part is (SIMD lane - self.re). + """ + var other_casted = other.cast[Self.dtype]() + return Self(other_casted - self.re, -self.im) + + fn __rsub__(self, other: ImaginaryUnit) -> Self: + """ + Returns the difference of the imaginary unit 1j and this ComplexSIMD instance. + + Args: + other: Imaginary unit (1j) from which this complex number is subtracted. + + Returns: + ComplexSIMD instance equal to 1j - self. + If self = a + bj, then result = -a + (1-b)j. + + Examples: + ```mojo + from numojo.prelude import * + + var z = ComplexSIMD[cf64, 1](3.0, 2.0) # 3 + 2j + var result = `1j` - z # -3 + (-1)j = -3 - 1j + print(result) # (-3 - 1j) + ``` + """ + return Self(-self.re, Self._broadcast(1) - self.im) + fn __mul__(self, other: Self) -> Self: """ - Multiplies two ComplexSIMD instances. + Returns the element-wise product of two ComplexSIMD instances. - Arguments: - other: The ComplexSIMD instance to multiply with. + Args: + other: Another ComplexSIMD instance. Returns: - Self: A new ComplexSIMD instance representing the product. + ComplexSIMD instance where each lane is the product of corresponding lanes, using complex multiplication: (a+bi)(c+di) = (ac - bd) + (ad + bc)i. """ return Self( self.re * other.re - self.im * other.im, self.re * other.im + self.im * other.re, ) + fn __mul__(self, other: Scalar[Self.dtype]) -> Self: + """ + Returns the product of this ComplexSIMD instance and a scalar. + + Args: + other: Scalar value to multiply with both real and imaginary parts. + + Returns: + ComplexSIMD instance where each lane is scaled by the scalar. + """ + var scalar_simd = Self._broadcast(other) + return Self(self.re * scalar_simd, self.im * scalar_simd) + + fn __mul__(self, other: SIMD[*_, size = Self.width]) -> Self: + """ + Returns the product of this ComplexSIMD instance and a SIMD vector. + + Args: + other: SIMD vector to multiply with both real and imaginary parts. + + Returns: + ComplexSIMD instance where each lane is scaled by the corresponding lane in the SIMD vector. + """ + var other_casted = other.cast[Self.dtype]() + return Self(self.re * other_casted, self.im * other_casted) + + fn __mul__(self, other: ImaginaryUnit) -> Self: + """ + Returns the product of this ComplexSIMD instance and the imaginary unit 1j. + + Args: + other: Imaginary unit (1j) to multiply with this complex number. + + Returns: + ComplexSIMD instance where each lane is multiplied by 1j. + If self = a + bj, then result = (a + bj) * 1j = -b + aj. + + Examples: + ```mojo + from numojo.prelude import * + + var z = ComplexSIMD[cf64, 1](3.0, 2.0) # 3 + 2j + var result = z * `1j` # -2 + 3j + print(result) # (-2 + 3j) + ``` + """ + return Self(-self.im, self.re) + fn __imul__(mut self, other: Self): """ - Performs in-place multiplication with another ComplexSIMD instance. + In-place complex multiplication with another ComplexSIMD instance. - Arguments: - other: The ComplexSIMD instance to multiply with. + Args: + other: Another ComplexSIMD instance. """ - var re = self.re * other.re - self.im * other.im + var new_re = self.re * other.re - self.im * other.im self.im = self.re * other.im + self.im * other.re - self.re = re + self.re = new_re + + fn __imul__(mut self, other: Scalar[Self.dtype]): + """ + In-place multiplication of this ComplexSIMD instance by a scalar. + + Args: + other: Scalar value to multiply with both real and imaginary parts. + """ + var scalar_simd = Self._broadcast(other) + self.re *= scalar_simd + self.im *= scalar_simd + + fn __imul__(mut self, other: SIMD[*_, size = Self.width]): + """ + In-place multiplication of this ComplexSIMD instance by a SIMD vector. + + Args: + other: SIMD vector to multiply with both real and imaginary parts. + """ + var other_casted = other.cast[Self.dtype]() + self.re *= other_casted + self.im *= other_casted + + fn __imul__(mut self, other: ImaginaryUnit): + """ + In-place multiplication of this ComplexSIMD instance by the imaginary unit 1j. + + Args: + other: Imaginary unit (1j) to multiply with this complex number. + + Examples: + ```mojo + from numojo.prelude import * + + var z = ComplexSIMD[cf64, 1](3.0, 2.0) # 3 + 2j + z *= `1j` # Now z = -2 + 3j + print(z) # (-2 + 3j) + ``` + """ + var new_re = -self.im + self.im = self.re + self.re = new_re + + fn __rmul__(self, other: Scalar[Self.dtype]) -> Self: + """ + Returns the product of a scalar and this ComplexSIMD instance. + + Args: + other: Scalar value to multiply with both real and imaginary parts. + + Returns: + ComplexSIMD instance where each lane is scaled by the scalar. + """ + var scalar_simd = Self._broadcast(other) + return Self(scalar_simd * self.re, scalar_simd * self.im) + + fn __rmul__(self, other: SIMD[*_, size = Self.width]) -> Self: + """ + Returns the product of a SIMD vector and this ComplexSIMD instance. + + Args: + other: SIMD vector to multiply with both real and imaginary parts. + + Returns: + ComplexSIMD instance where each lane is scaled by the corresponding lane in the SIMD vector. + """ + var other_casted = other.cast[Self.dtype]() + return Self(other_casted * self.re, other_casted * self.im) + + fn __rmul__(self, other: ImaginaryUnit) -> Self: + """ + Returns the product of the imaginary unit 1j and this ComplexSIMD instance. + + Args: + other: Imaginary unit (1j) to multiply with this complex number. + + Returns: + ComplexSIMD instance where each lane is multiplied by 1j. + If self = a + bj, then result = 1j * (a + bj) = -b + aj. + + Examples: + ```mojo + from numojo.prelude import * + + var z = ComplexSIMD[cf64, 1](3.0, 2.0) # 3 + 2j + var result = `1j` * z # -2 + 3j + print(result) # (-2 + 3j) + ``` + """ + return Self(-self.im, self.re) fn __truediv__(self, other: Self) -> Self: """ - Divides this ComplexSIMD instance by another. + Performs element-wise complex division of two ComplexSIMD instances. - Arguments: - other: The ComplexSIMD instance to divide by. + Args: + other: Another ComplexSIMD instance to divide by. Returns: - Self: A new ComplexSIMD instance representing the quotient. + ComplexSIMD instance where each lane is the result of dividing the corresponding lanes: + (a + bi) / (c + di) = [(ac + bd) / (c^2 + d^2)] + [(bc - ad) / (c^2 + d^2)]i + where a, b are self.re, self.im and c, d are other.re, other.im. """ var denom = other.re * other.re + other.im * other.im return Self( @@ -193,214 +687,1075 @@ struct ComplexSIMD[cdtype: ComplexDType, width: Int = 1]( (self.im * other.re - self.re * other.im) / denom, ) + fn __truediv__(self, other: Scalar[Self.dtype]) -> Self: + """ + Performs element-wise division of this ComplexSIMD instance by a scalar. + + Args: + other: Scalar value to divide both real and imaginary parts by. + + Returns: + ComplexSIMD instance where each lane is divided by the scalar. + """ + var scalar_simd = Self._broadcast(other) + return Self(self.re / scalar_simd, self.im / scalar_simd) + + fn __truediv__(self, other: SIMD[*_, size = Self.width]) -> Self: + """ + Performs element-wise division of this ComplexSIMD instance by a SIMD vector. + + Args: + other: SIMD vector to divide both real and imaginary parts by. + + Returns: + ComplexSIMD instance where each lane is divided by the corresponding lane in the SIMD vector. + """ + var other_casted = other.cast[Self.dtype]() + return Self(self.re / other_casted, self.im / other_casted) + + fn __truediv__(self, other: ImaginaryUnit) -> Self: + """ + Performs division of this ComplexSIMD instance by the imaginary unit 1j. + + Args: + other: Imaginary unit (1j) to divide this complex number by. + + Returns: + ComplexSIMD instance where each lane is divided by 1j. + If self = a + bj, then result = (a + bj) / 1j = b - aj. + + Examples: + ```mojo + from numojo.prelude import * + + var z = ComplexSIMD[cf64, 1](3.0, 2.0) # 3 + 2j + var result = z / `1j` # 2 - 3j + print(result) # (2 - 3j) + ``` + """ + return Self(self.im, -self.re) + fn __itruediv__(mut self, other: Self): """ - Performs in-place division by another ComplexSIMD instance. + Performs in-place element-wise complex division of self by another ComplexSIMD instance. - Arguments: - other: The ComplexSIMD instance to divide by. + Args: + other: Another ComplexSIMD instance to divide by. """ var denom = other.re * other.re + other.im * other.im - var re = (self.re * other.re + self.im * other.im) / denom + var new_re = (self.re * other.re + self.im * other.im) / denom self.im = (self.im * other.re - self.re * other.im) / denom - self.re = re + self.re = new_re - fn __pow__(self, other: Self) -> Self: + fn __itruediv__(mut self, other: Scalar[Self.dtype]): """ - Raises this ComplexSIMD instance to the power of another. + Performs in-place element-wise division of this ComplexSIMD instance by a scalar. - Arguments: - other: The ComplexSIMD instance to raise to the power of. + Args: + other: Scalar value to divide both real and imaginary parts by. + """ + var scalar_simd = Self._broadcast(other) + self.re /= scalar_simd + self.im /= scalar_simd + + fn __itruediv__(mut self, other: SIMD[*_, size = Self.width]): + """ + Performs in-place element-wise division of this ComplexSIMD instance by a SIMD vector. + + Args: + other: SIMD vector to divide both real and imaginary parts by. + """ + var other_casted = other.cast[Self.dtype]() + self.re /= other_casted + self.im /= other_casted + + fn __itruediv__(mut self, other: ImaginaryUnit): + """ + Performs in-place division of this ComplexSIMD instance by the imaginary unit 1j. + + Args: + other: Imaginary unit (1j) to divide this complex number by. + + Examples: + ```mojo + from numojo.prelude import * + + var z = ComplexSIMD[cf64, 1](3.0, 2.0) # 3 + 2j + z /= `1j` # Now z = 2 - 3j + print(z) # (2 - 3j) + ``` + """ + var new_re = self.im + self.im = -self.re + self.re = new_re + + fn __rtruediv__(self, other: Scalar[Self.dtype]) -> Self: + """ + Performs element-wise division of a scalar by this ComplexSIMD instance. + + Args: + other: Scalar value to be divided by this ComplexSIMD instance. Returns: - Self: A new ComplexSIMD instance representing the result. + ComplexSIMD instance where each lane is the result of dividing the scalar by the corresponding lane: + other / (a + bi) = [other * a / (a^2 + b^2)] + [-other * b / (a^2 + b^2)]i + where a, b are self.re, self.im. + """ + var denom = self.re * self.re + self.im * self.im + var scalar_simd = Self._broadcast(other) + return Self( + (scalar_simd * self.re) / denom, + (-scalar_simd * self.im) / denom, + ) + + fn __rtruediv__(self, other: SIMD[*_, size = Self.width]) -> Self: + """ + Performs element-wise division of a SIMD vector by this ComplexSIMD instance. + + Args: + other: SIMD vector to be divided by this ComplexSIMD instance. + + Returns: + ComplexSIMD instance where each lane is the result of dividing the corresponding lane in the SIMD vector by the corresponding lane in this ComplexSIMD: + other[i] / (a + bi) = [other[i] * a / (a^2 + b^2)] + [-other[i] * b / (a^2 + b^2)]i + where a, b are self.re, self.im. + """ + var denom = self.re * self.re + self.im * self.im + var other_casted = other.cast[Self.dtype]() + return Self( + (other_casted * self.re) / denom, + (-other_casted * self.im) / denom, + ) + + fn __rtruediv__(self, other: ImaginaryUnit) -> Self: + """ + Performs division of the imaginary unit 1j by this ComplexSIMD instance. + + Args: + other: Imaginary unit (1j) to be divided by this ComplexSIMD instance. + + Returns: + ComplexSIMD instance where each lane is the result of dividing 1j by the corresponding lane. + If self = a + bj, then result = 1j / (a + bj) = [b / (a² + b²)] + [-a / (a² + b²)]j. + + Examples: + ```mojo + from numojo.prelude import * + + var z = ComplexSIMD[cf64, 1](3.0, 4.0) # 3 + 4j + var result = `1j` / z # 1j / (3 + 4j) = 0.16 - 0.12j + print(result) # (0.16 - 0.12j) + ``` + """ + var denom = self.re * self.re + self.im * self.im + return Self( + self.im / denom, + -self.re / denom, + ) + + fn reciprocal(self) raises -> Self: + """ + Returns the element-wise reciprocal (1 / self) of the ComplexSIMD instance. + + Returns: + ComplexSIMD instance representing the reciprocal of each lane: + 1 / (a + bi) = (a / (a^2 + b^2)) + (-b / (a^2 + b^2)). + """ + var d = self.norm() + if d == 0: + raise Error( + "Cannot compute reciprocal of zero norm complex number." + ) + return Self(self.re / d, -self.im / d) + + # --- Power helpers --- + fn elem_pow(self, other: Self) -> Self: + """ + Raises each component of this ComplexSIMD to the power of the corresponding component in another ComplexSIMD. + + Args: + other: Another ComplexSIMD instance. + + Returns: + ComplexSIMD instance where each lane is (re^other.re, im^other.im). """ return Self(self.re**other.re, self.im**other.im) - fn __pow__(self, other: Scalar[Self.dtype]) -> Self: + fn elem_pow(self, exponent: Scalar[Self.dtype]) -> Self: """ - Raises this ComplexSIMD instance to the power of a scalar. + Raises each component of this ComplexSIMD to a scalar exponent. - Arguments: - other: The scalar to raise to the power of. + Args: + exponent: Scalar exponent to apply to both real and imaginary parts. Returns: - Self: A new ComplexSIMD instance representing the result. + ComplexSIMD instance where each lane is (re^exponent, im^exponent). """ - return Self(self.re**other, self.im**other) + return Self(self.re**exponent, self.im**exponent) + fn __pow__(self, n: Int) -> Self: + """ + Raises this ComplexSIMD to an integer. + + Args: + n: Integer exponent. + + Returns: + ComplexSIMD instance raised to the power n. + For negative n, returns the reciprocal of self raised to -n. + """ + if n == 0: + return Self.one() + var base = self + var exp = n + var result = Self.one() + var is_negative = exp < 0 + if is_negative: + exp = -exp + while exp > 0: + if (exp & 1) == 1: + result = result * base + base = base * base + exp >>= 1 + if is_negative: + return Self.one() / result + return result + + # --- Unary operators --- fn __pos__(self) -> Self: """ - Returns the ComplexSIMD instance itself. + Returns the positive value of this ComplexSIMD (identity operation). Returns: - Self: The ComplexSIMD instance itself. + The same ComplexSIMD instance. """ return self fn __neg__(self) -> Self: """ - Negates the ComplexSIMD instance. + Returns the negation of this ComplexSIMD. Returns: - Self: The negated ComplexSIMD instance. + ComplexSIMD instance with both real and imaginary parts negated. """ - return self * Self(-1, -1) + return Self(-self.re, -self.im) + + # --- Helpers --- + @staticmethod + @always_inline + fn _abs_simd( + x: SIMD[Self.dtype, Self.width] + ) -> SIMD[Self.dtype, Self.width]: + return sqrt(x * x) + # --- Equality --- fn __eq__(self, other: Self) -> Bool: """ - Checks if two ComplexSIMD instances are equal. - - Arguments: - self: The first ComplexSIMD instance. - other: The second ComplexSIMD instance to compare with. + Checks if two ComplexSIMD instances are exactly equal. Returns: - Bool: True if the instances are equal, False otherwise. + True if both the real and imaginary parts are equal for all lanes, otherwise False. """ return (self.re == other.re) and (self.im == other.im) + fn __eq__(self, other: ImaginaryUnit) -> Bool: + """ + Checks if this ComplexSIMD instance is equal to the imaginary unit 1j. + + Args: + other: Imaginary unit (1j) to compare with this ComplexSIMD instance. + + Returns: + True if the real part is 0 and the imaginary part is 1 for all lanes, otherwise False. + + Examples: + ```mojo + from numojo.prelude import * + + var z1 = ComplexSIMD[cf64, 1](0.0, 1.0) # 0 + 1j + var z2 = ComplexSIMD[cf64, 1](1.0, 1.0) # 1 + 1j + print(z1 == `1j`) # True + print(z2 == `1j`) # False + ``` + """ + return (self.re == Self._broadcast(0)) and ( + self.im == Self._broadcast(1) + ) + fn __ne__(self, other: Self) -> Bool: """ Checks if two ComplexSIMD instances are not equal. - Arguments: - self: The first ComplexSIMD instance. - other: The second ComplexSIMD instance to compare with. + Returns: + True if either the real or imaginary parts differ for any lane, otherwise False. + """ + return ~(self == other) + + fn __ne__(self, other: ImaginaryUnit) -> Bool: + """ + Checks if this ComplexSIMD instance is not equal to the imaginary unit 1j. + + Args: + other: Imaginary unit (1j) to compare with this ComplexSIMD instance. Returns: - Bool: True if the instances are not equal, False otherwise. + True if either the real part is not 0 or the imaginary part is not 1 for any lane, otherwise False. + + Examples: + ```mojo + from numojo.prelude import * + + var z1 = ComplexSIMD[cf64, 1](0.0, 1.0) # 0 + 1j + var z2 = ComplexSIMD[cf64, 1](1.0, 1.0) # 1 + 1j + print(z1 != `1j`) # False + print(z2 != `1j`) # True + ``` """ return ~(self == other) - fn __str__(self) -> String: + fn allclose( + self, + other: Self, + *, + rtol: Scalar[Self.dtype] = 1e-5, + atol: Scalar[Self.dtype] = 1e-8, + ) -> Bool: """ - Returns a string representation of the ComplexSIMD instance. + Checks if two ComplexSIMD instances are approximately equal within given tolerances. + + For each lane, compares the real and imaginary parts using the formula: + abs(a - b) <= atol + rtol * abs(b) + where a and b are the corresponding components of self and other. + + Args: + other: Another ComplexSIMD instance to compare against. + rtol: Relative tolerance. + atol: Absolute tolerance. Returns: - String: The string representation of the ComplexSIMD instance. + True if all lanes of both real and imaginary parts are within the specified tolerances, otherwise False. + + Note: + For SIMD width > 1, all lanes must satisfy the tolerance criteria. """ + # TODO: Optionally return a SIMD[Bool] mask instead of a single Bool. + var diff_re = Self._abs_simd(self.re - other.re) + var diff_im = Self._abs_simd(self.im - other.im) + var rtol_b = Self._broadcast(rtol) + var atol_b = Self._broadcast(atol) + var thresh_re = atol_b + rtol_b * Self._abs_simd(other.re) + var thresh_im = atol_b + rtol_b * Self._abs_simd(other.im) + var ok_re = diff_re <= thresh_re + var ok_im = diff_im <= thresh_im + return ok_re and ok_im + + # --- Representations --- + fn __str__(self) -> String: return String.write(self) fn write_to[W: Writer](self, mut writer: W): """ - Writes the ComplexSIMD instance to a writer. + Returns a string representation of the ComplexSIMD instance. - Arguments: - self: The ComplexSIMD instance to write. - writer: The writer to write to. + For width == 1, the format is: (re + im j). + For width > 1, the format is: [(re0 + im0 j), (re1 + im1 j), ...]. """ try: - writer.write(String("({} + {} j)").format(self.re, self.im)) + + @parameter + if Self.width == 1: + writer.write(String("({} + {} j)").format(self.re, self.im)) + else: + var s = String("[") + for i in range(0, Self.width): + if i > 0: + s += ", " + s += String("({} + {} j)").format(self.re[i], self.im[i]) + s += "]" + writer.write(s) except e: - writer.write("Cannot convert ComplexSIMD to string") + writer.write("<>") fn __repr__(self) raises -> String: """ - Returns a string representation of the ComplexSIMD instance. - - Returns: - String: The string representation of the ComplexSIMD instance. + Returns a string representation of the ComplexSIMD instance for debugging. `ComplexSIMD[dtype](re=, im=)`. """ - return String("ComplexSIMD[{}]({}, {})").format( + return String("ComplexSIMD[{}](re={}, im={})").format( String(Self.dtype), self.re, self.im ) - fn __getitem__(self, idx: Int) raises -> SIMD[Self.dtype, Self.width]: + # --- Indexing --- + fn __getitem__(self, idx: Int) raises -> ComplexScalar[Self.cdtype]: """ - Gets the real or imaginary part of the ComplexSIMD instance. + Returns the complex number at the specified lane index. - Arguments: - self: The ComplexSIMD instance. - idx: The index to access (0 for real, 1 for imaginary). + Args: + idx: SIMD lane index (0 to width-1). Returns: - SIMD[Self.dtype, 1]: The requested part of the ComplexSIMD instance. + ComplexScalar containing the complex number at that lane index. + + Raises: + Error if lane index is out of range for the SIMD width. + + Example: + ```mojo + from numojo.prelude import * + var c_simd = ComplexSIMD[cf32, 2](SIMD[f32, 2](1, 2), SIMD[f32, 2](3, 4)) + var c0 = c_simd[0] # 1 + 3j + var c1 = c_simd[1] # 2 + 4j + ``` """ - if idx == 0: - return self.re - elif idx == 1: - return self.im - else: - raise Error("Index out of range") + if idx < 0 or idx >= Self.width: + raise Error("Lane index out of range for SIMD width") + return ComplexScalar[Self.cdtype](self.re[idx], self.im[idx]) fn __setitem__( - mut self, idx: Int, value: SIMD[Self.dtype, Self.width] + mut self, idx: Int, value: ComplexScalar[Self.cdtype] ) raises: """ - Sets the real and imaginary parts of the ComplexSIMD instance. + Sets the complex scalar at the specified lane index. + + Args: + idx: SIMD lane index (0 to width-1). + value: ComplexScalar whose values will be assigned. + + Raises: + Error if lane index is out of range for the SIMD width. + + Example: + ```mojo + from numojo.prelude import * + var c_simd = nm.ComplexSIMD[cf32, 2](SIMD[f32, 2](1, 2), SIMD[f32, 2](3, 4)) # [(1 + 3j), (2 + 4j)] + c_simd[0] = nm.CScalar[cf32](5, 6) + print(c_simd) # [(1 + 3j), (2 + 4j)] becomes [(5 + 6j), (2 + 4j)] + ``` + """ + if idx < 0 or idx >= Self.width: + raise Error("Lane index out of range for SIMD width") + self.re[idx] = value.re + self.im[idx] = value.im - Arguments: - self: The ComplexSIMD instance to modify. - idx: The index to access (0 for real, 1 for imaginary). - value: The new value to set. + fn item[name: String](self, idx: Int) raises -> Scalar[Self.dtype]: """ - if idx == 0: - self.re = value - elif idx == 1: - self.im = value + Returns the scalar value for the specified lane index and component. + + Parameters: + name: Name of the component ('re' or 'im'). + + Args: + idx: Lane index to retrieve. + + Returns: + Scalar value of the specified component at the given lane index. + + Raises: + - Error if the component name is invalid. + - Error if lane index is out of range for the SIMD width. + + Example: + ```mojo + from numojo.prelude import * + var c_simd = nm.ComplexSIMD[cf32, 2](SIMD[f32, 2](1, 2), SIMD[f32, 2](3, 4)) # [(1 + 3j), (2 + 4j)] + var re0 = c_simd.item["re"](0) # 1.0 + var im1 = c_simd.item["im"](1) # 4.0 + ``` + """ + if idx < 0 or idx >= Self.width: + raise Error("Lane index out of range for SIMD width") + + @parameter + if name == "re": + return self.re[idx] + elif name == "im": + return self.im[idx] else: - raise Error("Index out of range") + raise Error("Invalid component name: {}".format(name)) - fn __setitem__(mut self, idx: Int, value: Self) raises: + fn itemset[ + name: String + ](mut self, idx: Int, val: Scalar[Self.dtype]) raises: """ - Sets the real and imaginary parts of the ComplexSIMD instance. + Sets the scalar value for the specified lane index and component. + + Parameters: + name: Name of the component ('re' or 'im'). + + Args: + idx: Lane index to set. + val: Scalar value to assign to the specified component. - Arguments: - self: The ComplexSIMD instance to modify. - idx: The index to access (0 for real, 1 for imaginary). - value: The new value to set. + Raises: + - Error if the component name is invalid. + - Error if lane index is out of range for the SIMD width. + + Example: + ```mojo + from numojo.prelude import * + var c_simd = nm.ComplexSIMD[cf32, 2](SIMD[f32, 2](1, 2), SIMD[f32, 2](3, 4)) # [(1 + 3j), (2 + 4j)] + c_simd.itemset["re"](0, 5.0) # Now first complex number is (5 + 3j) + c_simd.itemset["im"](1, 6.0) # Now second complex number is (2 + 6j) + ``` """ - if idx == 0: - self.re = value.re - elif idx == 1: - self.im = value.im + if idx < 0 or idx >= Self.width: + raise Error("Lane index out of range for SIMD width") + + @parameter + if name == "re": + self.re[idx] = val + elif name == "im": + self.im[idx] = val else: - raise Error("Index out of range") + raise Error("Invalid component name: {}".format(name)) - fn item(self, idx: Int) raises -> SIMD[Self.dtype, Self.width]: + fn real(self) -> SIMD[Self.dtype, Self.width]: """ - Gets the real or imaginary part of the ComplexSIMD instance. + Returns the real part(s) of the complex number(s). + + Returns: + SIMD vector containing the real components. """ - return self[idx] + return self.re - fn itemset(mut self, val: ComplexSIMD[cdtype, Self.width]): + fn imag(self) -> SIMD[Self.dtype, Self.width]: """ - Sets the real and imaginary parts of the ComplexSIMD instance. + Returns the imaginary part(s) of the complex number(s). - Arguments: - self: The ComplexSIMD instance to modify. - val: The new value for the real and imaginary parts. + Returns: + SIMD vector containing the imaginary components. """ - self.re = val.re - self.im = val.im + return self.im + # --- Magnitude / norm / conjugate --- fn __abs__(self) -> SIMD[Self.dtype, Self.width]: """ - Returns the magnitude of the ComplexSIMD instance. + Returns the magnitude (absolute value) of the complex number(s). + + Returns: + SIMD vector containing the magnitude for each lane: sqrt(re^2 + im^2). """ return sqrt(self.re * self.re + self.im * self.im) fn norm(self) -> SIMD[Self.dtype, Self.width]: """ - Returns the squared magnitude of the ComplexSIMD instance. + Returns the squared magnitude (norm) of the complex number(s). + + Returns: + SIMD vector containing the squared magnitude for each lane: re^2 + im^2. """ - return sqrt(self.re * self.re + self.im * self.im) + return self.re * self.re + self.im * self.im fn conj(self) -> Self: """ Returns the complex conjugate of the ComplexSIMD instance. + + Returns: + ComplexSIMD instance with the imaginary part negated: (re, -im). """ return Self(self.re, -self.im) - fn real(self) -> SIMD[Self.dtype, Self.width]: + +@register_passable("trivial") +struct ImaginaryUnit(Boolable, Stringable, Writable): + """ + Constant representing the imaginary unit complex number 0 + 1j. + + The ImaginaryUnit struct provides a convenient way to work with the imaginary unit + in complex arithmetic operations. It supports arithmetic operations with SIMD vectors, + scalars, and other complex numbers, enabling Python-like syntax for complex number creation. + + Examples: + ```mojo + from numojo.prelude import * + + # Create complex numbers using the imaginary unit + var z1 = 3 + 4 * `1j` # 3 + 4j + var z2 = `1j` * 2 # 0 + 2j + var z3 = 5.0 + `1j` # 5 + 1j + + # Powers of the imaginary unit + print(`1j` ** 0) # 1 + 0j + print(`1j` ** 1) # 0 + 1j + print(`1j` ** 2) # -1 + 0j + print(`1j` ** 3) # 0 - 1j + + # SIMD complex vectors + var c4 = SIMD[f32, 4](1.0) + `1j` * SIMD[f32, 4](2.0) # ComplexSIMD[cf32, 4] + var c5 = SIMD[f64, 2](3.0, 4.0) + `1j` # ComplexSIMD[cf64, 2] + var d = SIMD[f32, 2](1) + SIMD[f32, 2](2) * `1j` # creates [( 1 + 2 j) (1 + 2 j)] + + # Mathematical properties + var c6 = `1j` * `1j` # -1 (Scalar[f64]) + var c7 = `1j` ** 3 # (0 - 1j) (ComplexScalar[cf64]) + var c8 = (1 + `1j`) / `1j` # (1 - 1j) (ComplexScalar[cf64]) + ``` + """ + + fn __init__(out self): """ - Returns the real part of the ComplexSIMD instance. + Constructor for ImaginaryUnit. + + Creates an instance representing the imaginary unit 1j. """ - return self.re + pass - fn imag(self) -> SIMD[Self.dtype, Self.width]: + fn conj(self) -> ComplexSIMD[cf64, 1]: + """ + Returns the complex conjugate of the imaginary unit. + + Returns: + ComplexSIMD representing -1j (the conjugate of 1j). + + Examples: + ```mojo + from numojo.prelude import * + + var conj_i = `1j`.conj() # -1j + print(conj_i) # (0 - 1j) + ``` """ - Returns the imaginary part of the ComplexSIMD instance. + return -self + + # --- Arithmetic operators with SIMD and Scalar types --- + # Addition: 1j + SIMD -> ComplexSIMD + fn __add__[ + dtype: DType, width: Int + ](self, other: SIMD[dtype, width]) -> ComplexSIMD[ + ComplexDType(mlir_value=dtype._mlir_value), width + ]: """ - return self.im + Returns the sum of the imaginary unit 1j and a SIMD vector. + + Args: + other: SIMD vector to add to the imaginary unit. + + Returns: + ComplexSIMD where real part equals the SIMD vector and imaginary part is 1. + + Examples: + ```mojo + from numojo.prelude import * + + var vec = SIMD[DType.float32, 4](1.0, 2.0, 3.0, 4.0) + var result = `1j` + vec # [1+1j, 2+1j, 3+1j, 4+1j] + ``` + """ + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), width]( + other, SIMD[dtype, width](1) + ) + + # Addition: 1j + Scalar -> ComplexScalar + fn __add__[ + dtype: DType + ](self, other: Scalar[dtype]) -> ComplexScalar[ + ComplexDType(mlir_value=dtype._mlir_value) + ]: + """ + Returns the sum of the imaginary unit 1j and a scalar. + + Args: + other: Scalar to add to the imaginary unit. + + Returns: + ComplexScalar with real part equal to the scalar and imaginary part 1. + + Examples: + ```mojo + from numojo.prelude import * + + var result = `1j` + 3.5 # 3.5 + 1j + ``` + """ + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), 1]( + other, 1 + ) + + fn __add__(self, other: Int) -> ComplexScalar[ComplexDType.int]: + """ + Returns the sum of the imaginary unit 1j and an integer. + + Args: + other: Integer to add to the imaginary unit. + + Returns: + ComplexScalar with real part equal to the integer and imaginary part 1. + + Examples: + ```mojo + from numojo.prelude import * + + var result = `1j` + 5 # 5 + 1j + ``` + """ + return ComplexSIMD[ComplexDType.int, 1](other, 1) + + # SIMD + 1j -> ComplexSIMD + fn __radd__[ + dtype: DType, width: Int + ](self, other: SIMD[dtype, width]) -> ComplexSIMD[ + ComplexDType(mlir_value=dtype._mlir_value), width + ]: + """Returns the sum of a SIMD vector and the imaginary unit.""" + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), width]( + other, SIMD[dtype, width](1) + ) + + # Addition: Scalar + 1j -> ComplexScalar + fn __radd__[ + dtype: DType + ](self, other: Scalar[dtype]) -> ComplexScalar[ + ComplexDType(mlir_value=dtype._mlir_value) + ]: + """Returns the sum of a scalar and the imaginary unit.""" + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), 1]( + other, 1 + ) + + # Addition: Int + 1j -> ComplexScalar + fn __radd__(self, other: Int) -> ComplexScalar[ComplexDType.int]: + """Returns the sum of an integer and the imaginary unit.""" + return ComplexSIMD[ComplexDType.int, 1](other, 1) + + # Subtraction: 1j - SIMD -> ComplexSIMD + fn __sub__[ + dtype: DType, width: Int + ](self, other: SIMD[dtype, width]) -> ComplexSIMD[ + ComplexDType(mlir_value=dtype._mlir_value), width + ]: + """Returns the difference of the imaginary unit and a SIMD vector.""" + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), width]( + -other, SIMD[dtype, width](1) + ) + + # Subtraction: 1j - Scalar -> ComplexScalar + fn __sub__[ + dtype: DType + ](self, other: Scalar[dtype]) -> ComplexScalar[ + ComplexDType(mlir_value=dtype._mlir_value) + ]: + """Returns the difference of the imaginary unit and a scalar.""" + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), 1]( + -other, 1 + ) + + fn __sub__(self, other: Int) -> ComplexScalar[ComplexDType.int]: + """Returns the difference of the imaginary unit and an integer.""" + return ComplexSIMD[ComplexDType.int, 1](-other, 1) + + # Subtraction: SIMD - 1j -> ComplexSIMD + fn __rsub__[ + dtype: DType, width: Int + ](self, other: SIMD[dtype, width]) -> ComplexSIMD[ + ComplexDType(mlir_value=dtype._mlir_value), width + ]: + """Returns the difference of a SIMD vector and the imaginary unit.""" + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), width]( + other, SIMD[dtype, width](-1) + ) + + # Subtraction: Scalar - 1j -> ComplexScalar + fn __rsub__[ + dtype: DType + ](self, other: Scalar[dtype]) -> ComplexScalar[ + ComplexDType(mlir_value=dtype._mlir_value) + ]: + """Returns the difference of a scalar and the imaginary unit.""" + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), 1]( + other, -1 + ) + + # Subtraction: Int - 1j -> ComplexScalar + fn __rsub__(self, other: Int) -> ComplexScalar[ComplexDType.int]: + """Returns the difference of an integer and the imaginary unit.""" + return ComplexSIMD[ComplexDType.int, 1](other, -1) + + # Multiplication: 1j * SIMD -> ComplexSIMD + fn __mul__[ + dtype: DType, width: Int + ](self, other: SIMD[dtype, width]) -> ComplexSIMD[ + ComplexDType(mlir_value=dtype._mlir_value), width + ]: + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), width]( + SIMD[dtype, width](0), other + ) + + # Multiplication: 1j * Scalar -> ComplexScalar + fn __mul__[ + dtype: DType + ](self, other: Scalar[dtype]) -> ComplexScalar[ + ComplexDType(mlir_value=dtype._mlir_value) + ]: + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), 1]( + 0, other + ) + + # Multiplication: 1j * Int -> ComplexScalar + fn __mul__(self, other: Int) -> ComplexScalar[ComplexDType.int]: + return ComplexSIMD[ComplexDType.int, 1](0, other) + + fn __rmul__[ + dtype: DType, + width: Int, + ](self, other: SIMD[dtype, width]) -> ComplexSIMD[ + ComplexDType(mlir_value=dtype._mlir_value), width + ]: + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), width]( + SIMD[dtype, width](0), other + ) + + fn __rmul__[ + dtype: DType + ](self, other: Scalar[dtype]) -> ComplexScalar[ + ComplexDType(mlir_value=dtype._mlir_value) + ]: + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), 1]( + 0, other + ) + + # Multiplication: Scalar * 1j -> ComplexScalar + fn __rmul__(self, other: Int) -> ComplexScalar[ComplexDType.int]: + return ComplexSIMD[ComplexDType.int, 1](0, other) + + # Division: 1j / SIMD -> ComplexSIMD + fn __truediv__[ + dtype: DType, width: Int + ](self, other: SIMD[dtype, width]) -> ComplexSIMD[ + ComplexDType(mlir_value=dtype._mlir_value), width + ]: + """Returns the division of the imaginary unit by a SIMD vector.""" + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), width]( + SIMD[dtype, width](0), 1 / other + ) + + # Division: 1j / Scalar -> ComplexScalar + fn __truediv__[ + dtype: DType + ](self, other: Scalar[dtype]) -> ComplexScalar[ + ComplexDType(mlir_value=dtype._mlir_value) + ]: + """Returns the division of the imaginary unit by a scalar.""" + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), 1]( + 0, 1 / other + ) + + # Division: SIMD / 1j -> ComplexSIMD + fn __rtruediv__[ + dtype: DType, width: Int + ](self, other: SIMD[dtype, width]) -> ComplexSIMD[ + ComplexDType(mlir_value=dtype._mlir_value), width + ]: + """ + Returns the division of a SIMD vector by the imaginary unit. + + Args: + other: SIMD vector to be divided by the imaginary unit. + + Returns: + ComplexSIMD representing (0 - other j). + """ + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), width]( + SIMD[dtype, width](0), -other + ) + + # Division: Scalar / 1j -> ComplexScalar + fn __rtruediv__[ + dtype: DType + ](self, other: Scalar[dtype]) -> ComplexScalar[ + ComplexDType(mlir_value=dtype._mlir_value) + ]: + """ + Returns the division of a scalar by the imaginary unit. + + Args: + other: Scalar to be divided by the imaginary unit. + + Returns: + ComplexScalar representing (0 - other j). + """ + return ComplexSIMD[ComplexDType(mlir_value=dtype._mlir_value), 1]( + 0, -other + ) + + # Division: Int / 1j -> ComplexScalar + fn __rtruediv__(self, other: Int) -> ComplexScalar[ComplexDType.int]: + """ + Returns the division of an integer by the imaginary unit. + + Args: + other: Integer to be divided by the imaginary unit. + + Returns: + ComplexScalar representing (0 - other j). + """ + return ComplexSIMD[ComplexDType.int, 1](0, -other) + + # Self-operations: 1j with 1j + fn __mul__(self, other: ImaginaryUnit) -> Scalar[DType.float64]: + """ + Returns the product of the imaginary unit with itself: 1j * 1j = -1. + + Args: + other: Another imaginary unit to multiply with. + + Returns: + Scalar value -1, since (1j) * (1j) = -1. + + Examples: + ```mojo + from numojo.prelude import * + + var result = `1j` * `1j` # -1 + print(result) # -1 + ``` + """ + return -1 + + fn __add__(self, other: ImaginaryUnit) -> ComplexScalar[cf64]: + """ + Returns the sum of the imaginary unit with itself: 1j + 1j = 2j. + + Args: + other: Another imaginary unit to add. + + Returns: + ComplexScalar representing 2j. + + Examples: + ```mojo + from numojo.prelude import * + + var result = `1j` + `1j` # 2j + print(result) # (0 + 2j) + ``` + """ + return ComplexSIMD[cf64, 1](0, 2) + + fn __sub__(self, other: ImaginaryUnit) -> Scalar[DType.float64]: + """ + Returns the difference of the imaginary unit with itself: 1j - 1j = 0. + + Args: + other: Another imaginary unit to subtract. + + Returns: + Scalar value 0. + + Examples: + ```mojo + from numojo.prelude import * + + var result = `1j` - `1j` # 0 + print(result) # 0 + ``` + """ + return 0 + + fn __truediv__(self, other: ImaginaryUnit) -> Scalar[DType.float64]: + """ + Returns the division of the imaginary unit by itself: 1j / 1j = 1. + + Args: + other: Another imaginary unit to divide by. + + Returns: + Scalar value 1. + + Examples: + ```mojo + from numojo.prelude import * + + var result = `1j` / `1j` # 1 + print(result) # 1 + ``` + """ + return 1 + + fn __pow__(self, exponent: Int) -> ComplexScalar[cf64]: + """ + Returns the imaginary unit raised to an integer power. + + The powers of 1j cycle with period 4: + - 1j^0 = 1 + - 1j^1 = 1j + - 1j^2 = -1 + - 1j^3 = -1j + - 1j^4 = 1 (cycle repeats) + + Args: + exponent: Integer exponent. + + Returns: + ComplexScalar representing (1j) ** exponent. + + Examples: + ```mojo + from numojo.prelude import * + + print(`1j` ** 0) # (1 + 0j) + print(`1j` ** 1) # (0 + 1j) + print(`1j` ** 2) # (-1 + 0j) + print(`1j` ** 3) # (0 - 1j) + print(`1j` ** 4) # (1 + 0j) + ``` + """ + var remainder = exponent % 4 + if remainder == 0: + return ComplexSIMD[cf64, 1](1, 0) + elif remainder == 1: + return ComplexSIMD[cf64, 1](0, 1) + elif remainder == 2: + return ComplexSIMD[cf64, 1](-1, 0) + else: + return ComplexSIMD[cf64, 1](0, -1) + + fn __neg__(self) -> ComplexScalar[cf64]: + """ + Returns the negation of the imaginary unit: -1j. + + Returns: + ComplexScalar representing -1j. + + Examples: + ```mojo + from numojo.prelude import * + + var result = -`1j` # -1j + print(result) # (0 - 1j) + ``` + """ + return ComplexSIMD[cf64, 1](0, -1) + + fn __str__(self) -> String: + """ + Returns the string representation of the imaginary unit. + + Returns: + String representing the imaginary unit as "(0 + 1 j)". + """ + return "(0 + 1 j)" + + fn write_to[W: Writer](self, mut writer: W): + """ + Writes the string representation of the imaginary unit to a writer. + + Args: + writer: Writer instance to write the string representation to. + """ + writer.write("(0 + 1 j)") + + fn __bool__(self) -> Bool: + """ + Returns the boolean value of the imaginary unit. + + Returns: + Always True, since the imaginary unit 1j is non-zero. + + Examples: + ```mojo + from numojo.prelude import * + + if `1j`: + print("Imaginary unit is truthy") # This will execute + ``` + """ + return True diff --git a/numojo/core/ndarray.mojo b/numojo/core/ndarray.mojo index 0a872411..5e37e684 100644 --- a/numojo/core/ndarray.mojo +++ b/numojo/core/ndarray.mojo @@ -175,6 +175,15 @@ struct NDArray[dtype: DType = DType.float64]( Args: shape: Variadic shape. order: Memory order C or F. + + Example: + ```mojo + import numojo as nm + var a = nm.NDArray[nm.f32](nm.Shape(2,3), order="C") + ``` + + Note: + This constructor should not be used by users directly. Use factory functions in `nomojo.routines.creation` module instead. """ self.ndim = shape.ndim @@ -199,6 +208,15 @@ struct NDArray[dtype: DType = DType.float64]( Args: shape: List of shape. order: Memory order C or F. + + Example: + ```mojo + import numojo as nm + var a = nm.NDArray[nm.f32](List[Int](2,3), order="C") + ``` + + Note: + This constructor should not be used by users directly. Use factory functions in `numojo.routines.creation` module instead. """ self = Self(Shape(shape), order) @@ -215,6 +233,15 @@ struct NDArray[dtype: DType = DType.float64]( Args: shape: Variadic List of shape. order: Memory order C or F. + + Example: + ```mojo + from numojo.prelude import * + var A = nm.ComplexNDArray[cf32](VariadicList(2,3,4)) + ``` + + Note: + This constructor should not be used by users directly. Use factory functions in `numojo.routines.creation` module instead. """ self = Self(Shape(shape), order) @@ -226,12 +253,25 @@ struct NDArray[dtype: DType = DType.float64]( strides: List[Int], ) raises: """ - Extremely specific NDArray initializer. + Initialize a NDArray with a specific shape, offset, and strides. Args: - shape: List of shape. - offset: Offset value. - strides: List of strides. + shape: List of integers specifying the shape of the array. + offset: Integer offset into the underlying buffer. + strides: List of integers specifying the stride for each dimension. + + Notes: + - This constructor is intended for advanced use cases requiring precise control over memory layout. + - The resulting array is uninitialized and should be filled before use. + + Example: + ```mojo + from numojo.prelude import * + var shape = List[Int](2, 3) + var offset = 0 + var strides = List[Int](3, 1) + var arr = NDArray[f32](shape, offset, strides) + ``` """ self.shape = NDArrayShape(shape) self.ndim = self.shape.ndim @@ -253,16 +293,19 @@ struct NDArray[dtype: DType = DType.float64]( flags: Flags, ): """ - Constructs an extremely specific array, with value uninitialized. - The properties do not need to be compatible and are not checked. - For example, it can construct a 0-D array (numojo scalar). + Initialize a NDArray with explicit shape, strides, number of dimensions, size, and flags. This constructor creates an uninitialized NDArray with the provided properties. No compatibility checks are performed between shape, strides, ndim, size, or flags. This allows construction of arrays with arbitrary metadata, including 0-D arrays (scalars). Args: - shape: Shape of array. - strides: Strides of array. + shape: Shape of the array. + strides: Strides for each dimension. ndim: Number of dimensions. - size: Size of array. - flags: Flags of array. + size: Total number of elements. + flags: Memory layout flags. + + Notes: + - This constructor is intended for advanced or internal use cases requiring manual control. + - The resulting array is uninitialized; values must be set before use. + - No validation is performed on the consistency of the provided arguments. """ self.shape = shape @@ -282,14 +325,25 @@ struct NDArray[dtype: DType = DType.float64]( strides: NDArrayStrides, ) raises: """ - Initialize an NDArray view with given shape, buffer, offset, and strides. - ***Unsafe!*** This function is currently unsafe. Only for internal use. + Initialize a NDArray view with explicit shape, raw buffers, offset, and strides. + + This constructor creates a view over existing memory buffers for the real and imaginary parts, + using the provided shape, offset, and stride information. It is intended for advanced or internal + use cases where direct control over memory layout is required. + + ***Unsafe!*** This function is unsafe and should only be used internally. The caller is responsible + for ensuring that the buffers are valid and that the shape, offset, and strides are consistent. Args: - shape: Shape of the array. - buffer: Unsafe pointer to the buffer. - offset: Offset value. - strides: Strides of the array. + shape: NDArrayShape specifying the dimensions of the array. + buffer: Unsafe pointer to the buffer containing the real part data. + offset: Integer offset into the buffers. + strides: NDArrayStrides specifying the stride for each dimension. + + Notes: + - No validation is performed on the buffers or metadata. + - The resulting NDArray shares memory with the provided buffers. + - Incorrect usage may lead to undefined behavior. """ self.shape = shape self.strides = strides @@ -384,6 +438,32 @@ struct NDArray[dtype: DType = DType.float64]( # fn load[width: Int](self, *indices: Int) raises -> SIMD[dtype, width] # Load SIMD at coordinates # ===-------------------------------------------------------------------===# + @always_inline + fn normalize(self, idx: Int, dim: Int) -> Int: + """ + Normalize a potentially negative index to its positive equivalent + within the bounds of the given dimension. + + Args: + idx: The index to normalize. Can be negative to indicate indexing + from the end (e.g., -1 refers to the last element). + dim: The size of the dimension to normalize against. + + Returns: + The normalized index as a non-negative integer. + + Example: + ```mojo + from numojo.prelude import * + var mat = Matrix[f32](shape=(3, 4)) + var norm_idx = mat.normalize(-1, mat.shape[0]) # Normalize -1 to 2 + ``` + """ + var idx_norm = idx + if idx_norm < 0: + idx_norm = dim + idx_norm + return idx_norm + fn _getitem(self, *indices: Int) -> Scalar[dtype]: """ Get item at indices and bypass all boundary checks. @@ -395,17 +475,16 @@ struct NDArray[dtype: DType = DType.float64]( Returns: The element of the array at the indices. - Notes: - This function is unsafe and should be used only on internal use. - Examples: + ```mojo + import numojo as nm + from numojo.prelude import * + var A = nm.ones[f32](nm.Shape(2,3,4)) + print(A._getitem(1,2,3)) + ``` - ```mojo - import numojo as nm - from numojo.prelude import * - var A = nm.ones[f32](nm.Shape(2,3,4)) - print(A._getitem(1,2,3)) - ``` + Notes: + This function is unsafe and should be used only on internal use. """ var index_of_buffer: Int = 0 for i in range(self.ndim): @@ -423,9 +502,6 @@ struct NDArray[dtype: DType = DType.float64]( Returns: The element of the array at the indices. - Notes: - This function is unsafe and should be used only on internal use. - Examples: ```mojo @@ -434,6 +510,9 @@ struct NDArray[dtype: DType = DType.float64]( var A = nm.ones[f32](nm.Shape(2,3,4)) print(A._getitem(List[Int](1,2,3))) ``` + + Notes: + This function is unsafe and should be used only on internal use. """ var index_of_buffer: Int = 0 for i in range(self.ndim): @@ -452,10 +531,10 @@ struct NDArray[dtype: DType = DType.float64]( Examples: - ```console - >>>import numojo - >>>var a = numojo.arange(3)[0] - >>>print(a[]) # gets values of the 0-D array. + ```mojo + import numojo as nm + var a = nm.arange(3)[0] + print(a[]) # gets values of the 0-D array. ```. """ if self.ndim != 0: @@ -551,18 +630,15 @@ struct NDArray[dtype: DType = DType.float64]( IndexError: If `idx` is out of bounds after normalization. Notes: - Order preservation: The resulting copy preserves the source - array's memory order (C or F). Performance fast path: For - C-contiguous arrays the slice is a single contiguous block and is - copied with one `memcpy`. For F-contiguous or arbitrary - strided layouts a unified stride-based element loop is used. - (Future enhancement: return a non-owning view instead of + Order preservation: The resulting copy preserves the source array's memory order (C or F). Performance fast path: For C-contiguous arrays the slice is a single contiguous block and is + copied with one `memcpy`. For F-contiguous or arbitrary strided layouts a unified stride-based element loop is used. (Future enhancement: return a non-owning view instead of copying.) - Examples: + Example: ```mojo import numojo as nm - var a = nm.arange(0, 12, 1).reshape(nm.Shape(3, 4)) + from numojo.prelude import * + var a = nm.arange(0, 12, 1).reshape(Shape(3, 4)) print(a.shape) # (3,4) print(a[1].shape) # (4,) -- 1-D slice print(a[-1].shape) # (4,) -- negative index @@ -617,11 +693,11 @@ struct NDArray[dtype: DType = DType.float64]( count=block, ) return result^ - # (F-order or arbitrary stride layout) # TODO: Optimize this further (multi-axis unrolling / smarter linear index without div/mod) - self._copy_first_axis_slice(self, norm, result) - return result^ + else: + self._copy_first_axis_slice(self, norm, result) + return result^ # perhaps move these to a utility module fn _copy_first_axis_slice( @@ -707,7 +783,7 @@ struct NDArray[dtype: DType = DType.float64]( var narr: Self = self[slice_list^] return narr^ - fn _calculate_strides_efficient(self, shape: List[Int]) -> List[Int]: + fn _calculate_strides(self, shape: List[Int]) -> List[Int]: var strides = List[Int](capacity=len(shape)) if self.flags.C_CONTIGUOUS: # C_CONTIGUOUS @@ -748,7 +824,7 @@ struct NDArray[dtype: DType = DType.float64]( - This method supports advanced slicing similar to NumPy's multi-dimensional slicing. - The returned array shares data with the original array if possible. - Examples: + Example: ```mojo import numojo as nm var a = nm.arange(10).reshape(nm.Shape(2, 5)) @@ -805,7 +881,7 @@ struct NDArray[dtype: DType = DType.float64]( ncoefficients.append(1) # only C & F order are supported - var nstrides: List[Int] = self._calculate_strides_efficient( + var nstrides: List[Int] = self._calculate_strides( nshape, ) var narr: Self = Self(offset=noffset, shape=nshape, strides=nstrides) @@ -931,7 +1007,7 @@ struct NDArray[dtype: DType = DType.float64]( ncoefficients.append(1) # only C & F order are supported - var nstrides: List[Int] = self._calculate_strides_efficient( + var nstrides: List[Int] = self._calculate_strides( nshape, ) var narr: Self = Self(offset=noffset, shape=nshape, strides=nstrides) @@ -1230,8 +1306,8 @@ struct NDArray[dtype: DType = DType.float64]( # var shape = indices.shape.join(self.shape._pop(0)) var shape = indices.shape.join(self.shape._pop(0)) - var result = NDArray[dtype](shape) - var size_per_item = self.size // self.shape[0] + var result: NDArray[dtype] = NDArray[dtype](shape) + var size_per_item: Int = self.size // self.shape[0] # Fill in the values for i in range(indices.size): @@ -1534,8 +1610,7 @@ struct NDArray[dtype: DType = DType.float64]( ) ) - if index < 0: - index += self.size + index = self.normalize(index, self.size) if (index < 0) or (index >= self.size): raise Error( @@ -1665,8 +1740,7 @@ struct NDArray[dtype: DType = DType.float64]( ```. """ - if index < 0: - index += self.size + index = self.normalize(index, self.size) if index >= self.size: raise Error( @@ -1769,9 +1843,7 @@ struct NDArray[dtype: DType = DType.float64]( var indices_list: List[Int] = List[Int](capacity=self.ndim) for i in range(self.ndim): var idx_i = indices[i] - if idx_i < 0: - idx_i += self.shape[i] - elif idx_i >= self.shape[i]: + if idx_i < 0 or idx_i >= self.shape[i]: raise Error( IndexError( message=String( @@ -1788,6 +1860,7 @@ struct NDArray[dtype: DType = DType.float64]( ), ) ) + idx_i = self.normalize(idx_i, self.shape[i]) indices_list.append(idx_i) # indices_list already built above @@ -1895,8 +1968,7 @@ struct NDArray[dtype: DType = DType.float64]( ) var norm = idx - if norm < 0: - norm += self.shape[0] + norm = self.normalize(norm, self.shape[0]) if (norm < 0) or (norm >= self.shape[0]): raise Error( IndexError( @@ -1912,21 +1984,6 @@ struct NDArray[dtype: DType = DType.float64]( ) ) - if val.ndim != self.ndim - 1: - raise Error( - ValueError( - message=String( - "Value ndim {} incompatible with target slice ndim {}." - ).format(val.ndim, self.ndim - 1), - suggestion=String( - "Reshape or expand value to ndim {}." - ).format(self.ndim - 1), - location=String( - "NDArray.__setitem__(idx: Int, val: NDArray)" - ), - ) - ) - if self.shape[1:] != val.shape: var expected_shape: NDArrayShape = self.shape[1:] raise Error( @@ -1997,10 +2054,11 @@ struct NDArray[dtype: DType = DType.float64]( Examples: - ```console - >>> import numojo - >>> var A = numojo.random.rand[numojo.i16](2, 2, 2) - >>> A[numojo.Item(0, 1, 1)] = 10 + ```mojo + import numojo as nm + from numojo.prelude import * + var A = numojo.random.rand[numojo.i16](2, 2, 2) + A[Item(0, 1, 1)] = 10 ```. """ if index.__len__() != self.ndim: @@ -2035,8 +2093,7 @@ struct NDArray[dtype: DType = DType.float64]( ), ) ) - if index[i] < 0: - index[i] += self.shape[i] + index[i] = self.normalize(index[i], self.shape[i]) var idx: Int = _get_offset(index, self.strides) self._buf.ptr.store(idx, val) diff --git a/numojo/prelude.mojo b/numojo/prelude.mojo index 7edd1297..08a98589 100644 --- a/numojo/prelude.mojo +++ b/numojo/prelude.mojo @@ -26,7 +26,12 @@ from numojo.core.item import Item, item from numojo.core.matrix import Matrix from numojo.core.ndarray import NDArray from numojo.core.ndshape import Shape, NDArrayShape -from numojo.core.complex.complex_simd import ComplexSIMD, CScalar +from numojo.core.complex.complex_simd import ( + ComplexSIMD, + CScalar, + ComplexScalar, + `1j`, +) from numojo.core.complex.complex_ndarray import ComplexNDArray from numojo.core.complex.complex_dtype import ( ci8, diff --git a/numojo/routines/creation.mojo b/numojo/routines/creation.mojo index 770a89dd..1c5eaf90 100644 --- a/numojo/routines/creation.mojo +++ b/numojo/routines/creation.mojo @@ -933,7 +933,7 @@ fn geomspace[ ) var base: ComplexSIMD[cdtype] = stop / start var power: Scalar[dtype] = 1 / Scalar[dtype](num - 1) - var r: ComplexSIMD[cdtype] = base**power + var r: ComplexSIMD[cdtype] = base.elem_pow(power) for i in range(num): result.store[1]( i, @@ -947,7 +947,7 @@ fn geomspace[ ) var base: ComplexSIMD[cdtype] = stop / start var power: Scalar[dtype] = 1 / Scalar[dtype](num) - var r: ComplexSIMD[cdtype] = base**power + var r: ComplexSIMD[cdtype] = base.elem_pow(power) for i in range(num): result.store[1]( i, diff --git a/tests/core/test_complexSIMD.mojo b/tests/core/test_complexSIMD.mojo index 74f9bc66..b2036288 100644 --- a/tests/core/test_complexSIMD.mojo +++ b/tests/core/test_complexSIMD.mojo @@ -10,8 +10,12 @@ fn test_complex_init() raises: assert_equal(c1.im, 2.0, "init failed") var c2 = ComplexSIMD[cf32](c1) - assert_equal(c2.re, c1.re) - assert_equal(c2.im, c1.im) + assert_equal(c2.re, c1.re, "init failed") + assert_equal(c2.im, c1.im, "init failed") + + var c3 = ComplexSIMD[cf32, 2](1.0) + assert_equal(c3.re[0], 1.0, "init failed") + assert_equal(c3.im[0], 1.0, "init failed") fn test_complex_add() raises: