Raw File
# -*- coding: utf-8 -*-
"""Generating the training data.

This script generates the training data according to the config specifications.

To run this script, pass in the desired config file as argument::

    $ generate baobab/configs/tdlmc_diagonal_config.py --n_data 1000


import os, sys
import random
import argparse
import gc
from types import SimpleNamespace
from tqdm import tqdm
import numpy as np
import pandas as pd
# Lenstronomy modules
import lenstronomy
print("Lenstronomy path being used: {:s}".format(lenstronomy.__path__[0]))
from lenstronomy.LensModel.lens_model import LensModel
from lenstronomy.LightModel.light_model import LightModel
from lenstronomy.PointSource.point_source import PointSource
# Baobab modules
from baobab.configs import BaobabConfig
import baobab.bnn_priors as bnn_priors
from baobab.sim_utils import Imager, Selection

def parse_args():
    """Parse command-line arguments

    parser = argparse.ArgumentParser()
    parser.add_argument('config', help='Baobab config file path')
    parser.add_argument('--n_data', default=None, dest='n_data', type=int,
                        help='size of dataset to generate (overrides config file)')
    parser.add_argument('--dest_dir', default=None, dest='dest_dir', type=str,
                        help='destination for output folder (overrides config file)')
    args = parser.parse_args()
    # sys.argv rerouting for setuptools entry point
    if args is None:
        args = SimpleNamespace()
        args.config = sys.argv[0]
        args.n_data = sys.argv[1]
        args.dest_dir = sys.argv[2]
    return args

def main():
    args = parse_args()
    cfg = BaobabConfig.from_file(args.config)
    if args.n_data is not None:
        cfg.n_data = args.n_data
    if args.dest_dir is not None:
        cfg.destination_dir = args.dest_dir
    # Seed for reproducibility
    # Create data directory
    save_dir = cfg.out_dir
    if not os.path.exists(save_dir):
        print("Destination folder path: {:s}".format(save_dir))
        raise OSError("Destination folder already exists.")
    # Instantiate density models
    kwargs_model = dict(
                    lens_model_list=[cfg.bnn_omega.lens_mass.profile, cfg.bnn_omega.external_shear.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'])
    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'])
        lens_light_model = None
    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])
        ps_model = None
    # Instantiate Selection object
    selection = Selection(cfg.selection, cfg.components)
    # Instantiate Imager object
    if cfg.bnn_omega.kinematics.calculate_vel_disp or cfg.bnn_omega.time_delays.calculate_time_delays:
        for_cosmography = True
        for_cosmography = False
    imager = Imager(cfg.components, lens_mass_model, src_light_model, lens_light_model=lens_light_model, ps_model=ps_model, kwargs_numerics=cfg.numerics, min_magnification=cfg.selection.magnification.min, for_cosmography=for_cosmography, magnification_frac_err=cfg.bnn_omega.magnification.frac_error_sigma)
    # Initialize BNN prior
    if for_cosmography:
        kwargs_lens_eqn_solver = {'min_distance': 0.05, 'search_window': cfg.instrument['pixel_scale']*cfg.image['num_pix'], 'num_iter_max': 100}
        bnn_prior = getattr(bnn_priors, cfg.bnn_prior_class)(cfg.bnn_omega, cfg.components, kwargs_lens_eqn_solver)
        kwargs_lens_eqn_solver = {}
        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
        if selection.reject_initial(sample): # select on sampled model parameters
        # Generate the image
        img, img_features = imager.generate_image(sample, cfg.image.num_pix, cfg.survey_object_dict)
        if img is None: # select on stats computed while rendering the image
        # Save image file
        if cfg.image.squeeze_bandpass_dimension:
            img = np.squeeze(img)
        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: # Log model parameters
            for param_name, param_value in sample[comp].items():
                meta['{:s}_{:s}'.format(comp, param_name)] = param_value  
        if cfg.bnn_prior_class in ['EmpiricalBNNPrior', 'DiagonalCosmoBNNPrior']: # Log other stats
            for misc_name, misc_value in sample['misc'].items():
                meta['{:s}'.format(misc_name)] = misc_value
        if 'agn_light' in cfg.components:
            meta['x_image'] = img_features['x_image'].tolist()
            meta['y_image'] = img_features['y_image'].tolist()
            meta['n_img'] = len(img_features['y_image'])
            meta['magnification'] = img_features['magnification'].tolist()
            meta['measured_magnification'] = img_features['measured_magnification'].tolist()
        meta['img_filename'] = img_filename
        metadata = metadata.append(meta, ignore_index=True)
        # Export metadata.csv for the first time
        if current_idx == 0:
            metadata = metadata.reindex(sorted(metadata.columns), axis=1) # sort columns lexicographically
            metadata.to_csv(metadata_path, index=None) # export to csv
            metadata = pd.DataFrame() # init empty df for next checkpoint chunk
        # Export metadata every checkpoint interval
        if (current_idx + 1)%cfg.checkpoint_interval == 0:
            metadata.to_csv(metadata_path, index=None, mode='a', header=None) # export to csv
            metadata = pd.DataFrame() # init empty df for next checkpoint chunk
        # Update progress
        current_idx += 1
    # Export to csv
    metadata.to_csv(metadata_path, index=None, mode='a', header=None)
if __name__ == '__main__':
back to top