Fitting Methods#
Here we will explore the various fitting methods in AstroPhot. You have already encountered some of the methods, but here we will take a more systematic approach and discuss their strengths/weaknesses. Each method will be applied to the same problem with the same initial conditions so you can see how they operate.
%load_ext autoreload
%autoreload 2
%matplotlib inline
import torch
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
from scipy.stats import gaussian_kde as kde
from scipy.stats import norm
from tqdm import tqdm
from corner import corner
import astrophot as ap
Initializing model sky
Initializing model sersic_0
Initializing model sersic_1
Levenberg-Marquardt#
This fitter is identitied as ap.fit.LM and it employs a variant of the second order Newton’s method to converge very quickly to the local minimum. This is the generally accepted best algorithm for most use cases in \(\chi^2\) minimization. If you don’t know what to pick, start with this minimizer. The LM optimizer bridges the gap between first-order gradient descent and second order Newton’s method. When far from the minimum, Newton’s method is unstable and can give wildly wrong results, so LM takes gradient descent steps. However, near the minimum it switches to the Newton’s method which has “quadratic convergence” this means that it takes only a few iterations to converge to several decimal places. This can be represented as:
\((H + LI)h = g\)
Where H is the Hessian matrix of second derivatives, L is the damping parameter, I is the identity matrix, h is the step we will take in parameter space, and g is the gradient. We solve this linear system for h to get the next update step. The “L” scale parameter goes from L >> 1 which represents gradient descent to L << 1 which is Newton’s Method. When L >> 1 the hessian is effectively zero and we get \(h = g/L\) which is just gradient descent with \(1/L\) as the learning rate. In AstroPhot the damping parameter is treated somewhat differently, but the concept is the same.
LM can handle a lot of scenarios and converge to the minimum. Keep in mind, however, that it is seeking a local minimum, so it is best to start off the algorithm as close as possible to the best fit parameters. AstroPhot can automatically initialize, as discussed in other notebooks, but even that needs help sometimes (often in the form of a segmentation map).
The main drawback of LM is its memory consumption which goes as \(\mathcal{O}(PN)\) where P is the number of pixels and N is the number of parameters.
MODEL = initialize_model(target, False)
res_lm = ap.fit.LM(MODEL, verbose=1).fit()
print(res_lm.message)
Initializing model sky
Initializing model sersic_0
Initializing model sersic_1
==Starting LM fit for 'group' with 15 dynamic parameters and 9801 pixels==
Chi^2/DoF: 202.138, L: 1
/home/docs/checkouts/readthedocs.org/user_builds/astrophot/envs/v0.17.2/lib/python3.12/site-packages/torch/jit/_script.py:1488: DeprecationWarning: `torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`.
warnings.warn(
Chi^2/DoF: 166.515, L: 1
Chi^2/DoF: 15.1555, L: 1
Chi^2/DoF: 7.65442, L: 1
Chi^2/DoF: 5.80401, L: 1
Chi^2/DoF: 4.87884, L: 1
Chi^2/DoF: 4.32103, L: 1
Chi^2/DoF: 3.14634, L: 0.111
Chi^2/DoF: 2.80655, L: 0.111
Chi^2/DoF: 2.11155, L: 0.0123
Chi^2/DoF: 1.79243, L: 0.0123
Chi^2/DoF: 1.59195, L: 0.0123
Chi^2/DoF: 1.11131, L: 0.00137
Chi^2/DoF: 1.01572, L: 0.000152
Chi^2/DoF: 1.01388, L: 1.69e-05
Final Chi^2/DoF: 1.01388, L: 2.09e-07. Converged: success
success
Initializing model sky
Initializing model sersic_0
Initializing model sersic_1
Now that LM has found the \(\chi^2\) minimum, we can do a really neat trick. Since LM needs the hessian matrix, we have access to the hessian matrix at the minimum. This is in fact equal to the negative Fisher information matrix. If we take the matrix inverse of this matrix then we get the covariance matrix for a multivariate gaussian approximation of the \(\chi^2\) surface near the minimum. With the covariance matrix we can create a corner plot just like we would with an MCMC. We will see later that the MCMC methods (at least the ones which converge) produce very similar results!
param_names = list(MODEL.build_params_array_names())
set, sky = true_params()
fig, ax = ap.plots.covariance_matrix(
res_lm.covariance_matrix.detach().cpu().numpy(),
MODEL.get_values().detach().cpu().numpy(),
labels=param_names,
figsize=(20, 20),
reference_values=np.concatenate((sky, set.ravel())),
)
Iterative Fit (models)#
This iterative fitter is identified as ap.fit.Iter, this method is generally employed for large models where it is not feasible to hold all the relevant data in memory at once. The iterative fitter will cycle through the models in a GroupModel object and fit them one at a time to the image, using the residuals from the previous cycle. This can be a very robust way to deal with some fits, especially if the overlap between models is not too strong. It is however more dependent on good initialization than other methods like the Levenberg-Marquardt. Also, it is possible for the Iter method to get stuck in a local minimum under certain circumstances.
Note that while the Iterative fitter needs a GroupModel object to iterate over, it is not necessarily true that the sub models are ComponentModel objects, they could be GroupModel objects as well. In this way it is possible to cycle through and fit “clusters” of objects that are nearby, so long as it doesn’t consume too much memory.
By only fitting one model at a time it is possible to get caught in a local minimum, or to get out of a local minimum that a different fitter was stuck in. For this reason it can be good to mix-and-match the iterative optimizers so they can help each other get unstuck if a fit is very challenging.
MODEL = initialize_model(target, False)
res_iter = ap.fit.Iter(MODEL, verbose=1).fit()
Initializing model sky
Initializing model sersic_0
Initializing model sersic_1
Iterative Fit (Param)#
This iterative fitter is identified as ap.fit.IterParam, this is generally employed for large and interconnected models where it is not feasible to hold all the relevant data in memory at once. Unlike ap.fit.Iter which is intended to cycle through the sub models in a group model, this fitter iterates through parameters. The set of parameters which make up the model is broken into chunks and then fitting proceeds only on those chunks, rather than on all parameters simultaneously. For large models that have lots of interconnected/shared parameters, it doesn’t really make sense to cycle through one sub-model at a time as optimizing that model may throw another model that is sharing a parameter into a bad part of parameter space. Thus ap.fit.IterParam is safe to use on any AstroPhot model without concern for this issue, the fitter will industriously proceed to high likelihood solutions monotonically.
The tradeoff for this fitter is the same as for the other iterative fitter, if there are strong covariances in the likelihood structure then this fitter can take a long time to converge. The advantage here is that as the user you may take greater control over the combinations if you wish. The chunks argument can be set to an integer like 6 in which case, 6 parameters at a time will be fit (the last chunk may be smaller). Alternatively, the chunks parameter may be set to a tuple of numpy arrays, these should be boolean arrays that select the parameters for each chunk. For example, here is a possible chunks setup for a 7 parameter sersic model: ([1,1,0,0,0,0,0], [0,0,1,1,0,0,0], [0,0,0,0,1,1,1]) which makes three chunks to fit the x,y then q, PA then n, Re, Ie parameters. Note that you do not need to make the chunks exclusive, it is totally fine to have a parameter pop up in multiple chunks! Finally, there’s the order the chunks are fit in. This can either chunk_order="sequential" the default the chunks are fit in the order given, or chunk_order="random" where each iteration a new random order is decided for the chunks to be evaluated.
MODEL = initialize_model(target, False)
res_iterparam = ap.fit.IterParam(MODEL, chunks=5, verbose=1).fit()
Initializing model sky
Initializing model sersic_0
Initializing model sersic_1
==Starting LM fit for 'group' with 15 dynamic parameters and 9801 pixels==
Chi^2/DoF: 202.138, L: 1
Chi^2/DoF: 40.3008, L: 1
Chi^2/DoF: 20.1483, L: 1
Chi^2/DoF: 13.4895, L: 1
Chi^2/DoF: 10.0988, L: 1
Chi^2/DoF: 7.53482, L: 1
Chi^2/DoF: 5.2923, L: 1
Chi^2/DoF: 3.96364, L: 1
Chi^2/DoF: 2.99242, L: 1
Chi^2/DoF: 2.45525, L: 1
Chi^2/DoF: 2.07622, L: 1
Chi^2/DoF: 1.80351, L: 1
Chi^2/DoF: 1.60405, L: 1
Chi^2/DoF: 1.24465, L: 0.111
Chi^2/DoF: 1.0652, L: 0.00137
Chi^2/DoF: 1.02529, L: 0.00137
Chi^2/DoF: 1.01645, L: 0.00137
Chi^2/DoF: 1.01445, L: 0.000152
Chi^2/DoF: 1.01401, L: 0.000152
Chi^2/DoF: 1.01391, L: 0.000152
Chi^2/DoF: 1.01389, L: 0.000152
Final Chi^2/DoF: 1.01388, L: 0.000152. Converged: success
MODEL_init = initialize_model(target, False)
fig, axarr = plt.subplots(1, 4, figsize=(24, 5))
plt.subplots_adjust(wspace=0.1)
ap.plots.model_image(fig, axarr[0], MODEL_init)
axarr[0].set_title("Model before optimization")
ap.plots.residual_image(fig, axarr[1], MODEL_init, normalize_residuals=True)
axarr[1].set_title("Residuals before optimization")
ap.plots.model_image(fig, axarr[2], MODEL)
axarr[2].set_title("Model after optimization")
ap.plots.residual_image(fig, axarr[3], MODEL, normalize_residuals=True)
axarr[3].set_title("Residuals after optimization")
plt.show()
Initializing model sky
Initializing model sersic_0
Initializing model sersic_1
The ap.fit.IterParam fitter can also generate a covariance matrix of uncertainties, just keep in mind that it only evaluates the covariances for parameters in the same chunk.
param_names = list(MODEL.build_params_array_names())
set, sky = true_params()
fig, ax = ap.plots.covariance_matrix(
res_iterparam.covariance_matrix.detach().cpu().numpy(),
MODEL.get_values().detach().cpu().numpy(),
labels=param_names,
figsize=(20, 20),
reference_values=np.concatenate((sky, set.ravel())),
)
Scipy Minimize#
Any AstroPhot model becomes a function model(x) where x is a 1D tensor of
all the current dynamic parameters. This functional format is common for
external packages to use. AstroPhot includes a wrapper to access the
scipy.optimize.minimize minimizer list. AstroPhot will ensure the minimizers
respect the valid ranges set for each parameter.
Typically, the AstroPhot LM optimizer is faster and more accurate than the Scipy ones. The exact reason is unclear, but the Scipy minimizers are intended for very general use, while the LM optimizer is specifically optimized for gaussian log likelihoods.
In the case below, the minimizer thinks it has terminated successfully, although in fact it is quite far from the minimum. Consider this a lesson in trusting the “success” message from an optimizer. It turns out to be very challenging to identify if an optimizer is at a minimum, let alone the global minimum.
MODEL = initialize_model(target, False)
res_scipy = ap.fit.ScipyFit(MODEL, method="Powell", verbose=1).fit()
print(res_scipy.scipy_res)
Initializing model sky
Initializing model sersic_0
Initializing model sersic_1
Final 2NLL/DoF: 3.27593. Converged: success: True, message: Optimization terminated successfully.
message: Optimization terminated successfully.
success: True
status: 0
fun: 16029.110164752814
x: [ 2.297e+01 5.861e+01 ... 1.303e+01 4.921e+02]
nit: 26
direc: [[ 1.945e+00 1.968e-02 ... 2.035e-01 -5.820e+01]
[ 0.000e+00 1.000e+00 ... 0.000e+00 0.000e+00]
...
[ 0.000e+00 0.000e+00 ... 1.000e+00 0.000e+00]
[ 2.425e+00 7.320e-03 ... -9.140e-01 3.942e+01]]
nfev: 4189
Initializing model sky
Initializing model sersic_0
Initializing model sersic_1
Gradient Descent (Slalom)#
A gradient descent fitter uses local gradient information to determine the direction of increased likelihood in parameter space. The challenge with gradient descent is choosing a step size. The Slalom algorithm developed for AstroPhot uses a few samples along the gradient direction to determine a parabola which it can then jump to the minimum of. In some sense this is like a 1D version of the Levenberg-Marquardt algorithm and the 1 dimension it choses is that along the gradient (plus momentum).
It is also possible to access the PyTorch gradient descent algorithms like Adam through the AstroPhot wrapper ap.fit.Grad which perform gradient descent using various algorithm designed for machine learning. In general though, those algorithms perform better on stochastic gradient descent problems, not static problems like seen by AstroPhot. So Slalom tends to perform better.
As you see below, Slalom ends with a decent fit, though not good enough for perfect residuals like some other methods (Levenberg-Marquardt). This is typically the case. However, gradient descent can be very helpful for complex optimization tasks, because it is a slower optimization algorithm, it can be more stable in some circumstances. Try using it in cases where LM fails to get things back on track. Just make sure to finish off with an LM round to ensure you have settled into the minimum.
MODEL = initialize_model(target, False)
res_grad = ap.fit.Slalom(MODEL, verbose=1, momentum=0.1).fit()
Initializing model sky
Initializing model sersic_0
Initializing model sersic_1
iter: 0, step size: 3.662535e-01, posterior density: 5.116711e+05
iter: 10, step size: 3.200572e-02, posterior density: 1.076052e+05
iter: 20, step size: 3.636985e-01, posterior density: 6.932667e+04
iter: 30, step size: 9.491573e-02, posterior density: 4.954043e+04
iter: 40, step size: 7.715230e-03, posterior density: 3.508777e+04
iter: 50, step size: 7.691659e-03, posterior density: 3.078378e+04
iter: 60, step size: 8.186814e-03, posterior density: 2.880637e+04
iter: 70, step size: 1.040778e+00, posterior density: 1.249034e+04
iter: 80, step size: 4.185484e-03, posterior density: 8.939124e+03
iter: 90, step size: 6.185570e-03, posterior density: 8.454234e+03
iter: 100, step size: 9.386691e-03, posterior density: 8.193303e+03
iter: 110, step size: 1.394801e-03, posterior density: 7.879821e+03
iter: 120, step size: 2.629153e-03, posterior density: 7.644821e+03
iter: 130, step size: 1.106532e-03, posterior density: 7.453023e+03
iter: 140, step size: 9.470766e-04, posterior density: 7.310017e+03
iter: 150, step size: 1.226329e-03, posterior density: 7.220375e+03
iter: 160, step size: 5.506503e-04, posterior density: 7.009075e+03
iter: 170, step size: 2.057291e-03, posterior density: 6.964639e+03
iter: 180, step size: 7.426339e-04, posterior density: 6.868799e+03
iter: 190, step size: 2.125832e-03, posterior density: 6.832346e+03
iter: 200, step size: 5.556855e-04, posterior density: 6.756988e+03
iter: 210, step size: 4.663705e-04, posterior density: 6.705010e+03
iter: 220, step size: 5.489810e-04, posterior density: 6.592376e+03
iter: 230, step size: 1.171404e-03, posterior density: 6.564158e+03
iter: 240, step size: 4.643993e-04, posterior density: 6.550652e+03
iter: 250, step size: 6.572783e-04, posterior density: 6.522439e+03
iter: 260, step size: 2.718795e-04, posterior density: 6.518673e+03
iter: 270, step size: 2.571539e-04, posterior density: 6.506138e+03
iter: 280, step size: 2.480211e-04, posterior density: 6.494628e+03
iter: 290, step size: 2.374904e-04, posterior density: 6.478310e+03
iter: 300, step size: 6.344610e-04, posterior density: 6.464007e+03
Slalom Fitting complete in 50.4788544178009 sec with message: success
Initializing model sky
Initializing model sersic_0
Initializing model sersic_1
Metropolis Adjusted Langevin Algorithm (MALA)#
This is one of the simplest gradient based samplers, and is very powerful. The standard Metropolis Hastings algorithm will use a gaussian proposal distribution then use the Metropolis Hastings accept/reject stage. MALA uses gradient information to determine a better proposal distribution locally (while maintaining detailed balance) and then uses the Metropolis Hastings accept/reject stage. The ap.fit.MALA fitter object is just a basic wrapper over the ap.fit.func.mala function, so feel free to check it out if you want more details on this simple and powerful sampler!
MODEL = initialize_model(target, False)
# Use LM to start the sampler at a high likelihood location, no burn-in needed!
res1 = ap.fit.LM(MODEL, verbose=0).fit()
res_mala = ap.fit.MALA(
model=MODEL,
chains=4,
max_iter=300,
epsilon=8e-1,
likelihood="poisson",
mass_matrix=res1.covariance_matrix.detach().cpu().numpy(),
).fit()
chain_mala = res_mala.chain.reshape(-1, res_mala.chain.shape[-1])
Hamiltonian Monte-Carlo (HMC)#
The ap.fit.HMC takes a fixed number of steps at a fixed step size following Hamiltonian dynamics. This is in contrast to NUTS which attempts to optimally choose these parameters. The simplest way to think of HMC is as performing a number of MALA steps all in one go, so if leapfrog_steps = 10 then HMC is very similar to running MALA then taking every tenth step and adding it to the chain. HMC results will still have autocorrelation which will depend on the problem and choice of step parameters.
MODEL = initialize_model(target, False)
# Use LM to start the sampler at a high likelihood location, no burn-in needed!
res1 = ap.fit.LM(MODEL, verbose=0).fit()
# Run the HMC sampler
res_hmc = ap.fit.HMC(
MODEL,
warmup=1,
max_iter=150,
epsilon=1e-1,
leapfrog_steps=10,
inv_mass=res1.covariance_matrix,
).fit()
Metropolis Hastings#
This is the more standard MCMC algorithm using the Metropolis Hastngs accept step identified with ap.fit.MHMCMC. Under the hood, this is just a wrapper for the excellent emcee package, if you want to take advantage of more emcee features you can very easily use ap.fit.MHMCMC as a starting point. However, one should keep in mind that for large models it can take exceedingly long to actually converge to the posterior. Instead of waiting that long, we demonstrate the functionality with 100 steps (and 30 chains), but suggest using MALA for any real world problem. Still, if there is something MALA can’t handle (a function that isn’t differentiable) then MHMCMC can save the day (even if it takes all day to do it).
MODEL = initialize_model(target, False)
# Use LM to start the sampler at a high likelihood location, no burn-in needed!
print("running LM fit")
res1 = ap.fit.LM(MODEL, verbose=0).fit()
# Run the MHMCMC sampler
print("running MHMCMC sampling")
res_mh = ap.fit.MHMCMC(MODEL, verbose=1, max_iter=100).fit()
Initializing model sky
Initializing model sersic_0
Initializing model sersic_1
running LM fit
running MHMCMC sampling
# corner plot of the posterior
# note that, even 3000 samples is not enough to overcome the autocorrelation so the posterior has not converged.
param_names = list(MODEL.build_params_array_names())
set, sky = true_params()
fig = corner(
res_mh.chain,
labels=param_names,
truths=np.concatenate((sky, set.ravel())),
smooth=0.8,
)
