import numpy as np import galsim import os, sys import astropy.constants as cons from astropy.table import Table from scipy import interpolate from ObservationSim.MockObject._util import eObs, integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, getABMAG,convolveGaussXorders from ObservationSim.MockObject.SpecDisperser import SpecDisperser from ObservationSim.MockObject.MockObject import MockObject # import tracemalloc class Galaxy(MockObject): def __init__(self, param, rotation=None, logger=None): super().__init__(param, logger=logger) # self.thetaR = self.param["theta"] # self.bfrac = self.param["bfrac"] # self.hlr_disk = self.param["hlr_disk"] # self.hlr_bulge = self.param["hlr_bulge"] # Extract ellipticity components # self.e_disk = galsim.Shear(g=self.param["ell_disk"], beta=self.thetaR*galsim.degrees) # self.e_bulge = galsim.Shear(g=self.param["ell_bulge"], beta=self.thetaR*galsim.degrees) # self.e_total = galsim.Shear(g=self.param["ell_tot"], beta=self.thetaR*galsim.degrees) # self.e1_disk, self.e2_disk = self.e_disk.g1, self.e_disk.g2 # self.e1_bulge, self.e2_bulge = self.e_bulge.g1, self.e_bulge.g2 # self.e1_total, self.e2_total = self.e_total.g1, self.e_total.g2 if rotation is not None: self.rotateEllipticity(rotation) if not hasattr(self, "disk_sersic_idx"): self.disk_sersic_idx = 1. if not hasattr(self, "bulge_sersic_idx"): self.bulge_sersic_idx = 4. def unload_SED(self): """(Test) free up SED memory """ del self.sed def getGSObj_multiband(self, tel, psf_list, bandpass_list, filt, nphotons_tot=None, g1=0, g2=0, exptime=150., fd_shear=None): if len(psf_list) != len(bandpass_list): raise ValueError("!!!The number of PSF profiles and the number of bandpasses must be equal.") objs = [] if nphotons_tot == None: nphotons_tot = self.getElectronFluxFilt(filt, tel, exptime) # print("nphotons_tot = ", nphotons_tot) try: full = integrate_sed_bandpass(sed=self.sed, bandpass=filt.bandpass_full) except Exception as e: print(e) if self.logger: self.logger.error(e) return -1 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) return -1 ratio = sub/full if not (ratio == -1 or (ratio != ratio)): nphotons = ratio * nphotons_tot else: return -1 psf = psf_list[i] disk = galsim.Sersic(n=self.disk_sersic_idx, half_light_radius=self.hlr_disk, flux=1.0) disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk) disk = disk.shear(disk_shape) bulge = galsim.Sersic(n=self.bulge_sersic_idx, half_light_radius=self.hlr_bulge, flux=1.0) bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge) bulge = bulge.shear(bulge_shape) gal = self.bfrac * bulge + (1.0 - self.bfrac) * disk gal = gal.withFlux(nphotons) if fd_shear is not None: 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 fd_shear is not None: # gal = gal.shear(fd_shear) objs.append(gal) final = galsim.Sum(objs) return final 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) # print("nphotons_tot = ", nphotons_tot) try: full = integrate_sed_bandpass(sed=self.sed, bandpass=filt.bandpass_full) except Exception as e: print(e) if self.logger: self.logger.error(e) return 2, None nphotons_sum = 0 photons_list = [] xmax, ymax = 0, 0 # # [C6 TEST] # print('hlr_disk = %.4f, hlr_bulge = %.4f'%(self.hlr_disk, self.hlr_bulge)) # tracemalloc.start() big_galaxy = False if self.hlr_disk > 3.0 or self.hlr_bulge > 3.0: # Very big galaxy big_galaxy = True # (TEST) Galsim Parameters if self.getMagFilter(filt) <= 15 and (not big_galaxy): 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.real_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) real_wcs_local = self.real_wcs.local(self.real_pos) disk = galsim.Sersic(n=self.disk_sersic_idx, half_light_radius=self.hlr_disk, flux=1.0, gsparams=gsp) ###disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk) ###disk = disk.shear(disk_shape) bulge = galsim.Sersic(n=self.bulge_sersic_idx, half_light_radius=self.hlr_bulge, flux=1.0, gsparams=gsp) ###bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge) ###bulge = bulge.shear(bulge_shape) 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) # return False continue ratio = sub/full if not (ratio == -1 or (ratio != ratio)): nphotons = ratio * nphotons_tot else: continue nphotons_sum += nphotons # # [C6 TEST] # print("nphotons_sub-band_%d = %.2f"%(i, nphotons)) psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass, folding_threshold=folding_threshold) # disk = galsim.Sersic(n=self.disk_sersic_idx, half_light_radius=self.hlr_disk, flux=1.0, gsparams=gsp) # disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk) # disk = disk.shear(disk_shape) # bulge = galsim.Sersic(n=self.bulge_sersic_idx, half_light_radius=self.hlr_bulge, flux=1.0, gsparams=gsp) # bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge) # bulge = bulge.shear(bulge_shape) gal_temp = self.bfrac * bulge + (1.0 - self.bfrac) * disk gal_temp = gal_temp.withFlux(nphotons) if not big_galaxy: # Not apply PSF for very big galaxy gal_temp = galsim.Convolve(psf, gal_temp) if i == 0: gal = gal_temp else: gal = gal + gal_temp # (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(nphotons) # gal_shear = galsim.Shear(g1=g1, g2=g2) # gal = gal.shear(gal_shear) # 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) # # Use (explicit) stamps to draw # stamp = gal.drawImage(wcs=real_wcs_local, method='phot', offset=offset, save_photons=True) # xmax = max(xmax, stamp.xmax - stamp.xmin) # ymax = max(ymax, stamp.ymax - stamp.ymin) # photons = stamp.photons # photons.x += x_nominal # photons.y += y_nominal # photons_list.append(photons) # del gal # # [C6 TEST] # print('xmax = %d, ymax = %d '%(xmax, ymax)) # # Output memory usage # snapshot = tracemalloc.take_snapshot() # top_stats = snapshot.statistics('lineno') # for stat in top_stats[:10]: # print(stat) if fd_shear: g1 += fd_shear.g1 g2 += fd_shear.g2 gal_shear = galsim.Shear(g1=g1, g2=g2) gal = gal.shear(gal_shear) # if fd_shear is not None: # gal = gal.shear(fd_shear) stamp = gal.drawImage(wcs=real_wcs_local, method='phot', offset=offset, save_photons=True) photons = stamp.photons photons.x += x_nominal photons.y += y_nominal photons_list.append(photons) # stamp = galsim.ImageF(int(xmax * 1.1), int(ymax * 1.1)) stamp.wcs = real_wcs_local 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) stamp[bounds] = chip.img[bounds] if not big_galaxy: for i in range(len(photons_list)): if i == 0: chip.sensor.accumulate(photons_list[i], stamp) else: chip.sensor.accumulate(photons_list[i], stamp, resume=True) else: sensor = galsim.Sensor() for i in range(len(photons_list)): if i == 0: sensor.accumulate(photons_list[i], stamp) else: sensor.accumulate(photons_list[i], stamp, resume=True) del sensor chip.img[bounds] = stamp[bounds] chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin) else: # Return code 0: object photons missed this detector print("obj %s missed"%(self.id)) if self.logger: self.logger.info("obj %s missed"%(self.id)) return 0, pos_shear # # [C6 TEST] # print("nphotons_sum = ", nphotons_sum) del photons_list del stamp 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.real_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) real_wcs_local = self.real_wcs.local(self.real_pos) big_galaxy = False if self.hlr_disk > 3.0 or self.hlr_bulge > 3.0: # Very big galaxy big_galaxy = True if self.getMagFilter(filt) <= 15 and (not big_galaxy): 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 for i in range(len(bandpass_list)): bandpass = bandpass_list[i] psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass, folding_threshold=folding_threshold) disk = galsim.Sersic(n=self.disk_sersic_idx, half_light_radius=self.hlr_disk, flux=1.0, gsparams=gsp) disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk) disk = disk.shear(disk_shape) bulge = galsim.Sersic(n=self.bulge_sersic_idx, half_light_radius=self.hlr_bulge, flux=1.0, gsparams=gsp) bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge) bulge = bulge.shear(bulge_shape) gal = self.bfrac * bulge + (1.0 - self.bfrac) * disk # (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=real_wcs_local, offset=offset) 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=bandpass.blue_limit * 10, band_end=bandpass.red_limit * 10, conf=chip.sls_conf[0], isAlongY=0, flat_cube=flat_cube) self.addSLStoChipImage(sdp=sdp_p1, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=real_wcs_local) 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=bandpass.blue_limit * 10, band_end=bandpass.red_limit * 10, conf=chip.sls_conf[1], isAlongY=0, flat_cube=flat_cube) self.addSLStoChipImage(sdp=sdp_p2, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=real_wcs_local) 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=bandpass.blue_limit * 10, band_end=bandpass.red_limit * 10, conf=chip.sls_conf[1], isAlongY=0, flat_cube=flat_cube) self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=real_wcs_local) 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=bandpass.blue_limit * 10, band_end=bandpass.red_limit * 10, conf=chip.sls_conf[0], isAlongY=0, flat_cube=flat_cube) self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=real_wcs_local) del sdp # print(self.y_nominal, starImg.center.y, starImg.ymin) del psf return 1, pos_shear def getGSObj(self, psf, g1=0, g2=0, flux=None, filt=None, tel=None, exptime=150.): if flux == None: flux = self.getElectronFluxFilt(filt, tel, exptime) disk = galsim.Sersic(n=self.disk_sersic_idx, half_light_radius=self.hlr_disk, flux=1.0) disk_shape = galsim.Shear(g1=self.e1_disk, g2=self.e2_disk) disk = disk.shear(disk_shape) bulge = galsim.Sersic(n=self.bulge_sersic_idx, half_light_radius=self.hlr_bulge, flux=1.0) bulge_shape = galsim.Shear(g1=self.e1_bulge, g2=self.e2_bulge) bulge = bulge.shear(bulge_shape) gal = self.bfrac * bulge + (1.0 - self.bfrac) * disk gal = gal.withFlux(flux) gal_shear = galsim.Shear(g1=g1, g2=g2) gal = gal.shear(gal_shear) final = galsim.Convolve(psf, gal) return final def rotateEllipticity(self, rotation): if rotation == 1: self.e1_disk, self.e2_disk, self.e1_bulge, self.e2_bulge, self.e1_total, self.e2_total = -self.e2_disk, self.e1_disk, -self.e2_bulge, self.e1_bulge, -self.e2_total, self.e1_total if rotation == 2: self.e1_disk, self.e2_disk, self.e1_bulge, self.e2_bulge, self.e1_total, self.e2_total = -self.e1_disk, -self.e2_disk, -self.e1_bulge, -self.e2_bulge, -self.e1_total, -self.e2_total if rotation == 3: self.e1_disk, self.e2_disk, self.e1_bulge, self.e2_bulge, self.e1_total, self.e2_total = self.e2_disk, -self.e1_disk, self.e2_bulge, -self.e1_bulge, self.e2_total, -self.e1_total def drawObject(self, img, final, noise_level=0.0, flux=None, filt=None, tel=None, exptime=150.): """ Override the method in parent class Need to constrain the size of image stamp for extended objects """ isUpdated = True if flux == None: flux = self.getElectronFluxFilt(filt, tel, exptime) stamp = final.drawImage(wcs=self.localWCS, offset=self.offset) stamp_arr = stamp.array mask = (stamp_arr >= 0.001*noise_level) # why 0.001? err = int(np.sqrt(mask.sum())) if np.mod(err, 2) == 1: err += 1 # if err == 1: if err == 0: subSize = 16 # why 16? else: subSize = max([err, 16]) fluxRatio = flux / stamp_arr[mask].sum() final = final.withScaledFlux(fluxRatio) imgSub = galsim.ImageF(subSize, subSize) # Draw with FFT # stamp = final.drawImage(image=imgSub, wcs=self.localWCS, offset=self.offset) # Draw with Photon Shoot stamp = final.drawImage(image=imgSub, wcs=self.localWCS, method='phot', offset=self.offset) stamp.setCenter(self.x_nominal, self.y_nominal) if np.sum(np.isnan(stamp.array)) >= 1: stamp.setZero() bounds = stamp.bounds & img.bounds if bounds.area() == 0: isUpdated = False else: img[bounds] += stamp[bounds] return img, stamp, isUpdated def getObservedEll(self, g1=0, g2=0): e1_obs, e2_obs, e_obs, theta = eObs(self.e1_total, self.e2_total, g1, g2) return self.e1_total, self.e2_total, g1, g2, e1_obs, e2_obs