def fit_2dwave_solution(pixtab2d, deg=5):
# Transpose the input table cause it's easier and faster
# to iterate over rows than columns:
tab2d = pixtab2d.T
fit_table2d = np.zeros_like(tab2d)
col = np.arange(pixtab2d.shape[0])
for num, points in enumerate(tab2d):
# Median filter each column
med_col, mask = median_filter_data(points)
# Fit Cheb. poly to each filtered column
try:
cheb_polyfit = Chebyshev.fit(col[mask],
points[mask],
deg=deg,
domain=(col.min(), col.max()))
except:
print(
"Something went wrong in Chebyshev polynomial fitting. Here's the input:"
)
print("x:", col[mask])
print("y:", points[mask])
print("degree:", deg)
print("domain: %r %r" % (col.min(), col.max()))
raise
finally:
np.savetxt('pixtable_2d_dump.dat', pixtab2d)
# Insert back into the fit_table
fit_table2d[num] = cheb_polyfit(col)
# Transpose back to original orientation of the pixel table
return fit_table2d.T
def varPlot(fig, pca, x, bestFit):
lim_w = numpy.argmax(numpy.abs(bestFit), axis=0)
min_w = numpy.argmin(bestFit, axis=0)
max_w = numpy.argmax(bestFit, axis=0)
print "min: %s" % (min_w)
print "max: %s" % (max_w)
print "lim: %s" % (lim_w)
if x == None:
x = numpy.arange(bestFit.shape[0])
ii = numpy.argsort(x)
for i in range(pca.n_components):
if i == 0:
pRef = p = fig.add_subplot(pca.n_components,1,i+1)
else:
p = fig.add_subplot(pca.n_components,1,i+1, sharex=pRef)
p.locator_params(axis='y', nbins=4)
p.plot(x[ii], bestFit[ii,i], '+')
p.hlines(0, *p.get_xlim(), color='r', alpha=0.5)
t = T.fit(x[ii], bestFit[ii,i], 1)
pt = t(x[ii])
p.plot(x[ii], pt)
p.plot(x[min_w[i]], bestFit[min_w[i],i], 'r*')
p.plot(x[max_w[i]], bestFit[max_w[i],i], 'r*')
return min_w, max_w, lim_w
def set_level(self):
self.status_bar.showMessage('calculating levels...')
f = cv2.calcHist([self.foct_data.flatten()], [0], None, [256],
[0, 256])
f = f[:, 0].astype(int)
# fit vars
# only want to fit right of max peak
max_loc = f.argmax()
f_fmax = f[max_loc:]
x_max = np.linspace(max_loc, 255, num=len(f_fmax))
# fit a polynomial of order
order = 50
poly_fit = T.fit(x_max, f_fmax, order)
# get real roots
roots = np.real_if_close(poly_fit.deriv(2).roots())
roots = roots[np.isreal(roots)].real
# get roots to right of max peak
roots = roots[((roots >= max_loc) & (roots < 256))]
up_thresh = int(roots[((poly_fit(roots) / f.max()) < 0.15).argmax()])
# get lower thresh
r_f = np.arange(len(f)) > f.argmax()
up_f = (f / f.max() < 0.002)
low_thresh = (r_f & up_f).argmax()
self.setLevels(up_thresh, low_thresh)
self.status_bar.clearMessage()
self.status_bar.showMessage('Done!', 3000)
def get_fit(data, name='c', id_var=None):
"""Fit a curve to data for variable deviations."""
# Change to long format
try:
data = data.melt(id_vars=id_var)
except KeyError:
data = data.melt()
# Drop missing values
data = data.dropna()
# Take all x and y data
x_data = data.variable.values.astype(np.float)
y_data = data.value.values.astype(np.float)
# #Fit x and y data with 4th degree polynomial
# z = np.polyfit(x_data.astype(np.float), y_data.astype(np.float), 4)
# f = np.poly1d(z)
# x_curve = data.variable.unique()
# y_curve = f(x_curve)
# !!!! Chebyshev polynomial
c_fit = Chebyshev.fit(x_data, y_data, 5, domain=[0, 1])
x_curve = data.variable.unique()
y_curve = c_fit(x_curve)
# !!!! End
# Create DF from obtained fit
curve = pd.DataFrame(index=[name], columns=x_curve.astype(int), data=np.reshape(y_curve, (-1, len(y_curve))))
return curve
def fit(self):
pixvals, wavelengths = self.get_table_values()
mask = ~np.isnan(wavelengths)
order = int(self.poly_order.text())
if np.sum(~np.isnan(wavelengths)) < order:
msg = "Not enough data points to perform fit!\n"
msg += "Choose a lower polynomial order or identify more lines."
QMessageBox.critical(None, 'Not enough data to fit', msg)
else:
p_fit = Chebyshev.fit(pixvals[mask],
wavelengths[mask],
order,
domain=[self.pix.min(),
self.pix.max()])
wave_solution = p_fit(self.pix)
scatter = np.std(wavelengths[mask] - p_fit(pixvals[mask]))
scatter_label = r"$\sigma_{\lambda} = %.2f$ Å" % scatter
self.cheb_fit = p_fit
self._scatter = scatter
self._residuals = wavelengths - p_fit(pixvals)
if self._fit_ref is None:
fit_ref = self.ax2.plot(self.pix,
wave_solution,
color='RoyalBlue',
label=scatter_label)
self._fit_ref = fit_ref[0]
else:
self._fit_ref.set_ydata(wave_solution)
self._fit_ref.set_label(scatter_label)
self.update_plot()
self.ax2.legend(handlelength=0.5, frameon=False)
self.canvas.draw()
self.canvas.setFocus()
def fit_background_image(data, order_bg=3, xmin=0, xmax=None, kappa=10, fwhm_scale=3):
"""
Fit background in 2D spectral data. The background is fitted along the spatial rows by a Chebyshev polynomium.
Parameters
==========
data : np.array(M, N)
Data array holding the input image
order_bg : integer [default=3]
Order of the Chebyshev polynomium to fit the background
xmin, xmax : integer [default=0, None]
Mask out pixels below xmin and above xmax
fwhm_scale : float [default=3]
Number of FWHM below and above centroid of auto-detected trace
that will be masked out during fitting.
kappa : float [default=10]
Threshold for masking out cosmic rays etc.
Returns
=======
bg2D : np.array(M, N)
Background model of the 2D frame, same shape as input data.
"""
x = np.arange(data.shape[1])
if xmax is None:
xmax = len(x)
if xmax < 0:
xmax = len(x) + xmax
SPSF = np.nanmedian(data, 0)
noise = 1.5*mad(SPSF)
peaks, properties = signal.find_peaks(SPSF, prominence=kappa*noise, width=3)
mask = (x >= xmin) & (x <= xmax)
for num, center in enumerate(peaks):
width = properties['widths'][num]
x1 = center - width*fwhm_scale
x2 = center + width*fwhm_scale
obj = (x >= x1) * (x <= x2)
mask &= ~obj
bg2D = np.zeros_like(data)
for i, row in enumerate(data):
# Median filter the data to remove outliers:
med_row = median_filter(row, 15)
noise = mad(row)*1.4826
this_mask = mask * (np.abs(row - med_row) < 10*noise)
if np.sum(this_mask) > order_bg+1:
bg_model = Chebyshev.fit(x[this_mask], row[this_mask], order_bg, domain=[x.min(), x.max()])
bg2D[i] = bg_model(x)
return bg2D
def subtract_arc_background(arc2D, deg=5):
"""Subtract continuum background in arc line frames"""
if arc2D.dtype not in [np.float64, np.float32]:
arc2D = arc2D.astype(np.float64)
pix = np.arange(arc2D.shape[1])
bg2d = np.zeros_like(arc2D)
for i, row in enumerate(arc2D):
med1d, mask1d = median_filter_data(row)
cheb_fit = Chebyshev.fit(pix[mask1d], row[mask1d], deg=deg)
bg1d = cheb_fit(pix)
bg2d[i] = bg1d
bg_sub2d = arc2D - bg2d
return bg_sub2d, bg2d
def detrending(time, flux, err, pl):
fluxTrue = np.isfinite(flux)
index = np.where(fluxTrue == True)[0]
flux_nan = flux[index]
time = time[index]
err_nan = err[index]
time_nan = (time - min(time))
time_nan, flux_nan = zip(*sorted(zip(time_nan, flux_nan)))
time_nan = np.array(time_nan)
flux_nan = np.array(flux_nan)
p = T.fit(time_nan, flux_nan, pl)
flux_model = p(time_nan)
flux_detrended = (flux_nan - flux_model)
flux_detrended = flux_detrended + 1
return time_nan, flux_nan, flux_model, flux_detrended, err_nan
exa=p([1,2,3])
print(exa.deriv(1))
import matplotlib.pyplot as plt
from numpy.polynomial import Chebyshev as T
x = np.linspace(-1, 1, 100)
#x= np.linspace(-2,2,100)
for i in range(5):
ax = plt.plot(x, T.basis(i)(x), lw=2, label="$T_%d$"%i)
plt.legend()
plt.show()
np.random.seed(11)
x = np.linspace(0, 2*np.pi, 20)
y = np.sin(x) + np.random.normal(scale=.1, size=x.shape)
p = T.fit(x, y, 5)
plt.plot(x, y, 'o')
xx, yy = p.linspace()
plt.plot(xx, yy, lw=2)
p.domain
p.window
plt.show()
print("polyvalues finding ")
polyval(1,[1,2,3])
from scipy.stats import alpha
a = 3.57
mean, var, skew, kurt = alpha.stats(a, moments='mvsk')
def wavecal_1d(input_fname,
pixtable_fname,
*,
output,
order_wl=None,
log=False,
N_out=None,
linearize=True):
"""Apply wavelength calibration to 1D spectrum"""
msg = list()
pixtable = np.loadtxt(pixtable_fname)
msg.append(" - Loaded pixel table: %s" % pixtable_fname)
if order_wl is None:
order_wl, found_in_file = get_order_from_file(pixtable_fname)
if found_in_file:
msg.append(" - Loaded polynomial order from file: %i" %
order_wl)
else:
msg.append(" - Using default polynomial order: %i" %
order_wl)
else:
msg.append(" - Using polynomial order: %i" % order_wl)
if log:
linearize = True
# Load input data:
hdu_list = fits.open(input_fname)
msg.append(" - Loaded spectrum: %s" % input_fname)
output_hdu = fits.HDUList()
for hdu in hdu_list[1:]:
wl_unit = hdu.columns['WAVE'].unit
if wl_unit and wl_unit.lower() in ['angstrom', 'nm', 'a', 'aa']:
msg.append(
" [ERROR] - Spectrum is already wavelength calibrated.")
msg.append("")
output_msg = "\n".join(msg)
return output_msg
tab = hdu.data
hdr = hdu.header
pix = tab['WAVE']
flux1d = tab['FLUX']
err1d = tab['ERR']
if N_out is None:
N_out = len(pix)
else:
if N_out != len(pix):
linearize = True
# Fit wavelength solution
solution = Chebyshev.fit(pixtable[:, 0],
pixtable[:, 1],
deg=order_wl,
domain=[pix.min(), pix.max()])
wl = solution(pix)
res = np.std(pixtable[:, 1] - solution(pixtable[:, 0]))
msg.append(
" - Fitting wavelength solution with polynomium of order: %i"
% order_wl)
msg.append(
" - Standard deviation of wavelength residuals: %.3f Å" %
res)
hdr['CUNIT1'] = 'Angstrom'
hdr['WAVERES'] = (np.round(res, 2), "RMS of wavelength residuals")
if linearize:
if log:
msg.append(
" - Interpolating spectrum onto logarithmic grid")
wl_new = np.logspace(np.log10(wl.min()), np.log10(wl.max()),
N_out)
dv = np.diff(wl_new)[0] / wl_new[0] * 299792.
dlog = np.diff(np.log10(wl_new))[0]
msg.append(" - wavelength step: %.3f [log(Å)]" %
dlog)
msg.append(" - wavelength step: %.1f [km/s]" % dv)
else:
msg.append(
" - Interpolating spectrum onto linear grid")
wl_new = np.linspace(wl.min(), wl.max(), N_out)
dl = np.diff(wl_new)[0]
msg.append(" - wavelength step: %.3f [Å]" % dl)
if np.diff(wl)[0] < 0:
# Wavelengths are decreasing: Flip arrays
flux1d = flux1d[::-1]
wl = wl[::-1]
err1d = err1d[::-1]
interp_flux, interp_err = spectres.spectres(wl_new,
wl,
flux1d,
spec_errs=err1d,
verbose=False,
fill=0.)
final_wl = wl_new
else:
msg.append(
" - Using raw input grid, no interpolation used.")
msg.append("[WARNING] - Wavelength steps may not be constant!")
final_wl = wl
interp_flux = flux1d
interp_err = err1d
col_wl = fits.Column(name='WAVE',
array=final_wl,
format='D',
unit=hdr['CUNIT1'])
col_flux = fits.Column(name='FLUX',
array=interp_flux,
format='D',
unit=hdr['BUNIT'])
col_err = fits.Column(name='ERR',
array=interp_err,
format='D',
unit=hdr['BUNIT'])
output_tab = fits.BinTableHDU.from_columns([col_wl, col_flux, col_err],
header=hdr)
output_hdu.append(output_tab)
output_hdu.writeto(output, overwrite=True)
msg.append(" [OUTPUT] - Saving wavelength calibrated 1D spectrum: %s" %
output)
msg.append("")
output_msg = "\n".join(msg)
return output_msg
def apply_transform(img2D,
pix,
fit_table2d,
ref_table,
err2D=None,
mask2D=None,
header={},
order_wl=4,
log=False,
N_out=None,
interpolate=True):
"""
Apply 2D wavelength transformation to the input image
img2D : array, shape(M, N)
Input 2D spectrum oriented with dispersion along x-axis!
pix : array, shape (N)
Input pixel array along dispersion axis.
fit_table2d : array, shape(M, L)
Fitted wavelength solutions from corresponding arc-line frame
for L fitted reference lines.
ref_table : array, shape(L, 2)
Reference table for the given setup with two columns:
pixel and wavelength in Å
err2D : array, shape(M, N)
Associated error array, must be same shape as `img2D`
header : FITS Header or dict
The FITS header corresponding to `img2D`, or a dictionary
order_wl : int
Polynomial order used for wavelength fit as a function of input pixel
A Chebyshev polynomium is used to fit the solution for each row in img2D.
log : bool [default=False]
Use logarithmic binning in wavelength?
N_out : int
Number of output pixels along dispersion axis.
If `None` is given, the default is to use the same number
of pixels as in the input image.
interpolate : bool [default=True]
Interpolate the image onto new grid or use sub-pixel shifting
Returns
-------
img2D_tr : array, shape(M, N_out)
Transformed 2D spectrum with wavelength along x-axis.
wl : array, shape(N_out)
Coresponding wavelength array along x-axis.
hdr_tr : FITS Header or dict
Updated FITS Header or dictionary with wavelength information
"""
msg = list()
# Define wavelength grid at midpoint:
if N_out is None:
N_out = img2D.shape[1]
else:
if not interpolate:
N_out = img2D.shape[1]
msg.append("[WARNING] - Interpolation turned off!")
msg.append("[WARNING] - N_out was given: %i" % N_out)
msg.append(
"[WARNING] - Cannot change sampling without interpolating")
pix_in = pix
cen = fit_table2d.shape[0] // 2
ref_wl = ref_table[:, 1]
central_solution = Chebyshev.fit(fit_table2d[cen],
ref_wl,
deg=order_wl,
domain=[pix_in.min(),
pix_in.max()])
wl_central = central_solution(pix_in)
wl_residuals = np.std(ref_wl - central_solution(fit_table2d[cen]))
msg.append(" - Residuals of wavelength solution: %.2f Å" %
wl_residuals)
if log:
wl = np.logspace(np.log10(wl_central.min()),
np.log10(wl_central.max()), N_out)
hdr_tr = header.copy()
hdr_tr['CRPIX1'] = 1
hdr_tr['CDELT1'] = np.diff(np.log10(wl))[0]
hdr_tr['CRVAL1'] = np.log10(wl[0])
hdr_tr['CTYPE1'] = 'LOGLAM '
hdr_tr['CUNIT1'] = 'Angstrom'
msg.append(
" - Creating logarithmically sampled wavelength grid")
msg.append(" - Sampling: %.3f (logÅ/pix)" %
np.diff(np.log10(wl))[0])
else:
wl = np.linspace(wl_central.min(), wl_central.max(), N_out)
hdr_tr = header.copy()
hdr_tr['CRPIX1'] = 1
hdr_tr['CDELT1'] = np.diff(wl)[0]
hdr_tr['CRVAL1'] = wl[0]
hdr_tr['CTYPE1'] = 'LINEAR '
hdr_tr['CUNIT1'] = 'Angstrom'
msg.append(" - Creating linearly sampled wavelength grid")
msg.append(" - Sampling: %.3f (Å/pix)" % np.diff(wl)[0])
# Calculate the maximum curvature of the arc lines:
max_curvature = table2D_max_curvature(fit_table2d)
msg.append(" - Maximum curvature of arc lines: %.3f pixels" %
max_curvature)
if max_curvature < 0.1:
interpolate = False
msg.append(
" - Maximum curvature less than 1/10 pixel. No need to interpolate the data"
)
if interpolate:
img2D_tr = np.zeros((img2D.shape[0], N_out))
err2D_tr = np.zeros((img2D.shape[0], N_out))
mask2D_tr = np.zeros((img2D.shape[0], N_out))
for i, row in enumerate(img2D):
# - fit the chebyshev polynomium
solution_row = Chebyshev.fit(fit_table2d[i],
ref_wl,
deg=order_wl,
domain=[pix_in.min(),
pix_in.max()])
wl_row = solution_row(pix_in)
if np.diff(wl_row)[0] < 0:
# Wavelengths are decreasing: Flip arrays
row = row[::-1]
wl_row = wl_row[::-1]
flip_array = True
else:
flip_array = False
# -- interpolate the data onto the fixed wavelength grid
if err2D is not None:
err_row = err2D[i]
if flip_array:
err_row = err_row[::-1]
interp_row, interp_err = spectres.spectres(wl,
wl_row,
row,
spec_errs=err_row,
verbose=False,
fill=0.)
err2D_tr[i] = interp_err
else:
interp_row = spectres.spectres(wl,
wl_row,
row,
verbose=False,
fill=0.)
mask_row = mask2D[i]
mask_int = np.interp(wl, wl_row, mask_row)
mask2D_tr[i] = np.ceil(mask_int).astype(int)
img2D_tr[i] = interp_row
else:
img2D_tr = img2D
err2D_tr = err2D
mask2D_tr = mask2D
output_msg = "\n".join(msg)
return img2D_tr, err2D_tr, mask2D_tr, wl, hdr_tr, output_msg
def find_apertures(data,
scan_step=100,
align_deg=2,
degree=4,
mode='normal',
figpath='images'):
"""Find order locations for CFHT/ESPaDOnS data.
Args:
data ():
scan_step (int):
align_deg (int):
degree (int):
mode (normal):
figpath (str):
"""
# prepare for figures
if mode == 'debug':
dbgpath = 'debug'
if not os.path.exists(dbgpath):
os.mkdir(dbgpath)
plot_alignfit = True
plot_orderfit = True
plot_detection = True
figname_alignfit = lambda x1: os.path.join(
dbgpath, 'alignfit_{:04d}.png'.format(x1))
figname_orderfit = lambda iorder: os.path.join(
dbgpath, 'orderfit_{:03d}.png'.format(iorder))
figname_detection = os.path.join(dbgpath, 'order_detection.png')
else:
plot_alignfit = False
plot_orderfit = False
plot_detection = False
plot_orderalign = True
plot_allorders = True
figname_orderalign = os.path.join(figpath, 'order_alignment.png')
figname_allorders = os.path.join(figpath, 'order_all.png')
ny, nx = data.shape
allx = np.arange(nx)
def forward(x, p):
deg = len(p) - 1
res = p[0]
for i in range(deg):
res = res * x + p[i + 1]
return res
def forward_der(x, p):
deg = len(p) - 1
p_der = [(deg - i) * p[i] for i in range(deg)]
return forward(x, p_der)
def backward(y, p):
x = y
for ite in range(20):
dy = forward(x, p) - y
y_der = forward_der(x, p)
dx = dy / y_der
x = x - dx
if (np.abs(dx) < 1e-7).all():
break
return x
def fitfunc(p, interfunc, n):
#return p[-2]*interfunc(forward(np.arange(n), p[0:-2]))+p[-1]
return interfunc(forward(np.arange(n), p[0:-1])) + p[-1]
def resfunc(p, interfunc, flux0, mask=None):
res_lst = flux0 - fitfunc(p, interfunc, flux0.size)
if mask is None:
mask = np.ones_like(flux0, dtype=np.bool)
return res_lst[mask]
x0 = ny // 2
x_lst = {-1: [], 1: []}
param_lst = {-1: [], 1: []}
x1 = x0
direction = -1
icol = 0
all_order_param_lst = {}
all_aligned_x_lst = {}
if plot_orderalign:
fig0 = plt.figure(figsize=(12, 6), dpi=200)
ax0 = fig0.add_axes([0.07, 0.1, 0.4, 0.8])
ax1 = fig0.add_axes([0.53, 0.1, 0.4, 0.8])
while (True):
#flux1 = np.mean(logdata[x1-2:x1+3, :], axis=0)
flux1 = np.mean(data[x1 - 2:x1 + 3, :], axis=0)
# fix negative values in cross-section
negmask = flux1 < 0
if negmask.sum() > 0:
message = 'Negative values in Col {}: {}'.format(x1, allx[negmask])
f = intp.InterpolatedUnivariateSpline(allx[~negmask],
flux1[~negmask],
k=1,
ext='const')
flux1 = f(allx)
logflux1 = np.log(flux1)
if icol == 0:
logflux1_center = logflux1
if plot_orderalign:
ax0.plot(np.arange(nx), (logflux1 - 1) * 100 + x1,
color='C0',
lw=0.6)
ax1.plot(np.arange(nx), (logflux1 - 1) * 100 + x1,
color='C0',
lw=0.6)
all_order_param_lst[x1] = find_order_locations(flux1, x1)
all_aligned_x_lst[x1] = allx
else:
p0 = [0.0 for i in range(align_deg + 1)]
p0[-3] = 1.0
#p0 = [0.0 for i in range(deg+2)]
#p0[-4] = 1.0
interfunc = intp.InterpolatedUnivariateSpline(np.arange(
logflux1.size),
logflux1,
k=3,
ext=3)
mask = np.ones_like(logflux0, dtype=np.bool)
clipping = 5.
maxiter = 10
for i in range(maxiter):
param, _ = opt.leastsq(resfunc,
p0,
args=(interfunc, logflux0, mask))
res_lst = resfunc(param, interfunc, logflux0)
std = res_lst.std()
mask1 = res_lst < clipping * std
mask2 = res_lst > -clipping * std
new_mask = mask1 * mask2
if new_mask.sum() == mask.sum():
break
mask = new_mask
p0 = param
if plot_alignfit:
figalg = plt.figure(dpi=200)
axa1 = figalg.add_subplot(211)
axa2 = figalg.add_subplot(212)
axa1.plot(logflux0, lw=0.5, label='Template')
axa1.plot(logflux1, lw=0.5, label='Flux')
axa1.plot(fitfunc(param, interfunc, logflux0.size),
lw=0.5,
label='Shifted Flux')
axa2.plot(resfunc(param, interfunc, logflux0), lw=0.5)
axa1.set_xlim(0, nx - 1)
axa2.set_xlim(0, nx - 1)
axa1.set_ylim(1, 10)
axa1.legend(loc='lower center', ncol=3)
title = 'Order Alignment for Column {:04d}'.format(x1)
figalg.suptitle(title)
figname = figname_alignfit(x1)
figalg.savefig(figname)
plt.close(figalg)
message = 'savefig: "{}": {}'.format(figname, title)
logger.info(message)
param_lst[direction].append(param[0:-1])
#param_lst[direction].append(param[0:-2])
aligned_allx = allx.copy()
for param in param_lst[direction][::-1]:
aligned_allx = backward(aligned_allx, param)
if plot_orderalign:
ax0.plot(allx, (logflux1 - 1) * 100 + x1,
color='k',
alpha=0.2,
lw=0.6)
ax1.plot(aligned_allx, (logflux1 - 1) * 100 + x1,
color='k',
alpha=0.2,
lw=0.6)
all_order_param_lst[x1] = find_order_locations(flux1,
x1,
aligned_allx,
mode=mode)
all_aligned_x_lst[x1] = aligned_allx
x1 += direction * scan_step
if x1 <= 10:
# turn to the other direction
direction = +1
x1 = x0 + direction * scan_step
x_lst[direction].append(x1)
logflux0 = logflux1_center
icol += 1
continue
elif x1 >= ny - 20:
# scan ends
break
else:
x_lst[direction].append(x1)
logflux0 = logflux1
icol += 1
continue
if plot_orderalign:
title = 'Order Alignment'
fig0.suptitle(title)
fig0.savefig(figname_orderalign)
plt.close(fig0)
message = 'savefig: "{}": {}'.format(figname_orderalign, title)
logger.info(message)
aligned_bound_lst = []
all_aligned_order_param_lst = {}
for x1, order_param_lst in sorted(all_order_param_lst.items()):
aligned_x = all_aligned_x_lst[x1]
aligned_bound_lst.append(
(math.floor(aligned_x[0]), math.ceil(aligned_x[-1])))
f = intp.InterpolatedUnivariateSpline(allx, aligned_x, k=3)
# find aligned order param
aligned_order_param_lst = [(f(i1), f(i2), f(v1), f(v2), f(v3))
for i1, i2, v1, v2, v3 in order_param_lst]
all_aligned_order_param_lst[x1] = aligned_order_param_lst
aligned_peakAB_lst = []
aligned_peakA_lst = []
aligned_peakB_lst = []
for x1, aligned_order_param_lst in sorted(
all_aligned_order_param_lst.items()):
for _, _, newv1, newv2, newv3 in aligned_order_param_lst:
aligned_peakAB_lst.append(newv1)
aligned_peakA_lst.append(newv2)
aligned_peakB_lst.append(newv3)
minx = min(aligned_bound_lst, key=lambda item: item[0])[0]
maxx = max(aligned_bound_lst, key=lambda item: item[1])[1]
bins = np.arange(minx, maxx + 1, 1)
histAB, _ = np.histogram(aligned_peakAB_lst, bins=bins)
histA, _ = np.histogram(aligned_peakA_lst, bins=bins)
histB, _ = np.histogram(aligned_peakB_lst, bins=bins)
binx = bins[0:-1] + np.diff(bins) / 2
# find allsize_lst, which is the number of columns scanned in each
# cross-disp pixels
allsize_lst = np.zeros(maxx - minx)
for (x1, x2) in aligned_bound_lst:
xlst = np.ones(x2 - x1)
# add zeros in the beginning
xlst = np.insert(xlst, 0, [0] * (x1 - minx))
# add zeros in the end
xlst = np.append(xlst, [0] * (maxx - x2))
allsize_lst += xlst
# normalize the histogram
norm_histAB = histAB / allsize_lst
norm_histA = histA / allsize_lst
norm_histB = histB / allsize_lst
if plot_detection:
# fig5 is the order detection figure
fig5 = plt.figure(figsize=(10, 5), dpi=150)
ax51 = fig5.add_subplot(211)
ax52 = fig5.add_subplot(212)
ax51.fill_between(binx, histAB, color='C1', step='mid', alpha=0.6)
ax51.fill_between(binx, histA, color='C0', step='mid', alpha=0.6)
ax51.fill_between(binx, histB, color='C3', step='mid', alpha=0.6)
ax51.step(binx, allsize_lst)
ax52.fill_between(binx, norm_histAB, color='C1', step='mid', alpha=0.6)
ax52.fill_between(binx, norm_histA, color='C0', step='mid', alpha=0.6)
ax52.fill_between(binx, norm_histB, color='C3', step='mid', alpha=0.6)
y1, y2 = ax52.get_ylim()
# get group list
idx = np.where(norm_histAB > 1e-5)[0]
groupAB_lst = np.split(idx, np.where(np.diff(idx) > 2)[0] + 1)
centAB_lst = [
(binx[group] * norm_histAB[group]).sum() / (norm_histAB[group].sum())
for group in groupAB_lst
]
cumnAB_lst = [norm_histAB[group].sum() for group in groupAB_lst]
idx = np.where(norm_histA > 1e-5)[0]
groupA_lst = np.split(idx, np.where(np.diff(idx) > 2)[0] + 1)
centA_lst = [
(binx[group] * norm_histA[group]).sum() / (norm_histA[group].sum())
for group in groupA_lst
]
cumnA_lst = [norm_histA[group].sum() for group in groupA_lst]
idx = np.where(norm_histB > 1e-5)[0]
groupB_lst = np.split(idx, np.where(np.diff(idx) > 2)[0] + 1)
centB_lst = [
(binx[group] * norm_histB[group]).sum() / (norm_histB[group].sum())
for group in groupB_lst
]
cumnB_lst = [norm_histB[group].sum() for group in groupB_lst]
x1_lst = [x0]
for direction in [1, -1]:
for x1 in x_lst[direction]:
x1_lst.append(x1)
order_AB_lst = {}
order_A_lst = {}
order_B_lst = {}
iorder = 0
for group, cent, cumn, groupA, centA, cumnA, groupB, centB, cumnB in zip(
groupAB_lst,
centAB_lst,
cumnAB_lst,
groupA_lst,
centA_lst,
cumnA_lst,
groupB_lst,
centB_lst,
cumnB_lst,
):
if cumn < 0.3:
continue
xlst, yABlst, yAlst, yBlst = [], [], [], []
for x1 in x1_lst:
order_param_lst = all_order_param_lst[x1]
aligned_order_param_lst = all_aligned_order_param_lst[x1]
for (_, _, v1, v2, v3), (_, _, newv1, newv2,
newv3) in zip(order_param_lst,
aligned_order_param_lst):
if binx[group[0]] - 1 < newv1 < binx[group[-1]] + 1:
xlst.append(x1)
yABlst.append(v1)
yAlst.append(v2)
yBlst.append(v3)
break
xlst = np.array(xlst)
yABlst = np.array(yABlst)
yAlst = np.array(yAlst)
yBlst = np.array(yBlst)
# sort again
idx = xlst.argsort()
xlst = xlst[idx]
yABlst = yABlst[idx]
yAlst = yAlst[idx]
yBlst = yBlst[idx]
order_AB_lst[iorder] = (xlst, yABlst)
order_A_lst[iorder] = (xlst, yAlst)
order_B_lst[iorder] = (xlst, yBlst)
iorder += 1
# plot in order detection figure
if plot_detection:
for group in groupAB_lst:
i1, i2 = group[0], group[-1]
ax52.fill_betweenx([y1, y2],
binx[i1],
binx[i2],
color='C3',
alpha=0.1)
ax52.set_ylim(y1, y2)
ax51.set_xlim(minx, maxx)
ax52.set_xlim(minx, maxx)
fig5.savefig(figname_detection)
plt.close(fig5)
message = 'savefig: "{}": Order detections'.format(figname_detection)
logger.info(message)
aperture_set = ApertureSet(shape=(ny, nx))
aperture_set_A = ApertureSet(shape=(ny, nx))
aperture_set_B = ApertureSet(shape=(ny, nx))
# plot all order position in a single figure
if plot_allorders:
figall = plt.figure(figsize=(14, 7), dpi=200)
axall = figall.add_axes([0.1, 0.05, 0.8, 0.85])
######### fit order positions and pack them to ApertureSet ##########
for iorder in sorted(order_AB_lst.keys()):
xlst_AB, ylst_AB = order_AB_lst[iorder]
xlst_A, ylst_A = order_A_lst[iorder]
xlst_B, ylst_B = order_B_lst[iorder]
##################### Parse Center of Fiber A & B ################
fitmask = np.ones_like(xlst_AB, dtype=np.bool)
maxiter = 10
for nite in range(maxiter):
poly = Chebyshev.fit(xlst_AB[fitmask],
ylst_AB[fitmask],
deg=degree)
yres = ylst_AB - poly(xlst_AB)
std = yres[fitmask].std()
new_fitmask = (yres > -3 * std) * (yres < 3 * std)
if new_fitmask.sum() == fitmask.sum():
break
fitmask = new_fitmask
aperture_loc = ApertureLocation(direct='y', shape=(ny, nx))
aperture_loc.set_position(poly)
aperture_set[iorder] = aperture_loc
color = 'C1'
label = 'Fiber A+B'
if plot_allorders:
axall.scatter(xlst_AB,
ylst_AB,
s=15,
color='none',
edgecolor=color)
axall.scatter(xlst_AB[fitmask],
ylst_AB[fitmask],
s=15,
color=color)
# prepare newx, newy, and m (mask) for plotting a smooth line
newx = np.arange(0, ny)
newy = poly(newx)
m = (newy >= 0) * (newy < nx)
axall.plot(newx[m], newy[m], '-', color='C0', lw=0.7)
if plot_orderfit:
# plot position fitting of each order
# initialize fig
figm = plt.figure(dpi=150)
axm1 = figm.add_axes([0.1, 0.4, 0.8, 0.50])
axm2 = figm.add_axes([0.1, 0.1, 0.8, 0.25])
axm1.scatter(xlst_AB, ylst_AB, s=10, color='none', edgecolor=color)
axm1.scatter(xlst_AB[fitmask], ylst_AB[fitmask], s=10, color=color)
axm1.plot(newx[m], newy[m], '-', color=color, lw=0.7, label=label)
# plot fitting residuals
yres_AB = ylst_AB - poly(xlst_AB)
axm2.scatter(xlst_AB, yres_AB, s=10, color='none', edgecolor=color)
axm2.scatter(xlst_AB[fitmask], yres_AB[fitmask], s=10, color=color)
######################## Parse Fiber A ########################
fitmask = np.ones_like(xlst_A, dtype=np.bool)
maxiter = 10
for nite in range(maxiter):
poly = Chebyshev.fit(xlst_A[fitmask], ylst_A[fitmask], deg=degree)
yres = ylst_A - poly(xlst_A)
std = yres[fitmask].std()
new_fitmask = (yres > -3 * std) * (yres < 3 * std)
if new_fitmask.sum() == fitmask.sum():
break
fitmask = new_fitmask
aperture_loc = ApertureLocation(direct='y', shape=(ny, nx))
aperture_loc.set_position(poly)
aperture_set_A[iorder] = aperture_loc
color = 'C0'
label = 'Fiber A'
if plot_allorders:
axall.scatter(xlst_A, ylst_A, s=15, color='none', edgecolor=color)
axall.scatter(xlst_A[fitmask], ylst_A[fitmask], s=15, color=color)
# prepare newx, newy, and m (mask) for plotting a smooth line
newx = np.arange(0, ny)
newy = poly(newx)
m = (newy >= 0) * (newy < nx)
axall.plot(newx[m], newy[m], '-', color='C0', lw=0.7)
if plot_orderfit:
# plot position fitting of each order
axm1.scatter(xlst_A, ylst_A, s=10, color='none', edgecolor=color)
axm1.scatter(xlst_A[fitmask], ylst_A[fitmask], s=10, color=color)
axm1.plot(newx[m], newy[m], '-', color=color, lw=0.7, label=label)
# plot fitting residuals
yres_A = ylst_A - poly(xlst_A)
axm2.scatter(xlst_A, yres_A, s=10, color='none', edgecolor=color)
axm2.scatter(xlst_A[fitmask], yres_A[fitmask], s=10, color=color)
########################### Parse Fiber B #######################
fitmask = np.ones_like(xlst_B, dtype=np.bool)
maxiter = 10
for nite in range(maxiter):
poly = Chebyshev.fit(xlst_B[fitmask], ylst_B[fitmask], deg=degree)
yres = ylst_B - poly(xlst_B)
std = yres[fitmask].std()
new_fitmask = (yres > -3 * std) * (yres < 3 * std)
if new_fitmask.sum() == fitmask.sum():
break
fitmask = new_fitmask
aperture_loc = ApertureLocation(direct='y', shape=(ny, nx))
aperture_loc.set_position(poly)
aperture_set_B[iorder] = aperture_loc
color = 'C3'
label = 'Fiber B'
if plot_allorders:
axall.scatter(xlst_B, ylst_B, s=15, color='none', edgecolor=color)
axall.scatter(xlst_B[fitmask], ylst_B[fitmask], s=15, color=color)
# prepare newx, newy, and m (mask) for plotting a smooth line
newx = np.arange(0, ny)
newy = poly(newx)
m = (newy >= 0) * (newy < nx)
axall.plot(newx[m], newy[m], '-', color='C0', lw=0.7)
if plot_orderfit:
# plot position fitting of each order
axm1.scatter(xlst_B, ylst_B, s=10, color='none', edgecolor=color)
axm1.scatter(xlst_B[fitmask], ylst_B[fitmask], s=10, color=color)
axm1.plot(newx[m], newy[m], '-', color=color, lw=0.7, label=label)
# plot fitting residuals
yres_B = ylst_B - poly(xlst_B)
axm2.scatter(xlst_B, yres_B, s=10, color='none', edgecolor=color)
axm2.scatter(xlst_B[fitmask], yres_B[fitmask], s=10, color=color)
# decorate and save the order fit figure
axm1.set_xlim(0, ny - 1)
axm2.set_xlim(0, ny - 1)
legend = axm1.legend(loc='upper center', ncol=3)
title = 'Position fitting of Order {:03d}'.format(iorder)
figm.suptitle(title)
figname = figname_orderfit(iorder)
figm.savefig(figname)
plt.close(figm)
message = 'savefig: "{}": {}'.format(figname, title)
logger.info(message)
####################################################################
# decoration of all order figure
if plot_allorders:
axall.grid(True, ls='--', lw=0.5)
axall.set_axisbelow(True)
axall.set_xlim(0, ny - 1)
axall.set_ylim(0, nx - 1)
axall.set_aspect(1)
figall.savefig(figname_allorders)
plt.close(figall)
message = 'savefig: "{}": All order positions'.format(
figname_allorders)
logger.info(message)
return aperture_set, aperture_set_A, aperture_set_B
def create_2d_profile(img2D,
model_name='moffat',
dx=25,
width_scale=2,
xmin=None,
xmax=None,
ymin=None,
ymax=None,
order_center=3,
order_width=0,
w_cen=15,
kappa_cen=3.,
w_width=21,
kappa_width=3.):
"""
img2D : np.array(M, N)
Input image with dispersion along x-axis!
model_name : {'moffat' or 'gaussian' or 'tophat'}
Model type for the spectral PSF
dx : int [default=5]
Fit the trace for every dx column
width_scale : int [default=2]
The scaling factor of the FWHM used for the width of the tophat profile:
By default the flux is summed within 2*FWHM on either side of the centroid
xmin, xmax : int [default=None]
Minimum and maximum extent to fit along the dispersion axis
ymin, ymax : int [default=None]
Minimum and maximum extent to fit along the spatial axis
order_center : int [default=3]
Order of Chebyshev polynomium for the trace position
order_width : int [default=0]
Order of Chebyshev polynomium for the trace width
w_cen : int [default=15]
Kernel width of median filter for trace position
kappa_cen : float [default=3.0]
Threshold for median filtering. Reject outliers above: ±`kappa` * sigma,
where sigma is the robust standard deviation of the data points.
w_width : int [default=15]
Kernel width of median filter for trace width parameters
kappa_width : float [default=3.0]
Threshold for median filtering. Reject outliers above: ±`kappa` * sigma,
where sigma is the robust standard deviation of the data points.
Returns
-------
trace_models_2d : list(np.array(M, N))
List of trace models, one for each object identified in the image
trace_info : list
List of information dictionary for each trace:
The fitted position and width as well as the fit mask and the fitted values
"""
msg = list()
img2D = img2D.astype(np.float64)
x = np.arange(img2D.shape[1], dtype=np.float64)
y = np.arange(img2D.shape[0], dtype=np.float64)
if not xmin:
xmin = 0
if not xmax:
xmax = len(x)
if xmax < 0:
xmax = len(x) + xmax
if not ymin:
ymin = 0
if not ymax:
ymax = len(y)
if ymax < 0:
ymax = len(y) + ymax
if model_name == 'tophat':
# Fit the centroid using a Moffat profile, but the discard the with for the profile calculation
fit_values = fit_trace(img2D,
x,
y,
model_name='moffat',
dx=dx,
ymin=ymin,
ymax=ymax,
xmin=xmin,
xmax=xmax)
fwhm = fit_values[3]
if fwhm is None:
raise ValueError(
"FWHM of the spectral trace could not be determined! Maybe more than one object in slit..."
)
else:
fit_values = fit_trace(img2D,
x,
y,
model_name=model_name,
dx=dx,
ymin=ymin,
ymax=ymax,
xmin=xmin,
xmax=xmax)
x_binned, N_obj, trace_parameters, fwhm, fit_msg = fit_values
msg.append(fit_msg)
msg.append(
" - Creating 2D spectral profile from fitted parameters")
msg.append(" - Profile type: %s" % model_name)
msg.append(
" - Interpolating centroid using Chebyshev polynomium of degree: %i"
% order_center)
msg.append(
" - Interpolating profile width using Chebyshev polynomium of degree: %i"
% order_width)
trace_models_2d = list()
trace_info = list()
domain = [0, img2D.shape[1]]
for n in range(N_obj):
msg.append(" - Working on profile number %i" % (n + 1))
info_dict = dict()
info_dict['x_binned'] = x_binned
# Median filter
mu = np.array([p['mu_%i' % n] for p in trace_parameters])
mu_err = np.array([p['mu_%i' % n].stderr for p in trace_parameters])
mu_err[mu_err == 0] = 100.
w_mu = 1. / mu_err**2
mu_med, mask_mu = median_filter_data(mu, kappa_cen, w_cen)
mask_mu &= (x_binned > xmin) & (x_binned < xmax)
mu_fit = Chebyshev.fit(x_binned[mask_mu],
mu[mask_mu],
deg=order_center,
domain=domain,
w=w_mu[mask_mu])
info_dict['mu'] = mu
info_dict['mu_err'] = mu_err
info_dict['mask_mu'] = mask_mu
info_dict['fit_mu'] = mu_fit(x)
# Fit polynomium:
trace2D = np.zeros_like(img2D)
trace2D = np.zeros_like(img2D)
if model_name == 'gaussian':
# Median filter
sig = np.array([p['sig_%i' % n] for p in trace_parameters])
sig_err = np.array(
[p['sig_%i' % n].stderr for p in trace_parameters])
sig_err[sig_err == 0] = 100.
w_sig = 1. / sig_err**2
sig_med, mask_sig = median_filter_data(sig, kappa_width, w_width)
mask_sig &= (x_binned > xmin) & (x_binned < xmax)
sig_fit = Chebyshev.fit(x_binned[mask_sig],
sig[mask_sig],
deg=order_width,
domain=domain,
w=w_sig[mask_sig])
info_dict['sig'] = sig
info_dict['sig_err'] = sig_err
info_dict['mask_sig'] = mask_sig
info_dict['fit_sig'] = sig_fit(x)
for num, x_i in enumerate(x):
P_i = NN_gaussian(y, mu_fit(x_i), sig_fit(x_i), 0.)
P_i = P_i / np.sum(P_i)
trace2D[:, num] = P_i
trace_models_2d.append(trace2D)
elif model_name == 'moffat':
# Median filter
a = np.array([p['a_%i' % n] for p in trace_parameters])
a_med, mask_a = median_filter_data(a, kappa_width, w_width)
a_err = np.array([p['a_%i' % n].stderr for p in trace_parameters])
a_err[a_err == 0] = 100.
w_a = 1. / a_err**2
b = np.array([p['b_%i' % n] for p in trace_parameters])
b_med, mask_b = median_filter_data(b, kappa_width, w_width)
b_err = np.array([p['b_%i' % n].stderr for p in trace_parameters])
b_err[b_err == 0] = 100.
w_b = 1. / b_err**2
mask_a &= (x_binned > xmin) & (x_binned < xmax)
mask_b &= (x_binned > xmin) & (x_binned < xmax)
a_fit = Chebyshev.fit(x_binned[mask_a],
a[mask_a],
deg=order_width,
domain=domain,
w=w_a[mask_a])
b_fit = Chebyshev.fit(x_binned[mask_b],
b[mask_b],
deg=order_width,
domain=domain,
w=w_b[mask_b])
info_dict['a'] = a
info_dict['a_err'] = a_err
info_dict['mask_a'] = mask_a
info_dict['fit_a'] = a_fit(x)
info_dict['b'] = b
info_dict['b_err'] = b_err
info_dict['mask_b'] = mask_b
info_dict['fit_b'] = b_fit(x)
for num, x_i in enumerate(x):
P_i = NN_moffat(y, mu_fit(x_i), a_fit(x_i), b_fit(x_i), 0.)
P_i = P_i / np.sum(P_i)
trace2D[:, num] = P_i
trace_models_2d.append(trace2D)
elif model_name == 'tophat':
for num, x_i in enumerate(x):
center = mu_fit(x_i)
lower = int(center - width_scale * fwhm)
upper = int(center + width_scale * fwhm)
trace2D[lower:upper + 1, num] = 1 / (upper - lower + 1)
trace_models_2d.append(trace2D)
info_dict['fwhm'] = fwhm
trace_info.append(info_dict)
output_msg = "\n".join(msg)
return (trace_models_2d, trace_info, output_msg)
"""
profile_switch = 1 # 0: Couette Flow, 1 : Poiseuille Flow
alpha = 1.0 # streamwise wavenumber
beta = 0.0 # spanwise wavenumber
Re = 10000.0 # Reynold's number
N = 150 # Order of Chebyshev polynomials
""""""
""" Fit velocity profile """
if profile_switch == 0:
u = np.array([[-1., 1.], [-1., 1.]])
elif profile_switch == 1:
u = np.array([[-1., 0., 1.], [0., 1., 0.]])
U = Ch.fit(u[0, :], u[1, :], u.shape[1] - 1)
""""""
""" Define Orr-Sommerfeld operators """
k2 = alpha**2 + beta**2
jaRe = 1.0j * alpha * Re
def orr_somm_lhs(T):
return ((-U * k2 - U.deriv(2) - k2**2 / jaRe) * T +
(U + 2.0 * k2 / jaRe) * T.deriv(2) + (-1.0 / jaRe) * T.deriv(4))
def orr_somm_rhs(T):
return (T.deriv(2) - k2 * T)
from numpy.polynomial import Chebyshev as T x = 10*np.random.rand(100) y = 0.1*np.random.randn(100) + sin(x) plot(x,y,'.') fit = T.fit(x,y,10) xx, yy = fit.linspace(100) plot(xx,yy) fit_x = fit.deriv(1) fit_xx = fit.deriv(2) plot(*fit_x.linspace(100)) plot(*fit_xx.linspace(100)) fit = T.fit(x,y,50) xx, yy = fit.linspace(100)
import numpy as np import matplotlib.pyplot as plt from numpy.polynomial import Chebyshev as T np.random.seed(11) x = np.linspace(0, 2*np.pi, 20) y = np.sin(x) + np.random.normal(scale=.1, size=x.shape) p = T.fit(x, y, 5) plt.plot(x, y, 'o') # [<matplotlib.lines.Line2D object at 0x2136c10>] xx, yy = p.linspace() plt.plot(xx, yy, lw=2) # [<matplotlib.lines.Line2D object at 0x1cf2890>] p.domain # array([ 0. , 6.28318531]) p.window # array([-1., 1.]) plt.show()
import numpy as np import matplotlib.pyplot as plt from numpy.polynomial import Chebyshev as T np.random.seed() x = np.linspace(0, 2 * np.pi, 20) y = np.sin(x) + np.random.normal(scale=.2, size=x.shape) p = T.fit(x, y, 15) plt.plot(x, y, 'o') xx, yy = p.linspace() plt.plot(xx, yy, lw=2) plt.show()
