CODEWITHSUNDEEP
numpymatplotlib
user40465
user40465

Reputation: 421

Printing the equation of the best fit line

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 enter image description here

Upvotes: 6

Views: 8608

Answers (2)

user3521099
user3521099

Reputation:

Here's a one liner.

If fit is the poly1d object, while plotting the fitted line, just use label argument as bellow,

label='y=${}$'.format(''.join(['{}x^{}'.format(('{:.2f}'.format(j) if j<0 else '+{:.2f}'.format(j)),(len(fit.coef)-i-1)) for i,j in enumerate(fit.coef)]))

Upvotes: 0

HYRY
HYRY

Reputation: 97331

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:

enter image description here

Upvotes: 8

Related Questions