def setup(self): numpix = 10 self.ccd_gain = 4. self.pixel_scale = 0.13 self.read_noise = 10. kwargs_instrument = {'read_noise': self.read_noise, 'pixel_scale': self.pixel_scale, 'ccd_gain': self.ccd_gain} exposure_time = 100 sky_brightness = 20. self.magnitude_zero_point = 21. num_exposures = 2 seeing = 0.9 kwargs_observations = {'exposure_time': exposure_time, 'sky_brightness': sky_brightness, 'magnitude_zero_point': self.magnitude_zero_point, 'num_exposures': num_exposures, 'seeing': seeing, 'psf_type': 'GAUSSIAN', 'kernel_point_source': None} self.kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observations) self.api = DataAPI(numpix=numpix, data_count_unit='ADU', **self.kwargs_data) kwargs_observations = {'exposure_time': exposure_time, 'sky_brightness': sky_brightness, 'magnitude_zero_point': self.magnitude_zero_point, 'num_exposures': num_exposures, 'seeing': seeing, 'psf_type': 'PIXEL', 'kernel_point_source': np.ones((3, 3))} kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observations) self.api_pixel = DataAPI(numpix=numpix, data_count_unit='ADU', **kwargs_data) self.ra_at_xy_0 = 0.02 self.dec_at_xy_0 = 0.02 self.transform_pix2angle = [[-self.pixel_scale,0],[0,self.pixel_scale]] kwargs_pixel_grid = {'ra_at_xy_0':self.ra_at_xy_0,'dec_at_xy_0':self.dec_at_xy_0, 'transform_pix2angle':self.transform_pix2angle} self.api_pixel_grid = DataAPI(numpix=numpix, kwargs_pixel_grid=kwargs_pixel_grid, data_count_unit='ADU',**self.kwargs_data)
def test_psf_type(self): assert self.data_adu._psf_type == 'GAUSSIAN' kwargs_observations = { 'exposure_time': 1, 'sky_brightness': 1, 'magnitude_zero_point': self.magnitude_zero_point, 'num_exposures': 1, 'seeing': 1, 'psf_type': 'PIXEL' } kwargs_data = util.merge_dicts(self.kwargs_instrument, kwargs_observations) data_pixel = SingleBand(data_count_unit='ADU', **kwargs_data) assert data_pixel._psf_type == 'PIXEL' kwargs_observations = { 'exposure_time': 1, 'sky_brightness': 1, 'magnitude_zero_point': self.magnitude_zero_point, 'num_exposures': 1, 'seeing': 1, 'psf_type': 'NONE' } kwargs_data = util.merge_dicts(self.kwargs_instrument, kwargs_observations) data_pixel = SingleBand(data_count_unit='ADU', **kwargs_data) assert data_pixel._psf_type == 'NONE'
def test_raise(self): numpix = 10 self.ccd_gain = 4. self.pixel_scale = 0.13 self.read_noise = 10. kwargs_instrument = {'read_noise': self.read_noise, 'pixel_scale': self.pixel_scale, 'ccd_gain': self.ccd_gain} exposure_time = 100 sky_brightness = 20. magnitude_zero_point = 21. num_exposures = 2 seeing = 0.9 kwargs_observations = {'exposure_time': exposure_time, 'sky_brightness': sky_brightness, 'magnitude_zero_point': magnitude_zero_point, 'num_exposures': num_exposures, 'seeing': seeing, 'psf_type': 'wrong', 'kernel_point_source': None} kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observations) data_api = DataAPI(numpix=numpix, data_count_unit='ADU', **kwargs_data) print(data_api._psf_type) with self.assertRaises(ValueError): data_api = DataAPI(numpix=numpix, data_count_unit='ADU', **kwargs_data) psf_class = data_api.psf_class kwargs_observations = {'exposure_time': exposure_time, 'sky_brightness': sky_brightness, 'magnitude_zero_point': magnitude_zero_point, 'num_exposures': num_exposures, 'seeing': seeing, 'psf_type': 'PIXEL', 'kernel_point_source': None} kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observations) with self.assertRaises(ValueError): data_api = DataAPI(numpix=numpix, data_count_unit='ADU', **kwargs_data) psf_class = data_api.psf_class kwargs_data['kernel_point_source'] = np.ones((3, 3)) kwargs_pixel_grid = {'ra_at_xy_0':0.02,'dec_at_xy_0':0.02} with self.assertRaises(ValueError): data_api = DataAPI(numpix=numpix,kwargs_pixel_grid=kwargs_pixel_grid, **kwargs_data)
def setup(self): self.ccd_gain = 4. pixel_scale = 0.13 self.read_noise = 10. self.kwargs_instrument = { 'read_noise': self.read_noise, 'pixel_scale': pixel_scale, 'ccd_gain': self.ccd_gain } exposure_time = 100 sky_brightness = 20. self.magnitude_zero_point = 21. num_exposures = 2 seeing = 0.9 kwargs_observations = { 'exposure_time': exposure_time, 'sky_brightness': sky_brightness, 'magnitude_zero_point': self.magnitude_zero_point, 'num_exposures': num_exposures, 'seeing': seeing, 'psf_type': 'GAUSSIAN' } self.kwargs_data = util.merge_dicts(self.kwargs_instrument, kwargs_observations) self.data_adu = SingleBand(data_count_unit='ADU', **self.kwargs_data) self.data_e_ = SingleBand(data_count_unit='e-', **self.kwargs_data)
def test_raise(self): self.ccd_gain = 4. pixel_scale = 0.13 self.read_noise = 10. kwargs_instrument = {'read_noise': self.read_noise, 'pixel_scale': pixel_scale, 'ccd_gain': self.ccd_gain} exposure_time = 100 sky_brightness = 20. self.magnitude_zero_point = 21. num_exposures = 2 seeing = 0.9 kwargs_observations = {'exposure_time': exposure_time, 'sky_brightness': sky_brightness, 'magnitude_zero_point': self.magnitude_zero_point, 'num_exposures': num_exposures, 'seeing': seeing, 'psf_type': 'GAUSSIAN'} self.kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observations) with self.assertRaises(ValueError): SingleBand(data_count_unit='wrong', **self.kwargs_data) with self.assertRaises(ValueError): band = SingleBand(pixel_scale=1, exposure_time=1, magnitude_zero_point=1, read_noise=None, ccd_gain=None, sky_brightness=None, seeing=None, num_exposures=1, psf_type='GAUSSIAN', kernel_point_source=None, data_count_unit='ADU', background_noise=None) out = band.sky_brightness with self.assertRaises(ValueError): band = SingleBand(pixel_scale=1, exposure_time=1, magnitude_zero_point=1, read_noise=None, ccd_gain=None, sky_brightness=None, seeing=None, num_exposures=1, psf_type='GAUSSIAN', kernel_point_source=None, data_count_unit='ADU', background_noise=None) out = band.read_noise with self.assertRaises(ValueError): band = SingleBand(pixel_scale=1, exposure_time=1, magnitude_zero_point=1, read_noise=None, ccd_gain=None, sky_brightness=None, seeing=None, num_exposures=1, psf_type='GAUSSIAN', kernel_point_source=None, data_count_unit='ADU', background_noise=None) out = band.background_noise
def kwargs_single_band(self): """ :return: merged kwargs from camera and obs dicts """ kwargs = util.merge_dicts(self.camera, self.obs) return kwargs
def observation_constructor(instrument_name, observation_name): """ :param instrument_name: string, name of instrument referenced in this file :param observation_name: string, name of observation referenced in this file :return: instance of the SimulationAPI.data_type instance """ if instrument_name == 'LSST': kwargs_instrument = LSST_camera else: raise ValueError("instrument name %s not supported! Chose among %s " % (instrument_name, instrument_name_list)) if observation_name == 'LSST_g_band': kwargs_observation = LSST_g_band_obs elif observation_name == 'LSST_r_band': kwargs_observation = LSST_r_band_obs elif observation_name == 'LSST_i_band': kwargs_observation = LSST_i_band_obs else: raise ValueError('observation name %s not supported! Chose among %s' % (observation_name, observation_name_list)) kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observation) return kwargs_data
def test_raise(self): self.ccd_gain = 4. pixel_scale = 0.13 self.read_noise = 10. kwargs_instrument = { 'read_noise': self.read_noise, 'pixel_scale': pixel_scale, 'ccd_gain': self.ccd_gain } exposure_time = 100 sky_brightness = 20. self.magnitude_zero_point = 21. num_exposures = 2 seeing = 0.9 kwargs_observations = { 'exposure_time': exposure_time, 'sky_brightness': sky_brightness, 'magnitude_zero_point': self.magnitude_zero_point, 'num_exposures': num_exposures, 'seeing': seeing, 'psf_type': 'GAUSSIAN' } self.kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observations) with self.assertRaises(ValueError): SingleBand(data_count_unit='wrong', **self.kwargs_data)
def setup(self): numpix = 10 self.ccd_gain = 4. self.pixel_scale = 0.13 self.read_noise = 10. kwargs_instrument = { 'read_noise': self.read_noise, 'pixel_scale': self.pixel_scale, 'ccd_gain': self.ccd_gain } exposure_time = 100 sky_brightness = 20. self.magnitude_zero_point = 21. num_exposures = 2 seeing = 0.9 kwargs_observations = { 'exposure_time': exposure_time, 'sky_brightness': sky_brightness, 'magnitude_zero_point': self.magnitude_zero_point, 'num_exposures': num_exposures, 'seeing': seeing, 'psf_type': 'GAUSSIAN', 'psf_model': None } self.kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observations) self.api = DataAPI(numpix=numpix, data_count_unit='ADU', **self.kwargs_data) kwargs_observations = { 'exposure_time': exposure_time, 'sky_brightness': sky_brightness, 'magnitude_zero_point': self.magnitude_zero_point, 'num_exposures': num_exposures, 'seeing': seeing, 'psf_type': 'PIXEL', 'psf_model': np.ones((3, 3)) } kwargs_data = util.merge_dicts(kwargs_instrument, kwargs_observations) self.api_pixel = DataAPI(numpix=numpix, data_count_unit='ADU', **kwargs_data)
def test_kwargs_single_band(self): kwargs_g = util.merge_dicts(self.g.camera, self.g.obs) self.assertEqual(self.g.kwargs_single_band(), kwargs_g)
def test_kwargs_single_band(self): kwargs_F062 = util.merge_dicts(self.F062.camera, self.F062.obs) self.assertEqual(self.F062.kwargs_single_band(), kwargs_F062)
def main(): args = parse_args() cfg = BaobabConfig.from_file(args.config) if args.n_data is not None: cfg.n_data = args.n_data # Seed for reproducibility np.random.seed(cfg.seed) random.seed(cfg.seed) # Create data directory save_dir = cfg.out_dir if not os.path.exists(save_dir): os.makedirs(save_dir) print("Destination folder path: {:s}".format(save_dir)) print("Log path: {:s}".format(cfg.log_path)) cfg.export_log() else: raise OSError("Destination folder already exists.") # Instantiate PSF models psf_models = instantiate_PSF_models(cfg.psf, cfg.instrument.pixel_scale) n_psf = len(psf_models) # Instantiate density models kwargs_model = dict( lens_model_list=[ cfg.bnn_omega.lens_mass.profile, cfg.bnn_omega.external_shear.profile ], source_light_model_list=[cfg.bnn_omega.src_light.profile], ) lens_mass_model = LensModel( lens_model_list=kwargs_model['lens_model_list']) src_light_model = LightModel( light_model_list=kwargs_model['source_light_model_list']) lens_eq_solver = LensEquationSolver(lens_mass_model) lens_light_model = None ps_model = None if 'lens_light' in cfg.components: kwargs_model['lens_light_model_list'] = [ cfg.bnn_omega.lens_light.profile ] lens_light_model = LightModel( light_model_list=kwargs_model['lens_light_model_list']) if 'agn_light' in cfg.components: kwargs_model['point_source_model_list'] = [ cfg.bnn_omega.agn_light.profile ] ps_model = PointSource( point_source_type_list=kwargs_model['point_source_model_list'], fixed_magnification_list=[False]) # Instantiate Selection object selection = Selection(cfg.selection, cfg.components) # Initialize BNN prior bnn_prior = getattr(bnn_priors, cfg.bnn_prior_class)(cfg.bnn_omega, cfg.components) # Initialize empty metadata dataframe metadata = pd.DataFrame() metadata_path = os.path.join(save_dir, 'metadata.csv') current_idx = 0 # running idx of dataset pbar = tqdm(total=cfg.n_data) while current_idx < cfg.n_data: sample = bnn_prior.sample() # FIXME: sampling in batches # Selections on sampled parameters if selection.reject_initial(sample): continue psf_model = get_PSF_model(psf_models, n_psf, current_idx) # Instantiate the image maker data_api with detector and observation conditions kwargs_detector = util.merge_dicts(cfg.instrument, cfg.bandpass, cfg.observation) kwargs_detector.update(seeing=cfg.psf.fwhm, psf_type=cfg.psf.type, kernel_point_source=psf_model, background_noise=0.0) data_api = DataAPI(cfg.image.num_pix, **kwargs_detector) # Generate the image img, img_features = generate_image( sample, psf_model, data_api, lens_mass_model, src_light_model, lens_eq_solver, cfg.instrument.pixel_scale, cfg.image.num_pix, cfg.components, cfg.numerics, min_magnification=cfg.selection.magnification.min, lens_light_model=lens_light_model, ps_model=ps_model) if img is None: # couldn't make the magnification cut continue # Save image file img_filename = 'X_{0:07d}.npy'.format(current_idx) img_path = os.path.join(save_dir, img_filename) np.save(img_path, img) # Save labels meta = {} for comp in cfg.components: for param_name, param_value in sample[comp].items(): meta['{:s}_{:s}'.format(comp, param_name)] = param_value #if cfg.bnn_prior_class in ['DiagonalCosmoBNNPrior']: # if cfg.bnn_omega.time_delays.calculate_time_delays: # # Order time delays in increasing dec # unordered_td = sample['misc']['true_td'] # np array # increasing_dec_i = np.argsort(img_features['y_image']) # td = unordered_td[increasing_dec_i] # td = td[1:] - td[0] # take BCD - A # sample['misc']['true_td'] = list(td) # img_features['x_image'] = img_features['x_image'][increasing_dec_i] # img_features['y_image'] = img_features['y_image'][increasing_dec_i] if cfg.bnn_prior_class in [ 'EmpiricalBNNPrior', 'DiagonalCosmoBNNPrior' ]: for misc_name, misc_value in sample['misc'].items(): meta['{:s}'.format(misc_name)] = misc_value if 'agn_light' in cfg.components: x_image = np.zeros(4) y_image = np.zeros(4) n_img = len(img_features['x_image']) meta['n_img'] = n_img x_image[:n_img] = img_features['x_image'] y_image[:n_img] = img_features['y_image'] for i in range(4): meta['x_image_{:d}'.format(i)] = x_image[i] meta['y_image_{:d}'.format(i)] = y_image[i] meta['total_magnification'] = img_features['total_magnification'] meta['img_filename'] = img_filename metadata = metadata.append(meta, ignore_index=True) # Export metadata.csv for the first time if current_idx == 0: # Sort columns lexicographically metadata = metadata.reindex(sorted(metadata.columns), axis=1) # Export to csv metadata.to_csv(metadata_path, index=None) # Initialize empty dataframe for next checkpoint chunk metadata = pd.DataFrame() gc.collect() # Export metadata every checkpoint interval if (current_idx + 1) % cfg.checkpoint_interval == 0: # Export to csv metadata.to_csv(metadata_path, index=None, mode='a', header=None) # Initialize empty dataframe for next checkpoint chunk metadata = pd.DataFrame() gc.collect() # Update progress current_idx += 1 pbar.update(1) # Export to csv metadata.to_csv(metadata_path, index=None, mode='a', header=None) pbar.close()
def sim_single_image(sn_data, geo_dict, label): DECam = { 'read_noise': 10, # std of noise generated by read-out (in units of electrons) 'pixel_scale': 0.263, # scale (in arcseonds) of pixels 'ccd_gain': 4.5 # electrons/ADU (analog-to-digital unit). } obs_info = { 'exposure_time': 90, 'sky_brightness': sn_data['SKYMAG'], 'magnitude_zero_point': sn_data['ZPTMAG'], 'num_exposures': 1, 'seeing': sn_data['PSF'] / 3, 'psf_type': 'GAUSSIAN' } kwargs_model = { 'lens_model_list': [ 'SIE', 'SHEAR' ], # list of lens models to be used 'lens_light_model_list': [ 'SERSIC_ELLIPSE' ], # list of unlensed light models to be used 'source_light_model_list': [ 'SERSIC_ELLIPSE' ], # list of extended source models to be used 'point_source_model_list': [ 'SOURCE_POSITION' ], # list of point source models to be used } numpix = 64 kwargs_merged = util.merge_dicts(DECam, obs_info) kwargs_numerics = {'point_source_supersampling_factor': 1} sim = SimAPI(numpix=numpix, kwargs_single_band=kwargs_merged, kwargs_model=kwargs_model, kwargs_numerics=kwargs_numerics) imSim = sim.image_model_class x_grid, y_grid = util.make_grid(numPix=10, deltapix=1) flux = gauss.function(x_grid, y_grid, amp=1, sigma=2, e1=0.4, e2=0, center_x=0, center_y=0) image_gauss = util.array2image(flux) ## Set point source mag geo_dict['ps_kwargs']['magnitude'] = sn_data['ABMAG'] * 1.4 if label == 1: geo_dict = force_larger_ps_separation(geo_dict) elif label == 2: geo_dict = force_no_lensing_with_ps(geo_dict) elif label == 3: geo_dict = force_lensed_agn(geo_dict) elif label == 4: geo_dict = force_non_lensed_agn(geo_dict) else: print("Unexpected option passed as class label") sys.exit() #lens light kwargs_lens_light_mag = [geo_dict['lens_kwargs']] # source light kwargs_source_mag = [geo_dict['source_kwargs']] # point source kwargs_ps_mag = [geo_dict['ps_kwargs']] kwargs_lens_light, kwargs_source, kwargs_ps = sim.magnitude2amplitude( kwargs_lens_light_mag, kwargs_source_mag, kwargs_ps_mag) kwargs_lens = [ geo_dict['lens_model_kwargs_sie'], geo_dict['lens_model_kwargs_shear'] ] image = imSim.image(kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps) image += sim.noise_for_model(model=image) return np.log10(image + 2)
def test_kwargs_single_band(self): kwargs_VIS = util.merge_dicts(self.VIS.camera, self.VIS.obs) self.assertEqual(self.VIS.kwargs_single_band(), kwargs_VIS)
def test_kwargs_single_band(self): kwargs_TDLMC_F160W = util.merge_dicts(self.TDLMC_F160W.camera, self.TDLMC_F160W.obs) self.assertEqual(self.TDLMC_F160W.kwargs_single_band(), kwargs_TDLMC_F160W)
def main(): args = parse_args() cfg = Config.fromfile(args.config) if args.n_data is not None: cfg.n_data = args.n_data # Seed for reproducibility np.random.seed(cfg.seed) random.seed(cfg.seed) if not os.path.exists(cfg.out_dir): os.makedirs(cfg.out_dir) print("Destination folder: {:s}".format(cfg.out_dir)) else: raise OSError("Destination folder already exists.") # Instantiate PSF models psf_models = get_PSF_models(cfg.psf, cfg.instrument.pixel_scale) n_psf = len(psf_models) # Instantiate ImageData #kwargs_data = sim_util.data_configure_simple(**cfg.image) #image_data = ImageData(**kwargs_data) # Instantiate density models kwargs_model = dict( lens_model_list=[ cfg.bnn_omega.lens_mass.profile, cfg.bnn_omega.external_shear.profile ], source_light_model_list=[cfg.bnn_omega.src_light.profile], ) lens_mass_model = LensModel( lens_model_list=kwargs_model['lens_model_list']) src_light_model = LightModel( light_model_list=kwargs_model['source_light_model_list']) lens_eq_solver = LensEquationSolver(lens_mass_model) lens_light_model = None ps_model = None if 'lens_light' in cfg.components: kwargs_model['lens_light_model_list'] = [ cfg.bnn_omega.lens_light.profile ] lens_light_model = LightModel( light_model_list=kwargs_model['lens_light_model_list']) if 'agn_light' in cfg.components: kwargs_model['point_source_model_list'] = [ cfg.bnn_omega.agn_light.profile ] ps_model = PointSource( point_source_type_list=kwargs_model['point_source_model_list'], fixed_magnification_list=[False]) # Initialize BNN prior bnn_prior = getattr(bnn_priors, cfg.bnn_prior_class)(cfg.bnn_omega, cfg.components) # Initialize dataframe of labels param_list = [] for comp in cfg.components: param_list += [ '{:s}_{:s}'.format(comp, param) for param in bnn_prior.params[cfg.bnn_omega[comp]['profile']] ] if 'agn_light' in cfg.components: param_list += ['magnification_{:d}'.format(i) for i in range(4)] param_list += ['x_image_{:d}'.format(i) for i in range(4)] param_list += ['y_image_{:d}'.format(i) for i in range(4)] param_list += ['n_img'] param_list += ['img_path', 'total_magnification'] if cfg.bnn_prior_class == 'EmpiricalBNNPrior': param_list += [ 'z_lens', 'z_src', 'vel_disp_iso', 'R_eff_lens', 'R_eff_src', 'abmag_src' ] metadata = pd.DataFrame(columns=param_list) print("Starting simulation...") current_idx = 0 # running idx of dataset pbar = tqdm(total=cfg.n_data) while current_idx < cfg.n_data: psf_model = psf_models[current_idx % n_psf] sample = bnn_prior.sample() # FIXME: sampling in batches if sample['lens_mass']['theta_E'] < cfg.selection.theta_E.min: continue # Instantiate SimAPI (converts mag to amp and wraps around image model) kwargs_detector = util.merge_dicts(cfg.instrument, cfg.bandpass, cfg.observation) kwargs_detector.update( seeing=cfg.psf.fwhm, psf_type=cfg.psf.type, kernel_point_source=psf_model ) # keyword deprecation warning: I asked Simon to change this to data_api = DataAPI(cfg.image.num_pix, **kwargs_detector) image_data = data_api.data_class #sim_api = SimAPI(numpix=cfg.image.num_pix, # kwargs_single_band=kwargs_detector, # kwargs_model=kwargs_model, # kwargs_numerics=cfg.numerics) kwargs_lens_mass = [sample['lens_mass'], sample['external_shear']] kwargs_src_light = [sample['src_light']] kwargs_src_light = amp_to_mag_extended(kwargs_src_light, src_light_model, data_api) kwargs_lens_light = None kwargs_ps = None if 'agn_light' in cfg.components: x_image, y_image = lens_eq_solver.findBrightImage( sample['src_light']['center_x'], sample['src_light']['center_y'], kwargs_lens_mass, numImages=4, min_distance=cfg.instrument.pixel_scale, search_window=cfg.image.num_pix * cfg.instrument.pixel_scale) magnification = np.abs( lens_mass_model.magnification(x_image, y_image, kwargs=kwargs_lens_mass)) unlensed_mag = sample['agn_light']['magnitude'] # unlensed agn mag kwargs_unlensed_mag_ps = [{ 'ra_image': x_image, 'dec_image': y_image, 'magnitude': unlensed_mag }] # note unlensed magnitude kwargs_unlensed_amp_ps = amp_to_mag_point( kwargs_unlensed_mag_ps, ps_model, data_api) # note unlensed amp kwargs_ps = copy.deepcopy(kwargs_unlensed_amp_ps) for kw in kwargs_ps: kw.update(point_amp=kw['point_amp'] * magnification) else: kwargs_unlensed_amp_ps = None if 'lens_light' in cfg.components: kwargs_lens_light = [sample['lens_light']] kwargs_lens_light = amp_to_mag_extended(kwargs_lens_light, lens_light_model, data_api) # Instantiate image model image_model = ImageModel(image_data, psf_model, lens_mass_model, src_light_model, lens_light_model, ps_model, kwargs_numerics=cfg.numerics) # Compute magnification lensed_total_flux = get_lensed_total_flux(kwargs_lens_mass, kwargs_src_light, kwargs_lens_light, kwargs_ps, image_model) unlensed_total_flux = get_unlensed_total_flux(kwargs_src_light, src_light_model, kwargs_unlensed_amp_ps, ps_model) total_magnification = lensed_total_flux / unlensed_total_flux # Apply magnification cut if total_magnification < cfg.selection.magnification.min: continue # Generate image for export img = image_model.image(kwargs_lens_mass, kwargs_src_light, kwargs_lens_light, kwargs_ps) #kwargs_in_amp = sim_api.magnitude2amplitude(kwargs_lens_mass, kwargs_src_light, kwargs_lens_light, kwargs_ps) #imsim_api = sim_api.image_model_class #imsim_api.image(*kwargs_in_amp) # Add noise noise = data_api.noise_for_model(img, background_noise=True, poisson_noise=True, seed=None) img += noise # Save image file img_path = os.path.join(cfg.out_dir, 'X_{0:07d}.npy'.format(current_idx + 1)) np.save(img_path, img) # Save labels meta = {} for comp in cfg.components: for param_name, param_value in sample[comp].items(): meta['{:s}_{:s}'.format(comp, param_name)] = param_value if 'agn_light' in cfg.components: n_img = len(x_image) for i in range(n_img): meta['magnification_{:d}'.format(i)] = magnification[i] meta['x_image_{:d}'.format(i)] = x_image[i] meta['y_image_{:d}'.format(i)] = y_image[i] meta['n_img'] = n_img if cfg.bnn_prior_class == 'EmpiricalBNNPrior': for misc_name, misc_value in sample['misc'].items(): meta['{:s}'.format(misc_name)] = misc_value meta['total_magnification'] = total_magnification meta['img_path'] = img_path metadata = metadata.append(meta, ignore_index=True) # Update progress current_idx += 1 pbar.update(1) pbar.close() # Fix column ordering metadata = metadata[param_list] metadata_path = os.path.join(cfg.out_dir, 'metadata.csv') metadata.to_csv(metadata_path, index=None) print("Labels include: ", metadata.columns.values)