In [1]:
import os
import numpy as np
import time
import logging
from astropy.io import fits
from PyAstronomy import pyasl
#import matplotlib
#matplotlib.use('Qt5Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gs
from mpl_toolkits.mplot3d import Axes3D
from lmfit.models import GaussianModel, ConstantModel, SkewedGaussianModel, PolynomialModel
from scipy.interpolate import splrep, splev, interp1d
from scipy.ndimage import gaussian_filter
from scipy.stats import gaussian_kde
%matplotlib tk
In [2]:
def bissec(A, x):
n = len(A) - 1
if (x < A[0]):
return 0
elif (x > A[n]):
return n + 1
n1 = 0
n2 = n
while ((n2 - n1) > 1):
nm = round((n1 + n2) / 2)
if ((x - A[nm]) > 0):
n1 = nm
else:
n2 = nm
return n1
In [3]:
def measure_SNR(spec, region):
idx1 = (np.abs(spec[0] - region[0])).argmin()
idx2 = (np.abs(spec[0] - region[1])).argmin()
reg = np.array(spec[1][idx1:idx2])
if (len(reg) != 0) and (np.std(reg) > 1e-10):
return np.mean(reg) / np.std(reg)
else:
return 0
In [4]:
def deg2HMS(ra='', dec='', round=False):
RA, DEC, rs, ds = '', '', '', ''
if dec:
if str(dec)[0] == '-':
ds, dec = '-', abs(dec)
deg = int(dec)
decM = abs(int((dec - deg) * 60))
if round:
decS = int((abs((dec - deg) * 60) - decM) * 60)
else:
decS = (abs((dec - deg) * 60) - decM) * 60
DEC = '{0}{1}:{2}:{3}'.format(ds, deg, decM, decS)
if ra:
if str(ra)[0] == '-':
rs, ra = '-', abs(ra)
raH = int(ra / 15)
raM = int(((ra / 15) - raH) * 60)
if round:
raS = int(((((ra / 15) - raH) * 60) - raM) * 60)
else:
raS = ((((ra / 15) - raH) * 60) - raM) * 60
RA = '{0}{1}:{2}:{3}'.format(rs, raH, raM, raS)
if ra and dec:
return RA, DEC
else:
return RA or DEC
In [5]:
def zeroclean(spec, zl):
rej = []
for i, w in enumerate(spec[1]):
if w < zl:
rej.append(i)
specx = np.delete(spec[0], rej)
specy = np.delete(spec[1], rej)
return np.array([specx, specy])
In [100]:
def findcont(xn, yn, ptreg, ovl=0.3, sliord=1, kdeord=0.3, plax=None):
firstv = xn[0]
wavrange = xn[-1] - firstv
sliceno = wavrange / ptreg
slisize = wavrange / sliceno
points = [[], []]
sigmas = []
# logging.debug("slicing and fitting each slice")
for s in range(int(sliceno)+1):
if s != int(sliceno):
vi = firstv + s * slisize
vf = firstv + (s + 1) * slisize
i = bissec(xn, vi - (slisize * ovl))
f = bissec(xn, vf + (slisize * ovl))
slci = [xn[i:f], yn[i:f]]
else:
vf = xn[-1]
vi = xn[bissec(xn, xn[-1]-slisize)]
i = bissec(xn, vi - (slisize * ovl))
f = bissec(xn, vf + (slisize * ovl))
slci = [xn[i:f], yn[i:f]]
if len(slci[1]) > sliord+1:
contf = PolynomialModel(sliord)
pars_c = contf.make_params()
for p in range(sliord+1):
label = 'c' + str(p)
pars_c[label].set(value=1., vary=True)
contfit = contf.fit(slci[1], pars_c, x=slci[0])
contslc = contfit.eval(x=slci[0])
slc = [slci[0], slci[1]/contslc]
# ax.plot(slci[0], contslc, 'm-')
# md = np.median(slc[1])
# absor = min(slc[1])
# emiss = max(slc[1])
# freq, bin = np.histogram(slc[1], bins=50, range=(absor, high))
# rebin = [(bin[b + 1] + bin[b]) / 2 for b in range(len(bin) - 1)]
density = gaussian_kde(slc[1])
rebin = np.linspace(min(slc[1]), max(slc[1]), 200)
density.covariance_factor = lambda: .30
density._compute_covariance()
freq = density(rebin)
fmode = rebin[np.argmax(freq)]
# fsigma = (rebin[-1] - rebin[0])
sgauss = SkewedGaussianModel()
pars = sgauss.make_params()
pars['center'].set(value=fmode, vary=True, min=fmode-0.4, max=fmode+0.4)
pars['amplitude'].set(value=10., vary=True)
pars['sigma'].set(value=0.03, vary=True)
pars['gamma'].set(value=-1.0, vary=True)
out = sgauss.fit(freq, pars, x=rebin)
xrbn = np.linspace(rebin[0], rebin[-1], num=300)
yrbn = list(out.eval(x=xrbn))
xp = np.mean([vi, vf])
mode = xrbn[yrbn.index(max(yrbn))]*contfit.eval(x=xp)
points[0].append(xp)
points[1].append(mode)
try:
pass
# axz.plot(rebin, freq, 'k-', zs=xp)
# axz.plot(xrbn, yrbn, 'r-', zs=xp)
except:
pass
if plax != None:
rpt = plax.plot([xp], [mode], 'ro', picker=5)[0]
setattr(rpt, 'distpt', [rebin, freq])
setattr(rpt, 'dfitpt', [xrbn, yrbn])
finr = out.best_values
sigmas.append(finr['sigma'])
logging.debug("tracing spline")
spline = splrep(points[0], points[1], k=3)
sigma_function = splev(xn, splrep(points[0], sigmas, k=3))
return spline, sigma_function
In [93]:
class HoverShow:
def __init__(self, fig, axs):
self.fig = fig
self.axs = axs
pid = fig.canvas.mpl_connect('pick_event', self)
#cid = fig.canvas.mpl_connect('button_press_event', self)
def __call__(self, event):
if event.mouseevent.button == 1:
self.axs.clear()
dst = event.artist.distpt
fit = event.artist.dfitpt
self.axs.plot(dst[0], dst[1], 'k--', lw=0.7)
self.axs.plot(fit[0], fit[1], 'm-')
self.axs.set_xlim([min(dst[0])-0.05, max(dst[0])+0.05])
self.axs.set_ylim([min(dst[1])-0.1, max(dst[1])*1.1])
plt.draw()
In [98]:
def splinefit(spec, ptreg, zr=1.E-8, od=3, kdeord=0.3, noise=0.8):
spfig = plt.figure("Spline auto-fit", figsize=(10,5))
gd = gs.GridSpec(3, 4)
ax1 = spfig.add_subplot(gd[0, :-1])
ax2 = spfig.add_subplot(gd[1, :-1], sharex=ax1)
ax3 = spfig.add_subplot(gd[2, :], sharex=ax1)
ax4 = spfig.add_subplot(gd[:-1, -1])
# logging.debug("rejecting negative/zero values")
specr = zeroclean(spec, zr)
specx, specy = specr[0], specr[1]
# logging.debug("first polynomial guess of order %d" % od)
poly = PolynomialModel(od)
pars = poly.make_params()
for p in range(od+1):
label = 'c' + str(p)
pars[label].set(value=1., vary=True)
outcont = poly.fit(spec[1], pars, x=spec[0])
cont0 = outcont.eval(x=spec[0])
xfirst = specx
yfirst = specy / cont0
md1 = np.median(yfirst)
st1 = np.std(yfirst)
rej = []
# logging.debug("rejecting cosmic rays")
for i, w in enumerate(yfirst):
if w > md1 + 4 * st1:
rej.append(i)
xfirst = np.delete(xfirst, rej)
yfirst = np.delete(yfirst, rej)
ax1.plot(xfirst, yfirst, 'k-', linewidth=0.5)
# logging.debug("calculating first composite continuum")
firsteval, sigfunc = findcont(xfirst, yfirst, ptreg, kdeord=kdeord, plax=ax1)
firstcont = splev(xfirst, firsteval)
ax1.plot(xfirst, firstcont, 'g-')
xcut = []
ycut = []
# logging.debug("focusing on continuum points")
for n, i in enumerate(xfirst):
if abs(yfirst[n] - firstcont[n]) < sigfunc[n]*noise:
xcut.append(xfirst[n])
ycut.append(yfirst[n])
ax2.plot(xcut, ycut, 'k-', linewidth=0.5)
# logging.debug("calculating final composite continuum")
finaleval, sigf2 = findcont(xcut, ycut, ptreg, kdeord=kdeord, plax=ax2)
finalcont = splev(xfirst, finaleval)
ax2.plot(xfirst, finalcont, 'g-')
#axz.set_zlim3d(5100, 5200)
# logging.debug("normalizing")
yfin = yfirst / finalcont
ax3.plot(xfirst, yfin, 'b-', linewidth=0.5)
hovershow = HoverShow(spfig, ax4)
return np.array([xfirst, yfin])
In [106]:
def polyfit(spec, zr=1.E-8, sigma=2.0, dcft=0.21, od=4, itr=4, kde_cov=.2):
plfig = plt.figure("order %d polynomial fit"%od)
gd = gs.GridSpec(itr+1, 4)
specr = zeroclean(spec, zr)
specx, specy = specr[0], specr[1]
lpoly = np.polynomial.legendre.legfit(specx, specy, od, rcond=1E-60)
polycont = np.polynomial.legendre.Legendre(lpoly)
firstcont = polycont(specx)
yn = np.copy(specy) / firstcont
pl1 = plfig.add_subplot(gd[0, :-1])
pl1.plot(specx, specy, 'k-', linewidth=0.3)
pl1.plot(specx, firstcont, 'r-', linewidth=0.6)
for i in range(itr):
nitr = str(i + 1)
sigma -= i * dcft * sigma
md = np.median(yn)
density = gaussian_kde(yn)
rebin = np.linspace(min(yn), max(yn)*1.2, 200)
density.covariance_factor = lambda: kde_cov
density._compute_covariance()
freq = density(rebin)
gauss = SkewedGaussianModel()
pars = gauss.make_params()
pars['center'].set(value=md, vary=True)
pars['amplitude'].set(vary=True)
pars['sigma'].set(vary=True)
pars['gamma'].set(vary=True)
out = gauss.fit(freq, pars, x=rebin)
var = sigma * out.best_values['sigma']
xrbn = np.linspace(md-0.1, md+0.1, 200)
yrbn = list(out.eval(x=xrbn))
mode = xrbn[np.argmax(yrbn)]
pld = plfig.add_subplot(gd[i, -1])
pld.plot(rebin, freq, 'k-.', lw=0.7)
pld.plot(xrbn, yrbn, 'm-', lw=1.0)
pld.plot([mode-var/2, mode-var/2, mode+var/2, mode+var/2], [max(yrbn), 0, 0, max(yrbn)], 'r--', lw=0.5)
pld.set_xlim([mode-3*var, mode+3*var])
xnc = np.copy(specx)
ync = np.copy(specy)
mask = []
for n, w in enumerate(ync):
if (yn[n] > mode + var / 2) or (yn[n] < mode - var / 2):
mask.append(n)
xnc = np.delete(xnc, mask)
ync = np.delete(ync, mask)
pln = plfig.add_subplot(gd[i+1, :-1], sharex=pl1)
pln.plot(xnc, ync+0.5, 'r-', lw=0.3)
lpoly = np.polynomial.legendre.legfit(xnc, ync, od, rcond=1E-60)
polycont = np.polynomial.legendre.Legendre(lpoly)
contf = polycont(specx)
yn = specy / contf
pln.plot(specx, yn, 'k-', lw=0.3)
pln.plot(specx, contf+0.5, 'b-', lw=0.8)
pln.set_ylim([0, np.max(ync+0.5)*1.2])
return np.array([specx, yn])
In [9]:
def clean(spec, zr=1.E-8, rej=4.5):
specr = zeroclean(spec, zr)
specx, specy = specr[0], specr[1]
# md = np.median(spectra[1])
# n = int(len(spectra[0]) * 0.8)
# offset = round((len(spectra[0]) - n) / 2)
# absor = md - min(spectra[1][offset:n - offset])
# freq, bin = np.histogram(spectra[1], bins=30, range=(md - absor, md + absor))
# rebin = [(bin[b + 1] + bin[b]) / 2 for b in range(len(bin) - 1)]
density = gaussian_kde(specy)
rebin = np.linspace(min(specy), max(specy), 200)
density.covariance_factor = lambda: .30
density._compute_covariance()
freq = density(rebin)
gauss = SkewedGaussianModel()
pars = gauss.make_params()
pars['center'].set(vary=True)
pars['amplitude'].set(vary=True)
pars['sigma'].set(vary=True)
pars['gamma'].set(vary=True)
out = gauss.fit(freq, pars, x=rebin)
var = rej * out.best_values['sigma']
xrbn = np.linspace(rebin[0], rebin[-1], num=100)
yrbn = list(out.eval(x=xrbn))
mode = xrbn[yrbn.index(max(yrbn))]
mask = []
for i, w in enumerate(specy):
if w > mode + var / 2:
mask.append(i)
xnd = np.delete(specx, mask)
ynd = np.delete(specy, mask)
return np.array([xnd, ynd])
In [10]:
def pyfxcor(inspec, template, vmin=-200., vmax=200., resolution=1, edge_rej=800):
# global fig
rv, cc = pyasl.crosscorrRV(inspec[0], inspec[1], template[0], template[1], vmin, vmax, resolution, skipedge=edge_rej)
# ax4 = fig.add_subplot(414)
cen_gs = np.argmax(cc)
perfx, perfy = rv[cen_gs - 3:cen_gs + 4], cc[cen_gs - 3:cen_gs + 4]
# ax4.plot(rv, cc, 'b-')
curraxlim = plt.axis()
# logging.debug("fitting peak")
try:
gauss = ConstantModel() + GaussianModel()
pars = gauss.make_params()
pars['center'].set(value=rv[np.argmax(cc)], vary=True)
pars['amplitude'].set(value=max(cc)-min(cc), vary=True)
pars['sigma'].set(vary=True)
pars['c'].set(value=np.median(cc), vary=True)
out = gauss.fit(perfy, pars, x=perfx)
ct = out.best_values['center']
cterr = out.params['center'].stderr
except:
#plt.plot(rv, cc, 'b-')
#plt.show()
logging.warning("RV determination did not converge.\nSetting RV to 0.0")
return 0, 0
# out.plot_fit(ax=ax4, xlabel="RV", numpoints=100)
# plt.axis(curraxlim)
return ct, cterr
In [127]:
samp = np.genfromtxt("testspec_e.dat", unpack=True)
newspec = splinefit(samp, 20., zr=1.E-8, od=3, kdeord=0.3, noise=0.25)
plt.tight_layout()
#plt.show()
In [111]:
samp = np.genfromtxt("testspec_p2.dat", unpack=True)
newspec = polyfit(samp, zr=1.E-8, sigma=2.0, dcft=0.18, od=3, itr=4, kde_cov=.1)
plt.tight_layout()
plt.show()