from Config import ConfigDir, ReadConfig, ChipOutput from Config.Header import generatePrimaryHeader, generateExtensionHeader from Instrument import Telescope, Filter, FilterParam, FocalPlane, Chip from MockObject import Catalog, MockObject, Star, Galaxy, Quasar, calculateSkyMap_split_g from PSF import PSFGauss, PSFInterp, FieldDistortion from _util import makeSubDir, getShearFiled, makeSubDir_PointingList from astropy.time import Time as asTime from astropy.io import fits import numpy as np import mpi4py.MPI as MPI import galsim import os, sys import logging import psutil class Observation(object): def __init__(self, input_cat_dir=None, work_dir=None, data_dir=None): self.path_dict = ConfigDir(input_cat_dir, work_dir, data_dir) self.config = ReadConfig(self.path_dict["config_file"]) 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["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["field_dist"].lower() == "y": self.fd_model = FieldDistortion() 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., input_cat_name=None, shear_cat_file=None, cat_dir=None, sed_dir=None): if (ra_cen is None) or (dec_cen is None): ra_cen = self.config["ra_center"] dec_cen = self.config["dec_center"] if img_rot is None: img_rot = self.config["image_rot"] if self.config["psf_model"] == "Gauss": psf_model = PSFGauss(chip=chip) elif self.config["psf_model"] == "Interp": psf_model = PSFInterp(chip=chip) 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": skyfile = os.path.join(self.path_dict["data_dir"], 'skybackground/sky_emiss_hubble_50_50_A.dat') sky_map = calculateSkyMap_split_g(xLen=chip.npix_x, yLen=chip.npix_y, blueLimit=filt.blue_limit, redLimit=filt.red_limit, skyfn=skyfile, conf=chip.sls_conf, pixelSize=chip.pix_scale, isAlongY=0) # Load catalogues and templates self.cat = Catalog(config=self.config, chip=chip, cat_dir=cat_dir, sed_dir=sed_dir, pRa=ra_cen, pDec=dec_cen, rotation=img_rot, template_dir=self.path_dict["template_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 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 obj.load_SED( sed_path=sed_dir, survey_type=chip.survey_type, sed_templates=self.cat.tempSed_gal, normFilter=normF, target_filt=filt) elif obj.type == 'quasar': normF = chip.normF_galaxy obj.load_SED( sed_path=sed_dir, 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_method"] == "constant": if obj.type == 'star': g1, g2 = 0, 0 else: g1, g2 = self.g1_field, self.g2_field elif self.config["shear_method"] == "extra": # TODO: every object with individual shear from input catalog(s) g1, g2 = self.g1_field[j], self.g2_field[j] 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 chip.survey_type == "photometric": 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": 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 # =========================================================== chip.img = chip.addNoise(config=self.config, tel=self.tel, filt=filt, img=chip.img, sky_map=sky_map) chip.img = chip.addEffects(config=self.config, img=chip.img, chip_output=chip_output, filt=filt, pointing_ID=pointing_ID) 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"]) 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) chip.img = galsim.Image(chip.img.array, dtype=np.uint16) 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) 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) 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)) def runExposure(self, ra_cen=None, dec_cen=None, pointing_ID=0, img_rot=None, exptime=150., input_cat_name=None, shear_cat_file=None, oneChip=None): if (ra_cen == None) or (dec_cen == None): ra_cen = self.config["ra_center"] dec_cen = self.config["dec_center"] if img_rot == None: img_rot = self.config["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] # Just run one chip if oneChip is not None: if chip.chipID != oneChip: continue # Prepare output files chip_output = ChipOutput( config=self.config, focal_plane=self.focal_plane, chip=chip, filt=filt, exptime=exptime, 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, cat_dir=self.path_dict["cat_dir"], sed_dir=self.path_dict["SED_dir"]) print("finished running chip#%d..."%(chip.chipID), flush=True) def runExposure_MPI_PointingList(self, ra_cen=None, dec_cen=None, pRange=None, img_rot=None, exptime=150., input_cat_name=None, shear_cat_file=None): comm = MPI.COMM_WORLD ind_thread = comm.Get_rank() num_thread = comm.Get_size() nchips_per_fp = len(self.chip_list) 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 pointing_ID = pStart + ipoint 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] 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_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, cat_dir=self.path_dict["cat_dir"], sed_dir=self.path_dict["SED_dir"]) print("finished running chip#%d..."%(chip.chipID), flush=True)