ObservationSim.py

import os
import numpy as np
import mpi4py.MPI as MPI
import galsim
import logging
import psutil
from astropy.io import fits
from datetime import datetime

import traceback

from ObservationSim.Config import config_dir, ChipOutput
from ObservationSim.Config.Header import generatePrimaryHeader, generateExtensionHeader
from ObservationSim.Instrument import Telescope, Filter, FilterParam, FocalPlane, Chip
from ObservationSim.Instrument.Chip import Effects
from ObservationSim.MockObject import calculateSkyMap_split_g
from ObservationSim.PSF import PSFGauss, FieldDistortion, PSFInterp
from ObservationSim._util import get_shear_field, makeSubDir_PointingList
from ObservationSim.Astrometry.Astrometry_util import on_orbit_obs_position

class Observation(object):
    def __init__(self, config, Catalog, work_dir=None, data_dir=None):
        self.path_dict = config_dir(config=config, work_dir=work_dir, data_dir=data_dir)
        self.config = config
        self.tel = Telescope()
        self.focal_plane = FocalPlane(survey_type=self.config["obs_setting"]["survey_type"]) 
        self.filter_param = FilterParam() 
        self.chip_list = []
        self.filter_list = []
        self.all_filter = []
        self.Catalog = Catalog

        # 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

            # Make Chip & Filter lists
            chip = Chip(
                chipID=chipID, 
                config=self.config)
            filter_id, filter_type = chip.getChipFilter()
            filt = Filter(filter_id=filter_id, 
                filter_type=filter_type, 
                filter_param=self.filter_param)
            if not self.focal_plane.isIgnored(chipID=chipID):
                self.chip_list.append(chip)
                self.filter_list.append(filt)
            self.all_filter.append(filt)

        # Read catalog and shear(s)
        self.g1_field, self.g2_field, self.nshear = get_shear_field(config=self.config)

    def run_one_chip(self, chip, filt, pointing, chip_output, wcs_fp=None, psf_model=None, shear_cat_file=None, cat_dir=None, sed_dir=None):

        # print(':::::::::::::::::::Current Pointing Information::::::::::::::::::')
        # print("RA: %f, DEC; %f" % (pointing.ra, pointing.dec))
        # print("Time: %s" % datetime.fromtimestamp(pointing.timestamp).isoformat())
        # print("Exposure time: %f" % pointing.exp_time)
        # print("Satellite Position (x, y, z): (%f, %f, %f)" % (pointing.sat_x, pointing.sat_y, pointing.sat_z))
        # print("Satellite Velocity (x, y, z): (%f, %f, %f)" % (pointing.sat_vx, pointing.sat_vy, pointing.sat_vz))
        # print("Position Angle: %f" % pointing.img_pa.deg)
        # print('Chip : %d' % chip.chipID)
        # print(':::::::::::::::::::::::::::END:::::::::::::::::::::::::::::::::::')
        chip_output.logger.info(':::::::::::::::::::Current Pointing Information::::::::::::::::::')
        chip_output.logger.info("RA: %f, DEC; %f" % (pointing.ra, pointing.dec))
        chip_output.logger.info("Time: %s" % datetime.fromtimestamp(pointing.timestamp).isoformat())
        chip_output.logger.info("Exposure time: %f" % pointing.exp_time)
        chip_output.logger.info("Satellite Position (x, y, z): (%f, %f, %f)" % (pointing.sat_x, pointing.sat_y, pointing.sat_z))
        chip_output.logger.info("Satellite Velocity (x, y, z): (%f, %f, %f)" % (pointing.sat_vx, pointing.sat_vy, pointing.sat_vz))
        chip_output.logger.info("Position Angle: %f" % pointing.img_pa.deg)
        chip_output.logger.info('Chip : %d' % chip.chipID)
        chip_output.logger.info(':::::::::::::::::::::::::::END:::::::::::::::::::::::::::::::::::')

        if self.config["psf_setting"]["psf_model"] == "Gauss":
            psf_model = PSFGauss(chip=chip, psfRa=self.config["psf_setting"]["psf_rcont"])
        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)
            chip_output.logger.error("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 = get_shear_field(config=self.config, shear_cat_file=shear_cat_file)

        # Apply astrometric simulation for pointing
        if self.config["obs_setting"]["enable_astrometric_model"]:
            dt = datetime.fromtimestamp(pointing.timestamp)
            date_str = dt.date().isoformat()
            time_str = dt.time().isoformat()
            ra_cen, dec_cen = on_orbit_obs_position(
                input_ra_list=[pointing.ra],
                input_dec_list=[pointing.dec],
                input_pmra_list=[0.],
                input_pmdec_list=[0.],
                input_rv_list=[0.],
                input_parallax_list=[1e-9],
                input_nstars=1,
                input_x=pointing.sat_x,
                input_y=pointing.sat_y,
                input_z=pointing.sat_z,
                input_vx=pointing.sat_vx,
                input_vy=pointing.sat_vy,
                input_vz=pointing.sat_vz,
                input_epoch="J2015.5",
                input_date_str=date_str,
                input_time_str=time_str
            )
            ra_cen, dec_cen = ra_cen[0], dec_cen[0]
        else:
            ra_cen = pointing.ra
            dec_cen = pointing.dec

        # Get WCS for the focal plane
        if wcs_fp == None:
            wcs_fp = self.focal_plane.getTanWCS(ra_cen, dec_cen, pointing.img_pa, 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)
        elif chip.survey_type == "spectroscopic":
            # chip.loadSLSFLATCUBE(flat_fn='flat_cube.fits')
            flat_normal = np.ones_like(chip.img.array)
            if self.config["ins_effects"]["flat_fielding"] == True:
                # print("SLS flat preprocess,CHIP %d : Creating and applying Flat-Fielding"%chip.chipID, flush=True)
                # print(chip.img.bounds, flush=True)
                chip_output.logger.info("SLS flat preprocess,CHIP %d : Creating and applying Flat-Fielding"%chip.chipID)
                msg = str(chip.img.bounds)
                chip_output.logger.info(msg)
                flat_img = Effects.MakeFlatSmooth(
                    chip.img.bounds,
                    int(self.config["random_seeds"]["seed_flat"]))
                flat_normal = flat_normal * flat_img.array / np.mean(flat_img.array)
            if self.config["ins_effects"]["shutter_effect"] == True:
                # print("SLS flat preprocess,CHIP %d : Apply shutter effect"%chip.chipID, flush=True)
                chip_output.logger.info("SLS flat preprocess,CHIP %d : Apply shutter effect"%chip.chipID)
                shuttimg = Effects.ShutterEffectArr(chip.img, t_shutter=1.3, dist_bearing=735,
                                                    dt=1E-3)  # shutter effect normalized image for this chip
                flat_normal = flat_normal*shuttimg
                flat_normal = np.array(flat_normal,dtype='float32')
            sky_map = calculateSkyMap_split_g(
                skyMap=flat_normal,
                blueLimit=filt.blue_limit,
                redLimit=filt.red_limit,
                conf=chip.sls_conf,
                pixelSize=chip.pix_scale,
                isAlongY=0,
                flat_cube=chip.flat_cube)
            # sky_map = np.ones([9216, 9232])
            del flat_normal

        if pointing.pointing_type == 'MS':
            # Load catalogues and templates
            self.cat = self.Catalog(config=self.config, chip=chip, pointing=pointing, cat_dir=cat_dir, sed_dir=sed_dir, chip_output=chip_output)
            chip_output.create_output_file()
            self.nobj = len(self.cat.objs)

            for ifilt in range(len(self.all_filter)):
                temp_filter = self.all_filter[ifilt]
                # Update the limiting magnitude using exposure time in pointing
                temp_filter.update_limit_saturation_mags(exptime=pointing.exp_time, chip=chip)

                # Select cutting band filter for saturation/limiting magnitude
                if temp_filter.filter_type.lower() == self.config["obs_setting"]["cut_in_band"].lower():
                    cut_filter = temp_filter

            # Loop over objects
            missed_obj = 0
            bright_obj = 0
            dim_obj = 0
            for j in range(self.nobj):
                
                # (DEBUG)
                # if j >= 10:
                #     break

                obj = self.cat.objs[j]
                if obj.type == 'star' and self.config["run_option"]["galaxy_only"]:
                    continue
                elif obj.type == 'galaxy' and self.config["run_option"]["star_only"]:
                    continue
                elif obj.type == 'quasar' and self.config["run_option"]["star_only"]:
                    continue

                # load and convert SED; also caculate object's magnitude in all CSST bands
                try:
                    sed_data = self.cat.load_sed(obj)
                    norm_filt = self.cat.load_norm_filt(obj)
                    obj.sed, obj.param["mag_%s"%filt.filter_type], obj.param["flux_%s"%filt.filter_type] = self.cat.convert_sed(
                        mag=obj.param["mag_use_normal"],
                        sed=sed_data,
                        target_filt=filt, 
                        norm_filt=norm_filt,
                    )
                    _, obj.param["mag_%s"%cut_filter.filter_type], obj.param["flux_%s"%cut_filter.filter_type] = self.cat.convert_sed(
                        mag=obj.param["mag_use_normal"],
                        sed=sed_data,
                        target_filt=cut_filter, 
                        norm_filt=norm_filt,
                    )

                except Exception as e:
                    # print(e)
                    traceback.print_exc()
                    chip_output.logger.error(e)
                    continue
                
                # chip_output.logger.info("debug point #1")

                # Exclude very bright/dim objects (for now)
                # if filt.is_too_bright(mag=obj.getMagFilter(filt)):
                # if filt.is_too_bright(mag=obj.mag_use_normal):
                if cut_filter.is_too_bright(
                    mag=obj.param["mag_%s"%self.config["obs_setting"]["cut_in_band"].lower()],
                    margin=self.config["obs_setting"]["mag_sat_margin"]):
                    # 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),
                    margin=self.config["obs_setting"]["mag_lim_margin"]):
                # if cut_filter.is_too_dim(mag=obj.param["mag_%s"%self.config["obs_setting"]["cut_in_band"].lower()]):
                    # print("obj too dim!!", flush=True)
                    dim_obj += 1
                    obj.unload_SED()
                    # print(obj.getMagFilter(filt))
                    continue

                # chip_output.logger.info("debug point #2")

                if self.config["shear_setting"]["shear_type"] == "constant":
                    if obj.type == 'star':
                        obj.g1, obj.g2 = 0., 0.
                    else:
                        obj.g1, obj.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)
                        obj.g1, obj.g2 = self.g1_field[j], self.g2_field[j]
                    except:
                        # print("failed to load external shear.")
                        chip_output.logger.error("failed to load external shear.")
                        pass

                    # chip_output.logger.info("debug point #3")
                elif self.config["shear_setting"]["shear_type"] == "catalog":
                    pass
                else:
                    chip_output.logger.error("Unknown shear input")
                    raise ValueError("Unknown shear input")
                
                # chip_output.logger.info("debug point #4")
                header_wcs = generateExtensionHeader(
                    xlen=chip.npix_x,
                    ylen=chip.npix_y,
                    ra=ra_cen,
                    dec=dec_cen,
                    pa=pointing.img_pa.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')

                pos_img, offset, local_wcs, real_wcs = obj.getPosImg_Offset_WCS(img=chip.img, fdmodel=self.fd_model, chip=chip, verbose=False, img_header=header_wcs)
                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

                # chip_output.logger.info("debug point #5")

                # Draw object & update output catalog
                try:
                    # chip_output.logger.info("debug point #6")
                    # chip_output.logger.info("current filter type: %s"%filt.filter_type)
                    if self.config["run_option"]["out_cat_only"]:
                        isUpdated = True
                        obj.real_pos = obj.getRealPos(chip.img, global_x=obj.posImg.x, global_y=obj.posImg.y,
                                        img_real_wcs=obj.real_wcs)
                        pos_shear = 0.
                    elif chip.survey_type == "photometric" and not self.config["run_option"]["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=obj.g1, 
                            g2=obj.g2, 
                            exptime=pointing.exp_time
                            )
                    elif chip.survey_type == "spectroscopic" and not self.config["run_option"]["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=obj.g1, 
                            g2=obj.g2, 
                            exptime=pointing.exp_time,
                            normFilter=norm_filt)
                    # chip_output.logger.info("debug point #7")
                    if isUpdated:
                        # TODO: add up stats
                        # print("updating output catalog...")
                        chip_output.cat_add_obj(obj, pos_img, pos_shear)
                        pass
                    else:
                        # print("object omitted", flush=True)
                        continue
                except Exception as e:
                    # print(e)
                    traceback.print_exc()
                    chip_output.logger.error(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)
        chip_output.logger.info("check running:1: pointing-%d chip-%d pid-%d memory-%6.2fGB"%(pointing.id, chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) ))

        # Detector Effects
        # ===========================================================
        # whether to output zero, dark, flat calibration images.

        if not self.config["run_option"]["out_cat_only"]:
            chip.img = chip.addEffects(
                config=self.config, 
                img=chip.img, 
                chip_output=chip_output, 
                filt=filt, 
                ra_cen=pointing.ra, 
                dec_cen=pointing.dec,
                img_rot=pointing.img_pa,
                pointing_ID=pointing.id,
                timestamp_obs=pointing.timestamp,
                pointing_type=pointing.pointing_type,
                sky_map=sky_map, tel = self.tel,
                logger=chip_output.logger)
            
            if pointing.pointing_type == 'MS':
                datetime_obs = datetime.fromtimestamp(pointing.timestamp)
                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=pointing.ra, 
                    dec=pointing.dec, 
                    psize=chip.pix_scale, 
                    row_num=chip.rowID, 
                    col_num=chip.colID,
                    date=date_obs,
                    time_obs=time_obs,
                    exptime=pointing.exp_time,
                    im_type='SCI',
                    sat_pos=[pointing.sat_x, pointing.sat_y, pointing.sat_z],
                    sat_vel=[pointing.sat_vx, pointing.sat_vy, pointing.sat_vz])
                h_ext = generateExtensionHeader(
                    xlen=chip.npix_x, 
                    ylen=chip.npix_y, 
                    ra=pointing.ra, 
                    dec=pointing.dec, 
                    pa=pointing.img_pa.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))
                chip_output.logger.info("# objects that are too bright %d out of %d"%(bright_obj, self.nobj))
                chip_output.logger.info("# objects that are too dim %d out of %d"%(dim_obj, self.nobj))
                chip_output.logger.info("# 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)
        chip_output.logger.info("check running:2: pointing-%d chip-%d pid-%d memory-%6.2fGB"%(pointing.id, chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) ))

    def runExposure_MPI_PointingList(self, pointing_list, 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)

        for ipoint in range(len(pointing_list)):
            for ichip in range(nchips_per_fp):
                i = ipoint*nchips_per_fp + ichip
                pointing = pointing_list[ipoint]
                pointing_ID = pointing.id
                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 = 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=pointing.exp_time,
                    pointing_type=pointing.pointing_type,
                    pointing_ID=pointing_ID,  
                    subdir=sub_img_dir,
                    prefix=prefix)
                chip_output.logger.info("running pointing#%d, chip#%d, at PID#%d..."%(pointing_ID, chip.chipID, pid))
                self.run_one_chip(
                    chip=chip, 
                    filt=filt, 
                    chip_output=chip_output, 
                    pointing=pointing,
                    cat_dir=self.path_dict["cat_dir"])
                print("finished running chip#%d..."%(chip.chipID), flush=True)
                chip_output.logger.info("finished running chip#%d..."%(chip.chipID))
                for handler in chip_output.logger.handlers[:]:
                    chip_output.logger.removeHandler(handler)