Composable distributions for STROOPWAFEL parameter sampling.

A 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 (Uniform, PowerLaw, BrokenPowerLaw, TruncatedNormal) composed with a transform (Identity, Log10, Ln, Sin, 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 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.

class cosmic.sample.stroopwafel.distributions.BrokenPowerLaw(breaks, alphas, transform=None)

Bases: Distribution

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 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 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:

breakssequence of float

Internal breakpoints in sampling space, strictly increasing. The distribution has len(breaks) + 1 segments.

alphassequence 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.

transformTransform, optional

Map between physical and sampling space, by default Identity.

pdf(values, lo, hi)

Prior probability density at values (sampling space).

Parameters:

valuesnumpy.ndarray

(N,) array of values in sampling space.

lofloat

Lower bound in sampling space.

hifloat

Upper bound in sampling space.

Returns:

numpy.ndarray

(N,) array of prior densities, normalised over [lo, hi].

sample(n, lo, hi, rng=None)

Draw n samples in sampling space within [lo, hi].

Parameters:

nint

Number of samples to draw.

lofloat

Lower bound in sampling space.

hifloat

Upper bound in sampling space.

rngnumpy.random.Generator, optional

Random number generator, by default None.

Returns:

numpy.ndarray

(N,) array of samples in sampling space.

sigma(values, lo, hi, avg_density)

CDF-space step, identical in spirit to PowerLaw.sigma().

class cosmic.sample.stroopwafel.distributions.CosShift

Bases: 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 Transform.bounds() swaps the endpoints to keep lo <= hi.

to_physical(values)

Convert sampling-space values to physical space.

Parameters:

valuesnumpy.ndarray or float

Value(s) in sampling space.

Returns:

numpy.ndarray or float

Value(s) in physical space.

to_sampling(values)

Convert physical-space values to sampling space.

Parameters:

valuesnumpy.ndarray or float

Value(s) in physical space.

Returns:

numpy.ndarray or float

Value(s) in sampling space.

class cosmic.sample.stroopwafel.distributions.Distribution(transform=None)

Bases: object

A prior distribution, sampled and evaluated in sampling space.

Subclasses implement sample() and pdf(). 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); ParameterSpace handles that conversion via Transform.bounds().

Parameters:

transformTransform, optional

Map between physical and sampling space, by default Identity.

pdf(values, lo, hi)

Prior probability density at values (sampling space).

Parameters:

valuesnumpy.ndarray

(N,) array of values in sampling space.

lofloat

Lower bound in sampling space.

hifloat

Upper bound in sampling space.

Returns:

numpy.ndarray

(N,) array of prior densities, normalised over [lo, hi].

sample(n, lo, hi, rng=None)

Draw n samples in sampling space within [lo, hi].

Parameters:

nint

Number of samples to draw.

lofloat

Lower bound in sampling space.

hifloat

Upper bound in sampling space.

rngnumpy.random.Generator, optional

Random number generator, by default None.

Returns:

numpy.ndarray

(N,) array of samples in sampling space.

sigma(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. PowerLaw) override this with an exact CDF-space step.

Parameters:

valuesnumpy.ndarray

(K,) array of hit values in sampling space.

lofloat

Lower bound in sampling space.

hifloat

Upper bound in sampling space.

avg_densityfloat

Characteristic inter-sample spacing.

Returns:

numpy.ndarray

(K,) array of sigma values.

class cosmic.sample.stroopwafel.distributions.Identity

Bases: Transform

Identity transform: sampling space is physical space.

class cosmic.sample.stroopwafel.distributions.Ln

Bases: Transform

Sampling space is ln(physical) (physical must be positive).

to_physical(values)

Convert sampling-space values to physical space.

Parameters:

valuesnumpy.ndarray or float

Value(s) in sampling space.

Returns:

numpy.ndarray or float

Value(s) in physical space.

to_sampling(values)

Convert physical-space values to sampling space.

Parameters:

valuesnumpy.ndarray or float

Value(s) in physical space.

Returns:

numpy.ndarray or float

Value(s) in sampling space.

class cosmic.sample.stroopwafel.distributions.Log10

Bases: Transform

Sampling space is log10(physical) (physical must be positive).

to_physical(values)

Convert sampling-space values to physical space.

Parameters:

valuesnumpy.ndarray or float

Value(s) in sampling space.

Returns:

numpy.ndarray or float

Value(s) in physical space.

to_sampling(values)

Convert physical-space values to sampling space.

Parameters:

valuesnumpy.ndarray or float

Value(s) in physical space.

Returns:

numpy.ndarray or float

Value(s) in sampling space.

class cosmic.sample.stroopwafel.distributions.PowerLaw(alpha, transform=None)

Bases: Distribution

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 (Log10).

Parameters:

alphafloat

Power-law exponent.

transformTransform, optional

Map between physical and sampling space, by default Identity.

pdf(values, lo, hi)

Prior probability density at values (sampling space).

Parameters:

valuesnumpy.ndarray

(N,) array of values in sampling space.

lofloat

Lower bound in sampling space.

hifloat

Upper bound in sampling space.

Returns:

numpy.ndarray

(N,) array of prior densities, normalised over [lo, hi].

sample(n, lo, hi, rng=None)

Draw n samples in sampling space within [lo, hi].

Parameters:

nint

Number of samples to draw.

lofloat

Lower bound in sampling space.

hifloat

Upper bound in sampling space.

rngnumpy.random.Generator, optional

Random number generator, by default None.

Returns:

numpy.ndarray

(N,) array of samples in sampling space.

sigma(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).

class cosmic.sample.stroopwafel.distributions.Sin

Bases: Transform

Sampling space is sin(angle) for an angle in [-pi/2, pi/2].

to_physical(values)

Convert sampling-space values to physical space.

Parameters:

valuesnumpy.ndarray or float

Value(s) in sampling space.

Returns:

numpy.ndarray or float

Value(s) in physical space.

to_sampling(values)

Convert physical-space values to sampling space.

Parameters:

valuesnumpy.ndarray or float

Value(s) in physical space.

Returns:

numpy.ndarray or float

Value(s) in sampling space.

class cosmic.sample.stroopwafel.distributions.Transform

Bases: object

Map between physical space and sampling space.

The base class is the identity transform; subclasses override to_sampling() and to_physical(). bounds is derived automatically and handles non-monotonic-increasing maps (e.g. CosShift) by sorting the endpoints.

bounds(lo, hi)

Map a physical-space bound pair into sampling space.

Parameters:

lofloat

Lower bound in physical space.

hifloat

Upper bound in physical space.

Returns:

lo_samplingfloat

Lower bound in sampling space.

hi_samplingfloat

Upper bound in sampling space.

to_physical(values)

Convert sampling-space values to physical space.

Parameters:

valuesnumpy.ndarray or float

Value(s) in sampling space.

Returns:

numpy.ndarray or float

Value(s) in physical space.

to_sampling(values)

Convert physical-space values to sampling space.

Parameters:

valuesnumpy.ndarray or float

Value(s) in physical space.

Returns:

numpy.ndarray or float

Value(s) in sampling space.

class cosmic.sample.stroopwafel.distributions.TruncatedNormal(mu, scale, transform=None)

Bases: Distribution

Normal distribution truncated to [lo, hi] in sampling space.

Combined with 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:

mufloat

Mean of the underlying normal (in sampling space).

scalefloat

Standard deviation of the underlying normal (in sampling space).

transformTransform, optional

Map between physical and sampling space, by default Identity.

pdf(values, lo, hi)

Prior probability density at values (sampling space).

Parameters:

valuesnumpy.ndarray

(N,) array of values in sampling space.

lofloat

Lower bound in sampling space.

hifloat

Upper bound in sampling space.

Returns:

numpy.ndarray

(N,) array of prior densities, normalised over [lo, hi].

sample(n, lo, hi, rng=None)

Draw n samples in sampling space within [lo, hi].

Parameters:

nint

Number of samples to draw.

lofloat

Lower bound in sampling space.

hifloat

Upper bound in sampling space.

rngnumpy.random.Generator, optional

Random number generator, by default None.

Returns:

numpy.ndarray

(N,) array of samples in sampling space.

class cosmic.sample.stroopwafel.distributions.Uniform(transform=None)

Bases: Distribution

Uniform distribution on [lo, hi] in sampling space.

Combined with a transform this covers uniform (identity), flat_in_log (Log10), uniform_in_sine (Sin) and uniform_in_cosine (CosShift).

pdf(values, lo, hi)

Prior probability density at values (sampling space).

Parameters:

valuesnumpy.ndarray

(N,) array of values in sampling space.

lofloat

Lower bound in sampling space.

hifloat

Upper bound in sampling space.

Returns:

numpy.ndarray

(N,) array of prior densities, normalised over [lo, hi].

sample(n, lo, hi, rng=None)

Draw n samples in sampling space within [lo, hi].

Parameters:

nint

Number of samples to draw.

lofloat

Lower bound in sampling space.

hifloat

Upper bound in sampling space.

rngnumpy.random.Generator, optional

Random number generator, by default None.

Returns:

numpy.ndarray

(N,) array of samples in sampling space.

cosmic.sample.stroopwafel.distributions.get_distribution(dist)

Resolve a name or instance to a Distribution.

Parameters:

diststr or Distribution

Either a key in DISTRIBUTIONS or a distribution instance (returned unchanged).

Returns:

Distribution

The resolved distribution instance.

cosmic.sample.stroopwafel.distributions.register(name, distribution)

Register a distribution instance under name.

Parameters:

namestr

Key used to refer to the distribution (e.g. from a Parameter).

distributionDistribution

The distribution instance to register.