Filter Numpy Array with optional argument

Question

I am building a function which should prepare my data depending on the input. The variable x_imp contains indices on which features are important. However sometimes I still need all features so if 'x_imp = None' nothing should happen.

My solution was this (this is not the whole function just the inputs):

def get_train_data(x_cat, x_num,x_imp = None):
        x_cat = x_cat[:,x_imp]
        x_num = x_num[:,x_imp]
    return x_train

But this changes the shape of the data. For example if data.shape = (4, 5) then data[:,None].shape = (4, 1, 5)

How do I avoid this problem?

Answer 1

This happens because slicing by None is an alias for np.newaxis . Is there a reason not to just add an explicit if statement?

def get_train_data(x_cat, x_num,x_imp = None):
    if x_imp is not None:
        x_cat = x_cat[:,x_imp]
        x_num = x_num[:,x_imp]
    return x_train