healpix_geometry_analysis.problems.optax_optimizer

Classes

OptaxOptimizerProblem

Description of the optimization problem for optax

Module Contents

class OptaxOptimizerProblem

Bases: healpix_geometry_analysis.problems.base.BaseProblem

Description of the optimization problem for optax

Parameters:

geometry (TileGeometry) – Tile geometry object

initial_params(rng_key: jax.random.PRNGKey) → dict[str, object]

Sample initial parameter values

Parameters:

rng_key (jax.random.PRNGKey) – Random number generator key

Returns:

Initial parameter values, free parameters are sampled from the uniform distribution within their limits, and frozen parameters are set to their values.

Return type:

dict[str, object]

freeze_optimizer(optimizer)

Freeze parameters of the Optax optimizer

loss(params)

Loss function to minimize with optax