How to remove columns in Transformer function in Pipeline?

Question

I already used a custom transformation function in a scikit-learn pipeline. In this function I only added features to my data frame. It works great.

Below is a working example:

import pandas as pd
import numpy as np
from sklearn.pipeline import Pipeline
from sklearn.datasets import make_blobs

x, y = make_blobs(n_samples=300, n_features=2, centers=1)
x_train = pd.DataFrame(x[:150,:], columns=['x1','x2'])
x_test = pd.DataFrame(x[150:,:], columns=['x1','x2'])

class myTransformation(object) :
  def __init__(self, colname):
    self.colname = colname

  def transform(self, x) :
    dat = x.copy()
    squared = dat.loc[:,self.colname]**2
    squared.name = "%s_sqre"%self.colname
    dat.loc[:,squared.name] = squared
    dat.loc[:, self.colname+'_2'] = dat[self.colname]
    return dat

  def fit(self, dat, y=None) :
    return self

makePipe = Pipeline([('makeTransfo', myTransformation(colname="x2"))])
fittedPipe = makePipe.fit(x_train)
x_1 = fittedPipe.transform(x_train)
x_2 = fittedPipe.transform(x_test)

Now I would like to be able to add the ability to remove the equal columns in the data frames. For now, I have the following function:

def delSameCols(df) :
  cols = []
  for i in range(df.shape[1]) :
    for j in range(i+1, df.shape[1]) :
      if (df.iloc[:,i].dtype!='O') | (df.iloc[:,j].dtype!='O') :
        if np.array_equal(df.iloc[:,i],df.iloc[:,j]) :
          cols.append(df.columns[j])
  cols = list(set(cols))
  print( u'      -%s features removed'%len(cols) )
  return df.drop(cols, axis=1), cols

I have no idea how to deal with this/how to add a new function in the pipeline or directly in the existing function? Does anyone have any idea?

Answer 1

I succeeded in getting a satisfying solution. I posted an entire working script. What do you think about it? Especially the creation of an attribute (self.lstRemCols) not initialized in the init function?

import pandas as pd
import numpy as np
from sklearn.pipeline import Pipeline
from sklearn.datasets import make_blobs

class myTransfo(object) :
  def __init__(self, colname):
    self.colname = colname

  def transform(self, x) :
    dat = x.copy()
    squared = dat.loc[:,self.colname]**2
    squared.name = "%s_sqre"%self.colname
    dat.loc[:,squared.name] = squared
    dat.loc[:, self.colname+'_2'] = dat[self.colname]
    return dat

  def fit(self, dat, y=None) :
    return self

class removeSameCols(object) :
  def __init__(self) :
    pass

  def _delSameCols(self, df) :
    cols = []
    for i in range(df.shape[1]) :
      for j in range(i+1, df.shape[1]) :
        if (df.iloc[:,i].dtype!='O') | (df.iloc[:,j].dtype!='O') :
          if np.array_equal(df.iloc[:,i],df.iloc[:,j]) :
            cols.append(df.columns[j])
    cols = list(set(cols))
    print( u'      - %s features to be removed'%len(cols) )
    return cols

  def transform(self, x) :
    dat = x.copy()
    lstcols = list(set(dat.columns) - set(self.lstRemCols))
    return dat.loc[:, lstcols]

  def fit(self, x, y=None) :
    dat = x.copy()
    self.lstRemCols = self._delSameCols(dat)
    return self


x, y = make_blobs(n_samples=300, n_features=5)
x_train = pd.DataFrame(x[:150,:], columns=['x1','x2','x3','x4','x5'])
x_test = pd.DataFrame(x[150:,:], columns=['x1','x2','x3','x4','x5'])

makePipe2 = Pipeline([('makeCols', myTransfo(colname="x2")),
                      ('remCols', removeSameCols())])
makePipe2.fit(x_train)
x_1 = makePipe2.transform(x_train)
# test if only same columns in x_train are removed.
x_test.x4 = x_test.x5
x_2 = makePipe2.transform(x_test)