python 基线拟合

from numpy import random
import matplotlib.pyplot as plt
from astropy.modeling import models, fitting
from astropy.stats import sigma_clip
import numpy as np

#自己创建一个数据
x = np.array(range(600))
y = np.random.rand(600)+10
y = y+ x*0.01
y[10:15] = y[10:15]+100      #我们把着两个峰比作我们的信号
y[500:505] = y[500:505]+150 

plt.plot(x,y,label = 'before fit')     #如图一


#线性拟合
Order = 1   #把order改成2可以对基线做2次项的拟合

p1 = models.Polynomial1D(Order)
pfit = fitting.LinearLSQFitter()
sigma_clip_fit = fitting.FittingWithOutlierRemoval(pfit, sigma_clip, niter=3, sigma=3.0)   #可以剔除突出的信号，例如图上的两个峰
model_pI, mask_I = sigma_clip_fit(p1, x, y)
y_new  = y - model_pI(x)

plt.plot(x,y_new, label = 'after fit')
plt.legend()
plt.show()

这是拟合前后的图，可以看到不仅基线平直了，并且在0的水平线上