import os, sys import numpy as np import galsim import astropy.constants as cons from astropy.table import Table from scipy import interpolate from ObservationSim.MockObject.MockObject import MockObject from ObservationSim.MockObject.SpecDisperser import SpecDisperser from ObservationSim.MockObject._util import eObs, integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, getABMAG,convolveGaussXorders class Stamp(MockObject): def __init__(self, param, logger=None): super().__init__(param, logger=logger) def unload_SED(self): """(Test) free up SED memory """ del self.sed def drawObj_multiband(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0, exptime=150., fd_shear=None): if nphotons_tot == None: nphotons_tot = self.getElectronFluxFilt(filt, tel, exptime) try: full = integrate_sed_bandpass(sed=self.sed, bandpass=filt.bandpass_full) except Exception as e: print(e) self.logger.error(e) return False #nphotons_sum = 0 #photons_list = [] #xmax, ymax = 0, 0 if self.getMagFilter(filt) <= 15: folding_threshold = 5.e-4 else: folding_threshold = 5.e-3 gsp = galsim.GSParams(folding_threshold=folding_threshold) self.real_pos = self.getRealPos(chip.img, global_x=self.posImg.x, global_y=self.posImg.y, img_real_wcs=self.chip_wcs) x, y = self.real_pos.x + 0.5, self.real_pos.y + 0.5 x_nominal = int(np.floor(x + 0.5)) y_nominal = int(np.floor(y + 0.5)) dx = x - x_nominal dy = y - y_nominal offset = galsim.PositionD(dx, dy) chip_wcs_local = self.chip_wcs.local(self.real_pos) is_updated = 0 if fd_shear: g1 += fd_shear.g1 g2 += fd_shear.g2 gal_shear = galsim.Shear(g1=g1, g2=g2) for i in range(len(bandpass_list)): bandpass = bandpass_list[i] try: sub = integrate_sed_bandpass(sed=self.sed, bandpass=bandpass) except Exception as e: print(e) if self.logger: self.logger.error(e) continue ratio = sub/full if not (ratio == -1 or (ratio != ratio)): nphotons = ratio * nphotons_tot else: continue #nphotons_sum += nphotons psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass, folding_threshold=folding_threshold) _gal = self.param['image'] galImg= galsim.ImageF(_gal, scale=self.param['pixScale']) gal_temp= galsim.InterpolatedImage(galImg) gal_temp= gal_temp.shear(gal_shear) gal_temp= gal_temp.withFlux(nphotons) gal_temp= galsim.Convolve(psf, gal_temp) if i == 0: gal = gal_temp else: gal = gal + gal_temp stamp = gal.drawImage(wcs=chip_wcs_local, offset=offset) if np.sum(np.isnan(stamp.array)) > 0: # ERROR happens return 2, pos_shear stamp.setCenter(x_nominal, y_nominal) bounds = stamp.bounds & galsim.BoundsI(0, chip.npix_x - 1, 0, chip.npix_y - 1) if bounds.area() > 0: chip.img.setOrigin(0, 0) chip.img[bounds] += stamp[bounds] is_updated = 1 chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin) del stamp if is_updated == 0: print("fits obj %s missed"%(self.id)) if self.logger: self.logger.info("fits obj %s missed"%(self.id)) return 0, pos_shear return 1, pos_shear def drawObj_slitless(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0, exptime=150., normFilter=None, grating_split_pos=3685, fd_shear=None): if normFilter is not None: norm_thr_rang_ids = normFilter['SENSITIVITY'] > 0.001 sedNormFactor = getNormFactorForSpecWithABMAG(ABMag=self.param['mag_use_normal'], spectrum=self.sed, norm_thr=normFilter, sWave=np.floor(normFilter[norm_thr_rang_ids][0][0]), eWave=np.ceil(normFilter[norm_thr_rang_ids][-1][0])) if sedNormFactor == 0: return 2, None else: sedNormFactor = 1. normalSED = Table(np.array([self.sed['WAVELENGTH'], self.sed['FLUX'] * sedNormFactor]).T, names=('WAVELENGTH', 'FLUX')) self.real_pos = self.getRealPos(chip.img, global_x=self.posImg.x, global_y=self.posImg.y, img_real_wcs=self.chip_wcs) x, y = self.real_pos.x + 0.5, self.real_pos.y + 0.5 x_nominal = int(np.floor(x + 0.5)) y_nominal = int(np.floor(y + 0.5)) dx = x - x_nominal dy = y - y_nominal offset = galsim.PositionD(dx, dy) chip_wcs_local = self.chip_wcs.local(self.real_pos) if self.getMagFilter(filt) <= 15: folding_threshold = 5.e-4 else: folding_threshold = 5.e-3 gsp = galsim.GSParams(folding_threshold=folding_threshold) # nphotons_sum = 0 flat_cube = chip.flat_cube xOrderSigPlus = {'A':1.3909419820029296,'B':1.4760376591236062,'C':4.035447379743442,'D':5.5684364343742825,'E':16.260021029735388} grating_split_pos_chip = 0 + grating_split_pos branges = np.zeros([len(bandpass_list), 2]) # print(hasattr(psf_model, 'bandranges')) if hasattr(psf_model, 'bandranges'): if psf_model.bandranges is None: return 2, None if len(psf_model.bandranges) != len(bandpass_list): return 2, None branges = psf_model.bandranges else: for i in range(len(bandpass_list)): branges[i, 0] = bandpass_list[i].blue_limit * 10 branges[i, 1] = bandpass_list[i].red_limit * 10 for i in range(len(bandpass_list)): # bandpass = bandpass_list[i] brange = branges[i] # psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass, folding_threshold=folding_threshold) _gal = self.param['image'] galImg= galsim.ImageF(_gal, scale=self.param['pixScale']) gal = galsim.InterpolatedImage(galImg) # (TEST) Random knots # knots = galsim.RandomKnots(npoints=100, profile=disk) # kfrac = np.random.random()*(1.0 - self.bfrac) # gal = self.bfrac * bulge + (1.0 - self.bfrac - kfrac) * disk + kfrac * knots gal = gal.withFlux(tel.pupil_area * exptime) if fd_shear: g1 += fd_shear.g1 g2 += fd_shear.g2 gal_shear = galsim.Shear(g1=g1, g2=g2) gal = gal.shear(gal_shear) # gal = galsim.Convolve(psf, gal) # if not big_galaxy: # Not apply PSF for very big galaxy # gal = galsim.Convolve(psf, gal) # # if fd_shear is not None: # # gal = gal.shear(fd_shear) 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) 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] 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_p1 = galsim.Image(subImg_p1) star_p1.setOrigin(0, 0) 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, ycenter=ycenter_p1, origin=origin_p1, 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_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) subImg_p2 = starImg.array[:, subSlitPos+1:starImg.array.shape[1]] star_p2 = galsim.Image(subImg_p2) star_p2.setOrigin(0, 0) 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, ycenter=ycenter_p2, origin=origin_p2, 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_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, 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) del sdp elif grating_split_pos_chip>=gal_end[1]: sdp = SpecDisperser(orig_img=starImg, 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) del sdp # print(self.y_nominal, starImg.center.y, starImg.ymin) # del psf return 1, pos_shear