def test_basis_func(self):
p = Chebyshev([1, 2, 3])
assert_equal(self.as_latex(p),
r'$x \mapsto 1.0\,{T}_{0}(x) + 2.0\,{T}_{1}(x) + 3.0\,{T}_{2}(x)$')
# affine input - check no surplus parens are added
p = Chebyshev([1, 2, 3], domain=[-1, 0])
assert_equal(self.as_latex(p),
r'$x \mapsto 1.0\,{T}_{0}(1.0 + 2.0x) + 2.0\,{T}_{1}(1.0 + 2.0x) + 3.0\,{T}_{2}(1.0 + 2.0x)$')
def dlfParamsdz(z, pPhiStar, pMStar, pAlpha, pBeta):
dlog10phiStardz = T.deriv(T(pPhiStar))(1 + z)
dmStardz = T.deriv(T(pMStar))(1 + z)
dalphadz = T.deriv(T(pAlpha))(1 + z)
h, f0, z0, a, b = pBeta
zeta = np.log10((1.0 + z) / (1.0 + z0))
dbetadz = (-f0 * (a * 10.0**((a - 1) * zeta) /
(1.0 + z0) + b * 10.0**((b - 1) * zeta) / (1.0 + z0)) /
(10.0**(a * zeta) + 10.0**(b * zeta))**2)
return dlog10phiStardz, dmStardz, dalphadz, dbetadz
def get_Matrix(operator):
M = np.zeros((N + 1, N + 1), dtype=complex)
I = np.eye(N + 1)
y = np.cos(np.linspace(2, N - 1, N - 2, dtype=float) * pi / N)
for i in range(0, N + 1):
T = Ch(I[i, :])
M[i, 0] = T(-1.0)
M[i, 1] = T.deriv(1)(-1.0) * Dy(-1.0, 1)
M[i, 2:N] = operator(T, y)
M[i, N] = T(1.0)
#M[i,N] = T.deriv(1)(1.0)*Dy(1.0,1)
M = eliminate(M, 0)
M = eliminate(M, 1)
#M = eliminate(M,N-1)
M = eliminate(M, N)
return M[2:N, 2:N]
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 test_floordiv(Poly):
c1 = list(random((4, )) + .5)
c2 = list(random((3, )) + .5)
c3 = list(random((2, )) + .5)
p1 = Poly(c1)
p2 = Poly(c2)
p3 = Poly(c3)
p4 = p1 * p2 + p3
c4 = list(p4.coef)
assert_poly_almost_equal(p4 // p2, p1)
assert_poly_almost_equal(p4 // c2, p1)
assert_poly_almost_equal(c4 // p2, p1)
assert_poly_almost_equal(p4 // tuple(c2), p1)
assert_poly_almost_equal(tuple(c4) // p2, p1)
assert_poly_almost_equal(p4 // np.array(c2), p1)
assert_poly_almost_equal(np.array(c4) // p2, p1)
assert_poly_almost_equal(2 // p2, Poly([0]))
assert_poly_almost_equal(p2 // 2, 0.5 * p2)
assert_raises(TypeError, op.floordiv, p1, Poly([0],
domain=Poly.domain + 1))
assert_raises(TypeError, op.floordiv, p1, Poly([0],
window=Poly.window + 1))
if Poly is Polynomial:
assert_raises(TypeError, op.floordiv, p1, Chebyshev([0]))
else:
assert_raises(TypeError, op.floordiv, p1, Polynomial([0]))
def test_mod(Poly):
# This checks commutation, not numerical correctness
c1 = list(random((4, )) + .5)
c2 = list(random((3, )) + .5)
c3 = list(random((2, )) + .5)
p1 = Poly(c1)
p2 = Poly(c2)
p3 = Poly(c3)
p4 = p1 * p2 + p3
c4 = list(p4.coef)
assert_poly_almost_equal(p4 % p2, p3)
assert_poly_almost_equal(p4 % c2, p3)
assert_poly_almost_equal(c4 % p2, p3)
assert_poly_almost_equal(p4 % tuple(c2), p3)
assert_poly_almost_equal(tuple(c4) % p2, p3)
assert_poly_almost_equal(p4 % np.array(c2), p3)
assert_poly_almost_equal(np.array(c4) % p2, p3)
assert_poly_almost_equal(2 % p2, Poly([2]))
assert_poly_almost_equal(p2 % 2, Poly([0]))
assert_raises(TypeError, op.mod, p1, Poly([0], domain=Poly.domain + 1))
assert_raises(TypeError, op.mod, p1, Poly([0], window=Poly.window + 1))
if Poly is Polynomial:
assert_raises(TypeError, op.mod, p1, Chebyshev([0]))
else:
assert_raises(TypeError, op.mod, p1, Polynomial([0]))
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 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 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 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 test_split_rect_1():
# two functions from https://stackoverflow.com/a/9997374
def ccw(A,B,C):
return (C[1]-A[1]) * (B[0]-A[0]) > (B[1]-A[1]) * (C[0]-A[0])
# Return true if line segments AB and CD intersect
def intersect(A,B,C,D):
return ccw(A,C,D) != ccw(B,C,D) and ccw(A,B,C) != ccw(A,B,D)
mesh = om.read_trimesh('../data/rect.obj')
mesh_pts = mesh.points()
mesh_edges = mesh.ev_indices()
x = np.linspace(-1, 1, 101)
y = T.basis(5)(x)
# export Chebyshev polynomial
pts = np.zeros((x.size, 3))
pts[:, 0] = x
pts[:, 1] = y
mu.write_obj_lines('../data/curve_Chebyshev.obj', pts, [np.arange(pts.shape[0])])
edges = [np.array([0,0], 'i')]
ratios = [0.0]
for i in range(1, x.size-2):
A = np.array([x[i], y[i]])
B = np.array([x[i+1], y[i+1]])
temp_x = []
temp_e = []
temp_u = []
for e in mesh_edges:
C = mesh_pts[e[0], :2]
D = mesh_pts[e[1], :2]
if intersect(A, B, C, D):
# A = (x1, y1), B = (x2, y2)
# C = (x3, y3), D = (x4, y4)
u = ((A[0]-C[0])*(A[1]-B[1])-(A[1]-C[1])*(A[0]-B[0]))/((A[0]-B[0])*(C[1]-D[1])-(A[1]-B[1])*(C[0]-D[0]))
temp_e.append(e)
temp_u.append(u)
temp_x.append(C[0]*(1-u)+D[0]*u)
order = np.argsort(np.array(temp_x))
edges.extend([temp_e[i] for i in order])
ratios.extend([temp_u[i] for i in order])
edges.append(np.array([2,2], 'i'))
ratios.append(0.0)
edges = np.array(edges, 'i')
#print(ratios)
ratios = np.array(ratios, 'd')
pts0 = mesh_pts[edges[:, 0], :]
pts1 = mesh_pts[edges[:, 1], :]
pts = np.multiply(pts0, 1.-ratios[:, np.newaxis]) + \
np.multiply(pts1, ratios[:, np.newaxis])
mu.write_obj_lines('../data/curve.obj', pts, [np.arange(pts.shape[0])])
# print(on_edges, ratios)
mesh, curve_idx = mu.split_mesh_complete(mesh.points(),
mesh.face_vertex_indices(),
edges, ratios)
# print(curve_idx)
om.write_mesh('../data/mesh_split.obj', mesh)
def test_approximation(self):
def powx(x, p):
return x**p
x = np.linspace(0, 2, 10)
for deg in range(0, 10):
for t in range(0, deg + 1):
p = Chebyshev.interpolate(powx, deg, domain=[0, 2], args=(t, ))
assert_almost_equal(p(x), powx(x, t), decimal=12)
def chebyshev_approximation(fun, degree, R=4.0):
bases = chebyshev_bases(degree)
approx_coeffs = []
for i in range(len(bases)):
approx_coeffs.append(
quad(lambda x: fun(R * x) * bases[i](x) / np.sqrt(1 - x**2), -1,
1)[0])
std_coefs = Chebyshev(approx_coeffs).convert(kind=Polynomial).coef
return np.power(1 / R, range(degree + 1)) * std_coefs
def _create_Tderiv(self, k, n):
r"""Create the n'th derivative of the k'th Chebyshev polynomial."""
if not self._use_mp:
return ChebyshevT.basis(k).deriv(n)
else:
if n == 0:
return lambda x: mp.chebyt(k, x)
# Since the Chebyshev derivatives attain high values near the
# borders +-1 and usually are subtracted to obtain small values,
# minimizing their relative error (i.e. fixing the number of
# correct decimal places (dps)) is not enough to get precise
# results. Hence, the below `moredps` variable is set to higher
# values close to the border.
extradps = 5
pos = mp.fprod(
[mp.mpf(k**2 - p**2) / (2 * p + 1) for p in range(n)])
neg = (-1)**(k + n) * pos
if n == 1:
def dTk(x):
with mp.extradps(extradps):
x = clip(x, -1.0, 1.0)
if mp.almosteq(x, mp.one): return pos
if mp.almosteq(x, -mp.one): return neg
moredps = max(
0,
int(-math.log10(
min(abs(x - mp.one), abs(x + mp.one))) / 2))
moredps = min(moredps, 100)
with mp.extradps(moredps):
t = mp.acos(x)
return k * mp.sin(k * t) / mp.sin(t)
return dTk
if n == 2:
def ddTk(x):
with mp.extradps(extradps):
x = clip(x, -1.0, 1.0)
if mp.almosteq(x, mp.one): return pos
if mp.almosteq(x, -mp.one): return neg
moredps = max(
0,
int(-math.log10(
min(abs(x - mp.one), abs(x + mp.one))) * 1.5) +
2)
moredps = min(moredps, 100)
with mp.extradps(moredps):
t = mp.acos(x)
s = mp.sin(t)
return -k**2 * mp.cos(k * t) / s**2 + k * mp.cos(
t) * mp.sin(k * t) / s**3
return ddTk
raise NotImplementedError(
'Derivatives of order > 2 not implemented.')
def test_approximation(self):
def powx(x, p):
return x**p
x = np.linspace(0, 2, 10)
for deg in range(0, 10):
for t in range(0, deg + 1):
p = Chebyshev.interpolate(powx, deg, domain=[0, 2], args=(t,))
assert_almost_equal(p(x), powx(x, t), decimal=12)
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 chebyshev_approximation(fun, degree=4, is_zero=None):
bases = chebyshev_bases(degree)
approx_coeffs = []
for i in range(len(bases)):
if is_zero is None or not is_zero(i):
approx_coeffs.append(
quad(lambda x: fun(x) * bases[i](x) / np.sqrt(1 - x**2),
-1, 1)[0])
else:
approx_coeffs.append(0)
return Chebyshev(approx_coeffs).convert(kind=Polynomial).coef
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 test_divmod(Poly):
# This checks commutation, not numerical correctness
c1 = list(random((4, )) + .5)
c2 = list(random((3, )) + .5)
c3 = list(random((2, )) + .5)
p1 = Poly(c1)
p2 = Poly(c2)
p3 = Poly(c3)
p4 = p1 * p2 + p3
c4 = list(p4.coef)
quo, rem = divmod(p4, p2)
assert_poly_almost_equal(quo, p1)
assert_poly_almost_equal(rem, p3)
quo, rem = divmod(p4, c2)
assert_poly_almost_equal(quo, p1)
assert_poly_almost_equal(rem, p3)
quo, rem = divmod(c4, p2)
assert_poly_almost_equal(quo, p1)
assert_poly_almost_equal(rem, p3)
quo, rem = divmod(p4, tuple(c2))
assert_poly_almost_equal(quo, p1)
assert_poly_almost_equal(rem, p3)
quo, rem = divmod(tuple(c4), p2)
assert_poly_almost_equal(quo, p1)
assert_poly_almost_equal(rem, p3)
quo, rem = divmod(p4, np.array(c2))
assert_poly_almost_equal(quo, p1)
assert_poly_almost_equal(rem, p3)
quo, rem = divmod(np.array(c4), p2)
assert_poly_almost_equal(quo, p1)
assert_poly_almost_equal(rem, p3)
quo, rem = divmod(p2, 2)
assert_poly_almost_equal(quo, 0.5 * p2)
assert_poly_almost_equal(rem, Poly([0]))
quo, rem = divmod(2, p2)
assert_poly_almost_equal(quo, Poly([0]))
assert_poly_almost_equal(rem, Poly([2]))
assert_raises(TypeError, divmod, p1, Poly([0], domain=Poly.domain + 1))
assert_raises(TypeError, divmod, p1, Poly([0], window=Poly.window + 1))
if Poly is Polynomial:
assert_raises(TypeError, divmod, p1, Chebyshev([0]))
else:
assert_raises(TypeError, divmod, p1, Polynomial([0]))
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
def test_sub(Poly):
# This checks commutation, not numerical correctness
c1 = list(random((4, )) + .5)
c2 = list(random((3, )) + .5)
p1 = Poly(c1)
p2 = Poly(c2)
p3 = p1 - p2
assert_poly_almost_equal(p2 - p1, -p3)
assert_poly_almost_equal(p1 - c2, p3)
assert_poly_almost_equal(c2 - p1, -p3)
assert_poly_almost_equal(p1 - tuple(c2), p3)
assert_poly_almost_equal(tuple(c2) - p1, -p3)
assert_poly_almost_equal(p1 - np.array(c2), p3)
assert_poly_almost_equal(np.array(c2) - p1, -p3)
assert_raises(TypeError, op.sub, p1, Poly([0], domain=Poly.domain + 1))
assert_raises(TypeError, op.sub, p1, Poly([0], window=Poly.window + 1))
if Poly is Polynomial:
assert_raises(TypeError, op.sub, p1, Chebyshev([0]))
else:
assert_raises(TypeError, op.sub, p1, Polynomial([0]))
def check_add(Poly):
# This checks commutation, not numerical correctness
c1 = list(random((4,)) + .5)
c2 = list(random((3,)) + .5)
p1 = Poly(c1)
p2 = Poly(c2)
p3 = p1 + p2
assert_poly_almost_equal(p2 + p1, p3)
assert_poly_almost_equal(p1 + c2, p3)
assert_poly_almost_equal(c2 + p1, p3)
assert_poly_almost_equal(p1 + tuple(c2), p3)
assert_poly_almost_equal(tuple(c2) + p1, p3)
assert_poly_almost_equal(p1 + np.array(c2), p3)
assert_poly_almost_equal(np.array(c2) + p1, p3)
assert_raises(TypeError, p1.__add__, Poly([0], domain=Poly.domain + 1))
assert_raises(TypeError, p1.__add__, Poly([0], window=Poly.window + 1))
if Poly is Polynomial:
assert_raises(TypeError, p1.__add__, Chebyshev([0]))
else:
assert_raises(TypeError, p1.__add__, Polynomial([0]))
def test_mul(Poly):
c1 = list(random((4, )) + .5)
c2 = list(random((3, )) + .5)
p1 = Poly(c1)
p2 = Poly(c2)
p3 = p1 * p2
assert_poly_almost_equal(p2 * p1, p3)
assert_poly_almost_equal(p1 * c2, p3)
assert_poly_almost_equal(c2 * p1, p3)
assert_poly_almost_equal(p1 * tuple(c2), p3)
assert_poly_almost_equal(tuple(c2) * p1, p3)
assert_poly_almost_equal(p1 * np.array(c2), p3)
assert_poly_almost_equal(np.array(c2) * p1, p3)
assert_poly_almost_equal(p1 * 2, p1 * Poly([2]))
assert_poly_almost_equal(2 * p1, p1 * Poly([2]))
assert_raises(TypeError, op.mul, p1, Poly([0], domain=Poly.domain + 1))
assert_raises(TypeError, op.mul, p1, Poly([0], window=Poly.window + 1))
if Poly is Polynomial:
assert_raises(TypeError, op.mul, p1, Chebyshev([0]))
else:
assert_raises(TypeError, op.mul, p1, Polynomial([0]))
def test_dimensions(self):
for deg in range(1, 5):
assert_(Chebyshev.interpolate(self.f, deg).degree() == deg)
import matplotlib.pyplot as plt
from numpy.polynomial import Chebyshev as T
x = np.linspace(-2, 2, 100)
for i in range(6):
ax = plt.plot(x, T.basis(i)(x), lw=2, label="T_%d" % i)
# ...
plt.legend(loc="lower right")
# <matplotlib.legend.Legend object at 0x3b3ee10>
plt.show()
def T(n_):
return Chebyshev(np.append(np.zeros(n_), 1))
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 matplotlib.pyplot as plt from numpy.polynomial import Chebyshev as T x = np.linspace(-1, 1, 100) for i in range(6): ax = plt.plot(x, T.basis(i)(x), lw=2, label="T_%d"%i) # ... plt.legend(loc="upper left") # <matplotlib.legend.Legend object at 0x3b3ee10> plt.show()
def test_dimensions(self):
for deg in range(1, 5):
assert_(Chebyshev.interpolate(self.f, deg).degree() == deg)
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)
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
newdata =p([4,6])
print(newdata)
print(newdata.integ())
print(newdata.integ(1))
print(newdata.integ(lbnd=-1))
print(newdata.integ(lbnd=-1, k=1))
print("\n")
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()
#figure_width = 0.35*text_width
#figure_height = 0.75*figure_width #(figure_width / golden_ratio) #figure_width
figure_width = 0.75*text_width
figure_height = (figure_width / golden_ratio) #figure_width
figure_size = [figure_width, figure_height]
config.load_config_medium()
#fig = plt.figure(figsize=figure_size, dpi=100)
fig, axes = plt.subplots(nrows=1, ncols=1, sharex=True, sharey=True, squeeze=False, figsize=figure_size, dpi=100)
ax = axes[0][0]
x = np.linspace(-1, 1, 100)
for i in range(6):
ax.plot(x, T.basis(i)(x), lw=1.0, color=colors[i], label="$t_%d(x)$"%i)
#ax.plot(x, pow(x, i), lw=1.0, color=colors[i], label="$x^%d$"%i)
ax.set_xlabel(r"$x$")
ax.set_ylabel(r"$t_n(x)$", rotation=90)
ax.set_xlim(-1.1, 1.1)
ax.set_ylim(-1.1, 1.1)
ax.xaxis.labelpad = 5
ax.yaxis.labelpad = 2
# Shrink current axis's height by 10% on the bottom
#box = ax.get_position()
#ax.set_position([box.x0, box.y0 - 0.1*box.height, box.width, 0.9*box.height]) #left, bottom, width, height
plt.subplots_adjust(top=0.85, left=0.12, bottom=0.12, right=0.98) # Legend on top
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 chebyshev_approximation_alt(func, deg, a=-4, b=4):
return Chebyshev.interpolate(func, deg,
domain=[a, b]).convert(kind=Polynomial).coef
def chebyshev_basis(k):
for i in range(len(CHEBYSHEV_BASIS), k + 1):
coeffs = np.zeros(i + 1)
coeffs[-1] = (1. + np.sign(i)) / np.pi
CHEBYSHEV_BASIS.append(Chebyshev(coeffs))
return CHEBYSHEV_BASIS[k]
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
