In [ ]:
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
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
In [ ]:
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 [ ]:
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 [ ]:
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 [ ]:
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 [ ]:
def findcont(xn, yn, ptreg, ovl=0.3, sliord=1):
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[list(freq).index(max(freq))]
# fsigma = (rebin[-1] - rebin[0])
sgauss = SkewedGaussianModel()
pars = sgauss.make_params()
pars['center'].set(value=fmode, vary=True, min=fmode-0.3, max=fmode+0.3)
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
finr = out.best_values
sigmas.append(finr['sigma'])
logging.debug("tracing spline")
spline = splrep(points[0], points[1], k=3)
fitcont = splev(xn, spline)
# ax.plot(points[0], points[1], 'ro')
# ax.plot(xfirst, fitcont, 'g-')
sigma_function = splev(xn, splrep(points[0], sigmas, k=3))
return fitcont, sigma_function
In [ ]:
def splinefit(spec, ptreg, zr=1.E-8, od=3, noise=0.8):
# ax1 = fig.add_subplot(411)
# ax2 = fig.add_subplot(412, sharex=ax1)
# ax3 = fig.add_subplot(413, sharex=ax1)
# 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")
firstcont, sigfunc = findcont(xfirst, yfirst, ptreg)
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")
finalcont, sigf2 = findcont(xcut, ycut, ptreg)
# axz.set_zlim3d(5100, 5200)
# logging.debug("normalizing")
yfin = yfirst / finalcont
# ax3.plot(xfirst, yfin, 'b-', linewidth=0.5)
# meand = np.mean(spectra[1])
# ax3.plot(spectra[0], spectra[1]/meand + 2, 'r-', linewidth=0.5)
# ax3.plot(spectra[0], cont0/meand + 2, 'g-')
# ax3.set_ylim([0, max(cont0/meand + 2)*1.2])
return [xfirst, yfin]
In [ ]:
def polyfit(spec, zr=1.E-8, sigma=2.0, dcft=0.21, od=4, itr=4):
specr = zeroclean(spec, zr)
specx, specy = specr[0], specr[1]
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(specy, pars, x=specx)
firstcont = outcont.eval(x=specx)
xn = np.copy(specx)
yn = np.copy(specy) / firstcont
# pl1 = plt.subplot((itr + 1) * 100 + 11)
# pl1.plot(xn, spectra[1], 'k-', linewidth=0.3)
# pl1.plot(xn, firstcont, 'r-', linewidth=0.6)
# pl1.set_ylim([0, np.mean(firstcont)*1.5])
for i in range(itr):
nitr = str(i + 1)
sigma -= i * dcft * sigma
# md = np.median(yn)
# n = len([i for i in yn if i > 0.1])
# offset = (len(xn) - n) / 2
# absor = md - min(yn[offset:n - offset])
# freq, bin = np.histogram(yn, bins=50, range=(md - absor, md + absor))
# rebin = [(bin[b + 1] + bin[b]) / 2 for b in range(len(bin) - 1)]
density = gaussian_kde(yn)
rebin = np.linspace(min(yn), max(yn), 200)
density.covariance_factor = lambda: .30
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(rebin[0], rebin[-1], num=100)
yrbn = list(out.eval(x=xrbn))
mode = xrbn[yrbn.index(max(yrbn))]
ync = np.copy(xn)
xnc = np.copy(yn)
mask = []
for i, w in enumerate(ync):
if (ync[i] > mode + var / 2) or (ync[i] < mode - var / 2):
mask.append(i)
xnc = np.delete(xnc, mask)
ync = np.delete(ync, mask)
poly2 = PolynomialModel(od)
pars2 = poly2.make_params()
for p in range(od + 1):
label = 'c' + str(p)
pars2[label].set(value=1., vary=True)
outcont2 = poly2.fit(ync, pars2, x=xnc)
contf = outcont2.eval(x=specx)
yn = spec[1] / contf
# err = spec[2] / contf
# pln=plt.subplot(int((iter+1)*100+10+(i_+2)))
# pln.plot(specx, specy*(np.mean(contf)*0.8), 'k-', linewidth=0.3)
# pln.plot(specx, ync, 'r-', linewidth=0.3)
# pln.plot(specx, contf, 'b-', linewidth=0.6)
# pln.set_ylim([0, np.mean(contf)*1.2])
# plt.savefig(obj[0]+'_fit.png', dpi=300)
return np.array([specx, yn, err])
In [ ]:
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 [ ]:
def pyfxcor(inspec, template, vmin=-50., vmax=50., res=1, rej=800):
global fig
rv, cc = pyasl.crosscorrRV(inspec[0], inspec[1], template[0], template[1], vmin, vmax, res, skipedge=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['gamma'].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