ObservationSim.py

import os
import numpy as np
import mpi4py.MPI as MPI
import galsim
import psutil
import gc
from datetime import datetime

import traceback

from observation_sim.config import ChipOutput
from observation_sim.instruments import Telescope, Filter, FilterParam, FocalPlane, Chip
from observation_sim.instruments.chip import effects
from observation_sim.instruments.chip import chip_utils as chip_utils
from observation_sim.astrometry.Astrometry_util import on_orbit_obs_position
from observation_sim.sim_steps import SimSteps, SIM_STEP_TYPES


class Observation(object):
    def __init__(self, config, Catalog, work_dir=None, data_dir=None):
        self.config = config
        self.tel = Telescope()
        self.filter_param = FilterParam()
        self.Catalog = Catalog

    def prepare_chip_for_exposure(self, chip, ra_cen, dec_cen, pointing, wcs_fp=None, slsPSFOptim = False):
        # 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

        chip.slsPSFOptim = slsPSFOptim
        if chip.chipID in [1,2,3,4,5,10,21,26,27,28,29,30] and slsPSFOptim:
            chip.img_stack = {}
            for id1 in np.arange(2):
                gn = chip_utils.getChipSLSGratingID(chip.chipID)[id1]
                orders = {}
                # for id2 in ['-2','-1','0','1','2']:
                for id2 in ['0','1']:
                    o_n = "order"+id2
                    allbands = {}
                    for id3 in ['1','2','3','4']:
                        w_n = "w"+id3
                        allbands[w_n] = galsim.ImageF(chip.npix_x, chip.npix_y)
                        allbands[w_n].setOrigin(chip.bound.xmin, chip.bound.ymin)
                        allbands[w_n].wcs = wcs_fp
                    orders[o_n] = allbands      
                chip.img_stack[gn] = orders
        else:
            chip.img_stack = {}

        # Get random generators for this chip
        chip.rng_poisson, chip.poisson_noise = chip_utils.get_poisson(
            seed=int(self.config["random_seeds"]["seed_poisson"]) + pointing.id*30 + chip.chipID, sky_level=0.)

        # Get flat, shutter, and PRNU images
        chip.flat_img, _ = chip_utils.get_flat(
            img=chip.img, seed=int(self.config["random_seeds"]["seed_flat"]))
        if chip.chipID > 30:
            chip.shutter_img = np.ones_like(chip.img.array)
        else:
            chip.shutter_img = effects.ShutterEffectArr(
                chip.img, t_exp=pointing.exp_time, t_shutter=1.3, dist_bearing=735, dt=1E-3)
        chip.prnu_img = effects.PRNU_Img(xsize=chip.npix_x, ysize=chip.npix_y, sigma=0.01,
                                         seed=int(self.config["random_seeds"]["seed_prnu"]+chip.chipID))

        return chip

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

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

        # Apply astrometric simulation for pointing
        if self.config["obs_setting"]["enable_astrometric_model"]:
            dt = datetime.utcfromtimestamp(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="J2000",
                input_date_str=date_str,
                input_time_str=time_str
            )
            ra_cen, dec_cen = ra_cen[0], dec_cen[0]
            ra_offset, dec_offset = pointing.ra - ra_cen, pointing.dec - dec_cen
        else:
            ra_cen = pointing.ra
            dec_cen = pointing.dec
            ra_offset, dec_offset = 0., 0.
            
        slsPSFOpt = False
        # Prepare necessary chip properties for simulation
        chip = self.prepare_chip_for_exposure(chip, ra_cen, dec_cen, pointing, slsPSFOptim = slsPSFOpt)

        # Initialize SimSteps
        sim_steps = SimSteps(overall_config=self.config,
                             chip_output=chip_output,
                             all_filters=self.all_filters,
                             ra_offset=ra_offset,
                             dec_offset=dec_offset)

        for step in pointing.obs_param["call_sequence"]:
            if self.config["run_option"]["out_cat_only"]:
                if step != "scie_obs":
                    continue
            chip_output.Log_info("Starting simulation step: %s, calling function: %s" % (
                step, SIM_STEP_TYPES[step]))
            obs_param = pointing.obs_param["call_sequence"][step]
            step_name = SIM_STEP_TYPES[step]
            try:
                step_func = getattr(sim_steps, step_name)
                chip, filt, tel, pointing = step_func(
                    chip=chip,
                    filt=filt,
                    tel=self.tel,
                    pointing=pointing,
                    catalog=self.Catalog,
                    obs_param=obs_param)
                chip_output.Log_info("Finished simulation step: %s" % (step))
            except Exception as e:
                traceback.print_exc()
                chip_output.Log_error(e)
                chip_output.Log_error("Failed simulation on step: %s" % (step))
                break

        chip_output.Log_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)))
        del chip.img
        del chip.flat_img
        del chip.prnu_img
        del chip.shutter_img

    def runExposure_MPI_PointingList(self, pointing_list, chips=None):
        comm = MPI.COMM_WORLD
        ind_thread = comm.Get_rank()
        num_thread = comm.Get_size()

        process_counter = 0
        for ipoint in range(len(pointing_list)):
            # Construct chips & filters:
            pointing = pointing_list[ipoint]
            # pointing_ID = pointing.id
            pointing_ID = pointing.obs_id

            pointing.make_output_pointing_dir(
                overall_config=self.config, copy_obs_config=True)

            self.focal_plane = FocalPlane(
                chip_list=pointing.obs_param["run_chips"])
            # Make Chip & Filter lists
            self.chip_list = []
            self.filter_list = []
            self.all_filters = []
            for i in range(self.focal_plane.nchips):
                chipID = i + 1
                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_filters.append(filt)

            if chips is None:
                # Run all chips defined in configuration of this pointing
                run_chips = self.chip_list
                run_filts = self.filter_list
                nchips_per_fp = len(self.chip_list)
            else:
                # Only run a particular set of chips
                run_chips = []
                run_filts = []
                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)
                nchips_per_fp = len(chips)

            for ichip in range(nchips_per_fp):
                i_process = process_counter + ichip
                if i_process % num_thread != ind_thread:
                    continue
                pid = os.getpid()

                chip = run_chips[ichip]
                filt = run_filts[ichip]

                chip_output = ChipOutput(
                    config=self.config,
                    chip=chip,
                    filt=filt,
                    pointing=pointing
                )
                chip_output.Log_info("running pointing#%d, chip#%d, at PID#%d..." % (
                    int(pointing_ID), chip.chipID, pid))
                self.run_one_chip(
                    chip=chip,
                    filt=filt,
                    chip_output=chip_output,
                    pointing=pointing)
                chip_output.Log_info(
                    "finished running chip#%d..." % (chip.chipID))
                for handler in chip_output.logger.handlers[:]:
                    chip_output.logger.removeHandler(handler)
                gc.collect()
            process_counter += nchips_per_fp