Looping Linear Regression in Python

user_10610631 0
Viewed 537

df_tmp

zipcode.  var1   var2.  orders
Z1.       20.     19.    20
Z1.       50.     28.    10
Z1.       80.     45.    26
Z2.       19.     36.    10
Z2.       20.     90.    16
Z2.       40.     59.    18

I am using the following function for overall prediction -

#split x and y
x = df_tmp.drop(columns=['orders'])
y = df_tmp['orders']

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(x, y, test_size = 0.20, random_state = 0)

from sklearn import datasets, linear_model
from sklearn.linear_model import LinearRegression
import statsmodels.api as sm
from scipy import stats

X2 = sm.add_constant(X_train)
est = sm.OLS(y_train, X2).fit()
X_test = sm.add_constant(X_test) # add again the constant
y_test_predicted = est.predict(X_test)
y_test_predicted

However, I want to predict for every zipcode using a similar code. For-Looping might be extremely inefficient, looking for some suggestions on how this can be done

Thanks!

predict dataframe python loops
1 Answers

I would suggest creating a generator that contains the slices of the dataframe with different zipcodes, abstracting your modelling logic into a function and then mapping this onto this generator.

That will be much faster than using for loops.

Code here:

from sklearn.model_selection import train_test_split
import statsmodels.api as sm

zipcodes = df_tmp['zipcode.'].unique()
different_df_sections = (df_tmp.loc[df_tmp['zipcode.'] == x] for x in zipcodes)

def calculate_predictions(df_tmp):

    x = df_tmp.drop(columns=['orders'])
    y = df_tmp['orders']
    X_train, X_test, y_train, y_test = train_test_split(x, y, test_size = 0.20, random_state = 0)
    X2 = sm.add_constant(X_train)
    est = sm.OLS(y_train, X2).fit()
    X_test = sm.add_constant(X_test) # add again the constant
    y_test_predicted = est.predict(X_test)

    return y_test_predicted

different_predictions = map(calculate_predictions, different_df_sections)
all_predictions = [x for x in different_predictions]
user_134241850
Related