ObservationSim.py

from Config import ConfigDir, 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, FieldDistortion, PSFInterp
from _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 os, sys
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))
        
        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)