diff --git a/changelog.md b/changelog.md index a0447be77..9bdf8e7c9 100644 --- a/changelog.md +++ b/changelog.md @@ -1,6 +1,25 @@ # COSMIC Changelog ## Prepend only please! +## 4.2.0 + +This version, among other thins, introduces adaptive importance sampling to COSMIC. + +- Additions/changes + - **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 - Additions/changes diff --git a/docs/modules/sample.rst b/docs/modules/sample.rst index 2576cde9e..8bf6d50d8 100644 --- a/docs/modules/sample.rst +++ b/docs/modules/sample.rst @@ -29,6 +29,13 @@ Multidimensional sampler :no-inheritance-diagram: :no-heading: +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 5d1892b40..63e110011 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:: @@ -122,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.rst b/docs/pages/tutorials/adaptive.rst new file mode 100644 index 000000000..85ba4b838 --- /dev/null +++ b/docs/pages/tutorials/adaptive.rst @@ -0,0 +1,13 @@ +############################ +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/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 new file mode 100644 index 000000000..fa591b3ef --- /dev/null +++ b/docs/pages/tutorials/adaptive/basics.rst @@ -0,0 +1,381 @@ +.. _adaptive_basics: + +********************************************* +Adaptive importance sampling for rare systems +********************************************* + +The standard :ref:`independent ` and :ref:`multidimensional ` +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. + +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 +(`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 a flat Monte Carlo draw would require. +Example use cases include: + +* Double black holes (or neutron stars) merging within the Hubble time (GW sources) +* Long lived BH + stellar-companion systems + +If your target population is common the plain independent sampler is simpler and fast enough. + +How it works: think 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 + 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. + +Setup +===== + +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, 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 + + import numpy as np + from cosmic.sample.stroopwafel import ParameterSpace, Parameter + + 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'), + ]) + +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. + + +Complete the binary definition +------------------------------ + +Every binary handed to ``COSMIC`` is defined by **five** parameters: ``mass_1``, +``mass_2``, ``porb``, ``ecc``, and ``metallicity``. + +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. + +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): + +.. code-block:: python + + def derive_params(sampled): + """Provide binary parameters not drawn from the ParameterSpace.""" + return {'mass_2': sampled['mass_1'] * sampled['q']} + +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: + +.. code-block:: python + + params = ParameterSpace([ + Parameter('porb', 10**(0.15), 10**(5.5), dist='sana'), + ]) + + 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 +--------------------------- + +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: + +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(binary_params): + base_mask = default_reject(binary_params) + # additionally reject secondaries below 1 M_sun + base_mask |= (binary_params['mass_2'] < 1.0) + return base_mask + +Passing ``reject_systems`` is optional; you can instead pass ``None`` to skip physical rejection entirely. + + +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 +``(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`. +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. + +``merging_dco(kstar_1, kstar_2, max_merge_time=13.7)`` + 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 + + 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 NS-NS systems merging within the Hubble time + is_interesting_merging = merging_dco(kstar_1=[13], kstar_2=[13], max_merge_time=13.7) + + +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 + + 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'] < 10)) + | ((bpp['kstar_2'] == 14) & (bpp['kstar_1'] < 10)) + ) + & (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 + + +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 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. + +Examples +-------- + +Bound BH + BH binaries +^^^^^^^^^^^^^^^^^^^^^^ + +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. + +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 + +.. 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. + +.. code-block:: python + + 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'), + ]) + +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']} + + +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 + 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", + nproc=4, + n_generations=1, + 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}") + + +BH + star binaries surviving 100 Myr +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +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 + + import numpy as np + from cosmic.sample.stroopwafel import AdaptiveSampler + from cosmic.sample.stroopwafel.rejection import default_reject + + sampler = AdaptiveSampler( + parameter_space=params, # reuse from BHBH example + 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", + nproc=4, + n_generations=1, + seed=42, + ) + result = sampler.run() + + +Rules of thumb +============== + +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`` +-------------- + +``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 from the output DataFrames. + +``total_systems`` +----------------- + +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. You are probably safe with just 1 generation unless you have a very rare population. + +.. 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). + +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 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! 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 diff --git a/docs/pages/tutorials/adaptive/checkpoint.rst b/docs/pages/tutorials/adaptive/checkpoint.rst new file mode 100644 index 000000000..6f989fe52 --- /dev/null +++ b/docs/pages/tutorials/adaptive/checkpoint.rst @@ -0,0 +1,153 @@ +.. _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). + +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 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; +* 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 +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 + + 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 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, + SSEDict={'stellar_engine': 'sse'}, + is_interesting=any_dco(kstar_1=[14], kstar_2=[14]), + derive_params=derive_params, + reject_systems="default", + 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 + + 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``, ``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: + +.. code-block:: python + + sampler = AdaptiveSampler.from_checkpoint( + 'checkpoint.h5', nproc=16 + ) + +Any :class:`~cosmic.sample.stroopwafel.AdaptiveSampler` constructor argument may be +overridden this way (``parameter_space``, ``BSEDict``, ``SSEDict``, ``derive_params``, +``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 +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 stored in a checkpoint +============================== + +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 COSMIC output tables (``bpp``, ``bcm``, ``initC``, ``kick_info``) for the explored + systems; +* 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``, + ``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 +``from_checkpoint`` if you instead want a fresh stream). + + +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`` 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 new file mode 100644 index 000000000..66a5835c9 --- /dev/null +++ b/docs/pages/tutorials/adaptive/distributions.rst @@ -0,0 +1,212 @@ +.. _adaptive_distributions: + +********************************* +Distributions and custom priors +********************************* + +This tutorial assumes that you've already gone through :ref:`adaptive_basics`. + +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 +======================= + +You can select any of the following distributions 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'`` + - 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. + * - ``'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. + * - ``'disberg'`` + - log-normal + - Natal-kick magnitude, :math:`\ln v \sim \mathcal{N}(5.67, 0.59)`. + +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 +=========================== + +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`, + :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`, + :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, BrokenPowerLaw, TruncatedNormal, Log10, + ) + + 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. + + +Defining your own distribution +============================== + +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 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 + 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 — 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 + + 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, 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)`` + 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)) + +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:: + + 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 + +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. + +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 new file mode 100644 index 000000000..f95bd5154 --- /dev/null +++ b/docs/pages/tutorials/adaptive/outputs.rst @@ -0,0 +1,144 @@ +.. _adaptive_outputs: + +*************************************** +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 + +.. 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` (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 +---------------- + +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. + +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). + +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). + +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. + +.. 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 +==================================================== + +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, 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. + +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 + + import numpy as np + from cosmic.output import COSMICStroopOutput + + 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 population so that you can access the full evolution history if you need it. + +Wrap-up +======= + +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. + +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! 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" diff --git a/src/cosmic/data/cosmic-settings.json b/src/cosmic/data/cosmic-settings.json index 6f7936db8..e1588607f 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 (motivated by adaptive importance sampling)", + "version_added": "4.2.0" } ] }, diff --git a/src/cosmic/output.py b/src/cosmic/output.py index 024e87086..bcc5ec3e0 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 @@ -10,7 +11,8 @@ import warnings -__all__ = ['COSMICOutput', 'COSMICPopOutput', 'save_initC', 'load_initC'] +__all__ = ['COSMICOutput', 'COSMICPopOutput', 'COSMICStroopOutput', + 'STROOPWAFELCheckpoint', 'save_initC', 'load_initC'] kstar_translator = [ @@ -72,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') @@ -334,6 +336,436 @@ 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, + ) + + 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 + ---------- + 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 (n_samples, D) containing the drawn samples in physical space. + bin_nums : `numpy.ndarray` + 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] + probs = self.weights[hit_idx] / self.weights[hit_idx].sum() + + # draw a representative sample with replacement + rng = rng or np.random.default_rng() + bin_nums = rng.choice(hit_idx, size=n_samples, replace=True, p=probs) + + # `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 + + +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 + + 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``, ``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. + 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. + 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. + """ + + 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): + self.config = dict(config) + 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.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) + + # 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 + 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 + + # 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['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 + + @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']), + ) + + config = dill.loads(f['config'][:].tobytes()) + + meta = f['meta'] + 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']) + + # 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( + 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, + num_explored=num_explored, num_hits=num_hits, + num_hits_exploratory=num_hits_exploratory, + fraction_explored=fraction_explored, + prior_fraction_rejected=prior_fraction_rejected, + ) + + class COSMICPopOutput(): def __init__(self, file, label=None): # read in convergence tables and totals 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/__init__.py b/src/cosmic/sample/stroopwafel/__init__.py new file mode 100755 index 000000000..ec0189b11 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/__init__.py @@ -0,0 +1,14 @@ +"""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. + +""" +from .main import AdaptiveSampler +from .parameter_space import ParameterSpace, Parameter + +__all__ = ['AdaptiveSampler', 'ParameterSpace', 'Parameter'] diff --git a/src/cosmic/sample/stroopwafel/distributions.py b/src/cosmic/sample/stroopwafel/distributions.py new file mode 100644 index 000000000..676eac65a --- /dev/null +++ b/src/cosmic/sample/stroopwafel/distributions.py @@ -0,0 +1,511 @@ +"""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:`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 +``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 +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 + + +# --------------------------------------------------------------------------- +# 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 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. + + 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)``. + + 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': 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()), +} + + +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/main.py b/src/cosmic/sample/stroopwafel/main.py new file mode 100644 index 000000000..6e6dc0433 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/main.py @@ -0,0 +1,916 @@ +import os +from functools import partial +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, STROOPWAFELCheckpoint + +from .mixture_model import GaussianMixture +from .rejection import default_reject + + +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. + 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 + 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 uses :func:`~cosmic.sample.stroopwafel.rejection.default_reject`. + Pass None to skip physical rejection entirely. + nproc : `int`, optional + Number of CPU cores for COSMIC, by default 1 + kappa : `float`, optional + Gaussian width scaling factor, by default 1.0 + n_generations : `int`, optional + Number of refinement generations, by default 1 + 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 + 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. + 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 + #: 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, + 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): + self.param_space = parameter_space + self.total_systems = total_systems + self.batch_size = batch_size + self.bse_dict = BSEDict + self.sse_dict = SSEDict + self.derive_fn = derive_params + # 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 + 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) + self.min_active_fraction = min_active_fraction + self.min_entropy_change = min_entropy_change + + # 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) 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 = [] + + # 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. + + For multi-job workflows (e.g. SLURM) use :meth:`run_exploration` + and :meth:`run_refinement` instead. + + Returns + ------- + `COSMICStroopOutput` + Container holding all samples, weights, COSMIC output tables, + and associated metadata. + """ + self._explore() + + if not self.mc_only and self.num_hits > 0: + self._adapt() + self._refine() + + 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. + """ + 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 (the checkpoint is self-contained) + sampler = AdaptiveSampler.from_checkpoint('checkpoint.h5') + 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() + + #: Constructor arguments that may be overridden in :meth:`from_checkpoint`. + _OVERRIDABLE = frozenset({ + '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', + }) + + @classmethod + 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`. + **overrides + Any :class:`AdaptiveSampler` constructor argument + (``parameter_space``, ``total_systems``, ``batch_size``, + ``BSEDict``, ``is_interesting``, ``derive_params``, + ``reject_systems``, ``nproc``, ``kappa``, + ``n_generations``, ``only_save_hit_tables``, ``seed``). Passing + ``seed`` starts a fresh RNG instead of restoring the stored state. + + Returns + ------- + `AdaptiveSampler` + Ready to call :meth:`run_refinement`. + """ + if isinstance(checkpoint, (str, os.PathLike)): + checkpoint = STROOPWAFELCheckpoint.from_file(checkpoint) + + unknown = set(overrides) - cls._OVERRIDABLE + if unknown: + raise TypeError( + f"Unknown from_checkpoint override(s): {sorted(unknown)}. " + f"Allowed: {sorted(cls._OVERRIDABLE)}." + ) + + # 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 + + def _make_checkpoint(self): + """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) + 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, + 'SSEDict': self.sse_dict, + 'is_interesting': self.is_interesting_fn, + 'derive_params': self.derive_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, + '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, + } + + 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, + 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, + ) + + def _load_checkpoint(self, checkpoint): + """Restore 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] + + # ------------------------------------------------------------------ + # 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. + + 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 + # ------------------------------------------------------------------ + 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 / max(1.0 - self.prior_fraction_rejected, self.min_active_fraction) + )) + + # Sample from prior + samples, mask = self.param_space.sample(n_oversample, rng=self.rng) + + # Assemble binary parameters (sampled + derived) and reject + samples_phys = self.param_space.to_physical(samples) + 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 + + # 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_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_params) + + # Record which are hits + is_hit = np.zeros(len(selected), dtype=bool) + 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)) + 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) + 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) + rejected += np.sum(self._physical_reject_mask(phys)) + + 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, + min_active_fraction=self.min_active_fraction, + min_entropy_change=self.min_entropy_change + ) + + 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._physical_reject_mask, + 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) + 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] + 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_gauss_idx = valid_gauss_idx[indices[:n_take]] + 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_params) + + is_hit = np.zeros(n_take, dtype=bool) + 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)) + 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) + + self.num_hits += n_hits + self.finished += n_take + gen_finished += n_take + self._print_progress() + + # 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) + gen_priors = self.param_space.compute_prior(gen_samples) + + # Save current state in case we need to revert + saved_mixture = GaussianMixture( + 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( + gen_samples, gen_is_hit, gen_priors, + self.prior_fraction_rejected, + tolerance=1e-10, entropies=entropies + ) + + if should_revert: + self.mixture = saved_mixture + print(" Expectation-maximization (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 and assemble the final result. + + 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 + ------- + `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) + 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 / 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 + 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 / 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( + 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 + + # 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, + 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 + + 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. + df = pd.concat(frames, ignore_index=True) + df.index = df['bin_num'].values + return df + + return ( + _concat(bpp_list), + _concat(bcm_list), + _concat(initC_list), + _concat(kick_info_list), + ) + + # ------------------------------------------------------------------ + # COSMIC interface + # ------------------------------------------------------------------ + + # 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, binary_params): + """Evolve a batch of binaries with COSMIC and identify hits. + + Parameters + ---------- + 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 + ------- + 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. + bcm : `pandas.DataFrame` + COSMIC user-timestep output. + initC : `pandas.DataFrame` + COSMIC initial conditions output. + kick_info : `pandas.DataFrame` + COSMIC natal kick information output. + """ + n = len(binary_params['mass_1']) + + batch_initial = InitialBinaryTable.InitialBinaries( + 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=binary_params['metallicity'], + ) + + # 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 present_kick_cols: + for col in self._ALL_KICK_COLUMNS: + 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 + 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, + BSEDict=bse_dict_run, + SSEDict=self.sse_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, bcm, 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/meson.build b/src/cosmic/sample/stroopwafel/meson.build new file mode 100644 index 000000000..103abb233 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/meson.build @@ -0,0 +1,14 @@ +python_sources = [ + '__init__.py', + 'distributions.py', + 'main.py', + 'mixture_model.py', + 'parameter_space.py', + 'presets.py', + 'rejection.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..16d6c34b2 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/mixture_model.py @@ -0,0 +1,340 @@ +"""Vectorized Gaussian Mixture Model for adaptive importance sampling. + +Stores all mixture components as numpy arrays and provides vectorized +sampling, PDF evaluation, and expectation-maximization (EM) updates. +""" +import numpy as np +from scipy.stats import multivariate_normal, entropy as scipy_entropy + + +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, + 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): + 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=1.0, + min_active_fraction=0.01, min_entropy_change=0.01): + """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 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 + ------- + `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, + 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. + + 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 > 0.0: + 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 + 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, reject_mask_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. + reject_mask_fn : `callable` + Function ``(samples_physical) -> bool_mask`` returning True for + physically 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) + rejected += np.sum(reject_mask_fn(s_physical)) + + 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 (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 + ---------- + 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 / 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) + + 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): + 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] < self.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..05a7880a3 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/parameter_space.py @@ -0,0 +1,219 @@ +"""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 .distributions import Distribution, get_distribution + + +@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, always in physical space (e.g. solar masses for + ``'kroupa'``, days for ``'sana'``, km/s for ``'disberg'``). The + parameter's distribution maps this into sampling space via its + transform. + max_value : `float` + 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 + dist: "str | Distribution" = 'uniform' + + def __post_init__(self): + self.distribution = get_distribution(self.dist) + self.lo, self.hi = self.distribution.transform.bounds(self.min_value, self.max_value) + + +class ParameterSpace: + """An ordered collection of Parameters with vectorized operations. + + 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. The column order of every (N, D) array + produced by this class matches the order of this list. + """ + + def __init__(self, params): + # 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) + + 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): + col = p.distribution.sample(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] = p.distribution.transform.to_physical(samples[:, i]) + 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] = p.distribution.transform.to_sampling(samples[:, i]) + 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): + 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) + + 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): + 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 diff --git a/src/cosmic/sample/stroopwafel/presets.py b/src/cosmic/sample/stroopwafel/presets.py new file mode 100644 index 000000000..19f80f838 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/presets.py @@ -0,0 +1,97 @@ +"""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. + """ + # 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) + 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/rejection.py b/src/cosmic/sample/stroopwafel/rejection.py new file mode 100644 index 000000000..be6d91804 --- /dev/null +++ b/src/cosmic/sample/stroopwafel/rejection.py @@ -0,0 +1,68 @@ +"""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. +""" +from cosmic.sample.sampler.independent import Sample +from cosmic.utils import calc_Roche_radius, a_from_p + + +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 + stars are in contact at ZAMS, or either star overflows its Roche lobe + at periastron. The orbital separation is computed from the orbital + period via Kepler's third law, and both stars share the binary + metallicity. + + Parameters + ---------- + binary_params : `dict` + 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 + + Returns + ------- + `numpy.ndarray` + (N,) boolean mask where True indicates a rejected system. + """ + mass_1 = binary_params['mass_1'] + mass_2 = binary_params['mass_2'] + porb = binary_params['porb'] + ecc = binary_params['ecc'] + metallicity = binary_params['metallicity'] + + # compute separation from periods and masses + separation = a_from_p(p=porb, m1=mass_1, m2=mass_2) + + # get stellar radii at ZAMS (depends on the stellar engine in SSEDict) + sampler = Sample() + 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) + rl_1 = calc_Roche_radius(mass_1, mass_2, peri_sep) + rl_2 = calc_Roche_radius(mass_2, mass_1, peri_sep) + + rejected = ( + (mass_2 < min_secondary_mass) + | (separation <= (radius_1 + radius_2)) + | (radius_1 / rl_1 > 1) + | (radius_2 / rl_2 > 1) + ) + + return rejected diff --git a/src/cosmic/src/kick.f b/src/cosmic/src/kick.f index 396b46992..67fef22ea 100644 --- a/src/cosmic/src/kick.f +++ b/src/cosmic/src/kick.f @@ -212,6 +212,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) @@ -262,7 +268,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 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' ] diff --git a/src/cosmic/tests/test_stroopwafel.py b/src/cosmic/tests/test_stroopwafel.py new file mode 100644 index 000000000..8697c236b --- /dev/null +++ b/src/cosmic/tests/test_stroopwafel.py @@ -0,0 +1,726 @@ +"""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, AdaptiveSampler +from cosmic.sample.stroopwafel.distributions import ( + DISTRIBUTIONS, register, get_distribution, + Distribution, Uniform, PowerLaw, BrokenPowerLaw, TruncatedNormal, + Identity, Log10, Ln, Sin, CosShift, +) +from cosmic.sample.stroopwafel.mixture_model import GaussianMixture +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 + + +# -------------------------------------------------------------------------- +# 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='disberg'), + ]) + + +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', 'disberg'}, + ) + + def test_builtin_compositions(self): + 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['disberg'], TruncatedNormal) + self.assertIsInstance(DISTRIBUTIONS['disberg'].transform, Ln) + 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 [(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) + 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(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(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) + + 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 [(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) + 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 = 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 + 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(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))) + with self.assertRaises(ValueError): + 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 +# ========================================================================== +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_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): + + def test_wide_binary_not_rejected(self): + binary_params = { + 'mass_1': np.array([20.0]), + 'mass_2': np.array([10.0]), + 'porb': np.array([1000.0]), # days -> several AU, well detached + 'ecc': np.array([0.1]), + 'metallicity': np.array([0.014]), + } + self.assertFalse(default_reject(binary_params)[0]) + + def test_low_mass_secondary_rejected(self): + 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]), + } + 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() + + +# ========================================================================== +# 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 +# ========================================================================== +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))) + + +# ========================================================================== +# 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}, + '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, + '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 +# ========================================================================== +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) + + +# ========================================================================== +# 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()