Skip to content
ObservationSim.py 16.1 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.io import fits
Fang Yuedong's avatar
Fang Yuedong committed
from datetime import datetime
Fang Yuedong's avatar
Fang Yuedong committed
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)


Fang Yuedong's avatar
Fang Yuedong committed
    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', input_cat_name=None, shear_cat_file=None, cat_dir=None, sed_dir=None):
Fang Yuedong's avatar
Fang Yuedong committed

        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)

Fang Yuedong's avatar
Fang Yuedong committed
        if pointing_type == 'MS':
            # 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
Fang Yuedong's avatar
Fang Yuedong committed
                    obj.load_SED(
Fang Yuedong's avatar
Fang Yuedong committed
                        sed_path=sed_dir, 
Fang Yuedong's avatar
Fang Yuedong committed
                        survey_type=chip.survey_type, 
Fang Yuedong's avatar
Fang Yuedong committed
                        sed_templates=self.cat.tempSed_gal, 
Fang Yuedong's avatar
Fang Yuedong committed
                        normFilter=normF, 
Fang Yuedong's avatar
Fang Yuedong committed
                        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
Fang Yuedong's avatar
Fang Yuedong committed
                    obj.unload_SED()
Fang Yuedong's avatar
Fang Yuedong committed
                    # print(obj.getMagFilter(filt))
Fang Yuedong's avatar
Fang Yuedong committed
                    continue

Fang Yuedong's avatar
Fang Yuedong committed
                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
Fang Yuedong's avatar
Fang Yuedong committed

Fang Yuedong's avatar
Fang Yuedong committed
                # 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)
Fang Yuedong's avatar
Fang Yuedong committed
                    pass
Fang Yuedong's avatar
Fang Yuedong committed
                # Unload SED:
                obj.unload_SED()
                del obj
Fang Yuedong's avatar
Fang Yuedong committed

Fang Yuedong's avatar
Fang Yuedong committed
            del psf_model
            del self.cat
Fang Yuedong's avatar
Fang Yuedong committed

        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)
Fang Yuedong's avatar
Fang Yuedong committed
        # chip.img = chip.addEffects(config=self.config, img=chip.img, chip_output=chip_output, filt=filt, pointing_ID=pointing_ID)

        # 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)
Fang Yuedong's avatar
Fang Yuedong committed
        
Fang Yuedong's avatar
Fang Yuedong committed
        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)
            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))
Fang Yuedong's avatar
Fang Yuedong committed
        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)


Fang Yuedong's avatar
Fang Yuedong committed
    def runExposure(self, ra_cen=None, dec_cen=None, timestamp_obs=1621915200, pointing_ID=0, pointing_type='MS', img_rot=None, exptime=150., input_cat_name=None, shear_cat_file=None, oneChip=None):
Fang Yuedong's avatar
Fang Yuedong committed

        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,
Fang Yuedong's avatar
Fang Yuedong committed
                pointing_type=pointing_type,
Fang Yuedong's avatar
Fang Yuedong committed
                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, 
Fang Yuedong's avatar
Fang Yuedong committed
                timestamp_obs=timestamp_obs,
                pointing_type=pointing_type,
Fang Yuedong's avatar
Fang Yuedong committed
                cat_dir=self.path_dict["cat_dir"], 
                sed_dir=self.path_dict["SED_dir"])
            print("finished running chip#%d..."%(chip.chipID), flush=True)

Fang Yuedong's avatar
Fang Yuedong committed
    def runExposure_MPI_PointingList(self, ra_cen=None, dec_cen=None, pRange=None, timestamp_obs=[1621915200], pointing_type=['MS'], img_rot=None, exptime=150., input_cat_name=None, shear_cat_file=None):
Fang Yuedong's avatar
Fang Yuedong committed
        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

Fang Yuedong's avatar
Fang Yuedong committed
        # TEMP
        if len(timestamp_obs) == 1:
            timestamp_obs *= len(pRange)
            pointing_type *= len(pRange)

Fang Yuedong's avatar
Fang Yuedong committed
        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,
Fang Yuedong's avatar
Fang Yuedong committed
                    pointing_type=pointing_type[ipoint],
Fang Yuedong's avatar
Fang Yuedong committed
                    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, 
Fang Yuedong's avatar
Fang Yuedong committed
                    exptime=exptime,
                    timestamp_obs=timestamp_obs[ipoint],
                    pointing_type=pointing_type[ipoint],
Fang Yuedong's avatar
Fang Yuedong committed
                    cat_dir=self.path_dict["cat_dir"], 
                    sed_dir=self.path_dict["SED_dir"])
                print("finished running chip#%d..."%(chip.chipID), flush=True)