Printing the equation of the best fit line

Question

I have created the best fit lines for the dataset using the following code:

fig, ax = plt.subplots()
for dd,KK in DATASET.groupby('Z'):
  fit = polyfit(x,y,3)
  fit_fn = poly1d(fit)
  ax.plot(KK['x'],KK['y'],'o',KK['x'], fit_fn(KK['x']),'k',linewidth=4)
  ax.set_xlabel('x')
  ax.set_ylabel('y')

The graph displays the best fit line for each group of Z. I want print the equation of the best fit line on top of the line.Please suggest what can i do out here

Answer 1

So you need to write some function that convert a poly parameters array to a latex string, here is an example:

import pylab as pl
import numpy as np

x = np.random.randn(100)
y = 1 + 2 * x + 3 * x * x + np.random.randn(100) * 2

poly = pl.polyfit(x, y, 2)

def poly2latex(poly, variable="x", width=2):
    t = ["{0:0.{width}f}"]
    t.append(t[-1] + " {variable}")
    t.append(t[-1] + "^{1}")

    def f():
        for i, v in enumerate(reversed(poly)):
            idx = i if i < 2 else 2
            yield t[idx].format(v, i, variable=variable, width=width)

    return "${}$".format("+".join(f()))

pl.plot(x, y, "o", alpha=0.4)
x2 = np.linspace(-2, 2, 100)
y2 = np.polyval(poly, x2)
pl.plot(x2, y2, lw=2, color="r")
pl.text(x2[5], y2[5], poly2latex(poly), fontsize=16)

Here is the output: