amax parameter of iid_spin model missing for conversion_function on bilby sampler initialization

[mdmould]

Running bilby.run_sampler with the gwpopulation.hyperpe.HyperparameterLikelihood fails with the iid_spin model when sampling in parameters that require the gwpopulation.conversions.convert_to_beta_parameters conversion function.

This may be an issue with how bilby interacts with fixed parameters rather than an issue with gwpopulation itself, but I'll post the issue here with a reproducible example below.

Versions are bilby==2.2.3 and gwpopulation==1.0.0rc2.dev1+g37f1b4c (pip install git+https://github.com/ColmTalbot/gwpopulation, but the issue also occurs with gwpopulation==0.10.0).

import numpy as np
import pandas
import bilby
from gwpopulation.hyperpe import HyperparameterLikelihood
from gwpopulation.models.spin import iid_spin
from gwpopulation.conversions import convert_to_beta_parameters

posterior = pandas.DataFrame()
for key in 'a_1', 'a_2', 'cos_tilt_1', 'cos_tilt_2':
    posterior[key] = np.ones(1)

priors = bilby.prior.PriorDict(conversion_function = lambda args: convert_to_beta_parameters(args)[0])
priors['mu_chi'] = bilby.prior.Uniform(0, 1)
priors['sigma_chi'] = bilby.prior.Uniform(0, 1)
priors['xi_spin'] = bilby.prior.Uniform(0, 1)
priors['sigma_spin'] = bilby.prior.Uniform(0, 1)
priors['amax'] = 1

likelihood = HyperparameterLikelihood(
    posteriors = [posterior],
    hyper_prior = iid_spin,
    conversion_function = convert_to_beta_parameters,
    cupy = False,
)

result = bilby.run_sampler(likelihood, priors = priors)

This results in a key error for amax.

Following the traceback, the error is caused when the sampler is initialized and the _verify_parameters method runs, which tries to evaluate the likelihood but on only the set of parameters stored in search_parameter_keys. Looking at _initialise_parameters, it seems this does not contain parameters with delta function priors, and so the parameters for this initial likelihood call do not include amax. The error is ultimately triggered when these parameters are passed to the sampling routines in PriorDict, which attempt to call PriorDict.conversion_function, but the amax parameter for convert_to_beta_parameters is missing.

You can verify the issue is with the delta function by changing to, e.g., priors['amax'] = bilby.prior.Uniform(0, 1), which works fine.

A workaround is provided by gwpopulation_pipe.utils.prior_conversion when given as the conversion_function for the PriorDict above, as it manually populates amax* parameters to prevent this error.

[ColmTalbot]

I'm just seeing this issue now. For my reference, here is the traceback I get for this example.

I think a possible fix would be to include any parameters with delta function priors in the call to evaluate the constraints.

08:44 bilby INFO    : No prior values provided, defaulting to 1.
08:44 bilby INFO    : Running for label 'label', output will be saved to 'outdir'
08:44 bilby INFO    : Analysis priors:
08:44 bilby INFO    : mu_chi=Uniform(minimum=0, maximum=1, name=None, latex_label=None, unit=None, boundary=None)
08:44 bilby INFO    : sigma_chi=Uniform(minimum=0, maximum=1, name=None, latex_label=None, unit=None, boundary=None)
08:44 bilby INFO    : xi_spin=Uniform(minimum=0, maximum=1, name=None, latex_label=None, unit=None, boundary=None)
08:44 bilby INFO    : sigma_spin=Uniform(minimum=0, maximum=1, name=None, latex_label=None, unit=None, boundary=None)
08:44 bilby INFO    : amax=1
08:44 bilby INFO    : Analysis likelihood class: <class 'gwpopulation.hyperpe.HyperparameterLikelihood'>
08:44 bilby INFO    : Analysis likelihood noise evidence: nan
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[1], line 26
     17 priors['amax'] = 1
     19 likelihood = HyperparameterLikelihood(
     20     posteriors = [posterior],
     21     hyper_prior = iid_spin,
     22     conversion_function = convert_to_beta_parameters,
     23     cupy = False,
     24 )
---> 26 result = bilby.run_sampler(likelihood, priors = priors)

File ~/.conda/envs/o4-population/lib/python3.9/site-packages/bilby/core/sampler/__init__.py:190, in run_sampler(likelihood, priors, label, outdir, sampler, use_ratio, injection_parameters, conversion_function, plot, default_priors_file, clean, meta_data, save, gzip, result_class, npool, **kwargs)
    188 if sampler.lower() in IMPLEMENTED_SAMPLERS:
    189     sampler_class = IMPLEMENTED_SAMPLERS[sampler.lower()]
--> 190     sampler = sampler_class(
    191         likelihood,
    192         priors=priors,
    193         outdir=outdir,
    194         label=label,
    195         injection_parameters=injection_parameters,
    196         meta_data=meta_data,
    197         use_ratio=use_ratio,
    198         plot=plot,
    199         result_class=result_class,
    200         npool=npool,
    201         **kwargs,
    202     )
    203 else:
    204     print(IMPLEMENTED_SAMPLERS)

File ~/.conda/envs/o4-population/lib/python3.9/site-packages/bilby/core/sampler/dynesty.py:211, in Dynesty.__init__(self, likelihood, priors, outdir, label, use_ratio, plot, skip_import_verification, check_point, check_point_plot, n_check_point, check_point_delta_t, resume, nestcheck, exit_code, print_method, maxmcmc, nact, naccept, rejection_sample_posterior, proposals, **kwargs)
    209 self.print_method = print_method
    210 self._translate_kwargs(kwargs)
--> 211 super(Dynesty, self).__init__(
    212     likelihood=likelihood,
    213     priors=priors,
    214     outdir=outdir,
    215     label=label,
    216     use_ratio=use_ratio,
    217     plot=plot,
    218     skip_import_verification=skip_import_verification,
    219     exit_code=exit_code,
    220     **kwargs,
    221 )
    222 self.n_check_point = n_check_point
    223 self.check_point = check_point

File ~/.conda/envs/o4-population/lib/python3.9/site-packages/bilby/core/sampler/base_sampler.py:252, in Sampler.__init__(self, likelihood, priors, outdir, label, use_ratio, plot, skip_import_verification, injection_parameters, meta_data, result_class, likelihood_benchmark, soft_init, exit_code, npool, **kwargs)
    249 self.exit_code = exit_code
    251 if not soft_init:
--> 252     self._verify_parameters()
    253     self._time_likelihood()
    254     self._verify_use_ratio()

File ~/.conda/envs/o4-population/lib/python3.9/site-packages/bilby/core/sampler/base_sampler.py:417, in Sampler._verify_parameters(self)
    412 if self.priors.test_has_redundant_keys():
    413     raise IllegalSamplingSetError(
    414         "Your sampling set contains redundant parameters."
    415     )
--> 417 theta = self.priors.sample_subset_constrained_as_array(
    418     self.search_parameter_keys, size=1
    419 )[:, 0]
    420 try:
    421     self.log_likelihood(theta)

File ~/.conda/envs/o4-population/lib/python3.9/site-packages/bilby/core/prior/dict.py:398, in PriorDict.sample_subset_constrained_as_array(self, keys, size)
    383 def sample_subset_constrained_as_array(self, keys=iter([]), size=None):
    384     """Return an array of samples
    385 
    386     Parameters
   (...)
    396         An array of shape (len(key), size) of the samples (ordered by keys)
    397     """
--> 398     samples_dict = self.sample_subset_constrained(keys=keys, size=size)
    399     samples_dict = {key: np.atleast_1d(val) for key, val in samples_dict.items()}
    400     samples_list = [samples_dict[key] for key in keys]

File ~/.conda/envs/o4-population/lib/python3.9/site-packages/bilby/core/prior/dict.py:450, in PriorDict.sample_subset_constrained(self, keys, size)
    448     while True:
    449         sample = self.sample_subset(keys=keys, size=size)
--> 450         if self.evaluate_constraints(sample):
    451             return sample
    452 else:

File ~/.conda/envs/o4-population/lib/python3.9/site-packages/bilby/core/prior/dict.py:57, in PriorDict.evaluate_constraints(self, sample)
     56 def evaluate_constraints(self, sample):
---> 57     out_sample = self.conversion_function(sample)
     58     prob = 1
     59     for key in self:

Cell In[1], line 12, in <lambda>(args)
      9 for key in 'a_1', 'a_2', 'cos_tilt_1', 'cos_tilt_2':
     10     posterior[key] = np.ones(1)
---> 12 priors = bilby.prior.PriorDict(conversion_function = lambda args: convert_to_beta_parameters(args)[0])
     13 priors['mu_chi'] = bilby.prior.Uniform(0, 1)
     14 priors['sigma_chi'] = bilby.prior.Uniform(0, 1)

File ~/modules/o4-population/gwpopulation/gwpopulation/conversions.py:63, in convert_to_beta_parameters(parameters, remove)
     61     done = done or _done
     62 if not done:
---> 63     _ = _convert("")
     65 return converted, added_keys

File ~/modules/o4-population/gwpopulation/gwpopulation/conversions.py:48, in convert_to_beta_parameters.<locals>._convert(suffix)
     45 if mu in converted.keys() and sigma in converted.keys():
     46     done = True
     47     converted[alpha], converted[beta], _, = mu_var_max_to_alpha_beta_max(
---> 48         parameters[mu], parameters[sigma], parameters[amax]
     49     )
     50     if remove:
     51         added_keys.append(alpha)

KeyError: 'amax'

[mdmould]

This just popped up for me again because I was using the beta spin magnitude models. My general fix was to replace the conversion function from gwpopulation like:

def get_conversion(conversion = lambda pars: (pars, None), fixed = {}):
    return lambda pars: conversion({**pars, **fixed})

fixed = dict(amax = 1)
conversion_function = get_conversion(convert_to_beta_parameters, fixed)

I'm not sure where/how this would need to be added to make it work generally.

[ColmTalbot]

You may have moved on from this, but I think what you suggested is sensible. I think this is worth addressing in bilby, we should probably be picking up the fixed parameters. The difference between prior conversions, likelihood conversions, and post conversions is pretty annoying, but might be easily resolvable. I've also just removed the tracking of added keys as Bilby is moving away from storing parameter values inside the likelihood.

I think the code you have won't do what you expect anyway as you haven't specified (the presumably intended) constraints on the alpha and beta parameters.

import numpy as np
import pandas
import bilby
from gwpopulation.hyperpe import HyperparameterLikelihood
from gwpopulation.models.spin import iid_spin
from gwpopulation.conversions import convert_to_beta_parameters

posterior = pandas.DataFrame()
for key in 'a_1', 'a_2', 'cos_tilt_1', 'cos_tilt_2':
    posterior[key] = np.ones(1)


def get_conversion(conversion = lambda pars: (pars, None), fixed = {}):
    return lambda pars: conversion({**pars, **fixed})[0]


fixed = dict(amax = 1)
conversion_function = get_conversion(convert_to_beta_parameters, fixed)

priors = bilby.prior.PriorDict(conversion_function=conversion_function)
priors['mu_chi'] = bilby.prior.Uniform(0, 1)
priors['sigma_chi'] = bilby.prior.Uniform(0, 1)
priors["alpha_chi"] = bilby.prior.Constraint(0, np.inf)
priors["beta_chi"] = bilby.prior.Constraint(0, np.inf)
priors['xi_spin'] = bilby.prior.Uniform(0, 1)
priors['sigma_spin'] = bilby.prior.Uniform(0, 1)
priors['amax'] = bilby.prior.DeltaFunction(1)

likelihood = HyperparameterLikelihood(
    posteriors = [posterior],
    hyper_prior = iid_spin,
    conversion_function = convert_to_beta_parameters,
)

result = bilby.run_sampler(likelihood, priors = priors)

[mdmould]

Ah yes, the original intention is to convert to the alpha--beta parametrization when need to constrain a non-singular beta distribution, but I didn't include those constraints in the example.

When I've needed to do this I've either just used the above modified conversion function or hard coded anything necessary in custom population models.