python - 有没有办法在 matplotlib 图上动态编写数学公式？

Question

我正在用 Python 为我的实验室工作制作一个模板。总结一下它的目的，它是绘制数据点并使用 scipy curve_fit 拟合预定义的模型。通常我拟合多项式或指数曲线。我设法在绘图上动态打印拟合参数，但每次都必须手动输入相关方程。我想知道，是否有一种优雅的方式来动态地执行此操作？我读过关于 sympy 的内容，但暂时我做不到。

代码如下:

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from datetime import datetime

#two example functions
def f(x, p0, p1):
    return p0 * x + p1

def g(x, p0, p1):
    return p0 * np.exp(x * p1)

#example data
xval = np.array([0,1,2,3,4,5,6])
yval = np.array([0, 2,3.95,5.8,8.1, 10.2, 12.4])

#curve fitting
popt, pcov = curve_fit(f, xval, yval)

plt.rcParams.update({'font.size': 12})
plt.figure(figsize=(9,7))
plt.plot(xval, yval,'ko', label = 'Data points', markersize = 7)
plt.title('TITLE', fontsize = 15)
plt.grid()
plt.plot(xval, f(xval, *popt),'r-', label = 'Fit')
#printing the params on plot
for idx in range(len(popt)):
    plt.text(0.8,0.05+0.05*(idx+1), 'p'+str(idx)+' = {0:.5f}'.format(popt[idx]), transform=plt.gca().transAxes)

#manually writing the equation, that's what I want to print dynamically
plt.text(0.8, 0.05, '$y = p0 \cdot x + p1 $' , transform=plt.gca().transAxes)

plt.text(0.86, 1.01, datetime.today().strftime('%Y.%m.%d.'), transform=plt.gca().transAxes)
plt.text(0 ,1.01, 'NAME', transform=plt.gca().transAxes)
plt.ylabel('Y axis title')
plt.xlabel('X axis title')
plt.legend()
plt.show()

预期结果是:

如果我使用函数进行拟合 - 假设 g(x, p0, p1) 返回 p0 * np.exp(x * p1) 然后返回公式本身应该打印在绘图上，就像示例代码中的另一个公式一样:

plt.text(0.8, 0.05, '$y = p0 \cdot x + p1 $' , transform=plt.gca().transAxes) 

除非它是手动解决方案。

我非常感谢任何建议。

Answer 1

我认为你可以使用 sympy 包。它允许定义自定义变量、创建表达式然后对其求值。我不确定对性能有什么影响

这是经过更改的代码:

import numpy as np
import sympy
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from datetime import datetime


#two example functions
x, p0, p1 = sympy.var("x p0 p1")

f = p0 * x + p1 

g = p0 * sympy.exp(x*p1)

def partial_fun(sympy_expr):

    def res_fun(X, P0, P1):
        return np.array([sympy_expr.evalf(subs={x: x_, p0: P0, p1: P1}) for x_ in X], dtype=np.float)

    return res_fun

#example data
xval = np.array([0,1,2,3,4,5,6])
yval = np.array([0, 2,3.95,5.8,8.1, 10.2, 12.4])

#curve fitting
popt, pcov = curve_fit(partial_fun(f), xval, yval)

plt.rcParams.update({'font.size': 12})
plt.figure(figsize=(9,7))
plt.plot(xval, yval,'ko', label = 'Data points', markersize = 7)
plt.title('TITLE', fontsize = 15)
plt.grid()
plt.plot(xval, partial_fun(f)(xval, *popt),'r-', label = 'Fit')
#printing the params on plot
for idx in range(len(popt)):
    plt.text(0.8,0.05+0.05*(idx+1), 'p'+str(idx)+' = {0:.5f}'.format(popt[idx]), transform=plt.gca().transAxes)

#manually writing the equation, that's what I want to print dynamically
plt.text(0.8, 0.05, f'$y = {f} $' , transform=plt.gca().transAxes)

plt.text(0.86, 1.01, datetime.today().strftime('%Y.%m.%d.'), transform=plt.gca().transAxes)
plt.text(0 ,1.01, 'NAME', transform=plt.gca().transAxes)
plt.ylabel('Y axis title')
plt.xlabel('X axis title')
plt.legend()
plt.show()