How to loop through multiple polynomial fits changing the degree

Question

My code functions properly but I am repeating a block several times to vary the polynomial variable, degree. I assume this can and should be looped to allow quicker iterations, but I'm not sure how to do it. Prior to the code below I generate the train_test split which I keep for plotting.

After several iterations, I use np.vstack on the y_predictions to create a single array.

from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures

### degree 1 ####
poly1 = PolynomialFeatures(degree=1)
x1_poly = poly1.fit_transform(X_train)
                           
linreg1 = LinearRegression().fit(x1_poly, y_train)
pred_1= poly1.transform(x_prediction_data)
y1_poly_pred=linreg1.predict(pred_1)

### degree 3 #####
poly3 = PolynomialFeatures(degree=3)
x3_poly = poly3.fit_transform(X_train)
                                  
linreg3 = LinearRegression().fit(x3_poly, y_train)
pred_3= poly3.transform(x_prediction_data)
y3_poly_pred=linreg3.predict(pred_3)

#### ect... will make several other degree = 6, 9 ...

Answer 1

I would recommend collecting your answers in a dictionary, but I created a list for simplicity.

The code iterates over i, which is the degree of your polynomials. Trains the model, etc..., then collects its answers.

prediction_collector = []
for i in [1,3,6,9]:

    poly = PolynomialFeatures(degree=i)
    x_poly = poly.fit_transform(X_train)
                               
    linreg = LinearRegression().fit(x_poly, y_train)
    pred= poly.transform(x_prediction_data)
    y_poly_pred=linreg.predict(pred)

    # to collect the answer after each iteration/increase of degrees
    predictions_collector.append(y_poly_pred)