__author__ = 'dgilman'
from scipy.optimize import minimize
import numpy as np
from lenstronomy.Sampling.Samplers.pso import ParticleSwarmOptimizer
from lenstronomy.LensModel.QuadOptimizer.multi_plane_fast import MultiplaneFast
from lenstronomy.Sampling.Pool.pool import choose_pool
__all__ = ['Optimizer']
[docs]class Optimizer(object):
"""
class which executes the optimization routines. Currently implemented as a particle swarm optimization followed by
a downhill simplex routine.
Particle swarm optimizer is modified from the CosmoHammer particle swarm routine with different convergence criteria implemented.
"""
def __init__(self, x_image, y_image, lens_model_list, redshift_list, z_lens, z_source,
parameter_class, astropy_instance=None, numerical_alpha_class=None,
particle_swarm=True, re_optimize=False, re_optimize_scale=1.,
pso_convergence_mean=50000, foreground_rays=None,
tol_source=1e-5, tol_simplex_func=1e-3, simplex_n_iterations=400):
"""
:param x_image: x_image to fit (should be length 4)
:param y_image: y_image to fit (should be length 4)
:param lens_model_list: list of lens models for the system
:param redshift_list: list of lens redshifts for the system
:param z_lens: the main deflector redshift, the lens models being optimizer must be at this redshift
:param z_source: the source redshift
:param parameter_class: an instance of ParamClass (see documentation in QuadOptimizer.param_manager)
:param astropy_instance: an instance of astropy to pass to the lens model
:param numerical_alpha_class: a class to compute numerical deflection angles to pass to the lens model
:param particle_swarm: bool, whether or not to use a PSO fit first
:param re_optimize: bool, if True the initial spread of particles will be very tight
:param re_optimize_scale: float, controls how tight the initial spread of particles is
:param pso_convergence_mean: when to terminate the PSO fit
:param foreground_rays: (optional) can pass in pre-computed foreground light rays from a previous fit
so as to not waste time recomputing them
:param tol_source: sigma in the source plane chi^2
:param tol_simplex_func: tolerance for the downhill simplex optimization
:param simplex_n_iterations: number of iterations per dimension for the downhill simplex optimization
"""
self.fast_rayshooting = MultiplaneFast(x_image, y_image, z_lens, z_source,
lens_model_list, redshift_list, astropy_instance, parameter_class,
foreground_rays, tol_source, numerical_alpha_class)
self._tol_source = tol_source
self._pso_convergence_mean = pso_convergence_mean
self._param_class = parameter_class
self._tol_simplex_func = tol_simplex_func
self._simplex_n_iterations = simplex_n_iterations
self._particle_swarm = particle_swarm
self._re_optimize = re_optimize
self._re_optimize_scale = re_optimize_scale
[docs] def optimize(self, n_particles=50, n_iterations=250, verbose=False, threadCount=1):
"""
:param n_particles: number of PSO particles, will be ignored if self._particle_swarm is False
:param n_iterations: number of PSO iterations, will be ignored if self._particle_swarm is False
:param verbose: whether to print stuff
:param pool: instance of Pool with a map method for multiprocessing
:return: keyword arguments that map (x_image, y_image) to the same source coordinate (source_x, source_y)
"""
if self._particle_swarm:
if threadCount > 1:
pool = choose_pool(mpi=False, processes=threadCount)
else:
pool = None
kwargs = self._fit_pso(n_particles, n_iterations, pool, verbose)
else:
kwargs = self._param_class.kwargs_lens
kwargs_lens_final, source_penalty = self._fit_amoeba(kwargs, verbose)
args_lens_final = self._param_class.kwargs_to_args(kwargs_lens_final)
source_x_array, source_y_array = self.fast_rayshooting.ray_shooting_fast(args_lens_final)
source_x, source_y = np.mean(source_x_array), np.mean(source_y_array)
if verbose:
print('optimization done.')
print('Recovered source position: ', (source_x_array, source_y_array))
return kwargs_lens_final, [source_x, source_y]
def _fit_pso(self, n_particles, n_iterations, pool, verbose):
"""
Executes the PSO
"""
low_bounds, high_bounds = self._param_class.bounds(self._re_optimize, self._re_optimize_scale)
pso = ParticleSwarmOptimizer(self.fast_rayshooting.logL, low_bounds, high_bounds, n_particles,
pool, args=[self._tol_source])
best, info = pso.optimize(n_iterations, verbose, early_stop_tolerance=self._pso_convergence_mean)
if verbose:
print('PSO done... ')
print('source plane chi^2: ', self.fast_rayshooting.source_plane_chi_square(best))
print('total chi^2: ', self.fast_rayshooting.chi_square(best))
kwargs = self._param_class.args_to_kwargs(best)
return kwargs
def _fit_amoeba(self, kwargs, verbose):
"""
Executes the downhill simplex
"""
args_init = self._param_class.kwargs_to_args(kwargs)
options = {'adaptive': True, 'fatol': self._tol_simplex_func,
'maxiter': self._simplex_n_iterations * len(args_init)}
method = 'Nelder-Mead'
if verbose:
print('starting amoeba... ')
opt = minimize(self.fast_rayshooting.chi_square, x0=args_init,
method=method, options=options)
kwargs = self._param_class.args_to_kwargs(opt['x'])
source_penalty = opt['fun']
return kwargs, source_penalty