pep8 formatting (3515cc73) · Commits · csst-sims / csst_msc_sim

observation_sim/mock_objects/SpecDisperser/SpecDisperser.py

+143 −128

Original line number	Diff line number	Diff line
		@@ -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
		@@ -77,12 +90,10 @@ class SpecDisperser(object):
		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,8 +162,9 @@ 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)
		@@ -156,12 +173,10 @@ class SpecDisperser(object):
		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)
		@@ -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,8 +274,11 @@ 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()
		@@ -277,16 +294,18 @@ class SpecDisperser(object):
		# else:
		# beam_flat[k] = self.flat_cube[:, originOut_y + i, originOut_x + j]

		status = disperse.disperse_grism_object(self.thumb_img.astype(np.float32),
		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),
		modelf,
		x0,
		array(self.img_sh, dtype=int64),
		array(beam_sh, dtype=int64),
		beam_flat,
		lam_beam[lam_index][nonz])
		lam_beam[lam_index][nonz],
		)

		model = modelf.reshape(beam_sh)
		# n1 = np.sum(np.isinf(model))
		@@ -309,14 +328,12 @@ class SpecDisperser(object):
		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

		@@ -517,8 +533,7 @@ class aXeConf():
		u0 = dydx[1] + 2 * dydx[2] * (0)
		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

		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