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