Add trendline with equation in 2D array

Question

For understanding, I generated two 2D arrays. After that, I plot the scatter plot of that array. Now I want to draw a linear trendline and also the equation of that line. But I am doing something wrong. I really don't know how to get equation as well. My code is:

# Import Libraries
import numpy as np
import matplotlib.pyplot as plt

# Generate Random data 
a = np.zeros(shape=(8,8))
a[0] = [1,2,3,4,5,6,7,8]
a[1] = [1,0,3,4,0,6,6,8]
a[2] = [1,2,3,4,5,3,7,8]# Import Libraries
import numpy as np
import matplotlib.pyplot as plt

# Generate Random data 
a = np.zeros(shape=(8,8))
a[0] = [1,2,3,4,5,6,7,8]
a[1] = [1,0,3,4,0,6,6,8]
a[2] = [1,2,3,4,5,3,7,8]
a[3] = [1,2,3,0,5,6,7,8]
a[4] = [1,2,3,4,5,6,7,5]
a[5] = [1,2,3,4,5,6,7,8]
a[6] = [1,2,0,1,5,0,8,8]
a[7] = [1,2,3,4,5,6,7,8]

b = np.zeros(shape=(8,8))
b[0] = [1,1,3,4,5,6,7,8]
b[1] = [1,0,3,4,5,6,7,8]
b[2] = [1,2,3,4,5,6,5,6]
b[3] = [2,2,3,0,5,6,7,8]
b[4] = [1,2,3,8,8,6,7,8]
b[5] = [1,2,3,4,5,6,7,9]
b[6] = [1,2,6,4,5,0,7,8]
b[7] = [1,2,3,4,5,6,7,9]

# Draw scatterplot
plt.figure();
plt.title('Scatter plot')
plt.xlabel('a')
plt.ylabel('b')
plt.xlim(0, 10)
plt.ylim(0, 10)
plt.scatter(a, b)
plt.show()

# Add trendline with equation
#z = np.polyfit(a, b, 1)
#p = np.poly1d(z)
#plt.plot(a,p(a),"r--")
#print "y=%.6fx+(%.6f)"%(z[0],z[1]) #Dont know how it comes!

Thanks for help and suggestions!

Answer 1

The following code works:

plt.figure();
plt.suptitle('Scatter plot')
plt.xlabel('a')
plt.ylabel('b')
plt.scatter(a, b)

z = np.polyfit(a.flatten(), b.flatten(), 1)
p = np.poly1d(z)
plt.plot(a,p(a),"r--")
plt.title("y=%.6fx+%.6f"%(z[0],z[1])) 

plt.show()

np.polyfit, in your case, needs to have x and y as 1d arrays. I put the equation (y = coef x + b) as the title of the plot, but you can change that as you wish.

For instance, plt.text(8,1,"y=%.6fx+%.6f"%(z[0],z[1]), ha='right') instead of plt.title("y=%.6fx+%.6f"%(z[0],z[1])) would print your equation nicely in the lower right corner of your plot (right aligned, at the coordinates x=8, y=1)