From 1449476fc92f7633afd885b1f8cd5ff30b9c2630 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Fri, 10 Apr 2026 16:03:17 -0400 Subject: [PATCH 01/45] first draft --- src/cosmic/sample/stroopwafel/__init__.py | 19 + src/cosmic/sample/stroopwafel/constants.py | 25 + src/cosmic/sample/stroopwafel/engine.py | 460 ++++++++++++++++++ .../stroopwafel/examples/example_bhbh.py | 125 +++++ .../stroopwafel/examples/test_new_modules.py | 299 ++++++++++++ src/cosmic/sample/stroopwafel/io.py | 107 ++++ src/cosmic/sample/stroopwafel/meson.build | 19 + .../sample/stroopwafel/mixture_model.py | 316 ++++++++++++ .../sample/stroopwafel/parameter_space.py | 278 +++++++++++ src/cosmic/sample/stroopwafel/presets.py | 92 ++++ src/cosmic/sample/stroopwafel/priors.py | 171 +++++++ src/cosmic/sample/stroopwafel/rejection.py | 125 +++++ src/cosmic/sample/stroopwafel/result.py | 82 ++++ src/cosmic/sample/stroopwafel/samplers.py | 186 +++++++ src/cosmic/sample/stroopwafel/transforms.py | 87 ++++ 15 files changed, 2391 insertions(+) create mode 100755 src/cosmic/sample/stroopwafel/__init__.py create mode 100755 src/cosmic/sample/stroopwafel/constants.py create mode 100644 src/cosmic/sample/stroopwafel/engine.py create mode 100644 src/cosmic/sample/stroopwafel/examples/example_bhbh.py create mode 100644 src/cosmic/sample/stroopwafel/examples/test_new_modules.py create mode 100644 src/cosmic/sample/stroopwafel/io.py create mode 100644 src/cosmic/sample/stroopwafel/meson.build create mode 100644 src/cosmic/sample/stroopwafel/mixture_model.py create mode 100644 src/cosmic/sample/stroopwafel/parameter_space.py create mode 100644 src/cosmic/sample/stroopwafel/presets.py create mode 100644 src/cosmic/sample/stroopwafel/priors.py create mode 100644 src/cosmic/sample/stroopwafel/rejection.py create mode 100644 src/cosmic/sample/stroopwafel/result.py create mode 100644 src/cosmic/sample/stroopwafel/samplers.py create mode 100644 src/cosmic/sample/stroopwafel/transforms.py diff --git a/src/cosmic/sample/stroopwafel/__init__.py b/src/cosmic/sample/stroopwafel/__init__.py new file mode 100755 index 000000000..ce5731ed6 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/__init__.py @@ -0,0 +1,19 @@ +"""STROOPWAFEL adaptive importance sampling for COSMIC. + +Provides a vectorized reimplementation of the STROOPWAFEL algorithm +for efficiently sampling rare binary stellar evolution outcomes. The +three-phase pipeline (exploration → adaptation → refinement) places +Gaussian components at discovered hits, then importance-samples from +the resulting mixture to concentrate compute budget on interesting +regions of parameter space. + +Example +------- +>>> from cosmic.sample.stroopwafel import AdaptiveSampler, ParameterSpace, Parameter +>>> from cosmic.sample.stroopwafel.presets import merging_dco +""" +from .engine import AdaptiveSampler +from .parameter_space import ParameterSpace, Parameter +from .result import STROOPWAFELResult + +__all__ = ['AdaptiveSampler', 'ParameterSpace', 'Parameter', 'STROOPWAFELResult'] diff --git a/src/cosmic/sample/stroopwafel/constants.py b/src/cosmic/sample/stroopwafel/constants.py new file mode 100755 index 000000000..74ecfe6d8 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/constants.py @@ -0,0 +1,25 @@ +"""Constants used throughout the STROOPWAFEL adaptive sampling module.""" +ALPHA_IMF = -2.3 +SANA_G = -0.55 +SANA_ECC = -0.45 + +R_COEFF =\ + [[1.71535900, 0.62246212, -0.92557761, -1.16996966, -0.30631491],\ + [6.59778800, -0.42450044, -12.13339427, -10.73509484, -2.51487077],\ + [10.08855000, -7.11727086, -31.67119479, -24.24848322, -5.33608972],\ + [1.01249500, 0.32699690, -0.00923418, -0.03876858, -0.00412750],\ + [0.07490166, 0.02410413, 0.07233664, 0.03040467, 0.00197741],\ + [0.01077422, 0.00000000, 0.00000000, 0.00000000, 0.00000000],\ + [3.08223400, 0.94472050, -2.15200882, -2.49219496, -0.63848738],\ + [17.84778000, -7.45345690, -48.96066856, -40.05386135, -9.09331816],\ + [0.00022582, -0.00186899, 0.00388783, 0.00142402, -0.00007671]] + +MINIMUM_SECONDARY_MASS = 0.1 +R_SOL_TO_AU = 0.00465047 +METALLICITY_SOL = 0.0142 +ZSOL = 0.02 +REJECTION_SAMPLES_PER_BATCH = 1e4 +TOTAL_REJECTION_SAMPLES = 1e6 +NUM_GENERATIONS = 1 +MIN_ENTROPY_CHANGE = 0.01 +KAPPA = 1.0 diff --git a/src/cosmic/sample/stroopwafel/engine.py b/src/cosmic/sample/stroopwafel/engine.py new file mode 100644 index 000000000..9c0d57298 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/engine.py @@ -0,0 +1,460 @@ +"""AdaptiveSampler: the main STROOPWAFEL engine. + +Orchestrates the explore -> adapt -> refine -> weight calculation pipeline +using vectorized operations throughout. No Location objects are created. +""" +import os +import numpy as np +from scipy.stats import multivariate_normal + +from cosmic.sample.initialbinarytable import InitialBinaryTable +from cosmic.evolve import Evolve + +from .mixture_model import GaussianMixture +from .result import STROOPWAFELResult +from .constants import KAPPA, NUM_GENERATIONS + + +class AdaptiveSampler: + """Adaptive importance sampler for binary population synthesis with COSMIC. + + Parameters + ---------- + parameter_space : `ParameterSpace` + Defines the sampling dimensions and their distributions. + total_systems : `int` + Total number of systems to simulate across all phases. + batch_size : `int` + Number of systems evolved per COSMIC call. + BSEDict : `dict` + COSMIC binary stellar evolution parameters. + compute_derived : `callable` + Function with signature + ``(samples_physical, param_names) -> dict`` that computes + derived quantities (e.g., ``'mass_2'``, ``'separation'``). + reject_systems : `callable` + Function with signature + ``(samples_physical, derived, param_names) -> bool_mask`` + returning True for physically unacceptable systems. + is_interesting : `callable` + Function with signature ``(bpp) -> (n_hits, hit_bin_nums)`` + identifying systems of interest from COSMIC output. + output_path : `str`, optional + Directory for output files, by default ``'output'`` + nproc : `int`, optional + Number of CPU cores for COSMIC, by default 1 + kappa : `float`, optional + Gaussian width scaling factor, by default ``KAPPA`` + n_generations : `int`, optional + Number of refinement generations, by default ``NUM_GENERATIONS`` + mc_only : `bool`, optional + If True, only run exploration (standard Monte Carlo), by default + False + seed : `int` or None, optional + Random seed for reproducibility, by default None + """ + + def __init__(self, parameter_space, total_systems, batch_size, BSEDict, + compute_derived, reject_systems, is_interesting, + output_path='output', nproc=1, kappa=KAPPA, + n_generations=NUM_GENERATIONS, mc_only=False, seed=None): + self.param_space = parameter_space + self.total_systems = total_systems + self.batch_size = batch_size + self.bse_dict = BSEDict + self.compute_derived_fn = compute_derived + self.reject_fn = reject_systems + self.is_interesting_fn = is_interesting + self.output_path = output_path + self.nproc = nproc + self.kappa = kappa + self.n_generations = n_generations + self.mc_only = mc_only + self.rng = np.random.default_rng(seed) + + # State + self.num_explored = 0 + self.num_hits = 0 + self.fraction_explored = 1.0 + self.finished = 0 + self.prior_fraction_rejected = 0.0 + self.mixture = None + + # Accumulators for all samples (in sampling space) + self._all_samples = [] + self._all_is_hit = [] + self._all_generation = [] + self._all_gaussian_idx = [] + + def run(self): + """Run the full STROOPWAFEL pipeline. + + Returns + ------- + `STROOPWAFELResult` + Container holding all samples, weights, and metadata. + """ + os.makedirs(self.output_path, exist_ok=True) + + self._explore() + + if not self.mc_only and self.num_hits > 0: + self._adapt() + self._refine() + + result = self._compute_weights() + return result + + # ------------------------------------------------------------------ + # Exploration + # ------------------------------------------------------------------ + def _explore(self): + print("Exploration phase started") + + if not self.mc_only: + self.prior_fraction_rejected = self._estimate_prior_rejection_rate() + print(f" Prior rejection rate: {self.prior_fraction_rejected:.4f}") + + while self._should_continue_exploring(): + n_oversample = int(2 * np.ceil(self.batch_size / (1 - self.prior_fraction_rejected))) + + # Sample from prior + samples, mask = self.param_space.sample(n_oversample, rng=self.rng) + + # Transform to physical space for rejection checking + samples_phys = self.param_space.to_physical(samples) + derived = self.compute_derived_fn(samples_phys, self.param_space.names) + + # Physical rejection + phys_rejected = self.reject_fn(samples_phys, derived, self.param_space.names) + + # Combined mask: in bounds AND not physically rejected + valid = mask & ~phys_rejected + + # Select up to batch_size valid samples + valid_indices = np.where(valid)[0] + self.rng.shuffle(valid_indices) + selected = valid_indices[:self.batch_size] + + batch_samples = samples[selected] + batch_samples_phys = samples_phys[selected] + batch_derived = {k: v[selected] for k, v in derived.items()} + + # Evolve with COSMIC + n_hits, hit_bin_nums, bpp, initC, kick_info = self._evolve_batch( + batch_samples_phys, batch_derived + ) + + # Record which are hits + is_hit = np.zeros(len(selected), dtype=bool) + is_hit[hit_bin_nums] = True + + # Accumulate + self._all_samples.append(batch_samples) + self._all_is_hit.append(is_hit) + self._all_generation.append(np.zeros(len(selected), dtype=int)) + self._all_gaussian_idx.append(np.full(len(selected), -1, dtype=int)) + + self.num_hits += n_hits + self.finished += len(selected) + self.num_explored += len(selected) + self._update_fraction_explored() + + self._print_progress() + + self.num_hits_exploratory = self.num_hits + print(f"\nExploration done: {self.num_hits} hits / {self.num_explored} explored " + f"(rate={self.num_hits/max(1, self.num_explored):.6f}, " + f"f_expl={self.fraction_explored:.4f})") + + def _estimate_prior_rejection_rate(self, n_test=100000): + """Estimate fraction of prior samples that are physically rejected. + + Parameters + ---------- + n_test : `int`, optional + Number of samples to use for the estimate, by default 100000 + + Returns + ------- + `float` + Estimated fraction of samples rejected by bounds and physical + criteria combined. + """ + samples, mask = self.param_space.sample(n_test, rng=self.rng) + rejected = n_test - np.sum(mask) + + valid_samples = samples[mask] + if len(valid_samples) > 0: + phys = self.param_space.to_physical(valid_samples) + derived = self.compute_derived_fn(phys, self.param_space.names) + phys_rejected = self.reject_fn(phys, derived, self.param_space.names) + rejected += np.sum(phys_rejected) + + return rejected / n_test + + def _update_fraction_explored(self): + if self.num_hits == 0 or self.num_explored == 0: + return + u = 1.0 / (self.fraction_explored * self.total_systems) + r = self.num_hits / self.num_explored + num = r * (np.sqrt(1.0 - r) - np.sqrt(u)) + den = np.sqrt(1.0 - r) * (np.sqrt(u * (1.0 - r)) + r) + if den != 0: + self.fraction_explored = 1 - num / den + + def _should_continue_exploring(self): + if self.mc_only: + return self.num_explored < self.total_systems + return self.num_explored / self.total_systems < self.fraction_explored + + # ------------------------------------------------------------------ + # Adaptation + # ------------------------------------------------------------------ + def _adapt(self): + print("Adaptation phase started") + + # Gather all exploration hits in sampling space + all_samples = np.vstack(self._all_samples) + all_is_hit = np.concatenate(self._all_is_hit) + hit_samples = all_samples[all_is_hit] + + average_density_one_dim = 1.0 / np.power(self.num_explored, 1.0 / self.param_space.ndim) + + self.mixture = GaussianMixture.from_hits( + hit_samples, self.param_space, average_density_one_dim, kappa=self.kappa + ) + + print(f" Created {self.mixture.n_components} Gaussian components") + print("Adaptation phase finished") + + # ------------------------------------------------------------------ + # Refinement + # ------------------------------------------------------------------ + def _refine(self): + print("Refinement phase started") + entropies = [] + + for gen in range(self.n_generations): + # Estimate rejection rate for current mixture + dist_rejection_rate = self.mixture.compute_rejection_rate( + self.param_space, self.compute_derived_fn, self.reject_fn, + n_per_component=10000, rng=self.rng + ) + + n_per_gen = int((self.total_systems - self.num_explored) / self.n_generations) + gen_samples_list = [] + gen_is_hit_list = [] + gen_finished = 0 + + while gen_finished < n_per_gen and self.finished < self.total_systems: + # Sample from mixture + samples, mask, gauss_idx = self.mixture.sample( + self.batch_size, self.param_space, + consider_rejection=True, rng=self.rng + ) + + # Apply bounds and physical rejection + valid_samples = samples[mask] + valid_gauss_idx = gauss_idx[mask] + + if len(valid_samples) == 0: + continue + + phys = self.param_space.to_physical(valid_samples) + derived = self.compute_derived_fn(phys, self.param_space.names) + phys_rejected = self.reject_fn(phys, derived, self.param_space.names) + + keep = ~phys_rejected + valid_samples = valid_samples[keep] + valid_gauss_idx = valid_gauss_idx[keep] + phys = phys[keep] + derived = {k: v[keep] for k, v in derived.items()} + + # Trim to batch size + self.rng.shuffle(np.arange(len(valid_samples))) + n_take = min(len(valid_samples), self.batch_size) + batch_samples = valid_samples[:n_take] + batch_phys = phys[:n_take] + batch_gauss_idx = valid_gauss_idx[:n_take] + batch_derived = {k: v[:n_take] for k, v in derived.items()} + + # Evolve with COSMIC + n_hits, hit_bin_nums, bpp, initC, kick_info = self._evolve_batch( + batch_phys, batch_derived + ) + + is_hit = np.zeros(n_take, dtype=bool) + is_hit[hit_bin_nums] = True + + # Accumulate + self._all_samples.append(batch_samples) + self._all_is_hit.append(is_hit) + self._all_generation.append(np.full(n_take, gen + 1, dtype=int)) + self._all_gaussian_idx.append(batch_gauss_idx) + + gen_samples_list.append(batch_samples) + gen_is_hit_list.append(is_hit) + + self.num_hits += n_hits + self.finished += n_take + gen_finished += n_take + self._print_progress() + + # EM update (if not the last generation) + if gen < self.n_generations - 1 and len(gen_samples_list) > 0: + gen_samples = np.vstack(gen_samples_list) + gen_is_hit = np.concatenate(gen_is_hit_list) + gen_priors = self.param_space.compute_prior(gen_samples) + + # Save current state in case we need to revert + saved_mixture = GaussianMixture( + self.mixture.means.copy(), + self.mixture.covariances.copy(), + self.mixture.alphas.copy(), + self.mixture.rejection_rate + ) + + should_revert = self.mixture.update_em( + gen_samples, gen_is_hit, gen_priors, + self.prior_fraction_rejected, + tolerance=1e-10, entropies=entropies + ) + + if should_revert: + self.mixture = saved_mixture + print(" EM update reverted (insufficient entropy change)") + + n_refined = self.total_systems - self.num_explored + if n_refined > 0: + refine_hits = self.num_hits - self.num_hits_exploratory + print(f"\nRefinement done: {refine_hits} hits / {n_refined} refined " + f"(rate={refine_hits/max(1, n_refined):.6f})") + + # ------------------------------------------------------------------ + # Weight calculation + # ------------------------------------------------------------------ + def _compute_weights(self): + """Compute importance sampling weights for all samples. + + Builds the final ``STROOPWAFELResult`` using the formula + ``w(x) = π(x) / Q(x)`` where + ``Q = f_e·π + (1 − f_e)·q`` is a mixture of the prior and the + Gaussian proposal. + + Returns + ------- + `STROOPWAFELResult` + Container holding all samples, importance weights, and + associated metadata. + """ + all_samples = np.vstack(self._all_samples) + all_is_hit = np.concatenate(self._all_is_hit) + all_generation = np.concatenate(self._all_generation) + all_gaussian_idx = np.concatenate(self._all_gaussian_idx) + N = len(all_samples) + + # Prior probabilities + pi_norm = 1.0 / (1 - self.prior_fraction_rejected) + pi = self.param_space.compute_prior(all_samples) * pi_norm + + # Start with the exploration-phase contribution to denominator + fraction_explored = self.num_explored / float(N) + den = fraction_explored * pi + + # Add refinement-phase contributions from each generation's mixture + if self.mixture is not None and not self.mc_only: + q_norm = 1.0 / (1 - self.mixture.rejection_rate) + # Evaluate the mixture PDF incrementally (memory-efficient) + for k in range(self.mixture.n_components): + xPDF_k = multivariate_normal.pdf( + all_samples, self.mixture.means[k], self.mixture.covariances[k], + allow_singular=True + ) + den += (xPDF_k * self.mixture.alphas[k] + * (1 - fraction_explored) * q_norm) / self.n_generations + + weights = pi / den + + # Build result + result = STROOPWAFELResult( + samples=self.param_space.to_physical(all_samples), + param_names=self.param_space.names, + weights=weights, + is_hit=all_is_hit, + generation=all_generation, + gaussian_idx=all_gaussian_idx, + num_explored=self.num_explored, + num_hits=self.num_hits, + fraction_explored=self.fraction_explored, + ) + + print(f"\nTotal hits: {self.num_hits}") + print(f"Sum of weights: {np.sum(weights):.4f}") + print(f"Weighted hit rate: {result.hit_rate:.8f} +/- {result.hit_rate_uncertainty:.8f}") + + return result + + # ------------------------------------------------------------------ + # COSMIC interface + # ------------------------------------------------------------------ + def _evolve_batch(self, samples_physical, derived): + """Evolve a batch of binaries with COSMIC and identify hits. + + Parameters + ---------- + samples_physical : `numpy.ndarray` + (N, D) array of binary parameters in physical space. + derived : `dict` + Dictionary with keys ``'mass_2'``, ``'metallicity_1'``, etc., + each mapping to an (N,) array. + + Returns + ------- + n_hits : `int` + Number of systems classified as hits. + hit_bin_nums : `numpy.ndarray` + Integer array of 0-indexed positions within the batch that + are hits. + bpp : `pandas.DataFrame` + COSMIC binary population parameters output. + initC : `pandas.DataFrame` + COSMIC initial conditions output. + kick_info : `pandas.DataFrame` + COSMIC natal kick information output. + """ + n = len(samples_physical) + idx = {name: i for i, name in enumerate(self.param_space.names)} + + batch_initial = InitialBinaryTable.InitialBinaries( + m1=samples_physical[:, idx['mass_1']], + m2=derived['mass_2'], + porb=samples_physical[:, idx['porb']], + ecc=samples_physical[:, idx['ecc']], + tphysf=np.full(n, 13700.0), + kstar1=np.full(n, 1), + kstar2=np.full(n, 1), + metallicity=derived['metallicity_1'], + ) + + bpp, bcm, initC, kick_info = Evolve.evolve( + initialbinarytable=batch_initial, + BSEDict=self.bse_dict, + nproc=self.nproc, + ) + + # Apply user's hit identification + n_hits, hit_bin_nums = self.is_interesting_fn(bpp) + + return n_hits, hit_bin_nums, bpp, initC, kick_info + + # ------------------------------------------------------------------ + # Utilities + # ------------------------------------------------------------------ + def _print_progress(self): + pct = 100 * self.finished / self.total_systems + bar_len = 20 + filled = int(bar_len * self.finished // self.total_systems) + bar = '|' * filled + '-' * (bar_len - filled) + print(f'\r progress |{bar}| {pct:.1f}% complete ' + f'({self.finished}/{self.total_systems})', end='\r') diff --git a/src/cosmic/sample/stroopwafel/examples/example_bhbh.py b/src/cosmic/sample/stroopwafel/examples/example_bhbh.py new file mode 100644 index 000000000..3278e82fe --- /dev/null +++ b/src/cosmic/sample/stroopwafel/examples/example_bhbh.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python +"""Example: Find merging BH-BH binaries using the new vectorized STROOPWAFEL API. + +Equivalent to the old tests/BHBH_testing.py but using the new API. +""" +import os +import sys +import time +import argparse +import numpy as np + +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) # adds COSMIC-stroopwafel/ + +from stroopwafel import AdaptiveSampler, ParameterSpace, Parameter +from stroopwafel.presets import merging_dco +from stroopwafel.rejection import default_reject +from stroopwafel import io as swio + +# ------------------------------------------------------------------ +# Parse arguments +# ------------------------------------------------------------------ +parser = argparse.ArgumentParser() +parser.add_argument('--num_systems', type=int, default=10000) +parser.add_argument('--num_cores', type=int, default=1) +parser.add_argument('--num_per_core', type=int, default=100) +parser.add_argument('--mc_only', type=bool, default=False) +parser.add_argument('--output_dir', default='output/BHBH_new') +parser.add_argument('--model', default='fiducial') +parser.add_argument('--seed', type=int, default=None) +args = parser.parse_args() + +# ------------------------------------------------------------------ +# Physics models (same as old code) +# ------------------------------------------------------------------ +fiducial = { + 'xi': 1.0, 'bhflag': 1, 'neta': 0.5, 'windflag': 3, 'wdflag': 1, + 'alpha1': 1.0, 'pts1': 0.001, 'pts3': 0.02, 'pts2': 0.01, + 'epsnov': 0.001, 'hewind': 0.5, 'ck': 1000, 'bwind': 0.0, + 'lambdaf': 0.0, 'mxns': 3.0, 'beta': -1.0, 'tflag': 1, 'acc2': 1.5, + 'grflag': 1, 'remnantflag': 4, 'ceflag': 0, 'eddfac': 1.0, + 'ifflag': 0, 'bconst': 3000, 'sigma': 265.0, 'gamma': -2.0, + 'pisn': 45.0, + 'natal_kick_array': [[-100.0, -100.0, -100.0, -100.0, 0.0], + [-100.0, -100.0, -100.0, -100.0, 0.0]], + 'bhsigmafrac': 1.0, 'polar_kick_angle': 90, + 'qcrit_array': [0.0] * 16, + 'cekickflag': 2, 'cehestarflag': 0, 'cemergeflag': 0, + 'ecsn': 2.5, 'ecsn_mlow': 1.8, 'aic': 1, 'ussn': 0, + 'sigmadiv': -20.0, 'qcflag': 5, 'eddlimflag': 0, + 'fprimc_array': [2.0 / 21.0] * 16, + 'bhspinflag': 0, 'bhspinmag': 0.0, 'rejuv_fac': 1.0, + 'rejuvflag': 0, 'htpmb': 1, 'ST_cr': 1, 'ST_tide': 1, + 'bdecayfac': 1, 'rembar_massloss': 0.5, 'kickflag': 5, + 'zsun': 0.014, 'bhms_coll_flag': 0, 'don_lim': -1, + 'acc_lim': -1, 'rtmsflag': 0, 'wd_mass_lim': 1, +} + +BSEDict = fiducial # extend with model variants as needed + +# ------------------------------------------------------------------ +# Define parameter space +# ------------------------------------------------------------------ +params = ParameterSpace([ + Parameter('mass_1', 5.0, 150.0, sampler='kroupa', prior='kroupa'), + Parameter('q', 0.0, 1.0, sampler='uniform', prior='uniform'), + Parameter('porb', 0.15, 5.5, sampler='sana', prior='sana'), + Parameter('ecc', 1e-9, 0.99999999, sampler='sana_ecc', prior='sana_ecc'), + Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log', prior='flat_in_log'), +]) + +# ------------------------------------------------------------------ +# Define derived quantities (vectorized) +# ------------------------------------------------------------------ +def compute_derived(samples_physical, param_names): + """Compute derived quantities from sampled parameters.""" + idx = {name: i for i, name in enumerate(param_names)} + m1 = samples_physical[:, idx['mass_1']] + q = samples_physical[:, idx['q']] + porb = samples_physical[:, idx['porb']] + z = samples_physical[:, idx['metallicity']] + + mass_2 = m1 * q + separation = ((porb ** 2) * (m1 + mass_2)) ** (1.0 / 3.0) + + return { + 'mass_2': mass_2, + 'metallicity_1': z, + 'metallicity_2': z, + 'separation': separation, + } + +# ------------------------------------------------------------------ +# Hit definition: merging BH-BH binaries within Hubble time +# ------------------------------------------------------------------ +is_interesting = merging_dco(kstar_1=[14], kstar_2=[14], max_merge_time=13.7) + +# ------------------------------------------------------------------ +# Run +# ------------------------------------------------------------------ +if __name__ == '__main__': + start = time.time() + + sw = AdaptiveSampler( + parameter_space=params, + total_systems=args.num_systems, + batch_size=args.num_per_core, + BSEDict=BSEDict, + compute_derived=compute_derived, + reject_systems=default_reject, + is_interesting=is_interesting, + output_path=args.output_dir, + nproc=args.num_cores, + mc_only=args.mc_only, + seed=args.seed, + ) + + result = sw.run() + + # Save results + output_file = os.path.join(args.output_dir, 'stroopwafel_result.h5') + swio.save_result(output_file, result) + + elapsed = time.time() - start + print(f"\nTotal time: {elapsed:.1f}s") + print(f"Results saved to {output_file}") diff --git a/src/cosmic/sample/stroopwafel/examples/test_new_modules.py b/src/cosmic/sample/stroopwafel/examples/test_new_modules.py new file mode 100644 index 000000000..e797b792a --- /dev/null +++ b/src/cosmic/sample/stroopwafel/examples/test_new_modules.py @@ -0,0 +1,299 @@ +"""Tests for the new vectorized STROOPWAFEL modules. + +Run with: python -m pytest tests/test_new_modules.py -v +Or just: python tests/test_new_modules.py +""" +import sys +import os +import math +import numpy as np + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) # adds COSMIC-stroopwafel/ + +from stroopwafel import ParameterSpace, Parameter +from stroopwafel.samplers import SAMPLERS +from stroopwafel.priors import PRIORS +from stroopwafel.transforms import to_sampling_space, to_physical_space, transform_bounds +from stroopwafel.mixture_model import GaussianMixture +from stroopwafel.rejection import get_zams_radius, calculate_roche_lobe_radius, default_reject +from stroopwafel.result import STROOPWAFELResult +from stroopwafel.constants import ( + R_COEFF, ZSOL, R_SOL_TO_AU, ALPHA_IMF, SANA_G, SANA_ECC +) + + +def make_default_params(): + return ParameterSpace([ + Parameter('mass_1', 5.0, 150.0, sampler='kroupa', prior='kroupa'), + Parameter('q', 0.0, 1.0, sampler='uniform', prior='uniform'), + Parameter('porb', 0.15, 5.5, sampler='sana', prior='sana'), + Parameter('ecc', 1e-9, 0.99999999, sampler='sana_ecc', prior='sana_ecc'), + Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log', prior='flat_in_log'), + ]) + + +# ==================================================================== +# ParameterSpace tests +# ==================================================================== + +def test_parameter_space_names_sorted(): + params = make_default_params() + assert params.names == sorted(params.names) + assert params.ndim == 5 + + +def test_sampling_produces_valid_arrays(): + params = make_default_params() + rng = np.random.default_rng(42) + samples, mask = params.sample(1000, rng=rng) + assert samples.shape == (1000, 5) + assert mask.shape == (1000,) + assert mask.dtype == bool + + +def test_roundtrip_transform(): + params = make_default_params() + rng = np.random.default_rng(42) + samples, _ = params.sample(1000, rng=rng) + physical = params.to_physical(samples) + back = params.to_sampling(physical) + assert np.allclose(samples, back, atol=1e-12) + + +def test_prior_positive(): + params = make_default_params() + rng = np.random.default_rng(42) + samples, mask = params.sample(1000, rng=rng) + priors = params.compute_prior(samples[mask]) + assert np.all(priors > 0) + assert np.all(np.isfinite(priors)) + + +def test_in_bounds(): + params = make_default_params() + rng = np.random.default_rng(42) + samples, mask = params.sample(1000, rng=rng) + mask2 = params.in_bounds(samples) + np.testing.assert_array_equal(mask, mask2) + + +# ==================================================================== +# ZAMS radius / Roche lobe tests (vs old scalar implementation) +# ==================================================================== + +def old_get_zams_radius(mass, metallicity): + """Old scalar implementation for comparison.""" + metallicity_xi = math.log10(metallicity / ZSOL) + rc = [] + for coeff in R_COEFF: + value = 1; total = 0 + for series in coeff: + total += series * value + value *= metallicity_xi + rc.append(total) + top = (rc[0] * pow(mass, 2.5) + rc[1] * pow(mass, 6.5) + + rc[2] * pow(mass, 11) + rc[3] * pow(mass, 19) + + rc[4] * pow(mass, 19.5)) + bottom = (rc[5] + rc[6] * pow(mass, 2) + rc[7] * pow(mass, 8.5) + + pow(mass, 18.5) + rc[8] * pow(mass, 19.5)) + return (top / bottom) * R_SOL_TO_AU + + +def old_roche(m1, m2): + q = m1 / m2 + return 0.49 / (0.6 + pow(q, -2.0 / 3.0) * math.log(1.0 + pow(q, 1.0 / 3.0))) + + +def test_zams_radius_matches_old(): + masses = np.linspace(1, 100, 50) + mets = np.full(50, 0.014) + new = get_zams_radius(masses, mets) + old = np.array([old_get_zams_radius(m, z) for m, z in zip(masses, mets)]) + np.testing.assert_allclose(new, old, atol=1e-12) + + +def test_zams_radius_multiple_metallicities(): + masses = np.array([1.0, 10.0, 50.0]) + mets = np.array([0.001, 0.014, 0.03]) + new = get_zams_radius(masses, mets) + old = np.array([old_get_zams_radius(m, z) for m, z in zip(masses, mets)]) + np.testing.assert_allclose(new, old, atol=1e-12) + + +def test_roche_lobe_matches_old(): + m1 = np.array([5.0, 10.0, 20.0, 50.0, 100.0]) + m2 = np.array([3.0, 8.0, 10.0, 25.0, 50.0]) + new = calculate_roche_lobe_radius(m1, m2) + old = np.array([old_roche(a, b) for a, b in zip(m1, m2)]) + np.testing.assert_allclose(new, old, atol=1e-12) + + +# ==================================================================== +# Default rejection function tests +# ==================================================================== + +def test_default_reject_basic(): + param_names = ['ecc', 'mass_1', 'metallicity', 'porb', 'q'] + # Create a sample that should NOT be rejected (wide binary) + samples = np.array([[0.1, 20.0, 0.014, 100.0, 0.5]]) # reasonable binary + derived = { + 'mass_2': np.array([10.0]), + 'metallicity_1': np.array([0.014]), + 'metallicity_2': np.array([0.014]), + 'separation': np.array([50.0]), # AU + } + rejected = default_reject(samples, derived, param_names) + assert not rejected[0], "Wide binary should not be rejected" + + +def test_default_reject_low_mass(): + param_names = ['ecc', 'mass_1', 'metallicity', 'porb', 'q'] + samples = np.array([[0.1, 20.0, 0.014, 100.0, 0.001]]) + derived = { + 'mass_2': np.array([0.02]), # below minimum secondary mass + 'metallicity_1': np.array([0.014]), + 'metallicity_2': np.array([0.014]), + 'separation': np.array([50.0]), + } + rejected = default_reject(samples, derived, param_names) + assert rejected[0], "Low mass secondary should be rejected" + + +# ==================================================================== +# GaussianMixture tests +# ==================================================================== + +def test_mixture_from_hits(): + params = make_default_params() + rng = np.random.default_rng(42) + samples, _ = params.sample(1000, rng=rng) + # Use moderate samples as "hits" + hits = samples[400:410] # 10 hits + + avg_density = 1.0 / np.power(1000, 1.0 / params.ndim) + mixture = GaussianMixture.from_hits(hits, params, avg_density) + + assert mixture.n_components == 10 + assert mixture.means.shape == (10, 5) + assert mixture.covariances.shape == (10, 5, 5) + assert np.allclose(np.sum(mixture.alphas), 1.0) + assert np.allclose(mixture.alphas, 0.1) + + +def test_mixture_pdf_nonnegative(): + params = make_default_params() + rng = np.random.default_rng(42) + samples, _ = params.sample(1000, rng=rng) + hits = samples[400:410] + + avg_density = 1.0 / np.power(1000, 1.0 / params.ndim) + mixture = GaussianMixture.from_hits(hits, params, avg_density) + + pdf_vals = mixture.pdf(samples[:100]) + assert np.all(pdf_vals >= 0) + assert np.all(np.isfinite(pdf_vals)) + + +def test_mixture_sample(): + params = make_default_params() + rng = np.random.default_rng(42) + samples, _ = params.sample(1000, rng=rng) + hits = samples[400:410] + + avg_density = 1.0 / np.power(1000, 1.0 / params.ndim) + mixture = GaussianMixture.from_hits(hits, params, avg_density) + + msamp, mmask, midx = mixture.sample(2000, params, rng=rng) + assert msamp.shape[1] == 5 + assert len(mmask) == len(msamp) + assert len(midx) == len(msamp) + assert np.all(midx >= 0) + assert np.all(midx < 10) + + +# ==================================================================== +# STROOPWAFELResult tests +# ==================================================================== + +def test_result_hit_rate(): + result = STROOPWAFELResult() + result.weights = np.ones(100) + result.is_hit = np.zeros(100, dtype=bool) + result.is_hit[:10] = True + assert abs(result.hit_rate - 0.1) < 1e-10 + + +# ==================================================================== +# Prior comparison with old code +# ==================================================================== + +def test_kroupa_prior_matches(): + """Test that vectorized kroupa prior matches old scalar version.""" + from stroopwafel.priors import kroupa + lo, hi = 5.0, 150.0 + values = np.linspace(5.1, 149.9, 100) + vec_result = kroupa(values, lo, hi) + + # Old scalar + norm = (ALPHA_IMF + 1) / (hi**(ALPHA_IMF + 1) - lo**(ALPHA_IMF + 1)) + old_result = norm * values**ALPHA_IMF + np.testing.assert_allclose(vec_result, old_result, atol=1e-12) + + +def test_sana_prior_matches(): + from stroopwafel.priors import sana + lo, hi = 0.15, 5.5 + values = np.linspace(0.2, 5.4, 100) + vec_result = sana(values, lo, hi) + + norm = (SANA_G + 1) / (hi**(SANA_G + 1) - lo**(SANA_G + 1)) + old_result = norm * values**SANA_G + np.testing.assert_allclose(vec_result, old_result, atol=1e-12) + + +# ==================================================================== +# Performance benchmark +# ==================================================================== + +def bench_zams_radius(): + """Benchmark vectorized vs scalar ZAMS radius.""" + import time + N = 10000 + masses = np.random.uniform(1, 100, N) + mets = np.random.uniform(0.0001, 0.03, N) + + # Vectorized + start = time.time() + _ = get_zams_radius(masses, mets) + vec_time = time.time() - start + + # Scalar + start = time.time() + _ = [old_get_zams_radius(m, z) for m, z in zip(masses, mets)] + scalar_time = time.time() - start + + speedup = scalar_time / vec_time if vec_time > 0 else float('inf') + print(f"\nZAMS radius benchmark (N={N}):") + print(f" Vectorized: {vec_time*1000:.1f}ms") + print(f" Scalar: {scalar_time*1000:.1f}ms") + print(f" Speedup: {speedup:.0f}x") + + +if __name__ == '__main__': + # Run all tests + test_funcs = [v for k, v in sorted(globals().items()) if k.startswith('test_')] + passed = 0 + failed = 0 + for func in test_funcs: + try: + func() + print(f" [PASS] {func.__name__}") + passed += 1 + except Exception as e: + print(f" [FAIL] {func.__name__}: {e}") + failed += 1 + + print(f"\n{passed} passed, {failed} failed") + + if failed == 0: + bench_zams_radius() diff --git a/src/cosmic/sample/stroopwafel/io.py b/src/cosmic/sample/stroopwafel/io.py new file mode 100644 index 000000000..ebe79c9df --- /dev/null +++ b/src/cosmic/sample/stroopwafel/io.py @@ -0,0 +1,107 @@ +"""HDF5-based I/O for STROOPWAFEL state and results. + +Replaces the old CSV-based print_samples/read_samples with array-native storage. +""" +import pandas as pd +import h5py + + +def save_result(path, result): + """Save a `STROOPWAFELResult` to an HDF5 file. + + Parameters + ---------- + path : `str` + Output file path (created or overwritten). + result : `STROOPWAFELResult` + Result container to serialise. + """ + with h5py.File(path, 'w') as f: + f.create_dataset('samples', data=result.samples) + f.create_dataset('weights', data=result.weights) + f.create_dataset('is_hit', data=result.is_hit) + f.create_dataset('generation', data=result.generation) + f.create_dataset('gaussian_idx', data=result.gaussian_idx) + f.attrs['param_names'] = result.param_names + f.attrs['num_explored'] = result.num_explored + f.attrs['num_hits'] = result.num_hits + f.attrs['fraction_explored'] = result.fraction_explored + + +def save_cosmic_output(path, bpp_frames, initC_frames, kick_info_frames): + """Append COSMIC output DataFrames to an HDF5 file. + + Parameters + ---------- + path : `str` + Output file path (appended to if it already exists). + bpp_frames : `list` of `pandas.DataFrame` + Binary population parameter tables from each batch. + initC_frames : `list` of `pandas.DataFrame` + Initial conditions tables from each batch. + kick_info_frames : `list` of `pandas.DataFrame` + Natal kick information tables from each batch. + """ + if bpp_frames: + full_bpp = pd.concat(bpp_frames, ignore_index=True) + full_bpp.to_hdf(path, key='bpp', mode='a', format='table') + if initC_frames: + full_initC = pd.concat(initC_frames, ignore_index=True) + full_initC.to_hdf(path, key='initC', mode='a', format='table') + if kick_info_frames: + full_kicks = pd.concat(kick_info_frames, ignore_index=True) + full_kicks.to_hdf(path, key='kick_info', mode='a', format='table') + + +def save_mixture(path, mixture, generation): + """Save `GaussianMixture` state to an HDF5 file. + + Parameters + ---------- + path : `str` + Output file path (appended to if it already exists). + mixture : `GaussianMixture` + Mixture model to serialise. + generation : `int` + Refinement generation number, used as the HDF5 group key. + """ + with h5py.File(path, 'a') as f: + grp_name = f'mixture/gen_{generation}' + if grp_name in f: + del f[grp_name] + grp = f.create_group(grp_name) + grp.create_dataset('means', data=mixture.means) + grp.create_dataset('covariances', data=mixture.covariances) + grp.create_dataset('alphas', data=mixture.alphas) + grp.attrs['rejection_rate'] = mixture.rejection_rate + + +def load_mixture(path, generation): + """Load `GaussianMixture` state from an HDF5 file. + + Parameters + ---------- + path : `str` + File path to read from. + generation : `int` + Refinement generation number to load. + + Returns + ------- + `GaussianMixture` or None + The stored mixture model, or None if the requested generation + is not found or the file does not exist. + """ + from .mixture_model import GaussianMixture + + try: + with h5py.File(path, 'r') as f: + grp = f[f'mixture/gen_{generation}'] + return GaussianMixture( + means=grp['means'][:], + covariances=grp['covariances'][:], + alphas=grp['alphas'][:], + rejection_rate=grp.attrs['rejection_rate'], + ) + except (KeyError, FileNotFoundError): + return None diff --git a/src/cosmic/sample/stroopwafel/meson.build b/src/cosmic/sample/stroopwafel/meson.build new file mode 100644 index 000000000..87bd56fdf --- /dev/null +++ b/src/cosmic/sample/stroopwafel/meson.build @@ -0,0 +1,19 @@ +python_sources = [ + '__init__.py', + 'constants.py', + 'engine.py', + 'io.py', + 'mixture_model.py', + 'parameter_space.py', + 'presets.py', + 'priors.py', + 'rejection.py', + 'result.py', + 'samplers.py', + 'transforms.py', +] + +py3.install_sources( + python_sources, + subdir: 'cosmic/sample/stroopwafel' +) \ No newline at end of file diff --git a/src/cosmic/sample/stroopwafel/mixture_model.py b/src/cosmic/sample/stroopwafel/mixture_model.py new file mode 100644 index 000000000..1f3618a14 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/mixture_model.py @@ -0,0 +1,316 @@ +"""Vectorized Gaussian Mixture Model for adaptive importance sampling. + +Stores all mixture components as numpy arrays and provides vectorized +sampling, PDF evaluation, and EM updates. +""" +import numpy as np +from scipy.stats import multivariate_normal, entropy as scipy_entropy +from .constants import KAPPA, MIN_ENTROPY_CHANGE + + +class GaussianMixture: + """A mixture of K multivariate Gaussians. + + Parameters + ---------- + means : `numpy.ndarray` + (K, D) array of component means. + covariances : `numpy.ndarray` + (K, D, D) array of covariance matrices. + alphas : `numpy.ndarray` + (K,) array of mixture weights (must sum to 1). + rejection_rate : `float`, optional + Fraction of samples that fall outside bounds, by default 0.0 + """ + + def __init__(self, means, covariances, alphas, rejection_rate=0.0): + self.means = np.asarray(means) + self.covariances = np.asarray(covariances) + self.alphas = np.asarray(alphas, dtype=float) + self.rejection_rate = rejection_rate + + @property + def n_components(self): + return len(self.means) + + @property + def ndim(self): + return self.means.shape[1] + + @classmethod + def from_hits(cls, hit_samples, param_space, average_density_one_dim, kappa=KAPPA): + """Create a Gaussian mixture by placing one component at each hit. + + Parameters + ---------- + hit_samples : `numpy.ndarray` + (K, D) array of hit locations in sampling space. + param_space : `ParameterSpace` + Parameter space instance providing bounds and sigma computation. + average_density_one_dim : `float` + Characteristic inter-sample spacing, + ``1 / num_explored ** (1 / D)``. + kappa : `float`, optional + Width scaling factor for the Gaussian covariances, by default + ``KAPPA`` + + Returns + ------- + `GaussianMixture` + A new mixture with one component centred on each hit. + """ + K = len(hit_samples) + D = param_space.ndim + + # Compute sigma for each hit and dimension + sigmas = param_space.compute_sigma(hit_samples, average_density_one_dim) + + # Build diagonal covariance matrices, scaled by kappa + covariances = np.zeros((K, D, D)) + for k in range(K): + covariances[k] = np.diag(sigmas[k] ** 2) * kappa ** 2 + + # Equal mixture weights + alphas = np.full(K, 1.0 / K) + + return cls(hit_samples.copy(), covariances, alphas) + + def sample(self, n_total, param_space, consider_rejection=False, rng=None): + """Sample from the mixture distribution. + + Parameters + ---------- + n_total : `int` + Total number of samples desired. + param_space : `ParameterSpace` + Parameter space used for bounds checking. + consider_rejection : `bool`, optional + If True, oversample to account for the current rejection + rate, by default False + rng : `numpy.random.Generator`, optional + Random number generator, by default None + + Returns + ------- + samples : `numpy.ndarray` + (M, D) array of samples in sampling space. + mask : `numpy.ndarray` + (M,) boolean array indicating which samples are in bounds. + gaussian_indices : `numpy.ndarray` + (M,) integer array indicating which component generated each + sample. + """ + rng = rng or np.random.default_rng() + all_samples = [] + all_indices = [] + + for k in range(self.n_components): + n_k = int(np.ceil(n_total * self.alphas[k])) + if consider_rejection and self.rejection_rate < 1.0: + n_k = int(2 * np.ceil(n_k / (1 - self.rejection_rate))) + if n_k <= 0: + continue + s = rng.multivariate_normal(self.means[k], self.covariances[k], size=n_k) + all_samples.append(s) + all_indices.append(np.full(n_k, k, dtype=int)) + + if not all_samples: + D = param_space.ndim + return np.empty((0, D)), np.empty(0, dtype=bool), np.empty(0, dtype=int) + + samples = np.vstack(all_samples) + gaussian_indices = np.concatenate(all_indices) + mask = param_space.in_bounds(samples) + + return samples, mask, gaussian_indices + + def pdf(self, samples): + """Evaluate the mixture PDF at the given samples. + + Parameters + ---------- + samples : `numpy.ndarray` + (N, D) array in sampling space. + + Returns + ------- + `numpy.ndarray` + (N,) array of PDF values (weighted sum of components). + """ + N = len(samples) + result = np.zeros(N) + for k in range(self.n_components): + result += self.alphas[k] * multivariate_normal.pdf( + samples, self.means[k], self.covariances[k], allow_singular=True + ) + return result + + def component_pdfs(self, samples): + """Evaluate each component's PDF at the given samples. + + Parameters + ---------- + samples : `numpy.ndarray` + (N, D) array in sampling space. + + Returns + ------- + `numpy.ndarray` + (K, N) array where element [k, n] is the PDF of component k + evaluated at sample n. + """ + K = self.n_components + N = len(samples) + xPDF = np.empty((K, N)) + for k in range(K): + xPDF[k, :] = multivariate_normal.pdf( + samples, self.means[k], self.covariances[k], allow_singular=True + ) + return xPDF + + def compute_rejection_rate(self, param_space, compute_derived_fn, reject_fn, + n_per_component=10000, rng=None): + """Estimate the rejection rate of the mixture. + + Samples from each component, transforms to physical space, applies + rejection criteria, and computes the weighted rejection rate. + + Parameters + ---------- + param_space : `ParameterSpace` + Parameter space for bounds checking and coordinate transforms. + compute_derived_fn : `callable` + Function with signature + ``(samples_physical, param_names) -> dict`` that computes + derived quantities. + reject_fn : `callable` + Function with signature + ``(samples_physical, derived, param_names) -> bool_mask`` + returning True for rejected systems. + n_per_component : `int`, optional + Number of samples per component for the estimate, by default + 10000 + rng : `numpy.random.Generator`, optional + Random number generator, by default None + + Returns + ------- + `float` + Estimated rejection rate (also stored as + ``self.rejection_rate``). + """ + rng = rng or np.random.default_rng() + fractional_rejected = 0.0 + + for k in range(self.n_components): + n = n_per_component + s = rng.multivariate_normal(self.means[k], self.covariances[k], size=n) + + # Bounds rejection + bounds_mask = param_space.in_bounds(s) + rejected = n - np.sum(bounds_mask) + + # Physical rejection on in-bounds samples + s_valid = s[bounds_mask] + if len(s_valid) > 0: + s_physical = param_space.to_physical(s_valid) + derived = compute_derived_fn(s_physical, param_space.names) + phys_rejected = reject_fn(s_physical, derived, param_space.names) + rejected += np.sum(phys_rejected) + + fractional_rejected += rejected * self.alphas[k] / n + + self.rejection_rate = fractional_rejected + return self.rejection_rate + + def update_em(self, samples, is_hit, prior_probs, prior_fraction_rejected, + tolerance=1e-10, entropies=None): + """Perform one EM-like update of the mixture parameters. + + Parameters + ---------- + samples : `numpy.ndarray` + (N, D) array of samples in sampling space. + is_hit : `numpy.ndarray` + (N,) boolean or integer array indicating hits. + prior_probs : `numpy.ndarray` + (N,) array of prior probabilities for each sample. + prior_fraction_rejected : `float` + Fraction of prior samples that are physically rejected. + tolerance : `float`, optional + Minimum mixture weight to keep a component, by default 1e-10 + entropies : `list`, optional + List of previous entropy values (mutated in place for + convergence tracking), by default None + + Returns + ------- + `bool` + True if the entropy check triggers reversion to the previous + mixture state. + """ + pi_norm = 1.0 / (1 - prior_fraction_rejected) + q_norm = 1.0 / (1 - self.rejection_rate) + pi = prior_probs * pi_norm + is_hit = np.asarray(is_hit, dtype=float) + + N = len(samples) + K = self.n_components + + # Compute component PDFs: (K, N) + xPDF = self.component_pdfs(samples).T # -> (N, K) + qPDF = xPDF * self.alphas * q_norm # (N, K) + + # Responsibilities + qPDF_sum = np.sum(qPDF, axis=1) # (N,) + rho = qPDF / qPDF_sum[:, None] # (N, K) + + # Importance weights + gaussian_weights = (pi * is_hit) / qPDF_sum # (N,) + weight_sum = np.sum(gaussian_weights) + if weight_sum == 0: + return False + weights_normalized = (gaussian_weights / weight_sum)[:, None] # (N, 1) + + # Update alphas + new_alphas = np.sum(weights_normalized * rho, axis=0) # (K,) + + # Remove insignificant components + keep = new_alphas > tolerance + if not np.any(keep): + return False + + new_alphas = new_alphas[keep] + rho = rho[:, keep] + K_new = int(np.sum(keep)) + + # Update means + new_means = np.empty((K_new, self.ndim)) + for d in range(self.ndim): + new_means[:, d] = np.sum( + weights_normalized * samples[:, d:d+1] * rho, axis=0 + ) + new_means = new_means / new_alphas[:, None] + + # Update covariances + old_covs = self.covariances[keep] + new_covs = np.empty_like(old_covs) + for k in range(K_new): + distance = (new_means[k] - samples)[:, :, None] # (N, D, 1) + matrix = np.einsum('nij,nji->nij', distance, distance) # (N, D, D) + factor = weights_normalized[:, 0] * rho[:, k] # (N,) + new_covs[k] = np.sum(factor[:, None, None] * matrix, axis=0) / new_alphas[k] + + # Entropy check + entropy_change = np.exp(scipy_entropy(weights_normalized[:, 0])) / N + if entropies is not None: + if len(entropies) >= 1 and entropy_change - entropies[-1] < MIN_ENTROPY_CHANGE: + return True # Signal to revert + entropies.append(entropy_change) + + # Apply updates + self.means = new_means + self.covariances = new_covs + self.alphas = new_alphas + + return False diff --git a/src/cosmic/sample/stroopwafel/parameter_space.py b/src/cosmic/sample/stroopwafel/parameter_space.py new file mode 100644 index 000000000..ede8b1e5b --- /dev/null +++ b/src/cosmic/sample/stroopwafel/parameter_space.py @@ -0,0 +1,278 @@ +"""Vectorized parameter space definition. + +A `ParameterSpace` holds an ordered list of `Parameter` instances and provides +vectorized operations on (N, D) sample arrays including sampling, scale +transforms, prior evaluation, and bounds checking. +""" +import numpy as np +from dataclasses import dataclass + +from .samplers import SAMPLERS +from .priors import PRIORS +from .transforms import to_sampling_space, to_physical_space, transform_bounds +from .constants import ALPHA_IMF, SANA_G, SANA_ECC + + +@dataclass +class Parameter: + """A single dimension in the parameter space. + + Parameters + ---------- + name : `str` + Name of the parameter (used for column ordering). + min_value : `float` + Lower bound in physical space. + max_value : `float` + Upper bound in physical space. + sampler : `str`, optional + Name of the sampling distribution, by default ``'uniform'`` + prior : `str`, optional + Name of the prior distribution, by default ``'uniform'`` + """ + name: str + min_value: float + max_value: float + sampler: str = 'uniform' + prior: str = 'uniform' + + def __post_init__(self): + self.lo, self.hi = transform_bounds(self.min_value, self.max_value, self.sampler) + + +class ParameterSpace: + """An ordered collection of Parameters with vectorized operations. + + All methods operate on (N, D) numpy arrays where columns are ordered + alphabetically by parameter name. + + Parameters + ---------- + params : `list` of `Parameter` + List of parameter definitions. They will be sorted by name + internally. + """ + + def __init__(self, params): + # Sort by name for deterministic column ordering (same as old code) + self.params = sorted(params, key=lambda p: p.name) + self.names = [p.name for p in self.params] + self._name_to_idx = {p.name: i for i, p in enumerate(self.params)} + self.ndim = len(self.params) + + def idx(self, name): + """Get the column index for a parameter by name. + + Parameters + ---------- + name : `str` + Parameter name. + + Returns + ------- + `int` + Column index in the (N, D) sample arrays. + """ + return self._name_to_idx[name] + + def sample(self, n, rng=None): + """Draw n samples from the prior distribution. + + Parameters + ---------- + n : `int` + Number of samples to draw. + rng : `numpy.random.Generator`, optional + Random number generator, by default None + + Returns + ------- + samples : `numpy.ndarray` + (N, D) array of samples in sampling space. + mask : `numpy.ndarray` + (N,) boolean array indicating which samples are in bounds. + """ + rng = rng or np.random.default_rng() + samples = np.empty((n, self.ndim)) + mask = np.ones(n, dtype=bool) + + for i, p in enumerate(self.params): + sampler_fn = SAMPLERS[p.sampler] + col = sampler_fn(n, p.lo, p.hi, rng=rng) + samples[:, i] = col + mask &= (col >= p.lo) & (col <= p.hi) + + return samples, mask + + def in_bounds(self, samples): + """Check which rows of an (N, D) array are within parameter bounds. + + Parameters + ---------- + samples : `numpy.ndarray` + (N, D) array of samples in sampling space. + + Returns + ------- + `numpy.ndarray` + (N,) boolean mask where True means the sample is in bounds. + """ + mask = np.ones(len(samples), dtype=bool) + for i, p in enumerate(self.params): + mask &= (samples[:, i] >= p.lo) & (samples[:, i] <= p.hi) + return mask + + def to_physical(self, samples): + """Convert an (N, D) array from sampling space to physical space. + + Parameters + ---------- + samples : `numpy.ndarray` + (N, D) array in sampling space. + + Returns + ------- + `numpy.ndarray` + (N, D) array in physical space. + """ + result = samples.copy() + for i, p in enumerate(self.params): + result[:, i] = to_physical_space(samples[:, i], p.sampler) + return result + + def to_sampling(self, samples): + """Convert an (N, D) array from physical space to sampling space. + + Parameters + ---------- + samples : `numpy.ndarray` + (N, D) array in physical space. + + Returns + ------- + `numpy.ndarray` + (N, D) array in sampling space. + """ + result = samples.copy() + for i, p in enumerate(self.params): + result[:, i] = to_sampling_space(samples[:, i], p.sampler) + return result + + def compute_prior(self, samples): + """Compute the prior probability for each row. + + The joint prior is the product of the marginal priors across all + dimensions. + + Parameters + ---------- + samples : `numpy.ndarray` + (N, D) array in sampling space. + + Returns + ------- + `numpy.ndarray` + (N,) array of prior probabilities. + """ + log_prior = np.zeros(len(samples)) + for i, p in enumerate(self.params): + prior_fn = PRIORS[p.prior] + col_prior = prior_fn(samples[:, i], p.lo, p.hi) + # Clamp to avoid log(0) + log_prior += np.log(np.maximum(col_prior, 1e-300)) + return np.exp(log_prior) + + def compute_sigma(self, hit_samples, average_density_one_dim): + """Compute per-hit, per-dimension Gaussian widths (sigma). + + Parameters + ---------- + hit_samples : `numpy.ndarray` + (K, D) array of hit locations in sampling space. + average_density_one_dim : `float` + Characteristic inter-sample spacing, typically + ``1 / num_explored ** (1 / D)``. + + Returns + ------- + `numpy.ndarray` + (K, D) array of sigma values. + """ + K = len(hit_samples) + sigmas = np.empty((K, self.ndim)) + + for i, p in enumerate(self.params): + col = hit_samples[:, i] + if p.sampler in ('kroupa', 'sana', 'sana_ecc'): + sigmas[:, i] = self._sigma_power_law(p, col, average_density_one_dim) + else: + # uniform, flat_in_log, etc: sigma = avg_density / prior + prior_fn = PRIORS[p.prior] + prior_vals = prior_fn(col, p.lo, p.hi) + sigmas[:, i] = average_density_one_dim / prior_vals + return sigmas + + def _sigma_power_law(self, param, values, avg_density): + """Compute sigma for power-law distributions. + + Handles ``kroupa``, ``sana``, and ``sana_ecc`` samplers by mapping + each value to a normalised CDF position, stepping by + ``avg_density`` in CDF space, mapping back, and taking the + maximum of the two distances. + + Parameters + ---------- + param : `Parameter` + The parameter definition for this dimension. + values : `numpy.ndarray` + (K,) array of hit values in sampling space. + avg_density : `float` + Characteristic inter-sample spacing. + + Returns + ------- + `numpy.ndarray` + (K,) array of sigma values. + """ + if param.sampler == 'kroupa': + alpha = ALPHA_IMF + elif param.sampler == 'sana': + alpha = SANA_G + elif param.sampler == 'sana_ecc': + alpha = SANA_ECC + else: + raise ValueError(f"Unknown power-law sampler: {param.sampler}") + + a = alpha + 1 # e.g., -1.3 for kroupa + lo, hi = param.lo, param.hi + norm = a / (hi**a - lo**a) + + # Step 1: normalized inverse CDF position (matches old code) + # inv_X = (norm / X)^(1 / -alpha) for X = value, lo, hi + inv_exp = 1.0 / (-alpha) + inv_val = np.power(norm / values, inv_exp) + inv_lo = np.power(norm / lo, inv_exp) # scalar + inv_hi = np.power(norm / hi, inv_exp) # scalar + inv_normalized = (inv_val - inv_lo) / (inv_hi - inv_lo) + + # Step 2: step in CDF space, clamp to [0, 1] + inv_right = np.clip(inv_normalized + avg_density, 0, 1) + inv_left = np.clip(inv_normalized - avg_density, 0, 1) + + # Step 3: inverse_back maps CDF position back to parameter space + # inverse_back(inv) = norm / |inv * (inv_hi - inv_lo) + inv_lo|^(-alpha) + inv_range = inv_hi - inv_lo # scalar + + right_arg = np.abs(inv_right * inv_range + inv_lo) + left_arg = np.abs(inv_left * inv_range + inv_lo) + + # Avoid zero bases which would give inf when raised to -alpha (positive power) + right_arg = np.maximum(right_arg, 1e-300) + left_arg = np.maximum(left_arg, 1e-300) + + right_vals = norm / np.power(right_arg, -alpha) + left_vals = norm / np.power(left_arg, -alpha) + + right_dist = np.abs(right_vals - values) + left_dist = np.abs(left_vals - values) + return np.maximum(right_dist, left_dist) diff --git a/src/cosmic/sample/stroopwafel/presets.py b/src/cosmic/sample/stroopwafel/presets.py new file mode 100644 index 000000000..dfea320d1 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/presets.py @@ -0,0 +1,92 @@ +"""Preset hit-definition functions for common DCO selections. + +Each preset factory returns a callable with signature +``(bpp) -> (n_hits, hit_bin_nums)`` where ``hit_bin_nums`` are the +``bin_num`` values of systems classified as hits. +""" +import numpy as np + + +def merging_dco(kstar_1, kstar_2, max_merge_time=13.7): + """Create a hit function selecting merging double compact objects. + + Parameters + ---------- + kstar_1 : `list` of `int` + Allowed kstar types for star 1 (e.g., ``[14]`` for black holes). + kstar_2 : `list` of `int` + Allowed kstar types for star 2 (e.g., ``[14]`` for black holes). + max_merge_time : `float`, optional + Maximum merger time in Gyr, by default 13.7 (Hubble time) + + Returns + ------- + `callable` + Function with signature ``(bpp) -> (n_hits, hit_bin_nums)`` + where ``bpp`` is a `pandas.DataFrame` and ``hit_bin_nums`` is a + `numpy.ndarray` of bin_num values. + """ + from legwork import evol + import astropy.units as u + + k1_set = set(kstar_1) + k2_set = set(kstar_2) + + def is_interesting(bpp): + # Select rows matching the DCO type (either ordering) + pairs_mask = ( + (bpp.kstar_1.isin(k1_set) & bpp.kstar_2.isin(k2_set)) + | (bpp.kstar_1.isin(k2_set) & bpp.kstar_2.isin(k1_set)) + ) + # Must still be bound (sep > 0) + interesting_mask = pairs_mask & (bpp.sep > 0) + candidates = bpp.loc[interesting_mask].drop_duplicates(subset='bin_num', keep='first') + + if len(candidates) == 0: + return 0, np.array([], dtype=int) + + # Compute merger times using LEGWORK + merge_times = evol.get_t_merge_ecc( + ecc_i=candidates.ecc.values, + a_i=candidates.sep.values * u.Rsun, + m_1=candidates.mass_1.values * u.Msun, + m_2=candidates.mass_2.values * u.Msun, + ) + merge_mask = merge_times < (max_merge_time * u.Gyr) + hits = candidates[merge_mask] + + return len(hits), hits.bin_num.values + + return is_interesting + + +def any_dco(kstar_1, kstar_2): + """Create a hit function selecting DCOs regardless of merge time. + + Parameters + ---------- + kstar_1 : `list` of `int` + Allowed kstar types for star 1. + kstar_2 : `list` of `int` + Allowed kstar types for star 2. + + Returns + ------- + `callable` + Function with signature ``(bpp) -> (n_hits, hit_bin_nums)`` + where ``bpp`` is a `pandas.DataFrame` and ``hit_bin_nums`` is a + `numpy.ndarray` of bin_num values. + """ + k1_set = set(kstar_1) + k2_set = set(kstar_2) + + def is_interesting(bpp): + pairs_mask = ( + (bpp.kstar_1.isin(k1_set) & bpp.kstar_2.isin(k2_set)) + | (bpp.kstar_1.isin(k2_set) & bpp.kstar_2.isin(k1_set)) + ) + interesting_mask = pairs_mask & (bpp.sep > 0) + candidates = bpp.loc[interesting_mask].drop_duplicates(subset='bin_num', keep='first') + return len(candidates), candidates.bin_num.values + + return is_interesting diff --git a/src/cosmic/sample/stroopwafel/priors.py b/src/cosmic/sample/stroopwafel/priors.py new file mode 100644 index 000000000..8d312a730 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/priors.py @@ -0,0 +1,171 @@ +"""Vectorized prior probability density functions. + +Each function takes ``(values, lo, hi)`` where ``values`` is an (N,) ndarray +and ``lo``/``hi`` are the bounds in the sampling (transformed) space, and +returns an (N,) ndarray of prior probability densities. +""" +import numpy as np +from .constants import ALPHA_IMF, SANA_G, SANA_ECC + + +def uniform(values, lo, hi): + """Compute the uniform prior probability density. + + Parameters + ---------- + values : `numpy.ndarray` + (N,) array of sample values (unused, density is constant). + lo : `float` + Lower bound of the uniform distribution. + hi : `float` + Upper bound of the uniform distribution. + + Returns + ------- + `numpy.ndarray` + (N,) array of constant prior densities. + """ + return np.full(len(values), 1.0 / (hi - lo)) + + +def flat_in_log(values, lo, hi): + """Compute the flat-in-log prior probability density. + + The prior is uniform in log10 space; bounds are already in log10. + + Parameters + ---------- + values : `numpy.ndarray` + (N,) array of sample values in log10 space (unused, density is + constant). + lo : `float` + Lower bound in log10 space. + hi : `float` + Upper bound in log10 space. + + Returns + ------- + `numpy.ndarray` + (N,) array of constant prior densities. + """ + return np.full(len(values), 1.0 / (hi - lo)) + + +def kroupa(values, lo, hi): + """Compute the Kroupa IMF power-law prior probability density. + + Parameters + ---------- + values : `numpy.ndarray` + (N,) array of sample values. + lo : `float` + Lower bound of the distribution. + hi : `float` + Upper bound of the distribution. + + Returns + ------- + `numpy.ndarray` + (N,) array of prior densities. + """ + a = ALPHA_IMF + norm = (a + 1) / (hi**(a + 1) - lo**(a + 1)) + return norm * np.power(values, a) + + +def sana(values, lo, hi): + """Compute the Sana orbital period power-law prior probability density. + + Parameters + ---------- + values : `numpy.ndarray` + (N,) array of sample values. + lo : `float` + Lower bound of the distribution. + hi : `float` + Upper bound of the distribution. + + Returns + ------- + `numpy.ndarray` + (N,) array of prior densities. + """ + a = SANA_G + norm = (a + 1) / (hi**(a + 1) - lo**(a + 1)) + return norm * np.power(values, a) + + +def sana_ecc(values, lo, hi): + """Compute the Sana eccentricity power-law prior probability density. + + Parameters + ---------- + values : `numpy.ndarray` + (N,) array of sample values. + lo : `float` + Lower bound of the distribution. + hi : `float` + Upper bound of the distribution. + + Returns + ------- + `numpy.ndarray` + (N,) array of prior densities. + """ + a = SANA_ECC + norm = (a + 1) / (hi**(a + 1) - lo**(a + 1)) + return norm * np.power(values, a) + + +def uniform_in_sine(values, lo, hi): + """Compute the uniform-in-sine prior probability density. + + Parameters + ---------- + values : `numpy.ndarray` + (N,) array of sample values in sine space (unused, density is + constant). + lo : `float` + Lower bound in sine space. + hi : `float` + Upper bound in sine space. + + Returns + ------- + `numpy.ndarray` + (N,) array of constant prior densities. + """ + return np.full(len(values), 1.0 / (hi - lo)) + + +def uniform_in_cosine(values, lo, hi): + """Compute the uniform-in-cosine prior probability density. + + Parameters + ---------- + values : `numpy.ndarray` + (N,) array of sample values in cosine space (unused, density is + constant). + lo : `float` + Lower bound in cosine space. + hi : `float` + Upper bound in cosine space. + + Returns + ------- + `numpy.ndarray` + (N,) array of constant prior densities. + """ + return np.full(len(values), 1.0 / (hi - lo)) + + +# Registry mapping string names to functions +PRIORS = { + 'uniform': uniform, + 'flat_in_log': flat_in_log, + 'kroupa': kroupa, + 'sana': sana, + 'sana_ecc': sana_ecc, + 'uniform_in_sine': uniform_in_sine, + 'uniform_in_cosine': uniform_in_cosine, +} diff --git a/src/cosmic/sample/stroopwafel/rejection.py b/src/cosmic/sample/stroopwafel/rejection.py new file mode 100644 index 000000000..74b6f82cb --- /dev/null +++ b/src/cosmic/sample/stroopwafel/rejection.py @@ -0,0 +1,125 @@ +"""Vectorized rejection checking for binary star systems. + +Provides fully vectorized numpy operations for ZAMS radius, Roche lobe +radius, and physical rejection criteria used to discard unphysical binary +systems prior to evolution. +""" +import numpy as np +from .constants import R_COEFF, ZSOL, R_SOL_TO_AU, MINIMUM_SECONDARY_MASS + + +def get_zams_radius(mass, metallicity): + """Compute zero-age main sequence radius for arrays of stars. + + Parameters + ---------- + mass : `numpy.ndarray` + (N,) array of stellar masses in solar masses. + metallicity : `numpy.ndarray` + (N,) array of metallicities. + + Returns + ------- + `numpy.ndarray` + (N,) array of ZAMS radii in AU. + """ + mass = np.asarray(mass, dtype=float) + metallicity = np.asarray(metallicity, dtype=float) + + xi = np.log10(metallicity / ZSOL) + + # R_COEFF is a (9, 5) matrix of polynomial coefficients + # For each of the 9 radius coefficients, evaluate the polynomial in xi + R = np.array(R_COEFF) # (9, 5) + # Build Vandermonde matrix: xi^0, xi^1, xi^2, xi^3, xi^4 + powers = np.column_stack([xi**k for k in range(5)]) # (N, 5) + # rc[j, n] = sum_k R[j, k] * xi[n]^k + rc = R @ powers.T # (9, N) + + top = (rc[0] * mass**2.5 + rc[1] * mass**6.5 + rc[2] * mass**11 + + rc[3] * mass**19 + rc[4] * mass**19.5) + bottom = (rc[5] + rc[6] * mass**2 + rc[7] * mass**8.5 + + mass**18.5 + rc[8] * mass**19.5) + + return (top / bottom) * R_SOL_TO_AU + + +def calculate_roche_lobe_radius(mass1, mass2): + """Compute Roche lobe radius using the Eggleton (1983) approximation. + + Parameters + ---------- + mass1 : `numpy.ndarray` + (N,) array of masses of the star filling its Roche lobe. + mass2 : `numpy.ndarray` + (N,) array of companion masses. + + Returns + ------- + `numpy.ndarray` + (N,) array of Roche lobe radii (dimensionless, in units of + orbital separation). + """ + mass1 = np.asarray(mass1, dtype=float) + mass2 = np.asarray(mass2, dtype=float) + q = mass1 / mass2 + q_cbrt = np.power(q, 1.0 / 3.0) + return 0.49 / (0.6 + np.power(q, -2.0 / 3.0) * np.log(1.0 + q_cbrt)) + + +def default_reject(samples_physical, derived, param_names, + min_secondary_mass=MINIMUM_SECONDARY_MASS): + """Default rejection function for DCO progenitor systems. + + Rejects systems where the secondary mass is below the minimum, the + stars are in contact at ZAMS, or either star overflows its Roche lobe + at periastron. + + Parameters + ---------- + samples_physical : `numpy.ndarray` + (N, D) array of samples in physical space. + derived : `dict` + Dictionary with keys ``'mass_2'``, ``'metallicity_1'``, + ``'metallicity_2'``, and ``'separation'``, each mapping to an + (N,) array. + param_names : `list` of `str` + Sorted list of parameter names (used to find column indices). + min_secondary_mass : `float`, optional + Minimum allowed secondary mass in solar masses, by default + ``MINIMUM_SECONDARY_MASS`` + + Returns + ------- + `numpy.ndarray` + (N,) boolean mask where True indicates a rejected system. + """ + idx = {name: i for i, name in enumerate(param_names)} + + mass_1 = samples_physical[:, idx['mass_1']] + mass_2 = derived['mass_2'] + met_1 = derived['metallicity_1'] + met_2 = derived['metallicity_2'] + separation = derived['separation'] + ecc = samples_physical[:, idx['ecc']] + + # Compute ZAMS radii + radius_1 = get_zams_radius(mass_1, met_1) + radius_2 = get_zams_radius(mass_2, met_2) + + # Roche lobe radii at periastron + peri_sep = separation * (1 - ecc) + rl_1 = peri_sep * calculate_roche_lobe_radius(mass_1, mass_2) + rl_2 = peri_sep * calculate_roche_lobe_radius(mass_2, mass_1) + + roche_tracker_1 = radius_1 / rl_1 + roche_tracker_2 = radius_2 / rl_2 + + rejected = ( + (mass_2 < min_secondary_mass) + | (separation <= (radius_1 + radius_2)) + | (roche_tracker_1 > 1) + | (roche_tracker_2 > 1) + ) + + return rejected diff --git a/src/cosmic/sample/stroopwafel/result.py b/src/cosmic/sample/stroopwafel/result.py new file mode 100644 index 000000000..f933cdaf0 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/result.py @@ -0,0 +1,82 @@ +"""Container for STROOPWAFEL adaptive sampling results.""" +import numpy as np + + +class STROOPWAFELResult: + """Results from an AdaptiveSampler run. + + Parameters + ---------- + samples : `numpy.ndarray` + (N, D) array of all samples in physical space. + param_names : `list` of `str` + Parameter names defining the column ordering. + weights : `numpy.ndarray` + (N,) array of importance sampling weights. + is_hit : `numpy.ndarray` + (N,) boolean array indicating which samples are hits. + generation : `numpy.ndarray` + (N,) integer array (0 = exploration, 1+ = refinement). + gaussian_idx : `numpy.ndarray` + (N,) integer array (-1 = exploration, k = from Gaussian + component k). + num_explored : `int` + Number of systems simulated during the exploration phase. + num_hits : `int` + Total number of hits found across all phases. + fraction_explored : `float` + Adaptive exploration fraction. + bpp_frames : `list` of `pandas.DataFrame`, optional + COSMIC binary population parameter tables, by default None + initC_frames : `list` of `pandas.DataFrame`, optional + COSMIC initial conditions tables, by default None + kick_info_frames : `list` of `pandas.DataFrame`, optional + COSMIC natal kick information tables, by default None + """ + + def __init__(self, samples, param_names, weights, is_hit, generation, + gaussian_idx, num_explored, num_hits, fraction_explored, + bpp_frames=None, initC_frames=None, kick_info_frames=None): + self.samples = samples + self.param_names = param_names + self.weights = weights + self.is_hit = is_hit + self.generation = generation + self.gaussian_idx = gaussian_idx + self.num_explored = num_explored + self.num_hits = num_hits + self.fraction_explored = fraction_explored + self.bpp_frames = bpp_frames or [] + self.initC_frames = initC_frames or [] + self.kick_info_frames = kick_info_frames or [] + + @property + def hit_rate(self): + """Importance-weighted hit rate (sum of weights over hits / N). + + Returns + ------- + `float` + Weighted fraction of systems that are hits, or 0.0 if + weights or hit flags are not available. + """ + if self.weights is None or self.is_hit is None: + return 0.0 + return np.sum(self.weights[self.is_hit]) / len(self.weights) + + @property + def hit_rate_uncertainty(self): + """Standard error on the importance-weighted hit rate. + + Returns + ------- + `float` + Standard error ``std(w_hits) / sqrt(N)``, or 0.0 if fewer + than two hits are present. + """ + if self.weights is None or self.is_hit is None: + return 0.0 + w_hits = self.weights[self.is_hit] + if len(w_hits) < 2: + return 0.0 + return np.std(w_hits, ddof=1) / np.sqrt(len(self.weights)) diff --git a/src/cosmic/sample/stroopwafel/samplers.py b/src/cosmic/sample/stroopwafel/samplers.py new file mode 100644 index 000000000..9f7f44345 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/samplers.py @@ -0,0 +1,186 @@ +"""Vectorized sampling functions for each distribution type. + +Each function draws ``n`` samples between ``lo`` and ``hi`` bounds (in the +sampling/transformed space) and returns an (N,) ndarray. +""" +import numpy as np +from .constants import ALPHA_IMF, SANA_G, SANA_ECC + + +def uniform(n, lo, hi, rng=None): + """Draw samples from a uniform distribution. + + Parameters + ---------- + n : `int` + Number of samples to draw. + lo : `float` + Lower bound of the uniform distribution. + hi : `float` + Upper bound of the uniform distribution. + rng : `numpy.random.Generator`, optional + Random number generator, by default None + + Returns + ------- + `numpy.ndarray` + (N,) array of uniform samples. + """ + rng = rng or np.random.default_rng() + return rng.uniform(lo, hi, n) + + +def flat_in_log(n, lo, hi, rng=None): + """Sample uniformly in log10 space. + + Parameters + ---------- + n : `int` + Number of samples to draw. + lo : `float` + Lower bound (already in log10 space). + hi : `float` + Upper bound (already in log10 space). + rng : `numpy.random.Generator`, optional + Random number generator, by default None + + Returns + ------- + `numpy.ndarray` + (N,) array of samples in log10 space. + """ + rng = rng or np.random.default_rng() + return rng.uniform(lo, hi, n) + + +def kroupa(n, lo, hi, rng=None): + """Inverse CDF sampling from a Kroupa-like power law p(x) ~ x^alpha. + + Parameters + ---------- + n : `int` + Number of samples to draw. + lo : `float` + Lower bound of the distribution. + hi : `float` + Upper bound of the distribution. + rng : `numpy.random.Generator`, optional + Random number generator, by default None + + Returns + ------- + `numpy.ndarray` + (N,) array of power-law distributed samples. + """ + rng = rng or np.random.default_rng() + u = rng.uniform(0, 1, n) + a = ALPHA_IMF + 1 + return np.power(u * (hi**a - lo**a) + lo**a, 1.0 / a) + + +def sana(n, lo, hi, rng=None): + """Inverse CDF sampling from the Sana orbital period distribution. + + Parameters + ---------- + n : `int` + Number of samples to draw. + lo : `float` + Lower bound (in log10 space). + hi : `float` + Upper bound (in log10 space). + rng : `numpy.random.Generator`, optional + Random number generator, by default None + + Returns + ------- + `numpy.ndarray` + (N,) array of samples from the Sana period distribution. + """ + rng = rng or np.random.default_rng() + u = rng.uniform(0, 1, n) + a = SANA_G + 1 + return np.power(u * (hi**a - lo**a) + lo**a, 1.0 / a) + + +def sana_ecc(n, lo, hi, rng=None): + """Inverse CDF sampling from the Sana eccentricity distribution. + + Parameters + ---------- + n : `int` + Number of samples to draw. + lo : `float` + Lower bound of the eccentricity distribution. + hi : `float` + Upper bound of the eccentricity distribution. + rng : `numpy.random.Generator`, optional + Random number generator, by default None + + Returns + ------- + `numpy.ndarray` + (N,) array of samples from the Sana eccentricity distribution. + """ + rng = rng or np.random.default_rng() + u = rng.uniform(0, 1, n) + a = SANA_ECC + 1 + return np.power(u * (hi**a - lo**a) + lo**a, 1.0 / a) + + +def uniform_in_sine(n, lo, hi, rng=None): + """Sample uniformly in sine-transformed space. + + Parameters + ---------- + n : `int` + Number of samples to draw. + lo : `float` + Lower bound in sine space. + hi : `float` + Upper bound in sine space. + rng : `numpy.random.Generator`, optional + Random number generator, by default None + + Returns + ------- + `numpy.ndarray` + (N,) array of uniform samples in sine space. + """ + rng = rng or np.random.default_rng() + return rng.uniform(lo, hi, n) + + +def uniform_in_cosine(n, lo, hi, rng=None): + """Sample uniformly in cosine-transformed space. + + Parameters + ---------- + n : `int` + Number of samples to draw. + lo : `float` + Lower bound in cosine space. + hi : `float` + Upper bound in cosine space. + rng : `numpy.random.Generator`, optional + Random number generator, by default None + + Returns + ------- + `numpy.ndarray` + (N,) array of uniform samples in cosine space. + """ + rng = rng or np.random.default_rng() + return rng.uniform(lo, hi, n) + + +# Registry mapping string names to functions +SAMPLERS = { + 'uniform': uniform, + 'flat_in_log': flat_in_log, + 'kroupa': kroupa, + 'sana': sana, + 'sana_ecc': sana_ecc, + 'uniform_in_sine': uniform_in_sine, + 'uniform_in_cosine': uniform_in_cosine, +} diff --git a/src/cosmic/sample/stroopwafel/transforms.py b/src/cosmic/sample/stroopwafel/transforms.py new file mode 100644 index 000000000..b5c7120c9 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/transforms.py @@ -0,0 +1,87 @@ +"""Vectorized scale transformations between physical and sampling space. + +Some distributions (flat_in_log, sana, uniform_in_sine, etc.) sample in a +transformed space (e.g., log10). These functions convert between the two +representations. All functions operate on (N,) column arrays for a given +parameter. +""" +import numpy as np + + +def to_sampling_space(values, sampler_type): + """Convert physical-space values to the sampling (transformed) space. + + Parameters + ---------- + values : `numpy.ndarray` + (N,) array of values in physical space. + sampler_type : `str` + Name of the sampling distribution (e.g., ``'flat_in_log'``, + ``'uniform_in_sine'``). + + Returns + ------- + `numpy.ndarray` + (N,) array of values in sampling space. + """ + if sampler_type in ('flat_in_log', 'sana'): + return np.log10(values) + elif sampler_type == 'uniform_in_sine': + return np.sin(values) + elif sampler_type == 'uniform_in_cosine': + return np.cos(values + np.pi / 2) + return values + + +def to_physical_space(values, sampler_type): + """Convert sampling-space values back to physical space. + + Parameters + ---------- + values : `numpy.ndarray` + (N,) array of values in sampling space. + sampler_type : `str` + Name of the sampling distribution (e.g., ``'flat_in_log'``, + ``'uniform_in_sine'``). + + Returns + ------- + `numpy.ndarray` + (N,) array of values in physical space. + """ + if sampler_type in ('flat_in_log', 'sana'): + return np.power(10.0, values) + elif sampler_type == 'uniform_in_sine': + return np.arcsin(values) + elif sampler_type == 'uniform_in_cosine': + return np.arccos(values) - np.pi / 2 + return values + + +def transform_bounds(lo, hi, sampler_type): + """Get the bounds in sampling space for a given physical-space range. + + Parameters + ---------- + lo : `float` + Lower bound in physical space. + hi : `float` + Upper bound in physical space. + sampler_type : `str` + Name of the sampling distribution. + + Returns + ------- + lo_transformed : `float` + Lower bound in sampling space. + hi_transformed : `float` + Upper bound in sampling space. + """ + if sampler_type == 'flat_in_log': + return np.log10(lo), np.log10(hi) + elif sampler_type == 'uniform_in_sine': + return -1.0, 1.0 + elif sampler_type == 'uniform_in_cosine': + return -1.0, 1.0 + # sana, sana_ecc, kroupa, uniform: bounds stay as-is + return lo, hi From 18578715289f73920c6748bf029836c72c79e3ad Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Fri, 10 Apr 2026 16:08:49 -0400 Subject: [PATCH 02/45] remove useless constants --- src/cosmic/sample/stroopwafel/constants.py | 6 ------ src/cosmic/sample/stroopwafel/engine.py | 9 ++++----- src/cosmic/sample/stroopwafel/mixture_model.py | 7 +++---- src/cosmic/sample/stroopwafel/rejection.py | 7 +++---- 4 files changed, 10 insertions(+), 19 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/constants.py b/src/cosmic/sample/stroopwafel/constants.py index 74ecfe6d8..c04c86e66 100755 --- a/src/cosmic/sample/stroopwafel/constants.py +++ b/src/cosmic/sample/stroopwafel/constants.py @@ -14,12 +14,6 @@ [17.84778000, -7.45345690, -48.96066856, -40.05386135, -9.09331816],\ [0.00022582, -0.00186899, 0.00388783, 0.00142402, -0.00007671]] -MINIMUM_SECONDARY_MASS = 0.1 R_SOL_TO_AU = 0.00465047 -METALLICITY_SOL = 0.0142 ZSOL = 0.02 -REJECTION_SAMPLES_PER_BATCH = 1e4 -TOTAL_REJECTION_SAMPLES = 1e6 -NUM_GENERATIONS = 1 MIN_ENTROPY_CHANGE = 0.01 -KAPPA = 1.0 diff --git a/src/cosmic/sample/stroopwafel/engine.py b/src/cosmic/sample/stroopwafel/engine.py index 9c0d57298..1efa9f941 100644 --- a/src/cosmic/sample/stroopwafel/engine.py +++ b/src/cosmic/sample/stroopwafel/engine.py @@ -12,7 +12,6 @@ from .mixture_model import GaussianMixture from .result import STROOPWAFELResult -from .constants import KAPPA, NUM_GENERATIONS class AdaptiveSampler: @@ -44,9 +43,9 @@ class AdaptiveSampler: nproc : `int`, optional Number of CPU cores for COSMIC, by default 1 kappa : `float`, optional - Gaussian width scaling factor, by default ``KAPPA`` + Gaussian width scaling factor, by default 1.0 n_generations : `int`, optional - Number of refinement generations, by default ``NUM_GENERATIONS`` + Number of refinement generations, by default 1 mc_only : `bool`, optional If True, only run exploration (standard Monte Carlo), by default False @@ -56,8 +55,8 @@ class AdaptiveSampler: def __init__(self, parameter_space, total_systems, batch_size, BSEDict, compute_derived, reject_systems, is_interesting, - output_path='output', nproc=1, kappa=KAPPA, - n_generations=NUM_GENERATIONS, mc_only=False, seed=None): + output_path='output', nproc=1, kappa=1.0, + n_generations=1, mc_only=False, seed=None): self.param_space = parameter_space self.total_systems = total_systems self.batch_size = batch_size diff --git a/src/cosmic/sample/stroopwafel/mixture_model.py b/src/cosmic/sample/stroopwafel/mixture_model.py index 1f3618a14..4b44a463a 100644 --- a/src/cosmic/sample/stroopwafel/mixture_model.py +++ b/src/cosmic/sample/stroopwafel/mixture_model.py @@ -5,7 +5,7 @@ """ import numpy as np from scipy.stats import multivariate_normal, entropy as scipy_entropy -from .constants import KAPPA, MIN_ENTROPY_CHANGE +from .constants import MIN_ENTROPY_CHANGE class GaussianMixture: @@ -38,7 +38,7 @@ def ndim(self): return self.means.shape[1] @classmethod - def from_hits(cls, hit_samples, param_space, average_density_one_dim, kappa=KAPPA): + def from_hits(cls, hit_samples, param_space, average_density_one_dim, kappa=1.0): """Create a Gaussian mixture by placing one component at each hit. Parameters @@ -51,8 +51,7 @@ def from_hits(cls, hit_samples, param_space, average_density_one_dim, kappa=KAPP Characteristic inter-sample spacing, ``1 / num_explored ** (1 / D)``. kappa : `float`, optional - Width scaling factor for the Gaussian covariances, by default - ``KAPPA`` + Width scaling factor for the Gaussian covariances, by default 1.0 Returns ------- diff --git a/src/cosmic/sample/stroopwafel/rejection.py b/src/cosmic/sample/stroopwafel/rejection.py index 74b6f82cb..95e06cd7f 100644 --- a/src/cosmic/sample/stroopwafel/rejection.py +++ b/src/cosmic/sample/stroopwafel/rejection.py @@ -5,7 +5,7 @@ systems prior to evolution. """ import numpy as np -from .constants import R_COEFF, ZSOL, R_SOL_TO_AU, MINIMUM_SECONDARY_MASS +from .constants import R_COEFF, ZSOL, R_SOL_TO_AU def get_zams_radius(mass, metallicity): @@ -68,7 +68,7 @@ def calculate_roche_lobe_radius(mass1, mass2): def default_reject(samples_physical, derived, param_names, - min_secondary_mass=MINIMUM_SECONDARY_MASS): + min_secondary_mass=0.08): """Default rejection function for DCO progenitor systems. Rejects systems where the secondary mass is below the minimum, the @@ -86,8 +86,7 @@ def default_reject(samples_physical, derived, param_names, param_names : `list` of `str` Sorted list of parameter names (used to find column indices). min_secondary_mass : `float`, optional - Minimum allowed secondary mass in solar masses, by default - ``MINIMUM_SECONDARY_MASS`` + Minimum allowed secondary mass in solar masses, by default 0.08 Returns ------- From a471c1c96c4d2e7cd51fee6ee1d49344dbd54ce2 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Fri, 10 Apr 2026 19:17:15 -0400 Subject: [PATCH 03/45] final fixes for working stroopwafel implementation --- src/cosmic/sample/meson.build | 3 +- src/cosmic/sample/stroopwafel/constants.py | 27 ++++--- src/cosmic/sample/stroopwafel/engine.py | 9 ++- .../sample/stroopwafel/mixture_model.py | 11 +-- .../sample/stroopwafel/parameter_space.py | 71 ++++++++++--------- 5 files changed, 70 insertions(+), 51 deletions(-) diff --git a/src/cosmic/sample/meson.build b/src/cosmic/sample/meson.build index fa4e4fec3..96ecbb82b 100644 --- a/src/cosmic/sample/meson.build +++ b/src/cosmic/sample/meson.build @@ -10,4 +10,5 @@ py3.install_sources( ) subdir('cmc') -subdir('sampler') \ No newline at end of file +subdir('sampler') +subdir('stroopwafel') \ No newline at end of file diff --git a/src/cosmic/sample/stroopwafel/constants.py b/src/cosmic/sample/stroopwafel/constants.py index c04c86e66..f71ac2b02 100755 --- a/src/cosmic/sample/stroopwafel/constants.py +++ b/src/cosmic/sample/stroopwafel/constants.py @@ -3,17 +3,24 @@ SANA_G = -0.55 SANA_ECC = -0.45 -R_COEFF =\ - [[1.71535900, 0.62246212, -0.92557761, -1.16996966, -0.30631491],\ - [6.59778800, -0.42450044, -12.13339427, -10.73509484, -2.51487077],\ - [10.08855000, -7.11727086, -31.67119479, -24.24848322, -5.33608972],\ - [1.01249500, 0.32699690, -0.00923418, -0.03876858, -0.00412750],\ - [0.07490166, 0.02410413, 0.07233664, 0.03040467, 0.00197741],\ - [0.01077422, 0.00000000, 0.00000000, 0.00000000, 0.00000000],\ - [3.08223400, 0.94472050, -2.15200882, -2.49219496, -0.63848738],\ - [17.84778000, -7.45345690, -48.96066856, -40.05386135, -9.09331816],\ - [0.00022582, -0.00186899, 0.00388783, 0.00142402, -0.00007671]] +R_COEFF = [ + [1.71535900, 0.62246212, -0.92557761, -1.16996966, -0.30631491], + [6.59778800, -0.42450044, -12.13339427, -10.73509484, -2.51487077], + [10.08855000, -7.11727086, -31.67119479, -24.24848322, -5.33608972], + [1.01249500, 0.32699690, -0.00923418, -0.03876858, -0.00412750], + [0.07490166, 0.02410413, 0.07233664, 0.03040467, 0.00197741], + [0.01077422, 0.00000000, 0.00000000, 0.00000000, 0.00000000], + [3.08223400, 0.94472050, -2.15200882, -2.49219496, -0.63848738], + [17.84778000, -7.45345690, -48.96066856, -40.05386135, -9.09331816], + [0.00022582, -0.00186899, 0.00388783, 0.00142402, -0.00007671] +] R_SOL_TO_AU = 0.00465047 ZSOL = 0.02 MIN_ENTROPY_CHANGE = 0.01 + +# 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 array sizes or overflowing float64 normalisation constants. +MIN_ACTIVE_FRACTION = 0.01 diff --git a/src/cosmic/sample/stroopwafel/engine.py b/src/cosmic/sample/stroopwafel/engine.py index 1efa9f941..0eab96417 100644 --- a/src/cosmic/sample/stroopwafel/engine.py +++ b/src/cosmic/sample/stroopwafel/engine.py @@ -12,6 +12,7 @@ from .mixture_model import GaussianMixture from .result import STROOPWAFELResult +from .constants import MIN_ACTIVE_FRACTION class AdaptiveSampler: @@ -115,7 +116,9 @@ def _explore(self): print(f" Prior rejection rate: {self.prior_fraction_rejected:.4f}") while self._should_continue_exploring(): - n_oversample = int(2 * np.ceil(self.batch_size / (1 - self.prior_fraction_rejected))) + n_oversample = int(2 * np.ceil( + self.batch_size / max(1.0 - self.prior_fraction_rejected, MIN_ACTIVE_FRACTION) + )) # Sample from prior samples, mask = self.param_space.sample(n_oversample, rng=self.rng) @@ -354,7 +357,7 @@ def _compute_weights(self): N = len(all_samples) # Prior probabilities - pi_norm = 1.0 / (1 - self.prior_fraction_rejected) + pi_norm = 1.0 / max(1.0 - self.prior_fraction_rejected, MIN_ACTIVE_FRACTION) pi = self.param_space.compute_prior(all_samples) * pi_norm # Start with the exploration-phase contribution to denominator @@ -363,7 +366,7 @@ def _compute_weights(self): # Add refinement-phase contributions from each generation's mixture if self.mixture is not None and not self.mc_only: - q_norm = 1.0 / (1 - self.mixture.rejection_rate) + q_norm = 1.0 / max(1.0 - self.mixture.rejection_rate, MIN_ACTIVE_FRACTION) # Evaluate the mixture PDF incrementally (memory-efficient) for k in range(self.mixture.n_components): xPDF_k = multivariate_normal.pdf( diff --git a/src/cosmic/sample/stroopwafel/mixture_model.py b/src/cosmic/sample/stroopwafel/mixture_model.py index 4b44a463a..1fe6b6dc9 100644 --- a/src/cosmic/sample/stroopwafel/mixture_model.py +++ b/src/cosmic/sample/stroopwafel/mixture_model.py @@ -5,7 +5,7 @@ """ import numpy as np from scipy.stats import multivariate_normal, entropy as scipy_entropy -from .constants import MIN_ENTROPY_CHANGE +from .constants import MIN_ENTROPY_CHANGE, MIN_ACTIVE_FRACTION class GaussianMixture: @@ -105,8 +105,9 @@ def sample(self, n_total, param_space, consider_rejection=False, rng=None): for k in range(self.n_components): n_k = int(np.ceil(n_total * self.alphas[k])) - if consider_rejection and self.rejection_rate < 1.0: - n_k = int(2 * np.ceil(n_k / (1 - self.rejection_rate))) + if consider_rejection and self.rejection_rate > 0.0: + active = max(1.0 - self.rejection_rate, MIN_ACTIVE_FRACTION) + n_k = int(2 * np.ceil(n_k / active)) if n_k <= 0: continue s = rng.multivariate_normal(self.means[k], self.covariances[k], size=n_k) @@ -248,8 +249,8 @@ def update_em(self, samples, is_hit, prior_probs, prior_fraction_rejected, True if the entropy check triggers reversion to the previous mixture state. """ - pi_norm = 1.0 / (1 - prior_fraction_rejected) - q_norm = 1.0 / (1 - self.rejection_rate) + pi_norm = 1.0 / max(1.0 - prior_fraction_rejected, MIN_ACTIVE_FRACTION) + q_norm = 1.0 / max(1.0 - self.rejection_rate, MIN_ACTIVE_FRACTION) pi = prior_probs * pi_norm is_hit = np.asarray(is_hit, dtype=float) diff --git a/src/cosmic/sample/stroopwafel/parameter_space.py b/src/cosmic/sample/stroopwafel/parameter_space.py index ede8b1e5b..bb520aa60 100644 --- a/src/cosmic/sample/stroopwafel/parameter_space.py +++ b/src/cosmic/sample/stroopwafel/parameter_space.py @@ -215,10 +215,18 @@ def compute_sigma(self, hit_samples, average_density_one_dim): def _sigma_power_law(self, param, values, avg_density): """Compute sigma for power-law distributions. - Handles ``kroupa``, ``sana``, and ``sana_ecc`` samplers by mapping - each value to a normalised CDF position, stepping by - ``avg_density`` in CDF space, mapping back, and taking the - maximum of the two distances. + Maps each hit value to its CDF position, steps by ``avg_density`` + in CDF space, maps back to parameter space, and returns the + maximum of the two resulting distances. + + The CDF of a power-law p(x) ∝ x^α on [lo, hi] is + + F(x) = (x^a − lo^a) / (hi^a − lo^a), a = α + 1 + + so the inverse is F⁻¹(u) = (u·(hi^a − lo^a) + lo^a)^{1/a}. + All intermediate quantities stay in [lo^a, hi^a], avoiding the + catastrophic cancellation that occurred in the previous + ``inv_lo``-based formulation when ``lo`` was very small. Parameters ---------- @@ -233,6 +241,12 @@ def _sigma_power_law(self, param, values, avg_density): ------- `numpy.ndarray` (K,) array of sigma values. + + Raises + ------ + `ValueError` + If ``param.lo <= 0`` (the CDF formula requires a positive + lower bound) or if the sampler is not a recognised power law. """ if param.sampler == 'kroupa': alpha = ALPHA_IMF @@ -243,36 +257,29 @@ def _sigma_power_law(self, param, values, avg_density): else: raise ValueError(f"Unknown power-law sampler: {param.sampler}") - a = alpha + 1 # e.g., -1.3 for kroupa - lo, hi = param.lo, param.hi - norm = a / (hi**a - lo**a) - - # Step 1: normalized inverse CDF position (matches old code) - # inv_X = (norm / X)^(1 / -alpha) for X = value, lo, hi - inv_exp = 1.0 / (-alpha) - inv_val = np.power(norm / values, inv_exp) - inv_lo = np.power(norm / lo, inv_exp) # scalar - inv_hi = np.power(norm / hi, inv_exp) # scalar - inv_normalized = (inv_val - inv_lo) / (inv_hi - inv_lo) - - # Step 2: step in CDF space, clamp to [0, 1] - inv_right = np.clip(inv_normalized + avg_density, 0, 1) - inv_left = np.clip(inv_normalized - avg_density, 0, 1) + if param.lo <= 0: + raise ValueError( + f"Parameter '{param.name}' has lo={param.lo} <= 0 but uses " + f"the '{param.sampler}' power-law sampler. The CDF formula " + f"requires lo > 0. For 'sana' (period), pass the log10(P) " + f"lower bound directly, e.g. min_value=0.15 rather than " + f"min_value=np.log10(0.15)." + ) - # Step 3: inverse_back maps CDF position back to parameter space - # inverse_back(inv) = norm / |inv * (inv_hi - inv_lo) + inv_lo|^(-alpha) - inv_range = inv_hi - inv_lo # scalar + a = alpha + 1 # 0.55 for sana_ecc, 0.45 for sana, -1.3 for kroupa + lo_a = float(param.lo) ** a + hi_a = float(param.hi) ** a + range_a = hi_a - lo_a # always finite; no huge intermediates - right_arg = np.abs(inv_right * inv_range + inv_lo) - left_arg = np.abs(inv_left * inv_range + inv_lo) + # --- Forward CDF: F(x) = (x^a - lo_a) / range_a ---------------- + u = np.clip((np.power(values, a) - lo_a) / range_a, 0.0, 1.0) - # Avoid zero bases which would give inf when raised to -alpha (positive power) - right_arg = np.maximum(right_arg, 1e-300) - left_arg = np.maximum(left_arg, 1e-300) + # --- Step in CDF space ------------------------------------------ + u_right = np.clip(u + avg_density, 0.0, 1.0) + u_left = np.clip(u - avg_density, 0.0, 1.0) - right_vals = norm / np.power(right_arg, -alpha) - left_vals = norm / np.power(left_arg, -alpha) + # --- Inverse CDF: F^{-1}(u) = (u * range_a + lo_a)^{1/a} ------- + x_right = np.power(u_right * range_a + lo_a, 1.0 / a) + x_left = np.power(u_left * range_a + lo_a, 1.0 / a) - right_dist = np.abs(right_vals - values) - left_dist = np.abs(left_vals - values) - return np.maximum(right_dist, left_dist) + return np.maximum(np.abs(x_right - values), np.abs(x_left - values)) From aaf06833f5009d87737c5aa31cc2e9770b0247e6 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 11 May 2026 22:21:22 -0400 Subject: [PATCH 04/45] actually apply the shuffle properly --- src/cosmic/sample/stroopwafel/engine.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/engine.py b/src/cosmic/sample/stroopwafel/engine.py index 0eab96417..7dc9d2ffa 100644 --- a/src/cosmic/sample/stroopwafel/engine.py +++ b/src/cosmic/sample/stroopwafel/engine.py @@ -273,13 +273,14 @@ def _refine(self): phys = phys[keep] derived = {k: v[keep] for k, v in derived.items()} - # Trim to batch size - self.rng.shuffle(np.arange(len(valid_samples))) + # Trim to batch size with randomisation + indices = np.arange(len(valid_samples)) + self.rng.shuffle(indices) n_take = min(len(valid_samples), self.batch_size) - batch_samples = valid_samples[:n_take] - batch_phys = phys[:n_take] - batch_gauss_idx = valid_gauss_idx[:n_take] - batch_derived = {k: v[:n_take] for k, v in derived.items()} + batch_samples = valid_samples[indices[:n_take]] + batch_phys = phys[indices[:n_take]] + batch_gauss_idx = valid_gauss_idx[indices[:n_take]] + batch_derived = {k: v[indices[:n_take]] for k, v in derived.items()} # Evolve with COSMIC n_hits, hit_bin_nums, bpp, initC, kick_info = self._evolve_batch( From 6efd6aafb7556124287eac02cf4542af4ee3c558 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 11 May 2026 22:22:01 -0400 Subject: [PATCH 05/45] clean up einsum --- src/cosmic/sample/stroopwafel/mixture_model.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/mixture_model.py b/src/cosmic/sample/stroopwafel/mixture_model.py index 1fe6b6dc9..d8aeba186 100644 --- a/src/cosmic/sample/stroopwafel/mixture_model.py +++ b/src/cosmic/sample/stroopwafel/mixture_model.py @@ -296,12 +296,16 @@ def update_em(self, samples, is_hit, prior_probs, prior_fraction_rejected, old_covs = self.covariances[keep] new_covs = np.empty_like(old_covs) for k in range(K_new): - distance = (new_means[k] - samples)[:, :, None] # (N, D, 1) - matrix = np.einsum('nij,nji->nij', distance, distance) # (N, D, D) - factor = weights_normalized[:, 0] * rho[:, k] # (N,) - new_covs[k] = np.sum(factor[:, None, None] * matrix, axis=0) / new_alphas[k] - - # Entropy check + diff = samples - new_means[k] # (N, D) + outer = np.einsum('ni,nj->nij', diff, diff) # (N, D, D) + factor = weights_normalized[:, 0] * rho[:, k] # (N,) + new_covs[k] = np.sum(factor[:, None, None] * outer, axis=0) / new_alphas[k] + + # Normalised effective sample size: exp(H(w)) / N, where H is the + # Shannon entropy of the normalised importance weights. Ranges from + # 1/N (all weight on one sample) to 1 (uniform weights). We revert + # if this metric changes by less than MIN_ENTROPY_CHANGE between + # generations, which indicates the mixture has stopped improving. entropy_change = np.exp(scipy_entropy(weights_normalized[:, 0])) / N if entropies is not None: if len(entropies) >= 1 and entropy_change - entropies[-1] < MIN_ENTROPY_CHANGE: From d50f31cae172317c9aba424b40e09bf5ed254b4e Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 11 May 2026 22:22:32 -0400 Subject: [PATCH 06/45] clarify sana bounds --- src/cosmic/sample/stroopwafel/parameter_space.py | 11 +++++++++-- src/cosmic/sample/stroopwafel/samplers.py | 9 +++++++-- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/parameter_space.py b/src/cosmic/sample/stroopwafel/parameter_space.py index bb520aa60..e32081981 100644 --- a/src/cosmic/sample/stroopwafel/parameter_space.py +++ b/src/cosmic/sample/stroopwafel/parameter_space.py @@ -22,9 +22,16 @@ class Parameter: name : `str` Name of the parameter (used for column ordering). min_value : `float` - Lower bound in physical space. + Lower bound. For most samplers this is in physical space (e.g. + solar masses for ``'kroupa'``, eccentricity for ``'sana_ecc'``). + + **Exception:** ``'sana'`` (orbital period) operates in + log10(period / days) internally, so ``min_value`` and + ``max_value`` must be passed in log10 space. For example, to + span periods from ~1.4 d to ~316 000 d use + ``min_value=0.15, max_value=5.5`` (i.e. log10 of those values). max_value : `float` - Upper bound in physical space. + Upper bound (same unit convention as ``min_value``). sampler : `str`, optional Name of the sampling distribution, by default ``'uniform'`` prior : `str`, optional diff --git a/src/cosmic/sample/stroopwafel/samplers.py b/src/cosmic/sample/stroopwafel/samplers.py index 9f7f44345..d42a8df39 100644 --- a/src/cosmic/sample/stroopwafel/samplers.py +++ b/src/cosmic/sample/stroopwafel/samplers.py @@ -81,14 +81,19 @@ def kroupa(n, lo, hi, rng=None): def sana(n, lo, hi, rng=None): """Inverse CDF sampling from the Sana orbital period distribution. + The Sana et al. (2012) distribution is a power law in log10(period), + so this sampler operates entirely in log10(period / days) space. + ``lo`` and ``hi`` must therefore be given as log10 values (e.g. + ``lo=0.15, hi=5.5`` spans ~1.4 d to ~316 000 d). + Parameters ---------- n : `int` Number of samples to draw. lo : `float` - Lower bound (in log10 space). + Lower bound in log10(period / days). hi : `float` - Upper bound (in log10 space). + Upper bound in log10(period / days). rng : `numpy.random.Generator`, optional Random number generator, by default None From f2a189f07eaf9846f0f517dbd5d095b0cc1214a8 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 11 May 2026 22:22:54 -0400 Subject: [PATCH 07/45] explain cosine sampling --- src/cosmic/sample/stroopwafel/transforms.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/cosmic/sample/stroopwafel/transforms.py b/src/cosmic/sample/stroopwafel/transforms.py index b5c7120c9..2eae01271 100644 --- a/src/cosmic/sample/stroopwafel/transforms.py +++ b/src/cosmic/sample/stroopwafel/transforms.py @@ -29,6 +29,9 @@ def to_sampling_space(values, sampler_type): elif sampler_type == 'uniform_in_sine': return np.sin(values) elif sampler_type == 'uniform_in_cosine': + # Angles are measured from –π/2 to π/2 (e.g. declination-like + # coordinates), so the sampling variable is cos(θ + π/2) = –sin(θ). + # The round-trip is exact for θ ∈ [–π/2, π/2]. return np.cos(values + np.pi / 2) return values @@ -54,6 +57,7 @@ def to_physical_space(values, sampler_type): elif sampler_type == 'uniform_in_sine': return np.arcsin(values) elif sampler_type == 'uniform_in_cosine': + # Inverse of cos(θ + π/2): arccos(u) – π/2 return np.arccos(values) - np.pi / 2 return values @@ -83,5 +87,8 @@ def transform_bounds(lo, hi, sampler_type): return -1.0, 1.0 elif sampler_type == 'uniform_in_cosine': return -1.0, 1.0 - # sana, sana_ecc, kroupa, uniform: bounds stay as-is + # kroupa, uniform: bounds are already in physical space. + # sana, sana_ecc: bounds are passed by the caller in the native + # sampling space (log10(period) for sana, eccentricity for sana_ecc), + # so no further transformation is needed here. return lo, hi From a10609f9c6f2e1ca889ad04c6c86c653e3684063 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 11 May 2026 22:23:22 -0400 Subject: [PATCH 08/45] clean up tests --- .../stroopwafel/examples/test_new_modules.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/examples/test_new_modules.py b/src/cosmic/sample/stroopwafel/examples/test_new_modules.py index e797b792a..448b25d12 100644 --- a/src/cosmic/sample/stroopwafel/examples/test_new_modules.py +++ b/src/cosmic/sample/stroopwafel/examples/test_new_modules.py @@ -216,10 +216,20 @@ def test_mixture_sample(): # ==================================================================== def test_result_hit_rate(): - result = STROOPWAFELResult() - result.weights = np.ones(100) - result.is_hit = np.zeros(100, dtype=bool) - result.is_hit[:10] = True + weights = np.ones(100) + is_hit = np.zeros(100, dtype=bool) + is_hit[:10] = True + result = STROOPWAFELResult( + samples=np.zeros((100, 1)), + param_names=['mass_1'], + weights=weights, + is_hit=is_hit, + generation=np.zeros(100, dtype=int), + gaussian_idx=np.full(100, -1, dtype=int), + num_explored=100, + num_hits=10, + fraction_explored=1.0, + ) assert abs(result.hit_rate - 0.1) < 1e-10 From ed914d3451bfbd5441174872240e7cb6f55c423d Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 11 May 2026 22:23:31 -0400 Subject: [PATCH 09/45] add new docs page for stroopwafel sampling --- docs/pages/sample/adaptive.rst | 539 +++++++++++++++++++++++++++++++++ 1 file changed, 539 insertions(+) create mode 100644 docs/pages/sample/adaptive.rst diff --git a/docs/pages/sample/adaptive.rst b/docs/pages/sample/adaptive.rst new file mode 100644 index 000000000..0044700db --- /dev/null +++ b/docs/pages/sample/adaptive.rst @@ -0,0 +1,539 @@ +.. _adaptive: + +********************************************** +Adaptive importance sampling for rare systems +********************************************** + +The standard :ref:`independent ` and :ref:`multidimensional ` +samplers draw binary parameters from their prior distributions and evolve each system with +COSMIC. For most stellar outcomes this works well: common-envelope episodes, mass-transferring +binaries, and white dwarf systems all occur frequently enough that thousands of random draws +yield a workable sample. + +Some outcomes are extremely rare. Double black holes that merge within the Hubble time +form at rates of order 1-in-10,000 per binary evolved (or far less at high metallicity). +Sampling such populations with a flat Monte Carlo requires tens of millions of COSMIC calls +to accumulate even a few hundred systems — prohibitive for any serious parameter survey. + +COSMIC includes a vectorised implementation of the STROOPWAFEL algorithm +(`Broekgaarden et al. 2019 `_) that solves this +problem using *adaptive importance sampling*. The sampler first explores parameter space +to locate the progenitor regions of the target population, then concentrates its simulation +budget on those regions, and finally corrects the biased sampling with importance weights so +that any weighted statistic remains an unbiased estimator of the true prior-weighted +distribution. + + +When should I use this? +======================== + +Use STROOPWAFEL whenever you need a statistically representative sample of a rare binary +outcome and cannot afford the total binary count that flat Monte Carlo would require. +Typical use cases include: + +* Double black holes (or neutron stars) merging within the Hubble time (GW sources) +* BH + stellar-companion systems, such as X-ray binaries or Be/X-ray binaries +* Short-period post-common-envelope binaries that survive to become AM CVn systems +* Any binary channel with a formation efficiency ≲ 10\ :sup:`−3` per prior draw + +If your target population is common (≳ 1 % of all binaries evolving as the target), the +plain independent sampler is simpler and fast enough. The break-even point is roughly +where collecting 100 hits would require more than ~10,000 total evolutions. + + +How it works: the battleships analogy +====================================== + +STROOPWAFEL works in three phases that map neatly onto the board game Battleships. + +**Exploration — random fire** + Binaries are drawn at random from the prior distributions and evolved with COSMIC. + Every binary that produces the desired outcome is recorded as a *hit*. An adaptive + stopping criterion (based on the observed hit rate) decides when the remaining budget + is better spent on refinement than on further random exploration. + +**Adaptation — mark the ships** + One multivariate Gaussian component is placed at each hit location in parameter space. + The width of each Gaussian is derived from the local density of the prior (via the CDF + of the prior distribution) so that the proposal is appropriately broad regardless of + the parameter's scale. Together the Gaussians form a *mixture model* — a coarse map of + where progenitors live. + +**Refinement — concentrate fire** + New binaries are drawn from the Gaussian mixture instead of the broad prior. Because + the mixture is concentrated near known progenitor regions, a much larger fraction of the + simulated systems produce hits. Between refinement generations the mixture is + optionally updated via an expectation-maximisation (EM) step that shifts weight toward + the most productive Gaussians. + +**Weight calculation — correct the bias** + Every simulated system receives an importance weight + + .. math:: + + w(x) = \frac{\pi(x)}{Q(x)}, \qquad + Q(x) = f_e\,\pi(x) + (1 - f_e)\,q(x), + + where :math:`\pi(x)` is the prior probability density, :math:`q(x)` is the Gaussian + mixture density, and :math:`f_e` is the fraction of systems drawn from the prior + during exploration. Using these weights, any statistic computed on the hit population + is an unbiased estimator of the corresponding prior-weighted quantity. + + +Setting up the parameter space +================================ + +The :class:`~cosmic.sample.stroopwafel.ParameterSpace` class defines which binary +parameters are sampled and their distributions. Each +:class:`~cosmic.sample.stroopwafel.Parameter` specifies a name, physical-space bounds, a +sampling distribution, and a prior distribution. + +.. code-block:: python + + import numpy as np + from cosmic.sample.stroopwafel import ParameterSpace, Parameter + + params = ParameterSpace([ + Parameter('mass_1', 5.0, 150.0, sampler='kroupa', prior='kroupa'), + Parameter('q', 0.01, 1.0, sampler='uniform', prior='uniform'), + Parameter('porb', 0.15, 5.5, sampler='sana', prior='sana'), + Parameter('ecc', 1e-9, 0.9999, sampler='sana_ecc', prior='sana_ecc'), + Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log',prior='flat_in_log'), + ]) + +The available samplers and corresponding prior names are: + +.. list-table:: + :header-rows: 1 + :widths: 20 45 35 + + * - Sampler name + - Distribution + - Typical use + * - ``'uniform'`` + - Uniform between bounds + - Mass ratio, any flat prior + * - ``'kroupa'`` + - Kroupa (2001) power law (:math:`dN/dm_1 \propto m_1^{-2.3}`) + - Primary mass + * - ``'sana'`` + - Sana et al. (2012) period distribution + (:math:`dN/d\!\log P \propto (\log P)^{-0.55}`) + - Orbital period + * - ``'sana_ecc'`` + - Sana et al. (2012) eccentricity distribution + (:math:`dN/de \propto e^{-0.45}`) + - Orbital eccentricity + * - ``'flat_in_log'`` + - Uniform in :math:`\log_{10}` (Öpik's law) + - Metallicity, semi-major axis + +.. note:: + + The ``'sana'`` sampler operates in :math:`\log_{10}(P/\text{days})` space. The bounds + you supply are :math:`\log_{10}` values directly — **not** periods in days. The + Sana et al. (2012) fit is valid over :math:`\log_{10}(P) \in [0.15,\, 5.5]`, + corresponding to periods of roughly 1.4 to 316,000 days. + + The lower bound **must be positive** (i.e. :math:`\log_{10}(P_\text{min}) > 0`, + so :math:`P_\text{min} > 1` day). Do **not** pass ``np.log10(P_min)`` when that + value would be negative. + +Parameters are stored and returned in alphabetical order by name. Use +``params.names`` to inspect the column ordering of any sample array. + + +Defining derived quantities +============================ + +COSMIC requires ``mass_2``, ``separation``, and ``metallicity`` in addition to the +directly-sampled parameters. The ``compute_derived`` callback converts a ``(N, D)`` +array of physical-space samples and a sorted list of parameter names into a dictionary of +``(N,)`` arrays: + +.. code-block:: python + + def compute_derived(samples_physical, param_names): + """Convert sampled parameters to COSMIC inputs.""" + idx = {name: i for i, name in enumerate(param_names)} + + mass_1 = samples_physical[:, idx['mass_1']] + q = samples_physical[:, idx['q']] + porb = samples_physical[:, idx['porb']] # physical days after 10^x transform + z = samples_physical[:, idx['metallicity']] + + mass_2 = mass_1 * q + + # Kepler's third law: a [AU] from P [yr] and M [M_sun]; then convert to AU + separation = ((porb / 365.25) ** 2 * (mass_1 + mass_2)) ** (1.0 / 3.0) + + return { + 'mass_2': mass_2, + 'metallicity_1': z, + 'metallicity_2': z, + 'separation': separation, # AU, required by default_reject + } + +.. note:: + + The keys ``'mass_2'``, ``'metallicity_1'``, ``'metallicity_2'``, and ``'separation'`` + are expected by :func:`~cosmic.sample.stroopwafel.rejection.default_reject`. If you + supply a custom rejection function you are free to use different key names. + + +Defining the rejection function +================================ + +Before a batch is passed to COSMIC, unphysical systems are filtered out: stars already +overflowing their Roche lobes at ZAMS, binaries whose components are in contact, and +systems below the hydrogen-burning limit for the secondary. The built-in +:func:`~cosmic.sample.stroopwafel.rejection.default_reject` function performs all of these +checks: + +.. code-block:: python + + from cosmic.sample.stroopwafel.rejection import default_reject + +It expects the arrays produced by ``compute_derived`` and returns a boolean mask where +``True`` means the system is rejected. For most use cases this default is appropriate. +If you need extra cuts — for example, discarding systems with very low metallicity or +imposing a minimum primary mass — you can wrap it: + +.. code-block:: python + + def my_reject(samples_physical, derived, param_names): + base_mask = default_reject(samples_physical, derived, param_names) + # additionally reject secondaries below 1 M_sun + base_mask |= (derived['mass_2'] < 1.0) + return base_mask + + +Defining the hit criterion +============================ + +The ``is_interesting`` argument identifies which evolved systems count as hits. It receives +the COSMIC ``bpp`` DataFrame for the current batch and must return a tuple +``(n_hits, hit_bin_nums)`` where ``hit_bin_nums`` is an integer array of ``bin_num`` values +(0-indexed within the batch). + +STROOPWAFEL ships two preset factory functions in +:mod:`cosmic.sample.stroopwafel.presets`. + +``any_dco(kstar_1, kstar_2)`` + Selects all bound double compact objects whose stellar types match the supplied lists, + regardless of merger time. Use this for populations where you care about the DCO + existing rather than merging within the Hubble time. + +``merging_dco(kstar_1, kstar_2, max_merge_time=13.7)`` + Like ``any_dco`` but additionally requires the merger time (computed via LEGWORK) to be + less than ``max_merge_time`` Gyr. Suitable for gravitational-wave source studies. + +.. code-block:: python + + from cosmic.sample.stroopwafel.presets import any_dco, merging_dco + + # All bound BH-BH systems, no merger time cut + is_interesting = any_dco(kstar_1=[14], kstar_2=[14]) + + # Only BH-BH systems merging within the Hubble time + is_interesting_merging = merging_dco(kstar_1=[14], kstar_2=[14], max_merge_time=13.7) + +See :ref:`kstar-table` for the full list of stellar type codes. + +You can also write a fully custom hit function. For example, to find BH + stellar +companion systems (``kstar_2 ∈ 0–9``) that remain bound for at least 100 Myr after the +BH forms: + +.. code-block:: python + + import numpy as np + + _STELLAR_TYPES = set(range(10)) # kstar 0–9: MS through He-giant branch + + def bh_star_100myr(bpp): + """Hit: BH with a stellar companion bound for at least 100 Myr.""" + bh_star = bpp.loc[ + ( + ((bpp['kstar_1'] == 14) & bpp['kstar_2'].isin(_STELLAR_TYPES)) + | ((bpp['kstar_2'] == 14) & bpp['kstar_1'].isin(_STELLAR_TYPES)) + ) + & (bpp['sep'] > 0) + ] + if bh_star.empty: + return 0, np.array([], dtype=int) + + # Duration in the BH + star state for each binary + span = bh_star.groupby('bin_num')['tphys'].agg(lambda t: t.max() - t.min()) + hits = span.index[span >= 100.0].values + return len(hits), hits + + +Example 1: Bound BH + BH binaries +==================================== + +The following end-to-end example samples all bound BH-BH systems (no merger time +restriction) using a five-dimensional parameter space covering primary mass, mass ratio, +orbital period, eccentricity, and metallicity. + +.. code-block:: python + + import numpy as np + from cosmic.sample.stroopwafel import AdaptiveSampler, ParameterSpace, Parameter + from cosmic.sample.stroopwafel.presets import any_dco + from cosmic.sample.stroopwafel.rejection import default_reject + + # ------------------------------------------------------------------ + # BSE physics settings + # ------------------------------------------------------------------ + BSEDict = { + "pts1": 0.001, "pts2": 0.01, "pts3": 0.02, "zsun": 0.014, + "windflag": 3, "neta": 0.5, "bwind": 0.0, "hewind": 0.5, + "beta": 0.125, "xi": 0.5, "acc2": 1.5, "LBV_flag": 1, + "alpha1": 1.0, "lambdaf": 0.0, "ceflag": 1, "cekickflag": 2, + "cemergeflag": 1, "cehestarflag": 0, "qcflag": 5, + "qcrit_array": [0.0] * 16, + "kickflag": 5, "sigma": 265.0, "bhflag": 1, "bhsigmafrac": 1.0, + "sigmadiv": -20.0, "ecsn": 2.25, "ecsn_mlow": 1.6, "aic": 1, + "ussn": 1, "polar_kick_angle": 90.0, + "natal_kick_array": [[-100.0]*5, [-100.0]*5], + "remnantflag": 4, "mxns": 3.0, "rembar_massloss": 0.5, + "wd_mass_lim": 1, "grflag": 1, "eddfac": 10, "tflag": 1, + "ST_tide": 1, "ifflag": 1, "wdflag": 1, "epsnov": 0.001, + "bdecayfac": 1, "bconst": 3000, "ck": 1000, "htpmb": 1, + "ST_cr": 1, "rtmsflag": 0, + "fprimc_array": [2.0 / 21.0] * 16, + "mm_mu_ns": 400.0, "mm_mu_bh": 200.0, "pisn": -2, + } + + # ------------------------------------------------------------------ + # Parameter space + # ------------------------------------------------------------------ + params = ParameterSpace([ + Parameter('mass_1', 5.0, 150.0, sampler='kroupa', prior='kroupa'), + Parameter('q', 0.01, 1.0, sampler='uniform', prior='uniform'), + Parameter('porb', 0.15, 5.5, sampler='sana', prior='sana'), + Parameter('ecc', 1e-9, 0.9999, sampler='sana_ecc', prior='sana_ecc'), + Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log',prior='flat_in_log'), + ]) + + # ------------------------------------------------------------------ + # Derived quantities + # ------------------------------------------------------------------ + def compute_derived(samples_physical, param_names): + idx = {name: i for i, name in enumerate(param_names)} + mass_1 = samples_physical[:, idx['mass_1']] + q = samples_physical[:, idx['q']] + porb = samples_physical[:, idx['porb']] + z = samples_physical[:, idx['metallicity']] + mass_2 = mass_1 * q + separation = ((porb / 365.25) ** 2 * (mass_1 + mass_2)) ** (1.0 / 3.0) + return { + 'mass_2': mass_2, + 'metallicity_1': z, + 'metallicity_2': z, + 'separation': separation, + } + + # ------------------------------------------------------------------ + # Run + # ------------------------------------------------------------------ + sampler = AdaptiveSampler( + parameter_space=params, + total_systems=50_000, + batch_size=500, + BSEDict=BSEDict, + compute_derived=compute_derived, + reject_systems=default_reject, + is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), + output_path='output/bhbh', + nproc=4, + n_generations=3, + seed=42, + ) + + result = sampler.run() + + print(f"Total hits: {result.num_hits}") + print(f"Weighted hit rate: {result.hit_rate:.4e} ± {result.hit_rate_uncertainty:.4e}") + + +Example 2: BH + star binaries surviving 100 Myr +================================================== + +For outcomes that are less extreme but still rare — such as persistent BH + star systems — +STROOPWAFEL provides substantial efficiency gains over flat Monte Carlo. Using the same +parameter space, ``compute_derived``, and ``BSEDict`` as Example 1: + +.. code-block:: python + + import numpy as np + from cosmic.sample.stroopwafel import AdaptiveSampler + from cosmic.sample.stroopwafel.rejection import default_reject + + _STELLAR_TYPES = set(range(10)) # kstar 0–9: MS through He-giant + + def bh_star_100myr(bpp): + """Hit: BH + stellar companion bound for at least 100 Myr.""" + bh_star = bpp.loc[ + ( + ((bpp['kstar_1'] == 14) & bpp['kstar_2'].isin(_STELLAR_TYPES)) + | ((bpp['kstar_2'] == 14) & bpp['kstar_1'].isin(_STELLAR_TYPES)) + ) + & (bpp['sep'] > 0) + ] + if bh_star.empty: + return 0, np.array([], dtype=int) + span = bh_star.groupby('bin_num')['tphys'].agg(lambda t: t.max() - t.min()) + hits = span.index[span >= 100.0].values + return len(hits), hits + + sampler = AdaptiveSampler( + parameter_space=params, # reuse from Example 1 + total_systems=20_000, + batch_size=500, + BSEDict=BSEDict, # reuse from Example 1 + compute_derived=compute_derived, # reuse from Example 1 + reject_systems=default_reject, + is_interesting=bh_star_100myr, + output_path='output/bh_star', + nproc=4, + n_generations=2, + seed=42, + ) + + result = sampler.run() + +Because BH + star systems are more common than merging BH-BH pairs, a smaller total budget +is needed and fewer refinement generations are required before the mixture model is +well-constrained. + + +Choosing ``total_systems``, ``batch_size``, and ``n_generations`` +=================================================================== + +``batch_size`` +-------------- + +``batch_size`` sets how many systems are passed to +:meth:`~cosmic.evolve.Evolve.evolve` per call. + +* Aim for ``batch_size`` to be a multiple of ``nproc`` so that COSMIC distributes work + evenly across cores. +* Values of 200–1000 are typical. Batches smaller than ~50 increase Python overhead per + call; batches larger than ~5000 may cause memory pressure on the output DataFrames. +* A practical starting point is ``batch_size = 100 * nproc``. + +``total_systems`` +----------------- + +This is the total number of binary evolutions across all phases. + +* For **very rare events** (hit rate ≲ 10\ :sup:`−4`, e.g. merging BH-BH at near-solar + metallicity), start with ``total_systems`` in the range 100,000–500,000. The exploration + phase will find tens to hundreds of hits; refinement then multiplies that count many-fold. +* For **moderately rare events** (hit rate ~ 10\ :sup:`−3` to 10\ :sup:`−2`, e.g. any bound + BH-BH or long-lived BH + star), 20,000–50,000 systems is usually sufficient. +* As a rule of thumb, aim for at least ~30 hits during exploration before the adaptation + phase begins — fewer hits lead to a poorly-constrained Gaussian mixture. If exploration + ends with very few hits, increase ``total_systems`` and re-run. + +``n_generations`` +----------------- + +Each refinement generation uses an equal share of the remaining budget after exploration. +The EM step between generations can improve the mixture, but with diminishing returns: + +* ``n_generations = 1`` (the default) uses the mixture as constructed from exploration + hits, with no EM updates. This is a good starting point for any new target population. +* ``n_generations = 3`` gives a noticeable improvement for very rare populations with + complex progenitor structure. +* Beyond 5 generations the returns diminish rapidly. + +.. tip:: + + To run a plain Monte Carlo without any adaptation or refinement — useful as a baseline + or for common populations — pass ``mc_only=True``. The sampler will draw from the + prior only, and importance weights will equal the prior density divided by itself + (i.e. all weights are equal). + + +Working with results +====================== + +:meth:`~cosmic.sample.stroopwafel.engine.AdaptiveSampler.run` returns a +:class:`~cosmic.sample.stroopwafel.result.STROOPWAFELResult` object containing all +simulated systems, their importance weights, and summary statistics: + +.. code-block:: python + + import numpy as np + + # Shape of sample array and column ordering + print(result.samples.shape) # (N_total, D) + print(result.param_names) # sorted alphabetically, e.g. + # ['ecc', 'mass_1', 'metallicity', 'porb', 'q'] + + # Extract hits and their weights + hit_samples = result.samples[result.is_hit] # (N_hits, D) + hit_weights = result.weights[result.is_hit] # (N_hits,) + + # Normalise weights for the hit population + hit_weights_norm = hit_weights / hit_weights.sum() + + # Importance-weighted primary mass histogram + m1_col = result.param_names.index('mass_1') + m1_hits = hit_samples[:, m1_col] + hist, edges = np.histogram(m1_hits, bins=20, weights=hit_weights_norm) + + # Importance-weighted hit rate (fraction of prior draws producing a hit) + print(f"Hit rate: {result.hit_rate:.4e} ± {result.hit_rate_uncertainty:.4e}") + + # How the budget was spent + print(f"Explored: {result.num_explored} Total hits: {result.num_hits}") + +.. note:: + + ``result.samples`` stores samples in **physical space** (masses in M\ :sub:`☉`, + periods in days, etc.) in the alphabetically-sorted column order defined by + ``ParameterSpace``. Always use ``result.param_names`` to map column indices to + parameter names rather than relying on the order in which you defined the parameters. + + +Saving and loading results +============================ + +The full result can be saved to HDF5 for later analysis: + +.. code-block:: python + + from cosmic.sample.stroopwafel import io as swio + + swio.save_result('bhbh_result.h5', result) + +The file stores the sample array, importance weights, hit flags, generation labels, and +summary statistics as HDF5 datasets and attributes. + +To reload the data in a later session without re-running the sampler: + +.. code-block:: python + + import h5py + import numpy as np + + with h5py.File('bhbh_result.h5', 'r') as f: + samples = f['samples'][:] + weights = f['weights'][:] + is_hit = f['is_hit'][:] + param_names = list(f.attrs['param_names']) + num_hits = f.attrs['num_hits'] + + hits = samples[is_hit] + m1 = hits[:, param_names.index('mass_1')] + w = weights[is_hit] + print(f"Weighted mean BH primary mass: {np.average(m1, weights=w):.1f} M_sun") + +.. note:: + + The HDF5 file does **not** store the raw COSMIC ``bpp`` output tables. If you need + the evolutionary histories of the hit systems, re-evolve them with COSMIC using the + hit sample coordinates from ``result.samples[result.is_hit]`` and the same ``BSEDict``. From d68dd967159ad697b3dbe4a042d3c2ae595b9d70 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 11 May 2026 22:23:43 -0400 Subject: [PATCH 10/45] add it to the main page (the new docs page) --- docs/pages/runpop.rst | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/pages/runpop.rst b/docs/pages/runpop.rst index fff24e0c9..cf7803d97 100644 --- a/docs/pages/runpop.rst +++ b/docs/pages/runpop.rst @@ -31,6 +31,16 @@ We consider both cases in the guides below. sample/independent sample/multidim +For rare outcomes such as merging double black holes or persistent X-ray binary +systems, COSMIC also provides an adaptive importance sampler based on the +STROOPWAFEL algorithm that concentrates the simulation budget on progenitor +regions of parameter space. + +.. toctree:: + :maxdepth: 1 + + sample/adaptive + You can also use COSMIC to sample the initial conditions for a Globular Cluster (GC) using the ClusterMonteCarlo (CMC) software package. Check out the guide below for more information. From 13252b7ffb7e4789413a1d996aa79a2b4c7dfc78 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 19 May 2026 16:25:54 -0700 Subject: [PATCH 11/45] add kick params for sampling --- src/cosmic/sample/stroopwafel/constants.py | 7 + src/cosmic/sample/stroopwafel/engine.py | 31 +- .../sample/stroopwafel/examples/bh_star.py | 299 ++++++++++++++++++ src/cosmic/sample/stroopwafel/priors.py | 33 +- src/cosmic/sample/stroopwafel/samplers.py | 44 ++- src/cosmic/sample/stroopwafel/transforms.py | 9 + 6 files changed, 420 insertions(+), 3 deletions(-) create mode 100644 src/cosmic/sample/stroopwafel/examples/bh_star.py diff --git a/src/cosmic/sample/stroopwafel/constants.py b/src/cosmic/sample/stroopwafel/constants.py index f71ac2b02..21112f903 100755 --- a/src/cosmic/sample/stroopwafel/constants.py +++ b/src/cosmic/sample/stroopwafel/constants.py @@ -3,6 +3,13 @@ SANA_G = -0.55 SANA_ECC = -0.45 +# Log-normal natal kick distribution. +# The kick magnitude v [km/s] follows LogNormal(mu, sigma), meaning +# ln(v) ~ Normal(NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA). +# With mu=5.67, sigma=0.59 the median kick is exp(5.67) ≈ 291 km/s. +NATAL_KICK_LOG_MU = 5.67 # mean of ln(v_kick / km s⁻¹) +NATAL_KICK_LOG_SIGMA = 0.59 # std dev of ln(v_kick / km s⁻¹) + R_COEFF = [ [1.71535900, 0.62246212, -0.92557761, -1.16996966, -0.30631491], [6.59778800, -0.42450044, -12.13339427, -10.73509484, -2.51487077], diff --git a/src/cosmic/sample/stroopwafel/engine.py b/src/cosmic/sample/stroopwafel/engine.py index 7dc9d2ffa..6f1465ec8 100644 --- a/src/cosmic/sample/stroopwafel/engine.py +++ b/src/cosmic/sample/stroopwafel/engine.py @@ -401,6 +401,18 @@ def _compute_weights(self): # ------------------------------------------------------------------ # COSMIC interface # ------------------------------------------------------------------ + + # Natal kick parameters that can be sampled in the ParameterSpace and + # injected as per-binary columns in the InitialBinaryTable. When all + # eight are present, ``natal_kick_array`` is stripped from the BSEDict + # so COSMIC reads the per-binary values instead of global defaults. + # The ``randomseed`` columns required by COSMIC are filled with zeros + # (COSMIC ignores them when the kick values are already provided). + _KICK_PARAM_NAMES = ( + 'natal_kick_1', 'phi_1', 'theta_1', 'mean_anomaly_1', + 'natal_kick_2', 'phi_2', 'theta_2', 'mean_anomaly_2', + ) + def _evolve_batch(self, samples_physical, derived): """Evolve a batch of binaries with COSMIC and identify hits. @@ -440,9 +452,26 @@ def _evolve_batch(self, samples_physical, derived): metallicity=derived['metallicity_1'], ) + # If natal kick parameters were sampled, inject them as per-binary + # columns and omit the global natal_kick_array from the BSEDict. + # COSMIC uses per-binary values when all FLATTENED_NATAL_KICK_COLUMNS + # are present in the table AND natal_kick_array is absent from BSEDict. + param_names_set = set(self.param_space.names) + if all(k in param_names_set for k in self._KICK_PARAM_NAMES): + for col in self._KICK_PARAM_NAMES: + batch_initial[col] = samples_physical[:, idx[col]] + # randomseed columns are required by COSMIC's reshape but unused + # when kick values are provided directly. + batch_initial['randomseed_1'] = 0 + batch_initial['randomseed_2'] = 0 + bse_dict_run = {k: v for k, v in self.bse_dict.items() + if k != 'natal_kick_array'} + else: + bse_dict_run = self.bse_dict + bpp, bcm, initC, kick_info = Evolve.evolve( initialbinarytable=batch_initial, - BSEDict=self.bse_dict, + BSEDict=bse_dict_run, nproc=self.nproc, ) diff --git a/src/cosmic/sample/stroopwafel/examples/bh_star.py b/src/cosmic/sample/stroopwafel/examples/bh_star.py new file mode 100644 index 000000000..06c5da637 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/examples/bh_star.py @@ -0,0 +1,299 @@ +#!/usr/bin/env python +"""Example: black hole + normal-star binaries with STROOPWAFEL vs Monte Carlo. + +A "hit" is any binary where one component is a black hole (kstar=14) while +its companion is still a normal (non-degenerate) star (kstar < 10) and the +system remains bound (sep > 0). BH+star binaries are intrinsically rare +because they require a massive primary that forms a BH without disrupting +the binary, making them a good testbed for adaptive importance sampling. + +The script runs both methods on identical parameter spaces and system budgets, +then prints a side-by-side comparison of hit rates, uncertainties, and +wall-clock times. The key figure of merit is how many Monte Carlo systems +would be needed to match STROOPWAFEL's statistical precision. + +Usage +----- + python bh_star.py # full comparison (default 10 000 systems) + python bh_star.py --num_systems 50000 # larger run for a clearer signal + python bh_star.py --mc_only # MC baseline only + python bh_star.py --sw_only # STROOPWAFEL only +""" +import os +import sys +import time +import argparse +import numpy as np + +sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +sys.path.append(os.path.join(os.path.dirname(__file__), '../..')) + +from stroopwafel import AdaptiveSampler, ParameterSpace, Parameter +from stroopwafel.rejection import default_reject +from stroopwafel import io as swio + +# ------------------------------------------------------------------ +# CLI +# ------------------------------------------------------------------ +parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, +) +parser.add_argument('--num_systems', type=int, default=50000, + help='Total systems to evolve per method (default: 50000)') +parser.add_argument('--batch_size', type=int, default=1000, + help='Systems per COSMIC call (default: 1000)') +parser.add_argument('--num_cores', type=int, default=8, + help='CPU cores for COSMIC (default: 8)') +parser.add_argument('--n_generations', type=int, default=1, + help='STROOPWAFEL refinement generations (default: 1)') +parser.add_argument('--mc_only', action='store_true', + help='Run Monte Carlo baseline only, skip STROOPWAFEL') +parser.add_argument('--sw_only', action='store_true', + help='Run STROOPWAFEL only, skip Monte Carlo baseline') +parser.add_argument('--output_dir', default='output/bh_star', + help='Directory for output files (default: output/bh_star)') +parser.add_argument('--seed', type=int, default=117, + help='Base random seed; MC uses seed, STROOPWAFEL uses seed+1') +args = parser.parse_args() + +# ------------------------------------------------------------------ +# BSE physics +# ------------------------------------------------------------------ +BSEDict = { + 'xi': 1.0, 'bhflag': 1, 'neta': 0.5, 'windflag': 3, 'wdflag': 1, + 'alpha1': [1.0, 1.0], 'pts1': 0.001, 'pts3': 0.02, 'pts2': 0.01, + 'epsnov': 0.001, 'hewind': 0.5, 'ck': 1000, 'bwind': 0.0, + 'lambdaf': 0.0, 'mxns': 3.0, 'beta': -1.0, 'tflag': 1, 'acc2': 1.5, + 'grflag': 1, 'remnantflag': 4, 'ceflag': 0, 'eddfac': 1.0, + 'ifflag': 0, 'bconst': 3000, 'sigma': 265.0, 'gamma': -2.0, + 'pisn': 45.0, + # natal_kick_array is intentionally absent: kick parameters are sampled + # per-binary in the ParameterSpace below, and the engine injects them + # directly into the InitialBinaryTable for each COSMIC call. + 'bhsigmafrac': 1.0, 'polar_kick_angle': 90, + 'qcrit_array': [0.0] * 16, + 'cekickflag': 2, 'cehestarflag': 0, 'cemergeflag': 0, + 'ecsn': 2.5, 'ecsn_mlow': 1.8, 'aic': 1, 'ussn': 0, + 'sigmadiv': -20.0, 'qcflag': 5, 'eddlimflag': 0, + 'fprimc_array': [2.0 / 21.0] * 16, + 'bhspinflag': 0, 'bhspinmag': 0.0, 'rejuv_fac': 1.0, + 'rejuvflag': 0, 'htpmb': 1, 'ST_cr': 1, 'ST_tide': 1, + 'bdecayfac': 1, 'rembar_massloss': 0.5, 'kickflag': 5, + 'zsun': 0.014, 'bhms_coll_flag': 0, 'don_lim': -1, + 'acc_lim': [-1, -1], 'rtmsflag': 0, 'wd_mass_lim': 1, + "ppi_co_shift": 0.0, "ppi_extra_ml": 0.0, "fryer_mass_limit": 0, + "maltsev_mode": 0, "maltsev_fallback": 0.5, "maltsev_pf_prob": 0.1, + "mm_mu_ns": 800, "mm_mu_bh": 400, "LBV_flag": 1, + "fryer_fmix": 0.5, "fryer_mcrit_nsbh": 5.0, "smt_periastron_check": 0 +} + +# ------------------------------------------------------------------ +# Parameter space +# +# Orbital / stellar parameters (5 dimensions) +# mass_1 : primary mass [Msun], Kroupa IMF +# q : mass ratio m2/m1 ∈ [0.01, 1], uniform +# porb : log10(orbital period / days), Sana power law +# bounds 0.15 to 5.5 → periods ~1.4 d to ~316 000 d +# ecc : eccentricity, Sana power law +# metallicity : metallicity, log-uniform +# +# Natal kick parameters (8 dimensions, 4 per star) +# natal_kick_N : kick speed [km/s], log-normal (mu=5.67, sigma=0.59 in ln-space) +# physical bounds [0.1, 5000] km/s; median ≈ 291 km/s +# phi_N : kick elevation angle [degrees], uniform ∈ [−90, 90] +# theta_N : kick azimuthal angle [degrees], uniform ∈ [0, 360] +# mean_anomaly_N: orbital phase at kick [degrees], uniform ∈ [0, 360] +# +# Total: 13-dimensional parameter space. +# Note: ParameterSpace sorts parameters alphabetically, so the internal +# column order is fixed and independent of the order given here. +# ------------------------------------------------------------------ +params = ParameterSpace([ + # --- orbital / stellar --- + Parameter('mass_1', 5.0, 150.0, sampler='kroupa', prior='kroupa'), + Parameter('q', 0.01, 1.0, sampler='uniform', prior='uniform'), + Parameter('porb', 0.15, 5.5, sampler='sana', prior='sana'), + Parameter('ecc', 1e-9, 0.99999999, sampler='sana_ecc', prior='sana_ecc'), + Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log', prior='flat_in_log'), + # --- natal kick: star 1 --- + Parameter('natal_kick_1', 0.1, 5000.0, sampler='log_normal', prior='log_normal'), + Parameter('phi_1', -90.0, 90.0, sampler='uniform', prior='uniform'), + Parameter('theta_1', 0.0, 360.0, sampler='uniform', prior='uniform'), + Parameter('mean_anomaly_1', 0.0, 360.0, sampler='uniform', prior='uniform'), + # --- natal kick: star 2 --- + Parameter('natal_kick_2', 0.1, 5000.0, sampler='log_normal', prior='log_normal'), + Parameter('phi_2', -90.0, 90.0, sampler='uniform', prior='uniform'), + Parameter('theta_2', 0.0, 360.0, sampler='uniform', prior='uniform'), + Parameter('mean_anomaly_2', 0.0, 360.0, sampler='uniform', prior='uniform'), +]) + +# ------------------------------------------------------------------ +# Derived quantities +# ------------------------------------------------------------------ +def compute_derived(samples_physical, param_names): + """Compute mass_2, metallicities, and orbital separation from samples. + + Parameters + ---------- + samples_physical : numpy.ndarray + (N, D) array of samples in physical space. + param_names : list of str + Column labels for each dimension. + + Returns + ------- + dict + Keys: 'mass_2', 'metallicity_1', 'metallicity_2', 'separation'. + """ + idx = {name: i for i, name in enumerate(param_names)} + m1 = samples_physical[:, idx['mass_1']] + q = samples_physical[:, idx['q']] + porb = samples_physical[:, idx['porb']] # days (sana sampler output) + z = samples_physical[:, idx['metallicity']] + + mass_2 = m1 * q + + # Kepler's third law: a³ [AU³] = (P [yr])² · M [Msun] + # Convert period from days to years before applying. + separation = ((porb / 365.25) ** 2 * (m1 + mass_2)) ** (1.0 / 3.0) + + return { + 'mass_2': mass_2, + 'metallicity_1': z, + 'metallicity_2': z, + 'separation': separation, + } + +# ------------------------------------------------------------------ +# Hit definition: BH (kstar=14) + normal star (kstar 0–9), still bound +# ------------------------------------------------------------------ +_STAR_KSTARS = set(range(10)) # kstar 0–9: non-degenerate stars + +def is_bh_star(bpp): + """Identify binaries that are in a bound BH + normal-star phase for ≥ 100 Myr. + + Parameters + ---------- + bpp : pandas.DataFrame + COSMIC binary population parameters output. + + Returns + ------- + n_hits : int + Number of distinct binaries satisfying the criterion. + hit_bin_nums : numpy.ndarray + Integer bin_num values of those binaries. + """ + bh_star_mask = ( + ( (bpp.kstar_1 == 14) & bpp.kstar_2.isin(_STAR_KSTARS)) + | (bpp.kstar_1.isin(_STAR_KSTARS) & (bpp.kstar_2 == 14)) + ) & (bpp.sep > 0) + + bh_star_rows = bpp.loc[bh_star_mask] + if len(bh_star_rows) == 0: + return 0, np.array([], dtype=int) + + # Group by bin_num (the natural key) so that min/max tphys are aligned + # on the same index. drop_duplicates would give rows with different + # integer-row indices that pandas would NOT align correctly on subtraction. + phase_start = bh_star_rows.groupby('bin_num')['tphys'].min() # Series indexed by bin_num + phase_end = bh_star_rows.groupby('bin_num')['tphys'].max() # Series indexed by bin_num + duration = phase_end - phase_start # tphys in Myr + + long_enough_bin_nums = duration[duration >= 100.0].index.values + return len(long_enough_bin_nums), long_enough_bin_nums + +# ------------------------------------------------------------------ +# Helper: run one sampler and return (result, elapsed_seconds) +# ------------------------------------------------------------------ +def run_sampler(mc_only, seed): + label = "Monte Carlo" if mc_only else "STROOPWAFEL" + print(f"\n{'='*60}") + print(f" Running {label} (seed={seed})") + print(f"{'='*60}") + + sw = AdaptiveSampler( + parameter_space=params, + total_systems=args.num_systems, + batch_size=args.batch_size, + BSEDict=BSEDict, + compute_derived=compute_derived, + reject_systems=default_reject, + is_interesting=is_bh_star, + output_path=os.path.join(args.output_dir, 'mc' if mc_only else 'sw'), + nproc=args.num_cores, + n_generations=args.n_generations, + mc_only=mc_only, + seed=seed, + ) + + t0 = time.time() + result = sw.run() + elapsed = time.time() - t0 + + return result, elapsed + +# ------------------------------------------------------------------ +# Helper: summarise one result +# ------------------------------------------------------------------ +def print_summary(label, result, elapsed): + raw_hits = int(np.sum(result.is_hit)) + hit_rate = result.hit_rate + uncertainty = result.hit_rate_uncertainty + + print(f"\n--- {label} summary ---") + print(f" Systems evolved : {len(result.weights):,}") + print(f" Raw hits found : {raw_hits:,}") + print(f" Weighted hit rate: {hit_rate:.6e} ± {uncertainty:.6e}") + print(f" Wall-clock time : {elapsed:.1f} s") + +# ------------------------------------------------------------------ +# Helper: side-by-side comparison +# ------------------------------------------------------------------ +def print_comparison(mc_result, mc_elapsed, sw_result, sw_elapsed): + mc_rate = mc_result.hit_rate + mc_unc = mc_result.hit_rate_uncertainty + sw_rate = sw_result.hit_rate + sw_unc = sw_result.hit_rate_uncertainty + + print(f"\n{'='*60}") + print(" Efficiency comparison") + print(f"{'='*60}") + print(f"{'':30s} {'Monte Carlo':>15s} {'STROOPWAFEL':>15s}") + print(f" {'Systems evolved':<28s} {len(mc_result.weights):>15,} {len(sw_result.weights):>15,}") + print(f" {'Raw hits found':<28s} {int(np.sum(mc_result.is_hit)):>15,} {int(np.sum(sw_result.is_hit)):>15,}") + print(f" {'Weighted hit rate':<28s} {mc_rate:>15.4e} {sw_rate:>15.4e}") + print(f" {'Uncertainty (1σ)':<28s} {mc_unc:>15.4e} {sw_unc:>15.4e}") + print(f" {'Wall-clock time (s)':<28s} {mc_elapsed:>15.1f} {sw_elapsed:>15.1f}") + + if sw_unc > 0 and mc_unc > 0: + # How many MC systems would give the same precision as SW? + equivalent_mc = len(mc_result.weights) * (mc_unc / sw_unc) ** 2 + speedup = equivalent_mc / len(sw_result.weights) + print(f"\n STROOPWAFEL is ~{speedup:.1f}x more statistically efficient:") + print(f" Monte Carlo would need ~{equivalent_mc:,.0f} systems to match") + print(f" STROOPWAFEL's precision on {len(sw_result.weights):,} systems.") + +# ------------------------------------------------------------------ +# Main +# ------------------------------------------------------------------ +if __name__ == '__main__': + os.makedirs(args.output_dir, exist_ok=True) + + mc_result = mc_elapsed = None + sw_result = sw_elapsed = None + + if not args.sw_only: + mc_result, mc_elapsed = run_sampler(mc_only=True, seed=args.seed) + print_summary("Monte Carlo", mc_result, mc_elapsed) + swio.save_result(os.path.join(args.output_dir, 'mc_result.h5'), mc_result) + + if not args.mc_only: + sw_result, sw_elapsed = run_sampler(mc_only=False, seed=args.seed + 1) + print_summary("STROOPWAFEL", sw_result, sw_elapsed) + swio.save_result(os.path.join(args.output_dir, 'sw_result.h5'), sw_result) + + if mc_result is not None and sw_result is not None: + print_comparison(mc_result, mc_elapsed, sw_result, sw_elapsed) diff --git a/src/cosmic/sample/stroopwafel/priors.py b/src/cosmic/sample/stroopwafel/priors.py index 8d312a730..1305d05f3 100644 --- a/src/cosmic/sample/stroopwafel/priors.py +++ b/src/cosmic/sample/stroopwafel/priors.py @@ -5,7 +5,8 @@ returns an (N,) ndarray of prior probability densities. """ import numpy as np -from .constants import ALPHA_IMF, SANA_G, SANA_ECC +from scipy.stats import norm as _scipy_norm +from .constants import ALPHA_IMF, SANA_G, SANA_ECC, NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA def uniform(values, lo, hi): @@ -159,6 +160,35 @@ def uniform_in_cosine(values, lo, hi): return np.full(len(values), 1.0 / (hi - lo)) +def log_normal(values, lo, hi): + """Normal PDF in ln-space for a (truncated) log-normal natal kick prior. + + The kick magnitude v [km/s] follows LogNormal(``NATAL_KICK_LOG_MU``, + ``NATAL_KICK_LOG_SIGMA``), so the sampling-space variable ``x = ln(v)`` + has a (truncated) normal distribution. The returned density is the + normal PDF evaluated at ``values``, normalised so that it integrates to 1 + over the sampling-space interval ``[lo, hi]``. + + Parameters + ---------- + values : `numpy.ndarray` + (N,) array of sample values in ``ln(v / km s⁻¹)`` space. + lo : `float` + Lower bound in ``ln(v / km s⁻¹)`` space. + hi : `float` + Upper bound in ``ln(v / km s⁻¹)`` space. + + Returns + ------- + `numpy.ndarray` + (N,) array of prior probability densities. + """ + p_lo = _scipy_norm.cdf(lo, loc=NATAL_KICK_LOG_MU, scale=NATAL_KICK_LOG_SIGMA) + p_hi = _scipy_norm.cdf(hi, loc=NATAL_KICK_LOG_MU, scale=NATAL_KICK_LOG_SIGMA) + norm_factor = p_hi - p_lo # probability mass within bounds + return _scipy_norm.pdf(values, loc=NATAL_KICK_LOG_MU, scale=NATAL_KICK_LOG_SIGMA) / norm_factor + + # Registry mapping string names to functions PRIORS = { 'uniform': uniform, @@ -168,4 +198,5 @@ def uniform_in_cosine(values, lo, hi): 'sana_ecc': sana_ecc, 'uniform_in_sine': uniform_in_sine, 'uniform_in_cosine': uniform_in_cosine, + 'log_normal': log_normal, } diff --git a/src/cosmic/sample/stroopwafel/samplers.py b/src/cosmic/sample/stroopwafel/samplers.py index d42a8df39..837a779bf 100644 --- a/src/cosmic/sample/stroopwafel/samplers.py +++ b/src/cosmic/sample/stroopwafel/samplers.py @@ -4,7 +4,8 @@ sampling/transformed space) and returns an (N,) ndarray. """ import numpy as np -from .constants import ALPHA_IMF, SANA_G, SANA_ECC +from scipy.stats import norm as _scipy_norm +from .constants import ALPHA_IMF, SANA_G, SANA_ECC, NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA def uniform(n, lo, hi, rng=None): @@ -179,6 +180,46 @@ def uniform_in_cosine(n, lo, hi, rng=None): return rng.uniform(lo, hi, n) +def log_normal(n, lo, hi, rng=None): + """Inverse-CDF sampling from a (truncated) log-normal natal kick distribution. + + The kick magnitude v follows LogNormal(``NATAL_KICK_LOG_MU``, + ``NATAL_KICK_LOG_SIGMA``), so ``ln(v)`` is normally distributed. + ``lo`` and ``hi`` are bounds in ``ln(v)`` space (i.e. the natural log of + the physical bounds in km/s), and sampling is restricted to that range + via the truncated-normal inverse CDF. + + With ``NATAL_KICK_LOG_MU = 5.67`` and ``NATAL_KICK_LOG_SIGMA = 0.59``, + the median kick is ``exp(5.67) ≈ 291 km/s`` (Hobbs et al. 2005 / Fryer + et al.). Physical bounds of ``[0.1, 5000] km/s`` map to + ``lo ≈ –2.30``, ``hi ≈ 8.52`` in sampling space, which captures + essentially all of the probability mass. + + Parameters + ---------- + n : `int` + Number of samples to draw. + lo : `float` + Lower bound in ``ln(v / km s⁻¹)`` space. + hi : `float` + Upper bound in ``ln(v / km s⁻¹)`` space. + rng : `numpy.random.Generator`, optional + Random number generator, by default None + + Returns + ------- + `numpy.ndarray` + (N,) array of samples in ``ln(v)`` space. + """ + rng = rng or np.random.default_rng() + # Truncated-normal inverse CDF: map uniform draws to [CDF(lo), CDF(hi)], + # then apply the inverse normal CDF. + p_lo = _scipy_norm.cdf(lo, loc=NATAL_KICK_LOG_MU, scale=NATAL_KICK_LOG_SIGMA) + p_hi = _scipy_norm.cdf(hi, loc=NATAL_KICK_LOG_MU, scale=NATAL_KICK_LOG_SIGMA) + u = rng.uniform(p_lo, p_hi, n) + return _scipy_norm.ppf(u, loc=NATAL_KICK_LOG_MU, scale=NATAL_KICK_LOG_SIGMA) + + # Registry mapping string names to functions SAMPLERS = { 'uniform': uniform, @@ -188,4 +229,5 @@ def uniform_in_cosine(n, lo, hi, rng=None): 'sana_ecc': sana_ecc, 'uniform_in_sine': uniform_in_sine, 'uniform_in_cosine': uniform_in_cosine, + 'log_normal': log_normal, } diff --git a/src/cosmic/sample/stroopwafel/transforms.py b/src/cosmic/sample/stroopwafel/transforms.py index 2eae01271..8e204b62a 100644 --- a/src/cosmic/sample/stroopwafel/transforms.py +++ b/src/cosmic/sample/stroopwafel/transforms.py @@ -33,6 +33,9 @@ def to_sampling_space(values, sampler_type): # coordinates), so the sampling variable is cos(θ + π/2) = –sin(θ). # The round-trip is exact for θ ∈ [–π/2, π/2]. return np.cos(values + np.pi / 2) + elif sampler_type == 'log_normal': + # Sampling space is ln(v); physical space is v [km/s]. + return np.log(values) return values @@ -59,6 +62,9 @@ def to_physical_space(values, sampler_type): elif sampler_type == 'uniform_in_cosine': # Inverse of cos(θ + π/2): arccos(u) – π/2 return np.arccos(values) - np.pi / 2 + elif sampler_type == 'log_normal': + # Inverse of ln: exp(x) gives v [km/s]. + return np.exp(values) return values @@ -87,6 +93,9 @@ def transform_bounds(lo, hi, sampler_type): return -1.0, 1.0 elif sampler_type == 'uniform_in_cosine': return -1.0, 1.0 + elif sampler_type == 'log_normal': + # Physical bounds [lo, hi] in km/s → sampling bounds in ln(km/s). + return np.log(lo), np.log(hi) # kroupa, uniform: bounds are already in physical space. # sana, sana_ecc: bounds are passed by the caller in the native # sampling space (log10(period) for sana, eccentricity for sana_ecc), From cd2ff0e1a17fad42ed0fe3f724dc6b28f254229d Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 25 May 2026 16:04:28 -0400 Subject: [PATCH 12/45] add kick sampling --- src/cosmic/sample/stroopwafel/engine.py | 45 ++++++++++------ .../sample/stroopwafel/examples/bh_star.py | 53 +++++++++---------- 2 files changed, 53 insertions(+), 45 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/engine.py b/src/cosmic/sample/stroopwafel/engine.py index 6f1465ec8..4079f3226 100644 --- a/src/cosmic/sample/stroopwafel/engine.py +++ b/src/cosmic/sample/stroopwafel/engine.py @@ -402,16 +402,20 @@ def _compute_weights(self): # COSMIC interface # ------------------------------------------------------------------ - # Natal kick parameters that can be sampled in the ParameterSpace and - # injected as per-binary columns in the InitialBinaryTable. When all - # eight are present, ``natal_kick_array`` is stripped from the BSEDict - # so COSMIC reads the per-binary values instead of global defaults. - # The ``randomseed`` columns required by COSMIC are filled with zeros - # (COSMIC ignores them when the kick values are already provided). - _KICK_PARAM_NAMES = ( + # All natal kick columns that COSMIC accepts as per-binary overrides. + # Any subset may appear in the ParameterSpace; columns that are absent + # receive ``_KICK_SENTINEL`` (-100), which tells COSMIC to draw that + # component from its own prescription (kickflag / sigma in BSEDict). + # This lets you sample only the parameters that actually discriminate + # hits (e.g. natal_kick_1 for binary survival) while leaving irrelevant + # dimensions (angles, secondary kick) out of the parameter space + # entirely, avoiding the dimensionality cost they would otherwise impose + # on the mixture model. + _ALL_KICK_COLUMNS = frozenset([ 'natal_kick_1', 'phi_1', 'theta_1', 'mean_anomaly_1', 'natal_kick_2', 'phi_2', 'theta_2', 'mean_anomaly_2', - ) + ]) + _KICK_SENTINEL = -100.0 def _evolve_batch(self, samples_physical, derived): """Evolve a batch of binaries with COSMIC and identify hits. @@ -452,22 +456,29 @@ def _evolve_batch(self, samples_physical, derived): metallicity=derived['metallicity_1'], ) - # If natal kick parameters were sampled, inject them as per-binary - # columns and omit the global natal_kick_array from the BSEDict. - # COSMIC uses per-binary values when all FLATTENED_NATAL_KICK_COLUMNS - # are present in the table AND natal_kick_array is absent from BSEDict. + # If any natal kick columns were sampled, activate the per-binary + # injection path. Each column present in the ParameterSpace gets its + # sampled value; every other kick column gets _KICK_SENTINEL (-100) + # so COSMIC draws that component from its built-in prescription. + # natal_kick_array is stripped from the BSEDict copy so COSMIC reads + # the per-binary columns instead of the global default. param_names_set = set(self.param_space.names) - if all(k in param_names_set for k in self._KICK_PARAM_NAMES): - for col in self._KICK_PARAM_NAMES: - batch_initial[col] = samples_physical[:, idx[col]] - # randomseed columns are required by COSMIC's reshape but unused - # when kick values are provided directly. + sampled_kick_cols = self._ALL_KICK_COLUMNS & param_names_set + + if sampled_kick_cols: + for col in self._ALL_KICK_COLUMNS: + batch_initial[col] = (samples_physical[:, idx[col]] + if col in sampled_kick_cols + else self._KICK_SENTINEL) batch_initial['randomseed_1'] = 0 batch_initial['randomseed_2'] = 0 bse_dict_run = {k: v for k, v in self.bse_dict.items() if k != 'natal_kick_array'} else: bse_dict_run = self.bse_dict + if "natal_kick_array" not in bse_dict_run: + bse_dict_run['natal_kick_array'] = np.array([[-100, -100, -100, -100, 0], + [-100, -100, -100, -100, 0]]) bpp, bcm, initC, kick_info = Evolve.evolve( initialbinarytable=batch_initial, diff --git a/src/cosmic/sample/stroopwafel/examples/bh_star.py b/src/cosmic/sample/stroopwafel/examples/bh_star.py index 06c5da637..639ba3ab0 100644 --- a/src/cosmic/sample/stroopwafel/examples/bh_star.py +++ b/src/cosmic/sample/stroopwafel/examples/bh_star.py @@ -92,41 +92,38 @@ # Parameter space # # Orbital / stellar parameters (5 dimensions) -# mass_1 : primary mass [Msun], Kroupa IMF -# q : mass ratio m2/m1 ∈ [0.01, 1], uniform -# porb : log10(orbital period / days), Sana power law -# bounds 0.15 to 5.5 → periods ~1.4 d to ~316 000 d -# ecc : eccentricity, Sana power law -# metallicity : metallicity, log-uniform +# mass_1 : primary mass [Msun], Kroupa IMF +# q : mass ratio m2/m1 ∈ [0.01, 1], uniform +# porb : log10(orbital period / days), Sana power law +# bounds 0.15 to 5.5 → periods ~1.4 d to ~316 000 d +# ecc : eccentricity, Sana power law +# metallicity : metallicity, log-uniform # -# Natal kick parameters (8 dimensions, 4 per star) -# natal_kick_N : kick speed [km/s], log-normal (mu=5.67, sigma=0.59 in ln-space) -# physical bounds [0.1, 5000] km/s; median ≈ 291 km/s -# phi_N : kick elevation angle [degrees], uniform ∈ [−90, 90] -# theta_N : kick azimuthal angle [degrees], uniform ∈ [0, 360] -# mean_anomaly_N: orbital phase at kick [degrees], uniform ∈ [0, 360] +# Primary natal kick magnitude (1 dimension) +# natal_kick_1 : kick speed [km/s], log-normal (mu=5.67, sigma=0.59 in ln-space) +# physical bounds [0.1, 5000] km/s; median ≈ 291 km/s +# +# Total: 6-dimensional parameter space. +# +# Kick angles (phi_1, theta_1, mean_anomaly_1) and the secondary kick are +# intentionally excluded. Angles have flat hit-probability across their full +# range so they contribute no information to the mixture model while each +# extra dimension widens the Gaussians by N^(1/D_old - 1/D_new). The engine +# fills all omitted kick columns with the -100 sentinel so COSMIC draws those +# components from its own prescription (kickflag=5 / sigma=265 km/s). # -# Total: 13-dimensional parameter space. # Note: ParameterSpace sorts parameters alphabetically, so the internal # column order is fixed and independent of the order given here. # ------------------------------------------------------------------ params = ParameterSpace([ # --- orbital / stellar --- - Parameter('mass_1', 5.0, 150.0, sampler='kroupa', prior='kroupa'), - Parameter('q', 0.01, 1.0, sampler='uniform', prior='uniform'), - Parameter('porb', 0.15, 5.5, sampler='sana', prior='sana'), - Parameter('ecc', 1e-9, 0.99999999, sampler='sana_ecc', prior='sana_ecc'), - Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log', prior='flat_in_log'), - # --- natal kick: star 1 --- - Parameter('natal_kick_1', 0.1, 5000.0, sampler='log_normal', prior='log_normal'), - Parameter('phi_1', -90.0, 90.0, sampler='uniform', prior='uniform'), - Parameter('theta_1', 0.0, 360.0, sampler='uniform', prior='uniform'), - Parameter('mean_anomaly_1', 0.0, 360.0, sampler='uniform', prior='uniform'), - # --- natal kick: star 2 --- - Parameter('natal_kick_2', 0.1, 5000.0, sampler='log_normal', prior='log_normal'), - Parameter('phi_2', -90.0, 90.0, sampler='uniform', prior='uniform'), - Parameter('theta_2', 0.0, 360.0, sampler='uniform', prior='uniform'), - Parameter('mean_anomaly_2', 0.0, 360.0, sampler='uniform', prior='uniform'), + Parameter('mass_1', 5.0, 150.0, sampler='kroupa', prior='kroupa'), + Parameter('q', 0.01, 1.0, sampler='uniform', prior='uniform'), + Parameter('porb', 0.15, 5.5, sampler='sana', prior='sana'), + Parameter('ecc', 1e-9, 0.99999999, sampler='sana_ecc', prior='sana_ecc'), + Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log', prior='flat_in_log'), + # --- primary natal kick magnitude only --- + # Parameter('natal_kick_1', 0.1, 100.0, sampler='log_normal', prior='log_normal'), ]) # ------------------------------------------------------------------ From 15bccb5349643d53502344aac9138310a8b00c21 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 2 Jun 2026 10:41:39 -0400 Subject: [PATCH 13/45] create COSMICStroopOutput --- src/cosmic/output.py | 189 +++++++++++++++++++++- src/cosmic/sample/stroopwafel/__init__.py | 3 +- src/cosmic/sample/stroopwafel/engine.py | 94 +++++++++-- src/cosmic/sample/stroopwafel/io.py | 107 ------------ src/cosmic/sample/stroopwafel/meson.build | 2 - src/cosmic/sample/stroopwafel/result.py | 82 ---------- 6 files changed, 266 insertions(+), 211 deletions(-) delete mode 100644 src/cosmic/sample/stroopwafel/io.py delete mode 100644 src/cosmic/sample/stroopwafel/result.py diff --git a/src/cosmic/output.py b/src/cosmic/output.py index 024e87086..48bb418dc 100644 --- a/src/cosmic/output.py +++ b/src/cosmic/output.py @@ -10,7 +10,7 @@ import warnings -__all__ = ['COSMICOutput', 'COSMICPopOutput', 'save_initC', 'load_initC'] +__all__ = ['COSMICOutput', 'COSMICPopOutput', 'COSMICStroopOutput', 'save_initC', 'load_initC'] kstar_translator = [ @@ -334,6 +334,193 @@ def plot_distribution(self, x_col, y_col=None, c_col=None, when='final', return fig, ax +class COSMICStroopOutput(COSMICOutput): + """Results from a STROOPWAFEL adaptive importance-sampling run. + + Extends `COSMICOutput` with the sampled parameters, importance weights, + hit flags, and STROOPWAFEL bookkeeping arrays. + + The link between the numpy arrays and the COSMIC tables is ``bin_num``: + ``samples[bin_num]``, ``weights[bin_num]``, and ``is_hit[bin_num]`` all + correspond to the row(s) in ``bpp``/``bcm``/``initC``/``kick_info`` with + that ``bin_num``. Bin numbers are assigned sequentially (0-indexed) + across all batches so they can be used directly as array indices. + + Parameters + ---------- + bpp, bcm, initC, kick_info : `pandas.DataFrame` + COSMIC output tables (concatenated across all batches). + samples : `numpy.ndarray` + (N, D) array of sampled parameters in **physical** space. + param_names : `list` of `str` + Parameter names corresponding to the columns of ``samples``. + weights : `numpy.ndarray` + (N,) importance-sampling weights. + is_hit : `numpy.ndarray` + (N,) boolean array; True where the system satisfied the hit criterion. + generation : `numpy.ndarray` + (N,) integer array — 0 = exploration phase, 1+ = refinement generation. + gaussian_idx : `numpy.ndarray` + (N,) integer array — -1 = drawn from prior, k = drawn from Gaussian k. + num_explored : `int` + Number of systems evolved during the exploration phase. + num_hits : `int` + Total raw hit count across all phases. + fraction_explored : `float` + Fraction of total systems used for exploration. + label : `str`, optional + Human-readable label for the run, by default None + """ + + def __init__(self, bpp, bcm, initC, kick_info, + samples, param_names, weights, is_hit, + generation, gaussian_idx, + num_explored, num_hits, fraction_explored, + label=None): + super().__init__(bpp=bpp, bcm=bcm, initC=initC, kick_info=kick_info, label=label) + self.samples = np.asarray(samples) + self.param_names = list(param_names) + self.weights = np.asarray(weights) + self.is_hit = np.asarray(is_hit, dtype=bool) + self.generation = np.asarray(generation, dtype=int) + self.gaussian_idx = np.asarray(gaussian_idx, dtype=int) + self.num_explored = int(num_explored) + self.num_hits = int(num_hits) + self.fraction_explored = float(fraction_explored) + + def __repr__(self): + return ( + f'' + ) + + def __getitem__(self, key): + """Subset by bin_num, keeping all arrays aligned.""" + cosmic_subset = super().__getitem__(key) + bin_nums = cosmic_subset.initC['bin_num'].values + return COSMICStroopOutput( + bpp=cosmic_subset.bpp, + bcm=cosmic_subset.bcm, + initC=cosmic_subset.initC, + kick_info=cosmic_subset.kick_info, + samples=self.samples[bin_nums], + param_names=self.param_names, + weights=self.weights[bin_nums], + is_hit=self.is_hit[bin_nums], + generation=self.generation[bin_nums], + gaussian_idx=self.gaussian_idx[bin_nums], + num_explored=self.num_explored, + num_hits=int(self.is_hit[bin_nums].sum()), + fraction_explored=self.fraction_explored, + label=self.label, + ) + + # ------------------------------------------------------------------ + # Derived statistics + # ------------------------------------------------------------------ + + @property + def hit_rate(self): + """Importance-weighted hit rate: sum(w[is_hit]) / N. + + Returns + ------- + `float` + """ + if len(self.weights) == 0: + return 0.0 + return float(np.sum(self.weights[self.is_hit]) / len(self.weights)) + + @property + def hit_rate_uncertainty(self): + """Standard error on the importance-weighted hit rate. + + Returns ``std(w[is_hit], ddof=1) / sqrt(N)``, or 0.0 if fewer + than two hits are present. + + Returns + ------- + `float` + """ + w_hits = self.weights[self.is_hit] + if len(w_hits) < 2: + return 0.0 + return float(np.std(w_hits, ddof=1) / np.sqrt(len(self.weights))) + + # ------------------------------------------------------------------ + # I/O + # ------------------------------------------------------------------ + + def save(self, output_file): + """Save to an HDF5 file. + + Writes COSMIC tables via the parent ``save``, then appends + STROOPWAFEL arrays to a ``stroopwafel/`` group. + + Parameters + ---------- + output_file : `str` + Filename/path to create or overwrite. + """ + super().save(output_file) + with h5.File(output_file, 'a') as f: + grp = f.require_group('stroopwafel') + for name, arr in [ + ('samples', self.samples), + ('weights', self.weights), + ('is_hit', self.is_hit), + ('generation', self.generation), + ('gaussian_idx', self.gaussian_idx), + ]: + if name in grp: + del grp[name] + grp.create_dataset(name, data=arr) + grp.attrs['param_names'] = self.param_names + grp.attrs['num_explored'] = self.num_explored + grp.attrs['num_hits'] = self.num_hits + grp.attrs['fraction_explored'] = self.fraction_explored + + @classmethod + def from_file(cls, path, label=None): + """Load from an HDF5 file written by :meth:`save`. + + Parameters + ---------- + path : `str` + File path to read. + label : `str`, optional + Override the stored label, by default None + + Returns + ------- + `COSMICStroopOutput` + """ + cosmic = COSMICOutput(file=path, label=label) + with h5.File(path, 'r') as f: + grp = f['stroopwafel'] + samples = grp['samples'][:] + weights = grp['weights'][:] + is_hit = grp['is_hit'][:] + generation = grp['generation'][:] + gaussian_idx = grp['gaussian_idx'][:] + param_names = list(grp.attrs['param_names']) + num_explored = int(grp.attrs['num_explored']) + num_hits = int(grp.attrs['num_hits']) + fraction_explored = float(grp.attrs['fraction_explored']) + return cls( + bpp=cosmic.bpp, bcm=cosmic.bcm, + initC=cosmic.initC, kick_info=cosmic.kick_info, + samples=samples, param_names=param_names, + weights=weights, is_hit=is_hit, + generation=generation, gaussian_idx=gaussian_idx, + num_explored=num_explored, num_hits=num_hits, + fraction_explored=fraction_explored, + label=cosmic.label if label is None else label, + ) + + class COSMICPopOutput(): def __init__(self, file, label=None): # read in convergence tables and totals diff --git a/src/cosmic/sample/stroopwafel/__init__.py b/src/cosmic/sample/stroopwafel/__init__.py index ce5731ed6..e1f94d65a 100755 --- a/src/cosmic/sample/stroopwafel/__init__.py +++ b/src/cosmic/sample/stroopwafel/__init__.py @@ -14,6 +14,5 @@ """ from .engine import AdaptiveSampler from .parameter_space import ParameterSpace, Parameter -from .result import STROOPWAFELResult -__all__ = ['AdaptiveSampler', 'ParameterSpace', 'Parameter', 'STROOPWAFELResult'] +__all__ = ['AdaptiveSampler', 'ParameterSpace', 'Parameter'] diff --git a/src/cosmic/sample/stroopwafel/engine.py b/src/cosmic/sample/stroopwafel/engine.py index 4079f3226..0345690f4 100644 --- a/src/cosmic/sample/stroopwafel/engine.py +++ b/src/cosmic/sample/stroopwafel/engine.py @@ -5,13 +5,14 @@ """ import os import numpy as np +import pandas as pd from scipy.stats import multivariate_normal from cosmic.sample.initialbinarytable import InitialBinaryTable from cosmic.evolve import Evolve +from cosmic.output import COSMICStroopOutput from .mixture_model import GaussianMixture -from .result import STROOPWAFELResult from .constants import MIN_ACTIVE_FRACTION @@ -80,19 +81,26 @@ def __init__(self, parameter_space, total_systems, batch_size, BSEDict, self.prior_fraction_rejected = 0.0 self.mixture = None - # Accumulators for all samples (in sampling space) + # Accumulators for all samples (in sampling space) and COSMIC output. + # One entry per batch; _build_cosmic_output() renumbers bin_nums + # globally so that samples[bin_num] == the system with that bin_num. self._all_samples = [] self._all_is_hit = [] self._all_generation = [] self._all_gaussian_idx = [] + self._all_bpp = [] + self._all_bcm = [] + self._all_initC = [] + self._all_kick_info = [] def run(self): """Run the full STROOPWAFEL pipeline. Returns ------- - `STROOPWAFELResult` - Container holding all samples, weights, and metadata. + `COSMICStroopOutput` + Container holding all samples, weights, COSMIC output tables, + and associated metadata. """ os.makedirs(self.output_path, exist_ok=True) @@ -143,7 +151,7 @@ def _explore(self): batch_derived = {k: v[selected] for k, v in derived.items()} # Evolve with COSMIC - n_hits, hit_bin_nums, bpp, initC, kick_info = self._evolve_batch( + n_hits, hit_bin_nums, bpp, bcm, initC, kick_info = self._evolve_batch( batch_samples_phys, batch_derived ) @@ -151,11 +159,15 @@ def _explore(self): is_hit = np.zeros(len(selected), dtype=bool) is_hit[hit_bin_nums] = True - # Accumulate + # Accumulate samples and COSMIC output self._all_samples.append(batch_samples) self._all_is_hit.append(is_hit) self._all_generation.append(np.zeros(len(selected), dtype=int)) self._all_gaussian_idx.append(np.full(len(selected), -1, dtype=int)) + self._all_bpp.append(bpp) + self._all_bcm.append(bcm) + self._all_initC.append(initC) + self._all_kick_info.append(kick_info) self.num_hits += n_hits self.finished += len(selected) @@ -283,18 +295,22 @@ def _refine(self): batch_derived = {k: v[indices[:n_take]] for k, v in derived.items()} # Evolve with COSMIC - n_hits, hit_bin_nums, bpp, initC, kick_info = self._evolve_batch( + n_hits, hit_bin_nums, bpp, bcm, initC, kick_info = self._evolve_batch( batch_phys, batch_derived ) is_hit = np.zeros(n_take, dtype=bool) is_hit[hit_bin_nums] = True - # Accumulate + # Accumulate samples and COSMIC output self._all_samples.append(batch_samples) self._all_is_hit.append(is_hit) self._all_generation.append(np.full(n_take, gen + 1, dtype=int)) self._all_gaussian_idx.append(batch_gauss_idx) + self._all_bpp.append(bpp) + self._all_bcm.append(bcm) + self._all_initC.append(initC) + self._all_kick_info.append(kick_info) gen_samples_list.append(batch_samples) gen_is_hit_list.append(is_hit) @@ -338,18 +354,18 @@ def _refine(self): # Weight calculation # ------------------------------------------------------------------ def _compute_weights(self): - """Compute importance sampling weights for all samples. + """Compute importance sampling weights and assemble the final result. - Builds the final ``STROOPWAFELResult`` using the formula - ``w(x) = π(x) / Q(x)`` where + Uses the formula ``w(x) = π(x) / Q(x)`` where ``Q = f_e·π + (1 − f_e)·q`` is a mixture of the prior and the Gaussian proposal. Returns ------- - `STROOPWAFELResult` - Container holding all samples, importance weights, and - associated metadata. + `COSMICStroopOutput` + All samples, importance weights, COSMIC output tables, and + associated metadata. ``samples[bin_num]`` corresponds to + the COSMIC table row(s) with that ``bin_num``. """ all_samples = np.vstack(self._all_samples) all_is_hit = np.concatenate(self._all_is_hit) @@ -379,8 +395,11 @@ def _compute_weights(self): weights = pi / den - # Build result - result = STROOPWAFELResult( + # Concatenate COSMIC output with globally unique bin_nums + bpp, bcm, initC, kick_info = self._build_cosmic_output() + + result = COSMICStroopOutput( + bpp=bpp, bcm=bcm, initC=initC, kick_info=kick_info, samples=self.param_space.to_physical(all_samples), param_names=self.param_space.names, weights=weights, @@ -398,6 +417,45 @@ def _compute_weights(self): return result + def _build_cosmic_output(self): + """Concatenate per-batch COSMIC tables with globally unique bin_nums. + + Batch k receives bin_nums ``offset_k .. offset_k + N_k − 1`` where + ``offset_k = N_0 + … + N_{k-1}``. This ensures + ``samples[bin_num]`` and ``is_hit[bin_num]`` directly index the + system with that bin_num in the concatenated COSMIC tables. + + Returns + ------- + bpp, bcm, initC, kick_info : `pandas.DataFrame` + """ + offset = 0 + bpp_list, bcm_list, initC_list, kick_info_list = [], [], [], [] + for batch_samples, bpp, bcm, initC, kick_info in zip( + self._all_samples, + self._all_bpp, self._all_bcm, self._all_initC, self._all_kick_info, + ): + n = len(batch_samples) + bpp_list.append(bpp.assign(bin_num=bpp['bin_num'] + offset)) + bcm_list.append(bcm.assign(bin_num=bcm['bin_num'] + offset)) + initC_list.append(initC.assign(bin_num=initC['bin_num'] + offset)) + kick_info_list.append(kick_info.assign(bin_num=kick_info['bin_num'] + offset)) + offset += n + def _concat(frames): + # Concatenate and set bin_num as the index (drop=False keeps it + # as a column too, so filtering via df['bin_num'].isin(...) still + # works). For bpp/bcm the index is non-unique (multiple timestep + # rows per system); for initC/kick_info it is unique. + return (pd.concat(frames, ignore_index=True) + .set_index('bin_num', drop=False)) + + return ( + _concat(bpp_list), + _concat(bcm_list), + _concat(initC_list), + _concat(kick_info_list), + ) + # ------------------------------------------------------------------ # COSMIC interface # ------------------------------------------------------------------ @@ -437,6 +495,8 @@ def _evolve_batch(self, samples_physical, derived): are hits. bpp : `pandas.DataFrame` COSMIC binary population parameters output. + bcm : `pandas.DataFrame` + COSMIC user-timestep output. initC : `pandas.DataFrame` COSMIC initial conditions output. kick_info : `pandas.DataFrame` @@ -489,7 +549,7 @@ def _evolve_batch(self, samples_physical, derived): # Apply user's hit identification n_hits, hit_bin_nums = self.is_interesting_fn(bpp) - return n_hits, hit_bin_nums, bpp, initC, kick_info + return n_hits, hit_bin_nums, bpp, bcm, initC, kick_info # ------------------------------------------------------------------ # Utilities diff --git a/src/cosmic/sample/stroopwafel/io.py b/src/cosmic/sample/stroopwafel/io.py deleted file mode 100644 index ebe79c9df..000000000 --- a/src/cosmic/sample/stroopwafel/io.py +++ /dev/null @@ -1,107 +0,0 @@ -"""HDF5-based I/O for STROOPWAFEL state and results. - -Replaces the old CSV-based print_samples/read_samples with array-native storage. -""" -import pandas as pd -import h5py - - -def save_result(path, result): - """Save a `STROOPWAFELResult` to an HDF5 file. - - Parameters - ---------- - path : `str` - Output file path (created or overwritten). - result : `STROOPWAFELResult` - Result container to serialise. - """ - with h5py.File(path, 'w') as f: - f.create_dataset('samples', data=result.samples) - f.create_dataset('weights', data=result.weights) - f.create_dataset('is_hit', data=result.is_hit) - f.create_dataset('generation', data=result.generation) - f.create_dataset('gaussian_idx', data=result.gaussian_idx) - f.attrs['param_names'] = result.param_names - f.attrs['num_explored'] = result.num_explored - f.attrs['num_hits'] = result.num_hits - f.attrs['fraction_explored'] = result.fraction_explored - - -def save_cosmic_output(path, bpp_frames, initC_frames, kick_info_frames): - """Append COSMIC output DataFrames to an HDF5 file. - - Parameters - ---------- - path : `str` - Output file path (appended to if it already exists). - bpp_frames : `list` of `pandas.DataFrame` - Binary population parameter tables from each batch. - initC_frames : `list` of `pandas.DataFrame` - Initial conditions tables from each batch. - kick_info_frames : `list` of `pandas.DataFrame` - Natal kick information tables from each batch. - """ - if bpp_frames: - full_bpp = pd.concat(bpp_frames, ignore_index=True) - full_bpp.to_hdf(path, key='bpp', mode='a', format='table') - if initC_frames: - full_initC = pd.concat(initC_frames, ignore_index=True) - full_initC.to_hdf(path, key='initC', mode='a', format='table') - if kick_info_frames: - full_kicks = pd.concat(kick_info_frames, ignore_index=True) - full_kicks.to_hdf(path, key='kick_info', mode='a', format='table') - - -def save_mixture(path, mixture, generation): - """Save `GaussianMixture` state to an HDF5 file. - - Parameters - ---------- - path : `str` - Output file path (appended to if it already exists). - mixture : `GaussianMixture` - Mixture model to serialise. - generation : `int` - Refinement generation number, used as the HDF5 group key. - """ - with h5py.File(path, 'a') as f: - grp_name = f'mixture/gen_{generation}' - if grp_name in f: - del f[grp_name] - grp = f.create_group(grp_name) - grp.create_dataset('means', data=mixture.means) - grp.create_dataset('covariances', data=mixture.covariances) - grp.create_dataset('alphas', data=mixture.alphas) - grp.attrs['rejection_rate'] = mixture.rejection_rate - - -def load_mixture(path, generation): - """Load `GaussianMixture` state from an HDF5 file. - - Parameters - ---------- - path : `str` - File path to read from. - generation : `int` - Refinement generation number to load. - - Returns - ------- - `GaussianMixture` or None - The stored mixture model, or None if the requested generation - is not found or the file does not exist. - """ - from .mixture_model import GaussianMixture - - try: - with h5py.File(path, 'r') as f: - grp = f[f'mixture/gen_{generation}'] - return GaussianMixture( - means=grp['means'][:], - covariances=grp['covariances'][:], - alphas=grp['alphas'][:], - rejection_rate=grp.attrs['rejection_rate'], - ) - except (KeyError, FileNotFoundError): - return None diff --git a/src/cosmic/sample/stroopwafel/meson.build b/src/cosmic/sample/stroopwafel/meson.build index 87bd56fdf..f99e1caa2 100644 --- a/src/cosmic/sample/stroopwafel/meson.build +++ b/src/cosmic/sample/stroopwafel/meson.build @@ -2,13 +2,11 @@ python_sources = [ '__init__.py', 'constants.py', 'engine.py', - 'io.py', 'mixture_model.py', 'parameter_space.py', 'presets.py', 'priors.py', 'rejection.py', - 'result.py', 'samplers.py', 'transforms.py', ] diff --git a/src/cosmic/sample/stroopwafel/result.py b/src/cosmic/sample/stroopwafel/result.py deleted file mode 100644 index f933cdaf0..000000000 --- a/src/cosmic/sample/stroopwafel/result.py +++ /dev/null @@ -1,82 +0,0 @@ -"""Container for STROOPWAFEL adaptive sampling results.""" -import numpy as np - - -class STROOPWAFELResult: - """Results from an AdaptiveSampler run. - - Parameters - ---------- - samples : `numpy.ndarray` - (N, D) array of all samples in physical space. - param_names : `list` of `str` - Parameter names defining the column ordering. - weights : `numpy.ndarray` - (N,) array of importance sampling weights. - is_hit : `numpy.ndarray` - (N,) boolean array indicating which samples are hits. - generation : `numpy.ndarray` - (N,) integer array (0 = exploration, 1+ = refinement). - gaussian_idx : `numpy.ndarray` - (N,) integer array (-1 = exploration, k = from Gaussian - component k). - num_explored : `int` - Number of systems simulated during the exploration phase. - num_hits : `int` - Total number of hits found across all phases. - fraction_explored : `float` - Adaptive exploration fraction. - bpp_frames : `list` of `pandas.DataFrame`, optional - COSMIC binary population parameter tables, by default None - initC_frames : `list` of `pandas.DataFrame`, optional - COSMIC initial conditions tables, by default None - kick_info_frames : `list` of `pandas.DataFrame`, optional - COSMIC natal kick information tables, by default None - """ - - def __init__(self, samples, param_names, weights, is_hit, generation, - gaussian_idx, num_explored, num_hits, fraction_explored, - bpp_frames=None, initC_frames=None, kick_info_frames=None): - self.samples = samples - self.param_names = param_names - self.weights = weights - self.is_hit = is_hit - self.generation = generation - self.gaussian_idx = gaussian_idx - self.num_explored = num_explored - self.num_hits = num_hits - self.fraction_explored = fraction_explored - self.bpp_frames = bpp_frames or [] - self.initC_frames = initC_frames or [] - self.kick_info_frames = kick_info_frames or [] - - @property - def hit_rate(self): - """Importance-weighted hit rate (sum of weights over hits / N). - - Returns - ------- - `float` - Weighted fraction of systems that are hits, or 0.0 if - weights or hit flags are not available. - """ - if self.weights is None or self.is_hit is None: - return 0.0 - return np.sum(self.weights[self.is_hit]) / len(self.weights) - - @property - def hit_rate_uncertainty(self): - """Standard error on the importance-weighted hit rate. - - Returns - ------- - `float` - Standard error ``std(w_hits) / sqrt(N)``, or 0.0 if fewer - than two hits are present. - """ - if self.weights is None or self.is_hit is None: - return 0.0 - w_hits = self.weights[self.is_hit] - if len(w_hits) < 2: - return 0.0 - return np.std(w_hits, ddof=1) / np.sqrt(len(self.weights)) From 92c12e816d710279c579b353470445cbf50f6f8f Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 2 Jun 2026 10:41:46 -0400 Subject: [PATCH 14/45] account for output in tests --- .../stroopwafel/examples/test_new_modules.py | 23 ++++++++++++------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/examples/test_new_modules.py b/src/cosmic/sample/stroopwafel/examples/test_new_modules.py index 448b25d12..42420bfe6 100644 --- a/src/cosmic/sample/stroopwafel/examples/test_new_modules.py +++ b/src/cosmic/sample/stroopwafel/examples/test_new_modules.py @@ -16,7 +16,7 @@ from stroopwafel.transforms import to_sampling_space, to_physical_space, transform_bounds from stroopwafel.mixture_model import GaussianMixture from stroopwafel.rejection import get_zams_radius, calculate_roche_lobe_radius, default_reject -from stroopwafel.result import STROOPWAFELResult +from cosmic.output import COSMICStroopOutput from stroopwafel.constants import ( R_COEFF, ZSOL, R_SOL_TO_AU, ALPHA_IMF, SANA_G, SANA_ECC ) @@ -216,17 +216,24 @@ def test_mixture_sample(): # ==================================================================== def test_result_hit_rate(): - weights = np.ones(100) - is_hit = np.zeros(100, dtype=bool) + import pandas as pd + N = 100 + weights = np.ones(N) + is_hit = np.zeros(N, dtype=bool) is_hit[:10] = True - result = STROOPWAFELResult( - samples=np.zeros((100, 1)), + bin_nums = np.arange(N) + result = COSMICStroopOutput( + bpp=pd.DataFrame({'bin_num': bin_nums}), + bcm=pd.DataFrame({'bin_num': bin_nums}), + initC=pd.DataFrame({'bin_num': bin_nums}), + kick_info=pd.DataFrame({'bin_num': bin_nums}), + samples=np.zeros((N, 1)), param_names=['mass_1'], weights=weights, is_hit=is_hit, - generation=np.zeros(100, dtype=int), - gaussian_idx=np.full(100, -1, dtype=int), - num_explored=100, + generation=np.zeros(N, dtype=int), + gaussian_idx=np.full(N, -1, dtype=int), + num_explored=N, num_hits=10, fraction_explored=1.0, ) From 2d12d5d9c0cdc26e2de409b92dee2f0218ee9d04 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 2 Jun 2026 10:42:16 -0400 Subject: [PATCH 15/45] bh star working with new stuff --- src/cosmic/sample/stroopwafel/examples/bh_star.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/examples/bh_star.py b/src/cosmic/sample/stroopwafel/examples/bh_star.py index 639ba3ab0..fc9314fbc 100644 --- a/src/cosmic/sample/stroopwafel/examples/bh_star.py +++ b/src/cosmic/sample/stroopwafel/examples/bh_star.py @@ -61,7 +61,7 @@ # BSE physics # ------------------------------------------------------------------ BSEDict = { - 'xi': 1.0, 'bhflag': 1, 'neta': 0.5, 'windflag': 3, 'wdflag': 1, + 'xi': 1.0, 'bhflag': 4, 'neta': 0.5, 'windflag': 3, 'wdflag': 1, 'alpha1': [1.0, 1.0], 'pts1': 0.001, 'pts3': 0.02, 'pts2': 0.01, 'epsnov': 0.001, 'hewind': 0.5, 'ck': 1000, 'bwind': 0.0, 'lambdaf': 0.0, 'mxns': 3.0, 'beta': -1.0, 'tflag': 1, 'acc2': 1.5, @@ -123,7 +123,7 @@ Parameter('ecc', 1e-9, 0.99999999, sampler='sana_ecc', prior='sana_ecc'), Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log', prior='flat_in_log'), # --- primary natal kick magnitude only --- - # Parameter('natal_kick_1', 0.1, 100.0, sampler='log_normal', prior='log_normal'), + Parameter('natal_kick_1', 0.1, 100.0, sampler='log_normal', prior='log_normal'), ]) # ------------------------------------------------------------------ @@ -285,12 +285,12 @@ def print_comparison(mc_result, mc_elapsed, sw_result, sw_elapsed): if not args.sw_only: mc_result, mc_elapsed = run_sampler(mc_only=True, seed=args.seed) print_summary("Monte Carlo", mc_result, mc_elapsed) - swio.save_result(os.path.join(args.output_dir, 'mc_result.h5'), mc_result) + mc_result.save(os.path.join(args.output_dir, 'mc_result.h5')) if not args.mc_only: sw_result, sw_elapsed = run_sampler(mc_only=False, seed=args.seed + 1) print_summary("STROOPWAFEL", sw_result, sw_elapsed) - swio.save_result(os.path.join(args.output_dir, 'sw_result.h5'), sw_result) + sw_result.save(os.path.join(args.output_dir, 'sw_result.h5')) if mc_result is not None and sw_result is not None: print_comparison(mc_result, mc_elapsed, sw_result, sw_elapsed) From f29e1983fbc67add21eb1e85bcc46fde7f6e6e47 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 2 Jun 2026 10:46:54 -0400 Subject: [PATCH 16/45] we already have a roche radius function --- .../stroopwafel/examples/test_new_modules.py | 53 +------------------ src/cosmic/sample/stroopwafel/rejection.py | 28 ++-------- 2 files changed, 4 insertions(+), 77 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/examples/test_new_modules.py b/src/cosmic/sample/stroopwafel/examples/test_new_modules.py index 42420bfe6..ac7406afc 100644 --- a/src/cosmic/sample/stroopwafel/examples/test_new_modules.py +++ b/src/cosmic/sample/stroopwafel/examples/test_new_modules.py @@ -15,7 +15,7 @@ from stroopwafel.priors import PRIORS from stroopwafel.transforms import to_sampling_space, to_physical_space, transform_bounds from stroopwafel.mixture_model import GaussianMixture -from stroopwafel.rejection import get_zams_radius, calculate_roche_lobe_radius, default_reject +from stroopwafel.rejection import get_zams_radius, default_reject from cosmic.output import COSMICStroopOutput from stroopwafel.constants import ( R_COEFF, ZSOL, R_SOL_TO_AU, ALPHA_IMF, SANA_G, SANA_ECC @@ -77,57 +77,6 @@ def test_in_bounds(): np.testing.assert_array_equal(mask, mask2) -# ==================================================================== -# ZAMS radius / Roche lobe tests (vs old scalar implementation) -# ==================================================================== - -def old_get_zams_radius(mass, metallicity): - """Old scalar implementation for comparison.""" - metallicity_xi = math.log10(metallicity / ZSOL) - rc = [] - for coeff in R_COEFF: - value = 1; total = 0 - for series in coeff: - total += series * value - value *= metallicity_xi - rc.append(total) - top = (rc[0] * pow(mass, 2.5) + rc[1] * pow(mass, 6.5) - + rc[2] * pow(mass, 11) + rc[3] * pow(mass, 19) - + rc[4] * pow(mass, 19.5)) - bottom = (rc[5] + rc[6] * pow(mass, 2) + rc[7] * pow(mass, 8.5) - + pow(mass, 18.5) + rc[8] * pow(mass, 19.5)) - return (top / bottom) * R_SOL_TO_AU - - -def old_roche(m1, m2): - q = m1 / m2 - return 0.49 / (0.6 + pow(q, -2.0 / 3.0) * math.log(1.0 + pow(q, 1.0 / 3.0))) - - -def test_zams_radius_matches_old(): - masses = np.linspace(1, 100, 50) - mets = np.full(50, 0.014) - new = get_zams_radius(masses, mets) - old = np.array([old_get_zams_radius(m, z) for m, z in zip(masses, mets)]) - np.testing.assert_allclose(new, old, atol=1e-12) - - -def test_zams_radius_multiple_metallicities(): - masses = np.array([1.0, 10.0, 50.0]) - mets = np.array([0.001, 0.014, 0.03]) - new = get_zams_radius(masses, mets) - old = np.array([old_get_zams_radius(m, z) for m, z in zip(masses, mets)]) - np.testing.assert_allclose(new, old, atol=1e-12) - - -def test_roche_lobe_matches_old(): - m1 = np.array([5.0, 10.0, 20.0, 50.0, 100.0]) - m2 = np.array([3.0, 8.0, 10.0, 25.0, 50.0]) - new = calculate_roche_lobe_radius(m1, m2) - old = np.array([old_roche(a, b) for a, b in zip(m1, m2)]) - np.testing.assert_allclose(new, old, atol=1e-12) - - # ==================================================================== # Default rejection function tests # ==================================================================== diff --git a/src/cosmic/sample/stroopwafel/rejection.py b/src/cosmic/sample/stroopwafel/rejection.py index 95e06cd7f..e9acc7f06 100644 --- a/src/cosmic/sample/stroopwafel/rejection.py +++ b/src/cosmic/sample/stroopwafel/rejection.py @@ -6,6 +6,7 @@ """ import numpy as np from .constants import R_COEFF, ZSOL, R_SOL_TO_AU +from cosmic.utils import calc_Roche_radius def get_zams_radius(mass, metallicity): @@ -44,29 +45,6 @@ def get_zams_radius(mass, metallicity): return (top / bottom) * R_SOL_TO_AU -def calculate_roche_lobe_radius(mass1, mass2): - """Compute Roche lobe radius using the Eggleton (1983) approximation. - - Parameters - ---------- - mass1 : `numpy.ndarray` - (N,) array of masses of the star filling its Roche lobe. - mass2 : `numpy.ndarray` - (N,) array of companion masses. - - Returns - ------- - `numpy.ndarray` - (N,) array of Roche lobe radii (dimensionless, in units of - orbital separation). - """ - mass1 = np.asarray(mass1, dtype=float) - mass2 = np.asarray(mass2, dtype=float) - q = mass1 / mass2 - q_cbrt = np.power(q, 1.0 / 3.0) - return 0.49 / (0.6 + np.power(q, -2.0 / 3.0) * np.log(1.0 + q_cbrt)) - - def default_reject(samples_physical, derived, param_names, min_secondary_mass=0.08): """Default rejection function for DCO progenitor systems. @@ -108,8 +86,8 @@ def default_reject(samples_physical, derived, param_names, # Roche lobe radii at periastron peri_sep = separation * (1 - ecc) - rl_1 = peri_sep * calculate_roche_lobe_radius(mass_1, mass_2) - rl_2 = peri_sep * calculate_roche_lobe_radius(mass_2, mass_1) + rl_1 = calc_Roche_radius(mass_1, mass_2, peri_sep) + rl_2 = calc_Roche_radius(mass_2, mass_1, peri_sep) roche_tracker_1 = radius_1 / rl_1 roche_tracker_2 = radius_2 / rl_2 From 79178a255b2e9274d14a0bf4cea4b83e6632c425 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 2 Jun 2026 11:16:53 -0400 Subject: [PATCH 17/45] clearer error about legwork --- src/cosmic/sample/stroopwafel/presets.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/cosmic/sample/stroopwafel/presets.py b/src/cosmic/sample/stroopwafel/presets.py index dfea320d1..19f80f838 100644 --- a/src/cosmic/sample/stroopwafel/presets.py +++ b/src/cosmic/sample/stroopwafel/presets.py @@ -26,7 +26,12 @@ def merging_dco(kstar_1, kstar_2, max_merge_time=13.7): where ``bpp`` is a `pandas.DataFrame` and ``hit_bin_nums`` is a `numpy.ndarray` of bin_num values. """ - from legwork import evol + # check whether the user has LEGWORK installed, if not tell them they need it for this preset + try: + from legwork import evol + except ImportError: + raise ImportError("The 'merging_dco' preset requires the LEGWORK package. " + "Please install it with 'pip install legwork' to use this preset.") import astropy.units as u k1_set = set(kstar_1) From 8b726a06592abcc7edefbdd8ac525611df9c6070 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 2 Jun 2026 11:18:11 -0400 Subject: [PATCH 18/45] add temporary bhflag 4 --- src/cosmic/data/cosmic-settings.json | 5 +++++ src/cosmic/src/kick.f | 8 +++++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/src/cosmic/data/cosmic-settings.json b/src/cosmic/data/cosmic-settings.json index 6f7936db8..2d458bbc6 100644 --- a/src/cosmic/data/cosmic-settings.json +++ b/src/cosmic/data/cosmic-settings.json @@ -1092,6 +1092,11 @@ { "name": 3, "description": "BH natal kicks are not decreased compared to NS kicks and are drawn from the same Maxwellian distribution with dispersion = sigma set above" + }, + { + "name": 4, + "description": "Same as 1, but is also applied if the kick is provided directly in natal_kick_array", + "version_added": "4.1.0" } ] }, diff --git a/src/cosmic/src/kick.f b/src/cosmic/src/kick.f index 67299594a..386290af9 100644 --- a/src/cosmic/src/kick.f +++ b/src/cosmic/src/kick.f @@ -211,6 +211,12 @@ SUBROUTINE kick_pfahl(kw,m1,m1c,m1n,m2,ecc,sep,jorb,vk,sn,r2, * was passed. if(natal_kick_array(sn,1).ge.0.d0)then vk = natal_kick_array(sn,1) + + if(kw.eq.14.and.bhflag.eq.4)then + fallback = MIN(fallback,1.d0) + vk = MAX((1.d0-fallback)*vk,0.d0) + endif + vk2 = vk*vk * per supplied kick value we mimic a call to random number generator xx = RAN3(idum1) @@ -261,7 +267,7 @@ SUBROUTINE kick_pfahl(kw,m1,m1c,m1n,m2,ecc,sep,jorb,vk,sn,r2, if(kw.eq.14.and.bhflag.eq.0)then vk2 = 0.d0 vk = 0.d0 - elseif(kw.eq.14.and.bhflag.eq.1)then + elseif(kw.eq.14.and.(bhflag.eq.1.or.bhflag.eq.4))then fallback = MIN(fallback,1.d0) vk = MAX((1.d0-fallback)*vk,0.d0) vk2 = vk*vk From 48792258ee41e3407473234cb3ae4a3d9113fba5 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 2 Jun 2026 12:18:09 -0400 Subject: [PATCH 19/45] fix table indices --- src/cosmic/sample/stroopwafel/engine.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/engine.py b/src/cosmic/sample/stroopwafel/engine.py index 0345690f4..684ca9de8 100644 --- a/src/cosmic/sample/stroopwafel/engine.py +++ b/src/cosmic/sample/stroopwafel/engine.py @@ -446,8 +446,9 @@ def _concat(frames): # as a column too, so filtering via df['bin_num'].isin(...) still # works). For bpp/bcm the index is non-unique (multiple timestep # rows per system); for initC/kick_info it is unique. - return (pd.concat(frames, ignore_index=True) - .set_index('bin_num', drop=False)) + df = pd.concat(frames, ignore_index=True) + df.index = df['bin_num'].values + return df return ( _concat(bpp_list), From 4cd1967ba0640593fbe3a4574047c95127a132c7 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 2 Jun 2026 12:18:18 -0400 Subject: [PATCH 20/45] update changelog and version --- changelog.md | 6 ++++++ src/cosmic/_version.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/changelog.md b/changelog.md index 5fa722f81..4d9e9a841 100644 --- a/changelog.md +++ b/changelog.md @@ -1,6 +1,12 @@ # COSMIC Changelog ## Prepend only please! +## 4.2.0 + +- Additions/changes + - Adaptive importance sampling is now available through the STROOPWAFEL algorithm (https://arxiv.org/abs/1905.00910) + - This is accessed through ``cosmic.sample.stroopwafel.AdaptiveSampler`` + ## 4.1.0 - Additions/changes diff --git a/src/cosmic/_version.py b/src/cosmic/_version.py index 703970876..0fd7811c0 100644 --- a/src/cosmic/_version.py +++ b/src/cosmic/_version.py @@ -1 +1 @@ -__version__ = "4.1.0" +__version__ = "4.2.0" From db49faa4c2ba232014dd5db2dfef781cd717218f Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 2 Jun 2026 15:16:50 -0400 Subject: [PATCH 21/45] add checkpointing for slurm splitting and memory saving with only saving hit tables --- src/cosmic/output.py | 195 +++++++++++++++++- src/cosmic/sample/stroopwafel/engine.py | 251 +++++++++++++++++++++++- 2 files changed, 440 insertions(+), 6 deletions(-) diff --git a/src/cosmic/output.py b/src/cosmic/output.py index 48bb418dc..adab8eed4 100644 --- a/src/cosmic/output.py +++ b/src/cosmic/output.py @@ -10,7 +10,8 @@ import warnings -__all__ = ['COSMICOutput', 'COSMICPopOutput', 'COSMICStroopOutput', 'save_initC', 'load_initC'] +__all__ = ['COSMICOutput', 'COSMICPopOutput', 'COSMICStroopOutput', + 'STROOPWAFELCheckpoint', 'save_initC', 'load_initC'] kstar_translator = [ @@ -521,6 +522,198 @@ def from_file(cls, path, label=None): ) +class STROOPWAFELCheckpoint: + """Serialisable snapshot of `AdaptiveSampler` state after exploration. + + Saving this to disk decouples the exploration + adaptation phases from + the (typically more expensive) refinement phase, which is the key + enabler for multi-job SLURM workflows: + + .. code-block:: bash + + # Job 1 - exploration (embarrassingly parallel within the job) + python run_explore.py # writes checkpoint.h5 + + # Job 2 - refinement (can be a larger allocation) + python run_refine.py # reads checkpoint.h5, writes result.h5 + + Parameters + ---------- + mixture : `GaussianMixture` or None + Gaussian mixture fitted to exploration hits. ``None`` if no hits + were found or adaptation has not been run yet. + samples : `numpy.ndarray` + (N, D) array of explored samples in **sampling space** (the + internal transformed space used by the mixture model). Convert + to physical space with ``param_space.to_physical(samples)``. + is_hit, generation, gaussian_idx : `numpy.ndarray` + Per-sample flags and bookkeeping arrays (shapes (N,)). + bpp, bcm, initC, kick_info : `pandas.DataFrame` + COSMIC output from exploration, indexed by globally unique + ``bin_num`` so that ``samples[bin_num]`` gives the corresponding + physical parameters. + param_names : `list` of `str` + Parameter names for the D columns of ``samples``. + num_explored : `int` + Systems evolved during exploration. + num_hits : `int` + Raw hit count from exploration. + num_hits_exploratory : `int` + Same as ``num_hits`` (stored separately for use in ``_refine``). + fraction_explored : `float` + Adaptive fraction of total budget used for exploration. + prior_fraction_rejected : `float` + Estimated fraction of prior samples that fail physical rejection. + total_systems : `int` + Full system budget (exploration + refinement combined). + n_generations : `int` + Number of refinement generations originally requested. + """ + + def __init__(self, mixture, samples, is_hit, generation, gaussian_idx, + bpp, bcm, initC, kick_info, param_names, + num_explored, num_hits, num_hits_exploratory, + fraction_explored, prior_fraction_rejected, + total_systems, n_generations): + self.mixture = mixture + self.samples = np.asarray(samples) + self.is_hit = np.asarray(is_hit, dtype=bool) + self.generation = np.asarray(generation, dtype=int) + self.gaussian_idx = np.asarray(gaussian_idx, dtype=int) + self.bpp = bpp + self.bcm = bcm + self.initC = initC + self.kick_info = kick_info + self.param_names = list(param_names) + self.num_explored = int(num_explored) + self.num_hits = int(num_hits) + self.num_hits_exploratory = int(num_hits_exploratory) + self.fraction_explored = float(fraction_explored) + self.prior_fraction_rejected = float(prior_fraction_rejected) + self.total_systems = int(total_systems) + self.n_generations = int(n_generations) + + def __repr__(self): + adapted = self.mixture is not None + return ( + f'' + ) + + # ------------------------------------------------------------------ + # I/O + # ------------------------------------------------------------------ + + def save(self, path): + """Save to an HDF5 file. + + Parameters + ---------- + path : `str` + File path to create or overwrite. + """ + # COSMIC tables — use the same helpers as COSMICOutput so the file + # is readable by ordinary COSMIC tooling if needed. + self.bpp.reset_index(drop=True).to_hdf(path, key='bpp', mode='w') + self.bcm.reset_index(drop=True).to_hdf(path, key='bcm', mode='a') + save_initC(path, self.initC.reset_index(drop=True), + key='initC', settings_key='initC_settings') + self.kick_info.reset_index(drop=True).to_hdf(path, key='kick_info', mode='a') + + with h5.File(path, 'a') as f: + # Sample arrays + f.create_dataset('samples', data=self.samples) + f.create_dataset('is_hit', data=self.is_hit) + f.create_dataset('generation', data=self.generation) + f.create_dataset('gaussian_idx', data=self.gaussian_idx) + + # Mixture model (optional) + if self.mixture is not None: + mg = f.create_group('mixture') + mg.create_dataset('means', data=self.mixture.means) + mg.create_dataset('covariances', data=self.mixture.covariances) + mg.create_dataset('alphas', data=self.mixture.alphas) + mg.attrs['rejection_rate'] = self.mixture.rejection_rate + + # Scalar metadata + meta = f.require_group('meta') + meta.attrs['param_names'] = self.param_names + meta.attrs['num_explored'] = self.num_explored + meta.attrs['num_hits'] = self.num_hits + meta.attrs['num_hits_exploratory'] = self.num_hits_exploratory + meta.attrs['fraction_explored'] = self.fraction_explored + meta.attrs['prior_fraction_rejected'] = self.prior_fraction_rejected + meta.attrs['total_systems'] = self.total_systems + meta.attrs['n_generations'] = self.n_generations + + @classmethod + def from_file(cls, path): + """Load a checkpoint previously written by :meth:`save`. + + Parameters + ---------- + path : `str` + File path to read. + + Returns + ------- + `STROOPWAFELCheckpoint` + """ + # COSMIC tables + bpp = pd.read_hdf(path, key='bpp') + bcm = pd.read_hdf(path, key='bcm') + initC = load_initC(path, key='initC', settings_key='initC_settings') + kick_info = pd.read_hdf(path, key='kick_info') + + with h5.File(path, 'r') as f: + samples = f['samples'][:] + is_hit = f['is_hit'][:] + generation = f['generation'][:] + gaussian_idx = f['gaussian_idx'][:] + + # Mixture (may be absent) + mixture = None + if 'mixture' in f: + from cosmic.sample.stroopwafel.mixture_model import GaussianMixture + mg = f['mixture'] + mixture = GaussianMixture( + means=mg['means'][:], + covariances=mg['covariances'][:], + alphas=mg['alphas'][:], + rejection_rate=float(mg.attrs['rejection_rate']), + ) + + meta = f['meta'] + param_names = list(meta.attrs['param_names']) + num_explored = int(meta.attrs['num_explored']) + num_hits = int(meta.attrs['num_hits']) + num_hits_exploratory = int(meta.attrs['num_hits_exploratory']) + fraction_explored = float(meta.attrs['fraction_explored']) + prior_fraction_rejected = float(meta.attrs['prior_fraction_rejected']) + total_systems = int(meta.attrs['total_systems']) + n_generations = int(meta.attrs['n_generations']) + + # Re-apply the bin_num index so the same invariant holds as + # when the checkpoint was created. + for df in (bpp, bcm, initC, kick_info): + if 'bin_num' in df.columns: + df.set_index('bin_num', drop=False, inplace=True) + + return cls( + mixture=mixture, + samples=samples, is_hit=is_hit, + generation=generation, gaussian_idx=gaussian_idx, + bpp=bpp, bcm=bcm, initC=initC, kick_info=kick_info, + param_names=param_names, + num_explored=num_explored, num_hits=num_hits, + num_hits_exploratory=num_hits_exploratory, + fraction_explored=fraction_explored, + prior_fraction_rejected=prior_fraction_rejected, + total_systems=total_systems, n_generations=n_generations, + ) + + class COSMICPopOutput(): def __init__(self, file, label=None): # read in convergence tables and totals diff --git a/src/cosmic/sample/stroopwafel/engine.py b/src/cosmic/sample/stroopwafel/engine.py index 684ca9de8..01e5af234 100644 --- a/src/cosmic/sample/stroopwafel/engine.py +++ b/src/cosmic/sample/stroopwafel/engine.py @@ -10,7 +10,7 @@ from cosmic.sample.initialbinarytable import InitialBinaryTable from cosmic.evolve import Evolve -from cosmic.output import COSMICStroopOutput +from cosmic.output import COSMICStroopOutput, STROOPWAFELCheckpoint from .mixture_model import GaussianMixture from .constants import MIN_ACTIVE_FRACTION @@ -53,12 +53,21 @@ class AdaptiveSampler: False seed : `int` or None, optional Random seed for reproducibility, by default None + only_save_hit_tables : `bool`, 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. """ def __init__(self, parameter_space, total_systems, batch_size, BSEDict, compute_derived, reject_systems, is_interesting, output_path='output', nproc=1, kappa=1.0, - n_generations=1, mc_only=False, seed=None): + n_generations=1, mc_only=False, seed=None, + only_save_hit_tables=False): self.param_space = parameter_space self.total_systems = total_systems self.batch_size = batch_size @@ -71,6 +80,7 @@ def __init__(self, parameter_space, total_systems, batch_size, BSEDict, self.kappa = kappa self.n_generations = n_generations self.mc_only = mc_only + self.only_save_hit_tables = only_save_hit_tables self.rng = np.random.default_rng(seed) # State @@ -94,7 +104,10 @@ def __init__(self, parameter_space, total_systems, batch_size, BSEDict, self._all_kick_info = [] def run(self): - """Run the full STROOPWAFEL pipeline. + """Run the full STROOPWAFEL pipeline in a single call. + + For multi-job workflows (e.g. SLURM) use :meth:`run_exploration` + and :meth:`run_refinement` instead. Returns ------- @@ -110,8 +123,230 @@ def run(self): self._adapt() self._refine() - result = self._compute_weights() - return result + return self._compute_weights() + + # ------------------------------------------------------------------ + # Multi-job entry points + # ------------------------------------------------------------------ + + def run_exploration(self): + """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 + :meth:`run_refinement` (or :meth:`from_checkpoint`) to + continue on a different node or job. + """ + os.makedirs(self.output_path, exist_ok=True) + self._explore() + if not self.mc_only and self.num_hits > 0: + self._adapt() + return self._make_checkpoint() + + def run_refinement(self): + """Run the refinement phase and return the final result. + + Must be called after :meth:`run_exploration` **or** after + :meth:`from_checkpoint` has restored exploration state. Suitable + as the second SLURM job:: + + # job_refine.py + sampler = AdaptiveSampler.from_checkpoint( + 'checkpoint.h5', + parameter_space=params, batch_size=1000, + BSEDict=BSEDict, compute_derived=compute_derived, + reject_systems=default_reject, is_interesting=is_bh_star, + ) + result = sampler.run_refinement() + result.save('result.h5') + + Returns + ------- + `COSMICStroopOutput` + """ + if not self.mc_only and self.num_hits > 0 and self.mixture is not None: + self._refine() + return self._compute_weights() + + @classmethod + def from_checkpoint(cls, checkpoint, parameter_space, batch_size, BSEDict, + compute_derived, reject_systems, is_interesting, + output_path='output', nproc=1, seed=None, + total_systems=None, n_generations=None, + only_save_hit_tables=False): + """Create a sampler pre-loaded with state from a :class:`STROOPWAFELCheckpoint`. + + The callable arguments (``BSEDict``, ``compute_derived``, etc.) are + not stored in the checkpoint and must be supplied again. All other + state — explored samples, COSMIC output, mixture model, scalar + counters — is restored from the checkpoint. + + Parameters + ---------- + checkpoint : `STROOPWAFELCheckpoint` or `str` + A checkpoint object or path to an HDF5 file written by + :meth:`STROOPWAFELCheckpoint.save`. + parameter_space : `ParameterSpace` + Must use the same parameter names as when the checkpoint was + created; a `ValueError` is raised if they differ. + batch_size : `int` + Systems per COSMIC call for the refinement phase. + BSEDict : `dict` + COSMIC binary stellar evolution parameters. + compute_derived : `callable` + Same signature as for the original sampler. + reject_systems : `callable` + Same signature as for the original sampler. + is_interesting : `callable` + Same signature as for the original sampler. + output_path : `str`, optional + Directory for output files, by default ``'output'`` + nproc : `int`, optional + CPU cores for COSMIC, by default 1 + seed : `int` or None, optional + RNG seed for refinement sampling, by default None + total_systems : `int`, optional + Override the total system budget stored in the checkpoint. + n_generations : `int`, optional + Override the number of refinement generations stored in the + checkpoint. + + Returns + ------- + `AdaptiveSampler` + Ready to call :meth:`run_refinement`. + + Raises + ------ + ValueError + If ``parameter_space.names`` does not match the checkpoint's + ``param_names``. + """ + if isinstance(checkpoint, (str, os.PathLike)): + checkpoint = STROOPWAFELCheckpoint.from_file(checkpoint) + + if list(parameter_space.names) != list(checkpoint.param_names): + raise ValueError( + f"Parameter space mismatch.\n" + f" checkpoint : {checkpoint.param_names}\n" + f" provided : {list(parameter_space.names)}" + ) + + ts = total_systems if total_systems is not None else checkpoint.total_systems + ng = n_generations if n_generations is not None else checkpoint.n_generations + + sampler = cls( + parameter_space=parameter_space, + total_systems=ts, + batch_size=batch_size, + BSEDict=BSEDict, + compute_derived=compute_derived, + reject_systems=reject_systems, + is_interesting=is_interesting, + output_path=output_path, + nproc=nproc, + n_generations=ng, + mc_only=False, + seed=seed, + only_save_hit_tables=only_save_hit_tables, + ) + sampler._load_checkpoint(checkpoint) + return sampler + + def _make_checkpoint(self): + """Package current engine state into a `STROOPWAFELCheckpoint`.""" + all_samples = np.vstack(self._all_samples) + all_is_hit = np.concatenate(self._all_is_hit) + all_generation = np.concatenate(self._all_generation) + all_gaussian_idx = np.concatenate(self._all_gaussian_idx) + bpp, bcm, initC, kick_info = self._build_cosmic_output() + + return STROOPWAFELCheckpoint( + mixture=self.mixture, + samples=all_samples, + is_hit=all_is_hit, + generation=all_generation, + gaussian_idx=all_gaussian_idx, + bpp=bpp, bcm=bcm, initC=initC, kick_info=kick_info, + param_names=self.param_space.names, + num_explored=self.num_explored, + num_hits=self.num_hits, + num_hits_exploratory=getattr(self, 'num_hits_exploratory', self.num_hits), + fraction_explored=self.fraction_explored, + prior_fraction_rejected=self.prior_fraction_rejected, + total_systems=self.total_systems, + n_generations=self.n_generations, + ) + + def _load_checkpoint(self, checkpoint): + """Restore engine state from a `STROOPWAFELCheckpoint`. + + The checkpoint's samples, COSMIC output, and mixture are treated + as a single exploration "super-batch" with globally assigned + bin_nums (0..N_explore-1). Refinement batches added afterward + will receive bin_nums starting from N_explore. + """ + self.num_explored = checkpoint.num_explored + self.num_hits = checkpoint.num_hits + self.num_hits_exploratory = checkpoint.num_hits_exploratory + self.fraction_explored = checkpoint.fraction_explored + self.prior_fraction_rejected = checkpoint.prior_fraction_rejected + self.finished = checkpoint.num_explored + self.mixture = checkpoint.mixture + + # Treat the full checkpoint as a single pre-computed batch so that + # _build_cosmic_output() can extend it cleanly with refinement batches. + self._all_samples = [checkpoint.samples] + self._all_is_hit = [checkpoint.is_hit] + self._all_generation = [checkpoint.generation] + self._all_gaussian_idx = [checkpoint.gaussian_idx] + # The COSMIC tables already have globally assigned bin_nums from + # when the checkpoint was created; _build_cosmic_output() will add + # offset=0 for this "batch", leaving them unchanged. + self._all_bpp = [checkpoint.bpp] + self._all_bcm = [checkpoint.bcm] + self._all_initC = [checkpoint.initC] + self._all_kick_info = [checkpoint.kick_info] + + def _filter_tables(self, bpp, bcm, initC, kick_info, hit_bin_nums): + """Return COSMIC tables filtered to hit systems when requested. + + When ``only_save_hit_tables`` is False the tables are returned + unchanged. When True, only rows whose ``bin_num`` appears in + ``hit_bin_nums`` are kept; all other rows are discarded immediately, + reducing peak memory proportional to the miss rate. + + Parameters + ---------- + bpp, bcm, initC, kick_info : `pandas.DataFrame` + Raw per-batch COSMIC output with local (0-indexed) bin_nums. + hit_bin_nums : `numpy.ndarray` + Local bin_num values of the hit systems in this batch. + + Returns + ------- + bpp, bcm, initC, kick_info : `pandas.DataFrame` + Possibly filtered DataFrames. + """ + if not self.only_save_hit_tables: + return bpp, bcm, initC, kick_info + mask = hit_bin_nums # already local bin_num values + return ( + bpp[bpp['bin_num'].isin(mask)], + bcm[bcm['bin_num'].isin(mask)], + initC[initC['bin_num'].isin(mask)], + kick_info[kick_info['bin_num'].isin(mask)], + ) # ------------------------------------------------------------------ # Exploration @@ -160,6 +395,9 @@ def _explore(self): is_hit[hit_bin_nums] = True # Accumulate samples and COSMIC output + bpp, bcm, initC, kick_info = self._filter_tables( + bpp, bcm, initC, kick_info, hit_bin_nums + ) self._all_samples.append(batch_samples) self._all_is_hit.append(is_hit) self._all_generation.append(np.zeros(len(selected), dtype=int)) @@ -303,6 +541,9 @@ def _refine(self): is_hit[hit_bin_nums] = True # Accumulate samples and COSMIC output + bpp, bcm, initC, kick_info = self._filter_tables( + bpp, bcm, initC, kick_info, hit_bin_nums + ) self._all_samples.append(batch_samples) self._all_is_hit.append(is_hit) self._all_generation.append(np.full(n_take, gen + 1, dtype=int)) From 26dff3962b0084c61b0db5de537fc9709f93ea0c Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 2 Jun 2026 15:17:03 -0400 Subject: [PATCH 22/45] better upper limit for kick prior --- src/cosmic/sample/stroopwafel/examples/bh_star.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/examples/bh_star.py b/src/cosmic/sample/stroopwafel/examples/bh_star.py index fc9314fbc..0ad86347a 100644 --- a/src/cosmic/sample/stroopwafel/examples/bh_star.py +++ b/src/cosmic/sample/stroopwafel/examples/bh_star.py @@ -30,7 +30,6 @@ from stroopwafel import AdaptiveSampler, ParameterSpace, Parameter from stroopwafel.rejection import default_reject -from stroopwafel import io as swio # ------------------------------------------------------------------ # CLI @@ -123,7 +122,7 @@ Parameter('ecc', 1e-9, 0.99999999, sampler='sana_ecc', prior='sana_ecc'), Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log', prior='flat_in_log'), # --- primary natal kick magnitude only --- - Parameter('natal_kick_1', 0.1, 100.0, sampler='log_normal', prior='log_normal'), + Parameter('natal_kick_1', 0.1, 5000.0, sampler='log_normal', prior='log_normal'), ]) # ------------------------------------------------------------------ From 6c725032a8161c91db08cf605c617859360e54d8 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 2 Jun 2026 17:40:58 -0400 Subject: [PATCH 23/45] add stroopwafel to api docs and tutorials --- docs/modules/sample.rst | 6 + docs/pages/tutorials.rst | 41 +++-- docs/pages/tutorials/adaptive.rst | 12 ++ .../adaptive.rst => adaptive/basics.rst} | 167 +++++------------- docs/pages/tutorials/adaptive/checkpoint.rst | 8 + docs/pages/tutorials/adaptive/outputs.rst | 6 + 6 files changed, 109 insertions(+), 131 deletions(-) create mode 100644 docs/pages/tutorials/adaptive.rst rename docs/pages/tutorials/{sample/adaptive.rst => adaptive/basics.rst} (77%) create mode 100644 docs/pages/tutorials/adaptive/checkpoint.rst create mode 100644 docs/pages/tutorials/adaptive/outputs.rst diff --git a/docs/modules/sample.rst b/docs/modules/sample.rst index 2576cde9e..786240f8b 100644 --- a/docs/modules/sample.rst +++ b/docs/modules/sample.rst @@ -29,6 +29,12 @@ Multidimensional sampler :no-inheritance-diagram: :no-heading: +Adaptive importance sampler +--------------------------- + +.. automodapi:: cosmic.sample.stroopwafel + :no-heading: + CMC related functions --------------------- diff --git a/docs/pages/tutorials.rst b/docs/pages/tutorials.rst index 5d1892b40..e29b2b3a2 100644 --- a/docs/pages/tutorials.rst +++ b/docs/pages/tutorials.rst @@ -52,58 +52,77 @@ This page contains of tutorials that show you how to use ``COSMIC`` and become a :ref:`cmc_sampling` - .. grid-item-card:: .. container:: tutorial-card - .. rubric:: Converged populations + .. rubric:: Re-running simulations :class: tutorial-card-title .. rst-class:: tutorial-card-link - :ref:`fixedpop` + :ref:`rerun_rerun` + + .. rst-class:: tutorial-card-link + + :ref:`rerun_restart` .. grid-item-card:: .. container:: tutorial-card - .. rubric:: Re-running simulations + .. rubric:: Modifying timesteps :class: tutorial-card-title .. rst-class:: tutorial-card-link - :ref:`rerun_rerun` + :ref:`timesteps_resolution` .. rst-class:: tutorial-card-link - :ref:`rerun_restart` + :ref:`timesteps_modifiers` + .. grid-item-card:: .. container:: tutorial-card - .. rubric:: Modifying timesteps + .. rubric:: Converged populations :class: tutorial-card-title .. rst-class:: tutorial-card-link - :ref:`timesteps_resolution` + :ref:`fixedpop` + + .. grid-item-card:: + + .. container:: tutorial-card + + .. rubric:: Analysing simulations + :class: tutorial-card-title .. rst-class:: tutorial-card-link - :ref:`timesteps_modifiers` + :ref:`analysis_interface` .. grid-item-card:: .. container:: tutorial-card - .. rubric:: Analysing simulations + .. rubric:: Adaptive importance sampling :class: tutorial-card-title .. rst-class:: tutorial-card-link + + :ref:`adaptive_basics` - :ref:`analysis_interface` + .. rst-class:: tutorial-card-link + + :ref:`adaptive_outputs` + + .. rst-class:: tutorial-card-link + + :ref:`adaptive_checkpoint` .. grid-item-card:: diff --git a/docs/pages/tutorials/adaptive.rst b/docs/pages/tutorials/adaptive.rst new file mode 100644 index 000000000..a75c63032 --- /dev/null +++ b/docs/pages/tutorials/adaptive.rst @@ -0,0 +1,12 @@ +############################ +Adaptive importance sampling +############################ + +These tutorials will cover how to use the adaptive importance sampling method in ``COSMIC`` to efficiently sample from the parameter space and obtain results with fewer samples. + +.. toctree:: + :maxdepth: 2 + + adaptive/basics + adaptive/outputs + adaptive/checkpoint \ No newline at end of file diff --git a/docs/pages/tutorials/sample/adaptive.rst b/docs/pages/tutorials/adaptive/basics.rst similarity index 77% rename from docs/pages/tutorials/sample/adaptive.rst rename to docs/pages/tutorials/adaptive/basics.rst index 0044700db..2c5e207ad 100644 --- a/docs/pages/tutorials/sample/adaptive.rst +++ b/docs/pages/tutorials/adaptive/basics.rst @@ -1,21 +1,20 @@ -.. _adaptive: +.. _adaptive_basics: -********************************************** +********************************************* Adaptive importance sampling for rare systems -********************************************** +********************************************* The standard :ref:`independent ` and :ref:`multidimensional ` -samplers draw binary parameters from their prior distributions and evolve each system with -COSMIC. For most stellar outcomes this works well: common-envelope episodes, mass-transferring +samplers draw binary initial parameters from distributions that are defined for each sampler. These samples can then be evolved with COSMIC to find the final population. For many scenarios this works well: common-envelope episodes, mass-transferring binaries, and white dwarf systems all occur frequently enough that thousands of random draws yield a workable sample. -Some outcomes are extremely rare. Double black holes that merge within the Hubble time -form at rates of order 1-in-10,000 per binary evolved (or far less at high metallicity). -Sampling such populations with a flat Monte Carlo requires tens of millions of COSMIC calls -to accumulate even a few hundred systems — prohibitive for any serious parameter survey. +However, some outcomes are extremely rare. Binary black holes that merge within the Hubble time +form at very low rates (even more so for NS + NS mergers, especially at high metallicity). +Sampling such populations with regular Monte Carlo draws may require tens of millions of COSMIC calls +to accumulate even a few hundred systems — which may end up being prohibitive for a large parameter survey. -COSMIC includes a vectorised implementation of the STROOPWAFEL algorithm +``COSMIC`` includes a vectorised implementation of the ``STROOPWAFEL`` algorithm (`Broekgaarden et al. 2019 `_) that solves this problem using *adaptive importance sampling*. The sampler first explores parameter space to locate the progenitor regions of the target population, then concentrates its simulation @@ -25,30 +24,25 @@ distribution. When should I use this? -======================== +======================= -Use STROOPWAFEL whenever you need a statistically representative sample of a rare binary +Use ``STROOPWAFEL`` whenever you need a statistically representative sample of a rare binary outcome and cannot afford the total binary count that flat Monte Carlo would require. -Typical use cases include: +Example use cases include: * Double black holes (or neutron stars) merging within the Hubble time (GW sources) -* BH + stellar-companion systems, such as X-ray binaries or Be/X-ray binaries -* Short-period post-common-envelope binaries that survive to become AM CVn systems -* Any binary channel with a formation efficiency ≲ 10\ :sup:`−3` per prior draw +* Long lived BH + stellar-companion systems -If your target population is common (≳ 1 % of all binaries evolving as the target), the -plain independent sampler is simpler and fast enough. The break-even point is roughly -where collecting 100 hits would require more than ~10,000 total evolutions. +If your target population is common the plain independent sampler is simpler and fast enough. +How it works: think Battleships +=============================== -How it works: the battleships analogy -====================================== - -STROOPWAFEL works in three phases that map neatly onto the board game Battleships. +``STROOPWAFEL`` works in three phases that you can think of as a game of Battleships. You wouldn't continue to shoot randomly at the grid after you find a ship — instead, you would concentrate your fire around the hit location to sink it. Similarly, ``STROOPWAFEL`` first explores parameter space with random draws, then adapts a proposal distribution based on the hits it finds, and finally concentrates its sampling from that proposal. **Exploration — random fire** - Binaries are drawn at random from the prior distributions and evolved with COSMIC. - Every binary that produces the desired outcome is recorded as a *hit*. An adaptive + Binaries are drawn at random from the prior distributions and evolved with ``COSMIC``. + Every binary that produces the desired outcome is recorded as a *hit*. An adaptive stopping criterion (based on the observed hit rate) decides when the remaining budget is better spent on refinement than on further random exploration. @@ -56,7 +50,7 @@ STROOPWAFEL works in three phases that map neatly onto the board game Battleship One multivariate Gaussian component is placed at each hit location in parameter space. The width of each Gaussian is derived from the local density of the prior (via the CDF of the prior distribution) so that the proposal is appropriately broad regardless of - the parameter's scale. Together the Gaussians form a *mixture model* — a coarse map of + the parameter's scale. Together the Gaussians form a *mixture model* — a coarse map of where progenitors live. **Refinement — concentrate fire** @@ -76,12 +70,14 @@ STROOPWAFEL works in three phases that map neatly onto the board game Battleship where :math:`\pi(x)` is the prior probability density, :math:`q(x)` is the Gaussian mixture density, and :math:`f_e` is the fraction of systems drawn from the prior - during exploration. Using these weights, any statistic computed on the hit population + during exploration. Using these weights, any statistic computed on the hit population is an unbiased estimator of the corresponding prior-weighted quantity. +Setup +===== -Setting up the parameter space -================================ +Define the parameter space +-------------------------- The :class:`~cosmic.sample.stroopwafel.ParameterSpace` class defines which binary parameters are sampled and their distributions. Each @@ -143,8 +139,8 @@ Parameters are stored and returned in alphabetical order by name. Use ``params.names`` to inspect the column ordering of any sample array. -Defining derived quantities -============================ +Define any derived quantities +----------------------------- COSMIC requires ``mass_2``, ``separation``, and ``metallicity`` in addition to the directly-sampled parameters. The ``compute_derived`` callback converts a ``(N, D)`` @@ -181,8 +177,8 @@ array of physical-space samples and a sorted list of parameter names into a dict supply a custom rejection function you are free to use different key names. -Defining the rejection function -================================ +Choose a rejection function +--------------------------- Before a batch is passed to COSMIC, unphysical systems are filtered out: stars already overflowing their Roche lobes at ZAMS, binaries whose components are in contact, and @@ -208,8 +204,8 @@ imposing a minimum primary mass — you can wrap it: return base_mask -Defining the hit criterion -============================ +Identify what constitutes a hit +------------------------------- The ``is_interesting`` argument identifies which evolved systems count as hits. It receives the COSMIC ``bpp`` DataFrame for the current batch and must return a tuple @@ -268,8 +264,14 @@ BH forms: return len(hits), hits -Example 1: Bound BH + BH binaries -==================================== +Running the sampler +=================== + +Examples +======== + +Bound BH + BH binaries +---------------------- The following end-to-end example samples all bound BH-BH systems (no merger time restriction) using a five-dimensional parameter space covering primary mass, mass ratio, @@ -357,8 +359,8 @@ orbital period, eccentricity, and metallicity. print(f"Weighted hit rate: {result.hit_rate:.4e} ± {result.hit_rate_uncertainty:.4e}") -Example 2: BH + star binaries surviving 100 Myr -================================================== +BH + star binaries surviving 100 Myr +------------------------------------ For outcomes that are less extreme but still rare — such as persistent BH + star systems — STROOPWAFEL provides substantial efficiency gains over flat Monte Carlo. Using the same @@ -407,9 +409,10 @@ Because BH + star systems are more common than merging BH-BH pairs, a smaller to is needed and fewer refinement generations are required before the mixture model is well-constrained. +Rules of thumb +============== -Choosing ``total_systems``, ``batch_size``, and ``n_generations`` -=================================================================== +Choosing ``total_systems``, ``batch_size``, and ``n_generations`` is something of an art but there are some rules of thumb to get you started. The optimal settings depend on the rarity and complexity of the target population, the dimensionality of the parameter space, and your computational resources. ``batch_size`` -------------- @@ -457,83 +460,7 @@ The EM step between generations can improve the mixture, but with diminishing re (i.e. all weights are equal). -Working with results -====================== - -:meth:`~cosmic.sample.stroopwafel.engine.AdaptiveSampler.run` returns a -:class:`~cosmic.sample.stroopwafel.result.STROOPWAFELResult` object containing all -simulated systems, their importance weights, and summary statistics: - -.. code-block:: python - - import numpy as np - - # Shape of sample array and column ordering - print(result.samples.shape) # (N_total, D) - print(result.param_names) # sorted alphabetically, e.g. - # ['ecc', 'mass_1', 'metallicity', 'porb', 'q'] - - # Extract hits and their weights - hit_samples = result.samples[result.is_hit] # (N_hits, D) - hit_weights = result.weights[result.is_hit] # (N_hits,) - - # Normalise weights for the hit population - hit_weights_norm = hit_weights / hit_weights.sum() - - # Importance-weighted primary mass histogram - m1_col = result.param_names.index('mass_1') - m1_hits = hit_samples[:, m1_col] - hist, edges = np.histogram(m1_hits, bins=20, weights=hit_weights_norm) - - # Importance-weighted hit rate (fraction of prior draws producing a hit) - print(f"Hit rate: {result.hit_rate:.4e} ± {result.hit_rate_uncertainty:.4e}") - - # How the budget was spent - print(f"Explored: {result.num_explored} Total hits: {result.num_hits}") - -.. note:: - - ``result.samples`` stores samples in **physical space** (masses in M\ :sub:`☉`, - periods in days, etc.) in the alphabetically-sorted column order defined by - ``ParameterSpace``. Always use ``result.param_names`` to map column indices to - parameter names rather than relying on the order in which you defined the parameters. - - -Saving and loading results -============================ - -The full result can be saved to HDF5 for later analysis: - -.. code-block:: python - - from cosmic.sample.stroopwafel import io as swio - - swio.save_result('bhbh_result.h5', result) - -The file stores the sample array, importance weights, hit flags, generation labels, and -summary statistics as HDF5 datasets and attributes. - -To reload the data in a later session without re-running the sampler: - -.. code-block:: python - - import h5py - import numpy as np - - with h5py.File('bhbh_result.h5', 'r') as f: - samples = f['samples'][:] - weights = f['weights'][:] - is_hit = f['is_hit'][:] - param_names = list(f.attrs['param_names']) - num_hits = f.attrs['num_hits'] - - hits = samples[is_hit] - m1 = hits[:, param_names.index('mass_1')] - w = weights[is_hit] - print(f"Weighted mean BH primary mass: {np.average(m1, weights=w):.1f} M_sun") - -.. note:: +Saving your results +=================== - The HDF5 file does **not** store the raw COSMIC ``bpp`` output tables. If you need - the evolutionary histories of the hit systems, re-evolve them with COSMIC using the - hit sample coordinates from ``result.samples[result.is_hit]`` and the same ``BSEDict``. +TODO diff --git a/docs/pages/tutorials/adaptive/checkpoint.rst b/docs/pages/tutorials/adaptive/checkpoint.rst new file mode 100644 index 000000000..3112fcdf3 --- /dev/null +++ b/docs/pages/tutorials/adaptive/checkpoint.rst @@ -0,0 +1,8 @@ +.. _adaptive_checkpoint: + +****************************** +Saving and loading checkpoints +****************************** + +In this tutorial, we will cover how to save and load checkpoints during an adaptive importance sampling run in ``COSMIC``. This allows you to save the state of your simulation once the adaptation phase is complete, and then load it later to continue sampling from the same Gaussian mixture model (GMM) without having to redo the adaptation phase. This can be particularly useful if you want to run multiple sampling phases with the same adapted GMM (e.g. over multiple nodes on a cluster). + diff --git a/docs/pages/tutorials/adaptive/outputs.rst b/docs/pages/tutorials/adaptive/outputs.rst new file mode 100644 index 000000000..92be168c7 --- /dev/null +++ b/docs/pages/tutorials/adaptive/outputs.rst @@ -0,0 +1,6 @@ +.. _adaptive_outputs: + +*************************************** +Handling outputs from adaptive sampling +*************************************** + From 0358c81c19b1a7b204dc193b18cb32827bc3748d Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 2 Jun 2026 17:59:01 -0400 Subject: [PATCH 24/45] basics stroopwafel tutorial working --- docs/modules/sample.rst | 1 + docs/pages/tutorials.rst | 3 +- docs/pages/tutorials/adaptive/basics.rst | 154 +++++----------------- docs/pages/tutorials/sample.rst | 10 -- src/cosmic/sample/stroopwafel/__init__.py | 4 - 5 files changed, 36 insertions(+), 136 deletions(-) diff --git a/docs/modules/sample.rst b/docs/modules/sample.rst index 786240f8b..8bf6d50d8 100644 --- a/docs/modules/sample.rst +++ b/docs/modules/sample.rst @@ -33,6 +33,7 @@ Adaptive importance sampler --------------------------- .. automodapi:: cosmic.sample.stroopwafel + :no-inheritance-diagram: :no-heading: CMC related functions diff --git a/docs/pages/tutorials.rst b/docs/pages/tutorials.rst index e29b2b3a2..63e110011 100644 --- a/docs/pages/tutorials.rst +++ b/docs/pages/tutorials.rst @@ -141,8 +141,9 @@ This page contains of tutorials that show you how to use ``COSMIC`` and become a tutorials/evolve tutorials/sample - tutorials/convergence tutorials/rerun tutorials/timesteps + tutorials/convergence tutorials/analysis + tutorials/adaptive tutorials/misc \ No newline at end of file diff --git a/docs/pages/tutorials/adaptive/basics.rst b/docs/pages/tutorials/adaptive/basics.rst index 2c5e207ad..8fa84dffc 100644 --- a/docs/pages/tutorials/adaptive/basics.rst +++ b/docs/pages/tutorials/adaptive/basics.rst @@ -97,46 +97,7 @@ sampling distribution, and a prior distribution. Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log',prior='flat_in_log'), ]) -The available samplers and corresponding prior names are: - -.. list-table:: - :header-rows: 1 - :widths: 20 45 35 - - * - Sampler name - - Distribution - - Typical use - * - ``'uniform'`` - - Uniform between bounds - - Mass ratio, any flat prior - * - ``'kroupa'`` - - Kroupa (2001) power law (:math:`dN/dm_1 \propto m_1^{-2.3}`) - - Primary mass - * - ``'sana'`` - - Sana et al. (2012) period distribution - (:math:`dN/d\!\log P \propto (\log P)^{-0.55}`) - - Orbital period - * - ``'sana_ecc'`` - - Sana et al. (2012) eccentricity distribution - (:math:`dN/de \propto e^{-0.45}`) - - Orbital eccentricity - * - ``'flat_in_log'`` - - Uniform in :math:`\log_{10}` (Öpik's law) - - Metallicity, semi-major axis - -.. note:: - - The ``'sana'`` sampler operates in :math:`\log_{10}(P/\text{days})` space. The bounds - you supply are :math:`\log_{10}` values directly — **not** periods in days. The - Sana et al. (2012) fit is valid over :math:`\log_{10}(P) \in [0.15,\, 5.5]`, - corresponding to periods of roughly 1.4 to 316,000 days. - - The lower bound **must be positive** (i.e. :math:`\log_{10}(P_\text{min}) > 0`, - so :math:`P_\text{min} > 1` day). Do **not** pass ``np.log10(P_min)`` when that - value would be negative. - -Parameters are stored and returned in alphabetical order by name. Use -``params.names`` to inspect the column ordering of any sample array. +Parameters are stored and returned in alphabetical order by name. Use ``params.names`` to check the order and ``params.index('param_name')`` to get the index of a particular parameter. Define any derived quantities @@ -180,7 +141,7 @@ array of physical-space samples and a sorted list of parameter names into a dict Choose a rejection function --------------------------- -Before a batch is passed to COSMIC, unphysical systems are filtered out: stars already +Before a batch is passed to ``COSMIC``, unphysical systems are filtered out: stars already overflowing their Roche lobes at ZAMS, binaries whose components are in contact, and systems below the hydrogen-burning limit for the secondary. The built-in :func:`~cosmic.sample.stroopwafel.rejection.default_reject` function performs all of these @@ -208,11 +169,11 @@ Identify what constitutes a hit ------------------------------- The ``is_interesting`` argument identifies which evolved systems count as hits. It receives -the COSMIC ``bpp`` DataFrame for the current batch and must return a tuple +the ``COSMIC`` ``bpp`` DataFrame for the current batch and must return a tuple ``(n_hits, hit_bin_nums)`` where ``hit_bin_nums`` is an integer array of ``bin_num`` values (0-indexed within the batch). -STROOPWAFEL ships two preset factory functions in +The ``STROOPWAFEL`` sampler comes with two preset functions in :mod:`cosmic.sample.stroopwafel.presets`. ``any_dco(kstar_1, kstar_2)`` @@ -236,22 +197,18 @@ STROOPWAFEL ships two preset factory functions in See :ref:`kstar-table` for the full list of stellar type codes. -You can also write a fully custom hit function. For example, to find BH + stellar -companion systems (``kstar_2 ∈ 0–9``) that remain bound for at least 100 Myr after the -BH forms: +You can also write a fully custom hit function. For example, to find BH + stellar companion systems that remain bound for at least 100 Myr after the BH forms: .. code-block:: python import numpy as np - _STELLAR_TYPES = set(range(10)) # kstar 0–9: MS through He-giant branch - def bh_star_100myr(bpp): """Hit: BH with a stellar companion bound for at least 100 Myr.""" bh_star = bpp.loc[ ( - ((bpp['kstar_1'] == 14) & bpp['kstar_2'].isin(_STELLAR_TYPES)) - | ((bpp['kstar_2'] == 14) & bpp['kstar_1'].isin(_STELLAR_TYPES)) + ((bpp['kstar_1'] == 14) & (bpp['kstar_2'] < 10)) + | ((bpp['kstar_2'] == 14) & (bpp['kstar_1'] < 10)) ) & (bpp['sep'] > 0) ] @@ -267,16 +224,22 @@ BH forms: Running the sampler =================== +With all the pieces in place, you can run the sampler with :class:`~cosmic.sample.stroopwafel.AdaptiveSampler`. The most important arguments are the parameter space, the total number of systems to evolve, the batch size, the BSE physics settings, the derived quantity function, the rejection function, and the hit function. See the API documentation (:class:`~cosmic.sample.stroopwafel.AdaptiveSampler`) for a full list of options. + +The examples below demonstrate how you could go about this. + Examples -======== +-------- Bound BH + BH binaries ----------------------- +^^^^^^^^^^^^^^^^^^^^^^ The following end-to-end example samples all bound BH-BH systems (no merger time restriction) using a five-dimensional parameter space covering primary mass, mass ratio, orbital period, eccentricity, and metallicity. +.. include:: ../../../_generated/default_bsedict.rst + .. code-block:: python import numpy as np @@ -284,29 +247,6 @@ orbital period, eccentricity, and metallicity. from cosmic.sample.stroopwafel.presets import any_dco from cosmic.sample.stroopwafel.rejection import default_reject - # ------------------------------------------------------------------ - # BSE physics settings - # ------------------------------------------------------------------ - BSEDict = { - "pts1": 0.001, "pts2": 0.01, "pts3": 0.02, "zsun": 0.014, - "windflag": 3, "neta": 0.5, "bwind": 0.0, "hewind": 0.5, - "beta": 0.125, "xi": 0.5, "acc2": 1.5, "LBV_flag": 1, - "alpha1": 1.0, "lambdaf": 0.0, "ceflag": 1, "cekickflag": 2, - "cemergeflag": 1, "cehestarflag": 0, "qcflag": 5, - "qcrit_array": [0.0] * 16, - "kickflag": 5, "sigma": 265.0, "bhflag": 1, "bhsigmafrac": 1.0, - "sigmadiv": -20.0, "ecsn": 2.25, "ecsn_mlow": 1.6, "aic": 1, - "ussn": 1, "polar_kick_angle": 90.0, - "natal_kick_array": [[-100.0]*5, [-100.0]*5], - "remnantflag": 4, "mxns": 3.0, "rembar_massloss": 0.5, - "wd_mass_lim": 1, "grflag": 1, "eddfac": 10, "tflag": 1, - "ST_tide": 1, "ifflag": 1, "wdflag": 1, "epsnov": 0.001, - "bdecayfac": 1, "bconst": 3000, "ck": 1000, "htpmb": 1, - "ST_cr": 1, "rtmsflag": 0, - "fprimc_array": [2.0 / 21.0] * 16, - "mm_mu_ns": 400.0, "mm_mu_bh": 200.0, "pisn": -2, - } - # ------------------------------------------------------------------ # Parameter space # ------------------------------------------------------------------ @@ -349,7 +289,7 @@ orbital period, eccentricity, and metallicity. is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), output_path='output/bhbh', nproc=4, - n_generations=3, + n_generations=1, seed=42, ) @@ -360,11 +300,11 @@ orbital period, eccentricity, and metallicity. BH + star binaries surviving 100 Myr ------------------------------------- +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ For outcomes that are less extreme but still rare — such as persistent BH + star systems — STROOPWAFEL provides substantial efficiency gains over flat Monte Carlo. Using the same -parameter space, ``compute_derived``, and ``BSEDict`` as Example 1: +parameter space, ``compute_derived``, and ``BSEDict`` as the previous examples, we can simply swap out the hit function to find BH + star systems that remain bound for at least 100 Myr after the BH forms: .. code-block:: python @@ -372,34 +312,17 @@ parameter space, ``compute_derived``, and ``BSEDict`` as Example 1: from cosmic.sample.stroopwafel import AdaptiveSampler from cosmic.sample.stroopwafel.rejection import default_reject - _STELLAR_TYPES = set(range(10)) # kstar 0–9: MS through He-giant - - def bh_star_100myr(bpp): - """Hit: BH + stellar companion bound for at least 100 Myr.""" - bh_star = bpp.loc[ - ( - ((bpp['kstar_1'] == 14) & bpp['kstar_2'].isin(_STELLAR_TYPES)) - | ((bpp['kstar_2'] == 14) & bpp['kstar_1'].isin(_STELLAR_TYPES)) - ) - & (bpp['sep'] > 0) - ] - if bh_star.empty: - return 0, np.array([], dtype=int) - span = bh_star.groupby('bin_num')['tphys'].agg(lambda t: t.max() - t.min()) - hits = span.index[span >= 100.0].values - return len(hits), hits - sampler = AdaptiveSampler( - parameter_space=params, # reuse from Example 1 + parameter_space=params, # reuse from BHBH example total_systems=20_000, batch_size=500, - BSEDict=BSEDict, # reuse from Example 1 - compute_derived=compute_derived, # reuse from Example 1 + BSEDict=BSEDict, # reuse from BHBH example + compute_derived=compute_derived, # reuse from BHBH example reject_systems=default_reject, - is_interesting=bh_star_100myr, + is_interesting=bh_star_100myr, # we defined this earlier output_path='output/bh_star', nproc=4, - n_generations=2, + n_generations=1, seed=42, ) @@ -422,35 +345,19 @@ Choosing ``total_systems``, ``batch_size``, and ``n_generations`` is something o * Aim for ``batch_size`` to be a multiple of ``nproc`` so that COSMIC distributes work evenly across cores. -* Values of 200–1000 are typical. Batches smaller than ~50 increase Python overhead per +* Values of 200-1000 are typical. Batches smaller than ~50 increase Python overhead per call; batches larger than ~5000 may cause memory pressure on the output DataFrames. -* A practical starting point is ``batch_size = 100 * nproc``. ``total_systems`` ----------------- -This is the total number of binary evolutions across all phases. - -* For **very rare events** (hit rate ≲ 10\ :sup:`−4`, e.g. merging BH-BH at near-solar - metallicity), start with ``total_systems`` in the range 100,000–500,000. The exploration - phase will find tens to hundreds of hits; refinement then multiplies that count many-fold. -* For **moderately rare events** (hit rate ~ 10\ :sup:`−3` to 10\ :sup:`−2`, e.g. any bound - BH-BH or long-lived BH + star), 20,000–50,000 systems is usually sufficient. -* As a rule of thumb, aim for at least ~30 hits during exploration before the adaptation - phase begins — fewer hits lead to a poorly-constrained Gaussian mixture. If exploration - ends with very few hits, increase ``total_systems`` and re-run. +This is the total number of binary evolutions across all phases. It's hard to know how many you'll need without knowing the rarity of the target population. As a rule of thumb, aim for at least ~30 hits during exploration before the adaptation phase begins — fewer hits lead to a poorly-constrained Gaussian mixture. If exploration ends with very few hits, increase ``total_systems`` and re-run. ``n_generations`` ----------------- Each refinement generation uses an equal share of the remaining budget after exploration. -The EM step between generations can improve the mixture, but with diminishing returns: - -* ``n_generations = 1`` (the default) uses the mixture as constructed from exploration - hits, with no EM updates. This is a good starting point for any new target population. -* ``n_generations = 3`` gives a noticeable improvement for very rare populations with - complex progenitor structure. -* Beyond 5 generations the returns diminish rapidly. +The EM step between generations can improve the mixture, but with diminishing returns. You are probably safe with just 1 generation unless you have a very rare population. .. tip:: @@ -459,8 +366,13 @@ The EM step between generations can improve the mixture, but with diminishing re prior only, and importance weights will equal the prior density divided by itself (i.e. all weights are equal). - Saving your results =================== -TODO +Once you have your samples, you can save them to disk as an HDF5 file with the :meth:`~cosmic.output.COSMICSTROOPWAFELResult.save` method. This saves the parameter samples, derived quantities, and hit information in a compact format that can be loaded later for analysis. + +.. code-block:: python + + result.save('bhbh_samples.h5') + +We'll talk more about how to load and analyse these results in the :ref:`adaptive_outputs` tutorial next! \ No newline at end of file diff --git a/docs/pages/tutorials/sample.rst b/docs/pages/tutorials/sample.rst index cf7803d97..fff24e0c9 100644 --- a/docs/pages/tutorials/sample.rst +++ b/docs/pages/tutorials/sample.rst @@ -31,16 +31,6 @@ We consider both cases in the guides below. sample/independent sample/multidim -For rare outcomes such as merging double black holes or persistent X-ray binary -systems, COSMIC also provides an adaptive importance sampler based on the -STROOPWAFEL algorithm that concentrates the simulation budget on progenitor -regions of parameter space. - -.. toctree:: - :maxdepth: 1 - - sample/adaptive - You can also use COSMIC to sample the initial conditions for a Globular Cluster (GC) using the ClusterMonteCarlo (CMC) software package. Check out the guide below for more information. diff --git a/src/cosmic/sample/stroopwafel/__init__.py b/src/cosmic/sample/stroopwafel/__init__.py index e1f94d65a..f4b125798 100755 --- a/src/cosmic/sample/stroopwafel/__init__.py +++ b/src/cosmic/sample/stroopwafel/__init__.py @@ -7,10 +7,6 @@ the resulting mixture to concentrate compute budget on interesting regions of parameter space. -Example -------- ->>> from cosmic.sample.stroopwafel import AdaptiveSampler, ParameterSpace, Parameter ->>> from cosmic.sample.stroopwafel.presets import merging_dco """ from .engine import AdaptiveSampler from .parameter_space import ParameterSpace, Parameter From 373ca5064a5a3d0cef8850064fbc04bd0dcb664c Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 15 Jun 2026 12:40:58 -0400 Subject: [PATCH 25/45] working on tutorials --- docs/pages/tutorials/adaptive/outputs.rst | 83 +++++++++++++++++++++++ src/cosmic/output.py | 3 + 2 files changed, 86 insertions(+) diff --git a/docs/pages/tutorials/adaptive/outputs.rst b/docs/pages/tutorials/adaptive/outputs.rst index 92be168c7..e269e711c 100644 --- a/docs/pages/tutorials/adaptive/outputs.rst +++ b/docs/pages/tutorials/adaptive/outputs.rst @@ -4,3 +4,86 @@ Handling outputs from adaptive sampling *************************************** +This tutorial assumes that you've already gone through :ref:`adaptive_basics`. + +Reading your results from a file +================================ + +After you've finished running your adaptive sampling simulation, you will now have some results stored as :class:`~cosmic.output.COSMICStroopOutput` object. If you saved these results to a file, then you can reload them by running + +.. code-block:: python + + from cosmic.output import COSMICStroopOutput + + results = COSMICStroopOutput.from_file("path/to/your/file.h5") + +Understanding your outputs +========================== + +The :class:`~cosmic.output.COSMICStroopOutput` class stores all of the information you need to analyse your simulation. Let's step through some of the different attributes that you will need, and assume for the purposes of this guide that you have ``N`` samples, in ``D`` dimensions, with ``H`` hits. + +Sample information +------------------ + +Each :class:`~cosmic.output.COSMICStroopOutput` object contains a full record of every sample that was made during the simulation. + +- ``samples``: contains a every sampled point and as such has shape ``(N, D)`` +- ``param_names``: is a list of length ``D`` with names corresponding to each column in the ``samples`` array +- ``is_hit``: is an array of length ``N`` with boolean values for whether a sample was a hit (and as such will sum to ``H``) + +Hit details +----------- + +For the actual hits (i.e. the samples that you most care about), this class also stores the full evolution history. In particular, ``bpp``, ``bcm``, ``initC``, and ``kick_info`` all contain the usual ``COSMIC`` evolution tables (see :ref:`evolve_single` if you're not familiar). + +.. tip:: + + If you want to just explore your hits in particular, you can sub-select them by doing + + .. code-block:: python + + just_hits = results[results.is_hit] + + which masks the class just like you would with a :class:`~cosmic.output.COSMICOutput`. Be aware you likely still need the full population for access to the weights (we'll cover weights below). + +General metadata +---------------- + +The class also stores other metadata that you may find useful. These include: + +- ``num_explored``, which is the number of systems evolved during the exploration phase +- ``num_hits``, which is the total raw hit count across all phases +- ``fraction_explored``, which is the fraction of total systems used for exploration. + + +How to interpret adaptive sampling weights +========================================== + +So now for the important intuition part. ``COSMIC`` and ``STROOPWAFEL`` have now provided you with a sample of a rare population by preferentially sampling the parameter space that you've specified. However, we of course want to account for the fact that this *is* a rare population. This is where the weights come in. Each sample is assigned an adaptive importance sampling weight, which tells you how rare this sample is (smaller weights are rarer). + +.. admonition:: TODO + + Add maths + +This means that if you want to plot a true distribution of your sampled systems -- let's say the primary mass -- you need to use the weights in your plotting. + +.. code-block:: python + + import matplotlib.pyplot as plt + from cosmic.output import COSMICStroopOutput + + results = COSMICStroopOutput.from_file("YOUR_SIMULATION.h5") + primary_mass = results.samples[:, results.param_names.index("mass_1")] + + plt.hist(primary_mass, weights=results.weights, bins=50, density=True) + plt.xlabel(r"Primary mass, $m_1$ [$\rm M_\odot$]") + plt.ylabel("Probability density") + plt.show() + + +.. warning:: + + You should **always** apply your weights to any plot that you make. In histograms you can supply them directly. For scatter plots you could consider changing the size of your points, or using a 2D histogram or ``hexbin`` instead. + +Drawing a representative sample from your simulation +==================================================== \ No newline at end of file diff --git a/src/cosmic/output.py b/src/cosmic/output.py index adab8eed4..475d485d8 100644 --- a/src/cosmic/output.py +++ b/src/cosmic/output.py @@ -520,6 +520,9 @@ def from_file(cls, path, label=None): fraction_explored=fraction_explored, label=cosmic.label if label is None else label, ) + + def draw_representative_sample(self, sample_size): + raise NotImplementedError class STROOPWAFELCheckpoint: From ce9c10ab0e0f127efb27621cc3f36034a5bec77a Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Sat, 27 Jun 2026 22:45:14 -0400 Subject: [PATCH 26/45] consolidate into single Distributions class --- .../sample/stroopwafel/distributions.py | 421 ++++++++++++++++++ .../sample/stroopwafel/examples/bh_star.py | 12 +- .../stroopwafel/examples/example_bhbh.py | 10 +- src/cosmic/sample/stroopwafel/meson.build | 4 +- .../sample/stroopwafel/parameter_space.py | 139 ++---- src/cosmic/sample/stroopwafel/priors.py | 202 --------- src/cosmic/sample/stroopwafel/samplers.py | 233 ---------- src/cosmic/sample/stroopwafel/transforms.py | 103 ----- src/cosmic/tests/meson.build | 2 + 9 files changed, 468 insertions(+), 658 deletions(-) create mode 100644 src/cosmic/sample/stroopwafel/distributions.py delete mode 100644 src/cosmic/sample/stroopwafel/priors.py delete mode 100644 src/cosmic/sample/stroopwafel/samplers.py delete mode 100644 src/cosmic/sample/stroopwafel/transforms.py diff --git a/src/cosmic/sample/stroopwafel/distributions.py b/src/cosmic/sample/stroopwafel/distributions.py new file mode 100644 index 000000000..53429bf27 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/distributions.py @@ -0,0 +1,421 @@ +"""Composable distributions for STROOPWAFEL parameter sampling. + +A :class:`Distribution` bundles the four things STROOPWAFEL needs to know +about a parameter's prior, which used to be spread across ``samplers.py``, +``priors.py`` and ``transforms.py``: + +* how to **draw** samples (in sampling space), +* the prior **pdf** (in sampling space), +* the adaptive-sampling **sigma** (Gaussian kernel width), and +* the **transform** between physical and sampling space. + +Each distribution is a *base distribution* (:class:`Uniform`, +:class:`PowerLaw`, :class:`TruncatedNormal`) composed with a *transform* +(:class:`Identity`, :class:`Log10`, :class:`Ln`, :class:`Sin`, +:class:`CosShift`). The base operates entirely in sampling space; the +transform maps to and from the physical space the user specifies bounds in +and the simulation consumes. For example ``flat_in_log`` is +``Uniform(transform=Log10())`` and ``sana`` is +``PowerLaw(SANA_G, transform=Log10())``. + +To add a distribution, build an instance and either pass it straight to a +:class:`~cosmic.sample.stroopwafel.parameter_space.Parameter` or register it +under a name:: + + from cosmic.sample.stroopwafel.distributions import PowerLaw, register + register('my_imf', PowerLaw(-2.7)) + +All array-valued methods take and return (N,) ndarrays. +""" +import numpy as np +from scipy.stats import norm as _scipy_norm + +from .constants import ( + ALPHA_IMF, SANA_G, SANA_ECC, NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA, +) + + +# --------------------------------------------------------------------------- +# Transforms between physical and sampling space +# --------------------------------------------------------------------------- +class Transform: + """Map between physical space and sampling space. + + The base class is the identity transform; subclasses override + :meth:`to_sampling` and :meth:`to_physical`. ``bounds`` is derived + automatically and handles non-monotonic-increasing maps (e.g. + :class:`CosShift`) by sorting the endpoints. + """ + + def to_sampling(self, values): + """Convert physical-space values to sampling space. + + Parameters + ---------- + values : `numpy.ndarray` or `float` + Value(s) in physical space. + + Returns + ------- + `numpy.ndarray` or `float` + Value(s) in sampling space. + """ + return values + + def to_physical(self, values): + """Convert sampling-space values to physical space. + + Parameters + ---------- + values : `numpy.ndarray` or `float` + Value(s) in sampling space. + + Returns + ------- + `numpy.ndarray` or `float` + Value(s) in physical space. + """ + return values + + def bounds(self, lo, hi): + """Map a physical-space bound pair into sampling space. + + Parameters + ---------- + lo : `float` + Lower bound in physical space. + hi : `float` + Upper bound in physical space. + + Returns + ------- + lo_sampling : `float` + Lower bound in sampling space. + hi_sampling : `float` + Upper bound in sampling space. + """ + a = self.to_sampling(lo) + b = self.to_sampling(hi) + return (a, b) if a <= b else (b, a) + + +class Identity(Transform): + """Identity transform: sampling space is physical space.""" + + +class Log10(Transform): + """Sampling space is ``log10(physical)`` (physical must be positive).""" + + def to_sampling(self, values): + return np.log10(values) + + def to_physical(self, values): + return np.power(10.0, values) + + +class Ln(Transform): + """Sampling space is ``ln(physical)`` (physical must be positive).""" + + def to_sampling(self, values): + return np.log(values) + + def to_physical(self, values): + return np.exp(values) + + +class Sin(Transform): + """Sampling space is ``sin(angle)`` for an angle in ``[-pi/2, pi/2]``.""" + + def to_sampling(self, values): + return np.sin(values) + + def to_physical(self, values): + return np.arcsin(values) + + +class CosShift(Transform): + """Sampling space is ``cos(angle + pi/2) = -sin(angle)``. + + Suitable for declination-like coordinates measured from ``-pi/2`` to + ``pi/2``. The map is monotonically decreasing, so :meth:`Transform.bounds` + swaps the endpoints to keep ``lo <= hi``. + """ + + def to_sampling(self, values): + return np.cos(values + np.pi / 2) + + def to_physical(self, values): + return np.arccos(values) - np.pi / 2 + + +# --------------------------------------------------------------------------- +# Base distribution +# --------------------------------------------------------------------------- +class Distribution: + """A prior distribution, sampled and evaluated in sampling space. + + Subclasses implement :meth:`sample` and :meth:`pdf`. :meth:`sigma` has a + default implementation (``avg_density / pdf``) that power laws override. + All ``lo``/``hi`` arguments are bounds in *sampling* space (i.e. already + passed through ``self.transform``); :class:`ParameterSpace` handles that + conversion via :meth:`Transform.bounds`. + + Parameters + ---------- + transform : `Transform`, optional + Map between physical and sampling space, by default :class:`Identity`. + """ + + def __init__(self, transform=None): + self.transform = transform if transform is not None else Identity() + + def sample(self, n, lo, hi, rng=None): + """Draw ``n`` samples in sampling space within ``[lo, hi]``. + + Parameters + ---------- + n : `int` + Number of samples to draw. + lo : `float` + Lower bound in sampling space. + hi : `float` + Upper bound in sampling space. + rng : `numpy.random.Generator`, optional + Random number generator, by default None. + + Returns + ------- + `numpy.ndarray` + (N,) array of samples in sampling space. + """ + raise NotImplementedError + + def pdf(self, values, lo, hi): + """Prior probability density at ``values`` (sampling space). + + Parameters + ---------- + values : `numpy.ndarray` + (N,) array of values in sampling space. + lo : `float` + Lower bound in sampling space. + hi : `float` + Upper bound in sampling space. + + Returns + ------- + `numpy.ndarray` + (N,) array of prior densities, normalised over ``[lo, hi]``. + """ + raise NotImplementedError + + def sigma(self, values, lo, hi, avg_density): + """Per-sample Gaussian kernel width for adaptive refinement. + + The default places a kernel whose width is the local inter-sample + spacing, ``avg_density / pdf``. Distributions with closed-form CDFs + (e.g. :class:`PowerLaw`) override this with an exact CDF-space step. + + Parameters + ---------- + values : `numpy.ndarray` + (K,) array of hit values in sampling space. + lo : `float` + Lower bound in sampling space. + hi : `float` + Upper bound in sampling space. + avg_density : `float` + Characteristic inter-sample spacing. + + Returns + ------- + `numpy.ndarray` + (K,) array of sigma values. + """ + return avg_density / self.pdf(values, lo, hi) + + +# --------------------------------------------------------------------------- +# Concrete distributions +# --------------------------------------------------------------------------- +class Uniform(Distribution): + """Uniform distribution on ``[lo, hi]`` in sampling space. + + Combined with a transform this covers ``uniform`` (identity), + ``flat_in_log`` (:class:`Log10`), ``uniform_in_sine`` (:class:`Sin`) and + ``uniform_in_cosine`` (:class:`CosShift`). + """ + + def sample(self, n, lo, hi, rng=None): + rng = rng or np.random.default_rng() + return rng.uniform(lo, hi, n) + + def pdf(self, values, lo, hi): + return np.full(len(values), 1.0 / (hi - lo)) + + +class PowerLaw(Distribution): + r"""Power-law distribution ``p(x) \propto x^alpha`` on ``[lo, hi]``. + + Sampling uses the inverse CDF and the prior is the normalised power law, + both in sampling space. Combined with a transform this covers ``kroupa`` + and ``sana_ecc`` (identity) and ``sana`` (:class:`Log10`). + + Parameters + ---------- + alpha : `float` + Power-law exponent. + transform : `Transform`, optional + Map between physical and sampling space, by default :class:`Identity`. + """ + + def __init__(self, alpha, transform=None): + super().__init__(transform) + self.alpha = alpha + + def sample(self, n, lo, hi, rng=None): + rng = rng or np.random.default_rng() + u = rng.uniform(0, 1, n) + a = self.alpha + 1 + return np.power(u * (hi**a - lo**a) + lo**a, 1.0 / a) + + def pdf(self, values, lo, hi): + a = self.alpha + norm = (a + 1) / (hi**(a + 1) - lo**(a + 1)) + return norm * np.power(values, a) + + def sigma(self, values, lo, hi, avg_density): + """Exact CDF-space step for the power-law sigma. + + Maps each hit to its CDF position, steps by ``avg_density`` in CDF + space, maps back, and returns the larger of the two distances. The + CDF of ``p(x) \\propto x^alpha`` on ``[lo, hi]`` is + ``F(x) = (x^a - lo^a) / (hi^a - lo^a)`` with ``a = alpha + 1``, so the + inverse is ``F^{-1}(u) = (u*(hi^a - lo^a) + lo^a)^{1/a}``. Keeping + every intermediate in ``[lo^a, hi^a]`` avoids the catastrophic + cancellation seen with very small ``lo``. + + Raises + ------ + `ValueError` + If ``lo <= 0`` (the power-law variable must be positive over the + whole range). + """ + if lo <= 0: + raise ValueError( + f"PowerLaw.sigma requires a positive lower bound in sampling " + f"space, but got lo={lo}. When combined with a log transform " + f"(e.g. 'sana'), make sure the physical lower bound maps to a " + f"positive sampling-space value (for 'sana', period > 1 day)." + ) + a = self.alpha + 1 + lo_a = float(lo) ** a + hi_a = float(hi) ** a + range_a = hi_a - lo_a # always finite; no huge intermediates + + # Forward CDF, step in CDF space, then inverse CDF. + u = np.clip((np.power(values, a) - lo_a) / range_a, 0.0, 1.0) + u_right = np.clip(u + avg_density, 0.0, 1.0) + u_left = np.clip(u - avg_density, 0.0, 1.0) + x_right = np.power(u_right * range_a + lo_a, 1.0 / a) + x_left = np.power(u_left * range_a + lo_a, 1.0 / a) + + return np.maximum(np.abs(x_right - values), np.abs(x_left - values)) + + +class TruncatedNormal(Distribution): + """Normal distribution truncated to ``[lo, hi]`` in sampling space. + + Combined with :class:`Ln` this gives the ``log_normal`` natal-kick prior: + the kick magnitude ``v`` follows ``LogNormal(mu, scale)`` so ``ln(v)`` is + normally distributed, and sampling space is ``ln(v)``. + + Parameters + ---------- + mu : `float` + Mean of the underlying normal (in sampling space). + scale : `float` + Standard deviation of the underlying normal (in sampling space). + transform : `Transform`, optional + Map between physical and sampling space, by default :class:`Identity`. + """ + + def __init__(self, mu, scale, transform=None): + super().__init__(transform) + self.mu = mu + self.scale = scale + + def sample(self, n, lo, hi, rng=None): + rng = rng or np.random.default_rng() + # Truncated-normal inverse CDF: map uniform draws into + # [CDF(lo), CDF(hi)] then apply the inverse normal CDF. + p_lo = _scipy_norm.cdf(lo, loc=self.mu, scale=self.scale) + p_hi = _scipy_norm.cdf(hi, loc=self.mu, scale=self.scale) + u = rng.uniform(p_lo, p_hi, n) + return _scipy_norm.ppf(u, loc=self.mu, scale=self.scale) + + def pdf(self, values, lo, hi): + p_lo = _scipy_norm.cdf(lo, loc=self.mu, scale=self.scale) + p_hi = _scipy_norm.cdf(hi, loc=self.mu, scale=self.scale) + norm_factor = p_hi - p_lo # probability mass within bounds + return _scipy_norm.pdf(values, loc=self.mu, scale=self.scale) / norm_factor + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- +DISTRIBUTIONS = { + 'uniform': Uniform(), + 'flat_in_log': Uniform(transform=Log10()), + 'uniform_in_sine': Uniform(transform=Sin()), + 'uniform_in_cosine': Uniform(transform=CosShift()), + 'kroupa': PowerLaw(ALPHA_IMF), + 'sana': PowerLaw(SANA_G, transform=Log10()), + 'sana_ecc': PowerLaw(SANA_ECC), + 'log_normal': TruncatedNormal(NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA, transform=Ln()), +} + + +def register(name, distribution): + """Register a distribution instance under ``name``. + + Parameters + ---------- + name : `str` + Key used to refer to the distribution (e.g. from a + :class:`~cosmic.sample.stroopwafel.parameter_space.Parameter`). + distribution : `Distribution` + The distribution instance to register. + """ + if not isinstance(distribution, Distribution): + raise TypeError( + f"register() expects a Distribution instance, got " + f"{type(distribution).__name__}." + ) + DISTRIBUTIONS[name] = distribution + + +def get_distribution(dist): + """Resolve a name or instance to a :class:`Distribution`. + + Parameters + ---------- + dist : `str` or `Distribution` + Either a key in :data:`DISTRIBUTIONS` or a distribution instance + (returned unchanged). + + Returns + ------- + `Distribution` + The resolved distribution instance. + """ + if isinstance(dist, Distribution): + return dist + try: + return DISTRIBUTIONS[dist] + except KeyError: + raise KeyError( + f"Unknown distribution {dist!r}. Registered names: " + f"{sorted(DISTRIBUTIONS)}. Alternatively pass a Distribution " + f"instance directly." + ) from None diff --git a/src/cosmic/sample/stroopwafel/examples/bh_star.py b/src/cosmic/sample/stroopwafel/examples/bh_star.py index 0ad86347a..32e777b63 100644 --- a/src/cosmic/sample/stroopwafel/examples/bh_star.py +++ b/src/cosmic/sample/stroopwafel/examples/bh_star.py @@ -116,13 +116,13 @@ # ------------------------------------------------------------------ params = ParameterSpace([ # --- orbital / stellar --- - Parameter('mass_1', 5.0, 150.0, sampler='kroupa', prior='kroupa'), - Parameter('q', 0.01, 1.0, sampler='uniform', prior='uniform'), - Parameter('porb', 0.15, 5.5, sampler='sana', prior='sana'), - Parameter('ecc', 1e-9, 0.99999999, sampler='sana_ecc', prior='sana_ecc'), - Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log', prior='flat_in_log'), + Parameter('mass_1', 5.0, 150.0, dist='kroupa'), + Parameter('q', 0.01, 1.0, dist='uniform'), + Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), # ~1.4 d to ~316 000 d + Parameter('ecc', 1e-9, 0.99999999, dist='sana_ecc'), + Parameter('metallicity', 0.0001, 0.03, dist='flat_in_log'), # --- primary natal kick magnitude only --- - Parameter('natal_kick_1', 0.1, 5000.0, sampler='log_normal', prior='log_normal'), + Parameter('natal_kick_1', 0.1, 5000.0, dist='log_normal'), ]) # ------------------------------------------------------------------ diff --git a/src/cosmic/sample/stroopwafel/examples/example_bhbh.py b/src/cosmic/sample/stroopwafel/examples/example_bhbh.py index 3278e82fe..f84fa5234 100644 --- a/src/cosmic/sample/stroopwafel/examples/example_bhbh.py +++ b/src/cosmic/sample/stroopwafel/examples/example_bhbh.py @@ -61,11 +61,11 @@ # Define parameter space # ------------------------------------------------------------------ params = ParameterSpace([ - Parameter('mass_1', 5.0, 150.0, sampler='kroupa', prior='kroupa'), - Parameter('q', 0.0, 1.0, sampler='uniform', prior='uniform'), - Parameter('porb', 0.15, 5.5, sampler='sana', prior='sana'), - Parameter('ecc', 1e-9, 0.99999999, sampler='sana_ecc', prior='sana_ecc'), - Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log', prior='flat_in_log'), + Parameter('mass_1', 5.0, 150.0, dist='kroupa'), + Parameter('q', 0.0, 1.0, dist='uniform'), + Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), # ~1.4 d to ~316 000 d + Parameter('ecc', 1e-9, 0.99999999, dist='sana_ecc'), + Parameter('metallicity', 0.0001, 0.03, dist='flat_in_log'), ]) # ------------------------------------------------------------------ diff --git a/src/cosmic/sample/stroopwafel/meson.build b/src/cosmic/sample/stroopwafel/meson.build index f99e1caa2..691dc69aa 100644 --- a/src/cosmic/sample/stroopwafel/meson.build +++ b/src/cosmic/sample/stroopwafel/meson.build @@ -1,14 +1,12 @@ python_sources = [ '__init__.py', 'constants.py', + 'distributions.py', 'engine.py', 'mixture_model.py', 'parameter_space.py', 'presets.py', - 'priors.py', 'rejection.py', - 'samplers.py', - 'transforms.py', ] py3.install_sources( diff --git a/src/cosmic/sample/stroopwafel/parameter_space.py b/src/cosmic/sample/stroopwafel/parameter_space.py index e32081981..6462fcdd3 100644 --- a/src/cosmic/sample/stroopwafel/parameter_space.py +++ b/src/cosmic/sample/stroopwafel/parameter_space.py @@ -7,10 +7,7 @@ import numpy as np from dataclasses import dataclass -from .samplers import SAMPLERS -from .priors import PRIORS -from .transforms import to_sampling_space, to_physical_space, transform_bounds -from .constants import ALPHA_IMF, SANA_G, SANA_ECC +from .distributions import Distribution, get_distribution @dataclass @@ -22,29 +19,35 @@ class Parameter: name : `str` Name of the parameter (used for column ordering). min_value : `float` - Lower bound. For most samplers this is in physical space (e.g. - solar masses for ``'kroupa'``, eccentricity for ``'sana_ecc'``). - - **Exception:** ``'sana'`` (orbital period) operates in - log10(period / days) internally, so ``min_value`` and - ``max_value`` must be passed in log10 space. For example, to - span periods from ~1.4 d to ~316 000 d use - ``min_value=0.15, max_value=5.5`` (i.e. log10 of those values). + Lower bound, always in physical space (e.g. solar masses for + ``'kroupa'``, days for ``'sana'``, km/s for ``'log_normal'``). The + parameter's distribution maps this into sampling space via its + transform. max_value : `float` - Upper bound (same unit convention as ``min_value``). - sampler : `str`, optional - Name of the sampling distribution, by default ``'uniform'`` - prior : `str`, optional - Name of the prior distribution, by default ``'uniform'`` + Upper bound, in physical space (same convention as ``min_value``). + dist : `str` or `~cosmic.sample.stroopwafel.distributions.Distribution`, optional + The prior distribution, given either as a name registered in + :data:`~cosmic.sample.stroopwafel.distributions.DISTRIBUTIONS` + (e.g. ``'kroupa'``, ``'sana'``, ``'flat_in_log'``) or as a + :class:`~cosmic.sample.stroopwafel.distributions.Distribution` + instance for custom priors. By default ``'uniform'``. + + Attributes + ---------- + distribution : `Distribution` + The resolved distribution instance. + lo, hi : `float` + The bounds in sampling space (``min_value``/``max_value`` passed + through ``distribution.transform``). """ name: str min_value: float max_value: float - sampler: str = 'uniform' - prior: str = 'uniform' + dist: "str | Distribution" = 'uniform' def __post_init__(self): - self.lo, self.hi = transform_bounds(self.min_value, self.max_value, self.sampler) + self.distribution = get_distribution(self.dist) + self.lo, self.hi = self.distribution.transform.bounds(self.min_value, self.max_value) class ParameterSpace: @@ -104,8 +107,7 @@ def sample(self, n, rng=None): mask = np.ones(n, dtype=bool) for i, p in enumerate(self.params): - sampler_fn = SAMPLERS[p.sampler] - col = sampler_fn(n, p.lo, p.hi, rng=rng) + col = p.distribution.sample(n, p.lo, p.hi, rng=rng) samples[:, i] = col mask &= (col >= p.lo) & (col <= p.hi) @@ -144,7 +146,7 @@ def to_physical(self, samples): """ result = samples.copy() for i, p in enumerate(self.params): - result[:, i] = to_physical_space(samples[:, i], p.sampler) + result[:, i] = p.distribution.transform.to_physical(samples[:, i]) return result def to_sampling(self, samples): @@ -162,7 +164,7 @@ def to_sampling(self, samples): """ result = samples.copy() for i, p in enumerate(self.params): - result[:, i] = to_sampling_space(samples[:, i], p.sampler) + result[:, i] = p.distribution.transform.to_sampling(samples[:, i]) return result def compute_prior(self, samples): @@ -183,8 +185,7 @@ def compute_prior(self, samples): """ log_prior = np.zeros(len(samples)) for i, p in enumerate(self.params): - prior_fn = PRIORS[p.prior] - col_prior = prior_fn(samples[:, i], p.lo, p.hi) + col_prior = p.distribution.pdf(samples[:, i], p.lo, p.hi) # Clamp to avoid log(0) log_prior += np.log(np.maximum(col_prior, 1e-300)) return np.exp(log_prior) @@ -209,84 +210,10 @@ def compute_sigma(self, hit_samples, average_density_one_dim): sigmas = np.empty((K, self.ndim)) for i, p in enumerate(self.params): - col = hit_samples[:, i] - if p.sampler in ('kroupa', 'sana', 'sana_ecc'): - sigmas[:, i] = self._sigma_power_law(p, col, average_density_one_dim) - else: - # uniform, flat_in_log, etc: sigma = avg_density / prior - prior_fn = PRIORS[p.prior] - prior_vals = prior_fn(col, p.lo, p.hi) - sigmas[:, i] = average_density_one_dim / prior_vals + try: + sigmas[:, i] = p.distribution.sigma( + hit_samples[:, i], p.lo, p.hi, average_density_one_dim + ) + except ValueError as e: + raise ValueError(f"Parameter '{p.name}': {e}") from e return sigmas - - def _sigma_power_law(self, param, values, avg_density): - """Compute sigma for power-law distributions. - - Maps each hit value to its CDF position, steps by ``avg_density`` - in CDF space, maps back to parameter space, and returns the - maximum of the two resulting distances. - - The CDF of a power-law p(x) ∝ x^α on [lo, hi] is - - F(x) = (x^a − lo^a) / (hi^a − lo^a), a = α + 1 - - so the inverse is F⁻¹(u) = (u·(hi^a − lo^a) + lo^a)^{1/a}. - All intermediate quantities stay in [lo^a, hi^a], avoiding the - catastrophic cancellation that occurred in the previous - ``inv_lo``-based formulation when ``lo`` was very small. - - Parameters - ---------- - param : `Parameter` - The parameter definition for this dimension. - values : `numpy.ndarray` - (K,) array of hit values in sampling space. - avg_density : `float` - Characteristic inter-sample spacing. - - Returns - ------- - `numpy.ndarray` - (K,) array of sigma values. - - Raises - ------ - `ValueError` - If ``param.lo <= 0`` (the CDF formula requires a positive - lower bound) or if the sampler is not a recognised power law. - """ - if param.sampler == 'kroupa': - alpha = ALPHA_IMF - elif param.sampler == 'sana': - alpha = SANA_G - elif param.sampler == 'sana_ecc': - alpha = SANA_ECC - else: - raise ValueError(f"Unknown power-law sampler: {param.sampler}") - - if param.lo <= 0: - raise ValueError( - f"Parameter '{param.name}' has lo={param.lo} <= 0 but uses " - f"the '{param.sampler}' power-law sampler. The CDF formula " - f"requires lo > 0. For 'sana' (period), pass the log10(P) " - f"lower bound directly, e.g. min_value=0.15 rather than " - f"min_value=np.log10(0.15)." - ) - - a = alpha + 1 # 0.55 for sana_ecc, 0.45 for sana, -1.3 for kroupa - lo_a = float(param.lo) ** a - hi_a = float(param.hi) ** a - range_a = hi_a - lo_a # always finite; no huge intermediates - - # --- Forward CDF: F(x) = (x^a - lo_a) / range_a ---------------- - u = np.clip((np.power(values, a) - lo_a) / range_a, 0.0, 1.0) - - # --- Step in CDF space ------------------------------------------ - u_right = np.clip(u + avg_density, 0.0, 1.0) - u_left = np.clip(u - avg_density, 0.0, 1.0) - - # --- Inverse CDF: F^{-1}(u) = (u * range_a + lo_a)^{1/a} ------- - x_right = np.power(u_right * range_a + lo_a, 1.0 / a) - x_left = np.power(u_left * range_a + lo_a, 1.0 / a) - - return np.maximum(np.abs(x_right - values), np.abs(x_left - values)) diff --git a/src/cosmic/sample/stroopwafel/priors.py b/src/cosmic/sample/stroopwafel/priors.py deleted file mode 100644 index 1305d05f3..000000000 --- a/src/cosmic/sample/stroopwafel/priors.py +++ /dev/null @@ -1,202 +0,0 @@ -"""Vectorized prior probability density functions. - -Each function takes ``(values, lo, hi)`` where ``values`` is an (N,) ndarray -and ``lo``/``hi`` are the bounds in the sampling (transformed) space, and -returns an (N,) ndarray of prior probability densities. -""" -import numpy as np -from scipy.stats import norm as _scipy_norm -from .constants import ALPHA_IMF, SANA_G, SANA_ECC, NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA - - -def uniform(values, lo, hi): - """Compute the uniform prior probability density. - - Parameters - ---------- - values : `numpy.ndarray` - (N,) array of sample values (unused, density is constant). - lo : `float` - Lower bound of the uniform distribution. - hi : `float` - Upper bound of the uniform distribution. - - Returns - ------- - `numpy.ndarray` - (N,) array of constant prior densities. - """ - return np.full(len(values), 1.0 / (hi - lo)) - - -def flat_in_log(values, lo, hi): - """Compute the flat-in-log prior probability density. - - The prior is uniform in log10 space; bounds are already in log10. - - Parameters - ---------- - values : `numpy.ndarray` - (N,) array of sample values in log10 space (unused, density is - constant). - lo : `float` - Lower bound in log10 space. - hi : `float` - Upper bound in log10 space. - - Returns - ------- - `numpy.ndarray` - (N,) array of constant prior densities. - """ - return np.full(len(values), 1.0 / (hi - lo)) - - -def kroupa(values, lo, hi): - """Compute the Kroupa IMF power-law prior probability density. - - Parameters - ---------- - values : `numpy.ndarray` - (N,) array of sample values. - lo : `float` - Lower bound of the distribution. - hi : `float` - Upper bound of the distribution. - - Returns - ------- - `numpy.ndarray` - (N,) array of prior densities. - """ - a = ALPHA_IMF - norm = (a + 1) / (hi**(a + 1) - lo**(a + 1)) - return norm * np.power(values, a) - - -def sana(values, lo, hi): - """Compute the Sana orbital period power-law prior probability density. - - Parameters - ---------- - values : `numpy.ndarray` - (N,) array of sample values. - lo : `float` - Lower bound of the distribution. - hi : `float` - Upper bound of the distribution. - - Returns - ------- - `numpy.ndarray` - (N,) array of prior densities. - """ - a = SANA_G - norm = (a + 1) / (hi**(a + 1) - lo**(a + 1)) - return norm * np.power(values, a) - - -def sana_ecc(values, lo, hi): - """Compute the Sana eccentricity power-law prior probability density. - - Parameters - ---------- - values : `numpy.ndarray` - (N,) array of sample values. - lo : `float` - Lower bound of the distribution. - hi : `float` - Upper bound of the distribution. - - Returns - ------- - `numpy.ndarray` - (N,) array of prior densities. - """ - a = SANA_ECC - norm = (a + 1) / (hi**(a + 1) - lo**(a + 1)) - return norm * np.power(values, a) - - -def uniform_in_sine(values, lo, hi): - """Compute the uniform-in-sine prior probability density. - - Parameters - ---------- - values : `numpy.ndarray` - (N,) array of sample values in sine space (unused, density is - constant). - lo : `float` - Lower bound in sine space. - hi : `float` - Upper bound in sine space. - - Returns - ------- - `numpy.ndarray` - (N,) array of constant prior densities. - """ - return np.full(len(values), 1.0 / (hi - lo)) - - -def uniform_in_cosine(values, lo, hi): - """Compute the uniform-in-cosine prior probability density. - - Parameters - ---------- - values : `numpy.ndarray` - (N,) array of sample values in cosine space (unused, density is - constant). - lo : `float` - Lower bound in cosine space. - hi : `float` - Upper bound in cosine space. - - Returns - ------- - `numpy.ndarray` - (N,) array of constant prior densities. - """ - return np.full(len(values), 1.0 / (hi - lo)) - - -def log_normal(values, lo, hi): - """Normal PDF in ln-space for a (truncated) log-normal natal kick prior. - - The kick magnitude v [km/s] follows LogNormal(``NATAL_KICK_LOG_MU``, - ``NATAL_KICK_LOG_SIGMA``), so the sampling-space variable ``x = ln(v)`` - has a (truncated) normal distribution. The returned density is the - normal PDF evaluated at ``values``, normalised so that it integrates to 1 - over the sampling-space interval ``[lo, hi]``. - - Parameters - ---------- - values : `numpy.ndarray` - (N,) array of sample values in ``ln(v / km s⁻¹)`` space. - lo : `float` - Lower bound in ``ln(v / km s⁻¹)`` space. - hi : `float` - Upper bound in ``ln(v / km s⁻¹)`` space. - - Returns - ------- - `numpy.ndarray` - (N,) array of prior probability densities. - """ - p_lo = _scipy_norm.cdf(lo, loc=NATAL_KICK_LOG_MU, scale=NATAL_KICK_LOG_SIGMA) - p_hi = _scipy_norm.cdf(hi, loc=NATAL_KICK_LOG_MU, scale=NATAL_KICK_LOG_SIGMA) - norm_factor = p_hi - p_lo # probability mass within bounds - return _scipy_norm.pdf(values, loc=NATAL_KICK_LOG_MU, scale=NATAL_KICK_LOG_SIGMA) / norm_factor - - -# Registry mapping string names to functions -PRIORS = { - 'uniform': uniform, - 'flat_in_log': flat_in_log, - 'kroupa': kroupa, - 'sana': sana, - 'sana_ecc': sana_ecc, - 'uniform_in_sine': uniform_in_sine, - 'uniform_in_cosine': uniform_in_cosine, - 'log_normal': log_normal, -} diff --git a/src/cosmic/sample/stroopwafel/samplers.py b/src/cosmic/sample/stroopwafel/samplers.py deleted file mode 100644 index 837a779bf..000000000 --- a/src/cosmic/sample/stroopwafel/samplers.py +++ /dev/null @@ -1,233 +0,0 @@ -"""Vectorized sampling functions for each distribution type. - -Each function draws ``n`` samples between ``lo`` and ``hi`` bounds (in the -sampling/transformed space) and returns an (N,) ndarray. -""" -import numpy as np -from scipy.stats import norm as _scipy_norm -from .constants import ALPHA_IMF, SANA_G, SANA_ECC, NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA - - -def uniform(n, lo, hi, rng=None): - """Draw samples from a uniform distribution. - - Parameters - ---------- - n : `int` - Number of samples to draw. - lo : `float` - Lower bound of the uniform distribution. - hi : `float` - Upper bound of the uniform distribution. - rng : `numpy.random.Generator`, optional - Random number generator, by default None - - Returns - ------- - `numpy.ndarray` - (N,) array of uniform samples. - """ - rng = rng or np.random.default_rng() - return rng.uniform(lo, hi, n) - - -def flat_in_log(n, lo, hi, rng=None): - """Sample uniformly in log10 space. - - Parameters - ---------- - n : `int` - Number of samples to draw. - lo : `float` - Lower bound (already in log10 space). - hi : `float` - Upper bound (already in log10 space). - rng : `numpy.random.Generator`, optional - Random number generator, by default None - - Returns - ------- - `numpy.ndarray` - (N,) array of samples in log10 space. - """ - rng = rng or np.random.default_rng() - return rng.uniform(lo, hi, n) - - -def kroupa(n, lo, hi, rng=None): - """Inverse CDF sampling from a Kroupa-like power law p(x) ~ x^alpha. - - Parameters - ---------- - n : `int` - Number of samples to draw. - lo : `float` - Lower bound of the distribution. - hi : `float` - Upper bound of the distribution. - rng : `numpy.random.Generator`, optional - Random number generator, by default None - - Returns - ------- - `numpy.ndarray` - (N,) array of power-law distributed samples. - """ - rng = rng or np.random.default_rng() - u = rng.uniform(0, 1, n) - a = ALPHA_IMF + 1 - return np.power(u * (hi**a - lo**a) + lo**a, 1.0 / a) - - -def sana(n, lo, hi, rng=None): - """Inverse CDF sampling from the Sana orbital period distribution. - - The Sana et al. (2012) distribution is a power law in log10(period), - so this sampler operates entirely in log10(period / days) space. - ``lo`` and ``hi`` must therefore be given as log10 values (e.g. - ``lo=0.15, hi=5.5`` spans ~1.4 d to ~316 000 d). - - Parameters - ---------- - n : `int` - Number of samples to draw. - lo : `float` - Lower bound in log10(period / days). - hi : `float` - Upper bound in log10(period / days). - rng : `numpy.random.Generator`, optional - Random number generator, by default None - - Returns - ------- - `numpy.ndarray` - (N,) array of samples from the Sana period distribution. - """ - rng = rng or np.random.default_rng() - u = rng.uniform(0, 1, n) - a = SANA_G + 1 - return np.power(u * (hi**a - lo**a) + lo**a, 1.0 / a) - - -def sana_ecc(n, lo, hi, rng=None): - """Inverse CDF sampling from the Sana eccentricity distribution. - - Parameters - ---------- - n : `int` - Number of samples to draw. - lo : `float` - Lower bound of the eccentricity distribution. - hi : `float` - Upper bound of the eccentricity distribution. - rng : `numpy.random.Generator`, optional - Random number generator, by default None - - Returns - ------- - `numpy.ndarray` - (N,) array of samples from the Sana eccentricity distribution. - """ - rng = rng or np.random.default_rng() - u = rng.uniform(0, 1, n) - a = SANA_ECC + 1 - return np.power(u * (hi**a - lo**a) + lo**a, 1.0 / a) - - -def uniform_in_sine(n, lo, hi, rng=None): - """Sample uniformly in sine-transformed space. - - Parameters - ---------- - n : `int` - Number of samples to draw. - lo : `float` - Lower bound in sine space. - hi : `float` - Upper bound in sine space. - rng : `numpy.random.Generator`, optional - Random number generator, by default None - - Returns - ------- - `numpy.ndarray` - (N,) array of uniform samples in sine space. - """ - rng = rng or np.random.default_rng() - return rng.uniform(lo, hi, n) - - -def uniform_in_cosine(n, lo, hi, rng=None): - """Sample uniformly in cosine-transformed space. - - Parameters - ---------- - n : `int` - Number of samples to draw. - lo : `float` - Lower bound in cosine space. - hi : `float` - Upper bound in cosine space. - rng : `numpy.random.Generator`, optional - Random number generator, by default None - - Returns - ------- - `numpy.ndarray` - (N,) array of uniform samples in cosine space. - """ - rng = rng or np.random.default_rng() - return rng.uniform(lo, hi, n) - - -def log_normal(n, lo, hi, rng=None): - """Inverse-CDF sampling from a (truncated) log-normal natal kick distribution. - - The kick magnitude v follows LogNormal(``NATAL_KICK_LOG_MU``, - ``NATAL_KICK_LOG_SIGMA``), so ``ln(v)`` is normally distributed. - ``lo`` and ``hi`` are bounds in ``ln(v)`` space (i.e. the natural log of - the physical bounds in km/s), and sampling is restricted to that range - via the truncated-normal inverse CDF. - - With ``NATAL_KICK_LOG_MU = 5.67`` and ``NATAL_KICK_LOG_SIGMA = 0.59``, - the median kick is ``exp(5.67) ≈ 291 km/s`` (Hobbs et al. 2005 / Fryer - et al.). Physical bounds of ``[0.1, 5000] km/s`` map to - ``lo ≈ –2.30``, ``hi ≈ 8.52`` in sampling space, which captures - essentially all of the probability mass. - - Parameters - ---------- - n : `int` - Number of samples to draw. - lo : `float` - Lower bound in ``ln(v / km s⁻¹)`` space. - hi : `float` - Upper bound in ``ln(v / km s⁻¹)`` space. - rng : `numpy.random.Generator`, optional - Random number generator, by default None - - Returns - ------- - `numpy.ndarray` - (N,) array of samples in ``ln(v)`` space. - """ - rng = rng or np.random.default_rng() - # Truncated-normal inverse CDF: map uniform draws to [CDF(lo), CDF(hi)], - # then apply the inverse normal CDF. - p_lo = _scipy_norm.cdf(lo, loc=NATAL_KICK_LOG_MU, scale=NATAL_KICK_LOG_SIGMA) - p_hi = _scipy_norm.cdf(hi, loc=NATAL_KICK_LOG_MU, scale=NATAL_KICK_LOG_SIGMA) - u = rng.uniform(p_lo, p_hi, n) - return _scipy_norm.ppf(u, loc=NATAL_KICK_LOG_MU, scale=NATAL_KICK_LOG_SIGMA) - - -# Registry mapping string names to functions -SAMPLERS = { - 'uniform': uniform, - 'flat_in_log': flat_in_log, - 'kroupa': kroupa, - 'sana': sana, - 'sana_ecc': sana_ecc, - 'uniform_in_sine': uniform_in_sine, - 'uniform_in_cosine': uniform_in_cosine, - 'log_normal': log_normal, -} diff --git a/src/cosmic/sample/stroopwafel/transforms.py b/src/cosmic/sample/stroopwafel/transforms.py deleted file mode 100644 index 8e204b62a..000000000 --- a/src/cosmic/sample/stroopwafel/transforms.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Vectorized scale transformations between physical and sampling space. - -Some distributions (flat_in_log, sana, uniform_in_sine, etc.) sample in a -transformed space (e.g., log10). These functions convert between the two -representations. All functions operate on (N,) column arrays for a given -parameter. -""" -import numpy as np - - -def to_sampling_space(values, sampler_type): - """Convert physical-space values to the sampling (transformed) space. - - Parameters - ---------- - values : `numpy.ndarray` - (N,) array of values in physical space. - sampler_type : `str` - Name of the sampling distribution (e.g., ``'flat_in_log'``, - ``'uniform_in_sine'``). - - Returns - ------- - `numpy.ndarray` - (N,) array of values in sampling space. - """ - if sampler_type in ('flat_in_log', 'sana'): - return np.log10(values) - elif sampler_type == 'uniform_in_sine': - return np.sin(values) - elif sampler_type == 'uniform_in_cosine': - # Angles are measured from –π/2 to π/2 (e.g. declination-like - # coordinates), so the sampling variable is cos(θ + π/2) = –sin(θ). - # The round-trip is exact for θ ∈ [–π/2, π/2]. - return np.cos(values + np.pi / 2) - elif sampler_type == 'log_normal': - # Sampling space is ln(v); physical space is v [km/s]. - return np.log(values) - return values - - -def to_physical_space(values, sampler_type): - """Convert sampling-space values back to physical space. - - Parameters - ---------- - values : `numpy.ndarray` - (N,) array of values in sampling space. - sampler_type : `str` - Name of the sampling distribution (e.g., ``'flat_in_log'``, - ``'uniform_in_sine'``). - - Returns - ------- - `numpy.ndarray` - (N,) array of values in physical space. - """ - if sampler_type in ('flat_in_log', 'sana'): - return np.power(10.0, values) - elif sampler_type == 'uniform_in_sine': - return np.arcsin(values) - elif sampler_type == 'uniform_in_cosine': - # Inverse of cos(θ + π/2): arccos(u) – π/2 - return np.arccos(values) - np.pi / 2 - elif sampler_type == 'log_normal': - # Inverse of ln: exp(x) gives v [km/s]. - return np.exp(values) - return values - - -def transform_bounds(lo, hi, sampler_type): - """Get the bounds in sampling space for a given physical-space range. - - Parameters - ---------- - lo : `float` - Lower bound in physical space. - hi : `float` - Upper bound in physical space. - sampler_type : `str` - Name of the sampling distribution. - - Returns - ------- - lo_transformed : `float` - Lower bound in sampling space. - hi_transformed : `float` - Upper bound in sampling space. - """ - if sampler_type == 'flat_in_log': - return np.log10(lo), np.log10(hi) - elif sampler_type == 'uniform_in_sine': - return -1.0, 1.0 - elif sampler_type == 'uniform_in_cosine': - return -1.0, 1.0 - elif sampler_type == 'log_normal': - # Physical bounds [lo, hi] in km/s → sampling bounds in ln(km/s). - return np.log(lo), np.log(hi) - # kroupa, uniform: bounds are already in physical space. - # sana, sana_ecc: bounds are passed by the caller in the native - # sampling space (log10(period) for sana, eccentricity for sana_ecc), - # so no further transformation is needed here. - return lo, hi diff --git a/src/cosmic/tests/meson.build b/src/cosmic/tests/meson.build index 7cd621809..a56f6e0d8 100644 --- a/src/cosmic/tests/meson.build +++ b/src/cosmic/tests/meson.build @@ -1,7 +1,9 @@ python_sources = [ 'test_evolve.py', + 'test_kick.py', 'test_match.py', 'test_sample.py', + 'test_stroopwafel.py', 'test_utils.py' ] From 4feca691ce5303040d545c787aceab497f23f874 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Sat, 27 Jun 2026 22:45:32 -0400 Subject: [PATCH 27/45] move tests to main folder --- .../stroopwafel/examples/test_new_modules.py | 265 ----------- src/cosmic/tests/test_stroopwafel.py | 410 ++++++++++++++++++ 2 files changed, 410 insertions(+), 265 deletions(-) delete mode 100644 src/cosmic/sample/stroopwafel/examples/test_new_modules.py create mode 100644 src/cosmic/tests/test_stroopwafel.py diff --git a/src/cosmic/sample/stroopwafel/examples/test_new_modules.py b/src/cosmic/sample/stroopwafel/examples/test_new_modules.py deleted file mode 100644 index ac7406afc..000000000 --- a/src/cosmic/sample/stroopwafel/examples/test_new_modules.py +++ /dev/null @@ -1,265 +0,0 @@ -"""Tests for the new vectorized STROOPWAFEL modules. - -Run with: python -m pytest tests/test_new_modules.py -v -Or just: python tests/test_new_modules.py -""" -import sys -import os -import math -import numpy as np - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) # adds COSMIC-stroopwafel/ - -from stroopwafel import ParameterSpace, Parameter -from stroopwafel.samplers import SAMPLERS -from stroopwafel.priors import PRIORS -from stroopwafel.transforms import to_sampling_space, to_physical_space, transform_bounds -from stroopwafel.mixture_model import GaussianMixture -from stroopwafel.rejection import get_zams_radius, default_reject -from cosmic.output import COSMICStroopOutput -from stroopwafel.constants import ( - R_COEFF, ZSOL, R_SOL_TO_AU, ALPHA_IMF, SANA_G, SANA_ECC -) - - -def make_default_params(): - return ParameterSpace([ - Parameter('mass_1', 5.0, 150.0, sampler='kroupa', prior='kroupa'), - Parameter('q', 0.0, 1.0, sampler='uniform', prior='uniform'), - Parameter('porb', 0.15, 5.5, sampler='sana', prior='sana'), - Parameter('ecc', 1e-9, 0.99999999, sampler='sana_ecc', prior='sana_ecc'), - Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log', prior='flat_in_log'), - ]) - - -# ==================================================================== -# ParameterSpace tests -# ==================================================================== - -def test_parameter_space_names_sorted(): - params = make_default_params() - assert params.names == sorted(params.names) - assert params.ndim == 5 - - -def test_sampling_produces_valid_arrays(): - params = make_default_params() - rng = np.random.default_rng(42) - samples, mask = params.sample(1000, rng=rng) - assert samples.shape == (1000, 5) - assert mask.shape == (1000,) - assert mask.dtype == bool - - -def test_roundtrip_transform(): - params = make_default_params() - rng = np.random.default_rng(42) - samples, _ = params.sample(1000, rng=rng) - physical = params.to_physical(samples) - back = params.to_sampling(physical) - assert np.allclose(samples, back, atol=1e-12) - - -def test_prior_positive(): - params = make_default_params() - rng = np.random.default_rng(42) - samples, mask = params.sample(1000, rng=rng) - priors = params.compute_prior(samples[mask]) - assert np.all(priors > 0) - assert np.all(np.isfinite(priors)) - - -def test_in_bounds(): - params = make_default_params() - rng = np.random.default_rng(42) - samples, mask = params.sample(1000, rng=rng) - mask2 = params.in_bounds(samples) - np.testing.assert_array_equal(mask, mask2) - - -# ==================================================================== -# Default rejection function tests -# ==================================================================== - -def test_default_reject_basic(): - param_names = ['ecc', 'mass_1', 'metallicity', 'porb', 'q'] - # Create a sample that should NOT be rejected (wide binary) - samples = np.array([[0.1, 20.0, 0.014, 100.0, 0.5]]) # reasonable binary - derived = { - 'mass_2': np.array([10.0]), - 'metallicity_1': np.array([0.014]), - 'metallicity_2': np.array([0.014]), - 'separation': np.array([50.0]), # AU - } - rejected = default_reject(samples, derived, param_names) - assert not rejected[0], "Wide binary should not be rejected" - - -def test_default_reject_low_mass(): - param_names = ['ecc', 'mass_1', 'metallicity', 'porb', 'q'] - samples = np.array([[0.1, 20.0, 0.014, 100.0, 0.001]]) - derived = { - 'mass_2': np.array([0.02]), # below minimum secondary mass - 'metallicity_1': np.array([0.014]), - 'metallicity_2': np.array([0.014]), - 'separation': np.array([50.0]), - } - rejected = default_reject(samples, derived, param_names) - assert rejected[0], "Low mass secondary should be rejected" - - -# ==================================================================== -# GaussianMixture tests -# ==================================================================== - -def test_mixture_from_hits(): - params = make_default_params() - rng = np.random.default_rng(42) - samples, _ = params.sample(1000, rng=rng) - # Use moderate samples as "hits" - hits = samples[400:410] # 10 hits - - avg_density = 1.0 / np.power(1000, 1.0 / params.ndim) - mixture = GaussianMixture.from_hits(hits, params, avg_density) - - assert mixture.n_components == 10 - assert mixture.means.shape == (10, 5) - assert mixture.covariances.shape == (10, 5, 5) - assert np.allclose(np.sum(mixture.alphas), 1.0) - assert np.allclose(mixture.alphas, 0.1) - - -def test_mixture_pdf_nonnegative(): - params = make_default_params() - rng = np.random.default_rng(42) - samples, _ = params.sample(1000, rng=rng) - hits = samples[400:410] - - avg_density = 1.0 / np.power(1000, 1.0 / params.ndim) - mixture = GaussianMixture.from_hits(hits, params, avg_density) - - pdf_vals = mixture.pdf(samples[:100]) - assert np.all(pdf_vals >= 0) - assert np.all(np.isfinite(pdf_vals)) - - -def test_mixture_sample(): - params = make_default_params() - rng = np.random.default_rng(42) - samples, _ = params.sample(1000, rng=rng) - hits = samples[400:410] - - avg_density = 1.0 / np.power(1000, 1.0 / params.ndim) - mixture = GaussianMixture.from_hits(hits, params, avg_density) - - msamp, mmask, midx = mixture.sample(2000, params, rng=rng) - assert msamp.shape[1] == 5 - assert len(mmask) == len(msamp) - assert len(midx) == len(msamp) - assert np.all(midx >= 0) - assert np.all(midx < 10) - - -# ==================================================================== -# STROOPWAFELResult tests -# ==================================================================== - -def test_result_hit_rate(): - import pandas as pd - N = 100 - weights = np.ones(N) - is_hit = np.zeros(N, dtype=bool) - is_hit[:10] = True - bin_nums = np.arange(N) - result = COSMICStroopOutput( - bpp=pd.DataFrame({'bin_num': bin_nums}), - bcm=pd.DataFrame({'bin_num': bin_nums}), - initC=pd.DataFrame({'bin_num': bin_nums}), - kick_info=pd.DataFrame({'bin_num': bin_nums}), - samples=np.zeros((N, 1)), - param_names=['mass_1'], - weights=weights, - is_hit=is_hit, - generation=np.zeros(N, dtype=int), - gaussian_idx=np.full(N, -1, dtype=int), - num_explored=N, - num_hits=10, - fraction_explored=1.0, - ) - assert abs(result.hit_rate - 0.1) < 1e-10 - - -# ==================================================================== -# Prior comparison with old code -# ==================================================================== - -def test_kroupa_prior_matches(): - """Test that vectorized kroupa prior matches old scalar version.""" - from stroopwafel.priors import kroupa - lo, hi = 5.0, 150.0 - values = np.linspace(5.1, 149.9, 100) - vec_result = kroupa(values, lo, hi) - - # Old scalar - norm = (ALPHA_IMF + 1) / (hi**(ALPHA_IMF + 1) - lo**(ALPHA_IMF + 1)) - old_result = norm * values**ALPHA_IMF - np.testing.assert_allclose(vec_result, old_result, atol=1e-12) - - -def test_sana_prior_matches(): - from stroopwafel.priors import sana - lo, hi = 0.15, 5.5 - values = np.linspace(0.2, 5.4, 100) - vec_result = sana(values, lo, hi) - - norm = (SANA_G + 1) / (hi**(SANA_G + 1) - lo**(SANA_G + 1)) - old_result = norm * values**SANA_G - np.testing.assert_allclose(vec_result, old_result, atol=1e-12) - - -# ==================================================================== -# Performance benchmark -# ==================================================================== - -def bench_zams_radius(): - """Benchmark vectorized vs scalar ZAMS radius.""" - import time - N = 10000 - masses = np.random.uniform(1, 100, N) - mets = np.random.uniform(0.0001, 0.03, N) - - # Vectorized - start = time.time() - _ = get_zams_radius(masses, mets) - vec_time = time.time() - start - - # Scalar - start = time.time() - _ = [old_get_zams_radius(m, z) for m, z in zip(masses, mets)] - scalar_time = time.time() - start - - speedup = scalar_time / vec_time if vec_time > 0 else float('inf') - print(f"\nZAMS radius benchmark (N={N}):") - print(f" Vectorized: {vec_time*1000:.1f}ms") - print(f" Scalar: {scalar_time*1000:.1f}ms") - print(f" Speedup: {speedup:.0f}x") - - -if __name__ == '__main__': - # Run all tests - test_funcs = [v for k, v in sorted(globals().items()) if k.startswith('test_')] - passed = 0 - failed = 0 - for func in test_funcs: - try: - func() - print(f" [PASS] {func.__name__}") - passed += 1 - except Exception as e: - print(f" [FAIL] {func.__name__}: {e}") - failed += 1 - - print(f"\n{passed} passed, {failed} failed") - - if failed == 0: - bench_zams_radius() diff --git a/src/cosmic/tests/test_stroopwafel.py b/src/cosmic/tests/test_stroopwafel.py new file mode 100644 index 000000000..3de5fecaf --- /dev/null +++ b/src/cosmic/tests/test_stroopwafel.py @@ -0,0 +1,410 @@ +"""Unit tests for the STROOPWAFEL adaptive importance sampling module. + +Covers the composable distribution/transform design +(`cosmic.sample.stroopwafel.distributions`), the vectorized +`ParameterSpace`, physical rejection, the Gaussian mixture model, and the +`COSMICStroopOutput` result type. +""" + +__author__ = 'Tom Wagg ' + +import unittest +import numpy as np +import pandas as pd +from scipy.stats import norm, kstest +from scipy.integrate import trapezoid + +from cosmic.sample.stroopwafel import ParameterSpace, Parameter +from cosmic.sample.stroopwafel.distributions import ( + DISTRIBUTIONS, register, get_distribution, + Distribution, Uniform, PowerLaw, TruncatedNormal, + Identity, Log10, Ln, Sin, CosShift, +) +from cosmic.sample.stroopwafel.mixture_model import GaussianMixture +from cosmic.sample.stroopwafel.rejection import get_zams_radius, default_reject +from cosmic.sample.stroopwafel.constants import ( + ALPHA_IMF, SANA_G, SANA_ECC, NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA, +) +from cosmic.output import COSMICStroopOutput + +HALF_PI = np.pi / 2 +KS_PVALUE_MIN = 0.01 # samplers must not be rejected against their own CDF + + +# -------------------------------------------------------------------------- +# Theoretical CDFs used for goodness-of-fit tests +# -------------------------------------------------------------------------- +def _uniform_cdf(x, lo, hi): + return (x - lo) / (hi - lo) + + +def _powerlaw_cdf(x, alpha, lo, hi): + a = alpha + 1 + return (np.power(x, a) - lo**a) / (hi**a - lo**a) + + +def _truncnorm_cdf(x, mu, scale, lo, hi): + lo_c, hi_c = norm.cdf(lo, mu, scale), norm.cdf(hi, mu, scale) + return (norm.cdf(x, mu, scale) - lo_c) / (hi_c - lo_c) + + +def canonical_space(): + """The reference 6-D parameter space exercising every built-in dist.""" + return ParameterSpace([ + Parameter('mass_1', 5.0, 150.0, dist='kroupa'), + Parameter('q', 0.0, 1.0, dist='uniform'), + Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), + Parameter('ecc', 1e-9, 0.99999999, dist='sana_ecc'), + Parameter('metallicity', 1e-4, 0.03, dist='flat_in_log'), + Parameter('natal_kick_1', 0.1, 5000.0, dist='log_normal'), + ]) + + +class Triangular(Distribution): + """A user-defined distribution, used to test extensibility.""" + + def sample(self, n, lo, hi, rng=None): + rng = rng or np.random.default_rng() + return rng.triangular(lo, 0.5 * (lo + hi), hi, n) + + def pdf(self, values, lo, hi): + mode = 0.5 * (lo + hi) + height = 2.0 / (hi - lo) + return np.where(values < mode, + height * (values - lo) / (mode - lo), + height * (hi - values) / (hi - mode)) + + +# ========================================================================== +# Transforms +# ========================================================================== +class TestTransforms(unittest.TestCase): + + def test_identity(self): + t = Identity() + x = np.array([-3.0, 0.0, 7.5]) + np.testing.assert_array_equal(t.to_sampling(x), x) + np.testing.assert_array_equal(t.to_physical(x), x) + self.assertEqual(t.bounds(2.0, 9.0), (2.0, 9.0)) + + def test_log10(self): + t = Log10() + x = np.geomspace(1e-4, 1e3, 50) + np.testing.assert_allclose(t.to_physical(t.to_sampling(x)), x, rtol=1e-12) + lo, hi = t.bounds(1e-4, 1e3) + self.assertAlmostEqual(lo, -4.0) + self.assertAlmostEqual(hi, 3.0) + + def test_ln(self): + t = Ln() + x = np.geomspace(0.1, 5000.0, 50) + np.testing.assert_allclose(t.to_physical(t.to_sampling(x)), x, rtol=1e-12) + lo, hi = t.bounds(0.1, 5000.0) + self.assertAlmostEqual(lo, np.log(0.1)) + self.assertAlmostEqual(hi, np.log(5000.0)) + + def test_sin(self): + t = Sin() + x = np.linspace(-HALF_PI, HALF_PI, 50) + np.testing.assert_allclose(t.to_physical(t.to_sampling(x)), x, atol=1e-12) + self.assertEqual(t.bounds(-HALF_PI, HALF_PI), (-1.0, 1.0)) + + def test_cosshift_bounds_are_sorted(self): + # CosShift is monotonically decreasing, so bounds() must swap ends. + t = CosShift() + x = np.linspace(-HALF_PI, HALF_PI, 50) + np.testing.assert_allclose(t.to_physical(t.to_sampling(x)), x, atol=1e-12) + lo, hi = t.bounds(-HALF_PI, HALF_PI) + self.assertLess(lo, hi) + np.testing.assert_allclose([lo, hi], [-1.0, 1.0], atol=1e-12) + + +# ========================================================================== +# Distributions +# ========================================================================== +class TestDistributions(unittest.TestCase): + + def test_registry_contents(self): + self.assertEqual( + set(DISTRIBUTIONS), + {'uniform', 'flat_in_log', 'uniform_in_sine', 'uniform_in_cosine', + 'kroupa', 'sana', 'sana_ecc', 'log_normal'}, + ) + + def test_builtin_compositions(self): + self.assertIsInstance(DISTRIBUTIONS['kroupa'], PowerLaw) + self.assertIsInstance(DISTRIBUTIONS['kroupa'].transform, Identity) + self.assertIsInstance(DISTRIBUTIONS['sana'], PowerLaw) + self.assertIsInstance(DISTRIBUTIONS['sana'].transform, Log10) + self.assertIsInstance(DISTRIBUTIONS['flat_in_log'], Uniform) + self.assertIsInstance(DISTRIBUTIONS['flat_in_log'].transform, Log10) + self.assertIsInstance(DISTRIBUTIONS['log_normal'], TruncatedNormal) + self.assertIsInstance(DISTRIBUTIONS['log_normal'].transform, Ln) + self.assertEqual(DISTRIBUTIONS['kroupa'].alpha, ALPHA_IMF) + self.assertEqual(DISTRIBUTIONS['sana'].alpha, SANA_G) + self.assertEqual(DISTRIBUTIONS['sana_ecc'].alpha, SANA_ECC) + + def test_uniform_pdf_constant(self): + d = Uniform() + vals = np.linspace(2.0, 8.0, 100) + np.testing.assert_allclose(d.pdf(vals, 2.0, 8.0), 1.0 / 6.0) + + def test_powerlaw_pdf_matches_closed_form(self): + for alpha, lo, hi in [(ALPHA_IMF, 5.0, 150.0), + (SANA_G, 0.15, 5.5), + (SANA_ECC, 1e-9, 0.999)]: + vals = np.linspace(lo * 1.01, hi * 0.99, 100) + got = PowerLaw(alpha).pdf(vals, lo, hi) + a = alpha + 1 + expected = (a / (hi**a - lo**a)) * vals**alpha + np.testing.assert_allclose(got, expected, rtol=1e-12) + + def test_powerlaw_pdf_normalised(self): + d = PowerLaw(ALPHA_IMF) + lo, hi = 5.0, 150.0 + x = np.linspace(lo, hi, 200000) + self.assertAlmostEqual(trapezoid(d.pdf(x, lo, hi), x), 1.0, places=4) + + def test_truncnorm_pdf_normalised(self): + d = TruncatedNormal(NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA) + lo, hi = 2.0, 9.0 + x = np.linspace(lo, hi, 200000) + self.assertAlmostEqual(trapezoid(d.pdf(x, lo, hi), x), 1.0, places=4) + + def test_uniform_sampler_ks(self): + rng = np.random.default_rng(0) + s = Uniform().sample(20000, 2.0, 8.0, rng=rng) + self.assertTrue(np.all((s >= 2.0) & (s <= 8.0))) + p = kstest(s, lambda v: _uniform_cdf(v, 2.0, 8.0)).pvalue + self.assertGreater(p, KS_PVALUE_MIN) + + def test_powerlaw_sampler_ks(self): + for alpha, lo, hi, seed in [(ALPHA_IMF, 5.0, 150.0, 1), + (SANA_G, 0.15, 5.5, 2), + (SANA_ECC, 1e-9, 0.999, 3)]: + rng = np.random.default_rng(seed) + s = PowerLaw(alpha).sample(20000, lo, hi, rng=rng) + self.assertTrue(np.all((s >= lo) & (s <= hi))) + p = kstest(s, lambda v: _powerlaw_cdf(v, alpha, lo, hi)).pvalue + self.assertGreater(p, KS_PVALUE_MIN, msg=f"alpha={alpha}") + + def test_truncnorm_sampler_ks(self): + rng = np.random.default_rng(4) + mu, scale, lo, hi = NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA, 2.0, 9.0 + s = TruncatedNormal(mu, scale).sample(20000, lo, hi, rng=rng) + self.assertTrue(np.all((s >= lo) & (s <= hi))) + p = kstest(s, lambda v: _truncnorm_cdf(v, mu, scale, lo, hi)).pvalue + self.assertGreater(p, KS_PVALUE_MIN) + + def test_default_sigma_is_avg_density_over_pdf(self): + d = Uniform() + vals = np.linspace(2.0, 8.0, 50) + ad = 0.05 + np.testing.assert_allclose( + d.sigma(vals, 2.0, 8.0, ad), ad / d.pdf(vals, 2.0, 8.0), rtol=1e-12, + ) + + def test_powerlaw_sigma_positive_and_raises_on_nonpositive_lo(self): + d = PowerLaw(ALPHA_IMF) + vals = np.linspace(6.0, 140.0, 100) + sig = d.sigma(vals, 5.0, 150.0, 0.05) + self.assertTrue(np.all(sig > 0) and np.all(np.isfinite(sig))) + with self.assertRaises(ValueError): + d.sigma(vals, 0.0, 150.0, 0.05) + + +# ========================================================================== +# Registry / extensibility +# ========================================================================== +class TestRegistry(unittest.TestCase): + + def test_get_distribution_passthrough(self): + d = PowerLaw(-1.7) + self.assertIs(get_distribution(d), d) + + def test_get_distribution_unknown_raises(self): + with self.assertRaises(KeyError): + get_distribution('does_not_exist') + + def test_register_and_lookup(self): + self.addCleanup(lambda: DISTRIBUTIONS.pop('test_steep_imf', None)) + register('test_steep_imf', PowerLaw(-2.7)) + p = Parameter('m', 5.0, 100.0, dist='test_steep_imf') + self.assertEqual(p.distribution.alpha, -2.7) + + def test_register_rejects_non_instance(self): + with self.assertRaises(TypeError): + register('bad', PowerLaw) # a class, not an instance + + def test_custom_distribution_instance(self): + p = Parameter('x', 1.0, 10.0, dist=PowerLaw(-1.7)) + s = p.distribution.sample(500, p.lo, p.hi, rng=np.random.default_rng(0)) + self.assertTrue(np.all((s >= 1.0) & (s <= 10.0))) + + def test_custom_subclass_end_to_end(self): + # A user-defined Distribution composed with a transform, run through + # the full ParameterSpace pipeline. + ps = ParameterSpace([Parameter('t', 1.0, 100.0, dist=Triangular(transform=Log10()))]) + samples, mask = ps.sample(5000, rng=np.random.default_rng(1)) + phys = ps.to_physical(samples[mask]) + self.assertTrue(np.all((phys >= 1.0) & (phys <= 100.0))) + self.assertTrue(np.all(ps.compute_prior(samples[mask]) > 0)) + sig = ps.compute_sigma(samples[mask][:100], 0.05) + self.assertTrue(np.all(sig > 0) and np.all(np.isfinite(sig))) + + +# ========================================================================== +# ParameterSpace +# ========================================================================== +class TestParameterSpace(unittest.TestCase): + + def setUp(self): + self.ps = canonical_space() + + def test_names_sorted_and_ndim(self): + self.assertEqual(self.ps.names, sorted(self.ps.names)) + self.assertEqual(self.ps.ndim, 6) + + def test_idx(self): + for i, name in enumerate(self.ps.names): + self.assertEqual(self.ps.idx(name), i) + + def test_sample_shapes_and_mask(self): + samples, mask = self.ps.sample(1000, rng=np.random.default_rng(42)) + self.assertEqual(samples.shape, (1000, 6)) + self.assertEqual(mask.shape, (1000,)) + self.assertEqual(mask.dtype, bool) + + def test_roundtrip_transform(self): + samples, _ = self.ps.sample(1000, rng=np.random.default_rng(42)) + back = self.ps.to_sampling(self.ps.to_physical(samples)) + np.testing.assert_allclose(back, samples, atol=1e-12) + + def test_prior_positive_finite(self): + samples, mask = self.ps.sample(1000, rng=np.random.default_rng(42)) + priors = self.ps.compute_prior(samples[mask]) + self.assertTrue(np.all(priors > 0) and np.all(np.isfinite(priors))) + + def test_in_bounds_matches_sample_mask(self): + samples, mask = self.ps.sample(1000, rng=np.random.default_rng(42)) + np.testing.assert_array_equal(self.ps.in_bounds(samples), mask) + + def test_sana_bounds_linear_to_log10(self): + # sana bounds are given in linear (physical) days; the Log10 transform + # maps them into log10-period sampling space. + porb = next(p for p in self.ps.params if p.name == 'porb') + self.assertAlmostEqual(porb.lo, 0.15) + self.assertAlmostEqual(porb.hi, 5.5) + + def test_compute_sigma_positive_finite(self): + samples, mask = self.ps.sample(5000, rng=np.random.default_rng(7)) + hits = samples[mask][:200] + ad = 1.0 / np.power(5000, 1.0 / self.ps.ndim) + sig = self.ps.compute_sigma(hits, ad) + self.assertEqual(sig.shape, hits.shape) + self.assertTrue(np.all(sig > 0) and np.all(np.isfinite(sig))) + + def test_compute_sigma_reports_parameter_name_on_error(self): + # A sana period with a 1-day lower bound maps to log10(1)=0, which the + # power-law sigma cannot handle; the error should name the parameter. + ps = ParameterSpace([Parameter('porb', 1.0, 1000.0, dist='sana')]) + samples, mask = ps.sample(100, rng=np.random.default_rng(0)) + with self.assertRaisesRegex(ValueError, 'porb'): + ps.compute_sigma(samples[mask][:10], 0.05) + + +# ========================================================================== +# Rejection +# ========================================================================== +class TestRejection(unittest.TestCase): + + PARAM_NAMES = ['ecc', 'mass_1', 'metallicity', 'porb', 'q'] + + def test_get_zams_radius_positive_finite(self): + masses = np.array([1.0, 10.0, 30.0, 100.0]) + mets = np.array([0.02, 0.02, 0.001, 0.014]) + r = get_zams_radius(masses, mets) + self.assertEqual(r.shape, (4,)) + self.assertTrue(np.all(r > 0) and np.all(np.isfinite(r))) + + def test_wide_binary_not_rejected(self): + samples = np.array([[0.1, 20.0, 0.014, 100.0, 0.5]]) + derived = { + 'mass_2': np.array([10.0]), + 'metallicity_1': np.array([0.014]), + 'metallicity_2': np.array([0.014]), + 'separation': np.array([50.0]), + } + rejected = default_reject(samples, derived, self.PARAM_NAMES) + self.assertFalse(rejected[0]) + + def test_low_mass_secondary_rejected(self): + samples = np.array([[0.1, 20.0, 0.014, 100.0, 0.001]]) + derived = { + 'mass_2': np.array([0.02]), # below the 0.08 Msun minimum + 'metallicity_1': np.array([0.014]), + 'metallicity_2': np.array([0.014]), + 'separation': np.array([50.0]), + } + rejected = default_reject(samples, derived, self.PARAM_NAMES) + self.assertTrue(rejected[0]) + + +# ========================================================================== +# Gaussian mixture model +# ========================================================================== +class TestGaussianMixture(unittest.TestCase): + + def setUp(self): + self.ps = canonical_space() + rng = np.random.default_rng(42) + samples, mask = self.ps.sample(20000, rng=rng) + self.valid = samples[mask] + self.hits = self.valid[:10] + self.ad = 1.0 / np.power(20000, 1.0 / self.ps.ndim) + + def test_from_hits_shapes(self): + gm = GaussianMixture.from_hits(self.hits, self.ps, self.ad) + self.assertEqual(gm.n_components, 10) + self.assertEqual(gm.means.shape, (10, 6)) + self.assertEqual(gm.covariances.shape, (10, 6, 6)) + self.assertAlmostEqual(float(np.sum(gm.alphas)), 1.0) + + def test_pdf_nonnegative_finite(self): + gm = GaussianMixture.from_hits(self.hits, self.ps, self.ad) + pdf = gm.pdf(self.valid[:200]) + self.assertTrue(np.all(pdf >= 0) and np.all(np.isfinite(pdf))) + + def test_sample_shapes_and_idx_range(self): + gm = GaussianMixture.from_hits(self.hits, self.ps, self.ad) + msamp, mmask, midx = gm.sample(2000, self.ps, rng=np.random.default_rng(1)) + self.assertEqual(msamp.shape[1], 6) + self.assertEqual(len(mmask), len(msamp)) + self.assertEqual(len(midx), len(msamp)) + self.assertTrue(np.all((midx >= 0) & (midx < 10))) + + +# ========================================================================== +# COSMICStroopOutput +# ========================================================================== +class TestStroopOutput(unittest.TestCase): + + def test_hit_rate(self): + N = 100 + bin_nums = np.arange(N) + is_hit = np.zeros(N, dtype=bool) + is_hit[:10] = True + frame = pd.DataFrame({'bin_num': bin_nums}) + result = COSMICStroopOutput( + bpp=frame, bcm=frame, initC=frame, kick_info=frame, + samples=np.zeros((N, 1)), param_names=['mass_1'], + weights=np.ones(N), is_hit=is_hit, + generation=np.zeros(N, dtype=int), + gaussian_idx=np.full(N, -1, dtype=int), + num_explored=N, num_hits=10, fraction_explored=1.0, + ) + self.assertAlmostEqual(result.hit_rate, 0.1) + + +if __name__ == '__main__': + unittest.main() From c5bbe8e25b6bb0235470e10b58dcb1e39d9399b0 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Sun, 28 Jun 2026 21:39:35 -0400 Subject: [PATCH 28/45] docs overhaul --- docs/pages/tutorials/adaptive.rst | 3 +- docs/pages/tutorials/adaptive/basics.rst | 43 ++-- docs/pages/tutorials/adaptive/checkpoint.rst | 146 ++++++++++++ .../tutorials/adaptive/distributions.rst | 207 ++++++++++++++++++ docs/pages/tutorials/adaptive/outputs.rst | 65 +++++- 5 files changed, 443 insertions(+), 21 deletions(-) create mode 100644 docs/pages/tutorials/adaptive/distributions.rst diff --git a/docs/pages/tutorials/adaptive.rst b/docs/pages/tutorials/adaptive.rst index a75c63032..85ba4b838 100644 --- a/docs/pages/tutorials/adaptive.rst +++ b/docs/pages/tutorials/adaptive.rst @@ -6,7 +6,8 @@ These tutorials will cover how to use the adaptive importance sampling method in .. toctree:: :maxdepth: 2 - + adaptive/basics + adaptive/distributions adaptive/outputs adaptive/checkpoint \ No newline at end of file diff --git a/docs/pages/tutorials/adaptive/basics.rst b/docs/pages/tutorials/adaptive/basics.rst index 8fa84dffc..e9ef612fd 100644 --- a/docs/pages/tutorials/adaptive/basics.rst +++ b/docs/pages/tutorials/adaptive/basics.rst @@ -27,7 +27,7 @@ When should I use this? ======================= Use ``STROOPWAFEL`` whenever you need a statistically representative sample of a rare binary -outcome and cannot afford the total binary count that flat Monte Carlo would require. +outcome and cannot afford the total binary count that a flat Monte Carlo draw would require. Example use cases include: * Double black holes (or neutron stars) merging within the Hubble time (GW sources) @@ -81,8 +81,9 @@ Define the parameter space The :class:`~cosmic.sample.stroopwafel.ParameterSpace` class defines which binary parameters are sampled and their distributions. Each -:class:`~cosmic.sample.stroopwafel.Parameter` specifies a name, physical-space bounds, a -sampling distribution, and a prior distribution. +:class:`~cosmic.sample.stroopwafel.Parameter` specifies a name, physical-space bounds, and +a distribution. The distribution sets both how the parameter is drawn and its prior +probability density — in importance sampling these are one and the same. .. code-block:: python @@ -90,14 +91,22 @@ sampling distribution, and a prior distribution. from cosmic.sample.stroopwafel import ParameterSpace, Parameter params = ParameterSpace([ - Parameter('mass_1', 5.0, 150.0, sampler='kroupa', prior='kroupa'), - Parameter('q', 0.01, 1.0, sampler='uniform', prior='uniform'), - Parameter('porb', 0.15, 5.5, sampler='sana', prior='sana'), - Parameter('ecc', 1e-9, 0.9999, sampler='sana_ecc', prior='sana_ecc'), - Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log',prior='flat_in_log'), + Parameter('mass_1', 5.0, 150.0, dist='kroupa'), + Parameter('q', 0.01, 1.0, dist='uniform'), + Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), + Parameter('ecc', 1e-9, 0.9999, dist='sana_ecc'), + Parameter('metallicity', 0.0001, 0.03, dist='flat_in_log'), ]) -Parameters are stored and returned in alphabetical order by name. Use ``params.names`` to check the order and ``params.index('param_name')`` to get the index of a particular parameter. +The ``dist`` argument names one of the built-in distributions. Several common initial-distribution choices are available by default - the Kroupa IMF, Sana orbital periods and eccentricities, flat-in-log metallicity, and more — and lets you define your own just as easily. See :ref:`adaptive_distributions` for the full list and how to add custom distributions. + +Parameters are stored and returned in the order you provide them, which sets the column order of every sample array. Use ``params.names`` to check the order and ``params.idx('param_name')`` to get the index of a particular parameter. + +.. note:: + + Bounds are always given in **physical** space. The ``'sana'`` period is sampled in + :math:`\log_{10}(P / \mathrm{day})`, so the example writes its bounds as ``10**(0.15)`` + and ``10**(5.5)`` - the distribution applies the :math:`\log_{10}` transform internally. Define any derived quantities @@ -105,8 +114,8 @@ Define any derived quantities COSMIC requires ``mass_2``, ``separation``, and ``metallicity`` in addition to the directly-sampled parameters. The ``compute_derived`` callback converts a ``(N, D)`` -array of physical-space samples and a sorted list of parameter names into a dictionary of -``(N,)`` arrays: +array of physical-space samples and the list of parameter names (in column order) into a +dictionary of ``(N,)`` arrays: .. code-block:: python @@ -251,11 +260,11 @@ orbital period, eccentricity, and metallicity. # Parameter space # ------------------------------------------------------------------ params = ParameterSpace([ - Parameter('mass_1', 5.0, 150.0, sampler='kroupa', prior='kroupa'), - Parameter('q', 0.01, 1.0, sampler='uniform', prior='uniform'), - Parameter('porb', 0.15, 5.5, sampler='sana', prior='sana'), - Parameter('ecc', 1e-9, 0.9999, sampler='sana_ecc', prior='sana_ecc'), - Parameter('metallicity', 0.0001, 0.03, sampler='flat_in_log',prior='flat_in_log'), + Parameter('mass_1', 5.0, 150.0, dist='kroupa'), + Parameter('q', 0.01, 1.0, dist='uniform'), + Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), + Parameter('ecc', 1e-9, 0.9999, dist='sana_ecc'), + Parameter('metallicity', 0.0001, 0.03, dist='flat_in_log'), ]) # ------------------------------------------------------------------ @@ -369,7 +378,7 @@ The EM step between generations can improve the mixture, but with diminishing re Saving your results =================== -Once you have your samples, you can save them to disk as an HDF5 file with the :meth:`~cosmic.output.COSMICSTROOPWAFELResult.save` method. This saves the parameter samples, derived quantities, and hit information in a compact format that can be loaded later for analysis. +Once you have your samples, you can save them to disk as an HDF5 file with the :meth:`~cosmic.output.COSMICStroopOutput.save` method. This saves the parameter samples, derived quantities, and hit information in a compact format that can be loaded later for analysis. .. code-block:: python diff --git a/docs/pages/tutorials/adaptive/checkpoint.rst b/docs/pages/tutorials/adaptive/checkpoint.rst index 3112fcdf3..6691a7750 100644 --- a/docs/pages/tutorials/adaptive/checkpoint.rst +++ b/docs/pages/tutorials/adaptive/checkpoint.rst @@ -6,3 +6,149 @@ Saving and loading checkpoints In this tutorial, we will cover how to save and load checkpoints during an adaptive importance sampling run in ``COSMIC``. This allows you to save the state of your simulation once the adaptation phase is complete, and then load it later to continue sampling from the same Gaussian mixture model (GMM) without having to redo the adaptation phase. This can be particularly useful if you want to run multiple sampling phases with the same adapted GMM (e.g. over multiple nodes on a cluster). +This tutorial assumes that you've already gone through :ref:`adaptive_basics`. + +Why split a run in two? +======================= + +A call to :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run` performs all three phases +back to back: exploration, adaptation, and refinement. Splitting the run lets you stop +after adaptation — once the Gaussian mixture has been fitted to the exploration hits — and +resume the (usually much larger) refinement phase separately. This is useful when you want +to + +* run exploration and refinement as **separate cluster jobs**, perhaps with different + walltimes or allocations; +* fan a single adapted mixture out across **several refinement jobs** on different nodes; or +* simply **inspect** the mixture before committing compute to refinement. + +Instead of :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run`, you use the two +multi-job entry points: +:meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run_exploration` (which returns a +checkpoint) and :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run_refinement`. + + +Stage 1 — explore, adapt, and save a checkpoint +=============================================== + +Set up the sampler exactly as you would for a normal run (see :ref:`adaptive_basics`), but +call :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run_exploration` instead of +``run()``. This runs the exploration and adaptation phases and returns a +:class:`~cosmic.output.STROOPWAFELCheckpoint`, which you then save to disk. + +.. code-block:: python + + from cosmic.sample.stroopwafel import AdaptiveSampler, ParameterSpace, Parameter + from cosmic.sample.stroopwafel.presets import any_dco + from cosmic.sample.stroopwafel.rejection import default_reject + + params = ParameterSpace([ + Parameter('mass_1', 5.0, 150.0, dist='kroupa'), + Parameter('q', 0.01, 1.0, dist='uniform'), + Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), + Parameter('ecc', 1e-9, 0.9999, dist='sana_ecc'), + Parameter('metallicity', 0.0001, 0.03, dist='flat_in_log'), + ]) + + def compute_derived(samples_physical, param_names): + idx = {name: i for i, name in enumerate(param_names)} + mass_1 = samples_physical[:, idx['mass_1']] + q = samples_physical[:, idx['q']] + porb = samples_physical[:, idx['porb']] + z = samples_physical[:, idx['metallicity']] + mass_2 = mass_1 * q + separation = ((porb / 365.25) ** 2 * (mass_1 + mass_2)) ** (1.0 / 3.0) + return {'mass_2': mass_2, 'metallicity_1': z, + 'metallicity_2': z, 'separation': separation} + + sampler = AdaptiveSampler( + parameter_space=params, + total_systems=500_000, + batch_size=1000, + BSEDict=BSEDict, + compute_derived=compute_derived, + reject_systems=default_reject, + is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), + output_path='output/explore', + nproc=4, + seed=42, + ) + + checkpoint = sampler.run_exploration() # exploration + adaptation only + checkpoint.save('checkpoint.h5') + +.. note:: + + If exploration finds no hits there is nothing to adapt, and the checkpoint's mixture + will be ``None``. In that case you should increase ``total_systems`` and re-run + exploration before attempting refinement (see the rules of thumb in + :ref:`adaptive_basics`). + + +Stage 2 — load the checkpoint and refine +======================================== + +In a second script (or cluster job) rebuild the sampler with +:meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.from_checkpoint`, then call +:meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run_refinement`. + +.. code-block:: python + + from cosmic.sample.stroopwafel import AdaptiveSampler + from cosmic.sample.stroopwafel.presets import any_dco + from cosmic.sample.stroopwafel.rejection import default_reject + + # `params`, `compute_derived`, and `BSEDict` must be available again here — + # in practice, import them from a shared module used by both jobs. + + sampler = AdaptiveSampler.from_checkpoint( + 'checkpoint.h5', + parameter_space=params, + batch_size=1000, + BSEDict=BSEDict, + compute_derived=compute_derived, + reject_systems=default_reject, + is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), + output_path='output/refine', + nproc=4, + seed=7, + ) + + result = sampler.run_refinement() + result.save('result.h5') + +The ``result`` is an ordinary :class:`~cosmic.output.COSMICStroopOutput` — identical in form +to what a single :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run` would have produced — +so you can analyse it exactly as described in :ref:`adaptive_outputs`. + + +What is and isn't stored in a checkpoint +======================================== + +A checkpoint stores everything that is *derived from running COSMIC*: + +* the fitted Gaussian mixture model, +* every exploration sample (in the internal sampling space) and its hit/bookkeeping flags, +* the COSMIC output tables (``bpp``, ``bcm``, ``initC``, ``kick_info``) for the explored + systems, and +* the scalar counters needed to compute unbiased weights later (``num_explored``, + ``fraction_explored``, ``prior_fraction_rejected``, ``total_systems``, ...). + +It deliberately does **not** store your Python callables or physics settings — the +``BSEDict``, ``compute_derived``, ``reject_systems``, and ``is_interesting`` functions are +not serialisable in general, so you must supply them again to +:meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.from_checkpoint`. Only the parameter +**names** are stored, and they are checked against the ``parameter_space`` you provide; a +mismatch raises a ``ValueError`` to stop you from refining against the wrong setup. + + +Reusing one checkpoint for several refinement jobs +================================================== + +Because :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.from_checkpoint` reads the +checkpoint without modifying it, you can launch any number of independent refinement jobs +from the same ``checkpoint.h5`` — for example with different ``seed`` values on different +nodes — and each will draw a fresh, independent refinement sample from the shared mixture. +You can also override the budget stored in the checkpoint by passing ``total_systems`` or +``n_generations`` to :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.from_checkpoint`. + diff --git a/docs/pages/tutorials/adaptive/distributions.rst b/docs/pages/tutorials/adaptive/distributions.rst new file mode 100644 index 000000000..643a3a46f --- /dev/null +++ b/docs/pages/tutorials/adaptive/distributions.rst @@ -0,0 +1,207 @@ +.. _adaptive_distributions: + +********************************* +Distributions and custom priors +********************************* + +This tutorial assumes that you've already gone through :ref:`adaptive_basics`. + +Every :class:`~cosmic.sample.stroopwafel.Parameter` is given a *distribution*, which does +three jobs: it draws samples, it evaluates the prior probability density, and it defines +the adaptive-sampling kernel width used during refinement. ``COSMIC`` provides a small set +of built-in distributions covering the usual initial-condition choices, and a simple, +composable interface for defining your own. + + +Built-in distributions +======================= + +Pass any of the following names as the ``dist`` argument to a +:class:`~cosmic.sample.stroopwafel.Parameter`: + +.. list-table:: + :header-rows: 1 + :widths: 18 22 60 + + * - Name + - Sampled as + - Description + * - ``'uniform'`` + - uniform + - Flat between the bounds. Good default for mass ratio, etc. + * - ``'flat_in_log'`` + - uniform in :math:`\log_{10}` + - Flat in the log of the parameter (e.g. metallicity). + * - ``'kroupa'`` + - power law, :math:`\alpha = -2.3` + - Kroupa initial mass function for the primary mass. + * - ``'sana'`` + - power law in :math:`\log_{10} P`, :math:`\alpha = -0.55` + - Sana et al. (2012) orbital-period distribution. + * - ``'sana_ecc'`` + - power law, :math:`\alpha = -0.45` + - Sana et al. (2012) eccentricity distribution. + * - ``'uniform_in_sine'`` + - uniform in :math:`\sin\theta` + - Isotropic angle (e.g. inclination-like coordinates). + * - ``'uniform_in_cosine'`` + - uniform in :math:`\cos\theta` + - Isotropic angle for declination-like coordinates. + * - ``'log_normal'`` + - log-normal + - Natal-kick magnitude, :math:`\ln v \sim \mathcal{N}(5.67, 0.59)`. + +The full registry is available programmatically as +:data:`cosmic.sample.stroopwafel.distributions.DISTRIBUTIONS`. + + +How distributions are built +=========================== + +Internally each distribution is a **base distribution** composed with a **coordinate +transform**: + +* the base distribution (:class:`~cosmic.sample.stroopwafel.distributions.Uniform`, + :class:`~cosmic.sample.stroopwafel.distributions.PowerLaw`, or + :class:`~cosmic.sample.stroopwafel.distributions.TruncatedNormal`) handles sampling and + the density in the *sampling space*; and +* the transform (:class:`~cosmic.sample.stroopwafel.distributions.Log10`, + :class:`~cosmic.sample.stroopwafel.distributions.Ln`, etc.) maps between the physical + space you specify bounds in and that sampling space. + +This is why, for example, ``'flat_in_log'`` is just a uniform distribution paired with a +:math:`\log_{10}` transform, and ``'sana'`` is a power law paired with the same transform. +You can build the same objects yourself: + +.. code-block:: python + + from cosmic.sample.stroopwafel.distributions import ( + Uniform, PowerLaw, TruncatedNormal, Log10, + ) + + Uniform(transform=Log10()) # equivalent to 'flat_in_log' + PowerLaw(-0.55, transform=Log10()) # equivalent to 'sana' + PowerLaw(-2.3) # equivalent to 'kroupa' + +Bounds are always given to a :class:`~cosmic.sample.stroopwafel.Parameter` in **physical** +space; the transform converts them into sampling space automatically (and round-trips +samples back to physical space before they are evolved by ``COSMIC``). + + +Defining your own distribution +============================== + +There are three ways to use a custom distribution, in increasing order of effort. + +1. Pass a distribution instance directly +---------------------------------------- + +The quickest option is to tweak one of the built-in base distributions and hand the +instance straight to a :class:`~cosmic.sample.stroopwafel.Parameter` via ``dist``. No +registration required: + +.. code-block:: python + + from cosmic.sample.stroopwafel import Parameter + from cosmic.sample.stroopwafel.distributions import PowerLaw + + # A steeper-than-Kroupa IMF for the primary mass + Parameter('mass_1', 5.0, 150.0, dist=PowerLaw(-2.7)) + + +2. Register a named distribution +-------------------------------- + +If you want to reuse a distribution across several parameter spaces — or simply refer to it +by a memorable name — register it once with +:func:`~cosmic.sample.stroopwafel.distributions.register`: + +.. code-block:: python + + from cosmic.sample.stroopwafel import Parameter + from cosmic.sample.stroopwafel.distributions import register, PowerLaw + + register('imf_steep', PowerLaw(-2.7)) + + # ... anywhere later + Parameter('mass_1', 5.0, 150.0, dist='imf_steep') + + +3. Write a new distribution class +--------------------------------- + +For a genuinely new functional form, subclass +:class:`~cosmic.sample.stroopwafel.distributions.Distribution` and implement two methods: + +``sample(n, lo, hi, rng)`` + Draw ``n`` samples in **sampling space**, restricted to ``[lo, hi]``. + +``pdf(values, lo, hi)`` + Return the prior density at ``values``, normalised over ``[lo, hi]`` in sampling space. + +Both ``lo`` and ``hi`` are bounds in sampling space — the parameter space has already +applied the transform, so you do not need to worry about it here. The example below +implements a truncated exponential distribution: + +.. code-block:: python + + import numpy as np + from cosmic.sample.stroopwafel import Parameter + from cosmic.sample.stroopwafel.distributions import Distribution + + class Exponential(Distribution): + """p(x) ∝ exp(-x / scale), truncated to [lo, hi].""" + + def __init__(self, scale, transform=None): + super().__init__(transform) + self.scale = scale + + def sample(self, n, lo, hi, rng=None): + rng = rng or np.random.default_rng() + u = rng.uniform(0, 1, n) + c_lo, c_hi = np.exp(-lo / self.scale), np.exp(-hi / self.scale) + return -self.scale * np.log(c_lo - u * (c_lo - c_hi)) # inverse CDF + + def pdf(self, values, lo, hi): + norm = self.scale * (np.exp(-lo / self.scale) - np.exp(-hi / self.scale)) + return np.exp(-values / self.scale) / norm + + Parameter('some_param', 0.0, 10.0, dist=Exponential(scale=2.0)) + +That is all that is required — your distribution now works everywhere the built-ins do, and +can be combined with any transform (``dist=Exponential(2.0, transform=Log10())``). + +.. note:: + + During refinement, ``STROOPWAFEL`` places a Gaussian kernel at each hit whose width is + set by :meth:`~cosmic.sample.stroopwafel.distributions.Distribution.sigma`. The default + implementation, ``avg_density / pdf``, is appropriate for almost all distributions and + is inherited automatically — you only need to override it if you have an exact + closed-form CDF you would rather step through (as + :class:`~cosmic.sample.stroopwafel.distributions.PowerLaw` does). + + +Custom transforms +================= + +Transforms are just as extensible. A transform implements ``to_sampling`` (physical → +sampling) and ``to_physical`` (sampling → physical); the corresponding bound conversion is +derived automatically and handles decreasing maps by swapping the endpoints. The built-in +transforms are :class:`~cosmic.sample.stroopwafel.distributions.Identity`, +:class:`~cosmic.sample.stroopwafel.distributions.Log10`, +:class:`~cosmic.sample.stroopwafel.distributions.Ln`, +:class:`~cosmic.sample.stroopwafel.distributions.Sin`, and +:class:`~cosmic.sample.stroopwafel.distributions.CosShift`. To add your own, subclass +:class:`~cosmic.sample.stroopwafel.distributions.Transform`: + +.. code-block:: python + + import numpy as np + from cosmic.sample.stroopwafel.distributions import Transform + + class Sqrt(Transform): + def to_sampling(self, values): + return np.sqrt(values) + + def to_physical(self, values): + return values ** 2 diff --git a/docs/pages/tutorials/adaptive/outputs.rst b/docs/pages/tutorials/adaptive/outputs.rst index e269e711c..22cc33a61 100644 --- a/docs/pages/tutorials/adaptive/outputs.rst +++ b/docs/pages/tutorials/adaptive/outputs.rst @@ -55,15 +55,38 @@ The class also stores other metadata that you may find useful. These include: - ``num_hits``, which is the total raw hit count across all phases - ``fraction_explored``, which is the fraction of total systems used for exploration. +In addition, two derived properties summarise the weighted population: + +- ``hit_rate``, the importance-weighted hit rate (an unbiased estimate of the true rate) +- ``hit_rate_uncertainty``, the standard error on that rate + How to interpret adaptive sampling weights ========================================== So now for the important intuition part. ``COSMIC`` and ``STROOPWAFEL`` have now provided you with a sample of a rare population by preferentially sampling the parameter space that you've specified. However, we of course want to account for the fact that this *is* a rare population. This is where the weights come in. Each sample is assigned an adaptive importance sampling weight, which tells you how rare this sample is (smaller weights are rarer). -.. admonition:: TODO +Concretely, each system :math:`x` is assigned an importance weight + +.. math:: + + w(x) = \frac{\pi(x)}{Q(x)}, \qquad + Q(x) = f_e\,\pi(x) + (1 - f_e)\,q(x), + +where :math:`\pi(x)` is the prior (astrophysical) probability density, :math:`q(x)` is the +density of the adapted Gaussian mixture, and :math:`f_e` is the fraction of the budget spent +on exploration. The denominator :math:`Q(x)` is therefore the *actual* density from which +the system was drawn — a blend of the broad prior (during exploration) and the concentrated +mixture (during refinement). - Add maths +In words, the weight is the ratio of how often a system *should* appear under the prior to +how often it *actually* appeared under the combined sampling scheme. A hit discovered deep +inside an oversampled region picks up a small weight (we drew far more of them than nature +would), while a system drawn straight from the prior has a weight close to one. Summed over +the population these weights turn any statistic into an unbiased estimator of the +corresponding prior-weighted quantity; for example ``sum(weights[is_hit]) / N`` is an +unbiased estimate of the true hit rate, which is exactly what +:attr:`~cosmic.output.COSMICStroopOutput.hit_rate` returns. This means that if you want to plot a true distribution of your sampled systems -- let's say the primary mass -- you need to use the weights in your plotting. @@ -86,4 +109,40 @@ This means that if you want to plot a true distribution of your sampled systems You should **always** apply your weights to any plot that you make. In histograms you can supply them directly. For scatter plots you could consider changing the size of your points, or using a 2D histogram or ``hexbin`` instead. Drawing a representative sample from your simulation -==================================================== \ No newline at end of file +==================================================== + +Applying weights at plot time is the right approach for visualising distributions, but +sometimes you want an actual *set of systems* that is representative of the underlying +population — for example to pass a fixed number of binaries into a downstream calculation. +Because the simulation deliberately oversamples the rare region, you cannot take the hits at +face value; you need to resample them in proportion to their weights. + +This is a standard weighted bootstrap: draw indices from the hit population with probability +proportional to their weights, with replacement. + +.. code-block:: python + + import numpy as np + from cosmic.output import COSMICStroopOutput + + results = COSMICStroopOutput.from_file("YOUR_SIMULATION.h5") + + # Restrict to hits and normalise their weights into probabilities + hit_idx = np.where(results.is_hit)[0] + probs = results.weights[hit_idx] / results.weights[hit_idx].sum() + + # Draw a representative sample of, say, 5000 systems + rng = np.random.default_rng(42) + chosen = rng.choice(hit_idx, size=5000, replace=True, p=probs) + + representative = results.samples[chosen] # (5000, D), in physical space + +The resulting ``representative`` array is distributed according to the true (prior-weighted) +population, so it can be histogrammed or analysed **without** any further weighting. Because +the draw is made with replacement, the same underlying system can appear more than once — +this is expected, and is the price of turning a weighted sample into an unweighted one. + +.. note:: + + A :meth:`~cosmic.output.COSMICStroopOutput.draw_representative_sample` convenience method + that wraps this resampling is planned; until it lands, use the weighted bootstrap above. \ No newline at end of file From e2f31fa92b44aa8e78f75390eedad2088030fcd3 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Sun, 28 Jun 2026 21:40:07 -0400 Subject: [PATCH 29/45] no alphabetical --- src/cosmic/sample/stroopwafel/examples/bh_star.py | 2 -- src/cosmic/sample/stroopwafel/parameter_space.py | 12 ++++++------ src/cosmic/sample/stroopwafel/rejection.py | 2 +- src/cosmic/tests/test_stroopwafel.py | 4 ---- 4 files changed, 7 insertions(+), 13 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/examples/bh_star.py b/src/cosmic/sample/stroopwafel/examples/bh_star.py index 32e777b63..d072e7252 100644 --- a/src/cosmic/sample/stroopwafel/examples/bh_star.py +++ b/src/cosmic/sample/stroopwafel/examples/bh_star.py @@ -111,8 +111,6 @@ # fills all omitted kick columns with the -100 sentinel so COSMIC draws those # components from its own prescription (kickflag=5 / sigma=265 km/s). # -# Note: ParameterSpace sorts parameters alphabetically, so the internal -# column order is fixed and independent of the order given here. # ------------------------------------------------------------------ params = ParameterSpace([ # --- orbital / stellar --- diff --git a/src/cosmic/sample/stroopwafel/parameter_space.py b/src/cosmic/sample/stroopwafel/parameter_space.py index 6462fcdd3..3d4c67560 100644 --- a/src/cosmic/sample/stroopwafel/parameter_space.py +++ b/src/cosmic/sample/stroopwafel/parameter_space.py @@ -53,19 +53,19 @@ def __post_init__(self): class ParameterSpace: """An ordered collection of Parameters with vectorized operations. - All methods operate on (N, D) numpy arrays where columns are ordered - alphabetically by parameter name. + All methods operate on (N, D) numpy arrays whose columns follow the order + in which the parameters were supplied. Parameters ---------- params : `list` of `Parameter` - List of parameter definitions. They will be sorted by name - internally. + List of parameter definitions. The column order of every (N, D) array + produced by this class matches the order of this list. """ def __init__(self, params): - # Sort by name for deterministic column ordering (same as old code) - self.params = sorted(params, key=lambda p: p.name) + # Columns follow the order the user supplied (no reordering). + self.params = list(params) self.names = [p.name for p in self.params] self._name_to_idx = {p.name: i for i, p in enumerate(self.params)} self.ndim = len(self.params) diff --git a/src/cosmic/sample/stroopwafel/rejection.py b/src/cosmic/sample/stroopwafel/rejection.py index e9acc7f06..2082b8fda 100644 --- a/src/cosmic/sample/stroopwafel/rejection.py +++ b/src/cosmic/sample/stroopwafel/rejection.py @@ -62,7 +62,7 @@ def default_reject(samples_physical, derived, param_names, ``'metallicity_2'``, and ``'separation'``, each mapping to an (N,) array. param_names : `list` of `str` - Sorted list of parameter names (used to find column indices). + Parameter names in column order (used to find column indices). min_secondary_mass : `float`, optional Minimum allowed secondary mass in solar masses, by default 0.08 diff --git a/src/cosmic/tests/test_stroopwafel.py b/src/cosmic/tests/test_stroopwafel.py index 3de5fecaf..01611edb8 100644 --- a/src/cosmic/tests/test_stroopwafel.py +++ b/src/cosmic/tests/test_stroopwafel.py @@ -261,10 +261,6 @@ class TestParameterSpace(unittest.TestCase): def setUp(self): self.ps = canonical_space() - def test_names_sorted_and_ndim(self): - self.assertEqual(self.ps.names, sorted(self.ps.names)) - self.assertEqual(self.ps.ndim, 6) - def test_idx(self): for i, name in enumerate(self.ps.names): self.assertEqual(self.ps.idx(name), i) From a7eb2186403625b036221069cc7db75c4e4c2e1b Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Sun, 28 Jun 2026 22:06:20 -0400 Subject: [PATCH 30/45] make derived stuff more general --- docs/pages/tutorials/adaptive/basics.rst | 104 ++++----- docs/pages/tutorials/adaptive/checkpoint.rst | 24 +- src/cosmic/sample/stroopwafel/engine.py | 214 ++++++++++++------ .../sample/stroopwafel/examples/bh_star.py | 39 +--- .../stroopwafel/examples/example_bhbh.py | 29 +-- .../sample/stroopwafel/mixture_model.py | 19 +- src/cosmic/sample/stroopwafel/rejection.py | 47 ++-- src/cosmic/tests/test_stroopwafel.py | 93 ++++++-- 8 files changed, 329 insertions(+), 240 deletions(-) diff --git a/docs/pages/tutorials/adaptive/basics.rst b/docs/pages/tutorials/adaptive/basics.rst index e9ef612fd..373fcedf1 100644 --- a/docs/pages/tutorials/adaptive/basics.rst +++ b/docs/pages/tutorials/adaptive/basics.rst @@ -109,42 +109,44 @@ Parameters are stored and returned in the order you provide them, which sets the and ``10**(5.5)`` - the distribution applies the :math:`\log_{10}` transform internally. -Define any derived quantities ------------------------------ +Complete the binary definition +------------------------------ -COSMIC requires ``mass_2``, ``separation``, and ``metallicity`` in addition to the -directly-sampled parameters. The ``compute_derived`` callback converts a ``(N, D)`` -array of physical-space samples and the list of parameter names (in column order) into a -dictionary of ``(N,)`` arrays: +Every binary handed to ``COSMIC`` is defined by **five** parameters: ``mass_1``, +``mass_2``, ``porb``, ``ecc``, and ``metallicity``. -.. code-block:: python +Each one must be provided in **exactly one** of two ways: either it is sampled (a +:class:`~cosmic.sample.stroopwafel.Parameter` with that name in your ``ParameterSpace``) or +it is returned by an optional ``derive_params`` function. If any of the five is neither +sampled nor derived, :class:`~cosmic.sample.stroopwafel.AdaptiveSampler` raises an error as soon as it is constructed. - def compute_derived(samples_physical, param_names): - """Convert sampled parameters to COSMIC inputs.""" - idx = {name: i for i, name in enumerate(param_names)} +In the parameter space above we sampled ``mass_1``, ``porb``, ``ecc``, and ``metallicity`` +directly, but sampled the mass *ratio* ``q`` rather than ``mass_2``. ``derive_params`` +fills in the gap. It receives a dictionary mapping each sampled name to its ``(N,)`` array +of physical values and returns a dictionary of the remaining parameters (a scalar is +broadcast to all ``N`` binaries): - mass_1 = samples_physical[:, idx['mass_1']] - q = samples_physical[:, idx['q']] - porb = samples_physical[:, idx['porb']] # physical days after 10^x transform - z = samples_physical[:, idx['metallicity']] +.. code-block:: python - mass_2 = mass_1 * q + def derive_params(sampled): + """Provide binary parameters not drawn from the ParameterSpace.""" + return {'mass_2': sampled['mass_1'] * sampled['q']} - # Kepler's third law: a [AU] from P [yr] and M [M_sun]; then convert to AU - separation = ((porb / 365.25) ** 2 * (mass_1 + mass_2)) ** (1.0 / 3.0) +This design makes it easy to adaptively sample in just a few dimensions while holding the +rest fixed. For instance, to explore *only* orbital period you would sample ``porb`` and +fix everything else: - return { - 'mass_2': mass_2, - 'metallicity_1': z, - 'metallicity_2': z, - 'separation': separation, # AU, required by default_reject - } +.. code-block:: python -.. note:: + params = ParameterSpace([ + Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), + ]) - The keys ``'mass_2'``, ``'metallicity_1'``, ``'metallicity_2'``, and ``'separation'`` - are expected by :func:`~cosmic.sample.stroopwafel.rejection.default_reject`. If you - supply a custom rejection function you are free to use different key names. + def derive_params(sampled): + return {'mass_1': 30.0, 'mass_2': 25.0, 'ecc': 0.0, 'metallicity': 0.02} + +If all five parameters are sampled directly, ``derive_params`` is unnecessary and can be +omitted entirely. Choose a rejection function @@ -160,19 +162,23 @@ checks: from cosmic.sample.stroopwafel.rejection import default_reject -It expects the arrays produced by ``compute_derived`` and returns a boolean mask where -``True`` means the system is rejected. For most use cases this default is appropriate. -If you need extra cuts — for example, discarding systems with very low metallicity or -imposing a minimum primary mass — you can wrap it: +It receives the assembled binary parameters as a dictionary (``mass_1``, ``mass_2``, +``porb``, ``ecc``, ``metallicity``) and returns a boolean mask where ``True`` means the +system is rejected; the orbital separation is computed internally from ``porb``. For most +use cases this default is appropriate. If you need extra cuts — for example imposing a +minimum secondary mass — you can wrap it: .. code-block:: python - def my_reject(samples_physical, derived, param_names): - base_mask = default_reject(samples_physical, derived, param_names) + def my_reject(binary_params): + base_mask = default_reject(binary_params) # additionally reject secondaries below 1 M_sun - base_mask |= (derived['mass_2'] < 1.0) + base_mask |= (binary_params['mass_2'] < 1.0) return base_mask +Passing ``reject_systems`` is optional; omit it (or pass ``None``) to skip physical +rejection entirely. + Identify what constitutes a hit ------------------------------- @@ -233,7 +239,7 @@ You can also write a fully custom hit function. For example, to find BH + stell Running the sampler =================== -With all the pieces in place, you can run the sampler with :class:`~cosmic.sample.stroopwafel.AdaptiveSampler`. The most important arguments are the parameter space, the total number of systems to evolve, the batch size, the BSE physics settings, the derived quantity function, the rejection function, and the hit function. See the API documentation (:class:`~cosmic.sample.stroopwafel.AdaptiveSampler`) for a full list of options. +With all the pieces in place, you can run the sampler with :class:`~cosmic.sample.stroopwafel.AdaptiveSampler`. The most important arguments are the parameter space, the total number of systems to evolve, the batch size, the BSE physics settings, the hit function, the ``derive_params`` function (if needed), and the rejection function. See the API documentation (:class:`~cosmic.sample.stroopwafel.AdaptiveSampler`) for a full list of options. The examples below demonstrate how you could go about this. @@ -268,22 +274,10 @@ orbital period, eccentricity, and metallicity. ]) # ------------------------------------------------------------------ - # Derived quantities + # Complete the binary definition (mass_2 from the sampled mass ratio) # ------------------------------------------------------------------ - def compute_derived(samples_physical, param_names): - idx = {name: i for i, name in enumerate(param_names)} - mass_1 = samples_physical[:, idx['mass_1']] - q = samples_physical[:, idx['q']] - porb = samples_physical[:, idx['porb']] - z = samples_physical[:, idx['metallicity']] - mass_2 = mass_1 * q - separation = ((porb / 365.25) ** 2 * (mass_1 + mass_2)) ** (1.0 / 3.0) - return { - 'mass_2': mass_2, - 'metallicity_1': z, - 'metallicity_2': z, - 'separation': separation, - } + def derive_params(sampled): + return {'mass_2': sampled['mass_1'] * sampled['q']} # ------------------------------------------------------------------ # Run @@ -293,9 +287,9 @@ orbital period, eccentricity, and metallicity. total_systems=50_000, batch_size=500, BSEDict=BSEDict, - compute_derived=compute_derived, - reject_systems=default_reject, is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), + derive_params=derive_params, + reject_systems=default_reject, output_path='output/bhbh', nproc=4, n_generations=1, @@ -313,7 +307,7 @@ BH + star binaries surviving 100 Myr For outcomes that are less extreme but still rare — such as persistent BH + star systems — STROOPWAFEL provides substantial efficiency gains over flat Monte Carlo. Using the same -parameter space, ``compute_derived``, and ``BSEDict`` as the previous examples, we can simply swap out the hit function to find BH + star systems that remain bound for at least 100 Myr after the BH forms: +parameter space, ``derive_params``, and ``BSEDict`` as the previous examples, we can simply swap out the hit function to find BH + star systems that remain bound for at least 100 Myr after the BH forms: .. code-block:: python @@ -326,9 +320,9 @@ parameter space, ``compute_derived``, and ``BSEDict`` as the previous examples, total_systems=20_000, batch_size=500, BSEDict=BSEDict, # reuse from BHBH example - compute_derived=compute_derived, # reuse from BHBH example - reject_systems=default_reject, is_interesting=bh_star_100myr, # we defined this earlier + derive_params=derive_params, # reuse from BHBH example + reject_systems=default_reject, output_path='output/bh_star', nproc=4, n_generations=1, diff --git a/docs/pages/tutorials/adaptive/checkpoint.rst b/docs/pages/tutorials/adaptive/checkpoint.rst index 6691a7750..5663e007d 100644 --- a/docs/pages/tutorials/adaptive/checkpoint.rst +++ b/docs/pages/tutorials/adaptive/checkpoint.rst @@ -50,25 +50,17 @@ call :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run_exploration` instead Parameter('metallicity', 0.0001, 0.03, dist='flat_in_log'), ]) - def compute_derived(samples_physical, param_names): - idx = {name: i for i, name in enumerate(param_names)} - mass_1 = samples_physical[:, idx['mass_1']] - q = samples_physical[:, idx['q']] - porb = samples_physical[:, idx['porb']] - z = samples_physical[:, idx['metallicity']] - mass_2 = mass_1 * q - separation = ((porb / 365.25) ** 2 * (mass_1 + mass_2)) ** (1.0 / 3.0) - return {'mass_2': mass_2, 'metallicity_1': z, - 'metallicity_2': z, 'separation': separation} + def derive_params(sampled): + return {'mass_2': sampled['mass_1'] * sampled['q']} sampler = AdaptiveSampler( parameter_space=params, total_systems=500_000, batch_size=1000, BSEDict=BSEDict, - compute_derived=compute_derived, - reject_systems=default_reject, is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), + derive_params=derive_params, + reject_systems=default_reject, output_path='output/explore', nproc=4, seed=42, @@ -98,7 +90,7 @@ In a second script (or cluster job) rebuild the sampler with from cosmic.sample.stroopwafel.presets import any_dco from cosmic.sample.stroopwafel.rejection import default_reject - # `params`, `compute_derived`, and `BSEDict` must be available again here — + # `params`, `derive_params`, and `BSEDict` must be available again here — # in practice, import them from a shared module used by both jobs. sampler = AdaptiveSampler.from_checkpoint( @@ -106,9 +98,9 @@ In a second script (or cluster job) rebuild the sampler with parameter_space=params, batch_size=1000, BSEDict=BSEDict, - compute_derived=compute_derived, - reject_systems=default_reject, is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), + derive_params=derive_params, + reject_systems=default_reject, output_path='output/refine', nproc=4, seed=7, @@ -135,7 +127,7 @@ A checkpoint stores everything that is *derived from running COSMIC*: ``fraction_explored``, ``prior_fraction_rejected``, ``total_systems``, ...). It deliberately does **not** store your Python callables or physics settings — the -``BSEDict``, ``compute_derived``, ``reject_systems``, and ``is_interesting`` functions are +``BSEDict``, ``derive_params``, ``reject_systems``, and ``is_interesting`` functions are not serialisable in general, so you must supply them again to :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.from_checkpoint`. Only the parameter **names** are stored, and they are checked against the ``parameter_space`` you provide; a diff --git a/src/cosmic/sample/stroopwafel/engine.py b/src/cosmic/sample/stroopwafel/engine.py index 01e5af234..5d767df37 100644 --- a/src/cosmic/sample/stroopwafel/engine.py +++ b/src/cosmic/sample/stroopwafel/engine.py @@ -29,17 +29,25 @@ class AdaptiveSampler: Number of systems evolved per COSMIC call. BSEDict : `dict` COSMIC binary stellar evolution parameters. - compute_derived : `callable` - Function with signature - ``(samples_physical, param_names) -> dict`` that computes - derived quantities (e.g., ``'mass_2'``, ``'separation'``). - reject_systems : `callable` - Function with signature - ``(samples_physical, derived, param_names) -> bool_mask`` - returning True for physically unacceptable systems. is_interesting : `callable` Function with signature ``(bpp) -> (n_hits, hit_bin_nums)`` identifying systems of interest from COSMIC output. + derive_params : `callable`, 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 + :attr:`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_systems : `callable`, 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 + :func:`~cosmic.sample.stroopwafel.rejection.default_reject`. By + default None (no physical rejection). output_path : `str`, optional Directory for output files, by default ``'output'`` nproc : `int`, optional @@ -63,8 +71,13 @@ class AdaptiveSampler: By default False. """ + #: The parameters that fully define a binary for COSMIC. Each must be + #: either sampled (a Parameter in ``parameter_space``) or returned by + #: ``derive_params``. + REQUIRED_PARAMS = ('mass_1', 'mass_2', 'porb', 'ecc', 'metallicity') + def __init__(self, parameter_space, total_systems, batch_size, BSEDict, - compute_derived, reject_systems, is_interesting, + is_interesting, derive_params=None, reject_systems=None, output_path='output', nproc=1, kappa=1.0, n_generations=1, mc_only=False, seed=None, only_save_hit_tables=False): @@ -72,7 +85,7 @@ def __init__(self, parameter_space, total_systems, batch_size, BSEDict, self.total_systems = total_systems self.batch_size = batch_size self.bse_dict = BSEDict - self.compute_derived_fn = compute_derived + self.derive_fn = derive_params self.reject_fn = reject_systems self.is_interesting_fn = is_interesting self.output_path = output_path @@ -103,6 +116,9 @@ def __init__(self, parameter_space, total_systems, batch_size, BSEDict, self._all_initC = [] self._all_kick_info = [] + # Fail fast if the binary model is under-specified. + self._validate_param_coverage() + def run(self): """Run the full STROOPWAFEL pipeline in a single call. @@ -164,8 +180,8 @@ def run_refinement(self): sampler = AdaptiveSampler.from_checkpoint( 'checkpoint.h5', parameter_space=params, batch_size=1000, - BSEDict=BSEDict, compute_derived=compute_derived, - reject_systems=default_reject, is_interesting=is_bh_star, + BSEDict=BSEDict, is_interesting=is_bh_star, + derive_params=derive_params, reject_systems=default_reject, ) result = sampler.run_refinement() result.save('result.h5') @@ -180,13 +196,13 @@ def run_refinement(self): @classmethod def from_checkpoint(cls, checkpoint, parameter_space, batch_size, BSEDict, - compute_derived, reject_systems, is_interesting, + is_interesting, derive_params=None, reject_systems=None, output_path='output', nproc=1, seed=None, total_systems=None, n_generations=None, only_save_hit_tables=False): """Create a sampler pre-loaded with state from a :class:`STROOPWAFELCheckpoint`. - The callable arguments (``BSEDict``, ``compute_derived``, etc.) are + The callable arguments (``BSEDict``, ``derive_params``, etc.) are not stored in the checkpoint and must be supplied again. All other state — explored samples, COSMIC output, mixture model, scalar counters — is restored from the checkpoint. @@ -203,11 +219,11 @@ def from_checkpoint(cls, checkpoint, parameter_space, batch_size, BSEDict, Systems per COSMIC call for the refinement phase. BSEDict : `dict` COSMIC binary stellar evolution parameters. - compute_derived : `callable` + is_interesting : `callable` Same signature as for the original sampler. - reject_systems : `callable` + derive_params : `callable`, optional Same signature as for the original sampler. - is_interesting : `callable` + reject_systems : `callable`, optional Same signature as for the original sampler. output_path : `str`, optional Directory for output files, by default ``'output'`` @@ -250,9 +266,9 @@ def from_checkpoint(cls, checkpoint, parameter_space, batch_size, BSEDict, total_systems=ts, batch_size=batch_size, BSEDict=BSEDict, - compute_derived=compute_derived, - reject_systems=reject_systems, is_interesting=is_interesting, + derive_params=derive_params, + reject_systems=reject_systems, output_path=output_path, nproc=nproc, n_generations=ng, @@ -318,6 +334,84 @@ def _load_checkpoint(self, checkpoint): self._all_initC = [checkpoint.initC] self._all_kick_info = [checkpoint.kick_info] + # ------------------------------------------------------------------ + # Binary-parameter assembly + # ------------------------------------------------------------------ + + def _binary_params(self, samples_physical): + """Assemble the full set of binary parameters for a batch. + + Merges the sampled columns with whatever ``derive_params`` returns, + then checks that every member of :attr:`REQUIRED_PARAMS` is present. + + Parameters + ---------- + samples_physical : `numpy.ndarray` + (N, D) array of sampled values in physical space. + + Returns + ------- + `dict` + Maps parameter name to an (N,) array. Guaranteed to contain + every name in :attr:`REQUIRED_PARAMS`, plus any sampled extras + (e.g. natal-kick columns) or additional keys from + ``derive_params``. + + Raises + ------ + `ValueError` + If a required parameter is neither sampled nor returned by + ``derive_params``, or if ``derive_params`` returns an array of + the wrong length. + """ + n = len(samples_physical) + params = {name: samples_physical[:, i] + for i, name in enumerate(self.param_space.names)} + + derived_keys = [] + if self.derive_fn is not None: + for key, value in self.derive_fn(params).items(): + value = np.asarray(value, dtype=float) + if value.ndim == 0: + value = np.full(n, value) # broadcast fixed values + elif len(value) != n: + raise ValueError( + f"derive_params returned '{key}' with length " + f"{len(value)}, but the batch has {n} binaries." + ) + params[key] = value + derived_keys.append(key) + + missing = [p for p in self.REQUIRED_PARAMS if p not in params] + if missing: + raise ValueError( + f"Each binary must be defined by {list(self.REQUIRED_PARAMS)}, " + f"but {missing} were neither sampled nor returned by " + f"derive_params.\n" + f" sampled by ParameterSpace : {self.param_space.names}\n" + f" returned by derive_params : {derived_keys}\n" + f"Add each missing parameter to the ParameterSpace, or return " + f"it from derive_params (a fixed value is fine, " + f"e.g. {{'mass_1': 30.0}})." + ) + return params + + def _physical_reject_mask(self, samples_physical): + """Boolean mask of physically rejected systems for a set of samples.""" + if self.reject_fn is None: + return np.zeros(len(samples_physical), dtype=bool) + return self.reject_fn(self._binary_params(samples_physical)) + + def _validate_param_coverage(self): + """Fail fast (at construction) if the binary model is under-specified. + + Runs the sampled-plus-derived assembly on a couple of probe draws so + that a missing required parameter raises immediately rather than after + COSMIC evolution has already begun. + """ + probe, _ = self.param_space.sample(2, rng=np.random.default_rng(0)) + self._binary_params(self.param_space.to_physical(probe)) + def _filter_tables(self, bpp, bcm, initC, kick_info, hit_bin_nums): """Return COSMIC tables filtered to hit systems when requested. @@ -366,12 +460,11 @@ def _explore(self): # Sample from prior samples, mask = self.param_space.sample(n_oversample, rng=self.rng) - # Transform to physical space for rejection checking + # Assemble binary parameters (sampled + derived) and reject samples_phys = self.param_space.to_physical(samples) - derived = self.compute_derived_fn(samples_phys, self.param_space.names) - - # Physical rejection - phys_rejected = self.reject_fn(samples_phys, derived, self.param_space.names) + binary_params = self._binary_params(samples_phys) + phys_rejected = (self.reject_fn(binary_params) if self.reject_fn is not None + else np.zeros(n_oversample, dtype=bool)) # Combined mask: in bounds AND not physically rejected valid = mask & ~phys_rejected @@ -382,13 +475,10 @@ def _explore(self): selected = valid_indices[:self.batch_size] batch_samples = samples[selected] - batch_samples_phys = samples_phys[selected] - batch_derived = {k: v[selected] for k, v in derived.items()} + batch_params = {k: v[selected] for k, v in binary_params.items()} # Evolve with COSMIC - n_hits, hit_bin_nums, bpp, bcm, initC, kick_info = self._evolve_batch( - batch_samples_phys, batch_derived - ) + n_hits, hit_bin_nums, bpp, bcm, initC, kick_info = self._evolve_batch(batch_params) # Record which are hits is_hit = np.zeros(len(selected), dtype=bool) @@ -439,9 +529,7 @@ def _estimate_prior_rejection_rate(self, n_test=100000): valid_samples = samples[mask] if len(valid_samples) > 0: phys = self.param_space.to_physical(valid_samples) - derived = self.compute_derived_fn(phys, self.param_space.names) - phys_rejected = self.reject_fn(phys, derived, self.param_space.names) - rejected += np.sum(phys_rejected) + rejected += np.sum(self._physical_reject_mask(phys)) return rejected / n_test @@ -490,7 +578,7 @@ def _refine(self): for gen in range(self.n_generations): # Estimate rejection rate for current mixture dist_rejection_rate = self.mixture.compute_rejection_rate( - self.param_space, self.compute_derived_fn, self.reject_fn, + self.param_space, self._physical_reject_mask, n_per_component=10000, rng=self.rng ) @@ -514,28 +602,25 @@ def _refine(self): continue phys = self.param_space.to_physical(valid_samples) - derived = self.compute_derived_fn(phys, self.param_space.names) - phys_rejected = self.reject_fn(phys, derived, self.param_space.names) + binary_params = self._binary_params(phys) + phys_rejected = (self.reject_fn(binary_params) if self.reject_fn is not None + else np.zeros(len(phys), dtype=bool)) keep = ~phys_rejected valid_samples = valid_samples[keep] valid_gauss_idx = valid_gauss_idx[keep] - phys = phys[keep] - derived = {k: v[keep] for k, v in derived.items()} + binary_params = {k: v[keep] for k, v in binary_params.items()} # Trim to batch size with randomisation indices = np.arange(len(valid_samples)) self.rng.shuffle(indices) n_take = min(len(valid_samples), self.batch_size) batch_samples = valid_samples[indices[:n_take]] - batch_phys = phys[indices[:n_take]] batch_gauss_idx = valid_gauss_idx[indices[:n_take]] - batch_derived = {k: v[indices[:n_take]] for k, v in derived.items()} + batch_params = {k: v[indices[:n_take]] for k, v in binary_params.items()} # Evolve with COSMIC - n_hits, hit_bin_nums, bpp, bcm, initC, kick_info = self._evolve_batch( - batch_phys, batch_derived - ) + n_hits, hit_bin_nums, bpp, bcm, initC, kick_info = self._evolve_batch(batch_params) is_hit = np.zeros(n_take, dtype=bool) is_hit[hit_bin_nums] = True @@ -717,16 +802,15 @@ def _concat(frames): ]) _KICK_SENTINEL = -100.0 - def _evolve_batch(self, samples_physical, derived): + def _evolve_batch(self, binary_params): """Evolve a batch of binaries with COSMIC and identify hits. Parameters ---------- - samples_physical : `numpy.ndarray` - (N, D) array of binary parameters in physical space. - derived : `dict` - Dictionary with keys ``'mass_2'``, ``'metallicity_1'``, etc., - each mapping to an (N,) array. + binary_params : `dict` + Assembled binary parameters (see :meth:`_binary_params`). Must + contain :attr:`REQUIRED_PARAMS`; any natal-kick columns present + are injected per-binary. Each value is an (N,) array. Returns ------- @@ -744,33 +828,31 @@ def _evolve_batch(self, samples_physical, derived): kick_info : `pandas.DataFrame` COSMIC natal kick information output. """ - n = len(samples_physical) - idx = {name: i for i, name in enumerate(self.param_space.names)} + n = len(binary_params['mass_1']) batch_initial = InitialBinaryTable.InitialBinaries( - m1=samples_physical[:, idx['mass_1']], - m2=derived['mass_2'], - porb=samples_physical[:, idx['porb']], - ecc=samples_physical[:, idx['ecc']], + m1=binary_params['mass_1'], + m2=binary_params['mass_2'], + porb=binary_params['porb'], + ecc=binary_params['ecc'], tphysf=np.full(n, 13700.0), kstar1=np.full(n, 1), kstar2=np.full(n, 1), - metallicity=derived['metallicity_1'], + metallicity=binary_params['metallicity'], ) - # If any natal kick columns were sampled, activate the per-binary - # injection path. Each column present in the ParameterSpace gets its - # sampled value; every other kick column gets _KICK_SENTINEL (-100) - # so COSMIC draws that component from its built-in prescription. - # natal_kick_array is stripped from the BSEDict copy so COSMIC reads - # the per-binary columns instead of the global default. - param_names_set = set(self.param_space.names) - sampled_kick_cols = self._ALL_KICK_COLUMNS & param_names_set + # If any natal kick columns are present (sampled or derived), activate + # the per-binary injection path. Each column present gets its value; + # every other kick column gets _KICK_SENTINEL (-100) so COSMIC draws + # that component from its built-in prescription. natal_kick_array is + # stripped from the BSEDict copy so COSMIC reads the per-binary columns + # instead of the global default. + present_kick_cols = self._ALL_KICK_COLUMNS & set(binary_params) - if sampled_kick_cols: + if present_kick_cols: for col in self._ALL_KICK_COLUMNS: - batch_initial[col] = (samples_physical[:, idx[col]] - if col in sampled_kick_cols + batch_initial[col] = (binary_params[col] + if col in present_kick_cols else self._KICK_SENTINEL) batch_initial['randomseed_1'] = 0 batch_initial['randomseed_2'] = 0 diff --git a/src/cosmic/sample/stroopwafel/examples/bh_star.py b/src/cosmic/sample/stroopwafel/examples/bh_star.py index d072e7252..bc033f342 100644 --- a/src/cosmic/sample/stroopwafel/examples/bh_star.py +++ b/src/cosmic/sample/stroopwafel/examples/bh_star.py @@ -126,39 +126,24 @@ # ------------------------------------------------------------------ # Derived quantities # ------------------------------------------------------------------ -def compute_derived(samples_physical, param_names): - """Compute mass_2, metallicities, and orbital separation from samples. +def derive_params(sampled): + """Provide the secondary mass from the sampled primary mass and mass ratio. + + A binary is defined by {mass_1, mass_2, porb, ecc, metallicity}. Here + mass_1, porb, ecc, and metallicity are sampled directly, so only mass_2 + needs deriving. Parameters ---------- - samples_physical : numpy.ndarray - (N, D) array of samples in physical space. - param_names : list of str - Column labels for each dimension. + sampled : dict + Maps each sampled parameter name to its (N,) array of physical values. Returns ------- dict - Keys: 'mass_2', 'metallicity_1', 'metallicity_2', 'separation'. + ``{'mass_2': ...}`` -- the one required parameter not sampled here. """ - idx = {name: i for i, name in enumerate(param_names)} - m1 = samples_physical[:, idx['mass_1']] - q = samples_physical[:, idx['q']] - porb = samples_physical[:, idx['porb']] # days (sana sampler output) - z = samples_physical[:, idx['metallicity']] - - mass_2 = m1 * q - - # Kepler's third law: a³ [AU³] = (P [yr])² · M [Msun] - # Convert period from days to years before applying. - separation = ((porb / 365.25) ** 2 * (m1 + mass_2)) ** (1.0 / 3.0) - - return { - 'mass_2': mass_2, - 'metallicity_1': z, - 'metallicity_2': z, - 'separation': separation, - } + return {'mass_2': sampled['mass_1'] * sampled['q']} # ------------------------------------------------------------------ # Hit definition: BH (kstar=14) + normal star (kstar 0–9), still bound @@ -213,9 +198,9 @@ def run_sampler(mc_only, seed): total_systems=args.num_systems, batch_size=args.batch_size, BSEDict=BSEDict, - compute_derived=compute_derived, - reject_systems=default_reject, is_interesting=is_bh_star, + derive_params=derive_params, + reject_systems=default_reject, output_path=os.path.join(args.output_dir, 'mc' if mc_only else 'sw'), nproc=args.num_cores, n_generations=args.n_generations, diff --git a/src/cosmic/sample/stroopwafel/examples/example_bhbh.py b/src/cosmic/sample/stroopwafel/examples/example_bhbh.py index f84fa5234..e1d5320e5 100644 --- a/src/cosmic/sample/stroopwafel/examples/example_bhbh.py +++ b/src/cosmic/sample/stroopwafel/examples/example_bhbh.py @@ -69,25 +69,14 @@ ]) # ------------------------------------------------------------------ -# Define derived quantities (vectorized) +# Provide binary parameters not sampled directly (vectorized) +# +# A binary is defined by {mass_1, mass_2, porb, ecc, metallicity}. All but +# mass_2 are sampled, so derive mass_2 from the sampled mass ratio. # ------------------------------------------------------------------ -def compute_derived(samples_physical, param_names): - """Compute derived quantities from sampled parameters.""" - idx = {name: i for i, name in enumerate(param_names)} - m1 = samples_physical[:, idx['mass_1']] - q = samples_physical[:, idx['q']] - porb = samples_physical[:, idx['porb']] - z = samples_physical[:, idx['metallicity']] - - mass_2 = m1 * q - separation = ((porb ** 2) * (m1 + mass_2)) ** (1.0 / 3.0) - - return { - 'mass_2': mass_2, - 'metallicity_1': z, - 'metallicity_2': z, - 'separation': separation, - } +def derive_params(sampled): + """Return the one required parameter (mass_2) not sampled directly.""" + return {'mass_2': sampled['mass_1'] * sampled['q']} # ------------------------------------------------------------------ # Hit definition: merging BH-BH binaries within Hubble time @@ -105,9 +94,9 @@ def compute_derived(samples_physical, param_names): total_systems=args.num_systems, batch_size=args.num_per_core, BSEDict=BSEDict, - compute_derived=compute_derived, - reject_systems=default_reject, is_interesting=is_interesting, + derive_params=derive_params, + reject_systems=default_reject, output_path=args.output_dir, nproc=args.num_cores, mc_only=args.mc_only, diff --git a/src/cosmic/sample/stroopwafel/mixture_model.py b/src/cosmic/sample/stroopwafel/mixture_model.py index d8aeba186..1ffe5306a 100644 --- a/src/cosmic/sample/stroopwafel/mixture_model.py +++ b/src/cosmic/sample/stroopwafel/mixture_model.py @@ -168,7 +168,7 @@ def component_pdfs(self, samples): ) return xPDF - def compute_rejection_rate(self, param_space, compute_derived_fn, reject_fn, + def compute_rejection_rate(self, param_space, reject_mask_fn, n_per_component=10000, rng=None): """Estimate the rejection rate of the mixture. @@ -179,14 +179,11 @@ def compute_rejection_rate(self, param_space, compute_derived_fn, reject_fn, ---------- param_space : `ParameterSpace` Parameter space for bounds checking and coordinate transforms. - compute_derived_fn : `callable` - Function with signature - ``(samples_physical, param_names) -> dict`` that computes - derived quantities. - reject_fn : `callable` - Function with signature - ``(samples_physical, derived, param_names) -> bool_mask`` - returning True for rejected systems. + reject_mask_fn : `callable` + Function ``(samples_physical) -> bool_mask`` returning True for + physically rejected systems. The engine supplies one that + assembles binary parameters and applies the user's rejection + function. n_per_component : `int`, optional Number of samples per component for the estimate, by default 10000 @@ -214,9 +211,7 @@ def compute_rejection_rate(self, param_space, compute_derived_fn, reject_fn, s_valid = s[bounds_mask] if len(s_valid) > 0: s_physical = param_space.to_physical(s_valid) - derived = compute_derived_fn(s_physical, param_space.names) - phys_rejected = reject_fn(s_physical, derived, param_space.names) - rejected += np.sum(phys_rejected) + rejected += np.sum(reject_mask_fn(s_physical)) fractional_rejected += rejected * self.alphas[k] / n diff --git a/src/cosmic/sample/stroopwafel/rejection.py b/src/cosmic/sample/stroopwafel/rejection.py index 2082b8fda..af8c49708 100644 --- a/src/cosmic/sample/stroopwafel/rejection.py +++ b/src/cosmic/sample/stroopwafel/rejection.py @@ -45,24 +45,21 @@ def get_zams_radius(mass, metallicity): return (top / bottom) * R_SOL_TO_AU -def default_reject(samples_physical, derived, param_names, - min_secondary_mass=0.08): +def default_reject(binary_params, min_secondary_mass=0.08): """Default rejection function for DCO progenitor systems. Rejects systems where the secondary mass is below the minimum, the stars are in contact at ZAMS, or either star overflows its Roche lobe - at periastron. + at periastron. The orbital separation is computed from the orbital + period via Kepler's third law, and both stars share the binary + metallicity. Parameters ---------- - samples_physical : `numpy.ndarray` - (N, D) array of samples in physical space. - derived : `dict` - Dictionary with keys ``'mass_2'``, ``'metallicity_1'``, - ``'metallicity_2'``, and ``'separation'``, each mapping to an - (N,) array. - param_names : `list` of `str` - Parameter names in column order (used to find column indices). + binary_params : `dict` + Assembled binary parameters with keys ``'mass_1'``, ``'mass_2'`` + (solar masses), ``'porb'`` (days), ``'ecc'``, and ``'metallicity'``, + each an (N,) array. min_secondary_mass : `float`, optional Minimum allowed secondary mass in solar masses, by default 0.08 @@ -71,32 +68,30 @@ def default_reject(samples_physical, derived, param_names, `numpy.ndarray` (N,) boolean mask where True indicates a rejected system. """ - idx = {name: i for i, name in enumerate(param_names)} + mass_1 = binary_params['mass_1'] + mass_2 = binary_params['mass_2'] + porb = binary_params['porb'] + ecc = binary_params['ecc'] + metallicity = binary_params['metallicity'] - mass_1 = samples_physical[:, idx['mass_1']] - mass_2 = derived['mass_2'] - met_1 = derived['metallicity_1'] - met_2 = derived['metallicity_2'] - separation = derived['separation'] - ecc = samples_physical[:, idx['ecc']] + # Semi-major axis [AU] from Kepler's third law (porb in days, masses in + # solar masses): a^3 [AU^3] = (P [yr])^2 * M [Msun]. + separation = ((porb / 365.25) ** 2 * (mass_1 + mass_2)) ** (1.0 / 3.0) - # Compute ZAMS radii - radius_1 = get_zams_radius(mass_1, met_1) - radius_2 = get_zams_radius(mass_2, met_2) + # ZAMS radii [AU] (both stars share the binary metallicity) + radius_1 = get_zams_radius(mass_1, metallicity) + radius_2 = get_zams_radius(mass_2, metallicity) # Roche lobe radii at periastron peri_sep = separation * (1 - ecc) rl_1 = calc_Roche_radius(mass_1, mass_2, peri_sep) rl_2 = calc_Roche_radius(mass_2, mass_1, peri_sep) - roche_tracker_1 = radius_1 / rl_1 - roche_tracker_2 = radius_2 / rl_2 - rejected = ( (mass_2 < min_secondary_mass) | (separation <= (radius_1 + radius_2)) - | (roche_tracker_1 > 1) - | (roche_tracker_2 > 1) + | (radius_1 / rl_1 > 1) + | (radius_2 / rl_2 > 1) ) return rejected diff --git a/src/cosmic/tests/test_stroopwafel.py b/src/cosmic/tests/test_stroopwafel.py index 01611edb8..36db0e960 100644 --- a/src/cosmic/tests/test_stroopwafel.py +++ b/src/cosmic/tests/test_stroopwafel.py @@ -14,7 +14,7 @@ from scipy.stats import norm, kstest from scipy.integrate import trapezoid -from cosmic.sample.stroopwafel import ParameterSpace, Parameter +from cosmic.sample.stroopwafel import ParameterSpace, Parameter, AdaptiveSampler from cosmic.sample.stroopwafel.distributions import ( DISTRIBUTIONS, register, get_distribution, Distribution, Uniform, PowerLaw, TruncatedNormal, @@ -314,8 +314,6 @@ def test_compute_sigma_reports_parameter_name_on_error(self): # ========================================================================== class TestRejection(unittest.TestCase): - PARAM_NAMES = ['ecc', 'mass_1', 'metallicity', 'porb', 'q'] - def test_get_zams_radius_positive_finite(self): masses = np.array([1.0, 10.0, 30.0, 100.0]) mets = np.array([0.02, 0.02, 0.001, 0.014]) @@ -324,26 +322,85 @@ def test_get_zams_radius_positive_finite(self): self.assertTrue(np.all(r > 0) and np.all(np.isfinite(r))) def test_wide_binary_not_rejected(self): - samples = np.array([[0.1, 20.0, 0.014, 100.0, 0.5]]) - derived = { + binary_params = { + 'mass_1': np.array([20.0]), 'mass_2': np.array([10.0]), - 'metallicity_1': np.array([0.014]), - 'metallicity_2': np.array([0.014]), - 'separation': np.array([50.0]), + 'porb': np.array([1000.0]), # days -> several AU, well detached + 'ecc': np.array([0.1]), + 'metallicity': np.array([0.014]), } - rejected = default_reject(samples, derived, self.PARAM_NAMES) - self.assertFalse(rejected[0]) + self.assertFalse(default_reject(binary_params)[0]) def test_low_mass_secondary_rejected(self): - samples = np.array([[0.1, 20.0, 0.014, 100.0, 0.001]]) - derived = { - 'mass_2': np.array([0.02]), # below the 0.08 Msun minimum - 'metallicity_1': np.array([0.014]), - 'metallicity_2': np.array([0.014]), - 'separation': np.array([50.0]), + binary_params = { + 'mass_1': np.array([20.0]), + 'mass_2': np.array([0.02]), # below the 0.08 Msun minimum + 'porb': np.array([1000.0]), + 'ecc': np.array([0.1]), + 'metallicity': np.array([0.014]), + } + self.assertTrue(default_reject(binary_params)[0]) + + def test_contact_binary_rejected(self): + # A 0.5-day orbit puts two massive stars in contact at ZAMS. + binary_params = { + 'mass_1': np.array([30.0]), + 'mass_2': np.array([25.0]), + 'porb': np.array([0.5]), + 'ecc': np.array([0.0]), + 'metallicity': np.array([0.014]), } - rejected = default_reject(samples, derived, self.PARAM_NAMES) - self.assertTrue(rejected[0]) + self.assertTrue(default_reject(binary_params)[0]) + + +# ========================================================================== +# Binary-parameter model (sampled + derived coverage) +# ========================================================================== +class TestBinaryModel(unittest.TestCase): + + def _make(self, params, derive_params=None): + return AdaptiveSampler( + parameter_space=params, total_systems=10, batch_size=5, BSEDict={}, + is_interesting=lambda bpp: (0, np.array([], dtype=int)), + derive_params=derive_params, reject_systems=None, + ) + + def test_missing_required_params_raise_at_construction(self): + # Only porb is sampled; mass_1/mass_2/ecc/metallicity are undefined. + params = ParameterSpace([Parameter('porb', 10**(0.15), 10**(5.5), dist='sana')]) + with self.assertRaisesRegex(ValueError, 'mass_1'): + self._make(params) + + def test_derive_params_fills_missing_with_scalars(self): + # Sample only porb; fix the rest via scalar returns (broadcast to N). + params = ParameterSpace([Parameter('porb', 10**(0.15), 10**(5.5), dist='sana')]) + sampler = self._make(params, derive_params=lambda s: { + 'mass_1': 30.0, 'mass_2': 25.0, 'ecc': 0.0, 'metallicity': 0.02, + }) + phys = params.to_physical(params.sample(4, rng=np.random.default_rng(0))[0]) + bp = sampler._binary_params(phys) + for key in AdaptiveSampler.REQUIRED_PARAMS: + self.assertEqual(bp[key].shape, (4,)) + np.testing.assert_array_equal(bp['mass_1'], np.full(4, 30.0)) + + def test_all_required_sampled_needs_no_derive(self): + params = ParameterSpace([ + Parameter('mass_1', 5.0, 150.0, dist='kroupa'), + Parameter('mass_2', 1.0, 100.0, dist='uniform'), + Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), + Parameter('ecc', 1e-9, 0.99, dist='sana_ecc'), + Parameter('metallicity', 1e-4, 0.03, dist='flat_in_log'), + ]) + self._make(params) # must not raise + + def test_derive_params_wrong_length_raises(self): + params = ParameterSpace([Parameter('porb', 10**(0.15), 10**(5.5), dist='sana')]) + sampler_factory = lambda: self._make(params, derive_params=lambda s: { + 'mass_1': np.array([30.0]), # wrong length vs the 2-row probe + 'mass_2': 25.0, 'ecc': 0.0, 'metallicity': 0.02, + }) + with self.assertRaisesRegex(ValueError, 'length'): + sampler_factory() # ========================================================================== From b7457b9ab207491842491f2e2334d48d079044f4 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Sun, 28 Jun 2026 22:18:11 -0400 Subject: [PATCH 31/45] rename to main --- src/cosmic/sample/stroopwafel/__init__.py | 2 +- .../sample/stroopwafel/{engine.py => main.py} | 19 ++++++++----------- src/cosmic/sample/stroopwafel/meson.build | 2 +- .../sample/stroopwafel/mixture_model.py | 4 +--- 4 files changed, 11 insertions(+), 16 deletions(-) rename src/cosmic/sample/stroopwafel/{engine.py => main.py} (98%) diff --git a/src/cosmic/sample/stroopwafel/__init__.py b/src/cosmic/sample/stroopwafel/__init__.py index f4b125798..ec0189b11 100755 --- a/src/cosmic/sample/stroopwafel/__init__.py +++ b/src/cosmic/sample/stroopwafel/__init__.py @@ -8,7 +8,7 @@ regions of parameter space. """ -from .engine import AdaptiveSampler +from .main import AdaptiveSampler from .parameter_space import ParameterSpace, Parameter __all__ = ['AdaptiveSampler', 'ParameterSpace', 'Parameter'] diff --git a/src/cosmic/sample/stroopwafel/engine.py b/src/cosmic/sample/stroopwafel/main.py similarity index 98% rename from src/cosmic/sample/stroopwafel/engine.py rename to src/cosmic/sample/stroopwafel/main.py index 5d767df37..56defc481 100644 --- a/src/cosmic/sample/stroopwafel/engine.py +++ b/src/cosmic/sample/stroopwafel/main.py @@ -1,8 +1,3 @@ -"""AdaptiveSampler: the main STROOPWAFEL engine. - -Orchestrates the explore -> adapt -> refine -> weight calculation pipeline -using vectorized operations throughout. No Location objects are created. -""" import os import numpy as np import pandas as pd @@ -14,6 +9,7 @@ from .mixture_model import GaussianMixture from .constants import MIN_ACTIVE_FRACTION +from .rejection import default_reject class AdaptiveSampler: @@ -46,8 +42,9 @@ class AdaptiveSampler: physically unacceptable systems, where ``binary_params`` is the assembled dict of binary parameters (sampled columns merged with the output of ``derive_params``). See - :func:`~cosmic.sample.stroopwafel.rejection.default_reject`. By - default None (no physical rejection). + :func:`~cosmic.sample.stroopwafel.rejection.default_reject`. + By default uses :func:`~cosmic.sample.stroopwafel.rejection.default_reject`. + Pass None to skip physical rejection entirely. output_path : `str`, optional Directory for output files, by default ``'output'`` nproc : `int`, optional @@ -77,7 +74,7 @@ class AdaptiveSampler: REQUIRED_PARAMS = ('mass_1', 'mass_2', 'porb', 'ecc', 'metallicity') def __init__(self, parameter_space, total_systems, batch_size, BSEDict, - is_interesting, derive_params=None, reject_systems=None, + is_interesting, derive_params=None, reject_systems="default", output_path='output', nproc=1, kappa=1.0, n_generations=1, mc_only=False, seed=None, only_save_hit_tables=False): @@ -86,7 +83,7 @@ def __init__(self, parameter_space, total_systems, batch_size, BSEDict, self.batch_size = batch_size self.bse_dict = BSEDict self.derive_fn = derive_params - self.reject_fn = reject_systems + self.reject_fn = reject_systems if reject_systems != "default" else default_reject self.is_interesting_fn = is_interesting self.output_path = output_path self.nproc = nproc @@ -280,7 +277,7 @@ def from_checkpoint(cls, checkpoint, parameter_space, batch_size, BSEDict, return sampler def _make_checkpoint(self): - """Package current engine state into a `STROOPWAFELCheckpoint`.""" + """Package current state into a `STROOPWAFELCheckpoint`.""" all_samples = np.vstack(self._all_samples) all_is_hit = np.concatenate(self._all_is_hit) all_generation = np.concatenate(self._all_generation) @@ -305,7 +302,7 @@ def _make_checkpoint(self): ) def _load_checkpoint(self, checkpoint): - """Restore engine state from a `STROOPWAFELCheckpoint`. + """Restore state from a `STROOPWAFELCheckpoint`. The checkpoint's samples, COSMIC output, and mixture are treated as a single exploration "super-batch" with globally assigned diff --git a/src/cosmic/sample/stroopwafel/meson.build b/src/cosmic/sample/stroopwafel/meson.build index 691dc69aa..0af7964e8 100644 --- a/src/cosmic/sample/stroopwafel/meson.build +++ b/src/cosmic/sample/stroopwafel/meson.build @@ -2,7 +2,7 @@ python_sources = [ '__init__.py', 'constants.py', 'distributions.py', - 'engine.py', + 'main.py', 'mixture_model.py', 'parameter_space.py', 'presets.py', diff --git a/src/cosmic/sample/stroopwafel/mixture_model.py b/src/cosmic/sample/stroopwafel/mixture_model.py index 1ffe5306a..6366e860a 100644 --- a/src/cosmic/sample/stroopwafel/mixture_model.py +++ b/src/cosmic/sample/stroopwafel/mixture_model.py @@ -181,9 +181,7 @@ def compute_rejection_rate(self, param_space, reject_mask_fn, Parameter space for bounds checking and coordinate transforms. reject_mask_fn : `callable` Function ``(samples_physical) -> bool_mask`` returning True for - physically rejected systems. The engine supplies one that - assembles binary parameters and applies the user's rejection - function. + physically rejected systems. n_per_component : `int`, optional Number of samples per component for the estimate, by default 10000 From 30212c51e87ce559dde9a2401e64001473842248 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Sun, 28 Jun 2026 22:42:15 -0400 Subject: [PATCH 32/45] update basics tutorial --- docs/pages/tutorials/adaptive/basics.rst | 88 +++++++++++------------- 1 file changed, 39 insertions(+), 49 deletions(-) diff --git a/docs/pages/tutorials/adaptive/basics.rst b/docs/pages/tutorials/adaptive/basics.rst index 373fcedf1..61c8a273f 100644 --- a/docs/pages/tutorials/adaptive/basics.rst +++ b/docs/pages/tutorials/adaptive/basics.rst @@ -158,10 +158,6 @@ systems below the hydrogen-burning limit for the secondary. The built-in :func:`~cosmic.sample.stroopwafel.rejection.default_reject` function performs all of these checks: -.. code-block:: python - - from cosmic.sample.stroopwafel.rejection import default_reject - It receives the assembled binary parameters as a dictionary (``mass_1``, ``mass_2``, ``porb``, ``ecc``, ``metallicity``) and returns a boolean mask where ``True`` means the system is rejected; the orbital separation is computed internally from ``porb``. For most @@ -176,8 +172,7 @@ minimum secondary mass — you can wrap it: base_mask |= (binary_params['mass_2'] < 1.0) return base_mask -Passing ``reject_systems`` is optional; omit it (or pass ``None``) to skip physical -rejection entirely. +Passing ``reject_systems`` is optional; you can instead pass ``None`` to skip physical rejection entirely. Identify what constitutes a hit @@ -188,17 +183,20 @@ the ``COSMIC`` ``bpp`` DataFrame for the current batch and must return a tuple ``(n_hits, hit_bin_nums)`` where ``hit_bin_nums`` is an integer array of ``bin_num`` values (0-indexed within the batch). -The ``STROOPWAFEL`` sampler comes with two preset functions in -:mod:`cosmic.sample.stroopwafel.presets`. +The ``STROOPWAFEL`` sampler comes with two preset functions in :mod:`cosmic.sample.stroopwafel.presets`. +These functions focus on double compact objects (DCOs) and are suitable for gravitational-wave source studies: ``any_dco(kstar_1, kstar_2)`` Selects all bound double compact objects whose stellar types match the supplied lists, - regardless of merger time. Use this for populations where you care about the DCO - existing rather than merging within the Hubble time. + regardless of merger time. ``merging_dco(kstar_1, kstar_2, max_merge_time=13.7)`` - Like ``any_dco`` but additionally requires the merger time (computed via LEGWORK) to be - less than ``max_merge_time`` Gyr. Suitable for gravitational-wave source studies. + Like ``any_dco`` but additionally requires the merger time to be + less than ``max_merge_time`` Gyr. + +.. note:: + + The ``merging_dco`` function requires the `LEGWORK python package `_ to compute the merger time. If you do not have LEGWORK installed, ``merging_dco`` will raise an error. .. code-block:: python @@ -207,10 +205,9 @@ The ``STROOPWAFEL`` sampler comes with two preset functions in # All bound BH-BH systems, no merger time cut is_interesting = any_dco(kstar_1=[14], kstar_2=[14]) - # Only BH-BH systems merging within the Hubble time - is_interesting_merging = merging_dco(kstar_1=[14], kstar_2=[14], max_merge_time=13.7) + # Only NS-NS systems merging within the Hubble time + is_interesting_merging = merging_dco(kstar_1=[13], kstar_2=[13], max_merge_time=13.7) -See :ref:`kstar-table` for the full list of stellar type codes. You can also write a fully custom hit function. For example, to find BH + stellar companion systems that remain bound for at least 100 Myr after the BH forms: @@ -239,9 +236,9 @@ You can also write a fully custom hit function. For example, to find BH + stell Running the sampler =================== -With all the pieces in place, you can run the sampler with :class:`~cosmic.sample.stroopwafel.AdaptiveSampler`. The most important arguments are the parameter space, the total number of systems to evolve, the batch size, the BSE physics settings, the hit function, the ``derive_params`` function (if needed), and the rejection function. See the API documentation (:class:`~cosmic.sample.stroopwafel.AdaptiveSampler`) for a full list of options. +Now we can put it all together! You can run the sampler with the main :class:`~cosmic.sample.stroopwafel.AdaptiveSampler` class. The most important arguments are the parameter space, the total number of systems to evolve, the batch size, the BSE physics settings, the hit function, the ``derive_params`` function (if needed), and the rejection function. See the API documentation (:class:`~cosmic.sample.stroopwafel.AdaptiveSampler`) for a full list of options. -The examples below demonstrate how you could go about this. +Let's try this out with a few examples. Examples -------- @@ -249,22 +246,22 @@ Examples Bound BH + BH binaries ^^^^^^^^^^^^^^^^^^^^^^ -The following end-to-end example samples all bound BH-BH systems (no merger time -restriction) using a five-dimensional parameter space covering primary mass, mass ratio, -orbital period, eccentricity, and metallicity. +Let's imagine we want to sample the population of bound BH + BH binaries. We can use the same parameter space and ``derive_params`` function as above, and the built-in ``any_dco`` hit function to select all bound BH + BH systems. We can sample over a five-dimensional parameter space covering primary mass, mass ratio, orbital period, eccentricity, and metallicity. -.. include:: ../../../_generated/default_bsedict.rst +First we can import the necessary parts from the ``cosmic.sample.stroopwafel`` module, the preset hit function, and setup a BSEDict. .. code-block:: python import numpy as np from cosmic.sample.stroopwafel import AdaptiveSampler, ParameterSpace, Parameter from cosmic.sample.stroopwafel.presets import any_dco - from cosmic.sample.stroopwafel.rejection import default_reject - # ------------------------------------------------------------------ - # Parameter space - # ------------------------------------------------------------------ +.. include:: ../../../_generated/default_bsedict.rst + +Then we can define a simple parameter space, where we avoid sampling low-mass primaries since we know they +cannot produce a BH. + +.. code-block:: python params = ParameterSpace([ Parameter('mass_1', 5.0, 150.0, dist='kroupa'), Parameter('q', 0.01, 1.0, dist='uniform'), @@ -273,29 +270,29 @@ orbital period, eccentricity, and metallicity. Parameter('metallicity', 0.0001, 0.03, dist='flat_in_log'), ]) - # ------------------------------------------------------------------ - # Complete the binary definition (mass_2 from the sampled mass ratio) - # ------------------------------------------------------------------ +Since we only sampled the mass ratio ``q``, we need to derive the secondary mass from the primary mass and ``q``: + +.. code-block:: python def derive_params(sampled): return {'mass_2': sampled['mass_1'] * sampled['q']} - # ------------------------------------------------------------------ - # Run - # ------------------------------------------------------------------ + +And then it's just a matter of setting it going! + +.. code-block:: python sampler = AdaptiveSampler( parameter_space=params, - total_systems=50_000, - batch_size=500, + total_systems=50_000, # adjust this for more samples + batch_size=500, # adjust this to sample more or fewer systems per call to COSMIC BSEDict=BSEDict, is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), derive_params=derive_params, - reject_systems=default_reject, + reject_systems="default", output_path='output/bhbh', nproc=4, n_generations=1, seed=42, ) - result = sampler.run() print(f"Total hits: {result.num_hits}") @@ -305,9 +302,7 @@ orbital period, eccentricity, and metallicity. BH + star binaries surviving 100 Myr ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -For outcomes that are less extreme but still rare — such as persistent BH + star systems — -STROOPWAFEL provides substantial efficiency gains over flat Monte Carlo. Using the same -parameter space, ``derive_params``, and ``BSEDict`` as the previous examples, we can simply swap out the hit function to find BH + star systems that remain bound for at least 100 Myr after the BH forms: +Now let's repeat that whole scenario, but instead of BH + BH binaries we want to sample BH + stellar companion systems that remain bound for at least 100 Myr after the BH forms. We can use the same parameter space and ``derive_params`` function as above, but this time we will use the custom ``bh_star_100myr`` hit function defined earlier. .. code-block:: python @@ -317,23 +312,19 @@ parameter space, ``derive_params``, and ``BSEDict`` as the previous examples, we sampler = AdaptiveSampler( parameter_space=params, # reuse from BHBH example - total_systems=20_000, + total_systems=50_000, batch_size=500, BSEDict=BSEDict, # reuse from BHBH example is_interesting=bh_star_100myr, # we defined this earlier derive_params=derive_params, # reuse from BHBH example - reject_systems=default_reject, + reject_systems="default", output_path='output/bh_star', nproc=4, n_generations=1, seed=42, ) - result = sampler.run() -Because BH + star systems are more common than merging BH-BH pairs, a smaller total budget -is needed and fewer refinement generations are required before the mixture model is -well-constrained. Rules of thumb ============== @@ -343,13 +334,12 @@ Choosing ``total_systems``, ``batch_size``, and ``n_generations`` is something o ``batch_size`` -------------- -``batch_size`` sets how many systems are passed to -:meth:`~cosmic.evolve.Evolve.evolve` per call. +``batch_size`` sets how many systems are passed to :meth:`~cosmic.evolve.Evolve.evolve` per call. * Aim for ``batch_size`` to be a multiple of ``nproc`` so that COSMIC distributes work evenly across cores. * Values of 200-1000 are typical. Batches smaller than ~50 increase Python overhead per - call; batches larger than ~5000 may cause memory pressure on the output DataFrames. + call; batches larger than ~5000 may cause memory pressure from the output DataFrames. ``total_systems`` ----------------- @@ -372,10 +362,10 @@ The EM step between generations can improve the mixture, but with diminishing re Saving your results =================== -Once you have your samples, you can save them to disk as an HDF5 file with the :meth:`~cosmic.output.COSMICStroopOutput.save` method. This saves the parameter samples, derived quantities, and hit information in a compact format that can be loaded later for analysis. +Once you have your samples, you can save them to disk as an HDF5 file with the :meth:`~cosmic.output.COSMICStroopOutput.save` method. This saves the parameter samples, derived quantities, and hit information that can be loaded later for analysis. .. code-block:: python result.save('bhbh_samples.h5') -We'll talk more about how to load and analyse these results in the :ref:`adaptive_outputs` tutorial next! \ No newline at end of file +We'll talk more about how to load and analyse these results in the :ref:`adaptive_outputs` tutorial! But first, let's look at how to define custom distributions for your parameters in the next tutorial :ref:`adaptive_distributions`. \ No newline at end of file From ceaff95a3d7e0b074c4bc192e61cb0bce8986abc Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Sun, 28 Jun 2026 22:49:40 -0400 Subject: [PATCH 33/45] delete constants, move to definitions --- src/cosmic/sample/stroopwafel/constants.py | 10 -- .../sample/stroopwafel/distributions.py | 112 ++++++++++++++++-- 2 files changed, 101 insertions(+), 21 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/constants.py b/src/cosmic/sample/stroopwafel/constants.py index 21112f903..f3a04334e 100755 --- a/src/cosmic/sample/stroopwafel/constants.py +++ b/src/cosmic/sample/stroopwafel/constants.py @@ -1,14 +1,4 @@ """Constants used throughout the STROOPWAFEL adaptive sampling module.""" -ALPHA_IMF = -2.3 -SANA_G = -0.55 -SANA_ECC = -0.45 - -# Log-normal natal kick distribution. -# The kick magnitude v [km/s] follows LogNormal(mu, sigma), meaning -# ln(v) ~ Normal(NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA). -# With mu=5.67, sigma=0.59 the median kick is exp(5.67) ≈ 291 km/s. -NATAL_KICK_LOG_MU = 5.67 # mean of ln(v_kick / km s⁻¹) -NATAL_KICK_LOG_SIGMA = 0.59 # std dev of ln(v_kick / km s⁻¹) R_COEFF = [ [1.71535900, 0.62246212, -0.92557761, -1.16996966, -0.30631491], diff --git a/src/cosmic/sample/stroopwafel/distributions.py b/src/cosmic/sample/stroopwafel/distributions.py index 53429bf27..247e24dae 100644 --- a/src/cosmic/sample/stroopwafel/distributions.py +++ b/src/cosmic/sample/stroopwafel/distributions.py @@ -10,9 +10,10 @@ * the **transform** between physical and sampling space. Each distribution is a *base distribution* (:class:`Uniform`, -:class:`PowerLaw`, :class:`TruncatedNormal`) composed with a *transform* -(:class:`Identity`, :class:`Log10`, :class:`Ln`, :class:`Sin`, -:class:`CosShift`). The base operates entirely in sampling space; the +:class:`PowerLaw`, :class:`BrokenPowerLaw`, :class:`TruncatedNormal`) composed +with a *transform* (:class:`Identity`, :class:`Log10`, :class:`Ln`, +:class:`Sin`, :class:`CosShift`). The base operates entirely in sampling +space; the transform maps to and from the physical space the user specifies bounds in and the simulation consumes. For example ``flat_in_log`` is ``Uniform(transform=Log10())`` and ``sana`` is @@ -30,10 +31,6 @@ import numpy as np from scipy.stats import norm as _scipy_norm -from .constants import ( - ALPHA_IMF, SANA_G, SANA_ECC, NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA, -) - # --------------------------------------------------------------------------- # Transforms between physical and sampling space @@ -323,6 +320,99 @@ def sigma(self, values, lo, hi, avg_density): return np.maximum(np.abs(x_right - values), np.abs(x_left - values)) +class BrokenPowerLaw(Distribution): + r"""Continuous broken power law: ``p(x) \propto x^alpha`` with ``alpha`` + changing at fixed breakpoints. + + The density is continuous across every breakpoint. With the default + :class:`Identity` transform this gives the Kroupa IMF + (``breaks=[0.5]``, ``alphas=[-1.3, -2.3]``): shallower below 0.5 Msun, + steeper above. Over any window that contains no breakpoint it reduces + exactly to :class:`PowerLaw`, so sampling masses above 0.5 Msun behaves + identically to a single power law. + + Sampling and ``sigma`` use a piecewise inverse CDF; the normalisation, + sampling, and density are all computed over the requested ``[lo, hi]`` + window only (segments outside it are ignored). + + Parameters + ---------- + breaks : sequence of `float` + Internal breakpoints in sampling space, strictly increasing. The + distribution has ``len(breaks) + 1`` segments. + alphas : sequence of `float` + Power-law exponent for each segment (``len(breaks) + 1`` of them); + ``alphas[i]`` applies below ``breaks[i]`` and ``alphas[-1]`` above the + final break. None may equal exactly -1. + transform : `Transform`, optional + Map between physical and sampling space, by default :class:`Identity`. + """ + + def __init__(self, breaks, alphas, transform=None): + super().__init__(transform) + self.breaks = np.asarray(breaks, dtype=float) + self.alphas = np.asarray(alphas, dtype=float) + if self.alphas.shape != (self.breaks.size + 1,): + raise ValueError("alphas must have exactly one more entry than breaks.") + if np.any(np.diff(self.breaks) <= 0): + raise ValueError("breaks must be strictly increasing.") + if np.any(self.alphas == -1.0): + raise ValueError("BrokenPowerLaw does not support an exponent of exactly -1.") + # Per-segment continuity coefficients (the lowest segment is 1). + coeffs = np.ones(self.alphas.size) + for i in range(self.breaks.size): + coeffs[i + 1] = coeffs[i] * self.breaks[i] ** (self.alphas[i] - self.alphas[i + 1]) + self.coeffs = coeffs + + def _segments(self, lo, hi): + """Decompose ``[lo, hi]`` into segments split at the in-range breaks. + + Returns the segment edges and, per segment, the exponent, continuity + coefficient, and cumulative unnormalised integral (``cum[-1]`` is the + normalisation constant ``Z``). + """ + interior = self.breaks[(self.breaks > lo) & (self.breaks < hi)] + edges = np.concatenate(([lo], interior, [hi])) + piece = np.searchsorted(self.breaks, edges[:-1], side='right') + alpha = self.alphas[piece] + coeff = self.coeffs[piece] + a = alpha + 1.0 + seg_int = coeff * (edges[1:] ** a - edges[:-1] ** a) / a + cum = np.concatenate(([0.0], np.cumsum(seg_int))) + return edges, alpha, coeff, cum + + def pdf(self, values, lo, hi): + edges, alpha, coeff, cum = self._segments(lo, hi) + s = np.clip(np.searchsorted(edges, values, side='right') - 1, 0, alpha.size - 1) + return coeff[s] * np.power(values, alpha[s]) / cum[-1] + + def sample(self, n, lo, hi, rng=None): + rng = rng or np.random.default_rng() + return self._ppf(rng.uniform(0, 1, n), lo, hi) + + def sigma(self, values, lo, hi, avg_density): + """CDF-space step, identical in spirit to :meth:`PowerLaw.sigma`.""" + u = self._cdf(values, lo, hi) + x_right = self._ppf(np.clip(u + avg_density, 0.0, 1.0), lo, hi) + x_left = self._ppf(np.clip(u - avg_density, 0.0, 1.0), lo, hi) + return np.maximum(np.abs(x_right - values), np.abs(x_left - values)) + + def _cdf(self, values, lo, hi): + edges, alpha, coeff, cum = self._segments(lo, hi) + s = np.clip(np.searchsorted(edges, values, side='right') - 1, 0, alpha.size - 1) + a = alpha[s] + 1.0 + partial = coeff[s] * (np.power(values, a) - np.power(edges[s], a)) / a + return (cum[s] + partial) / cum[-1] + + def _ppf(self, u, lo, hi): + edges, alpha, coeff, cum = self._segments(lo, hi) + target = np.clip(u, 0.0, 1.0) * cum[-1] + s = np.clip(np.searchsorted(cum, target, side='right') - 1, 0, alpha.size - 1) + a = alpha[s] + 1.0 + partial = target - cum[s] + return np.power(partial * a / coeff[s] + np.power(edges[s], a), 1.0 / a) + + class TruncatedNormal(Distribution): """Normal distribution truncated to ``[lo, hi]`` in sampling space. @@ -369,10 +459,10 @@ def pdf(self, values, lo, hi): 'flat_in_log': Uniform(transform=Log10()), 'uniform_in_sine': Uniform(transform=Sin()), 'uniform_in_cosine': Uniform(transform=CosShift()), - 'kroupa': PowerLaw(ALPHA_IMF), - 'sana': PowerLaw(SANA_G, transform=Log10()), - 'sana_ecc': PowerLaw(SANA_ECC), - 'log_normal': TruncatedNormal(NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA, transform=Ln()), + 'kroupa': BrokenPowerLaw(breaks=[0.5], alphas=[-1.3, -2.3]), + 'sana': PowerLaw(-0.55, transform=Log10()), + 'sana_ecc': PowerLaw(-0.45), + 'disberg': TruncatedNormal(5.67, 0.59, transform=Ln()), } From 74b8109822f0ce114ed012cd8eddd976728bf65a Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Sun, 28 Jun 2026 23:07:22 -0400 Subject: [PATCH 34/45] simplify rejection, towards deleting consts --- docs/pages/tutorials/adaptive/basics.rst | 3 + .../tutorials/adaptive/distributions.rst | 19 +-- src/cosmic/sample/stroopwafel/constants.py | 14 --- .../sample/stroopwafel/distributions.py | 4 +- .../sample/stroopwafel/examples/bh_star.py | 2 +- .../sample/stroopwafel/parameter_space.py | 2 +- src/cosmic/sample/stroopwafel/rejection.py | 55 ++------- src/cosmic/tests/test_stroopwafel.py | 111 ++++++++++++++---- 8 files changed, 114 insertions(+), 96 deletions(-) diff --git a/docs/pages/tutorials/adaptive/basics.rst b/docs/pages/tutorials/adaptive/basics.rst index 61c8a273f..6c77a3236 100644 --- a/docs/pages/tutorials/adaptive/basics.rst +++ b/docs/pages/tutorials/adaptive/basics.rst @@ -262,6 +262,7 @@ Then we can define a simple parameter space, where we avoid sampling low-mass pr cannot produce a BH. .. code-block:: python + params = ParameterSpace([ Parameter('mass_1', 5.0, 150.0, dist='kroupa'), Parameter('q', 0.01, 1.0, dist='uniform'), @@ -273,6 +274,7 @@ cannot produce a BH. Since we only sampled the mass ratio ``q``, we need to derive the secondary mass from the primary mass and ``q``: .. code-block:: python + def derive_params(sampled): return {'mass_2': sampled['mass_1'] * sampled['q']} @@ -280,6 +282,7 @@ Since we only sampled the mass ratio ``q``, we need to derive the secondary mass And then it's just a matter of setting it going! .. code-block:: python + sampler = AdaptiveSampler( parameter_space=params, total_systems=50_000, # adjust this for more samples diff --git a/docs/pages/tutorials/adaptive/distributions.rst b/docs/pages/tutorials/adaptive/distributions.rst index 643a3a46f..d28a42ec1 100644 --- a/docs/pages/tutorials/adaptive/distributions.rst +++ b/docs/pages/tutorials/adaptive/distributions.rst @@ -33,8 +33,10 @@ Pass any of the following names as the ``dist`` argument to a - uniform in :math:`\log_{10}` - Flat in the log of the parameter (e.g. metallicity). * - ``'kroupa'`` - - power law, :math:`\alpha = -2.3` - - Kroupa initial mass function for the primary mass. + - broken power law (:math:`\alpha = -1.3` for :math:`m < 0.5\,M_\odot`, + :math:`-2.3` above) + - Kroupa initial mass function for the primary mass. Use a lower bound of + :math:`\geq 0.08\,M_\odot` (COSMIC cannot evolve lower-mass stars). * - ``'sana'`` - power law in :math:`\log_{10} P`, :math:`\alpha = -0.55` - Sana et al. (2012) orbital-period distribution. @@ -47,7 +49,7 @@ Pass any of the following names as the ``dist`` argument to a * - ``'uniform_in_cosine'`` - uniform in :math:`\cos\theta` - Isotropic angle for declination-like coordinates. - * - ``'log_normal'`` + * - ``'disberg'`` - log-normal - Natal-kick magnitude, :math:`\ln v \sim \mathcal{N}(5.67, 0.59)`. @@ -62,7 +64,8 @@ Internally each distribution is a **base distribution** composed with a **coordi transform**: * the base distribution (:class:`~cosmic.sample.stroopwafel.distributions.Uniform`, - :class:`~cosmic.sample.stroopwafel.distributions.PowerLaw`, or + :class:`~cosmic.sample.stroopwafel.distributions.PowerLaw`, + :class:`~cosmic.sample.stroopwafel.distributions.BrokenPowerLaw`, or :class:`~cosmic.sample.stroopwafel.distributions.TruncatedNormal`) handles sampling and the density in the *sampling space*; and * the transform (:class:`~cosmic.sample.stroopwafel.distributions.Log10`, @@ -76,12 +79,12 @@ You can build the same objects yourself: .. code-block:: python from cosmic.sample.stroopwafel.distributions import ( - Uniform, PowerLaw, TruncatedNormal, Log10, + Uniform, PowerLaw, BrokenPowerLaw, TruncatedNormal, Log10, ) - Uniform(transform=Log10()) # equivalent to 'flat_in_log' - PowerLaw(-0.55, transform=Log10()) # equivalent to 'sana' - PowerLaw(-2.3) # equivalent to 'kroupa' + Uniform(transform=Log10()) # equivalent to 'flat_in_log' + PowerLaw(-0.55, transform=Log10()) # equivalent to 'sana' + BrokenPowerLaw(breaks=[0.5], alphas=[-1.3, -2.3]) # equivalent to 'kroupa' Bounds are always given to a :class:`~cosmic.sample.stroopwafel.Parameter` in **physical** space; the transform converts them into sampling space automatically (and round-trips diff --git a/src/cosmic/sample/stroopwafel/constants.py b/src/cosmic/sample/stroopwafel/constants.py index f3a04334e..ae570b448 100755 --- a/src/cosmic/sample/stroopwafel/constants.py +++ b/src/cosmic/sample/stroopwafel/constants.py @@ -1,19 +1,5 @@ """Constants used throughout the STROOPWAFEL adaptive sampling module.""" -R_COEFF = [ - [1.71535900, 0.62246212, -0.92557761, -1.16996966, -0.30631491], - [6.59778800, -0.42450044, -12.13339427, -10.73509484, -2.51487077], - [10.08855000, -7.11727086, -31.67119479, -24.24848322, -5.33608972], - [1.01249500, 0.32699690, -0.00923418, -0.03876858, -0.00412750], - [0.07490166, 0.02410413, 0.07233664, 0.03040467, 0.00197741], - [0.01077422, 0.00000000, 0.00000000, 0.00000000, 0.00000000], - [3.08223400, 0.94472050, -2.15200882, -2.49219496, -0.63848738], - [17.84778000, -7.45345690, -48.96066856, -40.05386135, -9.09331816], - [0.00022582, -0.00186899, 0.00388783, 0.00142402, -0.00007671] -] - -R_SOL_TO_AU = 0.00465047 -ZSOL = 0.02 MIN_ENTROPY_CHANGE = 0.01 # Minimum value of (1 - rejection_rate) used in every oversampling and diff --git a/src/cosmic/sample/stroopwafel/distributions.py b/src/cosmic/sample/stroopwafel/distributions.py index 247e24dae..676eac65a 100644 --- a/src/cosmic/sample/stroopwafel/distributions.py +++ b/src/cosmic/sample/stroopwafel/distributions.py @@ -17,7 +17,7 @@ transform maps to and from the physical space the user specifies bounds in and the simulation consumes. For example ``flat_in_log`` is ``Uniform(transform=Log10())`` and ``sana`` is -``PowerLaw(SANA_G, transform=Log10())``. +``PowerLaw(-0.55, transform=Log10())``. To add a distribution, build an instance and either pass it straight to a :class:`~cosmic.sample.stroopwafel.parameter_space.Parameter` or register it @@ -416,7 +416,7 @@ def _ppf(self, u, lo, hi): class TruncatedNormal(Distribution): """Normal distribution truncated to ``[lo, hi]`` in sampling space. - Combined with :class:`Ln` this gives the ``log_normal`` natal-kick prior: + Combined with :class:`Ln` this gives the ``'disberg'`` natal-kick prior: the kick magnitude ``v`` follows ``LogNormal(mu, scale)`` so ``ln(v)`` is normally distributed, and sampling space is ``ln(v)``. diff --git a/src/cosmic/sample/stroopwafel/examples/bh_star.py b/src/cosmic/sample/stroopwafel/examples/bh_star.py index bc033f342..94fff9707 100644 --- a/src/cosmic/sample/stroopwafel/examples/bh_star.py +++ b/src/cosmic/sample/stroopwafel/examples/bh_star.py @@ -120,7 +120,7 @@ Parameter('ecc', 1e-9, 0.99999999, dist='sana_ecc'), Parameter('metallicity', 0.0001, 0.03, dist='flat_in_log'), # --- primary natal kick magnitude only --- - Parameter('natal_kick_1', 0.1, 5000.0, dist='log_normal'), + Parameter('natal_kick_1', 0.1, 5000.0, dist='disberg'), ]) # ------------------------------------------------------------------ diff --git a/src/cosmic/sample/stroopwafel/parameter_space.py b/src/cosmic/sample/stroopwafel/parameter_space.py index 3d4c67560..05a7880a3 100644 --- a/src/cosmic/sample/stroopwafel/parameter_space.py +++ b/src/cosmic/sample/stroopwafel/parameter_space.py @@ -20,7 +20,7 @@ class Parameter: Name of the parameter (used for column ordering). min_value : `float` Lower bound, always in physical space (e.g. solar masses for - ``'kroupa'``, days for ``'sana'``, km/s for ``'log_normal'``). The + ``'kroupa'``, days for ``'sana'``, km/s for ``'disberg'``). The parameter's distribution maps this into sampling space via its transform. max_value : `float` diff --git a/src/cosmic/sample/stroopwafel/rejection.py b/src/cosmic/sample/stroopwafel/rejection.py index af8c49708..4152d5c35 100644 --- a/src/cosmic/sample/stroopwafel/rejection.py +++ b/src/cosmic/sample/stroopwafel/rejection.py @@ -4,45 +4,8 @@ radius, and physical rejection criteria used to discard unphysical binary systems prior to evolution. """ -import numpy as np -from .constants import R_COEFF, ZSOL, R_SOL_TO_AU -from cosmic.utils import calc_Roche_radius - - -def get_zams_radius(mass, metallicity): - """Compute zero-age main sequence radius for arrays of stars. - - Parameters - ---------- - mass : `numpy.ndarray` - (N,) array of stellar masses in solar masses. - metallicity : `numpy.ndarray` - (N,) array of metallicities. - - Returns - ------- - `numpy.ndarray` - (N,) array of ZAMS radii in AU. - """ - mass = np.asarray(mass, dtype=float) - metallicity = np.asarray(metallicity, dtype=float) - - xi = np.log10(metallicity / ZSOL) - - # R_COEFF is a (9, 5) matrix of polynomial coefficients - # For each of the 9 radius coefficients, evaluate the polynomial in xi - R = np.array(R_COEFF) # (9, 5) - # Build Vandermonde matrix: xi^0, xi^1, xi^2, xi^3, xi^4 - powers = np.column_stack([xi**k for k in range(5)]) # (N, 5) - # rc[j, n] = sum_k R[j, k] * xi[n]^k - rc = R @ powers.T # (9, N) - - top = (rc[0] * mass**2.5 + rc[1] * mass**6.5 + rc[2] * mass**11 - + rc[3] * mass**19 + rc[4] * mass**19.5) - bottom = (rc[5] + rc[6] * mass**2 + rc[7] * mass**8.5 - + mass**18.5 + rc[8] * mass**19.5) - - return (top / bottom) * R_SOL_TO_AU +from cosmic.sample.sampler.independent import Sample +from cosmic.utils import calc_Roche_radius, a_from_p def default_reject(binary_params, min_secondary_mass=0.08): @@ -74,15 +37,15 @@ def default_reject(binary_params, min_secondary_mass=0.08): ecc = binary_params['ecc'] metallicity = binary_params['metallicity'] - # Semi-major axis [AU] from Kepler's third law (porb in days, masses in - # solar masses): a^3 [AU^3] = (P [yr])^2 * M [Msun]. - separation = ((porb / 365.25) ** 2 * (mass_1 + mass_2)) ** (1.0 / 3.0) + # compute separation from periods and masses + separation = a_from_p(p=porb, m1=mass_1, m2=mass_2) - # ZAMS radii [AU] (both stars share the binary metallicity) - radius_1 = get_zams_radius(mass_1, metallicity) - radius_2 = get_zams_radius(mass_2, metallicity) + # get stellar radii at ZAMS + sampler = Sample() + radius_1 = sampler.set_reff(mass=mass_1, metallicity=metallicity) + radius_2 = sampler.set_reff(mass=mass_2, metallicity=metallicity) - # Roche lobe radii at periastron + # roche lobe radii at periastron peri_sep = separation * (1 - ecc) rl_1 = calc_Roche_radius(mass_1, mass_2, peri_sep) rl_2 = calc_Roche_radius(mass_2, mass_1, peri_sep) diff --git a/src/cosmic/tests/test_stroopwafel.py b/src/cosmic/tests/test_stroopwafel.py index 36db0e960..ddb884676 100644 --- a/src/cosmic/tests/test_stroopwafel.py +++ b/src/cosmic/tests/test_stroopwafel.py @@ -17,16 +17,19 @@ from cosmic.sample.stroopwafel import ParameterSpace, Parameter, AdaptiveSampler from cosmic.sample.stroopwafel.distributions import ( DISTRIBUTIONS, register, get_distribution, - Distribution, Uniform, PowerLaw, TruncatedNormal, + Distribution, Uniform, PowerLaw, BrokenPowerLaw, TruncatedNormal, Identity, Log10, Ln, Sin, CosShift, ) from cosmic.sample.stroopwafel.mixture_model import GaussianMixture -from cosmic.sample.stroopwafel.rejection import get_zams_radius, default_reject -from cosmic.sample.stroopwafel.constants import ( - ALPHA_IMF, SANA_G, SANA_ECC, NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA, -) +from cosmic.sample.stroopwafel.rejection import default_reject +from scipy.integrate import cumulative_trapezoid from cosmic.output import COSMICStroopOutput +# Built-in distribution parameters, kept in sync with the DISTRIBUTIONS registry. +KROUPA_BREAK, KROUPA_ALPHA_LOW, KROUPA_ALPHA_HIGH = 0.5, -1.3, -2.3 +SANA_G, SANA_ECC = -0.55, -0.45 +DISBERG_MU, DISBERG_SIGMA = 5.67, 0.59 + HALF_PI = np.pi / 2 KS_PVALUE_MIN = 0.01 # samplers must not be rejected against their own CDF @@ -56,7 +59,7 @@ def canonical_space(): Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), Parameter('ecc', 1e-9, 0.99999999, dist='sana_ecc'), Parameter('metallicity', 1e-4, 0.03, dist='flat_in_log'), - Parameter('natal_kick_1', 0.1, 5000.0, dist='log_normal'), + Parameter('natal_kick_1', 0.1, 5000.0, dist='disberg'), ]) @@ -128,19 +131,21 @@ def test_registry_contents(self): self.assertEqual( set(DISTRIBUTIONS), {'uniform', 'flat_in_log', 'uniform_in_sine', 'uniform_in_cosine', - 'kroupa', 'sana', 'sana_ecc', 'log_normal'}, + 'kroupa', 'sana', 'sana_ecc', 'disberg'}, ) def test_builtin_compositions(self): - self.assertIsInstance(DISTRIBUTIONS['kroupa'], PowerLaw) + self.assertIsInstance(DISTRIBUTIONS['kroupa'], BrokenPowerLaw) self.assertIsInstance(DISTRIBUTIONS['kroupa'].transform, Identity) + np.testing.assert_array_equal(DISTRIBUTIONS['kroupa'].breaks, [KROUPA_BREAK]) + np.testing.assert_array_equal(DISTRIBUTIONS['kroupa'].alphas, + [KROUPA_ALPHA_LOW, KROUPA_ALPHA_HIGH]) self.assertIsInstance(DISTRIBUTIONS['sana'], PowerLaw) self.assertIsInstance(DISTRIBUTIONS['sana'].transform, Log10) self.assertIsInstance(DISTRIBUTIONS['flat_in_log'], Uniform) self.assertIsInstance(DISTRIBUTIONS['flat_in_log'].transform, Log10) - self.assertIsInstance(DISTRIBUTIONS['log_normal'], TruncatedNormal) - self.assertIsInstance(DISTRIBUTIONS['log_normal'].transform, Ln) - self.assertEqual(DISTRIBUTIONS['kroupa'].alpha, ALPHA_IMF) + self.assertIsInstance(DISTRIBUTIONS['disberg'], TruncatedNormal) + self.assertIsInstance(DISTRIBUTIONS['disberg'].transform, Ln) self.assertEqual(DISTRIBUTIONS['sana'].alpha, SANA_G) self.assertEqual(DISTRIBUTIONS['sana_ecc'].alpha, SANA_ECC) @@ -150,7 +155,7 @@ def test_uniform_pdf_constant(self): np.testing.assert_allclose(d.pdf(vals, 2.0, 8.0), 1.0 / 6.0) def test_powerlaw_pdf_matches_closed_form(self): - for alpha, lo, hi in [(ALPHA_IMF, 5.0, 150.0), + for alpha, lo, hi in [(KROUPA_ALPHA_HIGH, 5.0, 150.0), (SANA_G, 0.15, 5.5), (SANA_ECC, 1e-9, 0.999)]: vals = np.linspace(lo * 1.01, hi * 0.99, 100) @@ -160,13 +165,13 @@ def test_powerlaw_pdf_matches_closed_form(self): np.testing.assert_allclose(got, expected, rtol=1e-12) def test_powerlaw_pdf_normalised(self): - d = PowerLaw(ALPHA_IMF) + d = PowerLaw(KROUPA_ALPHA_HIGH) lo, hi = 5.0, 150.0 x = np.linspace(lo, hi, 200000) self.assertAlmostEqual(trapezoid(d.pdf(x, lo, hi), x), 1.0, places=4) def test_truncnorm_pdf_normalised(self): - d = TruncatedNormal(NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA) + d = TruncatedNormal(DISBERG_MU, DISBERG_SIGMA) lo, hi = 2.0, 9.0 x = np.linspace(lo, hi, 200000) self.assertAlmostEqual(trapezoid(d.pdf(x, lo, hi), x), 1.0, places=4) @@ -179,7 +184,7 @@ def test_uniform_sampler_ks(self): self.assertGreater(p, KS_PVALUE_MIN) def test_powerlaw_sampler_ks(self): - for alpha, lo, hi, seed in [(ALPHA_IMF, 5.0, 150.0, 1), + for alpha, lo, hi, seed in [(KROUPA_ALPHA_HIGH, 5.0, 150.0, 1), (SANA_G, 0.15, 5.5, 2), (SANA_ECC, 1e-9, 0.999, 3)]: rng = np.random.default_rng(seed) @@ -190,7 +195,7 @@ def test_powerlaw_sampler_ks(self): def test_truncnorm_sampler_ks(self): rng = np.random.default_rng(4) - mu, scale, lo, hi = NATAL_KICK_LOG_MU, NATAL_KICK_LOG_SIGMA, 2.0, 9.0 + mu, scale, lo, hi = DISBERG_MU, DISBERG_SIGMA, 2.0, 9.0 s = TruncatedNormal(mu, scale).sample(20000, lo, hi, rng=rng) self.assertTrue(np.all((s >= lo) & (s <= hi))) p = kstest(s, lambda v: _truncnorm_cdf(v, mu, scale, lo, hi)).pvalue @@ -205,7 +210,7 @@ def test_default_sigma_is_avg_density_over_pdf(self): ) def test_powerlaw_sigma_positive_and_raises_on_nonpositive_lo(self): - d = PowerLaw(ALPHA_IMF) + d = PowerLaw(KROUPA_ALPHA_HIGH) vals = np.linspace(6.0, 140.0, 100) sig = d.sigma(vals, 5.0, 150.0, 0.05) self.assertTrue(np.all(sig > 0) and np.all(np.isfinite(sig))) @@ -213,6 +218,71 @@ def test_powerlaw_sigma_positive_and_raises_on_nonpositive_lo(self): d.sigma(vals, 0.0, 150.0, 0.05) +# ========================================================================== +# Broken power law (Kroupa IMF) +# ========================================================================== +class TestBrokenPowerLaw(unittest.TestCase): + + KROUPA = BrokenPowerLaw(breaks=[KROUPA_BREAK], alphas=[KROUPA_ALPHA_LOW, KROUPA_ALPHA_HIGH]) + + def test_input_validation(self): + with self.assertRaises(ValueError): + BrokenPowerLaw(breaks=[0.5], alphas=[-2.3]) # wrong length + with self.assertRaises(ValueError): + BrokenPowerLaw(breaks=[0.5, 0.3], alphas=[-1, -2, -3]) # not increasing + with self.assertRaises(ValueError): + BrokenPowerLaw(breaks=[0.5], alphas=[-1.0, -2.3]) # exponent -1 + + def test_pdf_continuous_at_break(self): + lo, hi = 0.08, 100.0 + below = self.KROUPA.pdf(np.array([KROUPA_BREAK - 1e-6]), lo, hi)[0] + above = self.KROUPA.pdf(np.array([KROUPA_BREAK + 1e-6]), lo, hi)[0] + self.assertAlmostEqual(below, above, places=4) + + def test_pdf_normalised_across_break(self): + lo, hi = 0.08, 100.0 + x = np.geomspace(lo, hi, 200000) + self.assertAlmostEqual(trapezoid(self.KROUPA.pdf(x, lo, hi), x), 1.0, places=3) + + def test_slope_changes_at_break(self): + # Local log-log slope should match the segment exponents. + lo, hi = 0.08, 100.0 + def local_slope(x0): + x = np.array([x0 * 0.999, x0 * 1.001]) + p = self.KROUPA.pdf(x, lo, hi) + return np.diff(np.log(p))[0] / np.diff(np.log(x))[0] + self.assertAlmostEqual(local_slope(0.2), KROUPA_ALPHA_LOW, places=2) + self.assertAlmostEqual(local_slope(5.0), KROUPA_ALPHA_HIGH, places=2) + + def test_reduces_to_powerlaw_above_break(self): + # With no break in range the example regime (m1 > 5) is unchanged. + lo, hi = 5.0, 150.0 + pl = PowerLaw(KROUPA_ALPHA_HIGH) + vals = np.linspace(lo * 1.01, hi * 0.99, 200) + np.testing.assert_allclose(self.KROUPA.pdf(vals, lo, hi), + pl.pdf(vals, lo, hi), rtol=1e-10) + a = self.KROUPA.sample(2000, lo, hi, rng=np.random.default_rng(0)) + b = pl.sample(2000, lo, hi, rng=np.random.default_rng(0)) + np.testing.assert_allclose(a, b, rtol=1e-10) + + def test_sampler_follows_pdf(self): + # KS test against the (independently integrated) pdf, across the break. + lo, hi = 0.08, 50.0 + s = self.KROUPA.sample(40000, lo, hi, rng=np.random.default_rng(7)) + self.assertTrue(np.all((s >= lo) & (s <= hi))) + grid = np.geomspace(lo, hi, 40000) + cdf_vals = cumulative_trapezoid(self.KROUPA.pdf(grid, lo, hi), grid, initial=0) + cdf_vals /= cdf_vals[-1] + pvalue = kstest(s, lambda v: np.interp(v, grid, cdf_vals)).pvalue + self.assertGreater(pvalue, KS_PVALUE_MIN) + + def test_sigma_positive_finite_across_break(self): + lo, hi = 0.08, 100.0 + vals = np.array([0.1, 0.3, 0.5, 1.0, 10.0, 80.0]) + sig = self.KROUPA.sigma(vals, lo, hi, 0.02) + self.assertTrue(np.all(sig > 0) and np.all(np.isfinite(sig))) + + # ========================================================================== # Registry / extensibility # ========================================================================== @@ -314,13 +384,6 @@ def test_compute_sigma_reports_parameter_name_on_error(self): # ========================================================================== class TestRejection(unittest.TestCase): - def test_get_zams_radius_positive_finite(self): - masses = np.array([1.0, 10.0, 30.0, 100.0]) - mets = np.array([0.02, 0.02, 0.001, 0.014]) - r = get_zams_radius(masses, mets) - self.assertEqual(r.shape, (4,)) - self.assertTrue(np.all(r > 0) and np.all(np.isfinite(r))) - def test_wide_binary_not_rejected(self): binary_params = { 'mass_1': np.array([20.0]), From 36fe089371b7319f7f8bbde78e62a26da8818f81 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Sun, 28 Jun 2026 23:33:03 -0400 Subject: [PATCH 35/45] delete constants, move into attributes --- src/cosmic/sample/stroopwafel/constants.py | 9 ------- src/cosmic/sample/stroopwafel/meson.build | 1 - .../sample/stroopwafel/mixture_model.py | 26 +++++++++++++------ 3 files changed, 18 insertions(+), 18 deletions(-) delete mode 100755 src/cosmic/sample/stroopwafel/constants.py diff --git a/src/cosmic/sample/stroopwafel/constants.py b/src/cosmic/sample/stroopwafel/constants.py deleted file mode 100755 index ae570b448..000000000 --- a/src/cosmic/sample/stroopwafel/constants.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Constants used throughout the STROOPWAFEL adaptive sampling module.""" - -MIN_ENTROPY_CHANGE = 0.01 - -# 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 array sizes or overflowing float64 normalisation constants. -MIN_ACTIVE_FRACTION = 0.01 diff --git a/src/cosmic/sample/stroopwafel/meson.build b/src/cosmic/sample/stroopwafel/meson.build index 0af7964e8..103abb233 100644 --- a/src/cosmic/sample/stroopwafel/meson.build +++ b/src/cosmic/sample/stroopwafel/meson.build @@ -1,6 +1,5 @@ python_sources = [ '__init__.py', - 'constants.py', 'distributions.py', 'main.py', 'mixture_model.py', diff --git a/src/cosmic/sample/stroopwafel/mixture_model.py b/src/cosmic/sample/stroopwafel/mixture_model.py index 6366e860a..7b2ebc93e 100644 --- a/src/cosmic/sample/stroopwafel/mixture_model.py +++ b/src/cosmic/sample/stroopwafel/mixture_model.py @@ -5,7 +5,6 @@ """ import numpy as np from scipy.stats import multivariate_normal, entropy as scipy_entropy -from .constants import MIN_ENTROPY_CHANGE, MIN_ACTIVE_FRACTION class GaussianMixture: @@ -23,11 +22,14 @@ class GaussianMixture: Fraction of samples that fall outside bounds, by default 0.0 """ - def __init__(self, means, covariances, alphas, rejection_rate=0.0): + def __init__(self, means, covariances, alphas, rejection_rate=0.0, + min_active_fraction=0.01, min_entropy_change=0.01): self.means = np.asarray(means) self.covariances = np.asarray(covariances) self.alphas = np.asarray(alphas, dtype=float) self.rejection_rate = rejection_rate + self.min_active_fraction = min_active_fraction + self.min_entropy_change = min_entropy_change @property def n_components(self): @@ -38,7 +40,8 @@ def ndim(self): return self.means.shape[1] @classmethod - def from_hits(cls, hit_samples, param_space, average_density_one_dim, kappa=1.0): + def from_hits(cls, hit_samples, param_space, average_density_one_dim, kappa=1.0, + min_active_fraction=0.01, min_entropy_change=0.01): """Create a Gaussian mixture by placing one component at each hit. Parameters @@ -52,6 +55,12 @@ def from_hits(cls, hit_samples, param_space, average_density_one_dim, kappa=1.0) ``1 / num_explored ** (1 / D)``. kappa : `float`, optional Width scaling factor for the Gaussian covariances, by default 1.0 + min_active_fraction : `float`, optional + Minimum value of (1 - rejection_rate) used in oversampling and + normalisation calculations, by default 0.01 + min_entropy_change : `float`, optional + Minimum change in normalised effective sample size (entropy) to + avoid reverting to the previous mixture state, by default 0.01 Returns ------- @@ -72,7 +81,8 @@ def from_hits(cls, hit_samples, param_space, average_density_one_dim, kappa=1.0) # Equal mixture weights alphas = np.full(K, 1.0 / K) - return cls(hit_samples.copy(), covariances, alphas) + return cls(hit_samples.copy(), covariances, alphas, + min_active_fraction=min_active_fraction, min_entropy_change=min_entropy_change) def sample(self, n_total, param_space, consider_rejection=False, rng=None): """Sample from the mixture distribution. @@ -106,7 +116,7 @@ def sample(self, n_total, param_space, consider_rejection=False, rng=None): for k in range(self.n_components): n_k = int(np.ceil(n_total * self.alphas[k])) if consider_rejection and self.rejection_rate > 0.0: - active = max(1.0 - self.rejection_rate, MIN_ACTIVE_FRACTION) + active = max(1.0 - self.rejection_rate, self.min_active_fraction) n_k = int(2 * np.ceil(n_k / active)) if n_k <= 0: continue @@ -242,8 +252,8 @@ def update_em(self, samples, is_hit, prior_probs, prior_fraction_rejected, True if the entropy check triggers reversion to the previous mixture state. """ - pi_norm = 1.0 / max(1.0 - prior_fraction_rejected, MIN_ACTIVE_FRACTION) - q_norm = 1.0 / max(1.0 - self.rejection_rate, MIN_ACTIVE_FRACTION) + pi_norm = 1.0 / max(1.0 - prior_fraction_rejected, self.min_active_fraction) + q_norm = 1.0 / max(1.0 - self.rejection_rate, self.min_active_fraction) pi = prior_probs * pi_norm is_hit = np.asarray(is_hit, dtype=float) @@ -301,7 +311,7 @@ def update_em(self, samples, is_hit, prior_probs, prior_fraction_rejected, # generations, which indicates the mixture has stopped improving. entropy_change = np.exp(scipy_entropy(weights_normalized[:, 0])) / N if entropies is not None: - if len(entropies) >= 1 and entropy_change - entropies[-1] < MIN_ENTROPY_CHANGE: + if len(entropies) >= 1 and entropy_change - entropies[-1] < self.min_entropy_change: return True # Signal to revert entropies.append(entropy_change) From 75d241a0b76eca33aa26e40c440bb338c36a88f7 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Sun, 28 Jun 2026 23:34:34 -0400 Subject: [PATCH 36/45] make checkpointing save state better --- docs/pages/tutorials/adaptive/checkpoint.rst | 65 ++++---- src/cosmic/output.py | 58 ++++--- src/cosmic/sample/stroopwafel/main.py | 165 ++++++++++--------- src/cosmic/tests/test_stroopwafel.py | 73 ++++++++ 4 files changed, 229 insertions(+), 132 deletions(-) diff --git a/docs/pages/tutorials/adaptive/checkpoint.rst b/docs/pages/tutorials/adaptive/checkpoint.rst index 5663e007d..aeb55962f 100644 --- a/docs/pages/tutorials/adaptive/checkpoint.rst +++ b/docs/pages/tutorials/adaptive/checkpoint.rst @@ -82,56 +82,59 @@ Stage 2 — load the checkpoint and refine In a second script (or cluster job) rebuild the sampler with :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.from_checkpoint`, then call -:meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run_refinement`. +:meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run_refinement`. The checkpoint is +**self-contained**, so this needs nothing but the file: .. code-block:: python from cosmic.sample.stroopwafel import AdaptiveSampler - from cosmic.sample.stroopwafel.presets import any_dco - from cosmic.sample.stroopwafel.rejection import default_reject - # `params`, `derive_params`, and `BSEDict` must be available again here — - # in practice, import them from a shared module used by both jobs. + sampler = AdaptiveSampler.from_checkpoint('checkpoint.h5') + result = sampler.run_refinement() + result.save('result.h5') + +You do not need to re-import or re-specify the parameter space, ``BSEDict``, or any of the +callables — they were all saved into the checkpoint. If you *want* to change something for +the refinement phase (a common one is running on more cores, or with a larger budget than +exploration), pass it as a keyword override: + +.. code-block:: python sampler = AdaptiveSampler.from_checkpoint( - 'checkpoint.h5', - parameter_space=params, - batch_size=1000, - BSEDict=BSEDict, - is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), - derive_params=derive_params, - reject_systems=default_reject, - output_path='output/refine', - nproc=4, - seed=7, + 'checkpoint.h5', nproc=16, total_systems=1_000_000, ) - result = sampler.run_refinement() - result.save('result.h5') +Any :class:`~cosmic.sample.stroopwafel.AdaptiveSampler` constructor argument may be +overridden this way (``parameter_space``, ``BSEDict``, ``derive_params``, +``reject_systems``, ``is_interesting``, ``batch_size``, ``output_path``, ``nproc``, +``kappa``, ``n_generations``, ``only_save_hit_tables``, ``seed``). The ``result`` is an ordinary :class:`~cosmic.output.COSMICStroopOutput` — identical in form to what a single :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run` would have produced — so you can analyse it exactly as described in :ref:`adaptive_outputs`. -What is and isn't stored in a checkpoint -======================================== +What is stored in a checkpoint +============================== -A checkpoint stores everything that is *derived from running COSMIC*: +A checkpoint is a complete snapshot — it stores both the exploration *results* and the full +*configuration*, so refinement can resume with no further input: -* the fitted Gaussian mixture model, -* every exploration sample (in the internal sampling space) and its hit/bookkeeping flags, +* the fitted Gaussian mixture model; +* every exploration sample (in the internal sampling space) and its hit/bookkeeping flags; * the COSMIC output tables (``bpp``, ``bcm``, ``initC``, ``kick_info``) for the explored - systems, and + systems; * the scalar counters needed to compute unbiased weights later (``num_explored``, - ``fraction_explored``, ``prior_fraction_rejected``, ``total_systems``, ...). - -It deliberately does **not** store your Python callables or physics settings — the -``BSEDict``, ``derive_params``, ``reject_systems``, and ``is_interesting`` functions are -not serialisable in general, so you must supply them again to -:meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.from_checkpoint`. Only the parameter -**names** are stored, and they are checked against the ``parameter_space`` you provide; a -mismatch raises a ``ValueError`` to stop you from refining against the wrong setup. + ``fraction_explored``, ``prior_fraction_rejected``, ...); and +* the full configuration needed to rebuild the sampler: the parameter space, ``BSEDict``, + the ``derive_params`` / ``reject_systems`` / ``is_interesting`` callables, the remaining + scalar settings, and the live RNG state. + +The callables and parameter space are serialised with :mod:`dill` (a dependency COSMIC +already ships), so lambdas and closures — such as the ``is_interesting`` returned by +``any_dco(...)`` — round-trip correctly. Because the RNG state is stored too, refinement +continues the random stream seamlessly rather than restarting it (pass ``seed=`` to +``from_checkpoint`` if you instead want a fresh stream). Reusing one checkpoint for several refinement jobs diff --git a/src/cosmic/output.py b/src/cosmic/output.py index 475d485d8..8625d5f92 100644 --- a/src/cosmic/output.py +++ b/src/cosmic/output.py @@ -1,4 +1,5 @@ import json +import dill import pandas as pd import h5py as h5 from cosmic.evolve import Evolve @@ -540,8 +541,22 @@ class STROOPWAFELCheckpoint: # Job 2 - refinement (can be a larger allocation) python run_refine.py # reads checkpoint.h5, writes result.h5 + A checkpoint is **self-contained**: alongside the exploration data it + stores everything needed to rebuild the sampler — the parameter space, + ``BSEDict``, the ``derive_params``/``reject_systems``/``is_interesting`` + callables, the remaining scalar settings, and the live RNG state — so + :meth:`AdaptiveSampler.from_checkpoint` needs nothing but the file. The + callables and parameter space are serialised with :mod:`dill`. + Parameters ---------- + config : `dict` + Everything needed to reconstruct the :class:`AdaptiveSampler` for the + refinement phase: the constructor keyword arguments (``parameter_space``, + ``total_systems``, ``batch_size``, ``BSEDict``, ``is_interesting``, + ``derive_params``, ``reject_systems``, ``output_path``, ``nproc``, + ``kappa``, ``n_generations``, ``only_save_hit_tables``) plus the live + ``rng``. mixture : `GaussianMixture` or None Gaussian mixture fitted to exploration hits. ``None`` if no hits were found or adaptation has not been run yet. @@ -555,8 +570,6 @@ class STROOPWAFELCheckpoint: COSMIC output from exploration, indexed by globally unique ``bin_num`` so that ``samples[bin_num]`` gives the corresponding physical parameters. - param_names : `list` of `str` - Parameter names for the D columns of ``samples``. num_explored : `int` Systems evolved during exploration. num_hits : `int` @@ -567,17 +580,13 @@ class STROOPWAFELCheckpoint: Adaptive fraction of total budget used for exploration. prior_fraction_rejected : `float` Estimated fraction of prior samples that fail physical rejection. - total_systems : `int` - Full system budget (exploration + refinement combined). - n_generations : `int` - Number of refinement generations originally requested. """ - def __init__(self, mixture, samples, is_hit, generation, gaussian_idx, - bpp, bcm, initC, kick_info, param_names, + def __init__(self, config, mixture, samples, is_hit, generation, gaussian_idx, + bpp, bcm, initC, kick_info, num_explored, num_hits, num_hits_exploratory, - fraction_explored, prior_fraction_rejected, - total_systems, n_generations): + fraction_explored, prior_fraction_rejected): + self.config = dict(config) self.mixture = mixture self.samples = np.asarray(samples) self.is_hit = np.asarray(is_hit, dtype=bool) @@ -587,14 +596,16 @@ def __init__(self, mixture, samples, is_hit, generation, gaussian_idx, self.bcm = bcm self.initC = initC self.kick_info = kick_info - self.param_names = list(param_names) self.num_explored = int(num_explored) self.num_hits = int(num_hits) self.num_hits_exploratory = int(num_hits_exploratory) self.fraction_explored = float(fraction_explored) self.prior_fraction_rejected = float(prior_fraction_rejected) - self.total_systems = int(total_systems) - self.n_generations = int(n_generations) + + # Convenience views derived from the stored config. + self.param_names = list(config['parameter_space'].names) + self.total_systems = int(config['total_systems']) + self.n_generations = int(config['n_generations']) def __repr__(self): adapted = self.mixture is not None @@ -639,16 +650,21 @@ def save(self, path): mg.create_dataset('alphas', data=self.mixture.alphas) mg.attrs['rejection_rate'] = self.mixture.rejection_rate - # Scalar metadata + # Full reconstruction config (parameter space, BSEDict, callables, + # RNG, scalar settings) serialised with dill so closures/lambdas + # survive the round-trip. + blob = np.frombuffer(dill.dumps(self.config), dtype=np.uint8) + if 'config' in f: + del f['config'] + f.create_dataset('config', data=blob) + + # Progress-state metadata (everything else lives in `config`). meta = f.require_group('meta') - meta.attrs['param_names'] = self.param_names meta.attrs['num_explored'] = self.num_explored meta.attrs['num_hits'] = self.num_hits meta.attrs['num_hits_exploratory'] = self.num_hits_exploratory meta.attrs['fraction_explored'] = self.fraction_explored meta.attrs['prior_fraction_rejected'] = self.prior_fraction_rejected - meta.attrs['total_systems'] = self.total_systems - meta.attrs['n_generations'] = self.n_generations @classmethod def from_file(cls, path): @@ -687,15 +703,14 @@ def from_file(cls, path): rejection_rate=float(mg.attrs['rejection_rate']), ) + config = dill.loads(f['config'][:].tobytes()) + meta = f['meta'] - param_names = list(meta.attrs['param_names']) num_explored = int(meta.attrs['num_explored']) num_hits = int(meta.attrs['num_hits']) num_hits_exploratory = int(meta.attrs['num_hits_exploratory']) fraction_explored = float(meta.attrs['fraction_explored']) prior_fraction_rejected = float(meta.attrs['prior_fraction_rejected']) - total_systems = int(meta.attrs['total_systems']) - n_generations = int(meta.attrs['n_generations']) # Re-apply the bin_num index so the same invariant holds as # when the checkpoint was created. @@ -704,16 +719,15 @@ def from_file(cls, path): df.set_index('bin_num', drop=False, inplace=True) return cls( + config=config, mixture=mixture, samples=samples, is_hit=is_hit, generation=generation, gaussian_idx=gaussian_idx, bpp=bpp, bcm=bcm, initC=initC, kick_info=kick_info, - param_names=param_names, num_explored=num_explored, num_hits=num_hits, num_hits_exploratory=num_hits_exploratory, fraction_explored=fraction_explored, prior_fraction_rejected=prior_fraction_rejected, - total_systems=total_systems, n_generations=n_generations, ) diff --git a/src/cosmic/sample/stroopwafel/main.py b/src/cosmic/sample/stroopwafel/main.py index 56defc481..0fea86f19 100644 --- a/src/cosmic/sample/stroopwafel/main.py +++ b/src/cosmic/sample/stroopwafel/main.py @@ -8,7 +8,6 @@ from cosmic.output import COSMICStroopOutput, STROOPWAFELCheckpoint from .mixture_model import GaussianMixture -from .constants import MIN_ACTIVE_FRACTION from .rejection import default_reject @@ -66,6 +65,16 @@ class AdaptiveSampler: ``weights``, and ``is_hit`` arrays are unaffected — all systems are retained there for correct importance-weight calculation. By default False. + min_active_fraction : `float`, 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_change : `float`, optional + Minimum change in entropy required for a generation to be considered + as having made progress. By default 0.01. """ #: The parameters that fully define a binary for COSMIC. Each must be @@ -77,7 +86,8 @@ def __init__(self, parameter_space, total_systems, batch_size, BSEDict, is_interesting, derive_params=None, reject_systems="default", output_path='output', nproc=1, kappa=1.0, n_generations=1, mc_only=False, seed=None, - only_save_hit_tables=False): + only_save_hit_tables=False, + min_active_fraction=0.01, min_entropy_change=0.01): self.param_space = parameter_space self.total_systems = total_systems self.batch_size = batch_size @@ -92,6 +102,8 @@ def __init__(self, parameter_space, total_systems, batch_size, BSEDict, self.mc_only = mc_only self.only_save_hit_tables = only_save_hit_tables self.rng = np.random.default_rng(seed) + self.min_active_fraction = min_active_fraction + self.min_entropy_change = min_entropy_change # State self.num_explored = 0 @@ -173,13 +185,8 @@ def run_refinement(self): :meth:`from_checkpoint` has restored exploration state. Suitable as the second SLURM job:: - # job_refine.py - sampler = AdaptiveSampler.from_checkpoint( - 'checkpoint.h5', - parameter_space=params, batch_size=1000, - BSEDict=BSEDict, is_interesting=is_bh_star, - derive_params=derive_params, reject_systems=default_reject, - ) + # job_refine.py (the checkpoint is self-contained) + sampler = AdaptiveSampler.from_checkpoint('checkpoint.h5') result = sampler.run_refinement() result.save('result.h5') @@ -191,88 +198,72 @@ def run_refinement(self): self._refine() return self._compute_weights() + #: Constructor arguments that may be overridden in :meth:`from_checkpoint`. + _OVERRIDABLE = frozenset({ + 'parameter_space', 'total_systems', 'batch_size', 'BSEDict', + 'is_interesting', 'derive_params', 'reject_systems', 'output_path', + 'nproc', 'kappa', 'n_generations', 'only_save_hit_tables', 'seed', + }) + @classmethod - def from_checkpoint(cls, checkpoint, parameter_space, batch_size, BSEDict, - is_interesting, derive_params=None, reject_systems=None, - output_path='output', nproc=1, seed=None, - total_systems=None, n_generations=None, - only_save_hit_tables=False): - """Create a sampler pre-loaded with state from a :class:`STROOPWAFELCheckpoint`. - - The callable arguments (``BSEDict``, ``derive_params``, etc.) are - not stored in the checkpoint and must be supplied again. All other - state — explored samples, COSMIC output, mixture model, scalar - counters — is restored from the checkpoint. + def from_checkpoint(cls, checkpoint, **overrides): + """Rebuild a sampler from a checkpoint, ready for :meth:`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 ---------- checkpoint : `STROOPWAFELCheckpoint` or `str` A checkpoint object or path to an HDF5 file written by :meth:`STROOPWAFELCheckpoint.save`. - parameter_space : `ParameterSpace` - Must use the same parameter names as when the checkpoint was - created; a `ValueError` is raised if they differ. - batch_size : `int` - Systems per COSMIC call for the refinement phase. - BSEDict : `dict` - COSMIC binary stellar evolution parameters. - is_interesting : `callable` - Same signature as for the original sampler. - derive_params : `callable`, optional - Same signature as for the original sampler. - reject_systems : `callable`, optional - Same signature as for the original sampler. - output_path : `str`, optional - Directory for output files, by default ``'output'`` - nproc : `int`, optional - CPU cores for COSMIC, by default 1 - seed : `int` or None, optional - RNG seed for refinement sampling, by default None - total_systems : `int`, optional - Override the total system budget stored in the checkpoint. - n_generations : `int`, optional - Override the number of refinement generations stored in the - checkpoint. + **overrides + Any :class:`AdaptiveSampler` constructor argument + (``parameter_space``, ``total_systems``, ``batch_size``, + ``BSEDict``, ``is_interesting``, ``derive_params``, + ``reject_systems``, ``output_path``, ``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 :meth:`run_refinement`. - - Raises - ------ - ValueError - If ``parameter_space.names`` does not match the checkpoint's - ``param_names``. """ if isinstance(checkpoint, (str, os.PathLike)): checkpoint = STROOPWAFELCheckpoint.from_file(checkpoint) - if list(parameter_space.names) != list(checkpoint.param_names): - raise ValueError( - f"Parameter space mismatch.\n" - f" checkpoint : {checkpoint.param_names}\n" - f" provided : {list(parameter_space.names)}" + unknown = set(overrides) - cls._OVERRIDABLE + if unknown: + raise TypeError( + f"Unknown from_checkpoint override(s): {sorted(unknown)}. " + f"Allowed: {sorted(cls._OVERRIDABLE)}." ) - ts = total_systems if total_systems is not None else checkpoint.total_systems - ng = n_generations if n_generations is not None else checkpoint.n_generations - - sampler = cls( - parameter_space=parameter_space, - total_systems=ts, - batch_size=batch_size, - BSEDict=BSEDict, - is_interesting=is_interesting, - derive_params=derive_params, - reject_systems=reject_systems, - output_path=output_path, - nproc=nproc, - n_generations=ng, - mc_only=False, - seed=seed, - only_save_hit_tables=only_save_hit_tables, - ) + # Start from the stored config, apply overrides; refinement is never mc_only. + kwargs = dict(checkpoint.config) + stored_rng = kwargs.pop('rng', None) + kwargs.update(overrides) + kwargs['mc_only'] = False + + sampler = cls(**kwargs) + # Restore the exact RNG stream unless the caller asked for a new seed. + if 'seed' not in overrides and stored_rng is not None: + sampler.rng = stored_rng sampler._load_checkpoint(checkpoint) return sampler @@ -284,21 +275,37 @@ def _make_checkpoint(self): all_gaussian_idx = np.concatenate(self._all_gaussian_idx) bpp, bcm, initC, kick_info = self._build_cosmic_output() + # Everything needed to rebuild the sampler for refinement, so that + # from_checkpoint() needs nothing but the file. + config = { + 'parameter_space': self.param_space, + 'total_systems': self.total_systems, + 'batch_size': self.batch_size, + 'BSEDict': self.bse_dict, + 'is_interesting': self.is_interesting_fn, + 'derive_params': self.derive_fn, + 'reject_systems': self.reject_fn, + 'output_path': self.output_path, + 'nproc': self.nproc, + 'kappa': self.kappa, + 'n_generations': self.n_generations, + 'only_save_hit_tables': self.only_save_hit_tables, + 'rng': self.rng, + } + return STROOPWAFELCheckpoint( + config=config, mixture=self.mixture, samples=all_samples, is_hit=all_is_hit, generation=all_generation, gaussian_idx=all_gaussian_idx, bpp=bpp, bcm=bcm, initC=initC, kick_info=kick_info, - param_names=self.param_space.names, num_explored=self.num_explored, num_hits=self.num_hits, num_hits_exploratory=getattr(self, 'num_hits_exploratory', self.num_hits), fraction_explored=self.fraction_explored, prior_fraction_rejected=self.prior_fraction_rejected, - total_systems=self.total_systems, - n_generations=self.n_generations, ) def _load_checkpoint(self, checkpoint): @@ -451,7 +458,7 @@ def _explore(self): while self._should_continue_exploring(): n_oversample = int(2 * np.ceil( - self.batch_size / max(1.0 - self.prior_fraction_rejected, MIN_ACTIVE_FRACTION) + self.batch_size / max(1.0 - self.prior_fraction_rejected, self.min_active_fraction) )) # Sample from prior @@ -697,7 +704,7 @@ def _compute_weights(self): N = len(all_samples) # Prior probabilities - pi_norm = 1.0 / max(1.0 - self.prior_fraction_rejected, MIN_ACTIVE_FRACTION) + pi_norm = 1.0 / max(1.0 - self.prior_fraction_rejected, self.min_active_fraction) pi = self.param_space.compute_prior(all_samples) * pi_norm # Start with the exploration-phase contribution to denominator @@ -706,7 +713,7 @@ def _compute_weights(self): # Add refinement-phase contributions from each generation's mixture if self.mixture is not None and not self.mc_only: - q_norm = 1.0 / max(1.0 - self.mixture.rejection_rate, MIN_ACTIVE_FRACTION) + q_norm = 1.0 / max(1.0 - self.mixture.rejection_rate, self.min_active_fraction) # Evaluate the mixture PDF incrementally (memory-efficient) for k in range(self.mixture.n_components): xPDF_k = multivariate_normal.pdf( diff --git a/src/cosmic/tests/test_stroopwafel.py b/src/cosmic/tests/test_stroopwafel.py index ddb884676..26ccad8b2 100644 --- a/src/cosmic/tests/test_stroopwafel.py +++ b/src/cosmic/tests/test_stroopwafel.py @@ -500,6 +500,79 @@ def test_sample_shapes_and_idx_range(self): self.assertTrue(np.all((midx >= 0) & (midx < 10))) +# ========================================================================== +# Checkpoint (self-contained restore) +# ========================================================================== +class TestCheckpoint(unittest.TestCase): + + def _build(self): + from cosmic.output import STROOPWAFELCheckpoint + ps = ParameterSpace([ + Parameter('mass_1', 5.0, 150.0, dist='kroupa'), + Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), + ]) + rng = np.random.default_rng(123) + rng.uniform(size=10) # advance the stream (as exploration would) + config = { + 'parameter_space': ps, 'total_systems': 1000, 'batch_size': 50, + 'BSEDict': {'foo': 1}, + 'is_interesting': lambda bpp: (0, np.array([], dtype=int)), + 'derive_params': lambda s: {'mass_2': s['mass_1'] * 0.5, + 'ecc': 0.0, 'metallicity': 0.02}, + 'reject_systems': None, 'output_path': 'out', 'nproc': 2, + 'kappa': 1.0, 'n_generations': 1, 'only_save_hit_tables': False, + 'rng': rng, + } + N = 20 + frame = pd.DataFrame({'bin_num': np.arange(N)}) + ckpt = STROOPWAFELCheckpoint( + config=config, mixture=None, + samples=np.zeros((N, ps.ndim)), is_hit=np.zeros(N, dtype=bool), + generation=np.zeros(N, dtype=int), gaussian_idx=np.full(N, -1), + bpp=frame, bcm=frame, initC=frame, kick_info=frame, + num_explored=N, num_hits=0, num_hits_exploratory=0, + fraction_explored=1.0, prior_fraction_rejected=0.0, + ) + return ckpt, config + + def test_restores_config_and_state(self): + ckpt, config = self._build() + sampler = AdaptiveSampler.from_checkpoint(ckpt) + self.assertIs(sampler.param_space, config['parameter_space']) + self.assertIs(sampler.derive_fn, config['derive_params']) + self.assertEqual(sampler.nproc, 2) + self.assertEqual(sampler.batch_size, 50) + self.assertEqual(sampler.num_explored, 20) # progress state restored + self.assertIs(sampler.rng, config['rng']) # RNG stream restored + + def test_overrides_win(self): + ckpt, _ = self._build() + sampler = AdaptiveSampler.from_checkpoint(ckpt, nproc=16, total_systems=5000) + self.assertEqual(sampler.nproc, 16) + self.assertEqual(sampler.total_systems, 5000) + + def test_seed_override_starts_fresh_rng(self): + ckpt, config = self._build() + sampler = AdaptiveSampler.from_checkpoint(ckpt, seed=7) + self.assertIsNot(sampler.rng, config['rng']) + + def test_unknown_override_raises(self): + ckpt, _ = self._build() + with self.assertRaises(TypeError): + AdaptiveSampler.from_checkpoint(ckpt, not_a_real_arg=1) + + def test_config_survives_dill_roundtrip(self): + # save()/from_file() rely on dill for the callables + RNG. + import dill + _, config = self._build() + restored = dill.loads(dill.dumps(config)) + out = restored['derive_params']({'mass_1': np.array([10.0])}) + self.assertAlmostEqual(out['mass_2'][0], 5.0) + np.testing.assert_array_equal( + restored['rng'].uniform(size=3), config['rng'].uniform(size=3) + ) + + # ========================================================================== # COSMICStroopOutput # ========================================================================== From 38802bed38697cc41770e2039191f1468e04d5e8 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Sun, 28 Jun 2026 23:36:03 -0400 Subject: [PATCH 37/45] minor fixes for mins --- src/cosmic/sample/stroopwafel/main.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/src/cosmic/sample/stroopwafel/main.py b/src/cosmic/sample/stroopwafel/main.py index 0fea86f19..46cfbbb23 100644 --- a/src/cosmic/sample/stroopwafel/main.py +++ b/src/cosmic/sample/stroopwafel/main.py @@ -566,7 +566,9 @@ def _adapt(self): average_density_one_dim = 1.0 / np.power(self.num_explored, 1.0 / self.param_space.ndim) self.mixture = GaussianMixture.from_hits( - hit_samples, self.param_space, average_density_one_dim, kappa=self.kappa + hit_samples, self.param_space, average_density_one_dim, kappa=self.kappa, + min_active_fraction=self.min_active_fraction, + min_entropy_change=self.min_entropy_change ) print(f" Created {self.mixture.n_components} Gaussian components") @@ -658,10 +660,12 @@ def _refine(self): # Save current state in case we need to revert saved_mixture = GaussianMixture( - self.mixture.means.copy(), - self.mixture.covariances.copy(), - self.mixture.alphas.copy(), - self.mixture.rejection_rate + means=self.mixture.means.copy(), + covariances=self.mixture.covariances.copy(), + alphas=self.mixture.alphas.copy(), + rejection_rate=self.mixture.rejection_rate, + min_active_fraction=self.mixture.min_active_fraction, + min_entropy_change=self.mixture.min_entropy_change ) should_revert = self.mixture.update_em( From 582983ab498d660b51412c9bf890269e41a18153 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 29 Jun 2026 11:42:28 -0400 Subject: [PATCH 38/45] remove output path, spell out EM stuff --- src/cosmic/output.py | 30 +++++++++++++++++-- .../sample/stroopwafel/examples/bh_star.py | 1 - .../stroopwafel/examples/example_bhbh.py | 1 - src/cosmic/sample/stroopwafel/main.py | 20 +++++-------- .../sample/stroopwafel/mixture_model.py | 21 +++++++++++-- 5 files changed, 55 insertions(+), 18 deletions(-) diff --git a/src/cosmic/output.py b/src/cosmic/output.py index 8625d5f92..c2551404c 100644 --- a/src/cosmic/output.py +++ b/src/cosmic/output.py @@ -522,8 +522,34 @@ def from_file(cls, path, label=None): label=cosmic.label if label is None else label, ) - def draw_representative_sample(self, sample_size): - raise NotImplementedError + def draw_representative_sample(self, sample_size, rng=None): + """Draw a representative sample of hits from the explored systems. + + Parameters + ---------- + sample_size : `int` + Number of hits to draw. + rng : `numpy.random.Generator`, optional + Random number generator to use for sampling. If None, a new default generator is created. + + Returns + ------- + representative_sample : `numpy.ndarray` + Array of shape (sample_size, D) containing the drawn samples in physical space. + bin_nums : `numpy.ndarray` + Array of shape (sample_size,) containing the corresponding bin numbers for the drawn samples. + """ + # restrict to hits and normalise their weights into probabilities + hit_idx = np.where(self.is_hit)[0] + probs = self.weights[hit_idx] / self.weights[hit_idx].sum() + + # draw a representative sample with replacement + rng = rng or np.random.default_rng() + chosen = rng.choice(hit_idx, size=sample_size, replace=True, p=probs) + + representative_sample = self.samples[chosen] + bin_nums = self.initC['bin_num'].iloc[chosen].values + return representative_sample, bin_nums class STROOPWAFELCheckpoint: diff --git a/src/cosmic/sample/stroopwafel/examples/bh_star.py b/src/cosmic/sample/stroopwafel/examples/bh_star.py index 94fff9707..f168cee28 100644 --- a/src/cosmic/sample/stroopwafel/examples/bh_star.py +++ b/src/cosmic/sample/stroopwafel/examples/bh_star.py @@ -201,7 +201,6 @@ def run_sampler(mc_only, seed): is_interesting=is_bh_star, derive_params=derive_params, reject_systems=default_reject, - output_path=os.path.join(args.output_dir, 'mc' if mc_only else 'sw'), nproc=args.num_cores, n_generations=args.n_generations, mc_only=mc_only, diff --git a/src/cosmic/sample/stroopwafel/examples/example_bhbh.py b/src/cosmic/sample/stroopwafel/examples/example_bhbh.py index e1d5320e5..d241ee1be 100644 --- a/src/cosmic/sample/stroopwafel/examples/example_bhbh.py +++ b/src/cosmic/sample/stroopwafel/examples/example_bhbh.py @@ -97,7 +97,6 @@ def derive_params(sampled): is_interesting=is_interesting, derive_params=derive_params, reject_systems=default_reject, - output_path=args.output_dir, nproc=args.num_cores, mc_only=args.mc_only, seed=args.seed, diff --git a/src/cosmic/sample/stroopwafel/main.py b/src/cosmic/sample/stroopwafel/main.py index 46cfbbb23..d4880500a 100644 --- a/src/cosmic/sample/stroopwafel/main.py +++ b/src/cosmic/sample/stroopwafel/main.py @@ -44,8 +44,6 @@ class AdaptiveSampler: :func:`~cosmic.sample.stroopwafel.rejection.default_reject`. By default uses :func:`~cosmic.sample.stroopwafel.rejection.default_reject`. Pass None to skip physical rejection entirely. - output_path : `str`, optional - Directory for output files, by default ``'output'`` nproc : `int`, optional Number of CPU cores for COSMIC, by default 1 kappa : `float`, optional @@ -84,7 +82,7 @@ class AdaptiveSampler: def __init__(self, parameter_space, total_systems, batch_size, BSEDict, is_interesting, derive_params=None, reject_systems="default", - output_path='output', nproc=1, kappa=1.0, + nproc=1, 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): @@ -95,7 +93,6 @@ def __init__(self, parameter_space, total_systems, batch_size, BSEDict, self.derive_fn = derive_params self.reject_fn = reject_systems if reject_systems != "default" else default_reject self.is_interesting_fn = is_interesting - self.output_path = output_path self.nproc = nproc self.kappa = kappa self.n_generations = n_generations @@ -140,8 +137,6 @@ def run(self): Container holding all samples, weights, COSMIC output tables, and associated metadata. """ - os.makedirs(self.output_path, exist_ok=True) - self._explore() if not self.mc_only and self.num_hits > 0: @@ -172,7 +167,6 @@ def run_exploration(self): :meth:`run_refinement` (or :meth:`from_checkpoint`) to continue on a different node or job. """ - os.makedirs(self.output_path, exist_ok=True) self._explore() if not self.mc_only and self.num_hits > 0: self._adapt() @@ -201,7 +195,7 @@ def run_refinement(self): #: Constructor arguments that may be overridden in :meth:`from_checkpoint`. _OVERRIDABLE = frozenset({ 'parameter_space', 'total_systems', 'batch_size', 'BSEDict', - 'is_interesting', 'derive_params', 'reject_systems', 'output_path', + 'is_interesting', 'derive_params', 'reject_systems', 'nproc', 'kappa', 'n_generations', 'only_save_hit_tables', 'seed', }) @@ -235,7 +229,7 @@ def from_checkpoint(cls, checkpoint, **overrides): Any :class:`AdaptiveSampler` constructor argument (``parameter_space``, ``total_systems``, ``batch_size``, ``BSEDict``, ``is_interesting``, ``derive_params``, - ``reject_systems``, ``output_path``, ``nproc``, ``kappa``, + ``reject_systems``, ``nproc``, ``kappa``, ``n_generations``, ``only_save_hit_tables``, ``seed``). Passing ``seed`` starts a fresh RNG instead of restoring the stored state. @@ -285,7 +279,6 @@ def _make_checkpoint(self): 'is_interesting': self.is_interesting_fn, 'derive_params': self.derive_fn, 'reject_systems': self.reject_fn, - 'output_path': self.output_path, 'nproc': self.nproc, 'kappa': self.kappa, 'n_generations': self.n_generations, @@ -652,7 +645,9 @@ def _refine(self): gen_finished += n_take self._print_progress() - # EM update (if not the last generation) + # Expectation-maximization (EM) update: re-fit the mixture to this + # generation's hits before the next one (skipped after the final + # generation, since there is no subsequent generation to use it). if gen < self.n_generations - 1 and len(gen_samples_list) > 0: gen_samples = np.vstack(gen_samples_list) gen_is_hit = np.concatenate(gen_is_hit_list) @@ -676,7 +671,8 @@ def _refine(self): if should_revert: self.mixture = saved_mixture - print(" EM update reverted (insufficient entropy change)") + print(" Expectation-maximization (EM) update reverted " + "(insufficient entropy change)") n_refined = self.total_systems - self.num_explored if n_refined > 0: diff --git a/src/cosmic/sample/stroopwafel/mixture_model.py b/src/cosmic/sample/stroopwafel/mixture_model.py index 7b2ebc93e..16d6c34b2 100644 --- a/src/cosmic/sample/stroopwafel/mixture_model.py +++ b/src/cosmic/sample/stroopwafel/mixture_model.py @@ -1,7 +1,7 @@ """Vectorized Gaussian Mixture Model for adaptive importance sampling. Stores all mixture components as numpy arrays and provides vectorized -sampling, PDF evaluation, and EM updates. +sampling, PDF evaluation, and expectation-maximization (EM) updates. """ import numpy as np from scipy.stats import multivariate_normal, entropy as scipy_entropy @@ -228,7 +228,24 @@ def compute_rejection_rate(self, param_space, reject_mask_fn, def update_em(self, samples, is_hit, prior_probs, prior_fraction_rejected, tolerance=1e-10, entropies=None): - """Perform one EM-like update of the mixture parameters. + """Perform one (importance-weighted) expectation-maximization (EM) update. + + EM is the standard algorithm for fitting a mixture model. It alternates + an **E-step** -- computing each component's *responsibility* for every + sample (the posterior probability that the sample was drawn from that + component) -- with an **M-step** that re-estimates every component's + weight, mean, and covariance as the responsibility-weighted moments of + the samples. Here the samples additionally carry importance weights + (``prior_probs * is_hit / q``), so productive components (those near many + hits) gain weight and recentre on where the hits actually are; components + whose weight falls below ``tolerance`` are dropped. + + The update is only worth keeping if it improves the proposal. This is + measured by the normalised effective sample size ``exp(H(w)) / N`` (with + ``H`` the Shannon entropy of the normalised weights): if it fails to + increase by at least ``min_entropy_change`` over the previous generation, + the method signals a revert (see Returns) -- STROOPWAFEL's indication + that the mixture has stopped improving. Parameters ---------- From 5d5f3d86533d6959133b40b1ad4699d18ee3578d Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 29 Jun 2026 11:44:22 -0400 Subject: [PATCH 39/45] review docs --- docs/pages/tutorials/adaptive/basics.rst | 2 - .../tutorials/adaptive/distributions.rst | 47 +++++++++---------- docs/pages/tutorials/adaptive/outputs.rst | 42 ++++++++--------- 3 files changed, 42 insertions(+), 49 deletions(-) diff --git a/docs/pages/tutorials/adaptive/basics.rst b/docs/pages/tutorials/adaptive/basics.rst index 6c77a3236..e168895ec 100644 --- a/docs/pages/tutorials/adaptive/basics.rst +++ b/docs/pages/tutorials/adaptive/basics.rst @@ -291,7 +291,6 @@ And then it's just a matter of setting it going! is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), derive_params=derive_params, reject_systems="default", - output_path='output/bhbh', nproc=4, n_generations=1, seed=42, @@ -321,7 +320,6 @@ Now let's repeat that whole scenario, but instead of BH + BH binaries we want to is_interesting=bh_star_100myr, # we defined this earlier derive_params=derive_params, # reuse from BHBH example reject_systems="default", - output_path='output/bh_star', nproc=4, n_generations=1, seed=42, diff --git a/docs/pages/tutorials/adaptive/distributions.rst b/docs/pages/tutorials/adaptive/distributions.rst index d28a42ec1..4c3d5ea24 100644 --- a/docs/pages/tutorials/adaptive/distributions.rst +++ b/docs/pages/tutorials/adaptive/distributions.rst @@ -6,18 +6,16 @@ Distributions and custom priors This tutorial assumes that you've already gone through :ref:`adaptive_basics`. -Every :class:`~cosmic.sample.stroopwafel.Parameter` is given a *distribution*, which does -three jobs: it draws samples, it evaluates the prior probability density, and it defines -the adaptive-sampling kernel width used during refinement. ``COSMIC`` provides a small set -of built-in distributions covering the usual initial-condition choices, and a simple, -composable interface for defining your own. +When using adaptive sampling you need to define a distribution to use for each :class:`~cosmic.sample.stroopwafel.Parameter` that you sample. +The distribution performs three operations: it draws samples, it evaluates the prior probability density, and it defines +the adaptive-sampling kernel width used during refinement. +In this tutorial we'll cover the built-in distributions and show how to define your own custom distributions and transforms. Built-in distributions ======================= -Pass any of the following names as the ``dist`` argument to a -:class:`~cosmic.sample.stroopwafel.Parameter`: +You can select any of the following distributions as the ``dist`` argument to a :class:`~cosmic.sample.stroopwafel.Parameter`: .. list-table:: :header-rows: 1 @@ -53,15 +51,13 @@ Pass any of the following names as the ``dist`` argument to a - log-normal - Natal-kick magnitude, :math:`\ln v \sim \mathcal{N}(5.67, 0.59)`. -The full registry is available programmatically as -:data:`cosmic.sample.stroopwafel.distributions.DISTRIBUTIONS`. +If you ever want to access this, you can get the full list of registered distributions with :data:`cosmic.sample.stroopwafel.distributions.DISTRIBUTIONS`. How distributions are built =========================== -Internally each distribution is a **base distribution** composed with a **coordinate -transform**: +Under the hood, we set each distribution up as a combination of a base distribution and a coordinate transform. This allows you to mix and match options: * the base distribution (:class:`~cosmic.sample.stroopwafel.distributions.Uniform`, :class:`~cosmic.sample.stroopwafel.distributions.PowerLaw`, @@ -82,33 +78,31 @@ You can build the same objects yourself: Uniform, PowerLaw, BrokenPowerLaw, TruncatedNormal, Log10, ) - Uniform(transform=Log10()) # equivalent to 'flat_in_log' - PowerLaw(-0.55, transform=Log10()) # equivalent to 'sana' + Uniform(transform=Log10()) # equivalent to 'flat_in_log' + PowerLaw(-0.55, transform=Log10()) # equivalent to 'sana' BrokenPowerLaw(breaks=[0.5], alphas=[-1.3, -2.3]) # equivalent to 'kroupa' Bounds are always given to a :class:`~cosmic.sample.stroopwafel.Parameter` in **physical** -space; the transform converts them into sampling space automatically (and round-trips -samples back to physical space before they are evolved by ``COSMIC``). +space; the transform converts them into sampling space automatically. Defining your own distribution ============================== -There are three ways to use a custom distribution, in increasing order of effort. +Now let's say that you want to define your own distribution. You can do this in three ways - let's take a look at them in order of increasing complexity. 1. Pass a distribution instance directly ---------------------------------------- -The quickest option is to tweak one of the built-in base distributions and hand the -instance straight to a :class:`~cosmic.sample.stroopwafel.Parameter` via ``dist``. No -registration required: +The quickest option is to tweak one of the built-in base distributions like we did above and hand the +instance straight to a :class:`~cosmic.sample.stroopwafel.Parameter` via ``dist``. .. code-block:: python from cosmic.sample.stroopwafel import Parameter from cosmic.sample.stroopwafel.distributions import PowerLaw - # A steeper-than-Kroupa IMF for the primary mass + # a steeper-than-Kroupa IMF for the primary mass Parameter('mass_1', 5.0, 150.0, dist=PowerLaw(-2.7)) @@ -116,8 +110,8 @@ registration required: -------------------------------- If you want to reuse a distribution across several parameter spaces — or simply refer to it -by a memorable name — register it once with -:func:`~cosmic.sample.stroopwafel.distributions.register`: +by a memorable name — you can register it once with +:func:`~cosmic.sample.stroopwafel.distributions.register`. This then allows you to refer to it by name in any :class:`~cosmic.sample.stroopwafel.Parameter`: .. code-block:: python @@ -133,7 +127,7 @@ by a memorable name — register it once with 3. Write a new distribution class --------------------------------- -For a genuinely new functional form, subclass +For a genuinely new functional form, you'll need to create a new class that subclasses off :class:`~cosmic.sample.stroopwafel.distributions.Distribution` and implement two methods: ``sample(n, lo, hi, rng)`` @@ -171,7 +165,7 @@ implements a truncated exponential distribution: Parameter('some_param', 0.0, 10.0, dist=Exponential(scale=2.0)) -That is all that is required — your distribution now works everywhere the built-ins do, and +And this class defines everything we need, our distribution now works everywhere the built-ins do, and can be combined with any transform (``dist=Exponential(2.0, transform=Log10())``). .. note:: @@ -208,3 +202,8 @@ transforms are :class:`~cosmic.sample.stroopwafel.distributions.Identity`, def to_physical(self, values): return values ** 2 + +And that's everything you need to know about distributions and transforms in ``COSMIC``'s implementation of '``STROOPWAFEL``. +You can now define your own custom priors and use them in your adaptive sampling runs. + +Next, we'll look at how to analyse the outputs of an adaptive sampling run in :ref:`adaptive_outputs`. \ No newline at end of file diff --git a/docs/pages/tutorials/adaptive/outputs.rst b/docs/pages/tutorials/adaptive/outputs.rst index 22cc33a61..42eea549e 100644 --- a/docs/pages/tutorials/adaptive/outputs.rst +++ b/docs/pages/tutorials/adaptive/outputs.rst @@ -6,10 +6,14 @@ Handling outputs from adaptive sampling This tutorial assumes that you've already gone through :ref:`adaptive_basics`. +In this tutorial we're going to cover how to read in and interpret the outputs from an adaptive sampling run. +We'll also cover how to draw a representative sample from your simulation, and how to use the weights that are generated during the adaptive sampling process. + Reading your results from a file ================================ -After you've finished running your adaptive sampling simulation, you will now have some results stored as :class:`~cosmic.output.COSMICStroopOutput` object. If you saved these results to a file, then you can reload them by running +After you've finished running your adaptive sampling simulation, you will now have some results stored as :class:`~cosmic.output.COSMICStroopOutput` object. +If you saved these results to a file, then you can reload them by running .. code-block:: python @@ -20,7 +24,8 @@ After you've finished running your adaptive sampling simulation, you will now ha Understanding your outputs ========================== -The :class:`~cosmic.output.COSMICStroopOutput` class stores all of the information you need to analyse your simulation. Let's step through some of the different attributes that you will need, and assume for the purposes of this guide that you have ``N`` samples, in ``D`` dimensions, with ``H`` hits. +The :class:`~cosmic.output.COSMICStroopOutput` class stores all of the information you need to analyse your simulation. +Let's step through some of the different attributes that you will need, and assume for the purposes of this guide that you have ``N`` samples, in ``D`` dimensions, with ``H`` hits. Sample information ------------------ @@ -34,7 +39,8 @@ Each :class:`~cosmic.output.COSMICStroopOutput` object contains a full record of Hit details ----------- -For the actual hits (i.e. the samples that you most care about), this class also stores the full evolution history. In particular, ``bpp``, ``bcm``, ``initC``, and ``kick_info`` all contain the usual ``COSMIC`` evolution tables (see :ref:`evolve_single` if you're not familiar). +For the actual hits (i.e. the samples that you most care about), this class also stores the full evolution history. +In particular, ``bpp``, ``bcm``, ``initC``, and ``kick_info`` all contain the usual ``COSMIC`` evolution tables (see :ref:`evolve_single` if you're not familiar). .. tip:: @@ -44,7 +50,8 @@ For the actual hits (i.e. the samples that you most care about), this class also just_hits = results[results.is_hit] - which masks the class just like you would with a :class:`~cosmic.output.COSMICOutput`. Be aware you likely still need the full population for access to the weights (we'll cover weights below). + which masks the class just like you would with a :class:`~cosmic.output.COSMICOutput` (see :ref:`analysis_interface` if you're not familiar). + Be aware you likely still need the full population for access to the weights (we'll cover weights below). General metadata ---------------- @@ -113,12 +120,11 @@ Drawing a representative sample from your simulation Applying weights at plot time is the right approach for visualising distributions, but sometimes you want an actual *set of systems* that is representative of the underlying -population — for example to pass a fixed number of binaries into a downstream calculation. +population. For example, you may want a fixed number of binaries for further analysis that don't require any weights. Because the simulation deliberately oversamples the rare region, you cannot take the hits at face value; you need to resample them in proportion to their weights. -This is a standard weighted bootstrap: draw indices from the hit population with probability -proportional to their weights, with replacement. +We provide a convenience method for this in :class:`~cosmic.output.COSMICStroopOutput`, but the underlying procedure is simple. It is a standard weighted bootstrap: draw indices from the hit population with probability proportional to their weights, with replacement. .. code-block:: python @@ -126,23 +132,13 @@ proportional to their weights, with replacement. from cosmic.output import COSMICStroopOutput results = COSMICStroopOutput.from_file("YOUR_SIMULATION.h5") + representative_sample, bin_nums = results.draw_representative_sample(n_samples=1000) - # Restrict to hits and normalise their weights into probabilities - hit_idx = np.where(results.is_hit)[0] - probs = results.weights[hit_idx] / results.weights[hit_idx].sum() - - # Draw a representative sample of, say, 5000 systems - rng = np.random.default_rng(42) - chosen = rng.choice(hit_idx, size=5000, replace=True, p=probs) - - representative = results.samples[chosen] # (5000, D), in physical space +The resulting ``representative_sample`` array provides a set of 1000 systems with their parameters drawn from the underlying population, and ``bin_nums`` provides the corresponding indices into the original hit population so that you can access the full evolution history if you need it. -The resulting ``representative`` array is distributed according to the true (prior-weighted) -population, so it can be histogrammed or analysed **without** any further weighting. Because -the draw is made with replacement, the same underlying system can appear more than once — -this is expected, and is the price of turning a weighted sample into an unweighted one. +Wrap-up +======= -.. note:: +And that's everything you need to know about working with the outputs from your adaptive sampling run. You can now read in your results, understand the weights, and draw a representative sample from your simulation. - A :meth:`~cosmic.output.COSMICStroopOutput.draw_representative_sample` convenience method - that wraps this resampling is planned; until it lands, use the weighted bootstrap above. \ No newline at end of file +Finally, we're going to look at how to checkpoint your adaptive sampling run so that you can resume it later in :ref:`adaptive_checkpoint` - see you there! From db63a55745abfd74c631bae67ff4ff36ac6acefd Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 29 Jun 2026 11:47:36 -0400 Subject: [PATCH 40/45] finish off docs --- docs/pages/tutorials/adaptive/checkpoint.rst | 35 ++++++++++--------- .../tutorials/adaptive/distributions.rst | 3 ++ 2 files changed, 22 insertions(+), 16 deletions(-) diff --git a/docs/pages/tutorials/adaptive/checkpoint.rst b/docs/pages/tutorials/adaptive/checkpoint.rst index aeb55962f..9b8a0f83e 100644 --- a/docs/pages/tutorials/adaptive/checkpoint.rst +++ b/docs/pages/tutorials/adaptive/checkpoint.rst @@ -13,13 +13,13 @@ Why split a run in two? A call to :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run` performs all three phases back to back: exploration, adaptation, and refinement. Splitting the run lets you stop -after adaptation — once the Gaussian mixture has been fitted to the exploration hits — and -resume the (usually much larger) refinement phase separately. This is useful when you want +after adaptation — once the Gaussian mixture has been fit to the exploration hits — and +resume the (usually much larger) refinement phase separately. This is useful when you want to * run exploration and refinement as **separate cluster jobs**, perhaps with different walltimes or allocations; -* fan a single adapted mixture out across **several refinement jobs** on different nodes; or +* spread a single adapted mixture out across **several refinement jobs** on different nodes; or * simply **inspect** the mixture before committing compute to refinement. Instead of :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run`, you use the two @@ -40,7 +40,6 @@ call :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run_exploration` instead from cosmic.sample.stroopwafel import AdaptiveSampler, ParameterSpace, Parameter from cosmic.sample.stroopwafel.presets import any_dco - from cosmic.sample.stroopwafel.rejection import default_reject params = ParameterSpace([ Parameter('mass_1', 5.0, 150.0, dist='kroupa'), @@ -60,8 +59,7 @@ call :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run_exploration` instead BSEDict=BSEDict, is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), derive_params=derive_params, - reject_systems=default_reject, - output_path='output/explore', + reject_systems="default", nproc=4, seed=42, ) @@ -82,8 +80,7 @@ Stage 2 — load the checkpoint and refine In a second script (or cluster job) rebuild the sampler with :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.from_checkpoint`, then call -:meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run_refinement`. The checkpoint is -**self-contained**, so this needs nothing but the file: +:meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run_refinement`. .. code-block:: python @@ -101,12 +98,12 @@ exploration), pass it as a keyword override: .. code-block:: python sampler = AdaptiveSampler.from_checkpoint( - 'checkpoint.h5', nproc=16, total_systems=1_000_000, + 'checkpoint.h5', nproc=16 ) Any :class:`~cosmic.sample.stroopwafel.AdaptiveSampler` constructor argument may be overridden this way (``parameter_space``, ``BSEDict``, ``derive_params``, -``reject_systems``, ``is_interesting``, ``batch_size``, ``output_path``, ``nproc``, +``reject_systems``, ``is_interesting``, ``batch_size``, ``nproc``, ``kappa``, ``n_generations``, ``only_save_hit_tables``, ``seed``). The ``result`` is an ordinary :class:`~cosmic.output.COSMICStroopOutput` — identical in form @@ -130,10 +127,8 @@ A checkpoint is a complete snapshot — it stores both the exploration *results* the ``derive_params`` / ``reject_systems`` / ``is_interesting`` callables, the remaining scalar settings, and the live RNG state. -The callables and parameter space are serialised with :mod:`dill` (a dependency COSMIC -already ships), so lambdas and closures — such as the ``is_interesting`` returned by -``any_dco(...)`` — round-trip correctly. Because the RNG state is stored too, refinement -continues the random stream seamlessly rather than restarting it (pass ``seed=`` to +The callables and parameter space are serialised with :mod:`dill` so it lets you use general functions. +Because the RNG state is stored too, refinement continues the random stream seamlessly rather than restarting it (pass ``seed=`` to ``from_checkpoint`` if you instead want a fresh stream). @@ -142,8 +137,16 @@ Reusing one checkpoint for several refinement jobs Because :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.from_checkpoint` reads the checkpoint without modifying it, you can launch any number of independent refinement jobs -from the same ``checkpoint.h5`` — for example with different ``seed`` values on different -nodes — and each will draw a fresh, independent refinement sample from the shared mixture. +from the same ``checkpoint.h5`` and each will draw a fresh, independent refinement sample from the shared mixture. You can also override the budget stored in the checkpoint by passing ``total_systems`` or ``n_generations`` to :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.from_checkpoint`. +.. warning:: + + You should make sure that you change the ``seed`` when you launch multiple refinement jobs from the same checkpoint, otherwise they will all draw the same random numbers and produce identical results. The simplest way to do this is to pass ``seed=None`` to :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.from_checkpoint`, which will seed the RNG from the system clock. + + +Wrap-up +======= + +And that's all on adaptive sampling folks! You should now know everything you need to run an adaptive importance sampling simulation in ``COSMIC`` - enjoy exploring those rare populations! \ No newline at end of file diff --git a/docs/pages/tutorials/adaptive/distributions.rst b/docs/pages/tutorials/adaptive/distributions.rst index 4c3d5ea24..66a5835c9 100644 --- a/docs/pages/tutorials/adaptive/distributions.rst +++ b/docs/pages/tutorials/adaptive/distributions.rst @@ -203,6 +203,9 @@ transforms are :class:`~cosmic.sample.stroopwafel.distributions.Identity`, def to_physical(self, values): return values ** 2 +Wrap-up +======= + And that's everything you need to know about distributions and transforms in ``COSMIC``'s implementation of '``STROOPWAFEL``. You can now define your own custom priors and use them in your adaptive sampling runs. From cf1cdc2e057ac72c7d91bd759c703aeed6f02b9d Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 29 Jun 2026 12:03:11 -0400 Subject: [PATCH 41/45] check over representative sample --- docs/pages/tutorials/adaptive/outputs.rst | 2 +- src/cosmic/output.py | 25 +++++--- src/cosmic/sample/stroopwafel/main.py | 3 + src/cosmic/tests/test_stroopwafel.py | 76 ++++++++++++++++++++++- 4 files changed, 96 insertions(+), 10 deletions(-) diff --git a/docs/pages/tutorials/adaptive/outputs.rst b/docs/pages/tutorials/adaptive/outputs.rst index 42eea549e..f95bd5154 100644 --- a/docs/pages/tutorials/adaptive/outputs.rst +++ b/docs/pages/tutorials/adaptive/outputs.rst @@ -134,7 +134,7 @@ We provide a convenience method for this in :class:`~cosmic.output.COSMICStroopO results = COSMICStroopOutput.from_file("YOUR_SIMULATION.h5") representative_sample, bin_nums = results.draw_representative_sample(n_samples=1000) -The resulting ``representative_sample`` array provides a set of 1000 systems with their parameters drawn from the underlying population, and ``bin_nums`` provides the corresponding indices into the original hit population so that you can access the full evolution history if you need it. +The resulting ``representative_sample`` array provides a set of 1000 systems with their parameters drawn from the underlying population, and ``bin_nums`` provides the corresponding indices into the original population so that you can access the full evolution history if you need it. Wrap-up ======= diff --git a/src/cosmic/output.py b/src/cosmic/output.py index c2551404c..f30424a26 100644 --- a/src/cosmic/output.py +++ b/src/cosmic/output.py @@ -522,22 +522,29 @@ def from_file(cls, path, label=None): label=cosmic.label if label is None else label, ) - def draw_representative_sample(self, sample_size, rng=None): + def draw_representative_sample(self, n_samples, rng=None): """Draw a representative sample of hits from the explored systems. + Performs a weighted bootstrap: hits are drawn with replacement in + proportion to their importance weights, yielding a set of systems + distributed according to the true (prior-weighted) population that + can be analysed without any further weighting. + Parameters ---------- - sample_size : `int` + n_samples : `int` Number of hits to draw. rng : `numpy.random.Generator`, optional Random number generator to use for sampling. If None, a new default generator is created. - + Returns ------- representative_sample : `numpy.ndarray` - Array of shape (sample_size, D) containing the drawn samples in physical space. + Array of shape (n_samples, D) containing the drawn samples in physical space. bin_nums : `numpy.ndarray` - Array of shape (sample_size,) containing the corresponding bin numbers for the drawn samples. + Array of shape (n_samples,) containing the corresponding bin numbers, so the + full evolution history of each drawn system can be recovered from the + ``bpp``/``bcm``/``initC``/``kick_info`` tables (e.g. ``self.initC.loc[bin_num]``). """ # restrict to hits and normalise their weights into probabilities hit_idx = np.where(self.is_hit)[0] @@ -545,10 +552,12 @@ def draw_representative_sample(self, sample_size, rng=None): # draw a representative sample with replacement rng = rng or np.random.default_rng() - chosen = rng.choice(hit_idx, size=sample_size, replace=True, p=probs) + bin_nums = rng.choice(hit_idx, size=n_samples, replace=True, p=probs) - representative_sample = self.samples[chosen] - bin_nums = self.initC['bin_num'].iloc[chosen].values + # `samples` is indexed by bin_num (samples[bin_num] is the system with that + # bin_num), so the drawn indices are themselves the bin numbers — use them to + # label the systems and to look rows up in the COSMIC tables via `.loc`. + representative_sample = self.samples[bin_nums] return representative_sample, bin_nums diff --git a/src/cosmic/sample/stroopwafel/main.py b/src/cosmic/sample/stroopwafel/main.py index d4880500a..308621d64 100644 --- a/src/cosmic/sample/stroopwafel/main.py +++ b/src/cosmic/sample/stroopwafel/main.py @@ -197,6 +197,7 @@ def run_refinement(self): 'parameter_space', 'total_systems', 'batch_size', 'BSEDict', 'is_interesting', 'derive_params', 'reject_systems', 'nproc', 'kappa', 'n_generations', 'only_save_hit_tables', 'seed', + 'min_active_fraction', 'min_entropy_change', }) @classmethod @@ -283,6 +284,8 @@ def _make_checkpoint(self): 'kappa': self.kappa, 'n_generations': self.n_generations, 'only_save_hit_tables': self.only_save_hit_tables, + 'min_active_fraction': self.min_active_fraction, + 'min_entropy_change': self.min_entropy_change, 'rng': self.rng, } diff --git a/src/cosmic/tests/test_stroopwafel.py b/src/cosmic/tests/test_stroopwafel.py index 26ccad8b2..e929813af 100644 --- a/src/cosmic/tests/test_stroopwafel.py +++ b/src/cosmic/tests/test_stroopwafel.py @@ -519,8 +519,9 @@ def _build(self): 'is_interesting': lambda bpp: (0, np.array([], dtype=int)), 'derive_params': lambda s: {'mass_2': s['mass_1'] * 0.5, 'ecc': 0.0, 'metallicity': 0.02}, - 'reject_systems': None, 'output_path': 'out', 'nproc': 2, + 'reject_systems': None, 'nproc': 2, 'kappa': 1.0, 'n_generations': 1, 'only_save_hit_tables': False, + 'min_active_fraction': 0.01, 'min_entropy_change': 0.01, 'rng': rng, } N = 20 @@ -595,5 +596,78 @@ def test_hit_rate(self): self.assertAlmostEqual(result.hit_rate, 0.1) +# ========================================================================== +# draw_representative_sample +# ========================================================================== +class TestDrawRepresentativeSample(unittest.TestCase): + + PARAM_NAMES = ['mass_1', 'porb', 'ecc', 'metallicity'] + + def _make_output(self): + """Build an output whose initC holds ONLY the hit rows, in shuffled + order (as ``only_save_hit_tables=True`` produces). This makes the + initC row *position* differ from the ``bin_num``, so the test exercises + label-based (not positional) lookup of the returned bin numbers. + """ + N = 12 + # Each system gets uniquely identifiable physical parameters. + samples = np.column_stack([ + 10.0 + np.arange(N), # mass_1 + 100.0 + np.arange(N), # porb + 0.01 * np.arange(N), # ecc + 0.001 + 0.001 * np.arange(N), # metallicity + ]) + is_hit = np.zeros(N, dtype=bool) + is_hit[[1, 3, 5, 7, 9, 11]] = True + weights = np.linspace(0.1, 2.0, N) + + hit_bins = np.where(is_hit)[0] + order = hit_bins.copy() + np.random.default_rng(0).shuffle(order) # initC rows out of bin_num order + initC = pd.DataFrame({ + 'bin_num': order, + 'mass_1': samples[order, 0], + 'porb': samples[order, 1], + 'ecc': samples[order, 2], + 'metallicity': samples[order, 3], + 'mass_2': samples[order, 0] * 0.5, # an extra IC column + }).set_index('bin_num', drop=False) + minimal = pd.DataFrame({'bin_num': order}) + + out = COSMICStroopOutput( + bpp=minimal, bcm=minimal, initC=initC, kick_info=minimal, + samples=samples, param_names=self.PARAM_NAMES, weights=weights, + is_hit=is_hit, generation=np.zeros(N, dtype=int), + gaussian_idx=np.full(N, -1, dtype=int), + num_explored=N, num_hits=len(hit_bins), fraction_explored=1.0, + ) + return out, samples + + def test_bin_nums_reference_matching_initial_conditions(self): + out, samples = self._make_output() + rep, bin_nums = out.draw_representative_sample(50, rng=np.random.default_rng(3)) + + self.assertEqual(rep.shape, (50, len(self.PARAM_NAMES))) + self.assertEqual(bin_nums.shape, (50,)) + # Only hits may be drawn. + self.assertTrue(np.all(out.is_hit[bin_nums])) + + for params, b in zip(rep, bin_nums): + # Returned parameters are the sample row for that bin_num. + np.testing.assert_array_equal(params, samples[b]) + # The initC row for that bin_num has identical initial conditions. + ic = out.initC.loc[b] + for col in self.PARAM_NAMES: + self.assertEqual(ic[col], params[self.PARAM_NAMES.index(col)]) + + def test_draws_are_importance_weighted(self): + # With all weight on a single hit, every draw should be that system. + out, _ = self._make_output() + out.weights = np.zeros(len(out.weights)) + out.weights[7] = 1.0 # bin_num 7 is a hit + _, bin_nums = out.draw_representative_sample(100, rng=np.random.default_rng(1)) + self.assertTrue(np.all(bin_nums == 7)) + + if __name__ == '__main__': unittest.main() From 90f91624255d9037b9196d08c024c84e67e54279 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 29 Jun 2026 12:25:41 -0400 Subject: [PATCH 42/45] add SSEDict to everything --- docs/pages/tutorials/adaptive/basics.rst | 11 +- docs/pages/tutorials/adaptive/checkpoint.rst | 11 +- src/cosmic/output.py | 8 +- .../sample/stroopwafel/examples/bh_bh.py | 65 ++++++++++ .../sample/stroopwafel/examples/bh_star.py | 4 + .../stroopwafel/examples/example_bhbh.py | 113 ------------------ src/cosmic/sample/stroopwafel/main.py | 32 ++++- src/cosmic/sample/stroopwafel/rejection.py | 16 ++- src/cosmic/tests/test_stroopwafel.py | 53 ++++++++ 9 files changed, 181 insertions(+), 132 deletions(-) create mode 100644 src/cosmic/sample/stroopwafel/examples/bh_bh.py delete mode 100644 src/cosmic/sample/stroopwafel/examples/example_bhbh.py diff --git a/docs/pages/tutorials/adaptive/basics.rst b/docs/pages/tutorials/adaptive/basics.rst index e168895ec..fa591b3ef 100644 --- a/docs/pages/tutorials/adaptive/basics.rst +++ b/docs/pages/tutorials/adaptive/basics.rst @@ -236,7 +236,7 @@ You can also write a fully custom hit function. For example, to find BH + stell Running the sampler =================== -Now we can put it all together! You can run the sampler with the main :class:`~cosmic.sample.stroopwafel.AdaptiveSampler` class. The most important arguments are the parameter space, the total number of systems to evolve, the batch size, the BSE physics settings, the hit function, the ``derive_params`` function (if needed), and the rejection function. See the API documentation (:class:`~cosmic.sample.stroopwafel.AdaptiveSampler`) for a full list of options. +Now we can put it all together! You can run the sampler with the main :class:`~cosmic.sample.stroopwafel.AdaptiveSampler` class. The most important arguments are the parameter space, the total number of systems to evolve, the batch size, the BSE and SSE physics settings, the hit function, the ``derive_params`` function (if needed), and the rejection function. See the API documentation (:class:`~cosmic.sample.stroopwafel.AdaptiveSampler`) for a full list of options. Let's try this out with a few examples. @@ -258,6 +258,13 @@ First we can import the necessary parts from the ``cosmic.sample.stroopwafel`` m .. include:: ../../../_generated/default_bsedict.rst +COSMIC v4+ also expects an ``SSEDict`` of single stellar evolution settings (which selects +the stellar engine). We'll use the default ``sse`` engine here, you can swap in METISSE too if you like! + +.. code-block:: python + + SSEDict = {'stellar_engine': 'sse'} + Then we can define a simple parameter space, where we avoid sampling low-mass primaries since we know they cannot produce a BH. @@ -288,6 +295,7 @@ And then it's just a matter of setting it going! total_systems=50_000, # adjust this for more samples batch_size=500, # adjust this to sample more or fewer systems per call to COSMIC BSEDict=BSEDict, + SSEDict=SSEDict, is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), derive_params=derive_params, reject_systems="default", @@ -317,6 +325,7 @@ Now let's repeat that whole scenario, but instead of BH + BH binaries we want to total_systems=50_000, batch_size=500, BSEDict=BSEDict, # reuse from BHBH example + SSEDict=SSEDict, # reuse from BHBH example is_interesting=bh_star_100myr, # we defined this earlier derive_params=derive_params, # reuse from BHBH example reject_systems="default", diff --git a/docs/pages/tutorials/adaptive/checkpoint.rst b/docs/pages/tutorials/adaptive/checkpoint.rst index 9b8a0f83e..6f989fe52 100644 --- a/docs/pages/tutorials/adaptive/checkpoint.rst +++ b/docs/pages/tutorials/adaptive/checkpoint.rst @@ -57,6 +57,7 @@ call :meth:`~cosmic.sample.stroopwafel.AdaptiveSampler.run_exploration` instead total_systems=500_000, batch_size=1000, BSEDict=BSEDict, + SSEDict={'stellar_engine': 'sse'}, is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), derive_params=derive_params, reject_systems="default", @@ -90,8 +91,8 @@ In a second script (or cluster job) rebuild the sampler with result = sampler.run_refinement() result.save('result.h5') -You do not need to re-import or re-specify the parameter space, ``BSEDict``, or any of the -callables — they were all saved into the checkpoint. If you *want* to change something for +You do not need to re-import or re-specify the parameter space, ``BSEDict``, ``SSEDict``, or +any of the callables — they were all saved into the checkpoint. If you *want* to change something for the refinement phase (a common one is running on more cores, or with a larger budget than exploration), pass it as a keyword override: @@ -102,7 +103,7 @@ exploration), pass it as a keyword override: ) Any :class:`~cosmic.sample.stroopwafel.AdaptiveSampler` constructor argument may be -overridden this way (``parameter_space``, ``BSEDict``, ``derive_params``, +overridden this way (``parameter_space``, ``BSEDict``, ``SSEDict``, ``derive_params``, ``reject_systems``, ``is_interesting``, ``batch_size``, ``nproc``, ``kappa``, ``n_generations``, ``only_save_hit_tables``, ``seed``). @@ -124,8 +125,8 @@ A checkpoint is a complete snapshot — it stores both the exploration *results* * the scalar counters needed to compute unbiased weights later (``num_explored``, ``fraction_explored``, ``prior_fraction_rejected``, ...); and * the full configuration needed to rebuild the sampler: the parameter space, ``BSEDict``, - the ``derive_params`` / ``reject_systems`` / ``is_interesting`` callables, the remaining - scalar settings, and the live RNG state. + ``SSEDict``, the ``derive_params`` / ``reject_systems`` / ``is_interesting`` callables, the + remaining scalar settings, and the live RNG state. The callables and parameter space are serialised with :mod:`dill` so it lets you use general functions. Because the RNG state is stored too, refinement continues the random stream seamlessly rather than restarting it (pass ``seed=`` to diff --git a/src/cosmic/output.py b/src/cosmic/output.py index f30424a26..1fa913b26 100644 --- a/src/cosmic/output.py +++ b/src/cosmic/output.py @@ -588,10 +588,10 @@ class STROOPWAFELCheckpoint: config : `dict` Everything needed to reconstruct the :class:`AdaptiveSampler` for the refinement phase: the constructor keyword arguments (``parameter_space``, - ``total_systems``, ``batch_size``, ``BSEDict``, ``is_interesting``, - ``derive_params``, ``reject_systems``, ``output_path``, ``nproc``, - ``kappa``, ``n_generations``, ``only_save_hit_tables``) plus the live - ``rng``. + ``total_systems``, ``batch_size``, ``BSEDict``, ``SSEDict``, + ``is_interesting``, ``derive_params``, ``reject_systems``, ``nproc``, + ``kappa``, ``n_generations``, ``only_save_hit_tables``, + ``min_active_fraction``, ``min_entropy_change``) plus the live ``rng``. mixture : `GaussianMixture` or None Gaussian mixture fitted to exploration hits. ``None`` if no hits were found or adaptation has not been run yet. diff --git a/src/cosmic/sample/stroopwafel/examples/bh_bh.py b/src/cosmic/sample/stroopwafel/examples/bh_bh.py new file mode 100644 index 000000000..394fa61cb --- /dev/null +++ b/src/cosmic/sample/stroopwafel/examples/bh_bh.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python +import time +import argparse + +from cosmic.sample.stroopwafel import AdaptiveSampler, ParameterSpace, Parameter +from cosmic.sample.stroopwafel.presets import merging_dco + +import sys +sys.path.append("../../../../../docs") +import generate_default_bsedict +BSEDict = generate_default_bsedict.get_default_BSE_settings(to_python=True) + + +parser = argparse.ArgumentParser() +parser.add_argument('--num_systems', type=int, default=100_000) +parser.add_argument('--num_cores', type=int, default=1) +parser.add_argument('--num_per_core', type=int, default=1000) +parser.add_argument('--mc_only', type=bool, default=False) +parser.add_argument('--seed', type=int, default=None) +args = parser.parse_args() + +# COSMIC v4+ single stellar evolution settings (sse engine; swap for METISSE) +SSEDict = {'stellar_engine': 'sse'} + +# ------------------------------------------------------------------ +# Define parameter space +# ------------------------------------------------------------------ +params = ParameterSpace([ + Parameter('mass_1', 5.0, 150.0, dist='kroupa'), + Parameter('q', 0.0, 1.0, dist='uniform'), + Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), + Parameter('ecc', 1e-9, 0.99999999, dist='sana_ecc'), + Parameter('metallicity', 0.0001, 0.03, dist='flat_in_log'), + Parameter('natal_kick_1', 0.0, 5000.0, dist='disberg') +]) + +def derive_params(sampled): + """Return the one required parameter (mass_2) not sampled directly.""" + return {'mass_2': sampled['mass_1'] * sampled['q']} + +is_interesting = merging_dco(kstar_1=[14], kstar_2=[14], max_merge_time=13.7) + +if __name__ == '__main__': + start = time.time() + + sw = AdaptiveSampler( + parameter_space=params, + total_systems=args.num_systems, + batch_size=args.num_per_core, + BSEDict=BSEDict, + SSEDict=SSEDict, + is_interesting=is_interesting, + derive_params=derive_params, + reject_systems="default", + nproc=args.num_cores, + mc_only=args.mc_only, + seed=args.seed, + ) + + result = sw.run() + result.save("output/bh_bh/result.h5") + + elapsed = time.time() - start + print(f"\nTotal time: {elapsed:.1f}s") + print(f"Results saved to output/bh_bh/result.h5") diff --git a/src/cosmic/sample/stroopwafel/examples/bh_star.py b/src/cosmic/sample/stroopwafel/examples/bh_star.py index f168cee28..78182e243 100644 --- a/src/cosmic/sample/stroopwafel/examples/bh_star.py +++ b/src/cosmic/sample/stroopwafel/examples/bh_star.py @@ -87,6 +87,9 @@ "fryer_fmix": 0.5, "fryer_mcrit_nsbh": 5.0, "smt_periastron_check": 0 } +# COSMIC v4+ single stellar evolution settings (sse engine; swap for METISSE) +SSEDict = {'stellar_engine': 'sse'} + # ------------------------------------------------------------------ # Parameter space # @@ -198,6 +201,7 @@ def run_sampler(mc_only, seed): total_systems=args.num_systems, batch_size=args.batch_size, BSEDict=BSEDict, + SSEDict=SSEDict, is_interesting=is_bh_star, derive_params=derive_params, reject_systems=default_reject, diff --git a/src/cosmic/sample/stroopwafel/examples/example_bhbh.py b/src/cosmic/sample/stroopwafel/examples/example_bhbh.py deleted file mode 100644 index d241ee1be..000000000 --- a/src/cosmic/sample/stroopwafel/examples/example_bhbh.py +++ /dev/null @@ -1,113 +0,0 @@ -#!/usr/bin/env python -"""Example: Find merging BH-BH binaries using the new vectorized STROOPWAFEL API. - -Equivalent to the old tests/BHBH_testing.py but using the new API. -""" -import os -import sys -import time -import argparse -import numpy as np - -sys.path.append(os.path.join(os.path.dirname(__file__), '..')) # adds COSMIC-stroopwafel/ - -from stroopwafel import AdaptiveSampler, ParameterSpace, Parameter -from stroopwafel.presets import merging_dco -from stroopwafel.rejection import default_reject -from stroopwafel import io as swio - -# ------------------------------------------------------------------ -# Parse arguments -# ------------------------------------------------------------------ -parser = argparse.ArgumentParser() -parser.add_argument('--num_systems', type=int, default=10000) -parser.add_argument('--num_cores', type=int, default=1) -parser.add_argument('--num_per_core', type=int, default=100) -parser.add_argument('--mc_only', type=bool, default=False) -parser.add_argument('--output_dir', default='output/BHBH_new') -parser.add_argument('--model', default='fiducial') -parser.add_argument('--seed', type=int, default=None) -args = parser.parse_args() - -# ------------------------------------------------------------------ -# Physics models (same as old code) -# ------------------------------------------------------------------ -fiducial = { - 'xi': 1.0, 'bhflag': 1, 'neta': 0.5, 'windflag': 3, 'wdflag': 1, - 'alpha1': 1.0, 'pts1': 0.001, 'pts3': 0.02, 'pts2': 0.01, - 'epsnov': 0.001, 'hewind': 0.5, 'ck': 1000, 'bwind': 0.0, - 'lambdaf': 0.0, 'mxns': 3.0, 'beta': -1.0, 'tflag': 1, 'acc2': 1.5, - 'grflag': 1, 'remnantflag': 4, 'ceflag': 0, 'eddfac': 1.0, - 'ifflag': 0, 'bconst': 3000, 'sigma': 265.0, 'gamma': -2.0, - 'pisn': 45.0, - 'natal_kick_array': [[-100.0, -100.0, -100.0, -100.0, 0.0], - [-100.0, -100.0, -100.0, -100.0, 0.0]], - 'bhsigmafrac': 1.0, 'polar_kick_angle': 90, - 'qcrit_array': [0.0] * 16, - 'cekickflag': 2, 'cehestarflag': 0, 'cemergeflag': 0, - 'ecsn': 2.5, 'ecsn_mlow': 1.8, 'aic': 1, 'ussn': 0, - 'sigmadiv': -20.0, 'qcflag': 5, 'eddlimflag': 0, - 'fprimc_array': [2.0 / 21.0] * 16, - 'bhspinflag': 0, 'bhspinmag': 0.0, 'rejuv_fac': 1.0, - 'rejuvflag': 0, 'htpmb': 1, 'ST_cr': 1, 'ST_tide': 1, - 'bdecayfac': 1, 'rembar_massloss': 0.5, 'kickflag': 5, - 'zsun': 0.014, 'bhms_coll_flag': 0, 'don_lim': -1, - 'acc_lim': -1, 'rtmsflag': 0, 'wd_mass_lim': 1, -} - -BSEDict = fiducial # extend with model variants as needed - -# ------------------------------------------------------------------ -# Define parameter space -# ------------------------------------------------------------------ -params = ParameterSpace([ - Parameter('mass_1', 5.0, 150.0, dist='kroupa'), - Parameter('q', 0.0, 1.0, dist='uniform'), - Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), # ~1.4 d to ~316 000 d - Parameter('ecc', 1e-9, 0.99999999, dist='sana_ecc'), - Parameter('metallicity', 0.0001, 0.03, dist='flat_in_log'), -]) - -# ------------------------------------------------------------------ -# Provide binary parameters not sampled directly (vectorized) -# -# A binary is defined by {mass_1, mass_2, porb, ecc, metallicity}. All but -# mass_2 are sampled, so derive mass_2 from the sampled mass ratio. -# ------------------------------------------------------------------ -def derive_params(sampled): - """Return the one required parameter (mass_2) not sampled directly.""" - return {'mass_2': sampled['mass_1'] * sampled['q']} - -# ------------------------------------------------------------------ -# Hit definition: merging BH-BH binaries within Hubble time -# ------------------------------------------------------------------ -is_interesting = merging_dco(kstar_1=[14], kstar_2=[14], max_merge_time=13.7) - -# ------------------------------------------------------------------ -# Run -# ------------------------------------------------------------------ -if __name__ == '__main__': - start = time.time() - - sw = AdaptiveSampler( - parameter_space=params, - total_systems=args.num_systems, - batch_size=args.num_per_core, - BSEDict=BSEDict, - is_interesting=is_interesting, - derive_params=derive_params, - reject_systems=default_reject, - nproc=args.num_cores, - mc_only=args.mc_only, - seed=args.seed, - ) - - result = sw.run() - - # Save results - output_file = os.path.join(args.output_dir, 'stroopwafel_result.h5') - swio.save_result(output_file, result) - - elapsed = time.time() - start - print(f"\nTotal time: {elapsed:.1f}s") - print(f"Results saved to {output_file}") diff --git a/src/cosmic/sample/stroopwafel/main.py b/src/cosmic/sample/stroopwafel/main.py index 308621d64..6e6dc0433 100644 --- a/src/cosmic/sample/stroopwafel/main.py +++ b/src/cosmic/sample/stroopwafel/main.py @@ -1,4 +1,5 @@ import os +from functools import partial import numpy as np import pandas as pd from scipy.stats import multivariate_normal @@ -27,6 +28,13 @@ class AdaptiveSampler: is_interesting : `callable` Function with signature ``(bpp) -> (n_hits, hit_bin_nums)`` identifying systems of interest from COSMIC output. + SSEDict : `dict`, 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 :func:`~cosmic.sample.stroopwafel.rejection.default_reject` so the + ZAMS radii it computes via ``set_reff`` use the same engine. By + default None (COSMIC's ``sse`` engine). derive_params : `callable`, optional Function ``(sampled) -> dict`` that supplies any binary parameters not drawn from ``parameter_space``. ``sampled`` is a dict mapping @@ -81,8 +89,8 @@ class AdaptiveSampler: REQUIRED_PARAMS = ('mass_1', 'mass_2', 'porb', 'ecc', 'metallicity') def __init__(self, parameter_space, total_systems, batch_size, BSEDict, - is_interesting, derive_params=None, reject_systems="default", - nproc=1, kappa=1.0, + is_interesting, SSEDict=None, derive_params=None, + reject_systems="default", nproc=1, 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): @@ -90,8 +98,18 @@ def __init__(self, parameter_space, total_systems, batch_size, BSEDict, self.total_systems = total_systems self.batch_size = batch_size self.bse_dict = BSEDict + self.sse_dict = SSEDict self.derive_fn = derive_params - self.reject_fn = reject_systems if reject_systems != "default" else default_reject + # Keep the original spec so a checkpoint can rebind the default to the + # (possibly overridden) SSEDict on reload. + self._reject_spec = reject_systems + # The default rejection uses set_reff, which depends on the stellar + # engine, so bind this run's SSEDict into it -- whether requested via + # the "default" sentinel or by passing default_reject explicitly. + if reject_systems == "default" or reject_systems is default_reject: + self.reject_fn = partial(default_reject, SSEDict=SSEDict) + else: + self.reject_fn = reject_systems self.is_interesting_fn = is_interesting self.nproc = nproc self.kappa = kappa @@ -194,7 +212,7 @@ def run_refinement(self): #: Constructor arguments that may be overridden in :meth:`from_checkpoint`. _OVERRIDABLE = frozenset({ - 'parameter_space', 'total_systems', 'batch_size', 'BSEDict', + 'parameter_space', 'total_systems', 'batch_size', 'BSEDict', 'SSEDict', 'is_interesting', 'derive_params', 'reject_systems', 'nproc', 'kappa', 'n_generations', 'only_save_hit_tables', 'seed', 'min_active_fraction', 'min_entropy_change', @@ -277,9 +295,12 @@ def _make_checkpoint(self): 'total_systems': self.total_systems, 'batch_size': self.batch_size, 'BSEDict': self.bse_dict, + 'SSEDict': self.sse_dict, 'is_interesting': self.is_interesting_fn, 'derive_params': self.derive_fn, - 'reject_systems': self.reject_fn, + # store the original spec ('default'/callable/None) so reload can + # rebind 'default' to the (possibly overridden) SSEDict. + 'reject_systems': self._reject_spec, 'nproc': self.nproc, 'kappa': self.kappa, 'n_generations': self.n_generations, @@ -874,6 +895,7 @@ def _evolve_batch(self, binary_params): bpp, bcm, initC, kick_info = Evolve.evolve( initialbinarytable=batch_initial, BSEDict=bse_dict_run, + SSEDict=self.sse_dict, nproc=self.nproc, ) diff --git a/src/cosmic/sample/stroopwafel/rejection.py b/src/cosmic/sample/stroopwafel/rejection.py index 4152d5c35..be6d91804 100644 --- a/src/cosmic/sample/stroopwafel/rejection.py +++ b/src/cosmic/sample/stroopwafel/rejection.py @@ -8,7 +8,7 @@ from cosmic.utils import calc_Roche_radius, a_from_p -def default_reject(binary_params, min_secondary_mass=0.08): +def default_reject(binary_params, SSEDict=None, min_secondary_mass=0.08): """Default rejection function for DCO progenitor systems. Rejects systems where the secondary mass is below the minimum, the @@ -23,6 +23,14 @@ def default_reject(binary_params, min_secondary_mass=0.08): Assembled binary parameters with keys ``'mass_1'``, ``'mass_2'`` (solar masses), ``'porb'`` (days), ``'ecc'``, and ``'metallicity'``, each an (N,) array. + SSEDict : `dict`, optional + COSMIC single stellar evolution settings. The ZAMS radii are + obtained from ``Sample.set_reff``, which depends on the stellar + engine (e.g. ``sse`` vs ``METISSE``); passing the same ``SSEDict`` + used for evolution ensures the rejection radii are computed + consistently. By default None (COSMIC's ``sse`` engine). + :class:`~cosmic.sample.stroopwafel.main.AdaptiveSampler` wires its + own ``SSEDict`` into this argument automatically. min_secondary_mass : `float`, optional Minimum allowed secondary mass in solar masses, by default 0.08 @@ -40,10 +48,10 @@ def default_reject(binary_params, min_secondary_mass=0.08): # compute separation from periods and masses separation = a_from_p(p=porb, m1=mass_1, m2=mass_2) - # get stellar radii at ZAMS + # get stellar radii at ZAMS (depends on the stellar engine in SSEDict) sampler = Sample() - radius_1 = sampler.set_reff(mass=mass_1, metallicity=metallicity) - radius_2 = sampler.set_reff(mass=mass_2, metallicity=metallicity) + radius_1 = sampler.set_reff(mass=mass_1, metallicity=metallicity, SSEDict=SSEDict) + radius_2 = sampler.set_reff(mass=mass_2, metallicity=metallicity, SSEDict=SSEDict) # roche lobe radii at periastron peri_sep = separation * (1 - ecc) diff --git a/src/cosmic/tests/test_stroopwafel.py b/src/cosmic/tests/test_stroopwafel.py index e929813af..8697c236b 100644 --- a/src/cosmic/tests/test_stroopwafel.py +++ b/src/cosmic/tests/test_stroopwafel.py @@ -466,6 +466,58 @@ def test_derive_params_wrong_length_raises(self): sampler_factory() +# ========================================================================== +# SSEDict wiring +# ========================================================================== +class TestSSEDict(unittest.TestCase): + + SSE = {'stellar_engine': 'sse'} + + def _make(self, reject_systems): + params = ParameterSpace([ + Parameter('mass_1', 5.0, 150.0, dist='kroupa'), + Parameter('mass_2', 1.0, 100.0, dist='uniform'), + Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), + Parameter('ecc', 1e-9, 0.99, dist='sana_ecc'), + Parameter('metallicity', 1e-4, 0.03, dist='flat_in_log'), + ]) + return AdaptiveSampler( + parameter_space=params, total_systems=10, batch_size=5, BSEDict={}, + SSEDict=self.SSE, + is_interesting=lambda bpp: (0, np.array([], dtype=int)), + reject_systems=reject_systems, + ) + + def test_sse_dict_stored(self): + self.assertEqual(self._make("default").sse_dict, self.SSE) + + def test_sse_dict_bound_into_default_reject_via_sentinel(self): + sampler = self._make("default") + self.assertEqual(getattr(sampler.reject_fn, 'keywords', {}).get('SSEDict'), + self.SSE) + + def test_sse_dict_bound_when_default_reject_passed_explicitly(self): + sampler = self._make(default_reject) + self.assertEqual(getattr(sampler.reject_fn, 'keywords', {}).get('SSEDict'), + self.SSE) + + def test_custom_reject_not_rebound(self): + my_reject = lambda binary_params: np.zeros(len(binary_params['mass_1']), bool) + sampler = self._make(my_reject) + self.assertIs(sampler.reject_fn, my_reject) # left untouched + + def test_default_reject_accepts_ssedict_kwarg(self): + # default_reject runs with an explicit SSEDict and returns a mask. + bp = { + 'mass_1': np.array([20.0]), 'mass_2': np.array([10.0]), + 'porb': np.array([1000.0]), 'ecc': np.array([0.1]), + 'metallicity': np.array([0.014]), + } + mask = default_reject(bp, SSEDict=self.SSE) + self.assertEqual(mask.shape, (1,)) + self.assertFalse(bool(mask[0])) # a wide, detached binary survives + + # ========================================================================== # Gaussian mixture model # ========================================================================== @@ -519,6 +571,7 @@ def _build(self): 'is_interesting': lambda bpp: (0, np.array([], dtype=int)), 'derive_params': lambda s: {'mass_2': s['mass_1'] * 0.5, 'ecc': 0.0, 'metallicity': 0.02}, + 'SSEDict': {'stellar_engine': 'sse'}, 'reject_systems': None, 'nproc': 2, 'kappa': 1.0, 'n_generations': 1, 'only_save_hit_tables': False, 'min_active_fraction': 0.01, 'min_entropy_change': 0.01, From 920287da2534e34f95b5c9054fddd1a7a8caa996 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 29 Jun 2026 12:56:30 -0400 Subject: [PATCH 43/45] update changelog and version stuff --- changelog.md | 17 +++++++++++++++-- src/cosmic/data/cosmic-settings.json | 4 ++-- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/changelog.md b/changelog.md index 03f774e83..9bdf8e7c9 100644 --- a/changelog.md +++ b/changelog.md @@ -3,9 +3,22 @@ ## 4.2.0 +This version, among other thins, introduces adaptive importance sampling to COSMIC. + - Additions/changes - - Adaptive importance sampling is now available through the STROOPWAFEL algorithm (https://arxiv.org/abs/1905.00910) - - This is accessed through ``cosmic.sample.stroopwafel.AdaptiveSampler`` + - **Adaptive importance sampling via the STROOPWAFEL algorithm** ([Broekgaarden et al. 2019](https://arxiv.org/abs/1905.00910)) — a new vectorised module for efficiently sampling rare binary outcomes (e.g. merging double compact objects), where flat Monte Carlo would need millions of evolutions to collect a handful of systems. Lives in ``cosmic.sample.stroopwafel``: + - ``AdaptiveSampler`` runs the three-phase pipeline (exploration → adaptation → refinement) and returns importance-weighted results; pass ``mc_only=True`` for a plain Monte Carlo baseline. + - ``ParameterSpace`` / ``Parameter`` define the sampled dimensions. Each parameter takes a single composable ``dist`` (see below), and columns follow the order the parameters are supplied. + - A binary is defined by ``{mass_1, mass_2, porb, ecc, metallicity}``; each is either sampled or returned by a user ``derive_params`` callback, validated when the sampler is constructed. + - Composable distributions (``cosmic.sample.stroopwafel.distributions``): base distributions (``Uniform``, ``PowerLaw``, ``BrokenPowerLaw``, ``TruncatedNormal``) combined with coordinate transforms (``Identity``, ``Log10``, ``Ln``, ``Sin``, ``CosShift``). Built-ins include ``kroupa`` (a continuous broken power law, α=-1.3 below 0.5 Msun and -2.3 above), ``sana``, ``sana_ecc``, ``flat_in_log``, ``uniform``, ``uniform_in_sine``, ``uniform_in_cosine`` and ``disberg``. Define your own by passing a ``Distribution`` instance, calling ``register(...)``, or subclassing ``Distribution``. + - Physical rejection via ``default_reject`` (or a custom callback). It is ``SSEDict``-aware, so ZAMS radii are computed with the same stellar engine used for evolution. + - Built-in hit-definition presets ``any_dco`` and ``merging_dco`` (``cosmic.sample.stroopwafel.presets``), or supply any ``(bpp) -> (n_hits, hit_bin_nums)`` function. + - Results are returned as ``cosmic.output.COSMICStroopOutput``, holding the samples, importance weights, hit flags, and full COSMIC tables, with ``hit_rate``/``hit_rate_uncertainty`` properties, ``draw_representative_sample(...)`` (weighted bootstrap), and ``save``/``from_file``. + - Self-contained checkpointing: ``run_exploration()`` returns a ``STROOPWAFELCheckpoint`` (``.save(path)``), and ``AdaptiveSampler.from_checkpoint(path)`` rebuilds the sampler — parameter space, settings, callables (serialised with ``dill``), and RNG state — for ``run_refinement()`` with no re-specification. Any setting can be overridden as a keyword. + - ``bhflag = 4`` is added as another option, which applies fallback-modulation even when kicks are directly supplied + +- Documentation: + - New "Adaptive importance sampling" tutorial series under ``docs/pages/tutorials/adaptive/``: getting started (``basics``), defining and customising distributions (``distributions``), interpreting outputs and weights (``outputs``), and saving/resuming runs (``checkpoint``). ## 4.1.0 diff --git a/src/cosmic/data/cosmic-settings.json b/src/cosmic/data/cosmic-settings.json index 2d458bbc6..e1588607f 100644 --- a/src/cosmic/data/cosmic-settings.json +++ b/src/cosmic/data/cosmic-settings.json @@ -1095,8 +1095,8 @@ }, { "name": 4, - "description": "Same as 1, but is also applied if the kick is provided directly in natal_kick_array", - "version_added": "4.1.0" + "description": "Same as 1, but is also applied if the kick is provided directly in natal_kick_array (motivated by adaptive importance sampling)", + "version_added": "4.2.0" } ] }, From d3119683aa458a627559c14e56b74509c7cfbcb9 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Mon, 29 Jun 2026 13:14:28 -0400 Subject: [PATCH 44/45] remove old examples dir --- .../sample/stroopwafel/examples/bh_bh.py | 65 ---- .../sample/stroopwafel/examples/bh_star.py | 281 ------------------ 2 files changed, 346 deletions(-) delete mode 100644 src/cosmic/sample/stroopwafel/examples/bh_bh.py delete mode 100644 src/cosmic/sample/stroopwafel/examples/bh_star.py diff --git a/src/cosmic/sample/stroopwafel/examples/bh_bh.py b/src/cosmic/sample/stroopwafel/examples/bh_bh.py deleted file mode 100644 index 394fa61cb..000000000 --- a/src/cosmic/sample/stroopwafel/examples/bh_bh.py +++ /dev/null @@ -1,65 +0,0 @@ -#!/usr/bin/env python -import time -import argparse - -from cosmic.sample.stroopwafel import AdaptiveSampler, ParameterSpace, Parameter -from cosmic.sample.stroopwafel.presets import merging_dco - -import sys -sys.path.append("../../../../../docs") -import generate_default_bsedict -BSEDict = generate_default_bsedict.get_default_BSE_settings(to_python=True) - - -parser = argparse.ArgumentParser() -parser.add_argument('--num_systems', type=int, default=100_000) -parser.add_argument('--num_cores', type=int, default=1) -parser.add_argument('--num_per_core', type=int, default=1000) -parser.add_argument('--mc_only', type=bool, default=False) -parser.add_argument('--seed', type=int, default=None) -args = parser.parse_args() - -# COSMIC v4+ single stellar evolution settings (sse engine; swap for METISSE) -SSEDict = {'stellar_engine': 'sse'} - -# ------------------------------------------------------------------ -# Define parameter space -# ------------------------------------------------------------------ -params = ParameterSpace([ - Parameter('mass_1', 5.0, 150.0, dist='kroupa'), - Parameter('q', 0.0, 1.0, dist='uniform'), - Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), - Parameter('ecc', 1e-9, 0.99999999, dist='sana_ecc'), - Parameter('metallicity', 0.0001, 0.03, dist='flat_in_log'), - Parameter('natal_kick_1', 0.0, 5000.0, dist='disberg') -]) - -def derive_params(sampled): - """Return the one required parameter (mass_2) not sampled directly.""" - return {'mass_2': sampled['mass_1'] * sampled['q']} - -is_interesting = merging_dco(kstar_1=[14], kstar_2=[14], max_merge_time=13.7) - -if __name__ == '__main__': - start = time.time() - - sw = AdaptiveSampler( - parameter_space=params, - total_systems=args.num_systems, - batch_size=args.num_per_core, - BSEDict=BSEDict, - SSEDict=SSEDict, - is_interesting=is_interesting, - derive_params=derive_params, - reject_systems="default", - nproc=args.num_cores, - mc_only=args.mc_only, - seed=args.seed, - ) - - result = sw.run() - result.save("output/bh_bh/result.h5") - - elapsed = time.time() - start - print(f"\nTotal time: {elapsed:.1f}s") - print(f"Results saved to output/bh_bh/result.h5") diff --git a/src/cosmic/sample/stroopwafel/examples/bh_star.py b/src/cosmic/sample/stroopwafel/examples/bh_star.py deleted file mode 100644 index 78182e243..000000000 --- a/src/cosmic/sample/stroopwafel/examples/bh_star.py +++ /dev/null @@ -1,281 +0,0 @@ -#!/usr/bin/env python -"""Example: black hole + normal-star binaries with STROOPWAFEL vs Monte Carlo. - -A "hit" is any binary where one component is a black hole (kstar=14) while -its companion is still a normal (non-degenerate) star (kstar < 10) and the -system remains bound (sep > 0). BH+star binaries are intrinsically rare -because they require a massive primary that forms a BH without disrupting -the binary, making them a good testbed for adaptive importance sampling. - -The script runs both methods on identical parameter spaces and system budgets, -then prints a side-by-side comparison of hit rates, uncertainties, and -wall-clock times. The key figure of merit is how many Monte Carlo systems -would be needed to match STROOPWAFEL's statistical precision. - -Usage ------ - python bh_star.py # full comparison (default 10 000 systems) - python bh_star.py --num_systems 50000 # larger run for a clearer signal - python bh_star.py --mc_only # MC baseline only - python bh_star.py --sw_only # STROOPWAFEL only -""" -import os -import sys -import time -import argparse -import numpy as np - -sys.path.append(os.path.join(os.path.dirname(__file__), '..')) -sys.path.append(os.path.join(os.path.dirname(__file__), '../..')) - -from stroopwafel import AdaptiveSampler, ParameterSpace, Parameter -from stroopwafel.rejection import default_reject - -# ------------------------------------------------------------------ -# CLI -# ------------------------------------------------------------------ -parser = argparse.ArgumentParser( - description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter, -) -parser.add_argument('--num_systems', type=int, default=50000, - help='Total systems to evolve per method (default: 50000)') -parser.add_argument('--batch_size', type=int, default=1000, - help='Systems per COSMIC call (default: 1000)') -parser.add_argument('--num_cores', type=int, default=8, - help='CPU cores for COSMIC (default: 8)') -parser.add_argument('--n_generations', type=int, default=1, - help='STROOPWAFEL refinement generations (default: 1)') -parser.add_argument('--mc_only', action='store_true', - help='Run Monte Carlo baseline only, skip STROOPWAFEL') -parser.add_argument('--sw_only', action='store_true', - help='Run STROOPWAFEL only, skip Monte Carlo baseline') -parser.add_argument('--output_dir', default='output/bh_star', - help='Directory for output files (default: output/bh_star)') -parser.add_argument('--seed', type=int, default=117, - help='Base random seed; MC uses seed, STROOPWAFEL uses seed+1') -args = parser.parse_args() - -# ------------------------------------------------------------------ -# BSE physics -# ------------------------------------------------------------------ -BSEDict = { - 'xi': 1.0, 'bhflag': 4, 'neta': 0.5, 'windflag': 3, 'wdflag': 1, - 'alpha1': [1.0, 1.0], 'pts1': 0.001, 'pts3': 0.02, 'pts2': 0.01, - 'epsnov': 0.001, 'hewind': 0.5, 'ck': 1000, 'bwind': 0.0, - 'lambdaf': 0.0, 'mxns': 3.0, 'beta': -1.0, 'tflag': 1, 'acc2': 1.5, - 'grflag': 1, 'remnantflag': 4, 'ceflag': 0, 'eddfac': 1.0, - 'ifflag': 0, 'bconst': 3000, 'sigma': 265.0, 'gamma': -2.0, - 'pisn': 45.0, - # natal_kick_array is intentionally absent: kick parameters are sampled - # per-binary in the ParameterSpace below, and the engine injects them - # directly into the InitialBinaryTable for each COSMIC call. - 'bhsigmafrac': 1.0, 'polar_kick_angle': 90, - 'qcrit_array': [0.0] * 16, - 'cekickflag': 2, 'cehestarflag': 0, 'cemergeflag': 0, - 'ecsn': 2.5, 'ecsn_mlow': 1.8, 'aic': 1, 'ussn': 0, - 'sigmadiv': -20.0, 'qcflag': 5, 'eddlimflag': 0, - 'fprimc_array': [2.0 / 21.0] * 16, - 'bhspinflag': 0, 'bhspinmag': 0.0, 'rejuv_fac': 1.0, - 'rejuvflag': 0, 'htpmb': 1, 'ST_cr': 1, 'ST_tide': 1, - 'bdecayfac': 1, 'rembar_massloss': 0.5, 'kickflag': 5, - 'zsun': 0.014, 'bhms_coll_flag': 0, 'don_lim': -1, - 'acc_lim': [-1, -1], 'rtmsflag': 0, 'wd_mass_lim': 1, - "ppi_co_shift": 0.0, "ppi_extra_ml": 0.0, "fryer_mass_limit": 0, - "maltsev_mode": 0, "maltsev_fallback": 0.5, "maltsev_pf_prob": 0.1, - "mm_mu_ns": 800, "mm_mu_bh": 400, "LBV_flag": 1, - "fryer_fmix": 0.5, "fryer_mcrit_nsbh": 5.0, "smt_periastron_check": 0 -} - -# COSMIC v4+ single stellar evolution settings (sse engine; swap for METISSE) -SSEDict = {'stellar_engine': 'sse'} - -# ------------------------------------------------------------------ -# Parameter space -# -# Orbital / stellar parameters (5 dimensions) -# mass_1 : primary mass [Msun], Kroupa IMF -# q : mass ratio m2/m1 ∈ [0.01, 1], uniform -# porb : log10(orbital period / days), Sana power law -# bounds 0.15 to 5.5 → periods ~1.4 d to ~316 000 d -# ecc : eccentricity, Sana power law -# metallicity : metallicity, log-uniform -# -# Primary natal kick magnitude (1 dimension) -# natal_kick_1 : kick speed [km/s], log-normal (mu=5.67, sigma=0.59 in ln-space) -# physical bounds [0.1, 5000] km/s; median ≈ 291 km/s -# -# Total: 6-dimensional parameter space. -# -# Kick angles (phi_1, theta_1, mean_anomaly_1) and the secondary kick are -# intentionally excluded. Angles have flat hit-probability across their full -# range so they contribute no information to the mixture model while each -# extra dimension widens the Gaussians by N^(1/D_old - 1/D_new). The engine -# fills all omitted kick columns with the -100 sentinel so COSMIC draws those -# components from its own prescription (kickflag=5 / sigma=265 km/s). -# -# ------------------------------------------------------------------ -params = ParameterSpace([ - # --- orbital / stellar --- - Parameter('mass_1', 5.0, 150.0, dist='kroupa'), - Parameter('q', 0.01, 1.0, dist='uniform'), - Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), # ~1.4 d to ~316 000 d - Parameter('ecc', 1e-9, 0.99999999, dist='sana_ecc'), - Parameter('metallicity', 0.0001, 0.03, dist='flat_in_log'), - # --- primary natal kick magnitude only --- - Parameter('natal_kick_1', 0.1, 5000.0, dist='disberg'), -]) - -# ------------------------------------------------------------------ -# Derived quantities -# ------------------------------------------------------------------ -def derive_params(sampled): - """Provide the secondary mass from the sampled primary mass and mass ratio. - - A binary is defined by {mass_1, mass_2, porb, ecc, metallicity}. Here - mass_1, porb, ecc, and metallicity are sampled directly, so only mass_2 - needs deriving. - - Parameters - ---------- - sampled : dict - Maps each sampled parameter name to its (N,) array of physical values. - - Returns - ------- - dict - ``{'mass_2': ...}`` -- the one required parameter not sampled here. - """ - return {'mass_2': sampled['mass_1'] * sampled['q']} - -# ------------------------------------------------------------------ -# Hit definition: BH (kstar=14) + normal star (kstar 0–9), still bound -# ------------------------------------------------------------------ -_STAR_KSTARS = set(range(10)) # kstar 0–9: non-degenerate stars - -def is_bh_star(bpp): - """Identify binaries that are in a bound BH + normal-star phase for ≥ 100 Myr. - - Parameters - ---------- - bpp : pandas.DataFrame - COSMIC binary population parameters output. - - Returns - ------- - n_hits : int - Number of distinct binaries satisfying the criterion. - hit_bin_nums : numpy.ndarray - Integer bin_num values of those binaries. - """ - bh_star_mask = ( - ( (bpp.kstar_1 == 14) & bpp.kstar_2.isin(_STAR_KSTARS)) - | (bpp.kstar_1.isin(_STAR_KSTARS) & (bpp.kstar_2 == 14)) - ) & (bpp.sep > 0) - - bh_star_rows = bpp.loc[bh_star_mask] - if len(bh_star_rows) == 0: - return 0, np.array([], dtype=int) - - # Group by bin_num (the natural key) so that min/max tphys are aligned - # on the same index. drop_duplicates would give rows with different - # integer-row indices that pandas would NOT align correctly on subtraction. - phase_start = bh_star_rows.groupby('bin_num')['tphys'].min() # Series indexed by bin_num - phase_end = bh_star_rows.groupby('bin_num')['tphys'].max() # Series indexed by bin_num - duration = phase_end - phase_start # tphys in Myr - - long_enough_bin_nums = duration[duration >= 100.0].index.values - return len(long_enough_bin_nums), long_enough_bin_nums - -# ------------------------------------------------------------------ -# Helper: run one sampler and return (result, elapsed_seconds) -# ------------------------------------------------------------------ -def run_sampler(mc_only, seed): - label = "Monte Carlo" if mc_only else "STROOPWAFEL" - print(f"\n{'='*60}") - print(f" Running {label} (seed={seed})") - print(f"{'='*60}") - - sw = AdaptiveSampler( - parameter_space=params, - total_systems=args.num_systems, - batch_size=args.batch_size, - BSEDict=BSEDict, - SSEDict=SSEDict, - is_interesting=is_bh_star, - derive_params=derive_params, - reject_systems=default_reject, - nproc=args.num_cores, - n_generations=args.n_generations, - mc_only=mc_only, - seed=seed, - ) - - t0 = time.time() - result = sw.run() - elapsed = time.time() - t0 - - return result, elapsed - -# ------------------------------------------------------------------ -# Helper: summarise one result -# ------------------------------------------------------------------ -def print_summary(label, result, elapsed): - raw_hits = int(np.sum(result.is_hit)) - hit_rate = result.hit_rate - uncertainty = result.hit_rate_uncertainty - - print(f"\n--- {label} summary ---") - print(f" Systems evolved : {len(result.weights):,}") - print(f" Raw hits found : {raw_hits:,}") - print(f" Weighted hit rate: {hit_rate:.6e} ± {uncertainty:.6e}") - print(f" Wall-clock time : {elapsed:.1f} s") - -# ------------------------------------------------------------------ -# Helper: side-by-side comparison -# ------------------------------------------------------------------ -def print_comparison(mc_result, mc_elapsed, sw_result, sw_elapsed): - mc_rate = mc_result.hit_rate - mc_unc = mc_result.hit_rate_uncertainty - sw_rate = sw_result.hit_rate - sw_unc = sw_result.hit_rate_uncertainty - - print(f"\n{'='*60}") - print(" Efficiency comparison") - print(f"{'='*60}") - print(f"{'':30s} {'Monte Carlo':>15s} {'STROOPWAFEL':>15s}") - print(f" {'Systems evolved':<28s} {len(mc_result.weights):>15,} {len(sw_result.weights):>15,}") - print(f" {'Raw hits found':<28s} {int(np.sum(mc_result.is_hit)):>15,} {int(np.sum(sw_result.is_hit)):>15,}") - print(f" {'Weighted hit rate':<28s} {mc_rate:>15.4e} {sw_rate:>15.4e}") - print(f" {'Uncertainty (1σ)':<28s} {mc_unc:>15.4e} {sw_unc:>15.4e}") - print(f" {'Wall-clock time (s)':<28s} {mc_elapsed:>15.1f} {sw_elapsed:>15.1f}") - - if sw_unc > 0 and mc_unc > 0: - # How many MC systems would give the same precision as SW? - equivalent_mc = len(mc_result.weights) * (mc_unc / sw_unc) ** 2 - speedup = equivalent_mc / len(sw_result.weights) - print(f"\n STROOPWAFEL is ~{speedup:.1f}x more statistically efficient:") - print(f" Monte Carlo would need ~{equivalent_mc:,.0f} systems to match") - print(f" STROOPWAFEL's precision on {len(sw_result.weights):,} systems.") - -# ------------------------------------------------------------------ -# Main -# ------------------------------------------------------------------ -if __name__ == '__main__': - os.makedirs(args.output_dir, exist_ok=True) - - mc_result = mc_elapsed = None - sw_result = sw_elapsed = None - - if not args.sw_only: - mc_result, mc_elapsed = run_sampler(mc_only=True, seed=args.seed) - print_summary("Monte Carlo", mc_result, mc_elapsed) - mc_result.save(os.path.join(args.output_dir, 'mc_result.h5')) - - if not args.mc_only: - sw_result, sw_elapsed = run_sampler(mc_only=False, seed=args.seed + 1) - print_summary("STROOPWAFEL", sw_result, sw_elapsed) - sw_result.save(os.path.join(args.output_dir, 'sw_result.h5')) - - if mc_result is not None and sw_result is not None: - print_comparison(mc_result, mc_elapsed, sw_result, sw_elapsed) From 340328b5950dc577dff1c5d9aac2752ed7e11e38 Mon Sep 17 00:00:00 2001 From: Tom Wagg Date: Tue, 30 Jun 2026 15:10:44 -0400 Subject: [PATCH 45/45] fix double hyphen --- src/cosmic/output.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cosmic/output.py b/src/cosmic/output.py index 1fa913b26..bcc5ec3e0 100644 --- a/src/cosmic/output.py +++ b/src/cosmic/output.py @@ -74,7 +74,7 @@ def __init__(self, bpp=None, bcm=None, initC=None, kick_info=None, file=None, la self.bpp = pd.read_hdf(file, key=f'bpp{file_key_suffix}') self.bcm = pd.read_hdf(file, key=f'bcm{file_key_suffix}') self.initC = load_initC(file, key=f'initC{file_key_suffix}', - settings_key=f'initC_{file_key_suffix}_settings') + settings_key=f'initC{file_key_suffix}_settings') self.kick_info = pd.read_hdf(file, key=f'kick_info{file_key_suffix}') with h5.File(file, 'r') as f: file_version = f.attrs.get('COSMIC_version', 'unknown')