Merge branch 'develop' into milky_way_extinction (92db78dd) · Commits · csst-sims / csst_msc_sim

observation_sim/ObservationSim.py

+24 −3

Original line number	Diff line number	Diff line
		@@ -23,7 +23,7 @@ class Observation(object):
		self.filter_param = FilterParam()
		self.Catalog = Catalog

		def prepare_chip_for_exposure(self, chip, ra_cen, dec_cen, pointing, wcs_fp=None):
		def prepare_chip_for_exposure(self, chip, ra_cen, dec_cen, pointing, wcs_fp=None, slsPSFOptim = False):
		# Get WCS for the focal plane
		if wcs_fp == None:
		wcs_fp = self.focal_plane.getTanWCS(
		@@ -34,6 +34,26 @@ class Observation(object):
		chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin)
		chip.img.wcs = wcs_fp

		chip.slsPSFOptim = slsPSFOptim
		if chip.chipID in [1,2,3,4,5,10,21,26,27,28,29,30] and slsPSFOptim:
		chip.img_stack = {}
		for id1 in np.arange(2):
		gn = chip_utils.getChipSLSGratingID(chip.chipID)[id1]
		orders = {}
		# for id2 in ['-2','-1','0','1','2']:
		for id2 in ['0','1']:
		o_n = "order"+id2
		allbands = {}
		for id3 in ['1','2','3','4']:
		w_n = "w"+id3
		allbands[w_n] = galsim.ImageF(chip.npix_x, chip.npix_y)
		allbands[w_n].setOrigin(chip.bound.xmin, chip.bound.ymin)
		allbands[w_n].wcs = wcs_fp
		orders[o_n] = allbands
		chip.img_stack[gn] = orders
		else:
		chip.img_stack = {}

		# Get random generators for this chip
		chip.rng_poisson, chip.poisson_noise = chip_utils.get_poisson(
		seed=int(self.config["random_seeds"]["seed_poisson"]) + pointing.id*30 + chip.chipID, sky_level=0.)
		@@ -98,8 +118,9 @@ class Observation(object):
		dec_cen = pointing.dec
		ra_offset, dec_offset = 0., 0.

		slsPSFOpt = False
		# Prepare necessary chip properties for simulation
		chip = self.prepare_chip_for_exposure(chip, ra_cen, dec_cen, pointing)
		chip = self.prepare_chip_for_exposure(chip, ra_cen, dec_cen, pointing, slsPSFOptim = slsPSFOpt)

		# Initialize SimSteps
		sim_steps = SimSteps(overall_config=self.config,

observation_sim/mock_objects/Galaxy.py

+88 −39

Original line number	Diff line number	Diff line
		@@ -323,28 +323,73 @@ class Galaxy(MockObject):
		# # if fd_shear is not None:
		# # gal = gal.shear(fd_shear)

		starImg = gal.drawImage(
		wcs=chip_wcs_local, offset=offset, method='real_space')
		galImg_List = []
		try:
		pos_img_local = [0,0]
		x_start = chip.x_cen/chip.pix_size - chip.npix_x / 2.
		y_start = chip.y_cen/chip.pix_size - chip.npix_y / 2.
		pos_img_local[0] = pos_img.x - x_start
		pos_img_local[1] = pos_img.y - y_start
		nnx = 0
		nny = 0
		for order in ["A","B"]:
		psf, pos_shear = psf_model.get_PSF(
		chip, pos_img_local=pos_img_local, bandNo=i+1, galsimGSObject=True, g_order=order, grating_split_pos=grating_split_pos)
		star_p = galsim.Convolve(psf, gal)
		if nnx == 0:
		galImg = star_p.drawImage(wcs=chip_wcs_local, offset=offset)
		nnx = galImg.xmax - galImg.xmin + 1
		nny = galImg.ymax - galImg.ymin + 1
		else:
		galImg = star_p.drawImage(nx = nnx, ny = nny, wcs=chip_wcs_local, offset=offset)
		galImg.setOrigin(0, 0)
		# n1 = np.sum(np.isinf(galImg.array))
		# n2 = np.sum(np.isnan(galImg.array))
		# if n1>0 or n2 > 0:
		# print("DEBUG: Galaxy, inf:%d, nan:%d"%(n1, n2))
		if np.sum(np.isnan(galImg.array)) > 0:
		# ERROR happens
		return 2, pos_shear
		galImg_List.append(galImg)
		for order in ["C","D","E"]:
		galImg_List.append(galImg)
		except:
		psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img)
		star_p = galsim.Convolve(psf, gal)
		galImg = star_p.drawImage(wcs=chip_wcs_local, offset=offset)
		galImg.setOrigin(0, 0)
		if np.sum(np.isnan(galImg.array)) > 0:
		# ERROR happens
		return 2, pos_shear
		for order in ["A","B","C","D","E"]:
		galImg_List.append(galImg)

		# starImg = gal.drawImage(
		# wcs=chip_wcs_local, offset=offset, method='real_space')

		origin_star = [y_nominal - (starImg.center.y - starImg.ymin),
		x_nominal - (starImg.center.x - starImg.xmin)]
		starImg.setOrigin(0, 0)
		origin_star = [y_nominal - (galImg.center.y - galImg.ymin),
		x_nominal - (galImg.center.x - galImg.xmin)]
		galImg.setOrigin(0, 0)
		gal_origin = [origin_star[0], origin_star[1]]
		gal_end = [origin_star[0] + starImg.array.shape[0] -
		1, origin_star[1] + starImg.array.shape[1] - 1]
		gal_end = [origin_star[0] + galImg.array.shape[0] -
		1, origin_star[1] + galImg.array.shape[1] - 1]

		if gal_origin[1] < grating_split_pos_chip < gal_end[1]:
		subSlitPos = int(grating_split_pos_chip - gal_origin[1] + 1)
		# part img disperse

		subImg_p1 = starImg.array[:, 0:subSlitPos]
		star_p1s=[]
		for galImg in galImg_List:

		subImg_p1 = galImg.array[:, 0:subSlitPos]
		star_p1 = galsim.Image(subImg_p1)
		star_p1.setOrigin(0, 0)
		star_p1s.append(star_p1)
		origin_p1 = origin_star
		xcenter_p1 = min(x_nominal, grating_split_pos_chip-1) - 0
		ycenter_p1 = y_nominal-0

		sdp_p1 = SpecDisperser(orig_img=star_p1, xcenter=xcenter_p1,
		sdp_p1 = SpecDisperser(orig_img=star_p1s, xcenter=xcenter_p1,
		ycenter=ycenter_p1, origin=origin_p1,
		tar_spec=normalSED,
		band_start=brange[0], band_end=brange[1],
		@@ -352,21 +397,25 @@ class Galaxy(MockObject):
		isAlongY=0,
		flat_cube=flat_cube)

		# self.addSLStoChipImage(sdp=sdp_p1, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=chip_wcs_local)
		pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp_p1, chip=chip, pos_img_local=[xcenter_p1, ycenter_p1],
		psf_model=psf_model, bandNo=i + 1,
		grating_split_pos=grating_split_pos,
		local_wcs=chip_wcs_local, pos_img=pos_img)
		self.addSLStoChipImage(sdp=sdp_p1, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=chip_wcs_local)
		# pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp_p1, chip=chip, pos_img_local=[xcenter_p1, ycenter_p1],
		# psf_model=psf_model, bandNo=i + 1,
		# grating_split_pos=grating_split_pos,
		# local_wcs=chip_wcs_local, pos_img=pos_img)

		star_p2s=[]
		for galImg in galImg_List:

		subImg_p2 = starImg.array[:,
		subSlitPos+1:starImg.array.shape[1]]
		subImg_p2 = galImg.array[:,
		subSlitPos + 1:galImg.array.shape[1]]
		star_p2 = galsim.Image(subImg_p2)
		star_p2.setOrigin(0, 0)
		star_p2s.append(star_p2)
		origin_p2 = [origin_star[0], grating_split_pos_chip]
		xcenter_p2 = max(x_nominal, grating_split_pos_chip - 1) - 0
		ycenter_p2 = y_nominal - 0

		sdp_p2 = SpecDisperser(orig_img=star_p2, xcenter=xcenter_p2,
		sdp_p2 = SpecDisperser(orig_img=star_p2s, xcenter=xcenter_p2,
		ycenter=ycenter_p2, origin=origin_p2,
		tar_spec=normalSED,
		band_start=brange[0], band_end=brange[1],
		@@ -374,41 +423,41 @@ class Galaxy(MockObject):
		isAlongY=0,
		flat_cube=flat_cube)

		# self.addSLStoChipImage(sdp=sdp_p2, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=chip_wcs_local)
		pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp_p2, chip=chip, pos_img_local=[xcenter_p2, ycenter_p2],
		psf_model=psf_model, bandNo=i + 1,
		grating_split_pos=grating_split_pos,
		local_wcs=chip_wcs_local, pos_img=pos_img)
		self.addSLStoChipImage(sdp=sdp_p2, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=chip_wcs_local)
		# pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp_p2, chip=chip, pos_img_local=[xcenter_p2, ycenter_p2],
		# psf_model=psf_model, bandNo=i + 1,
		# grating_split_pos=grating_split_pos,
		# local_wcs=chip_wcs_local, pos_img=pos_img)

		del sdp_p1
		del sdp_p2
		elif grating_split_pos_chip <= gal_origin[1]:
		sdp = SpecDisperser(orig_img=starImg, xcenter=x_nominal - 0,
		sdp = SpecDisperser(orig_img=galImg_List, xcenter=x_nominal - 0,
		ycenter=y_nominal - 0, origin=origin_star,
		tar_spec=normalSED,
		band_start=brange[0], band_end=brange[1],
		conf=chip.sls_conf[1],
		isAlongY=0,
		flat_cube=flat_cube)
		# self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=chip_wcs_local)
		pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp, chip=chip, pos_img_local=[x_nominal, y_nominal],
		psf_model=psf_model, bandNo=i + 1,
		grating_split_pos=grating_split_pos,
		local_wcs=chip_wcs_local, pos_img=pos_img)
		self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=chip_wcs_local)
		# pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp, chip=chip, pos_img_local=[x_nominal, y_nominal],
		# psf_model=psf_model, bandNo=i + 1,
		# grating_split_pos=grating_split_pos,
		# local_wcs=chip_wcs_local, pos_img=pos_img)
		del sdp
		elif grating_split_pos_chip >= gal_end[1]:
		sdp = SpecDisperser(orig_img=starImg, xcenter=x_nominal - 0,
		sdp = SpecDisperser(orig_img=galImg_List, xcenter=x_nominal - 0,
		ycenter=y_nominal - 0, origin=origin_star,
		tar_spec=normalSED,
		band_start=brange[0], band_end=brange[1],
		conf=chip.sls_conf[0],
		isAlongY=0,
		flat_cube=flat_cube)
		# self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=chip_wcs_local)
		pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp, chip=chip, pos_img_local=[x_nominal, y_nominal],
		psf_model=psf_model, bandNo=i + 1,
		grating_split_pos=grating_split_pos,
		local_wcs=chip_wcs_local, pos_img=pos_img)
		self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=chip_wcs_local)
		# pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp, chip=chip, pos_img_local=[x_nominal, y_nominal],
		# psf_model=psf_model, bandNo=i + 1,
		# grating_split_pos=grating_split_pos,
		# local_wcs=chip_wcs_local, pos_img=pos_img)
		del sdp

		# print(self.y_nominal, starImg.center.y, starImg.ymin)

observation_sim/mock_objects/MockObject.py

+223 −112

File changed.

Preview size limit exceeded, changes collapsed.

observation_sim/mock_objects/SpecDisperser/SpecDisperser.py

+55 −11

Original line number	Diff line number	Diff line
		@@ -65,10 +65,30 @@ class SpecDisperser(object):
		# self.img_x = orig_img.shape[1]
		# self.img_y = orig_img.shape[0]

		self.thumb_img = np.abs(orig_img.array)
		self.thumb_x = orig_img.center.x
		self.thumb_y = orig_img.center.y
		self.img_sh = orig_img.array.shape
		# 5 orders, A, B ,
		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, 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)
		self.thumb_x = orig_img_one.center.x
		self.thumb_y = orig_img_one.center.y
		self.img_sh = orig_img_one.array.shape

		self.id = gid

		@@ -78,10 +98,13 @@ class SpecDisperser(object):
		self.isAlongY = isAlongY
		self.flat_cube = flat_cube
		if self.isAlongY == 1:
		self.thumb_img, self.thumb_x, self.thumb_y = rotate90(array_orig=self.thumb_img, xc=orig_img.center.x,
		yc=orig_img.center.y, isClockwise=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.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)

		self.img_sh = orig_img.array.T.shape
		self.img_sh = self.orig_img_orders[order].array.T.shape
		self.xcenter = ycenter
		self.ycenter = xcenter

		@@ -111,10 +134,16 @@ class SpecDisperser(object):

		def compute_spec(self, beam):

		# if beam == "B":
		# return self.thumb_img, self.origin[1], self.origin[0], None, None, None

		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
		self.thumb_y = self.orig_img_orders[beam].center.y
		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),
		@@ -169,7 +198,8 @@ class SpecDisperser(object):

		dyc = cast[int](np.floor(ytrace_beam+0.5))

		dypix = cast[int](np.floor(ytrace_beam - dyc[0] + x0[0] + 0.5))
		# dypix = cast[int](np.floor(ytrace_beam - dyc[0] + x0[0] + 0.5))
		dypix = dyc - dyc[0] + x0[0]

		frac_ids = yfrac_beam < 0

		@@ -248,7 +278,8 @@ class SpecDisperser(object):
		# 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],
		flat_index[nonz],
		yfrac_beam[nonz],
		sensitivity_beam[nonz],
		modelf, x0,
		array(self.img_sh,
		@@ -258,11 +289,24 @@ class SpecDisperser(object):
		lam_beam[lam_index][nonz])

		model = modelf.reshape(beam_sh)
		# n1 = np.sum(np.isinf(model))
		# n2 = np.sum(np.isnan(model))
		# n3 = np.sum(np.isinf(modelf))
		# n4 = np.sum(np.isnan(modelf))
		# 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))
		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):

observation_sim/mock_objects/SpecDisperser/disperse_c/disperse.pyx

+44 −5

Original line number	Diff line number	Diff line
		@@ -21,6 +21,29 @@ cdef extern from "math.h":
		double sqrt(double x)
		double exp(double x)


		def check_nan2D(np.ndarray[FTYPE_t, ndim=2] arr):
		cdef int i, j
		cdef int nrows = arr.shape[0]
		cdef int ncols = arr.shape[1]

		# 遍历数组的每个元素并检查是否存在 NaN
		for i in range(nrows):
		for j in range(ncols):
		if np.isnan(arr[i, j]) \| np.isinf(arr[i, j]):
		return True
		return False

		def check_nan1d(np.ndarray[DTYPE_t, ndim=1] arr):
		cdef int i
		cdef int n = arr.shape[0]

		# 遍历数组的每个元素并检查是否存在 NaN
		for i in range(n):
		if np.isnan(arr[i]) \| np.isinf(arr[i]):
		return True
		return False

		@cython.boundscheck(False)
		@cython.wraparound(False)
		@cython.embedsignature(True)
		@@ -54,6 +77,18 @@ def disperse_grism_object(np.ndarray[FTYPE_t, ndim=2] flam,
		nk = len(idxl)
		nl = len(full)


		#if check_nan2D(flam):
		# print("DEBUG: disperse, input Array 'flam' contains NaN.")
		#if check_nan1d(ysens):
		# print("DEBUG: disperse, input Array 'ysens' contains NaN.")
		#if check_nan1d(yfrac):
		# print("DEBUG: disperse, input Array 'yfrac' contains NaN.")
		#if check_nan1d(full):
		# print("DEBUG: disperse, input Array 'full' contains NaN before processing.")



		if (flat is not None):
		nlamb = len(wlambda)
		nflat = len(flat)
		@@ -95,14 +130,15 @@ def disperse_grism_object(np.ndarray[FTYPE_t, ndim=2] flam,

		else:
		for i in range(0-x0[1], x0[1]):
		if (x0[1]+i < 0) \| (x0[1]+i >= shd[1]):
		x_pos = x0[1]+i
		if (x_pos < 0) \| (x_pos >= shd[1]):
		continue

		for j in range(0-x0[0], x0[0]):
		if (x0[0]+j < 0) \| (x0[0]+j >= shd[0]):
		y_pos = x0[0]+j
		if (y_pos < 0) \| (y_pos >= shd[0]):
		continue

		fl_ij = flam[x0[0]+j, x0[1]+i] #/1.e-17
		fl_ij = flam[y_pos, x_pos] #/1.e-17
		if (fl_ij == 0):
		continue

		@@ -110,11 +146,14 @@ def disperse_grism_object(np.ndarray[FTYPE_t, ndim=2] flam,
		k1 = idxl[k]+j*shg[1]+i
		if (k1 >= 0) & (k1 < nl):
		full[k1] += ysens[k]fl_ij(1-yfrac[k])

		k2 = idxl[k]+(j+1)*shg[1]+i
		if (k2 >= 0) & (k2 < nl):
		full[k2] += ysens[k]fl_ijyfrac[k]

		#if (check_nan1d(full)):
		# print("DEBUG: disperse, output Array 'full' contains NaN after processing.+++++++++++++++++++++++++++")


		return True

		@cython.boundscheck(False)