diff --git a/observation_sim/mock_objects/SpecDisperser/SpecDisperser.py b/observation_sim/mock_objects/SpecDisperser/SpecDisperser.py index 2b2d7c29baa9e9965725da0cf92fac98b8495ae2..6936eb139856832dcd6dd2be4fb6d151f53e9d32 100644 --- a/observation_sim/mock_objects/SpecDisperser/SpecDisperser.py +++ b/observation_sim/mock_objects/SpecDisperser/SpecDisperser.py @@ -50,8 +50,21 @@ def rotate90(array_orig=None, xc=0, yc=0, isClockwise=0): class SpecDisperser(object): - def __init__(self, orig_img=None, xcenter=0, ycenter=0, origin=[100, 100], tar_spec=None, band_start=6200, - band_end=10000, isAlongY=0, conf='../param/CONF/csst.conf', gid=0, flat_cube=None, ignoreBeam=[]): + def __init__( + self, + orig_img=None, + xcenter=0, + ycenter=0, + origin=[100, 100], + tar_spec=None, + band_start=6200, + band_end=10000, + isAlongY=0, + conf="../param/CONF/csst.conf", + gid=0, + flat_cube=None, + ignoreBeam=[], + ): """ orig_img: normal image,galsim image xcenter, ycenter: the center of galaxy in orig_img @@ -66,23 +79,21 @@ class SpecDisperser(object): # self.img_y = orig_img.shape[0] # 5 orders, A, B , - orderName=["A","B","C","D","E"] + orderName = ["A", "B", "C", "D", "E"] self.orig_img_orders = OrderedDict() if isinstance(orig_img, list): orig_img_list = orig_img list_len = len(orig_img_list) if list_len < 5: - for i in np.arange(5-list_len): - orig_img_list.append(orig_img_list[list_len-1]) + for i in np.arange(5 - list_len): + orig_img_list.append(orig_img_list[list_len - 1]) for i, k in enumerate(orig_img_list): self.orig_img_orders[orderName[i]] = k - if isinstance(orig_img, galsim.Image): for i in np.arange(5): self.orig_img_orders[orderName[i]] = orig_img - orig_img_one = self.orig_img_orders["A"] self.thumb_img = np.abs(orig_img_one.array) @@ -99,8 +110,12 @@ class SpecDisperser(object): self.flat_cube = flat_cube if self.isAlongY == 1: for order in orderName: - self.orig_img_orders[order], self.thumb_x, self.thumb_y = rotate90(array_orig=self.orig_img_orders[order], xc=orig_img_one.center.x, - yc=orig_img_one.center.y, isClockwise=1) + self.orig_img_orders[order], self.thumb_x, self.thumb_y = rotate90( + array_orig=self.orig_img_orders[order], + xc=orig_img_one.center.x, + yc=orig_img_one.center.y, + isClockwise=1, + ) # self.thumb_img, self.thumb_x, self.thumb_y = rotate90(array_orig=self.thumb_img, xc=orig_img_one.center.x, # yc=orig_img_one.center.y, isClockwise=1) @@ -139,6 +154,7 @@ class SpecDisperser(object): from .disperse_c import interp from .disperse_c import disperse + # from MockObject.disperse_c import disperse self.thumb_img = np.abs(self.orig_img_orders[beam].array) self.thumb_x = self.orig_img_orders[beam].center.x @@ -146,22 +162,21 @@ class SpecDisperser(object): self.img_sh = self.orig_img_orders[beam].array.shape dx = self.grating_conf.dxlam[beam] xoff = 0 - ytrace_beam, lam_beam = self.grating_conf.get_beam_trace(x=self.xcenter, y=self.ycenter, dx=(dx + xoff), - beam=beam) + ytrace_beam, lam_beam = self.grating_conf.get_beam_trace( + x=self.xcenter, y=self.ycenter, dx=(dx + xoff), beam=beam + ) # Account for pixel centering of the trace - yfrac_beam = ytrace_beam - floor(ytrace_beam+0.5) + yfrac_beam = ytrace_beam - floor(ytrace_beam + 0.5) ysens = lam_beam * 0 lam_index = argsort(lam_beam) conf_sens = self.grating_conf.sens[beam] - lam_intep = np.linspace(self.band_start, self.band_end, int( - (self.band_end - self.band_start) / 0.1)) + lam_intep = np.linspace(self.band_start, self.band_end, int((self.band_end - self.band_start) / 0.1)) - thri = interpolate.interp1d( - conf_sens['WAVELENGTH'], conf_sens['SENSITIVITY']) - spci = interpolate.interp1d(self.spec['WAVELENGTH'], self.spec['FLUX']) + thri = interpolate.interp1d(conf_sens["WAVELENGTH"], conf_sens["SENSITIVITY"]) + spci = interpolate.interp1d(self.spec["WAVELENGTH"], self.spec["FLUX"]) beam_thr = thri(lam_intep) spec_sample = spci(lam_intep) @@ -179,14 +194,14 @@ class SpecDisperser(object): # # self.writerSensitivityFile(conffile = self.grating_conf_file, beam = beam, w = lam_beam[lam_index], sens = ysensitivity[lam_index]) - ysens[lam_index] = interp.interp_conserve_c(lam_beam[lam_index], lam_intep, bean_thr_spec, integrate=1, left=0, - right=0) + ysens[lam_index] = interp.interp_conserve_c( + lam_beam[lam_index], lam_intep, bean_thr_spec, integrate=1, left=0, right=0 + ) sensitivity_beam = ysens len_spec_x = len(dx) - len_spec_y = int( - abs(ceil(ytrace_beam[-1]) - floor(ytrace_beam[0])) + 1) + len_spec_y = int(abs(ceil(ytrace_beam[-1]) - floor(ytrace_beam[0])) + 1) beam_sh = (self.img_sh[0] + len_spec_y, self.img_sh[1] + len_spec_x) modelf = zeros(product(beam_sh), dtype=float) @@ -196,7 +211,7 @@ class SpecDisperser(object): dxpix = dx - dx[0] + x0[1] - dyc = cast[int](np.floor(ytrace_beam+0.5)) + dyc = cast[int](np.floor(ytrace_beam + 0.5)) # dypix = cast[int](np.floor(ytrace_beam - dyc[0] + x0[0] + 0.5)) dypix = dyc - dyc[0] + x0[0] @@ -204,7 +219,7 @@ class SpecDisperser(object): frac_ids = yfrac_beam < 0 dypix[frac_ids] = dypix[frac_ids] - 1 - yfrac_beam[frac_ids] = 1+yfrac_beam[frac_ids] + yfrac_beam[frac_ids] = 1 + yfrac_beam[frac_ids] flat_index = idx[dypix, dxpix] @@ -231,9 +246,8 @@ class SpecDisperser(object): else: beam_flat = zeros([len(modelf), len(self.flat_cube)]) - sub_flat_cube = zeros( - [len(self.flat_cube), beam_sh[0], beam_sh[1]]) - sub_flat_cube[0] = sub_flat_cube[0] + 1. + sub_flat_cube = zeros([len(self.flat_cube), beam_sh[0], beam_sh[1]]) + sub_flat_cube[0] = sub_flat_cube[0] + 1.0 overlap_flag = 1 @@ -260,33 +274,38 @@ class SpecDisperser(object): overlap_flag = 0 if overlap_flag == 1: - sub_flat_cube[:, beam_y_s-originOut_y:beam_y_e-originOut_y+1, beam_x_s-originOut_x:beam_x_e - - originOut_x+1] = self.flat_cube[:, beam_y_s:beam_y_e+1, beam_x_s:beam_x_e+1] + sub_flat_cube[ + :, + beam_y_s - originOut_y : beam_y_e - originOut_y + 1, + beam_x_s - originOut_x : beam_x_e - originOut_x + 1, + ] = self.flat_cube[:, beam_y_s : beam_y_e + 1, beam_x_s : beam_x_e + 1] for i in arange(0, len(self.flat_cube), 1): beam_flat[:, i] = sub_flat_cube[i].flatten() -# beam_flat = zeros([len(modelf), len(self.flat_cube)]) -# flat_sh = self.flat_cube[0].shape -# for i in arange(0, beam_sh[0], 1): -# for j in arange(0, beam_sh[1], 1): -# k = i * beam_sh[1] + j -# if originOut_y + i >= flat_sh[0] or originOut_y + i < 0 or originOut_x + j>= flat_sh[1] or originOut_x+j < 0: -# temp_bf = np.zeros_like(self.flat_cube[:, 0, 0]) -# temp_bf[0] = 1.0 -# beam_flat[k] = temp_bf -# else: -# beam_flat[k] = self.flat_cube[:, originOut_y + i, originOut_x + j] - - status = disperse.disperse_grism_object(self.thumb_img.astype(np.float32), - flat_index[nonz], - yfrac_beam[nonz], - sensitivity_beam[nonz], - modelf, x0, - array(self.img_sh, - dtype=int64), - array(beam_sh, dtype=int64), - beam_flat, - lam_beam[lam_index][nonz]) + # beam_flat = zeros([len(modelf), len(self.flat_cube)]) + # flat_sh = self.flat_cube[0].shape + # for i in arange(0, beam_sh[0], 1): + # for j in arange(0, beam_sh[1], 1): + # k = i * beam_sh[1] + j + # if originOut_y + i >= flat_sh[0] or originOut_y + i < 0 or originOut_x + j>= flat_sh[1] or originOut_x+j < 0: + # temp_bf = np.zeros_like(self.flat_cube[:, 0, 0]) + # temp_bf[0] = 1.0 + # beam_flat[k] = temp_bf + # else: + # beam_flat[k] = self.flat_cube[:, originOut_y + i, originOut_x + j] + + status = disperse.disperse_grism_object( + self.thumb_img.astype(np.float32), + flat_index[nonz], + yfrac_beam[nonz], + sensitivity_beam[nonz], + modelf, + x0, + array(self.img_sh, dtype=int64), + array(beam_sh, dtype=int64), + beam_flat, + lam_beam[lam_index][nonz], + ) model = modelf.reshape(beam_sh) # n1 = np.sum(np.isinf(model)) @@ -296,27 +315,25 @@ class SpecDisperser(object): # if n1>0 or n2 > 0: # print("DEBUG: SpecDisperser, inf:%d, nan:%d--------%d,%d"%(n1, n2, n3, n4)) # print(dypix) - # n1 = np.sum(np.isinf(self.thumb_img.astype(np.float32))) - # n2 = np.sum(np.isnan(self.thumb_img.astype(np.float32))) - # n3 = np.sum(np.isinf(yfrac_beam)) - # n4 = np.sum(np.isnan(yfrac_beam)) - # n5 = np.sum(np.isinf(sensitivity_beam)) - # n6 = np.sum(np.isnan(sensitivity_beam)) - # print("DEBUG: SpecDisperser, innput ---inf:%d, nan:%d, yfrac_beam:%d/%d, sensitivity_beam:%d/%d"%(n1, n2, n3, n4, n5, n6)) + # n1 = np.sum(np.isinf(self.thumb_img.astype(np.float32))) + # n2 = np.sum(np.isnan(self.thumb_img.astype(np.float32))) + # n3 = np.sum(np.isinf(yfrac_beam)) + # n4 = np.sum(np.isnan(yfrac_beam)) + # n5 = np.sum(np.isinf(sensitivity_beam)) + # n6 = np.sum(np.isnan(sensitivity_beam)) + # print("DEBUG: SpecDisperser, innput ---inf:%d, nan:%d, yfrac_beam:%d/%d, sensitivity_beam:%d/%d"%(n1, n2, n3, n4, n5, n6)) self.beam_flux[beam] = sum(modelf) if self.isAlongY == 1: model, _, _ = rotate90(array_orig=model, isClockwise=0) return model, originOut_x, originOut_y, dxpix, dypix, lam_beam, ysens - def writerSensitivityFile(self, conffile='', beam='', w=None, sens=None): - orders = {'A': '1st', 'B': '0st', 'C': '2st', 'D': '-1st', 'E': '-2st'} - sens_file_name = conffile[0:-5] + \ - '_sensitivity_' + orders[beam] + '.fits' + def writerSensitivityFile(self, conffile="", beam="", w=None, sens=None): + orders = {"A": "1st", "B": "0st", "C": "2st", "D": "-1st", "E": "-2st"} + sens_file_name = conffile[0:-5] + "_sensitivity_" + orders[beam] + ".fits" if not os.path.exists(sens_file_name) == True: - senstivity_out = Table( - array([w, sens]).T, names=('WAVELENGTH', 'SENSITIVITY')) - senstivity_out.write(sens_file_name, format='fits') + senstivity_out = Table(array([w, sens]).T, names=("WAVELENGTH", "SENSITIVITY")) + senstivity_out.write(sens_file_name, format="fits") """ @@ -324,8 +341,8 @@ Demonstrate aXe trace polynomials. """ -class aXeConf(): - def __init__(self, conf_file='WFC3.IR.G141.V2.5.conf'): +class aXeConf: + def __init__(self, conf_file="WFC3.IR.G141.V2.5.conf"): """Read an aXe-compatible configuration file Parameters @@ -340,17 +357,17 @@ class aXeConf(): self.count_beam_orders() # Global XOFF/YOFF offsets - if 'XOFF' in self.conf.keys(): - self.xoff = np.float(self.conf['XOFF']) + if "XOFF" in self.conf.keys(): + self.xoff = np.float(self.conf["XOFF"]) else: - self.xoff = 0. + self.xoff = 0.0 - if 'YOFF' in self.conf.keys(): - self.yoff = np.float(self.conf['YOFF']) + if "YOFF" in self.conf.keys(): + self.yoff = np.float(self.conf["YOFF"]) else: - self.yoff = 0. + self.yoff = 0.0 - def read_conf_file(self, conf_file='WFC3.IR.G141.V2.5.conf'): + def read_conf_file(self, conf_file="WFC3.IR.G141.V2.5.conf"): """Read an aXe config file, convert floats and arrays Parameters @@ -366,11 +383,11 @@ class aXeConf(): lines = open(conf_file).readlines() for line in lines: # empty / commented lines - if (line.startswith('#')) | (line.strip() == '') | ('"' in line): + if (line.startswith("#")) | (line.strip() == "") | ('"' in line): continue # split the line, taking out ; and # comments - spl = line.split(';')[0].split('#')[0].split() + spl = line.split(";")[0].split("#")[0].split() param = spl[0] if len(spl) > 2: value = np.cast[float](spl[1:]) @@ -385,22 +402,20 @@ class aXeConf(): return conf def count_beam_orders(self): - """Get the maximum polynomial order in DYDX or DLDP for each beam - """ + """Get the maximum polynomial order in DYDX or DLDP for each beam""" self.orders = {} - for beam in ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J']: + for beam in ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"]: order = 0 - while 'DYDX_{0:s}_{1:d}'.format(beam, order) in self.conf.keys(): + while "DYDX_{0:s}_{1:d}".format(beam, order) in self.conf.keys(): order += 1 - while 'DLDP_{0:s}_{1:d}'.format(beam, order) in self.conf.keys(): + while "DLDP_{0:s}_{1:d}".format(beam, order) in self.conf.keys(): order += 1 self.orders[beam] = order - 1 def get_beams(self): - """Get beam parameters and read sensitivity curves - """ + """Get beam parameters and read sensitivity curves""" import os from collections import OrderedDict from astropy.table import Table, Column @@ -413,12 +428,13 @@ class aXeConf(): for beam in self.orders: if self.orders[beam] > 0: self.beams.append(beam) - self.dxlam[beam] = np.arange(self.conf['BEAM{0}'.format(beam)].min(), - self.conf['BEAM{0}'.format(beam)].max(), dtype=int) - self.nx[beam] = int(self.dxlam[beam].max() - - self.dxlam[beam].min()) + 1 + self.dxlam[beam] = np.arange( + self.conf["BEAM{0}".format(beam)].min(), self.conf["BEAM{0}".format(beam)].max(), dtype=int + ) + self.nx[beam] = int(self.dxlam[beam].max() - self.dxlam[beam].min()) + 1 self.sens[beam] = Table.read( - '{0}/{1}'.format(os.path.dirname(self.conf_file), self.conf['SENSITIVITY_{0}'.format(beam)])) + "{0}/{1}".format(os.path.dirname(self.conf_file), self.conf["SENSITIVITY_{0}".format(beam)]) + ) # self.sens[beam].wave = np.cast[np.double](self.sens[beam]['WAVELENGTH']) # self.sens[beam].sens = np.cast[np.double](self.sens[beam]['SENSITIVITY']) @@ -505,7 +521,7 @@ class aXeConf(): # $\lambda = dldp_0 + dldp_1 dp + dldp_2 dp^2$ ... poly_order = len(dydx) - 1 - if (poly_order == 2): + if poly_order == 2: if np.abs(np.unique(dydx[2])).max() == 0: poly_order = 1 @@ -515,10 +531,9 @@ class aXeConf(): dp = np.sqrt(1 + dydx[1] ** 2) * (dx) elif poly_order == 2: # quadratic trace u0 = dydx[1] + 2 * dydx[2] * (0) - dp0 = (u0 * np.sqrt(1 + u0 ** 2) + np.arcsinh(u0)) / (4 * dydx[2]) + dp0 = (u0 * np.sqrt(1 + u0**2) + np.arcsinh(u0)) / (4 * dydx[2]) u = dydx[1] + 2 * dydx[2] * (dx) - dp = (u * np.sqrt(1 + u ** 2) + np.arcsinh(u)) / \ - (4 * dydx[2]) - dp0 + dp = (u * np.sqrt(1 + u**2) + np.arcsinh(u)) / (4 * dydx[2]) - dp0 else: # high order shape, numerical integration along trace # (this can be slow) @@ -530,11 +545,10 @@ class aXeConf(): dyfull += i * dydx[i] * (xfull - 0.5) ** (i - 1) # Integrate from 0 to dx / -dx - dpfull = xfull * 0. + dpfull = xfull * 0.0 lt0 = xfull < 0 if lt0.sum() > 1: - dpfull[lt0] = np.cumsum( - np.sqrt(1 + dyfull[lt0][::-1] ** 2))[::-1] + dpfull[lt0] = np.cumsum(np.sqrt(1 + dyfull[lt0][::-1] ** 2))[::-1] dpfull[lt0] *= -1 # @@ -548,7 +562,7 @@ class aXeConf(): return dp - def get_beam_trace(self, x=507, y=507, dx=0., beam='A'): + def get_beam_trace(self, x=507, y=507, dx=0.0, beam="A"): """Get an aXe beam trace for an input reference pixel and list of output x pixels `dx` Parameters @@ -578,18 +592,16 @@ class aXeConf(): NORDER = self.orders[beam] + 1 xi, yi = x - self.xoff, y - self.yoff - xoff_beam = self.field_dependent( - xi, yi, self.conf['XOFF_{0}'.format(beam)]) - yoff_beam = self.field_dependent( - xi, yi, self.conf['YOFF_{0}'.format(beam)]) + xoff_beam = self.field_dependent(xi, yi, self.conf["XOFF_{0}".format(beam)]) + yoff_beam = self.field_dependent(xi, yi, self.conf["YOFF_{0}".format(beam)]) # y offset of trace (DYDX) dydx = np.zeros(NORDER) # 0 #+1.e-80 dydx = [0] * NORDER for i in range(NORDER): - if 'DYDX_{0:s}_{1:d}'.format(beam, i) in self.conf.keys(): - coeffs = self.conf['DYDX_{0:s}_{1:d}'.format(beam, i)] + if "DYDX_{0:s}_{1:d}".format(beam, i) in self.conf.keys(): + coeffs = self.conf["DYDX_{0:s}_{1:d}".format(beam, i)] dydx[i] = self.field_dependent(xi, yi, coeffs) # $dy = dydx_0+dydx_1 dx+dydx_2 dx^2+$ ... @@ -603,14 +615,20 @@ class aXeConf(): dldp = [0] * NORDER for i in range(NORDER): - if 'DLDP_{0:s}_{1:d}'.format(beam, i) in self.conf.keys(): - coeffs = self.conf['DLDP_{0:s}_{1:d}'.format(beam, i)] + if "DLDP_{0:s}_{1:d}".format(beam, i) in self.conf.keys(): + coeffs = self.conf["DLDP_{0:s}_{1:d}".format(beam, i)] dldp[i] = self.field_dependent(xi, yi, coeffs) - self.eval_input = {'x': x, 'y': y, 'beam': beam, 'dx': dx} - self.eval_output = {'xi': xi, 'yi': yi, 'dldp': dldp, 'dydx': dydx, - 'xoff_beam': xoff_beam, 'yoff_beam': yoff_beam, - 'dy': dy} + self.eval_input = {"x": x, "y": y, "beam": beam, "dx": dx} + self.eval_output = { + "xi": xi, + "yi": yi, + "dldp": dldp, + "dydx": dydx, + "xoff_beam": xoff_beam, + "yoff_beam": yoff_beam, + "dy": dy, + } dp = self.evaluate_dp(dx - xoff_beam, dydx) # ## dp is the arc length along the trace @@ -648,13 +666,13 @@ class aXeConf(): # dp = np.interp(dx-xoff_beam, xfull, dpfull) # Evaluate dldp - lam = dp * 0. + lam = dp * 0.0 for i in range(NORDER): - lam += dldp[i] * dp ** i + lam += dldp[i] * dp**i return dy, lam - def show_beams(self, beams=['E', 'D', 'C', 'B', 'A']): + def show_beams(self, beams=["E", "D", "C", "B", "A"]): """ Make a demo plot of the beams of a given configuration file """ @@ -663,45 +681,42 @@ class aXeConf(): x0, x1 = 507, 507 dx = np.arange(-800, 1200) - if 'WFC3.UV' in self.conf_file: + if "WFC3.UV" in self.conf_file: x0, x1 = 2073, 250 dx = np.arange(-1200, 1200) - if 'G800L' in self.conf_file: + if "G800L" in self.conf_file: x0, x1 = 2124, 1024 dx = np.arange(-1200, 1200) s = 200 # marker size fig = plt.figure(figsize=[10, 3]) - plt.scatter(0, 0, marker='s', s=s, color='black', edgecolor='0.8', - label='Direct') + plt.scatter(0, 0, marker="s", s=s, color="black", edgecolor="0.8", label="Direct") for beam in beams: - if 'XOFF_{0}'.format(beam) not in self.conf.keys(): + if "XOFF_{0}".format(beam) not in self.conf.keys(): continue - xoff = self.field_dependent( - x0, x1, self.conf['XOFF_{0}'.format(beam)]) + xoff = self.field_dependent(x0, x1, self.conf["XOFF_{0}".format(beam)]) dy, lam = self.get_beam_trace(x0, x1, dx=dx, beam=beam) - xlim = self.conf['BEAM{0}'.format(beam)] + xlim = self.conf["BEAM{0}".format(beam)] ok = (dx >= xlim[0]) & (dx <= xlim[1]) - plt.scatter(dx[ok] + xoff, dy[ok], c=lam[ok] / 1.e4, marker='s', s=s, - alpha=0.5, edgecolor='None') - plt.text(np.median(dx[ok]), np.median(dy[ok]) + 1, beam, - ha='center', va='center', fontsize=14) - print('Beam {0}, lambda=({1:.1f} - {2:.1f})'.format(beam, - lam[ok].min(), lam[ok].max())) + plt.scatter(dx[ok] + xoff, dy[ok], c=lam[ok] / 1.0e4, marker="s", s=s, alpha=0.5, edgecolor="None") + plt.text(np.median(dx[ok]), np.median(dy[ok]) + 1, beam, ha="center", va="center", fontsize=14) + print("Beam {0}, lambda=({1:.1f} - {2:.1f})".format(beam, lam[ok].min(), lam[ok].max())) plt.grid() - plt.xlabel(r'$\Delta x$') - plt.ylabel(r'$\Delta y$') + plt.xlabel(r"$\Delta x$") + plt.ylabel(r"$\Delta y$") cb = plt.colorbar(pad=0.01, fraction=0.05) - cb.set_label(r'$\lambda\,(\mu\mathrm{m})$') + cb.set_label(r"$\lambda\,(\mu\mathrm{m})$") plt.title(self.conf_file) plt.tight_layout() - plt.savefig('{0}.pdf'.format(self.conf_file)) + plt.savefig("{0}.pdf".format(self.conf_file)) # def load_grism_config(conf_file): + + # """Load parameters from an aXe configuration file # Parameters