mirror of https://github.com/exaloop/codon.git
111 lines
4.4 KiB
Python
111 lines
4.4 KiB
Python
import numpy as np
|
|
from numpy import *
|
|
|
|
@test
|
|
def test_basic_indexing():
|
|
#Single element indexing
|
|
arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
|
|
assert (arr[[0, 1], [1, 2]] == np.array([2, 6])).all()
|
|
arr[[0, 2], [1, 2]] = [10, 11]
|
|
assert arr[0, 1] == 10 and arr[2, 2] == 11
|
|
assert arr[-1, -1] == 11
|
|
assert (arr[0] == np.array([1, 10, 3])).all()
|
|
|
|
#Slicing and striding
|
|
x = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
|
|
assert (x[1:7:2] == np.array([1, 3, 5])).all()
|
|
assert (x[-2:10] == np.array([8, 9])).all()
|
|
assert (x[-3:3:-1] == np.array([7, 6, 5, 4])).all()
|
|
assert (x[5:] == np.array([5, 6, 7, 8, 9])).all()
|
|
mx = np.array([[[1], [2], [3]], [[4], [5], [6]]])
|
|
assert (mx[1:2] == np.array([[[4], [5], [6]]])).all()
|
|
assert (arr[1:10:5, ::-1] == np.array([[6, 5, 4]])).all()
|
|
selected_rows = arr[1:3]
|
|
selected_columns = arr[:, 1]
|
|
assert (selected_rows == np.array([[4, 5, 6], [7, 8, 11]])).all()
|
|
assert (selected_columns == np.array([10, 5, 8])).all()
|
|
|
|
#Dimensional indexing tools
|
|
assert (mx[..., 0] == np.array([[1, 2, 3], [4, 5, 6]])).all()
|
|
assert (mx[:, np.newaxis, :, :].shape == (2, 1, 3, 1))
|
|
assert (mx[:, None, :, :].shape == (2, 1, 3, 1))
|
|
dx = np.arange(5)
|
|
assert (dx[:, np.newaxis] + dx[np.newaxis, :] == np.array([[0, 1, 2, 3, 4],
|
|
[1, 2, 3, 4, 5],
|
|
[2, 3, 4, 5, 6],
|
|
[3, 4, 5, 6, 7],
|
|
[4, 5, 6, 7,
|
|
8]])).all()
|
|
|
|
arr = np.array(42)
|
|
assert (arr[np.array(True)] == [42]).all()
|
|
assert arr[np.array(False)].size == 0
|
|
arr[np.array(False)] = -1
|
|
assert arr.item() == 42
|
|
arr[np.array(True)] = -1
|
|
assert arr.item() == -1
|
|
|
|
test_basic_indexing()
|
|
|
|
@test
|
|
def test_advanced_indexing():
|
|
#Integer array indexing
|
|
x = np.arange(10, 1, -1)
|
|
assert (x[np.array([3, 3, 1, 8])] == np.array([7, 7, 9, 2])).all()
|
|
assert (x[np.array([3, 3, -3, 8])] == np.array([7, 7, 4, 2])).all()
|
|
ax = np.array([[1, 2], [3, 4], [5, 6]])
|
|
assert (ax[np.array([1, -1])] == np.array([[3, 4], [5, 6]])).all()
|
|
try:
|
|
ax[np.array([3, 4])]
|
|
except IndexError as e:
|
|
assert True
|
|
y = np.arange(35).reshape(5, 7)
|
|
assert (y[np.array([0, 2, 4]),
|
|
np.array([0, 1, 2])] == np.array([0, 15, 30])).all()
|
|
try:
|
|
y[np.array([0, 2, 4]), np.array([0, 1])]
|
|
except IndexError as e:
|
|
assert True
|
|
assert (y[np.array([0, 2, 4]), 1] == np.array([1, 15, 29])).all()
|
|
assert (y[np.array([0, 2, 4])] == np.array([[0, 1, 2, 3, 4, 5, 6],
|
|
[14, 15, 16, 17, 18, 19, 20],
|
|
[28, 29, 30, 31, 32, 33,
|
|
34]])).all()
|
|
arr = np.arange(12).reshape(4, 3)
|
|
rows = np.array([0, 3])
|
|
columns = np.array([0, 2])
|
|
assert (rows[:, np.newaxis] == np.array([[0], [3]])).all()
|
|
assert (arr[rows[:, np.newaxis], columns] == np.array([[0, 2],
|
|
[9, 11]])).all()
|
|
|
|
#Boolean array indexing
|
|
x = np.array([[1., 2.], [np.nan, 3.], [np.nan, np.nan]])
|
|
assert (x[~np.isnan(x)] == np.array([1., 2., 3.])).all()
|
|
x = np.array([1., -1., -2., 3])
|
|
x[x < 0] += 20
|
|
assert (x == np.array([1., 19., 18., 3.])).all()
|
|
arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
|
|
assert (arr[arr > 5] == np.array([6, 7, 8, 9])).all()
|
|
mask = arr > 4
|
|
assert (arr[mask] == np.array([5, 6, 7, 8, 9])).all()
|
|
arr2 = np.arange(12).reshape(4, 3)
|
|
rows = (arr2.sum(-1) % 2) == 0
|
|
assert (rows == np.array([False, True, False, True])).all()
|
|
a = np.arange(30).reshape(2, 3, 5)
|
|
b = np.array([[True, True, False], [False, True, True]])
|
|
assert (a[b] == np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9],
|
|
[20, 21, 22, 23, 24], [25, 26, 27, 28,
|
|
29]])).all()
|
|
|
|
test_advanced_indexing()
|
|
|
|
@test
|
|
def test_combined_indexing():
|
|
x = np.arange(35).reshape(5, 7)
|
|
b = x > 20
|
|
assert (x[b[:, 5], 1:3] == np.array([[22, 23], [29, 30]])).all()
|
|
assert (np.array([42, 32])[()] == np.array([42, 32])).all()
|
|
assert np.array(42)[()] == 42
|
|
|
|
test_combined_indexing()
|