Разделить массив с помощью маскировки в Python

Question

У меня есть массив NumPy my_array размером 100x20.Я хочу создать функцию, которая получает в качестве входных данных двумерный массив my_arr и индекс x и возвращает два массива: один с размером 1x20 test_arr и один с 99x20 train_arr.Вектор test_arr будет соответствовать строке матрицы my_arr с индексом x, а train_arr будет содержать остальные строки.Я пытался найти решение, используя маскировку:

def split_train_test(my_arr, x):

   a = np.ma.array(my_arr, mask=False)
   a.mask[x, :] = True
   a = np.array(a.compressed())
   return a

Видимо, это не работает так, как я хотел.Как я могу вернуть массив numpy в результате и массивы train и test должным образом?

Answer 1

Вы также можете использовать логический индекс в качестве маски:

def split_train_test(my_arr, x):

    # define mask
    mask=np.zeros(my_arr.shape[0], dtype=bool)
    mask[x] = True # True only at index x, False elsewhere

    return my_arr[mask, :], my_arr[~mask, :]

Пример выполнения:

test_arr, train_arr = split_train_test(np.random.rand(100, 20), x=10)

print(test_arr.shape, train_arr.shape)
((1L, 20L), (99L, 20L))

---

РЕДАКТИРОВАТЬ:

Если кто-то ищет общий случай, когда более чем один элемент должен быть выделен для тестового массива (скажем, 80% -20% разделения), x может также принять массив:

my_arr = np.random.rand(100, 20)
x = np.random.choice(np.arange(my_arr.shape[0]), int(my_arr .shape[0]*0.8), replace=False)

test_arr, train_arr = split_train_test(my_arr, x)
print(test_arr.shape, train_arr.shape)
((80L, 20L), (20L, 20L))

Answer 2

Вы можете использовать простой индекс и numpy.delete для этого:

def split_train_test(my_arr, x):
    return np.delete(my_arr, x, 0), my_arr[x:x+1]

my_arr = np.arange(10).reshape(5,2)

train, test = split_train_test(my_arr, 2)

train
#array([[0, 1],
#       [2, 3],
#       [6, 7],
#       [8, 9]])

test
#array([[4, 5]])