AdaptiveSampler

class cosmic.sample.stroopwafel.AdaptiveSampler(parameter_space, total_systems, batch_size, BSEDict, is_interesting, SSEDict=None, derive_params=None, reject_systems='default', nproc=1, timestep_conditions=None, kappa=1.0, n_generations=1, mc_only=False, seed=None, only_save_hit_tables=False, min_active_fraction=0.01, min_entropy_change=0.01)

Bases: object

Adaptive importance sampler for binary population synthesis with COSMIC.

Parameters:

parameter_spaceParameterSpace

Defines the sampling dimensions and their distributions.

total_systemsint

Total number of systems to simulate across all phases.

batch_sizeint

Number of systems evolved per COSMIC call.

BSEDictdict

COSMIC binary stellar evolution parameters.

is_interestingcallable

Function with signature (bpp) -> (n_hits, hit_bin_nums) identifying systems of interest from COSMIC output.

SSEDictdict, optional

COSMIC single stellar evolution settings (required by COSMIC v4+; e.g. selecting the sse or METISSE stellar engine). Passed to every COSMIC evolution and, when reject_systems='default', wired into default_reject() so the ZAMS radii it computes via set_reff use the same engine. By default None (COSMIC’s sse engine).

derive_paramscallable, optional

Function (sampled) -> dict that supplies any binary parameters not drawn from parameter_space. sampled is a dict mapping each sampled parameter name to its (N,) array of physical values; the returned dict provides the remaining members of REQUIRED_PARAMS (a scalar is broadcast to all N binaries, e.g. {'mass_1': 30.0}). Every required parameter must be either sampled or returned here. May be None when all of them are sampled directly. By default None.

reject_systemscallable, optional

Function (binary_params) -> bool_mask returning True for physically unacceptable systems, where binary_params is the assembled dict of binary parameters (sampled columns merged with the output of derive_params). See default_reject(). By default uses default_reject(). Pass None to skip physical rejection entirely.

nprocint, optional

Number of CPU cores for COSMIC, by default 1

timestep_conditionsdict, optional

Dictionary of timestep conditions to pass to COSMIC, by default None.

kappafloat, optional

Gaussian width scaling factor, by default 1.0

n_generationsint, optional

Number of refinement generations, by default 1

mc_onlybool, optional

If True, only run exploration (standard Monte Carlo), by default False

seedint or None, optional

Random seed for reproducibility, by default None

only_save_hit_tablesbool, optional

If True, only rows corresponding to hit systems are kept in the accumulated bpp, bcm, initC, and kick_info tables. Non-hit rows are discarded immediately after each batch, reducing peak memory proportionally to the miss rate. The samples, weights, and is_hit arrays are unaffected — all systems are retained there for correct importance-weight calculation. By default False.

min_active_fractionfloat, optional

Minimum value of (1 - rejection_rate) used in every oversampling and normalisation calculation. Caps the oversampling multiplier at ×200 (= 2 / 0.01) and prevents near-zero denominators from producing astronomical (yes that was purposeful) array sizes or overflowing float64 normalisation constants. By default 0.01.

min_entropy_changefloat, optional

Minimum change in entropy required for a generation to be considered as having made progress. By default 0.01.

Attributes Summary

REQUIRED_PARAMS

The parameters that fully define a binary for COSMIC.

Methods Summary

from_checkpoint(checkpoint, **overrides)	Rebuild a sampler from a checkpoint, ready for run_refinement().
run()	Run the full STROOPWAFEL pipeline in a single call.
run_exploration()	Run the exploration and adaptation phases and return a checkpoint.
run_refinement()	Run the refinement phase and return the final result.

Attributes Documentation

REQUIRED_PARAMS = ('mass_1', 'mass_2', 'porb', 'ecc', 'metallicity')

The parameters that fully define a binary for COSMIC. Each must be either sampled (a Parameter in parameter_space) or returned by derive_params.

Methods Documentation

classmethod from_checkpoint(checkpoint, **overrides)

Rebuild a sampler from a checkpoint, ready for run_refinement().

The checkpoint is self-contained: by default every constructor argument — the parameter space, BSEDict, the derive_params/reject_systems/is_interesting callables, the scalar settings, and the RNG state — is restored from the file, so no re-specification is needed:

sampler = AdaptiveSampler.from_checkpoint('checkpoint.h5')
result = sampler.run_refinement()

Any constructor argument can be overridden by passing it as a keyword, which is useful for e.g. running refinement with more cores or a larger budget than exploration:

sampler = AdaptiveSampler.from_checkpoint(
    'checkpoint.h5', nproc=16, total_systems=1_000_000,
)

Parameters:

checkpointSTROOPWAFELCheckpoint or str

A checkpoint object or path to an HDF5 file written by STROOPWAFELCheckpoint.save().

**overrides

Any AdaptiveSampler constructor argument (parameter_space, total_systems, batch_size, BSEDict, is_interesting, derive_params, reject_systems, nproc, kappa, n_generations, only_save_hit_tables, seed). Passing seed starts a fresh RNG instead of restoring the stored state.

Returns:

AdaptiveSampler

Ready to call run_refinement().

run()

Run the full STROOPWAFEL pipeline in a single call.

For multi-job workflows (e.g. SLURM) use run_exploration() and run_refinement() instead.

Returns:

COSMICStroopOutput

Container holding all samples, weights, COSMIC output tables, and associated metadata.

run_exploration()

Run the exploration and adaptation phases and return a checkpoint.

Suitable as the first SLURM job in a two-stage workflow:

# job_explore.py
sampler = AdaptiveSampler(params, total_systems=500_000, ...)
ckpt = sampler.run_exploration()
ckpt.save('checkpoint.h5')

Returns:

STROOPWAFELCheckpoint

Serialisable snapshot containing the fitted mixture model, all exploration samples, and COSMIC output. Pass to run_refinement() (or from_checkpoint()) to continue on a different node or job.

run_refinement()

Run the refinement phase and return the final result.

Must be called after run_exploration() or after from_checkpoint() has restored exploration state. Suitable as the second SLURM job:

# job_refine.py  (the checkpoint is self-contained)
sampler = AdaptiveSampler.from_checkpoint('checkpoint.h5')
result = sampler.run_refinement()
result.save('result.h5')

Returns:

COSMICStroopOutput