jaxtronomy.Sampling.Samplers package

Submodules

jaxtronomy.Sampling.Samplers.optax module

class OptaxMinimizer(logL_func, args_mean, args_sigma, args_lower, args_upper, maxiter=500)

Bases: object

Gradient descent using Optax’s L-BFGS method.

__init__(logL_func, args_mean, args_sigma, args_lower, args_upper, maxiter=500)

Parameters:

• logL_func – callable function, usually the logL function from the likelihood module

• args_mean – array of args, to be used as the mean of a normal distribution obtained by using Param.kwargs2args

• args_sigma – array of args, to be used as the std of a normal distribution obtained by using Param.kwargs2args

• args_lower – array of args, to be used as the lower bound of a normal distribution obtained by using Param.kwargs2args

• args_upper – array of args, to be used as the upper bound of a normal distribution obtained by using Param.kwargs2args

• maxiter – int, maximum number of iterations to perform during minimization of the loss function

run(num_chains, tol, rng_seed)

Runs the gradient descent.

Parameters:

• num_chains – int, number of chains to run

• tol – float, when |logL[i] - logL[i-1]| < tol three times in a row, the gradient descent for that chain is stopped

• rng_seed – int, used to draw initial parameters from the prior distribution

run_single(init_params, tol)

Runs the gradient descent for a single chain.

Parameters:

• init_params – 1d array of floats, initial parameters for the loss function

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

Module contents