Skip to content
ObservationSim.py 14.2 KiB
Newer Older
Fang Yuedong's avatar
Fang Yuedong committed
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)