Single machine learning algorithm for multiple classes of data : One hot encoder

Question

I have data of the following kind:

   x1  x2  y
0  0   1  1
1  0   2  2
2  0   3  3
3  0   4  4
4  1   1  4
5  1   2  8
6  1   3  12
7  1   4  16

Is it possible to construct a single machine learning algorithm in python/scikit-learn by defining column x1 in such a way that a simple linear regression should give predict(x1=0, x2=5) = 5 and predict(x1=1, x2=5) = 20. My actual problem has multiple values of x1.

To illustrate the problem better: I have the following code with one hot encoder and it doesn't seem to give the accuracy of training the data separately.

import pandas as pd
from sklearn.linear_model import LinearRegression

# Dataframe with x1 = 0 and linear regression gives a slope of 1 as expected

df = pd.DataFrame(data=[{'x1': 0, 'x2': 1, 'y': 1},
                        {'x1': 0, 'x2': 2, 'y': 2},
                        {'x1': 0, 'x2': 3, 'y': 3},
                        {'x1': 0, 'x2': 4, 'y': 4}
                        ],
                  columns=['x1', 'x2', 'y'])

X = df[['x1', 'x2']]
y = df['y']
reg = LinearRegression().fit(X, y)
print(reg.predict(np.array([[0, 5]]))) # Output is 5 as expected

# Dataframe with x1 = 1 and linear regression gives a slope of 5 as expected

df = pd.DataFrame(data=[{'x1': 1, 'x2': 1, 'y': 4},
                        {'x1': 1, 'x2': 2, 'y': 8},
                        {'x1': 1, 'x2': 3, 'y': 12},
                        {'x1': 1, 'x2': 4, 'y': 16}
                        ],
                  columns=['x1', 'x2', 'y'])

X = df[['x1', 'x2']]
y = df['y']
reg = LinearRegression().fit(X, y)
print(reg.predict(np.array([[1, 5]]))) # Output is 20 as expected 

# Combine the two data frames x1 = 0 and x1 = 1 

df = pd.DataFrame(data=[{'x1': 0, 'x2': 1, 'y': 1},
                        {'x1': 0, 'x2': 2, 'y': 2},
                        {'x1': 0, 'x2': 3, 'y': 3},
                        {'x1': 0, 'x2': 4, 'y': 4},
                        {'x1': 1, 'x2': 1, 'y': 4},
                        {'x1': 1, 'x2': 2, 'y': 8},
                        {'x1': 1, 'x2': 3, 'y': 12},
                        {'x1': 1, 'x2': 4, 'y': 16}
                        ],
                  columns=['x1', 'x2', 'y'])

X = df[['x1', 'x2']]
y = df['y']
reg = LinearRegression().fit(X, y)
print(reg.predict(np.array([[0, 5]]))) # Output is 8.75 
print(reg.predict(np.array([[1, 5]]))) # Output is 16.25

# use one hot encoder

df = pd.get_dummies(df, columns=["x1"], prefix=["x1"])
X = df[['x1_0', 'x1_1', 'x2']]
y = df['y']
reg = LinearRegression().fit(X, y)
print(reg.predict(np.array([[1, 0, 5]]))) # Output is 8.75
print(reg.predict(np.array([[0, 1, 5]]))) # Output is 16.25

How can I use pandas and sklearn for the combined data to get the same accuracy using one machine learning model?

Answer 1

I know this is late, but the question got bumped by Community, so...

Your discussion and data suggest that what you want to produce is the model $$y = (3x_1+1)x_2,$$ but that is not a linear model and so linear regression will not find it. You can try any number of other nonlinear models, the best type depending on your real use-case. For instance,

1. If you really just want a linear model for each value of $x_1$, then it's probably best just to split the data along $x_1$ as you started with. You might need to examine your motivation behind wanting just "one machine learning model".

2. You could introduce a new feature, equal to $x_1 x_2$. This loses some information from option (1) when $x_1$ has more than two values, but it might be suitable.

3. You could use a tree-based model with linear regression done at each leaf, as in https://stats.stackexchange.com/questions/78563/regression-tree-algorithm-with-linear-regression-models-in-each-leaf . If the model fitting procedure decides to split only on $x_1$ and regress only on $x_2$, it mostly recovers option (1), though there's no reason to necessarily expect that.

Answer 2

You can have x1 as a categorical variable, convert it to dummy variables (one hot encoder) and then run linear regression (or any other algorithm).