An error occurred while loading the file. Please try again.
Galaxy.py 20.45 KiB
import numpy as np
import galsim
import gc
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)
    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.):
        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)
            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)
                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=1.0, 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=4.0, 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) gal_shear = galsim.Shear(g1=g1, g2=g2) gal = gal.shear(gal_shear) gal = galsim.Convolve(psf, gal) 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.): 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) self.logger.error(e) return False 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) 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) self.logger.error(e) # return False continue ratio = sub/full if not (ratio == -1 or (ratio != ratio)):
nphotons = ratio * nphotons_tot else: # return False 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=1.0, 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=4.0, 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(nphotons) gal_shear = galsim.Shear(g1=g1, g2=g2) gal = gal.shear(gal_shear) if self.hlr_disk < 10.0: # Not apply PSF for very big galaxy gal = galsim.Convolve(psf, gal) # Use (explicit) stamps to draw # stamp = gal.drawImage(wcs=self.localWCS, method='phot', offset=self.offset, save_photons=True) # xmax = max(xmax, stamp.xmax) # ymax = max(ymax, stamp.ymax) # photons = stamp.photons # photons.x += self.x_nominal # photons.y += self.y_nominal # photons_list.append(photons) 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) 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) # stamp = galsim.ImageF(int(xmax*1.1), int(ymax*1.1)) # stamp.wcs = self.localWCS # stamp.setCenter(self.x_nominal, self.y_nominal) # bounds = stamp.bounds & chip.img.bounds # 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) # # # print(stamp.array.sum()) # # chip.img[bounds] += stamp[bounds] # chip.img[bounds] = stamp[bounds] # # [C6 TEST] # print("nphotons_sum = ", nphotons_sum) del photons_list del stamp gc.collect() return True, 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): 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 False 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: # 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=1.0, 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=4.0, 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) gal_shear = galsim.Shear(g1=g1, g2=g2) gal = gal.shear(gal_shear) gal = galsim.Convolve(psf, gal) 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 True, 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=1.0, 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=4.0, 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