Skip to content
GitLab
Commit 3515cc73 authored by Zhang Xin's avatar Zhang Xin
Browse files

pep8 formatting

parent 8fd9723d
Pipeline #7111 failed with stage
in 0 seconds
Show whitespace changes
Inline Side-by-side
observation_sim/mock_objects/SpecDisperser/SpecDisperser.py
View file @ 3515cc73
......@@ -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),
# 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),
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
......@@ -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
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment