jaxtronomy.Workflow package

Submodules

jaxtronomy.Workflow.fitting_sequence module

class FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params, mpi=False, verbose=True)

Bases: object

Class to define a sequence of fitting applied, inherit the Fitting class this is a Workflow manager that allows to update model configurations before executing another step in the modelling The user can take this module as an example of how to create their own workflows or build their own around the FittingSequence.

__init__(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params, mpi=False, verbose=True)

Parameters:

• kwargs_data_joint – keyword argument specifying the data according to Likelihood

• kwargs_model – keyword arguments to describe all model components used in class_creator.create_class_instances()

• kwargs_constraints – keyword arguments of the Param() class to handle parameter constraints during the sampling (except upper and lower limits and sampling input mean and width)

• kwargs_likelihood – keyword arguments of the Likelihood() class to handle parameters and settings of the likelihood

• kwargs_params – setting of the sampling bounds and initial guess mean and spread. The argument is organized as: ‘lens_model’: [kwargs_init, kwargs_sigma, kwargs_fixed, kwargs_lower, kwargs_upper] ‘source_model’: [kwargs_init, kwargs_sigma, kwargs_fixed, kwargs_lower, kwargs_upper] ‘lens_light_model’: [kwargs_init, kwargs_sigma, kwargs_fixed, kwargs_lower, kwargs_upper] ‘point_source_model’: [kwargs_init, kwargs_sigma, kwargs_fixed, kwargs_lower, kwargs_upper] ‘extinction_model’: [kwargs_init, kwargs_sigma, kwargs_fixed, kwargs_lower, kwargs_upper] ‘special’: [kwargs_init, kwargs_sigma, kwargs_fixed, kwargs_lower, kwargs_upper] ‘tracer_source_model’: [kwargs_init, kwargs_sigma, kwargs_fixed, kwargs_lower, kwargs_upper]

• mpi – MPI option (bool), if True, will launch an MPI Pool job for the steps in the fitting sequence where possible

• verbose – bool, if True prints temporary results and indicators of the fitting process

property kwargs_fixed

Returns the updated kwargs_fixed from the update manager.

Returns:

list of fixed kwargs, see UpdateManager()

fit_sequence(fitting_list)

Parameters:

fitting_list – list of [[‘string’, {kwargs}], ..] with ‘string being the specific fitting option and kwargs being the arguments passed to this option

Returns:

fitting results

best_fit(bijective=False)

Parameters:

bijective – bool, if True, the mapping of image2source_plane and the mass_scaling parameterisation are inverted. If you do not use those options, there is no effect.

Returns:

best fit model of the current state of the FittingSequence class

update_state(kwargs_update)

Updates current best fit state to the input model keywords specified.

Parameters:

kwargs_update – format of kwargs_result

Returns:

None

best_fit_likelihood(verbose=False)

Returns the log likelihood of the best fit model of the current state of this class.

Parameters:

verbose – bool, if True, prints likelihood statements

Returns:

log likelihood, float

property bic

Bayesian information criterion (BIC) of the model.

Returns:

bic value, float

property param_class

Returns:

Param() class instance reflecting the current state of FittingSequence

property likelihood_class

Returns:

Likelihood() class instance reflecting the current state of FittingSequence

simplex(n_iterations, method='Nelder-Mead')

Downhill simplex optimization using the Nelder-Mead algorithm.

Parameters:

• n_iterations – maximum number of iterations to perform

• method – the optimization method used, see documentation in scipy.optimize.minimize

Returns:

result of the best fit

mcmc(n_burn, n_run, walkerRatio=None, n_walkers=None, sigma_scale=1, threadCount=1, vectorization_batch_size=None, init_samples=None, re_use_samples=True, sampler_type='emcee', progress=True, backend_filename=None, start_from_backend=False, **kwargs_zeus)

MCMC routine.

Parameters:

• n_burn – number of burn in iterations (will not be saved)

• n_run – number of MCMC iterations that are saved

• walkerRatio – ratio of walkers/number of free parameters

• n_walkers – integer, number of walkers of emcee (optional, if set, overwrites the walkerRatio input

• sigma_scale – scaling of the initial parameter spread relative to the width in the initial settings

• threadCount – number of CPU threads. If MPI option is set, threadCount=1

• vectorization_batch_size – int, only relevant for GPU, determines the number of particles/walkers whose logL will be computed simultaneously. None defaults to one particle at a time and 0 means to compute all particles/walkers simultaneously.

• init_samples – initial sample from where to start the MCMC process

• re_use_samples – bool, if True, re-uses the samples described in init_samples.nOtherwise starts from scratch.

• sampler_type – string, which MCMC sampler to be used. Options are ‘emcee’ and ‘zeus’

• progress – boolean, if True shows progress bar in EMCEE

• backend_filename (string) – name of the HDF5 file where sampling state is saved (through emcee backend engine)

• start_from_backend (bool) – if True, start from the state saved in backup_filename. O therwise, create a new backup file with name backup_filename (any already existing file is overwritten!).

• kwargs_zeus – zeus-specific kwargs

Returns:

list of output arguments, e.g. MCMC samples, parameter names, logL distances of all samples specified by the specific sampler used

optax(num_chains, maxiter, tolerance=0.01, sigma_scale=1, rng_seed=None)

Uses Optax L-BFGS gradient descent to find best-fit parameters.

Parameters:

• num_chains – int, number of minimization chains to run

• maxiter – int, maximum number of iterations during gradient descent process

• tolerance – float, if |logL[i] - logL[i-1]| < tol three times in a row, the gradient descent is stopped

• sigma_scale – scales the standard deviation of the prior distribution

• rng_seed – int, used to generate random initial starting point from the prior distribution. If None, a random seed will be used.

pso(n_particles, n_iterations, sigma_scale=1, print_key='PSO', threadCount=1, vectorization_batch_size=None)

Particle Swarm Optimization.

Parameters:

• n_particles – number of particles in the Particle Swarm Optimization

• n_iterations – number of iterations in the optimization process

• sigma_scale – scaling of the initial parameter spread relative to the width in the initial settings

• print_key – string, printed text when executing this routine

• threadCount – number of CPU threads. If MPI option is set, threadCount=1

• rng_seed – int, seed used for randomness in PSO. If None, a random seed is generated.

• vectorization_batch_size – int, only relevant for GPU, determines the number of particles/walkers whose logL will be computed simultaneously. None defaults to one particle at a time and 0 means to compute all particles/walkers simultaneously.

Returns:

result of the best fit, the PSO chain of the best fit parameter after each iteration [lnlikelihood, parameters, velocities], list of parameters in same order as in chain

nested_sampling(sampler_type='dynesty', kwargs_run={}, prior_type='uniform', width_scale=1, sigma_scale=1, output_basename='chain', remove_output_dir=True, dypolychord_dynamic_goal=0.8, polychord_settings={}, dypolychord_seed_increment=200, output_dir='nested_sampling_chains', dynesty_bound='multi', dynesty_sample='auto')

Run (Dynamic) Nested Sampling algorithms, depending on the type of algorithm.

Parameters:

• sampler_type – ‘MULTINEST’, ‘DYPOLYCHORD’, ‘DYNESTY’

• kwargs_run – keywords passed to the core sampling method

• prior_type – ‘uniform’ of ‘gaussian’, for converting the unit hypercube to param cube :param width_scale: scale the width (lower/upper limits) of the parameters space by this factor

• sigma_scale – if prior_type is ‘gaussian’, scale the gaussian sigma by this factor

• output_basename – name of the folder in which the core MultiNest/PolyChord code will save output files

• remove_output_dir – if True, the above folder is removed after completion

• dypolychord_dynamic_goal – dynamic goal for DyPolyChord (trade-off between evidence (0) and posterior (1) computation) :param polychord_settings: settings dictionary to send to pypolychord. Check dypolychord documentation for details.

• dypolychord_seed_increment – seed increment for dypolychord with MPI. Check dypolychord documentation for details.

• dynesty_bound – see https://dynesty.readthedocs.io

• sampler_type – ‘MULTINEST’, ‘DYPOLYCHORD’, ‘DYNESTY’

• kwargs_run – keywords passed to the core sampling method

• prior_type – ‘uniform’ of ‘gaussian’, for converting the unit hypercube to param cube

• width_scale – scale the width (lower/upper limits) of the parameters space by this factor

• sigma_scale – if prior_type is ‘gaussian’, scale the gaussian sigma by this factor

• output_basename – name of the folder in which the core MultiNest/PolyChord code will save output files

• remove_output_dir – if True, the above folder is removed after completion

• dypolychord_dynamic_goal – dynamic goal for DyPolyChord (trade-off between evidence (0) and posterior (1) computation)

• polychord_settings – settings dictionary to send to pypolychord. Check dypolychord documentation for details.

• dypolychord_seed_increment – seed increment for dypolychord with MPI. Check dypolychord documentation for details.

• dynesty_bound – see https://dynesty.readthedocs.io for details

• dynesty_sample – see https://dynesty.readthedocs.io for details

Returns:

list of output arguments : samples, mean inferred values, log- likelihood, log-evidence, error on log-evidence for each sample

psf_iteration(compute_bands=None, **kwargs_psf_iter)

Iterative PSF reconstruction.

Parameters:

• compute_bands – bool list, if multiple bands, this process can be limited to a subset of bands

• kwargs_psf_iter – keyword arguments as used or available in PSFIteration.update_iterative() definition

Returns:

0, updated PSF is stored in self.multi_band_list

align_images(n_particles=10, n_iterations=10, align_offset=True, align_rotation=False, threadCount=1, compute_bands=None, delta_shift=0.2, delta_rot=0.1)

Aligns the coordinate systems of different exposures within a fixed model parameterisation by executing a PSO with relative coordinate shifts as free parameters.

Parameters:

• n_particles – number of particles in the Particle Swarm Optimization

• n_iterations – number of iterations in the optimization process

• align_offset (boolean) – aligns shift in Ra and Dec

• align_rotation (boolean) – aligns coordinate rotation

• delta_shift – astrometric shift tolerance

• delta_rot – rotation angle tolerance [in radian]

• compute_bands – bool list, if multiple bands, this process can be limited to a subset of bands for which the coordinate system is being fit for best alignment to the model parameters

Returns:

0, updated coordinate system for the band(s)

update_settings(kwargs_model=None, kwargs_constraints=None, kwargs_likelihood=None, lens_add_fixed=None, source_add_fixed=None, lens_light_add_fixed=None, ps_add_fixed=None, special_add_fixed=None, tracer_source_add_fixed=None, lens_remove_fixed=None, source_remove_fixed=None, lens_light_remove_fixed=None, ps_remove_fixed=None, special_remove_fixed=None, tracer_source_remove_fixed=None, change_source_lower_limit=None, change_source_upper_limit=None, change_lens_lower_limit=None, change_lens_upper_limit=None, change_sigma_lens=None, change_sigma_source=None, change_sigma_lens_light=None)

Updates lenstronomy settings “on the fly”.

Parameters:

• kwargs_model – kwargs, specified keyword arguments overwrite the existing ones

• kwargs_constraints – kwargs, specified keyword arguments overwrite the existing ones

• kwargs_likelihood – kwargs, specified keyword arguments overwrite the existing ones

• lens_add_fixed – [[i_model, [‘param1’, ‘param2’,…], […]]

• source_add_fixed – [[i_model, [‘param1’, ‘param2’,…], […]]

• lens_light_add_fixed – [[i_model, [‘param1’, ‘param2’,…], […]]

• ps_add_fixed – [[i_model, [‘param1’, ‘param2’,…], […]]

• special_add_fixed – [‘param1’, ‘param2’,…]

• special_add_fixed – [‘param1’, ‘param2’,…]

• tracer_source_add_fixed – [[i_model, [‘param1’, ‘param2’,…], […]]

• lens_remove_fixed – [[i_model, [‘param1’, ‘param2’,…], […]]

• source_remove_fixed – [[i_model, [‘param1’, ‘param2’,…], […]]

• lens_light_remove_fixed – [[i_model, [‘param1’, ‘param2’,…], […]]

• ps_remove_fixed – [[i_model, [‘param1’, ‘param2’,…], […]]

• special_remove_fixed – [‘param1’, ‘param2’,…]

• special_remove_fixed – [‘param1’, ‘param2’,…]

• tracer_source_remove_fixed – [[i_model, [‘param1’, ‘param2’,…], […]]

• change_lens_lower_limit – [[i_model, [‘param_name1’, ‘param_name2’, …], [value1, value2, …]]]

• change_lens_upper_limit – [[i_model, [‘param_name1’, ‘param_name2’, …], [value1, value2, …]]]

• change_source_lower_limit – [[i_model, [‘param_name1’, ‘param_name2’, …], [value1, value2, …]]]

• change_source_upper_limit – [[i_model, [‘’param_name1’, ‘param_name2’, …], [value1, value2, …]]]

• change_sigma_lens – [[i_model, [‘param_name1’, ‘param_name2’, …], [value1, value2, …]]]

• change_sigma_source – [[i_model, [‘param_name1’, ‘param_name2’, …], [value1, value2, …]]]

• change_sigma_lens_light – [[i_model, [‘param_name1’, ‘param_name2’, …], [value1, value2, …]]]

Returns:

0, the settings are overwritten for the next fitting step to come

set_param_value(**kwargs)

Set a parameter to a specific value. kwargs are below.

Parameters:

• lens – [[i_model, [‘param1’, ‘param2’,…], […]]

• source – [[i_model, [‘param1’, ‘param2’,…], […]]

• lens_light – [[i_model, [‘param1’, ‘param2’,…], […]]

• ps – [[i_model, [‘param1’, ‘param2’,…], […]]

Returns:

0, the value of the param is overwritten

Return type:

fix_not_computed(free_bands)

Fixes lens model parameters of imaging bands/frames that are not computed and frees the parameters of the other lens models to the initial kwargs_fixed options.

Parameters:

free_bands – bool list of length of imaging bands in order of imaging bands, if False: set fixed lens model

Returns:

None

best_fit_from_samples(samples, logl)

Return best fit (max likelihood) value of samples in lenstronomy conventions.

Parameters:

• samples – samples of multi-dimensional parameter space

• logl – likelihood values for each sample

Returns:

kwargs_result in lenstronomy convention

property psf_iteration_memory

returns all the psf iterations performed in the FittingSequence It stores in a list of dictionaries: “sequence”: what PSF sequence it is (0, 1 etc) “band”: index of the imaging band that is being corrected “psf_before” kwargs_psf prior to the iteration “psf_after” kwargs_psf as a result of the iteration

Returns:

list of all psf corrections

Module contents