jmquintana79
4/14/2016 - 2:39 PM

## Regression lineal (or polynomial)

Regression lineal (or polynomial)

``````def polyfit(namex,x,namey, y, degree):
import numpy as np
import matplotlib.pyplot as plt

# initialize
results = {}

# calculate polynomial regression
coeffs = np.polyfit(x, y, degree)

# build model (function)
p = np.poly1d(coeffs)

# fit values, and mean
yhat = p(x)                      # or [p(z) for z in x]
ybar = np.sum(y)/len(y)          # or sum(y)/len(y)
ssreg = np.sum((yhat-ybar)**2)   # or sum([ (yihat - ybar)**2 for yihat in yhat])
sstot = np.sum((y - ybar)**2)    # or sum([ (yi - ybar)**2 for yi in y])

# Set regression info
results['polynomial'] = coeffs.tolist() # Polynomial Coefficients
results['R2'] = ssreg / sstot # R**2
results['correlation'] = np.corrcoef(x,y)[0,1]  # correlation index

## PLOT

# build object
fig, ax = plt.subplots()
# plot scatter
ax.scatter(x,y, color="blue")
# set title
plt.title("REGRESSION grade %s- R**2: %.3f Correlation: %.3f"%(degree,results['R2'],results['correlation']))
# axes labels
plt.xlabel(namex)
plt.ylabel(namey)
# plot grid axis
ax.xaxis.grid(True)
ax.yaxis.grid(True)
# plot fitted line
ax.plot(x,yhat, label='fit', color="red")
# display
plt.show()

# return info results
return results``````