def test_2d_orthopolynomial_in_compound_model(): """ Ensure that OrthoPolynomialBase (ie. Chebyshev2D & Legendre2D) models get evaluated & fitted correctly when part of a compound model. Regression test for https://github.com/astropy/astropy/pull/6085. """ y, x = np.mgrid[0:5, 0:5] z = x + y fitter = fitting.LevMarLSQFitter() simple_model = Chebyshev2D(2, 2) with pytest.warns(AstropyUserWarning, match='Model is linear in parameters'): simple_fit = fitter(simple_model, x, y, z) fitter = fitting.LevMarLSQFitter() # re-init to compare like with like compound_model = Identity(2) | Chebyshev2D(2, 2) compound_model.fittable = True compound_model.linear = True with pytest.warns(AstropyUserWarning, match='Model is linear in parameters'): compound_fit = fitter(compound_model, x, y, z) assert_allclose(simple_fit(x, y), compound_fit(x, y), atol=1e-15)
def fit_2dspec(xl, yl, zl, x_degree=4, y_degree=3, x_domain=None, y_domain=None): from astropy.modeling import fitting # Fit the data using astropy.modeling if x_domain is None: x_domain = [min(xl), max(xl)] # more room for y_domain?? if y_domain is None: #y_domain = [orders[0]-2, orders[-1]+2] y_domain = [min(yl), max(yl)] from astropy.modeling.polynomial import Chebyshev2D p_init = Chebyshev2D(x_degree=x_degree, y_degree=y_degree, x_domain=x_domain, y_domain=y_domain) f = fitting.LinearLSQFitter() p = f(p_init, xl, yl, zl) for i in [0]: dd = p(xl, yl) - zl m = np.abs(dd) < 3. * dd.std() p = f(p, xl[m], yl[m], zl[m]) return p, m
def test__fcache(): model = OrthoPolynomialBase(x_degree=2, y_degree=2) with pytest.raises(NotImplementedError) as err: model._fcache(np.asanyarray(1), np.asanyarray(1)) assert str(err.value) == "Subclasses should implement this" model = Hermite2D(x_degree=2, y_degree=2) assert model._fcache(np.asanyarray(1), np.asanyarray(1)) == { 0: np.asanyarray(1), 1: 2, 3: np.asanyarray(1), 4: 2, 2: 2.0, 5: -4.0 } model = Legendre2D(x_degree=2, y_degree=2) assert model._fcache(np.asanyarray(1), np.asanyarray(1)) == { 0: np.asanyarray(1), 1: np.asanyarray(1), 2: 1.0, 3: np.asanyarray(1), 4: np.asanyarray(1), 5: 1.0 } model = Chebyshev2D(x_degree=2, y_degree=2) assert model._fcache(np.asanyarray(1), np.asanyarray(1)) == { 0: np.asanyarray(1), 1: np.asanyarray(1), 2: 1.0, 3: np.asanyarray(1), 4: np.asanyarray(1), 5: 1.0 }
def test_fit_deriv_shape_error(): model = Hermite2D(x_degree=2, y_degree=2) with pytest.raises(ValueError) as err: model.fit_deriv(np.array([1, 2]), np.array([3, 4, 5])) assert str(err.value) == "x and y must have the same shape" model = Chebyshev2D(x_degree=2, y_degree=2) with pytest.raises(ValueError) as err: model.fit_deriv(np.array([1, 2]), np.array([3, 4, 5])) assert str(err.value) == "x and y must have the same shape" model = Legendre2D(x_degree=2, y_degree=2) with pytest.raises(ValueError) as err: model.fit_deriv(np.array([1, 2]), np.array([3, 4, 5])) assert str(err.value) == "x and y must have the same shape" model = Polynomial2D(degree=2) with pytest.raises(ValueError) as err: model.fit_deriv(np.array([1, 2]), np.array([3, 4, 5])) assert str(err.value) == "Expected x and y to be of equal size"
def convert_to_slitoffset_map(self, extractor, pm_list, n_slice_one_direction, slit_slice): xi = np.linspace(0, 2048, 128 + 1) from astropy.modeling import fitting from astropy.modeling.polynomial import Chebyshev2D x_domain = [0, 2048] y_domain = [0., 1.] p2_list = [] orders = extractor.orders_w_solutions for o in orders: oi = np.zeros_like(xi) + o shift_list = [] for p, m in pm_list[:n_slice_one_direction]: shift_list.append(p(xi, oi)) shift_list.append(np.zeros_like(xi)) for p, m in pm_list[n_slice_one_direction:]: shift_list.append(p(xi, oi)) p_init = Chebyshev2D(x_degree=1, y_degree=2, x_domain=x_domain, y_domain=y_domain) f = fitting.LinearLSQFitter() yi = 0.5 * (slit_slice[:-1] + slit_slice[1:]) xl, yl = np.meshgrid(xi, yi) zl = np.array(shift_list) p = f(p_init, xl, yl, zl) p2_list.append(p) return p2_list
def process_distortion_sky_band(utdate, refdate, band, obsids, config): from libs.products import ProductDB, PipelineStorage igr_path = IGRINSPath(config, utdate) igr_storage = PipelineStorage(igr_path) sky_filenames = igr_path.get_filenames(band, obsids) sky_basename = os.path.splitext(os.path.basename(sky_filenames[0]))[0] master_obsid = obsids[0] flaton_db_name = igr_path.get_section_filename_base( "PRIMARY_CALIB_PATH", "flat_on.db", ) flaton_db = ProductDB(flaton_db_name) # thar_db_name = igr_path.get_section_filename_base("PRIMARY_CALIB_PATH", # "thar.db", # ) # thar_db = ProductDB(thar_db_name) from libs.storage_descriptions import (COMBINED_IMAGE_DESC, ONED_SPEC_JSON_DESC) raw_spec_products = igr_storage.load( [COMBINED_IMAGE_DESC, ONED_SPEC_JSON_DESC], sky_basename) # raw_spec_products = PipelineProducts.load(sky_path.get_secondary_path("raw_spec")) from libs.storage_descriptions import SKY_WVLSOL_JSON_DESC wvlsol_products = igr_storage.load([SKY_WVLSOL_JSON_DESC], sky_basename)[SKY_WVLSOL_JSON_DESC] orders_w_solutions = wvlsol_products["orders"] wvl_solutions = wvlsol_products["wvl_sol"] ap = load_aperture2(igr_storage, band, master_obsid, flaton_db, raw_spec_products[ONED_SPEC_JSON_DESC]["orders"], orders_w_solutions) #orders_w_solutions = ap.orders if 1: # load reference data from libs.master_calib import load_sky_ref_data ref_utdate = config.get_value("REFDATE", utdate) sky_ref_data = load_sky_ref_data(ref_utdate, band) ohlines_db = sky_ref_data["ohlines_db"] ref_ohline_indices = sky_ref_data["ohline_indices"] orders_w_solutions = wvlsol_products["orders"] wvl_solutions = wvlsol_products["wvl_sol"] if 1: n_slice_one_direction = 2 n_slice = n_slice_one_direction * 2 + 1 i_center = n_slice_one_direction slit_slice = np.linspace(0., 1., n_slice + 1) slice_center = (slit_slice[i_center], slit_slice[i_center + 1]) slice_up = [(slit_slice[i_center+i], slit_slice[i_center+i+1]) \ for i in range(1, n_slice_one_direction+1)] slice_down = [(slit_slice[i_center-i-1], slit_slice[i_center-i]) \ for i in range(n_slice_one_direction)] d = raw_spec_products[COMBINED_IMAGE_DESC].data s_center = ap.extract_spectra_v2(d, slice_center[0], slice_center[1]) s_up, s_down = [], [] for s1, s2 in slice_up: s = ap.extract_spectra_v2(d, s1, s2) s_up.append(s) for s1, s2 in slice_down: s = ap.extract_spectra_v2(d, s1, s2) s_down.append(s) if 1: # now fit #ohline_indices = [ref_ohline_indices[o] for o in orders_w_solutions] if 0: def test_order(oi): ax = subplot(111) ax.plot(wvl_solutions[oi], s_center[oi]) #ax.plot(wvl_solutions[oi], raw_spec_products["specs"][oi]) o = orders[oi] line_indices = ref_ohline_indices[o] for li in line_indices: um = np.take(ohlines_db.um, li) intensity = np.take(ohlines_db.intensity, li) ax.vlines(um, ymin=0, ymax=-intensity) from libs.reidentify_ohlines import fit_ohlines, fit_ohlines_pixel def get_reidentified_lines_OH(orders_w_solutions, wvl_solutions, s_center): ref_pixel_list, reidentified_lines = \ fit_ohlines(ohlines_db, ref_ohline_indices, orders_w_solutions, wvl_solutions, s_center) reidentified_lines_map = dict( zip(orders_w_solutions, reidentified_lines)) return reidentified_lines_map, ref_pixel_list if band == "H": reidentified_lines_map, ref_pixel_list_oh = \ get_reidentified_lines_OH(orders_w_solutions, wvl_solutions, s_center) def refit_centroid(s_center, ref_pixel_list=ref_pixel_list_oh): centroids = fit_ohlines_pixel(s_center, ref_pixel_list) return centroids else: # band K reidentified_lines_map, ref_pixel_list_oh = \ get_reidentified_lines_OH(orders_w_solutions, wvl_solutions, s_center) import libs.master_calib as master_calib fn = "hitran_bootstrap_K_%s.json" % ref_utdate bootstrap_name = master_calib.get_master_calib_abspath(fn) import json bootstrap = json.load(open(bootstrap_name)) import libs.hitran as hitran r, ref_pixel_dict_hitrans = hitran.reidentify( wvl_solutions, s_center, bootstrap) # for i, s in r.items(): # ss = reidentified_lines_map[int(i)] # ss0 = np.concatenate([ss[0], s["pixel"]]) # ss1 = np.concatenate([ss[1], s["wavelength"]]) # reidentified_lines_map[int(i)] = (ss0, ss1) #reidentified_lines_map, ref_pixel_list def refit_centroid(s_center, ref_pixel_list=ref_pixel_list_oh, ref_pixel_dict_hitrans=ref_pixel_dict_hitrans): centroids_oh = fit_ohlines_pixel(s_center, ref_pixel_list) s_dict = dict(zip(orders_w_solutions, s_center)) centroids_dict_hitrans = hitran.fit_hitrans_pixel( s_dict, ref_pixel_dict_hitrans) centroids = [] for o, c_oh in zip(orders_w_solutions, centroids_oh): if o in centroids_dict_hitrans: c = np.concatenate( [c_oh, centroids_dict_hitrans[o]["pixel"]]) centroids.append(c) else: centroids.append(c_oh) return centroids # reidentified_lines_map = get_reidentified_lines(orders_w_solutions, # wvl_solutions, # s_center) if 1: # TODO: we should not need this, instead recycle from preivious step. fitted_centroid_center = refit_centroid(s_center) # fitted_centroid_center = fit_ohlines_pixel(s_center, # ref_pixel_list) d_shift_up = [] for s in s_up: # TODO: ref_pixel_list_filtered need to be updated with recent fit. fitted_centroid = refit_centroid(s) # fitted_centroid = fit_ohlines_pixel(s, # ref_pixel_list) d_shift = [ b - a for a, b in zip(fitted_centroid_center, fitted_centroid) ] d_shift_up.append(d_shift) d_shift_down = [] for s in s_down: # TODO: ref_pixel_list_filtered need to be updated with recent fit. fitted_centroid = refit_centroid(s) # fitted_centroid = fit_ohlines_pixel(s, # ref_pixel_list) #fitted_centroid_center, d_shift = [ b - a for a, b in zip(fitted_centroid_center, fitted_centroid) ] d_shift_down.append(d_shift) if 1: # now fit orders = orders_w_solutions x_domain = [0, 2048] y_domain = [orders[0] - 2, orders[-1] + 2] xl = np.concatenate(fitted_centroid_center) yl_ = [ o + np.zeros_like(x_) for o, x_ in zip(orders, fitted_centroid_center) ] yl = np.concatenate(yl_) from libs.ecfit import fit_2dspec, check_fit_simple zl_list = [np.concatenate(d_) for d_ \ in d_shift_down[::-1] + d_shift_up] pm_list = [] for zl in zl_list: p, m = fit_2dspec(xl, yl, zl, x_degree=1, y_degree=1, x_domain=x_domain, y_domain=y_domain) pm_list.append((p, m)) zz_std_list = [] for zl, (p, m) in zip(zl_list, pm_list): z_m = p(xl[m], yl[m]) zz = z_m - zl[m] zz_std_list.append(zz.std()) fig_list = [] from matplotlib.figure import Figure for zl, (p, m) in zip(zl_list, pm_list): fig = Figure() check_fit_simple(fig, xl[m], yl[m], zl[m], p, orders) fig_list.append(fig) if 1: xi = np.linspace(0, 2048, 128 + 1) from astropy.modeling import fitting from astropy.modeling.polynomial import Chebyshev2D x_domain = [0, 2048] y_domain = [0., 1.] p2_list = [] for o in orders: oi = np.zeros_like(xi) + o shift_list = [] for p, m in pm_list[:n_slice_one_direction]: shift_list.append(p(xi, oi)) shift_list.append(np.zeros_like(xi)) for p, m in pm_list[n_slice_one_direction:]: shift_list.append(p(xi, oi)) p_init = Chebyshev2D(x_degree=1, y_degree=2, x_domain=x_domain, y_domain=y_domain) f = fitting.LinearLSQFitter() yi = 0.5 * (slit_slice[:-1] + slit_slice[1:]) xl, yl = np.meshgrid(xi, yi) zl = np.array(shift_list) p = f(p_init, xl, yl, zl) p2_list.append(p) if 1: p2_dict = dict(zip(orders, p2_list)) # save order_map, etc order_map = ap.make_order_map() slitpos_map = ap.make_slitpos_map() order_map2 = ap.make_order_map(mask_top_bottom=True) slitoffset_map = np.empty_like(slitpos_map) slitoffset_map.fill(np.nan) wavelength_map = np.empty_like(slitpos_map) wavelength_map.fill(np.nan) from scipy.interpolate import interp1d for o, wvl in zip(ap.orders, wvl_solutions): xi = np.arange(0, 2048) xl, yl = np.meshgrid(xi, xi) msk = order_map == o xl_msk = xl[msk] slitoffset_map_msk = p2_dict[o](xl_msk, slitpos_map[msk]) slitoffset_map[msk] = slitoffset_map_msk wvl_interp1d = interp1d(xi, wvl, bounds_error=False) wavelength_map[msk] = wvl_interp1d(xl_msk - slitoffset_map_msk) from libs.storage_descriptions import (ORDERMAP_FITS_DESC, SLITPOSMAP_FITS_DESC, SLITOFFSET_FITS_DESC, WAVELENGTHMAP_FITS_DESC, ORDERMAP_MASKED_FITS_DESC) from libs.products import PipelineImage, PipelineProducts products = PipelineProducts("Distortion map") for desc, im in [(ORDERMAP_FITS_DESC, order_map), (SLITPOSMAP_FITS_DESC, slitpos_map), (SLITOFFSET_FITS_DESC, slitoffset_map), (WAVELENGTHMAP_FITS_DESC, wavelength_map), (ORDERMAP_MASKED_FITS_DESC, order_map2)]: products.add(desc, PipelineImage([], im)) igr_storage.store(products, mastername=sky_filenames[0], masterhdu=None) from libs.qa_helper import figlist_to_pngs sky_figs = igr_path.get_section_filename_base( "QA_PATH", "oh_distortion", "oh_distortion_" + sky_basename) print fig_list figlist_to_pngs(sky_figs, fig_list) if 0: # test x = np.arange(2048, dtype="d") oi = 10 o = orders[oi] yi = 0.5 * (slit_slice[:-1] + slit_slice[1:]) ax1 = subplot(211) s1 = s_up[-1][oi] s2 = s_down[-1][oi] ax1.plot(x, s1) ax1.plot(x, s2) ax2 = subplot(212, sharex=ax1, sharey=ax1) dx1 = p2_dict[o](x, yi[-1] + np.zeros_like(x)) ax2.plot(x - dx1, s1) dx2 = p2_dict[o](x, yi[0] + np.zeros_like(x)) ax2.plot(x - dx2, s2)
def trace_aperture_chebyshev(xy_list, domain=None): """ a list of (x_array, y_array). y_array must be a masked array """ import numpy.polynomial.chebyshev as cheb #for x, y in r["cent_bottom_list"]: # xy_list = r["cent_up_list"] # domain = [0, 2047] # if domain is None: # xmax = max(max(x) for x, y in xy_list) # xmin = min(min(x) for x, y in xy_list) # domain = [xmin, xmax] if domain is None: domain = [0, 2047] # we first fit the all traces with 2d chebyshev polynomials x_list, o_list, y_list = [], [], [] for o, (x, y) in enumerate(xy_list): if hasattr(y, "mask"): msk = ~y.mask & np.isfinite(y.data) y = y.data else: msk = np.isfinite(np.array(y, "d")) x1 = np.array(x)[msk] x_list.append(x1) o_list.append(np.zeros(len(x1))+o) y_list.append(np.array(y)[msk]) n_o = len(xy_list) from astropy.modeling import models, fitting from astropy.modeling.polynomial import Chebyshev2D x_degree, y_degree = 4, 5 p_init = Chebyshev2D(x_degree, y_degree, x_domain=domain, y_domain=[0, n_o-1]) fit_p = fitting.LinearLSQFitter() xxx, ooo, yyy = (np.concatenate(x_list), np.concatenate(o_list), np.concatenate(y_list)) p = fit_p(p_init, xxx, ooo, yyy) if 0: ax1 = subplot(121) for o, xy in enumerate(xy_list): ax1.plot(x_list[o], y_list[o] - p(x_list[o], o+np.zeros_like(x_list[o]))) for ii in range(3): # number of iteration mmm = np.abs(yyy - p(xxx, ooo)) < 1 # This need to be fixed with actual estimation of sigma. p = fit_p(p_init, xxx[mmm], ooo[mmm], yyy[mmm]) if 0: ax2=subplot(122, sharey=ax1) for o, xy in enumerate(xy_list): ax2=plot(x_list[o], y_list[o] - p(x_list[o], o+np.zeros_like(x_list[o]))) # Now we need to derive a 1d chebyshev for each order. While # there should be an analitical way, here we refit the trace for # each order using the result of 2d fit. xx = np.arange(domain[0], domain[1]) oo = np.zeros_like(xx) def _get_f(o0): y_m = p(xx, oo+o0) f = cheb.Chebyshev.fit(xx, y_m, x_degree, domain=domain) return f f_list = [] ooo = [o[0] for o in o_list] #for x, o in zip(x_list, o_list): f_list = [_get_f(o0) for o0 in ooo] def _get_f_old(next_orders, y_thresh): oi = next_orders.pop(0) y_m = p(xx, oo+oi) f = cheb.Chebyshev.fit(xx, y_m, x_degree, domain=domain) if next_orders: # if not the last order if np.all(y_thresh(y_m)): print "all negative at ", oi next_orders = next_orders[:1] return oi, f, next_orders def _get_f(next_orders, y_thresh): oi = next_orders.pop(0) y_m = p(xx, oo+oi) f = cheb.Chebyshev.fit(xx, y_m, x_degree, domain=domain) if np.all(y_thresh(y_m)): print "all negative at ", oi next_orders = [] return oi, f, next_orders # go down in order f_list_down = [] o_list_down = [] go_down_orders = [ooo[0]-_oi for _oi in range(1, 5)] while go_down_orders: oi, f, go_down_orders = _get_f(go_down_orders, y_thresh=lambda y_m: y_m < domain[0]) f_list_down.append(f) o_list_down.append(oi) f_list_up = [] o_list_up = [] go_up_orders = [ooo[-1]+_oi for _oi in range(1, 5)] while go_up_orders: oi, f, go_up_orders = _get_f(go_up_orders, y_thresh=lambda y_m: y_m > domain[-1]) f_list_up.append(f) o_list_up.append(oi) if 0: _get_f(next_orders) oi = go_down_orders.pop(0) y_m = p(xx, oo+ooo[0]-oi) f = cheb.Chebyshev.fit(xx, y_m, x_degree, domain=domain) if go_down_orders: # if not the last order if np.all(y_m < domain[0]): print "all negative at ", ooo[0]-oi go_down_orders = [oi+1] else: f_list_down.append(f) else: f_list_down.append(f) print o_list_down[::-1] + ooo + o_list_up return f_list, f_list_down[::-1] + f_list + f_list_up
def trace_aperture_chebyshev(xy_list, domain=None): """ a list of (x_array, y_array). y_array must be a masked array """ import numpy.polynomial.chebyshev as cheb #for x, y in r["cent_bottom_list"]: # xy_list = r["cent_up_list"] # domain = [0, 2047] # if domain is None: # xmax = max(max(x) for x, y in xy_list) # xmin = min(min(x) for x, y in xy_list) # domain = [xmin, xmax] if domain is None: domain = [0, 2047] # we first fit the all traces with 2d chebyshev polynomials x_list, o_list, y_list = [], [], [] for o, (x, y) in enumerate(xy_list): if hasattr(y, "mask"): msk = ~y.mask & np.isfinite(y.data) y = y.data else: msk = np.isfinite(np.array(y, "d")) x1 = np.array(x)[msk] x_list.append(x1) o_list.append(np.zeros(len(x1)) + o) y_list.append(np.array(y)[msk]) n_o = len(xy_list) from astropy.modeling import models, fitting from astropy.modeling.polynomial import Chebyshev2D x_degree, y_degree = 4, 5 p_init = Chebyshev2D(x_degree, y_degree, x_domain=domain, y_domain=[0, n_o - 1]) fit_p = fitting.LinearLSQFitter() xxx, ooo, yyy = (np.concatenate(x_list), np.concatenate(o_list), np.concatenate(y_list)) p = fit_p(p_init, xxx, ooo, yyy) if 0: ax1 = subplot(121) for o, xy in enumerate(xy_list): ax1.plot(x_list[o], y_list[o] - p(x_list[o], o + np.zeros_like(x_list[o]))) for ii in range(3): # number of iteration mmm = np.abs( yyy - p(xxx, ooo) ) < 1 # This need to be fixed with actual estimation of sigma. p = fit_p(p_init, xxx[mmm], ooo[mmm], yyy[mmm]) if 0: ax2 = subplot(122, sharey=ax1) for o, xy in enumerate(xy_list): ax2 = plot(x_list[o], y_list[o] - p(x_list[o], o + np.zeros_like(x_list[o]))) # Now we need to derive a 1d chebyshev for each order. While # there should be an analitical way, here we refit the trace for # each order using the result of 2d fit. f_list = [] for x, o in zip(x_list, o_list): y_m = p(x, o) f = cheb.Chebyshev.fit(x, y_m, x_degree, domain=domain) f_list.append(f) return f_list