Q:

Split a numpy array using masking in python

I have a numpy array my_array of size 100x20. I want to create a function that receives as an input a 2d numpy array my_arr and an index x and will return two arrays one with size 1x20 test_arr and one with 99x20 train_arr. The vector test_arr will correspond to the row of the matrix my_arr with the index x and the train_arr will contain the rest rows. I tried to follow a solution using masking:
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

Apparently this is not working as i wanted. How can i return a numpy array as a result and the train and test arrays properly?

A:

You can use simple index and numpy.delete for this:
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]])

