How to Assign Labels with Sklearn One Hot Encoder | by Vitalii Dodonov | Towards Data Science

Most machine learning algorithms don’t work with categorical data out of the box. It is a common practice to apply label encoder or one hot encoder to a categorical variable before using it in predictive modelling.
Sklearn library in Python is a go-to for most Machine Learning practitioners. It provides outstanding generalized functionality that allows us to perform complex operations with a single line of code. At the same time, some of its functionality is too general and may be a challenge to work with.
One hot encoding is a process of transforming a categorical variable into N binary columns where N is the number of unique values in the original column. For example, in my recent study about stock price behaviour during COVID-19
towardsdatascience.com
I used one hot encoding to break down industry sector classifications
into its binary column representation like so:
One hot transformation can be accomplished using the default sklearn package:
sklearn.preprocessing.OneHotEncoder# df = some DataFrameencoder = OneHotEncoder()encoder.fit_transform(df)
The output of the code above implementation looks like this:
It is correct, but it doesn’t provide labels, which makes it difficult to know the meaning behind the new columns. Without this knowledge, analyzing the output of predictive models would be impossible. Some guidance on label assignments is provided in the documentation but it is still a challenge to work with. Hence, there is a need for a better solution.
An elegant way to add labels to the One Hot Encoder output while keeping the “one line” implementation is to create a wrapper class that assigns labels during the transform() operation “inside the box”. Here is the code:
With this wrapper, the output loos like this:
I will now break down its implementation and explain how it works in detail.
from sklearn.preprocessing import OneHotEncoder as SklearnOneHotEncoderclass OneHotEncoder(SklearnOneHotEncoder):    def __init__(self, **kwargs):        super(OneHotEncoder, self).__init__(**kwargs)        self.fit_flag = False
This class inherits from the sklearn.preprocessing.OneHotEncoder, which means all the original functionality is preserved. I also add self.fit_flag property to keep track of whether the encoder has been fit or not.
def fit(self, X, **kwargs):    out = super().fit(X)    self.fit_flag = True    return out
The fit() method is essentially the original fit() method. I use super().fit(X) to access the original method from sklearn.preprocessing.OneHotEncoder and then update self.fit_flag to True before returning the output.
def transform(self, X, **kwargs):    sparse_matrix = super(OneHotEncoder, self).transform(X)    new_columns = self.get_new_columns(X=X)    d_out = pd.DataFrame(sparse_matrix.toarray(), columns=new_columns, index=X.index)    return d_outdef get_new_columns(self, X):    new_columns = []    for i, column in enumerate(X.columns):        j = 0        while j < len(self.categories_[i]):            new_columns.append(f'{column}_<{self.categories_[i][j]}>')            j += 1    return new_columns
Transform method is more involved because this is where I implement label assignments. The sparse_matrix variable is created first, storing the output from the original transform() method. I then call get_new_column() method which I wrote to access OneHotEncoder’s internal properties to retrieve class names.
The get_new_columns() method is essentially a nested iterator with two levels. At the first level, I iterate over columns in the original DataFrame. At the second level, I iterate over values in self.categories_[i] where self.categories_ belongs to the original OneHotEncoder and “i” is the ordinal index of columns in the table being transformed. At each step, I append column name to a list which will be returned once the iterations are complete. Eventually, the length of the new_columns variable is equivalent to the horizontal dimension of the sparse_matrix once it is converted to an array.
In the final step, I convert the sparse_matrix to an array using built-in functionality, then create a DataFrame with the column names returned by the get_new_columns() method.
def fit_transform(self, X, **kwargs):    self.fit(X)    return self.transform(X)
The last method consists of sequential execution of my modified fit() and transform() methods, returning a One Hot DataFrame with the assigned column names.
I hope you found this tutorial useful. Feel free to grab the code snippet above and use it in your own applications.