import os import sys import numpy as np import galsim import astropy.constants as cons from astropy.table import Table from scipy import interpolate from observation_sim.mock_objects.MockObject import MockObject from observation_sim.mock_objects.SpecDisperser import SpecDisperser from observation_sim.mock_objects._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 is 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