diff --git a/numojo/__init__.mojo b/numojo/__init__.mojo index 9488b874..e75f9416 100644 --- a/numojo/__init__.mojo +++ b/numojo/__init__.mojo @@ -12,7 +12,7 @@ alias __version__: String = "V0.8.0" from numojo.core.ndarray import NDArray from numojo.core.ndshape import NDArrayShape, Shape from numojo.core.ndstrides import NDArrayStrides, Strides -from numojo.core.item import Item, item +from numojo.core.item import Item from numojo.core.matrix import Matrix from numojo.core.complex.complex_simd import ComplexSIMD, CScalar from numojo.core.complex.complex_ndarray import ComplexNDArray diff --git a/numojo/core/__init__.mojo b/numojo/core/__init__.mojo index dcf4fc46..5d70c5d7 100644 --- a/numojo/core/__init__.mojo +++ b/numojo/core/__init__.mojo @@ -67,5 +67,4 @@ from .error import ( ) alias idx = Item -alias shape = NDArrayShape alias Shape = NDArrayShape diff --git a/numojo/core/array_methods.mojo b/numojo/core/array_methods.mojo new file mode 100644 index 00000000..114294cf --- /dev/null +++ b/numojo/core/array_methods.mojo @@ -0,0 +1,93 @@ +comptime newaxis: NewAxis = NewAxis() +comptime ellipsis: Ellipsis = Ellipsis() + + +struct Ellipsis(Stringable): + """ + Represents an ellipsis (`...`) used in array slicing to indicate the inclusion of all remaining dimensions. + This can be used to simplify slicing operations when the exact number of dimensions is not known or when you want to include all remaining dimensions without explicitly specifying them. + + Example: + ``` + from numojo.prelude import * + from numojo.routines.creation import arange + + var arr = arange(Shape(3, 4, 5, 6)) + sliced_arr = arr[nm.ellipsis, 2] # Equivalent to arr[:, :, :, 2] + ``` + """ + + fn __init__(out self): + """ + Initializes an Ellipsis instance. + """ + pass + + fn __repr__(self) -> String: + """ + Returns a string representation of the Ellipsis instance. + + Returns: + Str: The string "Ellipsis()". + """ + return "numojo.ellipsis()" + + fn __str__(self) -> String: + """ + Returns a string representation of the Ellipsis instance. + + Returns: + Str: The string "Ellipsis()". + """ + return "numojo.ellipsis()" + + fn __eq__(self, other: Self) -> Bool: + """ + Checks equality between two Ellipsis instances. + """ + return True + + fn __ne__(self, other: Self) -> Bool: + """ + Checks inequality between two Ellipsis instances. + """ + return False + + +# TODO: add an initializer with int field to specify number of new axes to add! +struct NewAxis(Stringable): + fn __init__(out self): + """ + Initializes a NewAxis instance. + """ + pass + + fn __repr__(self) -> String: + """ + Returns a string representation of the NewAxis instance. + + Returns: + Str: The string "NewAxis()". + """ + return "numojo.newaxis()" + + fn __str__(self) -> String: + """ + Returns a string representation of the NewAxis instance. + + Returns: + Str: The string "NewAxis()". + """ + return "numojo.newaxis()" + + fn __eq__(self, other: Self) -> Bool: + """ + Checks equality between two NewAxis instances. + """ + return True + + fn __ne__(self, other: Self) -> Bool: + """ + Checks inequality between two NewAxis instances. + """ + return False diff --git a/numojo/core/complex/complex_ndarray.mojo b/numojo/core/complex/complex_ndarray.mojo index c0518ceb..92c9e1f9 100644 --- a/numojo/core/complex/complex_ndarray.mojo +++ b/numojo/core/complex/complex_ndarray.mojo @@ -4285,7 +4285,7 @@ struct _ComplexNDArrayIter[ for offset in range(self.size_of_item): var remainder = offset - var item = Item(ndim=self.ndim, initialized=False) + var item: Item = Item(ndim=self.ndim) for i in range(self.ndim - 1, -1, -1): if i != self.dimension: @@ -4348,7 +4348,7 @@ struct _ComplexNDArrayIter[ for offset in range(self.size_of_item): var remainder = offset - var item = Item(ndim=self.ndim, initialized=False) + var item: Item = Item(ndim=self.ndim) for i in range(self.ndim - 1, -1, -1): if i != self.dimension: diff --git a/numojo/core/complex/complex_simd.mojo b/numojo/core/complex/complex_simd.mojo index 96e67c7b..bf85db1c 100644 --- a/numojo/core/complex/complex_simd.mojo +++ b/numojo/core/complex/complex_simd.mojo @@ -940,6 +940,60 @@ struct ComplexSIMD[cdtype: ComplexDType = ComplexDType.float64, width: Int = 1]( """ return Self(-self.re, -self.im) + fn __invert__( + self, + ) -> Self where cdtype == ComplexDType.bool or cdtype.is_integral(): + """ + Element-wise logical NOT operation on this ComplexSIMD instance. + + Returns: + ComplexSIMD instance where each lane is the logical NOT of the corresponding lane. + """ + return Self(~self.re, ~self.im) + + # --- Comparison operators --- + fn __and__( + self, other: Self + ) -> Self where cdtype == ComplexDType.bool or cdtype.is_integral(): + """ + Element-wise logical AND operation between two ComplexSIMD instances. + + Args: + other: Another ComplexSIMD instance. + + Returns: + True if both the real and imaginary parts are non-zero for all lanes, otherwise False. + """ + return Self(self.re & other.re, self.im & other.im) + + fn __or__( + self, other: Self + ) -> Self where cdtype == ComplexDType.bool or cdtype.is_integral(): + """ + Element-wise logical OR operation between two ComplexSIMD instances. + + Args: + other: Another ComplexSIMD instance. + + Returns: + True if either the real or imaginary part is non-zero for any lane, otherwise False. + """ + return Self(self.re | other.re, self.im | other.im) + + fn __xor__( + self, other: Self + ) -> Self where cdtype == ComplexDType.bool or cdtype.is_integral(): + """ + Element-wise logical XOR operation between two ComplexSIMD instances. + + Args: + other: Another ComplexSIMD instance. + + Returns: + True if exactly one of the real or imaginary parts is non-zero for any lane, otherwise False. + """ + return Self(self.re ^ other.re, self.im ^ other.im) + # --- Helpers --- @staticmethod @always_inline diff --git a/numojo/core/data_container.mojo b/numojo/core/data_container.mojo index 3f96a6cf..612179ff 100644 --- a/numojo/core/data_container.mojo +++ b/numojo/core/data_container.mojo @@ -11,6 +11,14 @@ from memory import UnsafePointer, LegacyUnsafePointer # temporary DataContainer to support transition from LegacyUnsafePointer to UnsafePointer. struct DataContainerNew[dtype: DType, origin: MutOrigin](ImplicitlyCopyable): + """ + DataContainer is managing a contiguous block of memory containing elements of type `Scalar[dtype]`, using an `UnsafePointer` with a specified mutability origin. It provides basic memory management and pointer access for low-level array operations. + + Type Parameters: + dtype: The data type of the elements stored in the container. + origin: The mutability origin for the pointer, controlling aliasing and mutation semantics. + """ + var ptr: UnsafePointer[Scalar[dtype], origin] fn __init__(out self, size: Int): diff --git a/numojo/core/item.mojo b/numojo/core/item.mojo index 57a33d42..251b3ab0 100644 --- a/numojo/core/item.mojo +++ b/numojo/core/item.mojo @@ -5,9 +5,9 @@ Implements Item type. """ from builtin.type_aliases import Origin -from builtin.int import index as index_int -from memory import memset_zero, memcpy -from memory import LegacyUnsafePointer as UnsafePointer +from builtin.int import index as convert_to_int +from memory import memcpy, memset_zero +from memory import UnsafePointer from memory import memcmp from os import abort from sys import simd_width_of @@ -18,25 +18,44 @@ from numojo.core.traits.indexer_collection_element import ( IndexerCollectionElement, ) -# simple alias for users. Use `Item` internally. -alias item = Item - @register_passable struct Item( ImplicitlyCopyable, Movable, Representable, Sized, Stringable, Writable ): """ - Specifies the indices of an item of an array. + Represents a multi-dimensional index for array access. + + The `Item` struct is used to specify the coordinates of an element within an N-dimensional array. + For example, `arr[Item(1, 2, 3)]` retrieves the element at position (1, 2, 3) in a 3D array. + + Each `Item` instance holds a sequence of integer indices, one for each dimension of the array. + This allows for precise and flexible indexing into arrays of arbitrary dimensionality. + + Example: + ```mojo + from numojo.prelude import * + import numojo as nm + var arr = nm.arange[f32](0, 27).reshape(Shape(3, 3, 3)) + var value = arr[Item(1, 2, 3)] # Accesses arr[1, 2, 3] + ``` + + Fields: + _buf: Pointer to the buffer storing the indices. + _ndim: Number of dimensions (length of the index tuple). """ # Aliases - alias _type: DType = DType.int + alias element_type: DType = DType.int + """The data type of the Item elements.""" + alias _origin: MutOrigin = MutOrigin.external + """Internal origin of the Item instance.""" # Fields - var _buf: UnsafePointer[Scalar[Self._type]] + var _buf: UnsafePointer[Scalar[Self.element_type], Self._origin] var ndim: Int + # add constraint for ndim >= 0 for Item instance. @always_inline("nodebug") fn __init__[T: Indexer](out self, *args: T): """Construct the tuple. @@ -47,10 +66,10 @@ struct Item( Args: args: Initial values. """ - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(args.__len__()) - self.ndim = args.__len__() - for i in range(args.__len__()): - self._buf[i] = index(args[i]) + self._buf = alloc[Scalar[Self.element_type]](len(args)) + self.ndim = len(args) + for i in range(len(args)): + (self._buf + i).init_pointee_copy(convert_to_int(args[i])) @always_inline("nodebug") fn __init__[T: IndexerCollectionElement](out self, args: List[T]) raises: @@ -63,9 +82,9 @@ struct Item( args: Initial values. """ self.ndim = len(args) - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) for i in range(self.ndim): - (self._buf + i).init_pointee_copy(index(args[i])) + (self._buf + i).init_pointee_copy(convert_to_int(args[i])) @always_inline("nodebug") fn __init__(out self, args: VariadicList[Int]) raises: @@ -75,53 +94,20 @@ struct Item( args: Initial values. """ self.ndim = len(args) - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](len(args)) for i in range(self.ndim): - (self._buf + i).init_pointee_copy(Int(args[i])) + (self._buf + i).init_pointee_copy(args[i]) - fn __init__( - out self, - *, - ndim: Int, - initialized: Bool, - ) raises: - """ - Construct Item with number of dimensions. - This method is useful when you want to create a Item with given ndim - without knowing the Item values. + @always_inline("nodebug") + fn __init__(out self, ndim: Int): + """Construct the Item with given length and initialize to zero. Args: - ndim: Number of dimensions. - initialized: Whether the shape is initialized. - If yes, the values will be set to 0. - If no, the values will be uninitialized. - - Raises: - Error: If the number of dimensions is negative. + ndim: The length of the tuple. """ - if ndim < 0: - raise Error( - IndexError( - message=String( - "Invalid ndim: got {}; must be >= 0." - ).format(ndim), - suggestion=String( - "Pass a non-negative dimension count when constructing" - " Item." - ), - location=String("Item.__init__(ndim: Int)"), - ) - ) - - if ndim == 0: - self.ndim = 0 - self._buf = UnsafePointer[Scalar[Self._type]]() - else: - self.ndim = ndim - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(ndim) - if initialized: - for i in range(ndim): - (self._buf + i).init_pointee_copy(0) + self.ndim = ndim + self._buf = alloc[Scalar[Self.element_type]](ndim) + memset_zero(self._buf, ndim) @always_inline("nodebug") fn __copyinit__(out self, other: Self): @@ -131,7 +117,7 @@ struct Item( other: The tuple to copy. """ self.ndim = other.ndim - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) memcpy(dest=self._buf, src=other._buf, count=self.ndim) @always_inline("nodebug") @@ -158,10 +144,10 @@ struct Item( Returns: The normalized index. """ - var normalized_idx: Int = index - if normalized_idx < 0: - normalized_idx += self.ndim - return normalized_idx + var norm_idx: Int = index + if norm_idx < 0: + norm_idx += self.ndim + return norm_idx @always_inline("nodebug") fn __getitem__[T: Indexer](self, idx: T) raises -> Int: @@ -176,12 +162,12 @@ struct Item( Returns: The value at the specified index. """ - var index: Int = index_int(idx) + var index: Int = convert_to_int(idx) if index >= self.ndim or index < -self.ndim: raise Error( IndexError( message=String("Index {} out of range [{} , {}).").format( - index_int(idx), -self.ndim, self.ndim + convert_to_int(idx), -self.ndim, self.ndim ), suggestion=String( "Use indices in [-ndim, ndim) (negative indices wrap)." @@ -189,7 +175,7 @@ struct Item( location=String("Item.__getitem__"), ) ) - var normalized_idx: Int = self.normalize_index(index_int(idx)) + var normalized_idx: Int = self.normalize_index(convert_to_int(idx)) return Int(self._buf[normalized_idx]) @always_inline("nodebug") @@ -220,11 +206,18 @@ struct Item( var length = updated_slice[2] if length <= 0: - var empty_result = Self(ndim=0, initialized=False) - return empty_result + raise Error( + ShapeError( + message="Provided slice results in an empty Item.", + suggestion=( + "Adjust slice parameters to obtain non-empty result." + ), + location="Item.__getitem__(self, slice_list: Slice)", + ) + ) - var result = Self(ndim=length, initialized=False) - var idx = start + var result: Item = Self(ndim=length) + var idx: Int = start for i in range(length): (result._buf + i).init_pointee_copy(self._buf[idx]) idx += step @@ -242,16 +235,12 @@ struct Item( idx: The index of the value to set. val: The value to set. """ - - var normalized_idx: Int = index_int(idx) - if normalized_idx < 0: - normalized_idx = index_int(idx) + self.ndim - - if normalized_idx < 0 or normalized_idx >= self.ndim: + var norm_idx: Int = self.normalize_index(convert_to_int(idx)) + if norm_idx < 0 or norm_idx >= self.ndim: raise Error( IndexError( message=String("Index {} out of range [{} , {}).").format( - index_int(idx), -self.ndim, self.ndim + convert_to_int(idx), -self.ndim, self.ndim ), suggestion=String( "Use indices in [-ndim, ndim) (negative indices wrap)." @@ -260,7 +249,7 @@ struct Item( ) ) - self._buf[normalized_idx] = index(val) + self._buf[norm_idx] = index(val) fn __iter__(self) raises -> _ItemIter: """Iterate over elements of the NDArray, returning copied value. @@ -278,7 +267,7 @@ struct Item( ) fn __repr__(self) -> String: - var result: String = "numojo.Item" + String(self) + var result: String = "Item" + String(self) return result fn __str__(self) -> String: @@ -286,18 +275,12 @@ struct Item( for i in range(self.ndim): result += String(self._buf[i]) if i < self.ndim - 1: - result += "," + result += ", " result += ")" return result fn write_to[W: Writer](self, mut writer: W): - writer.write( - "Item at index: " - + String(self) - + " " - + "Length: " - + String(self.ndim) - ) + writer.write("Coordinates: " + String(self) + " ") @always_inline("nodebug") fn __eq__(self, other: Self) -> Bool: @@ -348,16 +331,15 @@ struct Item( # ===-------------------------------------------------------------------===# # Other methods # ===-------------------------------------------------------------------===# - @always_inline("nodebug") - fn copy(read self) raises -> Self: + fn deep_copy(read self) raises -> Self: """ Returns a deep copy of the item. Returns: A new Item with the same values. """ - var res = Self(ndim=self.ndim, initialized=False) + var res: Item = Item(ndim=self.ndim) memcpy(dest=res._buf, src=self._buf, count=self.ndim) return res^ @@ -372,7 +354,7 @@ struct Item( Returns: A new item with the given axes swapped. """ - var res: Self = Self(ndim=self.ndim, initialized=False) + var res: Item = Item(ndim=self.ndim) memcpy(dest=res._buf, src=self._buf, count=self.ndim) res[axis1] = self[axis2] res[axis2] = self[axis1] @@ -402,7 +384,7 @@ struct Item( for i in range(len(others)): total_dims += others[i].ndim - var new_item: Self = Self(ndim=total_dims, initialized=False) + var new_item: Item = Item(ndim=total_dims) var index: UInt = 0 for i in range(self.ndim): @@ -436,11 +418,11 @@ struct Item( print(item._flip()) # Item: (3, 2, 1) ``` """ - var result: Self = Self(ndim=self.ndim, initialized=False) - memcpy(dest=result._buf, src=self._buf, count=self.ndim) - for i in range(result.ndim): - result._buf[i] = self._buf[self.ndim - 1 - i] - return result^ + var res: Item = Item(ndim=self.ndim) + memcpy(dest=res._buf, src=self._buf, count=self.ndim) + for i in range(res.ndim): + res._buf[i] = self._buf[self.ndim - 1 - i] + return res^ fn _move_axis_to_end(self, var axis: Int) raises -> Self: """ @@ -464,17 +446,17 @@ struct Item( if axis < 0: axis += self.ndim - var result: Self = Self(ndim=self.ndim, initialized=False) - memcpy(dest=result._buf, src=self._buf, count=self.ndim) + var res: Item = Item(ndim=self.ndim) + memcpy(dest=res._buf, src=self._buf, count=self.ndim) if axis == self.ndim - 1: - return result^ + return res^ - var value: Scalar[Self._type] = result._buf[axis] - for i in range(axis, result.ndim - 1): - result._buf[i] = result._buf[i + 1] - result._buf[result.ndim - 1] = value - return result^ + var value: Scalar[Self.element_type] = res._buf[axis] + for i in range(axis, res.ndim - 1): + res._buf[i] = res._buf[i + 1] + res._buf[res.ndim - 1] = value + return res^ fn _pop(self, axis: Int) raises -> Self: """ @@ -487,7 +469,7 @@ struct Item( Returns: A new item with the item at the given axis (index) dropped. """ - var res: Self = Self(ndim=self.ndim - 1, initialized=False) + var res: Item = Item(ndim=self.ndim - 1) memcpy(dest=res._buf, src=self._buf, count=axis) memcpy( dest=res._buf + axis, @@ -516,7 +498,7 @@ struct Item( ``` """ var total_dims: Int = self.ndim + len(values) - var new_item: Self = Self(ndim=total_dims, initialized=False) + var new_item: Item = Item(ndim=total_dims) var offset: UInt = 0 for i in range(self.ndim): @@ -601,7 +583,9 @@ struct Item( return (start, step, length) - fn load[width: Int = 1](self, idx: Int) raises -> SIMD[Self._type, width]: + fn load[ + width: Int = 1 + ](self, idx: Int) raises -> SIMD[Self.element_type, width]: """ Load a SIMD vector from the Item at the specified index. @@ -634,7 +618,7 @@ struct Item( fn store[ width: Int = 1 - ](self, idx: Int, value: SIMD[Self._type, width]) raises: + ](self, idx: Int, value: SIMD[Self.element_type, width]) raises: """ Store a SIMD vector into the Item at the specified index. @@ -663,7 +647,9 @@ struct Item( self._buf.store[width=width](idx, value) - fn unsafe_load[width: Int = 1](self, idx: Int) -> SIMD[Self._type, width]: + fn unsafe_load[ + width: Int = 1 + ](self, idx: Int) -> SIMD[Self.element_type, width]: """ Unsafely load a SIMD vector from the Item at the specified index. @@ -680,7 +666,7 @@ struct Item( fn unsafe_store[ width: Int = 1 - ](self, idx: Int, value: SIMD[Self._type, width]): + ](self, idx: Int, value: SIMD[Self.element_type, width]): """ Unsafely store a SIMD vector into the Item at the specified index. @@ -712,7 +698,7 @@ struct _ItemIter[ item: Item, length: Int, ): - self.index = 0 if forward else length + self.index = 0 if forward else length - 1 self.length = length self.item = item @@ -724,7 +710,7 @@ struct _ItemIter[ if forward: return self.index < self.length else: - return self.index > 0 + return self.index >= 0 fn __next__(mut self) raises -> Scalar[DType.int]: @parameter diff --git a/numojo/core/matrix.mojo b/numojo/core/matrix.mojo index 41e1e6b5..aa7d6d09 100644 --- a/numojo/core/matrix.mojo +++ b/numojo/core/matrix.mojo @@ -396,13 +396,22 @@ struct MatrixBase[ data: DataContainer[dtype, origin], ) where own_data == False: """ - Initialize Matrix that does not own the data. - The data is owned by another Matrix. + Initialize a non-owning `MatrixView`. + + This constructor creates a Matrix instance that acts as a view into an + existing data buffer. The view does not allocate or manage memory; it + references data owned by another Matrix. It is an unsafe operation and should not be called by users directly. Args: - shape: Shape of the view. - strides: Strides of the view. - data: DataContainer that holds the data buffer. + shape: A tuple representing the dimensions of the view as (rows, columns). + strides: A tuple representing the memory strides for accessing elements in the view. Strides determine how to traverse the data buffer to access elements in the matrix. + data: A DataContainer instance that holds the data buffer being referenced. + + Notes: + - This constructor is intended for internal use to create views into existing matrices! Users should not call this directly. + - The view does not own the data and relies on the lifetime of the + original data owner. + - Modifications to the view affect the original data by default. """ self.shape = shape self.strides = strides @@ -1633,6 +1642,37 @@ struct MatrixBase[ """ self._buf.ptr.store(idx, val) + fn store[width: Int = 1](self, idx: Int, val: SIMD[dtype, width]) raises: + """ + Store a SIMD element into the matrix at the specified linear index. + + Parameters: + width: The width of the SIMD element to store. Defaults to 1. + + Args: + idx: The linear index where the element will be stored. Negative indices are supported and follow Python conventions. + val: The SIMD element to store at the given index. + + Raises: + Error: If the provided index is out of bounds. + + Example: + ```mojo + from numojo.prelude import * + var mat = Matrix.ones(shape=(4, 4)) + var simd_element = SIMD[f64, 4](2.0, 2.0, 2.0, 2.0) + mat.store[4](2, simd_element) # Store a SIMD element of width 4 at index 2 + ``` + """ + if idx >= self.size or idx < -self.size: + raise Error( + String("Index {} exceed the matrix size {}").format( + idx, self.size + ) + ) + var idx_norm = self.normalize(idx, self.size) + self._buf.ptr.store[width=width](idx_norm, val) + # ===-------------------------------------------------------------------===# # Other dunders and auxiliary methods # ===-------------------------------------------------------------------===# diff --git a/numojo/core/ndarray.mojo b/numojo/core/ndarray.mojo index 5e37e684..b4831f7d 100644 --- a/numojo/core/ndarray.mojo +++ b/numojo/core/ndarray.mojo @@ -84,6 +84,7 @@ from numojo.core.error import ( ValueError, ArithmeticError, ) +from numojo.core.array_methods import ellipsis, newaxis # ===----------------------------------------------------------------------===# # === numojo routines (creation / io / logic) === @@ -5429,7 +5430,7 @@ struct _NDArrayIter[ for offset in range(self.size_of_item): var remainder = offset - var item = Item(ndim=self.ndim, initialized=False) + var item = Item(ndim=self.ndim) for i in range(self.ndim - 1, -1, -1): if i != self.dimension: @@ -5484,7 +5485,7 @@ struct _NDArrayIter[ for offset in range(self.size_of_item): var remainder = offset - var item = Item(ndim=self.ndim, initialized=False) + var item: Item = Item(ndim=self.ndim) for i in range(self.ndim - 1, -1, -1): if i != self.dimension: @@ -5652,7 +5653,7 @@ struct _NDAxisIter[ self.index -= 1 var remainder = current_index * self.size_of_item - var item = Item(ndim=self.ndim, initialized=False) + var item: Item = Item(ndim=self.ndim) if self.order == "C": for i in range(self.ndim): @@ -5714,7 +5715,7 @@ struct _NDAxisIter[ var elements: NDArray[dtype] = NDArray[dtype](Shape(self.size_of_item)) var remainder: Int = index * self.size_of_item - var item: Item = Item(ndim=self.ndim, initialized=True) + var item: Item = Item(ndim=self.ndim) if self.order == "C": for i in range(self.ndim): @@ -5781,7 +5782,7 @@ struct _NDAxisIter[ ) var remainder: Int = index * self.size_of_item - var item: Item = Item(ndim=self.ndim, initialized=True) + var item: Item = Item(ndim=self.ndim) for i in range(self.axis): item._buf[i] = remainder // self.strides_compatible[i] remainder %= self.strides_compatible[i] @@ -5893,7 +5894,7 @@ struct _NDIter[is_mutable: Bool, //, origin: Origin[is_mutable], dtype: DType]( self.index += 1 var remainder = current_index - var indices = Item(ndim=self.ndim, initialized=False) + var indices = Item(ndim=self.ndim) if self.order == "C": for i in range(self.ndim): @@ -5935,7 +5936,7 @@ struct _NDIter[is_mutable: Bool, //, origin: Origin[is_mutable], dtype: DType]( ) var remainder = index - var indices = Item(ndim=self.ndim, initialized=False) + var indices = Item(ndim=self.ndim) if self.order == "C": for i in range(self.ndim): diff --git a/numojo/core/ndshape.mojo b/numojo/core/ndshape.mojo index 65831193..1e3d8c53 100644 --- a/numojo/core/ndshape.mojo +++ b/numojo/core/ndshape.mojo @@ -9,7 +9,7 @@ Implements NDArrayShape type. """ from memory import memcpy, memcmp -from memory import LegacyUnsafePointer as UnsafePointer +from memory import UnsafePointer from numojo.core.error import IndexError, ShapeError, ValueError @@ -19,7 +19,7 @@ alias Shape = NDArrayShape @register_passable struct NDArrayShape( - ImplicitlyCopyable, Movable, Sized, Stringable & Representable, Writable + ImplicitlyCopyable, Movable, Representable, Sized, Stringable, Writable ): """ Presents the shape of `NDArray` type. @@ -29,13 +29,30 @@ struct NDArrayShape( The elements of the shape must be positive. The number of dimension and values of elements are checked upon creation of the shape. + + Example: + ```mojo + import numojo as nm + var shape1 = nm.Shape(2, 3, 4) + print(shape1) # Shape: (2,3,4) + var shape2 = nm.Shape([5, 6, 7]) + print(shape2) # Shape: (5,6,7) + + Fields: + _buf: UnsafePointer[Scalar[DType.int]] + Data buffer. + _ndim: Int + Number of dimensions of array. It must be larger than 0. """ # Aliases - alias _type: DType = DType.int + alias element_type: DType = DType.int + """The data type of the NDArrayShape elements.""" + alias _origin: MutOrigin = MutOrigin.external + """Internal origin of the NDArrayShape instance.""" # Fields - var _buf: UnsafePointer[Scalar[Self._type]] + var _buf: UnsafePointer[Scalar[Self.element_type], Self._origin] """Data buffer.""" var ndim: Int """Number of dimensions of array. It must be larger than 0.""" @@ -64,7 +81,7 @@ struct NDArrayShape( ) self.ndim = 1 - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(shape) + self._buf = alloc[Scalar[Self.element_type]](shape) self._buf.init_pointee_copy(shape) @always_inline("nodebug") @@ -90,7 +107,7 @@ struct NDArrayShape( ) self.ndim = len(shape) - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) for i in range(self.ndim): if shape[i] < 1: raise Error( @@ -132,7 +149,7 @@ struct NDArrayShape( ) ) self.ndim = len(shape) - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) for i in range(self.ndim): if shape[i] < 1: raise Error( @@ -187,7 +204,7 @@ struct NDArrayShape( ) ) self.ndim = len(shape) - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) for i in range(self.ndim): if shape[i] < 0: raise Error( @@ -233,7 +250,7 @@ struct NDArrayShape( ) self.ndim = len(shape) - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) for i in range(self.ndim): if shape[i] < 1: raise Error( @@ -292,7 +309,7 @@ struct NDArrayShape( ) self.ndim = len(shape) - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) for i in range(self.ndim): if shape[i] < 1: raise Error( @@ -341,7 +358,7 @@ struct NDArrayShape( ) self.ndim = len(shape) - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) for i in range(self.ndim): if shape[i] < 1: raise Error( @@ -388,7 +405,7 @@ struct NDArrayShape( shape: Another NDArrayShape to initialize from. """ self.ndim = shape.ndim - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(shape.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) memcpy(dest=self._buf, src=shape._buf, count=shape.ndim) for i in range(self.ndim): (self._buf + i).init_pointee_copy(shape._buf[i]) @@ -430,13 +447,15 @@ struct NDArrayShape( ) if ndim == 0: - # This is a 0darray (numojo scalar) + # This denotes a 0darray (numojo scalar) self.ndim = ndim - self._buf = UnsafePointer[Scalar[Self._type]]() - + self._buf = alloc[Scalar[Self.element_type]]( + 1 + ) # allocate 1 element to avoid null pointer + self._buf.init_pointee_copy(0) else: self.ndim = ndim - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(ndim) + self._buf = alloc[Scalar[Self.element_type]](ndim) if initialized: for i in range(ndim): (self._buf + i).init_pointee_copy(1) @@ -453,9 +472,9 @@ struct NDArrayShape( Example: ```mojo - import numojo as nm - var shape = nm.Shape(2, 3, 4) - var strides = nm.Strides.row_major(shape) + from numojo.prelude import * + var shape = Shape(2, 3, 4) + var strides = shape.row_major() print(strides) # Strides: (12, 4, 1) ``` """ @@ -473,9 +492,9 @@ struct NDArrayShape( Example: ```mojo - import numojo as nm - var shape = nm.Shape(2, 3, 4) - var strides = nm.Strides.col_major(shape) + from numojo.prelude import * + var shape = Shape(2, 3, 4) + var strides = shape.col_major() print(strides) # Strides: (1, 2, 6) ``` """ @@ -491,16 +510,34 @@ struct NDArrayShape( other: Another NDArrayShape to initialize from. """ self.ndim = other.ndim - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(other.ndim) - memcpy(dest=self._buf, src=other._buf, count=other.ndim) + if other.ndim == 0: + self._buf = alloc[Scalar[Self.element_type]](1) + self._buf.init_pointee_copy(0) + else: + self._buf = alloc[Scalar[Self.element_type]](other.ndim) + memcpy(dest=self._buf, src=other._buf, count=other.ndim) + + @always_inline("nodebug") + fn deep_copy(self) raises -> Self: + """ + Returns a deep copy of the NDArrayShape. + + Returns: + A new NDArrayShape with the same values and new origin. + """ + var res: NDArrayShape = NDArrayShape(ndim=self.ndim, initialized=True) + memcpy(dest=res._buf, src=self._buf, count=self.ndim) + return res^ fn __del__(deinit self): """ Destructor for NDArrayShape. - Frees the allocated memory for the data buffer. + Frees the allocated memory for the data buffer of the shape. + + Notes: + Even when ndim is 0, the buffer is still allocated with 1 element to avoid null pointer, so it needs to be freed here. """ - if self.ndim > 0: - self._buf.free() + self._buf.free() fn normalize_index(self, index: Int) -> Int: """ @@ -631,7 +668,7 @@ struct NDArrayShape( return result^ @always_inline("nodebug") - fn __setitem__(mut self, index: Int, val: Scalar[Self._type]) raises: + fn __setitem__(mut self, index: Int, val: Scalar[Self.element_type]) raises: """ Sets shape at specified index. @@ -678,7 +715,7 @@ struct NDArrayShape( Returns: String representation of the shape of the array. """ - return "numojo.Shape" + String(self) + return "numojo.Shape" + self.__str__() @always_inline("nodebug") fn __str__(self) -> String: @@ -698,7 +735,7 @@ struct NDArrayShape( fn write_to[W: Writer](self, mut writer: W): writer.write( - "Shape: " + String(self) + " " + "ndim: " + String(self.ndim) + "Shape: " + self.__str__() + " " + "ndim: " + String(self.ndim) ) @always_inline("nodebug") @@ -810,10 +847,10 @@ struct NDArrayShape( Returns: The total number of elements in the corresponding array. """ - var size: Scalar[Self._type] = 1 + var size_of_arr: Scalar[Self.element_type] = 1 for i in range(self.ndim): - size *= self._buf[i] - return Int(size) + size_of_arr *= self._buf[i] + return Int(size_of_arr) fn swapaxes(self, axis1: Int, axis2: Int) raises -> Self: """ @@ -933,7 +970,9 @@ struct NDArrayShape( ) return res^ - fn load[width: Int = 1](self, idx: Int) raises -> SIMD[Self._type, width]: + fn load[ + width: Int = 1 + ](self, idx: Int) raises -> SIMD[Self.element_type, width]: """ Load a SIMD vector from the Shape at the specified index. @@ -966,7 +1005,7 @@ struct NDArrayShape( fn store[ width: Int = 1 - ](self, idx: Int, value: SIMD[Self._type, width]) raises: + ](self, idx: Int, value: SIMD[Self.element_type, width]) raises: """ Store a SIMD vector into the Shape at the specified index. @@ -995,7 +1034,9 @@ struct NDArrayShape( self._buf.store[width=width](idx, value) - fn unsafe_load[width: Int = 1](self, idx: Int) -> SIMD[Self._type, width]: + fn unsafe_load[ + width: Int = 1 + ](self, idx: Int) -> SIMD[Self.element_type, width]: """ Unsafely load a SIMD vector from the Shape at the specified index. @@ -1012,7 +1053,7 @@ struct NDArrayShape( fn unsafe_store[ width: Int = 1 - ](self, idx: Int, value: SIMD[Self._type, width]): + ](self, idx: Int, value: SIMD[Self.element_type, width]): """ Unsafely store a SIMD vector into the Shape at the specified index. diff --git a/numojo/core/ndstrides.mojo b/numojo/core/ndstrides.mojo index 83f0db33..15589b6b 100644 --- a/numojo/core/ndstrides.mojo +++ b/numojo/core/ndstrides.mojo @@ -9,11 +9,10 @@ Implements NDArrayStrides type. """ from memory import memcmp, memcpy -from memory import LegacyUnsafePointer as UnsafePointer +from memory import UnsafePointer from numojo.core.error import IndexError, ValueError -alias strides = NDArrayStrides alias Strides = NDArrayStrides """An alias of the NDArrayStrides.""" @@ -31,14 +30,44 @@ struct NDArrayStrides( """ # Aliases - alias _type: DType = DType.int + alias element_type: DType = DType.int + """The data type of the NDArrayStrides elements.""" + alias _origin: MutOrigin = MutOrigin.external + """Internal origin of the NDArrayStrides instance.""" # Fields - var _buf: UnsafePointer[Scalar[Self._type]] + var _buf: UnsafePointer[Scalar[Self.element_type], Self._origin] """Data buffer.""" var ndim: Int """Number of dimensions of array. It must be larger than 0.""" + @always_inline("nodebug") + fn __init__(out self, shape: Int) raises: + """ + Initializes the NDArrayStrides with one dimension. + + Raises: + Error: If the shape is not positive. + + Args: + shape: Size of the array. + """ + + if shape < 1: + raise Error( + ShapeError( + message=String( + "Stride value must be positive, got {}." + ).format(shape), + suggestion="Use positive integers for stride value.", + location="NDArrayStrides.__init__(shape: Int)", + ) + ) + + self.ndim = 1 + self._buf = alloc[Scalar[Self.element_type]](shape) + self._buf.init_pointee_copy(shape) + @always_inline("nodebug") fn __init__(out self, *strides: Int) raises: """ @@ -62,7 +91,7 @@ struct NDArrayStrides( ) self.ndim = len(strides) - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) for i in range(self.ndim): (self._buf + i).init_pointee_copy(strides[i]) @@ -89,7 +118,7 @@ struct NDArrayStrides( ) self.ndim = len(strides) - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) for i in range(self.ndim): (self._buf + i).init_pointee_copy(strides[i]) @@ -118,7 +147,7 @@ struct NDArrayStrides( ) self.ndim = len(strides) - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) for i in range(self.ndim): (self._buf + i).init_pointee_copy(strides[i]) @@ -133,7 +162,7 @@ struct NDArrayStrides( """ self.ndim = strides.ndim - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(self.ndim) + self._buf = alloc[Scalar[Self.element_type]](self.ndim) memcpy(dest=self._buf, src=strides._buf, count=strides.ndim) @always_inline("nodebug") @@ -156,7 +185,7 @@ struct NDArrayStrides( """ self.ndim = shape.ndim - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(shape.ndim) + self._buf = alloc[Scalar[Self.element_type]](shape.ndim) if order == "C": var temp = 1 @@ -272,11 +301,12 @@ struct NDArrayStrides( if ndim == 0: # This is a 0darray (numojo scalar) self.ndim = ndim - self._buf = UnsafePointer[Scalar[Self._type]]() + self._buf = alloc[Scalar[Self.element_type]](1) + self._buf.init_pointee_copy(0) else: self.ndim = ndim - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(ndim) + self._buf = alloc[Scalar[Self.element_type]](ndim) if initialized: for i in range(ndim): (self._buf + i).init_pointee_copy(0) @@ -298,7 +328,8 @@ struct NDArrayStrides( Example: ```mojo import numojo as nm - var shape = nm.Shape(2, 3, 4) + from numojo.prelude import * + var shape = Shape(2, 3, 4) var strides = nm.Strides.row_major(shape) print(strides) # Strides: (12, 4, 1) ``` @@ -322,7 +353,8 @@ struct NDArrayStrides( Example: ```mojo import numojo as nm - var shape = nm.Shape(2, 3, 4) + from numojo.prelude import * + var shape = Shape(2, 3, 4) var strides = nm.Strides.col_major(shape) print(strides) # Strides: (1, 2, 6) ``` @@ -339,16 +371,36 @@ struct NDArrayStrides( other: Strides of the array. """ self.ndim = other.ndim - self._buf = UnsafePointer[Scalar[Self._type]]().alloc(other.ndim) - memcpy(dest=self._buf, src=other._buf, count=other.ndim) + if other.ndim == 0: + self._buf = alloc[Scalar[Self.element_type]](1) + self._buf.init_pointee_copy(0) + else: + self._buf = alloc[Scalar[Self.element_type]](other.ndim) + memcpy(dest=self._buf, src=other._buf, count=other.ndim) + + @always_inline("nodebug") + fn deep_copy(self) raises -> Self: + """ + Returns a deep copy of the NDArrayStride. + + Returns: + A new NDArrayStride with the same values and new origin. + """ + var res: NDArrayStrides = NDArrayStrides( + ndim=self.ndim, initialized=True + ) + memcpy(dest=res._buf, src=self._buf, count=self.ndim) + return res^ fn __del__(deinit self): """ Destructor for NDArrayStrides. Frees the allocated memory for the data buffer. + + Notes: + Even when ndim is 0, the buffer is still allocated with 1 element to avoid null pointer, so it needs to be freed here. """ - if self.ndim > 0: - self._buf.free() + self._buf.free() fn normalize_index(self, index: Int) -> Int: """ @@ -504,7 +556,7 @@ struct NDArrayStrides( return result^ @always_inline("nodebug") - fn __setitem__(mut self, index: Int, val: Scalar[Self._type]) raises: + fn __setitem__(mut self, index: Int, val: Scalar[Self.element_type]) raises: """ Sets stride at specified index. @@ -833,7 +885,9 @@ struct NDArrayStrides( ) return res - fn load[width: Int = 1](self, idx: Int) raises -> SIMD[Self._type, width]: + fn load[ + width: Int = 1 + ](self, idx: Int) raises -> SIMD[Self.element_type, width]: """ Load a SIMD vector from the Strides at the specified index. @@ -866,7 +920,7 @@ struct NDArrayStrides( fn store[ width: Int = 1 - ](self, idx: Int, value: SIMD[Self._type, width]) raises: + ](self, idx: Int, value: SIMD[Self.element_type, width]) raises: """ Store a SIMD vector into the Strides at the specified index. @@ -895,7 +949,9 @@ struct NDArrayStrides( self._buf.store[width=width](idx, value) - fn unsafe_load[width: Int = 1](self, idx: Int) -> SIMD[Self._type, width]: + fn unsafe_load[ + width: Int = 1 + ](self, idx: Int) -> SIMD[Self.element_type, width]: """ Unsafely load a SIMD vector from the Strides at the specified index. @@ -912,7 +968,7 @@ struct NDArrayStrides( fn unsafe_store[ width: Int = 1 - ](self, idx: Int, value: SIMD[Self._type, width]): + ](self, idx: Int, value: SIMD[Self.element_type, width]): """ Unsafely store a SIMD vector into the Strides at the specified index. diff --git a/numojo/core/utility.mojo b/numojo/core/utility.mojo index 4407a748..89585823 100644 --- a/numojo/core/utility.mojo +++ b/numojo/core/utility.mojo @@ -155,8 +155,8 @@ fn _transfer_offset(offset: Int, strides: NDArrayStrides) raises -> Int: The offset of the array of a flipped memory layout. """ - var remainder = offset - var indices = Item(ndim=len(strides), initialized=False) + var remainder: Int = offset + var indices: Item = Item(ndim=len(strides)) for i in range(len(strides)): indices[i] = remainder // strides[i] remainder %= strides[i] diff --git a/numojo/prelude.mojo b/numojo/prelude.mojo index 08a98589..69f475fe 100644 --- a/numojo/prelude.mojo +++ b/numojo/prelude.mojo @@ -22,7 +22,7 @@ from numojo.prelude import * import numojo as nm -from numojo.core.item import Item, item +from numojo.core.item import Item from numojo.core.matrix import Matrix from numojo.core.ndarray import NDArray from numojo.core.ndshape import Shape, NDArrayShape diff --git a/numojo/routines/creation.mojo b/numojo/routines/creation.mojo index 1c5eaf90..58d45258 100644 --- a/numojo/routines/creation.mojo +++ b/numojo/routines/creation.mojo @@ -48,6 +48,7 @@ from numojo.core.data_container import DataContainer # ===------------------------------------------------------------------------===# # Numerical ranges # ===------------------------------------------------------------------------===# +# FIXME: a lot of the creation routines uses ._buf.ptr directly. This should be changed to ._buf[idx] once we use the new DataContainer with correct origins. fn arange[ dtype: DType = DType.float64 ]( @@ -252,10 +253,13 @@ fn linspace[ print(arr2) # [0.0, 2.0, 4.0, 6.0, 8.0] # Parallel computation for large arrays - var large = nm.linspace[nm.f64](0.0, 1000.0, 10000, parallel=True) + var large = nm.linspace[nm.f64, parallel=True](0.0, 1000.0, 10000) ``` """ - constrained[not dtype.is_integral()]() + constrained[ + not dtype.is_integral(), + "numojo.linspace requires floating-point dtype.", + ]() @parameter if parallel: @@ -1236,7 +1240,7 @@ fn ones[ ```mojo import numojo as nm - var arr = nm.ones[nm.f64](Shape(2, 3)) + var arr = nm.ones[nm.f64](nm.Shape(2, 3)) # [[1, 1, 1], # [1, 1, 1]] ``` @@ -1389,7 +1393,7 @@ fn zeros[ ```mojo import numojo as nm - var arr = nm.zeros[nm.f64](Shape(2, 3)) + var arr = nm.zeros[nm.f64](nm.Shape(2, 3)) # [[0, 0, 0], # [0, 0, 0]] ``` @@ -1546,7 +1550,7 @@ fn full[ ```mojo import numojo as nm - var arr = nm.full[nm.f64](Shape(2, 3), fill_value=7.0) + var arr = nm.full[nm.f64](nm.Shape(2, 3), fill_value=7.0) # [[7, 7, 7], # [7, 7, 7]] ``` @@ -1654,7 +1658,7 @@ fn full[ import numojo as nm var val = nm.CScalar[nm.cf64](3.0, 4.0) - var arr = nm.full[nm.cf64](Shape(2, 2), fill_value=val) + var arr = nm.full[nm.cf64](nm.Shape(2, 2), fill_value=val) ``` """ var A = ComplexNDArray[cdtype](shape=shape, order=order) @@ -1772,7 +1776,7 @@ fn diag[ # [0, 0, 2]] # Extract diagonal from 2-D array - var mat = nm.ones[nm.f64](Shape(3, 3)) + var mat = nm.ones[nm.f64](nm.Shape(3, 3)) var d = nm.diag[nm.f64](mat) # [1, 1, 1] ``` @@ -1852,7 +1856,7 @@ fn diagflat[ ```mojo import numojo as nm - var v = nm.arange[nm.f64](4.0).reshape(Shape(2, 2)) # 2x2 array + var v = nm.arange[nm.f64](4.0).reshape(nm.Shape(2, 2)) # 2x2 array var d = nm.diagflat[nm.f64](v) # Flattens to [0,1,2,3] then creates diagonal # [[0, 0, 0, 0], # [0, 1, 0, 0], @@ -2592,9 +2596,10 @@ fn array[ import numojo as nm from numojo.prelude import * from python import Python + var np = Python.import_module("numpy") var np_arr = np.array([1, 2, 3, 4]) - A = nm.array[f32](real=np_arr, imag=np_arr, order="C") + A = nm.array[cf32](real=np_arr, imag=np_arr, order="C") ``` Parameters: @@ -2664,72 +2669,90 @@ fn array[ return A^ -# fn array[ -# dtype: DType = DType.float64 -# ](data: Tensor[dtype]) raises -> NDArray[dtype]: -# """ -# Create array from tensor. - -# Example: -# ```mojo -# import numojo as nm -# from tensor import Tensor, TensorShape -# from numojo.prelude import * - -# fn main() raises: -# height = 256 -# width = 256 -# channels = 3 -# image = Tensor[DType.float32].rand(TensorShape(height, width, channels)) -# print(image) -# print(nm.array(image)) -# ``` - -# Parameters: -# dtype: Datatype of the NDArray elements. - -# Args: -# data: Tensor. - -# Returns: -# NDArray. -# """ +fn meshgrid[ + dtype: DType = DType.float64, indexing: String = "xy" +](*arrays: NDArray[dtype]) raises -> List[NDArray[dtype]]: + """ + Generate coordinate matrices from coordinate vectors. -# return from_tensor(data) + Parameters: + dtype: Datatype of the NDArray elements. + indexing: Cartesian ('xy', default) or matrix ('ij') indexing of output. + Args: + arrays: 1-D input arrays representing the coordinates of a grid. -# fn array[ -# dtype: DType = DType.float64 -# ](real: Tensor[dtype], imag: Tensor[dtype]) raises -> ComplexNDArray[cdtype]: -# """ -# Create array from tensor. + Returns: + A list of N-D coordinate arrays for evaluating expressions on an N-D grid. -# Example: -# ```mojo -# import numojo as nm -# from tensor import Tensor, TensorShape -# from numojo.prelude import * + Examples: + ```mojo + from numojo.routines.creation import meshgrid, arange + from numojo.prelude import * -# fn main() raises: -# height = 256 -# width = 256 -# channels = 3 -# image = Tensor[DType.float32].rand(TensorShape(height, width, channels)) -# print(nm.array(real=image, imag=image)) -# ``` + var x = arange[f64](3.0) # [0, 1, 2] + var y = arange[f64](2.0) # [0, 1] + var grids = meshgrid[f64, indexing="xy"](x, y) + # grids[0]: [[0, 1, 2], + # [0, 1, 2]] + # grids[1]: [[0, 0, 0], + # [1, 1, 1]] + ``` + """ + var n: Int = len(arrays) + if n < 2: + raise Error("meshgrid requires at least two input arrays.") + for i in range(len(arrays)): + if arrays[i].ndim != 1: + raise Error("meshgrid only supports 1-D input arrays.") -# Parameters: -# dtype: Datatype of the NDArray elements. + var grids: List[NDArray[dtype]] = List[NDArray[dtype]](capacity=n) + var final_shape: List[Int] = List[Int](capacity=n) -# Args: -# real: Tensor. -# imag: Tensor. + @parameter + if indexing == "xy": + final_shape.append(arrays[1].size) + final_shape.append(arrays[0].size) + for i in range(2, n): + final_shape.append(arrays[i].size) + else: + for i in range(n): + final_shape.append(arrays[i].size) -# Returns: -# ComplexNDArray. -# """ + for i in range(n): + var grid: NDArray[dtype] = NDArray[dtype](Shape(final_shape)) + var broadcast_dim: Int = i -# return from_tensorC(real, imag) + @parameter + if indexing == "xy": + if i == 0: + broadcast_dim = 1 + elif i == 1: + broadcast_dim = 0 + + var dim_size: Int = arrays[i].size + var outer_size: Int = 1 + var inner_size: Int = 1 + + for j in range(broadcast_dim): + outer_size *= final_shape[j] + for j in range(broadcast_dim + 1, len(final_shape)): + inner_size *= final_shape[j] + + # for outer in range(outer_size): + @parameter + fn closure(outer: Int) -> None: + for k in range(dim_size): + for inner in range(inner_size): + var idx = ( + outer * dim_size * inner_size + k * inner_size + inner + ) + grid._buf.ptr[idx] = arrays[i]._buf.ptr[k] + + parallelize[closure](outer_size, outer_size) + grids.append(grid^) + + return grids^ # ===----------------------------------------------------------------------=== # diff --git a/numojo/routines/indexing.mojo b/numojo/routines/indexing.mojo index aefdd23b..b67a71b5 100644 --- a/numojo/routines/indexing.mojo +++ b/numojo/routines/indexing.mojo @@ -166,7 +166,7 @@ fn compress[ for offset in range(current_slice.size): var remainder: Int = count - var item: Item = Item(ndim=result.ndim, initialized=False) + var item: Item = Item(ndim=result.ndim) # First along the axis var j = normalized_axis diff --git a/numojo/routines/logic/comparison.mojo b/numojo/routines/logic/comparison.mojo index b3493fe3..4af3702e 100644 --- a/numojo/routines/logic/comparison.mojo +++ b/numojo/routines/logic/comparison.mojo @@ -11,6 +11,7 @@ import math import numojo.routines.math._math_funcs as _mf from numojo.core.ndarray import NDArray +from numojo.core.matrix import Matrix, MatrixBase # ===-------------------------------------a-----------------------------------===# @@ -306,3 +307,309 @@ fn not_equal[ return backend().math_func_compare_array_and_scalar[dtype, SIMD.ne]( array1, scalar ) + + +# TODO: Add backend to these functions. +fn allclose[ + dtype: DType +]( + a: NDArray[dtype], + b: NDArray[dtype], + rtol: Scalar[dtype] = 1e-5, + atol: Scalar[dtype] = 1e-8, + equal_nan: Bool = False, +) raises -> Bool: + """ + Determines whether two NDArrays are element-wise equal within a specified tolerance. + + This function compares each element of `a` and `b` and returns True if all corresponding elements satisfy the condition: + abs(a_i - b_i) <= atol + rtol * abs(b_i) + Optionally, if `equal_nan` is True, NaN values at the same positions are considered equal. + + Parameters: + dtype: The data type of the input NDArray. + + Args: + a: First NDArray to compare. + b: Second NDArray to compare. + rtol: Relative tolerance (default: 1e-5). The maximum allowed difference, relative to the magnitude of `b`. + atol: Absolute tolerance (default: 1e-8). The minimum absolute difference allowed. + equal_nan: If True, NaNs in the same position are considered equal (default: False). + + Returns: + True if all elements of `a` and `b` are equal within the specified tolerances, otherwise False. + + Example: + ```mojo + import numojo as nm + from numojo.routines.logic.comparison import allclose + var arr1 = nm.array[nm.f32]([1.0, 2.0, 3.0]) + var arr2 = nm.array[nm.f32]([1.0, 2.00001, 2.99999]) + print(allclose[nm.f32](arr1, arr2)) # Output: True + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Shape Mismatch error shapes must match for this function" + ), + location=( + "numojo.routines.logic.comparision.allclose(a: NDArray, b:" + " NDArray)" + ), + ) + ) + + for i in range(a.size): + val_a: Scalar[dtype] = a.load(i) + val_b: Scalar[dtype] = b.load(i) + if equal_nan and (math.isnan(val_a) and math.isnan(val_b)): + continue + if abs(val_a - val_b) <= atol + rtol * abs(val_b): + continue + else: + return False + + return True + + +fn isclose[ + dtype: DType +]( + a: NDArray[dtype], + b: NDArray[dtype], + rtol: Scalar[dtype] = 1e-5, + atol: Scalar[dtype] = 1e-8, + equal_nan: Bool = False, +) raises -> NDArray[DType.bool]: + """ + Performs element-wise comparison of two NDArrays to determine if their values are equal within a specified tolerance. + + For each element pair (a_i, b_i), the result is True if: + abs(a_i - b_i) <= atol + rtol * abs(b_i) + Optionally, if `equal_nan` is True, NaN values at the same positions are considered equal. + + Parameters: + dtype: The data type of the input NDArray. + + Args: + a: First NDArray to compare. + b: Second NDArray to compare. + rtol: Relative tolerance (default: 1e-5). The maximum allowed difference, relative to the magnitude of `b`. + atol: Absolute tolerance (default: 1e-8). The minimum absolute difference allowed. + equal_nan: If True, NaNs in the same position are considered equal (default: False). + + Returns: + An NDArray of bools, where each element is True if the corresponding elements of `a` and `b` are equal within the specified tolerances, otherwise False. + + Example: + ```mojo + import numojo as nm + from numojo.routines.logic.comparison import isclose + var arr1 = nm.array[nm.f32]([1.0, 2.0, 3.0]) + var arr2 = nm.array[nm.f32]([1.0, 2.00001, 2.99999]) + print(isclose[nm.f32](arr1, arr2)) # Output: [True, True, True] + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Shape Mismatch error shapes must match for this function" + ), + location=( + "numojo.routines.logic.comparision.isclose(a: Scalar, b:" + " Scalar)" + ), + ) + ) + + var res: NDArray[DType.bool] = NDArray[DType.bool](a.shape) + for i in range(a.size): + val_a: Scalar[dtype] = a.load(i) + val_b: Scalar[dtype] = b.load(i) + if equal_nan and (math.isnan(val_a) and math.isnan(val_b)): + res.store(i, True) + continue + if abs(val_a - val_b) <= atol + rtol * abs(val_b): + res.store(i, True) + continue + else: + res.store(i, False) + + return res^ + + +fn allclose[ + dtype: DType +]( + a: Matrix[dtype], + b: Matrix[dtype], + rtol: Scalar[dtype] = 1e-5, + atol: Scalar[dtype] = 1e-8, + equal_nan: Bool = False, +) raises -> Bool: + """ + Determines whether two Matrix are element-wise equal within a specified tolerance. + + This function compares each element of `a` and `b` and returns True if all corresponding elements satisfy the condition: + abs(a_i - b_i) <= atol + rtol * abs(b_i) + Optionally, if `equal_nan` is True, NaN values at the same positions are considered equal. + + Parameters: + dtype: The data type of the input Matrix. + + Args: + a: First Matrix to compare. + b: Second Matrix to compare. + rtol: Relative tolerance (default: 1e-5). The maximum allowed difference, relative to the magnitude of `b`. + atol: Absolute tolerance (default: 1e-8). The minimum absolute difference allowed. + equal_nan: If True, NaNs in the same position are considered equal (default: False). + + Returns: + True if all elements of `a` and `b` are equal within the specified tolerances, otherwise False. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.comparison import allclose + var mat1 = Matrix.rand[f32]((2, 2)) + var mat2 = Matrix.rand[f32]((2, 2)) + print(allclose[f32](mat1, mat2)) # Output: True + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Shape Mismatch error shapes must match for this function" + ), + location=( + "numojo.routines.logic.comparision.allclose(a: NDArray, b:" + " NDArray)" + ), + ) + ) + + for i in range(a.size): + val_a: Scalar[dtype] = a.load(i) + val_b: Scalar[dtype] = b.load(i) + if equal_nan and (math.isnan(val_a) and math.isnan(val_b)): + continue + if abs(val_a - val_b) <= atol + rtol * abs(val_b): + continue + else: + return False + + return True + + +fn isclose[ + dtype: DType +]( + a: MatrixBase[dtype, **_], + b: MatrixBase[dtype, **_], + rtol: Scalar[dtype] = 1e-5, + atol: Scalar[dtype] = 1e-8, + equal_nan: Bool = False, +) raises -> Matrix[DType.bool]: + """ + Performs element-wise comparison of two Matrix to determine if their values are equal within a specified tolerance. + + For each element pair (a_i, b_i), the result is True if: + abs(a_i - b_i) <= atol + rtol * abs(b_i) + Optionally, if `equal_nan` is True, NaN values at the same positions are considered equal. + + Parameters: + dtype: The data type of the input Matrix. + + Args: + a: First Matrix to compare. + b: Second Matrix to compare. + rtol: Relative tolerance (default: 1e-5). The maximum allowed difference, relative to the magnitude of `b`. + atol: Absolute tolerance (default: 1e-8). The minimum absolute difference allowed. + equal_nan: If True, NaNs in the same position are considered equal (default: False). + + Returns: + An NDArray of bools, where each element is True if the corresponding elements of `a` and `b` are equal within the specified tolerances, otherwise False. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.comparison import isclose + var mat1 = Matrix.rand[f32]((2, 2)) + var mat2 = Matrix.rand[f32]((2, 2)) + print(isclose[f32](mat1, mat2)) + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Shape Mismatch error shapes must match for this function" + ), + location=( + "numojo.routines.logic.comparision.isclose(a: Scalar, b:" + " Scalar)" + ), + ) + ) + + var res: Matrix[DType.bool] = Matrix[DType.bool](a.shape) + for i in range(a.size): + val_a: Scalar[dtype] = a.load(i) + val_b: Scalar[dtype] = b.load(i) + if equal_nan and (math.isnan(val_a) and math.isnan(val_b)): + res._store_idx(i, val=True) + continue + if abs(val_a - val_b) <= atol + rtol * abs(val_b): + res._store_idx(i, val=True) + continue + else: + res._store_idx(i, val=False) + + return res^ + + +# TODO: define the allclose, isclose with correct behaviour for ComplexNDArray. + + +fn array_equal[ + dtype: DType +](array1: NDArray[dtype], array2: NDArray[dtype]) raises -> Bool: + """ + Checks if two NDArrays are equal element-wise and shape-wise. + + Parameters: + dtype: The dtype of the input NDArray. + + Args: + array1: First NDArray to compare. + array2: Second NDArray to compare. + + Returns: + True if the two NDArrays are equal element-wise and shape-wise, False otherwise. + + Examples: + ```mojo + from numojo.prelude import * + import numojo as nm + from numojo.routines.logic.comparison import array_equal + + var arr = nm.arange[i32](0, 10) + var arr2 = nm.arange[i32](0, 10) + print(array_equal[i32](arr, arr2)) # Output: True + ``` + """ + if array1.shape != array2.shape: + return False + + for i in range(array1.size): + if array1.load(i) != array2.load(i): + return False + + return True + + +# TODO: define array_equiv with correct broadcast semantics. diff --git a/numojo/routines/logic/contents.mojo b/numojo/routines/logic/contents.mojo index 5ffb971c..0cdaec9d 100644 --- a/numojo/routines/logic/contents.mojo +++ b/numojo/routines/logic/contents.mojo @@ -7,6 +7,7 @@ Implements Checking routines: currently not SIMD due to bool bit packing issue import math +from utils.numerics import neg_inf, inf import numojo.routines.math._math_funcs as _mf from numojo.core.ndarray import NDArray @@ -29,6 +30,7 @@ from numojo.core.ndarray import NDArray # return backend().math_func_is[dtype, math.is_odd](array) +# FIXME: Make all SIMD vectorized operations once bool bit-packing issue is resolved. fn isinf[ dtype: DType, backend: _mf.Backend = _mf.Vectorized ](array: NDArray[dtype]) raises -> NDArray[DType.bool]: @@ -97,3 +99,91 @@ fn isnan[ for i in range(result_array.size): result_array.store(i, math.isnan(array.load(i))) return result_array^ + + +# TODO: Optimize the following functions by implementing a generic backend or just using vectorized operations. +# TODO: Implement the same for complex ndarray. +fn isneginf[ + dtype: DType, backend: _mf.Backend = _mf.Vectorized +](array: NDArray[dtype]) raises -> NDArray[DType.bool]: + """ + Checks if each element of the input array is negative infinity. + + Parameters: + dtype: DType - Data type of the input array. + backend: _mf.Backend - Backend to use for the operation. Defaults to _mf.Vectorized. + + Args: + array: NDArray[dtype] - Input array to check. + + Returns: + NDArray[DType.bool] - A array of the same shape as `array` with True for negative infinite elements and False for others. + """ + var result_array: NDArray[DType.bool] = NDArray[DType.bool](array.shape) + for i in range(result_array.size): + result_array.store(i, neg_inf[dtype]() == array.load(i)) + return result_array^ + + +fn isposinf[ + dtype: DType, backend: _mf.Backend = _mf.Vectorized +](array: NDArray[dtype]) raises -> NDArray[DType.bool]: + """ + Checks if each element of the input array is positive infinity. + Parameters: + dtype: DType - Data type of the input array. + backend: _mf.Backend - Backend to use for the operation. Defaults to _mf.Vectorized. + + Args: + array: NDArray[dtype] - Input array to check. + + Returns: + NDArray[DType.bool] - A array of the same shape as `array` with True for positive infinite elements and False for others. + """ + var result_array: NDArray[DType.bool] = NDArray[DType.bool](array.shape) + for i in range(result_array.size): + result_array.store(i, inf[dtype]() == array.load(i)) + return result_array^ + + +fn isneginf[ + dtype: DType, backend: _mf.Backend = _mf.Vectorized +](matrix: Matrix[dtype]) raises -> Matrix[DType.bool]: + """ + Checks if each element of the input Matrix is negative infinity. + + Parameters: + dtype: DType - Data type of the input Matrix. + backend: _mf.Backend - Backend to use for the operation. Defaults to _mf.Vectorized. + + Args: + matrix: Matrix[dtype] - Input Matrix to check. + + Returns: + Matrix[DType.bool] - A array of the same shape as `array` with True for negative infinite elements and False for others. + """ + var result_array: Matrix[DType.bool] = Matrix[DType.bool](matrix.shape) + for i in range(result_array.size): + result_array.store(i, neg_inf[dtype]() == matrix.load(i)) + return result_array^ + + +fn isposinf[ + dtype: DType, backend: _mf.Backend = _mf.Vectorized +](matrix: Matrix[dtype]) raises -> Matrix[DType.bool]: + """ + Checks if each element of the input Matrix is positive infinity. + Parameters: + dtype: DType - Data type of the input Matrix. + backend: _mf.Backend - Backend to use for the operation. Defaults to _mf.Vectorized. + + Args: + matrix: Matrix[dtype] - Input Matrix to check. + + Returns: + Matrix[DType.bool] - A array of the same shape as `array` with True for positive infinite elements and False for others. + """ + var result_array: Matrix[DType.bool] = Matrix[DType.bool](matrix.shape) + for i in range(result_array.size): + result_array.store(i, inf[dtype]() == matrix.load(i)) + return result_array^ diff --git a/numojo/routines/logic/logical_ops.mojo b/numojo/routines/logic/logical_ops.mojo new file mode 100644 index 00000000..c4b13b3b --- /dev/null +++ b/numojo/routines/logic/logical_ops.mojo @@ -0,0 +1,548 @@ +# ===----------------------------------------------------------------------=== # +# Logical Operations Module +# ===----------------------------------------------------------------------=== # +from numojo.core.error import ShapeError + + +# TODO: add `where` argument support to logical operations +# FIXME: Make all SIMD vectorized operations once bool bit-packing issue is resolved. +# ===----------------------------------------------------------------------=== # +# NDArray operations +# ===----------------------------------------------------------------------=== # +fn logical_and[ + dtype: DType +](a: NDArray[dtype], b: NDArray[dtype]) raises -> NDArray[DType.bool] where ( + dtype == DType.bool or dtype.is_integral() +): + """ + Element-wise logical AND operation between two arrays. + + Args: + a: First input array. + b: Second input array. + + Returns: + An array containing the result of the logical AND operation. + + Raises: + - ShapeError: If the input arrays do not have the same shape. + + Notes: + - Supports only boolean and integral data types. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.logical_ops import logical_and + + var a = nm.arange(0, 10) + var b = nm.arange(5, 15) + var result = logical_and(a > 3, b < 10) + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Input arrays must have the same shape for logical AND" + " operation." + ), + location="numojo.routines.logic.logical_and", + ) + ) + var res: NDArray[DType.bool] = NDArray[DType.bool](a.shape) + for i in range(res.size): + res.store(i, Scalar[DType.bool](a.load(i) & b.load(i))) + return res^ + + +fn logical_or[ + dtype: DType +](a: NDArray[dtype], b: NDArray[dtype]) raises -> NDArray[DType.bool] where ( + dtype == DType.bool or dtype.is_integral() +): + """ + Element-wise logical OR operation between two arrays. + + Args: + a: First input array. + b: Second input array. + + Returns: + An array containing the result of the logical OR operation. + + Raises: + - ShapeError: If the input arrays do not have the same shape. + + Notes: + - Supports only boolean and integral data types. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.logical_ops import logical_or + + var a = nm.arange(0, 10) + var b = nm.arange(5, 15) + var result = logical_or(a < 3, b > 10) + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Input arrays must have the same shape for logical OR" + " operation." + ), + location="numojo.routines.logic.logical_or", + ) + ) + var res: NDArray[DType.bool] = NDArray[DType.bool](a.shape) + for i in range(res.size): + res.store(i, Scalar[DType.bool](a.load(i) | b.load(i))) + return res^ + + +fn logical_not[ + dtype: DType +](a: NDArray[dtype]) raises -> NDArray[DType.bool] where ( + dtype == DType.bool or dtype.is_integral() +): + """ + Element-wise logical NOT operation on an array. + + Args: + a: Input array. + + Returns: + An array containing the result of the logical NOT operation. + + Raises: + - ShapeError: If the input array is not of a supported data type. + + Notes: + - Supports only boolean and integral data types. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.logical_ops import logical_not + + var a = nm.arange(0, 10) + var result = logical_not(a < 5) + ``` + """ + var res: NDArray[DType.bool] = NDArray[DType.bool](a.shape) + for i in range(res.size): + res.store(i, Scalar[DType.bool](~a.load(i))) + return res^ + + +fn logical_xor[ + dtype: DType +](a: NDArray[dtype], b: NDArray[dtype]) raises -> NDArray[DType.bool] where ( + dtype == DType.bool or dtype.is_integral() +): + """ + Element-wise logical XOR operation between two arrays. + + Args: + a: First input array. + b: Second input array. + + Returns: + An array containing the result of the logical XOR operation. + + Raises: + - ShapeError: If the input arrays do not have the same shape. + + Notes: + - Supports only boolean and integral data types. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.logical_ops import logical_xor + + var a = nm.arange(0, 10) + var b = nm.arange(5, 15) + var result = logical_xor(a > 3, b < 10) + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Input arrays must have the same shape for logical XOR" + " operation." + ), + location="numojo.routines.logic.logical_xor", + ) + ) + var res: NDArray[DType.bool] = NDArray[DType.bool](a.shape) + for i in range(res.size): + res.store(i, Scalar[DType.bool](a.load(i) ^ b.load(i))) + return res^ + + +# ===----------------------------------------------------------------------=== # +# ComplexNDArray operations +# ===----------------------------------------------------------------------=== # +fn logical_and[ + cdtype: ComplexDType +]( + a: ComplexNDArray[cdtype], b: ComplexNDArray[cdtype] +) raises -> ComplexNDArray[cdtype] where ( + cdtype == ComplexDType.bool or cdtype.is_integral() +): + """ + Element-wise logical AND operation between two arrays. + + Args: + a: First input array. + b: Second input array. + + Returns: + An array containing the result of the logical AND operation. + + Raises: + - ShapeError: If the input arrays do not have the same shape. + + Notes: + - Supports only boolean and integral data types. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.logical_ops import logical_and + + var a = nm.arange(0, 10) + var b = nm.arange(5, 15) + var result = logical_and(a > 3, b < 10) + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Input arrays must have the same shape for logical AND" + " operation." + ), + location="numojo.routines.logic.logical_and", + ) + ) + var res: ComplexNDArray[cdtype] = ComplexNDArray[cdtype](a.shape) + for i in range(res.size): + res.store(i, a.load(i) & b.load(i)) + return res^ + + +fn logical_or[ + cdtype: ComplexDType +]( + a: ComplexNDArray[cdtype], b: ComplexNDArray[cdtype] +) raises -> ComplexNDArray[cdtype] where ( + cdtype == ComplexDType.bool or cdtype.is_integral() +): + """ + Element-wise logical OR operation between two arrays. + + Args: + a: First input array. + b: Second input array. + + Returns: + An array containing the result of the logical OR operation. + + Raises: + - ShapeError: If the input arrays do not have the same shape. + + Notes: + - Supports only boolean and integral data types. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.logical_ops import logical_or + + var a = nm.arange(0, 10) + var b = nm.arange(5, 15) + var result = logical_or(a < 3, b > 10) + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Input arrays must have the same shape for logical OR" + " operation." + ), + location="numojo.routines.logic.logical_or", + ) + ) + var res: ComplexNDArray[cdtype] = ComplexNDArray[cdtype](a.shape) + for i in range(res.size): + res.store(i, a.load(i) | b.load(i)) + return res^ + + +fn logical_not[ + cdtype: ComplexDType +](a: ComplexNDArray[cdtype]) raises -> ComplexNDArray[cdtype] where ( + cdtype == ComplexDType.bool or cdtype.is_integral() +): + """ + Element-wise logical NOT operation on an array. + + Args: + a: Input array. + + Returns: + An array containing the result of the logical NOT operation. + + Raises: + - ShapeError: If the input array is not of a supported data type. + + Notes: + - Supports only boolean and integral data types. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.logical_ops import logical_not + + var a = nm.arange(0, 10) + var result = logical_not(a < 5) + ``` + """ + var res: ComplexNDArray[cdtype] = ComplexNDArray[cdtype](a.shape) + for i in range(res.size): + res.store(i, ~a.load(i)) + return res^ + + +fn logical_xor[ + cdtype: ComplexDType +]( + a: ComplexNDArray[cdtype], b: ComplexNDArray[cdtype] +) raises -> ComplexNDArray[cdtype] where ( + cdtype == ComplexDType.bool or cdtype.is_integral() +): + """ + Element-wise logical XOR operation between two arrays. + + Args: + a: First input array. + b: Second input array. + + Returns: + An array containing the result of the logical XOR operation. + + Raises: + - ShapeError: If the input arrays do not have the same shape. + + Notes: + - Supports only boolean and integral data types. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.logical_ops import logical_xor + + var a = nm.arange(0, 10) + var b = nm.arange(5, 15) + var result = logical_xor(a > 3, b < 10) + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Input arrays must have the same shape for logical XOR" + " operation." + ), + location="numojo.routines.logic.logical_xor", + ) + ) + var res: ComplexNDArray[cdtype] = ComplexNDArray[cdtype](a.shape) + for i in range(res.size): + res.store(i, a.load(i) ^ b.load(i)) + return res^ + + +# ===----------------------------------------------------------------------=== # +# Matrix operations +# ===----------------------------------------------------------------------=== # +fn logical_and[ + dtype: DType +](a: Matrix[dtype], b: Matrix[dtype]) raises -> Matrix[DType.bool] where ( + dtype == DType.bool or dtype.is_integral() +): + """ + Element-wise logical AND operation between two matrices. + + Args: + a: First input matrix. + b: Second input matrix. + + Returns: + A matrix containing the result of the logical AND operation. + + Raises: + - ShapeError: If the input matrices do not have the same shape. + + Notes: + - Supports only boolean and integral data types. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.logical_ops import logical_and + + var a = Matrix.rand[i32]((2, 5)) + var b = Matrix.rand[i32]((2, 5)) + var result = logical_and(a > 3, b < 10) + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Input matrices must have the same shape for logical AND" + " operation." + ), + location="numojo.routines.logic.logical_and", + ) + ) + var res: Matrix[DType.bool] = Matrix[DType.bool](a.shape) + for i in range(res.size): + res._buf.store(i, Scalar[DType.bool](a.load(i) & b.load(i))) + return res^ + + +fn logical_or[ + dtype: DType +](a: Matrix[dtype], b: Matrix[dtype]) raises -> Matrix[DType.bool] where ( + dtype == DType.bool or dtype.is_integral() +): + """ + Element-wise logical OR operation between two matrices. + + Args: + a: First input matrix. + b: Second input matrix. + + Returns: + A matrix containing the result of the logical OR operation. + + Raises: + - ShapeError: If the input matrices do not have the same shape. + + Notes: + - Supports only boolean and integral data types. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.logical_ops import logical_or + var a = Matrix.rand[i32]((2, 5)) + var b = Matrix.rand[i32]((2, 5)) + var result = logical_or(a < 3, b > 10) + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Input matrices must have the same shape for logical OR" + " operation." + ), + location="numojo.routines.logic.logical_or", + ) + ) + var res: Matrix[DType.bool] = Matrix[DType.bool](a.shape) + for i in range(res.size): + res._buf.store(i, Scalar[DType.bool](a.load(i) | b.load(i))) + return res^ + + +fn logical_not[ + dtype: DType +]( + a: Matrix[dtype], +) raises -> Matrix[DType.bool] where ( + dtype == DType.bool or dtype.is_integral() +): + """ + Element-wise logical NOT operation on a matrix. + + Args: + a: Input matrix. + + Returns: + A matrix containing the result of the logical NOT operation. + + Raises: + - ShapeError: If the input matrix is not of a supported data type. + + Notes: + - Supports only boolean and integral data types. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.logical_ops import logical_not + var a = Matrix.rand[i32]((2, 5)) + var result = logical_not(a < 5) + ``` + """ + var res: Matrix[DType.bool] = Matrix[DType.bool](a.shape) + for i in range(res.size): + res._buf.store(i, Scalar[DType.bool](~a.load(i))) + return res^ + + +fn logical_xor[ + dtype: DType +](a: Matrix[dtype], b: Matrix[dtype]) raises -> Matrix[DType.bool] where ( + dtype == DType.bool or dtype.is_integral() +): + """ + Element-wise logical XOR operation between two matrices. + + Args: + a: First input matrix. + b: Second input matrix. + + Returns: + A matrix containing the result of the logical XOR operation. + + Raises: + - ShapeError: If the input matrices do not have the same shape. + + Notes: + - Supports only boolean and integral data types. + + Example: + ```mojo + from numojo.prelude import * + from numojo.routines.logic.logical_ops import logical_xor + var a = Matrix.rand[i32]((2, 5)) + var b = Matrix.rand[i32]((2, 5)) + var result = logical_xor(a > 3, b < 10) + ``` + """ + if a.shape != b.shape: + raise Error( + ShapeError( + message=( + "Input matrices must have the same shape for logical XOR" + " operation." + ), + location="numojo.routines.logic.logical_xor", + ) + ) + var res: Matrix[DType.bool] = Matrix[DType.bool](a.shape) + for i in range(res.size): + res._buf.store(i, Scalar[DType.bool](a.load(i) ^ b.load(i))) + return res^ diff --git a/numojo/routines/logic/truth.mojo b/numojo/routines/logic/truth.mojo index b2a97d49..4213e35a 100644 --- a/numojo/routines/logic/truth.mojo +++ b/numojo/routines/logic/truth.mojo @@ -76,7 +76,7 @@ fn all[ raise Error(String("The axis can either be 1 or 0!")) -fn allt(array: NDArray[DType.bool]) raises -> Scalar[DType.bool]: +fn all(array: NDArray[DType.bool]) raises -> Scalar[DType.bool]: """ If all True. diff --git a/numojo/routines/manipulation.mojo b/numojo/routines/manipulation.mojo index 0d485f12..48701797 100644 --- a/numojo/routines/manipulation.mojo +++ b/numojo/routines/manipulation.mojo @@ -431,7 +431,7 @@ fn broadcast_to[ # Iterate all items in the new array and fill in correct values. for offset in range(b.size): var remainder = offset - var indices = Item(ndim=b.ndim, initialized=False) + var indices = Item(ndim=b.ndim) for i in range(b.ndim): indices[i] = remainder // b.strides[i] @@ -555,7 +555,7 @@ fn _broadcast_back_to[ # Iterate all items in the new array and fill in correct values. for offset in range(b.size): var remainder = offset - var indices = Item(ndim=b.ndim, initialized=False) + var indices = Item(ndim=b.ndim) for i in range(b.ndim): indices[i] = remainder // b.strides[i] diff --git a/numojo/routines/math/misc.mojo b/numojo/routines/math/misc.mojo index 610b43ce..23c7729e 100644 --- a/numojo/routines/math/misc.mojo +++ b/numojo/routines/math/misc.mojo @@ -20,6 +20,7 @@ import numojo.routines.math._math_funcs as _mf from numojo.core.ndarray import NDArray +# TODO: Implement same routines for Matrix. fn cbrt[ dtype: DType, backend: _mf.Backend = _mf.Vectorized ](array: NDArray[dtype]) raises -> NDArray[dtype]: