From ae370761a3dece3b5fa51cb46aa4e9311e663f6a Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sat, 1 Feb 2025 21:08:13 -0500 Subject: [PATCH 01/25] implement fitting sequence --- jaxtronomy/Workflow/fitting_sequence.py | 958 +++++++++++++++++ test/test_Workflow/test_fitting_sequence.py | 1037 +++++++++++++++++++ 2 files changed, 1995 insertions(+) create mode 100644 jaxtronomy/Workflow/fitting_sequence.py create mode 100644 test/test_Workflow/test_fitting_sequence.py diff --git a/jaxtronomy/Workflow/fitting_sequence.py b/jaxtronomy/Workflow/fitting_sequence.py new file mode 100644 index 0000000..b0ca737 --- /dev/null +++ b/jaxtronomy/Workflow/fitting_sequence.py @@ -0,0 +1,958 @@ +import copy + +from jaxtronomy.Sampling.likelihood import LikelihoodModule +from jaxtronomy.ImSim.MultiBand.single_band_multi_model import SingleBandMultiModel +from jaxtronomy.Sampling.Samplers.jaxopt_minimizer import JaxoptMinimizer + +from lenstronomy.Workflow.psf_fitting import PsfFitting +from lenstronomy.Workflow.alignment_matching import AlignmentFitting +from lenstronomy.Workflow.flux_calibration import FluxCalibration +from lenstronomy.Workflow.multi_band_manager import MultiBandUpdateManager +from lenstronomy.Sampling.sampler import Sampler +from lenstronomy.Sampling.Samplers.multinest_sampler import MultiNestSampler +from lenstronomy.Sampling.Samplers.polychord_sampler import DyPolyChordSampler +from lenstronomy.Sampling.Samplers.dynesty_sampler import DynestySampler +from lenstronomy.Sampling.Samplers.nautilus_sampler import NautilusSampler +from lenstronomy.Sampling.Samplers.cobaya_sampler import CobayaSampler + +import jax +import numpy as np +import lenstronomy.Util.analysis_util as analysis_util + +jax.config.update("jax_enable_x64", True) + + +__all__ = ["FittingSequence"] + + +class FittingSequence(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.""" + + def __init__( + self, + kwargs_data_joint, + kwargs_model, + kwargs_constraints, + kwargs_likelihood, + kwargs_params, + mpi=False, + verbose=True, + ): + """ + + :param kwargs_data_joint: keyword argument specifying the data according to LikelihoodModule + :param kwargs_model: keyword arguments to describe all model components used in + class_creator.create_class_instances() + :param 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) + :param kwargs_likelihood: keyword arguments of the Likelihood() class to handle parameters and settings of the + likelihood + :param 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] + :param mpi: MPI option (bool), if True, will launch an MPI Pool job for the steps in the fitting sequence where + possible + :param verbose: bool, if True prints temporary results and indicators of the fitting process + """ + self.kwargs_data_joint = kwargs_data_joint + self.multi_band_list = kwargs_data_joint.get("multi_band_list", []) + self.multi_band_type = kwargs_data_joint.get("multi_band_type", "single-band") + self._verbose = verbose + self._mpi = mpi + self._updateManager = MultiBandUpdateManager( + kwargs_model, + kwargs_constraints, + kwargs_likelihood, + kwargs_params, + num_bands=len(self.multi_band_list), + ) + self._mcmc_init_samples = None + self._psf_iteration_memory = [] + self._psf_iteration_index = 0 # index of the sequence of the PSF iteration (how many times it is being run) + + @property + def kwargs_fixed(self): + """Returns the updated kwargs_fixed from the update manager. + + :return: list of fixed kwargs, see UpdateManager() + """ + return self._updateManager.fixed_kwargs + + def fit_sequence(self, fitting_list): + """ + + :param fitting_list: list of [['string', {kwargs}], ..] with 'string being the specific fitting option and + kwargs being the arguments passed to this option + :return: fitting results + """ + chain_list = [] + for i, fitting in enumerate(fitting_list): + fitting_type = fitting[0] + kwargs = fitting[1] + + if fitting_type in [ + "PSO", + "SIMPLEX", + "MCMC", + "emcee", + "zeus", + "Cobaya", + "dynesty", + "dyPolyChord", + "MultiNest", + "nested_sampling", + "Nautilus", + ]: + self._updateManager.check_initial_state() + if fitting_type == "restart": + self._updateManager.set_init_state() + self._updateManager.check_initial_state() + + elif fitting_type == "update_settings": + self.update_settings(**kwargs) + + elif fitting_type == "set_param_value": + self.set_param_value(**kwargs) + + elif fitting_type == "fix_not_computed": + self.fix_not_computed(**kwargs) + + elif fitting_type == "psf_iteration": + raise ValueError("psf_iteration not supported in jaxtronomy") + #self.psf_iteration(**kwargs) + + elif fitting_type == "align_images": + self.align_images(**kwargs) + + elif fitting_type == "calibrate_images": + self.flux_calibration(**kwargs) + + elif fitting_type == "PSO": + kwargs_result, chain, param = self.pso(**kwargs) + self._updateManager.update_param_state(**kwargs_result) + + chain_list.append([fitting_type, chain, param]) + + elif fitting_type == "SIMPLEX": + kwargs_result = self.simplex(**kwargs) + self._updateManager.update_param_state(**kwargs_result) + chain_list.append([fitting_type, kwargs_result]) + + elif fitting_type in ["MCMC", "emcee", "zeus"]: + if fitting_type == "MCMC": + print("MCMC selected. Sampling with default option emcee.") + fitting_type = "emcee" + if "init_samples" not in kwargs: + kwargs["init_samples"] = self._mcmc_init_samples + elif kwargs["init_samples"] is None: + kwargs["init_samples"] = self._mcmc_init_samples + mcmc_output = self.mcmc(**kwargs, sampler_type=fitting_type) + kwargs_result = self._result_from_mcmc(mcmc_output) + self._updateManager.update_param_state(**kwargs_result) + chain_list.append(mcmc_output) + + elif fitting_type == "Cobaya": + print("Using the Metropolis--Hastings MCMC sampler in Cobaya.") + param_class = self.param_class + kwargs_temp = self._updateManager.parameter_state + mean_start = param_class.kwargs2args(**kwargs_temp) + kwargs_sigma = self._updateManager.sigma_kwargs + sigma_start = np.array(param_class.kwargs2args(**kwargs_sigma)) + # pass the likelihood and starting info to the sampler + sampler = CobayaSampler(self.likelihoodModule, mean_start, sigma_start) + # run the sampler + updated_info, sampler_type, best_fit_values = sampler.run(**kwargs) + # change the best-fit values returned by cobaya into lenstronomy kwargs format + best_fit_kwargs = self.param_class.args2kwargs( + best_fit_values, bijective=True + ) + # collect the products + mh_output = [updated_info, sampler_type, best_fit_kwargs] + # append the products to the chain list + chain_list.append(mh_output) + + elif fitting_type in [ + "dynesty", + "dyPolyChord", + "MultiNest", + "nested_sampling", + ]: + if fitting_type == "nested_sampling": + print( + "Nested sampling selected. Sampling with default option dynesty." + ) + fitting_type = "dynesty" + ns_output = self.nested_sampling(**kwargs, sampler_type=fitting_type) + chain_list.append(ns_output) + + elif fitting_type == "Nautilus": + # do importance nested sampling with Nautilus + nautilus = NautilusSampler( + likelihood_module=self.likelihoodModule, mpi=self._mpi, **kwargs + ) + points, log_w, log_l, log_z = nautilus.run(**kwargs) + chain_list.append([points, log_w, log_l, log_z]) + if kwargs.get("verbose", False): + print(len(points), "number of points sampled") + kwargs_result = self.best_fit_from_samples(points, log_l) + self._updateManager.update_param_state(**kwargs_result) + + elif fitting_type == "Jaxopt": + args_history, logL_history, kwargs_result = self.jaxopt(**kwargs) + self._updateManager.update_param_state(**kwargs_result) + chain_list.append([fitting_type, args_history, logL_history, kwargs_result]) + else: + raise ValueError( + "fitting_sequence {} is not supported. Please use: 'PSO', 'SIMPLEX', " + "'MCMC' or 'emcee', 'zeus', 'Cobaya', " + "'dynesty', 'dyPolyChord', 'Multinest', 'Nautilus, '" + "'psf_iteration', 'restart', 'update_settings', 'calibrate_images' or " + "'align_images'".format(fitting_type) + ) + + return chain_list + + def best_fit(self, bijective=False): + """ + + :param 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. + :return: best fit model of the current state of the FittingSequence class + """ + + return self._updateManager.best_fit(bijective=bijective) + + def update_state(self, kwargs_update): + """Updates current best fit state to the input model keywords specified. + + :param kwargs_update: format of kwargs_result + :return: None + """ + self._updateManager.update_param_state(**kwargs_update) + + def best_fit_likelihood(self, verbose=False): + """Returns the log likelihood of the best fit model of the current state of this + class. + + :param verbose: bool, if True, prints likelihood statements + :return: log likelihood, float + """ + kwargs_result = self.best_fit(bijective=True) + param_class = self.param_class + likelihoodModule = self.likelihoodModule + logL = likelihoodModule.logL( + param_class.kwargs2args(**kwargs_result), verbose=verbose + ) + return logL + + @property + def bic(self): + """Bayesian information criterion (BIC) of the model. + + :return: bic value, float + """ + num_data = self.likelihoodModule.num_data + num_param_nonlinear = self.param_class.num_param()[0] + num_param_linear = self.param_class.num_param_linear() + num_param = num_param_nonlinear + num_param_linear + bic = analysis_util.bic_model( + self.best_fit_likelihood(verbose=False), num_data, num_param + ) + return bic + + @property + def param_class(self): + """ + + :return: Param() class instance reflecting the current state of FittingSequence + """ + return self._updateManager.param_class + + + @property + def likelihoodModule(self): + """ + + :return: Likelihood() class instance reflecting the current state of FittingSequence + """ + kwargs_model = self._updateManager.kwargs_model + kwargs_likelihood = self._updateManager.kwargs_likelihood + likelihoodModule = LikelihoodModule( + self.kwargs_data_joint, kwargs_model, self.param_class, **kwargs_likelihood + ) + return likelihoodModule + + + def simplex(self, n_iterations, method="Nelder-Mead"): + """Downhill simplex optimization using the Nelder-Mead algorithm. + + :param n_iterations: maximum number of iterations to perform + :param method: the optimization method used, see documentation in + scipy.optimize.minimize + :return: result of the best fit + """ + + param_class = self.param_class + kwargs_temp = self._updateManager.parameter_state + init_pos = param_class.kwargs2args(**kwargs_temp) + sampler = Sampler(likelihoodModule=self.likelihoodModule) + result = sampler.simplex(init_pos, n_iterations, method) + + kwargs_result = param_class.args2kwargs(result, bijective=True) + return kwargs_result + + def mcmc( + self, + n_burn, + n_run, + walkerRatio=None, + n_walkers=None, + sigma_scale=1, + threadCount=1, + init_samples=None, + re_use_samples=True, + sampler_type="emcee", + progress=True, + backend_filename=None, + start_from_backend=False, + **kwargs_zeus + ): + """MCMC routine. + + :param n_burn: number of burn in iterations (will not be saved) + :param n_run: number of MCMC iterations that are saved + :param walkerRatio: ratio of walkers/number of free parameters + :param n_walkers: integer, number of walkers of emcee (optional, if set, overwrites the walkerRatio input + :param sigma_scale: scaling of the initial parameter spread relative to the width in the initial settings + :param threadCount: number of CPU threads. If MPI option is set, threadCount=1 + :param init_samples: initial sample from where to start the MCMC process + :param re_use_samples: bool, if True, re-uses the samples described in init_samples.nOtherwise starts from + scratch. + :param sampler_type: string, which MCMC sampler to be used. Options are 'emcee' and 'zeus' + :param progress: boolean, if True shows progress bar in EMCEE + :param backend_filename: name of the HDF5 file where sampling state is saved (through emcee backend engine) + :type backend_filename: string + :param start_from_backend: 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!). + :type start_from_backend: bool + :param kwargs_zeus: zeus-specific kwargs + :return: list of output arguments, e.g. MCMC samples, parameter names, logL distances of all samples specified + by the specific sampler used + """ + param_class = self.param_class + # run PSO + mcmc_class = Sampler(likelihoodModule=self.likelihoodModule) + kwargs_temp = self._updateManager.parameter_state + mean_start = param_class.kwargs2args(**kwargs_temp) + kwargs_sigma = self._updateManager.sigma_kwargs + sigma_start = np.array(param_class.kwargs2args(**kwargs_sigma)) * sigma_scale + num_param, param_list = param_class.num_param() + if n_walkers is None: + if walkerRatio is None: + raise ValueError( + "MCMC sampler needs either n_walkers or walkerRatio as input argument" + ) + n_walkers = num_param * walkerRatio + # run MCMC + if init_samples is not None and re_use_samples is True: + num_samples, num_param_prev = np.shape(init_samples) + if num_param_prev == num_param: + print("re-using previous samples to initialize the next MCMC run.") + idxs = np.random.choice(len(init_samples), n_walkers) + initpos = init_samples[idxs] + else: + raise ValueError( + "Can not re-use previous MCMC samples as number of parameters have changed!" + ) + else: + initpos = None + + if sampler_type == "zeus": + # check if zeus is specified, if not default to emcee + samples, dist = mcmc_class.mcmc_zeus( + n_walkers, + n_run, + n_burn, + mean_start, + sigma_start, + mpi=self._mpi, + threadCount=threadCount, + progress=progress, + initpos=initpos, + backend_filename=backend_filename, + **kwargs_zeus + ) + output = [sampler_type, samples, param_list, dist] + else: + # sample with emcee + samples, dist = mcmc_class.mcmc_emcee( + n_walkers, + n_run, + n_burn, + mean_start, + sigma_start, + mpi=self._mpi, + threadCount=threadCount, + progress=progress, + initpos=initpos, + backend_filename=backend_filename, + start_from_backend=start_from_backend, + ) + output = [sampler_type, samples, param_list, dist] + + self._mcmc_init_samples = samples # overwrites previous samples to continue from there in the next MCMC run + return output + + def jaxopt( + self, + method="BFGS", + maxiter=100, + ): + """Uses the Jaxopt Scipy Minimizer + + :param method: string, options are BFGS, Nelder-Mead, Powell, CG, BFGS, Newton-CG, + L-BFGS-B, TNC, COBYLA, SLSQP, trust-constr, dogleg, trust-ncg, trust-exact, trust-krylov + :param maxiter: int, number of iterations to perform during minimization of the loss function + """ + print(f"Running {method} minimization with {maxiter} iterations:") + param_class = self.param_class + likelihood_module=self.likelihoodModule + + # Coonverts kwargs to args for the mean, sigma, lower, and upper parameter values + kwargs_temp = self._updateManager.parameter_state + args_mean = param_class.kwargs2args(**kwargs_temp) + + kwargs_sigma = self._updateManager.sigma_kwargs + args_sigma = param_class.kwargs2args(**kwargs_sigma) + + kwargs_lower = self._updateManager._lower_kwargs + args_lower = param_class.kwargs2args(*kwargs_lower) + + kwargs_upper = self._updateManager._upper_kwargs + args_upper = param_class.kwargs2args(*kwargs_upper) + + # Initialize the solver class + minimizer = JaxoptMinimizer( + method=method, + logL_func=likelihood_module.logL, + args_mean=args_mean, + args_sigma=args_sigma, + args_lower=args_lower, + args_upper=args_upper, + maxiter=maxiter + ) + + # Runs the minimizer and prints results + args_result, final_logL = minimizer.run(args_mean) + kwargs_result = param_class.args2kwargs(args_result) + print("best fit log_likelihood:", final_logL) + print("kwargs_result:", kwargs_result) + + return minimizer.parameter_history, minimizer.logL_history, kwargs_result + + + def pso( + self, n_particles, n_iterations, sigma_scale=1, print_key="PSO", threadCount=1 + ): + """Particle Swarm Optimization. + + :param n_particles: number of particles in the Particle Swarm Optimization + :param n_iterations: number of iterations in the optimization process + :param sigma_scale: scaling of the initial parameter spread relative to the + width in the initial settings + :param print_key: string, printed text when executing this routine + :param threadCount: number of CPU threads. If MPI option is set, threadCount=1 + :return: 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 + """ + + param_class = self.param_class + kwargs_temp = self._updateManager.parameter_state + init_pos = param_class.kwargs2args(**kwargs_temp) + kwargs_sigma = self._updateManager.sigma_kwargs + sigma_start = param_class.kwargs2args(**kwargs_sigma) + lower_start = np.array(init_pos) - np.array(sigma_start) * sigma_scale + upper_start = np.array(init_pos) + np.array(sigma_start) * sigma_scale + + num_param, param_list = param_class.num_param() + # run PSO + sampler = Sampler(likelihoodModule=self.likelihoodModule) + result, chain = sampler.pso( + n_particles, + n_iterations, + lower_start, + upper_start, + init_pos=init_pos, + threadCount=threadCount, + mpi=self._mpi, + print_key=print_key, + verbose=self._verbose, + ) + kwargs_result = param_class.args2kwargs(result, bijective=True) + return kwargs_result, chain, param_list + + def nested_sampling( + self, + 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. + + :param sampler_type: 'MULTINEST', 'DYPOLYCHORD', 'DYNESTY' + :param kwargs_run: keywords passed to the core sampling method + :param 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 + :param sigma_scale: if prior_type is 'gaussian', scale the gaussian sigma by + this factor + :param output_basename: name of the folder in which the core + MultiNest/PolyChord code will save output files + :param remove_output_dir: if True, the above folder is removed after completion + :param 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. + :param dypolychord_seed_increment: seed increment for dypolychord with MPI. + Check dypolychord documentation for details. + :param dynesty_bound: see https://dynesty.readthedocs.io + :param sampler_type: 'MULTINEST', 'DYPOLYCHORD', 'DYNESTY' + :param kwargs_run: keywords passed to the core sampling method + :param 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 + :param sigma_scale: if prior_type is 'gaussian', scale the gaussian sigma by + this factor + :param output_basename: name of the folder in which the core MultiNest/PolyChord + code will save output files + :param remove_output_dir: if True, the above folder is removed after completion + :param 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. + :param dypolychord_seed_increment: seed increment for dypolychord with MPI. + Check dypolychord documentation for details. + :param dynesty_bound: see https://dynesty.readthedocs.io for details + :param dynesty_sample: see https://dynesty.readthedocs.io for details + :return: list of output arguments : samples, mean inferred values, log- + likelihood, log-evidence, error on log-evidence for each sample + """ + mean_start, sigma_start = self._prepare_sampling(prior_type) + + if sampler_type == "dyPolyChord": + if "resume_dyn_run" in kwargs_run and kwargs_run["resume_dyn_run"] is True: + resume_dyn_run = True + else: + resume_dyn_run = False + sampler = DyPolyChordSampler( + self.likelihoodModule, + prior_type=prior_type, + prior_means=mean_start, + prior_sigmas=sigma_start, + width_scale=width_scale, + sigma_scale=sigma_scale, + output_dir=output_dir, + output_basename=output_basename, + polychord_settings=polychord_settings, + remove_output_dir=remove_output_dir, + resume_dyn_run=resume_dyn_run, + use_mpi=self._mpi, + ) + samples, means, logZ, logZ_err, logL, results_object = sampler.run( + dypolychord_dynamic_goal, kwargs_run + ) + + elif sampler_type == "MultiNest": + sampler = MultiNestSampler( + self.likelihoodModule, + prior_type=prior_type, + prior_means=mean_start, + prior_sigmas=sigma_start, + width_scale=width_scale, + sigma_scale=sigma_scale, + output_dir=output_dir, + output_basename=output_basename, + remove_output_dir=remove_output_dir, + use_mpi=self._mpi, + ) + samples, means, logZ, logZ_err, logL, results_object = sampler.run( + kwargs_run + ) + else: + sampler = DynestySampler( + self.likelihoodModule, + prior_type=prior_type, + prior_means=mean_start, + prior_sigmas=sigma_start, + width_scale=width_scale, + sigma_scale=sigma_scale, + bound=dynesty_bound, + sample=dynesty_sample, + use_mpi=self._mpi, + ) + samples, means, logZ, logZ_err, logL, results_object = sampler.run( + kwargs_run + ) + + # update current best fit values + self._update_state(samples[-1]) + + output = [ + sampler_type, + samples, + sampler.param_names, + logL, + logZ, + logZ_err, + results_object, + ] + + return output + + def psf_iteration(self, compute_bands=None, **kwargs_psf_iter): + """Iterative PSF reconstruction. + + :param compute_bands: bool list, if multiple bands, this process can be limited + to a subset of bands + :param kwargs_psf_iter: keyword arguments as used or available in + PSFIteration.update_iterative() definition + :return: 0, updated PSF is stored in self.multi_band_list + """ + kwargs_model = self._updateManager.kwargs_model + kwargs_likelihood = self._updateManager.kwargs_likelihood + likelihood_mask_list = kwargs_likelihood.get("image_likelihood_mask_list", None) + kwargs_pixelbased = kwargs_likelihood.get("kwargs_pixelbased", None) + kwargs_temp = self.best_fit(bijective=False) + if compute_bands is None: + compute_bands = [True] * len(self.multi_band_list) + + for band_index in range(len(self.multi_band_list)): + if compute_bands[band_index] is True: + kwargs_psf = self.multi_band_list[band_index][1] + kwargs_psf_before = copy.deepcopy(kwargs_psf) + image_model = SingleBandMultiModel( + self.multi_band_list, + kwargs_model, + likelihood_mask_list=likelihood_mask_list, + band_index=band_index, + kwargs_pixelbased=kwargs_pixelbased, + ) + psf_iter = PsfFitting(image_model_class=image_model) + kwargs_psf = psf_iter.update_iterative( + kwargs_psf, kwargs_params=kwargs_temp, **kwargs_psf_iter + ) + self.multi_band_list[band_index][1] = kwargs_psf + self._psf_iteration_memory.append( + { + "sequence": self._psf_iteration_index, + "band": band_index, + "psf_before": kwargs_psf_before, + "psf_after": kwargs_psf, + } + ) + self._psf_iteration_index += 1 + return 0 + + def align_images( + self, + 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. + + :param n_particles: number of particles in the Particle Swarm Optimization + :param n_iterations: number of iterations in the optimization process + :param align_offset: aligns shift in Ra and Dec + :type align_offset: boolean + :param align_rotation: aligns coordinate rotation + :type align_rotation: boolean + :param delta_shift: astrometric shift tolerance + :param delta_rot: rotation angle tolerance [in radian] + :param 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 + :return: 0, updated coordinate system for the band(s) + """ + kwargs_model = self._updateManager.kwargs_model + kwargs_likelihood = self._updateManager.kwargs_likelihood + likelihood_mask_list = kwargs_likelihood.get("image_likelihood_mask_list", None) + kwargs_temp = self.best_fit(bijective=False) + if compute_bands is None: + compute_bands = [True] * len(self.multi_band_list) + + for i in range(len(self.multi_band_list)): + if compute_bands[i] is True: + alignmentFitting = AlignmentFitting( + self.multi_band_list, + kwargs_model, + kwargs_temp, + band_index=i, + likelihood_mask_list=likelihood_mask_list, + align_offset=align_offset, + align_rotation=align_rotation, + ) + + kwargs_data, chain = alignmentFitting.pso( + n_particles=n_particles, + n_iterations=n_iterations, + delta_shift=delta_shift, + delta_rot=delta_rot, + threadCount=threadCount, + mpi=self._mpi, + print_key="Alignment fitting for band %s ..." % i, + ) + print("Align completed for band %s." % i) + print( + "ra_shift: %s, dec_shift: %s, phi_rot: %s" + % ( + kwargs_data.get("ra_shift", 0), + kwargs_data.get("dec_shift", 0), + kwargs_data.get("phi_rot", 0), + ) + ) + self.multi_band_list[i][0] = kwargs_data + return 0 + + def flux_calibration( + self, + n_particles=10, + n_iterations=10, + threadCount=1, + calibrate_bands=None, + scaling_lower_limit=0, + scaling_upper_limit=1000, + ): + """Calibrates flux_scaling between multiple images. This routine only works in + 'join-linear' model when fluxes are meant to be identical for different bands. + + :param n_particles: number of particles in the Particle Swarm Optimization + :param n_iterations: number of iterations in the optimization process + :param calibrate_bands: state which bands the flux calibration is applied to + :type calibrate_bands: list of booleans of length of the imaging bands + :param threadCount: number of CPU threads. If MPI option is set, threadCount=1 + :type threadCount: integer + :param scaling_lower_limit: lower limit of flux_scaling + :param scaling_upper_limit: upper limit of flux scaling + :return: 0, updated coordinate system for the band(s) + """ + kwargs_model = self._updateManager.kwargs_model + kwargs_temp = self.best_fit(bijective=False) + multi_band_type = self.kwargs_data_joint.get("multi_band_type", "multi-linear") + kwargs_imaging = self.likelihoodModule.kwargs_imaging + + calibration_fitting = FluxCalibration( + kwargs_imaging=kwargs_imaging, + kwargs_model=kwargs_model, + kwargs_params=kwargs_temp, + calibrate_bands=calibrate_bands, + ) + + multi_band_list, chain = calibration_fitting.pso( + n_particles=n_particles, + n_iterations=n_iterations, + threadCount=threadCount, + mpi=self._mpi, + scaling_lower_limit=scaling_lower_limit, + scaling_upper_limit=scaling_upper_limit, + ) + self.multi_band_list = multi_band_list + return 0 + + def update_settings( + self, + 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". + + :param kwargs_model: kwargs, specified keyword arguments overwrite the existing ones + :param kwargs_constraints: kwargs, specified keyword arguments overwrite the existing ones + :param kwargs_likelihood: kwargs, specified keyword arguments overwrite the existing ones + :param lens_add_fixed: [[i_model, ['param1', 'param2',...], [...]] + :param source_add_fixed: [[i_model, ['param1', 'param2',...], [...]] + :param lens_light_add_fixed: [[i_model, ['param1', 'param2',...], [...]] + :param ps_add_fixed: [[i_model, ['param1', 'param2',...], [...]] + :param special_add_fixed: ['param1', 'param2',...] + :param special_add_fixed: ['param1', 'param2',...] + :param tracer_source_add_fixed: [[i_model, ['param1', 'param2',...], [...]] + :param lens_remove_fixed: [[i_model, ['param1', 'param2',...], [...]] + :param source_remove_fixed: [[i_model, ['param1', 'param2',...], [...]] + :param lens_light_remove_fixed: [[i_model, ['param1', 'param2',...], [...]] + :param ps_remove_fixed: [[i_model, ['param1', 'param2',...], [...]] + :param special_remove_fixed: ['param1', 'param2',...] + :param special_remove_fixed: ['param1', 'param2',...] + :param tracer_source_remove_fixed: [[i_model, ['param1', 'param2',...], [...]] + :param change_lens_lower_limit: [[i_model, ['param_name1', 'param_name2', ...], [value1, value2, ...]]] + :param change_lens_upper_limit: [[i_model, ['param_name1', 'param_name2', ...], [value1, value2, ...]]] + :param change_source_lower_limit: [[i_model, ['param_name1', 'param_name2', ...], [value1, value2, ...]]] + :param change_source_upper_limit: [[i_model, [''param_name1', 'param_name2', ...], [value1, value2, ...]]] + :param change_sigma_lens: [[i_model, ['param_name1', 'param_name2', ...], [value1, value2, ...]]] + :param change_sigma_source: [[i_model, ['param_name1', 'param_name2', ...], [value1, value2, ...]]] + :param change_sigma_lens_light: [[i_model, ['param_name1', 'param_name2', ...], [value1, value2, ...]]] + :return: 0, the settings are overwritten for the next fitting step to come + """ + self._updateManager.update_options( + kwargs_model, kwargs_constraints, kwargs_likelihood + ) + self._updateManager.update_fixed( + lens_add_fixed, + source_add_fixed, + lens_light_add_fixed, + ps_add_fixed, + special_add_fixed, + tracer_source_add_fixed, + lens_remove_fixed, + source_remove_fixed, + lens_light_remove_fixed, + ps_remove_fixed, + special_remove_fixed, + tracer_source_remove_fixed, + ) + self._updateManager.update_limits( + change_source_lower_limit, + change_source_upper_limit, + change_lens_lower_limit, + change_lens_upper_limit, + ) + self._updateManager.update_sigmas( + change_sigma_lens=change_sigma_lens, + change_sigma_source=change_sigma_source, + change_sigma_lens_light=change_sigma_lens_light, + ) + return 0 + + def set_param_value(self, **kwargs): + """Set a parameter to a specific value. `kwargs` are below. + + :param lens: [[i_model, ['param1', 'param2',...], [...]] + :type lens: + :param source: [[i_model, ['param1', 'param2',...], [...]] + :type source: + :param lens_light: [[i_model, ['param1', 'param2',...], [...]] + :type lens_light: + :param ps: [[i_model, ['param1', 'param2',...], [...]] + :type ps: + :return: 0, the value of the param is overwritten + :rtype: + """ + self._updateManager.update_param_value(**kwargs) + + def fix_not_computed(self, 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. + + :param free_bands: bool list of length of imaging bands in order of imaging + bands, if False: set fixed lens model + :return: None + """ + self._updateManager.fix_not_computed(free_bands=free_bands) + + def _prepare_sampling(self, prior_type): + if prior_type == "gaussian": + mean_start = self.param_class.kwargs2args( + **self._updateManager.parameter_state + ) + sigma_start = self.param_class.kwargs2args( + **self._updateManager.sigma_kwargs + ) + mean_start = np.array(mean_start) + sigma_start = np.array(sigma_start) + else: + mean_start, sigma_start = None, None + return mean_start, sigma_start + + def _update_state(self, result): + """ + + :param result: array of parameters being sampled (e.g. result of MCMC chain) + :return: None, updates the parameter state of the class instance + """ + kwargs_result = self.param_class.args2kwargs(result, bijective=True) + self._updateManager.update_param_state(**kwargs_result) + + def _result_from_mcmc(self, mcmc_output): + """ + + :param mcmc_output: list returned by self.mcmc() + :return: kwargs_result like returned by self.pso(), from best logL MCMC sample + """ + _, samples, _, logL_values = mcmc_output + return self.best_fit_from_samples(samples, logL_values) + + def best_fit_from_samples(self, samples, logl): + """Return best fit (max likelihood) value of samples in lenstronomy conventions. + + :param samples: samples of multi-dimensional parameter space + :param logl: likelihood values for each sample + :return: kwargs_result in lenstronomy convention + """ + # get index of best logL sample + best_fit_index = np.argmax(logl) + best_fit_sample = samples[best_fit_index, :] + best_fit_result = best_fit_sample.tolist() + # get corresponding kwargs + kwargs_result = self.param_class.args2kwargs(best_fit_result, bijective=True) + return kwargs_result + + @property + def psf_iteration_memory(self): + """ + 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 + + :return: list of all psf corrections + """ + return self._psf_iteration_memory diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py new file mode 100644 index 0000000..f18a5f3 --- /dev/null +++ b/test/test_Workflow/test_fitting_sequence.py @@ -0,0 +1,1037 @@ +__author__ = "sibirrer" + +import copy + +import pytest +import numpy.testing as npt +import numpy as np +from lenstronomy.Util import simulation_util as sim_util, util +from jaxtronomy.ImSim.image_model import ImageModel +from jaxtronomy.PointSource.point_source import PointSource +from jaxtronomy.LensModel.lens_model import LensModel +from jaxtronomy.LightModel.light_model import LightModel +from jaxtronomy.Workflow.fitting_sequence import FittingSequence +from jaxtronomy.Data.imaging_data import ImageData +from lenstronomy.Data.psf import PSF + + +class TestFittingSequence(object): + """Test the fitting sequences.""" + + def setup_method(self): + # data specifics + sigma_bkg = 0.05 # background noise per pixel + exp_time = 100 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) + numPix = 10 # cutout pixel size + deltaPix = 0.05 # pixel size in arcsec (area per pixel = deltaPix**2) + fwhm = 0.5 # full width half max of PSF + + # PSF specification + + self.kwargs_data = sim_util.data_configure_simple( + numPix, deltaPix, exp_time, sigma_bkg + ) + data_class = ImageData(**self.kwargs_data) + kwargs_psf_gaussian = { + "psf_type": "GAUSSIAN", + "fwhm": fwhm, + "pixel_size": deltaPix, + "truncation": 3, + } + psf_gaussian = PSF(**kwargs_psf_gaussian) + self.kwargs_psf = { + "psf_type": "PIXEL", + "kernel_point_source": psf_gaussian.kernel_point_source, + } + psf_class = PSF(**self.kwargs_psf) + # 'EXTERNAL_SHEAR': external shear + kwargs_shear = { + "gamma1": 0.01, + "gamma2": 0.01, + } # gamma_ext: shear strength, psi_ext: shear angle (in radian) + kwargs_spemd = { + "theta_E": 1.0, + "gamma": 1.8, + "center_x": 0, + "center_y": 0, + "e1": 0.1, + "e2": 0.1, + } + + lens_model_list = ["EPL", "SHEAR"] + self.kwargs_lens = [kwargs_spemd, kwargs_shear] + lens_model_class = LensModel(lens_model_list=lens_model_list) + kwargs_sersic = { + "amp": 21.0, + "R_sersic": 0.1, + "n_sersic": 2, + "center_x": 0, + "center_y": 0, + } + # 'SERSIC_ELLIPSE': elliptical Sersic profile + kwargs_sersic_ellipse = { + "amp": 13.0, + "R_sersic": 0.6, + "n_sersic": 3, + "center_x": 0, + "center_y": 0, + "e1": 0.1, + "e2": 0.1, + } + + lens_light_model_list = ["SERSIC"] + self.kwargs_lens_light = [kwargs_sersic] + lens_light_model_class = LightModel(light_model_list=lens_light_model_list) + source_model_list = ["SERSIC_ELLIPSE"] + self.kwargs_source = [kwargs_sersic_ellipse] + source_model_class = LightModel(light_model_list=source_model_list) + self.kwargs_ps = [ + {"ra_source": 0.0, "dec_source": 0.0, "source_amp": 1.0} + ] # quasar point source position in the source plane and intrinsic brightness + point_source_list = []#["SOURCE_POSITION"] + point_source_class = PointSource( + point_source_type_list=point_source_list, #fixed_magnification_list=[True] + ) + kwargs_numerics = { + "supersampling_factor": 3, + "supersampling_convolution": True, + "compute_mode": "regular", + #"point_source_supersampling_factor": 1, + } + imageModel = ImageModel( + data_class, + psf_class, + lens_model_class, + source_model_class, + lens_light_model_class, + point_source_class, + kwargs_numerics=kwargs_numerics, + ) + image_sim = sim_util.simulate_simple( + imageModel, + self.kwargs_lens, + self.kwargs_source, + self.kwargs_lens_light, + #self.kwargs_psf, + point_source_add=False + ) + + data_class.update_data(image_sim) + self.data_class = data_class + self.psf_class = psf_class + self.kwargs_data["image_data"] = image_sim + self.kwargs_model = { + "lens_model_list": lens_model_list, + "source_light_model_list": source_model_list, + "lens_light_model_list": lens_light_model_list, + "point_source_model_list": point_source_list, + "fixed_magnification_list": [False], + "index_lens_model_list": [[0, 1]], + "point_source_frame_list": [[0]], + } + self.kwargs_numerics = kwargs_numerics + + #num_source_model = len(source_model_list) + + self.kwargs_constraints = { + #"num_point_source_list": [4], + #"image_plane_source_list": [False] * num_source_model, + "linear_solver": False + } + + self.kwargs_likelihood = { + # This is false by default anyways + "check_positive_flux": False, + } + + lens_sigma = [ + { + "theta_E": 0.1, + "gamma": 0.1, + "e1": 0.1, + "e2": 0.1, + "center_x": 0.1, + "center_y": 0.1, + }, + {"gamma1": 0.1, "gamma2": 0.1}, + ] + lens_lower = [ + { + "theta_E": 0.0, + "gamma": 1.5, + "center_x": -2, + "center_y": -2, + "e1": -0.4, + "e2": -0.4, + }, + {"gamma1": -0.3, "gamma2": -0.3}, + ] + lens_upper = [ + { + "theta_E": 10.0, + "gamma": 2.5, + "center_x": 2, + "center_y": 2, + "e1": 0.4, + "e2": 0.4, + }, + {"gamma1": 0.3, "gamma2": 0.3}, + ] + source_sigma = [ + { + "amp": 0.1, + "R_sersic": 0.05, + "n_sersic": 0.5, + "center_x": 0.1, + "center_y": 0.1, + "e1": 0.1, + "e2": 0.1, + } + ] + source_lower = [ + { + "amp": 0, + "R_sersic": 0.01, + "n_sersic": 0.5, + "center_x": -2, + "center_y": -2, + "e1": -0.4, + "e2": -0.4, + } + ] + source_upper = [ + { + "amp": 100, + "R_sersic": 10, + "n_sersic": 5.5, + "center_x": 2, + "center_y": 2, + "e1": 0.4, + "e2": 0.4, + } + ] + + lens_light_sigma = [ + {"amp": 0.1, "R_sersic": 0.05, "n_sersic": 0.5, "center_x": 0.1, "center_y": 0.1} + ] + lens_light_lower = [ + {"amp": 0, "R_sersic": 0.01, "n_sersic": 0.5, "center_x": -2, "center_y": -2} + ] + lens_light_upper = [ + {"amp": 100, "R_sersic": 10, "n_sersic": 5.5, "center_x": 2, "center_y": 2} + ] + ps_sigma = [{"ra_source": 1, "dec_source": 1, "point_amp": 1}] + + lens_param = ( + self.kwargs_lens, + lens_sigma, + [{}, {"ra_0": 0, "dec_0": 0}], + lens_lower, + lens_upper, + ) + source_param = ( + self.kwargs_source, + source_sigma, + [{}], + source_lower, + source_upper, + ) + lens_light_param = ( + self.kwargs_lens_light, + lens_light_sigma, + [{"center_x": 0}], + lens_light_lower, + lens_light_upper, + ) + ps_param = self.kwargs_ps, ps_sigma, [{}], self.kwargs_ps, self.kwargs_ps + + self.kwargs_params = { + "lens_model": lens_param, + "source_model": source_param, + "lens_light_model": lens_light_param, + "point_source_model": ps_param, + # 'special': special_param + } + image_band = [self.kwargs_data, self.kwargs_psf, self.kwargs_numerics] + multi_band_list = [image_band] + self.kwargs_data_joint = { + "multi_band_list": multi_band_list, + "multi_band_type": "single-band", + } + + def test_simulationAPI_image(self): + npt.assert_almost_equal(self.data_class.data[4, 4], 0.1, decimal=0) + + def test_simulationAPI_psf(self): + npt.assert_almost_equal( + np.sum(self.psf_class.kernel_point_source), 1, decimal=6 + ) + + def test_fitting_sequence(self): + fittingSequence = FittingSequence( + self.kwargs_data_joint, + self.kwargs_model, + self.kwargs_constraints, + self.kwargs_likelihood, + self.kwargs_params, + ) + + kwargs_result = fittingSequence.best_fit(bijective=False) + lens_temp = kwargs_result["kwargs_lens"] + npt.assert_almost_equal( + lens_temp[0]["theta_E"], self.kwargs_lens[0]["theta_E"], decimal=2 + ) + + logL = fittingSequence.best_fit_likelihood() + print(logL, "test") + # print(lens_temp, source_temp, lens_light_temp, ps_temp, special_temp) + assert logL < 0 + bic = fittingSequence.bic + assert bic > 0 + # npt.assert_almost_equal(bic, 20000000220.29376, decimal=-4) + + # npt.assert_almost_equal(logL, -10000000061.792593, decimal=-4) + + n_p = 2 + n_i = 2 + fitting_list = [] + + kwargs_pso = {"sigma_scale": 1, "n_particles": n_p, "n_iterations": n_i} + fitting_list.append(["PSO", kwargs_pso]) + kwargs_align = {"delta_shift": 0.2, "n_particles": 2, "n_iterations": 2} + fitting_list.append(["align_images", kwargs_align]) + kwargs_psf_iter = { + "num_iter": 2, + "psf_iter_factor": 0.5, + "stacking_method": "mean", + "new_procedure": False, + } + #fitting_list.append(["psf_iteration", kwargs_psf_iter]) + fitting_list.append(["restart", None]) + fitting_list.append(["fix_not_computed", {"free_bands": [True]}]) + n_sersic_overwrite = 4 + kwargs_update = { + "lens_light_add_fixed": [[0, ["n_sersic"], [n_sersic_overwrite]]], + "lens_light_remove_fixed": [[0, ["center_x"]]], + "change_source_lower_limit": [[0, ["n_sersic"], [0.1]]], + "change_source_upper_limit": [[0, ["n_sersic"], [10]]], + } + fitting_list.append(["update_settings", kwargs_update]) + + chain_list = fittingSequence.fit_sequence(fitting_list) + ( + lens_fixed, + source_fixed, + lens_light_fixed, + ps_fixed, + special_fixed, + extinction_fixed, + tracer_source_fixed, + ) = fittingSequence._updateManager.fixed_kwargs + kwargs_result = fittingSequence.best_fit(bijective=False) + npt.assert_almost_equal( + kwargs_result["kwargs_lens"][0]["theta_E"], + self.kwargs_lens[0]["theta_E"], + decimal=1, + ) + npt.assert_almost_equal( + fittingSequence._updateManager._lens_light_fixed[0]["n_sersic"], + n_sersic_overwrite, + decimal=8, + ) + npt.assert_almost_equal(lens_light_fixed[0]["n_sersic"], 4, decimal=-1) + assert fittingSequence._updateManager._lower_kwargs[1][0]["n_sersic"] == 0.1 + assert fittingSequence._updateManager._upper_kwargs[1][0]["n_sersic"] == 10 + + # test 'set_param_value' fitting sequence + fitting_list = [ + ["set_param_value", {"lens": [[1, ["gamma1"], [0.013]]]}], + ["set_param_value", {"lens_light": [[0, ["center_x"], [0.009]]]}], + ["set_param_value", {"source": [[0, ["n_sersic"], [2.993]]]}], + #["set_param_value", {"ps": [[0, ["ra_source"], [0.007]]]}], + ] + + fittingSequence.fit_sequence(fitting_list) + + kwargs_set = fittingSequence._updateManager.parameter_state + assert kwargs_set["kwargs_lens"][1]["gamma1"] == 0.013 + assert kwargs_set["kwargs_lens_light"][0]["center_x"] == 0.009 + assert kwargs_set["kwargs_source"][0]["n_sersic"] == 2.993 + #assert kwargs_set["kwargs_ps"][0]["ra_source"] == 0.007 + + from unittest import TestCase + + t = TestCase() + with t.assertRaises(ValueError): + fitting_list_two = [] + fitting_list_two.append(["fake_mcmc_method", kwargs_pso]) + fittingSequence.fit_sequence(fitting_list_two) + + with t.assertRaises(ValueError): + # should raise a value error for n_walkers = walkerRatio = None + fitting_list_three = [] + kwargs_test = {"n_burn": 10, "n_run": 10} + fitting_list_three.append(["emcee", kwargs_test]) + fittingSequence.fit_sequence(fitting_list_three) + + fitting_list4 = [["psf_iteration", kwargs_psf_iter]] + with t.assertRaises(ValueError): + fittingSequence.fit_sequence(fitting_list4) + + #psf_iteration_list = fittingSequence.psf_iteration_memory + #assert len(psf_iteration_list) == 1 + #assert "sequence" in psf_iteration_list[0] + #assert "band" in psf_iteration_list[0] + #assert "psf_before" in psf_iteration_list[0] + #assert "psf_after" in psf_iteration_list[0] + + def test_cobaya(self): + np.random.seed(42) + + # make a basic lens model to fit + # data specifics + sigma_bkg = 0.05 # background noise per pixel + exp_time = 100 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) + numPix = 10 # cutout pixel size + deltaPix = 0.05 # pixel size in arcsec (area per pixel = deltaPix**2) + fwhm = 0.5 # full width half max of PSF + + # PSF specification + kwargs_data = sim_util.data_configure_simple( + numPix, deltaPix, exp_time, sigma_bkg + ) + data_class = ImageData(**kwargs_data) + kwargs_psf_gaussian = { + "psf_type": "GAUSSIAN", + "fwhm": fwhm, + "pixel_size": deltaPix, + "truncation": 3, + } + psf_gaussian = PSF(**kwargs_psf_gaussian) + + # make a lens + lens_model_list = ["SIS"] + kwargs_lens = [{"theta_E": 1.5, "center_x": 0.0, "center_y": 0.0}] + lens_model_class = LensModel(lens_model_list=lens_model_list) + + # make a source + source_model_list = ["SERSIC"] + kwargs_source = [ + { + "amp": 1.0, + "R_sersic": 0.3, + "n_sersic": 3.0, + "center_x": 0.1, + "center_y": 0.1, + } + ] + source_model_class = LightModel(light_model_list=source_model_list) + + kwargs_numerics = { + "supersampling_factor": 1, + "supersampling_convolution": False, + } + + imageModel = ImageModel( + data_class, + psf_gaussian, + lens_model_class, + source_model_class, + kwargs_numerics=kwargs_numerics, + ) + image_sim = sim_util.simulate_simple(imageModel, kwargs_lens, kwargs_source, point_source_add=False) + + data_class.update_data(image_sim) + + kwargs_data["image_data"] = image_sim + + kwargs_model = { + "lens_model_list": lens_model_list, + "source_light_model_list": source_model_list, + } + + lens_fixed = [{"center_x": 0.0, "center_y": 0.0}] + lens_sigma = [{"theta_E": 0.01, "center_x": 0.1, "center_y": 0.1}] + lens_lower = [{"theta_E": 0.1, "center_x": -10, "center_y": -10}] + lens_upper = [{"theta_E": 3.0, "center_x": 10, "center_y": 10}] + + source_fixed = [{}] + source_sigma = [ + {"amp": 0.1, "R_sersic": 0.01, "n_sersic": 0.01, "center_x": 0.01, "center_y": 0.01} + ] + source_lower = [ + {"amp": 0, "R_sersic": 0.01, "n_sersic": 0.5, "center_x": -1, "center_y": -1} + ] + source_upper = [ + {"amp": 100, "R_sersic": 1.0, "n_sersic": 6.0, "center_x": 1, "center_y": 1} + ] + + lens_param = [kwargs_lens, lens_sigma, lens_fixed, lens_lower, lens_upper] + source_param = [ + kwargs_source, + source_sigma, + source_fixed, + source_lower, + source_upper, + ] + + kwargs_params = {"lens_model": lens_param, "source_model": source_param} + + multi_band_list = [[kwargs_data, kwargs_psf_gaussian, kwargs_numerics]] + + kwargs_data_joint = { + "multi_band_list": multi_band_list, + "multi_band_type": "single-band", + } + + fittingSequence = FittingSequence( + kwargs_data_joint, + kwargs_model, + self.kwargs_constraints, + self.kwargs_likelihood, + kwargs_params, + ) + + kwargs_cobaya = { + "proposal_widths": [0.001, 0.001, 0.001, 0.001, 0.001, 0.001], + "Rminus1_stop": 100, # does this need to be large? can we run in test mode? + "force_overwrite": True, + } + + chain_list = fittingSequence.fit_sequence([["Cobaya", kwargs_cobaya]]) + + def test_zeus(self): + np.random.seed(42) + # we make a very basic lens+source model to feed to check zeus can be run through fitting sequence + # we don't use the kwargs defined in setup() as those are modified during the tests; using unique kwargs here is safer + + # data specifics + sigma_bkg = 0.05 # background noise per pixel + exp_time = 100 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) + numPix = 10 # cutout pixel size + deltaPix = 0.05 # pixel size in arcsec (area per pixel = deltaPix**2) + fwhm = 0.5 # full width half max of PSF + + # PSF specification + + kwargs_data = sim_util.data_configure_simple( + numPix, deltaPix, exp_time, sigma_bkg + ) + data_class = ImageData(**kwargs_data) + kwargs_psf_gaussian = { + "psf_type": "GAUSSIAN", + "fwhm": fwhm, + "pixel_size": deltaPix, + "truncation": 3, + } + psf_gaussian = PSF(**kwargs_psf_gaussian) + + # make a lens + lens_model_list = ["EPL"] + kwargs_epl = { + "theta_E": 0.6, + "gamma": 2.6, + "center_x": 0.0, + "center_y": 0.0, + "e1": 0.1, + "e2": 0.1, + } + kwargs_lens = [kwargs_epl] + lens_model_class = LensModel(lens_model_list=lens_model_list) + + # make a source + source_model_list = ["SERSIC_ELLIPSE"] + kwargs_sersic_ellipse = { + "amp": 1.0, + "R_sersic": 0.6, + "n_sersic": 3, + "center_x": 0.0, + "center_y": 0.0, + "e1": 0.1, + "e2": 0.1, + } + kwargs_source = [kwargs_sersic_ellipse] + source_model_class = LightModel(light_model_list=source_model_list) + + kwargs_numerics = { + "supersampling_factor": 1, + "supersampling_convolution": False, + } + + imageModel = ImageModel( + data_class, + psf_gaussian, + lens_model_class, + source_model_class, + kwargs_numerics=kwargs_numerics, + ) + image_sim = sim_util.simulate_simple(imageModel, kwargs_lens, kwargs_source, point_source_add=False) + + data_class.update_data(image_sim) + + kwargs_data["image_data"] = image_sim + + kwargs_model = { + "lens_model_list": lens_model_list, + "source_light_model_list": source_model_list, + } + + lens_fixed = [{}] + lens_sigma = [ + { + "theta_E": 0.1, + "gamma": 0.1, + "e1": 0.1, + "e2": 0.1, + "center_x": 0.1, + "center_y": 0.1, + } + ] + lens_lower = [ + { + "theta_E": 0.0, + "gamma": 1.5, + "center_x": -2, + "center_y": -2, + "e1": -0.4, + "e2": -0.4, + } + ] + lens_upper = [ + { + "theta_E": 10.0, + "gamma": 2.5, + "center_x": 2, + "center_y": 2, + "e1": 0.4, + "e2": 0.4, + } + ] + + source_fixed = [{}] + source_sigma = [ + { + "amp": 0.1, + "R_sersic": 0.05, + "n_sersic": 0.5, + "center_x": 0.1, + "center_y": 0.1, + "e1": 0.1, + "e2": 0.1, + } + ] + source_lower = [ + { + "amp": 0, + "R_sersic": 0.01, + "n_sersic": 0.5, + "center_x": -2, + "center_y": -2, + "e1": -0.4, + "e2": -0.4, + } + ] + source_upper = [ + { + "amp": 100, + "R_sersic": 10, + "n_sersic": 5.5, + "center_x": 2, + "center_y": 2, + "e1": 0.4, + "e2": 0.4, + } + ] + + lens_param = [kwargs_lens, lens_sigma, lens_fixed, lens_lower, lens_upper] + source_param = [ + kwargs_source, + source_sigma, + source_fixed, + source_lower, + source_upper, + ] + + kwargs_params = {"lens_model": lens_param, "source_model": source_param} + + multi_band_list = [[kwargs_data, kwargs_psf_gaussian, kwargs_numerics]] + + kwargs_data_joint = { + "multi_band_list": multi_band_list, + "multi_band_type": "single-band", + } + + fittingSequence = FittingSequence( + kwargs_data_joint, + kwargs_model, + self.kwargs_constraints, + self.kwargs_likelihood, + kwargs_params, + ) + + fitting_list = [] + + kwargs_zeus = { + "n_burn": 2, + "n_run": 2, + "walkerRatio": 4, + "backend_filename": "test_mcmc_zeus.h5", + } + + fitting_list.append(["zeus", kwargs_zeus]) + + chain_list = fittingSequence.fit_sequence(fitting_list) + + def test_multinest(self): + # Nested sampler tests + # further decrease the parameter space for nested samplers to run faster + + fittingSequence = FittingSequence( + self.kwargs_data_joint, + self.kwargs_model, + self.kwargs_constraints, + self.kwargs_likelihood, + self.kwargs_params, + ) + fitting_list = [] + kwargs_update = { + #"ps_add_fixed": [[0, ["ra_source", "dec_source"], [0, 0]]], + "lens_light_add_fixed": [ + [0, ["n_sersic", "R_sersic", "center_x", "center_y"], [4, 0.1, 0, 0]] + ], + "source_add_fixed": [ + [ + 0, + ["R_sersic", "e1", "e2", "center_x", "center_y"], + [0.6, 0.1, 0.1, 0, 0], + ] + ], + "lens_add_fixed": [ + [ + 0, + ["gamma", "theta_E", "e1", "e2", "center_x", "center_y"], + [1.8, 1.0, 0.1, 0.1, 0, 0], + ], + [1, ["gamma1", "gamma2"], [0.01, 0.01]], + ], + "change_source_lower_limit": [[0, ["n_sersic"], [2.9]]], + "change_source_upper_limit": [[0, ["n_sersic"], [3.1]]], + } + fitting_list.append(["update_settings", kwargs_update]) + kwargs_multinest = { + "kwargs_run": { + "n_live_points": 10, + "evidence_tolerance": 0.5, + "sampling_efficiency": 0.8, # 1 for posterior-only, 0 for evidence-only + "importance_nested_sampling": False, + "multimodal": True, + "const_efficiency_mode": False, # reduce sampling_efficiency to 5% when True + }, + "remove_output_dir": True, + } + + fitting_list.append(["MultiNest", kwargs_multinest]) + + chain_list2 = fittingSequence.fit_sequence(fitting_list) + kwargs_fixed = fittingSequence._updateManager.fixed_kwargs + npt.assert_almost_equal(kwargs_fixed[0][1]["gamma1"], 0.01, decimal=2) + assert fittingSequence._updateManager._lower_kwargs[1][0]["n_sersic"] == 2.9 + assert fittingSequence._updateManager._upper_kwargs[1][0]["n_sersic"] == 3.1 + + kwargs_test = {"kwargs_lens": 1} + fittingSequence.update_state(kwargs_test) + kwargs_out = fittingSequence.best_fit(bijective=True) + assert kwargs_out["kwargs_lens"] == 1 + + def test_dynesty(self): + np.random.seed(42) + kwargs_params = copy.deepcopy(self.kwargs_params) + kwargs_params["lens_model"][0][0]["theta_E"] += 0.01 + kwargs_params["lens_model"][0][0]["gamma"] += 0.01 + fittingSequence = FittingSequence( + self.kwargs_data_joint, + self.kwargs_model, + self.kwargs_constraints, + self.kwargs_likelihood, + kwargs_params, + ) + + fitting_list = [] + kwargs_dynesty = { + "kwargs_run": { + "dlogz_init": 0.01, + "nlive_init": 20, + "nlive_batch": 20, + "maxbatch": 1, + }, + } + + fitting_list.append(["dynesty", kwargs_dynesty]) + + chain_list = fittingSequence.fit_sequence(fitting_list) + + def test_nautilus(self): + np.random.seed(42) + kwargs_params = copy.deepcopy(self.kwargs_params) + fittingSequence = FittingSequence( + self.kwargs_data_joint, + self.kwargs_model, + self.kwargs_constraints, + self.kwargs_likelihood, + kwargs_params, + ) + + fitting_list = [] + kwargs_nautilus = { + "prior_type": "uniform", + "verbose": True, + "f_live": 1.0, + "n_eff": 0.0, + "n_live": 2, + "seed": 42, + } + + fitting_list.append(["Nautilus", kwargs_nautilus]) + chain_list = fittingSequence.fit_sequence(fitting_list) + + def test_dypolychord(self): + fittingSequence = FittingSequence( + self.kwargs_data_joint, + self.kwargs_model, + self.kwargs_constraints, + self.kwargs_likelihood, + self.kwargs_params, + ) + fitting_list = [] + kwargs_dypolychord = { + "kwargs_run": { + "ninit": 8, + "nlive_const": 10, + #'seed_increment': 1, + "resume_dyn_run": False, + #'init_step': 10, + }, + "polychord_settings": { + "seed": 1, + #'num_repeats': 20 + }, + "dypolychord_dynamic_goal": 0.8, # 1 for posterior-only, 0 for evidence-only + "remove_output_dir": True, + } + + fitting_list.append(["dyPolyChord", kwargs_dypolychord]) + + chain_list = fittingSequence.fit_sequence(fitting_list) + + def test_minimizer(self): + n_p = 2 + n_i = 2 + + fitting_list = [] + + kwargs_simplex = {"n_iterations": n_i, "method": "Nelder-Mead"} + fitting_list.append(["SIMPLEX", kwargs_simplex]) + kwargs_simplex = {"n_iterations": n_i, "method": "Powell"} + fitting_list.append(["SIMPLEX", kwargs_simplex]) + kwargs_pso = {"sigma_scale": 1, "n_particles": n_p, "n_iterations": n_i} + fitting_list.append(["PSO", kwargs_pso]) + kwargs_mcmc = {"sigma_scale": 1, "n_burn": 1, "n_run": 1, "n_walkers": 10} + fitting_list.append(["emcee", kwargs_mcmc]) + kwargs_mcmc["re_use_samples"] = True + kwargs_mcmc["init_samples"] = np.array( + [[np.random.normal(1, 0.001)] for i in range(100)] + ) + fitting_list.append(["emcee", kwargs_mcmc]) + + def custom_likelihood(kwargs_lens, **kwargs): + theta_E = kwargs_lens[0]["theta_E"] + return -((theta_E - 1.0) ** 2) / 0.1**2 / 2 + + kwargs_likelihood = {"custom_logL_addition": custom_likelihood} + + kwargs_data_joint = {"multi_band_list": [], "multi_band_type": "single-band",} + kwargs_model = {"lens_model_list": ["SIS"]} + lens_param = ( + [{"theta_E": 1, "center_x": 0, "center_y": 0}], + [{"theta_E": 0.1, "center_x": 0.1, "center_y": 0.1}], + [{"center_x": 0, "center_y": 0}], + [{"theta_E": 0, "center_x": -10, "center_y": -10}], + [{"theta_E": 10, "center_x": 10, "center_y": 10}], + ) + + kwargs_params = {"lens_model": lens_param} + fittingSequence = FittingSequence( + kwargs_data_joint, + kwargs_model, + self.kwargs_constraints, + kwargs_likelihood, + kwargs_params, + ) + args = fittingSequence.param_class.kwargs2args( + kwargs_lens=[{"theta_E": 1, "center_x": 0, "center_y": 0}] + ) + kwargs_result = fittingSequence.param_class.args2kwargs(args) + print(kwargs_result) + print(args, "test args") + chain_list = fittingSequence.fit_sequence(fitting_list) + kwargs_result = fittingSequence.best_fit(bijective=False) + npt.assert_almost_equal( + kwargs_result["kwargs_lens"][0]["theta_E"], 1, decimal=2 + ) + + def test_jaxopt_minimizer(self): + # data specifics + background_rms = .005 # background noise per pixel + exp_time = 100. # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) + numPix = 60 # cutout pixel size per axis + pixel_scale = 0.11 # pixel size in arcsec (area per pixel = pixel_scale**2) + fwhm = 0.5 # full width at half maximum of PSF + + # lensing quantities + lens_model_list = ['EPL', 'SHEAR'] + kwargs_epl = {'theta_E': 0.66, 'gamma': 1.7, 'e1': .07, 'e2': -0.03, 'center_x': 0.05, 'center_y': 0.1} # parameters of the deflector lens model + kwargs_shear = {'gamma1': 0.0, 'gamma2': -0.05} # shear values to the source plane + + kwargs_lens = [kwargs_epl, kwargs_shear] + + lens_model_class = LensModel(lens_model_list) + + # Sersic parameters in the initial simulation for the source + kwargs_sersic = {'amp': 16., 'R_sersic': 0.1, 'n_sersic': 1., 'e1': -0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.} + source_model_list = ['SERSIC_ELLIPSE'] + kwargs_source = [kwargs_sersic] + + source_model_class = LightModel(source_model_list) + + + kwargs_sersic_lens = {'amp': 16., 'R_sersic': 0.6, 'n_sersic': 2., 'e1': -0.1, 'e2': 0.1, 'center_x': 0.05, 'center_y': 0.} + + lens_light_model_list = ['SERSIC_ELLIPSE'] + kwargs_lens_light = [kwargs_sersic_lens] + + kwargs_truth = {'kwargs_lens': kwargs_lens, + 'kwargs_source': kwargs_source, + 'kwargs_lens_light': kwargs_lens_light} + + + lens_light_model_class = LightModel(lens_light_model_list) + # generate the coordinate grid and image properties (we only read out the relevant lines we need) + _, _, ra_at_xy_0, dec_at_xy_0, _, _, Mpix2coord, _ = util.make_grid_with_coordtransform(numPix=numPix, deltapix=pixel_scale, center_ra=0, center_dec=0, subgrid_res=1, inverse=False) + + + kwargs_data = {'background_rms': background_rms, # rms of background noise + 'exposure_time': exp_time, # exposure time (or a map per pixel) + 'ra_at_xy_0': ra_at_xy_0, # RA at (0,0) pixel + 'dec_at_xy_0': dec_at_xy_0, # DEC at (0,0) pixel + 'transform_pix2angle': Mpix2coord, # matrix to translate shift in pixel in shift in relative RA/DEC (2x2 matrix). Make sure it's units are arcseconds or the angular un + 'image_data': np.zeros((numPix, numPix)) # 2d data vector, here initialized with zeros as place holders that get's overwritten once a simulated image with noise is cre + } + + data_class = ImageData(**kwargs_data) + # generate the psf variables + #kwargs_psf = {'psf_type': psf_type, 'fwhm': fwhm, 'pixel_size': pixel_scale, 'truncation': 3} + + # if you are using a PSF estimate from e.g. a star in the FoV of your exposure, you can set + kwargs_psf_gaussian = {'psf_type': 'GAUSSIAN', 'pixel_size': pixel_scale, 'fwhm': fwhm} + + + psf_class = PSF(**kwargs_psf_gaussian) + kwargs_numerics = {'supersampling_factor': 1, 'supersampling_convolution': False, 'convolution_type': 'fft_static'} + + imageModel = ImageModel(data_class, psf_class, lens_model_class=lens_model_class, + source_model_class=source_model_class, lens_light_model_class=lens_light_model_class, + kwargs_numerics=kwargs_numerics) + + # generate image + image_model = imageModel.image(kwargs_lens, kwargs_source, kwargs_lens_light=kwargs_lens_light, kwargs_ps=None, point_source_add=False) + kwargs_data['image_data'] = image_model + + data_class.update_data(image_model) + + fixed_lens = [] + kwargs_lens_init = [] + kwargs_lens_sigma = [] + kwargs_lower_lens = [] + kwargs_upper_lens = [] + + fixed_lens.append({}) + kwargs_lens_init.append({'center_x': 0.1, 'center_y': 0.2, + 'e1': 0.13, 'e2': -0.06, 'theta_E': 0.6, }) + kwargs_lens_sigma.append({'theta_E': 0.1, 'e1': 0.1, 'e2': 0.1, + 'center_x': 0.1, 'center_y': 0.1}) + kwargs_lower_lens.append({'theta_E': 0.01, 'e1': -0.5, 'e2': -0.5, 'center_x': -10., 'center_y': -10.}) + kwargs_upper_lens.append({'theta_E': 10., 'e1': 0.5, 'e2': 0.5, 'center_x': 10., 'center_y': 10.}) + + + kwargs_lens_init[0]['gamma'] = 2.1 + kwargs_lens_sigma[0]['gamma'] = 0.1 + kwargs_lower_lens[0]['gamma'] = 1.5 + kwargs_upper_lens[0]['gamma'] = 2.5 + + fixed_lens.append({'ra_0': 0, 'dec_0': 0}) + kwargs_lens_init.append({'gamma1': 0.1, 'gamma2': 0.04}) + kwargs_lens_sigma.append({'gamma1': 0.01, 'gamma2': 0.01}) + kwargs_lower_lens.append({'gamma1': -0.2, 'gamma2': -0.2}) + kwargs_upper_lens.append({'gamma1': 0.2, 'gamma2': 0.2}) + + lens_params = [kwargs_lens_init, kwargs_lens_sigma, fixed_lens, kwargs_lower_lens, kwargs_upper_lens] + + + fixed_source = [] + kwargs_source_init = [] + kwargs_source_sigma = [] + kwargs_lower_source = [] + kwargs_upper_source = [] + + + fixed_source.append({}) + kwargs_source_init.append({'R_sersic': 0.2, 'n_sersic': 2.5, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.3, 'center_y': 0., 'amp': 26.}) + kwargs_source_sigma.append({'n_sersic': 1, 'R_sersic': 0.01, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.2, 'center_y': 0.2, 'amp': 10}) + kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.001, 'n_sersic': .5, 'center_x': -10., 'center_y': -10., 'amp': 0.}) + kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 10, 'n_sersic': 5., 'center_x': 10, 'center_y': 10, 'amp': 100}) + + source_params = [kwargs_source_init, kwargs_source_sigma, fixed_source, kwargs_lower_source, kwargs_upper_source] + + + fixed_lens_light = [] + kwargs_lens_light_init = [] + kwargs_lens_light_sigma = [] + kwargs_lower_lens_light = [] + kwargs_upper_lens_light = [] + + + fixed_lens_light.append({}) + kwargs_lens_light_init.append({'R_sersic': 0.5, 'n_sersic': 2., 'e1': 0.1, 'e2': 0.3, 'center_x': 0.1, 'center_y': 0., 'amp': 7.}) + kwargs_lens_light_sigma.append({'n_sersic': 0.01, 'R_sersic': 0.03, 'e1': 0.5, 'e2': 0.5, 'center_x': 0.01, 'center_y': 0.01, 'amp': 10}) + kwargs_lower_lens_light.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.001, 'n_sersic': .5, 'center_x': -10., 'center_y': -10., 'amp': 0.}) + kwargs_upper_lens_light.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 10., 'n_sersic': 5., 'center_x': 10., 'center_y': 10., 'amp': 100.}) + + lens_light_params = [kwargs_lens_light_init, kwargs_lens_light_sigma, fixed_lens_light, kwargs_lower_lens_light, kwargs_upper_lens_light] + + kwargs_params = {'lens_model': lens_params, + 'source_model': source_params, + 'lens_light_model': lens_light_params} + + kwargs_model = {'lens_model_list': lens_model_list, 'source_light_model_list': source_model_list, 'lens_light_model_list': lens_light_model_list} + + multi_band_list = [[kwargs_data, kwargs_psf_gaussian, kwargs_numerics]] + + kwargs_data_joint = {'multi_band_list': multi_band_list, + 'multi_band_type': 'single-band' + } + + fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, self.kwargs_constraints, self.kwargs_likelihood, kwargs_params) + + # options are BFGS(), Nelder-Mead, Powell, CG, Newton-CG, L-BFGS-B, TNC(), COBYLA, SLSQP!, trust-constr!, dogleg!, trust-ncg!, trust-exact!, trust-krylov! + fitting_kwargs_list_jaxopt = [['Jaxopt', {'method': "BFGS", 'maxiter': 200}]] + chain_list = fitting_seq.fit_sequence(fitting_kwargs_list_jaxopt) + fitting_type, args_history, logL_history, kwargs_result = chain_list[0] + + assert fitting_type == "Jaxopt" + assert len(args_history) == len(logL_history) + npt.assert_almost_equal(kwargs_result['kwargs_lens'][0]['theta_E'], 0.66, decimal=3) + + + +if __name__ == "__main__": + pytest.main() From 7d3e9e98a46207f9a2a77c775eea9a0648818d82 Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sat, 1 Feb 2025 21:08:49 -0500 Subject: [PATCH 02/25] rename exp_map to exposure_map --- jaxtronomy/Data/image_noise.py | 6 +++--- jaxtronomy/Data/imaging_data.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/jaxtronomy/Data/image_noise.py b/jaxtronomy/Data/image_noise.py index 674a67d..96cced6 100644 --- a/jaxtronomy/Data/image_noise.py +++ b/jaxtronomy/Data/image_noise.py @@ -49,7 +49,7 @@ def __init__( ) else: # make sure no negative exposure values are present no dividing by zero - self.exp_map = jnp.where( + self.exposure_map = jnp.where( exposure_time <= 10 ** (-10), 10 ** (-10), exposure_time ) @@ -88,7 +88,7 @@ def __init__( self.C_D = covariance_matrix( self.data, self.background_rms, - self.exp_map, + self.exposure_map, ) if gradient_boost_factor is not None: @@ -107,7 +107,7 @@ def C_D_model(self, model): if self._noise_map is not None: return self._noise_map**2 else: - return covariance_matrix(model, self.background_rms, self.exp_map) + return covariance_matrix(model, self.background_rms, self.exposure_map) @export diff --git a/jaxtronomy/Data/imaging_data.py b/jaxtronomy/Data/imaging_data.py index 532bee0..6345fd8 100644 --- a/jaxtronomy/Data/imaging_data.py +++ b/jaxtronomy/Data/imaging_data.py @@ -153,7 +153,7 @@ def update_data(self, image_data): self.C_D = covariance_matrix( self.data, self.background_rms, - self.exp_map, + self.exposure_map, ) def _log_likelihood(self, model, mask, additional_error_map=0): From 0c73ff45ab1f41ec86742d74308675955a5bef8d Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sat, 1 Feb 2025 21:09:05 -0500 Subject: [PATCH 03/25] init file --- test/test_Workflow/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 test/test_Workflow/__init__.py diff --git a/test/test_Workflow/__init__.py b/test/test_Workflow/__init__.py new file mode 100644 index 0000000..e69de29 From 0b35daa43c8ec2db1ea68f5dded6b49c5934e4a3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 2 Feb 2025 02:09:29 +0000 Subject: [PATCH 04/25] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- jaxtronomy/Workflow/fitting_sequence.py | 37 +- test/test_Workflow/test_fitting_sequence.py | 415 +++++++++++++++----- 2 files changed, 332 insertions(+), 120 deletions(-) diff --git a/jaxtronomy/Workflow/fitting_sequence.py b/jaxtronomy/Workflow/fitting_sequence.py index b0ca737..acf8011 100644 --- a/jaxtronomy/Workflow/fitting_sequence.py +++ b/jaxtronomy/Workflow/fitting_sequence.py @@ -128,7 +128,7 @@ def fit_sequence(self, fitting_list): elif fitting_type == "psf_iteration": raise ValueError("psf_iteration not supported in jaxtronomy") - #self.psf_iteration(**kwargs) + # self.psf_iteration(**kwargs) elif fitting_type == "align_images": self.align_images(**kwargs) @@ -209,7 +209,9 @@ def fit_sequence(self, fitting_list): elif fitting_type == "Jaxopt": args_history, logL_history, kwargs_result = self.jaxopt(**kwargs) self._updateManager.update_param_state(**kwargs_result) - chain_list.append([fitting_type, args_history, logL_history, kwargs_result]) + chain_list.append( + [fitting_type, args_history, logL_history, kwargs_result] + ) else: raise ValueError( "fitting_sequence {} is not supported. Please use: 'PSO', 'SIMPLEX', " @@ -277,7 +279,6 @@ def param_class(self): """ return self._updateManager.param_class - @property def likelihoodModule(self): """ @@ -290,7 +291,6 @@ def likelihoodModule(self): self.kwargs_data_joint, kwargs_model, self.param_class, **kwargs_likelihood ) return likelihoodModule - def simplex(self, n_iterations, method="Nelder-Mead"): """Downhill simplex optimization using the Nelder-Mead algorithm. @@ -324,7 +324,7 @@ def mcmc( progress=True, backend_filename=None, start_from_backend=False, - **kwargs_zeus + **kwargs_zeus, ): """MCMC routine. @@ -389,7 +389,7 @@ def mcmc( progress=progress, initpos=initpos, backend_filename=backend_filename, - **kwargs_zeus + **kwargs_zeus, ) output = [sampler_type, samples, param_list, dist] else: @@ -417,26 +417,28 @@ def jaxopt( method="BFGS", maxiter=100, ): - """Uses the Jaxopt Scipy Minimizer - - :param method: string, options are BFGS, Nelder-Mead, Powell, CG, BFGS, Newton-CG, - L-BFGS-B, TNC, COBYLA, SLSQP, trust-constr, dogleg, trust-ncg, trust-exact, trust-krylov - :param maxiter: int, number of iterations to perform during minimization of the loss function + """Uses the Jaxopt Scipy Minimizer. + + :param method: string, options are BFGS, Nelder-Mead, Powell, CG, BFGS, Newton- + CG, L-BFGS-B, TNC, COBYLA, SLSQP, trust-constr, dogleg, trust-ncg, trust- + exact, trust-krylov + :param maxiter: int, number of iterations to perform during minimization of the + loss function """ print(f"Running {method} minimization with {maxiter} iterations:") param_class = self.param_class - likelihood_module=self.likelihoodModule + likelihood_module = self.likelihoodModule # Coonverts kwargs to args for the mean, sigma, lower, and upper parameter values kwargs_temp = self._updateManager.parameter_state args_mean = param_class.kwargs2args(**kwargs_temp) - + kwargs_sigma = self._updateManager.sigma_kwargs args_sigma = param_class.kwargs2args(**kwargs_sigma) - + kwargs_lower = self._updateManager._lower_kwargs args_lower = param_class.kwargs2args(*kwargs_lower) - + kwargs_upper = self._updateManager._upper_kwargs args_upper = param_class.kwargs2args(*kwargs_upper) @@ -448,7 +450,7 @@ def jaxopt( args_sigma=args_sigma, args_lower=args_lower, args_upper=args_upper, - maxiter=maxiter + maxiter=maxiter, ) # Runs the minimizer and prints results @@ -456,9 +458,8 @@ def jaxopt( kwargs_result = param_class.args2kwargs(args_result) print("best fit log_likelihood:", final_logL) print("kwargs_result:", kwargs_result) - - return minimizer.parameter_history, minimizer.logL_history, kwargs_result + return minimizer.parameter_history, minimizer.logL_history, kwargs_result def pso( self, n_particles, n_iterations, sigma_scale=1, print_key="PSO", threadCount=1 diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index f18a5f3..55f9664 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -88,15 +88,15 @@ def setup_method(self): self.kwargs_ps = [ {"ra_source": 0.0, "dec_source": 0.0, "source_amp": 1.0} ] # quasar point source position in the source plane and intrinsic brightness - point_source_list = []#["SOURCE_POSITION"] + point_source_list = [] # ["SOURCE_POSITION"] point_source_class = PointSource( - point_source_type_list=point_source_list, #fixed_magnification_list=[True] + point_source_type_list=point_source_list, # fixed_magnification_list=[True] ) kwargs_numerics = { "supersampling_factor": 3, "supersampling_convolution": True, "compute_mode": "regular", - #"point_source_supersampling_factor": 1, + # "point_source_supersampling_factor": 1, } imageModel = ImageModel( data_class, @@ -112,8 +112,8 @@ def setup_method(self): self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, - #self.kwargs_psf, - point_source_add=False + # self.kwargs_psf, + point_source_add=False, ) data_class.update_data(image_sim) @@ -131,11 +131,11 @@ def setup_method(self): } self.kwargs_numerics = kwargs_numerics - #num_source_model = len(source_model_list) + # num_source_model = len(source_model_list) self.kwargs_constraints = { - #"num_point_source_list": [4], - #"image_plane_source_list": [False] * num_source_model, + # "num_point_source_list": [4], + # "image_plane_source_list": [False] * num_source_model, "linear_solver": False } @@ -212,10 +212,22 @@ def setup_method(self): ] lens_light_sigma = [ - {"amp": 0.1, "R_sersic": 0.05, "n_sersic": 0.5, "center_x": 0.1, "center_y": 0.1} + { + "amp": 0.1, + "R_sersic": 0.05, + "n_sersic": 0.5, + "center_x": 0.1, + "center_y": 0.1, + } ] lens_light_lower = [ - {"amp": 0, "R_sersic": 0.01, "n_sersic": 0.5, "center_x": -2, "center_y": -2} + { + "amp": 0, + "R_sersic": 0.01, + "n_sersic": 0.5, + "center_x": -2, + "center_y": -2, + } ] lens_light_upper = [ {"amp": 100, "R_sersic": 10, "n_sersic": 5.5, "center_x": 2, "center_y": 2} @@ -306,7 +318,7 @@ def test_fitting_sequence(self): "stacking_method": "mean", "new_procedure": False, } - #fitting_list.append(["psf_iteration", kwargs_psf_iter]) + # fitting_list.append(["psf_iteration", kwargs_psf_iter]) fitting_list.append(["restart", None]) fitting_list.append(["fix_not_computed", {"free_bands": [True]}]) n_sersic_overwrite = 4 @@ -348,7 +360,7 @@ def test_fitting_sequence(self): ["set_param_value", {"lens": [[1, ["gamma1"], [0.013]]]}], ["set_param_value", {"lens_light": [[0, ["center_x"], [0.009]]]}], ["set_param_value", {"source": [[0, ["n_sersic"], [2.993]]]}], - #["set_param_value", {"ps": [[0, ["ra_source"], [0.007]]]}], + # ["set_param_value", {"ps": [[0, ["ra_source"], [0.007]]]}], ] fittingSequence.fit_sequence(fitting_list) @@ -357,7 +369,7 @@ def test_fitting_sequence(self): assert kwargs_set["kwargs_lens"][1]["gamma1"] == 0.013 assert kwargs_set["kwargs_lens_light"][0]["center_x"] == 0.009 assert kwargs_set["kwargs_source"][0]["n_sersic"] == 2.993 - #assert kwargs_set["kwargs_ps"][0]["ra_source"] == 0.007 + # assert kwargs_set["kwargs_ps"][0]["ra_source"] == 0.007 from unittest import TestCase @@ -378,12 +390,12 @@ def test_fitting_sequence(self): with t.assertRaises(ValueError): fittingSequence.fit_sequence(fitting_list4) - #psf_iteration_list = fittingSequence.psf_iteration_memory - #assert len(psf_iteration_list) == 1 - #assert "sequence" in psf_iteration_list[0] - #assert "band" in psf_iteration_list[0] - #assert "psf_before" in psf_iteration_list[0] - #assert "psf_after" in psf_iteration_list[0] + # psf_iteration_list = fittingSequence.psf_iteration_memory + # assert len(psf_iteration_list) == 1 + # assert "sequence" in psf_iteration_list[0] + # assert "band" in psf_iteration_list[0] + # assert "psf_before" in psf_iteration_list[0] + # assert "psf_after" in psf_iteration_list[0] def test_cobaya(self): np.random.seed(42) @@ -439,7 +451,9 @@ def test_cobaya(self): source_model_class, kwargs_numerics=kwargs_numerics, ) - image_sim = sim_util.simulate_simple(imageModel, kwargs_lens, kwargs_source, point_source_add=False) + image_sim = sim_util.simulate_simple( + imageModel, kwargs_lens, kwargs_source, point_source_add=False + ) data_class.update_data(image_sim) @@ -457,10 +471,22 @@ def test_cobaya(self): source_fixed = [{}] source_sigma = [ - {"amp": 0.1, "R_sersic": 0.01, "n_sersic": 0.01, "center_x": 0.01, "center_y": 0.01} + { + "amp": 0.1, + "R_sersic": 0.01, + "n_sersic": 0.01, + "center_x": 0.01, + "center_y": 0.01, + } ] source_lower = [ - {"amp": 0, "R_sersic": 0.01, "n_sersic": 0.5, "center_x": -1, "center_y": -1} + { + "amp": 0, + "R_sersic": 0.01, + "n_sersic": 0.5, + "center_x": -1, + "center_y": -1, + } ] source_upper = [ {"amp": 100, "R_sersic": 1.0, "n_sersic": 6.0, "center_x": 1, "center_y": 1} @@ -565,7 +591,9 @@ def test_zeus(self): source_model_class, kwargs_numerics=kwargs_numerics, ) - image_sim = sim_util.simulate_simple(imageModel, kwargs_lens, kwargs_source, point_source_add=False) + image_sim = sim_util.simulate_simple( + imageModel, kwargs_lens, kwargs_source, point_source_add=False + ) data_class.update_data(image_sim) @@ -695,7 +723,7 @@ def test_multinest(self): ) fitting_list = [] kwargs_update = { - #"ps_add_fixed": [[0, ["ra_source", "dec_source"], [0, 0]]], + # "ps_add_fixed": [[0, ["ra_source", "dec_source"], [0, 0]]], "lens_light_add_fixed": [ [0, ["n_sersic", "R_sersic", "center_x", "center_y"], [4, 0.1, 0, 0]] ], @@ -849,7 +877,10 @@ def custom_likelihood(kwargs_lens, **kwargs): kwargs_likelihood = {"custom_logL_addition": custom_likelihood} - kwargs_data_joint = {"multi_band_list": [], "multi_band_type": "single-band",} + kwargs_data_joint = { + "multi_band_list": [], + "multi_band_type": "single-band", + } kwargs_model = {"lens_model_list": ["SIS"]} lens_param = ( [{"theta_E": 1, "center_x": 0, "center_y": 0}], @@ -881,70 +912,125 @@ def custom_likelihood(kwargs_lens, **kwargs): def test_jaxopt_minimizer(self): # data specifics - background_rms = .005 # background noise per pixel - exp_time = 100. # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) + background_rms = 0.005 # background noise per pixel + exp_time = 100.0 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) numPix = 60 # cutout pixel size per axis pixel_scale = 0.11 # pixel size in arcsec (area per pixel = pixel_scale**2) fwhm = 0.5 # full width at half maximum of PSF # lensing quantities - lens_model_list = ['EPL', 'SHEAR'] - kwargs_epl = {'theta_E': 0.66, 'gamma': 1.7, 'e1': .07, 'e2': -0.03, 'center_x': 0.05, 'center_y': 0.1} # parameters of the deflector lens model - kwargs_shear = {'gamma1': 0.0, 'gamma2': -0.05} # shear values to the source plane + lens_model_list = ["EPL", "SHEAR"] + kwargs_epl = { + "theta_E": 0.66, + "gamma": 1.7, + "e1": 0.07, + "e2": -0.03, + "center_x": 0.05, + "center_y": 0.1, + } # parameters of the deflector lens model + kwargs_shear = { + "gamma1": 0.0, + "gamma2": -0.05, + } # shear values to the source plane kwargs_lens = [kwargs_epl, kwargs_shear] lens_model_class = LensModel(lens_model_list) # Sersic parameters in the initial simulation for the source - kwargs_sersic = {'amp': 16., 'R_sersic': 0.1, 'n_sersic': 1., 'e1': -0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.} - source_model_list = ['SERSIC_ELLIPSE'] + kwargs_sersic = { + "amp": 16.0, + "R_sersic": 0.1, + "n_sersic": 1.0, + "e1": -0.1, + "e2": 0.1, + "center_x": 0.1, + "center_y": 0.0, + } + source_model_list = ["SERSIC_ELLIPSE"] kwargs_source = [kwargs_sersic] source_model_class = LightModel(source_model_list) + kwargs_sersic_lens = { + "amp": 16.0, + "R_sersic": 0.6, + "n_sersic": 2.0, + "e1": -0.1, + "e2": 0.1, + "center_x": 0.05, + "center_y": 0.0, + } - kwargs_sersic_lens = {'amp': 16., 'R_sersic': 0.6, 'n_sersic': 2., 'e1': -0.1, 'e2': 0.1, 'center_x': 0.05, 'center_y': 0.} - - lens_light_model_list = ['SERSIC_ELLIPSE'] + lens_light_model_list = ["SERSIC_ELLIPSE"] kwargs_lens_light = [kwargs_sersic_lens] - kwargs_truth = {'kwargs_lens': kwargs_lens, - 'kwargs_source': kwargs_source, - 'kwargs_lens_light': kwargs_lens_light} - + kwargs_truth = { + "kwargs_lens": kwargs_lens, + "kwargs_source": kwargs_source, + "kwargs_lens_light": kwargs_lens_light, + } lens_light_model_class = LightModel(lens_light_model_list) # generate the coordinate grid and image properties (we only read out the relevant lines we need) - _, _, ra_at_xy_0, dec_at_xy_0, _, _, Mpix2coord, _ = util.make_grid_with_coordtransform(numPix=numPix, deltapix=pixel_scale, center_ra=0, center_dec=0, subgrid_res=1, inverse=False) - + _, _, ra_at_xy_0, dec_at_xy_0, _, _, Mpix2coord, _ = ( + util.make_grid_with_coordtransform( + numPix=numPix, + deltapix=pixel_scale, + center_ra=0, + center_dec=0, + subgrid_res=1, + inverse=False, + ) + ) - kwargs_data = {'background_rms': background_rms, # rms of background noise - 'exposure_time': exp_time, # exposure time (or a map per pixel) - 'ra_at_xy_0': ra_at_xy_0, # RA at (0,0) pixel - 'dec_at_xy_0': dec_at_xy_0, # DEC at (0,0) pixel - 'transform_pix2angle': Mpix2coord, # matrix to translate shift in pixel in shift in relative RA/DEC (2x2 matrix). Make sure it's units are arcseconds or the angular un - 'image_data': np.zeros((numPix, numPix)) # 2d data vector, here initialized with zeros as place holders that get's overwritten once a simulated image with noise is cre - } + kwargs_data = { + "background_rms": background_rms, # rms of background noise + "exposure_time": exp_time, # exposure time (or a map per pixel) + "ra_at_xy_0": ra_at_xy_0, # RA at (0,0) pixel + "dec_at_xy_0": dec_at_xy_0, # DEC at (0,0) pixel + "transform_pix2angle": Mpix2coord, # matrix to translate shift in pixel in shift in relative RA/DEC (2x2 matrix). Make sure it's units are arcseconds or the angular un + "image_data": np.zeros( + (numPix, numPix) + ), # 2d data vector, here initialized with zeros as place holders that get's overwritten once a simulated image with noise is cre + } data_class = ImageData(**kwargs_data) # generate the psf variables - #kwargs_psf = {'psf_type': psf_type, 'fwhm': fwhm, 'pixel_size': pixel_scale, 'truncation': 3} + # kwargs_psf = {'psf_type': psf_type, 'fwhm': fwhm, 'pixel_size': pixel_scale, 'truncation': 3} # if you are using a PSF estimate from e.g. a star in the FoV of your exposure, you can set - kwargs_psf_gaussian = {'psf_type': 'GAUSSIAN', 'pixel_size': pixel_scale, 'fwhm': fwhm} - + kwargs_psf_gaussian = { + "psf_type": "GAUSSIAN", + "pixel_size": pixel_scale, + "fwhm": fwhm, + } psf_class = PSF(**kwargs_psf_gaussian) - kwargs_numerics = {'supersampling_factor': 1, 'supersampling_convolution': False, 'convolution_type': 'fft_static'} + kwargs_numerics = { + "supersampling_factor": 1, + "supersampling_convolution": False, + "convolution_type": "fft_static", + } - imageModel = ImageModel(data_class, psf_class, lens_model_class=lens_model_class, - source_model_class=source_model_class, lens_light_model_class=lens_light_model_class, - kwargs_numerics=kwargs_numerics) + imageModel = ImageModel( + data_class, + psf_class, + lens_model_class=lens_model_class, + source_model_class=source_model_class, + lens_light_model_class=lens_light_model_class, + kwargs_numerics=kwargs_numerics, + ) # generate image - image_model = imageModel.image(kwargs_lens, kwargs_source, kwargs_lens_light=kwargs_lens_light, kwargs_ps=None, point_source_add=False) - kwargs_data['image_data'] = image_model + image_model = imageModel.image( + kwargs_lens, + kwargs_source, + kwargs_lens_light=kwargs_lens_light, + kwargs_ps=None, + point_source_add=False, + ) + kwargs_data["image_data"] = image_model data_class.update_data(image_model) @@ -955,27 +1041,49 @@ def test_jaxopt_minimizer(self): kwargs_upper_lens = [] fixed_lens.append({}) - kwargs_lens_init.append({'center_x': 0.1, 'center_y': 0.2, - 'e1': 0.13, 'e2': -0.06, 'theta_E': 0.6, }) - kwargs_lens_sigma.append({'theta_E': 0.1, 'e1': 0.1, 'e2': 0.1, - 'center_x': 0.1, 'center_y': 0.1}) - kwargs_lower_lens.append({'theta_E': 0.01, 'e1': -0.5, 'e2': -0.5, 'center_x': -10., 'center_y': -10.}) - kwargs_upper_lens.append({'theta_E': 10., 'e1': 0.5, 'e2': 0.5, 'center_x': 10., 'center_y': 10.}) - - - kwargs_lens_init[0]['gamma'] = 2.1 - kwargs_lens_sigma[0]['gamma'] = 0.1 - kwargs_lower_lens[0]['gamma'] = 1.5 - kwargs_upper_lens[0]['gamma'] = 2.5 - - fixed_lens.append({'ra_0': 0, 'dec_0': 0}) - kwargs_lens_init.append({'gamma1': 0.1, 'gamma2': 0.04}) - kwargs_lens_sigma.append({'gamma1': 0.01, 'gamma2': 0.01}) - kwargs_lower_lens.append({'gamma1': -0.2, 'gamma2': -0.2}) - kwargs_upper_lens.append({'gamma1': 0.2, 'gamma2': 0.2}) - - lens_params = [kwargs_lens_init, kwargs_lens_sigma, fixed_lens, kwargs_lower_lens, kwargs_upper_lens] + kwargs_lens_init.append( + { + "center_x": 0.1, + "center_y": 0.2, + "e1": 0.13, + "e2": -0.06, + "theta_E": 0.6, + } + ) + kwargs_lens_sigma.append( + {"theta_E": 0.1, "e1": 0.1, "e2": 0.1, "center_x": 0.1, "center_y": 0.1} + ) + kwargs_lower_lens.append( + { + "theta_E": 0.01, + "e1": -0.5, + "e2": -0.5, + "center_x": -10.0, + "center_y": -10.0, + } + ) + kwargs_upper_lens.append( + {"theta_E": 10.0, "e1": 0.5, "e2": 0.5, "center_x": 10.0, "center_y": 10.0} + ) + kwargs_lens_init[0]["gamma"] = 2.1 + kwargs_lens_sigma[0]["gamma"] = 0.1 + kwargs_lower_lens[0]["gamma"] = 1.5 + kwargs_upper_lens[0]["gamma"] = 2.5 + + fixed_lens.append({"ra_0": 0, "dec_0": 0}) + kwargs_lens_init.append({"gamma1": 0.1, "gamma2": 0.04}) + kwargs_lens_sigma.append({"gamma1": 0.01, "gamma2": 0.01}) + kwargs_lower_lens.append({"gamma1": -0.2, "gamma2": -0.2}) + kwargs_upper_lens.append({"gamma1": 0.2, "gamma2": 0.2}) + + lens_params = [ + kwargs_lens_init, + kwargs_lens_sigma, + fixed_lens, + kwargs_lower_lens, + kwargs_upper_lens, + ] fixed_source = [] kwargs_source_init = [] @@ -983,15 +1091,59 @@ def test_jaxopt_minimizer(self): kwargs_lower_source = [] kwargs_upper_source = [] - fixed_source.append({}) - kwargs_source_init.append({'R_sersic': 0.2, 'n_sersic': 2.5, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.3, 'center_y': 0., 'amp': 26.}) - kwargs_source_sigma.append({'n_sersic': 1, 'R_sersic': 0.01, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.2, 'center_y': 0.2, 'amp': 10}) - kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.001, 'n_sersic': .5, 'center_x': -10., 'center_y': -10., 'amp': 0.}) - kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 10, 'n_sersic': 5., 'center_x': 10, 'center_y': 10, 'amp': 100}) - - source_params = [kwargs_source_init, kwargs_source_sigma, fixed_source, kwargs_lower_source, kwargs_upper_source] + kwargs_source_init.append( + { + "R_sersic": 0.2, + "n_sersic": 2.5, + "e1": 0.1, + "e2": 0.1, + "center_x": 0.3, + "center_y": 0.0, + "amp": 26.0, + } + ) + kwargs_source_sigma.append( + { + "n_sersic": 1, + "R_sersic": 0.01, + "e1": 0.1, + "e2": 0.1, + "center_x": 0.2, + "center_y": 0.2, + "amp": 10, + } + ) + kwargs_lower_source.append( + { + "e1": -0.5, + "e2": -0.5, + "R_sersic": 0.001, + "n_sersic": 0.5, + "center_x": -10.0, + "center_y": -10.0, + "amp": 0.0, + } + ) + kwargs_upper_source.append( + { + "e1": 0.5, + "e2": 0.5, + "R_sersic": 10, + "n_sersic": 5.0, + "center_x": 10, + "center_y": 10, + "amp": 100, + } + ) + source_params = [ + kwargs_source_init, + kwargs_source_sigma, + fixed_source, + kwargs_lower_source, + kwargs_upper_source, + ] fixed_lens_light = [] kwargs_lens_light_init = [] @@ -999,38 +1151,97 @@ def test_jaxopt_minimizer(self): kwargs_lower_lens_light = [] kwargs_upper_lens_light = [] - fixed_lens_light.append({}) - kwargs_lens_light_init.append({'R_sersic': 0.5, 'n_sersic': 2., 'e1': 0.1, 'e2': 0.3, 'center_x': 0.1, 'center_y': 0., 'amp': 7.}) - kwargs_lens_light_sigma.append({'n_sersic': 0.01, 'R_sersic': 0.03, 'e1': 0.5, 'e2': 0.5, 'center_x': 0.01, 'center_y': 0.01, 'amp': 10}) - kwargs_lower_lens_light.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.001, 'n_sersic': .5, 'center_x': -10., 'center_y': -10., 'amp': 0.}) - kwargs_upper_lens_light.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 10., 'n_sersic': 5., 'center_x': 10., 'center_y': 10., 'amp': 100.}) + kwargs_lens_light_init.append( + { + "R_sersic": 0.5, + "n_sersic": 2.0, + "e1": 0.1, + "e2": 0.3, + "center_x": 0.1, + "center_y": 0.0, + "amp": 7.0, + } + ) + kwargs_lens_light_sigma.append( + { + "n_sersic": 0.01, + "R_sersic": 0.03, + "e1": 0.5, + "e2": 0.5, + "center_x": 0.01, + "center_y": 0.01, + "amp": 10, + } + ) + kwargs_lower_lens_light.append( + { + "e1": -0.5, + "e2": -0.5, + "R_sersic": 0.001, + "n_sersic": 0.5, + "center_x": -10.0, + "center_y": -10.0, + "amp": 0.0, + } + ) + kwargs_upper_lens_light.append( + { + "e1": 0.5, + "e2": 0.5, + "R_sersic": 10.0, + "n_sersic": 5.0, + "center_x": 10.0, + "center_y": 10.0, + "amp": 100.0, + } + ) - lens_light_params = [kwargs_lens_light_init, kwargs_lens_light_sigma, fixed_lens_light, kwargs_lower_lens_light, kwargs_upper_lens_light] + lens_light_params = [ + kwargs_lens_light_init, + kwargs_lens_light_sigma, + fixed_lens_light, + kwargs_lower_lens_light, + kwargs_upper_lens_light, + ] + + kwargs_params = { + "lens_model": lens_params, + "source_model": source_params, + "lens_light_model": lens_light_params, + } - kwargs_params = {'lens_model': lens_params, - 'source_model': source_params, - 'lens_light_model': lens_light_params} - - kwargs_model = {'lens_model_list': lens_model_list, 'source_light_model_list': source_model_list, 'lens_light_model_list': lens_light_model_list} + kwargs_model = { + "lens_model_list": lens_model_list, + "source_light_model_list": source_model_list, + "lens_light_model_list": lens_light_model_list, + } multi_band_list = [[kwargs_data, kwargs_psf_gaussian, kwargs_numerics]] - kwargs_data_joint = {'multi_band_list': multi_band_list, - 'multi_band_type': 'single-band' - } + kwargs_data_joint = { + "multi_band_list": multi_band_list, + "multi_band_type": "single-band", + } - fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, self.kwargs_constraints, self.kwargs_likelihood, kwargs_params) + fitting_seq = FittingSequence( + kwargs_data_joint, + kwargs_model, + self.kwargs_constraints, + self.kwargs_likelihood, + kwargs_params, + ) # options are BFGS(), Nelder-Mead, Powell, CG, Newton-CG, L-BFGS-B, TNC(), COBYLA, SLSQP!, trust-constr!, dogleg!, trust-ncg!, trust-exact!, trust-krylov! - fitting_kwargs_list_jaxopt = [['Jaxopt', {'method': "BFGS", 'maxiter': 200}]] + fitting_kwargs_list_jaxopt = [["Jaxopt", {"method": "BFGS", "maxiter": 200}]] chain_list = fitting_seq.fit_sequence(fitting_kwargs_list_jaxopt) fitting_type, args_history, logL_history, kwargs_result = chain_list[0] assert fitting_type == "Jaxopt" assert len(args_history) == len(logL_history) - npt.assert_almost_equal(kwargs_result['kwargs_lens'][0]['theta_E'], 0.66, decimal=3) - + npt.assert_almost_equal( + kwargs_result["kwargs_lens"][0]["theta_E"], 0.66, decimal=3 + ) if __name__ == "__main__": From 190fbccbe5ab41b4e7e5a4e05dba0734a755776b Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sat, 1 Feb 2025 21:36:03 -0500 Subject: [PATCH 05/25] change import statement --- test/test_Workflow/test_fitting_sequence.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index f18a5f3..f9deb47 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -6,11 +6,13 @@ import numpy.testing as npt import numpy as np from lenstronomy.Util import simulation_util as sim_util, util -from jaxtronomy.ImSim.image_model import ImageModel -from jaxtronomy.PointSource.point_source import PointSource + from jaxtronomy.LensModel.lens_model import LensModel from jaxtronomy.LightModel.light_model import LightModel +from lenstronomy.PointSource.point_source import PointSource + from jaxtronomy.Workflow.fitting_sequence import FittingSequence +from jaxtronomy.ImSim.image_model import ImageModel from jaxtronomy.Data.imaging_data import ImageData from lenstronomy.Data.psf import PSF From d6fc9d3ce71f72c25441953562c91f8cf57415fa Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sat, 1 Feb 2025 22:05:45 -0500 Subject: [PATCH 06/25] add test requirements for fitting sequence --- test_requirements.txt | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 test_requirements.txt diff --git a/test_requirements.txt b/test_requirements.txt new file mode 100644 index 0000000..3b98d2c --- /dev/null +++ b/test_requirements.txt @@ -0,0 +1,4 @@ +dynesty +zeus-mcmc +cobaya +nautilus-sampler>=0.7 \ No newline at end of file From 481ba1a9641384c3dd723d48b200724a75f2eb40 Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sat, 1 Feb 2025 22:19:35 -0500 Subject: [PATCH 07/25] added emcee test requirement --- test_requirements.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test_requirements.txt b/test_requirements.txt index 3b98d2c..d5d4e2b 100644 --- a/test_requirements.txt +++ b/test_requirements.txt @@ -1,4 +1,5 @@ dynesty zeus-mcmc cobaya -nautilus-sampler>=0.7 \ No newline at end of file +nautilus-sampler>=0.7 +emcee>=3.0.0 \ No newline at end of file From 332dbfe5f96d29c64d14386b6114e3d8e1ca7a20 Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sun, 2 Feb 2025 01:28:46 -0500 Subject: [PATCH 08/25] fixed tests --- test/test_Workflow/test_fitting_sequence.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index 9804365..b7780f7 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -432,7 +432,7 @@ def test_cobaya(self): source_model_list = ["SERSIC"] kwargs_source = [ { - "amp": 1.0, + "amp": 10.0, "R_sersic": 0.3, "n_sersic": 3.0, "center_x": 0.1, @@ -467,9 +467,9 @@ def test_cobaya(self): } lens_fixed = [{"center_x": 0.0, "center_y": 0.0}] - lens_sigma = [{"theta_E": 0.01, "center_x": 0.1, "center_y": 0.1}] - lens_lower = [{"theta_E": 0.1, "center_x": -10, "center_y": -10}] - lens_upper = [{"theta_E": 3.0, "center_x": 10, "center_y": 10}] + lens_sigma = [{"theta_E": 0.1, "center_x": 0.1, "center_y": 0.1}] + lens_lower = [{"theta_E": 1.0, "center_x": -10, "center_y": -10}] + lens_upper = [{"theta_E": 2.0, "center_x": 10, "center_y": 10}] source_fixed = [{}] source_sigma = [ @@ -491,7 +491,7 @@ def test_cobaya(self): } ] source_upper = [ - {"amp": 100, "R_sersic": 1.0, "n_sersic": 6.0, "center_x": 1, "center_y": 1} + {"amp": 20, "R_sersic": 1.0, "n_sersic": 6.0, "center_x": 1, "center_y": 1} ] lens_param = [kwargs_lens, lens_sigma, lens_fixed, lens_lower, lens_upper] @@ -521,7 +521,7 @@ def test_cobaya(self): ) kwargs_cobaya = { - "proposal_widths": [0.001, 0.001, 0.001, 0.001, 0.001, 0.001], + "proposal_widths": [0.01, 0.01, 0.01, 0.01, 0.01, 0.01], "Rminus1_stop": 100, # does this need to be large? can we run in test mode? "force_overwrite": True, } @@ -570,7 +570,7 @@ def test_zeus(self): # make a source source_model_list = ["SERSIC_ELLIPSE"] kwargs_sersic_ellipse = { - "amp": 1.0, + "amp": 10.0, "R_sersic": 0.6, "n_sersic": 3, "center_x": 0.0, @@ -663,7 +663,7 @@ def test_zeus(self): ] source_upper = [ { - "amp": 100, + "amp": 20, "R_sersic": 10, "n_sersic": 5.5, "center_x": 2, From 24d68eb10177aa4ec895642dec82706a3256c228 Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sun, 2 Feb 2025 02:16:47 -0500 Subject: [PATCH 09/25] changed tests --- test/test_Workflow/test_fitting_sequence.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index b7780f7..6640465 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -454,7 +454,7 @@ def test_cobaya(self): kwargs_numerics=kwargs_numerics, ) image_sim = sim_util.simulate_simple( - imageModel, kwargs_lens, kwargs_source, point_source_add=False + imageModel, kwargs_lens, kwargs_source, point_source_add=False, no_noise=True ) data_class.update_data(image_sim) @@ -483,7 +483,7 @@ def test_cobaya(self): ] source_lower = [ { - "amp": 0, + "amp": 9.5, "R_sersic": 0.01, "n_sersic": 0.5, "center_x": -1, @@ -491,7 +491,7 @@ def test_cobaya(self): } ] source_upper = [ - {"amp": 20, "R_sersic": 1.0, "n_sersic": 6.0, "center_x": 1, "center_y": 1} + {"amp": 10.5, "R_sersic": 1.0, "n_sersic": 6.0, "center_x": 1, "center_y": 1} ] lens_param = [kwargs_lens, lens_sigma, lens_fixed, lens_lower, lens_upper] @@ -522,7 +522,7 @@ def test_cobaya(self): kwargs_cobaya = { "proposal_widths": [0.01, 0.01, 0.01, 0.01, 0.01, 0.01], - "Rminus1_stop": 100, # does this need to be large? can we run in test mode? + "Rminus1_stop": 100, "force_overwrite": True, } @@ -594,7 +594,7 @@ def test_zeus(self): kwargs_numerics=kwargs_numerics, ) image_sim = sim_util.simulate_simple( - imageModel, kwargs_lens, kwargs_source, point_source_add=False + imageModel, kwargs_lens, kwargs_source, point_source_add=False, no_noise=True ) data_class.update_data(image_sim) @@ -652,7 +652,7 @@ def test_zeus(self): ] source_lower = [ { - "amp": 0, + "amp": 9.5, "R_sersic": 0.01, "n_sersic": 0.5, "center_x": -2, @@ -663,7 +663,7 @@ def test_zeus(self): ] source_upper = [ { - "amp": 20, + "amp": 10.5, "R_sersic": 10, "n_sersic": 5.5, "center_x": 2, From cf366c636b4acab2aff344483233d4c25f11cb67 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 2 Feb 2025 07:17:00 +0000 Subject: [PATCH 10/25] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- test/test_Workflow/test_fitting_sequence.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index 6640465..4e77b52 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -454,7 +454,11 @@ def test_cobaya(self): kwargs_numerics=kwargs_numerics, ) image_sim = sim_util.simulate_simple( - imageModel, kwargs_lens, kwargs_source, point_source_add=False, no_noise=True + imageModel, + kwargs_lens, + kwargs_source, + point_source_add=False, + no_noise=True, ) data_class.update_data(image_sim) @@ -491,7 +495,13 @@ def test_cobaya(self): } ] source_upper = [ - {"amp": 10.5, "R_sersic": 1.0, "n_sersic": 6.0, "center_x": 1, "center_y": 1} + { + "amp": 10.5, + "R_sersic": 1.0, + "n_sersic": 6.0, + "center_x": 1, + "center_y": 1, + } ] lens_param = [kwargs_lens, lens_sigma, lens_fixed, lens_lower, lens_upper] @@ -594,7 +604,11 @@ def test_zeus(self): kwargs_numerics=kwargs_numerics, ) image_sim = sim_util.simulate_simple( - imageModel, kwargs_lens, kwargs_source, point_source_add=False, no_noise=True + imageModel, + kwargs_lens, + kwargs_source, + point_source_add=False, + no_noise=True, ) data_class.update_data(image_sim) From ca4b61aca7c31b5b97b137d730fd7604f449f0fb Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sun, 2 Feb 2025 14:39:28 -0500 Subject: [PATCH 11/25] fixed compatibility issues with non-jax samplers --- jaxtronomy/Sampling/likelihood.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/jaxtronomy/Sampling/likelihood.py b/jaxtronomy/Sampling/likelihood.py index e6ca979..e50eb60 100644 --- a/jaxtronomy/Sampling/likelihood.py +++ b/jaxtronomy/Sampling/likelihood.py @@ -10,6 +10,7 @@ import jax from jax import jit, lax, numpy as jnp from functools import partial +import numpy as np __all__ = ["LikelihoodModule"] @@ -32,7 +33,7 @@ def __init__( kwargs_model, param_class, image_likelihood=True, - check_bounds=False, + check_bounds=True, astrometric_likelihood=False, image_position_likelihood=False, source_position_likelihood=False, @@ -417,8 +418,10 @@ def effective_num_data_points(self, **kwargs): num_param, param_names = self.param.num_param() return self.num_data - num_param + # This function should be used to convert the jax type to a normal float + # Required for samplers e.g. Cobaya which do not work with jax types def likelihood(self, a): - return self.logL(a) + return np.float64(self.logL(a)) def negativelogL(self, a): """For minimizer function, the negative value of the logl value is requested. From eee8c1c4f1eb30133b79744310f106ae8bca2a5e Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sun, 2 Feb 2025 14:40:03 -0500 Subject: [PATCH 12/25] changed tests --- test/test_Workflow/test_fitting_sequence.py | 22 +++++++++------------ 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index 4e77b52..b3d75b4 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -144,6 +144,7 @@ def setup_method(self): self.kwargs_likelihood = { # This is false by default anyways "check_positive_flux": False, + "source_marg": False } lens_sigma = [ @@ -432,7 +433,7 @@ def test_cobaya(self): source_model_list = ["SERSIC"] kwargs_source = [ { - "amp": 10.0, + "amp": 1.0, "R_sersic": 0.3, "n_sersic": 3.0, "center_x": 0.1, @@ -472,13 +473,12 @@ def test_cobaya(self): lens_fixed = [{"center_x": 0.0, "center_y": 0.0}] lens_sigma = [{"theta_E": 0.1, "center_x": 0.1, "center_y": 0.1}] - lens_lower = [{"theta_E": 1.0, "center_x": -10, "center_y": -10}] - lens_upper = [{"theta_E": 2.0, "center_x": 10, "center_y": 10}] + lens_lower = [{"theta_E": 0.1, "center_x": -10, "center_y": -10}] + lens_upper = [{"theta_E": 3.0, "center_x": 10, "center_y": 10}] - source_fixed = [{}] + source_fixed = [{"amp": 1.0}] source_sigma = [ { - "amp": 0.1, "R_sersic": 0.01, "n_sersic": 0.01, "center_x": 0.01, @@ -487,7 +487,6 @@ def test_cobaya(self): ] source_lower = [ { - "amp": 9.5, "R_sersic": 0.01, "n_sersic": 0.5, "center_x": -1, @@ -496,7 +495,6 @@ def test_cobaya(self): ] source_upper = [ { - "amp": 10.5, "R_sersic": 1.0, "n_sersic": 6.0, "center_x": 1, @@ -531,8 +529,9 @@ def test_cobaya(self): ) kwargs_cobaya = { - "proposal_widths": [0.01, 0.01, 0.01, 0.01, 0.01, 0.01], + "proposal_widths": [0.001, 0.001, 0.001, 0.001, 0.001], "Rminus1_stop": 100, + "max_tries": 1000, "force_overwrite": True, } @@ -580,7 +579,7 @@ def test_zeus(self): # make a source source_model_list = ["SERSIC_ELLIPSE"] kwargs_sersic_ellipse = { - "amp": 10.0, + "amp": 1.0, "R_sersic": 0.6, "n_sersic": 3, "center_x": 0.0, @@ -652,10 +651,9 @@ def test_zeus(self): } ] - source_fixed = [{}] + source_fixed = [{"amp": 1.0}] source_sigma = [ { - "amp": 0.1, "R_sersic": 0.05, "n_sersic": 0.5, "center_x": 0.1, @@ -666,7 +664,6 @@ def test_zeus(self): ] source_lower = [ { - "amp": 9.5, "R_sersic": 0.01, "n_sersic": 0.5, "center_x": -2, @@ -677,7 +674,6 @@ def test_zeus(self): ] source_upper = [ { - "amp": 10.5, "R_sersic": 10, "n_sersic": 5.5, "center_x": 2, From b4c2479b6f908b6c9020e8f851d05768021c539c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 2 Feb 2025 19:40:15 +0000 Subject: [PATCH 13/25] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- test/test_Workflow/test_fitting_sequence.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index b3d75b4..eb4c66a 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -144,7 +144,7 @@ def setup_method(self): self.kwargs_likelihood = { # This is false by default anyways "check_positive_flux": False, - "source_marg": False + "source_marg": False, } lens_sigma = [ From 48fad9ca63e9d754ed0970b13e8cf66ba2f312c2 Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sun, 2 Feb 2025 15:09:36 -0500 Subject: [PATCH 14/25] removed psf_iteration and calibrate_images --- jaxtronomy/Workflow/fitting_sequence.py | 92 +-------------------- test/test_Workflow/test_fitting_sequence.py | 11 +-- 2 files changed, 5 insertions(+), 98 deletions(-) diff --git a/jaxtronomy/Workflow/fitting_sequence.py b/jaxtronomy/Workflow/fitting_sequence.py index acf8011..93ca450 100644 --- a/jaxtronomy/Workflow/fitting_sequence.py +++ b/jaxtronomy/Workflow/fitting_sequence.py @@ -128,13 +128,12 @@ def fit_sequence(self, fitting_list): elif fitting_type == "psf_iteration": raise ValueError("psf_iteration not supported in jaxtronomy") - # self.psf_iteration(**kwargs) elif fitting_type == "align_images": self.align_images(**kwargs) elif fitting_type == "calibrate_images": - self.flux_calibration(**kwargs) + raise ValueError("calibrate_images not supported in jaxtronomy") elif fitting_type == "PSO": kwargs_result, chain, param = self.pso(**kwargs) @@ -630,50 +629,6 @@ def nested_sampling( return output - def psf_iteration(self, compute_bands=None, **kwargs_psf_iter): - """Iterative PSF reconstruction. - - :param compute_bands: bool list, if multiple bands, this process can be limited - to a subset of bands - :param kwargs_psf_iter: keyword arguments as used or available in - PSFIteration.update_iterative() definition - :return: 0, updated PSF is stored in self.multi_band_list - """ - kwargs_model = self._updateManager.kwargs_model - kwargs_likelihood = self._updateManager.kwargs_likelihood - likelihood_mask_list = kwargs_likelihood.get("image_likelihood_mask_list", None) - kwargs_pixelbased = kwargs_likelihood.get("kwargs_pixelbased", None) - kwargs_temp = self.best_fit(bijective=False) - if compute_bands is None: - compute_bands = [True] * len(self.multi_band_list) - - for band_index in range(len(self.multi_band_list)): - if compute_bands[band_index] is True: - kwargs_psf = self.multi_band_list[band_index][1] - kwargs_psf_before = copy.deepcopy(kwargs_psf) - image_model = SingleBandMultiModel( - self.multi_band_list, - kwargs_model, - likelihood_mask_list=likelihood_mask_list, - band_index=band_index, - kwargs_pixelbased=kwargs_pixelbased, - ) - psf_iter = PsfFitting(image_model_class=image_model) - kwargs_psf = psf_iter.update_iterative( - kwargs_psf, kwargs_params=kwargs_temp, **kwargs_psf_iter - ) - self.multi_band_list[band_index][1] = kwargs_psf - self._psf_iteration_memory.append( - { - "sequence": self._psf_iteration_index, - "band": band_index, - "psf_before": kwargs_psf_before, - "psf_after": kwargs_psf, - } - ) - self._psf_iteration_index += 1 - return 0 - def align_images( self, n_particles=10, @@ -742,51 +697,6 @@ def align_images( self.multi_band_list[i][0] = kwargs_data return 0 - def flux_calibration( - self, - n_particles=10, - n_iterations=10, - threadCount=1, - calibrate_bands=None, - scaling_lower_limit=0, - scaling_upper_limit=1000, - ): - """Calibrates flux_scaling between multiple images. This routine only works in - 'join-linear' model when fluxes are meant to be identical for different bands. - - :param n_particles: number of particles in the Particle Swarm Optimization - :param n_iterations: number of iterations in the optimization process - :param calibrate_bands: state which bands the flux calibration is applied to - :type calibrate_bands: list of booleans of length of the imaging bands - :param threadCount: number of CPU threads. If MPI option is set, threadCount=1 - :type threadCount: integer - :param scaling_lower_limit: lower limit of flux_scaling - :param scaling_upper_limit: upper limit of flux scaling - :return: 0, updated coordinate system for the band(s) - """ - kwargs_model = self._updateManager.kwargs_model - kwargs_temp = self.best_fit(bijective=False) - multi_band_type = self.kwargs_data_joint.get("multi_band_type", "multi-linear") - kwargs_imaging = self.likelihoodModule.kwargs_imaging - - calibration_fitting = FluxCalibration( - kwargs_imaging=kwargs_imaging, - kwargs_model=kwargs_model, - kwargs_params=kwargs_temp, - calibrate_bands=calibrate_bands, - ) - - multi_band_list, chain = calibration_fitting.pso( - n_particles=n_particles, - n_iterations=n_iterations, - threadCount=threadCount, - mpi=self._mpi, - scaling_lower_limit=scaling_lower_limit, - scaling_upper_limit=scaling_upper_limit, - ) - self.multi_band_list = multi_band_list - return 0 - def update_settings( self, kwargs_model=None, diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index eb4c66a..82d3830 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -392,13 +392,9 @@ def test_fitting_sequence(self): fitting_list4 = [["psf_iteration", kwargs_psf_iter]] with t.assertRaises(ValueError): fittingSequence.fit_sequence(fitting_list4) - - # psf_iteration_list = fittingSequence.psf_iteration_memory - # assert len(psf_iteration_list) == 1 - # assert "sequence" in psf_iteration_list[0] - # assert "band" in psf_iteration_list[0] - # assert "psf_before" in psf_iteration_list[0] - # assert "psf_after" in psf_iteration_list[0] + fitting_list5 = [["calibrate_images", {}]] + with t.assertRaises(ValueError): + fittingSequence.fit_sequence(fitting_list5) def test_cobaya(self): np.random.seed(42) @@ -536,6 +532,7 @@ def test_cobaya(self): } chain_list = fittingSequence.fit_sequence([["Cobaya", kwargs_cobaya]]) + assert fittingSequence.kwargs_fixed == (lens_fixed, source_fixed, [], [], {}, [], []) def test_zeus(self): np.random.seed(42) From 27b4026ad1fcc201adda12088ee4595610976e25 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 2 Feb 2025 20:10:16 +0000 Subject: [PATCH 15/25] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- test/test_Workflow/test_fitting_sequence.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index 82d3830..c6df1f2 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -532,7 +532,15 @@ def test_cobaya(self): } chain_list = fittingSequence.fit_sequence([["Cobaya", kwargs_cobaya]]) - assert fittingSequence.kwargs_fixed == (lens_fixed, source_fixed, [], [], {}, [], []) + assert fittingSequence.kwargs_fixed == ( + lens_fixed, + source_fixed, + [], + [], + {}, + [], + [], + ) def test_zeus(self): np.random.seed(42) From bd8175cb912dab88c4e33be100913c655aeeefc4 Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sun, 2 Feb 2025 15:16:48 -0500 Subject: [PATCH 16/25] changed docstrings --- jaxtronomy/Workflow/fitting_sequence.py | 9 +++++---- test/test_Workflow/test_fitting_sequence.py | 5 ++++- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/jaxtronomy/Workflow/fitting_sequence.py b/jaxtronomy/Workflow/fitting_sequence.py index 93ca450..34d9725 100644 --- a/jaxtronomy/Workflow/fitting_sequence.py +++ b/jaxtronomy/Workflow/fitting_sequence.py @@ -413,14 +413,15 @@ def mcmc( def jaxopt( self, - method="BFGS", + method="TNC", maxiter=100, ): """Uses the Jaxopt Scipy Minimizer. - :param method: string, options are BFGS, Nelder-Mead, Powell, CG, BFGS, Newton- - CG, L-BFGS-B, TNC, COBYLA, SLSQP, trust-constr, dogleg, trust-ncg, trust- - exact, trust-krylov + :param method: string. Otions are BFGS and TNC. + Other options such as Nelder-Mead, Powell, CG, Newton-CG, L-BFGS-B, COBYLA, + SLSQP, trust-constr, dogleg, trust-ncg, trust-exact, trust-krylov + either do not work yet or do not perform as well as BFGS and TNC :param maxiter: int, number of iterations to perform during minimization of the loss function """ diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index c6df1f2..6296148 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -1249,7 +1249,10 @@ def test_jaxopt_minimizer(self): kwargs_params, ) - # options are BFGS(), Nelder-Mead, Powell, CG, Newton-CG, L-BFGS-B, TNC(), COBYLA, SLSQP!, trust-constr!, dogleg!, trust-ncg!, trust-exact!, trust-krylov! + # options are BFGS and TNC + # Other options such as Nelder-Mead, Powell, CG, Newton-CG, L-BFGS-B, COBYLA, + # SLSQP, trust-constr, dogleg, trust-ncg, trust-exact, trust-krylov + # either do not work yet or do not perform as well as BFGS and TNC fitting_kwargs_list_jaxopt = [["Jaxopt", {"method": "BFGS", "maxiter": 200}]] chain_list = fitting_seq.fit_sequence(fitting_kwargs_list_jaxopt) fitting_type, args_history, logL_history, kwargs_result = chain_list[0] From 3802c113fe65986ae9e7d4f52205f2e9b409c71d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 2 Feb 2025 20:17:00 +0000 Subject: [PATCH 17/25] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- jaxtronomy/Workflow/fitting_sequence.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/jaxtronomy/Workflow/fitting_sequence.py b/jaxtronomy/Workflow/fitting_sequence.py index 34d9725..28f51db 100644 --- a/jaxtronomy/Workflow/fitting_sequence.py +++ b/jaxtronomy/Workflow/fitting_sequence.py @@ -418,10 +418,10 @@ def jaxopt( ): """Uses the Jaxopt Scipy Minimizer. - :param method: string. Otions are BFGS and TNC. - Other options such as Nelder-Mead, Powell, CG, Newton-CG, L-BFGS-B, COBYLA, - SLSQP, trust-constr, dogleg, trust-ncg, trust-exact, trust-krylov - either do not work yet or do not perform as well as BFGS and TNC + :param method: string. Otions are BFGS and TNC. Other options such as Nelder- + Mead, Powell, CG, Newton-CG, L-BFGS-B, COBYLA, SLSQP, trust-constr, dogleg, + trust-ncg, trust-exact, trust-krylov either do not work yet or do not + perform as well as BFGS and TNC :param maxiter: int, number of iterations to perform during minimization of the loss function """ From ef7ccf952779cd29467619394ef5cf571ffb96d0 Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sun, 2 Feb 2025 15:18:36 -0500 Subject: [PATCH 18/25] added test --- test/test_Workflow/test_fitting_sequence.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index 6296148..d6b2710 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -982,12 +982,6 @@ def test_jaxopt_minimizer(self): lens_light_model_list = ["SERSIC_ELLIPSE"] kwargs_lens_light = [kwargs_sersic_lens] - kwargs_truth = { - "kwargs_lens": kwargs_lens, - "kwargs_source": kwargs_source, - "kwargs_lens_light": kwargs_lens_light, - } - lens_light_model_class = LightModel(lens_light_model_list) # generate the coordinate grid and image properties (we only read out the relevant lines we need) _, _, ra_at_xy_0, dec_at_xy_0, _, _, Mpix2coord, _ = ( @@ -1262,6 +1256,7 @@ def test_jaxopt_minimizer(self): npt.assert_almost_equal( kwargs_result["kwargs_lens"][0]["theta_E"], 0.66, decimal=3 ) + npt.assert_almost_equal(logL_history[-1], 0, decimal=-2) if __name__ == "__main__": From be20d7e5860678e7083d9163fdd7edf7ddf63377 Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sun, 2 Feb 2025 15:43:09 -0500 Subject: [PATCH 19/25] added test for coverage --- test/test_Workflow/test_fitting_sequence.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index d6b2710..dfa4e95 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -811,7 +811,7 @@ def test_dynesty(self): }, } - fitting_list.append(["dynesty", kwargs_dynesty]) + fitting_list.append(["nested_sampling", kwargs_dynesty]) chain_list = fittingSequence.fit_sequence(fitting_list) @@ -883,10 +883,11 @@ def test_minimizer(self): kwargs_mcmc = {"sigma_scale": 1, "n_burn": 1, "n_run": 1, "n_walkers": 10} fitting_list.append(["emcee", kwargs_mcmc]) kwargs_mcmc["re_use_samples"] = True + # Change the number of parameters from 1 to 2; this should raise an error kwargs_mcmc["init_samples"] = np.array( - [[np.random.normal(1, 0.001)] for i in range(100)] + [[np.random.normal(1, 0.001)] * 2 for i in range(100)] ) - fitting_list.append(["emcee", kwargs_mcmc]) + fitting_list.append(["MCMC", kwargs_mcmc]) def custom_likelihood(kwargs_lens, **kwargs): theta_E = kwargs_lens[0]["theta_E"] @@ -915,6 +916,10 @@ def custom_likelihood(kwargs_lens, **kwargs): kwargs_likelihood, kwargs_params, ) + npt.assert_raises(ValueError, fittingSequence.fit_sequence, fitting_list) + kwargs_mcmc["init_samples"] = None + fitting_list[-1][1] = kwargs_mcmc + args = fittingSequence.param_class.kwargs2args( kwargs_lens=[{"theta_E": 1, "center_x": 0, "center_y": 0}] ) From ce00e7e0d97c6de1697ce429ec119bb480e9e8fd Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sun, 2 Feb 2025 15:43:23 -0500 Subject: [PATCH 20/25] removed unused function and imports --- jaxtronomy/Workflow/fitting_sequence.py | 23 +---------------------- 1 file changed, 1 insertion(+), 22 deletions(-) diff --git a/jaxtronomy/Workflow/fitting_sequence.py b/jaxtronomy/Workflow/fitting_sequence.py index 28f51db..7a14109 100644 --- a/jaxtronomy/Workflow/fitting_sequence.py +++ b/jaxtronomy/Workflow/fitting_sequence.py @@ -1,12 +1,7 @@ -import copy - from jaxtronomy.Sampling.likelihood import LikelihoodModule -from jaxtronomy.ImSim.MultiBand.single_band_multi_model import SingleBandMultiModel from jaxtronomy.Sampling.Samplers.jaxopt_minimizer import JaxoptMinimizer -from lenstronomy.Workflow.psf_fitting import PsfFitting from lenstronomy.Workflow.alignment_matching import AlignmentFitting -from lenstronomy.Workflow.flux_calibration import FluxCalibration from lenstronomy.Workflow.multi_band_manager import MultiBandUpdateManager from lenstronomy.Sampling.sampler import Sampler from lenstronomy.Sampling.Samplers.multinest_sampler import MultiNestSampler @@ -76,8 +71,6 @@ def __init__( num_bands=len(self.multi_band_list), ) self._mcmc_init_samples = None - self._psf_iteration_memory = [] - self._psf_iteration_index = 0 # index of the sequence of the PSF iteration (how many times it is being run) @property def kwargs_fixed(self): @@ -853,18 +846,4 @@ def best_fit_from_samples(self, samples, logl): best_fit_result = best_fit_sample.tolist() # get corresponding kwargs kwargs_result = self.param_class.args2kwargs(best_fit_result, bijective=True) - return kwargs_result - - @property - def psf_iteration_memory(self): - """ - 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 - - :return: list of all psf corrections - """ - return self._psf_iteration_memory + return kwargs_result \ No newline at end of file From ecc4225b47eb8173876c8e204e75f0a67cab6bf6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 2 Feb 2025 20:43:35 +0000 Subject: [PATCH 21/25] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- jaxtronomy/Workflow/fitting_sequence.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jaxtronomy/Workflow/fitting_sequence.py b/jaxtronomy/Workflow/fitting_sequence.py index 7a14109..5bef9d4 100644 --- a/jaxtronomy/Workflow/fitting_sequence.py +++ b/jaxtronomy/Workflow/fitting_sequence.py @@ -846,4 +846,4 @@ def best_fit_from_samples(self, samples, logl): best_fit_result = best_fit_sample.tolist() # get corresponding kwargs kwargs_result = self.param_class.args2kwargs(best_fit_result, bijective=True) - return kwargs_result \ No newline at end of file + return kwargs_result From 91b49ec3a0d5c38d51447cce87ec20ab3ae50646 Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sun, 2 Feb 2025 16:05:33 -0500 Subject: [PATCH 22/25] added more tests for coverage --- test/test_Workflow/test_fitting_sequence.py | 6 +++++- test_requirements.txt | 4 +++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index dfa4e95..7a0a9db 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -866,6 +866,11 @@ def test_dypolychord(self): fitting_list.append(["dyPolyChord", kwargs_dypolychord]) + kwargs_dypolychord2 = copy.deepcopy(kwargs_dypolychord) + kwargs_dypolychord2["kwargs_run"]["resume_dyn_run"] = True + kwargs_dypolychord2["prior_type"] = "gaussian" + fitting_list.append(["dyPolyChord", kwargs_dypolychord2]) + chain_list = fittingSequence.fit_sequence(fitting_list) def test_minimizer(self): @@ -918,7 +923,6 @@ def custom_likelihood(kwargs_lens, **kwargs): ) npt.assert_raises(ValueError, fittingSequence.fit_sequence, fitting_list) kwargs_mcmc["init_samples"] = None - fitting_list[-1][1] = kwargs_mcmc args = fittingSequence.param_class.kwargs2args( kwargs_lens=[{"theta_E": 1, "center_x": 0, "center_y": 0}] diff --git a/test_requirements.txt b/test_requirements.txt index d5d4e2b..303e870 100644 --- a/test_requirements.txt +++ b/test_requirements.txt @@ -2,4 +2,6 @@ dynesty zeus-mcmc cobaya nautilus-sampler>=0.7 -emcee>=3.0.0 \ No newline at end of file +emcee>=3.0.0 +nestcheck +dyPolyChord @ https://github.com/ejhigson/dyPolyChord/archive/refs/heads/master.zip \ No newline at end of file From 816ceb272a427fd7397c9a04a7200321691323f9 Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sun, 2 Feb 2025 16:16:52 -0500 Subject: [PATCH 23/25] moved some requirements to test_requirements --- requirements.txt | 4 +--- test_requirements.txt | 5 ++++- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/requirements.txt b/requirements.txt index 75c8eff..18d31b9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,4 @@ lenstronomy @ https://github.com/lenstronomy/lenstronomy/archive/refs/heads/main numpy>=1.17.0 scipy>=0.19.1 jax>=0.4.12 -jaxlib>=0.4.12 -numpyro>=0.15.3 -jaxopt>=0.8.3 \ No newline at end of file +jaxlib>=0.4.12 \ No newline at end of file diff --git a/test_requirements.txt b/test_requirements.txt index 303e870..bed4b23 100644 --- a/test_requirements.txt +++ b/test_requirements.txt @@ -1,7 +1,10 @@ +# samplers for FittingSequence +numpyro>=0.15.3 +jaxopt>=0.8.3 dynesty zeus-mcmc cobaya nautilus-sampler>=0.7 emcee>=3.0.0 -nestcheck +nestcheck @ https://github.com/ejhigson/nestcheck/archive/refs/heads/master.zip dyPolyChord @ https://github.com/ejhigson/dyPolyChord/archive/refs/heads/master.zip \ No newline at end of file From ca9d29695941d9b4f76ce7583456023669e35749 Mon Sep 17 00:00:00 2001 From: ahuang314 Date: Sun, 2 Feb 2025 16:36:26 -0500 Subject: [PATCH 24/25] removed resume_dyn_run feature in dyPolyChord sampling --- jaxtronomy/Workflow/fitting_sequence.py | 3 ++- test/test_Workflow/test_fitting_sequence.py | 2 ++ test_requirements.txt | 4 +--- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/jaxtronomy/Workflow/fitting_sequence.py b/jaxtronomy/Workflow/fitting_sequence.py index 5bef9d4..dd929b6 100644 --- a/jaxtronomy/Workflow/fitting_sequence.py +++ b/jaxtronomy/Workflow/fitting_sequence.py @@ -555,7 +555,8 @@ def nested_sampling( if sampler_type == "dyPolyChord": if "resume_dyn_run" in kwargs_run and kwargs_run["resume_dyn_run"] is True: - resume_dyn_run = True + #resume_dyn_run = True + raise ValueError("resume_dyn_run feature doesn't work") else: resume_dyn_run = False sampler = DyPolyChordSampler( diff --git a/test/test_Workflow/test_fitting_sequence.py b/test/test_Workflow/test_fitting_sequence.py index 7a0a9db..4eb7cee 100644 --- a/test/test_Workflow/test_fitting_sequence.py +++ b/test/test_Workflow/test_fitting_sequence.py @@ -870,7 +870,9 @@ def test_dypolychord(self): kwargs_dypolychord2["kwargs_run"]["resume_dyn_run"] = True kwargs_dypolychord2["prior_type"] = "gaussian" fitting_list.append(["dyPolyChord", kwargs_dypolychord2]) + npt.assert_raises(ValueError, fittingSequence.fit_sequence, fitting_list) + kwargs_dypolychord2["kwargs_run"]["resume_dyn_run"] = False chain_list = fittingSequence.fit_sequence(fitting_list) def test_minimizer(self): diff --git a/test_requirements.txt b/test_requirements.txt index bed4b23..395b55d 100644 --- a/test_requirements.txt +++ b/test_requirements.txt @@ -5,6 +5,4 @@ dynesty zeus-mcmc cobaya nautilus-sampler>=0.7 -emcee>=3.0.0 -nestcheck @ https://github.com/ejhigson/nestcheck/archive/refs/heads/master.zip -dyPolyChord @ https://github.com/ejhigson/dyPolyChord/archive/refs/heads/master.zip \ No newline at end of file +emcee>=3.0.0 \ No newline at end of file From 29a406e8f2dd348fcd2c651594e645cf6e5fc2b3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 2 Feb 2025 21:36:38 +0000 Subject: [PATCH 25/25] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- jaxtronomy/Workflow/fitting_sequence.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jaxtronomy/Workflow/fitting_sequence.py b/jaxtronomy/Workflow/fitting_sequence.py index dd929b6..c7a4283 100644 --- a/jaxtronomy/Workflow/fitting_sequence.py +++ b/jaxtronomy/Workflow/fitting_sequence.py @@ -555,7 +555,7 @@ def nested_sampling( if sampler_type == "dyPolyChord": if "resume_dyn_run" in kwargs_run and kwargs_run["resume_dyn_run"] is True: - #resume_dyn_run = True + # resume_dyn_run = True raise ValueError("resume_dyn_run feature doesn't work") else: resume_dyn_run = False