sheap.MasterSampler.Samplers.Utils.numpyro_utils module

NumPyro Model Utilities

Note

This require cleaning

Helper functions to construct NumPyro models from sheap parameters, including dictionary conversion, handling of tied parameters, and safe initialization near constraints.

Main Features

  • Apply arithmetic ties between parameters.

  • Convert flat parameter arrays to dictionaries for NumPyro.

  • Initialize values safely away from hard bounds.

  • Build a full NumPyro model for flux fitting with uncertainty.

Notes

  • Dependencies are expressed as tuples (tag, target_idx, src_idx, op, val) where op is one of {'+','-','*','/'}.

  • Parameters tied via dependencies are excluded from sampling and reconstructed after free parameters are sampled.

NumPyro_apply_arithmetic_ties(params, dependencies)[source]

Apply arithmetic ties to update dependent parameters in place. #TODO this is already in other place? :param params: Dictionary of parameter values keyed by "theta_{i}". :type params: dict :param dependencies: List of constraints of the form:

(tag, target_idx, src_idx, op, val).

Returns:

Updated parameter dictionary with tied targets assigned.

Return type:

dict

Parameters:

  • params (Dict[str, float])

  • dependencies (list of tuple)

Notes

  • Supported operations are +, -, *, /.

  • Example: if theta_3 = theta_1 * 2, then dependency is ("arith", 3, 1, "*", 2.0).

make_numpyro_model(name_list, wl, flux, sigma, constraints, params_i, theta_to_sheap, fixed_params, dependencies, model_func)[source]

Build a NumPyro probabilistic model for spectrum fitting.

Parameters:

  • name_list (list of str) – Names of NumPyro parameters (e.g. ["theta_0", "theta_1", ...]).

  • wl (jnp.ndarray) – Wavelength grid, shape (n_pix,).

  • flux (jnp.ndarray) – Observed flux, shape (n_pix,).

  • sigma (jnp.ndarray) – Flux uncertainties, shape (n_pix,).

  • constraints (list of tuple) – Bounds for each parameter as (low, high).

  • params_i (array-like) – Initial parameter values.

  • theta_to_sheap (dict) – Mapping from NumPyro parameter name to SHEAP parameter name.

  • fixed_params (dict) – Dictionary of fixed parameter values keyed by SHEAP names.

  • dependencies (list of tuple) – List of ties (see apply_arithmetic_ties()).

  • model_func (Callable) – Spectral model function with signature model_func(wl, theta).

Returns:

  • numpyro_model (Callable) – A function defining the NumPyro model.

  • init_value (dict) – Dictionary of initial parameter values for init_to_value.

Notes

The model: .. math:

\\theta_i \sim \mathrm{Uniform}(low_i, high_i)

f(\\lambda) = \\mathrm{model\_func}(\\lambda, \\theta)

y(\\lambda) \sim \mathcal{N}(f(\\lambda), \sigma)

  • Dependent (tied) parameters are excluded from sampling and reconstructed afterward.

  • Fixed parameters are held at their provided values.

NumPyro_params_to_dict(params, dependencies, constraints=None, eps=0.01)[source]

Convert flat parameters into a dictionary for NumPyro initialization.

Excludes dependent parameters and applies jitter to avoid initializing exactly at bounds.

Parameters:

  • params (array-like) – Flat parameter values (e.g. from a minimizer).

  • dependencies (list of tuple) – Ties of the form (tag, target_idx, src_idx, op, val).

  • constraints (dict or list, optional) – Parameter bounds. Either a dict mapping theta_i -> (low, high) or a list aligned with params.

  • eps (float, default=1e-2) – Minimum offset from bounds to avoid initialization at the edge.

Returns:

Dictionary of free parameters, keyed as theta_i.

Return type:

dict

Notes

  • If a parameter value is invalid (NaN, inf, outside bounds), it is reset to the midpoint of its allowed range.

  • Exact zeros are nudged to eps to avoid degenerate starts.