diff --git a/numojo/core/ndarray.mojo b/numojo/core/ndarray.mojo index 69ad7e1..4d737f7 100644 --- a/numojo/core/ndarray.mojo +++ b/numojo/core/ndarray.mojo @@ -478,13 +478,12 @@ struct NDArray[dtype: DType = DType.float64]( var size_at_dim: Int = self.ndshape[i] slice_list.append(Slice(0, size_at_dim)) - # self.__setitem__(slice_list=slice_list, val=val) var n_slices: Int = len(slice_list) var ndims: Int = 0 var count: Int = 0 var spec: List[Int] = List[Int]() for i in range(n_slices): - self._adjust_slice_(slice_list[i], self.ndshape[i]) + # self._adjust_slice_(slice_list[i], self.ndshape[i]) if ( slice_list[i].start.value() >= self.ndshape[i] or slice_list[i].end.value() > self.ndshape[i] @@ -611,43 +610,11 @@ struct NDArray[dtype: DType = DType.float64]( Example: `arr[1:3, 2:4]` returns the corresponding sliced array (2 x 2). """ - print("slices: ", slices[0], slices[1], slices[2]) var n_slices: Int = len(slices) var ndims: Int = 0 var count: Int = 0 var spec: List[Int] = List[Int]() - var slice_list: List[Slice] = List[Slice]() - for i in range(n_slices): - var start: Int = 0 - var end: Int = 0 - if slices[i].start is None and slices[i].end is None: - start = 0 - end = self.ndshape[i] - temp = Slice( - start=Optional(start), - end=Optional(end), - step=Optional(slices[i].step), - ) - slice_list.append(temp) - if slices[i].start is None and slices[i].end is not None: - start = 0 - temp = Slice( - start=Optional(start), - end=Optional(slices[i].end.value()), - step=Optional(slices[i].step), - ) - slice_list.append(temp) - if slices[i].start is not None and slices[i].end is None: - end = self.ndshape[i] - temp = Slice( - start=Optional(slices[i].start.value()), - end=Optional(end), - step=Optional(slices[i].step), - ) - slice_list.append(temp) - if slices[i].start is not None and slices[i].end is not None: - slice_list.append(slices[i]) - + var slice_list: List[Slice] = self._adjust_slice_(slices) for i in range(n_slices): if ( slice_list[i].start.value() >= self.ndshape[i] @@ -867,24 +834,50 @@ struct NDArray[dtype: DType = DType.float64]( var idx: Int = _get_index(index, self.coefficient) return self.data.load[width=1](idx) - fn _adjust_slice_(self, inout span: Slice, dim: Int): + fn _adjust_slice_(self, slice_list: List[Slice]) raises -> List[Slice]: """ Adjusts the slice values to lie within 0 and dim. """ - if span.start or span.end: - var start = int(span.start.value()) - var end = int(span.end.value()) - if start < 0: - start = dim + start - if not span.end: - end = dim - elif end < 0: - end = dim + end - if end > dim: - end = dim - if end < start: - start = 0 - end = 0 + var n_slices: Int = slice_list.__len__() + var slices = List[Slice]() + for i in range(n_slices): + if i >= self.ndim: + raise Error("Error: Number of slices exceeds array dimensions") + + var start: Int = 0 + var end: Int = self.ndshape[i] + var step: Int = 1 + if slice_list[i].start is not None: + start = slice_list[i].start.value() + if start < 0: + # start += self.ndshape[i] + raise Error( + "Error: Negative indexing in slices not supported" + " currently" + ) + + if slice_list[i].end is not None: + end = slice_list[i].end.value() + if end < 0: + # end += self.ndshape[i] + 1 + raise Error( + "Error: Negative indexing in slices not supported" + " currently" + ) + + step = slice_list[i].step + if step == 0: + raise Error("Error: Slice step cannot be zero") + + slices.append( + Slice( + start=Optional(start), + end=Optional(end), + step=Optional(step), + ) + ) + + return slices^ fn __getitem__(self, owned *slices: Slice) raises -> Self: """ @@ -908,7 +901,7 @@ struct NDArray[dtype: DType = DType.float64]( var narr: Self = self[slice_list] return narr - fn __getitem__(self, owned slices: List[Slice]) raises -> Self: + fn __getitem__(self, owned slice_list: List[Slice]) raises -> Self: """ Retreive slices of an array from list of slices. @@ -916,15 +909,16 @@ struct NDArray[dtype: DType = DType.float64]( `arr[1:3, 2:4]` returns the corresponding sliced array (2 x 2). """ - var n_slices: Int = slices.__len__() + var n_slices: Int = slice_list.__len__() if n_slices > self.ndim or n_slices < self.ndim: raise Error("Error: No of slices do not match shape") var ndims: Int = 0 var spec: List[Int] = List[Int]() var count: Int = 0 + + var slices: List[Slice] = self._adjust_slice_(slice_list) for i in range(slices.__len__()): - self._adjust_slice_(slices[i], self.ndshape[i]) if ( slices[i].start.value() >= self.ndshape[i] or slices[i].end.value() > self.ndshape[i] diff --git a/test.mojo b/test.mojo index 499cf43..ea50ea8 100644 --- a/test.mojo +++ b/test.mojo @@ -218,15 +218,15 @@ fn test_bool_masks2() raises: print(temp3.ndshape, temp3.stride, temp3.ndshape.ndsize) -fn test_creation_routines() raises: - var x = linspace[numojo.f32](0.0, 60.0, 60) - var y = ones[numojo.f32](3, 2) - var z = logspace[numojo.f32](-3, 0, 60) - var w = arange[f32](0.0, 24.0, step=1) - print(x) - print(y) - print(z) - print(w) +# fn test_creation_routines() raises: +# var x = linspace[numojo.f32](0.0, 60.0, 60) +# var y = ones[numojo.f32](shape(3, 2)) +# var z = logspace[numojo.f32](-3, 0, 60) +# var w = arange[f32](0.0, 24.0, step=1) +# print(x) +# print(y) +# print(z) +# print(w) fn test_slicing() raises: @@ -267,6 +267,33 @@ fn test_slicing() raises: print(slicedy3) # print("Time taken: ", (time.now() - start)/1e9/10) + # var np = Python.import_module("numpy") + # y = nm.arange[nm.f32](0.0, 24.0, step=1) + # y.reshape(2, 3, 4, order="C") + # np_y = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4, order="C") + # print(y) + # print(np_y) + # print() + # # Test slicing + # slicedy = y[:, :, 1:2] + # print("slicedy: ", slicedy) + # np_slicedy = np.take( + # np.take( + # np.take(np_y, np.arange(0, 2), axis=0), np.arange(0, 3), axis=1 + # ), + # np.arange(1, 2), + # axis=2, + # ) + # print("np_slicedy: ", np_slicedy) + # np_slicedy = np.squeeze( + # np_slicedy, axis=2 + # ) # Remove the dimension with size 1 + # var np_arr = slicedy.to_numpy() + # print() + # print(np_arr) + # print(np_slicedy) + # print(np.all(np.equal(np_arr, np_slicedy))) + fn test_rand_funcs[ dtype: DType = DType.float64 @@ -331,18 +358,19 @@ def test_solve(): fn test_setter() raises: print("Testing setter") - # var A = NDArray[i16](2, 3, 2, fill=Scalar[i16](1)) - # var B = NDArray[i16](3, 2, fill=Scalar[i16](2)) - # A[0] = B - # print(A) + var A = nm.full[i16](3, 3, 3, fill_value=1) + var B = nm.full[i16](3, 3, fill_value=2) + A[0] = B + print(A) - var A = ndarray[i16](3, 3, 3, fill=Scalar[i16](1)) - print("1: ", A) - var D = nm.random.rand[i16](3, 3, min=0, max=100) - A[1] = D # sets the elements of A[1:2, :, :] with the array `D` - print("2: ", A) - A[:, 0:1, :] = D # sets the elements of A[:, 0:1, :] with the array `D` - print("3: ", A) + var A1 = nm.full[i16](3, 4, 5, fill_value=1) + print("A1: ", A1) + var D1 = nm.random.rand[i16](3, 5, min=0, max=100) + A1[:, 0:1, :] = D1 # sets the elements of A[:, 0:1, :] with the array `D` + print("A3: ", A1) + var D = nm.random.rand[i16](4, 5, min=0, max=100) + A1[1] = D # sets the elements of A[1:2, :, :] with the array `D` + print("A2: ", A1) fn main() raises: @@ -359,118 +387,3 @@ fn main() raises: # test_solve() # test_linalg() test_setter() - - -# var x = numojo.full[numojo.f32](3, 2, fill_value=16.0) -# var x = numojo.NDArray[numojo.f32](data=List[SIMD[numojo.f32, 1]](1,2,3,4,5,6,7,8,9,10,11,12), shape=List[Int](2,3,2), -# order="F") -# print(x) -# print(x.stride) -# var y = numojo.NDArray[numojo.f32](data=List[SIMD[numojo.f32, 1]](1,2,3,4,5,6,7,8,9,10,11,12), shape=List[Int](2,3,2), -# order="C") -# print(y) -# print(y.stride) -# print() -# var summed = numojo.stats.sum(x,0) -# print(summed) -# print(numojo.stats.mean(x,0)) -# print(numojo.stats.cumprod(x)) - -# var maxval = x.max(axis=0) -# print(maxval) - - -# var array = nj.NDArray[nj.f64](10,10) -# for i in range(array.size()): -# array[i] = i -# # for i in range(10): -# # for j in range(10): -# # print(array[i, j]) -# var res = array.sum(axis=0) -# print(res) - -# var arr2 = numojo.NDArray[numojo.f32](data=List[SIMD[numojo.f32, 1]](1.0, 2.0, 4.0, 7.0, 11.0, 16.0), -# shape=List[Int](6)) -# var np = Python.import_module("numpy") -# var np_arr = numojo.to_numpy(arr2) -# print(np_arr) -# var result = numojo.math.calculus.differentiation.gradient[numojo.f32](arr2, spacing=1.0) -# print(result) -# print(arr1.any()) -# print(arr1.all()) -# print(arr1.argmax()) -# print(arr1.argmin()) -# print(arr1.astype[numojo.i16]()) -# print(arr1.flatten(inplace=True)) -# print(r.ndshape, r.stride, r.ndshape.ndsize) -# var t0 = time.now() -# var res = numojo.math.linalg.matmul_tiled_unrolled_parallelized[numojo.f32](arr, arr1) -# print((time.now()-t0)/1e9) -# var res = numojo.math.linalg.matmul_tiled_unrolled_parallelized[numojo.f32](arr, arr1) -# print(res) -# print(arr) -# print("2x3x1") -# var sliced = arr[:, :, 1:2] -# print(sliced) - -# print("1x3x4") -# var sliced1 = arr[::2, :] -# print(sliced1) - -# print("1x3x1") -# var sliced2 = arr[1:2, :, 2:3] -# print(sliced2) - -# var result = numojo.NDArray(3, 3) -# numojo.math.linalg.dot[t10=3, t11=3, t21=3, dtype=numojo.f32](result, arr, arr1) -# print(result) - - -# fn main() raises: -# var size:VariadicList[Int] = VariadicList[Int](16,128,256,512,1024) -# alias size1: StaticIntTuple[5] = StaticIntTuple[5](16,128,256,512,1024) -# var times:List[Float64] = List[Float64]() -# alias type:DType = DType.float64 -# measure_time[type, size1](size, times) - -# fn measure_time[dtype:DType, size1: StaticIntTuple[5]](size:VariadicList[Int], inout times:List[Float64]) raises: - -# for i in range(size.__len__()): -# var arr1 = numojo.NDArray[dtype](size[i], size[i]) -# var arr2 = numojo.NDArray[dtype](size[i], size[i]) -# var arr_mul = numojo.NDArray[dtype](size[i], size[i]) - -# var t0 = time.now() -# @parameter -# for i in range(50): -# numojo.math.linalg.dot[t10=size1[i], t11=size1[i], t21=size1[i], dtype=dtype](arr_mul, arr1, arr2) -# # var arr_mul = numojo.math.linalg.matmul_parallelized[dtype](arr1, arr2) -# # var arr_mul = numojo.math.linalg.matmul_tiled_unrolled_parallelized[dtype](arr1, arr2) -# keep(arr_mul.unsafe_ptr()) -# times.append(((time.now()-t0)/1e9)/50) - -# for i in range(size.__len__()): -# print(times[i]) - -# fn main() raises: -# alias type:DType = DType.float16 -# measure_time[type]() - -# fn measure_time[dtype:DType]() raises: -# var size:VariadicList[Int] = VariadicList[Int](16,128,256,512,1024) -# alias size1: StaticIntTuple[5] = StaticIntTuple[5](16,128,256,512,1024) - -# var n = 4 -# alias m = 4 -# var arr1 = numojo.NDArray[dtype](size[n], size[n]) -# var arr2 = numojo.NDArray[dtype](size[n], size[n]) -# var arr_mul = numojo.NDArray[dtype](size[n], size[n]) - -# var t0 = time.now() - -# for _ in range(50): -# numojo.math.linalg.dot[t10=size1[m], t11=size1[m], t21=size1[m], dtype=dtype](arr_mul, arr1, arr2) -# # var arr_mul = numojo.math.linalg.matmul_parallelized[dtype](arr1, arr2) -# # var arr_mul = numojo.math.linalg.matmul_tiled_unrolled_parallelized[dtype](arr1, arr2) -# keep(arr_mul.unsafe_ptr()) -# print(((time.now()-t0)/1e9)/50) diff --git a/tests/test_slicing.mojo b/tests/test_slicing.mojo index 2025814..7035a0a 100644 --- a/tests/test_slicing.mojo +++ b/tests/test_slicing.mojo @@ -1,55 +1,123 @@ -# import numojo as nm -# from numojo import * -# from testing.testing import assert_true, assert_almost_equal, assert_equal -# from utils_for_test import check, check_is_close +import numojo as nm +from numojo import * +from testing.testing import assert_true, assert_almost_equal, assert_equal +from utils_for_test import check, check_is_close -# def test_slicing(): +def test_slicing_getter1(): + var np = Python.import_module("numpy") + + # Test C-order array slicing + nm_arr = nm.arange[nm.f32](0.0, 24.0, step=1) + nm_arr.reshape(2, 3, 4, order="C") + np_arr = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) + + # Test case 1: Slicing all dimensions + nm_slice1 = nm_arr[:, :, 1:2] + np_sliced1 = np.take( + np.take( + np.take(np_arr, np.arange(0, 2), axis=0), np.arange(0, 3), axis=1 + ), + np.arange(1, 2), + axis=2, + ) + np_sliced1 = np.squeeze(np_sliced1, axis=2) + check(nm_slice1, np_sliced1, "3D array slicing (C-order) [:, :, 1:2]") + + +def test_slicing_getter2(): + var np = Python.import_module("numpy") + + # Test C-order array slicing + nm_arr = nm.arange[nm.f32](0.0, 24.0, step=1) + nm_arr.reshape(2, 3, 4, order="C") + np_arr = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) + + # Test case 2: Slicing with start and end indices + nm_slice2 = nm_arr[0:1, 1:3, 2:4] + np_sliced2 = np.take( + np.take( + np.take(np_arr, np.arange(0, 1), axis=0), np.arange(1, 3), axis=1 + ), + np.arange(2, 4), + axis=2, + ) + check(nm_slice2, np_sliced2, "3D array slicing (C-order) [0:1, 1:3, 2:4]") + + +def test_slicing_getter3(): + var np = Python.import_module("numpy") + + # Test C-order array slicing + nm_arr = nm.arange[nm.f32](0.0, 24.0, step=1) + nm_arr.reshape(2, 3, 4, order="C") + np_arr = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) + + # Test case 3: Slicing with mixed start, end, and step values + nm_slice3 = nm_arr[1:, 0:2, ::2] + np_sliced3 = np.take( + np.take( + np.take(np_arr, np.arange(1, np_arr.shape[0]), axis=0), + np.arange(0, 2), + axis=1, + ), + np.arange(0, np_arr.shape[2], 2), + axis=2, + ) + check(nm_slice3, np_sliced3, "3D array slicing (C-order) [1:, 0:2, ::2]") + + +def test_slicing_getter4(): + var np = Python.import_module("numpy") + + # Test C-order array slicing + nm_arr = nm.arange[nm.f32](0.0, 24.0, step=1) + nm_arr.reshape(2, 3, 4, order="C") + np_arr = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) + + # Test case 4: Slicing with step + nm_slice4 = nm_arr[::2, ::2, ::2] + np_sliced4 = np.take( + np.take( + np.take(np_arr, np.arange(0, np_arr.shape[0], 2), axis=0), + np.arange(0, np_arr.shape[1], 2), + axis=1, + ), + np.arange(0, np_arr.shape[2], 2), + axis=2, + ) + check(nm_slice4, np_sliced4, "3D array slicing (C-order) [::2, ::2, ::2]") + + +def test_slicing_getter5(): + var np = Python.import_module("numpy") + + # Test C-order array slicing + nm_arr = nm.arange[nm.f32](0.0, 24.0, step=1) + nm_arr.reshape(2, 3, 4, order="C") + np_arr = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) + + # Test case 5: Slicing with combination of integer and slices + nm_slice5 = nm_arr[1:2, :, 1:3] + np_sliced5 = np.take( + np.take(np_arr[1], np.arange(0, np_arr.shape[1]), axis=0), + np.arange(1, 3), + axis=1, + ) + check(nm_slice5, np_sliced5, "3D array slicing (C-order) [1, :, 1:3]") + + +# def test_slicing_setter1(): # var np = Python.import_module("numpy") # # Test C-order array slicing -# w = nm.arange[nm.f32](0.0, 24.0, step=1) -# w.reshape(2, 3, 4, order="C") -# np_w = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) - -# # Test basic indexing -# check_is_close( -# w[0, 0, 0], np_w[0, 0, 0], "3D array indexing (C-order) [0,0,0]" -# ) -# check_is_close( -# w[1, 2, 3], np_w[1, 2, 3], "3D array indexing (C-order) [1,2,3]" -# ) - -# # Test slicing -# slicedw = w[0:1, :, 1:2] -# # var py_list: PythonObject = [0:1, :, 1:2] -# np_slicedw = np_w[PythonObject0:1, :, 1:2] -# check_is_close(slicedw, np_slicedw, "3D array slicing (C-order)") - -# # Test F-order array slicing -# y = nm.arange[nm.f32](0.0, 24.0, step=1) -# y.reshape(2, 3, 4, order="F") -# np_y = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4, order="F") - -# # Test basic indexing -# check_is_close( -# y[0, 0, 0], np_y[0, 0, 0], "3D array indexing (F-order) [0,0,0]" -# ) -# check_is_close( -# y[1, 2, 3], np_y[1, 2, 3], "3D array indexing (F-order) [1,2,3]" -# ) - -# # Test slicing -# slicedy = y[:, :, 1:2] -# np_slicedy = np_y[:, :, 1:2] -# check_is_close(slicedy, np_slicedy, "3D array slicing (F-order)") - -# # Test integer array -# z = nm.arange[nm.i32](0, 24, step=1) -# z.reshape(2, 3, 4, order="C") -# np_z = np.arange(0, 24, dtype=np.int32).reshape(2, 3, 4, order="C") - -# # Test slicing for integer array -# slicedz = z[1:2, 0:2, :] -# np_slicedz = np_z[1:2, 0:2, :] -# check(slicedz, np_slicedz, "3D integer array slicing (C-order)") +# nm_arr = nm.arange[nm.f32](0.0, 24.0, step=1) +# nm_arr.reshape(2, 3, 4, order="C") +# np_arr = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) + +# # Test case 2: Setting a slice with another array +# nm_set_arr = nm.full[nm.f32](2, 2, fill_value=50.0) +# np_set_arr = np.full((1, 2, 2), 50, dtype=np.float32) +# nm_arr[1:2, 1:3, 2:4] = nm_set_arr +# np.put(np_arr, np.ravel_multi_index((np.arange(1, 2), np.arange(1, 3), np.arange(2, 4)), np_arr.shape), np_set_arr.flatten()) +# check(nm_arr, np_arr, "3D array slice setting (C-order) [1:2, 1:3, 2:4] = array")