petitRADTRANS.retrieval.sampler#

Classes#

`SamplerContext`	Lightweight container that carries everything a sampler needs from a Retrieval.
`Sampler`	Abstract base class for samplers.
`BlackJAXSamplingResults`
`NumPyroSamplingResults`
`BlackJAXSampler`	Abstract base class for samplers.
`BlackJAXHMCSampler`	Abstract base class for samplers.
`BlackJAXNUTSSampler`	Abstract base class for samplers.
`NumPyroSampler`	Abstract base class for samplers.
`NumPyroHMCSampler`	Abstract base class for samplers.
`NumPyroNUTSSampler`	Abstract base class for samplers.
`PymultinestSampler`	Sampler class for pymultinest.
`UltranestSampler`	Abstract base class for samplers.
`DynestySampler`	Abstract base class for samplers.
`JAXNSSampler`	Sampler class for JAXNS.
`JAXNSShardedStaticNestedSampler`	Sampler class for JAXNS.

Functions#

`_require_jax`()
`_require_jaxns`()
`_require_blackjax`()
`_require_numpyro`()
`_require_arviz`()
`_require_matplotlib_pyplot`()
`_require_pymultinest`()
`_require_ultranest`()
`_require_dynesty`()
`_check_differentiable_support`(→ None)	Raise if the retrieval has only legacy (non-differentiable) model groups.
`_build_device_likelihood_wrapper`(context)	Return a PyMultiNest-compatible log-likelihood that dispatches to a JAX device.
`_build_unconstrained_mcmc_interface`(context)	Build the unconstrained-space logdensity and transforms for gradient-based samplers.
`_diagonal_inverse_mass_from_hessian`(logdensity, ...)	Diagonal inverse mass matrix from the MAP Hessian of the log-density.
`_dense_inverse_mass_from_hessian`(logdensity, map_position)	Dense inverse mass matrix from the MAP Hessian of the log-density.

Module Contents#

petitRADTRANS.retrieval.sampler.lax = None#

petitRADTRANS.retrieval.sampler.random = None#

petitRADTRANS.retrieval.sampler.tree = None#

petitRADTRANS.retrieval.sampler.vmap = None#

petitRADTRANS.retrieval.sampler.tree_util = None#

petitRADTRANS.retrieval.sampler.jit = None#

petitRADTRANS.retrieval.sampler.local_device_count = None#

petitRADTRANS.retrieval.sampler.make_mesh = None#

petitRADTRANS.retrieval.sampler.shard_map = None#

petitRADTRANS.retrieval.sampler.set_mesh = None#

petitRADTRANS.retrieval.sampler.devices = None#

petitRADTRANS.retrieval.sampler.default_device = None#

petitRADTRANS.retrieval.sampler.device_get = None#

petitRADTRANS.retrieval.sampler.clear_caches = None#

petitRADTRANS.retrieval.sampler.jnp = None#

petitRADTRANS.retrieval.sampler.NamedSharding = None#

petitRADTRANS.retrieval.sampler.P = None#

petitRADTRANS.retrieval.sampler.tfp = None#

petitRADTRANS.retrieval.sampler.hessian = None#

petitRADTRANS.retrieval.sampler._require_jax()#

petitRADTRANS.retrieval.sampler.jaxns = None#

petitRADTRANS.retrieval.sampler.Model = None#

petitRADTRANS.retrieval.sampler.NestedSampler = None#

petitRADTRANS.retrieval.sampler.TerminationCondition = None#

petitRADTRANS.retrieval.sampler.save_results = None#

petitRADTRANS.retrieval.sampler.load_results = None#

petitRADTRANS.retrieval.sampler.jaxns_resample = None#

petitRADTRANS.retrieval.sampler.ShardedStaticNestedSampler = None#

petitRADTRANS.retrieval.sampler.UniformSampler = None#

petitRADTRANS.retrieval.sampler._JAXNS_TERMINATION_FIELDS = ('ess', 'evidence_uncert', 'live_evidence_frac', 'dlogZ', 'max_samples',...#

petitRADTRANS.retrieval.sampler._require_jaxns()#

petitRADTRANS.retrieval.sampler.blackjax = None#

petitRADTRANS.retrieval.sampler._require_blackjax()#

petitRADTRANS.retrieval.sampler.numpyro = None#

petitRADTRANS.retrieval.sampler.NumPyroHMC = None#

petitRADTRANS.retrieval.sampler.NumPyroNUTS = None#

petitRADTRANS.retrieval.sampler.NumPyroMCMC = None#

petitRADTRANS.retrieval.sampler.az = None#

petitRADTRANS.retrieval.sampler.plt = None#

petitRADTRANS.retrieval.sampler._require_numpyro()#

petitRADTRANS.retrieval.sampler._require_arviz()#

petitRADTRANS.retrieval.sampler._require_matplotlib_pyplot()#

petitRADTRANS.retrieval.sampler.pymultinest = None#

petitRADTRANS.retrieval.sampler._require_pymultinest()#

petitRADTRANS.retrieval.sampler.ultranest = None#

petitRADTRANS.retrieval.sampler.MLFriends = None#

petitRADTRANS.retrieval.sampler.RobustEllipsoidRegion = None#

petitRADTRANS.retrieval.sampler.dynesty = None#

petitRADTRANS.retrieval.sampler.DynamicNestedSampler = None#

petitRADTRANS.retrieval.sampler.DynestyPool = None#

petitRADTRANS.retrieval.sampler.dyplot = None#

petitRADTRANS.retrieval.sampler.dyfunc = None#

petitRADTRANS.retrieval.sampler._require_ultranest()#

petitRADTRANS.retrieval.sampler._require_dynesty()#

class petitRADTRANS.retrieval.sampler.SamplerContext#

Lightweight container that carries everything a sampler needs from a Retrieval.

Built once by Retrieval._build_sampler_context() and passed to each sampler’s prepare() class method so that sampler implementations never depend on the full Retrieval object.

parameter_layout: petitRADTRANS.retrieval.runtime.ParameterLayout#

runtime: petitRADTRANS.retrieval.runtime.RetrievalRuntime#

configuration: Any#

output_directory: str#

retrieval_name: str#

uncertainties_mode: str = 'default'#

print_log_likelihood_for_debugging: bool = False#

log_likelihood_func: Any = None#

prior_func: Any = None#

prior_ultranest_func: Any = None#

free_parameter_names: list = []#

n_free_parameters: int = 0#

likelihood_devices: list | None = None#

petitRADTRANS.retrieval.sampler._check_differentiable_support(context: SamplerContext, sampler_label: str) → None#: Raise if the retrieval has only legacy (non-differentiable) model groups.

petitRADTRANS.retrieval.sampler._build_device_likelihood_wrapper(context: SamplerContext)#

Return a PyMultiNest-compatible log-likelihood that dispatches to a JAX device.

Each MPI rank r selects context.likelihood_devices[r % len(likelihood_devices)], so N ranks stripe naturally across M devices. Requires differentiable model groups.

Returns (device_likelihood_func, target_device).

Notes#

The user is responsible for ensuring that JAX processes do not unexpectedly contend on the same device (e.g. set JAX_PLATFORMS=cuda,cpu or CUDA_VISIBLE_DEVICES per MPI process for true isolation when using multiple GPUs).

petitRADTRANS.retrieval.sampler._build_unconstrained_mcmc_interface(context: SamplerContext)#

Build the unconstrained-space logdensity and transforms for gradient-based samplers.

Returns (free_parameter_names, initial_position, logdensity, transform_positions, log_jacobian).

class petitRADTRANS.retrieval.sampler.Sampler(log_likelihood_func, prior_func, output_directory, retrieval_name)#

Bases: abc.ABC

Abstract base class for samplers.

log_likelihood_func#

prior_func#

output_directory#

retrieval_name#

abstractmethod run_sampler(**kwargs)#

classmethod prepare(context: SamplerContext, **sampler_kwargs)#

Abstractmethod:

Build a sampler instance, run-kwargs, and summary parameters.

Every concrete sampler should override this to assemble its own run kwargs from the common SamplerContext instead of relying on the Retrieval class.

Returns:: (sampler_instance, run_kwargs, summary_parameters)

post_run(results)#

Optional hook called after run_sampler.

Handles pretty-printing, results extraction, and any other sampler-specific finalization.

Returns:: (final_results, extra_finalize_kwargs) where the extra kwargs are forwarded to Retrieval._finalize_sampler_run.

plot_diagnostics(results=None)#: Optional hook to persist sampler-specific diagnostics plots.

_get_samples_archive_path()#

_load_sample_archive_results()#

load_results()#

static _extract_figure(plot_output)#

classmethod _save_plot_figure(plot_output, output_file)#

static _summary_scalar(value)#

get_summary(free_parameter_names=None)#

static _posterior_summary_from_matrix(samples, parameter_names)#

static _posterior_summary_from_mapping(samples_by_parameter, free_parameter_names=None)#

class petitRADTRANS.retrieval.sampler.BlackJAXSamplingResults#

algorithm: str#

samples: numpy.ndarray#

unconstrained_samples: numpy.ndarray#

logdensity: numpy.ndarray#

model_log_likelihood: numpy.ndarray | None#

warmup_parameters: dict#

acceptance_rate: numpy.ndarray | None = None#

is_divergent: numpy.ndarray | None = None#

num_integration_steps: numpy.ndarray | None = None#

output_file: str | None = None#

metadata_file: str | None = None#

class petitRADTRANS.retrieval.sampler.NumPyroSamplingResults#

algorithm: str#

samples: numpy.ndarray#

unconstrained_samples: numpy.ndarray#

logdensity: numpy.ndarray#

model_log_likelihood: numpy.ndarray | None#

warmup_parameters: dict#

acceptance_rate: numpy.ndarray | None = None#

is_divergent: numpy.ndarray | None = None#

num_integration_steps: numpy.ndarray | None = None#

output_file: str | None = None#

metadata_file: str | None = None#

class petitRADTRANS.retrieval.sampler.BlackJAXSampler(log_likelihood_func, prior_func, output_directory, retrieval_name)#

Bases: Sampler

Abstract base class for samplers.

results = None#

classmethod _prepare_blackjax(context: SamplerContext, algorithm_name: str, **sampler_kwargs)#: Shared preparation logic for all BlackJAX sampler variants.

abstract property algorithm_name#

abstract property algorithm_factory#

static _to_serializable(value)#

_get_output_prefix()#

_save_results(parameter_names, samples, unconstrained_samples, logdensity, model_log_likelihood, acceptance_rate, is_divergent, num_integration_steps, warmup_parameters, metadata)#

run_sampler(initial_position, parameter_names, num_samples=1000, num_warmup=1000, initial_step_size=1.0, target_acceptance_rate=0.8, is_mass_matrix_diagonal=True, progress_bar=False, seed=42654, use_jit=True, log_jacobian_fn=None, num_chains=1, **blackjax_kwargs)#

get_results()#

load_results()#

pretty_print_results()#

get_summary(free_parameter_names=None)#

class petitRADTRANS.retrieval.sampler.BlackJAXHMCSampler(log_likelihood_func, prior_func, output_directory, retrieval_name)#

Bases: BlackJAXSampler

Abstract base class for samplers.

property algorithm_name#

property algorithm_factory#

classmethod prepare(context: SamplerContext, **sampler_kwargs)#

Build a sampler instance, run-kwargs, and summary parameters.

Every concrete sampler should override this to assemble its own run kwargs from the common SamplerContext instead of relying on the Retrieval class.

Returns:: (sampler_instance, run_kwargs, summary_parameters)

class petitRADTRANS.retrieval.sampler.BlackJAXNUTSSampler(log_likelihood_func, prior_func, output_directory, retrieval_name)#

Bases: BlackJAXSampler

Abstract base class for samplers.

property algorithm_name#

property algorithm_factory#

classmethod prepare(context: SamplerContext, **sampler_kwargs)#

Build a sampler instance, run-kwargs, and summary parameters.

Every concrete sampler should override this to assemble its own run kwargs from the common SamplerContext instead of relying on the Retrieval class.

Returns:: (sampler_instance, run_kwargs, summary_parameters)

petitRADTRANS.retrieval.sampler._diagonal_inverse_mass_from_hessian(logdensity, map_position)#

Diagonal inverse mass matrix from the MAP Hessian of the log-density.

The optimal HMC metric is the posterior precision; the inverse mass matrix is therefore an estimate of the posterior covariance. Here we use the cheap diagonal approximation 1 / precision_ii at the MAP. logdensity has a maximum at the MAP, so -diag(Hessian) is the (positive) precision diagonal. Any non-finite or non-positive entry falls back to 1.0, leaving NUTS’ own warmup adaptation to refine it. Returns a 1-D NumPy array (the diagonal metric expected when is_mass_matrix_diagonal=True).

The diagonal is computed matrix-free, ONE coordinate at a time, via Hessian-vector products: diag(H)_i = (H e_i)_i where H e_i is the forward-over-reverse product jvp(grad(logdensity), e_i). This avoids ever materialising the dense Hessian. jax.hessian would vectorise all n directions simultaneously, holding n reverse-mode tapes of the (very large, full-resolution radiative-transfer) forward model at once – which for the lbl line-by-line models needs hundreds of GB and OOM-kills a whole node. Each HVP here keeps only a SINGLE such tape resident (~the cost of one gradient, which the MAP optimiser already evaluates comfortably), so peak memory is ~1x a gradient instead of ~n x.

petitRADTRANS.retrieval.sampler._dense_inverse_mass_from_hessian(logdensity, map_position)#

Dense inverse mass matrix from the MAP Hessian of the log-density.

The dense analogue of _diagonal_inverse_mass_from_hessian(): the inverse mass matrix is the full posterior covariance, estimated as the inverse of the precision -Hessian at the MAP. The (symmetrised) precision is eigen-floored to a small positive value so the returned covariance is symmetric positive-definite even where the curvature is flat or slightly indefinite; if the Hessian is non-finite the whole matrix falls back to the identity, leaving NUTS’ warmup to adapt it. Returns an (n, n) NumPy array (the dense metric used when is_mass_matrix_diagonal=False).

class petitRADTRANS.retrieval.sampler.NumPyroSampler(log_likelihood_func, prior_func, output_directory, retrieval_name)#

Bases: Sampler

Abstract base class for samplers.

results = None#

mcmc = None#

parameter_names = None#

classmethod _prepare_numpyro(context: SamplerContext, algorithm_name: str, **sampler_kwargs)#: Shared preparation logic for all NumPyro sampler variants.

abstract property algorithm_name#

abstract property kernel_class#

static _to_serializable(value)#

_get_output_prefix()#

static _prepare_kernel_kwargs(algorithm_name, log_likelihood_func, initial_step_size, is_mass_matrix_diagonal, target_acceptance_rate, numpyro_kwargs)#

_save_results(parameter_names, samples, unconstrained_samples, logdensity, model_log_likelihood, acceptance_rate, is_divergent, num_integration_steps, warmup_parameters, metadata)#

_checkpoint_path()#

_save_checkpoint(payload)#: Atomically persist a checkpoint (write to .tmp then os.replace).

_load_checkpoint()#

_remove_checkpoint()#

static _concat_or_none(chunks)#

_run_checkpointed_nuts(*, kernel, initial_position, rng_key, num_warmup, num_samples, thinning, progress_bar, extra_fields, checkpoint_interval)#

Run a single NUTS chain in blocks, checkpointing after each block.

Collects num_samples draws in chunks of checkpoint_interval. After every block the post-warmup MCMC state (mcmc.last_state, as NumPy) and all draws so far are pickled to <prefix>_checkpoint.pkl. If that file is present on entry (a resubmitted job after a timeout), sampling resumes from it – warmup is NOT repeated, because NumPyro’s post_warmup_state carries the frozen step size / mass matrix and the chain position forward.

Returns the same tuple the non-checkpointed single-device path produces: (unconstrained_samples, potential_energy, acceptance_rate, is_divergent, num_integration_steps, warmup_parameters).

run_sampler(initial_position, parameter_names, num_samples=1000, num_warmup=1000, initial_step_size=1.0, target_acceptance_rate=0.8, is_mass_matrix_diagonal=True, progress_bar=False, seed=42654, use_jit=True, log_jacobian_fn=None, num_chains=1, thinning=1, chain_method='vectorized', checkpoint_interval=0, **numpyro_kwargs)#

get_results()#

load_results()#

pretty_print_results()#

plot_diagnostics(results=None)#: Optional hook to persist sampler-specific diagnostics plots.

get_summary(free_parameter_names=None)#

class petitRADTRANS.retrieval.sampler.NumPyroHMCSampler(log_likelihood_func, prior_func, output_directory, retrieval_name)#

Bases: NumPyroSampler

Abstract base class for samplers.

property algorithm_name#

property kernel_class#

classmethod prepare(context: SamplerContext, **sampler_kwargs)#

Build a sampler instance, run-kwargs, and summary parameters.

Every concrete sampler should override this to assemble its own run kwargs from the common SamplerContext instead of relying on the Retrieval class.

Returns:: (sampler_instance, run_kwargs, summary_parameters)

class petitRADTRANS.retrieval.sampler.NumPyroNUTSSampler(log_likelihood_func, prior_func, output_directory, retrieval_name)#

Bases: NumPyroSampler

Abstract base class for samplers.

property algorithm_name#

property kernel_class#

classmethod prepare(context: SamplerContext, **sampler_kwargs)#

Build a sampler instance, run-kwargs, and summary parameters.

Every concrete sampler should override this to assemble its own run kwargs from the common SamplerContext instead of relying on the Retrieval class.

Returns:: (sampler_instance, run_kwargs, summary_parameters)

class petitRADTRANS.retrieval.sampler.PymultinestSampler(log_likelihood_func, prior_func, output_directory, retrieval_name)#

Bases: Sampler

Sampler class for pymultinest.

seed = -1#

analyzer = None#

_n_dims = None#

outputfiles_basename#

classmethod prepare(context: SamplerContext, **sampler_kwargs)#

Build a sampler instance, run-kwargs, and summary parameters.

Every concrete sampler should override this to assemble its own run kwargs from the common SamplerContext instead of relying on the Retrieval class.

Returns:: (sampler_instance, run_kwargs, summary_parameters)

post_run(results)#

Optional hook called after run_sampler.

Handles pretty-printing, results extraction, and any other sampler-specific finalization.

Returns:: (final_results, extra_finalize_kwargs) where the extra kwargs are forwarded to Retrieval._finalize_sampler_run.

run_sampler(n_dims, **kwargs)#

Args:

importance_nested_sampling: bool: If True, Multinest will use Importance Nested Sampling (INS). Read http://arxiv.org/abs/1306.2144 for more details on INS. Please read the MultiNest README file before using the INS in MultiNest v3.0.
n_params: int: Total no. of parameters, should be equal to ndims in most cases but if you need to store some additional parameters with the actual parameters then you need to pass them through the likelihood routine.
sampling_efficiency: float: defines the sampling efficiency. 0.8 and 0.3 are recommended for parameter estimation & evidence evalutation respectively. use ‘parameter’ or ‘model’ to select the respective default values
mode_tolerance: float: MultiNest can find multiple modes & also specify which samples belong to which mode. It might be desirable to have separate samples & mode statistics for modes with local log-evidence value greater than a particular value in which case Ztol should be set to that value. If there isn’t any particularly interesting Ztol value, then Ztol should be set to a very large negative number (e.g. -1e90).
evidence_tolerance: float: A value of 0.5 should give good enough accuracy.
n_clustering_params: int: If mmodal is T, MultiNest will attempt to separate out the modes. Mode separation is done through a clustering algorithm. Mode separation can be done on all the parameters (in which case nCdims should be set to ndims) & it can also be done on a subset of parameters (in which case nCdims < ndims) which might be advantageous as clustering is less accurate as the dimensionality increases. If nCdims < ndims then mode separation is done on the first nCdims parameters.
null_log_evidence: float: If mmodal is T, MultiNest can find multiple modes & also specify which samples belong to which mode. It might be desirable to have separate samples & mode statistics for modes with local log-evidence value greater than a particular value in which case nullZ should be set to that value. If there isn’t any particulrly interesting nullZ value, then nullZ should be set to a very large negative number (e.g. -1.d90).
init_MPI: bool: initialize MPI routines?, relevant only if compiling with MPI To run pymultinest with MPI, you need mpi4py installed. Then, the libmultinest_mpi library is loaded when you run with mpiexec or similar. init_MPI should be set to False, because importing mpi4py initialises MPI already.
log_zero: float: points with loglike < logZero will be ignored by MultiNest
max_iter: int: maximum number of iterations. 0 is unlimited.
write_output: bool: write output files? This is required for analysis.
dump_callback: callable: a callback function for dumping the current status
use_MPI: bool: if True (default), if run with mpiexec and mpi4py installed, use multinest MPI library. if False, use only a single processor.
resumebool: Continue existing retrieval. If FALSE THIS WILL OVERWRITE YOUR EXISTING RETRIEVAL.
seedint: Random number generator seed, -ve value for seed from the system clock (for reproducibility)

Returns:

analyzer

load_results()#

pretty_print(free_parameter_names, prefix)#

get_summary(free_parameter_names=None)#

class petitRADTRANS.retrieval.sampler.UltranestSampler(log_likelihood_func, prior_func, output_directory, retrieval_name)#

Bases: Sampler

Abstract base class for samplers.

results = None#

_get_log_dir()#

classmethod _resolve_vectorized(context: SamplerContext, requested: bool) → bool#

classmethod prepare(context: SamplerContext, **sampler_kwargs)#

Build a sampler instance, run-kwargs, and summary parameters.

Every concrete sampler should override this to assemble its own run kwargs from the common SamplerContext instead of relying on the Retrieval class.

Returns:: (sampler_instance, run_kwargs, summary_parameters)

load_results()#

run_sampler(parameter_names, **kwargs)#

get_results()#

pretty_print_results()#

get_summary(free_parameter_names=None)#

class petitRADTRANS.retrieval.sampler.DynestySampler(log_likelihood_func, prior_func, output_directory, retrieval_name)#

Bases: Sampler

Abstract base class for samplers.

sampler = None#

results = None#

raw_results = None#

parameter_names = []#

_run_summary#

_dynesty_pool = None#

_owns_dynesty_pool = False#

static _default_use_pool() → dict[str, bool]#

classmethod _configure_dynesty_pool(log_likelihood_func, prior_func, *, pool, pool_njobs)#

_close_dynesty_pool()#

classmethod _build_dynesty_interface(context: SamplerContext, *, use_jit: bool = True, pool_safe: bool = False)#

classmethod prepare(context: SamplerContext, **sampler_kwargs)#

Build a sampler instance, run-kwargs, and summary parameters.

Every concrete sampler should override this to assemble its own run kwargs from the common SamplerContext instead of relying on the Retrieval class.

Returns:: (sampler_instance, run_kwargs, summary_parameters)

post_run(results)#

Optional hook called after run_sampler.

Handles pretty-printing, results extraction, and any other sampler-specific finalization.

Returns:: (final_results, extra_finalize_kwargs) where the extra kwargs are forwarded to Retrieval._finalize_sampler_run.

static _normalize_weights(weights)#

_get_output_prefix()#

_save_results(parameter_names, samples, log_likelihood, weighted_samples, weighted_log_likelihood, weighted_log_weights, metadata)#

_build_equal_weight_posterior(raw_results, seed)#

static _summarize_ncall(raw_results)#

run_sampler(init_kwargs, **run_kwargs)#

get_results()#

load_results()#

pretty_print_results()#

plot_diagnostics(results=None)#: Optional hook to persist sampler-specific diagnostics plots.

get_summary(free_parameter_names=None)#

class petitRADTRANS.retrieval.sampler.JAXNSSampler(log_likelihood_func, prior_func, output_directory, retrieval_name)#

Bases: Sampler

Sampler class for JAXNS.

jaxns_parameters#

seed = -1#

sampler = None#

results = None#

raw_results = None#

parameter_names = []#

offload_results_to_cpu = True#

classmethod _build_jaxns_interface(context: SamplerContext)#: Build JAXNS prior_model generator and log_likelihood wrapper.

static _coerce_termination_condition_mapping(values, source_name: str) → dict[str, Any]#

classmethod _build_termination_condition(sampler_kwargs: dict[str, Any]) → tuple[Any, dict[str, Any]]#

classmethod prepare(context: SamplerContext, **sampler_kwargs)#

Build a sampler instance, run-kwargs, and summary parameters.

Every concrete sampler should override this to assemble its own run kwargs from the common SamplerContext instead of relying on the Retrieval class.

Returns:: (sampler_instance, run_kwargs, summary_parameters)

_get_output_prefix()#

_get_raw_results_path()#

_get_equal_weight_samples_path()#

_get_parameter_names_path()#

_get_evidence_path()#

_write_evidence_sidecar(processed_results)#

Persist scalar evidence/diagnostics next to the equal-weight samples.

JAXNS stores the evidence (log Z, ESS, …) only in memory and in its native raw-results file. Writing a tiny JSON sidecar lets the dependency -light, JAXNS-free multi-run combiner (petitRADTRANS.retrieval.multichain.combine_nested_chains()) read each run’s evidence to combine independent nested-sampling runs.

_get_cpu_device()#

_materialize_raw_results(termination_reason, state)#

_to_numpy_tree(pytree)#

static _get_named_samples_mapping(raw_results)#

static _samples_mapping_to_matrix(samples_by_parameter, parameter_names=None)#

_save_equal_weight_samples(parameter_names, samples, log_likelihood)#

_build_processed_results(raw_results, *, equal_weight_samples=None, equal_weight_log_likelihood=None, output_file=None, params_file=None, raw_results_file=None)#

_build_equal_weight_results(raw_results)#

post_run(results)#: JAXNS post-run: get_results is already called in run_sampler, pass stats=results.

run_sampler(use_jit=True, parameters={}, **jaxns_kwargs)#: Run the JAXNS sampler.

get_results()#

load_results()#

pretty_print_results()#

plot_diagnostics(results=None)#: Optional hook to persist sampler-specific diagnostics plots.

get_summary(free_parameter_names=None)#

class petitRADTRANS.retrieval.sampler.JAXNSShardedStaticNestedSampler(log_likelihood_func, prior_func, output_directory, retrieval_name)#

Bases: JAXNSSampler

Sampler class for JAXNS.

log_likelihood_func#

prior_function#

output_directory#

retrieval_name#

jaxns_parameters#

seed = -1#

sampler = None#

results = None#

run_sampler(use_jit=True, **jaxns_kwargs)#: Run the JAXNS sampler.

get_results()#

load_results()#

pretty_print_results()#

`lax`
`random`
`tree`
`vmap`
`tree_util`
`jit`
`local_device_count`
`make_mesh`
`shard_map`
`set_mesh`
`devices`
`default_device`
`device_get`
`clear_caches`
`jnp`
`NamedSharding`
`P`
`tfp`
`hessian`
`jaxns`
`Model`
`NestedSampler`
`TerminationCondition`
`save_results`
`load_results`
`jaxns_resample`
`ShardedStaticNestedSampler`
`UniformSampler`
`_JAXNS_TERMINATION_FIELDS`
`blackjax`
`numpyro`
`NumPyroHMC`
`NumPyroNUTS`
`NumPyroMCMC`
`az`
`plt`
`pymultinest`
`ultranest`
`MLFriends`
`RobustEllipsoidRegion`
`dynesty`
`DynamicNestedSampler`
`DynestyPool`
`dyplot`
`dyfunc`

petitRADTRANS.retrieval.sampler#

Attributes#

Classes#

Functions#

Module Contents#

Notes#