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 ObservationSim.Config import ChipOutput
from ObservationSim.Instrument import Telescope, Filter, FilterParam, FocalPlane, Chip
from ObservationSim.Instrument.Chip import Effects
from ObservationSim.Instrument.Chip import ChipUtils as chip_utils
from ObservationSim.Astrometry.Astrometry_util import on_orbit_obs_position
from ObservationSim.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):
        # 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

        # 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_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]
        else:
            ra_cen = pointing.ra
            dec_cen = pointing.dec

        # Prepare necessary chip properties for simulation
        chip = self.prepare_chip_for_exposure(chip, ra_cen, dec_cen, pointing)

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

        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

    def runExposure_MPI_PointingList(self, pointing_list, chips=None, use_mpi=False):
        if use_mpi:
            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 (defined in the overall config file)
                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 use_mpi:
                    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