import sys import os from ObservationSim.Config import ConfigDir, ChipOutput from ObservationSim.Config.Header import generatePrimaryHeader, generateExtensionHeader from ObservationSim.Instrument import Telescope, Filter, FilterParam, FocalPlane, Chip from ObservationSim.MockObject import Catalog, MockObject, Star, Galaxy, Quasar, calculateSkyMap_split_g from ObservationSim.PSF import PSFGauss, FieldDistortion, PSFInterp from ObservationSim._util import getShearFiled, makeSubDir_PointingList from astropy.io import fits from datetime import datetime import numpy as np import mpi4py.MPI as MPI import galsim import logging import psutil class Observation(object): def __init__(self, config, work_dir=None, data_dir=None): self.path_dict = ConfigDir(config=config, work_dir=work_dir, data_dir=data_dir) # self.config = ReadConfig(self.path_dict["config_file"]) self.config = config self.tel = Telescope(optEffCurve_path=self.path_dict["mirror_file"]) # Currently the default values are hard coded in self.focal_plane = FocalPlane(survey_type=self.config["obs_setting"]["survey_type"]) # Currently the default values are hard coded in self.filter_param = FilterParam(filter_dir=self.path_dict["filter_dir"]) # Currently the default values are hard coded in self.chip_list = [] self.filter_list = [] # if we want to apply field distortion? if self.config["ins_effects"]["field_dist"] == True: self.fd_model = FieldDistortion(fdModel_path=self.path_dict["fd_path"]) else: self.fd_model = None # Construct chips & filters: nchips = self.focal_plane.nchip_x*self.focal_plane.nchip_y for i in range(nchips): chipID = i + 1 if self.focal_plane.isIgnored(chipID=chipID): continue # Make Chip & Filter lists chip = Chip(chipID, ccdEffCurve_dir=self.path_dict["ccd_dir"], CRdata_dir=self.path_dict["CRdata_dir"], normalize_dir=self.path_dict["normalize_dir"], sls_dir=self.path_dict["sls_dir"], config=self.config) # currently there is no config file for chips filter_id, filter_type = chip.getChipFilter() filt = Filter(filter_id=filter_id, filter_type=filter_type, filter_param=self.filter_param, ccd_bandpass=chip.effCurve) self.chip_list.append(chip) self.filter_list.append(filt) # Read catalog and shear(s) self.g1_field, self.g2_field, self.nshear = getShearFiled(config=self.config) def runOneChip(self, chip, filt, chip_output, wcs_fp=None, psf_model=None, pointing_ID=0, ra_cen=None, dec_cen=None, img_rot=None, exptime=150., timestamp_obs=1621915200, pointing_type="MS", shear_cat_file=None, cat_dir=None, sed_dir=None): if (ra_cen is None) or (dec_cen is None): ra_cen = self.config["obs_setting"]["ra_center"] dec_cen = self.config["obs_setting"]["dec_center"] if img_rot is None: img_rot = self.config["obs_setting"]["image_rot"] if self.config["psf_setting"]["psf_model"] == "Gauss": psf_model = PSFGauss(chip=chip) elif self.config["psf_setting"]["psf_model"] == "Interp": psf_model = PSFInterp(chip=chip, PSF_data_file=self.path_dict["psf_dir"]) else: print("unrecognized PSF model type!!", flush=True) # Get (extra) shear fields if shear_cat_file is not None: self.g1_field, self.g2_field, self.nshear = getShearFiled(config=self.config, shear_cat_file=shear_cat_file) # Get WCS for the focal plane if wcs_fp == None: wcs_fp = self.focal_plane.getTanWCS(ra_cen, dec_cen, img_rot, chip.pix_scale) # Create chip Image chip.img = galsim.ImageF(chip.npix_x, chip.npix_y) chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin) chip.img.wcs = wcs_fp if chip.survey_type == "photometric": sky_map = None elif chip.survey_type == "spectroscopic": sky_map = calculateSkyMap_split_g(xLen=chip.npix_x, yLen=chip.npix_y, blueLimit=filt.blue_limit, redLimit=filt.red_limit, skyfn=self.path_dict["sky_file"], conf=chip.sls_conf, pixelSize=chip.pix_scale, isAlongY=0) if pointing_type == 'MS': # Load catalogues and templates self.cat = Catalog(config=self.config, chip=chip, cat_dir=cat_dir, sed_dir=sed_dir) self.nobj = len(self.cat.objs) # Loop over objects missed_obj = 0 bright_obj = 0 dim_obj = 0 for j in range(self.nobj): # if j >= 20: # break obj = self.cat.objs[j] # Load SED if obj.type == 'star': normF = chip.normF_star if self.config["galaxy_only"]: continue try: obj.load_SED( survey_type=chip.survey_type, normFilter=normF, target_filt=filt, sed_lib=self.cat.tempSED_star) except Exception as e: print(e) continue elif obj.type == 'galaxy': # or obj.type == quasar normF = chip.normF_galaxy if self.config["star_only"]: continue obj.load_SED( survey_type=chip.survey_type, sed_templates=self.cat.tempSed_gal, normFilter=normF, target_filt=filt) elif obj.type == 'quasar': normF = chip.normF_galaxy if self.config["star_only"]: continue obj.load_SED( survey_type=chip.survey_type, sed_templates=self.cat.tempSed_gal, normFilter=normF, target_filt=filt) # Exclude very bright/dim objects (for now) if filt.is_too_bright(mag=obj.getMagFilter(filt)): # print("obj too birght!!", flush=True) if obj.type != 'galaxy': bright_obj += 1 obj.unload_SED() continue if filt.is_too_dim(mag=obj.getMagFilter(filt)): # print("obj too dim!!", flush=True) dim_obj += 1 obj.unload_SED() # print(obj.getMagFilter(filt)) continue if self.config["shear_setting"]["shear_type"] == "constant": if obj.type == 'star': g1, g2 = 0, 0 else: g1, g2 = self.g1_field, self.g2_field elif self.config["shear_setting"]["shear_type"] == "extra": try: # TODO: every object with individual shear from input catalog(s) g1, g2 = self.g1_field[j], self.g2_field[j] except: print("failed to load external shear.") pass pos_img, offset, local_wcs = obj.getPosImg_Offset_WCS(img=chip.img, fdmodel=self.fd_model, chip=chip, verbose=False) if pos_img.x == -1 or pos_img.y == -1: # Exclude object which is outside the chip area (after field distortion) # print("obj missed!!") missed_obj += 1 obj.unload_SED() continue # Draw object & update output catalog try: if self.config["out_cat_only"]: isUpdated = True if chip.survey_type == "photometric" and not self.config["out_cat_only"]: isUpdated, pos_shear = obj.drawObj_multiband( tel=self.tel, pos_img=pos_img, psf_model=psf_model, bandpass_list=filt.bandpass_sub_list, filt=filt, chip=chip, g1=g1, g2=g2, exptime=exptime) elif chip.survey_type == "spectroscopic" and not self.config["out_cat_only"]: isUpdated, pos_shear = obj.drawObj_slitless( tel=self.tel, pos_img=pos_img, psf_model=psf_model, bandpass_list=filt.bandpass_sub_list, filt=filt, chip=chip, g1=g1, g2=g2, exptime=exptime, normFilter=normF) if isUpdated: # TODO: add up stats chip_output.cat_add_obj(obj, pos_img, pos_shear, g1, g2) pass else: # print("object omitted", flush=True) continue except Exception as e: print(e) pass # Unload SED: obj.unload_SED() del obj del psf_model del self.cat print("check running:1: pointing-{:} chip-{:} pid-{:} memory-{:6.2}GB".format(pointing_ID, chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) ), flush=True) # Detector Effects # =========================================================== # whether to output zero, dark, flat calibration images. chip.img = chip.addEffects( config=self.config, img=chip.img, chip_output=chip_output, filt=filt, ra_cen=ra_cen, dec_cen=dec_cen, img_rot=img_rot, pointing_ID=pointing_ID, timestamp_obs=timestamp_obs, pointing_type=pointing_type, sky_map=sky_map, tel = self.tel) if pointing_type == 'MS': datetime_obs = datetime.fromtimestamp(timestamp_obs) date_obs = datetime_obs.strftime("%y%m%d") time_obs = datetime_obs.strftime("%H%M%S") h_prim = generatePrimaryHeader( xlen=chip.npix_x, ylen=chip.npix_y, pointNum = str(pointing_ID), ra=ra_cen, dec=dec_cen, psize=chip.pix_scale, row_num=chip.rowID, col_num=chip.colID, # date=self.config["date_obs"], # time_obs=self.config["time_obs"], date=date_obs, time_obs=time_obs, im_type='MS') h_ext = generateExtensionHeader( xlen=chip.npix_x, ylen=chip.npix_y, ra=ra_cen, dec=dec_cen, pa=img_rot.deg, gain=chip.gain, readout=chip.read_noise, dark=chip.dark_noise, saturation=90000, psize=chip.pix_scale, row_num=chip.rowID, col_num=chip.colID, extName='raw') chip.img = galsim.Image(chip.img.array, dtype=np.uint16) # chip.img = galsim.Image(chip.img.array, dtype=np.uint32) hdu1 = fits.PrimaryHDU(header=h_prim) hdu2 = fits.ImageHDU(chip.img.array, header=h_ext) hdu1 = fits.HDUList([hdu1, hdu2]) fname = os.path.join(chip_output.subdir, h_prim['FILENAME'] + '.fits') hdu1.writeto(fname, output_verify='ignore', overwrite=True) print("# objects that are too bright %d out of %d"%(bright_obj, self.nobj)) print("# objects that are too dim %d out of %d"%(dim_obj, self.nobj)) print("# objects that are missed %d out of %d"%(missed_obj, self.nobj)) del chip.img print("check running:2: pointing-{:} chip-{:} pid-{:} memory-{:6.2}GB".format(pointing_ID, chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) ), flush=True) def runExposure(self, ra_cen=None, dec_cen=None, timestamp_obs=1621915200, pointing_ID=0, pointing_type='MS', img_rot=None, exptime=150., shear_cat_file=None, chips=None): if (ra_cen == None) or (dec_cen == None): ra_cen = self.config["obs_setting"]["ra_center"] dec_cen = self.config["obs_setting"]["dec_center"] if img_rot == None: img_rot = self.config["obs_setting"]["image_rot"] sub_img_dir, prefix = makeSubDir_PointingList(path_dict=self.path_dict, config=self.config, pointing_ID=pointing_ID) # Loop over chips for i in range(len(self.chip_list)): chip = self.chip_list[i] filt = self.filter_list[i] # Only run a particular set of chips if chips is not None: if chip.chipID not in chips: continue # Prepare output files chip_output = ChipOutput( config=self.config, focal_plane=self.focal_plane, chip=chip, filt=filt, exptime=exptime, pointing_type=pointing_type, pointing_ID=pointing_ID, subdir=sub_img_dir, prefix=prefix) self.runOneChip( chip=chip, filt=filt, chip_output=chip_output, pointing_ID = pointing_ID, ra_cen=ra_cen, dec_cen=dec_cen, img_rot=img_rot, exptime=exptime, timestamp_obs=timestamp_obs, pointing_type=pointing_type, cat_dir=self.path_dict["cat_dir"]) print("finished running chip#%d..."%(chip.chipID), flush=True) def runExposure_MPI_PointingList(self, ra_cen=None, dec_cen=None, pRange=None, timestamp_obs=np.array([1621915200]), pointing_type=np.array(['MS']), img_rot=None, exptime=150., shear_cat_file=None, chips=None, use_mpi=False): if use_mpi: comm = MPI.COMM_WORLD ind_thread = comm.Get_rank() num_thread = comm.Get_size() if chips is None: nchips_per_fp = len(self.chip_list) run_chips = self.chip_list run_filts = self.filter_list else: # Only run a particular set of chips run_chips = [] run_filts = [] nchips_per_fp = len(chips) for ichip in range(len(self.chip_list)): chip = self.chip_list[ichip] filt = self.filter_list[ichip] if chip.chipID in chips: run_chips.append(chip) run_filts.append(filt) # TEMP if len(timestamp_obs) == 1: timestamp_obs = np.tile(timestamp_obs, len(ra_cen)) pointing_type = np.tile(pointing_type, len(ra_cen)) if pRange is not None: timestamp_obs = timestamp_obs[pRange] pointing_type = pointing_type[pRange] ra_cen = ra_cen[pRange] dec_cen = dec_cen[pRange] # The Starting pointing ID if pRange is not None: pStart = pRange[0] else: pStart = 0 for ipoint in range(len(ra_cen)): for ichip in range(nchips_per_fp): i = ipoint*nchips_per_fp + ichip if pRange is None: pointing_ID = pStart + ipoint else: pointing_ID = pRange[ipoint] if use_mpi: if i % num_thread != ind_thread: continue pid = os.getpid() sub_img_dir, prefix = makeSubDir_PointingList(path_dict=self.path_dict, config=self.config, pointing_ID=pointing_ID) # chip = self.chip_list[ichip] # filt = self.filter_list[ichip] chip = run_chips[ichip] filt = run_filts[ichip] print("running pointing#%d, chip#%d, at PID#%d..."%(pointing_ID, chip.chipID, pid), flush=True) chip_output = ChipOutput( config=self.config, focal_plane=self.focal_plane, chip=chip, filt=filt, exptime=exptime, pointing_type=pointing_type[ipoint], pointing_ID=pointing_ID, subdir=sub_img_dir, prefix=prefix) self.runOneChip( chip=chip, filt=filt, chip_output=chip_output, pointing_ID = pointing_ID, ra_cen=ra_cen[ipoint], dec_cen=dec_cen[ipoint], img_rot=img_rot, exptime=exptime, timestamp_obs=timestamp_obs[ipoint], pointing_type=pointing_type[ipoint], cat_dir=self.path_dict["cat_dir"]) print("finished running chip#%d..."%(chip.chipID), flush=True)