petitRADTRANS.retrieval.multichain

Native multi-chain retrieval support.

This module holds the building blocks for running and combining multiple chains of the JAX-based retrieval samplers (NumPyro, BlackJAX, JAXNS). It is intentionally dependency-light (pure NumPy, no JAX import) so that the combine/diagnostics step can run as a lightweight task, separately from the heavy sampler processes, and so that it is straightforward to unit test.

See multichain_plan.md (same directory) for the full design. This file implements:

• shared primitives: deterministic per-chain seeding, chain/combined naming, sampler-family and output-directory mapping.

• the MCMC chain combiner (NumPyro + BlackJAX share one on-disk format)

and convergence diagnostics (split-R-hat and effective sample size), plus combine_mcmc_chains().

• process orchestration helpers: CPU partitioning, the auto execution policy, memory estimates, and run_chains_in_processes() (spawn pool + affinity + mpi-block), driven by Retrieval.run_multichain.

• the nested (JAXNS) multi-run combiner: combine_evidence() (inverse-variance log Z + between-run scatter) and combine_nested_chains() (evidence-weighted posterior).

• resume helpers (existing_chain_ids()) and the SLURM job-array launcher generator (render_chain_worker_script(), render_slurm_array_scripts()), driven by Retrieval.write_chain_launcher.

Classes

ChainExecution	Allowed values for the chain_execution policy selector.
ChainSpec	Identity of a single chain within a multi-chain run.
MCMCChainArrays	One chain's arrays, loaded from a *_samples.npz archive.
CombinedMCMC	Result of combining MCMC chains.
MCMCChainCombiner	Combine NumPyro/BlackJAX chains into one pooled posterior + diagnostics.
ProcessRunResult	Outcome of running chains across worker processes.
NestedChainArrays	One nested-sampling run: equal-weight samples + scalar evidence.
CombinedNested
NestedChainCombiner	Combine independent JAXNS runs into one evidence-weighted posterior.

Functions

derive_chain_seed(→ int)	Return the seed for chain chain_id given a base seed.
chain_retrieval_name(→ str)	Per-chain retrieval name used to namespace on-disk artifacts.
combined_retrieval_name(→ str)	Retrieval name under which the combined posterior is written.
backend_output_subdir(→ str)	Return the out_<Backend> sub-directory for a sampler type.
sampler_family(→ str)	Return 'mcmc', 'nested' or 'other' for a sampler type.
compute_mcmc_diagnostics(→ dict)	Compute split-R-hat and ESS per parameter for grouped chains.
combine_mcmc_chains(→ MCMCCombineSummary)	Combine per-chain MCMC artifacts on disk into one pooled posterior.
partition_cpus(→ list[Optional[list[int]]])	Split a CPU id list into n_groups contiguous, disjoint groups.
available_memory_gb(→ Optional[float])	Best-effort available memory in GiB from /proc/meminfo; else None.
estimate_peak_memory_gb(→ Optional[float])	Estimate peak resident memory for a given execution mode.
resolve_chain_execution(→ str)	Resolve auto to a concrete execution mode (pure policy).
run_chains_in_processes(→ ProcessRunResult)	Run one chain per worker process and collect successes/failures.
combine_evidence(→ dict)	Combine independent log-evidence estimates of the same posterior.
combine_nested_chains(→ NestedCombineSummary)	Combine per-run JAXNS artifacts on disk into one evidence-weighted posterior.
chain_artifact_path(→ str)	Path of the per-chain *_samples.npz artifact (used for resume checks).
existing_chain_ids(→ list[int])	Chain ids in 0..num_chains-1 whose per-chain artifact already exists.
render_chain_worker_script(→ str)	Render the auto-generated per-task worker script.
render_slurm_array_scripts(, env_setup_lines, str])	Render the SLURM array sbatch, the combine sbatch, and the submit driver.

Module Contents

class petitRADTRANS.retrieval.multichain.ChainExecution

Allowed values for the chain_execution policy selector.

AUTO = 'auto'

VECTORISED = 'vectorised'

SHARDED = 'sharded'

SEQUENTIAL = 'sequential'

PROCESSES = 'processes'

ALL

classmethod validate(value: str) → str

petitRADTRANS.retrieval.multichain.derive_chain_seed(base_seed: int, chain_id: int) → int

Return the seed for chain chain_id given a base seed.

Deterministic and reproducible from an integer, and shared by every execution mode so that in-process chain c and out-of-process chain c are identical for the same base seed. Kept deliberately simple (offset by the chain index); swap in jax.random.fold_in here if stronger decorrelation is ever required, but keep it the single source of truth.

petitRADTRANS.retrieval.multichain.chain_retrieval_name(base_name: str, chain_id: int) → str

Per-chain retrieval name used to namespace on-disk artifacts.

petitRADTRANS.retrieval.multichain.combined_retrieval_name(base_name: str) → str

Retrieval name under which the combined posterior is written.

petitRADTRANS.retrieval.multichain.backend_output_subdir(sampler_type: str) → str

Return the out_<Backend> sub-directory for a sampler type.

petitRADTRANS.retrieval.multichain.sampler_family(sampler_type: str) → str

Return 'mcmc', 'nested' or 'other' for a sampler type.

class petitRADTRANS.retrieval.multichain.ChainSpec

Identity of a single chain within a multi-chain run.

chain_id: int

seed: int

retrieval_name: str

classmethod build(base_name: str, base_seed: int, chain_id: int) → ChainSpec

class petitRADTRANS.retrieval.multichain.MCMCChainArrays

One chain’s arrays, loaded from a *_samples.npz archive.

chain_id: int

samples: numpy.ndarray

unconstrained_samples: numpy.ndarray | None = None

logdensity: numpy.ndarray | None = None

model_log_likelihood: numpy.ndarray | None = None

acceptance_rate: numpy.ndarray | None = None

is_divergent: numpy.ndarray | None = None

num_integration_steps: numpy.ndarray | None = None

property n_draws: int

property n_dim: int

class petitRADTRANS.retrieval.multichain.CombinedMCMC

Result of combining MCMC chains.

parameter_names: list[str]

pooled: dict[str, numpy.ndarray]

chain_index: numpy.ndarray

grouped_for_diagnostics: numpy.ndarray

n_per_chain: list[int]

diagnostics: dict

property n_total: int

class petitRADTRANS.retrieval.multichain.MCMCChainCombiner(rhat_threshold: float = 1.01)

Combine NumPyro/BlackJAX chains into one pooled posterior + diagnostics.

rhat_threshold

combine(chains: Sequence[MCMCChainArrays], parameter_names: Sequence[str]) → CombinedMCMC

petitRADTRANS.retrieval.multichain.compute_mcmc_diagnostics(grouped: numpy.ndarray, parameter_names: Sequence[str], *, divergences: dict | None = None, rhat_threshold: float = 1.01, n_per_chain: Sequence[int] | None = None) → dict

Compute split-R-hat and ESS per parameter for grouped chains.

Parameters

grouped:

Array of shape (n_chains, n_draws, n_dim).

parameter_names:

Names matching the n_dim axis.

petitRADTRANS.retrieval.multichain.combine_mcmc_chains(output_directory: os.PathLike | str, sampler_type: str, base_name: str, *, num_chains: int | None = None, write: bool = True, on_missing: str = 'warn', rhat_threshold: float = 1.01) → MCMCCombineSummary

Combine per-chain MCMC artifacts on disk into one pooled posterior.

Reads out_<Backend>/<base_name>_chain<c>_samples.npz (+ _params.json), pools the draws, computes diagnostics, and (if write) writes the combined posterior under <base_name>_combined in the same on-disk format that Retrieval._get_samples_dict reads, plus a _diagnostics.json sidecar.

Parameters

num_chains:

If given, expect chains 0..num_chains-1 (and report any missing). If None, discover whatever _chain<c>_ files exist.

on_missing:

"warn" (default) combines the survivors and warns; "raise" raises if any expected chain is missing.

petitRADTRANS.retrieval.multichain.partition_cpus(cpus: Sequence[int], n_groups: int) → list[list[int] | None]

Split a CPU id list into n_groups contiguous, disjoint groups.

Mirrors the affinity partitioning used by the launcher scripts. The remainder is spread over the first groups. If there are fewer CPUs than groups, pinning is not meaningful, so every group is None (let the OS schedule).

petitRADTRANS.retrieval.multichain.available_memory_gb() → float | None

Best-effort available memory in GiB from /proc/meminfo; else None.

petitRADTRANS.retrieval.multichain.estimate_peak_memory_gb(execution: str, num_chains: int, per_chain_memory_gb: float | None, max_workers: int | None = None) → float | None

Estimate peak resident memory for a given execution mode.

processes: min(max_workers, num_chains) x per_chain. sequential: per_chain (one resident model at a time). vectorised/sharded: num_chains x per_chain (chains co-resident). Returns None if per_chain_memory_gb is unknown.

petitRADTRANS.retrieval.multichain.resolve_chain_execution(requested: str, *, num_chains: int, num_devices: int = 1, gpu_available: bool = False, per_chain_memory_gb: float | None = None) → str

Resolve auto to a concrete execution mode (pure policy).

Policy (see multichain_plan.md S5):

• one chain -> sequential;

• GPU(s) present: sharded (one chain per device) when num_chains <= num_devices, else vectorised (batch on the devices);

• CPU only: processes (real parallelism; forced CPU “devices” share cores and do not help).

class petitRADTRANS.retrieval.multichain.ProcessRunResult

Outcome of running chains across worker processes.

succeeded_chain_ids: list[int]

failed: dict[int, str]

property all_succeeded: bool

petitRADTRANS.retrieval.multichain.run_chains_in_processes(chain_runner, chain_ids: Sequence[int], cpu_groups: Sequence[Sequence[int] | None], *, on_chain_failure: str = 'warn', max_workers: int | None = None, block_mpi: bool = True, executor_factory=None) → ProcessRunResult

Run one chain per worker process and collect successes/failures.

Parameters

chain_runner:

Module-level callable chain_runner(chain_id, cores) -> Any that builds a Retrieval and runs a single chain. Must be picklable (top-level) for the default spawn executor.

chain_ids, cpu_groups:

Parallel sequences; cpu_groups[i] is the core list (or None) for chain_ids[i].

on_chain_failure:

"warn" (default) records the failure and continues; "raise" raises after collecting results.

executor_factory:

Optional callable(max_workers) -> Executor used instead of a spawn ProcessPoolExecutor (for testing with a thread/inline executor).

class petitRADTRANS.retrieval.multichain.NestedChainArrays

One nested-sampling run: equal-weight samples + scalar evidence.

chain_id: int

sample_matrix: numpy.ndarray

log_likelihood: numpy.ndarray

parameter_names: list

log_Z_mean: float | None = None

log_Z_uncert: float | None = None

ess: float | None = None

total_num_likelihood_evaluations: float | None = None

log_efficiency: float | None = None

property n_samples: int

property n_dim: int

petitRADTRANS.retrieval.multichain.combine_evidence(log_z: Sequence[float | None], log_z_uncert: Sequence[float | None]) → dict

Combine independent log-evidence estimates of the same posterior.

Uses inverse-variance weighting in log space when per-run uncertainties are available (falling back to an unweighted mean otherwise), and reports a combined uncertainty that is the larger of the formal statistical error and the between-run scatter (standard error of the mean). The between-run scatter is the key signal: if runs disagree by more than their individual error bars, it dominates and flags non-convergence.

class petitRADTRANS.retrieval.multichain.CombinedNested

parameter_names: list

sample_matrix: numpy.ndarray

log_likelihood: numpy.ndarray

evidence: dict

diagnostics: dict

class petitRADTRANS.retrieval.multichain.NestedChainCombiner(seed: int = 0)

Combine independent JAXNS runs into one evidence-weighted posterior.

seed = 0

combine(chains: Sequence[NestedChainArrays]) → CombinedNested

petitRADTRANS.retrieval.multichain.combine_nested_chains(output_directory: os.PathLike | str, sampler_type: str, base_name: str, *, num_chains: int | None = None, write: bool = True, on_missing: str = 'warn', seed: int = 0) → NestedCombineSummary

Combine per-run JAXNS artifacts on disk into one evidence-weighted posterior.

Reads out_JAXNS/<base>_chain<c>_samples.npz (+ _evidence.json), weights each run by its evidence, resamples to an equal-weight posterior, and writes <base>_combined in the JAXNS on-disk format that get_samples reads, plus <base>_combined_evidence.json and <base>_combined_diagnostics.json.

petitRADTRANS.retrieval.multichain.chain_artifact_path(output_directory: os.PathLike | str, sampler_type: str, base_name: str, chain_id: int) → str

Path of the per-chain *_samples.npz artifact (used for resume checks).

petitRADTRANS.retrieval.multichain.existing_chain_ids(output_directory: os.PathLike | str, sampler_type: str, base_name: str, num_chains: int) → list[int]

Chain ids in 0..num_chains-1 whose per-chain artifact already exists.

petitRADTRANS.retrieval.multichain.render_chain_worker_script(builder_module: str, builder_function: str, output_directory: str, num_chains: int, chain_seed: int, retrieval_kwargs: dict | None, run_kwargs: dict | None, resume: bool) → str

Render the auto-generated per-task worker script.

The worker runs one chain (chain id from SLURM_ARRAY_TASK_ID or argv) or, in combine mode, merges the chains. builder_module/builder_function must be importable (not __main__); retrieval_kwargs/run_kwargs must be Python literals (load arrays inside the builder, not here).

petitRADTRANS.retrieval.multichain.render_slurm_array_scripts(*, base_name: str, worker_filename: str, num_chains: int, python: str = 'python', cpus_per_task: int = 8, partition: str | None = None, time: str | None = '24:00:00', mem: str | None = None, gres: str | None = None, account: str | None = None, log_dir: str = 'out', job_name: str | None = None, combine_dependency: str = 'afterany', extra_sbatch_lines: Sequence[str] = (), env_setup_lines: Sequence[str] = ()) → dict[str, str]

Render the SLURM array sbatch, the combine sbatch, and the submit driver.

Returns a mapping {filename: content}. The array job runs one chain per task (SLURM_ARRAY_TASK_ID); the combine job depends on it (afterany by default, so survivors are merged even if some tasks fail).