data_pipe.py

import numpy as np
from sklearn.model_selection import train_test_split

config = {
           'data': data,
           'train_test_ratio': 0.2 
          }

def feature_selection(data):
  """
    feature_selection function.
    It takes data array and returns the feature selected.
    
    Arguments:
      data np.array of shape (n_training_example, n_features)
      
    Returns:
      data_reduced np.array of shape (n_training_example, n_reduced_features)
  """
  
  data_reduced = data # In this case we keep all the features. Change this according to your analysis.
  
  assert data_reduced.shape[0] == data.shape[0], "Data leaking!"
  return data_reduced

def import_data(input=config):
  """
    import_data function.
    It makes use of sklearn.model_selection.train_test_split.
    
    Arguments:
      input dict containing the following variables
      data dict of np.arrays
            data.data is the array made of feature vectors rows.
            data.target is the array of target values.
      
      train_test_ratio float
            the ratio between train and test set sizes. 
            default: 0.2
      
    Returns:
      tuple of four np.arrays (X_train, X_test, Y_train, Y_test) of shape
        -  X_train (n_training_examples, n_features)
        -  X_test  (n_test_examples, n_features)
        -  Y_train (n_training_examples, )
        -  Y_test  (n_test_examples, )
  """
  data, train_test_ratio = input['data'], input['train_test_ratio']
  X = feature_selection(data.data)
  Y = data.target
  X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size = train_test_ratio, random_state=42)
  
  assert X_train.shape[1] == X_test.shape[1], "Train and test shapes do not correspond!"
  return X_train, X_test, Y_train, Y_test