Skip to content
ObservationSim.py 16.5 KiB
Newer Older
Fang Yuedong's avatar
Fang Yuedong committed
import numpy as np
import mpi4py.MPI as MPI
import galsim
import psutil
Fang Yuedong's avatar
Fang Yuedong committed
import gc
from datetime import datetime

Fang Yuedong's avatar
Fang Yuedong committed
import traceback

from ObservationSim.Config import config_dir, ChipOutput
from ObservationSim.Instrument import Telescope, Filter, FilterParam, FocalPlane, Chip
from ObservationSim.Instrument.Chip import Effects
Fang Yuedong's avatar
Fang Yuedong committed
from ObservationSim.Instrument.Chip import ChipUtils as chip_utils
from ObservationSim._util import makeSubDir_PointingList
from ObservationSim.Astrometry.Astrometry_util import on_orbit_obs_position
from ObservationSim.sim_steps import SimSteps, SIM_STEP_TYPES
Fang Yuedong's avatar
Fang Yuedong committed

Fang Yuedong's avatar
Fang Yuedong committed
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)
Fang Yuedong's avatar
Fang Yuedong committed
        self.config = config
        self.tel = Telescope()
        self.filter_param = FilterParam() 
        self.Catalog = Catalog
Fang Yuedong's avatar
Fang Yuedong committed

    def prepare_chip_for_exposure(self, chip, ra_cen, dec_cen, pointing, wcs_fp=None):
Fang Yuedong's avatar
Fang Yuedong committed
        # 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)
Fang Yuedong's avatar
Fang Yuedong committed

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

    def run_one_chip(self, chip, filt, pointing, chip_output, wcs_fp=None, psf_model=None, cat_dir=None, sed_dir=None):
Fang Yuedong's avatar
Fang Yuedong committed

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

        # Apply astrometric simulation for pointing
        if self.config["obs_setting"]["enable_astrometric_model"]:
Fang Yuedong's avatar
Fang Yuedong committed
            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,
Fang Yuedong's avatar
Fang Yuedong committed
                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
Fang Yuedong's avatar
Fang Yuedong committed
        chip = self.prepare_chip_for_exposure(chip, ra_cen, dec_cen, pointing)
Fang Yuedong's avatar
Fang Yuedong committed
        # Load catalogues
        # self.cat = self.Catalog(config=self.config, chip=chip, pointing=pointing, cat_dir=cat_dir, sed_dir=sed_dir, chip_output=chip_output, filt=filt)
Fang Yuedong's avatar
Fang Yuedong committed
        # Initialize SimSteps
        sim_steps = SimSteps(overall_config=self.config, chip_output=chip_output, all_filters=self.all_filters)
Fang Yuedong's avatar
Fang Yuedong committed
        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]))
Fang Yuedong's avatar
Fang Yuedong committed
            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, 
Fang Yuedong's avatar
Fang Yuedong committed
                    pointing=pointing, 
                    catalog=self.Catalog, 
Fang Yuedong's avatar
Fang Yuedong committed
                    obs_param=obs_param)
                chip_output.Log_info("Finished simulation step: %s"%(step))
Fang Yuedong's avatar
Fang Yuedong committed
            except Exception as e:
                traceback.print_exc()
                chip_output.Log_error(e)
                chip_output.Log_error("Failed simulation on step: %s"%(step))
                break
Fang Yuedong's avatar
Fang Yuedong committed

Fang Yuedong's avatar
Fang Yuedong committed
        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) ))
Zhang Xin's avatar
Zhang Xin committed
        del chip.img

    # def run_one_chip_calibration(self, chip, filt, pointing, chip_output, skyback_level = 20000, sky_level_filt = 'g', wcs_fp=None, psf_model=None, cat_dir=None, sed_dir=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
    #     pf_map = np.zeros_like(chip.img.array)
    #     if self.config["obs_setting"]["LED_TYPE"] is not None:
    #         if len(self.config["obs_setting"]["LED_TYPE"]) != 0:
    #             print("LED OPEN--------")

    #             led_obj = FlatLED(chip, filt)

    #             led_flat = led_obj.drawObj_LEDFlat(led_type_list=self.config["obs_setting"]["LED_TYPE"], exp_t_list=self.config["obs_setting"]["LED_TIME"])
    #             pf_map = led_flat
    #     # whether to output zero, dark, flat calibration images.
    #     expTime = self.config["obs_setting"]["exp_time"]
    #     norm_scaler = skyback_level/expTime / self.filter_param.param[sky_level_filt][5]

    #     if skyback_level == 0:
    #         self.config["ins_effects"]["shutter_effect"] = False

    #     if chip.survey_type == "photometric":
    #         sky_map = np.ones_like(chip.img.array) * norm_scaler * self.filter_param.param[chip.filter_type][5] / self.tel.pupil_area
    #     elif chip.survey_type == "spectroscopic":
    #         flat_normal = np.ones_like(chip.img.array)
    #         if self.config["ins_effects"]["flat_fielding"] == True:
    #             chip_output.Log_info("SLS flat preprocess,CHIP %d : Creating and applying Flat-Fielding" % chip.chipID)
    #             msg = str(chip.img.bounds)
    #             chip_output.Log_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:
    #             chip_output.Log_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 = sky_map * norm_scaler

    #     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,
    #         exptime=self.config["obs_setting"]["exp_time"],
    #         pointing_ID=pointing.id,
    #         timestamp_obs=pointing.timestamp,
    #         pointing_type=pointing.pointing_type,
    #         sky_map=sky_map, tel=self.tel,
    #         post_flash_map=pf_map,
    #         logger=chip_output.logger)

    #     datetime_obs = datetime.utcfromtimestamp(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,
    #         pixel_scale=chip.pix_scale,
    #         date=date_obs,
    #         time_obs=time_obs,
    #         exptime=self.config["obs_setting"]["exp_time"],
    #         im_type='DARKPF',
    #         sat_pos=[pointing.sat_x, pointing.sat_y, pointing.sat_z],
    #         sat_vel=[pointing.sat_vx, pointing.sat_vy, pointing.sat_vz],
    #         chip_name=str(chip.chipID).rjust(2, '0'))
    #     h_ext = generateExtensionHeader(
    #         chip=chip,
    #         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,
    #         pixel_scale=chip.pix_scale,
    #         pixel_size=chip.pix_size,
    #         xcen=chip.x_cen,
    #         ycen=chip.y_cen,
    #         extName='SCI',
    #         timestamp=pointing.timestamp,
    #         exptime=self.config["obs_setting"]["exp_time"],
    #         readoutTime=chip.readout_time)

    #     chip.img = galsim.Image(chip.img.array, dtype=np.uint16)
    #     hdu1 = fits.PrimaryHDU(header=h_prim)
    #     hdu1.add_checksum()
    #     hdu1.header.comments['CHECKSUM'] = 'HDU checksum'
    #     hdu1.header.comments['DATASUM'] = 'data unit checksum'
    #     hdu2 = fits.ImageHDU(chip.img.array, header=h_ext)
    #     hdu2.add_checksum()
    #     hdu2.header.comments['XTENSION'] = 'extension type'
    #     hdu2.header.comments['CHECKSUM'] = 'HDU checksum'
    #     hdu2.header.comments['DATASUM'] = 'data unit checksum'
    #     hdu1 = fits.HDUList([hdu1, hdu2])
    #     fname = os.path.join(chip_output.subdir, h_prim['FILENAME'] + '.fits')
    #     hdu1.writeto(fname, output_verify='ignore', overwrite=True)

    #         # chip_output.Log_info("# objects that are too bright %d out of %d" % (bright_obj, self.nobj))
    #         # chip_output.Log_info("# objects that are too dim %d out of %d" % (dim_obj, self.nobj))
    #         # chip_output.Log_info("# objects that are missed %d out of %d" % (missed_obj, self.nobj))
    #     del chip.img

    #     chip_output.Log_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)))
Fang Yuedong's avatar
Fang Yuedong committed
    def runExposure_MPI_PointingList(self, pointing_list, chips=None, use_mpi=False):
Fang Yuedong's avatar
Fang Yuedong committed
        if use_mpi:
            comm = MPI.COMM_WORLD
            ind_thread = comm.Get_rank()
            num_thread = comm.Get_size()
Fang Yuedong's avatar
Fang Yuedong committed

Fang Yuedong's avatar
Fang Yuedong committed
        process_counter = 0
        for ipoint in range(len(pointing_list)):
Fang Yuedong's avatar
Fang Yuedong committed
            # Construct chips & filters:
            pointing = pointing_list[ipoint]
            pointing_ID = pointing.id
            self.focal_plane = FocalPlane(chip_list=pointing.obs_param["run_chips"])
            # Make Chip & Filter lists
            self.chip_list = []
            self.filter_list = []
            self.all_filters = []
Fang Yuedong's avatar
Fang Yuedong committed
            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)
            
Fang Yuedong's avatar
Fang Yuedong committed
            for ichip in range(nchips_per_fp):
Fang Yuedong's avatar
Fang Yuedong committed
                i_process = process_counter + ichip
Fang Yuedong's avatar
Fang Yuedong committed
                if use_mpi:
Fang Yuedong's avatar
Fang Yuedong committed
                    if i_process % num_thread != ind_thread:
Fang Yuedong's avatar
Fang Yuedong committed
                        continue
Fang Yuedong's avatar
Fang Yuedong committed
                pid = os.getpid()

                sub_img_dir, prefix = makeSubDir_PointingList(path_dict=self.path_dict, config=self.config, pointing_ID=pointing_ID)

Fang Yuedong's avatar
Fang Yuedong committed
                chip = run_chips[ichip]
                filt = run_filts[ichip]
Fang Yuedong's avatar
Fang Yuedong committed
                # chip_output.Log_info("running pointing#%d, chip#%d, at PID#%d..."%(pointing_ID, chip.chipID, pid))
Fang Yuedong's avatar
Fang Yuedong committed
                chip_output = ChipOutput(
                    config=self.config, 
                    focal_plane=self.focal_plane, 
                    chip=chip, 
                    filt=filt,  
                    exptime=pointing.exp_time,
                    pointing_type=pointing.pointing_type,
Fang Yuedong's avatar
Fang Yuedong committed
                    pointing_ID=pointing_ID,  
                    subdir=sub_img_dir,
                    prefix=prefix,
                    ra_cen=pointing.ra,
                    dec_cen=pointing.dec)
Fang Yuedong's avatar
Fang Yuedong committed
                chip_output.Log_info("running pointing#%d, chip#%d, at PID#%d..."%(pointing_ID, chip.chipID, pid))
Fang Yuedong's avatar
Fang Yuedong committed
                self.run_one_chip(
                    chip=chip,
                    filt=filt,
                    chip_output=chip_output,
                    pointing=pointing)
                # if self.config["obs_setting"]["survey_type"] == "CALIBRATION":
                #     self.run_one_chip_calibration(chip=chip,
                #                                filt=filt,
                #                                chip_output=chip_output,
                #                                pointing=pointing,
                #                                skyback_level = self.config["obs_setting"]["FLAT_LEVEL"],
                #                                sky_level_filt = self.config["obs_setting"]["FLAT_LEVEL_FIL"])
                # else:
                #     self.run_one_chip(
                #         chip=chip,
                #         filt=filt,
                #         chip_output=chip_output,
                #         pointing=pointing)
Fang Yuedong's avatar
Fang Yuedong committed
                chip_output.Log_info("finished running chip#%d..."%(chip.chipID))
                for handler in chip_output.logger.handlers[:]:
                    chip_output.logger.removeHandler(handler)
Fang Yuedong's avatar
Fang Yuedong committed
                gc.collect()
Fang Yuedong's avatar
Fang Yuedong committed
            process_counter += nchips_per_fp