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): 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