如何用Python对数据集分段拟合：负区间线性+正区间余弦函数

实现步骤 #

1. 数据拆分：根据指定区间筛选对应数据子集
2. 线性拟合：对x∈[-0.65,-0.049]的子集用一次多项式拟合
3. 余弦拟合：对x∈[0.049,1.85]的子集用周期约0.4的余弦函数拟合
4. 结果可视化：绘制原始数据与拟合曲线对比

完整代码 #

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

# 给定数据集
x = np.array([ 1.85      ,  1.74990234,  1.65004883,  1.54995117,  1.45009766,
            1.35      ,  1.24990234,  1.15004883,  1.04995117,  0.95009766,
            0.85      ,  0.74990234,  0.65004883,  0.54995117,  0.45009766,
            0.35      ,  0.24990234,  0.15004883,  0.04995117, -0.04990234,
           -0.15      , -0.25009766, -0.34995117, -0.45004883, -0.54990234,
           -0.65      ])
y = np.array([-1.9764485 , -2.51604791, -2.47880877, -1.65158466, -0.57161459,
       -1.73628227, -1.67898043, -2.63570217, -0.51512979, -0.73739313,
       -1.37574256, -1.96644884, -0.57822485, -0.57741358, -1.14066144,
       -1.64367465, -1.4994257 , -0.11035846, -0.36308983,  0.81232546,
       -1.22662672, -0.36636166, -0.91063433,  0.52619598, -0.53499829,
        0.19404025])

# 1. 拆分数据集
# 线性拟合区间：x ∈ [-0.65, -0.049]
mask_linear = (x >= -0.65) & (x <= -0.049)
x_linear = x[mask_linear]
y_linear = y[mask_linear]

# 余弦拟合区间：x ∈ [0.049, 1.85]
mask_cos = (x >= 0.049) & (x <= 1.85)
x_cos = x[mask_cos]
y_cos = y[mask_cos]

# 2. 线性拟合（一次多项式）
p_linear = np.polyfit(x_linear, y_linear, 1)
def linear_func(x):
    return p_linear[0] * x + p_linear[1]

# 3. 余弦拟合（周期约0.4）
# 定义余弦函数，固定周期为0.4，也可放开让模型拟合周期
def cos_func(x, a, ph, off):
    period = 0.4
    return a * np.cos((2 * np.pi / period) * (x - ph)) + off

# 初始参数猜测：振幅、相位偏移、偏移量
guess_cos = [1.0, 0.1, -1.0]
popt_cos, _ = curve_fit(cos_func, x_cos, y_cos, p0=guess_cos)

# 4. 生成拟合曲线用于可视化
x_fit_linear = np.linspace(-0.65, -0.049, 100)
y_fit_linear = linear_func(x_fit_linear)

x_fit_cos = np.linspace(0.049, 1.85, 200)
y_fit_cos = cos_func(x_fit_cos, *popt_cos)

# 绘图展示结果
plt.figure(figsize=(10, 6))
plt.scatter(x, y, label='原始数据', c='blue', s=50)
plt.plot(x_fit_linear, y_fit_linear, label=f'线性拟合: y={p_linear[0]:.3f}x + {p_linear[1]:.3f}', c='red', linewidth=2)
plt.plot(x_fit_cos, y_fit_cos, label=f'余弦拟合: y={popt_cos[0]:.3f}cos((2π/0.4)(x-{popt_cos[1]:.3f})) + {popt_cos[2]:.3f}', c='green', linewidth=2)
plt.xlabel('x', fontsize=12)
plt.ylabel('y', fontsize=12)
plt.legend(fontsize=10)
plt.grid(alpha=0.3)
plt.show()

# 输出拟合参数
print("线性拟合参数：斜率 =", p_linear[0], "截距 =", p_linear[1])
print("余弦拟合参数：振幅 =", popt_cos[0], "相位偏移 =", popt_cos[1], "偏移量 =", popt_cos[2])

原代码问题修正 #

• 原代码筛选x>0.5，不符合需求的x≥0.049区间
• 原代码未处理线性拟合区间的逻辑
• 原代码存在语法错误：i_fit = y[i][ind[0]:ind[-1]])中i未定义，且多了一个闭合括号

内容的提问来源于stack exchange，提问作者Krystal