python - 如何对 pcolormesh 图进行线性回归

Question

我正在尝试运行最小二乘算法来将一条线拟合到我的数据中。

这段代码和我的问题类似； 'z' 是一个 50 in 50 矩阵，例如，我想通过功率大于 0.25 的数据拟合一条线。 (按幂表示右边的颜色条)

我想知道如何才能通过电源拟合最佳线？在我的数据中，x 轴是时间，y 轴是频率，z 是功率。此外，我的数据更均匀，因此超过 0.25 的幂几乎是线性的。

import matplotlib.pyplot as plt
import numpy as np
from scipy.ndimage.filters import gaussian_filter

# Generate data for the plot
x = np.linspace(0, 1, 51)
y = np.linspace(0, 1, 51)
r = np.random.RandomState(42)
z = gaussian_filter(r.random_sample([50, 50]), sigma=5, mode='wrap')
z -= np.min(z)
z /= np.max(z)

# Generate the plot
fig, ax = plt.subplots()
cmap = ax.pcolormesh(x, y, z)
fig.colorbar(cmap)
plt.show(fig)

我只需要一些关于如何绘制最佳拟合线的指导。

Answer 1

您首先需要找到最佳拟合线，并将其绘制到同一轴上:

thresh = 0.25

# location where z > thresh
# x,y,z are the same as yours
xc, yc = np.where(z > thresh)

# find best fit line
a, b = np.polyfit(x[xc],y[yc], 1)

# Generate the plot
fig, ax = plt.subplots()

# show color mesh
cmap = ax.pcolormesh(x, y, z, vmin=thresh)

# plot the line
ax.plot(x,a*x + b, c='w')

# set the y limit to avoid line outside of the mesh
ax.set_ylim(y.min(), y.max())

fig.colorbar(cmap)
plt.show(fig)

输出: