From 89274ce4c1971658a9c6245c03c00aeeaf1b7d50 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Fri, 18 Sep 2026 12:04:20 -0400 Subject: [PATCH 1/3] Make sure the ROQ waveform generator wrapper is picklable --- bilby/gw/likelihood/roq.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/bilby/gw/likelihood/roq.py b/bilby/gw/likelihood/roq.py index 61b0ca95a..60c50d65f 100644 --- a/bilby/gw/likelihood/roq.py +++ b/bilby/gw/likelihood/roq.py @@ -22,18 +22,31 @@ class _WaveformGeneratorWrapper: """ def __init__(self, waveform_generator, update_basis): - object.__setattr__(self, '_wrapped', waveform_generator) - object.__setattr__(self, '_update_basis', update_basis) + self._wrapped = waveform_generator + self._update_basis = update_basis def frequency_domain_strain(self, parameters=None, **kwargs): self._update_basis(parameters) return self._wrapped.frequency_domain_strain(parameters, **kwargs) + def __setstate__(self, d): + # since we overrode __setattr__ just doing self.__dict__ = d + # won't work + for k, v in d.items(): + setattr(self, k, v) + def __getattr__(self, name): return getattr(self._wrapped, name) def __setattr__(self, name, value): - setattr(self._wrapped, name, value) + if name in ["_wrapped", "_update_basis"]: + object.__setattr__(self, name, value) + elif not hasattr(self, "_wrapped"): + raise AttributeError( + "ROQ waveform generator wrapper not properly initialized" + ) + else: + setattr(self._wrapped, name, value) class ROQGravitationalWaveTransient(GravitationalWaveTransient): From d2069f7217fec63356543b711bb4a045e67a1000 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 5 Oct 2026 09:29:12 -0400 Subject: [PATCH 2/3] Add test of pickleability --- test/gw/likelihood/roq_test.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/test/gw/likelihood/roq_test.py b/test/gw/likelihood/roq_test.py index 34fa5504f..d16db2f17 100644 --- a/test/gw/likelihood/roq_test.py +++ b/test/gw/likelihood/roq_test.py @@ -8,6 +8,7 @@ import array_api_compat as aac import bilby +import dill import h5py import numpy as np import pytest @@ -285,6 +286,14 @@ def test_create_roq_weights_fails_due_to_min_frequency(self): priors=self.priors, ) + def test_roq_likelihood_pickles(self): + # the backend waveform generator is not pickleable... + self.roq.waveform_generator = self.roq.waveform_generator.wfg + dill.pickles(self.roq) + self.roq.waveform_generator = BackendWaveformGenerator( + self.roq.waveform_generator, self.xp + ) + def test_create_roq_weights_fails_due_to_duration(self): ifos = bilby.gw.detector.InterferometerList(["H1"]) ifos.set_strain_data_from_power_spectral_densities( From 3c259d09954daefb5aa5afbd679bd8128665f263 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Thu, 5 Nov 2026 12:31:34 -0500 Subject: [PATCH 3/3] STY: apply ruff formatting --- bilby/__init__.py | 3 +- bilby/bilby_mcmc/chain.py | 2 +- bilby/bilby_mcmc/flows.py | 1 - bilby/bilby_mcmc/proposals.py | 1 - bilby/bilby_mcmc/sampler.py | 3 - bilby/compat/pytrees/__init__.py | 2 +- bilby/compat/pytrees/likelihood.py | 34 +- bilby/compat/pytrees/prior.py | 12 +- bilby/compat/utils.py | 16 +- bilby/core/fisher.py | 24 +- bilby/core/grid.py | 160 +- bilby/core/likelihood.py | 174 +- bilby/core/prior/analytical.py | 582 ++++--- bilby/core/prior/base.py | 110 +- bilby/core/prior/conditional.py | 120 +- bilby/core/prior/dict.py | 67 +- bilby/core/prior/interpolated.py | 83 +- bilby/core/prior/joint.py | 42 +- bilby/core/prior/slabspike.py | 42 +- bilby/core/result.py | 1017 +++++++----- bilby/core/sampler/base_sampler.py | 1 - bilby/core/sampler/dynesty.py | 3 +- bilby/core/sampler/fake_sampler.py | 4 +- bilby/core/sampler/proposal.py | 2 +- bilby/core/series.py | 72 +- bilby/core/utils/calculus.py | 14 +- bilby/core/utils/cmd.py | 89 +- bilby/core/utils/colors.py | 8 +- bilby/core/utils/constants.py | 2 +- bilby/core/utils/conversion.py | 2 +- bilby/core/utils/counter.py | 3 +- bilby/core/utils/docs.py | 2 + bilby/core/utils/entry_points.py | 5 +- bilby/core/utils/env.py | 10 +- bilby/core/utils/introspection.py | 29 +- bilby/core/utils/io.py | 56 +- bilby/core/utils/log.py | 58 +- bilby/core/utils/meta_data.py | 4 +- bilby/core/utils/plotting.py | 7 +- bilby/core/utils/random.py | 13 +- bilby/core/utils/samples.py | 29 +- bilby/core/utils/series.py | 32 +- bilby/gw/__init__.py | 15 +- bilby/gw/compat/__init__.py | 2 +- bilby/gw/compat/cython.py | 4 +- bilby/gw/compat/pytrees.py | 21 +- bilby/gw/conversion.py | 1446 +++++++++++------ bilby/gw/cosmology.py | 9 +- bilby/gw/detector/__init__.py | 250 +-- bilby/gw/detector/calibration.py | 165 +- bilby/gw/detector/geometry.py | 87 +- bilby/gw/detector/interferometer.py | 535 +++--- bilby/gw/detector/networks.py | 106 +- bilby/gw/detector/psd.py | 231 ++- bilby/gw/detector/strain_data.py | 458 ++++-- bilby/gw/eos/__init__.py | 3 +- bilby/gw/eos/eos.py | 425 +++-- bilby/gw/eos/tov_solver.py | 94 +- bilby/gw/geometry.py | 22 +- bilby/gw/likelihood/__init__.py | 11 +- bilby/gw/likelihood/base.py | 553 ++++--- bilby/gw/likelihood/basic.py | 49 +- bilby/gw/likelihood/multiband.py | 410 +++-- bilby/gw/likelihood/relative.py | 178 +- bilby/gw/likelihood/roq.py | 877 ++++++---- bilby/gw/prior.py | 764 ++++++--- bilby/gw/result.py | 621 ++++--- bilby/gw/sampler/proposal.py | 18 +- bilby/gw/source.py | 1134 +++++++++---- bilby/gw/utils.py | 583 +++++-- bilby/gw/waveform_generator.py | 326 ++-- bilby/hyper/compat/pytrees.py | 8 +- bilby/hyper/likelihood.py | 57 +- bilby/hyper/model.py | 5 +- cli_bilby/bilby_result.py | 5 +- cli_bilby/plot_multiple_posteriors.py | 49 +- docs/conf.py | 92 +- .../linear_regression_bilby_mcmc.py | 1 + .../linear_regression_pymc.py | 1 + examples/core_examples/gaussian_example.py | 1 + examples/core_examples/grid_example.py | 1 + .../core_examples/hyper_parameter_example.py | 1 + examples/core_examples/linear_regression.py | 1 + .../core_examples/linear_regression_grid.py | 1 + .../linear_regression_unknown_noise.py | 1 + .../linear_regression_with_Fisher.py | 1 + examples/core_examples/logo/sample_logo.py | 3 +- .../core_examples/occam_factor_example.py | 4 +- examples/core_examples/radioactive_decay.py | 1 + .../gw_examples/data_examples/GW150914.py | 1 + .../gw_examples/data_examples/GW190425.py | 1 + .../binary_neutron_star_example.py | 1 - .../injection_examples/bns_eos_example.py | 1 - .../bns_polytrope_eos_example.py | 1 - .../bns_spectral_pca_eos_example.py | 1 - .../change_sampled_parameters.py | 1 + .../create_your_own_source_model.py | 1 + .../custom_proposal_example.py | 1 + .../fake_sampler_example.py | 6 +- .../injection_examples/jax_fast_tutorial.py | 1 + .../marginalized_likelihood.py | 1 + .../injection_examples/non_tensor.py | 1 + .../injection_examples/plot_skymap.py | 1 + .../plot_time_domain_data.py | 1 + .../injection_examples/relative_binning.py | 3 +- .../injection_examples/reproduce_mpa1_eos.py | 1 + .../sine_gaussian_example.py | 1 + .../standard_15d_cbc_tutorial.py | 1 + .../supernova_example/supernova_example.py | 1 + examples/tutorials/compare_samplers.ipynb | 2 +- examples/tutorials/conditional_priors.ipynb | 88 +- .../fitting_with_x_and_y_errors.ipynb | 12 +- scripts/check_subject_format.py | 4 +- test/bilby_mcmc/test_chain.py | 36 +- test/bilby_mcmc/test_proposals.py | 17 +- test/bilby_mcmc/test_sampler.py | 15 +- test/check_author_list.py | 31 +- test/check_subject_format_test.py | 58 +- test/conftest.py | 13 +- test/core/grid_test.py | 121 +- test/core/likelihood_test.py | 112 +- test/core/prior/analytical_test.py | 40 +- test/core/prior/base_test.py | 21 +- test/core/prior/conditional_test.py | 24 +- test/core/prior/dict_test.py | 116 +- test/core/prior/joint_test.py | 20 +- test/core/prior/prior_test.py | 118 +- test/core/prior/slabspike_test.py | 107 +- test/core/result_test.py | 182 ++- test/core/sampler/base_sampler_test.py | 27 +- test/core/sampler/dynesty_test.py | 16 +- test/core/sampler/general_sampler_tests.py | 3 +- .../core/sampler/implemented_samplers_test.py | 10 +- test/core/sampler/proposal_test.py | 67 +- test/core/series_test.py | 12 +- test/core/utils_test.py | 84 +- test/gw/conversion_test.py | 440 ++--- test/gw/cosmology_test.py | 8 +- test/gw/detector/calibration_test.py | 22 +- test/gw/detector/geometry_test.py | 24 +- test/gw/detector/injection_test.py | 4 +- test/gw/detector/interferometer_test.py | 143 +- test/gw/detector/networks_test.py | 24 +- test/gw/detector/psd_test.py | 24 +- test/gw/detector/strain_data_test.py | 21 +- test/gw/eos/eos_test.py | 86 +- test/gw/likelihood/base_test.py | 28 +- test/gw/likelihood/basic_test.py | 12 +- test/gw/likelihood/jit_test.py | 20 +- test/gw/likelihood/marginalization_test.py | 107 +- test/gw/likelihood/multiband_test.py | 241 +-- test/gw/likelihood/relative_binning_test.py | 85 +- test/gw/likelihood/roq_test.py | 326 ++-- test/gw/likelihood/utils.py | 5 +- test/gw/plot_test.py | 32 +- test/gw/prior_test.py | 47 +- test/gw/result_test.py | 26 +- test/gw/source_test.py | 86 +- test/gw/utils_test.py | 26 +- test/gw/waveform_generator_test.py | 185 ++- test/hyper/hyper_pe_test.py | 13 +- test/import_test.py | 11 +- test/integration/example_test.py | 4 +- test/integration/noise_realisation_test.py | 2 +- test/integration/other_test.py | 2 +- .../integration/sample_from_the_prior_test.py | 2 +- test/integration/sampler_run_test.py | 17 +- test/integration/test_waveforms.py | 9 +- 168 files changed, 10379 insertions(+), 5656 deletions(-) diff --git a/bilby/__init__.py b/bilby/__init__.py index 1c8a20222..387467e04 100644 --- a/bilby/__init__.py +++ b/bilby/__init__.py @@ -14,6 +14,7 @@ https://bilby-dev.github.io/bilby/installation.html. """ + import logging # Avoid configuring logging for applications that import bilby. @@ -29,4 +30,4 @@ try: from ._version import version as __version__ except ModuleNotFoundError: # development mode - __version__ = 'unknown' + __version__ = "unknown" diff --git a/bilby/bilby_mcmc/chain.py b/bilby/bilby_mcmc/chain.py index 4b78a3044..b09944510 100644 --- a/bilby/bilby_mcmc/chain.py +++ b/bilby/bilby_mcmc/chain.py @@ -386,7 +386,7 @@ def plot(self, outdir=".", label="label", priors=None, all_samples=None): position_indexes = np.arange(self.position + 1) # Plot the traceplots - for (start, stop, thin, color, alpha, ms) in plot_setups: + for start, stop, thin, color, alpha, ms in plot_setups: for ax, key in zip(axes[:, 0], self.keys): xx = position_indexes[start:stop:thin] / K yy = self.get_1d_array(key)[start:stop:thin] diff --git a/bilby/bilby_mcmc/flows.py b/bilby/bilby_mcmc/flows.py index b08ea3a93..26f7883ba 100644 --- a/bilby/bilby_mcmc/flows.py +++ b/bilby/bilby_mcmc/flows.py @@ -44,7 +44,6 @@ def __init__( batch_norm_between_layers=False, random_permutation=True, ): - if use_volume_preserving: coupling_constructor = AdditiveCouplingTransform else: diff --git a/bilby/bilby_mcmc/proposals.py b/bilby/bilby_mcmc/proposals.py index 09215f17e..e40304c76 100644 --- a/bilby/bilby_mcmc/proposals.py +++ b/bilby/bilby_mcmc/proposals.py @@ -134,7 +134,6 @@ def apply_reflective_boundary(self, key, val): return minimum + width * val_normalised_reflected def __call__(self, chain, likelihood=None, priors=None): - if getattr(self, "needs_likelihood_and_priors", False): sample, log_factor = self.propose(chain, likelihood, priors) else: diff --git a/bilby/bilby_mcmc/sampler.py b/bilby/bilby_mcmc/sampler.py index a0529c816..bcdfffc17 100644 --- a/bilby/bilby_mcmc/sampler.py +++ b/bilby/bilby_mcmc/sampler.py @@ -183,7 +183,6 @@ def __init__( normalize_prior=True, **kwargs, ): - super(Bilby_MCMC, self).__init__( likelihood=likelihood, priors=priors, @@ -304,7 +303,6 @@ def setup_chain_set(self): self.init_ptsampler() def init_ptsampler(self): - logger.info(f"Initializing BilbyPTMCMCSampler with:\n{self.get_setup_string()}") self.ptsampler = BilbyPTMCMCSampler( convergence_inputs=self.convergence_inputs, @@ -643,7 +641,6 @@ def get_initial_betas(self): return betas def setup_sampler_dictionary(self, convergence_inputs, proposal_cycle): - betas = self.get_initial_betas() logger.info( f"Initializing BilbyPTMCMCSampler with:" diff --git a/bilby/compat/pytrees/__init__.py b/bilby/compat/pytrees/__init__.py index 04e632b07..ad126b00b 100644 --- a/bilby/compat/pytrees/__init__.py +++ b/bilby/compat/pytrees/__init__.py @@ -2,4 +2,4 @@ from . import likelihood from . import prior -from . import utils \ No newline at end of file +from . import utils diff --git a/bilby/compat/pytrees/likelihood.py b/bilby/compat/pytrees/likelihood.py index dd1820c69..d4bf72046 100644 --- a/bilby/compat/pytrees/likelihood.py +++ b/bilby/compat/pytrees/likelihood.py @@ -128,7 +128,7 @@ def joint_likelihood_unflatten(aux_data, flat) -> JointLikelihood: def analytical_multidimensional_covariant_gaussian_flatten( - likelihood: AnalyticalMultidimensionalCovariantGaussian + likelihood: AnalyticalMultidimensionalCovariantGaussian, ): _, aux_data = likelihood_flatten(likelihood) children = (likelihood.cov, likelihood.mean, likelihood.sigma) @@ -143,12 +143,14 @@ def analytical_multidimensional_covariant_gaussian_unflatten( likelihood.cov = cov likelihood.mean = mean likelihood.sigma = sigma - likelihood.logpdf = partial(jax.scipy.stats.multivariate_normal.logpdf, mean=mean, cov=cov) + likelihood.logpdf = partial( + jax.scipy.stats.multivariate_normal.logpdf, mean=mean, cov=cov + ) return likelihood def analytical_multidimensional_bimodal_covariant_gaussian_flatten( - likelihood: AnalyticalMultidimensionalBimodalCovariantGaussian + likelihood: AnalyticalMultidimensionalBimodalCovariantGaussian, ): _, aux_data = likelihood_flatten(likelihood) children = (likelihood.cov, likelihood.mean_1, likelihood.mean_2, likelihood.sigma) @@ -164,8 +166,12 @@ def analytical_multidimensional_bimodal_covariant_gaussian_unflatten( likelihood.mean_1 = mean_1 likelihood.mean_2 = mean_2 likelihood.sigma = sigma - likelihood.logpdf_1 = partial(jax.scipy.stats.multivariate_normal.logpdf, mean=mean_1, cov=cov) - likelihood.logpdf_2 = partial(jax.scipy.stats.multivariate_normal.logpdf, mean=mean_2, cov=cov) + likelihood.logpdf_1 = partial( + jax.scipy.stats.multivariate_normal.logpdf, mean=mean_1, cov=cov + ) + likelihood.logpdf_2 = partial( + jax.scipy.stats.multivariate_normal.logpdf, mean=mean_2, cov=cov + ) return likelihood @@ -173,9 +179,21 @@ def analytical_multidimensional_bimodal_covariant_gaussian_unflatten( (Likelihood, likelihood_flatten, likelihood_unflatten), (GaussianLikelihood, gaussian_likelihood_flatten, gaussian_likelihood_unflatten), (ZeroLikelihood, zero_likelihood_flatten, zero_likelihood_unflatten), - (Analytical1DLikelihood, analytical_1d_likelihood_flatten, analytical_1d_likelihood_unflatten), - (PoissonLikelihood, analytical_1d_likelihood_flatten, analytical_1d_likelihood_unflatten), - (ExponentialLikelihood, analytical_1d_likelihood_flatten, analytical_1d_likelihood_unflatten), + ( + Analytical1DLikelihood, + analytical_1d_likelihood_flatten, + analytical_1d_likelihood_unflatten, + ), + ( + PoissonLikelihood, + analytical_1d_likelihood_flatten, + analytical_1d_likelihood_unflatten, + ), + ( + ExponentialLikelihood, + analytical_1d_likelihood_flatten, + analytical_1d_likelihood_unflatten, + ), (StudentTLikelihood, student_t_likelihood_flatten, student_t_likelihood_unflatten), (Multinomial, multinomial_flatten, multinomial_unflatten), (JointLikelihood, joint_likelihood_flatten, joint_likelihood_unflatten), diff --git a/bilby/compat/pytrees/prior.py b/bilby/compat/pytrees/prior.py index 79da25d0f..c2be2e2ab 100644 --- a/bilby/compat/pytrees/prior.py +++ b/bilby/compat/pytrees/prior.py @@ -69,9 +69,7 @@ def conditional_unflatten(aux_data, children) -> ConditionalBasePrior: def dict_flatten(prior_dict: PriorDict): prior_dict.convert_floats_to_delta_functions() - children = ( - {k: v for k, v in prior_dict.items()}, - ) + children = ({k: v for k, v in prior_dict.items()},) aux_data = ( prior_dict.__class__, prior_dict.conversion_function, @@ -141,8 +139,12 @@ def dirichlet_dict_unflatten(aux_data, children): register_pytree_node(PriorDict, dict_flatten, dict_unflatten) -register_pytree_node(ConditionalPriorDict, conditional_dict_flatten, conditional_dict_unflatten) -register_pytree_node(DirichletPriorDict, dirichlet_dict_flatten, dirichlet_dict_unflatten) +register_pytree_node( + ConditionalPriorDict, conditional_dict_flatten, conditional_dict_unflatten +) +register_pytree_node( + DirichletPriorDict, dirichlet_dict_flatten, dirichlet_dict_unflatten +) register_pytree_node(CBCPriorDict, conditional_dict_flatten, conditional_dict_unflatten) register_pytree_node(BBHPriorDict, conditional_dict_flatten, conditional_dict_unflatten) register_pytree_node(BNSPriorDict, conditional_dict_flatten, conditional_dict_unflatten) diff --git a/bilby/compat/utils.py b/bilby/compat/utils.py index 4876cedbf..8dc9bd3d5 100644 --- a/bilby/compat/utils.py +++ b/bilby/compat/utils.py @@ -102,10 +102,12 @@ def array_module(arr): # FIXME: remove direct import of orng to avoid hard dependency import orng + if isinstance(arr, orng.RandomGenerator): match arr.backend: case "jax": import jax.numpy as jnp + return jnp case _: return np @@ -116,7 +118,9 @@ def array_module(arr): except TypeError: if isinstance(arr, dict): try: - return aac.array_namespace(*[val for val in arr.values() if not isinstance(val, str)]) + return aac.array_namespace( + *[val for val in arr.values() if not isinstance(val, str)] + ) except TypeError: return np elif arr.__class__.__module__ == "builtins" and isinstance(arr, Iterable): @@ -167,9 +171,10 @@ def promote_to_array(args, xp, skip=None): else: skip = len(args) - skip if not aac.is_numpy_namespace(xp): - args = tuple( - xp.asarray(arg, device=BILBY_DEVICE) for arg in args[:skip] - ) + args[skip:] + args = ( + tuple(xp.asarray(arg, device=BILBY_DEVICE) for arg in args[:skip]) + + args[skip:] + ) return args @@ -191,6 +196,7 @@ def xp_wrap(func, no_xp=False): function The decorated function. """ + def parse_args_kwargs_for_xp(*args, xp=None, **kwargs): if not no_xp and xp is None: try: @@ -217,10 +223,12 @@ def parse_args_kwargs_for_xp(*args, xp=None, **kwargs): sig = inspect.signature(func) if any(name in sig.parameters for name in ("self", "cls")): + def wrapped(self, *args, xp=None, **kwargs): args, kwargs = parse_args_kwargs_for_xp(*args, xp=xp, **kwargs) return func(self, *args, **kwargs) else: + def wrapped(*args, xp=None, **kwargs): args, kwargs = parse_args_kwargs_for_xp(*args, xp=xp, **kwargs) return func(*args, **kwargs) diff --git a/bilby/core/fisher.py b/bilby/core/fisher.py index 7651a9d0c..2509b4d62 100644 --- a/bilby/core/fisher.py +++ b/bilby/core/fisher.py @@ -5,8 +5,10 @@ class FisherMatrixPosteriorEstimator(object): - def __init__(self, likelihood, priors, parameters=None, fd_eps=1e-6, n_prior_samples=100): - """ A class to estimate posteriors using the Fisher Matrix approach + def __init__( + self, likelihood, priors, parameters=None, fd_eps=1e-6, n_prior_samples=100 + ): + """A class to estimate posteriors using the Fisher Matrix approach Parameters ---------- @@ -35,7 +37,9 @@ def __init__(self, likelihood, priors, parameters=None, fd_eps=1e-6, n_prior_sam self.prior_samples = [ priors.sample_subset(self.parameter_names) for _ in range(n_prior_samples) ] - self.prior_bounds = [(priors[key].minimum, priors[key].maximum) for key in self.parameter_names] + self.prior_bounds = [ + (priors[key].minimum, priors[key].maximum) for key in self.parameter_names + ] self.prior_width_dict = {} for key in self.parameter_names: @@ -91,13 +95,13 @@ def get_finite_difference_xx(self, sample, ii): p = self.shift_sample_x(sample, ii, 1) m = self.shift_sample_x(sample, ii, -1) - dx = .5 * (p[ii] - m[ii]) + dx = 0.5 * (p[ii] - m[ii]) loglp = self.log_likelihood(p) logl = self.log_likelihood(sample) loglm = self.log_likelihood(m) - return (loglp - 2 * logl + loglm) / dx ** 2 + return (loglp - 2 * logl + loglm) / dx**2 def get_finite_difference_xy(self, sample, ii, jj): # Sample grid @@ -106,8 +110,8 @@ def get_finite_difference_xy(self, sample, ii, jj): mp = self.shift_sample_xy(sample, ii, -1, jj, 1) mm = self.shift_sample_xy(sample, ii, -1, jj, -1) - dx = .5 * (pp[ii] - mm[ii]) - dy = .5 * (pp[jj] - mm[jj]) + dx = 0.5 * (pp[ii] - mm[ii]) + dy = 0.5 * (pp[jj] - mm[jj]) loglpp = self.log_likelihood(pp) loglpm = self.log_likelihood(pm) @@ -117,7 +121,6 @@ def get_finite_difference_xy(self, sample, ii, jj): return (loglpp - loglpm - loglmp + loglmm) / (4 * dx * dy) def shift_sample_x(self, sample, x_key, x_coef): - vx = sample[x_key] dvx = self.fd_eps * self.prior_width_dict[x_key] @@ -127,7 +130,6 @@ def shift_sample_x(self, sample, x_key, x_coef): return shift_sample def shift_sample_xy(self, sample, x_key, x_coef, y_key, y_coef): - vx = sample[x_key] vy = sample[y_key] @@ -140,7 +142,7 @@ def shift_sample_xy(self, sample, x_key, x_coef, y_key, y_coef): return shift_sample def get_maximum_likelihood_sample(self, initial_sample=None): - """ A method to attempt optimization of the maximum likelihood + """A method to attempt optimization of the maximum likelihood This uses a simple scipy optimization approach, starting from a number of draws from the prior to avoid problems with local optimization. @@ -158,7 +160,7 @@ def get_maximum_likelihood_sample(self, initial_sample=None): def neg_log_like(x, self, T=1): sample = {key: val for key, val in zip(self.parameter_names, x)} - return - 1 / T * self.log_likelihood(sample) + return -1 / T * self.log_likelihood(sample) out = minimize( neg_log_like, diff --git a/bilby/core/grid.py b/bilby/core/grid.py index 0256a5177..add4ec467 100644 --- a/bilby/core/grid.py +++ b/bilby/core/grid.py @@ -6,15 +6,19 @@ from .prior import Prior, PriorDict from .utils import ( - logtrapzexp, check_directory_exists_and_if_not_mkdir, logger, - BilbyJsonEncoder, load_json, move_old_file + logtrapzexp, + check_directory_exists_and_if_not_mkdir, + logger, + BilbyJsonEncoder, + load_json, + move_old_file, ) from .result import FileMovedError from ..compat.utils import array_module, copy_array def grid_file_name(outdir, label, gzip=False): - """ Returns the standard filename used for a grid file + """Returns the standard filename used for a grid file Parameters ========== @@ -30,16 +34,21 @@ def grid_file_name(outdir, label, gzip=False): str: File name of the output file """ if gzip: - return os.path.join(outdir, '{}_grid.json.gz'.format(label)) + return os.path.join(outdir, "{}_grid.json.gz".format(label)) else: - return os.path.join(outdir, '{}_grid.json'.format(label)) + return os.path.join(outdir, "{}_grid.json".format(label)) class Grid(object): - def __init__( - self, likelihood=None, priors=None, grid_size=101, - save=False, label='no_label', outdir='.', gzip=False, + self, + likelihood=None, + priors=None, + grid_size=101, + save=False, + label="no_label", + outdir=".", + gzip=False, xp=None, ): """ @@ -84,9 +93,12 @@ def __init__( # evaluate the prior on the grid points if self.n_dims > 0: self._ln_prior = self.priors.ln_prob( - {key: self.mesh_grid[i].flatten() for i, key in - enumerate(self.parameter_names)}, axis=0).reshape( - self.mesh_grid[0].shape) + { + key: self.mesh_grid[i].flatten() + for i, key in enumerate(self.parameter_names) + }, + axis=0, + ).reshape(self.mesh_grid[0].shape) else: self._ln_prior = xp.asarray(0.0) self._ln_likelihood = None @@ -197,8 +209,7 @@ def _marginalize_single(self, log_array, name, non_marg_names=None): """ if name not in self.parameter_names: - raise ValueError("'{}' is not a recognised " - "parameter".format(name)) + raise ValueError("'{}' is not a recognised parameter".format(name)) if non_marg_names is None: non_marg_names = list(self.parameter_names) @@ -215,13 +226,15 @@ def _marginalize_single(self, log_array, name, non_marg_names=None): out = logtrapzexp(log_array, dx=dx, xp=xp) elif aac.is_torch_namespace(xp): # https://discuss.pytorch.org/t/apply-a-function-along-an-axis/130440 - out = xp.stack([ - logtrapzexp(x_i, dx=dx, xp=xp) for x_i in xp.unbind(log_array, dim=axis) - ], dim=min(axis, log_array.ndim - 2)) - else: - out = xp.apply_along_axis( - logtrapzexp, axis, log_array, dx + out = xp.stack( + [ + logtrapzexp(x_i, dx=dx, xp=xp) + for x_i in xp.unbind(log_array, dim=axis) + ], + dim=min(axis, log_array.ndim - 2), ) + else: + out = xp.apply_along_axis(logtrapzexp, axis, log_array, dx) else: # no marginalisation required, just remove the singleton dimension z = log_array.shape @@ -260,8 +273,9 @@ def marginalize_ln_likelihood(self, parameters=None, not_parameters=None): array-like: The marginalized ln likelihood. """ - return self.marginalize(self.ln_likelihood, parameters=parameters, - not_parameters=not_parameters) + return self.marginalize( + self.ln_likelihood, parameters=parameters, not_parameters=not_parameters + ) def marginalize_ln_posterior(self, parameters=None, not_parameters=None): """ @@ -281,8 +295,9 @@ def marginalize_ln_posterior(self, parameters=None, not_parameters=None): array-like: The marginalized ln posterior. """ - return self.marginalize(self.ln_posterior, parameters=parameters, - not_parameters=not_parameters) + return self.marginalize( + self.ln_posterior, parameters=parameters, not_parameters=not_parameters + ) def marginalize_likelihood(self, parameters=None, not_parameters=None): """ @@ -302,8 +317,9 @@ def marginalize_likelihood(self, parameters=None, not_parameters=None): array-like: The marginalized likelihood. """ - ln_like = self.marginalize(self.ln_likelihood, parameters=parameters, - not_parameters=not_parameters) + ln_like = self.marginalize( + self.ln_likelihood, parameters=parameters, not_parameters=not_parameters + ) xp = aac.get_namespace(ln_like) # NOTE: the output will not be properly normalised return xp.exp(ln_like - xp.max(ln_like)) @@ -326,8 +342,9 @@ def marginalize_posterior(self, parameters=None, not_parameters=None): array-like: The marginalized posterior. """ - ln_post = self.marginalize(self.ln_posterior, parameters=parameters, - not_parameters=not_parameters) + ln_post = self.marginalize( + self.ln_posterior, parameters=parameters, not_parameters=not_parameters + ) # NOTE: the output will not be properly normalised xp = aac.get_namespace(ln_post) return xp.exp(ln_post - xp.max(ln_post)) @@ -340,7 +357,10 @@ def _evaluate(self): else: from jax import vmap self._ln_likelihood = vmap(self.likelihood.log_likelihood)( - {key: self.mesh_grid[i].flatten() for i, key in enumerate(self.parameter_names)} + { + key: self.mesh_grid[i].flatten() + for i, key in enumerate(self.parameter_names) + } ).reshape(self.mesh_grid[0].shape) else: @@ -351,10 +371,15 @@ def _evaluate(self): def _evaluate_recursion(self, dimension, parameters): if dimension == self.n_dims: xp = aac.get_namespace(self.mesh_grid[0]) - current_point = tuple([[xp.where( - parameters[name] == - self.sample_points[name])[0].item()] for name in self.parameter_names]) - self._ln_likelihood[current_point] = self.likelihood.log_likelihood(parameters) + current_point = tuple( + [ + [xp.where(parameters[name] == self.sample_points[name])[0].item()] + for name in self.parameter_names + ] + ) + self._ln_likelihood[current_point] = self.likelihood.log_likelihood( + parameters + ) else: name = self.parameter_names[dimension] for ii in range(self._ln_likelihood.shape[dimension]): @@ -366,17 +391,20 @@ def _get_sample_points(self, grid_size, *, xp=np): if isinstance(self.priors[key], Prior): if isinstance(grid_size, int): self.sample_points[key] = self.priors[key].rescale( - xp.linspace(0, 1, grid_size)) + xp.linspace(0, 1, grid_size) + ) elif isinstance(grid_size, list): if isinstance(grid_size[ii], int): self.sample_points[key] = self.priors[key].rescale( - xp.linspace(0, 1, grid_size[ii])) + xp.linspace(0, 1, grid_size[ii]) + ) else: self.sample_points[key] = xp.asarray(grid_size[ii]) elif isinstance(grid_size, dict): if isinstance(grid_size[key], int): self.sample_points[key] = self.priors[key].rescale( - xp.linspace(0, 1, grid_size[key])) + xp.linspace(0, 1, grid_size[key]) + ) else: self.sample_points[key] = xp.asarray(grid_size[key]) else: @@ -384,15 +412,22 @@ def _get_sample_points(self, grid_size, *, xp=np): # set the mesh of points self.mesh_grid = xp.meshgrid( - *(self.sample_points[key] for key in self.parameter_names), - indexing='ij') + *(self.sample_points[key] for key in self.parameter_names), indexing="ij" + ) def _get_save_data_dictionary(self): # This list defines all the parameters saved in the grid object save_attrs = [ - 'label', 'outdir', 'parameter_names', 'n_dims', 'priors', - 'sample_points', 'ln_likelihood', 'ln_evidence', - 'ln_noise_evidence'] + "label", + "outdir", + "parameter_names", + "n_dims", + "priors", + "sample_points", + "ln_likelihood", + "ln_evidence", + "ln_noise_evidence", + ] dictionary = dict() for attr in save_attrs: try: @@ -408,14 +443,15 @@ def _safe_outdir_creation(self, outdir=None, caller_func=None): try: check_directory_exists_and_if_not_mkdir(outdir) except PermissionError: - raise FileMovedError("Can not write in the out directory.\n" - "Did you move the here file from another system?\n" - "Try calling " + caller_func.__name__ + " with the 'outdir' " - "keyword argument, e.g. " + caller_func.__name__ + "(outdir='.')") + raise FileMovedError( + "Can not write in the out directory.\n" + "Did you move the here file from another system?\n" + "Try calling " + caller_func.__name__ + " with the 'outdir' " + "keyword argument, e.g. " + caller_func.__name__ + "(outdir='.')" + ) return outdir - def save_to_file(self, filename=None, overwrite=False, outdir=None, - gzip=False): + def save_to_file(self, filename=None, overwrite=False, outdir=None, gzip=False): """ Writes the Grid to a file. @@ -445,22 +481,25 @@ def save_to_file(self, filename=None, overwrite=False, outdir=None, try: dictionary["priors"] = dictionary["priors"]._get_json_dict() - if gzip or (os.path.splitext(filename)[-1] == '.gz'): + if gzip or (os.path.splitext(filename)[-1] == ".gz"): import gzip + # encode to a string - json_str = json.dumps(dictionary, cls=BilbyJsonEncoder).encode('utf-8') - with gzip.GzipFile(filename, 'w') as file: + json_str = json.dumps(dictionary, cls=BilbyJsonEncoder).encode("utf-8") + with gzip.GzipFile(filename, "w") as file: file.write(json_str) else: - with open(filename, 'w') as file: + with open(filename, "w") as file: json.dump(dictionary, file, indent=2, cls=BilbyJsonEncoder) except Exception as e: - logger.error("\n\n Saving the data has failed with the " - "following message:\n {} \n\n".format(e)) + logger.error( + "\n\n Saving the data has failed with the " + "following message:\n {} \n\n".format(e) + ) @classmethod def read(cls, filename=None, outdir=None, label=None, gzip=False, xp=None): - """ Read in a saved .json grid file + """Read in a saved .json grid file Parameters ========== @@ -496,13 +535,18 @@ def read(cls, filename=None, outdir=None, label=None, gzip=False, xp=None): if os.path.isfile(filename): dictionary = load_json(filename, gzip) try: - grid = cls(likelihood=None, priors=dictionary['priors'], - grid_size=dictionary['sample_points'], - label=dictionary['label'], outdir=dictionary['outdir'], xp=xp) + grid = cls( + likelihood=None, + priors=dictionary["priors"], + grid_size=dictionary["sample_points"], + label=dictionary["label"], + outdir=dictionary["outdir"], + xp=xp, + ) # set the likelihood - grid._ln_likelihood = dictionary['ln_likelihood'] - grid.ln_noise_evidence = dictionary['ln_noise_evidence'] + grid._ln_likelihood = dictionary["ln_likelihood"] + grid.ln_noise_evidence = dictionary["ln_noise_evidence"] return grid except TypeError as e: diff --git a/bilby/core/likelihood.py b/bilby/core/likelihood.py index 0778a333c..47283b4d5 100644 --- a/bilby/core/likelihood.py +++ b/bilby/core/likelihood.py @@ -5,13 +5,15 @@ from array_api_compat import is_array_api_obj from scipy.special import gammaln, xlogy -from .utils import infer_parameters_from_function, infer_args_from_function_except_n_args +from .utils import ( + infer_parameters_from_function, + infer_args_from_function_except_n_args, +) from ..compat.patches import multivariate_logpdf from ..compat.utils import BackendNotImplementedError, array_module class Likelihood: - def __init__(self): """Empty likelihood class to be subclassed by other likelihoods""" self._meta_data = None @@ -59,7 +61,7 @@ def log_likelihood_ratio(self, parameters): @property def meta_data(self): - return getattr(self, '_meta_data', None) + return getattr(self, "_meta_data", None) @meta_data.setter def meta_data(self, meta_data): @@ -74,7 +76,7 @@ def marginalized_parameters(self): class ZeroLikelihood(Likelihood): - """ A special test-only class which already returns zero likelihood + """A special test-only class which already returns zero likelihood Parameters ========== @@ -123,30 +125,32 @@ def __init__(self, x, y, func, **kwargs): self.kwargs = kwargs def __repr__(self): - return self.__class__.__name__ + '(x={}, y={}, func={})'.format(self.x, self.y, self.func.__name__) + return self.__class__.__name__ + "(x={}, y={}, func={})".format( + self.x, self.y, self.func.__name__ + ) @property def func(self): - """ Make func read-only """ + """Make func read-only""" return self._func def model_parameters(self, parameters): - """ This sets up the function only parameters (i.e. not sigma for the GaussianLikelihood) """ + """This sets up the function only parameters (i.e. not sigma for the GaussianLikelihood)""" return {key: parameters[key] for key in self.function_keys} @property def function_keys(self): - """ Makes function_keys read_only """ + """Makes function_keys read_only""" return self._function_keys @property def n(self): - """ The number of data points """ + """The number of data points""" return len(self.x) @property def x(self): - """ The independent variable. Setter assures that single numbers will be converted to arrays internally """ + """The independent variable. Setter assures that single numbers will be converted to arrays internally""" return self._x @x.setter @@ -157,7 +161,7 @@ def x(self, x): @property def y(self): - """ The dependent variable. Setter assures that single numbers will be converted to arrays internally """ + """The dependent variable. Setter assures that single numbers will be converted to arrays internally""" return self._y @y.setter @@ -167,8 +171,10 @@ def y(self, y): self._y = y def residual(self, parameters): - """ Residual of the function against the data. """ - return self.y - self.func(self.x, **self.model_parameters(parameters=parameters), **self.kwargs) + """Residual of the function against the data.""" + return self.y - self.func( + self.x, **self.model_parameters(parameters=parameters), **self.kwargs + ) class GaussianLikelihood(Analytical1DLikelihood): @@ -200,13 +206,16 @@ def __init__(self, x, y, func, sigma=None, **kwargs): def log_likelihood(self, parameters): xp = array_module(self.x) sigma = parameters.get("sigma", self.sigma) - log_l = xp.sum(- (self.residual(parameters) / sigma)**2 / 2 - - xp.log(xp.asarray(2 * np.pi * sigma**2)) / 2) + log_l = xp.sum( + -((self.residual(parameters) / sigma) ** 2) / 2 + - xp.log(xp.asarray(2 * np.pi * sigma**2)) / 2 + ) return log_l def __repr__(self): - return self.__class__.__name__ + '(x={}, y={}, func={}, sigma={})' \ - .format(self.x, self.y, self.func.__name__, self.sigma) + return self.__class__.__name__ + "(x={}, y={}, func={}, sigma={})".format( + self.x, self.y, self.func.__name__, self.sigma + ) @property def sigma(self): @@ -227,7 +236,7 @@ def sigma(self, sigma): elif len(sigma) == self.n: self._sigma = sigma else: - raise ValueError('Sigma must be either float or array-like x.') + raise ValueError("Sigma must be either float or array-like x.") class PoissonLikelihood(Analytical1DLikelihood): @@ -256,15 +265,17 @@ def __init__(self, x, y, func, **kwargs): super(PoissonLikelihood, self).__init__(x=x, y=y, func=func, **kwargs) def log_likelihood(self, parameters): - rate = self.func(self.x, **self.model_parameters(parameters=parameters), **self.kwargs) + rate = self.func( + self.x, **self.model_parameters(parameters=parameters), **self.kwargs + ) if not is_array_api_obj(rate): raise ValueError( "Poisson rate function returns wrong value type! " - "Is {} when it should be numpy.ndarray".format(type(rate))) + "Is {} when it should be numpy.ndarray".format(type(rate)) + ) xp = aac.get_namespace(rate) - if not aac.is_jax_namespace(xp) and xp.any(rate < 0.): - raise ValueError(("Poisson rate function returns a negative", - " value!")) + if not aac.is_jax_namespace(xp) and xp.any(rate < 0.0): + raise ValueError(("Poisson rate function returns a negative", " value!")) rate = xp.maximum(rate, xp.asarray(0.0)) return xp.sum(-rate + self.y * xp.log(rate) - gammaln(self.y + 1)) @@ -273,7 +284,7 @@ def __repr__(self): @property def y(self): - """ Property assures that y-value is a positive integer. """ + """Property assures that y-value is a positive integer.""" return self._y @y.setter @@ -283,7 +294,9 @@ def y(self, y): xp = aac.get_namespace(y) # check array is a non-negative integer array # torch doesn't support checking dtype kind - if (not aac.is_torch_namespace(xp) and y.dtype.kind not in 'ui') or xp.any(y < 0): + if (not aac.is_torch_namespace(xp) and y.dtype.kind not in "ui") or xp.any( + y < 0 + ): raise ValueError("Data must be non-negative integers") self._y = y @@ -308,9 +321,13 @@ def __init__(self, x, y, func, **kwargs): super(ExponentialLikelihood, self).__init__(x=x, y=y, func=func, **kwargs) def log_likelihood(self, parameters): - mu = self.func(self.x, **self.model_parameters(parameters=parameters), **self.kwargs) + mu = self.func( + self.x, **self.model_parameters(parameters=parameters), **self.kwargs + ) xp = array_module(mu) - val = xp.nan_to_num(xp.log(mu) + (self.y / mu), nan=np.inf, posinf=np.inf, neginf=-np.inf) + val = xp.nan_to_num( + xp.log(mu) + (self.y / mu), nan=np.inf, posinf=np.inf, neginf=-np.inf + ) return -xp.sum(val) def __repr__(self): @@ -318,7 +335,7 @@ def __repr__(self): @property def y(self): - """ Property assures that y-value is positive. """ + """Property assures that y-value is positive.""" return self._y @y.setter @@ -368,32 +385,42 @@ def __init__(self, x, y, func, nu=None, sigma=1, **kwargs): def log_likelihood(self, parameters): nu = parameters.get("nu", self.nu) - if nu <= 0.: - raise ValueError("Number of degrees of freedom for Student's " - "t-likelihood must be positive") + if nu <= 0.0: + raise ValueError( + "Number of degrees of freedom for Student's " + "t-likelihood must be positive" + ) xp = array_module(self.x) - log_l = xp.sum(xp.nan_to_num( - - (nu + 1) * xp.log1p(self.lam * self.residual(parameters=parameters)**2 / nu) / 2 - + xp.log(xp.asarray(self.lam / (nu * np.pi))) / 2 - + gammaln((nu + 1) / 2) - gammaln(nu / 2), - nan=np.inf, posinf=np.inf, neginf=-np.inf - )) + log_l = xp.sum( + xp.nan_to_num( + -(nu + 1) + * xp.log1p(self.lam * self.residual(parameters=parameters) ** 2 / nu) + / 2 + + xp.log(xp.asarray(self.lam / (nu * np.pi))) / 2 + + gammaln((nu + 1) / 2) + - gammaln(nu / 2), + nan=np.inf, + posinf=np.inf, + neginf=-np.inf, + ) + ) return log_l def __repr__(self): - base_string = '(x={}, y={}, func={}, nu={}, sigma={})' + base_string = "(x={}, y={}, func={}, nu={}, sigma={})" return self.__class__.__name__ + base_string.format( - self.x, self.y, self.func.__name__, self.nu, self.sigma) + self.x, self.y, self.func.__name__, self.nu, self.sigma + ) @property def lam(self): - """ Converts 'scale' to 'precision' """ - return 1. / self.sigma ** 2 + """Converts 'scale' to 'precision'""" + return 1.0 / self.sigma**2 @property def nu(self): - """ This checks if nu or sigma have been set in parameters. If so, those + """This checks if nu or sigma have been set in parameters. If so, those values will be used. Otherwise, the attribute nu is used. The logic is that if nu is not in parameters the attribute is used which was given at init (i.e. the known nu as a float).""" @@ -443,8 +470,7 @@ def log_likelihood(self, parameters): ======= float """ - probs = [parameters[self.base + str(ii)] - for ii in range(self.n - 1)] + probs = [parameters[self.base + str(ii)] for ii in range(self.n - 1)] probs.append(1 - sum(probs)) return self._multinomial_ln_pdf(probs=probs) @@ -461,7 +487,8 @@ def _multinomial_ln_pdf(self, probs): """Lifted from scipy.stats.multinomial._logpdf""" xp = array_module(self.data) ln_prob = gammaln(self._total + 1) + xp.sum( - xlogy(self.data, probs) - gammaln(self.data + 1), axis=-1) + xlogy(self.data, probs) - gammaln(self.data + 1), axis=-1 + ) return ln_prob @@ -562,7 +589,7 @@ def __init__(self, *likelihoods): @property def likelihoods(self): - """ The list of likelihoods """ + """The list of likelihoods""" return self._likelihoods @likelihoods.setter @@ -572,12 +599,16 @@ def likelihoods(self, likelihoods): if all(isinstance(likelihood, Likelihood) for likelihood in likelihoods): self._likelihoods = list(likelihoods) else: - raise ValueError('Try setting the JointLikelihood like this\n' - 'JointLikelihood(first_likelihood, second_likelihood, ...)') + raise ValueError( + "Try setting the JointLikelihood like this\n" + "JointLikelihood(first_likelihood, second_likelihood, ...)" + ) elif isinstance(likelihoods, Likelihood): self._likelihoods = [likelihoods] else: - raise ValueError('Input likelihood is not a list of tuple. You need to set multiple likelihoods.') + raise ValueError( + "Input likelihood is not a list of tuple. You need to set multiple likelihoods." + ) def log_likelihood(self, parameters): """ @@ -592,20 +623,29 @@ def log_likelihood(self, parameters): ======= float """ - return sum([likelihood.log_likelihood(parameters=parameters) for likelihood in self.likelihoods]) + return sum( + [ + likelihood.log_likelihood(parameters=parameters) + for likelihood in self.likelihoods + ] + ) def noise_log_likelihood(self): - """ This is just the sum of the noise likelihoods of all parts of the joint likelihood""" - return sum([likelihood.noise_log_likelihood() for likelihood in self.likelihoods]) + """This is just the sum of the noise likelihoods of all parts of the joint likelihood""" + return sum( + [likelihood.noise_log_likelihood() for likelihood in self.likelihoods] + ) def function_to_celerite_mean_model(func): from celerite.modeling import Model as CeleriteModel + return _function_to_gp_model(func, CeleriteModel) def function_to_george_mean_model(func): from celerite.modeling import Model as GeorgeModel + return _function_to_gp_model(func, GeorgeModel) @@ -624,7 +664,6 @@ def compute_gradient(self, *args, **kwargs): class _GPLikelihood(Likelihood): - def __init__(self, kernel, mean_model, t, y, yerr=1e-6, gp_class=None): """ Basic Gaussian Process likelihood interface for `celerite` and `george`. @@ -654,7 +693,12 @@ def __init__(self, kernel, mean_model, t, y, yerr=1e-6, gp_class=None): self.y = np.array(y) self.yerr = np.array(yerr) self.GPClass = gp_class - self.gp = self.GPClass(kernel=self.kernel, mean=self.mean_model, fit_mean=True, fit_white_noise=True) + self.gp = self.GPClass( + kernel=self.kernel, + mean=self.mean_model, + fit_mean=True, + fit_white_noise=True, + ) self.gp.compute(self.t, yerr=self.yerr) super().__init__() @@ -676,7 +720,6 @@ def set_parameters(self, parameters): class CeleriteLikelihood(_GPLikelihood): - def __init__(self, kernel, mean_model, t, y, yerr=1e-6): """ Basic Gaussian Process likelihood interface for `celerite` and `george`. @@ -698,7 +741,15 @@ def __init__(self, kernel, mean_model, t, y, yerr=1e-6): applies for all y-values. Default is 1e-6, effectively assuming that no y-errors are present. """ import celerite - super().__init__(kernel=kernel, mean_model=mean_model, t=t, y=y, yerr=yerr, gp_class=celerite.GP) + + super().__init__( + kernel=kernel, + mean_model=mean_model, + t=t, + y=y, + yerr=yerr, + gp_class=celerite.GP, + ) def log_likelihood(self, parameters): """ @@ -716,7 +767,6 @@ def log_likelihood(self, parameters): class GeorgeLikelihood(_GPLikelihood): - def __init__(self, kernel, mean_model, t, y, yerr=1e-6): """ Basic Gaussian Process likelihood interface for `celerite` and `george`. @@ -738,7 +788,15 @@ def __init__(self, kernel, mean_model, t, y, yerr=1e-6): applies for all y-values. Default is 1e-6, effectively assuming that no y-errors are present. """ import george - super().__init__(kernel=kernel, mean_model=mean_model, t=t, y=y, yerr=yerr, gp_class=george.GP) + + super().__init__( + kernel=kernel, + mean_model=mean_model, + t=t, + y=y, + yerr=yerr, + gp_class=george.GP, + ) def log_likelihood(self, parameters): """ diff --git a/bilby/core/prior/analytical.py b/bilby/core/prior/analytical.py index 21be9f3da..99c0b0f4a 100644 --- a/bilby/core/prior/analytical.py +++ b/bilby/core/prior/analytical.py @@ -20,7 +20,6 @@ class DeltaFunction(Prior): - _leaves = ["peak", "_minimum", "_maximum", "least_recently_sampled"] def __init__(self, peak, name=None, latex_label=None, unit=None): @@ -38,8 +37,14 @@ def __init__(self, peak, name=None, latex_label=None, unit=None): See superclass """ - super(DeltaFunction, self).__init__(name=name, latex_label=latex_label, unit=unit, - minimum=peak, maximum=peak, check_range_nonzero=False) + super(DeltaFunction, self).__init__( + name=name, + latex_label=latex_label, + unit=unit, + minimum=peak, + maximum=peak, + check_range_nonzero=False, + ) self.peak = peak self._is_fixed = True self.least_recently_sampled = peak @@ -55,7 +60,7 @@ def rescale(self, val, *, xp=None): ======= float: Rescaled probability, equivalent to peak """ - return self.peak * val ** 0 + return self.peak * val**0 def prob(self, val, *, xp=None): """Return the prior probability of val @@ -69,7 +74,7 @@ def prob(self, val, *, xp=None): Union[float, array_like]: np.inf if val = peak, 0 otherwise """ - at_peak = (val == self.peak) + at_peak = val == self.peak # coerce bool to float for some array backends return at_peak * 1.0 @@ -78,11 +83,18 @@ def cdf(self, val, *, xp=None): class PowerLaw(Prior): - _leaves = ["alpha"] - def __init__(self, alpha, minimum, maximum, name=None, latex_label=None, - unit=None, boundary=None): + def __init__( + self, + alpha, + minimum, + maximum, + name=None, + latex_label=None, + unit=None, + boundary=None, + ): """Power law with bounds and alpha, spectral index Parameters @@ -102,9 +114,14 @@ def __init__(self, alpha, minimum, maximum, name=None, latex_label=None, boundary: str See superclass """ - super(PowerLaw, self).__init__(name=name, latex_label=latex_label, - minimum=minimum, maximum=maximum, unit=unit, - boundary=boundary) + super(PowerLaw, self).__init__( + name=name, + latex_label=latex_label, + minimum=minimum, + maximum=maximum, + unit=unit, + boundary=boundary, + ) self.alpha = alpha @xp_wrap @@ -127,12 +144,12 @@ def rescale(self, val, *, xp=None): maximum = xp.asarray(self.maximum) pow = xp.asarray(self.alpha + 1) - with np.errstate(divide='ignore', invalid='ignore'): + with np.errstate(divide="ignore", invalid="ignore"): safe_pow = xp.where(pow == 0, xp.ones(pow.shape), pow) base_x = ( - minimum ** safe_pow + val * (maximum ** safe_pow - minimum ** safe_pow) - ) ** (1. / safe_pow) + minimum**safe_pow + val * (maximum**safe_pow - minimum**safe_pow) + ) ** (1.0 / safe_pow) safe_min = xp.where(minimum <= 0, xp.ones_like(minimum), minimum) safe_max = xp.where(maximum <= 0, xp.ones_like(maximum), maximum) @@ -156,13 +173,13 @@ def prob(self, val, *, xp=None): maximum = xp.asarray(self.maximum) pow = xp.asarray(self.alpha + 1) - with np.errstate(divide='ignore', invalid='ignore'): + with np.errstate(divide="ignore", invalid="ignore"): norm = xp.where( pow != 0, - pow / (maximum ** pow - minimum ** pow), + pow / (maximum**pow - minimum**pow), 1 / xp.log(maximum / minimum), ) - return xp.nan_to_num(val ** self.alpha * norm) * self.is_in_prior_range(val) + return xp.nan_to_num(val**self.alpha * norm) * self.is_in_prior_range(val) @xp_wrap def ln_prob(self, val, *, xp=None): @@ -181,13 +198,13 @@ def ln_prob(self, val, *, xp=None): maximum = xp.asarray(self.maximum) pow = xp.asarray(self.alpha + 1) - with np.errstate(divide='ignore', invalid='ignore'): + with np.errstate(divide="ignore", invalid="ignore"): normalising = xp.where( pow != 0, - pow / (maximum ** pow - minimum ** pow), + pow / (maximum**pow - minimum**pow), 1 / xp.log(maximum / minimum), ) - ln_in_range = xp.log(1. * self.is_in_prior_range(val)) + ln_in_range = xp.log(1.0 * self.is_in_prior_range(val)) ln_p = self.alpha * xp.nan_to_num(xp.log(val)) + xp.log(normalising) return ln_p + ln_in_range @@ -198,12 +215,11 @@ def cdf(self, val, *, xp=None): maximum = xp.asarray(self.maximum) pow = xp.asarray(self.alpha + 1) - with np.errstate(divide='ignore', invalid='ignore'): + with np.errstate(divide="ignore", invalid="ignore"): safe_pow = xp.where(pow == 0, xp.ones(pow.shape), pow) - base_cdf = ( - (val ** safe_pow - minimum ** safe_pow) - / (maximum ** safe_pow - minimum ** safe_pow) + base_cdf = (val**safe_pow - minimum**safe_pow) / ( + maximum**safe_pow - minimum**safe_pow ) safe_min = xp.where(minimum <= 0, xp.ones_like(minimum), minimum) @@ -216,9 +232,9 @@ def cdf(self, val, *, xp=None): class Uniform(Prior): - - def __init__(self, minimum, maximum, name=None, latex_label=None, - unit=None, boundary=None): + def __init__( + self, minimum, maximum, name=None, latex_label=None, unit=None, boundary=None + ): """Uniform prior with bounds Parameters @@ -236,9 +252,14 @@ def __init__(self, minimum, maximum, name=None, latex_label=None, boundary: str See superclass """ - super(Uniform, self).__init__(name=name, latex_label=latex_label, - minimum=minimum, maximum=maximum, unit=unit, - boundary=boundary) + super(Uniform, self).__init__( + name=name, + latex_label=latex_label, + minimum=minimum, + maximum=maximum, + unit=unit, + boundary=boundary, + ) def rescale(self, val, *, xp=None): """ @@ -268,7 +289,9 @@ def prob(self, val, *, xp=None): ======= float: Prior probability of val """ - return ((val >= self.minimum) & (val <= self.maximum)) / (self.maximum - self.minimum) + return ((val >= self.minimum) & (val <= self.maximum)) / ( + self.maximum - self.minimum + ) @xp_wrap def ln_prob(self, val, *, xp=None): @@ -292,9 +315,9 @@ def cdf(self, val, *, xp=None): class LogUniform(PowerLaw): - - def __init__(self, minimum, maximum, name=None, latex_label=None, - unit=None, boundary=None): + def __init__( + self, minimum, maximum, name=None, latex_label=None, unit=None, boundary=None + ): """Log-Uniform prior with bounds Parameters @@ -312,16 +335,27 @@ def __init__(self, minimum, maximum, name=None, latex_label=None, boundary: str See superclass """ - super(LogUniform, self).__init__(name=name, latex_label=latex_label, unit=unit, - minimum=minimum, maximum=maximum, alpha=-1, boundary=boundary) + super(LogUniform, self).__init__( + name=name, + latex_label=latex_label, + unit=unit, + minimum=minimum, + maximum=maximum, + alpha=-1, + boundary=boundary, + ) if self.minimum <= 0: - logger.warning('You specified a uniform-in-log prior with minimum={}'.format(self.minimum)) + logger.warning( + "You specified a uniform-in-log prior with minimum={}".format( + self.minimum + ) + ) class SymmetricLogUniform(Prior): - - def __init__(self, minimum, maximum, name=None, latex_label=None, - unit=None, boundary=None): + def __init__( + self, minimum, maximum, name=None, latex_label=None, unit=None, boundary=None + ): """Symmetric Log-Uniform distributions with bounds This is identical to a Log-Uniform distribution, but mirrored about @@ -344,9 +378,14 @@ def __init__(self, minimum, maximum, name=None, latex_label=None, boundary: str See superclass """ - super(SymmetricLogUniform, self).__init__(name=name, latex_label=latex_label, - minimum=minimum, maximum=maximum, unit=unit, - boundary=boundary) + super(SymmetricLogUniform, self).__init__( + name=name, + latex_label=latex_label, + minimum=minimum, + maximum=maximum, + unit=unit, + boundary=boundary, + ) @xp_wrap def rescale(self, val, *, xp=None): @@ -373,7 +412,9 @@ def rescale(self, val, *, xp=None): return ( xp.sign(2 * val - 1) * self.minimum - * xp.exp(xp.abs(2 * val - 1) * xp.log(xp.asarray(self.maximum / self.minimum))) + * xp.exp( + xp.abs(2 * val - 1) * xp.log(xp.asarray(self.maximum / self.minimum)) + ) ) @xp_wrap @@ -389,8 +430,9 @@ def prob(self, val, *, xp=None): float: Prior probability of val """ val = xp.abs(val) - return (xp.nan_to_num(0.5 / val / xp.log(xp.asarray(self.maximum / self.minimum))) * - self.is_in_prior_range(val)) + return xp.nan_to_num( + 0.5 / val / xp.log(xp.asarray(self.maximum / self.minimum)) + ) * self.is_in_prior_range(val) @xp_wrap def ln_prob(self, val, *, xp=None): @@ -405,21 +447,29 @@ def ln_prob(self, val, *, xp=None): float: """ - return xp.nan_to_num(- xp.log(2 * xp.abs(val)) - xp.log(xp.log(xp.asarray(self.maximum / self.minimum)))) + return xp.nan_to_num( + -xp.log(2 * xp.abs(val)) + - xp.log(xp.log(xp.asarray(self.maximum / self.minimum))) + ) @xp_wrap def cdf(self, val, *, xp=None): - asymmetric = ( - xp.log(xp.maximum(xp.abs(val) / self.minimum, xp.asarray(1))) - / xp.log(xp.asarray(self.maximum / self.minimum)) - ) + asymmetric = xp.log( + xp.maximum(xp.abs(val) / self.minimum, xp.asarray(1)) + ) / xp.log(xp.asarray(self.maximum / self.minimum)) return xp.clip(0.5 * (1 + xp.sign(val) * asymmetric), 0, 1) class Cosine(Prior): - - def __init__(self, minimum=-np.pi / 2, maximum=np.pi / 2, name=None, - latex_label=None, unit=None, boundary=None): + def __init__( + self, + minimum=-np.pi / 2, + maximum=np.pi / 2, + name=None, + latex_label=None, + unit=None, + boundary=None, + ): """Cosine prior with bounds Parameters @@ -437,8 +487,14 @@ def __init__(self, minimum=-np.pi / 2, maximum=np.pi / 2, name=None, boundary: str See superclass """ - super(Cosine, self).__init__(minimum=minimum, maximum=maximum, name=name, - latex_label=latex_label, unit=unit, boundary=boundary) + super(Cosine, self).__init__( + minimum=minimum, + maximum=maximum, + name=name, + latex_label=latex_label, + unit=unit, + boundary=boundary, + ) @xp_wrap def rescale(self, val, *, xp=None): @@ -466,9 +522,8 @@ def prob(self, val, *, xp=None): @xp_wrap def cdf(self, val, *, xp=None): - _cdf = ( - (xp.sin(val) - xp.sin(xp.asarray(self.minimum))) / - (xp.sin(xp.asarray(self.maximum)) - xp.sin(xp.asarray(self.minimum))) + _cdf = (xp.sin(val) - xp.sin(xp.asarray(self.minimum))) / ( + xp.sin(xp.asarray(self.maximum)) - xp.sin(xp.asarray(self.minimum)) ) _cdf *= val >= self.minimum _cdf *= val <= self.maximum @@ -477,9 +532,15 @@ def cdf(self, val, *, xp=None): class Sine(Prior): - - def __init__(self, minimum=0, maximum=np.pi, name=None, - latex_label=None, unit=None, boundary=None): + def __init__( + self, + minimum=0, + maximum=np.pi, + name=None, + latex_label=None, + unit=None, + boundary=None, + ): """Sine prior with bounds Parameters @@ -497,8 +558,14 @@ def __init__(self, minimum=0, maximum=np.pi, name=None, boundary: str See superclass """ - super(Sine, self).__init__(minimum=minimum, maximum=maximum, name=name, - latex_label=latex_label, unit=unit, boundary=boundary) + super(Sine, self).__init__( + minimum=minimum, + maximum=maximum, + name=name, + latex_label=latex_label, + unit=unit, + boundary=boundary, + ) @xp_wrap def rescale(self, val, *, xp=None): @@ -526,9 +593,8 @@ def prob(self, val, *, xp=None): @xp_wrap def cdf(self, val, *, xp=None): - _cdf = ( - (xp.cos(val) - xp.cos(xp.asarray(self.minimum))) - / (xp.cos(xp.asarray(self.maximum)) - xp.cos(xp.asarray(self.minimum))) + _cdf = (xp.cos(val) - xp.cos(xp.asarray(self.minimum))) / ( + xp.cos(xp.asarray(self.maximum)) - xp.cos(xp.asarray(self.minimum)) ) _cdf *= val >= self.minimum _cdf *= val <= self.maximum @@ -537,10 +603,11 @@ def cdf(self, val, *, xp=None): class Gaussian(Prior): - _leaves = ["mu", "sigma"] - def __init__(self, mu, sigma, name=None, latex_label=None, unit=None, boundary=None): + def __init__( + self, mu, sigma, name=None, latex_label=None, unit=None, boundary=None + ): """Gaussian prior with mean mu and width sigma Parameters @@ -558,7 +625,9 @@ def __init__(self, mu, sigma, name=None, latex_label=None, unit=None, boundary=N boundary: str See superclass """ - super(Gaussian, self).__init__(name=name, latex_label=latex_label, unit=unit, boundary=boundary) + super(Gaussian, self).__init__( + name=name, latex_label=latex_label, unit=unit, boundary=boundary + ) self.mu = mu self.sigma = sigma @@ -573,7 +642,7 @@ def rescale(self, val, *, xp=None): This maps to the inverse CDF. This has been analytically solved for this case. """ - return self.mu + erfinv(2 * val - 1) * 2 ** 0.5 * self.sigma + return self.mu + erfinv(2 * val - 1) * 2**0.5 * self.sigma @xp_wrap def prob(self, val, *, xp=None): @@ -587,7 +656,11 @@ def prob(self, val, *, xp=None): ======= Union[float, array_like]: Prior probability of val """ - return xp.exp(-(self.mu - val) ** 2 / (2 * self.sigma ** 2)) / (2 * np.pi) ** 0.5 / self.sigma + return ( + xp.exp(-((self.mu - val) ** 2) / (2 * self.sigma**2)) + / (2 * np.pi) ** 0.5 + / self.sigma + ) @xp_wrap def ln_prob(self, val, *, xp=None): @@ -601,22 +674,33 @@ def ln_prob(self, val, *, xp=None): ======= Union[float, array_like]: Prior probability of val """ - return -0.5 * ((self.mu - val) ** 2 / self.sigma ** 2 + xp.log(xp.asarray(2 * np.pi * self.sigma ** 2))) + return -0.5 * ( + (self.mu - val) ** 2 / self.sigma**2 + + xp.log(xp.asarray(2 * np.pi * self.sigma**2)) + ) def cdf(self, val, *, xp=None): - return (1 - erf((self.mu - val) / 2 ** 0.5 / self.sigma)) / 2 + return (1 - erf((self.mu - val) / 2**0.5 / self.sigma)) / 2 class Normal(Gaussian): - """A synonym for the Gaussian distribution. """ + """A synonym for the Gaussian distribution.""" class TruncatedGaussian(Prior): - _leaves = ["mu", "sigma"] - def __init__(self, mu, sigma, minimum, maximum, name=None, - latex_label=None, unit=None, boundary=None): + def __init__( + self, + mu, + sigma, + minimum, + maximum, + name=None, + latex_label=None, + unit=None, + boundary=None, + ): """Truncated Gaussian prior with mean mu and width sigma https://en.wikipedia.org/wiki/Truncated_normal_distribution @@ -640,21 +724,29 @@ def __init__(self, mu, sigma, minimum, maximum, name=None, boundary: str See superclass """ - super(TruncatedGaussian, self).__init__(name=name, latex_label=latex_label, unit=unit, - minimum=minimum, maximum=maximum, boundary=boundary) + super(TruncatedGaussian, self).__init__( + name=name, + latex_label=latex_label, + unit=unit, + minimum=minimum, + maximum=maximum, + boundary=boundary, + ) self.mu = mu self.sigma = sigma @property def normalisation(self): - """ Calculates the proper normalisation of the truncated Gaussian + """Calculates the proper normalisation of the truncated Gaussian Returns ======= float: Proper normalisation of the truncated Gaussian """ - return (erf((self.maximum - self.mu) / 2 ** 0.5 / self.sigma) - erf( - (self.minimum - self.mu) / 2 ** 0.5 / self.sigma)) / 2 + return ( + erf((self.maximum - self.mu) / 2**0.5 / self.sigma) + - erf((self.minimum - self.mu) / 2**0.5 / self.sigma) + ) / 2 @xp_wrap def rescale(self, val, *, xp=None): @@ -663,8 +755,15 @@ def rescale(self, val, *, xp=None): This maps to the inverse CDF. This has been analytically solved for this case. """ - return erfinv(2 * val * self.normalisation + erf( - (self.minimum - self.mu) / 2 ** 0.5 / self.sigma)) * 2 ** 0.5 * self.sigma + self.mu + return ( + erfinv( + 2 * val * self.normalisation + + erf((self.minimum - self.mu) / 2**0.5 / self.sigma) + ) + * 2**0.5 + * self.sigma + + self.mu + ) @xp_wrap def prob(self, val, *, xp=None): @@ -678,12 +777,23 @@ def prob(self, val, *, xp=None): ======= float: Prior probability of val """ - return xp.exp(-(self.mu - val) ** 2 / (2 * self.sigma ** 2)) / (2 * np.pi) ** 0.5 \ - / self.sigma / self.normalisation * self.is_in_prior_range(val) + return ( + xp.exp(-((self.mu - val) ** 2) / (2 * self.sigma**2)) + / (2 * np.pi) ** 0.5 + / self.sigma + / self.normalisation + * self.is_in_prior_range(val) + ) def cdf(self, val, *, xp=None): - _cdf = (erf((val - self.mu) / 2 ** 0.5 / self.sigma) - erf( - (self.minimum - self.mu) / 2 ** 0.5 / self.sigma)) / 2 / self.normalisation + _cdf = ( + ( + erf((val - self.mu) / 2**0.5 / self.sigma) + - erf((self.minimum - self.mu) / 2**0.5 / self.sigma) + ) + / 2 + / self.normalisation + ) _cdf *= val >= self.minimum _cdf *= val <= self.maximum _cdf += val > self.maximum @@ -695,7 +805,6 @@ class TruncatedNormal(TruncatedGaussian): class HalfGaussian(TruncatedGaussian): - _leaves = ["sigma", "mu"] def __init__(self, sigma, name=None, latex_label=None, unit=None, boundary=None): @@ -714,9 +823,16 @@ def __init__(self, sigma, name=None, latex_label=None, unit=None, boundary=None) boundary: str See superclass """ - super(HalfGaussian, self).__init__(mu=0., sigma=sigma, minimum=0., maximum=np.inf, - name=name, latex_label=latex_label, - unit=unit, boundary=boundary) + super(HalfGaussian, self).__init__( + mu=0.0, + sigma=sigma, + minimum=0.0, + maximum=np.inf, + name=name, + latex_label=latex_label, + unit=unit, + boundary=boundary, + ) class HalfNormal(HalfGaussian): @@ -724,10 +840,11 @@ class HalfNormal(HalfGaussian): class LogNormal(Prior): - _leaves = ["mu", "sigma"] - def __init__(self, mu, sigma, name=None, latex_label=None, unit=None, boundary=None): + def __init__( + self, mu, sigma, name=None, latex_label=None, unit=None, boundary=None + ): """Log-normal prior with mean mu and width sigma https://en.wikipedia.org/wiki/Log-normal_distribution @@ -747,11 +864,18 @@ def __init__(self, mu, sigma, name=None, latex_label=None, unit=None, boundary=N boundary: str See superclass """ - super(LogNormal, self).__init__(name=name, minimum=0., latex_label=latex_label, - unit=unit, boundary=boundary) + super(LogNormal, self).__init__( + name=name, + minimum=0.0, + latex_label=latex_label, + unit=unit, + boundary=boundary, + ) - if sigma <= 0.: - raise ValueError("For the LogGaussian prior the standard deviation must be positive") + if sigma <= 0.0: + raise ValueError( + "For the LogGaussian prior the standard deviation must be positive" + ) self.mu = mu self.sigma = sigma @@ -763,7 +887,7 @@ def rescale(self, val, *, xp=None): This maps to the inverse CDF. This has been analytically solved for this case. """ - return xp.exp(self.mu + (2 * self.sigma ** 2)**0.5 * erfinv(2 * val - 1)) + return xp.exp(self.mu + (2 * self.sigma**2) ** 0.5 * erfinv(2 * val - 1)) @xp_wrap def prob(self, val, *, xp=None): @@ -792,17 +916,34 @@ def ln_prob(self, val, *, xp=None): Union[float, array_like]: Prior probability of val """ with np.errstate(divide="ignore", invalid="ignore"): - return xp.nan_to_num(( - -(xp.log(xp.maximum(val, xp.asarray(self.minimum))) - self.mu) ** 2 / self.sigma ** 2 / 2 - - xp.log((2 * np.pi)**0.5 * val * self.sigma) - ) + xp.log(val > self.minimum), nan=-xp.inf, neginf=-xp.inf, posinf=-xp.inf) + return xp.nan_to_num( + ( + -( + (xp.log(xp.maximum(val, xp.asarray(self.minimum))) - self.mu) + ** 2 + ) + / self.sigma**2 + / 2 + - xp.log((2 * np.pi) ** 0.5 * val * self.sigma) + ) + + xp.log(val > self.minimum), + nan=-xp.inf, + neginf=-xp.inf, + posinf=-xp.inf, + ) @xp_wrap def cdf(self, val, *, xp=None): with np.errstate(divide="ignore"): - return 0.5 + erf( - (xp.log(xp.maximum(val, xp.asarray(self.minimum))) - self.mu) / self.sigma / np.sqrt(2) - ) / 2 + return ( + 0.5 + + erf( + (xp.log(xp.maximum(val, xp.asarray(self.minimum))) - self.mu) + / self.sigma + / np.sqrt(2) + ) + / 2 + ) class LogGaussian(LogNormal): @@ -810,7 +951,6 @@ class LogGaussian(LogNormal): class Exponential(Prior): - _leaves = ["mu"] def __init__(self, mu, name=None, latex_label=None, unit=None, boundary=None): @@ -829,8 +969,13 @@ def __init__(self, mu, name=None, latex_label=None, unit=None, boundary=None): boundary: str See superclass """ - super(Exponential, self).__init__(name=name, minimum=0., latex_label=latex_label, - unit=unit, boundary=boundary) + super(Exponential, self).__init__( + name=name, + minimum=0.0, + latex_label=latex_label, + unit=unit, + boundary=boundary, + ) self.mu = mu @xp_wrap @@ -870,20 +1015,31 @@ def ln_prob(self, val, *, xp=None): Union[float, array_like]: Prior probability of val """ with np.errstate(divide="ignore"): - return -val / self.mu - xp.log(xp.asarray(self.mu)) + xp.log(val >= self.minimum) + return ( + -val / self.mu + - xp.log(xp.asarray(self.mu)) + + xp.log(val >= self.minimum) + ) @xp_wrap def cdf(self, val, *, xp=None): with np.errstate(divide="ignore", invalid="ignore", over="ignore"): - return xp.maximum(1. - xp.exp(-val / self.mu), xp.asarray(0.0)) + return xp.maximum(1.0 - xp.exp(-val / self.mu), xp.asarray(0.0)) class StudentT(Prior): - _leaves = ["df", "mu", "scale"] - def __init__(self, df, mu=0., scale=1., name=None, latex_label=None, - unit=None, boundary=None): + def __init__( + self, + df, + mu=0.0, + scale=1.0, + name=None, + latex_label=None, + unit=None, + boundary=None, + ): """Student's t-distribution prior with number of degrees of freedom df, mean mu and scale @@ -906,10 +1062,14 @@ def __init__(self, df, mu=0., scale=1., name=None, latex_label=None, boundary: str See superclass """ - super(StudentT, self).__init__(name=name, latex_label=latex_label, unit=unit, boundary=boundary) + super(StudentT, self).__init__( + name=name, latex_label=latex_label, unit=unit, boundary=boundary + ) - if df <= 0. or scale <= 0.: - raise ValueError("For the StudentT prior the number of degrees of freedom and scale must be positive") + if df <= 0.0 or scale <= 0.0: + raise ValueError( + "For the StudentT prior the number of degrees of freedom and scale must be positive" + ) self.df = df self.mu = mu @@ -960,8 +1120,11 @@ def ln_prob(self, val, *, xp=None): Union[float, array_like]: Prior probability of val """ return ( - gammaln(0.5 * (self.df + 1)) - gammaln(0.5 * self.df) - - xp.log(xp.asarray((np.pi * self.df)**0.5 * self.scale)) - (self.df + 1) / 2 + gammaln(0.5 * (self.df + 1)) + - gammaln(0.5 * self.df) + - xp.log(xp.asarray((np.pi * self.df) ** 0.5 * self.scale)) + - (self.df + 1) + / 2 * xp.log(1 + ((val - self.mu) / self.scale) ** 2 / self.df) ) @@ -970,11 +1133,19 @@ def cdf(self, val, *, xp=None): class Beta(Prior): - _leaves = ["alpha", "beta"] - def __init__(self, alpha, beta, minimum=0, maximum=1, name=None, - latex_label=None, unit=None, boundary=None): + def __init__( + self, + alpha, + beta, + minimum=0, + maximum=1, + name=None, + latex_label=None, + unit=None, + boundary=None, + ): """Beta distribution https://en.wikipedia.org/wiki/Beta_distribution @@ -1001,10 +1172,16 @@ def __init__(self, alpha, beta, minimum=0, maximum=1, name=None, boundary: str See superclass """ - super(Beta, self).__init__(minimum=minimum, maximum=maximum, name=name, - latex_label=latex_label, unit=unit, boundary=boundary) + super(Beta, self).__init__( + minimum=minimum, + maximum=maximum, + name=name, + latex_label=latex_label, + unit=unit, + boundary=boundary, + ) - if alpha <= 0. or beta <= 0.: + if alpha <= 0.0 or beta <= 0.0: raise ValueError("alpha and beta must both be positive values") self.alpha = alpha @@ -1072,16 +1249,17 @@ def cdf(self, val, *, xp=None): betainc( xp.asarray(self.alpha), xp.asarray(self.beta), - (val - self.minimum) / (self.maximum - self.minimum) + (val - self.minimum) / (self.maximum - self.minimum), ) ) + (val > self.maximum) class Logistic(Prior): - _leaves = ["mu", "scale"] - def __init__(self, mu, scale, name=None, latex_label=None, unit=None, boundary=None): + def __init__( + self, mu, scale, name=None, latex_label=None, unit=None, boundary=None + ): """Logistic distribution https://en.wikipedia.org/wiki/Logistic_distribution @@ -1101,9 +1279,11 @@ def __init__(self, mu, scale, name=None, latex_label=None, unit=None, boundary=N boundary: str See superclass """ - super(Logistic, self).__init__(name=name, latex_label=latex_label, unit=unit, boundary=boundary) + super(Logistic, self).__init__( + name=name, latex_label=latex_label, unit=unit, boundary=boundary + ) - if scale <= 0.: + if scale <= 0.0: raise ValueError("For the Logistic prior the scale must be positive") self.mu = mu @@ -1118,7 +1298,9 @@ def rescale(self, val, *, xp=None): """ with np.errstate(divide="ignore"): val = xp.asarray(val) - return self.mu + self.scale * xp.log(xp.maximum(val / (1 - val), xp.asarray(0))) + return self.mu + self.scale * xp.log( + xp.maximum(val / (1 - val), xp.asarray(0)) + ) @xp_wrap def prob(self, val, *, xp=None): @@ -1147,19 +1329,23 @@ def ln_prob(self, val, *, xp=None): Union[float, array_like]: Prior probability of val """ with np.errstate(over="ignore"): - return -(val - self.mu) / self.scale -\ - 2. * xp.log1p(xp.exp(-(val - self.mu) / self.scale)) - xp.log(xp.asarray(self.scale)) + return ( + -(val - self.mu) / self.scale + - 2.0 * xp.log1p(xp.exp(-(val - self.mu) / self.scale)) + - xp.log(xp.asarray(self.scale)) + ) @xp_wrap def cdf(self, val, *, xp=None): - return 1. / (1. + xp.exp(-(val - self.mu) / self.scale)) + return 1.0 / (1.0 + xp.exp(-(val - self.mu) / self.scale)) class Cauchy(Prior): - _leaves = ["alpha", "beta"] - def __init__(self, alpha, beta, name=None, latex_label=None, unit=None, boundary=None): + def __init__( + self, alpha, beta, name=None, latex_label=None, unit=None, boundary=None + ): """Cauchy distribution https://en.wikipedia.org/wiki/Cauchy_distribution @@ -1179,9 +1365,11 @@ def __init__(self, alpha, beta, name=None, latex_label=None, unit=None, boundary boundary: str See superclass """ - super(Cauchy, self).__init__(name=name, latex_label=latex_label, unit=unit, boundary=boundary) + super(Cauchy, self).__init__( + name=name, latex_label=latex_label, unit=unit, boundary=boundary + ) - if beta <= 0.: + if beta <= 0.0: raise ValueError("For the Cauchy prior the scale must be positive") self.alpha = alpha @@ -1209,7 +1397,7 @@ def prob(self, val, *, xp=None): ======= Union[float, array_like]: Prior probability of val """ - return 1. / self.beta / np.pi / (1. + ((val - self.alpha) / self.beta) ** 2) + return 1.0 / self.beta / np.pi / (1.0 + ((val - self.alpha) / self.beta) ** 2) @xp_wrap def ln_prob(self, val, *, xp=None): @@ -1223,7 +1411,9 @@ def ln_prob(self, val, *, xp=None): ======= Union[float, array_like]: Log prior probability of val """ - return - xp.log(xp.asarray(self.beta * np.pi)) - xp.log(1. + ((val - self.alpha) / self.beta) ** 2) + return -xp.log(xp.asarray(self.beta * np.pi)) - xp.log( + 1.0 + ((val - self.alpha) / self.beta) ** 2 + ) @xp_wrap def cdf(self, val, *, xp=None): @@ -1235,10 +1425,11 @@ class Lorentzian(Cauchy): class Gamma(Prior): - _leaves = ["k", "theta"] - def __init__(self, k, theta=1., name=None, latex_label=None, unit=None, boundary=None): + def __init__( + self, k, theta=1.0, name=None, latex_label=None, unit=None, boundary=None + ): """Gamma distribution https://en.wikipedia.org/wiki/Gamma_distribution @@ -1258,8 +1449,13 @@ def __init__(self, k, theta=1., name=None, latex_label=None, unit=None, boundary boundary: str See superclass """ - super(Gamma, self).__init__(name=name, minimum=0., latex_label=latex_label, - unit=unit, boundary=boundary) + super(Gamma, self).__init__( + name=name, + minimum=0.0, + latex_label=latex_label, + unit=unit, + boundary=boundary, + ) if k <= 0 or theta <= 0: raise ValueError("For the Gamma prior the shape and scale must be positive") @@ -1310,18 +1506,21 @@ def ln_prob(self, val, *, xp=None): """ with np.errstate(divide="ignore"): ln_prob = ( - xlogy(xp.asarray(self.k - 1), val) - val / self.theta - - xlogy(self.k, self.theta) - gammaln(self.k) + xlogy(xp.asarray(self.k - 1), val) + - val / self.theta + - xlogy(self.k, self.theta) + - gammaln(self.k) ) + xp.log(val >= self.minimum) return xp.nan_to_num(ln_prob, nan=-xp.inf, neginf=-xp.inf, posinf=xp.inf) @xp_wrap def cdf(self, val, *, xp=None): - return gammainc(xp.asarray(self.k), xp.maximum(val, xp.asarray(self.minimum)) / self.theta) + return gammainc( + xp.asarray(self.k), xp.maximum(val, xp.asarray(self.minimum)) / self.theta + ) class ChiSquared(Gamma): - _leaves = ["nu", "theta"] def __init__(self, nu, name=None, latex_label=None, unit=None, boundary=None): @@ -1344,10 +1543,18 @@ def __init__(self, nu, name=None, latex_label=None, unit=None, boundary=None): """ if nu <= 0 or not isinstance(nu, int): - raise ValueError("For the ChiSquared prior the number of degrees of freedom must be a positive integer") + raise ValueError( + "For the ChiSquared prior the number of degrees of freedom must be a positive integer" + ) - super(ChiSquared, self).__init__(name=name, k=nu / 2., theta=2., - latex_label=latex_label, unit=unit, boundary=boundary) + super(ChiSquared, self).__init__( + name=name, + k=nu / 2.0, + theta=2.0, + latex_label=latex_label, + unit=unit, + boundary=boundary, + ) @property def nu(self): @@ -1355,15 +1562,13 @@ def nu(self): @nu.setter def nu(self, nu): - self.k = nu / 2. + self.k = nu / 2.0 class FermiDirac(Prior): - _leaves = ["sigma", "mu", "r", "expr"] - def __init__(self, sigma, mu=None, r=None, name=None, latex_label=None, - unit=None): + def __init__(self, sigma, mu=None, r=None, name=None, latex_label=None, unit=None): """A Fermi-Dirac type prior, with a fixed lower boundary at zero (see, e.g. Section 2.3.5 of [1]_). The probability distribution is defined by Equation 22 of [1]_. @@ -1391,13 +1596,17 @@ def __init__(self, sigma, mu=None, r=None, name=None, latex_label=None, .. [1] M. Pitkin, M. Isi, J. Veitch & G. Woan, `arXiv:1705.08978v1 `_, 2017. """ - super(FermiDirac, self).__init__(name=name, latex_label=latex_label, unit=unit, minimum=0.) + super(FermiDirac, self).__init__( + name=name, latex_label=latex_label, unit=unit, minimum=0.0 + ) self.sigma = sigma if mu is None and r is None: - raise ValueError("For the Fermi-Dirac prior either a 'mu' value or 'r' " - "value must be given.") + raise ValueError( + "For the Fermi-Dirac prior either a 'mu' value or 'r' " + "value must be given." + ) if r is None and mu is not None: self.mu = mu @@ -1406,9 +1615,10 @@ def __init__(self, sigma, mu=None, r=None, name=None, latex_label=None, self.r = r self.mu = self.sigma * self.r - if self.r <= 0. or self.sigma <= 0.: - raise ValueError("For the Fermi-Dirac prior the values of sigma and r " - "must be positive.") + if self.r <= 0.0 or self.sigma <= 0.0: + raise ValueError( + "For the Fermi-Dirac prior the values of sigma and r must be positive." + ) xp = array_module((mu, sigma, r)) self.expr = xp.exp(self.r) @@ -1431,7 +1641,11 @@ def rescale(self, val, *, xp=None): .. [1] M. Pitkin, M. Isi, J. Veitch & G. Woan, `arXiv:1705.08978v1 `_, 2017. """ - inv = -1 / self.expr + (1 + self.expr)**-val + (1 + self.expr)**-val / self.expr + inv = ( + -1 / self.expr + + (1 + self.expr) ** -val + + (1 + self.expr) ** -val / self.expr + ) return -self.sigma * xp.log(xp.maximum(inv, xp.asarray(0))) @xp_wrap @@ -1447,7 +1661,7 @@ def prob(self, val, *, xp=None): float: Prior probability of val """ return ( - (xp.exp((val - self.mu) / self.sigma) + 1)**-1 + (xp.exp((val - self.mu) / self.sigma) + 1) ** -1 / (self.sigma * xp.log1p(xp.asarray(self.expr))) * (val >= self.minimum) ) @@ -1489,9 +1703,9 @@ def cdf(self, val, *, xp=None): `_, 2017. """ result = ( - (xp.logaddexp(xp.asarray(0.0), -xp.asarray(self.r)) - xp.logaddexp(-val / self.sigma, -xp.asarray(self.r))) - / xp.logaddexp(xp.asarray(0.0), xp.asarray(self.r)) - ) + xp.logaddexp(xp.asarray(0.0), -xp.asarray(self.r)) + - xp.logaddexp(-val / self.sigma, -xp.asarray(self.r)) + ) / xp.logaddexp(xp.asarray(0.0), xp.asarray(self.r)) return xp.clip(result, 0, 1) @@ -1533,8 +1747,13 @@ def __init__( # Small delta added to help with MCMC walking maximum = xp.max(values) * (1 + 1e-15) super(WeightedDiscreteValues, self).__init__( - name=name, latex_label=latex_label, minimum=minimum, - maximum=maximum, unit=unit, boundary=boundary) + name=name, + latex_label=latex_label, + minimum=minimum, + maximum=maximum, + unit=unit, + boundary=boundary, + ) self.nvalues = nvalues sorter = xp.argsort(values) self._values_array = values[sorter] @@ -1684,8 +1903,9 @@ def pytree_unflatten(cls, aux_data, children): class DiscreteValues(WeightedDiscreteValues): - def __init__(self, values, name=None, latex_label=None, - unit=None, boundary="periodic"): + def __init__( + self, values, name=None, latex_label=None, unit=None, boundary="periodic" + ): """An equal-weighted discrete-valued prior Parameters @@ -1885,17 +2105,13 @@ def cdf(self, val, *, xp=None): float: prior cumulative probability at val """ - return ( + return (val > self.mode) + (val > self.minimum) * (val <= self.maximum) / ( + self.scale + ) * ( (val > self.mode) - + (val > self.minimum) - * (val <= self.maximum) - / (self.scale) - * ( - (val > self.mode) - * (self.rescaled_maximum - val) ** 2.0 - / (self.mode - self.rescaled_maximum) - + (val <= self.mode) - * (val - self.rescaled_minimum) ** 2.0 - / (self.mode - self.rescaled_minimum) - ) + * (self.rescaled_maximum - val) ** 2.0 + / (self.mode - self.rescaled_maximum) + + (val <= self.mode) + * (val - self.rescaled_minimum) ** 2.0 + / (self.mode - self.rescaled_minimum) ) diff --git a/bilby/core/prior/base.py b/bilby/core/prior/base.py index 723b31459..6f81b9b3f 100644 --- a/bilby/core/prior/base.py +++ b/bilby/core/prior/base.py @@ -23,9 +23,17 @@ class Prior(object): _default_latex_labels = {} _leaves = [] - def __init__(self, name=None, latex_label=None, unit=None, minimum=-np.inf, - maximum=np.inf, check_range_nonzero=True, boundary=None): - """ Implements a Prior object + def __init__( + self, + name=None, + latex_label=None, + unit=None, + minimum=-np.inf, + maximum=np.inf, + check_range_nonzero=True, + boundary=None, + ): + """Implements a Prior object Parameters ========== @@ -134,7 +142,7 @@ def __eq__(self, other): if aac.is_array_api_obj(this_dict[key]): if not np.array_equal(this_dict[key], other_dict[key]): return False - elif isinstance(this_dict[key], type(scipy.stats.beta(1., 1.))): + elif isinstance(this_dict[key], type(scipy.stats.beta(1.0, 1.0))): continue else: if not this_dict[key] == other_dict[key]: @@ -155,6 +163,7 @@ def sample(self, size=None, *, random_state=None): """ from ..utils import random + rng = random.resolve_random_state(random_state) if isinstance(size, (int, np.integer)): @@ -198,12 +207,14 @@ def prob(self, val, *, xp=None): @xp_wrap def cdf(self, val, *, xp=None): - """ Generic method to calculate CDF, can be overwritten in subclass """ + """Generic method to calculate CDF, can be overwritten in subclass""" from scipy.integrate import cumulative_trapezoid + if np.any(np.isinf([self.minimum, self.maximum])): raise ValueError( "Unable to use the generic CDF calculation for priors with" - "infinite support") + "infinite support" + ) x = xp.linspace(self.minimum, self.maximum, 1000) pdf = self.prob(x, xp=xp) cdf = cumulative_trapezoid(pdf, x, initial=0) @@ -223,7 +234,7 @@ def ln_prob(self, val, *, xp=None): np.nan """ - with np.errstate(divide='ignore'): + with np.errstate(divide="ignore"): return xp.log(self.prob(val, xp=xp)) def is_in_prior_range(self, val): @@ -264,8 +275,11 @@ def _native(value): return value.item() return value - args = ', '.join( - [f'{key}={repr(_native(instantiation_dict[key]))}' for key in instantiation_dict] + args = ", ".join( + [ + f"{key}={repr(_native(instantiation_dict[key]))}" + for key in instantiation_dict + ] ) if "bilby.core.prior" in prior_module: return f"{prior_name}({args})" @@ -316,7 +330,7 @@ def unit(self, unit): @property def latex_label_with_unit(self): - """ If a unit is specified, returns a string of the latex label and unit """ + """If a unit is specified, returns a string of the latex label and unit""" if self.unit is not None: return "{} [{}]".format(self.latex_label, self.unit) else: @@ -353,8 +367,10 @@ def boundary(self): @boundary.setter def boundary(self, boundary): - if boundary not in ['periodic', 'reflective', None]: - raise ValueError('{} is not a valid setting for prior boundaries'.format(boundary)) + if boundary not in ["periodic", "reflective", None]: + raise ValueError( + "{} is not a valid setting for prior boundaries".format(boundary) + ) self._boundary = boundary @property @@ -381,19 +397,22 @@ def from_repr(cls, string): def _from_repr(cls, string): subclass_args = infer_args_from_method(cls.__init__) - string = string.replace(' ', '') + string = string.replace(" ", "") kwargs = cls._split_repr(string) for key in kwargs: val = kwargs[key] if key not in subclass_args and not hasattr(cls, "reference_params"): - raise AttributeError('Unknown argument {} for class {}'.format( - key, cls.__name__)) + raise AttributeError( + "Unknown argument {} for class {}".format(key, cls.__name__) + ) else: kwargs[key] = cls._parse_argument_string(val) - if key in ["condition_func", "conversion_function"] and isinstance(kwargs[key], str): + if key in ["condition_func", "conversion_function"] and isinstance( + kwargs[key], str + ): if "." in kwargs[key]: - module = '.'.join(kwargs[key].split('.')[:-1]) - name = kwargs[key].split('.')[-1] + module = ".".join(kwargs[key].split(".")[:-1]) + name = kwargs[key].split(".")[-1] else: module = __name__ name = kwargs[key] @@ -403,30 +422,29 @@ def _from_repr(cls, string): @classmethod def _split_repr(cls, string): subclass_args = infer_args_from_method(cls.__init__) - args = string.split(',') + args = string.split(",") remove = list() for ii, key in enumerate(args): - for paren_pair in ['()', '{}', '[]']: + for paren_pair in ["()", "{}", "[]"]: if paren_pair[0] in key: jj = ii while paren_pair[1] not in args[jj]: jj += 1 - args[ii] = ','.join([args[ii], args[jj]]).strip() + args[ii] = ",".join([args[ii], args[jj]]).strip() remove.append(jj) remove.reverse() for ii in remove: del args[ii] kwargs = dict() for ii, arg in enumerate(args): - if '=' not in arg: - logger.debug( - 'Reading priors with non-keyword arguments is dangerous!') + if "=" not in arg: + logger.debug("Reading priors with non-keyword arguments is dangerous!") key = subclass_args[ii] val = arg else: - split_arg = arg.split('=') + split_arg = arg.split("=") key = split_arg[0] - val = '='.join(split_arg[1:]) + val = "=".join(split_arg[1:]) kwargs[key] = val return kwargs @@ -471,20 +489,22 @@ def _parse_argument_string(cls, val): TypeError: If val cannot be parsed as described above. """ - if val == 'None': + if val == "None": val = None - elif re.sub(r'\'.*\'', '', val) in ['r', 'u']: + elif re.sub(r"\'.*\'", "", val) in ["r", "u"]: val = val[2:-1] elif val.startswith("'") and val.endswith("'"): val = val.strip("'") - elif '(' in val and not val.startswith(("[", "{")): - other_cls = val.split('(')[0] - vals = '('.join(val.split('(')[1:])[:-1] + elif "(" in val and not val.startswith(("[", "{")): + other_cls = val.split("(")[0] + vals = "(".join(val.split("(")[1:])[:-1] if "." in other_cls: - module = '.'.join(other_cls.split('.')[:-1]) - other_cls = other_cls.split('.')[-1] + module = ".".join(other_cls.split(".")[:-1]) + other_cls = other_cls.split(".")[-1] else: - module = __name__.replace('.' + os.path.basename(__file__).replace('.py', ''), '') + module = __name__.replace( + "." + os.path.basename(__file__).replace(".py", ""), "" + ) other_cls = getattr(import_module(module), other_cls) val = other_cls.from_repr(vals) else: @@ -492,13 +512,12 @@ def _parse_argument_string(cls, val): val = eval(val, dict(), dict(np=np, inf=np.inf, pi=np.pi)) except NameError: if "." in val: - module = '.'.join(val.split('.')[:-1]) - func = val.split('.')[-1] + module = ".".join(val.split(".")[:-1]) + func = val.split(".")[-1] new_val = getattr(import_module(module), func, val) if val == new_val: raise TypeError( - "Cannot evaluate prior, " - f"failed to parse argument {val}" + f"Cannot evaluate prior, failed to parse argument {val}" ) else: val = new_val @@ -506,13 +525,16 @@ def _parse_argument_string(cls, val): class Constraint(Prior): - _leaves = ["_minimum", "_maximum"] - def __init__(self, minimum, maximum, name=None, latex_label=None, - unit=None): - super(Constraint, self).__init__(minimum=minimum, maximum=maximum, name=name, - latex_label=latex_label, unit=unit) + def __init__(self, minimum, maximum, name=None, latex_label=None, unit=None): + super(Constraint, self).__init__( + minimum=minimum, + maximum=maximum, + name=name, + latex_label=latex_label, + unit=unit, + ) self._is_fixed = True def prob(self, val, *, xp=None): @@ -520,4 +542,4 @@ def prob(self, val, *, xp=None): class PriorException(Exception): - """ General base class for all prior exceptions """ + """General base class for all prior exceptions""" diff --git a/bilby/core/prior/conditional.py b/bilby/core/prior/conditional.py index 0dc3bb934..0f8647dcc 100644 --- a/bilby/core/prior/conditional.py +++ b/bilby/core/prior/conditional.py @@ -2,23 +2,47 @@ from .base import Prior, PriorException from .interpolated import Interped -from .analytical import DeltaFunction, PowerLaw, Uniform, LogUniform, \ - SymmetricLogUniform, Cosine, Sine, Gaussian, TruncatedGaussian, HalfGaussian, \ - LogNormal, Exponential, StudentT, Beta, Logistic, Cauchy, Gamma, ChiSquared, FermiDirac +from .analytical import ( + DeltaFunction, + PowerLaw, + Uniform, + LogUniform, + SymmetricLogUniform, + Cosine, + Sine, + Gaussian, + TruncatedGaussian, + HalfGaussian, + LogNormal, + Exponential, + StudentT, + Beta, + Logistic, + Cauchy, + Gamma, + ChiSquared, + FermiDirac, +) from ..utils import infer_args_from_method, infer_parameters_from_function from ...compat.utils import xp_wrap def conditional_prior_factory(prior_class): class ConditionalPrior(prior_class): - _leaves = prior_class._leaves - __name__ = 'Conditional{}'.format(prior_class.__name__) - __qualname__ = 'Conditional{}'.format(prior_class.__qualname__) - - def __init__(self, condition_func, name=None, latex_label=None, unit=None, - boundary=None, **reference_params): + __name__ = "Conditional{}".format(prior_class.__name__) + __qualname__ = "Conditional{}".format(prior_class.__qualname__) + + def __init__( + self, + condition_func, + name=None, + latex_label=None, + unit=None, + boundary=None, + **reference_params, + ): """ Parameters @@ -55,12 +79,20 @@ def condition_func(reference_params, y): This differs on the `prior_class`, for example for the Gaussian prior this is `mu` and `sigma`. """ - if 'boundary' in infer_args_from_method(super(ConditionalPrior, self).__init__): - super(ConditionalPrior, self).__init__(name=name, latex_label=latex_label, - unit=unit, boundary=boundary, **reference_params) + if "boundary" in infer_args_from_method( + super(ConditionalPrior, self).__init__ + ): + super(ConditionalPrior, self).__init__( + name=name, + latex_label=latex_label, + unit=unit, + boundary=boundary, + **reference_params, + ) else: - super(ConditionalPrior, self).__init__(name=name, latex_label=latex_label, - unit=unit, **reference_params) + super(ConditionalPrior, self).__init__( + name=name, latex_label=latex_label, unit=unit, **reference_params + ) self._required_variables = None self.condition_func = condition_func @@ -82,6 +114,7 @@ def sample(self, size=None, *, random_state=None, **required_variables): """ from ..utils import random + rng = random.resolve_random_state(random_state) if isinstance(size, int | np.integer): @@ -183,15 +216,19 @@ class depending on the required variables it depends on. """ required_variables.pop("xp", None) if sorted(list(required_variables)) == sorted(self.required_variables): - parameters = self.condition_func(self.reference_params.copy(), **required_variables) + parameters = self.condition_func( + self.reference_params.copy(), **required_variables + ) for key, value in parameters.items(): setattr(self, key, value) elif len(required_variables) == 0: return else: - raise IllegalRequiredVariablesException("Expected kwargs for {}. Got kwargs for {} instead." - .format(self.required_variables, - list(required_variables.keys()))) + raise IllegalRequiredVariablesException( + "Expected kwargs for {}. Got kwargs for {} instead.".format( + self.required_variables, list(required_variables.keys()) + ) + ) @property def reference_params(self): @@ -212,11 +249,13 @@ def condition_func(self, condition_func): self._condition_func = lambda reference_params: reference_params else: self._condition_func = condition_func - self._required_variables = infer_parameters_from_function(self.condition_func) + self._required_variables = infer_parameters_from_function( + self.condition_func + ) @property def required_variables(self): - """ The required variables to pass into the condition function. """ + """The required variables to pass into the condition function.""" return self._required_variables def get_instantiation_dict(self): @@ -245,12 +284,18 @@ def __repr__(self): """ prior_name = self.__class__.__name__ instantiation_dict = self.get_instantiation_dict() - instantiation_dict["condition_func"] = ".".join([ - instantiation_dict["condition_func"].__module__, - instantiation_dict["condition_func"].__name__ - ]) - args = ', '.join(['{}={}'.format(key, repr(instantiation_dict[key])) - for key in instantiation_dict]) + instantiation_dict["condition_func"] = ".".join( + [ + instantiation_dict["condition_func"].__module__, + instantiation_dict["condition_func"].__name__, + ] + ) + args = ", ".join( + [ + "{}={}".format(key, repr(instantiation_dict[key])) + for key in instantiation_dict + ] + ) return "{}({})".format(prior_name, args) return ConditionalPrior @@ -390,23 +435,22 @@ class DirichletElement(ConditionalBeta): def __init__(self, order, n_dimensions, label): """ """ super(DirichletElement, self).__init__( - minimum=0, maximum=1, alpha=1, beta=n_dimensions - order - 1, + minimum=0, + maximum=1, + alpha=1, + beta=n_dimensions - order - 1, name=label + str(order), - condition_func=self.dirichlet_condition + condition_func=self.dirichlet_condition, ) self.label = label self.n_dimensions = n_dimensions self.order = order - self._required_variables = [ - label + str(ii) for ii in range(order) - ] - self.__class__.__name__ = 'DirichletElement' - self.__class__.__qualname__ = 'DirichletElement' + self._required_variables = [label + str(ii) for ii in range(order)] + self.__class__.__name__ = "DirichletElement" + self.__class__.__qualname__ = "DirichletElement" def dirichlet_condition(self, reference_parms, **kwargs): - remaining = 1 - sum( - [kwargs[self.label + str(ii)] for ii in range(self.order)] - ) + remaining = 1 - sum([kwargs[self.label + str(ii)] for ii in range(self.order)]) return dict(minimum=reference_parms["minimum"], maximum=remaining) def __repr__(self): @@ -427,8 +471,8 @@ def pytree_unflatten(cls, aux_data, children): class ConditionalPriorException(PriorException): - """ General base class for all conditional prior exceptions """ + """General base class for all conditional prior exceptions""" class IllegalRequiredVariablesException(ConditionalPriorException): - """ Exception class for exceptions relating to handling the required variables. """ + """Exception class for exceptions relating to handling the required variables.""" diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py index 90f75f894..9582b9d86 100644 --- a/bilby/core/prior/dict.py +++ b/bilby/core/prior/dict.py @@ -64,8 +64,9 @@ def __hash__(self): def __setitem__(self, key, value): if not isinstance(value, (Prior, int, float)): raise TypeError( - "Unable to parse prior, bad entry: {} " - "= {} of type {}".format(key, value, type(value)) + "Unable to parse prior, bad entry: {} = {} of type {}".format( + key, value, type(value) + ) ) super().__setitem__(key, value) @@ -319,8 +320,9 @@ def from_dictionary(self, dictionary): continue else: raise TypeError( - "Unable to parse prior, bad entry: {} " - "= {} of type {}".format(key, val, type(val)) + "Unable to parse prior, bad entry: {} = {} of type {}".format( + key, val, type(val) + ) ) self.update(dictionary) @@ -356,9 +358,14 @@ def fill_priors(self, likelihood=None, default_priors_file=None): """ if likelihood is not None: - warn("Filling priors from likelihood parameters is deprecated", FutureWarning) + warn( + "Filling priors from likelihood parameters is deprecated", FutureWarning + ) if default_priors_file is not None: - warn("Setting default priors from a defaults file is deprecated", FutureWarning) + warn( + "Setting default priors from a defaults file is deprecated", + FutureWarning, + ) self.convert_floats_to_delta_functions() @@ -477,7 +484,9 @@ def check_efficiency(n_tested, n_valid): return efficiency = n_valid / float(n_tested) if n_tested >= 1e3 and efficiency < 1e-3: - logger.warning("Prior sampling efficiency is very low, please verify its validity.") + logger.warning( + "Prior sampling efficiency is very low, please verify its validity." + ) efficiency_warning_was_issued = True n_tested_samples, n_valid_samples = 0, 0 @@ -568,7 +577,9 @@ def prob(self, sample, *, normalized=True, xp=None, **kwargs): """ if xp is None: xp = array_module(sample.values()) - prob = xp.prod(xp.stack([self[key].prob(sample[key], xp=xp) for key in sample]), **kwargs) + prob = xp.prod( + xp.stack([self[key].prob(sample[key], xp=xp) for key in sample]), **kwargs + ) return self.check_prob(sample, prob, normalized=normalized, xp=xp) @@ -614,7 +625,10 @@ def ln_prob(self, sample, axis=None, *, normalized=True, xp=None): elif xp is None: # assume input is a dataframe xp = array_module(sample.values) - ln_prob = xp.sum(xp.stack([self[key].ln_prob(sample[key], xp=xp) for key in sample]), axis=axis) + ln_prob = xp.sum( + xp.stack([self[key].ln_prob(sample[key], xp=xp) for key in sample]), + axis=axis, + ) return self.check_ln_prob(sample, ln_prob, normalized=normalized, xp=xp) def check_ln_prob(self, sample, ln_prob, normalized=True, *, xp=None): @@ -669,7 +683,9 @@ def rescale(self, keys, theta, *, xp=None): if xp is None: xp = array_module(theta) - return xp.asarray([self[key].rescale(sample, xp=xp) for key, sample in zip(keys, theta)]) + return xp.asarray( + [self[key].rescale(sample, xp=xp) for key, sample in zip(keys, theta)] + ) def test_redundancy(self, key, disable_logging=False): """Empty redundancy test, should be overwritten in subclasses""" @@ -806,7 +822,9 @@ def sample_subset(self, keys=iter([]), size=None, *, random_state=None): rvars = { key: value[i] for key, value in required_variables.items() } - samples[key][i] = subset_dict[key].sample(**rvars, random_state=random_state) + samples[key][i] = subset_dict[key].sample( + **rvars, random_state=random_state + ) subset_dict[key].least_recently_sampled = samples[key] else: logger.debug("{} not a known prior.".format(key)) @@ -847,10 +865,12 @@ def prob(self, sample, *, normalized=True, xp=None, **kwargs): self._prepare_evaluation(*zip(*sample.items())) if xp is None: xp = array_module(sample.values()) - res = xp.asarray([ - self[key].prob(sample[key], **self.get_required_variables(key), xp=xp) - for key in sample - ]) + res = xp.asarray( + [ + self[key].prob(sample[key], **self.get_required_variables(key), xp=xp) + for key in sample + ] + ) prob = xp.prod(res, **kwargs) return self.check_prob(sample, prob, normalized=normalized, xp=xp) @@ -875,10 +895,14 @@ def ln_prob(self, sample, *, axis=None, normalized=True, xp=None): self._prepare_evaluation(*zip(*sample.items())) if xp is None: xp = array_module(sample.values()) - res = xp.asarray([ - self[key].ln_prob(sample[key], **self.get_required_variables(key), xp=xp) - for key in sample - ]) + res = xp.asarray( + [ + self[key].ln_prob( + sample[key], **self.get_required_variables(key), xp=xp + ) + for key in sample + ] + ) ln_prob = xp.sum(res, axis=axis) return self.check_ln_prob(sample, ln_prob, normalized=normalized, xp=xp) @@ -921,7 +945,10 @@ def rescale(self, keys, theta): theta[index], **self.get_required_variables(key) ) self[key].least_recently_sampled = result[key] - if isinstance(self[key], JointPrior) and self[key].dist.distname not in joint: + if ( + isinstance(self[key], JointPrior) + and self[key].dist.distname not in joint + ): joint[self[key].dist.distname] = [key] elif isinstance(self[key], JointPrior): joint[self[key].dist.distname].append(key) diff --git a/bilby/core/prior/interpolated.py b/bilby/core/prior/interpolated.py index 2ddd3f7e4..17f71b39a 100644 --- a/bilby/core/prior/interpolated.py +++ b/bilby/core/prior/interpolated.py @@ -10,7 +10,6 @@ class Interped(Prior): - _leaves = [ "xx", "_yy", @@ -25,8 +24,17 @@ class Interped(Prior): "_all_interpolated", ] - def __init__(self, xx, yy, minimum=np.nan, maximum=np.nan, name=None, - latex_label=None, unit=None, boundary=None): + def __init__( + self, + xx, + yy, + minimum=np.nan, + maximum=np.nan, + name=None, + latex_label=None, + unit=None, + boundary=None, + ): """Creates an interpolated prior function from arrays of xx and yy=p(xx) Parameters @@ -71,8 +79,14 @@ def __init__(self, xx, yy, minimum=np.nan, maximum=np.nan, name=None, self._all_interpolated = interp1d(x=xx, y=yy, bounds_error=False, fill_value=0) minimum = float(np.nanmax(np.array((min(xx), minimum)))) maximum = float(np.nanmin(np.array((max(xx), maximum)))) - super(Interped, self).__init__(name=name, latex_label=latex_label, unit=unit, - minimum=minimum, maximum=maximum, boundary=boundary) + super(Interped, self).__init__( + name=name, + latex_label=latex_label, + unit=unit, + minimum=minimum, + maximum=maximum, + boundary=boundary, + ) self._update_instance() def __eq__(self, other): @@ -127,9 +141,11 @@ def minimum(self): @minimum.setter def minimum(self, minimum): if minimum < self.min_limit: - raise ValueError('Minimum cannot be set below {}.'.format(round(self.min_limit, 2))) + raise ValueError( + "Minimum cannot be set below {}.".format(round(self.min_limit, 2)) + ) self._minimum = minimum - if '_maximum' in self.__dict__ and self._maximum < np.inf: + if "_maximum" in self.__dict__ and self._maximum < np.inf: self._update_instance() @property @@ -150,9 +166,11 @@ def maximum(self): @maximum.setter def maximum(self, maximum): if maximum > self.max_limit: - raise ValueError('Maximum cannot be set above {}.'.format(round(self.max_limit, 2))) + raise ValueError( + "Maximum cannot be set above {}.".format(round(self.max_limit, 2)) + ) self._maximum = maximum - if '_minimum' in self.__dict__ and self._minimum < np.inf: + if "_minimum" in self.__dict__ and self._minimum < np.inf: self._update_instance() @property @@ -171,7 +189,9 @@ def yy(self): @yy.setter def yy(self, yy): self._yy = yy - self._all_interpolated = interp1d(x=self.xx, y=self._yy, bounds_error=False, fill_value=0) + self._all_interpolated = interp1d( + x=self.xx, y=self._yy, bounds_error=False, fill_value=0 + ) self._update_instance() def _update_instance(self): @@ -182,21 +202,37 @@ def _update_instance(self): def _initialize_attributes(self): from scipy.integrate import cumulative_trapezoid + if trapezoid(self._yy, self.xx) != 1: - logger.debug('Supplied PDF for {} is not normalised, normalising.'.format(self.name)) + logger.debug( + "Supplied PDF for {} is not normalised, normalising.".format(self.name) + ) self._yy /= trapezoid(self._yy, self.xx) self.YY = cumulative_trapezoid(self._yy, self.xx, initial=0) # Need last element of cumulative distribution to be exactly one. self.YY = xpx.at(self.YY, -1).set(1) - self.probability_density = interp1d(x=self.xx, y=self._yy, bounds_error=False, fill_value=0) - self.cumulative_distribution = interp1d(x=self.xx, y=self.YY, bounds_error=False, fill_value=(0, 1)) - self.inverse_cumulative_distribution = interp1d(x=self.YY, y=self.xx, bounds_error=True) + self.probability_density = interp1d( + x=self.xx, y=self._yy, bounds_error=False, fill_value=0 + ) + self.cumulative_distribution = interp1d( + x=self.xx, y=self.YY, bounds_error=False, fill_value=(0, 1) + ) + self.inverse_cumulative_distribution = interp1d( + x=self.YY, y=self.xx, bounds_error=True + ) class FromFile(Interped): - - def __init__(self, file_name, minimum=None, maximum=None, name=None, - latex_label=None, unit=None, boundary=None): + def __init__( + self, + file_name, + minimum=None, + maximum=None, + name=None, + latex_label=None, + unit=None, + boundary=None, + ): """Creates an interpolated prior function from arrays of xx and yy=p(xx) extracted from a file Parameters @@ -220,9 +256,16 @@ def __init__(self, file_name, minimum=None, maximum=None, name=None, try: self.file_name = file_name xx, yy = np.genfromtxt(self.file_name).T - super(FromFile, self).__init__(xx=xx, yy=yy, minimum=minimum, - maximum=maximum, name=name, latex_label=latex_label, - unit=unit, boundary=boundary) + super(FromFile, self).__init__( + xx=xx, + yy=yy, + minimum=minimum, + maximum=maximum, + name=name, + latex_label=latex_label, + unit=unit, + boundary=boundary, + ) except IOError: logger.warning("Can't load {}.".format(self.file_name)) logger.warning("Format should be:") diff --git a/bilby/core/prior/joint.py b/bilby/core/prior/joint.py index 34fffc740..1fbb67c6f 100644 --- a/bilby/core/prior/joint.py +++ b/bilby/core/prior/joint.py @@ -419,7 +419,7 @@ def __init__( self.covs = [] self.corrcoefs = [] self.sigmas = [] - self.logprodsigmas = [] # log of product of sigmas, needed for "standard" multivariate normal + self.logprodsigmas = [] # log of product of sigmas, needed for "standard" multivariate normal self.weights = [] self.eigvalues = [] self.eigvectors = [] @@ -526,17 +526,13 @@ def add_mode(self, mus=None, sigmas=None, corrcoef=None, cov=None, weight=1.0): self.corrcoefs.append(np.asarray(corrcoef)) if len(self.corrcoefs[-1].shape) != 2: - raise ValueError( - "Correlation coefficient matrix must be a 2d array." - ) + raise ValueError("Correlation coefficient matrix must be a 2d array.") if ( self.corrcoefs[-1].shape[0] != self.corrcoefs[-1].shape[1] or self.corrcoefs[-1].shape[0] != self.num_vars ): - raise ValueError( - "Correlation coefficient matrix shape is inconsistent" - ) + raise ValueError("Correlation coefficient matrix shape is inconsistent") # check matrix is symmetric if not np.allclose(self.corrcoefs[-1], self.corrcoefs[-1].T): @@ -575,15 +571,11 @@ def add_mode(self, mus=None, sigmas=None, corrcoef=None, cov=None, weight=1.0): self.eigvalues.append(evals) self.eigvectors.append(evecs) except Exception as e: - raise RuntimeError( - "Problem getting eigenvalues and vectors: {}".format(e) - ) + raise RuntimeError("Problem getting eigenvalues and vectors: {}".format(e)) # check eigenvalues are positive if np.any(self.eigvalues[-1] <= 0.0): - raise ValueError( - "Correlation coefficient matrix is not positive definite" - ) + raise ValueError("Correlation coefficient matrix is not positive definite") self.sqeigvalues.append(np.sqrt(self.eigvalues[-1])) # set the weights @@ -607,7 +599,9 @@ def add_mode(self, mus=None, sigmas=None, corrcoef=None, cov=None, weight=1.0): # multivariate_normal(mean=mus, cov=cov).logpdf(x) # = multivariate_normal(mean=0, cov=corrcoefs).logpdf((x - mus)/sigmas) - logprodsigmas self.mvn.append( - scipy.stats.multivariate_normal(mean=np.zeros(self.num_vars), cov=self.corrcoefs[-1]) + scipy.stats.multivariate_normal( + mean=np.zeros(self.num_vars), cov=self.corrcoefs[-1] + ) ) @xp_wrap @@ -623,11 +617,13 @@ def _rescale(self, samp, *, xp=None, **kwargs): else: mode = np.argwhere(self.cumweights - random.rng.uniform(0, 1) > 0)[0][0] - samp = erfinv(2.0 * samp - 1) * 2.0 ** 0.5 + samp = erfinv(2.0 * samp - 1) * 2.0**0.5 # rotate and scale to the multivariate normal shape samp = xp.asarray(self.mus[mode]) + xp.asarray(self.sigmas[mode]) * xp.einsum( - "ij,kj->ik", samp * self.sqeigvalues[mode], xp.asarray(self.eigvectors[mode]) + "ij,kj->ik", + samp * self.sqeigvalues[mode], + xp.asarray(self.eigvectors[mode]), ) return samp @@ -686,7 +682,10 @@ def _ln_prob(self, samp, lnprob, outbounds, *, xp=None): # self.mvn[i] is a "standard" multivariate normal distribution; see add_mode() z = (samp[j] - self.mus[i]) / self.sigmas[i] lnprob = xpx.at(lnprob, j).set( - xp.logaddexp(lnprob[j], self.mvn[i].logpdf(z) - xp.asarray(self.logprodsigmas[i])) + xp.logaddexp( + lnprob[j], + self.mvn[i].logpdf(z) - xp.asarray(self.logprodsigmas[i]), + ) ) # set out-of-bounds values to -inf @@ -804,9 +803,9 @@ def rescale(self, val, *, xp=None, **kwargs): if self.dist.filled_rescale(): # print(self.dist.rescale_parameters) - values = xp.stack([ - xp.asarray(val) for val in self.dist.rescale_parameters.values() - ]).T + values = xp.stack( + [xp.asarray(val) for val in self.dist.rescale_parameters.values()] + ).T # values = xp.asarray(list(self.dist.rescale_parameters.values())).T samples = self.dist.rescale(values, **kwargs) self.dist.reset_rescale() @@ -880,8 +879,7 @@ def ln_prob(self, val, *, xp=None): shapes.add(xp.asarray(v).shape) if len(shapes) > 1: raise ValueError( - "Each parameter must have the same " - "number of requested values." + "Each parameter must have the same number of requested values." ) lnp = self.dist.ln_prob(xp.stack(values).T) diff --git a/bilby/core/prior/slabspike.py b/bilby/core/prior/slabspike.py index f4f132d1e..80dcfefd6 100644 --- a/bilby/core/prior/slabspike.py +++ b/bilby/core/prior/slabspike.py @@ -6,8 +6,14 @@ class SlabSpikePrior(Prior): - - _leaves = ["slab", "_spike_location", "_spike_height", "inverse_cdf_below_spike", "_minimum", "_maximum"] + _leaves = [ + "slab", + "_spike_location", + "_spike_height", + "inverse_cdf_below_spike", + "_minimum", + "_maximum", + ] def __init__(self, slab, spike_location=None, spike_height=0): """'Slab-and-spike' prior, see e.g. https://arxiv.org/abs/1812.07259 @@ -33,15 +39,25 @@ def __init__(self, slab, spike_location=None, spike_height=0): """ self.slab = slab - super().__init__(name=self.slab.name, latex_label=self.slab.latex_label, unit=self.slab.unit, - minimum=self.slab.minimum, maximum=self.slab.maximum, - check_range_nonzero=self.slab.check_range_nonzero, boundary=self.slab.boundary) + super().__init__( + name=self.slab.name, + latex_label=self.slab.latex_label, + unit=self.slab.unit, + minimum=self.slab.minimum, + maximum=self.slab.maximum, + check_range_nonzero=self.slab.check_range_nonzero, + boundary=self.slab.boundary, + ) self.spike_location = spike_location self.spike_height = spike_height try: - self.inverse_cdf_below_spike = self._find_inverse_cdf_fraction_before_spike() + self.inverse_cdf_below_spike = ( + self._find_inverse_cdf_fraction_before_spike() + ) except Exception as e: - logger.warning("Disregard the following warning when running tests:\n {}".format(e)) + logger.warning( + "Disregard the following warning when running tests:\n {}".format(e) + ) @property def spike_location(self): @@ -52,7 +68,9 @@ def spike_location(self, spike_loc): if spike_loc is None: spike_loc = self.minimum if not self.minimum <= spike_loc <= self.maximum: - raise ValueError("Spike location {} not within prior domain ".format(spike_loc)) + raise ValueError( + "Spike location {} not within prior domain ".format(spike_loc) + ) self._spike_location = spike_loc @property @@ -64,11 +82,13 @@ def spike_height(self, spike_height): if 0 <= spike_height <= 1: self._spike_height = spike_height else: - raise ValueError("Spike height must be between 0 and 1, but is {}".format(spike_height)) + raise ValueError( + "Spike height must be between 0 and 1, but is {}".format(spike_height) + ) @property def slab_fraction(self): - """ Relative prior weight of the slab. """ + """Relative prior weight of the slab.""" return 1 - self.spike_height def _find_inverse_cdf_fraction_before_spike(self): @@ -157,7 +177,7 @@ def ln_prob(self, val, *, xp=None): @xp_wrap def cdf(self, val, *, xp=None): - """ Return the CDF of the prior. + """Return the CDF of the prior. This calls to the slab CDF and adds a discrete step at the spike location. diff --git a/bilby/core/result.py b/bilby/core/result.py index c593d989c..305e24fd0 100644 --- a/bilby/core/result.py +++ b/bilby/core/result.py @@ -16,12 +16,17 @@ from . import utils from .utils import ( - logger, infer_parameters_from_function, + logger, + infer_parameters_from_function, check_directory_exists_and_if_not_mkdir, - latex_plot_format, safe_save_figure, - BilbyJsonEncoder, load_json, - move_old_file, get_version_information, - decode_bilby_json, docstring, + latex_plot_format, + safe_save_figure, + BilbyJsonEncoder, + load_json, + move_old_file, + get_version_information, + decode_bilby_json, + docstring, recursively_save_dict_contents_to_group, recursively_load_dict_contents_from_group, recursively_decode_bilby_json, @@ -39,8 +44,8 @@ def __eval_l(likelihood, params): return likelihood.log_likelihood() -def result_file_name(outdir, label, extension='json', gzip=False): - """ Returns the standard filename used for a result file +def result_file_name(outdir, label, extension="json", gzip=False): + """Returns the standard filename used for a result file Parameters ========== @@ -57,19 +62,19 @@ def result_file_name(outdir, label, extension='json', gzip=False): ======= str: File name of the output file """ - if extension == 'pickle': - extension = 'pkl' - if extension in ['json', 'hdf5', 'pkl']: - if extension == 'json' and gzip: - return os.path.join(outdir, '{}_result.{}.gz'.format(label, extension)) + if extension == "pickle": + extension = "pkl" + if extension in ["json", "hdf5", "pkl"]: + if extension == "json" and gzip: + return os.path.join(outdir, "{}_result.{}.gz".format(label, extension)) else: - return os.path.join(outdir, '{}_result.{}'.format(label, extension)) + return os.path.join(outdir, "{}_result.{}".format(label, extension)) else: raise ValueError("Extension type {} not understood".format(extension)) def _determine_file_name(filename, outdir, label, extension, gzip): - """ Helper method to determine the filename """ + """Helper method to determine the filename""" if filename is not None: if isinstance(filename, os.PathLike): # convert PathLike object to string @@ -83,8 +88,15 @@ def _determine_file_name(filename, outdir, label, extension, gzip): return result_file_name(outdir, label, extension, gzip) -def read_in_result(filename=None, outdir=None, label=None, extension=None, gzip=False, result_class=None): - """ Reads in a stored bilby result object +def read_in_result( + filename=None, + outdir=None, + label=None, + extension=None, + gzip=False, + result_class=None, +): + """Reads in a stored bilby result object Parameters ========== @@ -102,9 +114,7 @@ def read_in_result(filename=None, outdir=None, label=None, extension=None, gzip= but objects which inherit from this class can be given providing additional methods. """ - filename = _determine_file_name( - filename, outdir, label, extension or "json", gzip - ) + filename = _determine_file_name(filename, outdir, label, extension or "json", gzip) if result_class is None: result_class = Result @@ -115,19 +125,21 @@ def read_in_result(filename=None, outdir=None, label=None, extension=None, gzip= if extension is None: ext = os.path.splitext(filename)[1][1:] extension = ext if ext else None - if extension == 'gz': # gzipped file - extension = os.path.splitext(os.path.splitext(filename)[0])[1].lstrip('.') + if extension == "gz": # gzipped file + extension = os.path.splitext(os.path.splitext(filename)[0])[1].lstrip(".") if extension is None: - raise ValueError("No filetype extension provided and could not be inferred from filename") + raise ValueError( + "No filetype extension provided and could not be inferred from filename" + ) extension = extension.lower() read_functions = { - 'json': result_class.from_json, - 'hdf5': result_class.from_hdf5, - 'h5': result_class.from_hdf5, - 'pkl': result_class.from_pickle, - 'pickle': result_class.from_pickle, + "json": result_class.from_json, + "hdf5": result_class.from_hdf5, + "h5": result_class.from_hdf5, + "pkl": result_class.from_pickle, + "pickle": result_class.from_pickle, } if extension not in read_functions: raise ValueError( @@ -158,7 +170,7 @@ def read_in_result(filename=None, outdir=None, label=None, extension=None, gzip= def read_in_result_list(filename_list, invalid="warning"): - """ Read in a set of results + """Read in a set of results Parameters ========== @@ -174,9 +186,8 @@ def read_in_result_list(filename_list, invalid="warning"): """ results_list = [] for filename in filename_list: - if ( - not os.path.exists(filename) - and os.path.exists(f"{os.path.splitext(filename)[0]}.pkl") + if not os.path.exists(filename) and os.path.exists( + f"{os.path.splitext(filename)[0]}.pkl" ): pickle_path = f"{os.path.splitext(filename)[0]}.pkl" logger.warning( @@ -196,9 +207,17 @@ def read_in_result_list(filename_list, invalid="warning"): def get_weights_for_reweighting( - result, new_likelihood=None, new_prior=None, old_likelihood=None, - old_prior=None, resume_file=None, n_checkpoint=5000, npool=1, pool=None): - """ Calculate the weights for reweight() + result, + new_likelihood=None, + new_prior=None, + old_likelihood=None, + old_prior=None, + resume_file=None, + n_checkpoint=5000, + npool=1, + pool=None, +): + """Calculate the weights for reweight() See bilby.core.result.reweight() for help with the inputs @@ -231,11 +250,15 @@ def get_weights_for_reweighting( starting_index = 0 if (resume_file is not None) and os.path.exists(resume_file): - old_log_likelihood_array, old_log_prior_array, new_log_likelihood_array, new_log_prior_array = \ - np.genfromtxt(resume_file) + ( + old_log_likelihood_array, + old_log_prior_array, + new_log_likelihood_array, + new_log_prior_array, + ) = np.genfromtxt(resume_file) starting_index = np.argmin(np.abs(old_log_likelihood_array)) - logger.info(f'Checkpoint resuming from {starting_index}.') + logger.info(f"Checkpoint resuming from {starting_index}.") elif resume_file is not None: basedir = os.path.split(resume_file)[0] check_directory_exists_and_if_not_mkdir(basedir) @@ -254,30 +277,35 @@ def eval_pool(this_logl): chunksize = max(100, n // (2 * npool)) map_fn = partial(my_pool.imap, chunksize=chunksize) - log_l = list(tqdm( - map_fn(this_logl.log_likelihood, dict_samples[starting_index:]), - desc='Computing likelihoods', - total=n, - )) + log_l = list( + tqdm( + map_fn(this_logl.log_likelihood, dict_samples[starting_index:]), + desc="Computing likelihoods", + total=n, + ) + ) return log_l if old_likelihood is None: - old_log_likelihood_array[starting_index:] = \ - result.posterior["log_likelihood"][starting_index:].to_numpy() + old_log_likelihood_array[starting_index:] = result.posterior["log_likelihood"][ + starting_index: + ].to_numpy() else: old_log_likelihood_array[starting_index:] = eval_pool(old_likelihood) if new_likelihood is None: # Don't perform likelihood reweighting (i.e. likelihood isn't updated) - new_log_likelihood_array[starting_index:] = old_log_likelihood_array[starting_index:] + new_log_likelihood_array[starting_index:] = old_log_likelihood_array[ + starting_index: + ] else: new_log_likelihood_array[starting_index:] = eval_pool(new_likelihood) # Compute priors - for ii, sample in enumerate(tqdm(dict_samples[starting_index:], - desc='Computing priors', - total=n), - start=starting_index): + for ii, sample in enumerate( + tqdm(dict_samples[starting_index:], desc="Computing priors", total=n), + start=starting_index, + ): # prior calculation needs to not have prior or likelihood keys ln_prior = sample.pop("log_prior", np.nan) if "log_likelihood" in sample: @@ -299,19 +327,35 @@ def eval_pool(this_logl): if (ii % (n_checkpoint) == 0) and (resume_file is not None): checkpointed_index = np.argmin(np.abs(old_log_likelihood_array)) - logger.info(f'Checkpointing with {checkpointed_index} samples') + logger.info(f"Checkpointing with {checkpointed_index} samples") np.savetxt( resume_file, - [old_log_likelihood_array, old_log_prior_array, new_log_likelihood_array, new_log_prior_array]) + [ + old_log_likelihood_array, + old_log_prior_array, + new_log_likelihood_array, + new_log_prior_array, + ], + ) ln_weights = ( - new_log_likelihood_array + new_log_prior_array - old_log_likelihood_array - old_log_prior_array) + new_log_likelihood_array + + new_log_prior_array + - old_log_likelihood_array + - old_log_prior_array + ) - return ln_weights, new_log_likelihood_array, new_log_prior_array, old_log_likelihood_array, old_log_prior_array + return ( + ln_weights, + new_log_likelihood_array, + new_log_prior_array, + old_log_likelihood_array, + old_log_prior_array, + ) def rejection_sample(posterior, weights): - """ Perform rejection sampling on a posterior using weights + """Perform rejection sampling on a posterior using weights Parameters ========== @@ -330,11 +374,22 @@ def rejection_sample(posterior, weights): return posterior[keep] -def reweight(result, label=None, new_likelihood=None, new_prior=None, - old_likelihood=None, old_prior=None, conversion_function=None, npool=1, - verbose_output=False, resume_file=None, n_checkpoint=5000, - use_nested_samples=False, pool=None): - """ Reweight a result to a new likelihood/prior using rejection sampling +def reweight( + result, + label=None, + new_likelihood=None, + new_prior=None, + old_likelihood=None, + old_prior=None, + conversion_function=None, + npool=1, + verbose_output=False, + resume_file=None, + n_checkpoint=5000, + use_nested_samples=False, + pool=None, +): + """Reweight a result to a new likelihood/prior using rejection sampling Parameters ========== @@ -392,13 +447,23 @@ def reweight(result, label=None, new_likelihood=None, new_prior=None, nposterior = len(result.posterior) logger.info("Reweighting posterior with {} samples".format(nposterior)) - ln_weights, new_log_likelihood_array, new_log_prior_array, old_log_likelihood_array, old_log_prior_array =\ - get_weights_for_reweighting( - result, new_likelihood=new_likelihood, new_prior=new_prior, - old_likelihood=old_likelihood, old_prior=old_prior, - resume_file=resume_file, n_checkpoint=n_checkpoint, - npool=npool, pool=pool, - ) + ( + ln_weights, + new_log_likelihood_array, + new_log_prior_array, + old_log_likelihood_array, + old_log_prior_array, + ) = get_weights_for_reweighting( + result, + new_likelihood=new_likelihood, + new_prior=new_prior, + old_likelihood=old_likelihood, + old_prior=old_prior, + resume_file=resume_file, + n_checkpoint=n_checkpoint, + npool=npool, + pool=pool, + ) if use_nested_samples: ln_weights += np.log(result.posterior["weights"]) @@ -411,7 +476,9 @@ def reweight(result, label=None, new_likelihood=None, new_prior=None, result.posterior = rejection_sample(result.posterior, weights=weights) result.posterior = result.posterior.reset_index(drop=True) - logger.info("Rejection sampling resulted in {} samples".format(len(result.posterior))) + logger.info( + "Rejection sampling resulted in {} samples".format(len(result.posterior)) + ) result.meta_data["reweighted_using_rejection_sampling"] = True if use_nested_samples: @@ -428,7 +495,10 @@ def reweight(result, label=None, new_likelihood=None, new_prior=None, parameters = inspect.signature(conversion_function).parameters kwargs = dict() for key, value in [ - ("likelihood", new_likelihood), ("priors", new_prior), ("npool", npool), ("pool", pool) + ("likelihood", new_likelihood), + ("priors", new_prior), + ("npool", npool), + ("pool", pool), ]: if key in parameters: kwargs[key] = value @@ -441,28 +511,52 @@ def reweight(result, label=None, new_likelihood=None, new_prior=None, result.label += "_reweighted" if verbose_output: - return result, weights, new_log_likelihood_array, \ - new_log_prior_array, old_log_likelihood_array, old_log_prior_array + return ( + result, + weights, + new_log_likelihood_array, + new_log_prior_array, + old_log_likelihood_array, + old_log_prior_array, + ) else: return result class Result(object): - def __init__(self, label='no_label', outdir='.', sampler=None, - search_parameter_keys=None, fixed_parameter_keys=None, - constraint_parameter_keys=None, priors=None, - sampler_kwargs=None, injection_parameters=None, - meta_data=None, posterior=None, samples=None, - nested_samples=None, log_evidence=np.nan, - log_evidence_err=np.nan, information_gain=np.nan, - log_noise_evidence=np.nan, log_bayes_factor=np.nan, - log_likelihood_evaluations=None, - log_prior_evaluations=None, sampling_time=None, nburn=None, - num_likelihood_evaluations=None, walkers=None, - max_autocorrelation_time=None, use_ratio=None, - parameter_labels=None, parameter_labels_with_unit=None, - version=None): - """ A class to store the results of the sampling run + def __init__( + self, + label="no_label", + outdir=".", + sampler=None, + search_parameter_keys=None, + fixed_parameter_keys=None, + constraint_parameter_keys=None, + priors=None, + sampler_kwargs=None, + injection_parameters=None, + meta_data=None, + posterior=None, + samples=None, + nested_samples=None, + log_evidence=np.nan, + log_evidence_err=np.nan, + information_gain=np.nan, + log_noise_evidence=np.nan, + log_bayes_factor=np.nan, + log_likelihood_evaluations=None, + log_prior_evaluations=None, + sampling_time=None, + nburn=None, + num_likelihood_evaluations=None, + walkers=None, + max_autocorrelation_time=None, + use_ratio=None, + parameter_labels=None, + parameter_labels_with_unit=None, + version=None, + ): + """A class to store the results of the sampling run Parameters ========== @@ -578,8 +672,9 @@ def __init__(self, label='no_label', outdir='.', sampler=None, @staticmethod @docstring(_load_doctstring.format(format="pickle")) def from_pickle(filename=None, outdir=None, label=None): - filename = _determine_file_name(filename, outdir, label, 'hdf5', False) + filename = _determine_file_name(filename, outdir, label, "hdf5", False) import dill + with open(filename, "rb") as ff: return dill.load(ff) @@ -587,15 +682,16 @@ def from_pickle(filename=None, outdir=None, label=None): @docstring(_load_doctstring.format(format="hdf5")) def from_hdf5(cls, filename=None, outdir=None, label=None): import h5py - filename = _determine_file_name(filename, outdir, label, 'hdf5', False) + + filename = _determine_file_name(filename, outdir, label, "hdf5", False) with h5py.File(filename, "r") as ff: - data = recursively_load_dict_contents_from_group(ff, '/') + data = recursively_load_dict_contents_from_group(ff, "/") data["posterior"] = pd.DataFrame(data["posterior"]) data["priors"] = PriorDict._get_from_json_dict( json.loads(data["priors"], object_hook=decode_bilby_json) ) try: - cls = getattr(import_module(data['__module__']), data['__name__']) + cls = getattr(import_module(data["__module__"]), data["__name__"]) except ImportError: logger.debug( "Module {}.{} not found".format(data["__module__"], data["__name__"]) @@ -612,7 +708,7 @@ def from_hdf5(cls, filename=None, outdir=None, label=None): def from_json(cls, filename=None, outdir=None, label=None, gzip=False): from json.decoder import JSONDecodeError - filename = _determine_file_name(filename, outdir, label, 'json', gzip) + filename = _determine_file_name(filename, outdir, label, "json", gzip) if os.path.isfile(filename): try: @@ -629,22 +725,28 @@ def from_json(cls, filename=None, outdir=None, label=None, gzip=False): raise IOError("No result '{}' found".format(filename)) def __str__(self): - """Print a summary """ - if getattr(self, 'posterior', None) is not None: - if getattr(self, 'log_noise_evidence', None) is not None: - return ("nsamples: {:d}\n" - "ln_noise_evidence: {:6.3f}\n" - "ln_evidence: {:6.3f} +/- {:6.3f}\n" - "ln_bayes_factor: {:6.3f} +/- {:6.3f}\n" - .format(len(self.posterior), self.log_noise_evidence, self.log_evidence, - self.log_evidence_err, self.log_bayes_factor, - self.log_evidence_err)) + """Print a summary""" + if getattr(self, "posterior", None) is not None: + if getattr(self, "log_noise_evidence", None) is not None: + return ( + "nsamples: {:d}\n" + "ln_noise_evidence: {:6.3f}\n" + "ln_evidence: {:6.3f} +/- {:6.3f}\n" + "ln_bayes_factor: {:6.3f} +/- {:6.3f}\n".format( + len(self.posterior), + self.log_noise_evidence, + self.log_evidence, + self.log_evidence_err, + self.log_bayes_factor, + self.log_evidence_err, + ) + ) else: - return ("nsamples: {:d}\n" - "ln_evidence: {:6.3f} +/- {:6.3f}\n" - .format(len(self.posterior), self.log_evidence, self.log_evidence_err)) + return "nsamples: {:d}\nln_evidence: {:6.3f} +/- {:6.3f}\n".format( + len(self.posterior), self.log_evidence, self.log_evidence_err + ) else: - return '' + return "" @property def meta_data(self): @@ -662,7 +764,7 @@ def priors(self): if self._priors is not None: return self._priors else: - raise ValueError('Result object has no priors') + raise ValueError("Result object has no priors") @priors.setter def priors(self, priors): @@ -672,12 +774,14 @@ def priors(self, priors): else: self._priors = PriorDict(priors) if self.parameter_labels is None: - self.parameter_labels = [self.priors[k].latex_label for k in - self.search_parameter_keys] + self.parameter_labels = [ + self.priors[k].latex_label for k in self.search_parameter_keys + ] if self.parameter_labels_with_unit is None: self.parameter_labels_with_unit = [ - self.priors[k].latex_label_with_unit for k in - self.search_parameter_keys] + self.priors[k].latex_label_with_unit + for k in self.search_parameter_keys + ] elif priors is None: self._priors = priors self.parameter_labels = self.search_parameter_keys @@ -687,7 +791,7 @@ def priors(self, priors): @property def samples(self): - """ An array of samples """ + """An array of samples""" if self._samples is not None: return self._samples else: @@ -699,7 +803,7 @@ def samples(self, samples): @property def num_likelihood_evaluations(self): - """ number of likelihood evaluations """ + """number of likelihood evaluations""" if self._num_likelihood_evaluations is not None: return self._num_likelihood_evaluations else: @@ -711,7 +815,7 @@ def num_likelihood_evaluations(self, num_likelihood_evaluations): @property def nested_samples(self): - """" An array of unweighted samples """ + """ " An array of unweighted samples""" if self._nested_samples is not None: return self._nested_samples else: @@ -723,7 +827,7 @@ def nested_samples(self, nested_samples): @property def walkers(self): - """" An array of the ensemble walkers """ + """ " An array of the ensemble walkers""" if self._walkers is not None: return self._walkers else: @@ -735,7 +839,7 @@ def walkers(self, walkers): @property def nburn(self): - """" An array of the ensemble walkers """ + """ " An array of the ensemble walkers""" if self._nburn is not None: return self._nburn else: @@ -747,7 +851,7 @@ def nburn(self, nburn): @property def posterior(self): - """ A pandas data frame of the posterior """ + """A pandas data frame of the posterior""" if self._posterior is not None: return self._posterior else: @@ -795,22 +899,42 @@ def version(self): @version.setter def version(self, version): if version is None: - self._version = 'bilby={}'.format(utils.get_version_information()) + self._version = "bilby={}".format(utils.get_version_information()) else: self._version = version def _get_save_data_dictionary(self): # This list defines all the parameters saved in the result object save_attrs = [ - 'label', 'outdir', 'sampler', 'log_evidence', 'log_evidence_err', - 'log_noise_evidence', 'log_bayes_factor', 'priors', 'posterior', - 'injection_parameters', 'meta_data', 'search_parameter_keys', - 'fixed_parameter_keys', 'constraint_parameter_keys', - 'sampling_time', 'sampler_kwargs', 'use_ratio', 'information_gain', - 'log_likelihood_evaluations', 'log_prior_evaluations', - 'num_likelihood_evaluations', 'samples', 'nested_samples', - 'walkers', 'nburn', 'parameter_labels', 'parameter_labels_with_unit', - 'version'] + "label", + "outdir", + "sampler", + "log_evidence", + "log_evidence_err", + "log_noise_evidence", + "log_bayes_factor", + "priors", + "posterior", + "injection_parameters", + "meta_data", + "search_parameter_keys", + "fixed_parameter_keys", + "constraint_parameter_keys", + "sampling_time", + "sampler_kwargs", + "use_ratio", + "information_gain", + "log_likelihood_evaluations", + "log_prior_evaluations", + "num_likelihood_evaluations", + "samples", + "nested_samples", + "walkers", + "nburn", + "parameter_labels", + "parameter_labels_with_unit", + "version", + ] dictionary = dict() for attr in save_attrs: try: @@ -820,8 +944,9 @@ def _get_save_data_dictionary(self): pass return dictionary - def save_to_file(self, filename=None, overwrite=False, outdir=None, - extension=None, gzip=False): + def save_to_file( + self, filename=None, overwrite=False, outdir=None, extension=None, gzip=False + ): """ Writes the Result to a file. @@ -872,7 +997,7 @@ def save_to_file(self, filename=None, overwrite=False, outdir=None, if extension is None: logger.info("No extension given, defaulting to JSON.") - extension = 'json' + extension = "json" if extension is True: message = ( @@ -881,7 +1006,7 @@ def save_to_file(self, filename=None, overwrite=False, outdir=None, "This behaviour is deprecated and will be removed. " ) logger.warning(message) - extension = 'json' + extension = "json" outdir = _outdir if outdir is None else outdir outdir = self._safe_outdir_creation(outdir, self.save_to_file) @@ -896,31 +1021,37 @@ def save_to_file(self, filename=None, overwrite=False, outdir=None, dictionary = self._get_save_data_dictionary() # Convert callable sampler_kwargs to strings - if dictionary.get('sampler_kwargs', None) is not None: - for key in dictionary['sampler_kwargs']: - if hasattr(dictionary['sampler_kwargs'][key], '__call__'): - dictionary['sampler_kwargs'][key] = str(dictionary['sampler_kwargs']) + if dictionary.get("sampler_kwargs", None) is not None: + for key in dictionary["sampler_kwargs"]: + if hasattr(dictionary["sampler_kwargs"][key], "__call__"): + dictionary["sampler_kwargs"][key] = str( + dictionary["sampler_kwargs"] + ) try: # convert priors to JSON dictionary for both JSON and hdf5 files - if extension == 'json': + if extension == "json": dictionary["priors"] = dictionary["priors"]._get_json_dict() if gzip: import gzip + # encode to a string - json_str = json.dumps(dictionary, cls=BilbyJsonEncoder).encode('utf-8') - with gzip.GzipFile(output_path, 'w') as file: + json_str = json.dumps(dictionary, cls=BilbyJsonEncoder).encode( + "utf-8" + ) + with gzip.GzipFile(output_path, "w") as file: file.write(json_str) else: - with open(output_path, 'w') as file: + with open(output_path, "w") as file: json.dump(dictionary, file, indent=2, cls=BilbyJsonEncoder) - elif extension in ['hdf5', 'h5']: + elif extension in ["hdf5", "h5"]: import h5py + dictionary["__module__"] = self.__module__ dictionary["__name__"] = self.__class__.__name__ - with h5py.File(output_path, 'w') as h5file: - recursively_save_dict_contents_to_group(h5file, '/', dictionary) - elif extension in ['pkl', 'pickle']: + with h5py.File(output_path, "w") as h5file: + recursively_save_dict_contents_to_group(h5file, "/", dictionary) + elif extension in ["pkl", "pickle"]: safe_file_dump(self, output_path, "dill") else: raise ValueError(f"Extension type {extension} not understood") @@ -933,7 +1064,7 @@ def save_to_file(self, filename=None, overwrite=False, outdir=None, ) def save_posterior_samples(self, filename=None, outdir=None, label=None): - """ Saves posterior samples to a file + """Saves posterior samples to a file Generates a .dat file containing the posterior samples and auxiliary data saved in the posterior. Note, strings in the posterior are @@ -953,7 +1084,7 @@ def save_posterior_samples(self, filename=None, outdir=None, label=None): if label is None: label = self.label outdir = self._safe_outdir_creation(outdir, self.save_posterior_samples) - filename = '{}/{}_posterior_samples.dat'.format(outdir, label) + filename = "{}/{}_posterior_samples.dat".format(outdir, label) else: outdir = os.path.dirname(filename) self._safe_outdir_creation(outdir, self.save_posterior_samples) @@ -969,10 +1100,10 @@ def save_posterior_samples(self, filename=None, outdir=None, label=None): df.loc[:, key + "_angle"] = np.angle(complex_term) logger.info("Writing samples file to {}".format(filename)) - df.to_csv(filename, index=False, header=True, sep=' ') + df.to_csv(filename, index=False, header=True, sep=" ") def get_latex_labels_from_parameter_keys(self, keys): - """ Returns a list of latex_labels corresponding to the given keys + """Returns a list of latex_labels corresponding to the given keys Parameters ========== @@ -993,9 +1124,7 @@ def get_latex_labels_from_parameter_keys(self, keys): label = key else: label = None - logger.debug( - 'key {} not a parameter label or latex label'.format(key) - ) + logger.debug("key {} not a parameter label or latex label".format(key)) if label is None: label = key.replace("_", " ") latex_labels.append(label) @@ -1003,26 +1132,25 @@ def get_latex_labels_from_parameter_keys(self, keys): @property def covariance_matrix(self): - """ The covariance matrix of the samples the posterior """ + """The covariance matrix of the samples the posterior""" samples = self.posterior[self.search_parameter_keys].values return np.cov(samples.T) @property def posterior_volume(self): - """ The posterior volume """ + """The posterior volume""" if self.covariance_matrix.ndim == 0: return np.sqrt(self.covariance_matrix) else: - return 1 / np.sqrt(np.abs(np.linalg.det( - 1 / self.covariance_matrix))) + return 1 / np.sqrt(np.abs(np.linalg.det(1 / self.covariance_matrix))) @staticmethod def prior_volume(priors): - """ The prior volume, given a set of priors """ + """The prior volume, given a set of priors""" return np.prod([priors[k].maximum - priors[k].minimum for k in priors]) def occam_factor(self, priors): - """ The Occam factor, + """The Occam factor, See Chapter 28, `Mackay "Information Theory, Inference, and Learning Algorithms" `_ Cambridge @@ -1033,7 +1161,7 @@ def occam_factor(self, priors): @property def bayesian_model_dimensionality(self): - """ Characterises how many parameters are effectively constraint by the data + """Characterises how many parameters are effectively constraint by the data See @@ -1041,12 +1169,15 @@ def bayesian_model_dimensionality(self): ======= float: The model dimensionality """ - return 2 * (np.mean(self.posterior['log_likelihood']**2) - - np.mean(self.posterior['log_likelihood'])**2) + return 2 * ( + np.mean(self.posterior["log_likelihood"] ** 2) + - np.mean(self.posterior["log_likelihood"]) ** 2 + ) - def get_one_dimensional_median_and_error_bar(self, key, fmt='.2f', - quantiles=(0.16, 0.84)): - """ Calculate the median and error bar for a given key + def get_one_dimensional_median_and_error_bar( + self, key, fmt=".2f", quantiles=(0.16, 0.84) + ): + """Calculate the median and error bar for a given key Parameters ========== @@ -1064,7 +1195,7 @@ def get_one_dimensional_median_and_error_bar(self, key, fmt='.2f', An object with attributes, median, lower, upper and string """ - summary = namedtuple('summary', ['median', 'lower', 'upper', 'string']) + summary = namedtuple("summary", ["median", "lower", "upper", "string"]) if len(quantiles) != 2: raise ValueError("quantiles must be of length 2") @@ -1078,15 +1209,27 @@ def get_one_dimensional_median_and_error_bar(self, key, fmt='.2f', fmt = "{{0:{0}}}".format(fmt).format string_template = r"${{{0}}}_{{-{1}}}^{{+{2}}}$" summary.string = string_template.format( - fmt(summary.median), fmt(summary.minus), fmt(summary.plus)) + fmt(summary.median), fmt(summary.minus), fmt(summary.plus) + ) return summary @latex_plot_format - def plot_single_density(self, key, prior=None, cumulative=False, - title=None, truth=None, save=True, - file_base_name=None, bins=50, label_fontsize=16, - title_fontsize=16, quantiles=(0.16, 0.84), dpi=300): - """ Plot a 1D marginal density, either probability or cumulative. + def plot_single_density( + self, + key, + prior=None, + cumulative=False, + title=None, + truth=None, + save=True, + file_base_name=None, + bins=50, + label_fontsize=16, + title_fontsize=16, + quantiles=(0.16, 0.84), + dpi=300, + ): + """Plot a 1D marginal density, either probability or cumulative. Parameters ========== @@ -1127,52 +1270,68 @@ def plot_single_density(self, key, prior=None, cumulative=False, A matplotlib figure object """ import matplotlib.pyplot as plt - logger.info('Plotting {} marginal distribution'.format(key)) + + logger.info("Plotting {} marginal distribution".format(key)) label = self.get_latex_labels_from_parameter_keys([key])[0] label = sanity_check_labels([label])[0] fig, ax = plt.subplots() try: - ax.hist(self.posterior[key].values, bins=bins, density=True, - histtype='step', cumulative=cumulative) + ax.hist( + self.posterior[key].values, + bins=bins, + density=True, + histtype="step", + cumulative=cumulative, + ) except ValueError as e: logger.info( - 'Failed to generate 1d plot for {}, error message: {}' - .format(key, e)) + "Failed to generate 1d plot for {}, error message: {}".format(key, e) + ) return ax.set_xlabel(label, fontsize=label_fontsize) if truth is not None: - ax.axvline(truth, ls='-', color='orange') + ax.axvline(truth, ls="-", color="orange") summary = self.get_one_dimensional_median_and_error_bar( - key, quantiles=quantiles) - ax.axvline(summary.median - summary.minus, ls='--', color='C0') - ax.axvline(summary.median + summary.plus, ls='--', color='C0') + key, quantiles=quantiles + ) + ax.axvline(summary.median - summary.minus, ls="--", color="C0") + ax.axvline(summary.median + summary.plus, ls="--", color="C0") if title: ax.set_title(summary.string, fontsize=title_fontsize) if isinstance(prior, Prior): theta = np.linspace(ax.get_xlim()[0], ax.get_xlim()[1], 300) if cumulative is False: - ax.plot(theta, prior.prob(theta), color='C2') + ax.plot(theta, prior.prob(theta), color="C2") else: - ax.plot(theta, prior.cdf(theta), color='C2') + ax.plot(theta, prior.cdf(theta), color="C2") if save: fig.tight_layout() if cumulative: - file_name = file_base_name + key + '_cdf' + file_name = file_base_name + key + "_cdf" else: - file_name = file_base_name + key + '_pdf' + file_name = file_base_name + key + "_pdf" safe_save_figure(fig=fig, filename=file_name, dpi=dpi) plt.close(fig) else: return fig - def plot_marginals(self, parameters=None, priors=None, titles=True, - file_base_name=None, bins=50, label_fontsize=16, - title_fontsize=16, quantiles=(0.16, 0.84), dpi=300, - outdir=None): - """ Plot 1D marginal distributions + def plot_marginals( + self, + parameters=None, + priors=None, + titles=True, + file_base_name=None, + bins=50, + label_fontsize=16, + title_fontsize=16, + quantiles=(0.16, 0.84), + dpi=300, + outdir=None, + ): + """Plot 1D marginal distributions Parameters ========== @@ -1224,17 +1383,17 @@ def plot_marginals(self, parameters=None, priors=None, titles=True, if file_base_name is None: outdir = self._safe_outdir_creation(outdir, self.plot_marginals) - file_base_name = '{}/{}_1d/'.format(outdir, self.label) + file_base_name = "{}/{}_1d/".format(outdir, self.label) check_directory_exists_and_if_not_mkdir(file_base_name) if priors is True: - priors = getattr(self, 'priors', dict()) + priors = getattr(self, "priors", dict()) elif isinstance(priors, dict): pass elif priors in [False, None]: priors = dict() else: - raise ValueError('Input priors={} not understood'.format(priors)) + raise ValueError("Input priors={} not understood".format(priors)) for i, key in enumerate(plot_parameter_keys): if not isinstance(self.posterior[key].values[0], float): @@ -1243,15 +1402,32 @@ def plot_marginals(self, parameters=None, priors=None, titles=True, truth = truths.get(key, None) for cumulative in [False, True]: self.plot_single_density( - key, prior=prior, cumulative=cumulative, title=titles, - truth=truth, save=True, file_base_name=file_base_name, - bins=bins, label_fontsize=label_fontsize, dpi=dpi, - title_fontsize=title_fontsize, quantiles=quantiles) + key, + prior=prior, + cumulative=cumulative, + title=titles, + truth=truth, + save=True, + file_base_name=file_base_name, + bins=bins, + label_fontsize=label_fontsize, + dpi=dpi, + title_fontsize=title_fontsize, + quantiles=quantiles, + ) @latex_plot_format - def plot_corner(self, parameters=None, priors=None, titles=True, save=True, - filename=None, dpi=300, **kwargs): - """ Plot a corner-plot + def plot_corner( + self, + parameters=None, + priors=None, + titles=True, + save=True, + filename=None, + dpi=300, + **kwargs, + ): + """Plot a corner-plot Parameters ========== @@ -1305,16 +1481,22 @@ def plot_corner(self, parameters=None, priors=None, titles=True, save=True, return defaults_kwargs = dict( - bins=50, smooth=0.9, - title_kwargs=dict(fontsize=16), color='#0072C1', - truth_color='tab:orange', quantiles=[0.16, 0.84], - levels=(1 - np.exp(-0.5), 1 - np.exp(-2), 1 - np.exp(-9 / 2.)), - plot_density=False, plot_datapoints=True, fill_contours=True, - max_n_ticks=3) + bins=50, + smooth=0.9, + title_kwargs=dict(fontsize=16), + color="#0072C1", + truth_color="tab:orange", + quantiles=[0.16, 0.84], + levels=(1 - np.exp(-0.5), 1 - np.exp(-2), 1 - np.exp(-9 / 2.0)), + plot_density=False, + plot_datapoints=True, + fill_contours=True, + max_n_ticks=3, + ) - if 'lionize' in kwargs and kwargs['lionize'] is True: - defaults_kwargs['truth_color'] = 'tab:blue' - defaults_kwargs['color'] = '#FF8C00' + if "lionize" in kwargs and kwargs["lionize"] is True: + defaults_kwargs["truth_color"] = "tab:blue" + defaults_kwargs["color"] = "#FF8C00" label_kwargs_defaults = dict(fontsize=16) hist_kwargs_defaults = dict(density=True) @@ -1332,27 +1514,25 @@ def plot_corner(self, parameters=None, priors=None, titles=True, save=True, kwargs["hist_kwargs"] = hist_kwargs_defaults # Handle if truths was passed in - if 'truth' in kwargs: - kwargs['truths'] = kwargs.pop('truth') + if "truth" in kwargs: + kwargs["truths"] = kwargs.pop("truth") if "truths" in kwargs: - truths = kwargs.get('truths') + truths = kwargs.get("truths") if isinstance(parameters, list) and isinstance(truths, list): if len(parameters) != len(truths): - raise ValueError( - "Length of parameters and truths don't match") + raise ValueError("Length of parameters and truths don't match") elif isinstance(truths, dict) and parameters is None: - parameters = kwargs.pop('truths') + parameters = kwargs.pop("truths") elif isinstance(truths, bool): pass elif truths is None: kwargs["truths"] = False else: - raise ValueError( - "Combination of parameters and truths not understood") + raise ValueError("Combination of parameters and truths not understood") # If injection parameters where stored, use these as parameter values # but do not overwrite input parameters (or truths) - cond1 = getattr(self, 'injection_parameters', None) is not None + cond1 = getattr(self, "injection_parameters", None) is not None cond2 = parameters is None cond3 = bool(kwargs.get("truths", True)) if cond1 and cond2 and cond3: @@ -1365,26 +1545,26 @@ def plot_corner(self, parameters=None, priors=None, titles=True, save=True, # parameters to plot and the values as truths. if isinstance(parameters, dict): plot_parameter_keys = list(parameters.keys()) - kwargs['truths'] = list(parameters.values()) + kwargs["truths"] = list(parameters.values()) elif parameters is None: plot_parameter_keys = self.search_parameter_keys else: plot_parameter_keys = list(parameters) # Get latex formatted strings for the plot labels - kwargs['labels'] = kwargs.get( - 'labels', self.get_latex_labels_from_parameter_keys( - plot_parameter_keys)) + kwargs["labels"] = kwargs.get( + "labels", self.get_latex_labels_from_parameter_keys(plot_parameter_keys) + ) kwargs["labels"] = sanity_check_labels(kwargs["labels"]) # Unless already set, set the range to include all samples # This prevents ValueErrors being raised for parameters with no range - kwargs['range'] = kwargs.get('range', [1] * len(plot_parameter_keys)) + kwargs["range"] = kwargs.get("range", [1] * len(plot_parameter_keys)) # Remove truths if it is a bool - if isinstance(kwargs.get('truths'), bool): - kwargs.pop('truths') + if isinstance(kwargs.get("truths"), bool): + kwargs.pop("truths") # Create the data array to plot and pass everything to corner xs = self.posterior[plot_parameter_keys].values @@ -1392,40 +1572,50 @@ def plot_corner(self, parameters=None, priors=None, titles=True, save=True, fig = corner.corner(xs, **kwargs) else: ax = kwargs.get("ax", plt.subplot()) - ax.hist(xs, bins=kwargs["bins"], color=kwargs["color"], - histtype="step", **kwargs["hist_kwargs"]) + ax.hist( + xs, + bins=kwargs["bins"], + color=kwargs["color"], + histtype="step", + **kwargs["hist_kwargs"], + ) ax.set_xlabel(kwargs["labels"][0]) fig = plt.gcf() axes = fig.get_axes() # Add the titles - if titles and kwargs.get('quantiles', None) is not None: + if titles and kwargs.get("quantiles", None) is not None: for i, par in enumerate(plot_parameter_keys): ax = axes[i + i * len(plot_parameter_keys)] - if ax.title.get_text() == '': - ax.set_title(self.get_one_dimensional_median_and_error_bar( - par, quantiles=kwargs['quantiles']).string, - **kwargs['title_kwargs']) + if ax.title.get_text() == "": + ax.set_title( + self.get_one_dimensional_median_and_error_bar( + par, quantiles=kwargs["quantiles"] + ).string, + **kwargs["title_kwargs"], + ) # Add priors to the 1D plots if priors is True: - priors = getattr(self, 'priors', False) + priors = getattr(self, "priors", False) if isinstance(priors, dict): for i, par in enumerate(plot_parameter_keys): ax = axes[i + i * len(plot_parameter_keys)] theta = np.linspace(ax.get_xlim()[0], ax.get_xlim()[1], 300) - ax.plot(theta, priors[par].prob(theta), color='C2') + ax.plot(theta, priors[par].prob(theta), color="C2") elif priors in [False, None]: pass else: - raise ValueError('Input priors={} not understood'.format(priors)) + raise ValueError("Input priors={} not understood".format(priors)) if save: if filename is None: - outdir = self._safe_outdir_creation(kwargs.get('outdir'), self.plot_corner) - filename = '{}/{}_corner.png'.format(outdir, self.label) - logger.debug('Saving corner plot to {}'.format(filename)) + outdir = self._safe_outdir_creation( + kwargs.get("outdir"), self.plot_corner + ) + filename = "{}/{}_corner.png".format(outdir, self.label) + logger.debug("Saving corner plot to {}".format(filename)) safe_save_figure(fig=fig, filename=filename, dpi=dpi) plt.close(fig) @@ -1433,9 +1623,10 @@ def plot_corner(self, parameters=None, priors=None, titles=True, save=True, @latex_plot_format def plot_walkers(self, **kwargs): - """ Method to plot the trace of the walkers in an ensemble MCMC plot """ + """Method to plot the trace of the walkers in an ensemble MCMC plot""" import matplotlib.pyplot as plt - if hasattr(self, 'walkers') is False: + + if hasattr(self, "walkers") is False: logger.warning("Cannot plot_walkers as no walkers are saved") return @@ -1448,28 +1639,46 @@ def plot_walkers(self, **kwargs): walkers = self.walkers[:, :, :] parameter_labels = sanity_check_labels(self.parameter_labels) for i, ax in enumerate(axes): - ax.plot(idxs[:self.nburn + 1], walkers[:, :self.nburn + 1, i].T, - lw=0.1, color='r') + ax.plot( + idxs[: self.nburn + 1], + walkers[:, : self.nburn + 1, i].T, + lw=0.1, + color="r", + ) ax.set_ylabel(parameter_labels[i]) for i, ax in enumerate(axes): - ax.plot(idxs[self.nburn:], walkers[:, self.nburn:, i].T, lw=0.1, - color='k') + ax.plot( + idxs[self.nburn :], walkers[:, self.nburn :, i].T, lw=0.1, color="k" + ) ax.set_ylabel(parameter_labels[i]) fig.tight_layout() - outdir = self._safe_outdir_creation(kwargs.get('outdir'), self.plot_walkers) - filename = '{}/{}_walkers.png'.format(outdir, self.label) - logger.debug('Saving walkers plot to {}'.format('filename')) + outdir = self._safe_outdir_creation(kwargs.get("outdir"), self.plot_walkers) + filename = "{}/{}_walkers.png".format(outdir, self.label) + logger.debug("Saving walkers plot to {}".format("filename")) safe_save_figure(fig=fig, filename=filename) plt.close(fig) @latex_plot_format - def plot_with_data(self, model, x, y, ndraws=1000, npoints=1000, - xlabel=None, ylabel=None, data_label='data', - data_fmt='o', draws_label=None, filename=None, - maxl_label='max likelihood', dpi=300, outdir=None): - """ Generate a figure showing the data and fits to the data + def plot_with_data( + self, + model, + x, + y, + ndraws=1000, + npoints=1000, + xlabel=None, + ylabel=None, + data_label="data", + data_fmt="o", + draws_label=None, + filename=None, + maxl_label="max likelihood", + dpi=300, + outdir=None, + ): + """Generate a figure showing the data and fits to the data Parameters ========== @@ -1507,20 +1716,24 @@ def plot_with_data(self, model, x, y, ndraws=1000, npoints=1000, xsmooth = np.linspace(np.min(x), np.max(x), npoints) fig, ax = plt.subplots() - logger.info('Plotting {} draws'.format(ndraws)) + logger.info("Plotting {} draws".format(ndraws)) for _ in range(ndraws): - s = model_posterior.sample().to_dict('records')[0] - ax.plot(xsmooth, model(xsmooth, **s), alpha=0.25, lw=0.1, color='r', - label=draws_label) + s = model_posterior.sample().to_dict("records")[0] + ax.plot( + xsmooth, + model(xsmooth, **s), + alpha=0.25, + lw=0.1, + color="r", + label=draws_label, + ) try: if all(~np.isnan(self.posterior.log_likelihood)): - logger.info('Plotting maximum likelihood') + logger.info("Plotting maximum likelihood") s = model_posterior.iloc[self.posterior.log_likelihood.idxmax()] - ax.plot(xsmooth, model(xsmooth, **s), lw=1, color='k', - label=maxl_label) + ax.plot(xsmooth, model(xsmooth, **s), lw=1, color="k", label=maxl_label) except (AttributeError, TypeError): - logger.debug( - "No log likelihood values stored, unable to plot max") + logger.debug("No log likelihood values stored, unable to plot max") ax.plot(x, y, data_fmt, markersize=2, label=data_label) @@ -1536,7 +1749,7 @@ def plot_with_data(self, model, x, y, ndraws=1000, npoints=1000, fig.tight_layout() if filename is None: outdir = self._safe_outdir_creation(outdir, self.plot_with_data) - filename = '{}/{}_plot_with_data'.format(outdir, self.label) + filename = "{}/{}_plot_with_data".format(outdir, self.label) safe_save_figure(fig=fig, filename=filename, dpi=dpi) plt.close(fig) @@ -1545,15 +1758,17 @@ def _add_prior_fixed_values_to_posterior(posterior, priors): if priors is None: return posterior for key in priors: - if isinstance(priors[key], DeltaFunction) and \ - not isinstance(priors[key], ConditionalDeltaFunction): + if isinstance(priors[key], DeltaFunction) and not isinstance( + priors[key], ConditionalDeltaFunction + ): posterior[key] = priors[key].peak elif isinstance(priors[key], float): posterior[key] = priors[key] return posterior - def samples_to_posterior(self, likelihood=None, priors=None, - conversion_function=None, npool=1, pool=None): + def samples_to_posterior( + self, likelihood=None, priors=None, conversion_function=None, npool=1, pool=None + ): """ Convert array of samples to posterior (a Pandas data frame) @@ -1570,23 +1785,26 @@ def samples_to_posterior(self, likelihood=None, priors=None, should take the data_frame, likelihood and prior as arguments. """ - data_frame = pd.DataFrame( - self.samples, columns=self.search_parameter_keys) - data_frame = self._add_prior_fixed_values_to_posterior( - data_frame, priors) - data_frame['log_likelihood'] = getattr( - self, 'log_likelihood_evaluations', np.nan) + data_frame = pd.DataFrame(self.samples, columns=self.search_parameter_keys) + data_frame = self._add_prior_fixed_values_to_posterior(data_frame, priors) + data_frame["log_likelihood"] = getattr( + self, "log_likelihood_evaluations", np.nan + ) if self.log_prior_evaluations is None and priors is not None: - data_frame['log_prior'] = priors.ln_prob( - dict(data_frame[self.search_parameter_keys]), axis=0) + data_frame["log_prior"] = priors.ln_prob( + dict(data_frame[self.search_parameter_keys]), axis=0 + ) else: - data_frame['log_prior'] = self.log_prior_evaluations + data_frame["log_prior"] = self.log_prior_evaluations if conversion_function is not None: parameters = inspect.signature(conversion_function).parameters kwargs = dict() for key, value in [ - ("likelihood", likelihood), ("priors", priors), ("npool", npool), ("pool", pool) + ("likelihood", likelihood), + ("priors", priors), + ("npool", npool), + ("pool", pool), ]: if key in parameters: kwargs[key] = value @@ -1608,8 +1826,9 @@ def calculate_prior_values(self, priors): if isinstance(priors[key], DeltaFunction): continue else: - self.prior_values[key]\ - = priors[key].prob(self.posterior[key].values) + self.prior_values[key] = priors[key].prob( + self.posterior[key].values + ) def get_all_injection_credible_levels(self, keys=None, weights=None): """ @@ -1637,9 +1856,11 @@ def get_all_injection_credible_levels(self, keys=None, weights=None): "Result object has no 'injection_parameters'. " "Cannot compute credible levels." ) - credible_levels = {key: self.get_injection_credible_level(key, weights=weights) - for key in keys - if isinstance(self.injection_parameters.get(key, None), float)} + credible_levels = { + key: self.get_injection_credible_level(key, weights=weights) + for key in keys + if isinstance(self.injection_parameters.get(key, None), float) + } return credible_levels def get_injection_credible_level(self, parameter, weights=None): @@ -1665,23 +1886,26 @@ def get_injection_credible_level(self, parameter, weights=None): raise ( TypeError, "Result object has no 'injection_parameters'. " - "Cannot copmute credible levels." + "Cannot copmute credible levels.", ) if weights is None: weights = np.ones(len(self.posterior)) - if parameter in self.posterior and\ - parameter in self.injection_parameters: - credible_level =\ - sum(np.array(self.posterior[parameter].values < - self.injection_parameters[parameter]) * weights) / (sum(weights)) + if parameter in self.posterior and parameter in self.injection_parameters: + credible_level = sum( + np.array( + self.posterior[parameter].values + < self.injection_parameters[parameter] + ) + * weights + ) / (sum(weights)) return credible_level else: return np.nan def _check_attribute_match_to_other_object(self, name, other_object): - """ Check attribute name exists in other_object and is the same + """Check attribute name exists in other_object and is the same Parameters ========== @@ -1697,7 +1921,7 @@ def _check_attribute_match_to_other_object(self, name, other_object): """ a = getattr(self, name, False) b = getattr(other_object, name, False) - logger.debug('Checking {} value: {}=={}'.format(name, a, b)) + logger.debug("Checking {} value: {}=={}".format(name, a, b)) if (a is not False) and (b is not False): type_a = type(a) type_b = type(b) @@ -1713,7 +1937,7 @@ def _check_attribute_match_to_other_object(self, name, other_object): @property def kde(self): - """ Kernel density estimate built from the stored posterior + """Kernel density estimate built from the stored posterior Uses `scipy.stats.gaussian_kde` to generate the kernel density """ @@ -1721,11 +1945,12 @@ def kde(self): return self._kde else: self._kde = scipy.stats.gaussian_kde( - self.posterior[self.search_parameter_keys].values.T) + self.posterior[self.search_parameter_keys].values.T + ) return self._kde def posterior_probability(self, sample): - """ Calculate the posterior probability for a new sample + """Calculate the posterior probability for a new sample This queries a Kernel Density Estimate of the posterior to calculate the posterior probability density for the new sample. @@ -1745,8 +1970,9 @@ def posterior_probability(self, sample): """ if isinstance(sample, dict): sample = [sample] - ordered_sample = [[s[key] for key in self.search_parameter_keys] - for s in sample] + ordered_sample = [ + [s[key] for key in self.search_parameter_keys] for s in sample + ] return self.kde(ordered_sample) def _safe_outdir_creation(self, outdir=None, caller_func=None): @@ -1755,49 +1981,60 @@ def _safe_outdir_creation(self, outdir=None, caller_func=None): try: utils.check_directory_exists_and_if_not_mkdir(outdir) except PermissionError: - raise FileMovedError("Can not write in the out directory.\n" - "Did you move the here file from another system?\n" - "Try calling " + caller_func.__name__ + " with the 'outdir' " - "keyword argument, e.g. " + caller_func.__name__ + "(outdir='.')") + raise FileMovedError( + "Can not write in the out directory.\n" + "Did you move the here file from another system?\n" + "Try calling " + caller_func.__name__ + " with the 'outdir' " + "keyword argument, e.g. " + caller_func.__name__ + "(outdir='.')" + ) return outdir def get_weights_by_new_prior(self, old_prior, new_prior, prior_names=None): - """ Calculate a list of sample weights based on the ratio of new to old priors + """Calculate a list of sample weights based on the ratio of new to old priors - Parameters - ========== - old_prior: PriorDict, - The prior used in the generation of the original samples. + Parameters + ========== + old_prior: PriorDict, + The prior used in the generation of the original samples. - new_prior: PriorDict, - The prior to use to reweight the samples. + new_prior: PriorDict, + The prior to use to reweight the samples. - prior_names: list - A list of the priors to include in the ratio during reweighting. + prior_names: list + A list of the priors to include in the ratio during reweighting. - Returns - ======= - weights: array-like, - A list of sample weights. + Returns + ======= + weights: array-like, + A list of sample weights. - """ + """ weights = [] # Shared priors - these will form a ratio if prior_names is not None: - shared_parameters = {key: self.posterior[key] for key in new_prior if - key in old_prior and key in prior_names} + shared_parameters = { + key: self.posterior[key] + for key in new_prior + if key in old_prior and key in prior_names + } else: - shared_parameters = {key: self.posterior[key] for key in new_prior if key in old_prior} - parameters = [{key: self.posterior[key][i] for key in shared_parameters.keys()} - for i in range(len(self.posterior))] + shared_parameters = { + key: self.posterior[key] for key in new_prior if key in old_prior + } + parameters = [ + {key: self.posterior[key][i] for key in shared_parameters.keys()} + for i in range(len(self.posterior)) + ] for i in range(len(self.posterior)): weight = 1 for prior_key in shared_parameters.keys(): val = self.posterior[prior_key][i] weight *= new_prior.evaluate_constraints(parameters[i]) - weight *= new_prior[prior_key].prob(val) / old_prior[prior_key].prob(val) + weight *= new_prior[prior_key].prob(val) / old_prior[prior_key].prob( + val + ) weights.append(weight) @@ -1837,14 +2074,10 @@ def to_arviz(self, prior=None): posdict[key] = self.posterior[key].values if "log_likelihood" in posdict: - loglikedict = { - "log_likelihood": posdict.pop("log_likelihood") - } + loglikedict = {"log_likelihood": posdict.pop("log_likelihood")} else: if self.log_likelihood_evaluations is not None: - loglikedict = { - "log_likelihood": self.log_likelihood_evaluations - } + loglikedict = {"log_likelihood": self.log_likelihood_evaluations} else: loglikedict = None @@ -1866,7 +2099,7 @@ def to_arviz(self, prior=None): az_version = packaging.version.parse(importlib.metadata.version("arviz")) version = { "inference_library": "bilby: {}".format(self.sampler), - "inference_library_version": get_version_information() + "inference_library_version": get_version_information(), } if az_version < packaging.version.parse("1"): azdata = az.from_dict(**az_data_dict) @@ -1887,9 +2120,8 @@ def to_arviz(self, prior=None): class ResultList(list): - def __init__(self, results=None, consistency_level="warning"): - """ A class to store a list of :class:`bilby.core.result.Result` objects + """A class to store a list of :class:`bilby.core.result.Result` objects from equivalent runs on the same data. This provides methods for outputting combined results. @@ -1959,7 +2191,7 @@ def combine(self, shuffle=False, consistency_level="error"): result = copy(self[0]) if result.label is not None: - result.label += '_combined' + result.label += "_combined" self.check_consistent_sampler() self.check_consistent_data() @@ -2006,17 +2238,22 @@ def _combine_nested_sampled_runs(self, result): The result object with the combined evidences. """ from scipy.special import logsumexp + self.check_nested_samples() # Combine evidences log_evidences = np.array([res.log_evidence for res in self]) - result.log_evidence = logsumexp(log_evidences, b=1. / len(self)) + result.log_evidence = logsumexp(log_evidences, b=1.0 / len(self)) result.log_bayes_factor = result.log_evidence - result.log_noise_evidence # Propagate uncertainty in combined evidence - log_errs = [res.log_evidence_err for res in self if np.isfinite(res.log_evidence_err)] + log_errs = [ + res.log_evidence_err for res in self if np.isfinite(res.log_evidence_err) + ] if len(log_errs) > 0: - result.log_evidence_err = 0.5 * logsumexp(2 * np.array(log_errs), b=1. / len(self)) + result.log_evidence_err = 0.5 * logsumexp( + 2 * np.array(log_errs), b=1.0 / len(self) + ) else: result.log_evidence_err = np.nan @@ -2024,7 +2261,7 @@ def _combine_nested_sampled_runs(self, result): result_weights = np.exp(log_evidences - np.max(log_evidences)) posteriors = list() for res, frac in zip(self, result_weights): - selected_samples = (random.rng.uniform(size=len(res.posterior)) < frac) + selected_samples = random.rng.uniform(size=len(res.posterior)) < frac posteriors.append(res.posterior[selected_samples]) # remove original nested_samples @@ -2056,13 +2293,17 @@ def _combine_mcmc_sampled_runs(self, result): # Combine evidences log_evidences = np.array([res.log_evidence for res in self]) - result.log_evidence = logsumexp(log_evidences, b=1. / len(self)) + result.log_evidence = logsumexp(log_evidences, b=1.0 / len(self)) result.log_bayes_factor = result.log_evidence - result.log_noise_evidence # Propagate uncertainty in combined evidence - log_errs = [res.log_evidence_err for res in self if np.isfinite(res.log_evidence_err)] + log_errs = [ + res.log_evidence_err for res in self if np.isfinite(res.log_evidence_err) + ] if len(log_errs) > 0: - result.log_evidence_err = 0.5 * logsumexp(2 * np.array(log_errs), b=1. / len(self)) + result.log_evidence_err = 0.5 * logsumexp( + 2 * np.array(log_errs), b=1.0 / len(self) + ) else: result.log_evidence_err = np.nan @@ -2088,17 +2329,26 @@ def _error_or_warning_consistency(self, msg): elif self.consistency_level == "ignore": pass else: - raise ValueError(f"Input consistency_level {self.consistency_level} not understood") + raise ValueError( + f"Input consistency_level {self.consistency_level} not understood" + ) def check_consistent_priors(self): for res in self: for p in self[0].priors.keys(): - if not self[0].priors[p] == res.priors[p] or len(self[0].priors) != len(res.priors): + if not self[0].priors[p] == res.priors[p] or len(self[0].priors) != len( + res.priors + ): msg = "Inconsistent priors between results" self._error_or_warning_consistency(msg) def check_consistent_parameters(self): - if not np.all([set(self[0].search_parameter_keys) == set(res.search_parameter_keys) for res in self]): + if not np.all( + [ + set(self[0].search_parameter_keys) == set(res.search_parameter_keys) + for res in self + ] + ): msg = "Inconsistent parameters between results" self._error_or_warning_consistency(msg) @@ -2120,9 +2370,18 @@ def check_consistent_sampler(self): @latex_plot_format -def plot_multiple(results, filename=None, labels=None, colours=None, - save=True, evidences=False, corner_labels=None, linestyles=None, - fig=None, **kwargs): +def plot_multiple( + results, + filename=None, + labels=None, + colours=None, + save=True, + evidences=False, + corner_labels=None, + linestyles=None, + fig=None, + **kwargs, +): """Generate a corner plot overlaying two sets of results Parameters @@ -2169,22 +2428,22 @@ def plot_multiple(results, filename=None, labels=None, colours=None, import matplotlib.pyplot as plt import matplotlib.lines as mpllines - kwargs['show_titles'] = False - kwargs['truths'] = None + kwargs["show_titles"] = False + kwargs["truths"] = None if corner_labels is not None: - kwargs['labels'] = corner_labels - default_filename = '{}/{}'.format(results[0].outdir, 'combined') + kwargs["labels"] = corner_labels + default_filename = "{}/{}".format(results[0].outdir, "combined") lines = [] default_labels = [] for i, result in enumerate(results): if colours: c = colours[i] else: - c = 'C{}'.format(i) + c = "C{}".format(i) if linestyles is not None: linestyle = linestyles[i] else: - linestyle = 'solid' + linestyle = "solid" hist_kwargs = kwargs.get("hist_kwargs", dict()) contour_kwargs = kwargs.get("contour_kwargs", dict()) hist_kwargs["color"] = c @@ -2194,7 +2453,7 @@ def plot_multiple(results, filename=None, labels=None, colours=None, kwargs["hist_kwargs"] = hist_kwargs kwargs["contour_kwargs"] = contour_kwargs fig = result.plot_corner(fig=fig, save=False, color=c, **kwargs) - default_filename += '_{}'.format(result.label) + default_filename += "_{}".format(result.label) lines.append(mpllines.Line2D([0], [0], color=c, linestyle=linestyle)) default_labels.append(result.label) @@ -2210,9 +2469,9 @@ def plot_multiple(results, filename=None, labels=None, colours=None, if evidences: if np.isnan(results[0].log_bayes_factor): - template = r'{label} $\mathrm{{ln}}(Z)={lnz:1.3g}$' + template = r"{label} $\mathrm{{ln}}(Z)={lnz:1.3g}$" else: - template = r'{label} $\mathrm{{ln}}(B)={lnbf:1.3g}$' + template = r"{label} $\mathrm{{ln}}(B)={lnbf:1.3g}$" labels = [ template.format( label=label, @@ -2235,10 +2494,19 @@ def plot_multiple(results, filename=None, labels=None, colours=None, @latex_plot_format -def make_pp_plot(results, filename=None, save=True, confidence_interval=[0.68, 0.95, 0.997], - lines=None, legend_fontsize='x-small', keys=None, title=True, - confidence_interval_alpha=0.1, weight_list=None, - **kwargs): +def make_pp_plot( + results, + filename=None, + save=True, + confidence_interval=[0.68, 0.95, 0.997], + lines=None, + legend_fontsize="x-small", + keys=None, + title=True, + confidence_interval_alpha=0.1, + weight_list=None, + **kwargs, +): """ Make a P-P plot for a set of runs with injected signals. @@ -2304,27 +2572,34 @@ def make_pp_plot(results, filename=None, save=True, confidence_interval=[0.68, 0 if isinstance(confidence_interval, float): confidence_interval = [confidence_interval] if isinstance(confidence_interval_alpha, float): - confidence_interval_alpha = [confidence_interval_alpha] * len(confidence_interval) + confidence_interval_alpha = [confidence_interval_alpha] * len( + confidence_interval + ) elif len(confidence_interval_alpha) != len(confidence_interval): raise ValueError( - "confidence_interval_alpha must have the same length as confidence_interval") + "confidence_interval_alpha must have the same length as confidence_interval" + ) for ci, alpha in zip(confidence_interval, confidence_interval_alpha): - edge_of_bound = (1. - ci) / 2. + edge_of_bound = (1.0 - ci) / 2.0 lower = scipy.stats.binom.ppf(1 - edge_of_bound, N, x_values) / N upper = scipy.stats.binom.ppf(edge_of_bound, N, x_values) / N # The binomial point percent function doesn't always return 0 @ 0, # so set those bounds explicitly to be sure lower[0] = 0 upper[0] = 0 - ax.fill_between(x_values, lower, upper, alpha=alpha, color='k') + ax.fill_between(x_values, lower, upper, alpha=alpha, color="k") pvalues = [] logger.info("Key: KS-test p-value") for ii, key in enumerate(credible_levels): - pp = np.array([sum(credible_levels[key].values < xx) / - len(credible_levels) for xx in x_values]) - pvalue = scipy.stats.kstest(credible_levels[key], 'uniform').pvalue + pp = np.array( + [ + sum(credible_levels[key].values < xx) / len(credible_levels) + for xx in x_values + ] + ) + pvalue = scipy.stats.kstest(credible_levels[key], "uniform").pvalue pvalues.append(pvalue) logger.info("{}: {}".format(key, pvalue)) @@ -2335,16 +2610,18 @@ def make_pp_plot(results, filename=None, save=True, confidence_interval=[0.68, 0 label = "{} ({:2.3f})".format(name, pvalue) plt.plot(x_values, pp, lines[ii], label=label, **kwargs) - Pvals = namedtuple('pvals', ['combined_pvalue', 'pvalues', 'names']) - pvals = Pvals(combined_pvalue=scipy.stats.combine_pvalues(pvalues)[1], - pvalues=pvalues, - names=list(credible_levels.keys())) - logger.info( - "Combined p-value: {}".format(pvals.combined_pvalue)) + Pvals = namedtuple("pvals", ["combined_pvalue", "pvalues", "names"]) + pvals = Pvals( + combined_pvalue=scipy.stats.combine_pvalues(pvalues)[1], + pvalues=pvalues, + names=list(credible_levels.keys()), + ) + logger.info("Combined p-value: {}".format(pvals.combined_pvalue)) if title: - ax.set_title("N={}, p-value={:2.4f}".format( - len(results), pvals.combined_pvalue)) + ax.set_title( + "N={}, p-value={:2.4f}".format(len(results), pvals.combined_pvalue) + ) ax.set_xlabel("C.I.") ax.set_ylabel("Fraction of events in C.I.") ax.legend(handlelength=2, labelspacing=0.25, fontsize=legend_fontsize) @@ -2353,14 +2630,14 @@ def make_pp_plot(results, filename=None, save=True, confidence_interval=[0.68, 0 fig.tight_layout() if save: if filename is None: - filename = 'outdir/pp.png' + filename = "outdir/pp.png" safe_save_figure(fig=fig, filename=filename, dpi=500) return fig, pvals def sanity_check_labels(labels): - """ Check labels for plotting to remove matplotlib errors """ + """Check labels for plotting to remove matplotlib errors""" for ii, lab in enumerate(labels): if "_" in lab and "$" not in lab: lab = lab.replace("_", "-") @@ -2369,16 +2646,16 @@ def sanity_check_labels(labels): class ResultError(Exception): - """ Base exception for all Result related errors """ + """Base exception for all Result related errors""" class ResultListError(ResultError): - """ For Errors occurring during combining results. """ + """For Errors occurring during combining results.""" class FileMovedError(ResultError): - """ Exceptions that occur when files have been moved """ + """Exceptions that occur when files have been moved""" class FileLoadError(ResultError): - """ Exceptions that occur when files cannot be loaded """ + """Exceptions that occur when files cannot be loaded""" diff --git a/bilby/core/sampler/base_sampler.py b/bilby/core/sampler/base_sampler.py index a0a943406..b67caabf0 100644 --- a/bilby/core/sampler/base_sampler.py +++ b/bilby/core/sampler/base_sampler.py @@ -717,7 +717,6 @@ def write_current_state_and_exit(self, signum=None, frame=None): can be caught as a :code:`SystemExit`. """ if self.npool in (1, None) or getattr(self, "pool", None) is not None: - self._interrupted = True self._interruption_signum = signum diff --git a/bilby/core/sampler/dynesty.py b/bilby/core/sampler/dynesty.py index 7a9bb554d..80d60167b 100644 --- a/bilby/core/sampler/dynesty.py +++ b/bilby/core/sampler/dynesty.py @@ -814,8 +814,7 @@ def write_current_state(self): logger.info(f"Written checkpoint file {self.resume_file}") else: logger.warning( - "Cannot write pickle resume file! " - "Job will not resume if interrupted." + "Cannot write pickle resume file! Job will not resume if interrupted." ) self.sampler.pool = self.pool if self.sampler.pool is not None: diff --git a/bilby/core/sampler/fake_sampler.py b/bilby/core/sampler/fake_sampler.py index e36b9a1c8..b791acb7f 100644 --- a/bilby/core/sampler/fake_sampler.py +++ b/bilby/core/sampler/fake_sampler.py @@ -35,7 +35,7 @@ def __init__( injection_parameters=None, meta_data=None, result_class=None, - **kwargs + **kwargs, ): super(FakeSampler, self).__init__( likelihood=likelihood, @@ -48,7 +48,7 @@ def __init__( injection_parameters=None, meta_data=None, result_class=None, - **kwargs + **kwargs, ) self._read_parameter_list_from_file(sample_file) self.result.outdir = outdir diff --git a/bilby/core/sampler/proposal.py b/bilby/core/sampler/proposal.py index d23d19b4c..23115dd2a 100644 --- a/bilby/core/sampler/proposal.py +++ b/bilby/core/sampler/proposal.py @@ -205,7 +205,7 @@ def __init__( random_number_generator=random.rng.uniform, n_points=3, priors=None, - **random_number_generator_args + **random_number_generator_args, ): """ An ensemble walk diff --git a/bilby/core/series.py b/bilby/core/series.py index d8a8b3439..981e18628 100644 --- a/bilby/core/series.py +++ b/bilby/core/series.py @@ -2,18 +2,17 @@ class CoupledTimeAndFrequencySeries: - def __init__(self, duration=None, sampling_frequency=None, start_time=0): - """ A waveform generator - - Parameters - ========== - sampling_frequency: float, optional - The sampling frequency - duration: float, optional - Time duration of data - start_time: float, optional - Starting time of the time array + """A waveform generator + + Parameters + ========== + sampling_frequency: float, optional + The sampling frequency + duration: float, optional + Time duration of data + start_time: float, optional + Starting time of the time array """ self._duration = duration self._sampling_frequency = sampling_frequency @@ -24,12 +23,16 @@ def __init__(self, duration=None, sampling_frequency=None, start_time=0): self._time_array = None def __repr__(self): - return self.__class__.__name__ + '(duration={}, sampling_frequency={}, start_time={})'\ - .format(self.duration, self.sampling_frequency, self.start_time) + return ( + self.__class__.__name__ + + "(duration={}, sampling_frequency={}, start_time={})".format( + self.duration, self.sampling_frequency, self.start_time + ) + ) @property def frequency_array(self): - """ Frequency array for the waveforms. Automatically updates if sampling_frequency or duration are updated. + """Frequency array for the waveforms. Automatically updates if sampling_frequency or duration are updated. Returns ======= @@ -38,12 +41,15 @@ def frequency_array(self): if not self._frequency_array_updated: if self.sampling_frequency is not None and self.duration is not None: self._frequency_array = utils.create_frequency_series( - sampling_frequency=self.sampling_frequency, - duration=self.duration) + sampling_frequency=self.sampling_frequency, duration=self.duration + ) else: - raise ValueError('Can not calculate a frequency series without a ' - 'legitimate sampling_frequency ({}) or duration ({})' - .format(self.sampling_frequency, self.duration)) + raise ValueError( + "Can not calculate a frequency series without a " + "legitimate sampling_frequency ({}) or duration ({})".format( + self.sampling_frequency, self.duration + ) + ) self._frequency_array_updated = True return self._frequency_array @@ -51,13 +57,16 @@ def frequency_array(self): @frequency_array.setter def frequency_array(self, frequency_array): self._frequency_array = frequency_array - self._sampling_frequency, self._duration = \ - utils.get_sampling_frequency_and_duration_from_frequency_array(frequency_array) + self._sampling_frequency, self._duration = ( + utils.get_sampling_frequency_and_duration_from_frequency_array( + frequency_array + ) + ) self._frequency_array_updated = True @property def time_array(self): - """ Time array for the waveforms. Automatically updates if sampling_frequency or duration are updated. + """Time array for the waveforms. Automatically updates if sampling_frequency or duration are updated. Returns ======= @@ -69,11 +78,15 @@ def time_array(self): self._time_array = utils.create_time_series( sampling_frequency=self.sampling_frequency, duration=self.duration, - starting_time=self.start_time) + starting_time=self.start_time, + ) else: - raise ValueError('Can not calculate a time series without a ' - 'legitimate sampling_frequency ({}) or duration ({})' - .format(self.sampling_frequency, self.duration)) + raise ValueError( + "Can not calculate a time series without a " + "legitimate sampling_frequency ({}) or duration ({})".format( + self.sampling_frequency, self.duration + ) + ) self._time_array_updated = True return self._time_array @@ -81,14 +94,15 @@ def time_array(self): @time_array.setter def time_array(self, time_array): self._time_array = time_array - self._sampling_frequency, self._duration = \ + self._sampling_frequency, self._duration = ( utils.get_sampling_frequency_and_duration_from_time_array(time_array) + ) self._start_time = time_array[0] self._time_array_updated = True @property def duration(self): - """ Allows one to set the time duration and automatically updates the frequency and time array. + """Allows one to set the time duration and automatically updates the frequency and time array. Returns ======= @@ -105,7 +119,7 @@ def duration(self, duration): @property def sampling_frequency(self): - """ Allows one to set the sampling frequency and automatically updates the frequency and time array. + """Allows one to set the sampling frequency and automatically updates the frequency and time array. Returns ======= diff --git a/bilby/core/utils/calculus.py b/bilby/core/utils/calculus.py index 9d3515cc1..31e2e7958 100644 --- a/bilby/core/utils/calculus.py +++ b/bilby/core/utils/calculus.py @@ -222,7 +222,6 @@ def interp1d(x, y, kind="linear", axis=-1, bounds_error=None, fill_value=np.nan) class BoundedRectBivariateSpline(RectBivariateSpline): - def __init__(self, x, y, z, bbox=[None] * 4, kx=3, ky=3, s=0, fill_value=None): self.x_min, self.x_max, self.y_min, self.y_max = bbox if self.x_min is None: @@ -285,9 +284,16 @@ class WrappedInterp1d(_interp1d): makes sure that the output is a float if the input is a float or int. """ - def __init__(self, x, y, kind="linear", axis=-1, bounds_error=False, fill_value=np.nan): + def __init__( + self, x, y, kind="linear", axis=-1, bounds_error=False, fill_value=np.nan + ): super().__init__( - x=x, y=y, kind=kind, axis=axis, bounds_error=bounds_error, fill_value=fill_value + x=x, + y=y, + kind=kind, + axis=axis, + bounds_error=bounds_error, + fill_value=fill_value, ) self.kind = kind @@ -322,4 +328,4 @@ def round_up_to_power_of_two(x): float: next power of two """ - return 2**math.ceil(np.log2(x)) + return 2 ** math.ceil(np.log2(x)) diff --git a/bilby/core/utils/cmd.py b/bilby/core/utils/cmd.py index 817b1d6e4..30484d35d 100644 --- a/bilby/core/utils/cmd.py +++ b/bilby/core/utils/cmd.py @@ -7,7 +7,7 @@ def set_up_command_line_arguments(): - """ Sets up command line arguments that can be used to modify how scripts are run. + """Sets up command line arguments that can be used to modify how scripts are run. Returns ======= @@ -54,29 +54,64 @@ def set_up_command_line_arguments(): try: parser = argparse.ArgumentParser( description="Command line interface for bilby scripts", - add_help=False, allow_abbrev=False) + add_help=False, + allow_abbrev=False, + ) except TypeError: parser = argparse.ArgumentParser( - description="Command line interface for bilby scripts", - add_help=False) - parser.add_argument("-v", "--verbose", action="store_true", - help=("Increase output verbosity [logging.DEBUG]." + - " Overridden by script level settings")) - parser.add_argument("-q", "--quiet", action="store_true", - help=("Decrease output verbosity [logging.WARNING]." + - " Overridden by script level settings")) - parser.add_argument("-c", "--clean", action="store_true", - help="Force clean data, never use cached data") - parser.add_argument("-u", "--use-cached", action="store_true", - help="Force cached data and do not check its validity") - parser.add_argument("--sampler-help", nargs='?', default=False, - const='None', help="Print help for given sampler") - parser.add_argument("--bilby-test-mode", action="store_true", - help=("Used for testing only: don't run full PE, but" - " just check nothing breaks")) - parser.add_argument("--bilby-zero-likelihood-mode", action="store_true", - help=("Used for testing only: don't run full PE, but" - " just check nothing breaks")) + description="Command line interface for bilby scripts", add_help=False + ) + parser.add_argument( + "-v", + "--verbose", + action="store_true", + help=( + "Increase output verbosity [logging.DEBUG]." + + " Overridden by script level settings" + ), + ) + parser.add_argument( + "-q", + "--quiet", + action="store_true", + help=( + "Decrease output verbosity [logging.WARNING]." + + " Overridden by script level settings" + ), + ) + parser.add_argument( + "-c", + "--clean", + action="store_true", + help="Force clean data, never use cached data", + ) + parser.add_argument( + "-u", + "--use-cached", + action="store_true", + help="Force cached data and do not check its validity", + ) + parser.add_argument( + "--sampler-help", + nargs="?", + default=False, + const="None", + help="Print help for given sampler", + ) + parser.add_argument( + "--bilby-test-mode", + action="store_true", + help=( + "Used for testing only: don't run full PE, but just check nothing breaks" + ), + ) + parser.add_argument( + "--bilby-zero-likelihood-mode", + action="store_true", + help=( + "Used for testing only: don't run full PE, but just check nothing breaks" + ), + ) args, unknown_args = parser.parse_known_args() if args.quiet: args.log_level = logging.WARNING @@ -101,19 +136,19 @@ def run_commandline(cl, log_level=20, raise_error=True, return_output=True): """ - logger.log(log_level, 'Now executing: ' + cl) + logger.log(log_level, "Now executing: " + cl) if return_output: try: out = subprocess.check_output( - cl, stderr=subprocess.STDOUT, shell=True, - universal_newlines=True) + cl, stderr=subprocess.STDOUT, shell=True, universal_newlines=True + ) except subprocess.CalledProcessError as e: - logger.log(log_level, 'Execution failed: {}'.format(e.output)) + logger.log(log_level, "Execution failed: {}".format(e.output)) if raise_error: raise else: out = 0 - os.system('\n') + os.system("\n") return out else: process = subprocess.Popen(cl, shell=True) diff --git a/bilby/core/utils/colors.py b/bilby/core/utils/colors.py index 7ce34e66c..53c604ee6 100644 --- a/bilby/core/utils/colors.py +++ b/bilby/core/utils/colors.py @@ -1,5 +1,5 @@ class tcolors: - KEY = '\033[93m' - VALUE = '\033[91m' - HIGHLIGHT = '\033[95m' - END = '\033[0m' + KEY = "\033[93m" + VALUE = "\033[91m" + HIGHLIGHT = "\033[95m" + END = "\033[0m" diff --git a/bilby/core/utils/constants.py b/bilby/core/utils/constants.py index 8dbec27da..45f3cd03d 100644 --- a/bilby/core/utils/constants.py +++ b/bilby/core/utils/constants.py @@ -1,7 +1,7 @@ # Constants: values taken from LAL 505df9dd2e69b4812f1e8eee3a6d468ba7f80674 speed_of_light = 299792458.0 # m/s -parsec = 3.085677581491367e+16 # m +parsec = 3.085677581491367e16 # m solar_mass = 1.988409870698050731911960804878414216e30 # Kg radius_of_earth = 6378136.6 # m gravitational_constant = 6.6743e-11 # m^3 kg^-1 s^-2 diff --git a/bilby/core/utils/conversion.py b/bilby/core/utils/conversion.py index c89e2a9ec..5f020dc84 100644 --- a/bilby/core/utils/conversion.py +++ b/bilby/core/utils/conversion.py @@ -2,7 +2,7 @@ def ra_dec_to_theta_phi(ra, dec, gmst): - """ Convert from RA and DEC to polar coordinates on celestial sphere + """Convert from RA and DEC to polar coordinates on celestial sphere Parameters ========== diff --git a/bilby/core/utils/counter.py b/bilby/core/utils/counter.py index 29428057e..3218e4a3e 100644 --- a/bilby/core/utils/counter.py +++ b/bilby/core/utils/counter.py @@ -11,8 +11,9 @@ class Counter(object): initalval : int, 0 number to start counting from """ + def __init__(self, initval=0): - self.val = multiprocessing.RawValue('i', initval) + self.val = multiprocessing.RawValue("i", initval) self.lock = multiprocessing.Lock() def increment(self): diff --git a/bilby/core/utils/docs.py b/bilby/core/utils/docs.py index a7e326e8b..37c6e983b 100644 --- a/bilby/core/utils/docs.py +++ b/bilby/core/utils/docs.py @@ -12,10 +12,12 @@ def docstring(docstr, sep="\n"): sep: str Separation character for appending the existing docstring. """ + def _decorator(func): if func.__doc__ is None: func.__doc__ = docstr else: func.__doc__ = sep.join([func.__doc__, docstr]) return func + return _decorator diff --git a/bilby/core/utils/entry_points.py b/bilby/core/utils/entry_points.py index 305fc5704..80765c30d 100644 --- a/bilby/core/utils/entry_points.py +++ b/bilby/core/utils/entry_points.py @@ -1,4 +1,5 @@ import sys + if sys.version_info < (3, 10): from importlib_metadata import entry_points else: @@ -13,6 +14,4 @@ def get_entry_points(group): group: str Entry points you wish to query """ - return { - custom.name: custom for custom in entry_points(group=group) - } + return {custom.name: custom for custom in entry_points(group=group)} diff --git a/bilby/core/utils/env.py b/bilby/core/utils/env.py index fc3ba4e1d..7df35bc43 100644 --- a/bilby/core/utils/env.py +++ b/bilby/core/utils/env.py @@ -1,16 +1,12 @@ - - def string_to_boolean(value): """Convert a string to a boolean. Supports True/False (case-insensitive), and 1/0. """ value = value.strip().lower() - if value in ['true', '1']: + if value in ["true", "1"]: return True - elif value in ['false', '0']: + elif value in ["false", "0"]: return False else: - raise ValueError( - f"Invalid value for boolean: {value}" - ) + raise ValueError(f"Invalid value for boolean: {value}") diff --git a/bilby/core/utils/introspection.py b/bilby/core/utils/introspection.py index c4c424d17..a806fcb07 100644 --- a/bilby/core/utils/introspection.py +++ b/bilby/core/utils/introspection.py @@ -2,7 +2,7 @@ def infer_parameters_from_function(func): - """ Infers the arguments of a function + """Infers the arguments of a function (except the first arg which is assumed to be the dep. variable). Throws out `*args` and `**kwargs` type arguments @@ -36,7 +36,7 @@ def infer_parameters_from_function(func): def infer_args_from_method(method): - """ Infers all arguments of a method except for `self` + """Infers all arguments of a method except for `self` Throws out `*args` and `**kwargs` type arguments. @@ -51,7 +51,7 @@ def infer_args_from_method(method): def infer_args_from_function_except_n_args(func, n=1): - """ Inspects a function to find its arguments, and ignoring the + """Inspects a function to find its arguments, and ignoring the first n of these, returns a list of arguments from the function's signature. @@ -94,12 +94,10 @@ def infer_args_from_function_except_n_args(func, n=1): """ parameters = [ - parameter.name for parameter in - inspect.signature(func).parameters.values() - if parameter.kind not in ( - inspect.Parameter.VAR_POSITIONAL, - inspect.Parameter.VAR_KEYWORD - ) + parameter.name + for parameter in inspect.signature(func).parameters.values() + if parameter.kind + not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD) ] del parameters[:n] @@ -111,8 +109,9 @@ def _infer_args_from_function_except_for_first_arg(func): def get_dict_with_properties(obj): - property_names = [p for p in dir(obj.__class__) - if isinstance(getattr(obj.__class__, p), property)] + property_names = [ + p for p in dir(obj.__class__) if isinstance(getattr(obj.__class__, p), property) + ] dict_with_properties = obj.__dict__.copy() for key in property_names: dict_with_properties[key] = getattr(obj, key) @@ -152,7 +151,11 @@ def __init__(self, container_instance_name, property_name): self.container_instance_name = container_instance_name def __get__(self, instance, owner): - return getattr(getattr(instance, self.container_instance_name), self.property_name) + return getattr( + getattr(instance, self.container_instance_name), self.property_name + ) def __set__(self, instance, value): - setattr(getattr(instance, self.container_instance_name), self.property_name, value) + setattr( + getattr(instance, self.container_instance_name), self.property_name, value + ) diff --git a/bilby/core/utils/io.py b/bilby/core/utils/io.py index efe86c62e..99220366a 100644 --- a/bilby/core/utils/io.py +++ b/bilby/core/utils/io.py @@ -16,7 +16,7 @@ def check_directory_exists_and_if_not_mkdir(directory): - """ Checks if the given directory exists and creates it if it does not exist + """Checks if the given directory exists and creates it if it does not exist Parameters ========== @@ -56,7 +56,9 @@ def default(self, obj): return encode_astropy_cosmology(obj) if isinstance(obj, units.Quantity): return encode_astropy_quantity(obj) - if isinstance(obj, (units.PrefixUnit, units.UnitBase, units.FunctionUnitBase)): + if isinstance( + obj, (units.PrefixUnit, units.UnitBase, units.FunctionUnitBase) + ): return encode_astropy_unit(obj) except ImportError: logger.debug("Cannot import astropy, cannot write cosmological priors") @@ -96,10 +98,7 @@ def default(self, obj): "__name__": obj.__name__, } if isinstance(obj, (timedelta)): - return { - "__timedelta__": True, - "__total_seconds__": obj.total_seconds() - } + return {"__timedelta__": True, "__total_seconds__": obj.total_seconds()} return obj.isoformat() return json.JSONEncoder.default(self, obj) @@ -319,10 +318,12 @@ def decode_numpy_random_generator(dct): f"Original error: {e}" ) from e # Convert the state and inc integers back to integers - dct["bit_generator_state"]["state"]["state"] = \ - int(dct["bit_generator_state"]["state"]["state"]) - dct["bit_generator_state"]["state"]["inc"] = \ - int(dct["bit_generator_state"]["state"]["inc"]) + dct["bit_generator_state"]["state"]["state"] = int( + dct["bit_generator_state"]["state"]["state"] + ) + dct["bit_generator_state"]["state"]["inc"] = int( + dct["bit_generator_state"]["state"]["inc"] + ) generator = np.random.Generator(bit_generator()) generator.bit_generator.state = dct["bit_generator_state"] @@ -362,11 +363,11 @@ def decode_bilby_json(dct): cls = getattr(import_module(dct["__module__"]), dct["__name__"]) obj = cls(**dct["kwargs"]) except (AttributeError, ValueError) as e: - if type(e).__name__ == 'AttributeError': + if type(e).__name__ == "AttributeError": warning_message = "Unknown prior class for parameter {}, defaulting to base Prior object".format( dct["kwargs"]["name"] ) - elif type(e).__name__ == 'ValueError': + elif type(e).__name__ == "ValueError": warning_message = ( f"Unable to load prior {cls} with arguments {dct['kwargs']}, " "defaulting to base Prior object" @@ -375,7 +376,14 @@ def decode_bilby_json(dct): from ..prior import Prior for key in list(dct["kwargs"].keys()): - if key not in ["name", "latex_label", "unit", "minimum", "maximum", "boundary"]: + if key not in [ + "name", + "latex_label", + "unit", + "minimum", + "maximum", + "boundary", + ]: dct["kwargs"].pop(key) cls = Prior obj = cls(**dct["kwargs"]) @@ -516,9 +524,9 @@ def encode_for_hdf5(key, item): item = complex(item) if isinstance(item, np.ndarray): # Numpy's wide unicode strings are not supported by hdf5 - if item.dtype.kind == 'U': - logger.debug(f'converting dtype {item.dtype} for hdf5') - item = np.array(item, dtype='S') + if item.dtype.kind == "U": + logger.debug(f"converting dtype {item.dtype} for hdf5") + item = np.array(item, dtype="S") elif aac.is_array_api_obj(item): # temporarily dump all arrays as numpy arrays, we should figure ou # how to properly deserialize them @@ -549,7 +557,7 @@ def encode_for_hdf5(key, item): else: output.append(str(value).encode("utf-8")) else: - raise ValueError(f'Cannot save {key}: {type(item)} type') + raise ValueError(f"Cannot save {key}: {type(item)} type") elif isinstance(item, PriorDict): output = json.dumps(item._get_json_dict()) elif isinstance(item, pd.DataFrame): @@ -558,7 +566,9 @@ def encode_for_hdf5(key, item): output = encode_astropy_cosmology(item) elif units is not None and isinstance(item, units.Quantity): output = encode_astropy_quantity(item) - elif units is not None and isinstance(item, (units.PrefixUnit, units.UnitBase, units.FunctionUnitBase)): + elif units is not None and isinstance( + item, (units.PrefixUnit, units.UnitBase, units.FunctionUnitBase) + ): output = encode_astropy_unit(item) elif lal is not None and isinstance(item, lal.Dict): output = encode_lal_dict(item) @@ -575,7 +585,7 @@ def encode_for_hdf5(key, item): elif isinstance(item, datetime.timedelta): output = item.total_seconds() else: - raise ValueError(f'Cannot save {key}: {type(item)} type') + raise ValueError(f"Cannot save {key}: {type(item)} type") return output @@ -588,7 +598,9 @@ def decode_hdf5_dict(output): """ if ("__function__" in output) or ("__class__" in output): default = ".".join([output["__module__"], output["__name__"]]) - output = getattr(import_module(output["__module__"]), output["__name__"], default) + output = getattr( + import_module(output["__module__"]), output["__name__"], default + ) elif "__cosmology__" in output: output = decode_astropy_cosmology(output) elif "__astropy_quantity__" in output: @@ -666,7 +678,7 @@ def recursively_save_dict_contents_to_group(h5file, path, dic): def safe_file_dump(data, filename, module): - """ Safely dump data to a .pickle file + """Safely dump data to a .pickle file Parameters ========== @@ -686,7 +698,7 @@ def safe_file_dump(data, filename, module): def move_old_file(filename, overwrite=False): - """ Moves or removes an old file. + """Moves or removes an old file. Parameters ========== diff --git a/bilby/core/utils/log.py b/bilby/core/utils/log.py index a6bc1683c..33587e9e5 100644 --- a/bilby/core/utils/log.py +++ b/bilby/core/utils/log.py @@ -5,11 +5,11 @@ import sys from importlib import metadata -logger = logging.getLogger('bilby') +logger = logging.getLogger("bilby") -def setup_logger(outdir='.', label=None, log_level='INFO', print_version=False): - """ Setup logging output: call at the start of the script to use +def setup_logger(outdir=".", label=None, log_level="INFO", print_version=False): + """Setup logging output: call at the start of the script to use Parameters ========== @@ -27,28 +27,34 @@ def setup_logger(outdir='.', label=None, log_level='INFO', print_version=False): try: level = getattr(logging, log_level.upper()) except AttributeError: - raise ValueError('log_level {} not understood'.format(log_level)) + raise ValueError("log_level {} not understood".format(log_level)) else: level = int(log_level) - logger = logging.getLogger('bilby') + logger = logging.getLogger("bilby") logger.propagate = False logger.setLevel(level) if not any([isinstance(h, logging.StreamHandler) for h in logger.handlers]): stream_handler = logging.StreamHandler() - stream_handler.setFormatter(logging.Formatter( - '%(asctime)s %(name)s %(levelname)-8s: %(message)s', datefmt='%H:%M')) + stream_handler.setFormatter( + logging.Formatter( + "%(asctime)s %(name)s %(levelname)-8s: %(message)s", datefmt="%H:%M" + ) + ) stream_handler.setLevel(level) logger.addHandler(stream_handler) if not any([isinstance(h, logging.FileHandler) for h in logger.handlers]): if label: Path(outdir).mkdir(parents=True, exist_ok=True) - log_file = '{}/{}.log'.format(outdir, label) + log_file = "{}/{}.log".format(outdir, label) file_handler = logging.FileHandler(log_file) - file_handler.setFormatter(logging.Formatter( - '%(asctime)s %(levelname)-8s: %(message)s', datefmt='%H:%M')) + file_handler.setFormatter( + logging.Formatter( + "%(asctime)s %(levelname)-8s: %(message)s", datefmt="%H:%M" + ) + ) file_handler.setLevel(level) logger.addHandler(file_handler) @@ -58,11 +64,12 @@ def setup_logger(outdir='.', label=None, log_level='INFO', print_version=False): if print_version: version = get_version_information() - logger.info('Running bilby version: {}'.format(version)) + logger.info("Running bilby version: {}".format(version)) def get_version_information(): from bilby import __version__ + return __version__ @@ -106,12 +113,9 @@ def env_package_list(as_dataframe=False): conda_detected = (Path(prefix) / "conda-meta").is_dir() if conda_detected: try: - pkgs = json.loads(subprocess.check_output([ - "conda", - "list", - "--prefix", prefix, - "--json" - ])) + pkgs = json.loads( + subprocess.check_output(["conda", "list", "--prefix", prefix, "--json"]) + ) except (FileNotFoundError, subprocess.CalledProcessError): # When a conda env is in use but conda is unavailable conda_detected = False @@ -126,15 +130,23 @@ def env_package_list(as_dataframe=False): modules = loaded_modules_dict() pkgs = [{"name": x, "version": y} for x, y in modules.items()] else: - pkgs = json.loads(subprocess.check_output([ - sys.executable, - "-m", "pip", - "list", "installed", - "--format", "json", - ])) + pkgs = json.loads( + subprocess.check_output( + [ + sys.executable, + "-m", + "pip", + "list", + "installed", + "--format", + "json", + ] + ) + ) # convert to recarray for storage if as_dataframe: from pandas import DataFrame + return DataFrame(pkgs) return pkgs diff --git a/bilby/core/utils/meta_data.py b/bilby/core/utils/meta_data.py index 858ecf7a3..3e84305cd 100644 --- a/bilby/core/utils/meta_data.py +++ b/bilby/core/utils/meta_data.py @@ -44,6 +44,4 @@ def __new__(cls, *args, **kwargs): return cls._instance -global_meta_data = GlobalMetaData({ - "rng": random.rng -}) +global_meta_data = GlobalMetaData({"rng": random.rng}) diff --git a/bilby/core/utils/plotting.py b/bilby/core/utils/plotting.py index 1f2689dd5..8337990ba 100644 --- a/bilby/core/utils/plotting.py +++ b/bilby/core/utils/plotting.py @@ -35,6 +35,7 @@ def latex_plot_format(func): latex_plot_format-wrapped plotting function and will be set directly. """ + @functools.wraps(func) def wrapper_decorator(*args, **kwargs): import matplotlib.pyplot as plt @@ -52,7 +53,7 @@ def wrapper_decorator(*args, **kwargs): if bilby_mathdefault == 1: logger.debug("Setting mathdefault in the rcParams") - rcParams['text.latex.preamble'] = r'\providecommand{\mathdefault}[1][]{}' + rcParams["text.latex.preamble"] = r"\providecommand{\mathdefault}[1][]{}" logger.debug("Using BILBY_STYLE={}".format(bilby_style)) if bilby_style.lower() == "none": @@ -79,8 +80,8 @@ def wrapper_decorator(*args, **kwargs): return func(*args, **kwargs) else: logger.debug( - "Environment variable BILBY_STYLE={} not used" - .format(bilby_style) + "Environment variable BILBY_STYLE={} not used".format(bilby_style) ) return func(*args, **kwargs) + return wrapper_decorator diff --git a/bilby/core/utils/random.py b/bilby/core/utils/random.py index b8aac288d..ecf394b63 100644 --- a/bilby/core/utils/random.py +++ b/bilby/core/utils/random.py @@ -24,6 +24,7 @@ Do not import :code:`rng` directly from :code:`bilby.core.utils.random` since it will not be seeded correctly. """ + import sys import warnings @@ -49,6 +50,7 @@ class Generator: It should not be used directly, instead use :code:`random.rng` to generate random numbers. See the documentation for more details. """ + rng = default_rng() """Random number generator. @@ -81,7 +83,7 @@ def seed(seed): # Warn if the original rng object (i.e., pre-seed) still exists elsewhere for module in sys.modules.values(): - if not module or not hasattr(module, '__dict__'): + if not module or not hasattr(module, "__dict__"): continue rng_obj = module.__dict__.get("rng") if rng_obj is _original_rng: @@ -89,7 +91,7 @@ def seed(seed): "Detected that `rng` was likely imported directly before calling `seed()`. " "This means the imported reference will not reflect the newly seeded generator. " "Use `from bilby.core.utils import random` and access `random.rng` instead.", - RuntimeWarning + RuntimeWarning, ) break @@ -147,6 +149,7 @@ def _resolve_numpy_generator(random_state): return _resolve_numpy_generator(random_state) import orng + if isinstance(random_state, (np.random.Generator, orng.RandomGenerator)): return random_state elif aac.is_jax_array(random_state): @@ -181,9 +184,13 @@ def random_array_module(random_state): return np elif aac.is_jax_array(random_state) or getattr(random_state, "backend") == "jax": import jax.numpy as jnp + return jnp - elif aac.is_torch_array(random_state) or getattr(random_state, "backend") == "torch": + elif ( + aac.is_torch_array(random_state) or getattr(random_state, "backend") == "torch" + ): import torch + return torch else: return np diff --git a/bilby/core/utils/samples.py b/bilby/core/utils/samples.py index 93fdac0ac..db71d2924 100644 --- a/bilby/core/utils/samples.py +++ b/bilby/core/utils/samples.py @@ -4,7 +4,7 @@ class SamplesSummary(object): - """ Object to store a set of samples and calculate summary statistics + """Object to store a set of samples and calculate summary statistics Parameters ========== @@ -17,7 +17,8 @@ class SamplesSummary(object): The default confidence interval level, defaults t0 0.9 """ - def __init__(self, samples, average='median', confidence_level=.9): + + def __init__(self, samples, average="median", confidence_level=0.9): self.samples = samples self.average = average self.confidence_level = confidence_level @@ -43,14 +44,14 @@ def confidence_level(self, confidence_level): @property def average(self): - if self._average == 'mean': + if self._average == "mean": return self.mean - elif self._average == 'median': + elif self._average == "median": return self.median @average.setter def average(self, average): - allowed_averages = ['mean', 'median'] + allowed_averages = ["mean", "median"] if average in allowed_averages: self._average = average else: @@ -66,37 +67,37 @@ def mean(self): @property def _lower_level(self): - """ The credible interval lower quantile value """ - return (1 - self.confidence_level) / 2. + """The credible interval lower quantile value""" + return (1 - self.confidence_level) / 2.0 @property def _upper_level(self): - """ The credible interval upper quantile value """ - return (1 + self.confidence_level) / 2. + """The credible interval upper quantile value""" + return (1 + self.confidence_level) / 2.0 @property def lower_absolute_credible_interval(self): - """ Absolute lower value of the credible interval """ + """Absolute lower value of the credible interval""" return np.quantile(self.samples, self._lower_level, axis=0) @property def upper_absolute_credible_interval(self): - """ Absolute upper value of the credible interval """ + """Absolute upper value of the credible interval""" return np.quantile(self.samples, self._upper_level, axis=0) @property def lower_relative_credible_interval(self): - """ Relative (to average) lower value of the credible interval """ + """Relative (to average) lower value of the credible interval""" return self.lower_absolute_credible_interval - self.average @property def upper_relative_credible_interval(self): - """ Relative (to average) upper value of the credible interval """ + """Relative (to average) upper value of the credible interval""" return self.upper_absolute_credible_interval - self.average def kish_log_effective_sample_size(ln_weights): - """ Calculate the Kish effective sample size from the natural-log weights + """Calculate the Kish effective sample size from the natural-log weights See https://en.wikipedia.org/wiki/Effective_sample_size for details diff --git a/bilby/core/utils/series.py b/bilby/core/utils/series.py index 44f80e053..d3cbcc297 100644 --- a/bilby/core/utils/series.py +++ b/bilby/core/utils/series.py @@ -30,7 +30,7 @@ def get_sampling_frequency(time_array): if np.ptp(np.diff(time_array)) > tol: raise ValueError("Your time series was not evenly sampled") else: - return np.round(1. / (time_array[1] - time_array[0]), decimals=_TOL) + return np.round(1.0 / (time_array[1] - time_array[0]), decimals=_TOL) def get_sampling_frequency_and_duration_from_time_array(time_array): @@ -84,11 +84,13 @@ def get_sampling_frequency_and_duration_from_frequency_array(frequency_array): delta_freq = frequency_array[1] - frequency_array[0] duration = np.round(1 / delta_freq, decimals=_TOL) - sampling_frequency = np.round(2 * (number_of_frequencies - 1) / duration, decimals=14) + sampling_frequency = np.round( + 2 * (number_of_frequencies - 1) / duration, decimals=14 + ) return sampling_frequency, duration -def create_time_series(sampling_frequency, duration, starting_time=0.): +def create_time_series(sampling_frequency, duration, starting_time=0.0): """ Parameters @@ -113,7 +115,7 @@ def create_time_series(sampling_frequency, duration, starting_time=0.): def create_frequency_series(sampling_frequency, duration): - """ Create a frequency series with the correct length and spacing. + """Create a frequency series with the correct length and spacing. Parameters ========== @@ -135,7 +137,7 @@ def create_frequency_series(sampling_frequency, duration): def _check_legal_sampling_frequency_and_duration(sampling_frequency, duration): - """ By convention, sampling_frequency and duration have to multiply to an integer + """By convention, sampling_frequency and duration have to multiply to an integer This will check if the product of both parameters multiplies reasonably close to an integer. @@ -147,13 +149,12 @@ def _check_legal_sampling_frequency_and_duration(sampling_frequency, duration): """ num = sampling_frequency * duration - if abs(num % 1) > 10**(-_TOL): + if abs(num % 1) > 10 ** (-_TOL): raise IllegalDurationAndSamplingFrequencyException( - '\nYour sampling frequency and duration must multiply to a number' - 'up to (tol = {}) decimals close to an integer number. ' - '\nBut sampling_frequency={} and duration={} multiply to {}'.format( - _TOL, sampling_frequency, duration, - sampling_frequency * duration + "\nYour sampling frequency and duration must multiply to a number" + "up to (tol = {}) decimals close to an integer number. " + "\nBut sampling_frequency={} and duration={} multiply to {}".format( + _TOL, sampling_frequency, duration, sampling_frequency * duration ) ) @@ -199,7 +200,7 @@ def safe_white_noise(number_of_samples: int, duration: float, *, random_state=No def create_white_noise(sampling_frequency, duration, random_state=None): - """ Create white_noise which is then coloured by a given PSD + """Create white_noise which is then coloured by a given PSD Parameters ========== @@ -218,7 +219,7 @@ def create_white_noise(sampling_frequency, duration, random_state=None): def nfft(time_domain_strain, sampling_frequency): - """ Perform an FFT while keeping track of the frequency bins. Assumes input + """Perform an FFT while keeping track of the frequency bins. Assumes input time series is real (positive frequencies only). Parameters @@ -240,13 +241,14 @@ def nfft(time_domain_strain, sampling_frequency): frequency_domain_strain /= sampling_frequency frequency_array = xp.linspace( - 0, sampling_frequency / 2, len(frequency_domain_strain)) + 0, sampling_frequency / 2, len(frequency_domain_strain) + ) return frequency_domain_strain, frequency_array def infft(frequency_domain_strain, sampling_frequency): - """ Inverse FFT for use in conjunction with nfft. + """Inverse FFT for use in conjunction with nfft. Parameters ========== diff --git a/bilby/gw/__init__.py b/bilby/gw/__init__.py index cd09bc6f6..5f901e700 100644 --- a/bilby/gw/__init__.py +++ b/bilby/gw/__init__.py @@ -1,7 +1,16 @@ -from . import (conversion, cosmology, detector, eos, likelihood, prior, - result, source, utils, waveform_generator) +from . import ( + conversion, + cosmology, + detector, + eos, + likelihood, + prior, + result, + source, + utils, + waveform_generator, +) from .waveform_generator import WaveformGenerator, LALCBCWaveformGenerator from .likelihood import GravitationalWaveTransient from .detector import calibration from . import compat - diff --git a/bilby/gw/compat/__init__.py b/bilby/gw/compat/__init__.py index 245124a71..4d18be2bf 100644 --- a/bilby/gw/compat/__init__.py +++ b/bilby/gw/compat/__init__.py @@ -14,4 +14,4 @@ try: from .torch import n_leap_seconds except ModuleNotFoundError: - pass \ No newline at end of file + pass diff --git a/bilby/gw/compat/cython.py b/bilby/gw/compat/cython.py index 9d0a69af0..ca0ccdd9e 100644 --- a/bilby/gw/compat/cython.py +++ b/bilby/gw/compat/cython.py @@ -56,7 +56,9 @@ def time_delay_geocentric(detector1: ArrayLike, detector2: ArrayLike, ra, dec, t @dispatch(precedence=1) -def time_delay_from_geocenter(detector1: ArrayLike, ra: Real, dec: Real, time: Real | ArrayLike): +def time_delay_from_geocenter( + detector1: ArrayLike, ra: Real, dec: Real, time: Real | ArrayLike +): return _geometry.time_delay_from_geocenter(detector1, ra, dec, time) diff --git a/bilby/gw/compat/pytrees.py b/bilby/gw/compat/pytrees.py index 22e73c9a8..44c30277e 100644 --- a/bilby/gw/compat/pytrees.py +++ b/bilby/gw/compat/pytrees.py @@ -79,9 +79,7 @@ def interferometer_flatten(ifo: Interferometer): ifo.calibration_model, ifo.reference_time, ) - aux_data = ( - ifo.name, - ) + aux_data = (ifo.name,) return children, aux_data @@ -144,6 +142,7 @@ def psd_flatten(psd: PowerSpectralDensity): def psd_unflatten(aux_data, children) -> PowerSpectralDensity: import jax.numpy as jnp + psd = PowerSpectralDensity.__new__(PowerSpectralDensity) psd._cache = children[0] @@ -174,9 +173,7 @@ def geometry_flatten(geometry: InterferometerGeometry): geometry._y, geometry._detector_tensor, ) - aux_data = ( - geometry.length, - ) + aux_data = (geometry.length,) return children, aux_data @@ -224,9 +221,7 @@ def cubic_spline_flatten(cubic_spline: CubicSpline): cubic_spline.delta_log_spline_points, cubic_spline.nodes_to_spline_coefficients, ) - aux_data += ( - cubic_spline.n_points, - ) + aux_data += (cubic_spline.n_points,) return children, aux_data @@ -252,9 +247,7 @@ def precomputed_flatten(precomputed: Precomputed): precomputed.frequency_array, precomputed.parameters, ) - aux_data += ( - precomputed.label, - ) + aux_data += (precomputed.label,) return children, aux_data @@ -279,9 +272,7 @@ def notch_list_unflatten(aux_data, children) -> NotchList: def wfg_flatten(wfg: WaveformGenerator): - children = ( - wfg._times_and_frequencies, - ) + children = (wfg._times_and_frequencies,) aux_data = ( wfg.frequency_domain_source_model, wfg.time_domain_source_model, diff --git a/bilby/gw/conversion.py b/bilby/gw/conversion.py index b55d49379..488216b57 100644 --- a/bilby/gw/conversion.py +++ b/bilby/gw/conversion.py @@ -11,23 +11,32 @@ from pandas import DataFrame, Series from scipy.stats import norm -from .utils import (lalsim_SimNeutronStarEOS4ParamSDGammaCheck, - lalsim_SimNeutronStarEOS4ParameterSpectralDecomposition, - lalsim_SimNeutronStarEOS4ParamSDViableFamilyCheck, - lalsim_SimNeutronStarEOS3PieceDynamicPolytrope, - lalsim_SimNeutronStarEOS3PieceCausalAnalytic, - lalsim_SimNeutronStarEOS3PDViableFamilyCheck, - lalsim_CreateSimNeutronStarFamily, - lalsim_SimNeutronStarEOSMaxPseudoEnthalpy, - lalsim_SimNeutronStarEOSSpeedOfSoundGeometerized, - lalsim_SimNeutronStarFamMinimumMass, - lalsim_SimNeutronStarMaximumMass, - lalsim_SimNeutronStarRadius, - lalsim_SimNeutronStarLoveNumberK2) +from .utils import ( + lalsim_SimNeutronStarEOS4ParamSDGammaCheck, + lalsim_SimNeutronStarEOS4ParameterSpectralDecomposition, + lalsim_SimNeutronStarEOS4ParamSDViableFamilyCheck, + lalsim_SimNeutronStarEOS3PieceDynamicPolytrope, + lalsim_SimNeutronStarEOS3PieceCausalAnalytic, + lalsim_SimNeutronStarEOS3PDViableFamilyCheck, + lalsim_CreateSimNeutronStarFamily, + lalsim_SimNeutronStarEOSMaxPseudoEnthalpy, + lalsim_SimNeutronStarEOSSpeedOfSoundGeometerized, + lalsim_SimNeutronStarFamMinimumMass, + lalsim_SimNeutronStarMaximumMass, + lalsim_SimNeutronStarRadius, + lalsim_SimNeutronStarLoveNumberK2, +) from ..compat.utils import array_module from ..core.likelihood import MarginalizedLikelihoodReconstructionError -from ..core.utils import logger, solar_mass, gravitational_constant, speed_of_light, command_line_args, safe_file_dump +from ..core.utils import ( + logger, + solar_mass, + gravitational_constant, + speed_of_light, + command_line_args, + safe_file_dump, +) from ..core.utils.parallel import bilby_pool, sampling_convenience_dump from ..core.prior import DeltaFunction from .utils import lalsim_SimInspiralTransformPrecessingNewInitialConditions @@ -47,6 +56,7 @@ def redshift_to_comoving_distance(redshift, cosmology=None): def luminosity_distance_to_redshift(distance, cosmology=None): from astropy import units + cosmology = get_cosmology(cosmology) if isinstance(distance, Series): distance = distance.values @@ -55,6 +65,7 @@ def luminosity_distance_to_redshift(distance, cosmology=None): def comoving_distance_to_redshift(distance, cosmology=None): from astropy import units + cosmology = get_cosmology(cosmology) if isinstance(distance, Series): distance = distance.values @@ -104,8 +115,17 @@ def luminosity_distance_to_comoving_distance(distance, cosmology=None): def bilby_to_lalsimulation_spins( - theta_jn, phi_jl, tilt_1, tilt_2, phi_12, a_1, a_2, mass_1, mass_2, - reference_frequency, phase + theta_jn, + phi_jl, + tilt_1, + tilt_2, + phi_12, + a_1, + a_2, + mass_1, + mass_2, + reference_frequency, + phase, ): """ Convert from Bilby spin parameters to lalsimulation ones. @@ -153,9 +173,19 @@ def bilby_to_lalsimulation_spins( iota = theta_jn else: from numbers import Number + args = ( - theta_jn, phi_jl, tilt_1, tilt_2, phi_12, a_1, a_2, mass_1, - mass_2, reference_frequency, phase + theta_jn, + phi_jl, + tilt_1, + tilt_2, + phi_12, + a_1, + a_2, + mass_1, + mass_2, + reference_frequency, + phase, ) float_inputs = all([isinstance(arg, Number) for arg in args]) if float_inputs: @@ -208,29 +238,35 @@ def convert_to_lal_binary_black_hole_parameters(parameters): converted_parameters = parameters.copy() original_keys = list(converted_parameters.keys()) xp = array_module(parameters) - if 'luminosity_distance' not in original_keys: - if 'redshift' in converted_parameters.keys(): - converted_parameters['luminosity_distance'] = \ - redshift_to_luminosity_distance(parameters['redshift']) - elif 'comoving_distance' in converted_parameters.keys(): - converted_parameters['luminosity_distance'] = \ + if "luminosity_distance" not in original_keys: + if "redshift" in converted_parameters.keys(): + converted_parameters["luminosity_distance"] = ( + redshift_to_luminosity_distance(parameters["redshift"]) + ) + elif "comoving_distance" in converted_parameters.keys(): + converted_parameters["luminosity_distance"] = ( comoving_distance_to_luminosity_distance( - parameters['comoving_distance']) + parameters["comoving_distance"] + ) + ) for key in original_keys: - if key[-7:] == '_source': - if 'redshift' not in converted_parameters.keys(): - converted_parameters['redshift'] =\ - luminosity_distance_to_redshift( - parameters['luminosity_distance']) + if key[-7:] == "_source": + if "redshift" not in converted_parameters.keys(): + converted_parameters["redshift"] = luminosity_distance_to_redshift( + parameters["luminosity_distance"] + ) converted_parameters[key[:-7]] = converted_parameters[key] * ( - 1 + converted_parameters['redshift']) + 1 + converted_parameters["redshift"] + ) # we do not require the component masses be added if no mass parameters are present - converted_parameters = generate_component_masses(converted_parameters, require_add=False) + converted_parameters = generate_component_masses( + converted_parameters, require_add=False + ) - for idx in ['1', '2']: - key = 'chi_{}'.format(idx) + for idx in ["1", "2"]: + key = "chi_{}".format(idx) if key in original_keys: if "chi_{}_in_plane".format(idx) in original_keys: converted_parameters["a_{}".format(idx)] = ( @@ -242,10 +278,12 @@ def convert_to_lal_binary_black_hole_parameters(parameters): / converted_parameters[f"a_{idx}"] ) elif "a_{}".format(idx) not in original_keys: - converted_parameters['a_{}'.format(idx)] = abs( - converted_parameters[key]) - converted_parameters['cos_tilt_{}'.format(idx)] = \ - xp.sign(xp.asarray(converted_parameters[key])) + converted_parameters["a_{}".format(idx)] = abs( + converted_parameters[key] + ) + converted_parameters["cos_tilt_{}".format(idx)] = xp.sign( + xp.asarray(converted_parameters[key]) + ) else: with np.errstate(invalid="raise"): try: @@ -264,8 +302,8 @@ def convert_to_lal_binary_black_hole_parameters(parameters): if key not in converted_parameters: converted_parameters[key] = 0.0 - for angle in ['tilt_1', 'tilt_2', 'theta_jn']: - cos_angle = str('cos_' + angle) + for angle in ["tilt_1", "tilt_2", "theta_jn"]: + cos_angle = str("cos_" + angle) if cos_angle in converted_parameters.keys(): with np.errstate(invalid="ignore"): converted_parameters[angle] = xp.arccos(converted_parameters[cos_angle]) @@ -276,9 +314,11 @@ def convert_to_lal_binary_black_hole_parameters(parameters): converted_parameters["delta_phase"] - xp.sign(xp.cos(converted_parameters["theta_jn"])) * converted_parameters["psi"], - 2 * np.pi) - added_keys = [key for key in converted_parameters.keys() - if key not in original_keys] + 2 * np.pi, + ) + added_keys = [ + key for key in converted_parameters.keys() if key not in original_keys + ] return converted_parameters, added_keys @@ -313,275 +353,387 @@ def convert_to_lal_binary_neutron_star_parameters(parameters): """ converted_parameters = parameters.copy() original_keys = list(converted_parameters.keys()) - converted_parameters, added_keys =\ - convert_to_lal_binary_black_hole_parameters(converted_parameters) - - if not any([key in converted_parameters for key in - ['lambda_1', 'lambda_2', - 'lambda_tilde', 'delta_lambda_tilde', 'lambda_symmetric', - 'eos_polytrope_gamma_0', 'eos_spectral_pca_gamma_0', 'eos_v1']]): - converted_parameters['lambda_1'] = 0 - converted_parameters['lambda_2'] = 0 - added_keys = added_keys + ['lambda_1', 'lambda_2'] + converted_parameters, added_keys = convert_to_lal_binary_black_hole_parameters( + converted_parameters + ) + + if not any( + [ + key in converted_parameters + for key in [ + "lambda_1", + "lambda_2", + "lambda_tilde", + "delta_lambda_tilde", + "lambda_symmetric", + "eos_polytrope_gamma_0", + "eos_spectral_pca_gamma_0", + "eos_v1", + ] + ] + ): + converted_parameters["lambda_1"] = 0 + converted_parameters["lambda_2"] = 0 + added_keys = added_keys + ["lambda_1", "lambda_2"] return converted_parameters, added_keys - if 'delta_lambda_tilde' in converted_parameters.keys(): - converted_parameters['lambda_1'], converted_parameters['lambda_2'] =\ + if "delta_lambda_tilde" in converted_parameters.keys(): + converted_parameters["lambda_1"], converted_parameters["lambda_2"] = ( lambda_tilde_delta_lambda_tilde_to_lambda_1_lambda_2( - converted_parameters['lambda_tilde'], - parameters['delta_lambda_tilde'], converted_parameters['mass_1'], - converted_parameters['mass_2']) - elif 'lambda_tilde' in converted_parameters.keys(): - converted_parameters['lambda_1'], converted_parameters['lambda_2'] =\ + converted_parameters["lambda_tilde"], + parameters["delta_lambda_tilde"], + converted_parameters["mass_1"], + converted_parameters["mass_2"], + ) + ) + elif "lambda_tilde" in converted_parameters.keys(): + converted_parameters["lambda_1"], converted_parameters["lambda_2"] = ( lambda_tilde_to_lambda_1_lambda_2( - converted_parameters['lambda_tilde'], - converted_parameters['mass_1'], converted_parameters['mass_2']) - if 'lambda_2' not in converted_parameters.keys() and 'lambda_1' in converted_parameters.keys(): - converted_parameters['lambda_2'] =\ - converted_parameters['lambda_1']\ - * converted_parameters['mass_1']**5\ - / converted_parameters['mass_2']**5 - elif 'lambda_2' in converted_parameters.keys() and converted_parameters['lambda_2'] is None: - converted_parameters['lambda_2'] =\ - converted_parameters['lambda_1']\ - * converted_parameters['mass_1']**5\ - / converted_parameters['mass_2']**5 - elif 'eos_spectral_pca_gamma_0' in converted_parameters.keys(): # FIXME: This is a clunky way to do this + converted_parameters["lambda_tilde"], + converted_parameters["mass_1"], + converted_parameters["mass_2"], + ) + ) + if ( + "lambda_2" not in converted_parameters.keys() + and "lambda_1" in converted_parameters.keys() + ): + converted_parameters["lambda_2"] = ( + converted_parameters["lambda_1"] + * converted_parameters["mass_1"] ** 5 + / converted_parameters["mass_2"] ** 5 + ) + elif ( + "lambda_2" in converted_parameters.keys() + and converted_parameters["lambda_2"] is None + ): + converted_parameters["lambda_2"] = ( + converted_parameters["lambda_1"] + * converted_parameters["mass_1"] ** 5 + / converted_parameters["mass_2"] ** 5 + ) + elif ( + "eos_spectral_pca_gamma_0" in converted_parameters.keys() + ): # FIXME: This is a clunky way to do this converted_parameters = generate_source_frame_parameters(converted_parameters) float_eos_params = {} max_len = 1 - eos_keys = ['eos_spectral_pca_gamma_0', - 'eos_spectral_pca_gamma_1', - 'eos_spectral_pca_gamma_2', - 'eos_spectral_pca_gamma_3', - 'mass_1_source', 'mass_2_source'] + eos_keys = [ + "eos_spectral_pca_gamma_0", + "eos_spectral_pca_gamma_1", + "eos_spectral_pca_gamma_2", + "eos_spectral_pca_gamma_3", + "mass_1_source", + "mass_2_source", + ] for key in eos_keys: try: - if (len(converted_parameters[key]) > max_len): + if len(converted_parameters[key]) > max_len: max_len = len(converted_parameters[key]) except TypeError: float_eos_params[key] = converted_parameters[key] - if len(float_eos_params) == len(eos_keys): # case where all eos params are floats (pinned) + if len(float_eos_params) == len( + eos_keys + ): # case where all eos params are floats (pinned) g0, g1, g2, g3 = spectral_pca_to_spectral( - converted_parameters['eos_spectral_pca_gamma_0'], - converted_parameters['eos_spectral_pca_gamma_1'], - converted_parameters['eos_spectral_pca_gamma_2'], - converted_parameters['eos_spectral_pca_gamma_3']) - converted_parameters['lambda_1'], converted_parameters['lambda_2'], converted_parameters['eos_check'] = \ - spectral_params_to_lambda_1_lambda_2( - g0, g1, g2, g3, converted_parameters['mass_1_source'], converted_parameters['mass_2_source']) - elif len(float_eos_params) < len(eos_keys): # case where some or none of the eos params are floats (pinned) + converted_parameters["eos_spectral_pca_gamma_0"], + converted_parameters["eos_spectral_pca_gamma_1"], + converted_parameters["eos_spectral_pca_gamma_2"], + converted_parameters["eos_spectral_pca_gamma_3"], + ) + ( + converted_parameters["lambda_1"], + converted_parameters["lambda_2"], + converted_parameters["eos_check"], + ) = spectral_params_to_lambda_1_lambda_2( + g0, + g1, + g2, + g3, + converted_parameters["mass_1_source"], + converted_parameters["mass_2_source"], + ) + elif len(float_eos_params) < len( + eos_keys + ): # case where some or none of the eos params are floats (pinned) for key in float_eos_params.keys(): converted_parameters[key] = np.ones(max_len) * converted_parameters[key] - g0pca = converted_parameters['eos_spectral_pca_gamma_0'] - g1pca = converted_parameters['eos_spectral_pca_gamma_1'] - g2pca = converted_parameters['eos_spectral_pca_gamma_2'] - g3pca = converted_parameters['eos_spectral_pca_gamma_3'] - m1s = converted_parameters['mass_1_source'] - m2s = converted_parameters['mass_2_source'] + g0pca = converted_parameters["eos_spectral_pca_gamma_0"] + g1pca = converted_parameters["eos_spectral_pca_gamma_1"] + g2pca = converted_parameters["eos_spectral_pca_gamma_2"] + g3pca = converted_parameters["eos_spectral_pca_gamma_3"] + m1s = converted_parameters["mass_1_source"] + m2s = converted_parameters["mass_2_source"] all_lambda_1 = list() all_lambda_2 = list() all_eos_check = list() - for (g_0pca, g_1pca, g_2pca, g_3pca, m1_s, m2_s) in zip(g0pca, g1pca, g2pca, g3pca, m1s, m2s): - g_0, g_1, g_2, g_3 = spectral_pca_to_spectral(g_0pca, g_1pca, g_2pca, g_3pca) - lambda_1, lambda_2, eos_check = \ - spectral_params_to_lambda_1_lambda_2(g_0, g_1, g_2, g_3, m1_s, m2_s) + for g_0pca, g_1pca, g_2pca, g_3pca, m1_s, m2_s in zip( + g0pca, g1pca, g2pca, g3pca, m1s, m2s + ): + g_0, g_1, g_2, g_3 = spectral_pca_to_spectral( + g_0pca, g_1pca, g_2pca, g_3pca + ) + lambda_1, lambda_2, eos_check = spectral_params_to_lambda_1_lambda_2( + g_0, g_1, g_2, g_3, m1_s, m2_s + ) all_lambda_1.append(lambda_1) all_lambda_2.append(lambda_2) all_eos_check.append(eos_check) - converted_parameters['lambda_1'] = np.array(all_lambda_1) - converted_parameters['lambda_2'] = np.array(all_lambda_2) - converted_parameters['eos_check'] = np.array(all_eos_check) + converted_parameters["lambda_1"] = np.array(all_lambda_1) + converted_parameters["lambda_2"] = np.array(all_lambda_2) + converted_parameters["eos_check"] = np.array(all_eos_check) for key in float_eos_params.keys(): converted_parameters[key] = float_eos_params[key] - elif 'eos_polytrope_gamma_0' and 'eos_polytrope_log10_pressure_1' in converted_parameters.keys(): + elif ( + "eos_polytrope_gamma_0" + and "eos_polytrope_log10_pressure_1" in converted_parameters.keys() + ): converted_parameters = generate_source_frame_parameters(converted_parameters) float_eos_params = {} max_len = 1 - eos_keys = ['eos_polytrope_gamma_0', - 'eos_polytrope_gamma_1', - 'eos_polytrope_gamma_2', - 'eos_polytrope_log10_pressure_1', - 'eos_polytrope_log10_pressure_2', - 'mass_1_source', 'mass_2_source'] + eos_keys = [ + "eos_polytrope_gamma_0", + "eos_polytrope_gamma_1", + "eos_polytrope_gamma_2", + "eos_polytrope_log10_pressure_1", + "eos_polytrope_log10_pressure_2", + "mass_1_source", + "mass_2_source", + ] for key in eos_keys: try: - if (len(converted_parameters[key]) > max_len): + if len(converted_parameters[key]) > max_len: max_len = len(converted_parameters[key]) except TypeError: float_eos_params[key] = converted_parameters[key] - if len(float_eos_params) == len(eos_keys): # case where all eos params are floats (pinned) - converted_parameters['lambda_1'], converted_parameters['lambda_2'], converted_parameters['eos_check'] = \ - polytrope_or_causal_params_to_lambda_1_lambda_2( - converted_parameters['eos_polytrope_gamma_0'], - converted_parameters['eos_polytrope_log10_pressure_1'], - converted_parameters['eos_polytrope_gamma_1'], - converted_parameters['eos_polytrope_log10_pressure_2'], - converted_parameters['eos_polytrope_gamma_2'], - converted_parameters['mass_1_source'], - converted_parameters['mass_2_source'], - causal=0) - elif len(float_eos_params) < len(eos_keys): # case where some or none are floats (pinned) + if len(float_eos_params) == len( + eos_keys + ): # case where all eos params are floats (pinned) + ( + converted_parameters["lambda_1"], + converted_parameters["lambda_2"], + converted_parameters["eos_check"], + ) = polytrope_or_causal_params_to_lambda_1_lambda_2( + converted_parameters["eos_polytrope_gamma_0"], + converted_parameters["eos_polytrope_log10_pressure_1"], + converted_parameters["eos_polytrope_gamma_1"], + converted_parameters["eos_polytrope_log10_pressure_2"], + converted_parameters["eos_polytrope_gamma_2"], + converted_parameters["mass_1_source"], + converted_parameters["mass_2_source"], + causal=0, + ) + elif len(float_eos_params) < len( + eos_keys + ): # case where some or none are floats (pinned) for key in float_eos_params.keys(): converted_parameters[key] = np.ones(max_len) * converted_parameters[key] - pg0 = converted_parameters['eos_polytrope_gamma_0'] - pg1 = converted_parameters['eos_polytrope_gamma_1'] - pg2 = converted_parameters['eos_polytrope_gamma_2'] - logp1 = converted_parameters['eos_polytrope_log10_pressure_1'] - logp2 = converted_parameters['eos_polytrope_log10_pressure_2'] - m1s = converted_parameters['mass_1_source'] - m2s = converted_parameters['mass_2_source'] + pg0 = converted_parameters["eos_polytrope_gamma_0"] + pg1 = converted_parameters["eos_polytrope_gamma_1"] + pg2 = converted_parameters["eos_polytrope_gamma_2"] + logp1 = converted_parameters["eos_polytrope_log10_pressure_1"] + logp2 = converted_parameters["eos_polytrope_log10_pressure_2"] + m1s = converted_parameters["mass_1_source"] + m2s = converted_parameters["mass_2_source"] all_lambda_1 = np.empty(0) all_lambda_2 = np.empty(0) all_eos_check = np.empty(0, dtype=bool) - for (pg_0, pg_1, pg_2, logp_1, logp_2, m1_s, m2_s) in zip(pg0, pg1, pg2, logp1, logp2, m1s, m2s): - lambda_1, lambda_2, eos_check = \ + for pg_0, pg_1, pg_2, logp_1, logp_2, m1_s, m2_s in zip( + pg0, pg1, pg2, logp1, logp2, m1s, m2s + ): + lambda_1, lambda_2, eos_check = ( polytrope_or_causal_params_to_lambda_1_lambda_2( - pg_0, logp_1, pg_1, logp_2, pg_2, m1_s, m2_s, causal=0) + pg_0, logp_1, pg_1, logp_2, pg_2, m1_s, m2_s, causal=0 + ) + ) all_lambda_1 = np.append(all_lambda_1, lambda_1) all_lambda_2 = np.append(all_lambda_2, lambda_2) all_eos_check = np.append(all_eos_check, eos_check) - converted_parameters['lambda_1'] = all_lambda_1 - converted_parameters['lambda_2'] = all_lambda_2 - converted_parameters['eos_check'] = all_eos_check + converted_parameters["lambda_1"] = all_lambda_1 + converted_parameters["lambda_2"] = all_lambda_2 + converted_parameters["eos_check"] = all_eos_check for key in float_eos_params.keys(): converted_parameters[key] = float_eos_params[key] - elif 'eos_polytrope_gamma_0' and 'eos_polytrope_scaled_pressure_ratio' in converted_parameters.keys(): + elif ( + "eos_polytrope_gamma_0" + and "eos_polytrope_scaled_pressure_ratio" in converted_parameters.keys() + ): converted_parameters = generate_source_frame_parameters(converted_parameters) float_eos_params = {} max_len = 1 - eos_keys = ['eos_polytrope_gamma_0', - 'eos_polytrope_gamma_1', - 'eos_polytrope_gamma_2', - 'eos_polytrope_scaled_pressure_ratio', - 'eos_polytrope_scaled_pressure_2', - 'mass_1_source', 'mass_2_source'] + eos_keys = [ + "eos_polytrope_gamma_0", + "eos_polytrope_gamma_1", + "eos_polytrope_gamma_2", + "eos_polytrope_scaled_pressure_ratio", + "eos_polytrope_scaled_pressure_2", + "mass_1_source", + "mass_2_source", + ] for key in eos_keys: try: - if (len(converted_parameters[key]) > max_len): + if len(converted_parameters[key]) > max_len: max_len = len(converted_parameters[key]) except TypeError: float_eos_params[key] = converted_parameters[key] - if len(float_eos_params) == len(eos_keys): # case where all eos params are floats (pinned) + if len(float_eos_params) == len( + eos_keys + ): # case where all eos params are floats (pinned) logp1, logp2 = log_pressure_reparameterization_conversion( - converted_parameters['eos_polytrope_scaled_pressure_ratio'], - converted_parameters['eos_polytrope_scaled_pressure_2']) - converted_parameters['lambda_1'], converted_parameters['lambda_2'], converted_parameters['eos_check'] = \ - polytrope_or_causal_params_to_lambda_1_lambda_2( - converted_parameters['eos_polytrope_gamma_0'], - logp1, - converted_parameters['eos_polytrope_gamma_1'], - logp2, - converted_parameters['eos_polytrope_gamma_2'], - converted_parameters['mass_1_source'], - converted_parameters['mass_2_source'], - causal=0) - elif len(float_eos_params) < len(eos_keys): # case where some or none are floats (pinned) + converted_parameters["eos_polytrope_scaled_pressure_ratio"], + converted_parameters["eos_polytrope_scaled_pressure_2"], + ) + ( + converted_parameters["lambda_1"], + converted_parameters["lambda_2"], + converted_parameters["eos_check"], + ) = polytrope_or_causal_params_to_lambda_1_lambda_2( + converted_parameters["eos_polytrope_gamma_0"], + logp1, + converted_parameters["eos_polytrope_gamma_1"], + logp2, + converted_parameters["eos_polytrope_gamma_2"], + converted_parameters["mass_1_source"], + converted_parameters["mass_2_source"], + causal=0, + ) + elif len(float_eos_params) < len( + eos_keys + ): # case where some or none are floats (pinned) for key in float_eos_params.keys(): converted_parameters[key] = np.ones(max_len) * converted_parameters[key] - pg0 = converted_parameters['eos_polytrope_gamma_0'] - pg1 = converted_parameters['eos_polytrope_gamma_1'] - pg2 = converted_parameters['eos_polytrope_gamma_2'] - scaledratio = converted_parameters['eos_polytrope_scaled_pressure_ratio'] - scaled_p2 = converted_parameters['eos_polytrope_scaled_pressure_2'] - m1s = converted_parameters['mass_1_source'] - m2s = converted_parameters['mass_2_source'] + pg0 = converted_parameters["eos_polytrope_gamma_0"] + pg1 = converted_parameters["eos_polytrope_gamma_1"] + pg2 = converted_parameters["eos_polytrope_gamma_2"] + scaledratio = converted_parameters["eos_polytrope_scaled_pressure_ratio"] + scaled_p2 = converted_parameters["eos_polytrope_scaled_pressure_2"] + m1s = converted_parameters["mass_1_source"] + m2s = converted_parameters["mass_2_source"] all_lambda_1 = np.empty(0) all_lambda_2 = np.empty(0) all_eos_check = np.empty(0, dtype=bool) - for (pg_0, pg_1, pg_2, scaled_ratio, scaled_p_2, m1_s, - m2_s) in zip(pg0, pg1, pg2, scaledratio, scaled_p2, m1s, m2s): - logp_1, logp_2 = log_pressure_reparameterization_conversion(scaled_ratio, scaled_p_2) - lambda_1, lambda_2, eos_check = \ + for pg_0, pg_1, pg_2, scaled_ratio, scaled_p_2, m1_s, m2_s in zip( + pg0, pg1, pg2, scaledratio, scaled_p2, m1s, m2s + ): + logp_1, logp_2 = log_pressure_reparameterization_conversion( + scaled_ratio, scaled_p_2 + ) + lambda_1, lambda_2, eos_check = ( polytrope_or_causal_params_to_lambda_1_lambda_2( - pg_0, logp_1, pg_1, logp_2, pg_2, m1_s, m2_s, causal=0) + pg_0, logp_1, pg_1, logp_2, pg_2, m1_s, m2_s, causal=0 + ) + ) all_lambda_1 = np.append(all_lambda_1, lambda_1) all_lambda_2 = np.append(all_lambda_2, lambda_2) all_eos_check = np.append(all_eos_check, eos_check) - converted_parameters['lambda_1'] = all_lambda_1 - converted_parameters['lambda_2'] = all_lambda_2 - converted_parameters['eos_check'] = all_eos_check + converted_parameters["lambda_1"] = all_lambda_1 + converted_parameters["lambda_2"] = all_lambda_2 + converted_parameters["eos_check"] = all_eos_check for key in float_eos_params.keys(): converted_parameters[key] = float_eos_params[key] - elif 'eos_v1' in converted_parameters.keys(): + elif "eos_v1" in converted_parameters.keys(): converted_parameters = generate_source_frame_parameters(converted_parameters) float_eos_params = {} max_len = 1 - eos_keys = ['eos_v1', - 'eos_v2', - 'eos_v3', - 'eos_log10_pressure1_cgs', - 'eos_log10_pressure2_cgs', - 'mass_1_source', 'mass_2_source'] + eos_keys = [ + "eos_v1", + "eos_v2", + "eos_v3", + "eos_log10_pressure1_cgs", + "eos_log10_pressure2_cgs", + "mass_1_source", + "mass_2_source", + ] for key in eos_keys: try: - if (len(converted_parameters[key]) > max_len): + if len(converted_parameters[key]) > max_len: max_len = len(converted_parameters[key]) except TypeError: float_eos_params[key] = converted_parameters[key] - if len(float_eos_params) == len(eos_keys): # case where all eos params are floats (pinned) - converted_parameters['lambda_1'], converted_parameters['lambda_2'], converted_parameters['eos_check'] = \ - polytrope_or_causal_params_to_lambda_1_lambda_2( - converted_parameters['eos_v1'], - converted_parameters['eos_log10_pressure1_cgs'], - converted_parameters['eos_v2'], - converted_parameters['eos_log10_pressure2_cgs'], - converted_parameters['eos_v3'], - converted_parameters['mass_1_source'], - converted_parameters['mass_2_source'], - causal=1) - elif len(float_eos_params) < len(eos_keys): # case where some or none are floats (pinned) + if len(float_eos_params) == len( + eos_keys + ): # case where all eos params are floats (pinned) + ( + converted_parameters["lambda_1"], + converted_parameters["lambda_2"], + converted_parameters["eos_check"], + ) = polytrope_or_causal_params_to_lambda_1_lambda_2( + converted_parameters["eos_v1"], + converted_parameters["eos_log10_pressure1_cgs"], + converted_parameters["eos_v2"], + converted_parameters["eos_log10_pressure2_cgs"], + converted_parameters["eos_v3"], + converted_parameters["mass_1_source"], + converted_parameters["mass_2_source"], + causal=1, + ) + elif len(float_eos_params) < len( + eos_keys + ): # case where some or none are floats (pinned) for key in float_eos_params.keys(): converted_parameters[key] = np.ones(max_len) * converted_parameters[key] - v1 = converted_parameters['eos_v1'] - v2 = converted_parameters['eos_v2'] - v3 = converted_parameters['eos_v3'] - logp1 = converted_parameters['eos_log10_pressure1_cgs'] - logp2 = converted_parameters['eos_log10_pressure2_cgs'] - m1s = converted_parameters['mass_1_source'] - m2s = converted_parameters['mass_2_source'] + v1 = converted_parameters["eos_v1"] + v2 = converted_parameters["eos_v2"] + v3 = converted_parameters["eos_v3"] + logp1 = converted_parameters["eos_log10_pressure1_cgs"] + logp2 = converted_parameters["eos_log10_pressure2_cgs"] + m1s = converted_parameters["mass_1_source"] + m2s = converted_parameters["mass_2_source"] all_lambda_1 = np.empty(0) all_lambda_2 = np.empty(0) all_eos_check = np.empty(0, dtype=bool) - for (v_1, v_2, v_3, logp_1, logp_2, m1_s, m2_s) in zip(v1, v2, v3, logp1, logp2, m1s, m2s): - lambda_1, lambda_2, eos_check = \ + for v_1, v_2, v_3, logp_1, logp_2, m1_s, m2_s in zip( + v1, v2, v3, logp1, logp2, m1s, m2s + ): + lambda_1, lambda_2, eos_check = ( polytrope_or_causal_params_to_lambda_1_lambda_2( - v_1, logp_1, v_2, logp_2, v_3, m1_s, m2_s, causal=1) + v_1, logp_1, v_2, logp_2, v_3, m1_s, m2_s, causal=1 + ) + ) all_lambda_1 = np.append(all_lambda_1, lambda_1) all_lambda_2 = np.append(all_lambda_2, lambda_2) all_eos_check = np.append(all_eos_check, eos_check) - converted_parameters['lambda_1'] = all_lambda_1 - converted_parameters['lambda_2'] = all_lambda_2 - converted_parameters['eos_check'] = all_eos_check + converted_parameters["lambda_1"] = all_lambda_1 + converted_parameters["lambda_2"] = all_lambda_2 + converted_parameters["eos_check"] = all_eos_check for key in float_eos_params.keys(): converted_parameters[key] = float_eos_params[key] - elif 'lambda_symmetric' in converted_parameters.keys(): - if 'lambda_antisymmetric' in converted_parameters.keys(): - converted_parameters['lambda_1'], converted_parameters['lambda_2'] =\ + elif "lambda_symmetric" in converted_parameters.keys(): + if "lambda_antisymmetric" in converted_parameters.keys(): + converted_parameters["lambda_1"], converted_parameters["lambda_2"] = ( lambda_symmetric_lambda_antisymmetric_to_lambda_1_lambda_2( - converted_parameters['lambda_symmetric'], - converted_parameters['lambda_antisymmetric']) - elif 'mass_ratio' in converted_parameters.keys(): - if 'binary_love_uniform' in converted_parameters.keys(): - converted_parameters['lambda_1'], converted_parameters['lambda_2'] =\ + converted_parameters["lambda_symmetric"], + converted_parameters["lambda_antisymmetric"], + ) + ) + elif "mass_ratio" in converted_parameters.keys(): + if "binary_love_uniform" in converted_parameters.keys(): + converted_parameters["lambda_1"], converted_parameters["lambda_2"] = ( binary_love_lambda_symmetric_to_lambda_1_lambda_2_manual_marginalisation( - converted_parameters['binary_love_uniform'], - converted_parameters['lambda_symmetric'], - converted_parameters['mass_ratio']) + converted_parameters["binary_love_uniform"], + converted_parameters["lambda_symmetric"], + converted_parameters["mass_ratio"], + ) + ) else: - converted_parameters['lambda_1'], converted_parameters['lambda_2'] =\ + converted_parameters["lambda_1"], converted_parameters["lambda_2"] = ( binary_love_lambda_symmetric_to_lambda_1_lambda_2_automatic_marginalisation( - converted_parameters['lambda_symmetric'], - converted_parameters['mass_ratio']) + converted_parameters["lambda_symmetric"], + converted_parameters["mass_ratio"], + ) + ) - added_keys = [key for key in converted_parameters.keys() - if key not in original_keys] + added_keys = [ + key for key in converted_parameters.keys() if key not in original_keys + ] return converted_parameters, added_keys -def log_pressure_reparameterization_conversion(scaled_pressure_ratio, scaled_pressure_2): - ''' +def log_pressure_reparameterization_conversion( + scaled_pressure_ratio, scaled_pressure_2 +): + """ Converts the reparameterization joining pressures from (scaled_pressure_ratio,scaled_pressure_2) to (log10_pressure_1,log10_pressure_2). This reparameterization with a triangular prior (with mode = max) on scaled_pressure_2 @@ -602,8 +754,8 @@ def log_pressure_reparameterization_conversion(scaled_pressure_ratio, scaled_pre log10_pressure_1, log10_pressure_2: float joining pressures in the original parameterization - ''' - minimum_pressure = 33. + """ + minimum_pressure = 33.0 log10_pressure_1 = (scaled_pressure_ratio * scaled_pressure_2) + minimum_pressure log10_pressure_2 = minimum_pressure + scaled_pressure_2 @@ -611,7 +763,7 @@ def log_pressure_reparameterization_conversion(scaled_pressure_ratio, scaled_pre def spectral_pca_to_spectral(gamma_pca_0, gamma_pca_1, gamma_pca_2, gamma_pca_3): - ''' + """ Change of basis on parameter space from an efficient space to sample in (sample space) to the space used in spectral eos model (model space). @@ -630,28 +782,33 @@ def spectral_pca_to_spectral(gamma_pca_0, gamma_pca_1, gamma_pca_2, gamma_pca_3) converted_gamma_parameters: np.array() array of gamma_0, gamma_1, gamma_2, gamma_3 in model space - ''' + """ xp = array_module(gamma_pca_0) - sampled_pca_gammas = xp.asarray([gamma_pca_0, gamma_pca_1, gamma_pca_2, gamma_pca_3]) + sampled_pca_gammas = xp.asarray( + [gamma_pca_0, gamma_pca_1, gamma_pca_2, gamma_pca_3] + ) transformation_matrix = xp.asarray( [ [0.43801, -0.76705, 0.45143, 0.12646], [-0.53573, 0.17169, 0.67968, 0.47070], [0.52660, 0.31255, -0.19454, 0.76626], - [-0.49379, -0.53336, -0.54444, 0.41868] + [-0.49379, -0.53336, -0.54444, 0.41868], ] ) model_space_mean = xp.asarray([0.89421, 0.33878, -0.07894, 0.00393]) model_space_standard_deviation = xp.asarray([0.35700, 0.25769, 0.05452, 0.00312]) - converted_gamma_parameters = \ - model_space_mean + model_space_standard_deviation * (transformation_matrix @ sampled_pca_gammas) + converted_gamma_parameters = model_space_mean + model_space_standard_deviation * ( + transformation_matrix @ sampled_pca_gammas + ) return converted_gamma_parameters -def spectral_params_to_lambda_1_lambda_2(gamma_0, gamma_1, gamma_2, gamma_3, mass_1_source, mass_2_source): - ''' +def spectral_params_to_lambda_1_lambda_2( + gamma_0, gamma_1, gamma_2, gamma_3, mass_1_source, mass_2_source +): + """ Converts from the 4 spectral decomposition parameters and the source masses to the tidal deformability parameters. @@ -670,26 +827,46 @@ def spectral_params_to_lambda_1_lambda_2(gamma_0, gamma_1, gamma_2, gamma_3, mas whether or not the equation of state is viable / if eos_check = False, lambdas are 0 and the sample is rejected. - ''' + """ eos_check = True - if lalsim_SimNeutronStarEOS4ParamSDGammaCheck(gamma_0, gamma_1, gamma_2, gamma_3) != 0: + if ( + lalsim_SimNeutronStarEOS4ParamSDGammaCheck(gamma_0, gamma_1, gamma_2, gamma_3) + != 0 + ): lambda_1 = 0.0 lambda_2 = 0.0 eos_check = False else: - eos = lalsim_SimNeutronStarEOS4ParameterSpectralDecomposition(gamma_0, gamma_1, gamma_2, gamma_3) - if lalsim_SimNeutronStarEOS4ParamSDViableFamilyCheck(gamma_0, gamma_1, gamma_2, gamma_3) != 0: + eos = lalsim_SimNeutronStarEOS4ParameterSpectralDecomposition( + gamma_0, gamma_1, gamma_2, gamma_3 + ) + if ( + lalsim_SimNeutronStarEOS4ParamSDViableFamilyCheck( + gamma_0, gamma_1, gamma_2, gamma_3 + ) + != 0 + ): lambda_1 = 0.0 lambda_2 = 0.0 eos_check = False else: - lambda_1, lambda_2, eos_check = neutron_star_family_physical_check(eos, mass_1_source, mass_2_source) + lambda_1, lambda_2, eos_check = neutron_star_family_physical_check( + eos, mass_1_source, mass_2_source + ) return lambda_1, lambda_2, eos_check def polytrope_or_causal_params_to_lambda_1_lambda_2( - param1, log10_pressure1_cgs, param2, log10_pressure2_cgs, param3, mass_1_source, mass_2_source, causal): + param1, + log10_pressure1_cgs, + param2, + log10_pressure2_cgs, + param3, + mass_1_source, + mass_2_source, + causal, +): """ Converts parameters from sampled dynamic piecewise polytrope parameters to component tidal deformablity parameters. @@ -729,17 +906,38 @@ def polytrope_or_causal_params_to_lambda_1_lambda_2( else: if causal == 0: eos = lalsim_SimNeutronStarEOS3PieceDynamicPolytrope( - param1, log10_pressure1_cgs - 1., param2, log10_pressure2_cgs - 1., param3) + param1, + log10_pressure1_cgs - 1.0, + param2, + log10_pressure2_cgs - 1.0, + param3, + ) else: eos = lalsim_SimNeutronStarEOS3PieceCausalAnalytic( - param1, log10_pressure1_cgs - 1., param2, log10_pressure2_cgs - 1., param3) - if lalsim_SimNeutronStarEOS3PDViableFamilyCheck( - param1, log10_pressure1_cgs - 1., param2, log10_pressure2_cgs - 1., param3, causal) != 0: + param1, + log10_pressure1_cgs - 1.0, + param2, + log10_pressure2_cgs - 1.0, + param3, + ) + if ( + lalsim_SimNeutronStarEOS3PDViableFamilyCheck( + param1, + log10_pressure1_cgs - 1.0, + param2, + log10_pressure2_cgs - 1.0, + param3, + causal, + ) + != 0 + ): lambda_1 = 0.0 lambda_2 = 0.0 eos_check = False else: - lambda_1, lambda_2, eos_check = neutron_star_family_physical_check(eos, mass_1_source, mass_2_source) + lambda_1, lambda_2, eos_check = neutron_star_family_physical_check( + eos, mass_1_source, mass_2_source + ) return lambda_1, lambda_2, eos_check @@ -768,10 +966,16 @@ def neutron_star_family_physical_check(eos, mass_1_source, mass_2_source): eos_check = True family = lalsim_CreateSimNeutronStarFamily(eos) max_pseudo_enthalpy = lalsim_SimNeutronStarEOSMaxPseudoEnthalpy(eos) - max_speed_of_sound = lalsim_SimNeutronStarEOSSpeedOfSoundGeometerized(max_pseudo_enthalpy, eos) + max_speed_of_sound = lalsim_SimNeutronStarEOSSpeedOfSoundGeometerized( + max_pseudo_enthalpy, eos + ) min_mass = lalsim_SimNeutronStarFamMinimumMass(family) / solar_mass max_mass = lalsim_SimNeutronStarMaximumMass(family) / solar_mass - if max_speed_of_sound <= 1.1 and min_mass <= mass_1_source <= max_mass and min_mass <= mass_2_source <= max_mass: + if ( + max_speed_of_sound <= 1.1 + and min_mass <= mass_1_source <= max_mass + and min_mass <= mass_2_source <= max_mass + ): lambda_1 = lambda_from_mass_and_family(mass_1_source, family) lambda_2 = lambda_from_mass_and_family(mass_2_source, family) else: @@ -801,9 +1005,11 @@ def lambda_from_mass_and_family(mass_i, family): """ radius = lalsim_SimNeutronStarRadius(mass_i * solar_mass, family) love_number_k2 = lalsim_SimNeutronStarLoveNumberK2(mass_i * solar_mass, family) - mass_geometrized = mass_i * solar_mass * gravitational_constant / speed_of_light ** 2. + mass_geometrized = ( + mass_i * solar_mass * gravitational_constant / speed_of_light**2.0 + ) compactness = mass_geometrized / radius - lambda_i = (2. / 3.) * love_number_k2 / compactness ** 5. + lambda_i = (2.0 / 3.0) * love_number_k2 / compactness**5.0 return lambda_i @@ -888,11 +1094,11 @@ def chirp_mass_and_mass_ratio_to_component_masses(chirp_mass, mass_ratio): mass_2: float Mass of the lighter object """ - total_mass = chirp_mass_and_mass_ratio_to_total_mass(chirp_mass=chirp_mass, - mass_ratio=mass_ratio) - mass_1, mass_2 = ( - total_mass_and_mass_ratio_to_component_masses( - total_mass=total_mass, mass_ratio=mass_ratio) + total_mass = chirp_mass_and_mass_ratio_to_total_mass( + chirp_mass=chirp_mass, mass_ratio=mass_ratio + ) + mass_1, mass_2 = total_mass_and_mass_ratio_to_component_masses( + total_mass=total_mass, mass_ratio=mass_ratio ) return mass_1, mass_2 @@ -912,8 +1118,8 @@ def symmetric_mass_ratio_to_mass_ratio(symmetric_mass_ratio): Mass ratio of the binary """ - temp = (1 / symmetric_mass_ratio / 2 - 1) - return temp - (temp ** 2 - 1) ** 0.5 + temp = 1 / symmetric_mass_ratio / 2 - 1 + return temp - (temp**2 - 1) ** 0.5 def chirp_mass_and_total_mass_to_symmetric_mass_ratio(chirp_mass, total_mass): @@ -960,9 +1166,9 @@ def chirp_mass_and_primary_mass_to_mass_ratio(chirp_mass, mass_1): Mass ratio (mass_2/mass_1) of the binary """ a = (chirp_mass / mass_1) ** 5 - t0 = (9 * a + 3**0.5 * (27 * a ** 2 - 4 * a ** 3)**0.5)**(1 / 3) - t1 = (2)**(1 / 3) * 3 ** (2 / 3) - t2 = (2 / 3)**(1 / 3) * a + t0 = (9 * a + 3**0.5 * (27 * a**2 - 4 * a**3) ** 0.5) ** (1 / 3) + t1 = (2) ** (1 / 3) * 3 ** (2 / 3) + t2 = (2 / 3) ** (1 / 3) * a return t2 / t0 + t0 / t1 @@ -986,7 +1192,7 @@ def chirp_mass_and_mass_ratio_to_total_mass(chirp_mass, mass_ratio): """ with np.errstate(invalid="ignore"): - return chirp_mass * (1 + mass_ratio) ** 1.2 / mass_ratio ** 0.6 + return chirp_mass * (1 + mass_ratio) ** 1.2 / mass_ratio**0.6 def component_masses_to_chirp_mass(mass_1, mass_2): @@ -1088,10 +1294,11 @@ def mass_1_and_chirp_mass_to_mass_ratio(mass_1, chirp_mass): Mass ratio of the binary """ temp = (chirp_mass / mass_1) ** 5 - mass_ratio = (2 / 3 / (3 ** 0.5 * (27 * temp ** 2 - 4 * temp ** 3) ** 0.5 + - 9 * temp)) ** (1 / 3) * temp + \ - ((3 ** 0.5 * (27 * temp ** 2 - 4 * temp ** 3) ** 0.5 + - 9 * temp) / (2 * 3 ** 2)) ** (1 / 3) + mass_ratio = ( + 2 / 3 / (3**0.5 * (27 * temp**2 - 4 * temp**3) ** 0.5 + 9 * temp) + ) ** (1 / 3) * temp + ( + (3**0.5 * (27 * temp**2 - 4 * temp**3) ** 0.5 + 9 * temp) / (2 * 3**2) + ) ** (1 / 3) return mass_ratio @@ -1145,9 +1352,14 @@ def lambda_1_lambda_2_to_lambda_tilde(lambda_1, lambda_2, mass_1, mass_2): eta = component_masses_to_symmetric_mass_ratio(mass_1, mass_2) lambda_plus = lambda_1 + lambda_2 lambda_minus = lambda_1 - lambda_2 - lambda_tilde = 8 / 13 * ( - (1 + 7 * eta - 31 * eta**2) * lambda_plus + - (1 - 4 * eta)**0.5 * (1 + 9 * eta - 11 * eta**2) * lambda_minus) + lambda_tilde = ( + 8 + / 13 + * ( + (1 + 7 * eta - 31 * eta**2) * lambda_plus + + (1 - 4 * eta) ** 0.5 * (1 + 9 * eta - 11 * eta**2) * lambda_minus + ) + ) return lambda_tilde @@ -1177,15 +1389,23 @@ def lambda_1_lambda_2_to_delta_lambda_tilde(lambda_1, lambda_2, mass_1, mass_2): eta = component_masses_to_symmetric_mass_ratio(mass_1, mass_2) lambda_plus = lambda_1 + lambda_2 lambda_minus = lambda_1 - lambda_2 - delta_lambda_tilde = 1 / 2 * ( - (1 - 4 * eta) ** 0.5 * (1 - 13272 / 1319 * eta + 8944 / 1319 * eta**2) * - lambda_plus + (1 - 15910 / 1319 * eta + 32850 / 1319 * eta ** 2 + - 3380 / 1319 * eta ** 3) * lambda_minus) + delta_lambda_tilde = ( + 1 + / 2 + * ( + (1 - 4 * eta) ** 0.5 + * (1 - 13272 / 1319 * eta + 8944 / 1319 * eta**2) + * lambda_plus + + (1 - 15910 / 1319 * eta + 32850 / 1319 * eta**2 + 3380 / 1319 * eta**3) + * lambda_minus + ) + ) return delta_lambda_tilde def lambda_tilde_delta_lambda_tilde_to_lambda_1_lambda_2( - lambda_tilde, delta_lambda_tilde, mass_1, mass_2): + lambda_tilde, delta_lambda_tilde, mass_1, mass_2 +): """ Convert from dominant tidal terms to individual tidal parameters. @@ -1211,28 +1431,33 @@ def lambda_tilde_delta_lambda_tilde_to_lambda_1_lambda_2( """ eta = component_masses_to_symmetric_mass_ratio(mass_1, mass_2) - coefficient_1 = (1 + 7 * eta - 31 * eta**2) - coefficient_2 = (1 - 4 * eta)**0.5 * (1 + 9 * eta - 11 * eta**2) - coefficient_3 = (1 - 4 * eta)**0.5 *\ - (1 - 13272 / 1319 * eta + 8944 / 1319 * eta**2) - coefficient_4 = (1 - 15910 / 1319 * eta + 32850 / 1319 * eta**2 + - 3380 / 1319 * eta**3) - lambda_1 =\ - (13 * lambda_tilde / 8 * (coefficient_3 - coefficient_4) - - 2 * delta_lambda_tilde * (coefficient_1 - coefficient_2))\ - / ((coefficient_1 + coefficient_2) * (coefficient_3 - coefficient_4) - - (coefficient_1 - coefficient_2) * (coefficient_3 + coefficient_4)) - lambda_2 =\ - (13 * lambda_tilde / 8 * (coefficient_3 + coefficient_4) - - 2 * delta_lambda_tilde * (coefficient_1 + coefficient_2)) \ - / ((coefficient_1 - coefficient_2) * (coefficient_3 + coefficient_4) - - (coefficient_1 + coefficient_2) * (coefficient_3 - coefficient_4)) + coefficient_1 = 1 + 7 * eta - 31 * eta**2 + coefficient_2 = (1 - 4 * eta) ** 0.5 * (1 + 9 * eta - 11 * eta**2) + coefficient_3 = (1 - 4 * eta) ** 0.5 * ( + 1 - 13272 / 1319 * eta + 8944 / 1319 * eta**2 + ) + coefficient_4 = ( + 1 - 15910 / 1319 * eta + 32850 / 1319 * eta**2 + 3380 / 1319 * eta**3 + ) + lambda_1 = ( + 13 * lambda_tilde / 8 * (coefficient_3 - coefficient_4) + - 2 * delta_lambda_tilde * (coefficient_1 - coefficient_2) + ) / ( + (coefficient_1 + coefficient_2) * (coefficient_3 - coefficient_4) + - (coefficient_1 - coefficient_2) * (coefficient_3 + coefficient_4) + ) + lambda_2 = ( + 13 * lambda_tilde / 8 * (coefficient_3 + coefficient_4) + - 2 * delta_lambda_tilde * (coefficient_1 + coefficient_2) + ) / ( + (coefficient_1 - coefficient_2) * (coefficient_3 + coefficient_4) + - (coefficient_1 + coefficient_2) * (coefficient_3 - coefficient_4) + ) return lambda_1, lambda_2 -def lambda_tilde_to_lambda_1_lambda_2( - lambda_tilde, mass_1, mass_2): +def lambda_tilde_to_lambda_1_lambda_2(lambda_tilde, mass_1, mass_2): """ Convert from dominant tidal term to individual tidal parameters assuming lambda_1 * mass_1**5 = lambda_2 * mass_2**5. @@ -1257,9 +1482,15 @@ def lambda_tilde_to_lambda_1_lambda_2( """ eta = component_masses_to_symmetric_mass_ratio(mass_1, mass_2) q = mass_2 / mass_1 - lambda_1 = 13 / 8 * lambda_tilde / ( - (1 + 7 * eta - 31 * eta**2) * (1 + q**-5) + - (1 - 4 * eta)**0.5 * (1 + 9 * eta - 11 * eta**2) * (1 - q**-5)) + lambda_1 = ( + 13 + / 8 + * lambda_tilde + / ( + (1 + 7 * eta - 31 * eta**2) * (1 + q**-5) + + (1 - 4 * eta) ** 0.5 * (1 + 9 * eta - 11 * eta**2) * (1 - q**-5) + ) + ) lambda_2 = lambda_1 / q**5 return lambda_1, lambda_2 @@ -1282,7 +1513,7 @@ def lambda_1_lambda_2_to_lambda_symmetric(lambda_1, lambda_2): lambda_symmetric: float Symmetric tidal parameter. """ - lambda_symmetric = (lambda_2 + lambda_1) / 2. + lambda_symmetric = (lambda_2 + lambda_1) / 2.0 return lambda_symmetric @@ -1304,11 +1535,13 @@ def lambda_1_lambda_2_to_lambda_antisymmetric(lambda_1, lambda_2): lambda_antisymmetric: float Antisymmetric tidal parameter. """ - lambda_antisymmetric = (lambda_2 - lambda_1) / 2. + lambda_antisymmetric = (lambda_2 - lambda_1) / 2.0 return lambda_antisymmetric -def lambda_symmetric_lambda_antisymmetric_to_lambda_1(lambda_symmetric, lambda_antisymmetric): +def lambda_symmetric_lambda_antisymmetric_to_lambda_1( + lambda_symmetric, lambda_antisymmetric +): """ Convert from symmetric and antisymmetric tidal terms to lambda_1 @@ -1330,7 +1563,9 @@ def lambda_symmetric_lambda_antisymmetric_to_lambda_1(lambda_symmetric, lambda_a return lambda_1 -def lambda_symmetric_lambda_antisymmetric_to_lambda_2(lambda_symmetric, lambda_antisymmetric): +def lambda_symmetric_lambda_antisymmetric_to_lambda_2( + lambda_symmetric, lambda_antisymmetric +): """ Convert from symmetric and antisymmetric tidal terms to lambda_2 @@ -1352,7 +1587,9 @@ def lambda_symmetric_lambda_antisymmetric_to_lambda_2(lambda_symmetric, lambda_a return lambda_2 -def lambda_symmetric_lambda_antisymmetric_to_lambda_1_lambda_2(lambda_symmetric, lambda_antisymmetric): +def lambda_symmetric_lambda_antisymmetric_to_lambda_1_lambda_2( + lambda_symmetric, lambda_antisymmetric +): """ Convert from symmetric and antisymmetric tidal terms to lambda_1 and lambda_2 @@ -1372,13 +1609,18 @@ def lambda_symmetric_lambda_antisymmetric_to_lambda_1_lambda_2(lambda_symmetric, lambda_2: float Tidal parameter of less massive neutron star. """ - lambda_1 = lambda_symmetric_lambda_antisymmetric_to_lambda_1(lambda_symmetric, lambda_antisymmetric) - lambda_2 = lambda_symmetric_lambda_antisymmetric_to_lambda_2(lambda_symmetric, lambda_antisymmetric) + lambda_1 = lambda_symmetric_lambda_antisymmetric_to_lambda_1( + lambda_symmetric, lambda_antisymmetric + ) + lambda_2 = lambda_symmetric_lambda_antisymmetric_to_lambda_2( + lambda_symmetric, lambda_antisymmetric + ) return lambda_1, lambda_2 -def binary_love_fit_lambda_symmetric_mass_ratio_to_lambda_antisymmetric(lambda_symmetric, mass_ratio): - +def binary_love_fit_lambda_symmetric_mass_ratio_to_lambda_antisymmetric( + lambda_symmetric, mass_ratio +): """ Convert from symmetric tidal terms and mass ratio to antisymmetric tidal terms using BinaryLove relations. @@ -1410,13 +1652,13 @@ def binary_love_fit_lambda_symmetric_mass_ratio_to_lambda_antisymmetric(lambda_s lambda_symmetric_m3o5 = lambda_symmetric_m2o5 * lambda_symmetric_m1o5 q = mass_ratio - q2 = mass_ratio ** 2 + q2 = mass_ratio**2 # Eqn.2 from CHZ, incorporating the dependence on mass ratio n_polytropic = 0.743 # average polytropic index for the EoSs included in the fit - q_for_Fnofq = q ** (10. / (3. - n_polytropic)) - Fnofq = (1. - q_for_Fnofq) / (1. + q_for_Fnofq) + q_for_Fnofq = q ** (10.0 / (3.0 - n_polytropic)) + Fnofq = (1.0 - q_for_Fnofq) / (1.0 + q_for_Fnofq) # b_ij and c_ij coefficients are given in Table I of CHZ @@ -1437,23 +1679,34 @@ def binary_love_fit_lambda_symmetric_mass_ratio_to_lambda_antisymmetric(lambda_s # Eqn 1 from CHZ, giving the lambda_antisymmetric_fitOnly # not yet accounting for the uncertainty in the fit - numerator = 1.0 + \ - (b11 * q * lambda_symmetric_m1o5) + (b12 * q2 * lambda_symmetric_m1o5) + \ - (b21 * q * lambda_symmetric_m2o5) + (b22 * q2 * lambda_symmetric_m2o5) + \ - (b31 * q * lambda_symmetric_m3o5) + (b32 * q2 * lambda_symmetric_m3o5) + numerator = ( + 1.0 + + (b11 * q * lambda_symmetric_m1o5) + + (b12 * q2 * lambda_symmetric_m1o5) + + (b21 * q * lambda_symmetric_m2o5) + + (b22 * q2 * lambda_symmetric_m2o5) + + (b31 * q * lambda_symmetric_m3o5) + + (b32 * q2 * lambda_symmetric_m3o5) + ) - denominator = 1.0 + \ - (c11 * q * lambda_symmetric_m1o5) + (c12 * q2 * lambda_symmetric_m1o5) + \ - (c21 * q * lambda_symmetric_m2o5) + (c22 * q2 * lambda_symmetric_m2o5) + \ - (c31 * q * lambda_symmetric_m3o5) + (c32 * q2 * lambda_symmetric_m3o5) + denominator = ( + 1.0 + + (c11 * q * lambda_symmetric_m1o5) + + (c12 * q2 * lambda_symmetric_m1o5) + + (c21 * q * lambda_symmetric_m2o5) + + (c22 * q2 * lambda_symmetric_m2o5) + + (c31 * q * lambda_symmetric_m3o5) + + (c32 * q2 * lambda_symmetric_m3o5) + ) lambda_antisymmetric_fitOnly = Fnofq * lambda_symmetric * numerator / denominator return lambda_antisymmetric_fitOnly -def binary_love_lambda_symmetric_to_lambda_1_lambda_2_manual_marginalisation(binary_love_uniform, - lambda_symmetric, mass_ratio): +def binary_love_lambda_symmetric_to_lambda_1_lambda_2_manual_marginalisation( + binary_love_uniform, lambda_symmetric, mass_ratio +): """ Convert from symmetric tidal terms to lambda_1 and lambda_2 using BinaryLove relations @@ -1482,13 +1735,16 @@ def binary_love_lambda_symmetric_to_lambda_1_lambda_2_manual_marginalisation(bin lambda_2: float Tidal parameter of less massive neutron star. """ - lambda_antisymmetric_fitOnly = binary_love_fit_lambda_symmetric_mass_ratio_to_lambda_antisymmetric(lambda_symmetric, - mass_ratio) + lambda_antisymmetric_fitOnly = ( + binary_love_fit_lambda_symmetric_mass_ratio_to_lambda_antisymmetric( + lambda_symmetric, mass_ratio + ) + ) - lambda_symmetric_sqrt = lambda_symmetric ** 0.5 + lambda_symmetric_sqrt = lambda_symmetric**0.5 q = mass_ratio - q2 = mass_ratio ** 2 + q2 = mass_ratio**2 # mu_i and sigma_i coefficients are given in Table II of CHZ @@ -1497,7 +1753,7 @@ def binary_love_lambda_symmetric_to_lambda_1_lambda_2_manual_marginalisation(bin mu_3 = 0.5168637 mu_4 = -11.2765281 mu_5 = 14.9499544 - mu_6 = - 4.6638851 + mu_6 = -4.6638851 sigma_1 = -0.0000739 sigma_2 = 0.0103778 @@ -1511,46 +1767,49 @@ def binary_love_lambda_symmetric_to_lambda_1_lambda_2_manual_marginalisation(bin # uncertainty in the mean of the lambdaS residual fit, # using coefficients mu_1, mu_2 and mu_3 from Table II of CHZ - lambda_antisymmetric_lambda_symmetric_meanCorr = \ - (mu_1 / (lambda_symmetric * lambda_symmetric)) + \ - (mu_2 / lambda_symmetric) + mu_3 + lambda_antisymmetric_lambda_symmetric_meanCorr = ( + (mu_1 / (lambda_symmetric * lambda_symmetric)) + + (mu_2 / lambda_symmetric) + + mu_3 + ) # Eqn 8 from CHZ, correction on fit for lambdaA caused by # uncertainty in the standard deviation of lambdaS residual fit, # using coefficients sigma_1, sigma_2, sigma_3 and sigma_4 from Table II - lambda_antisymmetric_lambda_symmetric_stdCorr = \ - (sigma_1 * lambda_symmetric * lambda_symmetric_sqrt) + \ - (sigma_2 * lambda_symmetric) + \ - (sigma_3 * lambda_symmetric_sqrt) + sigma_4 + lambda_antisymmetric_lambda_symmetric_stdCorr = ( + (sigma_1 * lambda_symmetric * lambda_symmetric_sqrt) + + (sigma_2 * lambda_symmetric) + + (sigma_3 * lambda_symmetric_sqrt) + + sigma_4 + ) # Eqn 7, correction on fit for lambdaA caused by # uncertainty in the mean of the q residual fit, # using coefficients mu_4, mu_5 and mu_6 from Table II - lambda_antisymmetric_mass_ratio_meanCorr = \ - (mu_4 * q2) + (mu_5 * q) + mu_6 + lambda_antisymmetric_mass_ratio_meanCorr = (mu_4 * q2) + (mu_5 * q) + mu_6 # Eqn 9 from CHZ, correction on fit for lambdaA caused by # uncertainty in the standard deviation of the q residual fit, # using coefficients sigma_5, sigma_6 and sigma_7 from Table II - lambda_antisymmetric_mass_ratio_stdCorr = \ - (sigma_5 * q2) + (sigma_6 * q) + sigma_7 + lambda_antisymmetric_mass_ratio_stdCorr = (sigma_5 * q2) + (sigma_6 * q) + sigma_7 # Eqn 4 from CHZ, averaging the corrections from the # mean of the residual fits - lambda_antisymmetric_meanCorr = \ - (lambda_antisymmetric_lambda_symmetric_meanCorr + - lambda_antisymmetric_mass_ratio_meanCorr) / 2. + lambda_antisymmetric_meanCorr = ( + lambda_antisymmetric_lambda_symmetric_meanCorr + + lambda_antisymmetric_mass_ratio_meanCorr + ) / 2.0 # Eqn 5 from CHZ, averaging the corrections from the # standard deviations of the residual fits lambda_antisymmetric_stdCorr = ( - lambda_antisymmetric_lambda_symmetric_stdCorr ** 2 - + lambda_antisymmetric_mass_ratio_stdCorr ** 2 + lambda_antisymmetric_lambda_symmetric_stdCorr**2 + + lambda_antisymmetric_mass_ratio_stdCorr**2 ) ** 0.5 # Draw a correction on the fit from a @@ -1558,17 +1817,23 @@ def binary_love_lambda_symmetric_to_lambda_1_lambda_2_manual_marginalisation(bin # this is done by sampling a percent point function (inverse cdf) # through a U{0,1} variable called binary_love_uniform - lambda_antisymmetric_scatter = norm.ppf(binary_love_uniform, loc=0., - scale=lambda_antisymmetric_stdCorr) + lambda_antisymmetric_scatter = norm.ppf( + binary_love_uniform, loc=0.0, scale=lambda_antisymmetric_stdCorr + ) # Add the correction of the residual mean # and the Gaussian scatter to the lambda_antisymmetric_fitOnly value - lambda_antisymmetric = lambda_antisymmetric_fitOnly + \ - (lambda_antisymmetric_meanCorr + lambda_antisymmetric_scatter) + lambda_antisymmetric = lambda_antisymmetric_fitOnly + ( + lambda_antisymmetric_meanCorr + lambda_antisymmetric_scatter + ) - lambda_1 = lambda_symmetric_lambda_antisymmetric_to_lambda_1(lambda_symmetric, lambda_antisymmetric) - lambda_2 = lambda_symmetric_lambda_antisymmetric_to_lambda_2(lambda_symmetric, lambda_antisymmetric) + lambda_1 = lambda_symmetric_lambda_antisymmetric_to_lambda_1( + lambda_symmetric, lambda_antisymmetric + ) + lambda_2 = lambda_symmetric_lambda_antisymmetric_to_lambda_2( + lambda_symmetric, lambda_antisymmetric + ) # The BinaryLove model is only physically valid where # lambda_2 > lambda_1 as it assumes mass_1 > mass_2 @@ -1591,7 +1856,9 @@ def binary_love_lambda_symmetric_to_lambda_1_lambda_2_manual_marginalisation(bin return lambda_1, lambda_2 -def binary_love_lambda_symmetric_to_lambda_1_lambda_2_automatic_marginalisation(lambda_symmetric, mass_ratio): +def binary_love_lambda_symmetric_to_lambda_1_lambda_2_automatic_marginalisation( + lambda_symmetric, mass_ratio +): """ Convert from symmetric tidal terms to lambda_1 and lambda_2 using BinaryLove relations @@ -1626,26 +1893,29 @@ def binary_love_lambda_symmetric_to_lambda_1_lambda_2_automatic_marginalisation( binary_love_uniform = random.rng.uniform(0, 1, len(lambda_symmetric)) - lambda_1, lambda_2 = binary_love_lambda_symmetric_to_lambda_1_lambda_2_manual_marginalisation( - binary_love_uniform, lambda_symmetric, mass_ratio) + lambda_1, lambda_2 = ( + binary_love_lambda_symmetric_to_lambda_1_lambda_2_manual_marginalisation( + binary_love_uniform, lambda_symmetric, mass_ratio + ) + ) return lambda_1, lambda_2 -def _generate_all_cbc_parameters(sample, defaults, base_conversion, - likelihood=None, priors=None, npool=1, pool=None): +def _generate_all_cbc_parameters( + sample, defaults, base_conversion, likelihood=None, priors=None, npool=1, pool=None +): """Generate all cbc parameters, helper function for BBH/BNS""" output_sample = sample.copy() waveform_defaults = defaults for key in waveform_defaults: try: - output_sample[key] = \ - likelihood.waveform_generator.waveform_arguments[key] + output_sample[key] = likelihood.waveform_generator.waveform_arguments[key] except (KeyError, AttributeError): default = waveform_defaults[key] output_sample[key] = default - logger.debug('Assuming {} = {}'.format(key, default)) + logger.debug("Assuming {} = {}".format(key, default)) output_sample = fill_from_fixed_priors(output_sample, priors) output_sample, _ = base_conversion(output_sample) @@ -1663,13 +1933,17 @@ def _generate_all_cbc_parameters(sample, defaults, base_conversion, output_sample["time_jitter"] = 0.0 compute_per_detector_log_likelihoods( - samples=output_sample, likelihood=likelihood, npool=npool, pool=pool) + samples=output_sample, likelihood=likelihood, npool=npool, pool=pool + ) - marginalized_parameters = getattr(likelihood, "_marginalized_parameters", list()) + marginalized_parameters = getattr( + likelihood, "_marginalized_parameters", list() + ) if len(marginalized_parameters) > 0: try: generate_posterior_samples_from_marginalized_likelihood( - samples=output_sample, likelihood=likelihood, npool=npool, pool=pool) + samples=output_sample, likelihood=likelihood, npool=npool, pool=pool + ) except MarginalizedLikelihoodReconstructionError as e: logger.warning( "Marginalised parameter reconstruction failed with message " @@ -1685,42 +1959,49 @@ def _generate_all_cbc_parameters(sample, defaults, base_conversion, for par in marginalized_parameters: name = misnamed_marginalizations.get(par, par) if ( - getattr(likelihood, f'{name}_marginalization', False) + getattr(likelihood, f"{name}_marginalization", False) and par in likelihood.priors ): priors[par] = likelihood.priors[par] - if ( - not getattr(likelihood, "reference_frame", "sky") == "sky" - or "geocent" not in getattr(likelihood, "time_reference", "geocent") - ): + if not getattr( + likelihood, "reference_frame", "sky" + ) == "sky" or "geocent" not in getattr(likelihood, "time_reference", "geocent"): try: generate_sky_frame_parameters( samples=output_sample, likelihood=likelihood ) except TypeError: logger.info( - "Failed to generate sky frame parameters for type {}" - .format(type(output_sample)) + "Failed to generate sky frame parameters for type {}".format( + type(output_sample) + ) ) compute_snrs(output_sample, likelihood=likelihood, npool=npool, pool=pool) # Remove the time jitter if it was added if added_time_jitter: output_sample.pop("time_jitter") - for key, func in zip(["mass", "spin", "source frame"], [ - generate_mass_parameters, generate_spin_parameters, - generate_source_frame_parameters]): + for key, func in zip( + ["mass", "spin", "source frame"], + [ + generate_mass_parameters, + generate_spin_parameters, + generate_source_frame_parameters, + ], + ): try: output_sample = func(output_sample) except KeyError as e: logger.info( - "Generation of {} parameters failed with message {}".format( - key, e)) + "Generation of {} parameters failed with message {}".format(key, e) + ) return output_sample -def generate_all_bbh_parameters(sample, likelihood=None, priors=None, npool=1, pool=None): +def generate_all_bbh_parameters( + sample, likelihood=None, priors=None, npool=1, pool=None +): """ From either a single sample or a set of samples fill in all missing BBH parameters, in place. @@ -1737,16 +2018,25 @@ def generate_all_bbh_parameters(sample, likelihood=None, priors=None, npool=1, p Dictionary of prior objects, used to fill in non-sampled parameters. """ waveform_defaults = { - 'reference_frequency': 50.0, 'waveform_approximant': 'IMRPhenomPv2', - 'minimum_frequency': 20.0} + "reference_frequency": 50.0, + "waveform_approximant": "IMRPhenomPv2", + "minimum_frequency": 20.0, + } output_sample = _generate_all_cbc_parameters( - sample, defaults=waveform_defaults, + sample, + defaults=waveform_defaults, base_conversion=convert_to_lal_binary_black_hole_parameters, - likelihood=likelihood, priors=priors, npool=npool, pool=pool) + likelihood=likelihood, + priors=priors, + npool=npool, + pool=pool, + ) return output_sample -def generate_all_bns_parameters(sample, likelihood=None, priors=None, npool=1, pool=None): +def generate_all_bns_parameters( + sample, likelihood=None, priors=None, npool=1, pool=None +): """ From either a single sample or a set of samples fill in all missing BNS parameters, in place. @@ -1768,17 +2058,23 @@ def generate_all_bns_parameters(sample, likelihood=None, priors=None, npool=1, p If given, perform generation (where possible) using a multiprocessing pool """ waveform_defaults = { - 'reference_frequency': 50.0, 'waveform_approximant': 'TaylorF2', - 'minimum_frequency': 20.0} + "reference_frequency": 50.0, + "waveform_approximant": "TaylorF2", + "minimum_frequency": 20.0, + } output_sample = _generate_all_cbc_parameters( - sample, defaults=waveform_defaults, + sample, + defaults=waveform_defaults, base_conversion=convert_to_lal_binary_neutron_star_parameters, - likelihood=likelihood, priors=priors, npool=npool, pool=pool) + likelihood=likelihood, + priors=priors, + npool=npool, + pool=pool, + ) try: output_sample = generate_tidal_parameters(output_sample) except KeyError as e: - logger.debug( - "Generation of tidal parameters failed with message {}".format(e)) + logger.debug("Generation of tidal parameters failed with message {}".format(e)) return output_sample @@ -1839,7 +2135,7 @@ def fill_from_fixed_priors(sample, priors): def generate_component_masses(sample, require_add=False, source=False): - """" + """ " Add the component masses to the dataframe/dictionary We add: mass_1, mass_2 @@ -1868,11 +2164,13 @@ def generate_component_masses(sample, require_add=False, source=False): Returns dict : the updated dictionary """ + def check_and_return_quietly(require_add, sample): if require_add: raise KeyError("Insufficient mass parameters in input dictionary") else: return sample + output_sample = sample.copy() if source: @@ -1890,23 +2188,21 @@ def check_and_return_quietly(require_add, sample): if mass_2_key in sample.keys(): return output_sample if total_mass_key in sample.keys(): - output_sample[mass_2_key] = output_sample[total_mass_key] - ( - output_sample[mass_1_key] + output_sample[mass_2_key] = ( + output_sample[total_mass_key] - (output_sample[mass_1_key]) ) return output_sample elif "mass_ratio" in sample.keys(): pass elif "symmetric_mass_ratio" in sample.keys(): - output_sample["mass_ratio"] = ( - symmetric_mass_ratio_to_mass_ratio( - output_sample["symmetric_mass_ratio"]) + output_sample["mass_ratio"] = symmetric_mass_ratio_to_mass_ratio( + output_sample["symmetric_mass_ratio"] ) elif chirp_mass_key in sample.keys(): - output_sample["mass_ratio"] = ( - mass_1_and_chirp_mass_to_mass_ratio( - mass_1=output_sample[mass_1_key], - chirp_mass=output_sample[chirp_mass_key]) + output_sample["mass_ratio"] = mass_1_and_chirp_mass_to_mass_ratio( + mass_1=output_sample[mass_1_key], + chirp_mass=output_sample[chirp_mass_key], ) else: return check_and_return_quietly(require_add, sample) @@ -1927,21 +2223,19 @@ def check_and_return_quietly(require_add, sample): elif "mass_ratio" in sample.keys(): pass elif "symmetric_mass_ratio" in sample.keys(): - output_sample["mass_ratio"] = ( - symmetric_mass_ratio_to_mass_ratio( - output_sample["symmetric_mass_ratio"]) + output_sample["mass_ratio"] = symmetric_mass_ratio_to_mass_ratio( + output_sample["symmetric_mass_ratio"] ) elif chirp_mass_key in sample.keys(): - output_sample["mass_ratio"] = ( - mass_2_and_chirp_mass_to_mass_ratio( - mass_2=output_sample[mass_2_key], - chirp_mass=output_sample[chirp_mass_key]) + output_sample["mass_ratio"] = mass_2_and_chirp_mass_to_mass_ratio( + mass_2=output_sample[mass_2_key], + chirp_mass=output_sample[chirp_mass_key], ) else: check_and_return_quietly(require_add, sample) - output_sample[mass_1_key] = 1 / output_sample["mass_ratio"] * ( - output_sample[mass_2_key] + output_sample[mass_1_key] = ( + 1 / output_sample["mass_ratio"] * (output_sample[mass_2_key]) ) return output_sample @@ -1951,19 +2245,18 @@ def check_and_return_quietly(require_add, sample): if "mass_ratio" in sample.keys(): pass # We have everything we need already elif "symmetric_mass_ratio" in sample.keys(): - output_sample["mass_ratio"] = ( - symmetric_mass_ratio_to_mass_ratio( - output_sample["symmetric_mass_ratio"]) + output_sample["mass_ratio"] = symmetric_mass_ratio_to_mass_ratio( + output_sample["symmetric_mass_ratio"] ) elif chirp_mass_key in sample.keys(): output_sample["symmetric_mass_ratio"] = ( chirp_mass_and_total_mass_to_symmetric_mass_ratio( chirp_mass=output_sample[chirp_mass_key], - total_mass=output_sample[total_mass_key]) + total_mass=output_sample[total_mass_key], + ) ) - output_sample["mass_ratio"] = ( - symmetric_mass_ratio_to_mass_ratio( - output_sample["symmetric_mass_ratio"]) + output_sample["mass_ratio"] = symmetric_mass_ratio_to_mass_ratio( + output_sample["symmetric_mass_ratio"] ) else: return check_and_return_quietly(require_add, sample) @@ -1972,27 +2265,24 @@ def check_and_return_quietly(require_add, sample): if "mass_ratio" in sample.keys(): pass elif "symmetric_mass_ratio" in sample.keys(): - output_sample["mass_ratio"] = ( - symmetric_mass_ratio_to_mass_ratio( - sample["symmetric_mass_ratio"]) + output_sample["mass_ratio"] = symmetric_mass_ratio_to_mass_ratio( + sample["symmetric_mass_ratio"] ) else: return check_and_return_quietly(require_add, sample) - output_sample[total_mass_key] = ( - chirp_mass_and_mass_ratio_to_total_mass( - chirp_mass=output_sample[chirp_mass_key], - mass_ratio=output_sample["mass_ratio"]) + output_sample[total_mass_key] = chirp_mass_and_mass_ratio_to_total_mass( + chirp_mass=output_sample[chirp_mass_key], + mass_ratio=output_sample["mass_ratio"], ) # We haven't matched any of the criteria if total_mass_key not in output_sample.keys() or ( - "mass_ratio" not in output_sample.keys()): + "mass_ratio" not in output_sample.keys() + ): return check_and_return_quietly(require_add, sample) - mass_1, mass_2 = ( - total_mass_and_mass_ratio_to_component_masses( - total_mass=output_sample[total_mass_key], - mass_ratio=output_sample["mass_ratio"]) + mass_1, mass_2 = total_mass_and_mass_ratio_to_component_masses( + total_mass=output_sample[total_mass_key], mass_ratio=output_sample["mass_ratio"] ) output_sample[mass_1_key] = mass_1 output_sample[mass_2_key] = mass_2 @@ -2030,35 +2320,33 @@ def generate_mass_parameters(sample, source=False): output_sample = intermediate_sample.copy() if source: - mass_1_key = 'mass_1_source' - mass_2_key = 'mass_2_source' - total_mass_key = 'total_mass_source' - chirp_mass_key = 'chirp_mass_source' + mass_1_key = "mass_1_source" + mass_2_key = "mass_2_source" + total_mass_key = "total_mass_source" + chirp_mass_key = "chirp_mass_source" else: - mass_1_key = 'mass_1' - mass_2_key = 'mass_2' - total_mass_key = 'total_mass' - chirp_mass_key = 'chirp_mass' + mass_1_key = "mass_1" + mass_2_key = "mass_2" + total_mass_key = "total_mass" + chirp_mass_key = "chirp_mass" if chirp_mass_key not in output_sample.keys(): - output_sample[chirp_mass_key] = ( - component_masses_to_chirp_mass(output_sample[mass_1_key], - output_sample[mass_2_key]) + output_sample[chirp_mass_key] = component_masses_to_chirp_mass( + output_sample[mass_1_key], output_sample[mass_2_key] ) if total_mass_key not in output_sample.keys(): - output_sample[total_mass_key] = ( - component_masses_to_total_mass(output_sample[mass_1_key], - output_sample[mass_2_key]) + output_sample[total_mass_key] = component_masses_to_total_mass( + output_sample[mass_1_key], output_sample[mass_2_key] ) - if 'symmetric_mass_ratio' not in output_sample.keys(): - output_sample['symmetric_mass_ratio'] = ( - component_masses_to_symmetric_mass_ratio(output_sample[mass_1_key], - output_sample[mass_2_key]) + if "symmetric_mass_ratio" not in output_sample.keys(): + output_sample["symmetric_mass_ratio"] = ( + component_masses_to_symmetric_mass_ratio( + output_sample[mass_1_key], output_sample[mass_2_key] + ) ) - if 'mass_ratio' not in output_sample.keys(): - output_sample['mass_ratio'] = ( - component_masses_to_mass_ratio(output_sample[mass_1_key], - output_sample[mass_2_key]) + if "mass_ratio" not in output_sample.keys(): + output_sample["mass_ratio"] = component_masses_to_mass_ratio( + output_sample[mass_1_key], output_sample[mass_2_key] ) return output_sample @@ -2086,27 +2374,29 @@ def generate_spin_parameters(sample): output_sample = generate_component_spins(output_sample) xp = array_module(sample) - output_sample['chi_eff'] = (output_sample['spin_1z'] + - output_sample['spin_2z'] * - output_sample['mass_ratio']) /\ - (1 + output_sample['mass_ratio']) + output_sample["chi_eff"] = ( + output_sample["spin_1z"] + + output_sample["spin_2z"] * output_sample["mass_ratio"] + ) / (1 + output_sample["mass_ratio"]) - output_sample['chi_1_in_plane'] = ( - output_sample['spin_1x'] ** 2 + output_sample['spin_1y'] ** 2 + output_sample["chi_1_in_plane"] = ( + output_sample["spin_1x"] ** 2 + output_sample["spin_1y"] ** 2 ) ** 0.5 - output_sample['chi_2_in_plane'] = ( - output_sample['spin_2x'] ** 2 + output_sample['spin_2y'] ** 2 + output_sample["chi_2_in_plane"] = ( + output_sample["spin_2x"] ** 2 + output_sample["spin_2y"] ** 2 ) ** 0.5 - output_sample['chi_p'] = xp.maximum( - output_sample['chi_1_in_plane'], - (4 * output_sample['mass_ratio'] + 3) / - (3 * output_sample['mass_ratio'] + 4) * output_sample['mass_ratio'] * - output_sample['chi_2_in_plane']) + output_sample["chi_p"] = xp.maximum( + output_sample["chi_1_in_plane"], + (4 * output_sample["mass_ratio"] + 3) + / (3 * output_sample["mass_ratio"] + 4) + * output_sample["mass_ratio"] + * output_sample["chi_2_in_plane"], + ) try: - output_sample['cos_tilt_1'] = xp.cos(output_sample['tilt_1']) - output_sample['cos_tilt_2'] = xp.cos(output_sample['tilt_2']) + output_sample["cos_tilt_1"] = xp.cos(output_sample["tilt_1"]) + output_sample["cos_tilt_2"] = xp.cos(output_sample["tilt_2"]) except KeyError: pass @@ -2131,39 +2421,59 @@ def generate_component_spins(sample): """ output_sample = sample.copy() - spin_conversion_parameters =\ - ['theta_jn', 'phi_jl', 'tilt_1', 'tilt_2', 'phi_12', 'a_1', 'a_2', - 'mass_1', 'mass_2', 'reference_frequency', 'phase'] + spin_conversion_parameters = [ + "theta_jn", + "phi_jl", + "tilt_1", + "tilt_2", + "phi_12", + "a_1", + "a_2", + "mass_1", + "mass_2", + "reference_frequency", + "phase", + ] if all(key in output_sample.keys() for key in spin_conversion_parameters): xp = array_module(output_sample["theta_jn"]) ( - output_sample['iota'], output_sample['spin_1x'], - output_sample['spin_1y'], output_sample['spin_1z'], - output_sample['spin_2x'], output_sample['spin_2y'], - output_sample['spin_2z'] + output_sample["iota"], + output_sample["spin_1x"], + output_sample["spin_1y"], + output_sample["spin_1z"], + output_sample["spin_2x"], + output_sample["spin_2y"], + output_sample["spin_2z"], ) = xp.vectorize(bilby_to_lalsimulation_spins)( - output_sample['theta_jn'], output_sample['phi_jl'], - output_sample['tilt_1'], output_sample['tilt_2'], - output_sample['phi_12'], output_sample['a_1'], output_sample['a_2'], - output_sample['mass_1'] * solar_mass, - output_sample['mass_2'] * solar_mass, - output_sample['reference_frequency'], output_sample['phase'] + output_sample["theta_jn"], + output_sample["phi_jl"], + output_sample["tilt_1"], + output_sample["tilt_2"], + output_sample["phi_12"], + output_sample["a_1"], + output_sample["a_2"], + output_sample["mass_1"] * solar_mass, + output_sample["mass_2"] * solar_mass, + output_sample["reference_frequency"], + output_sample["phase"], + ) + + output_sample["phi_1"] = xp.fmod( + 2 * np.pi + xp.arctan2(output_sample["spin_1y"], output_sample["spin_1x"]), + 2 * np.pi, + ) + output_sample["phi_2"] = xp.fmod( + 2 * np.pi + xp.arctan2(output_sample["spin_2y"], output_sample["spin_2x"]), + 2 * np.pi, ) - output_sample['phi_1'] =\ - xp.fmod(2 * np.pi + xp.arctan2( - output_sample['spin_1y'], output_sample['spin_1x']), 2 * np.pi) - output_sample['phi_2'] =\ - xp.fmod(2 * np.pi + xp.arctan2( - output_sample['spin_2y'], output_sample['spin_2x']), 2 * np.pi) - - elif 'chi_1' in output_sample and 'chi_2' in output_sample: - output_sample['spin_1x'] = 0 - output_sample['spin_1y'] = 0 - output_sample['spin_1z'] = output_sample['chi_1'] - output_sample['spin_2x'] = 0 - output_sample['spin_2y'] = 0 - output_sample['spin_2z'] = output_sample['chi_2'] + elif "chi_1" in output_sample and "chi_2" in output_sample: + output_sample["spin_1x"] = 0 + output_sample["spin_1y"] = 0 + output_sample["spin_1z"] = output_sample["chi_1"] + output_sample["spin_2x"] = 0 + output_sample["spin_2y"] = 0 + output_sample["spin_2z"] = output_sample["chi_2"] else: logger.debug("Component spin extraction failed.") @@ -2188,14 +2498,18 @@ def generate_tidal_parameters(sample): """ output_sample = sample.copy() - output_sample['lambda_tilde'] =\ - lambda_1_lambda_2_to_lambda_tilde( - output_sample['lambda_1'], output_sample['lambda_2'], - output_sample['mass_1'], output_sample['mass_2']) - output_sample['delta_lambda_tilde'] = \ - lambda_1_lambda_2_to_delta_lambda_tilde( - output_sample['lambda_1'], output_sample['lambda_2'], - output_sample['mass_1'], output_sample['mass_2']) + output_sample["lambda_tilde"] = lambda_1_lambda_2_to_lambda_tilde( + output_sample["lambda_1"], + output_sample["lambda_2"], + output_sample["mass_1"], + output_sample["mass_2"], + ) + output_sample["delta_lambda_tilde"] = lambda_1_lambda_2_to_delta_lambda_tilde( + output_sample["lambda_1"], + output_sample["lambda_2"], + output_sample["mass_1"], + output_sample["mass_2"], + ) return output_sample @@ -2214,15 +2528,18 @@ def generate_source_frame_parameters(sample): """ output_sample = sample.copy() - output_sample['redshift'] =\ - luminosity_distance_to_redshift(output_sample['luminosity_distance']) - output_sample['comoving_distance'] =\ - redshift_to_comoving_distance(output_sample['redshift']) + output_sample["redshift"] = luminosity_distance_to_redshift( + output_sample["luminosity_distance"] + ) + output_sample["comoving_distance"] = redshift_to_comoving_distance( + output_sample["redshift"] + ) - for key in ['mass_1', 'mass_2', 'chirp_mass', 'total_mass']: + for key in ["mass_1", "mass_2", "chirp_mass", "total_mass"]: if key in output_sample: - output_sample['{}_source'.format(key)] =\ - output_sample[key] / (1 + output_sample['redshift']) + output_sample["{}_source".format(key)] = output_sample[key] / ( + 1 + output_sample["redshift"] + ) return output_sample @@ -2242,24 +2559,37 @@ def compute_snrs(sample, likelihood, npool=1, pool=None): """ if likelihood is not None: if isinstance(sample, dict): - signal_polarizations = likelihood.waveform_generator.frequency_domain_strain(sample.copy()) + signal_polarizations = ( + likelihood.waveform_generator.frequency_domain_strain(sample.copy()) + ) for ifo in likelihood.interferometers: - per_detector_snr = likelihood.calculate_snrs(signal_polarizations, ifo, parameters=sample) - sample['{}_matched_filter_snr'.format(ifo.name)] =\ + per_detector_snr = likelihood.calculate_snrs( + signal_polarizations, ifo, parameters=sample + ) + sample["{}_matched_filter_snr".format(ifo.name)] = ( per_detector_snr.complex_matched_filter_snr - sample['{}_optimal_snr'.format(ifo.name)] = \ - per_detector_snr.optimal_snr_squared.real ** 0.5 + ) + sample["{}_optimal_snr".format(ifo.name)] = ( + per_detector_snr.optimal_snr_squared.real**0.5 + ) else: from tqdm.auto import tqdm - logger.info('Computing SNRs for every sample.') + + logger.info("Computing SNRs for every sample.") fill_args = [(ii, row) for ii, row in sample.iterrows()] - with bilby_pool(likelihood=likelihood, priors=None, npool=npool, pool=pool) as _pool: + with bilby_pool( + likelihood=likelihood, priors=None, npool=npool, pool=pool + ) as _pool: if _pool is not None: - new_samples = _pool.map(_compute_snrs, tqdm(fill_args, file=sys.stdout)) + new_samples = _pool.map( + _compute_snrs, tqdm(fill_args, file=sys.stdout) + ) else: sampling_convenience_dump.likelihood = likelihood - new_samples = [_compute_snrs(xx) for xx in tqdm(fill_args, file=sys.stdout)] + new_samples = [ + _compute_snrs(xx) for xx in tqdm(fill_args, file=sys.stdout) + ] for ii, ifo in enumerate(likelihood.interferometers): snr_updates = dict() @@ -2272,7 +2602,7 @@ def compute_snrs(sample, likelihood, npool=1, pool=None): for k, v in snr_updates.items(): sample[k] = v else: - logger.debug('Not computing SNRs.') + logger.debug("Not computing SNRs.") def _compute_snrs(args): @@ -2285,13 +2615,17 @@ def _compute_snrs(args): ) snrs = list() for ifo in likelihood.interferometers: - snrs.append(likelihood.calculate_snrs( - signal_polarizations, ifo, return_array=False, parameters=sample - )) + snrs.append( + likelihood.calculate_snrs( + signal_polarizations, ifo, return_array=False, parameters=sample + ) + ) return snrs -def compute_per_detector_log_likelihoods(samples, likelihood, npool=1, block=10, pool=None): +def compute_per_detector_log_likelihoods( + samples, likelihood, npool=1, block=10, pool=None +): """ Calculate the log likelihoods in each detector. @@ -2313,7 +2647,7 @@ def compute_per_detector_log_likelihoods(samples, likelihood, npool=1, block=10, """ if likelihood is not None: if not callable(likelihood.compute_per_detector_log_likelihood): - logger.debug('Not computing per-detector log likelihoods.') + logger.debug("Not computing per-detector log likelihoods.") return samples if isinstance(samples, dict): @@ -2321,10 +2655,12 @@ def compute_per_detector_log_likelihoods(samples, likelihood, npool=1, block=10, return samples elif not isinstance(samples, DataFrame): - raise ValueError("Unable to handle input samples of type {}".format(type(samples))) + raise ValueError( + "Unable to handle input samples of type {}".format(type(samples)) + ) from tqdm.auto import tqdm - logger.info('Computing per-detector log likelihoods.') + logger.info("Computing per-detector log likelihoods.") # Initialize cache dict cached_samples_dict = dict() @@ -2335,7 +2671,9 @@ def compute_per_detector_log_likelihoods(samples, likelihood, npool=1, block=10, fill_args = [(ii, row) for ii, row in samples.iterrows()] ii = 0 pbar = tqdm(total=len(samples), file=sys.stdout) - with bilby_pool(likelihood=likelihood, priors=None, npool=npool, pool=pool) as _pool: + with bilby_pool( + likelihood=likelihood, priors=None, npool=npool, pool=pool + ) as _pool: while ii < len(samples): if ii in cached_samples_dict: ii += block @@ -2348,10 +2686,12 @@ def compute_per_detector_log_likelihoods(samples, likelihood, npool=1, block=10, map_fn = map sampling_convenience_dump.likelihood = likelihood - subset_samples = list(map_fn( - _compute_per_detector_log_likelihoods, - fill_args[ii: ii + block], - )) + subset_samples = list( + map_fn( + _compute_per_detector_log_likelihoods, + fill_args[ii : ii + block], + ) + ) cached_samples_dict[ii] = subset_samples @@ -2360,17 +2700,22 @@ def compute_per_detector_log_likelihoods(samples, likelihood, npool=1, block=10, pbar.close() new_samples = np.concatenate( - [np.array(val) for key, val in cached_samples_dict.items() if key != "_samples"] + [ + np.array(val) + for key, val in cached_samples_dict.items() + if key != "_samples" + ] ) - for ii, key in \ - enumerate([f'{ifo.name}_log_likelihood' for ifo in likelihood.interferometers]): + for ii, key in enumerate( + [f"{ifo.name}_log_likelihood" for ifo in likelihood.interferometers] + ): samples[key] = new_samples[:, ii] return samples else: - logger.debug('Not computing per-detector log likelihoods.') + logger.debug("Not computing per-detector log likelihoods.") def _compute_per_detector_log_likelihoods(args): @@ -2379,12 +2724,19 @@ def _compute_per_detector_log_likelihoods(args): _, sample = args sample = dict(sample).copy() new_sample = likelihood.compute_per_detector_log_likelihood(sample) - return tuple((new_sample[key] for key in - [f'{ifo.name}_log_likelihood' for ifo in likelihood.interferometers])) + return tuple( + ( + new_sample[key] + for key in [ + f"{ifo.name}_log_likelihood" for ifo in likelihood.interferometers + ] + ) + ) def generate_posterior_samples_from_marginalized_likelihood( - samples, likelihood, npool=1, block=10, use_cache=True, pool=None): + samples, likelihood, npool=1, block=10, use_cache=True, pool=None +): """ Reconstruct the distance posterior from a run which used a likelihood which explicitly marginalised over time/distance/phase. @@ -2420,13 +2772,17 @@ def generate_posterior_samples_from_marginalized_likelihood( if isinstance(samples, dict): return samples elif not isinstance(samples, DataFrame): - raise ValueError("Unable to handle input samples of type {}".format(type(samples))) + raise ValueError( + "Unable to handle input samples of type {}".format(type(samples)) + ) from tqdm.auto import tqdm - logger.info('Reconstructing marginalised parameters.') + logger.info("Reconstructing marginalised parameters.") try: - cache_filename = f"{likelihood.outdir}/.{likelihood.label}_generate_posterior_cache.pickle" + cache_filename = ( + f"{likelihood.outdir}/.{likelihood.label}_generate_posterior_cache.pickle" + ) except AttributeError: logger.warning("Likelihood has no outdir and label attribute: caching disabled") use_cache = False @@ -2440,13 +2796,19 @@ def generate_posterior_samples_from_marginalized_likelihood( cached_samples_dict = None # Check the samples are identical between the cache and current - if (cached_samples_dict is not None) and (cached_samples_dict["_samples"].equals(samples)): + if (cached_samples_dict is not None) and ( + cached_samples_dict["_samples"].equals(samples) + ): # Calculate reconstruction percentage and print a log message nsamples_converted = np.sum( - [len(val) for key, val in cached_samples_dict.items() if key != "_samples"] + [ + len(val) + for key, val in cached_samples_dict.items() + if key != "_samples" + ] ) perc = 100 * nsamples_converted / len(cached_samples_dict["_samples"]) - logger.info(f'Using cached reconstruction with {perc:0.1f}% converted.') + logger.info(f"Using cached reconstruction with {perc:0.1f}% converted.") else: logger.info("Cached samples dict out of date, ignoring") cached_samples_dict = dict(_samples=samples) @@ -2462,7 +2824,9 @@ def generate_posterior_samples_from_marginalized_likelihood( fill_args = [(ii, row, seed) for (ii, row), seed in zip(samples.iterrows(), seeds)] ii = 0 pbar = tqdm(total=len(samples), file=sys.stdout) - with bilby_pool(likelihood=likelihood, priors=None, npool=npool, pool=pool) as _pool: + with bilby_pool( + likelihood=likelihood, priors=None, npool=npool, pool=pool + ) as _pool: while ii < len(samples): if ii in cached_samples_dict: ii += block @@ -2470,9 +2834,9 @@ def generate_posterior_samples_from_marginalized_likelihood( continue if _pool is not None: - subset_samples = _pool.map(fill_sample, fill_args[ii: ii + block]) + subset_samples = _pool.map(fill_sample, fill_args[ii : ii + block]) else: - subset_samples = list(map(fill_sample, fill_args[ii: ii + block])) + subset_samples = list(map(fill_sample, fill_args[ii : ii + block])) cached_samples_dict[ii] = subset_samples @@ -2501,7 +2865,7 @@ def generate_sky_frame_parameters(samples, likelihood): raise ValueError from tqdm.auto import tqdm - logger.info('Generating sky frame parameters.') + logger.info("Generating sky frame parameters.") new_samples = list() for ii in tqdm(range(len(samples)), file=sys.stdout): sample = dict(samples.iloc[ii]).copy() @@ -2519,7 +2883,9 @@ def fill_sample(args): likelihood = sampling_convenience_dump.likelihood marginalized_parameters = getattr(likelihood, "_marginalized_parameters", list()) sample = dict(sample).copy() - new_sample = likelihood.generate_posterior_sample_from_marginalized_likelihood(sample) + new_sample = likelihood.generate_posterior_sample_from_marginalized_likelihood( + sample + ) return tuple((new_sample[key] for key in marginalized_parameters)) @@ -2555,13 +2921,17 @@ def identity_map_generation(sample, likelihood=None, priors=None, npool=1, pool= if likelihood is not None: compute_per_detector_log_likelihoods( - samples=output_sample, likelihood=likelihood, npool=npool, pool=pool) + samples=output_sample, likelihood=likelihood, npool=npool, pool=pool + ) - marginalized_parameters = getattr(likelihood, "_marginalized_parameters", list()) + marginalized_parameters = getattr( + likelihood, "_marginalized_parameters", list() + ) if len(marginalized_parameters) > 0: try: generate_posterior_samples_from_marginalized_likelihood( - samples=output_sample, likelihood=likelihood, npool=npool, pool=pool) + samples=output_sample, likelihood=likelihood, npool=npool, pool=pool + ) except MarginalizedLikelihoodReconstructionError as e: logger.warning( "Marginalised parameter reconstruction failed with message " @@ -2569,7 +2939,11 @@ def identity_map_generation(sample, likelihood=None, priors=None, npool=1, pool= "interpretation.".format(e) ) - if ("ra" in output_sample.keys() and "dec" in output_sample.keys() and "psi" in output_sample.keys()): + if ( + "ra" in output_sample.keys() + and "dec" in output_sample.keys() + and "psi" in output_sample.keys() + ): compute_snrs(output_sample, likelihood, npool=npool, pool=pool) else: logger.info( diff --git a/bilby/gw/cosmology.py b/bilby/gw/cosmology.py index 6606fdd85..df547e2e7 100644 --- a/bilby/gw/cosmology.py +++ b/bilby/gw/cosmology.py @@ -9,6 +9,7 @@ def _set_default_cosmology(): from astropy import cosmology as cosmo from ..core.utils.meta_data import global_meta_data + global DEFAULT_COSMOLOGY, COSMOLOGY if DEFAULT_COSMOLOGY is None: DEFAULT_COSMOLOGY = cosmo.Planck15 @@ -27,6 +28,7 @@ def get_available_cosmologies(): A tuple of strings with the names of the available cosmologies. """ from astropy.cosmology.realizations import available + return (*available, "Planck15_LAL") @@ -51,6 +53,7 @@ def get_cosmology(cosmology=None): Cosmology instance """ from astropy import cosmology as cosmo + _set_default_cosmology() if cosmology is None: cosmology = DEFAULT_COSMOLOGY @@ -79,8 +82,8 @@ def get_cosmology(cosmology=None): else: cosmology = getattr(cosmo, cosmology) elif isinstance(cosmology, dict): - if 'Ode0' in cosmology.keys(): - if 'w0' in cosmology.keys(): + if "Ode0" in cosmology.keys(): + if "w0" in cosmology.keys(): cosmology = cosmo.wCDM(**cosmology) else: cosmology = cosmo.LambdaCDM(**cosmology) @@ -105,6 +108,7 @@ def set_cosmology(cosmology=None): class. """ from ..core.utils.meta_data import global_meta_data + cosmology = get_cosmology(cosmology) global COSMOLOGY, DEFAULT_COSMOLOGY DEFAULT_COSMOLOGY = cosmology @@ -125,4 +129,5 @@ def z_at_value(func, fval, **kwargs): for detailed documentation. """ from astropy.cosmology import z_at_value + return z_at_value(func=func, fval=fval, **kwargs).value diff --git a/bilby/gw/detector/__init__.py b/bilby/gw/detector/__init__.py index 51ff98279..90e1492be 100644 --- a/bilby/gw/detector/__init__.py +++ b/bilby/gw/detector/__init__.py @@ -5,8 +5,9 @@ from .psd import * from .strain_data import * + def get_safe_signal_duration(mass_1, mass_2, a_1, a_2, tilt_1, tilt_2, flow=10): - """ Calculate the safe signal duration, given the parameters + """Calculate the safe signal duration, given the parameters Parameters ========== @@ -24,15 +25,27 @@ def get_safe_signal_duration(mass_1, mass_2, a_1, a_2, tilt_1, tilt_2, flow=10): """ from lal import MSUN_SI from lalsimulation import SimInspiralChirpTimeBound + chirp_time = SimInspiralChirpTimeBound( - flow, mass_1 * MSUN_SI, mass_2 * MSUN_SI, - a_1 * np.cos(tilt_1), a_2 * np.cos(tilt_2)) - return max(2**(int(np.log2(chirp_time)) + 1), 4) + flow, + mass_1 * MSUN_SI, + mass_2 * MSUN_SI, + a_1 * np.cos(tilt_1), + a_2 * np.cos(tilt_2), + ) + return max(2 ** (int(np.log2(chirp_time)) + 1), 4) def inject_signal_into_gwpy_timeseries( - data, waveform_generator, parameters, det, power_spectral_density=None, outdir=None, label=None): - """ Inject a signal into a gwpy timeseries + data, + waveform_generator, + parameters, + det, + power_spectral_density=None, + outdir=None, + label=None, +): + """Inject a signal into a gwpy timeseries Parameters ========== @@ -73,13 +86,17 @@ def inject_signal_into_gwpy_timeseries( ifo.strain_data.set_from_gwpy_timeseries(data) parameters_check, _ = convert_to_lal_binary_black_hole_parameters(parameters) - parameters_check = {key: parameters_check[key] for key in - ['mass_1', 'mass_2', 'a_1', 'a_2', 'tilt_1', 'tilt_2']} + parameters_check = { + key: parameters_check[key] + for key in ["mass_1", "mass_2", "a_1", "a_2", "tilt_1", "tilt_2"] + } safe_time = get_safe_signal_duration(**parameters_check) if data.duration.value < safe_time: ValueError( - "Injecting a signal with safe-duration {} longer than the data {}" - .format(safe_time, data.duration.value)) + "Injecting a signal with safe-duration {} longer than the data {}".format( + safe_time, data.duration.value + ) + ) waveform_polarizations = waveform_generator.time_domain_strain(parameters) @@ -87,24 +104,31 @@ def inject_signal_into_gwpy_timeseries( for mode in waveform_polarizations.keys(): det_response = ifo.antenna_response( - parameters['ra'], parameters['dec'], parameters['geocent_time'], - parameters['psi'], mode) + parameters["ra"], + parameters["dec"], + parameters["geocent_time"], + parameters["psi"], + mode, + ) signal += waveform_polarizations[mode] * det_response time_shift = ifo.time_delay_from_geocenter( - parameters['ra'], parameters['dec'], parameters['geocent_time']) + parameters["ra"], parameters["dec"], parameters["geocent_time"] + ) - dt = parameters['geocent_time'] + time_shift - data.times[0].value + dt = parameters["geocent_time"] + time_shift - data.times[0].value n_roll = dt * data.sample_rate.value n_roll = int(np.round(n_roll)) signal_shifted = TimeSeries( - data=np.roll(signal, n_roll), times=data.times, unit=data.unit) + data=np.roll(signal, n_roll), times=data.times, unit=data.unit + ) signal_and_data = data.inject(signal_shifted) if outdir is not None and label is not None: fig = Plot(signal_shifted) - fig.savefig('{}/{}_{}_time_domain_injected_signal'.format( - outdir, ifo.name, label)) + fig.savefig( + "{}/{}_{}_time_domain_injected_signal".format(outdir, ifo.name, label) + ) meta_data = dict(name=det) frequency_domain_signal, _ = utils.nfft( @@ -113,10 +137,12 @@ def inject_signal_into_gwpy_timeseries( frequency_domain_data, _ = utils.nfft( signal_and_data.value, waveform_generator.sampling_frequency ) - meta_data['optimal_SNR'] = ( - np.sqrt(ifo.optimal_snr_squared(signal=frequency_domain_signal)).real) + meta_data["optimal_SNR"] = np.sqrt( + ifo.optimal_snr_squared(signal=frequency_domain_signal) + ).real from ..utils import matched_filter_snr - meta_data['matched_filter_SNR'] = matched_filter_snr( + + meta_data["matched_filter_SNR"] = matched_filter_snr( frequency_domain_signal, frequency_domain_data, ifo.power_spectral_density_array, @@ -125,19 +151,29 @@ def inject_signal_into_gwpy_timeseries( meta_data["parameters"] = parameters logger.info("Injected signal in {}:".format(ifo.name)) - logger.info(" optimal SNR = {:.2f}".format(meta_data['optimal_SNR'])) - logger.info(" matched filter SNR = {:.2f}".format(meta_data['matched_filter_SNR'])) + logger.info(" optimal SNR = {:.2f}".format(meta_data["optimal_SNR"])) + logger.info(" matched filter SNR = {:.2f}".format(meta_data["matched_filter_SNR"])) for key in parameters: - logger.info(' {} = {}'.format(key, parameters[key])) + logger.info(" {} = {}".format(key, parameters[key])) return signal_and_data, meta_data def get_interferometer_with_fake_noise_and_injection( - name, injection_parameters, injection_polarizations=None, - waveform_generator=None, sampling_frequency=4096, duration=4, - start_time=None, outdir='outdir', label=None, plot=True, save=True, - zero_noise=False, raise_error=True): + name, + injection_parameters, + injection_polarizations=None, + waveform_generator=None, + sampling_frequency=4096, + duration=4, + start_time=None, + outdir="outdir", + label=None, + plot=True, + save=True, + zero_noise=False, + raise_error=True, +): """ Helper function to obtain an Interferometer instance with appropriate power spectral density and data, given an center_time. @@ -186,27 +222,33 @@ def get_interferometer_with_fake_noise_and_injection( utils.check_directory_exists_and_if_not_mkdir(outdir) if start_time is None: - start_time = injection_parameters['geocent_time'] + 2 - duration + start_time = injection_parameters["geocent_time"] + 2 - duration interferometer = get_empty_interferometer(name) interferometer.power_spectral_density = PowerSpectralDensity.from_aligo() if zero_noise: interferometer.set_strain_data_from_zero_noise( - sampling_frequency=sampling_frequency, duration=duration, - start_time=start_time) + sampling_frequency=sampling_frequency, + duration=duration, + start_time=start_time, + ) else: interferometer.set_strain_data_from_power_spectral_density( - sampling_frequency=sampling_frequency, duration=duration, - start_time=start_time) + sampling_frequency=sampling_frequency, + duration=duration, + start_time=start_time, + ) injection_polarizations = interferometer.inject_signal( parameters=injection_parameters, injection_polarizations=injection_polarizations, waveform_generator=waveform_generator, - raise_error=raise_error) + raise_error=raise_error, + ) signal = interferometer.get_detector_response( - injection_polarizations, injection_parameters) + injection_polarizations, injection_parameters + ) if plot: interferometer.plot_data(signal=signal, outdir=outdir, label=label) @@ -218,10 +260,18 @@ def get_interferometer_with_fake_noise_and_injection( def load_data_from_cache_file( - cache_file, start_time, segment_duration, psd_duration, psd_start_time, - channel_name=None, sampling_frequency=4096, roll_off=0.2, - overlap=0, outdir=None): - """ Helper routine to generate an interferometer from a cache file + cache_file, + start_time, + segment_duration, + psd_duration, + psd_start_time, + channel_name=None, + sampling_frequency=4096, + roll_off=0.2, + overlap=0, + outdir=None, +): + """Helper routine to generate an interferometer from a cache file Parameters ========== @@ -250,71 +300,84 @@ def load_data_from_cache_file( appropriate data in the cache file and a PSD. """ import lal + data_set = False psd_set = False - with open(cache_file, 'r') as ff: + with open(cache_file, "r") as ff: lines = ff.readlines() - if len(lines)>1: - raise ValueError('This method cannot handle cache files with' - ' multiple frames. Use `load_data_by_channel_name' - ' instead.') + if len(lines) > 1: + raise ValueError( + "This method cannot handle cache files with" + " multiple frames. Use `load_data_by_channel_name" + " instead." + ) else: line = lines[0] cache = lal.utils.cache.CacheEntry(line) - data_in_cache = ( - (cache.segment[0].gpsSeconds < start_time) & - (cache.segment[1].gpsSeconds > start_time + segment_duration)) - psd_in_cache = ( - (cache.segment[0].gpsSeconds < psd_start_time) & - (cache.segment[1].gpsSeconds > psd_start_time + psd_duration)) - ifo = get_empty_interferometer( - "{}1".format(cache.observatory)) + data_in_cache = (cache.segment[0].gpsSeconds < start_time) & ( + cache.segment[1].gpsSeconds > start_time + segment_duration + ) + psd_in_cache = (cache.segment[0].gpsSeconds < psd_start_time) & ( + cache.segment[1].gpsSeconds > psd_start_time + psd_duration + ) + ifo = get_empty_interferometer("{}1".format(cache.observatory)) if not data_in_cache: - raise ValueError('The specified data segment does not exist in' - ' this frame.') + raise ValueError( + "The specified data segment does not exist in this frame." + ) if not psd_in_cache: - raise ValueError('The specified PSD data segment does not exist' - ' in this frame.') + raise ValueError( + "The specified PSD data segment does not exist in this frame." + ) if (not data_set) & data_in_cache: ifo.set_strain_data_from_frame_file( frame_file=cache.path, sampling_frequency=sampling_frequency, duration=segment_duration, start_time=start_time, - channel=channel_name, buffer_time=0) + channel=channel_name, + buffer_time=0, + ) data_set = True if (not psd_set) & psd_in_cache: - ifo.power_spectral_density = \ - PowerSpectralDensity.from_frame_file( - cache.path, - psd_start_time=psd_start_time, - psd_duration=psd_duration, - fft_length=segment_duration, - sampling_frequency=sampling_frequency, - roll_off=roll_off, - overlap=overlap, - channel=channel_name, - name=cache.observatory, - outdir=outdir, - analysis_segment_start_time=start_time) + ifo.power_spectral_density = PowerSpectralDensity.from_frame_file( + cache.path, + psd_start_time=psd_start_time, + psd_duration=psd_duration, + fft_length=segment_duration, + sampling_frequency=sampling_frequency, + roll_off=roll_off, + overlap=overlap, + channel=channel_name, + name=cache.observatory, + outdir=outdir, + analysis_segment_start_time=start_time, + ) psd_set = True if data_set and psd_set: return ifo elif not data_set: - raise ValueError('Data not loaded for {}'.format(ifo.name)) + raise ValueError("Data not loaded for {}".format(ifo.name)) elif not psd_set: - raise ValueError('PSD not created for {}'.format(ifo.name)) + raise ValueError("PSD not created for {}".format(ifo.name)) def load_data_by_channel_name( - channel_name, start_time, segment_duration, psd_duration, psd_start_time, - sampling_frequency=4096, roll_off=0.2, - overlap=0, outdir=None): - """ Helper routine to generate an interferometer from a channel name - This function creates an empty interferometer specified in the name - of the channel. It calls `ifo.set_strain_data_from_channel_name` to - set the data and PSD in the interferometer using data retrieved from + channel_name, + start_time, + segment_duration, + psd_duration, + psd_start_time, + sampling_frequency=4096, + roll_off=0.2, + overlap=0, + outdir=None, +): + """Helper routine to generate an interferometer from a channel name + This function creates an empty interferometer specified in the name + of the channel. It calls `ifo.set_strain_data_from_channel_name` to + set the data and PSD in the interferometer using data retrieved from the specified channel using gwpy.TimeSeries.get() Parameters @@ -342,7 +405,7 @@ def load_data_by_channel_name( appropriate data fetched from the specified channel and a PSD. """ try: - det = channel_name.split(':')[-2] + det = channel_name.split(":")[-2] except IndexError: raise IndexError("Channel name must be of the format `IFO:Channel`") ifo = get_empty_interferometer(det) @@ -351,18 +414,19 @@ def load_data_by_channel_name( channel=channel_name, sampling_frequency=sampling_frequency, duration=segment_duration, - start_time=start_time) - - ifo.power_spectral_density = \ - PowerSpectralDensity.from_channel_name( - channel=channel_name, - psd_start_time=psd_start_time, - psd_duration=psd_duration, - fft_length=segment_duration, - sampling_frequency=sampling_frequency, - roll_off=roll_off, - overlap=overlap, - name=det, - outdir=outdir, - analysis_segment_start_time=start_time) + start_time=start_time, + ) + + ifo.power_spectral_density = PowerSpectralDensity.from_channel_name( + channel=channel_name, + psd_start_time=psd_start_time, + psd_duration=psd_duration, + fft_length=segment_duration, + sampling_frequency=sampling_frequency, + roll_off=roll_off, + overlap=overlap, + name=det, + outdir=outdir, + analysis_segment_start_time=start_time, + ) return ifo diff --git a/bilby/gw/detector/calibration.py b/bilby/gw/detector/calibration.py index c904c707d..056147d84 100644 --- a/bilby/gw/detector/calibration.py +++ b/bilby/gw/detector/calibration.py @@ -39,6 +39,7 @@ :code:`"data"` convention. """ + import copy import os @@ -77,7 +78,11 @@ def _check_calibration_correction_type(correction_type): def read_calibration_file( - filename, frequency_array, number_of_response_curves, starting_index=0, correction_type=None + filename, + frequency_array, + number_of_response_curves, + starting_index=0, + correction_type=None, ): r""" Function to read the hdf5 files from the calibration group containing the @@ -112,24 +117,30 @@ def read_calibration_file( """ import h5py - correction_type = _check_calibration_correction_type(correction_type=correction_type) + correction_type = _check_calibration_correction_type( + correction_type=correction_type + ) logger.info(f"Reading calibration draws from {filename}") - with h5py.File(filename, 'r') as calibration_file: + with h5py.File(filename, "r") as calibration_file: try: - dr = calibration_file['deltaR'] + dr = calibration_file["deltaR"] except KeyError: raise KeyError(f"File {filename} does not contain 'deltaR' group.") # slice draws according to starting_index and number_of_response_curves - calibration_amplitude = dr['draws_amp_rel'][starting_index: starting_index + number_of_response_curves] - calibration_phase = dr['draws_phase'][starting_index: starting_index + number_of_response_curves] - calibration_frequencies = dr['freq'][:] + calibration_amplitude = dr["draws_amp_rel"][ + starting_index : starting_index + number_of_response_curves + ] + calibration_phase = dr["draws_phase"][ + starting_index : starting_index + number_of_response_curves + ] + calibration_frequencies = dr["freq"][:] # read parameter draws if present; stored under CalParams/table as a structured dataset parameter_draws = None - if 'CalParams' in calibration_file and 'table' in calibration_file['CalParams']: - table_ds = calibration_file['CalParams']['table'] + if "CalParams" in calibration_file and "table" in calibration_file["CalParams"]: + table_ds = calibration_file["CalParams"]["table"] try: rec = np.array(table_ds) # convert structured array to DataFrame @@ -141,8 +152,12 @@ def read_calibration_file( # combine into complex responses and interpolate to requested frequency array calibration_draws = calibration_amplitude * np.exp(1j * calibration_phase) calibration_draws = interp1d( - calibration_frequencies, calibration_draws, kind='cubic', - bounds_error=False, fill_value=1)(frequency_array) + calibration_frequencies, + calibration_draws, + kind="cubic", + bounds_error=False, + fill_value=1, + )(frequency_array) if correction_type == "data": calibration_draws = 1 / calibration_draws @@ -150,7 +165,11 @@ def read_calibration_file( def write_calibration_file( - filename, frequency_array, calibration_draws, calibration_parameter_draws=None, correction_type=None + filename, + frequency_array, + calibration_draws, + calibration_parameter_draws=None, + correction_type=None, ): """ Function to write the generated response curves to file @@ -178,7 +197,9 @@ def write_calibration_file( """ import h5py - correction_type = _check_calibration_correction_type(correction_type=correction_type) + correction_type = _check_calibration_correction_type( + correction_type=correction_type + ) # if the correction is specified as mapping data -> template, invert when writing if correction_type == "data": @@ -186,31 +207,34 @@ def write_calibration_file( logger.info(f"Writing calibration draws to {filename}") # overwrite/create the file - with h5py.File(filename, 'w') as calibration_file: - deltaR_group = calibration_file.create_group('deltaR') + with h5py.File(filename, "w") as calibration_file: + deltaR_group = calibration_file.create_group("deltaR") # Save output: amplitude and phase arrays amp = np.abs(calibration_draws) phase = np.angle(calibration_draws) - deltaR_group.create_dataset('draws_amp_rel', data=amp, dtype=np.float64, compression='gzip') - deltaR_group.create_dataset('draws_phase', data=phase, dtype=np.float64, compression='gzip') - deltaR_group.create_dataset('freq', data=frequency_array, dtype=np.float64) + deltaR_group.create_dataset( + "draws_amp_rel", data=amp, dtype=np.float64, compression="gzip" + ) + deltaR_group.create_dataset( + "draws_phase", data=phase, dtype=np.float64, compression="gzip" + ) + deltaR_group.create_dataset("freq", data=frequency_array, dtype=np.float64) # Save calibration parameter draws (DataFrame) if provided. Store as a structured # dataset under CalParams/table so it can be read back into a DataFrame. if calibration_parameter_draws is not None: - cp_group = calibration_file.create_group('CalParams') + cp_group = calibration_file.create_group("CalParams") # convert DataFrame to a numpy recarray / structured array rec = calibration_parameter_draws.to_records(index=False) # create dataset - cp_group.create_dataset('table', data=rec, compression='gzip') + cp_group.create_dataset("table", data=rec, compression="gzip") class Recalibrate(object): + name = "none" - name = 'none' - - def __init__(self, prefix='recalib_'): + def __init__(self, prefix="recalib_"): """ Base calibration object. This applies no transformation @@ -223,7 +247,7 @@ def __init__(self, prefix='recalib_'): self.prefix = prefix def __repr__(self): - return self.__class__.__name__ + '(prefix=\'{}\')'.format(self.prefix) + return self.__class__.__name__ + "(prefix='{}')".format(self.prefix) def get_calibration_factor(self, frequency_array, **params): """Apply calibration model @@ -247,16 +271,20 @@ def get_calibration_factor(self, frequency_array, **params): return xp.ones_like(frequency_array) def set_calibration_parameters(self, **params): - self.params.update({key[len(self.prefix):]: params[key] for key in params - if self.prefix in key}) + self.params.update( + { + key[len(self.prefix) :]: params[key] + for key in params + if self.prefix in key + } + ) def __eq__(self, other): return self.__dict__ == other.__dict__ class CubicSpline(Recalibrate): - - name = 'cubic_spline' + name = "cubic_spline" def __init__(self, prefix, minimum_frequency, maximum_frequency, n_points): """ @@ -280,17 +308,22 @@ def __init__(self, prefix, minimum_frequency, maximum_frequency, n_points): """ super(CubicSpline, self).__init__(prefix=prefix) if n_points < 4: - raise ValueError('Cubic spline calibration requires at least 4 spline nodes.') + raise ValueError( + "Cubic spline calibration requires at least 4 spline nodes." + ) self.n_points = n_points self.minimum_frequency = minimum_frequency self.maximum_frequency = maximum_frequency self._log_spline_points = np.linspace( - np.log10(minimum_frequency), np.log10(maximum_frequency), n_points) + np.log10(minimum_frequency), np.log10(maximum_frequency), n_points + ) @property def delta_log_spline_points(self): if not hasattr(self, "_delta_log_spline_points"): - self._delta_log_spline_points = self._log_spline_points[1] - self._log_spline_points[0] + self._delta_log_spline_points = ( + self._log_spline_points[1] - self._log_spline_points[0] + ) return self._delta_log_spline_points @property @@ -329,13 +362,22 @@ def log_spline_points(self): return self._log_spline_points def __repr__(self): - return self.__class__.__name__ + '(prefix=\'{}\', minimum_frequency={}, maximum_frequency={}, n_points={})'\ - .format(self.prefix, self.minimum_frequency, self.maximum_frequency, self.n_points) + return ( + self.__class__.__name__ + + "(prefix='{}', minimum_frequency={}, maximum_frequency={}, n_points={})".format( + self.prefix, + self.minimum_frequency, + self.maximum_frequency, + self.n_points, + ) + ) def _evaluate_spline(self, kind, a, b, c, d, previous_nodes): """Evaluate Eq. (1) in https://dcc.ligo.org/LIGO-T2300140""" xp = array_module(self.params[f"{kind}_0"]) - parameters = xp.asarray([self.params[f"{kind}_{ii}"] for ii in range(self.n_points)]) + parameters = xp.asarray( + [self.params[f"{kind}_{ii}"] for ii in range(self.n_points)] + ) next_nodes = previous_nodes + 1 nodes = xp.asarray(self.nodes_to_spline_coefficients) spline_coefficients = nodes.dot(parameters) @@ -365,11 +407,16 @@ def get_calibration_factor(self, frequency_array, *, xp=np, **params): calibration_factor : array-like The factor to multiply the strain by. """ - log10f_per_deltalog10f = xp.nan_to_num( - xp.log10(frequency_array) - xp.asarray(self.log_spline_points[0]), - neginf=0.0, - ) / self.delta_log_spline_points - previous_nodes = xp.clip(xp.astype(log10f_per_deltalog10f, int), min=0, max=self.n_points - 2) + log10f_per_deltalog10f = ( + xp.nan_to_num( + xp.log10(frequency_array) - xp.asarray(self.log_spline_points[0]), + neginf=0.0, + ) + / self.delta_log_spline_points + ) + previous_nodes = xp.clip( + xp.astype(log10f_per_deltalog10f, int), min=0, max=self.n_points - 2 + ) b = log10f_per_deltalog10f - previous_nodes a = 1 - b c = (a**3 - a) / 6 @@ -379,13 +426,14 @@ def get_calibration_factor(self, frequency_array, *, xp=np, **params): delta_amplitude = self._evaluate_spline("amplitude", a, b, c, d, previous_nodes) delta_phase = self._evaluate_spline("phase", a, b, c, d, previous_nodes) - calibration_factor = (1 + delta_amplitude) * (2 + 1j * delta_phase) / (2 - 1j * delta_phase) + calibration_factor = ( + (1 + delta_amplitude) * (2 + 1j * delta_phase) / (2 - 1j * delta_phase) + ) return xp.nan_to_num(calibration_factor) class Precomputed(Recalibrate): - name = "precomputed" def __init__(self, label, curves, frequency_array, parameters=None): @@ -448,7 +496,7 @@ def constant_uncertainty_spline( frequency_array=frequency_array, n_points=n_nodes, parameters=parameters, - n_curves=n_curves + n_curves=n_curves, ) return cls( label=label, @@ -485,7 +533,9 @@ def from_envelope_file( ) @classmethod - def from_calibration_file(cls, label, filename, frequency_array, n_curves, starting_index=0): + def from_calibration_file( + cls, label, filename, frequency_array, n_curves, starting_index=0 + ): curves, parameters = read_calibration_file( filename=filename, frequency_array=frequency_array, @@ -565,17 +615,21 @@ def _generate_calibration_draws(interferometer, priors, n_curves): parameters = pd.DataFrame(calibration_priors.sample(n_curves)) - draws = np.array(curves_from_spline_and_prior( - parameters=parameters, - label=name, - n_points=interferometer.calibration_model.n_points, - frequency_array=frequencies, - n_curves=n_curves, - )) + draws = np.array( + curves_from_spline_and_prior( + parameters=parameters, + label=name, + n_points=interferometer.calibration_model.n_points, + frequency_array=frequencies, + n_curves=n_curves, + ) + ) return draws, parameters -def curves_from_spline_and_prior(parameters, label, n_points, frequency_array, n_curves): +def curves_from_spline_and_prior( + parameters, label, n_points, frequency_array, n_curves +): spline = CubicSpline( prefix=f"recalib_{label}_", minimum_frequency=frequency_array[0], @@ -584,8 +638,9 @@ def curves_from_spline_and_prior(parameters, label, n_points, frequency_array, n ) curves = list() for ii in range(n_curves): - curves.append(spline.get_calibration_factor( - frequency_array=frequency_array, - **parameters.iloc[ii] - )) + curves.append( + spline.get_calibration_factor( + frequency_array=frequency_array, **parameters.iloc[ii] + ) + ) return curves diff --git a/bilby/gw/detector/geometry.py b/bilby/gw/detector/geometry.py index f2461761e..b39e0530e 100644 --- a/bilby/gw/detector/geometry.py +++ b/bilby/gw/detector/geometry.py @@ -6,8 +6,17 @@ class InterferometerGeometry(object): - def __init__(self, length, latitude, longitude, elevation, xarm_azimuth, yarm_azimuth, - xarm_tilt=0., yarm_tilt=0.): + def __init__( + self, + length, + latitude, + longitude, + elevation, + xarm_azimuth, + yarm_azimuth, + xarm_tilt=0.0, + yarm_tilt=0.0, + ): """ Instantiate an Interferometer object. @@ -52,22 +61,39 @@ def __init__(self, length, latitude, longitude, elevation, xarm_azimuth, yarm_az self._detector_tensor = None def __eq__(self, other): - for attribute in ['length', 'latitude', 'longitude', 'elevation', - 'xarm_azimuth', 'yarm_azimuth', 'xarm_tilt', 'yarm_tilt']: + for attribute in [ + "length", + "latitude", + "longitude", + "elevation", + "xarm_azimuth", + "yarm_azimuth", + "xarm_tilt", + "yarm_tilt", + ]: if not getattr(self, attribute) == getattr(other, attribute): return False return True def __repr__(self): - return self.__class__.__name__ + '(length={}, latitude={}, longitude={}, elevation={}, ' \ - 'xarm_azimuth={}, yarm_azimuth={}, xarm_tilt={}, yarm_tilt={})' \ - .format(float(self.length), float(self.latitude), float(self.longitude), - float(self.elevation), float(self.xarm_azimuth), float(self.yarm_azimuth), float(self.xarm_tilt), - float(self.yarm_tilt)) + return ( + self.__class__.__name__ + + "(length={}, latitude={}, longitude={}, elevation={}, " + "xarm_azimuth={}, yarm_azimuth={}, xarm_tilt={}, yarm_tilt={})".format( + float(self.length), + float(self.latitude), + float(self.longitude), + float(self.elevation), + float(self.xarm_azimuth), + float(self.yarm_azimuth), + float(self.xarm_tilt), + float(self.yarm_tilt), + ) + ) @property def latitude(self): - """ Saves latitude in rad internally. Updates related quantities if set to a different value. + """Saves latitude in rad internally. Updates related quantities if set to a different value. Returns ======= @@ -94,7 +120,7 @@ def latitude_radians(self): @property def longitude(self): - """ Saves longitude in rad internally. Updates related quantities if set to a different value. + """Saves longitude in rad internally. Updates related quantities if set to a different value. Returns ======= @@ -121,7 +147,7 @@ def longitude_radians(self): @property def elevation(self): - """ Updates related quantities if set to a different values. + """Updates related quantities if set to a different values. Returns ======= @@ -136,7 +162,7 @@ def elevation(self, elevation): @property def xarm_azimuth(self): - """ Saves the x-arm azimuth in rad internally. Updates related quantities if set to a different values. + """Saves the x-arm azimuth in rad internally. Updates related quantities if set to a different values. Returns ======= @@ -152,7 +178,7 @@ def xarm_azimuth(self, xarm_azimuth): @property def yarm_azimuth(self): - """ Saves the y-arm azimuth in rad internally. Updates related quantities if set to a different values. + """Saves the y-arm azimuth in rad internally. Updates related quantities if set to a different values. Returns ======= @@ -168,7 +194,7 @@ def yarm_azimuth(self, yarm_azimuth): @property def xarm_tilt(self): - """ Updates related quantities if set to a different values. + """Updates related quantities if set to a different values. Returns ======= @@ -184,7 +210,7 @@ def xarm_tilt(self, xarm_tilt): @property def yarm_tilt(self): - """ Updates related quantities if set to a different values. + """Updates related quantities if set to a different values. Returns ======= @@ -200,7 +226,7 @@ def yarm_tilt(self, yarm_tilt): @property def vertex(self): - """ Position of the IFO vertex in geocentric coordinates in meters. + """Position of the IFO vertex in geocentric coordinates in meters. Is automatically updated if related quantities are modified. @@ -212,14 +238,15 @@ def vertex(self): xp = self.array_backend if not self._vertex_updated: - self._vertex = gwutils.get_vertex_position_geocentric(xp.asarray([self._latitude, self._longitude, - self.elevation])) + self._vertex = gwutils.get_vertex_position_geocentric( + xp.asarray([self._latitude, self._longitude, self.elevation]) + ) self._vertex_updated = True return self._vertex @property def x(self): - """ A unit vector along the x-arm + """A unit vector along the x-arm Is automatically updated if related quantities are modified. @@ -229,14 +256,14 @@ def x(self): """ if not self._x_updated: - self._x = self.unit_vector_along_arm('x') + self._x = self.unit_vector_along_arm("x") self._x_updated = True self._detector_tensor_updated = False return self._x @property def y(self): - """ A unit vector along the y-arm + """A unit vector along the y-arm Is automatically updated if related quantities are modified. @@ -246,7 +273,7 @@ def y(self): """ if not self._y_updated: - self._y = self.unit_vector_along_arm('y') + self._y = self.unit_vector_along_arm("y") self._y_updated = True self._detector_tensor_updated = False return self._y @@ -292,19 +319,13 @@ def unit_vector_along_arm(self, arm): ValueError: If arm is neither 'x' nor 'y' """ - if arm == 'x': + if arm == "x": return calculate_arm( - self._xarm_tilt, - self._xarm_azimuth, - self._longitude, - self._latitude + self._xarm_tilt, self._xarm_azimuth, self._longitude, self._latitude ) - elif arm == 'y': + elif arm == "y": return calculate_arm( - self._yarm_tilt, - self._yarm_azimuth, - self._longitude, - self._latitude + self._yarm_tilt, self._yarm_azimuth, self._longitude, self._latitude ) else: raise ValueError("Arm must either be 'x' or 'y'.") diff --git a/bilby/gw/detector/interferometer.py b/bilby/gw/detector/interferometer.py index 6fbb8d1f8..2a241384a 100644 --- a/bilby/gw/detector/interferometer.py +++ b/bilby/gw/detector/interferometer.py @@ -18,36 +18,50 @@ class Interferometer(object): - """Class for the Interferometer """ - - length = PropertyAccessor('geometry', 'length') - latitude = PropertyAccessor('geometry', 'latitude') - latitude_radians = PropertyAccessor('geometry', 'latitude_radians') - longitude = PropertyAccessor('geometry', 'longitude') - longitude_radians = PropertyAccessor('geometry', 'longitude_radians') - elevation = PropertyAccessor('geometry', 'elevation') - x = PropertyAccessor('geometry', 'x') - y = PropertyAccessor('geometry', 'y') - xarm_azimuth = PropertyAccessor('geometry', 'xarm_azimuth') - yarm_azimuth = PropertyAccessor('geometry', 'yarm_azimuth') - xarm_tilt = PropertyAccessor('geometry', 'xarm_tilt') - yarm_tilt = PropertyAccessor('geometry', 'yarm_tilt') - vertex = PropertyAccessor('geometry', 'vertex') - detector_tensor = PropertyAccessor('geometry', 'detector_tensor') - - duration = PropertyAccessor('strain_data', 'duration') - sampling_frequency = PropertyAccessor('strain_data', 'sampling_frequency') - start_time = PropertyAccessor('strain_data', 'start_time') - frequency_array = PropertyAccessor('strain_data', 'frequency_array') - time_array = PropertyAccessor('strain_data', 'time_array') - minimum_frequency = PropertyAccessor('strain_data', 'minimum_frequency') - maximum_frequency = PropertyAccessor('strain_data', 'maximum_frequency') - frequency_mask = PropertyAccessor('strain_data', 'frequency_mask') - frequency_domain_strain = PropertyAccessor('strain_data', 'frequency_domain_strain') - time_domain_strain = PropertyAccessor('strain_data', 'time_domain_strain') - - def __init__(self, name, power_spectral_density, minimum_frequency, maximum_frequency, length, latitude, longitude, - elevation, xarm_azimuth, yarm_azimuth, xarm_tilt=0., yarm_tilt=0., calibration_model=Recalibrate()): + """Class for the Interferometer""" + + length = PropertyAccessor("geometry", "length") + latitude = PropertyAccessor("geometry", "latitude") + latitude_radians = PropertyAccessor("geometry", "latitude_radians") + longitude = PropertyAccessor("geometry", "longitude") + longitude_radians = PropertyAccessor("geometry", "longitude_radians") + elevation = PropertyAccessor("geometry", "elevation") + x = PropertyAccessor("geometry", "x") + y = PropertyAccessor("geometry", "y") + xarm_azimuth = PropertyAccessor("geometry", "xarm_azimuth") + yarm_azimuth = PropertyAccessor("geometry", "yarm_azimuth") + xarm_tilt = PropertyAccessor("geometry", "xarm_tilt") + yarm_tilt = PropertyAccessor("geometry", "yarm_tilt") + vertex = PropertyAccessor("geometry", "vertex") + detector_tensor = PropertyAccessor("geometry", "detector_tensor") + + duration = PropertyAccessor("strain_data", "duration") + sampling_frequency = PropertyAccessor("strain_data", "sampling_frequency") + start_time = PropertyAccessor("strain_data", "start_time") + frequency_array = PropertyAccessor("strain_data", "frequency_array") + time_array = PropertyAccessor("strain_data", "time_array") + minimum_frequency = PropertyAccessor("strain_data", "minimum_frequency") + maximum_frequency = PropertyAccessor("strain_data", "maximum_frequency") + frequency_mask = PropertyAccessor("strain_data", "frequency_mask") + frequency_domain_strain = PropertyAccessor("strain_data", "frequency_domain_strain") + time_domain_strain = PropertyAccessor("strain_data", "time_domain_strain") + + def __init__( + self, + name, + power_spectral_density, + minimum_frequency, + maximum_frequency, + length, + latitude, + longitude, + elevation, + xarm_azimuth, + yarm_azimuth, + xarm_tilt=0.0, + yarm_tilt=0.0, + calibration_model=Recalibrate(), + ): """ Instantiate an Interferometer object. @@ -82,40 +96,60 @@ def __init__(self, name, power_spectral_density, minimum_frequency, maximum_freq Calibration model, this applies the calibration correction to the template, the default model applies no correction. """ - self.geometry = InterferometerGeometry(length, latitude, longitude, elevation, - xarm_azimuth, yarm_azimuth, xarm_tilt, yarm_tilt) + self.geometry = InterferometerGeometry( + length, + latitude, + longitude, + elevation, + xarm_azimuth, + yarm_azimuth, + xarm_tilt, + yarm_tilt, + ) self.name = name self.power_spectral_density = power_spectral_density self.calibration_model = calibration_model self.strain_data = InterferometerStrainData( - minimum_frequency=minimum_frequency, - maximum_frequency=maximum_frequency) + minimum_frequency=minimum_frequency, maximum_frequency=maximum_frequency + ) self.meta_data = dict(name=name) self.reference_time = None def __eq__(self, other): - if self.name == other.name and \ - self.geometry == other.geometry and \ - self.power_spectral_density.__eq__(other.power_spectral_density) and \ - self.calibration_model == other.calibration_model and \ - self.strain_data == other.strain_data: + if ( + self.name == other.name + and self.geometry == other.geometry + and self.power_spectral_density.__eq__(other.power_spectral_density) + and self.calibration_model == other.calibration_model + and self.strain_data == other.strain_data + ): return True return False def __repr__(self): - return self.__class__.__name__ + '(name=\'{}\', power_spectral_density={}, minimum_frequency={}, ' \ - 'maximum_frequency={}, length={}, latitude={}, longitude={}, elevation={}, ' \ - 'xarm_azimuth={}, yarm_azimuth={}, xarm_tilt={}, yarm_tilt={})' \ - .format(self.name, self.power_spectral_density, float(self.strain_data.minimum_frequency), - float(self.strain_data.maximum_frequency), float(self.geometry.length), - float(self.geometry.latitude), float(self.geometry.longitude), - float(self.geometry.elevation), float(self.geometry.xarm_azimuth), - float(self.geometry.yarm_azimuth), float(self.geometry.xarm_tilt), - float(self.geometry.yarm_tilt)) + return ( + self.__class__.__name__ + + "(name='{}', power_spectral_density={}, minimum_frequency={}, " + "maximum_frequency={}, length={}, latitude={}, longitude={}, elevation={}, " + "xarm_azimuth={}, yarm_azimuth={}, xarm_tilt={}, yarm_tilt={})".format( + self.name, + self.power_spectral_density, + float(self.strain_data.minimum_frequency), + float(self.strain_data.maximum_frequency), + float(self.geometry.length), + float(self.geometry.latitude), + float(self.geometry.longitude), + float(self.geometry.elevation), + float(self.geometry.xarm_azimuth), + float(self.geometry.yarm_azimuth), + float(self.geometry.xarm_tilt), + float(self.geometry.yarm_tilt), + ) + ) def set_strain_data_from_gwpy_timeseries(self, time_series, *, xp=None): - """ Set the `Interferometer.strain_data` from a gwpy TimeSeries + """Set the `Interferometer.strain_data` from a gwpy TimeSeries Parameters ========== @@ -129,9 +163,14 @@ def set_strain_data_from_gwpy_timeseries(self, time_series, *, xp=None): self.strain_data.set_from_gwpy_timeseries(time_series=time_series, xp=xp) def set_strain_data_from_frequency_domain_strain( - self, frequency_domain_strain, sampling_frequency=None, - duration=None, start_time=0, frequency_array=None): - """ Set the `Interferometer.strain_data` from a numpy array + self, + frequency_domain_strain, + sampling_frequency=None, + duration=None, + start_time=0, + frequency_array=None, + ): + """Set the `Interferometer.strain_data` from a numpy array Parameters ========== @@ -150,12 +189,16 @@ def set_strain_data_from_frequency_domain_strain( """ self.strain_data.set_from_frequency_domain_strain( frequency_domain_strain=frequency_domain_strain, - sampling_frequency=sampling_frequency, duration=duration, - start_time=start_time, frequency_array=frequency_array) + sampling_frequency=sampling_frequency, + duration=duration, + start_time=start_time, + frequency_array=frequency_array, + ) def set_strain_data_from_power_spectral_density( - self, sampling_frequency, duration, start_time=0, *, random_state=None): - """ Set the `Interferometer.strain_data` from a power spectal density + self, sampling_frequency, duration, start_time=0, *, random_state=None + ): + """Set the `Interferometer.strain_data` from a power spectal density This uses the `interferometer.power_spectral_density` object to set the `strain_data` to a noise realization. See @@ -172,13 +215,25 @@ def set_strain_data_from_power_spectral_density( """ self.strain_data.set_from_power_spectral_density( - self.power_spectral_density, sampling_frequency=sampling_frequency, - duration=duration, start_time=start_time, random_state=random_state) + self.power_spectral_density, + sampling_frequency=sampling_frequency, + duration=duration, + start_time=start_time, + random_state=random_state, + ) def set_strain_data_from_frame_file( - self, frame_file, sampling_frequency, duration, start_time=0, - channel=None, buffer_time=1, *, xp=None): - """ Set the `Interferometer.strain_data` from a frame file + self, + frame_file, + sampling_frequency, + duration, + start_time=0, + channel=None, + buffer_time=1, + *, + xp=None, + ): + """Set the `Interferometer.strain_data` from a frame file Parameters ========== @@ -201,12 +256,18 @@ def set_strain_data_from_frame_file( """ self.strain_data.set_from_frame_file( - frame_file=frame_file, sampling_frequency=sampling_frequency, - duration=duration, start_time=start_time, - channel=channel, buffer_time=buffer_time, xp=xp) + frame_file=frame_file, + sampling_frequency=sampling_frequency, + duration=duration, + start_time=start_time, + channel=channel, + buffer_time=buffer_time, + xp=xp, + ) def set_strain_data_from_channel_name( - self, channel, sampling_frequency, duration, start_time=0, *, xp=None): + self, channel, sampling_frequency, duration, start_time=0, *, xp=None + ): """ Set the `Interferometer.strain_data` by fetching from given channel using strain_data.set_from_channel_name() @@ -227,11 +288,15 @@ def set_strain_data_from_channel_name( """ self.strain_data.set_from_channel_name( - channel=channel, sampling_frequency=sampling_frequency, - duration=duration, start_time=start_time, xp=xp) + channel=channel, + sampling_frequency=sampling_frequency, + duration=duration, + start_time=start_time, + xp=xp, + ) def set_strain_data_from_csv(self, filename, *, xp=None): - """ Set the `Interferometer.strain_data` from a csv file + """Set the `Interferometer.strain_data` from a csv file Parameters ========== @@ -245,8 +310,9 @@ def set_strain_data_from_csv(self, filename, *, xp=None): self.strain_data.set_from_csv(filename, xp=xp) def set_strain_data_from_zero_noise( - self, sampling_frequency, duration, start_time=0): - """ Set the `Interferometer.strain_data` to zero noise + self, sampling_frequency, duration, start_time=0 + ): + """Set the `Interferometer.strain_data` to zero noise Parameters ========== @@ -260,8 +326,10 @@ def set_strain_data_from_zero_noise( """ self.strain_data.set_from_zero_noise( - sampling_frequency=sampling_frequency, duration=duration, - start_time=start_time) + sampling_frequency=sampling_frequency, + duration=duration, + start_time=start_time, + ) def antenna_response(self, ra, dec, time, psi, mode): """ @@ -293,14 +361,18 @@ def antenna_response(self, ra, dec, time, psi, mode): """ if mode in ["plus", "cross", "x", "y", "breathing", "longitudinal"]: polarization_tensor = get_polarization_tensor(ra, dec, time, psi, mode) - return three_by_three_matrix_contraction(self.geometry.detector_tensor, polarization_tensor) + return three_by_three_matrix_contraction( + self.geometry.detector_tensor, polarization_tensor + ) elif mode == self.name: return 1 else: return 0 - def get_detector_response(self, waveform_polarizations, parameters, frequencies=None): - """ Get the detector response for a particular waveform + def get_detector_response( + self, waveform_polarizations, parameters, frequencies=None + ): + """Get the detector response for a particular waveform Parameters ========== @@ -340,20 +412,23 @@ def get_detector_response(self, waveform_polarizations, parameters, frequencies= signal = {} for mode in waveform_polarizations.keys(): det_response = self.antenna_response( - parameters['ra'], - parameters['dec'], + parameters["ra"], + parameters["dec"], antenna_time, - parameters['psi'], mode) + parameters["psi"], + mode, + ) signal[mode] = waveform_polarizations[mode] * mask * det_response signal_ifo = sum(signal.values()) time_shift = self.time_delay_from_geocenter( - parameters['ra'], parameters['dec'], parameters['geocent_time']) + parameters["ra"], parameters["dec"], parameters["geocent_time"] + ) # Be careful to first subtract the two GPS times which are ~1e9 sec. # And then add the time_shift which varies at ~1e-5 sec - dt_geocent = parameters['geocent_time'] - self.strain_data.start_time + dt_geocent = parameters["geocent_time"] - self.strain_data.start_time dt = dt_geocent + time_shift xp = array_module(signal_ifo) @@ -361,13 +436,13 @@ def get_detector_response(self, waveform_polarizations, parameters, frequencies= signal_ifo = signal_ifo * xp.exp(-1j * 2 * np.pi * dt * frequencies) signal_ifo *= self.calibration_model.get_calibration_factor( - frequencies, prefix=f'recalib_{self.name}_', xp=xp, **parameters + frequencies, prefix=f"recalib_{self.name}_", xp=xp, **parameters ) return signal_ifo def check_signal_duration(self, parameters, raise_error=True): - """ Check that the signal with the given parameters fits in the data + """Check that the signal with the given parameters fits in the data Parameters ========== @@ -387,7 +462,9 @@ def check_signal_duration(self, parameters, raise_error=True): if ("mass_1" not in parameters) and ("mass_2" not in parameters): if raise_error: - raise AttributeError("Unable to check signal duration as mass not given") + raise AttributeError( + "Unable to check signal duration as mass not given" + ) else: return @@ -408,9 +485,14 @@ def check_signal_duration(self, parameters, raise_error=True): else: logger.warning(msg) - def inject_signal(self, parameters, injection_polarizations=None, - waveform_generator=None, raise_error=True): - """ General signal injection method. + def inject_signal( + self, + parameters, + injection_polarizations=None, + waveform_generator=None, + raise_error=True, + ): + """General signal injection method. Provide the injection parameters and either the injection polarizations or the waveform generator to inject a signal into the detector. Defaults to the injection polarizations is both are given. @@ -450,17 +532,20 @@ def inject_signal(self, parameters, injection_polarizations=None, if injection_polarizations is None and waveform_generator is None: raise ValueError( "inject_signal needs one of waveform_generator or " - "injection_polarizations.") + "injection_polarizations." + ) elif injection_polarizations is not None: - self.inject_signal_from_waveform_polarizations(parameters=parameters, - injection_polarizations=injection_polarizations) + self.inject_signal_from_waveform_polarizations( + parameters=parameters, injection_polarizations=injection_polarizations + ) elif waveform_generator is not None: - injection_polarizations = self.inject_signal_from_waveform_generator(parameters=parameters, - waveform_generator=waveform_generator) + injection_polarizations = self.inject_signal_from_waveform_generator( + parameters=parameters, waveform_generator=waveform_generator + ) return injection_polarizations def inject_signal_from_waveform_generator(self, parameters, waveform_generator): - """ Inject a signal using a waveform generator and a set of parameters. + """Inject a signal using a waveform generator and a set of parameters. Alternative to `inject_signal` and `inject_signal_from_waveform_polarizations` Parameters @@ -482,14 +567,16 @@ def inject_signal_from_waveform_generator(self, parameters, waveform_generator): The internally generated injection parameters """ - injection_polarizations = \ - waveform_generator.frequency_domain_strain(parameters) - self.inject_signal_from_waveform_polarizations(parameters=parameters, - injection_polarizations=injection_polarizations) + injection_polarizations = waveform_generator.frequency_domain_strain(parameters) + self.inject_signal_from_waveform_polarizations( + parameters=parameters, injection_polarizations=injection_polarizations + ) return injection_polarizations - def inject_signal_from_waveform_polarizations(self, parameters, injection_polarizations): - """ Inject a signal into the detector from a dict of waveform polarizations. + def inject_signal_from_waveform_polarizations( + self, parameters, injection_polarizations + ): + """Inject a signal into the detector from a dict of waveform polarizations. Alternative to `inject_signal` and `inject_signal_from_waveform_generator`. Parameters @@ -501,29 +588,35 @@ def inject_signal_from_waveform_polarizations(self, parameters, injection_polari `waveform_generator.frequency_domain_strain()`. """ - if not self.strain_data.time_within_data(parameters['geocent_time']): + if not self.strain_data.time_within_data(parameters["geocent_time"]): logger.warning( - 'Injecting signal outside segment, start_time={}, merger time={}.' - .format(self.strain_data.start_time, parameters['geocent_time'])) + "Injecting signal outside segment, start_time={}, merger time={}.".format( + self.strain_data.start_time, parameters["geocent_time"] + ) + ) signal_ifo = self.get_detector_response(injection_polarizations, parameters) self.strain_data.frequency_domain_strain += signal_ifo - self.meta_data['optimal_SNR'] = ( - self.optimal_snr_squared(signal=signal_ifo)).real ** 0.5 - self.meta_data['matched_filter_SNR'] = ( - self.matched_filter_snr(signal=signal_ifo)) - self.meta_data['parameters'] = parameters + self.meta_data["optimal_SNR"] = ( + self.optimal_snr_squared(signal=signal_ifo) + ).real ** 0.5 + self.meta_data["matched_filter_SNR"] = self.matched_filter_snr( + signal=signal_ifo + ) + self.meta_data["parameters"] = parameters logger.info("Injected signal in {}:".format(self.name)) - logger.info(" optimal SNR = {:.2f}".format(self.meta_data['optimal_SNR'])) - logger.info(" matched filter SNR = {:.2f}".format(self.meta_data['matched_filter_SNR'])) + logger.info(" optimal SNR = {:.2f}".format(self.meta_data["optimal_SNR"])) + logger.info( + " matched filter SNR = {:.2f}".format(self.meta_data["matched_filter_SNR"]) + ) for key in parameters: - logger.info(' {} = {}'.format(key, parameters[key])) + logger.info(" {} = {}".format(key, parameters[key])) @property def amplitude_spectral_density_array(self): - """ Returns the amplitude spectral density (ASD) given we know a power spectral density (PSD) + """Returns the amplitude spectral density (ASD) given we know a power spectral density (PSD) Returns ======= @@ -536,7 +629,7 @@ def amplitude_spectral_density_array(self): @property def power_spectral_density_array(self): - """ Returns the power spectral density (PSD) + """Returns the power spectral density (PSD) This accounts for whether the data in the interferometer has been windowed. @@ -550,8 +643,10 @@ def power_spectral_density_array(self): ) def unit_vector_along_arm(self, arm): - logger.warning("This method has been moved and will be removed in the future." - "Use Interferometer.geometry.unit_vector_along_arm instead.") + logger.warning( + "This method has been moved and will be removed in the future." + "Use Interferometer.geometry.unit_vector_along_arm instead." + ) return self.geometry.unit_vector_along_arm(arm) def time_delay_from_geocenter(self, ra, dec, time): @@ -588,9 +683,15 @@ def vertex_position_geocentric(self): """ xp = self.array_backend - return gwutils.get_vertex_position_geocentric(xp.asarray([self.geometry.latitude_radians, - self.geometry.longitude_radians, - self.geometry.elevation])) + return gwutils.get_vertex_position_geocentric( + xp.asarray( + [ + self.geometry.latitude_radians, + self.geometry.longitude_radians, + self.geometry.elevation, + ] + ) + ) def optimal_snr_squared(self, signal): """ @@ -608,8 +709,11 @@ def optimal_snr_squared(self, signal): xp = array_module(signal) return gwutils.optimal_snr_squared( signal=signal * mask, - power_spectral_density=xp.where(mask, self.power_spectral_density_array, xp.inf), - duration=self.strain_data.duration) + power_spectral_density=xp.where( + mask, self.power_spectral_density_array, xp.inf + ), + duration=self.strain_data.duration, + ) def inner_product(self, signal): """ @@ -628,8 +732,11 @@ def inner_product(self, signal): return gwutils.noise_weighted_inner_product( aa=signal * mask, bb=self.strain_data.frequency_domain_strain * mask, - power_spectral_density=xp.where(mask, self.power_spectral_density_array, xp.inf), - duration=self.strain_data.duration) + power_spectral_density=xp.where( + mask, self.power_spectral_density_array, xp.inf + ), + duration=self.strain_data.duration, + ) def template_template_inner_product(self, signal_1, signal_2): """A noise weighted inner product between two templates, using this ifo's PSD. @@ -648,8 +755,11 @@ def template_template_inner_product(self, signal_1, signal_2): return gwutils.noise_weighted_inner_product( aa=signal_1[self.strain_data.frequency_mask], bb=signal_2[self.strain_data.frequency_mask], - power_spectral_density=self.power_spectral_density_array[self.strain_data.frequency_mask], - duration=self.strain_data.duration) + power_spectral_density=self.power_spectral_density_array[ + self.strain_data.frequency_mask + ], + duration=self.strain_data.duration, + ) def matched_filter_snr(self, signal): """ @@ -666,11 +776,16 @@ def matched_filter_snr(self, signal): """ return gwutils.matched_filter_snr( signal=signal[self.strain_data.frequency_mask], - frequency_domain_strain=self.strain_data.frequency_domain_strain[self.strain_data.frequency_mask], - power_spectral_density=self.power_spectral_density_array[self.strain_data.frequency_mask], - duration=self.strain_data.duration) + frequency_domain_strain=self.strain_data.frequency_domain_strain[ + self.strain_data.frequency_mask + ], + power_spectral_density=self.power_spectral_density_array[ + self.strain_data.frequency_mask + ], + duration=self.strain_data.duration, + ) - def whiten_frequency_series(self, frequency_series : np.array) -> np.array: + def whiten_frequency_series(self, frequency_series: np.array) -> np.array: """Whitens a frequency series with the noise properties of the detector .. math:: @@ -689,11 +804,12 @@ def whiten_frequency_series(self, frequency_series : np.array) -> np.array: frequency_series : np.array The frequency series, whitened by the ASD """ - return frequency_series / (self.amplitude_spectral_density_array * (self.duration / 4)**0.5) + return frequency_series / ( + self.amplitude_spectral_density_array * (self.duration / 4) ** 0.5 + ) def get_whitened_time_series_from_whitened_frequency_series( - self, - whitened_frequency_series : np.array + self, whitened_frequency_series: np.array ) -> np.array: """Gets the whitened time series from a whitened frequency series. @@ -726,7 +842,8 @@ def get_whitened_time_series_from_whitened_frequency_series( whitened_time_series = ( xp.fft.irfft(whitened_frequency_series) - * self.frequency_mask.sum()**0.5 / frequency_window_factor + * self.frequency_mask.sum() ** 0.5 + / frequency_window_factor ) return whitened_time_series @@ -757,10 +874,12 @@ def whitened_time_domain_strain(self) -> np.array: array_like The whitened data in the time domain """ - return self.get_whitened_time_series_from_whitened_frequency_series(self.whitened_frequency_domain_strain) + return self.get_whitened_time_series_from_whitened_frequency_series( + self.whitened_frequency_domain_strain + ) def save_data(self, outdir, label=None): - """ Creates save files for interferometer data in plain text format. + """Creates save files for interferometer data in plain text format. Saves two files: the frequency domain strain data with three columns [f, real part of h(f), imaginary part of h(f)], and the amplitude spectral density with two columns [f, ASD(f)]. @@ -776,65 +895,92 @@ def save_data(self, outdir, label=None): """ if label is None: - filename_asd = '{}/{}_asd.dat'.format(outdir, self.name) - filename_data = '{}/{}_frequency_domain_data.dat'.format(outdir, self.name) + filename_asd = "{}/{}_asd.dat".format(outdir, self.name) + filename_data = "{}/{}_frequency_domain_data.dat".format(outdir, self.name) else: - filename_asd = '{}/{}_{}_asd.dat'.format(outdir, self.name, label) - filename_data = '{}/{}_{}_frequency_domain_data.dat'.format(outdir, self.name, label) - np.savetxt(filename_data, - np.array( - [self.strain_data.frequency_array, - self.strain_data.frequency_domain_strain.real, - self.strain_data.frequency_domain_strain.imag]).T, - header='f real_h(f) imag_h(f)') - np.savetxt(filename_asd, - np.array( - [self.strain_data.frequency_array, - self.amplitude_spectral_density_array]).T, - header='f h(f)') - - def plot_data(self, signal=None, outdir='.', label=None): + filename_asd = "{}/{}_{}_asd.dat".format(outdir, self.name, label) + filename_data = "{}/{}_{}_frequency_domain_data.dat".format( + outdir, self.name, label + ) + np.savetxt( + filename_data, + np.array( + [ + self.strain_data.frequency_array, + self.strain_data.frequency_domain_strain.real, + self.strain_data.frequency_domain_strain.imag, + ] + ).T, + header="f real_h(f) imag_h(f)", + ) + np.savetxt( + filename_asd, + np.array( + [ + self.strain_data.frequency_array, + self.amplitude_spectral_density_array, + ] + ).T, + header="f h(f)", + ) + + def plot_data(self, signal=None, outdir=".", label=None): import matplotlib.pyplot as plt + if utils.command_line_args.bilby_test_mode: return fig, ax = plt.subplots() df = self.strain_data.frequency_array[1] - self.strain_data.frequency_array[0] asd = gwutils.asd_from_freq_series( - freq_data=self.strain_data.frequency_domain_strain, df=df) - - ax.loglog(self.strain_data.frequency_array[self.strain_data.frequency_mask], - asd[self.strain_data.frequency_mask], - color='C0', label=self.name) - ax.loglog(self.strain_data.frequency_array[self.strain_data.frequency_mask], - self.amplitude_spectral_density_array[self.strain_data.frequency_mask], - color='C1', lw=1.0, label=self.name + ' ASD') + freq_data=self.strain_data.frequency_domain_strain, df=df + ) + + ax.loglog( + self.strain_data.frequency_array[self.strain_data.frequency_mask], + asd[self.strain_data.frequency_mask], + color="C0", + label=self.name, + ) + ax.loglog( + self.strain_data.frequency_array[self.strain_data.frequency_mask], + self.amplitude_spectral_density_array[self.strain_data.frequency_mask], + color="C1", + lw=1.0, + label=self.name + " ASD", + ) if signal is not None: - signal_asd = gwutils.asd_from_freq_series( - freq_data=signal, df=df) + signal_asd = gwutils.asd_from_freq_series(freq_data=signal, df=df) - ax.loglog(self.strain_data.frequency_array[self.strain_data.frequency_mask], - signal_asd[self.strain_data.frequency_mask], - color='C2', - label='Signal') + ax.loglog( + self.strain_data.frequency_array[self.strain_data.frequency_mask], + signal_asd[self.strain_data.frequency_mask], + color="C2", + label="Signal", + ) ax.grid(True) - ax.set_ylabel(r'Strain [strain/$\sqrt{\rm Hz}$]') - ax.set_xlabel(r'Frequency [Hz]') - ax.legend(loc='best') + ax.set_ylabel(r"Strain [strain/$\sqrt{\rm Hz}$]") + ax.set_xlabel(r"Frequency [Hz]") + ax.legend(loc="best") fig.tight_layout() if label is None: - fig.savefig( - '{}/{}_frequency_domain_data.png'.format(outdir, self.name)) + fig.savefig("{}/{}_frequency_domain_data.png".format(outdir, self.name)) else: fig.savefig( - '{}/{}_{}_frequency_domain_data.png'.format( - outdir, self.name, label)) + "{}/{}_{}_frequency_domain_data.png".format(outdir, self.name, label) + ) plt.close(fig) def plot_time_domain_data( - self, outdir='.', label=None, bandpass_frequencies=(50, 250), - notches=None, start_end=None, t0=None): - """ Plots the strain data in the time domain + self, + outdir=".", + label=None, + bandpass_frequencies=(50, 250), + notches=None, + start_end=None, + t0=None, + ): + """Plots the strain data in the time domain Parameters ========== @@ -860,16 +1006,20 @@ def plot_time_domain_data( if notches is None: notches = list() timeseries = TimeSeries( - data=self.strain_data.time_domain_strain, times=self.strain_data.time_array) + data=self.strain_data.time_domain_strain, times=self.strain_data.time_array + ) zpks = [] if bandpass_frequencies is not None: - zpks.append(bandpass( - bandpass_frequencies[0], bandpass_frequencies[1], - self.strain_data.sampling_frequency)) + zpks.append( + bandpass( + bandpass_frequencies[0], + bandpass_frequencies[1], + self.strain_data.sampling_frequency, + ) + ) if notches is not None: for line in notches: - zpks.append(notch( - line, self.strain_data.sampling_frequency)) + zpks.append(notch(line, self.strain_data.sampling_frequency)) if len(zpks) > 0: zpk = concatenate_zpks(*zpks) strain = timeseries.filter(zpk, filtfilt=False) @@ -880,14 +1030,14 @@ def plot_time_domain_data( if t0: x = self.strain_data.time_array - t0 - xlabel = 'GPS time [s] - {}'.format(t0) + xlabel = "GPS time [s] - {}".format(t0) else: x = self.strain_data.time_array - xlabel = 'GPS time [s]' + xlabel = "GPS time [s]" ax.plot(x, strain) ax.set_xlabel(xlabel) - ax.set_ylabel('Strain') + ax.set_ylabel("Strain") if start_end is not None: ax.set_xlim(*start_end) @@ -895,16 +1045,16 @@ def plot_time_domain_data( fig.tight_layout() if label is None: - fig.savefig( - '{}/{}_time_domain_data.png'.format(outdir, self.name)) + fig.savefig("{}/{}_time_domain_data.png".format(outdir, self.name)) else: fig.savefig( - '{}/{}_{}_time_domain_data.png'.format(outdir, self.name, label)) + "{}/{}_{}_time_domain_data.png".format(outdir, self.name, label) + ) plt.close(fig) @staticmethod def _filename_from_outdir_label_extension(outdir, label, extension="h5"): - return os.path.join(outdir, label + f'.{extension}') + return os.path.join(outdir, label + f".{extension}") _save_ifo_docstring = """ Save the object to a {format} file @@ -927,22 +1077,25 @@ def _filename_from_outdir_label_extension(outdir, label, extension="h5"): """ - @docstring(_save_ifo_docstring.format( - format="pickle", extra=".. versionadded:: 1.1.0" - )) + @docstring( + _save_ifo_docstring.format(format="pickle", extra=".. versionadded:: 1.1.0") + ) def to_pickle(self, outdir="outdir", label=None): utils.check_directory_exists_and_if_not_mkdir(outdir) - filename = self._filename_from_outdir_label_extension(outdir, label, extension="pkl") + filename = self._filename_from_outdir_label_extension( + outdir, label, extension="pkl" + ) safe_file_dump(self, filename, "dill") @classmethod @docstring(_load_docstring.format(format="pickle")) def from_pickle(cls, filename=None): import dill + with open(filename, "rb") as ff: res = dill.load(ff) if res.__class__ != cls: - raise TypeError('The loaded object is not an Interferometer') + raise TypeError("The loaded object is not an Interferometer") return res @property diff --git a/bilby/gw/detector/networks.py b/bilby/gw/detector/networks.py index 1ef5d6179..013ad8358 100644 --- a/bilby/gw/detector/networks.py +++ b/bilby/gw/detector/networks.py @@ -145,7 +145,7 @@ def inject_signal( waveform_generator=None, raise_error=True, ): - """ Inject a signal into noise in each of the three detectors. + """Inject a signal into noise in each of the three detectors. Parameters ========== @@ -218,8 +218,13 @@ def plot_data(self, signal=None, outdir=".", label=None): interferometer.plot_data(signal=signal, outdir=outdir, label=label) def plot_time_domain_data( - self, outdir=".", label=None, bandpass_frequencies=(50, 250), - notches=None, start_end=None, t0=None + self, + outdir=".", + label=None, + bandpass_frequencies=(50, 250), + notches=None, + start_end=None, + t0=None, ): """Plots the strain data in the time domain for each of the interfeormeters @@ -251,7 +256,7 @@ def plot_time_domain_data( bandpass_frequencies=bandpass_frequencies, notches=notches, start_end=start_end, - t0=t0 + t0=t0, ) @property @@ -382,7 +387,6 @@ def __init__( maximum_frequency = [maximum_frequency] * 3 for ii in range(3): - self.append( Interferometer( "{}{}".format(name, ii + 1), @@ -400,8 +404,15 @@ def __init__( ) ) - latitude, longitude, elevation, xarm_azimuth, \ - yarm_azimuth, xarm_tilt, yarm_tilt = self._get_next_vertex_parameters(ii) + ( + latitude, + longitude, + elevation, + xarm_azimuth, + yarm_azimuth, + xarm_tilt, + yarm_tilt, + ) = self._get_next_vertex_parameters(ii) @xp_wrap def _get_next_vertex_parameters(self, current_index, *, xp): @@ -441,10 +452,18 @@ def _get_next_vertex_parameters(self, current_index, *, xp): unit_vector_x = current_ifo.geometry.unit_vector_along_arm("x") unit_vector_y = current_ifo.geometry.unit_vector_along_arm("y") - vertex_geocentric = get_vertex_position_geocentric(xp.asarray([current_ifo.latitude_radians, - current_ifo.longitude_radians, - current_ifo.elevation])) - next_vertex_geocentric = vertex_geocentric + current_ifo.length * 1000 * unit_vector_x + vertex_geocentric = get_vertex_position_geocentric( + xp.asarray( + [ + current_ifo.latitude_radians, + current_ifo.longitude_radians, + current_ifo.elevation, + ] + ) + ) + next_vertex_geocentric = ( + vertex_geocentric + current_ifo.length * 1000 * unit_vector_x + ) next_vertex_ellipsoid = get_vertex_position_ellipsoid(next_vertex_geocentric) next_latitude_rad, next_longitude_rad, next_elevation = next_vertex_ellipsoid @@ -455,29 +474,48 @@ def _get_next_vertex_parameters(self, current_index, *, xp): next_unit_vector_x = rotation.apply(unit_vector_x) next_unit_vector_y = rotation.apply(unit_vector_y) - next_local_normal_vector = xp.asarray([xp.cos(next_latitude_rad) * xp.cos(next_longitude_rad), - xp.cos(next_latitude_rad) * xp.sin(next_longitude_rad), - xp.sin(next_latitude_rad)]) - next_local_north_vector = xp.asarray([-xp.sin(next_latitude_rad) * xp.cos(next_longitude_rad), - -xp.sin(next_latitude_rad) * xp.sin(next_longitude_rad), - xp.cos(next_latitude_rad)]) - next_local_east_vector = xp.asarray([-xp.sin(next_longitude_rad), - xp.cos(next_longitude_rad), 0]) - - next_xarm_tilt_rad = xp.arcsin(xp.dot(next_unit_vector_x, next_local_normal_vector)) - next_yarm_tilt_rad = xp.arcsin(xp.dot(next_unit_vector_y, next_local_normal_vector)) - next_xarm_azimuth_rad = xp.arctan2(xp.dot(next_unit_vector_x, next_local_north_vector), - xp.dot(next_unit_vector_x, next_local_east_vector)) - next_yarm_azimuth_rad = xp.arctan2(xp.dot(next_unit_vector_y, next_local_north_vector), - xp.dot(next_unit_vector_y, next_local_east_vector)) - - return (xp.rad2deg(next_latitude_rad), - xp.rad2deg(next_longitude_rad), - next_elevation, - xp.rad2deg(next_xarm_azimuth_rad), - xp.rad2deg(next_yarm_azimuth_rad), - next_xarm_tilt_rad, - next_yarm_tilt_rad) + next_local_normal_vector = xp.asarray( + [ + xp.cos(next_latitude_rad) * xp.cos(next_longitude_rad), + xp.cos(next_latitude_rad) * xp.sin(next_longitude_rad), + xp.sin(next_latitude_rad), + ] + ) + next_local_north_vector = xp.asarray( + [ + -xp.sin(next_latitude_rad) * xp.cos(next_longitude_rad), + -xp.sin(next_latitude_rad) * xp.sin(next_longitude_rad), + xp.cos(next_latitude_rad), + ] + ) + next_local_east_vector = xp.asarray( + [-xp.sin(next_longitude_rad), xp.cos(next_longitude_rad), 0] + ) + + next_xarm_tilt_rad = xp.arcsin( + xp.dot(next_unit_vector_x, next_local_normal_vector) + ) + next_yarm_tilt_rad = xp.arcsin( + xp.dot(next_unit_vector_y, next_local_normal_vector) + ) + next_xarm_azimuth_rad = xp.arctan2( + xp.dot(next_unit_vector_x, next_local_north_vector), + xp.dot(next_unit_vector_x, next_local_east_vector), + ) + next_yarm_azimuth_rad = xp.arctan2( + xp.dot(next_unit_vector_y, next_local_north_vector), + xp.dot(next_unit_vector_y, next_local_east_vector), + ) + + return ( + xp.rad2deg(next_latitude_rad), + xp.rad2deg(next_longitude_rad), + next_elevation, + xp.rad2deg(next_xarm_azimuth_rad), + xp.rad2deg(next_yarm_azimuth_rad), + next_xarm_tilt_rad, + next_yarm_tilt_rad, + ) _LEGACY_DETECTOR_NAMES = { diff --git a/bilby/gw/detector/psd.py b/bilby/gw/detector/psd.py index 864dbd380..b1d117164 100644 --- a/bilby/gw/detector/psd.py +++ b/bilby/gw/detector/psd.py @@ -12,9 +12,14 @@ class PowerSpectralDensity(object): - - def __init__(self, frequency_array=None, psd_array=None, asd_array=None, - psd_file=None, asd_file=None): + def __init__( + self, + frequency_array=None, + psd_array=None, + asd_array=None, + psd_file=None, + asd_file=None, + ): """ Instantiate a new PowerSpectralDensity object. @@ -50,8 +55,7 @@ def __init__(self, frequency_array=None, psd_array=None, asd_array=None, Interpolated function of the PSD """ - self._cache = dict( - frequency_array=np.array([]), psd_array=None, asd_array=None) + self._cache = dict(frequency_array=np.array([]), psd_array=None, asd_array=None) self.frequency_array = np.array(frequency_array) if psd_array is not None: self.psd_array = psd_array @@ -62,30 +66,37 @@ def __init__(self, frequency_array=None, psd_array=None, asd_array=None, def _update_cache(self, frequency_array): psd_array = self.power_spectral_density_interpolated(frequency_array) - self._cache['psd_array'] = psd_array - self._cache['asd_array'] = psd_array**0.5 - self._cache['frequency_array'] = frequency_array + self._cache["psd_array"] = psd_array + self._cache["asd_array"] = psd_array**0.5 + self._cache["frequency_array"] = frequency_array def __eq__(self, other): - if self.psd_file == other.psd_file \ - and self.asd_file == other.asd_file \ - and np.array_equal(self.frequency_array, other.frequency_array) \ - and np.array_equal(self.psd_array, other.psd_array) \ - and np.array_equal(self.asd_array, other.asd_array): + if ( + self.psd_file == other.psd_file + and self.asd_file == other.asd_file + and np.array_equal(self.frequency_array, other.frequency_array) + and np.array_equal(self.psd_array, other.psd_array) + and np.array_equal(self.asd_array, other.asd_array) + ): return True return False def __repr__(self): if self.asd_file is not None or self.psd_file is not None: - return self.__class__.__name__ + '(psd_file=\'{}\', asd_file=\'{}\')' \ - .format(self.psd_file, self.asd_file) + return self.__class__.__name__ + "(psd_file='{}', asd_file='{}')".format( + self.psd_file, self.asd_file + ) else: - return self.__class__.__name__ + '(frequency_array={}, psd_array={}, asd_array={})' \ - .format(self.frequency_array, self.psd_array, self.asd_array) + return ( + self.__class__.__name__ + + "(frequency_array={}, psd_array={}, asd_array={})".format( + self.frequency_array, self.psd_array, self.asd_array + ) + ) @staticmethod def from_amplitude_spectral_density_file(asd_file): - """ Set the amplitude spectral density from a given file + """Set the amplitude spectral density from a given file Parameters ========== @@ -97,7 +108,7 @@ def from_amplitude_spectral_density_file(asd_file): @staticmethod def from_power_spectral_density_file(psd_file): - """ Set the power spectral density from a given file + """Set the power spectral density from a given file Parameters ========== @@ -108,11 +119,20 @@ def from_power_spectral_density_file(psd_file): return PowerSpectralDensity(psd_file=psd_file) @staticmethod - def from_frame_file(frame_file, psd_start_time, psd_duration, - fft_length=4, sampling_frequency=4096, roll_off=0.2, - overlap=0, channel=None, name=None, outdir=None, - analysis_segment_start_time=None): - """ Generate power spectral density from a frame file + def from_frame_file( + frame_file, + psd_start_time, + psd_duration, + fft_length=4, + sampling_frequency=4096, + roll_off=0.2, + overlap=0, + channel=None, + name=None, + outdir=None, + analysis_segment_start_time=None, + ): + """Generate power spectral density from a frame file Parameters ========== @@ -143,19 +163,37 @@ def from_frame_file(frame_file, psd_start_time, psd_duration, """ strain = InterferometerStrainData(roll_off=roll_off) strain.set_from_frame_file( - frame_file, start_time=psd_start_time, duration=psd_duration, - channel=channel, sampling_frequency=sampling_frequency) + frame_file, + start_time=psd_start_time, + duration=psd_duration, + channel=channel, + sampling_frequency=sampling_frequency, + ) frequency_array, psd_array = strain.create_power_spectral_density( - fft_length=fft_length, name=name, outdir=outdir, overlap=overlap, - analysis_segment_start_time=analysis_segment_start_time) - return PowerSpectralDensity(frequency_array=frequency_array, psd_array=psd_array) + fft_length=fft_length, + name=name, + outdir=outdir, + overlap=overlap, + analysis_segment_start_time=analysis_segment_start_time, + ) + return PowerSpectralDensity( + frequency_array=frequency_array, psd_array=psd_array + ) @staticmethod - def from_channel_name(channel, psd_start_time, psd_duration, - fft_length=4, sampling_frequency=4096, roll_off=0.2, - overlap=0, name=None, outdir=None, - analysis_segment_start_time=None): - """ Generate power spectral density from a given channel name + def from_channel_name( + channel, + psd_start_time, + psd_duration, + fft_length=4, + sampling_frequency=4096, + roll_off=0.2, + overlap=0, + name=None, + outdir=None, + analysis_segment_start_time=None, + ): + """Generate power spectral density from a given channel name by loading data using `strain_data.set_from_channel_name` Parameters @@ -186,26 +224,42 @@ def from_channel_name(channel, psd_start_time, psd_duration, """ strain = InterferometerStrainData(roll_off=roll_off) strain.set_from_channel_name( - channel, duration=psd_duration, start_time=psd_start_time, - sampling_frequency=sampling_frequency) + channel, + duration=psd_duration, + start_time=psd_start_time, + sampling_frequency=sampling_frequency, + ) frequency_array, psd_array = strain.create_power_spectral_density( - fft_length=fft_length, name=name, outdir=outdir, overlap=overlap, - analysis_segment_start_time=analysis_segment_start_time) - return PowerSpectralDensity(frequency_array=frequency_array, psd_array=psd_array) + fft_length=fft_length, + name=name, + outdir=outdir, + overlap=overlap, + analysis_segment_start_time=analysis_segment_start_time, + ) + return PowerSpectralDensity( + frequency_array=frequency_array, psd_array=psd_array + ) @staticmethod def from_amplitude_spectral_density_array(frequency_array, asd_array): - return PowerSpectralDensity(frequency_array=frequency_array, asd_array=asd_array) + return PowerSpectralDensity( + frequency_array=frequency_array, asd_array=asd_array + ) @staticmethod def from_power_spectral_density_array(frequency_array, psd_array): - return PowerSpectralDensity(frequency_array=frequency_array, psd_array=psd_array) + return PowerSpectralDensity( + frequency_array=frequency_array, psd_array=psd_array + ) @staticmethod def from_aligo(): - logger.info("No power spectral density provided, using aLIGO," - "zero detuning, high power.") - return PowerSpectralDensity.from_power_spectral_density_file(psd_file='aLIGO_ZERO_DET_high_P_psd.txt') + logger.info( + "No power spectral density provided, using aLIGO,zero detuning, high power." + ) + return PowerSpectralDensity.from_power_spectral_density_file( + psd_file="aLIGO_ZERO_DET_high_P_psd.txt" + ) @property def psd_array(self): @@ -215,7 +269,7 @@ def psd_array(self): def psd_array(self, psd_array): self._check_frequency_array_matches_density_array(psd_array) self._psd_array = np.array(psd_array) - self._asd_array = psd_array ** 0.5 + self._asd_array = psd_array**0.5 self._interpolate_power_spectral_density() @property @@ -226,38 +280,40 @@ def asd_array(self): def asd_array(self, asd_array): self._check_frequency_array_matches_density_array(asd_array) self._asd_array = np.array(asd_array) - self._psd_array = asd_array ** 2 + self._psd_array = asd_array**2 self._interpolate_power_spectral_density() def _check_frequency_array_matches_density_array(self, density_array): if len(self.frequency_array) != len(density_array): - raise ValueError('Provided spectral density does not match frequency array. Not updating.\n' - 'Length spectral density {}\n Length frequency array {}\n' - .format(density_array, self.frequency_array)) + raise ValueError( + "Provided spectral density does not match frequency array. Not updating.\n" + "Length spectral density {}\n Length frequency array {}\n".format( + density_array, self.frequency_array + ) + ) def _interpolate_power_spectral_density(self): """Interpolate the loaded power spectral density so it can be resampled - for arbitrary frequency arrays. + for arbitrary frequency arrays. """ - self._power_spectral_density_interpolated = interp1d(self.frequency_array, - self.psd_array, - bounds_error=False, - fill_value=np.inf) + self._power_spectral_density_interpolated = interp1d( + self.frequency_array, self.psd_array, bounds_error=False, fill_value=np.inf + ) self._update_cache(self.frequency_array) def get_power_spectral_density_array(self, frequency_array): if aac.is_jax_array(frequency_array): return self.power_spectral_density_interpolated(frequency_array) - if not np.array_equal(frequency_array, self._cache['frequency_array']): + if not np.array_equal(frequency_array, self._cache["frequency_array"]): self._update_cache(frequency_array=frequency_array) - return self._cache['psd_array'] + return self._cache["psd_array"] def get_amplitude_spectral_density_array(self, frequency_array): if aac.is_jax_array(frequency_array): - return self.power_spectral_density_interpolated(frequency_array)**0.5 - if not np.array_equal(frequency_array, self._cache['frequency_array']): + return self.power_spectral_density_interpolated(frequency_array) ** 0.5 + if not np.array_equal(frequency_array, self._cache["frequency_array"]): self._update_cache(frequency_array=frequency_array) - return self._cache['asd_array'] + return self._cache["asd_array"] @property def power_spectral_density_interpolated(self): @@ -279,8 +335,14 @@ def _check_file_was_asd_file(self): if min(self.asd_array) < 1e-30: logger.warning("You specified an amplitude spectral density file.") logger.warning("{} WARNING {}".format("*" * 30, "*" * 30)) - logger.warning("The minimum of the provided curve is {:.2e}.".format(min(self.asd_array))) - logger.warning("You may have intended to provide this as a power spectral density.") + logger.warning( + "The minimum of the provided curve is {:.2e}.".format( + min(self.asd_array) + ) + ) + logger.warning( + "You may have intended to provide this as a power spectral density." + ) @property def psd_file(self): @@ -298,8 +360,14 @@ def _check_file_was_psd_file(self): if min(self.psd_array) > 1e-30: logger.warning("You specified a power spectral density file.") logger.warning("{} WARNING {}".format("*" * 30, "*" * 30)) - logger.warning("The minimum of the provided curve is {:.2e}.".format(min(self.psd_array))) - logger.warning("You may have intended to provide this as an amplitude spectral density.") + logger.warning( + "The minimum of the provided curve is {:.2e}.".format( + min(self.psd_array) + ) + ) + logger.warning( + "You may have intended to provide this as an amplitude spectral density." + ) @staticmethod def _validate_file_name(file): @@ -330,23 +398,26 @@ def _validate_file_name(file): logger.debug("PSD file {} exists".format(file)) return file else: - file_in_default_directory = ( - os.path.join(os.path.dirname(__file__), 'noise_curves', file)) + file_in_default_directory = os.path.join( + os.path.dirname(__file__), "noise_curves", file + ) if os.path.isfile(file_in_default_directory): logger.debug("PSD file {} exists in default dir.".format(file)) return file_in_default_directory else: raise ValueError( - "Unable to locate PSD file {} locally or in the default dir" - .format(file)) + "Unable to locate PSD file {} locally or in the default dir".format( + file + ) + ) return file def _import_amplitude_spectral_density(self): - """ Automagically load an amplitude spectral density curve """ + """Automagically load an amplitude spectral density curve""" self.frequency_array, self.asd_array = np.genfromtxt(self.asd_file).T def _import_power_spectral_density(self): - """ Automagically load a power spectral density curve """ + """Automagically load a power spectral density curve""" self.frequency_array, self.psd_array = np.genfromtxt(self.psd_file).T def get_noise_realisation(self, number_of_samples, duration, *, random_state=None): @@ -366,13 +437,17 @@ def get_noise_realisation(self, number_of_samples, duration, *, random_state=Non array_like: frequencies related to the frequency domain strain """ - white_noise, frequencies = utils.safe_white_noise(number_of_samples, duration, random_state=random_state) + white_noise, frequencies = utils.safe_white_noise( + number_of_samples, duration, random_state=random_state + ) with np.errstate(invalid="ignore"): - frequency_domain_strain = self._power_spectral_density_interpolated(frequencies) ** 0.5 * white_noise + frequency_domain_strain = ( + self._power_spectral_density_interpolated(frequencies) ** 0.5 + * white_noise + ) xp = aac.array_namespace(frequency_domain_strain) - out_of_bounds = ( - (frequencies < self.frequency_array.min()) - | (frequencies > self.frequency_array.max()) + out_of_bounds = (frequencies < self.frequency_array.min()) | ( + frequencies > self.frequency_array.max() ) frequency_domain_strain = xpx.at(frequency_domain_strain, out_of_bounds).set(0j) return xp.nan_to_num(frequency_domain_strain), xp.asarray(frequencies) @@ -383,7 +458,7 @@ def array_backend(self): @array_backend.setter def array_backend(self, xp): - """ Set the array backend for the cached arrays + """Set the array backend for the cached arrays Parameters ========== @@ -394,7 +469,7 @@ def array_backend(self, xp): self.frequency_array = xp.asarray(self.frequency_array) self._asd_array = xp.asarray(self._asd_array) self._psd_array = xp.asarray(self._psd_array) - self._cache['frequency_array'] = xp.asarray(self._cache['frequency_array']) - self._cache['psd_array'] = xp.asarray(self._cache['psd_array']) - self._cache['asd_array'] = xp.asarray(self._cache['asd_array']) + self._cache["frequency_array"] = xp.asarray(self._cache["frequency_array"]) + self._cache["psd_array"] = xp.asarray(self._cache["psd_array"]) + self._cache["asd_array"] = xp.asarray(self._cache["asd_array"]) self._interpolate_power_spectral_density() diff --git a/bilby/gw/detector/strain_data.py b/bilby/gw/detector/strain_data.py index fd63bbaa7..d545aa308 100644 --- a/bilby/gw/detector/strain_data.py +++ b/bilby/gw/detector/strain_data.py @@ -9,17 +9,24 @@ class InterferometerStrainData(object): - """ Strain data for an interferometer """ - - duration = PropertyAccessor('_times_and_frequencies', 'duration') - sampling_frequency = PropertyAccessor('_times_and_frequencies', 'sampling_frequency') - start_time = PropertyAccessor('_times_and_frequencies', 'start_time') - frequency_array = PropertyAccessor('_times_and_frequencies', 'frequency_array') - time_array = PropertyAccessor('_times_and_frequencies', 'time_array') - - def __init__(self, minimum_frequency=0, maximum_frequency=np.inf, - roll_off=0.2, notch_list=None): - """ Initiate an InterferometerStrainData object + """Strain data for an interferometer""" + + duration = PropertyAccessor("_times_and_frequencies", "duration") + sampling_frequency = PropertyAccessor( + "_times_and_frequencies", "sampling_frequency" + ) + start_time = PropertyAccessor("_times_and_frequencies", "start_time") + frequency_array = PropertyAccessor("_times_and_frequencies", "frequency_array") + time_array = PropertyAccessor("_times_and_frequencies", "time_array") + + def __init__( + self, + minimum_frequency=0, + maximum_frequency=np.inf, + roll_off=0.2, + notch_list=None, + ): + """Initiate an InterferometerStrainData object The initialised object contains no data, this should be added using one of the `set_from..` methods. @@ -53,22 +60,26 @@ def __init__(self, minimum_frequency=0, maximum_frequency=np.inf, self._channel = None def __eq__(self, other): - if self.minimum_frequency == other.minimum_frequency \ - and self.maximum_frequency == other.maximum_frequency \ - and self.roll_off == other.roll_off \ - and self.window_factor == other.window_factor \ - and self.sampling_frequency == other.sampling_frequency \ - and self.duration == other.duration \ - and self.start_time == other.start_time \ - and np.array_equal(self.time_array, other.time_array) \ - and np.array_equal(self.frequency_array, other.frequency_array) \ - and np.array_equal(self.frequency_domain_strain, other.frequency_domain_strain) \ - and np.array_equal(self.time_domain_strain, other.time_domain_strain): + if ( + self.minimum_frequency == other.minimum_frequency + and self.maximum_frequency == other.maximum_frequency + and self.roll_off == other.roll_off + and self.window_factor == other.window_factor + and self.sampling_frequency == other.sampling_frequency + and self.duration == other.duration + and self.start_time == other.start_time + and np.array_equal(self.time_array, other.time_array) + and np.array_equal(self.frequency_array, other.frequency_array) + and np.array_equal( + self.frequency_domain_strain, other.frequency_domain_strain + ) + and np.array_equal(self.time_domain_strain, other.time_domain_strain) + ): return True return False def time_within_data(self, time): - """ Check if time is within the data span + """Check if time is within the data span Parameters ========== @@ -101,7 +112,7 @@ def minimum_frequency(self, minimum_frequency): @property def maximum_frequency(self): - """ Force the maximum frequency be less than the Nyquist frequency """ + """Force the maximum frequency be less than the Nyquist frequency""" if self.sampling_frequency is not None: xp = array_module(self._maximum_frequency) self._maximum_frequency = xp.minimum( @@ -120,7 +131,7 @@ def notch_list(self): @notch_list.setter def notch_list(self, notch_list): - """ Set the notch_list + """Set the notch_list Parameters ========== @@ -141,7 +152,7 @@ def notch_list(self, notch_list): @property def frequency_mask(self): - """ Masking array for limiting the frequency band. + """Masking array for limiting the frequency band. Returns ======= @@ -150,8 +161,9 @@ def frequency_mask(self): """ if not self._frequency_mask_updated: frequency_array = self._times_and_frequencies.frequency_array - mask = ((frequency_array >= self.minimum_frequency) & - (frequency_array <= self.maximum_frequency)) + mask = (frequency_array >= self.minimum_frequency) & ( + frequency_array <= self.maximum_frequency + ) for notch in self.notch_list: mask[notch.get_idxs(frequency_array)] = False self._frequency_mask = mask @@ -188,22 +200,24 @@ def time_domain_window(self, roll_off=None, alpha=None): Window function over time array """ from scipy.signal.windows import tukey + if roll_off is not None: self.roll_off = roll_off elif alpha is not None: self.roll_off = alpha * self.duration / 2 window = tukey(len(self._time_domain_strain), alpha=self.alpha) - self.window_factor = np.mean(window ** 2) + self.window_factor = np.mean(window**2) return window @property def time_domain_strain(self): - """ The time domain strain, in units of strain """ + """The time domain strain, in units of strain""" if self._time_domain_strain is not None: return self._time_domain_strain elif self._frequency_domain_strain is not None: self._time_domain_strain = utils.infft( - self.frequency_domain_strain, self.sampling_frequency) + self.frequency_domain_strain, self.sampling_frequency + ) return self._time_domain_strain else: @@ -211,7 +225,7 @@ def time_domain_strain(self): @property def frequency_domain_strain(self): - """ Returns the frequency domain strain + """Returns the frequency domain strain This is the frequency domain strain normalised to units of strain / Hz, obtained by a one-sided Fourier transform of the @@ -220,14 +234,19 @@ def frequency_domain_strain(self): if self._frequency_domain_strain is not None: return self._frequency_domain_strain * self.frequency_mask elif self._time_domain_strain is not None: - logger.debug("Generating frequency domain strain from given time " - "domain strain.") - logger.debug("Applying a tukey window with alpha={}, roll off={}".format( - self.alpha, self.roll_off)) + logger.debug( + "Generating frequency domain strain from given time domain strain." + ) + logger.debug( + "Applying a tukey window with alpha={}, roll off={}".format( + self.alpha, self.roll_off + ) + ) # self.low_pass_filter() window = self.time_domain_window() self._frequency_domain_strain, self.frequency_array = utils.nfft( - self._time_domain_strain * window, self.sampling_frequency) + self._time_domain_strain * window, self.sampling_frequency + ) return self._frequency_domain_strain * self.frequency_mask else: raise ValueError("frequency domain strain data not yet set") @@ -235,7 +254,9 @@ def frequency_domain_strain(self): @frequency_domain_strain.setter def frequency_domain_strain(self, frequency_domain_strain): if not len(self.frequency_array) == len(frequency_domain_strain): - raise ValueError("The frequency_array and the set strain have different lengths") + raise ValueError( + "The frequency_array and the set strain have different lengths" + ) self._frequency_domain_strain = frequency_domain_strain self._time_domain_strain = None @@ -249,8 +270,10 @@ def to_gwpy_timeseries(self): raise ModuleNotFoundError("Cannot output strain data as gwpy TimeSeries") return TimeSeries( - self.time_domain_strain, sample_rate=self.sampling_frequency, - t0=self.start_time, channel=self.channel + self.time_domain_strain, + sample_rate=self.sampling_frequency, + t0=self.start_time, + channel=self.channel, ) def to_pycbc_timeseries(self): @@ -265,8 +288,9 @@ def to_pycbc_timeseries(self): raise ModuleNotFoundError("Cannot output strain data as PyCBC TimeSeries") return TimeSeries( - self.time_domain_strain, delta_t=(1. / self.sampling_frequency), - epoch=LIGOTimeGPS(self.start_time) + self.time_domain_strain, + delta_t=(1.0 / self.sampling_frequency), + epoch=LIGOTimeGPS(self.start_time), ) def to_lal_timeseries(self): @@ -279,8 +303,12 @@ def to_lal_timeseries(self): raise ModuleNotFoundError("Cannot output strain data as PyCBC TimeSeries") lal_data = CreateREAL8TimeSeries( - "", LIGOTimeGPS(self.start_time), 0, 1 / self.sampling_frequency, - SecondUnit, len(self.time_domain_strain) + "", + LIGOTimeGPS(self.start_time), + 0, + 1 / self.sampling_frequency, + SecondUnit, + len(self.time_domain_strain), ) lal_data.data.data[:] = self.time_domain_strain @@ -293,13 +321,15 @@ def to_gwpy_frequencyseries(self): try: from gwpy.frequencyseries import FrequencySeries except ModuleNotFoundError: - raise ModuleNotFoundError("Cannot output strain data as gwpy FrequencySeries") + raise ModuleNotFoundError( + "Cannot output strain data as gwpy FrequencySeries" + ) return FrequencySeries( self.frequency_domain_strain, frequencies=self.frequency_array, epoch=self.start_time, - channel=self.channel + channel=self.channel, ) def to_pycbc_frequencyseries(self): @@ -316,7 +346,7 @@ def to_pycbc_frequencyseries(self): return FrequencySeries( self.frequency_domain_strain, delta_f=1 / self.duration, - epoch=LIGOTimeGPS(self.start_time) + epoch=LIGOTimeGPS(self.start_time), ) def to_lal_frequencyseries(self): @@ -334,39 +364,49 @@ def to_lal_frequencyseries(self): self.frequency_array[0], 1 / self.duration, SecondUnit, - len(self.frequency_domain_strain) + len(self.frequency_domain_strain), ) lal_data.data.data[:] = self.frequency_domain_strain return lal_data def low_pass_filter(self, filter_freq=None): - """ Low pass filter the data """ + """Low pass filter the data""" from gwpy.signal.filter_design import lowpass from gwpy.timeseries import TimeSeries if filter_freq is None: - logger.debug( - "Setting low pass filter_freq using given maximum frequency") + logger.debug("Setting low pass filter_freq using given maximum frequency") filter_freq = self.maximum_frequency if 2 * filter_freq >= self.sampling_frequency: logger.info( "Low pass filter frequency of {}Hz requested, this is equal" - " or greater than the Nyquist frequency so no filter applied" - .format(filter_freq)) + " or greater than the Nyquist frequency so no filter applied".format( + filter_freq + ) + ) return - logger.debug("Applying low pass filter with filter frequency {}".format(filter_freq)) + logger.debug( + "Applying low pass filter with filter frequency {}".format(filter_freq) + ) bp = lowpass(filter_freq, self.sampling_frequency) - strain = TimeSeries(self.time_domain_strain, sample_rate=self.sampling_frequency) + strain = TimeSeries( + self.time_domain_strain, sample_rate=self.sampling_frequency + ) strain = strain.filter(bp, filtfilt=True) self._time_domain_strain = strain.value def create_power_spectral_density( - self, fft_length, overlap=0, name='unknown', outdir=None, - analysis_segment_start_time=None): - """ Use the time domain strain to generate a power spectral density + self, + fft_length, + overlap=0, + name="unknown", + outdir=None, + analysis_segment_start_time=None, + ): + """Use the time domain strain to generate a power spectral density This create a Tukey-windowed power spectral density and writes it to a PSD file. @@ -399,13 +439,16 @@ def create_power_spectral_density( if analysis_segment_start_time is not None: analysis_segment_end_time = analysis_segment_start_time + fft_length - inside = (analysis_segment_start_time > self.time_array[0] + - analysis_segment_end_time < self.time_array[-1]) + inside = ( + analysis_segment_start_time + > self.time_array[0] + analysis_segment_end_time + < self.time_array[-1] + ) if inside: logger.info("Removing analysis segment data from the PSD data") - idxs = ( - (self.time_array < analysis_segment_start_time) + - (self.time_array > analysis_segment_end_time)) + idxs = (self.time_array < analysis_segment_start_time) + ( + self.time_array > analysis_segment_end_time + ) data = data[idxs] # WARNING this line can cause issues if the data is non-contiguous @@ -413,63 +456,87 @@ def create_power_spectral_density( psd_alpha = 2 * self.roll_off / fft_length logger.info( "Tukey window PSD data with alpha={}, roll off={}".format( - psd_alpha, self.roll_off)) + psd_alpha, self.roll_off + ) + ) psd = strain.psd( - fftlength=fft_length, overlap=overlap, window=('tukey', psd_alpha)) + fftlength=fft_length, overlap=overlap, window=("tukey", psd_alpha) + ) if outdir: - psd_file = '{}/{}_PSD_{}_{}.txt'.format(outdir, name, self.start_time, self.duration) - with open('{}'.format(psd_file), 'w+') as opened_file: + psd_file = "{}/{}_PSD_{}_{}.txt".format( + outdir, name, self.start_time, self.duration + ) + with open("{}".format(psd_file), "w+") as opened_file: for f, p in zip(psd.frequencies.value, psd.value): - opened_file.write('{} {}\n'.format(f, p)) + opened_file.write("{} {}\n".format(f, p)) return psd.frequencies.value, psd.value def _infer_time_domain_dependence( - self, start_time, sampling_frequency, duration, time_array): - """ Helper function to figure out if the time_array, or - sampling_frequency and duration where given - """ - self._infer_dependence(domain='time', array=time_array, duration=duration, - sampling_frequency=sampling_frequency, start_time=start_time) + self, start_time, sampling_frequency, duration, time_array + ): + """Helper function to figure out if the time_array, or + sampling_frequency and duration where given + """ + self._infer_dependence( + domain="time", + array=time_array, + duration=duration, + sampling_frequency=sampling_frequency, + start_time=start_time, + ) def _infer_frequency_domain_dependence( - self, start_time, sampling_frequency, duration, frequency_array): - """ Helper function to figure out if the frequency_array, or - sampling_frequency and duration where given - """ - - self._infer_dependence(domain='frequency', array=frequency_array, - duration=duration, sampling_frequency=sampling_frequency, start_time=start_time) + self, start_time, sampling_frequency, duration, frequency_array + ): + """Helper function to figure out if the frequency_array, or + sampling_frequency and duration where given + """ + + self._infer_dependence( + domain="frequency", + array=frequency_array, + duration=duration, + sampling_frequency=sampling_frequency, + start_time=start_time, + ) - def _infer_dependence(self, domain, array, duration, sampling_frequency, start_time): + def _infer_dependence( + self, domain, array, duration, sampling_frequency, start_time + ): if (sampling_frequency is not None) and (duration is not None): if array is not None: raise ValueError( - "You have given the sampling_frequency, duration, and " - "an array") + "You have given the sampling_frequency, duration, and an array" + ) pass elif array is not None: - if domain == 'time': + if domain == "time": self.time_array = array - elif domain == 'frequency': + elif domain == "frequency": self.frequency_array = array self.start_time = start_time return elif sampling_frequency is None or duration is None: - raise ValueError( - "You must provide both sampling_frequency and duration") + raise ValueError("You must provide both sampling_frequency and duration") else: - raise ValueError( - "Insufficient information given to set arrays") - self._times_and_frequencies = CoupledTimeAndFrequencySeries(duration=duration, - sampling_frequency=sampling_frequency, - start_time=start_time) + raise ValueError("Insufficient information given to set arrays") + self._times_and_frequencies = CoupledTimeAndFrequencySeries( + duration=duration, + sampling_frequency=sampling_frequency, + start_time=start_time, + ) def set_from_time_domain_strain( - self, time_domain_strain, sampling_frequency=None, duration=None, - start_time=0, time_array=None): - """ Set the strain data from a time domain strain array + self, + time_domain_strain, + sampling_frequency=None, + duration=None, + start_time=0, + time_array=None, + ): + """Set the strain data from a time domain strain array This sets the time_domain_strain attribute, the frequency_domain_strain is automatically calculated after a low-pass filter and Tukey window @@ -490,12 +557,14 @@ def set_from_time_domain_strain( given. """ - self._infer_time_domain_dependence(start_time=start_time, - sampling_frequency=sampling_frequency, - duration=duration, - time_array=time_array) + self._infer_time_domain_dependence( + start_time=start_time, + sampling_frequency=sampling_frequency, + duration=duration, + time_array=time_array, + ) - logger.debug('Setting data using provided time_domain_strain') + logger.debug("Setting data using provided time_domain_strain") if np.shape(time_domain_strain) == np.shape(self.time_array): self._time_domain_strain = time_domain_strain self._frequency_domain_strain = None @@ -503,7 +572,7 @@ def set_from_time_domain_strain( raise ValueError("Data times do not match time array") def set_from_gwpy_timeseries(self, time_series, *, xp=np): - """ Set the strain data from a gwpy TimeSeries + """Set the strain data from a gwpy TimeSeries This sets the time_domain_strain attribute, the frequency_domain_strain is automatically calculated after a low-pass filter and Tukey window @@ -519,16 +588,18 @@ def set_from_gwpy_timeseries(self, time_series, *, xp=np): """ from gwpy.timeseries import TimeSeries - logger.debug('Setting data using provided gwpy TimeSeries object') + + logger.debug("Setting data using provided gwpy TimeSeries object") if not isinstance(time_series, TimeSeries): raise ValueError("Input time_series is not a gwpy TimeSeries") duration = xp.asarray(time_series.duration.value) sampling_frequency = xp.asarray(time_series.sample_rate.value) start_time = xp.asarray(time_series.epoch.value) - self._times_and_frequencies = \ - CoupledTimeAndFrequencySeries(duration=duration, - sampling_frequency=sampling_frequency, - start_time=start_time) + self._times_and_frequencies = CoupledTimeAndFrequencySeries( + duration=duration, + sampling_frequency=sampling_frequency, + start_time=start_time, + ) self._time_domain_strain = xp.asarray(time_series.value) self._frequency_domain_strain = None self._channel = time_series.channel @@ -538,9 +609,17 @@ def channel(self): return self._channel def set_from_open_data( - self, name, start_time, duration=4, outdir='outdir', cache=True, - *, xp=None, **kwargs): - """ Set the strain data from open LOSC data + self, + name, + start_time, + duration=4, + outdir="outdir", + cache=True, + *, + xp=None, + **kwargs, + ): + """Set the strain data from open LOSC data This sets the time_domain_strain attribute, the frequency_domain_strain is automatically calculated after a low-pass filter and Tukey window @@ -567,8 +646,13 @@ def set_from_open_data( """ timeseries = gwutils.get_open_strain_data( - name, float(start_time), float(start_time + duration), outdir=outdir, cache=cache, - **kwargs) + name, + float(start_time), + float(start_time + duration), + outdir=outdir, + cache=cache, + **kwargs, + ) if xp is None: xp = array_module((duration, start_time)) @@ -576,7 +660,7 @@ def set_from_open_data( self.set_from_gwpy_timeseries(timeseries, xp=xp) def set_from_csv(self, filename, xp=None): - """ Set the strain data from a csv file + """Set the strain data from a csv file Parameters ========== @@ -588,13 +672,19 @@ def set_from_csv(self, filename, xp=None): """ from gwpy.timeseries import TimeSeries - timeseries = TimeSeries.read(filename, format='csv') + + timeseries = TimeSeries.read(filename, format="csv") self.set_from_gwpy_timeseries(timeseries, xp=xp) def set_from_frequency_domain_strain( - self, frequency_domain_strain, sampling_frequency=None, - duration=None, start_time=0, frequency_array=None): - """ Set the `frequency_domain_strain` from a numpy array + self, + frequency_domain_strain, + sampling_frequency=None, + duration=None, + start_time=0, + frequency_array=None, + ): + """Set the `frequency_domain_strain` from a numpy array Parameters ========== @@ -612,12 +702,14 @@ def set_from_frequency_domain_strain( """ - self._infer_frequency_domain_dependence(start_time=start_time, - sampling_frequency=sampling_frequency, - duration=duration, - frequency_array=frequency_array) + self._infer_frequency_domain_dependence( + start_time=start_time, + sampling_frequency=sampling_frequency, + duration=duration, + frequency_array=frequency_array, + ) - logger.debug('Setting data using provided frequency_domain_strain') + logger.debug("Setting data using provided frequency_domain_strain") if np.shape(frequency_domain_strain) == np.shape(self.frequency_array): self._frequency_domain_strain = frequency_domain_strain self.window_factor = 1 @@ -625,9 +717,15 @@ def set_from_frequency_domain_strain( raise ValueError("Data frequencies do not match frequency_array") def set_from_power_spectral_density( - self, power_spectral_density, sampling_frequency, duration, - start_time=0, *, random_state=None): - """ Set the `frequency_domain_strain` by generating a noise realisation + self, + power_spectral_density, + sampling_frequency, + duration, + start_time=0, + *, + random_state=None, + ): + """Set the `frequency_domain_strain` by generating a noise realisation Parameters ========== @@ -642,27 +740,34 @@ def set_from_power_spectral_density( """ - self._times_and_frequencies = CoupledTimeAndFrequencySeries(duration=duration, - sampling_frequency=sampling_frequency, - start_time=start_time) + self._times_and_frequencies = CoupledTimeAndFrequencySeries( + duration=duration, + sampling_frequency=sampling_frequency, + start_time=start_time, + ) logger.debug( - 'Setting data using noise realization from provided' - 'power_spectal_density') - frequency_domain_strain, frequency_array = \ + "Setting data using noise realization from providedpower_spectal_density" + ) + frequency_domain_strain, frequency_array = ( power_spectral_density.get_noise_realisation( - self.frequency_array.shape[0], duration, random_state=random_state) + self.frequency_array.shape[0], duration, random_state=random_state + ) + ) xp = aac.array_namespace(frequency_domain_strain) self._frequency_array = xp.asarray(self.frequency_array) self._time_array = xp.asarray(self.time_array) - if self.duration == duration and frequency_array.shape == self.frequency_array.shape: + if ( + self.duration == duration + and frequency_array.shape == self.frequency_array.shape + ): self._frequency_domain_strain = frequency_domain_strain else: raise ValueError("Data frequencies do not match frequency_array") def set_from_zero_noise(self, sampling_frequency, duration, start_time=0): - """ Set the `frequency_domain_strain` to zero noise + """Set the `frequency_domain_strain` to zero noise This sets the `strain_data` to an array of (complex) zeros, while also making sure the frequency and time arrays match @@ -679,18 +784,29 @@ def set_from_zero_noise(self, sampling_frequency, duration, start_time=0): """ - self._times_and_frequencies = CoupledTimeAndFrequencySeries(duration=duration, - sampling_frequency=sampling_frequency, - start_time=start_time) - logger.debug('Setting zero noise data') + self._times_and_frequencies = CoupledTimeAndFrequencySeries( + duration=duration, + sampling_frequency=sampling_frequency, + start_time=start_time, + ) + logger.debug("Setting zero noise data") xp = aac.get_namespace(self.frequency_array) - self._frequency_domain_strain = xp.zeros_like(self.frequency_array, - dtype=complex) + self._frequency_domain_strain = xp.zeros_like( + self.frequency_array, dtype=complex + ) def set_from_frame_file( - self, frame_file, sampling_frequency, duration, start_time=0, - channel=None, buffer_time=1, *, xp=None): - """ Set the `frequency_domain_strain` from a frame fiile + self, + frame_file, + sampling_frequency, + duration, + start_time=0, + channel=None, + buffer_time=1, + *, + xp=None, + ): + """Set the `frequency_domain_strain` from a frame fiile Parameters ========== @@ -715,22 +831,30 @@ def set_from_frame_file( """ self._times_and_frequencies = CoupledTimeAndFrequencySeries( - duration=duration, sampling_frequency=sampling_frequency, - start_time=start_time) + duration=duration, + sampling_frequency=sampling_frequency, + start_time=start_time, + ) - logger.info('Reading data from frame file {}'.format(frame_file)) + logger.info("Reading data from frame file {}".format(frame_file)) strain = gwutils.read_frame_file( - frame_file, start_time=start_time, end_time=start_time + duration, - buffer_time=buffer_time, channel=channel, - resample=sampling_frequency) + frame_file, + start_time=start_time, + end_time=start_time + duration, + buffer_time=buffer_time, + channel=channel, + resample=sampling_frequency, + ) if xp is None: xp = aac.get_namespace(self.frequency_array) self.set_from_gwpy_timeseries(strain, xp=xp) - def set_from_channel_name(self, channel, duration, start_time, sampling_frequency, *, xp=None): - """ Set the `frequency_domain_strain` by fetching from given channel + def set_from_channel_name( + self, channel, duration, start_time, sampling_frequency, *, xp=None + ): + """Set the `frequency_domain_strain` by fetching from given channel using gwpy.TimesSeries.get(), which dynamically accesses either frames on disk, or a remote NDS2 server to find and return data. This function also verifies that the specified channel is given in the correct format. @@ -752,15 +876,18 @@ def set_from_channel_name(self, channel, duration, start_time, sampling_frequenc """ from gwpy.timeseries import TimeSeries - channel_comp = channel.split(':') + + channel_comp = channel.split(":") if len(channel_comp) != 2: - raise IndexError('Channel name must have format `IFO:Channel`') + raise IndexError("Channel name must have format `IFO:Channel`") self._times_and_frequencies = CoupledTimeAndFrequencySeries( - duration=duration, sampling_frequency=sampling_frequency, - start_time=start_time) + duration=duration, + sampling_frequency=sampling_frequency, + start_time=start_time, + ) - logger.info('Fetching data using channel {}'.format(channel)) + logger.info("Fetching data using channel {}".format(channel)) strain = TimeSeries.get(channel, start_time, start_time + duration) strain = strain.resample(sampling_frequency) @@ -772,7 +899,7 @@ def set_from_channel_name(self, channel, duration, start_time, sampling_frequenc class Notch(object): def __init__(self, minimum_frequency, maximum_frequency): - """ A notch object storing the maximum and minimum frequency of the notch + """A notch object storing the maximum and minimum frequency of the notch Parameters ========== @@ -785,12 +912,13 @@ def __init__(self, minimum_frequency, maximum_frequency): self.minimum_frequency = minimum_frequency self.maximum_frequency = maximum_frequency else: - msg = ("Your notch minimum_frequency {} and maximum_frequency {} are invalid" - .format(minimum_frequency, maximum_frequency)) + msg = "Your notch minimum_frequency {} and maximum_frequency {} are invalid".format( + minimum_frequency, maximum_frequency + ) raise ValueError(msg) def get_idxs(self, frequency_array): - """ Get a boolean mask for the frequencies in frequency_array in the notch + """Get a boolean mask for the frequencies in frequency_array in the notch Parameters ========== @@ -803,12 +931,12 @@ def get_idxs(self, frequency_array): An array of booleans which are True for frequencies in the notch """ - lower = (frequency_array > self.minimum_frequency) - upper = (frequency_array < self.maximum_frequency) + lower = frequency_array > self.minimum_frequency + upper = frequency_array < self.maximum_frequency return lower & upper def check_frequency(self, freq): - """ Check if freq is inside the notch + """Check if freq is inside the notch Parameters ========== @@ -829,7 +957,7 @@ def check_frequency(self, freq): class NotchList(list): def __init__(self, notch_list): - """ A list of notches + """A list of notches Parameters ========== @@ -852,7 +980,7 @@ def __init__(self, notch_list): raise ValueError(msg) def check_frequency(self, freq): - """ Check if freq is inside the notch list + """Check if freq is inside the notch list Parameters ========== diff --git a/bilby/gw/eos/__init__.py b/bilby/gw/eos/__init__.py index 36eba397f..a32b2e781 100644 --- a/bilby/gw/eos/__init__.py +++ b/bilby/gw/eos/__init__.py @@ -1,3 +1,2 @@ from .tov_solver import IntegrateTOV -from .eos import (SpectralDecompositionEOS, - EOSFamily, TabularEOS) +from .eos import SpectralDecompositionEOS, EOSFamily, TabularEOS diff --git a/bilby/gw/eos/eos.py b/bilby/gw/eos/eos.py index 2e962aa0c..c4a294df0 100644 --- a/bilby/gw/eos/eos.py +++ b/bilby/gw/eos/eos.py @@ -6,27 +6,39 @@ from ...core import utils C_SI = utils.speed_of_light # m/s -C_CGS = C_SI * 100. +C_CGS = C_SI * 100.0 G_SI = utils.gravitational_constant # m^3 kg^-1 s^-2 MSUN_SI = utils.solar_mass # Kg # Stores conversions from geometerized to cgs or si unit systems -conversion_dict = {'pressure': {'cgs': C_SI ** 4. / G_SI * 10., 'si': C_SI ** 4. / G_SI, 'geom': 1.}, - 'energy_density': {'cgs': C_SI ** 4. / G_SI * 10., 'si': C_SI ** 4. / G_SI, 'geom': 1.}, - 'density': {'cgs': C_SI ** 2. / G_SI / 1000., 'si': C_SI ** 2. / G_SI, 'geom': 1.}, - 'pseudo_enthalpy': {'dimensionless': 1.}, - 'mass': {'g': C_SI ** 2. / G_SI * 1000, 'kg': C_SI ** 2. / G_SI, 'geom': 1., - 'm_sol': C_SI ** 2. / G_SI / MSUN_SI}, - 'radius': {'cm': 100., 'm': 1., 'km': .001}, - 'tidal_deformability': {'geom': 1.}} +conversion_dict = { + "pressure": {"cgs": C_SI**4.0 / G_SI * 10.0, "si": C_SI**4.0 / G_SI, "geom": 1.0}, + "energy_density": { + "cgs": C_SI**4.0 / G_SI * 10.0, + "si": C_SI**4.0 / G_SI, + "geom": 1.0, + }, + "density": {"cgs": C_SI**2.0 / G_SI / 1000.0, "si": C_SI**2.0 / G_SI, "geom": 1.0}, + "pseudo_enthalpy": {"dimensionless": 1.0}, + "mass": { + "g": C_SI**2.0 / G_SI * 1000, + "kg": C_SI**2.0 / G_SI, + "geom": 1.0, + "m_sol": C_SI**2.0 / G_SI / MSUN_SI, + }, + "radius": {"cm": 100.0, "m": 1.0, "km": 0.001}, + "tidal_deformability": {"geom": 1.0}, +} # construct dictionary of pre-shipped EOS pressure denstity table -path_to_eos_tables = os.path.join(os.path.dirname(__file__), 'eos_tables') +path_to_eos_tables = os.path.join(os.path.dirname(__file__), "eos_tables") list_of_eos_tables = os.listdir(path_to_eos_tables) -valid_eos_files = [i for i in list_of_eos_tables if 'LAL' in i] -valid_eos_file_paths = [os.path.join(path_to_eos_tables, filename) for filename in valid_eos_files] -valid_eos_names = [i.split('_', maxsplit=1)[-1].strip('.dat') for i in valid_eos_files] +valid_eos_files = [i for i in list_of_eos_tables if "LAL" in i] +valid_eos_file_paths = [ + os.path.join(path_to_eos_tables, filename) for filename in valid_eos_files +] +valid_eos_names = [i.split("_", maxsplit=1)[-1].strip(".dat") for i in valid_eos_files] valid_eos_dict = dict(zip(valid_eos_names, valid_eos_file_paths)) @@ -68,8 +80,10 @@ def __init__(self, eos, sampling_flag=False, warning_flag=False): elif isinstance(eos, np.ndarray): table = eos else: - raise ValueError("eos provided is invalid type please supply a str name, str path to ASCII file, " - "or a numpy array") + raise ValueError( + "eos provided is invalid type please supply a str name, str path to ASCII file, " + "or a numpy array" + ) table = self.__remove_leading_zero(table) @@ -83,20 +97,27 @@ def __init__(self, eos, sampling_flag=False, warning_flag=False): self.warning_flag = True else: integrand = self.pressure / (self.energy_density + self.pressure) - self.pseudo_enthalpy = cumulative_trapezoid(integrand, np.log(self.pressure), initial=0) + integrand[0] + self.pseudo_enthalpy = ( + cumulative_trapezoid(integrand, np.log(self.pressure), initial=0) + + integrand[0] + ) - self.interp_energy_density_from_pressure = CubicSpline(np.log10(self.pressure), - np.log10(self.energy_density), - ) + self.interp_energy_density_from_pressure = CubicSpline( + np.log10(self.pressure), + np.log10(self.energy_density), + ) - self.interp_energy_density_from_pseudo_enthalpy = CubicSpline(np.log10(self.pseudo_enthalpy), - np.log10(self.energy_density)) + self.interp_energy_density_from_pseudo_enthalpy = CubicSpline( + np.log10(self.pseudo_enthalpy), np.log10(self.energy_density) + ) - self.interp_pressure_from_pseudo_enthalpy = CubicSpline(np.log10(self.pseudo_enthalpy), - np.log10(self.pressure)) + self.interp_pressure_from_pseudo_enthalpy = CubicSpline( + np.log10(self.pseudo_enthalpy), np.log10(self.pressure) + ) - self.interp_pseudo_enthalpy_from_energy_density = CubicSpline(np.log10(self.energy_density), - np.log10(self.pseudo_enthalpy)) + self.interp_pseudo_enthalpy_from_energy_density = CubicSpline( + np.log10(self.energy_density), np.log10(self.pseudo_enthalpy) + ) self.__construct_all_tables() @@ -110,13 +131,13 @@ def __remove_leading_zero(self, table): loglog interpolation breaks if the first entries are 0s """ - if table[0, 0] == 0. or table[0, 1] == 0.: + if table[0, 0] == 0.0 or table[0, 1] == 0.0: return table[1:, :] else: return table - def energy_from_pressure(self, pressure, interp_type='CubicSpline'): + def energy_from_pressure(self, pressure, interp_type="CubicSpline"): """ Find value of energy_from_pressure as in lalsimulation, return e = K * p**(3./5.) below min pressure @@ -137,26 +158,32 @@ def energy_from_pressure(self, pressure, interp_type='CubicSpline'): indices_greater_than_min = np.nonzero(pressure >= self.minimum_pressure) # We do this special for less than min pressure - energy_returned[indices_less_than_min] = 10 ** (np.log10(self.energy_density[0]) + - (3. / 5.) * (np.log10(pressure[indices_less_than_min]) - - np.log10(self.pressure[0]))) + energy_returned[indices_less_than_min] = 10 ** ( + np.log10(self.energy_density[0]) + + (3.0 / 5.0) + * (np.log10(pressure[indices_less_than_min]) - np.log10(self.pressure[0])) + ) - if interp_type == 'CubicSpline': + if interp_type == "CubicSpline": energy_returned[indices_greater_than_min] = ( - 10. ** self.interp_energy_density_from_pressure(np.log10(pressure[indices_greater_than_min]))) - elif interp_type == 'linear': - energy_returned[indices_greater_than_min] = np.interp(pressure[indices_greater_than_min], - self.pressure, - self.energy_density) + 10.0 + ** self.interp_energy_density_from_pressure( + np.log10(pressure[indices_greater_than_min]) + ) + ) + elif interp_type == "linear": + energy_returned[indices_greater_than_min] = np.interp( + pressure[indices_greater_than_min], self.pressure, self.energy_density + ) else: - raise ValueError('Interpolation scheme must be linear or CubicSpline') + raise ValueError("Interpolation scheme must be linear or CubicSpline") if energy_returned.size == 1: return energy_returned[0] else: return energy_returned - def pressure_from_pseudo_enthalpy(self, pseudo_enthalpy, interp_type='CubicSpline'): + def pressure_from_pseudo_enthalpy(self, pseudo_enthalpy, interp_type="CubicSpline"): """ Find p(h) as in lalsimulation, return p = K * h**(5./2.) below min enthalpy @@ -169,29 +196,46 @@ def pressure_from_pseudo_enthalpy(self, pseudo_enthalpy, interp_type='CubicSplin """ pseudo_enthalpy = np.atleast_1d(pseudo_enthalpy) pressure_returned = np.zeros(pseudo_enthalpy.size) - indices_less_than_min = np.nonzero(pseudo_enthalpy < self.minimum_pseudo_enthalpy) - indices_greater_than_min = np.nonzero(pseudo_enthalpy >= self.minimum_pseudo_enthalpy) - - pressure_returned[indices_less_than_min] = 10. ** (np.log10(self.pressure[0]) + - 2.5 * (np.log10(pseudo_enthalpy[indices_less_than_min]) - - np.log10(self.pseudo_enthalpy[0]))) - - if interp_type == 'CubicSpline': + indices_less_than_min = np.nonzero( + pseudo_enthalpy < self.minimum_pseudo_enthalpy + ) + indices_greater_than_min = np.nonzero( + pseudo_enthalpy >= self.minimum_pseudo_enthalpy + ) + + pressure_returned[indices_less_than_min] = 10.0 ** ( + np.log10(self.pressure[0]) + + 2.5 + * ( + np.log10(pseudo_enthalpy[indices_less_than_min]) + - np.log10(self.pseudo_enthalpy[0]) + ) + ) + + if interp_type == "CubicSpline": pressure_returned[indices_greater_than_min] = ( - 10. ** self.interp_pressure_from_pseudo_enthalpy(np.log10(pseudo_enthalpy[indices_greater_than_min]))) - elif interp_type == 'linear': - pressure_returned[indices_greater_than_min] = np.interp(pseudo_enthalpy[indices_greater_than_min], - self.pseudo_enthalpy, - self.pressure) + 10.0 + ** self.interp_pressure_from_pseudo_enthalpy( + np.log10(pseudo_enthalpy[indices_greater_than_min]) + ) + ) + elif interp_type == "linear": + pressure_returned[indices_greater_than_min] = np.interp( + pseudo_enthalpy[indices_greater_than_min], + self.pseudo_enthalpy, + self.pressure, + ) else: - raise ValueError('Interpolation scheme must be linear or CubicSpline') + raise ValueError("Interpolation scheme must be linear or CubicSpline") if pressure_returned.size == 1: return pressure_returned[0] else: return pressure_returned - def energy_density_from_pseudo_enthalpy(self, pseudo_enthalpy, interp_type='CubicSpline'): + def energy_density_from_pseudo_enthalpy( + self, pseudo_enthalpy, interp_type="CubicSpline" + ): """ Find energy_density_from_pseudo_enthalpy(pseudo_enthalpy) as in lalsimulation, return e = K * h**(3./2.) below min enthalpy @@ -204,28 +248,43 @@ def energy_density_from_pseudo_enthalpy(self, pseudo_enthalpy, interp_type='Cubi """ pseudo_enthalpy = np.atleast_1d(pseudo_enthalpy) energy_returned = np.zeros(pseudo_enthalpy.size) - indices_less_than_min = np.nonzero(pseudo_enthalpy < self.minimum_pseudo_enthalpy) - indices_greater_than_min = np.nonzero(pseudo_enthalpy >= self.minimum_pseudo_enthalpy) - - energy_returned[indices_less_than_min] = 10 ** (np.log10(self.energy_density[0]) + - 1.5 * (np.log10(pseudo_enthalpy[indices_less_than_min]) - - np.log10(self.pseudo_enthalpy[0]))) - if interp_type == 'CubicSpline': + indices_less_than_min = np.nonzero( + pseudo_enthalpy < self.minimum_pseudo_enthalpy + ) + indices_greater_than_min = np.nonzero( + pseudo_enthalpy >= self.minimum_pseudo_enthalpy + ) + + energy_returned[indices_less_than_min] = 10 ** ( + np.log10(self.energy_density[0]) + + 1.5 + * ( + np.log10(pseudo_enthalpy[indices_less_than_min]) + - np.log10(self.pseudo_enthalpy[0]) + ) + ) + if interp_type == "CubicSpline": x = np.log10(pseudo_enthalpy[indices_greater_than_min]) - energy_returned[indices_greater_than_min] = 10 ** self.interp_energy_density_from_pseudo_enthalpy(x) - elif interp_type == 'linear': - energy_returned[indices_greater_than_min] = np.interp(pseudo_enthalpy[indices_greater_than_min], - self.pseudo_enthalpy, - self.energy_density) + energy_returned[indices_greater_than_min] = ( + 10 ** self.interp_energy_density_from_pseudo_enthalpy(x) + ) + elif interp_type == "linear": + energy_returned[indices_greater_than_min] = np.interp( + pseudo_enthalpy[indices_greater_than_min], + self.pseudo_enthalpy, + self.energy_density, + ) else: - raise ValueError('Interpolation scheme must be linear or CubicSpline') + raise ValueError("Interpolation scheme must be linear or CubicSpline") if energy_returned.size == 1: return energy_returned[0] else: return energy_returned - def pseudo_enthalpy_from_energy_density(self, energy_density, interp_type='CubicSpline'): + def pseudo_enthalpy_from_energy_density( + self, energy_density, interp_type="CubicSpline" + ): """ Find h(epsilon) as in lalsimulation, return h = K * e**(2./3.) below min enthalpy @@ -239,28 +298,39 @@ def pseudo_enthalpy_from_energy_density(self, energy_density, interp_type='Cubic energy_density = np.atleast_1d(energy_density) pseudo_enthalpy_returned = np.zeros(energy_density.size) indices_less_than_min = np.nonzero(energy_density < self.minimum_energy_density) - indices_greater_than_min = np.nonzero(energy_density >= self.minimum_energy_density) - - pseudo_enthalpy_returned[indices_less_than_min] = 10 ** (np.log10(self.pseudo_enthalpy[0]) + (2. / 3.) * - (np.log10(energy_density[indices_less_than_min]) - - np.log10(self.energy_density[0]))) - - if interp_type == 'CubicSpline': + indices_greater_than_min = np.nonzero( + energy_density >= self.minimum_energy_density + ) + + pseudo_enthalpy_returned[indices_less_than_min] = 10 ** ( + np.log10(self.pseudo_enthalpy[0]) + + (2.0 / 3.0) + * ( + np.log10(energy_density[indices_less_than_min]) + - np.log10(self.energy_density[0]) + ) + ) + + if interp_type == "CubicSpline": x = np.log10(energy_density[indices_greater_than_min]) - pseudo_enthalpy_returned[indices_greater_than_min] = 10**self.interp_pseudo_enthalpy_from_energy_density(x) - elif interp_type == 'linear': - pseudo_enthalpy_returned[indices_greater_than_min] = np.interp(energy_density[indices_greater_than_min], - self.energy_density, - self.pseudo_enthalpy) + pseudo_enthalpy_returned[indices_greater_than_min] = ( + 10 ** self.interp_pseudo_enthalpy_from_energy_density(x) + ) + elif interp_type == "linear": + pseudo_enthalpy_returned[indices_greater_than_min] = np.interp( + energy_density[indices_greater_than_min], + self.energy_density, + self.pseudo_enthalpy, + ) else: - raise ValueError('Interpolation scheme must be linear or CubicSpline') + raise ValueError("Interpolation scheme must be linear or CubicSpline") if pseudo_enthalpy_returned.size == 1: return pseudo_enthalpy_returned[0] else: return pseudo_enthalpy_returned - def dedh(self, pseudo_enthalpy, rel_dh=1e-5, interp_type='CubicSpline'): + def dedh(self, pseudo_enthalpy, rel_dh=1e-5, interp_type="CubicSpline"): """ Value of [depsilon/dh](p) @@ -276,12 +346,16 @@ def dedh(self, pseudo_enthalpy, rel_dh=1e-5, interp_type='CubicSpline'): # step size=fraction of value dh = pseudo_enthalpy * rel_dh - eps_upper = self.energy_density_from_pseudo_enthalpy(pseudo_enthalpy + dh, interp_type=interp_type) - eps_lower = self.energy_density_from_pseudo_enthalpy(pseudo_enthalpy - dh, interp_type=interp_type) + eps_upper = self.energy_density_from_pseudo_enthalpy( + pseudo_enthalpy + dh, interp_type=interp_type + ) + eps_lower = self.energy_density_from_pseudo_enthalpy( + pseudo_enthalpy - dh, interp_type=interp_type + ) - return (eps_upper - eps_lower) / (2. * dh) + return (eps_upper - eps_lower) / (2.0 * dh) - def dedp(self, pressure, rel_dp=1e-5, interp_type='CubicSpline'): + def dedp(self, pressure, rel_dp=1e-5, interp_type="CubicSpline"): """ Find value of [depsilon/dp](p) @@ -300,9 +374,9 @@ def dedp(self, pressure, rel_dp=1e-5, interp_type='CubicSpline'): eps_upper = self.energy_from_pressure(pressure + dp, interp_type=interp_type) eps_lower = self.energy_from_pressure(pressure - dp, interp_type=interp_type) - return (eps_upper - eps_lower) / (2. * dp) + return (eps_upper - eps_lower) / (2.0 * dp) - def velocity_from_pseudo_enthalpy(self, pseudo_enthalpy, interp_type='CubicSpline'): + def velocity_from_pseudo_enthalpy(self, pseudo_enthalpy, interp_type="CubicSpline"): """ Returns the speed of sound in geometerized units in the neutron star at the specified pressure. @@ -316,7 +390,9 @@ def velocity_from_pseudo_enthalpy(self, pseudo_enthalpy, interp_type='CubicSplin :return v_s (`float`): Speed of sound at `pseudo-enthalpy` in geometerized units. """ - pressure = self.pressure_from_pseudo_enthalpy(pseudo_enthalpy, interp_type=interp_type) + pressure = self.pressure_from_pseudo_enthalpy( + pseudo_enthalpy, interp_type=interp_type + ) return self.dedp(pressure, interp_type=interp_type) ** -0.5 def check_causality(self): @@ -399,53 +475,67 @@ def plot(self, rep, xlim=None, ylim=None, units=None): import matplotlib.pyplot as plt # Set data based on specified representation - varnames = rep.split('-') + varnames = rep.split("-") - assert varnames[0] != varnames[ - 1], 'Cannot plot the same variable against itself. Please choose another representation' + assert varnames[0] != varnames[1], ( + "Cannot plot the same variable against itself. Please choose another representation" + ) # Correspondence of rep parameter, data, and latex symbol # rep_dict = {'energy_density': [self.epsilon, r'$\epsilon$'], # 'pressure': [self.p, r'$p$'], 'pseudo_enthalpy': [pseudo_enthalpy, r'$h$']} # FIXME: The second element in these arrays should be tex labels, but tex's not working rn - rep_dict = {'energy_density': [self.energy_density, 'energy_density'], - 'pressure': [self.pressure, 'pressure'], - 'pseudo_enthalpy': [self.pseudo_enthalpy, 'pseudo_enthalpy']} + rep_dict = { + "energy_density": [self.energy_density, "energy_density"], + "pressure": [self.pressure, "pressure"], + "pseudo_enthalpy": [self.pseudo_enthalpy, "pseudo_enthalpy"], + } xname = varnames[1] yname = varnames[0] # Set units - eos_default_units = {'pressure': 'cgs', 'energy_density': 'cgs', 'density': 'cgs', - 'pseudo_enthalpy': 'dimensionless'} + eos_default_units = { + "pressure": "cgs", + "energy_density": "cgs", + "density": "cgs", + "pseudo_enthalpy": "dimensionless", + } if units is None: - units = [eos_default_units[yname], eos_default_units[xname]] # Default unit system is cgs + units = [ + eos_default_units[yname], + eos_default_units[xname], + ] # Default unit system is cgs elif isinstance(units, str): - units = [units, units] # If only one unit system given, use for both + units = [units, units] # If only one unit system given, use for both xunits = units[1] yunits = units[0] # Ensure valid units - if xunits not in list(conversion_dict[xname].keys()) or yunits not in list(conversion_dict[yname].keys()): - s = ''' + if xunits not in list(conversion_dict[xname].keys()) or yunits not in list( + conversion_dict[yname].keys() + ): + s = """ Invalid unit system. Valid variable-unit pairs are: p: {p_units} e: {e_units} rho: {rho_units} h: {h_units}. - '''.format(p_units=list(conversion_dict['pressure'].keys()), - e_units=list(conversion_dict['energy_density'].keys()), - rho_units=list(conversion_dict['density'].keys()), - h_units=list(conversion_dict['pseudo_enthalpy'].keys())) + """.format( + p_units=list(conversion_dict["pressure"].keys()), + e_units=list(conversion_dict["energy_density"].keys()), + rho_units=list(conversion_dict["density"].keys()), + h_units=list(conversion_dict["pseudo_enthalpy"].keys()), + ) raise ValueError(s) xdat = rep_dict[xname][0] * conversion_dict[xname][xunits] ydat = rep_dict[yname][0] * conversion_dict[yname][yunits] - xlabel = rep_dict[varnames[1]][1].replace('_', ' ') - ylabel = rep_dict[varnames[0]][1].replace('_', ' ') + '(' + xlabel + ')' + xlabel = rep_dict[varnames[1]][1].replace("_", " ") + ylabel = rep_dict[varnames[0]][1].replace("_", " ") + "(" + xlabel + ")" # Determine plot ranges. Currently shows 10% wider than actual data range. if xlim is None: @@ -468,7 +558,7 @@ def plot(self, rep, xlim=None, ylim=None, units=None): def spectral_adiabatic_index(gammas, x): arg = 0 for i in range(len(gammas)): - arg += gammas[i] * x ** i + arg += gammas[i] * x**i return np.exp(arg) @@ -501,7 +591,16 @@ class SpectralDecompositionEOS(TabularEOS): number of points in pressure-energy density data. """ - def __init__(self, gammas, p0=3.01e33, e0=2.03e14, xmax=None, npts=100, sampling_flag=False, warning_flag=False): + def __init__( + self, + gammas, + p0=3.01e33, + e0=2.03e14, + xmax=None, + npts=100, + sampling_flag=False, + warning_flag=False, + ): self.warning_flag = warning_flag self.gammas = gammas self.p0 = p0 @@ -523,35 +622,38 @@ def __init__(self, gammas, p0=3.01e33, e0=2.03e14, xmax=None, npts=100, sampling else: self.e_pdat = self.__construct_e_of_p_table() - super().__init__(self.e_pdat, sampling_flag=self.sampling_flag, warning_flag=self.warning_flag) + super().__init__( + self.e_pdat, + sampling_flag=self.sampling_flag, + warning_flag=self.warning_flag, + ) def __determine_xmax(self, a_max=6.0): highest_order_gamma = np.abs(self.gammas[-1])[0] expansion_order = float(len(self.gammas) - 1) - xmax = (np.log(a_max) / highest_order_gamma) ** (1. / expansion_order) + xmax = (np.log(a_max) / highest_order_gamma) ** (1.0 / expansion_order) return xmax def __mu_integrand(self, x): - - return 1. / spectral_adiabatic_index(self.gammas, x) + return 1.0 / spectral_adiabatic_index(self.gammas, x) def mu(self, x): from scipy.integrate import quad + return np.exp(-quad(self.__mu_integrand, 0, x)[0]) def __eps_integrand(self, x): - return np.exp(x) * self.mu(x) / spectral_adiabatic_index(self.gammas, x) def energy_density(self, x, eps0): from scipy.integrate import quad + quad_result, quad_err = quad(self.__eps_integrand, 0, x) - eps_of_x = (eps0 * C_CGS ** 2.) / self.mu(x) + self.p0 / self.mu(x) * quad_result + eps_of_x = (eps0 * C_CGS**2.0) / self.mu(x) + self.p0 / self.mu(x) * quad_result return eps_of_x def __construct_a_of_x_table(self): - xdat = np.linspace(0, self.xmax, num=self.npts) # Generate adiabatic index points until a point is out of prior range @@ -568,7 +670,7 @@ def __construct_a_of_x_table(self): xmax_new = xdat[i - 1] # If EOS is too short, set prior to 0, else resample the function and set new xmax - if xmax_new < 4. or i == 0: + if xmax_new < 4.0 or i == 0: self.warning_flag = True else: xdat = np.linspace(0, xmax_new, num=self.npts) @@ -580,7 +682,6 @@ def __construct_a_of_x_table(self): self.adat = adat def __construct_e_of_p_table(self): - """ Creates p, epsilon table for a given set of spectral parameters """ @@ -598,11 +699,13 @@ def __construct_e_of_p_table(self): # convert eos to geometrized units in *m^-2* # IMPORTANT - eos_vals = eos_vals * 0.1 * G_SI / C_SI ** 4 + eos_vals = eos_vals * 0.1 * G_SI / C_SI**4 # doing as those before me have done and using SLY4 as low density region # SLY4 in geometrized units - low_density_path = os.path.join(os.path.dirname(__file__), 'eos_tables', 'LALSimNeutronStarEOS_SLY4.dat') + low_density_path = os.path.join( + os.path.dirname(__file__), "eos_tables", "LALSimNeutronStarEOS_SLY4.dat" + ) low_density = np.loadtxt(low_density_path) cutoff = eos_vals[0, :] @@ -636,8 +739,10 @@ class EOSFamily(object): The mass-radius and mass-k2 data should be populated here via the TOV solver upon object construction. """ + def __init__(self, eos, npts=500): from scipy.optimize import minimize_scalar + self.eos = eos # FIXME: starting_energy_density is set somewhat arbitrarily @@ -645,9 +750,9 @@ def __init__(self, eos, npts=500): ending_energy_density = max(self.eos.energy_density) log_starting_energy_density = np.log(starting_energy_density) log_ending_energy_density = np.log(ending_energy_density) - log_energy_density_grid = np.linspace(log_starting_energy_density, - log_ending_energy_density, - num=npts) + log_energy_density_grid = np.linspace( + log_starting_energy_density, log_ending_energy_density, num=npts + ) energy_density_grid = np.exp(log_energy_density_grid) # Generate m, r, and k2 lists @@ -671,10 +776,16 @@ def __init__(self, eos, npts=500): # Now replace with point with interpolated maximum mass # This is done by interpolating the last three points and then # minimizing the negative of the interpolated function - x = [energy_density_grid[i - 2], energy_density_grid[i - 1], energy_density_grid[i]] + x = [ + energy_density_grid[i - 2], + energy_density_grid[i - 1], + energy_density_grid[i], + ] y = [mass[i - 2], mass[i - 1], mass[i]] - f = interp1d(x, y, kind='quadratic', bounds_error=False, fill_value='extrapolate') + f = interp1d( + x, y, kind="quadratic", bounds_error=False, fill_value="extrapolate" + ) res = minimize_scalar(lambda x: -f(x)) @@ -692,8 +803,10 @@ def __init__(self, eos, npts=500): # with these quantities, then convert to SI. # Calculating dimensionless lambda values from k2, radii, and mass - tidal_deformability = [2. / 3. * k2 * r ** 5. / m ** 5. for k2, r, m in - zip(k2love_number, radius, mass)] + tidal_deformability = [ + 2.0 / 3.0 * k2 * r**5.0 / m**5.0 + for k2, r, m in zip(k2love_number, radius, mass) + ] # As a last resort, if highest mass is still smaller than second # to last point, remove the last point from each array @@ -707,16 +820,16 @@ def __init__(self, eos, npts=500): self.radius = np.array(radius) self.k2love_number = np.array(k2love_number) self.tidal_deformability = np.array(tidal_deformability) - self.maximum_mass = mass[-1] * conversion_dict['mass']['m_sol'] + self.maximum_mass = mass[-1] * conversion_dict["mass"]["m_sol"] def radius_from_mass(self, m): """ :param m: mass of neutron star in solar masses :return: radius of neutron star in meters """ - f = CubicSpline(self.mass, self.radius, bc_type='natural', extrapolate=True) + f = CubicSpline(self.mass, self.radius, bc_type="natural", extrapolate=True) - mass_converted_to_geom = m * MSUN_SI * G_SI / C_SI ** 2. + mass_converted_to_geom = m * MSUN_SI * G_SI / C_SI**2.0 return f(mass_converted_to_geom) def k2_from_mass(self, m): @@ -724,9 +837,11 @@ def k2_from_mass(self, m): :param m: mass of neutron star in solar masses. :return: dimensionless second tidal love number. """ - f = CubicSpline(self.mass, self.k2love_number, bc_type='natural', extrapolate=True) + f = CubicSpline( + self.mass, self.k2love_number, bc_type="natural", extrapolate=True + ) - m_geom = m * MSUN_SI * G_SI / C_SI ** 2. + m_geom = m * MSUN_SI * G_SI / C_SI**2.0 return f(m_geom) def lambda_from_mass(self, m): @@ -742,9 +857,9 @@ def lambda_from_mass(self, m): r = self.radius_from_mass(m) k = self.k2_from_mass(m) - m_geom = m * MSUN_SI * G_SI / C_SI ** 2. + m_geom = m * MSUN_SI * G_SI / C_SI**2.0 c = m_geom / r - lmbda = (2. / 3.) * k / c ** 5. + lmbda = (2.0 / 3.0) * k / c**5.0 return lmbda @@ -784,22 +899,34 @@ def plot(self, rep, xlim=None, ylim=None, units=None): import matplotlib.pyplot as plt # Set data based on specified representation - varnames = rep.split('-') + varnames = rep.split("-") - assert varnames[0] != varnames[ - 1], 'Cannot plot the same variable against itself. Please choose another representation' + assert varnames[0] != varnames[1], ( + "Cannot plot the same variable against itself. Please choose another representation" + ) # Correspondence of rep parameter, data, and latex symbol - rep_dict = {'mass': [self.mass, r'$M$'], 'radius': [self.radius, r'$R$'], 'k2': [self.k2love_number, r'$k_2$'], - 'tidal_deformability': [self.tidal_deformability, r'$l$']} + rep_dict = { + "mass": [self.mass, r"$M$"], + "radius": [self.radius, r"$R$"], + "k2": [self.k2love_number, r"$k_2$"], + "tidal_deformability": [self.tidal_deformability, r"$l$"], + } xname = varnames[1] yname = varnames[0] # Set units - fam_default_units = {'mass': 'm_sol', 'radius': 'km', 'tidal_deformability': 'geom'} + fam_default_units = { + "mass": "m_sol", + "radius": "km", + "tidal_deformability": "geom", + } if units is None: - units = [fam_default_units[yname], fam_default_units[xname]] # Default unit system is cgs + units = [ + fam_default_units[yname], + fam_default_units[xname], + ] # Default unit system is cgs elif isinstance(units, str): units = [units, units] # If only one unit system given, use for both @@ -807,22 +934,26 @@ def plot(self, rep, xlim=None, ylim=None, units=None): yunits = units[0] # Ensure valid units - if xunits not in list(conversion_dict[xname].keys()) or yunits not in list(conversion_dict[yname].keys()): - s = ''' + if xunits not in list(conversion_dict[xname].keys()) or yunits not in list( + conversion_dict[yname].keys() + ): + s = """ Invalid unit system. Valid variable-unit pairs are: m: {m_units} r: {r_units} l: {l_units}. - '''.format(m_units=list(conversion_dict['mass'].keys()), - r_units=list(conversion_dict['radius'].keys()), - l_units=list(conversion_dict['tidal_deformability'].keys())) + """.format( + m_units=list(conversion_dict["mass"].keys()), + r_units=list(conversion_dict["radius"].keys()), + l_units=list(conversion_dict["tidal_deformability"].keys()), + ) raise ValueError(s) xdat = rep_dict[varnames[1]][0] * conversion_dict[xname][xunits] ydat = rep_dict[varnames[0]][0] * conversion_dict[yname][yunits] - xlabel = rep_dict[varnames[1]][1].replace('_', ' ') - ylabel = rep_dict[varnames[0]][1].replace('_', ' ') + '(' + xlabel + ')' + xlabel = rep_dict[varnames[1]][1].replace("_", " ") + ylabel = rep_dict[varnames[0]][1].replace("_", " ") + "(" + xlabel + ")" # Determine plot ranges. Currently shows 10% wider than actual data range. if xlim is None: diff --git a/bilby/gw/eos/tov_solver.py b/bilby/gw/eos/tov_solver.py index 0135086ef..11c77d75e 100644 --- a/bilby/gw/eos/tov_solver.py +++ b/bilby/gw/eos/tov_solver.py @@ -4,8 +4,7 @@ class IntegrateTOV: - """Class that given an initial pressure a mass radius value and a k2-love number - """ + """Class that given an initial pressure a mass radius value and a k2-love number""" def __init__(self, eos, eps_0): self.eos = eos @@ -18,8 +17,8 @@ def __init__(self, eos, eps_0): mass0, radius0 = self.__mass_radius_cent(pseudo_enthalpy0, self.pseudo_enthalpy) # k2 integration starting vals - H0 = radius0 ** 2 - B0 = 2. * radius0 + H0 = radius0**2 + B0 = 2.0 * radius0 self.y = np.array([mass0, radius0, H0, B0]) @@ -34,12 +33,26 @@ def __mass_radius_cent(self, pseudo_enthalpy0, pseudo_enthalpy): p_c = self.eos.pressure_from_pseudo_enthalpy(pseudo_enthalpy0) depsdh_c = self.eos.dedh(pseudo_enthalpy0) - radius = ((3. * (pseudo_enthalpy0 - pseudo_enthalpy)) / (2. * np.pi * (eps_c + 3. * p_c))) ** 0.5 \ - * (1. - 0.25 * (eps_c - 3. * p_c - (3. / 5.) * depsdh_c) * - ((pseudo_enthalpy0 - pseudo_enthalpy) / (eps_c + 3. * p_c))) - - mass = (4. * np.pi) / 3. * eps_c * radius ** 3 * (1. - (3. / 5.) * - depsdh_c * (pseudo_enthalpy0 - pseudo_enthalpy) / eps_c) + radius = ( + (3.0 * (pseudo_enthalpy0 - pseudo_enthalpy)) + / (2.0 * np.pi * (eps_c + 3.0 * p_c)) + ) ** 0.5 * ( + 1.0 + - 0.25 + * (eps_c - 3.0 * p_c - (3.0 / 5.0) * depsdh_c) + * ((pseudo_enthalpy0 - pseudo_enthalpy) / (eps_c + 3.0 * p_c)) + ) + + mass = ( + (4.0 * np.pi) + / 3.0 + * eps_c + * radius**3 + * ( + 1.0 + - (3.0 / 5.0) * depsdh_c * (pseudo_enthalpy0 - pseudo_enthalpy) / eps_c + ) + ) return mass, radius @@ -59,25 +72,31 @@ def __tov_eqns(self, h, y): r = y[1] H = y[2] B = y[3] - eps = self.eos.energy_density_from_pseudo_enthalpy(h, interp_type='CubicSpline') - p = self.eos.pressure_from_pseudo_enthalpy(h, interp_type='CubicSpline') - depsdp = self.eos.dedp(p, interp_type='CubicSpline') + eps = self.eos.energy_density_from_pseudo_enthalpy(h, interp_type="CubicSpline") + p = self.eos.pressure_from_pseudo_enthalpy(h, interp_type="CubicSpline") + depsdp = self.eos.dedp(p, interp_type="CubicSpline") - dmdh = (- (4. * np.pi * eps * r ** 3 * (r - 2. * m)) / - (m + 4. * np.pi * r ** 3 * p)) + dmdh = -(4.0 * np.pi * eps * r**3 * (r - 2.0 * m)) / ( + m + 4.0 * np.pi * r**3 * p + ) - drdh = -(r * (r - 2. * m)) / (m + 4. * np.pi * r ** 3 * p) + drdh = -(r * (r - 2.0 * m)) / (m + 4.0 * np.pi * r**3 * p) dHdh = B * drdh # taken from Damour & Nagar - e_lam = (1. - 2. * m / r) ** (-1) - - C1 = 2. / r + e_lam * (2. * m / r ** 2. + 4. * np.pi * r * (p - eps)) - C0 = (e_lam * (- 6. / r ** 2. + 4. * np.pi * (eps + p) * - depsdp + 4. * np.pi * (5. * eps + 9. * p)) - - (2. * (m + 4. * np.pi * r ** 3. * p) / - (r ** 2. - 2. * m * r)) ** 2.) + e_lam = (1.0 - 2.0 * m / r) ** (-1) + + C1 = 2.0 / r + e_lam * (2.0 * m / r**2.0 + 4.0 * np.pi * r * (p - eps)) + C0 = ( + e_lam + * ( + -6.0 / r**2.0 + + 4.0 * np.pi * (eps + p) * depsdp + + 4.0 * np.pi * (5.0 * eps + 9.0 * p) + ) + - (2.0 * (m + 4.0 * np.pi * r**3.0 * p) / (r**2.0 - 2.0 * m * r)) ** 2.0 + ) dBdh = -(C1 * B + C0 * H) * drdh @@ -96,13 +115,19 @@ def __calc_k2(self, R, Beta, H, C): y = (R * Beta) / H - num = ((8. / 5.) * (1. - 2. * C) ** 2 * - C ** 5 * (2. * C * (y - 1.) - y + 2.)) - denom = (2. * C * (4. * (y + 1.) * C ** 4 + (6. * y - 4.) * C ** 3 + - (26. - 22. * y) * C ** 2 + 3. * (5. * y - 8.) * - C - 3. * y + 6.) - 3. * (1. - 2 * C) ** 2 * - (2. * C * (y - 1.) - y + 2.) * - np.log(1. / (1. - 2. * C))) + num = ( + (8.0 / 5.0) * (1.0 - 2.0 * C) ** 2 * C**5 * (2.0 * C * (y - 1.0) - y + 2.0) + ) + denom = 2.0 * C * ( + 4.0 * (y + 1.0) * C**4 + + (6.0 * y - 4.0) * C**3 + + (26.0 - 22.0 * y) * C**2 + + 3.0 * (5.0 * y - 8.0) * C + - 3.0 * y + + 6.0 + ) - 3.0 * (1.0 - 2 * C) ** 2 * (2.0 * C * (y - 1.0) - y + 2.0) * np.log( + 1.0 / (1.0 - 2.0 * C) + ) return num / denom @@ -116,8 +141,13 @@ def integrate_TOV(self): rel_err = 1e-4 abs_err = 0.0 - result = solve_ivp(self.__tov_eqns, (self.pseudo_enthalpy, 1e-16), self.y, rtol=rel_err, - atol=abs_err) + result = solve_ivp( + self.__tov_eqns, + (self.pseudo_enthalpy, 1e-16), + self.y, + rtol=rel_err, + atol=abs_err, + ) m_fin = result.y[0, -1] r_fin = result.y[1, -1] H_fin = result.y[2, -1] diff --git a/bilby/gw/geometry.py b/bilby/gw/geometry.py index fa4be3d86..4f8a69022 100644 --- a/bilby/gw/geometry.py +++ b/bilby/gw/geometry.py @@ -44,7 +44,9 @@ def antenna_response(detector_tensor, ra, dec, time, psi, mode): The antenna response (scalar or array depending on input). """ xp = array_module(detector_tensor) - polarization_tensor = get_polarization_tensor(*promote_to_array((ra, dec, time, psi), xp), mode) + polarization_tensor = get_polarization_tensor( + *promote_to_array((ra, dec, time, psi), xp), mode + ) return three_by_three_matrix_contraction(detector_tensor, polarization_tensor) @@ -153,14 +155,16 @@ def get_polarization_tensor(ra, dec, time, psi, mode): -xp.sin(theta) * xp.ones_like(phi), ] ) - v = xp.asarray([ - -xp.sin(phi), xp.cos(phi), xp.zeros_like(phi) - ]) * xp.ones_like(theta) - omega = xp.asarray([ - xp.sin(xp.pi - theta) * xp.cos(xp.pi + phi), - xp.sin(xp.pi - theta) * xp.sin(xp.pi + phi), - xp.cos(xp.pi - theta) * xp.ones_like(phi), - ]) + v = xp.asarray([-xp.sin(phi), xp.cos(phi), xp.zeros_like(phi)]) * xp.ones_like( + theta + ) + omega = xp.asarray( + [ + xp.sin(xp.pi - theta) * xp.cos(xp.pi + phi), + xp.sin(xp.pi - theta) * xp.sin(xp.pi + phi), + xp.cos(xp.pi - theta) * xp.ones_like(phi), + ] + ) m = -u * xp.sin(psi) - v * xp.cos(psi) n = -u * xp.cos(psi) + v * xp.sin(psi) if xp.__name__ == "mlx.core": diff --git a/bilby/gw/likelihood/__init__.py b/bilby/gw/likelihood/__init__.py index d752f77a9..4a5bd5b7d 100644 --- a/bilby/gw/likelihood/__init__.py +++ b/bilby/gw/likelihood/__init__.py @@ -9,7 +9,7 @@ def get_binary_black_hole_likelihood(interferometers): - """ A wrapper to quickly set up a likelihood for BBH parameter estimation + """A wrapper to quickly set up a likelihood for BBH parameter estimation Parameters ========== @@ -27,8 +27,9 @@ def get_binary_black_hole_likelihood(interferometers): duration=interferometers.duration, sampling_frequency=interferometers.sampling_frequency, frequency_domain_source_model=lal_binary_black_hole, - waveform_arguments={'waveform_approximant': 'IMRPhenomPv2', - 'reference_frequency': 50}) + waveform_arguments={ + "waveform_approximant": "IMRPhenomPv2", + "reference_frequency": 50, + }, + ) return GravitationalWaveTransient(interferometers, waveform_generator) - - diff --git a/bilby/gw/likelihood/base.py b/bilby/gw/likelihood/base.py index 8a7598680..b9962dedd 100644 --- a/bilby/gw/likelihood/base.py +++ b/bilby/gw/likelihood/base.py @@ -1,4 +1,3 @@ - import os import copy @@ -16,7 +15,7 @@ class GravitationalWaveTransient(Likelihood): - """ A gravitational-wave transient likelihood object + """A gravitational-wave transient likelihood object This is the usual likelihood object to use for transient gravitational wave parameter estimation. It computes the log-likelihood in the frequency @@ -143,18 +142,27 @@ def snrs_as_sample(self) -> dict: The dictionary of SNRs labelled accordingly """ return { - "matched_filter_snr" : self.complex_matched_filter_snr, - "optimal_snr" : self.optimal_snr_squared.real ** 0.5 + "matched_filter_snr": self.complex_matched_filter_snr, + "optimal_snr": self.optimal_snr_squared.real**0.5, } def __init__( - self, interferometers, waveform_generator, time_marginalization=False, - distance_marginalization=False, phase_marginalization=False, calibration_marginalization=False, priors=None, - distance_marginalization_lookup_table=None, calibration_lookup_table=None, - number_of_response_curves=1000, starting_index=0, jitter_time=True, reference_frame="sky", - time_reference="geocenter" + self, + interferometers, + waveform_generator, + time_marginalization=False, + distance_marginalization=False, + phase_marginalization=False, + calibration_marginalization=False, + priors=None, + distance_marginalization_lookup_table=None, + calibration_lookup_table=None, + number_of_response_curves=1000, + starting_index=0, + jitter_time=True, + reference_frame="sky", + time_reference="geocenter", ): - self.waveform_generator = waveform_generator super(GravitationalWaveTransient, self).__init__() self.interferometers = InterferometerList(interferometers) @@ -181,83 +189,100 @@ def __init__( self.reference_ifo = None if self.time_marginalization: - self._check_marginalized_prior_is_set(key='geocent_time') + self._check_marginalized_prior_is_set(key="geocent_time") self._setup_time_marginalization() - priors['geocent_time'] = float(self.interferometers.start_time) + priors["geocent_time"] = float(self.interferometers.start_time) if self.jitter_time: - priors['time_jitter'] = Uniform( - minimum=- self._delta_tc / 2, + priors["time_jitter"] = Uniform( + minimum=-self._delta_tc / 2, maximum=self._delta_tc / 2, - boundary='periodic', + boundary="periodic", name="time_jitter", - latex_label="$t_j$" + latex_label="$t_j$", ) - self._marginalized_parameters.append('geocent_time') + self._marginalized_parameters.append("geocent_time") elif self.jitter_time: logger.debug( "Time jittering requested with non-time-marginalised " - "likelihood, ignoring.") + "likelihood, ignoring." + ) self.jitter_time = False if self.phase_marginalization: - self._check_marginalized_prior_is_set(key='phase') - priors['phase'] = float(0) - self._marginalized_parameters.append('phase') + self._check_marginalized_prior_is_set(key="phase") + priors["phase"] = float(0) + self._marginalized_parameters.append("phase") if self.distance_marginalization: self._lookup_table_filename = None - self._check_marginalized_prior_is_set(key='luminosity_distance') + self._check_marginalized_prior_is_set(key="luminosity_distance") self._distance_array = np.linspace( - self.priors['luminosity_distance'].minimum, - self.priors['luminosity_distance'].maximum, int(1e4)) + self.priors["luminosity_distance"].minimum, + self.priors["luminosity_distance"].maximum, + int(1e4), + ) self.distance_prior_array = np.array( - [self.priors['luminosity_distance'].prob(distance) - for distance in self._distance_array]) - self._ref_dist = self.priors['luminosity_distance'].rescale(0.5) - self._setup_distance_marginalization( - distance_marginalization_lookup_table) - for key in ['redshift', 'comoving_distance']: + [ + self.priors["luminosity_distance"].prob(distance) + for distance in self._distance_array + ] + ) + self._ref_dist = self.priors["luminosity_distance"].rescale(0.5) + self._setup_distance_marginalization(distance_marginalization_lookup_table) + for key in ["redshift", "comoving_distance"]: if key in priors: del priors[key] - priors['luminosity_distance'] = float(self._ref_dist) - self._marginalized_parameters.append('luminosity_distance') + priors["luminosity_distance"] = float(self._ref_dist) + self._marginalized_parameters.append("luminosity_distance") if self.calibration_marginalization: self.number_of_response_curves = number_of_response_curves self.starting_index = starting_index self._setup_calibration_marginalization(calibration_lookup_table, priors) - self._marginalized_parameters.append('recalib_index') + self._marginalized_parameters.append("recalib_index") def __repr__(self): - return self.__class__.__name__ + '(interferometers={},\n\twaveform_generator={},\n\ttime_marginalization={}, ' \ - 'distance_marginalization={}, phase_marginalization={}, ' \ - 'calibration_marginalization={}, priors={})' \ - .format(self.interferometers, self.waveform_generator, self.time_marginalization, - self.distance_marginalization, self.phase_marginalization, self.calibration_marginalization, - self.priors) + return ( + self.__class__.__name__ + + "(interferometers={},\n\twaveform_generator={},\n\ttime_marginalization={}, " + "distance_marginalization={}, phase_marginalization={}, " + "calibration_marginalization={}, priors={})".format( + self.interferometers, + self.waveform_generator, + self.time_marginalization, + self.distance_marginalization, + self.phase_marginalization, + self.calibration_marginalization, + self.priors, + ) + ) def _check_set_duration_and_sampling_frequency_of_waveform_generator(self): - """ Check the waveform_generator has the same duration and + """Check the waveform_generator has the same duration and sampling_frequency as the interferometers. If they are unset, then set them, if they differ, raise an error """ - attributes = ['duration', 'sampling_frequency', 'start_time'] + attributes = ["duration", "sampling_frequency", "start_time"] for attribute in attributes: wfg_attr = getattr(self.waveform_generator, attribute) ifo_attr = getattr(self.interferometers, attribute) if wfg_attr is None: logger.debug( "The waveform_generator {} is None. Setting from the " - "provided interferometers.".format(attribute)) + "provided interferometers.".format(attribute) + ) elif wfg_attr != ifo_attr: logger.debug( "The waveform_generator {} is not equal to that of the " "provided interferometers. Overwriting the " - "waveform_generator.".format(attribute)) + "waveform_generator.".format(attribute) + ) setattr(self.waveform_generator, attribute, ifo_attr) - def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array=True, parameters): + def calculate_snrs( + self, waveform_polarizations, interferometer, *, return_array=True, parameters + ): """ Compute the snrs @@ -286,8 +311,10 @@ def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array ) _mask = interferometer.frequency_mask - if 'recalib_index' in parameters: - signal[_mask] *= self.calibration_draws[interferometer.name][int(parameters['recalib_index'])] + if "recalib_index" in parameters: + signal[_mask] *= self.calibration_draws[interferometer.name][ + int(parameters["recalib_index"]) + ] d_inner_h = interferometer.inner_product(signal=signal) optimal_snr_squared = interferometer.optimal_snr_squared(signal=signal) @@ -303,23 +330,29 @@ def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array d_inner_h_array = None optimal_snr_squared_array = None elif self.time_marginalization and self.calibration_marginalization: - d_inner_h_integrand = xp.tile( - interferometer.frequency_domain_strain.conj() * signal / - interferometer.power_spectral_density_array, (self.number_of_response_curves, 1)).T + interferometer.frequency_domain_strain.conj() + * signal + / interferometer.power_spectral_density_array, + (self.number_of_response_curves, 1), + ).T d_inner_h_integrand[_mask] *= self.calibration_draws[interferometer.name].T - d_inner_h_array = 4 / self.waveform_generator.duration * xp.fft.fft( - d_inner_h_integrand[0:-1], axis=0 - ).T + d_inner_h_array = ( + 4 + / self.waveform_generator.duration + * xp.fft.fft(d_inner_h_integrand[0:-1], axis=0).T + ) optimal_snr_squared_integrand = ( - normalization * xp.abs(signal)**2 / interferometer.power_spectral_density_array + normalization + * xp.abs(signal) ** 2 + / interferometer.power_spectral_density_array ) optimal_snr_squared_array = xp.dot( optimal_snr_squared_integrand[_mask], - self.calibration_abs_draws[interferometer.name].T + self.calibration_abs_draws[interferometer.name].T, ) elif self.time_marginalization and not self.calibration_marginalization: @@ -329,20 +362,26 @@ def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array / interferometer.power_spectral_density_array[0:-1] ) - elif self.calibration_marginalization and ('recalib_index' not in parameters): + elif self.calibration_marginalization and ("recalib_index" not in parameters): d_inner_h_integrand = ( - normalization * - interferometer.frequency_domain_strain.conj() * signal + normalization + * interferometer.frequency_domain_strain.conj() + * signal / interferometer.power_spectral_density_array ) - d_inner_h_array = xp.dot(d_inner_h_integrand[_mask], self.calibration_draws[interferometer.name].T) + d_inner_h_array = xp.dot( + d_inner_h_integrand[_mask], + self.calibration_draws[interferometer.name].T, + ) optimal_snr_squared_integrand = ( - normalization * xp.abs(signal)**2 / interferometer.power_spectral_density_array + normalization + * xp.abs(signal) ** 2 + / interferometer.power_spectral_density_array ) optimal_snr_squared_array = xp.dot( optimal_snr_squared_integrand[_mask], - self.calibration_abs_draws[interferometer.name].T + self.calibration_abs_draws[interferometer.name].T, ) return self._CalculatedSNRs( @@ -358,30 +397,30 @@ def _check_marginalized_prior_is_set(self, key): raise ValueError( "Cannot use marginalized likelihood for {}: prior is fixed".format(key) ) - if key not in self.priors or not isinstance( - self.priors[key], Prior): - if key == 'geocent_time': + if key not in self.priors or not isinstance(self.priors[key], Prior): + if key == "geocent_time": logger.warning( - 'Prior not provided for geocent time, using the full segment.' + "Prior not provided for geocent time, using the full segment." ) self.priors[key] = Uniform( self.interferometers.start_time, - self.interferometers.start_time + self.interferometers.duration) - elif key == 'luminosity_distance': - for key in ['redshift', 'comoving_distance']: + self.interferometers.start_time + self.interferometers.duration, + ) + elif key == "luminosity_distance": + for key in ["redshift", "comoving_distance"]: if key in self.priors: if not isinstance(self.priors[key], Cosmological): raise TypeError( "To marginalize over {}, the prior must be specified as a " "subclass of bilby.gw.prior.Cosmological.".format(key) ) - self.priors['luminosity_distance'] = self.priors[key].get_corresponding_prior( - 'luminosity_distance' - ) + self.priors["luminosity_distance"] = self.priors[ + key + ].get_corresponding_prior("luminosity_distance") del self.priors[key] else: logger.warning( - 'Prior not provided for {}, using the BBH default.'.format(key) + "Prior not provided for {}, using the BBH default.".format(key) ) self.priors[key] = BBHPriorDict()[key] @@ -393,9 +432,16 @@ def priors(self): def priors(self, priors): if priors is not None: self._prior = priors.copy() - elif any([self.time_marginalization, self.phase_marginalization, - self.distance_marginalization]): - raise ValueError("You can't use a marginalized likelihood without specifying a priors") + elif any( + [ + self.time_marginalization, + self.phase_marginalization, + self.distance_marginalization, + ] + ): + raise ValueError( + "You can't use a marginalized likelihood without specifying a priors" + ) else: self._prior = None @@ -403,11 +449,15 @@ def _calculate_noise_log_likelihood(self): log_l = 0 for interferometer in self.interferometers: mask = interferometer.frequency_mask - log_l -= abs(noise_weighted_inner_product( - interferometer.frequency_domain_strain[mask], - interferometer.frequency_domain_strain[mask], - interferometer.power_spectral_density_array[mask], - self.waveform_generator.duration) / 2) + log_l -= abs( + noise_weighted_inner_product( + interferometer.frequency_domain_strain[mask], + interferometer.frequency_domain_strain[mask], + interferometer.power_spectral_density_array[mask], + self.waveform_generator.duration, + ) + / 2 + ) return log_l def noise_log_likelihood(self): @@ -419,13 +469,14 @@ def noise_log_likelihood(self): def log_likelihood_ratio(self, parameters): parameters = copy.deepcopy(parameters) parameters.update(self.get_sky_frame_parameters(parameters)) - waveform_polarizations = \ - self.waveform_generator.frequency_domain_strain(parameters) + waveform_polarizations = self.waveform_generator.frequency_domain_strain( + parameters + ) if waveform_polarizations is None: return np.nan_to_num(-np.inf) if self.time_marginalization and self.jitter_time: - parameters['geocent_time'] += parameters['time_jitter'] + parameters["geocent_time"] += parameters["time_jitter"] total_snrs = self._CalculatedSNRs() @@ -441,12 +492,11 @@ def log_likelihood_ratio(self, parameters): log_l = self.compute_log_likelihood_from_snrs(total_snrs, parameters=parameters) if self.time_marginalization and self.jitter_time: - parameters['geocent_time'] -= parameters['time_jitter'] + parameters["geocent_time"] -= parameters["time_jitter"] return log_l.real def compute_log_likelihood_from_snrs(self, total_snrs, parameters): - if self.calibration_marginalization: log_l = self.calibration_marginalized_likelihood( d_inner_h_calibration_array=total_snrs.d_inner_h_array, @@ -463,13 +513,15 @@ def compute_log_likelihood_from_snrs(self, total_snrs, parameters): elif self.distance_marginalization: log_l = self.distance_marginalized_likelihood( - d_inner_h=total_snrs.d_inner_h, h_inner_h=total_snrs.optimal_snr_squared, + d_inner_h=total_snrs.d_inner_h, + h_inner_h=total_snrs.optimal_snr_squared, parameters=parameters, ) elif self.phase_marginalization: log_l = self.phase_marginalized_likelihood( - d_inner_h=total_snrs.d_inner_h, h_inner_h=total_snrs.optimal_snr_squared) + d_inner_h=total_snrs.d_inner_h, h_inner_h=total_snrs.optimal_snr_squared + ) else: log_l = np.real(total_snrs.d_inner_h) - total_snrs.optimal_snr_squared / 2 @@ -478,11 +530,12 @@ def compute_log_likelihood_from_snrs(self, total_snrs, parameters): def compute_per_detector_log_likelihood(self, parameters): parameters.update(self.get_sky_frame_parameters(parameters)) - waveform_polarizations = \ - self.waveform_generator.frequency_domain_strain(parameters) + waveform_polarizations = self.waveform_generator.frequency_domain_strain( + parameters + ) if self.time_marginalization and self.jitter_time: - parameters['geocent_time'] += parameters['time_jitter'] + parameters["geocent_time"] += parameters["time_jitter"] for interferometer in self.interferometers: per_detector_snr = self.calculate_snrs( @@ -491,11 +544,14 @@ def compute_per_detector_log_likelihood(self, parameters): parameters=parameters, ) - parameters['{}_log_likelihood'.format(interferometer.name)] = \ - self.compute_log_likelihood_from_snrs(per_detector_snr, parameters=parameters) + parameters["{}_log_likelihood".format(interferometer.name)] = ( + self.compute_log_likelihood_from_snrs( + per_detector_snr, parameters=parameters + ) + ) if self.time_marginalization and self.jitter_time: - parameters['geocent_time'] -= parameters['time_jitter'] + parameters["geocent_time"] -= parameters["time_jitter"] return parameters.copy() @@ -518,31 +574,38 @@ def generate_posterior_sample_from_marginalized_likelihood(self, parameters): """ if len(self._marginalized_parameters) > 0: signal_polarizations = copy.deepcopy( - self.waveform_generator.frequency_domain_strain( - parameters)) + self.waveform_generator.frequency_domain_strain(parameters) + ) else: return parameters if self.calibration_marginalization: - new_calibration = self.generate_calibration_sample_from_marginalized_likelihood( - signal_polarizations=signal_polarizations, parameters=parameters) - parameters['recalib_index'] = new_calibration + new_calibration = ( + self.generate_calibration_sample_from_marginalized_likelihood( + signal_polarizations=signal_polarizations, parameters=parameters + ) + ) + parameters["recalib_index"] = new_calibration if self.time_marginalization: new_time = self.generate_time_sample_from_marginalized_likelihood( - signal_polarizations=signal_polarizations, parameters=parameters) - parameters['geocent_time'] = new_time + signal_polarizations=signal_polarizations, parameters=parameters + ) + parameters["geocent_time"] = new_time if self.distance_marginalization: new_distance = self.generate_distance_sample_from_marginalized_likelihood( - signal_polarizations=signal_polarizations, parameters=parameters) - parameters['luminosity_distance'] = new_distance + signal_polarizations=signal_polarizations, parameters=parameters + ) + parameters["luminosity_distance"] = new_distance if self.phase_marginalization: new_phase = self.generate_phase_sample_from_marginalized_likelihood( - signal_polarizations=signal_polarizations, parameters=parameters) - parameters['phase'] = new_phase + signal_polarizations=signal_polarizations, parameters=parameters + ) + parameters["phase"] = new_phase return parameters.copy() def generate_calibration_sample_from_marginalized_likelihood( - self, signal_polarizations=None, *, parameters): + self, signal_polarizations=None, *, parameters + ): """ Generate a single sample from the posterior distribution for the set of calibration response curves when explicitly marginalizing over the calibration uncertainty. @@ -559,12 +622,13 @@ def generate_calibration_sample_from_marginalized_likelihood( """ from ...core.utils import random - if 'recalib_index' in parameters: - parameters.pop('recalib_index') + if "recalib_index" in parameters: + parameters.pop("recalib_index") parameters.update(self.get_sky_frame_parameters(parameters)) if signal_polarizations is None: - signal_polarizations = \ - self.waveform_generator.frequency_domain_strain(parameters) + signal_polarizations = self.waveform_generator.frequency_domain_strain( + parameters + ) log_like = self.get_calibration_log_likelihoods( signal_polarizations=signal_polarizations, parameters=parameters @@ -573,12 +637,15 @@ def generate_calibration_sample_from_marginalized_likelihood( calibration_post = np.exp(log_like - max(log_like)) calibration_post /= np.sum(calibration_post) - new_calibration = random.rng.choice(self.number_of_response_curves, p=calibration_post) + new_calibration = random.rng.choice( + self.number_of_response_curves, p=calibration_post + ) return new_calibration def generate_time_sample_from_marginalized_likelihood( - self, signal_polarizations=None, *, parameters): + self, signal_polarizations=None, *, parameters + ): """ Generate a single sample from the posterior distribution for coalescence time when using a likelihood which explicitly marginalises over time. @@ -599,15 +666,18 @@ def generate_time_sample_from_marginalized_likelihood( """ parameters.update(self.get_sky_frame_parameters(parameters)) if self.jitter_time: - parameters['geocent_time'] += parameters['time_jitter'] + parameters["geocent_time"] += parameters["time_jitter"] if signal_polarizations is None: - signal_polarizations = \ - self.waveform_generator.frequency_domain_strain(parameters) + signal_polarizations = self.waveform_generator.frequency_domain_strain( + parameters + ) times = create_time_series( sampling_frequency=16384, - starting_time=parameters['geocent_time'] - self.waveform_generator.start_time, - duration=self.waveform_generator.duration) + starting_time=parameters["geocent_time"] + - self.waveform_generator.start_time, + duration=self.waveform_generator.duration, + ) times = times % self.waveform_generator.duration times += self.waveform_generator.start_time @@ -637,16 +707,17 @@ def generate_time_sample_from_marginalized_likelihood( if self.distance_marginalization: time_log_like = self.distance_marginalized_likelihood( - d_inner_h, h_inner_h, parameters=parameters) + d_inner_h, h_inner_h, parameters=parameters + ) elif self.phase_marginalization: time_log_like = ln_i0(abs(d_inner_h)) - h_inner_h.real / 2 else: - time_log_like = (d_inner_h.real - h_inner_h.real / 2) + time_log_like = d_inner_h.real - h_inner_h.real / 2 - time_prior_array = self.priors['geocent_time'].prob(times) + time_prior_array = self.priors["geocent_time"].prob(times) time_post = np.exp(time_log_like - max(time_log_like)) * time_prior_array - keep = (time_post > max(time_post) / 1000) + keep = time_post > max(time_post) / 1000 if sum(keep) < 3: keep[1:-1] = keep[1:-1] | keep[2:] | keep[:-2] time_post = time_post[keep] @@ -656,7 +727,8 @@ def generate_time_sample_from_marginalized_likelihood( return new_time def generate_distance_sample_from_marginalized_likelihood( - self, signal_polarizations=None, *, parameters): + self, signal_polarizations=None, *, parameters + ): """ Generate a single sample from the posterior distribution for luminosity distance when using a likelihood which explicitly marginalises over @@ -678,31 +750,33 @@ def generate_distance_sample_from_marginalized_likelihood( """ parameters.update(self.get_sky_frame_parameters(parameters)) if signal_polarizations is None: - signal_polarizations = \ - self.waveform_generator.frequency_domain_strain(parameters) + signal_polarizations = self.waveform_generator.frequency_domain_strain( + parameters + ) d_inner_h, h_inner_h = self._calculate_inner_products( signal_polarizations, parameters=parameters ) d_inner_h_dist = ( - d_inner_h * parameters['luminosity_distance'] / self._distance_array + d_inner_h * parameters["luminosity_distance"] / self._distance_array ) h_inner_h_dist = ( - h_inner_h * parameters['luminosity_distance']**2 / self._distance_array**2 + h_inner_h * parameters["luminosity_distance"] ** 2 / self._distance_array**2 ) if self.phase_marginalization: distance_log_like = ln_i0(abs(d_inner_h_dist)) - h_inner_h_dist.real / 2 else: - distance_log_like = (d_inner_h_dist.real - h_inner_h_dist.real / 2) + distance_log_like = d_inner_h_dist.real - h_inner_h_dist.real / 2 - distance_post = (np.exp(distance_log_like - max(distance_log_like)) * - self.distance_prior_array) + distance_post = ( + np.exp(distance_log_like - max(distance_log_like)) + * self.distance_prior_array + ) - new_distance = Interped( - self._distance_array, distance_post).sample() + new_distance = Interped(self._distance_array, distance_post).sample() self._rescale_signal(signal_polarizations, new_distance) return new_distance @@ -712,7 +786,8 @@ def _calculate_inner_products(self, signal_polarizations, parameters): h_inner_h = 0 for interferometer in self.interferometers: per_detector_snr = self.calculate_snrs( - signal_polarizations, interferometer, parameters=parameters) + signal_polarizations, interferometer, parameters=parameters + ) d_inner_h += per_detector_snr.d_inner_h h_inner_h += per_detector_snr.optimal_snr_squared @@ -736,7 +811,8 @@ def _compute_full_waveform(self, signal_polarizations, interferometer, parameter return interferometer.get_detector_response(signal_polarizations, parameters) def generate_phase_sample_from_marginalized_likelihood( - self, signal_polarizations=None, *, parameters): + self, signal_polarizations=None, *, parameters + ): r""" Generate a single sample from the posterior distribution for phase when using a likelihood which explicitly marginalises over phase. @@ -759,8 +835,9 @@ def generate_phase_sample_from_marginalized_likelihood( """ parameters.update(self.get_sky_frame_parameters(parameters)) if signal_polarizations is None: - signal_polarizations = \ - self.waveform_generator.frequency_domain_strain(parameters) + signal_polarizations = self.waveform_generator.frequency_domain_strain( + parameters + ) d_inner_h, h_inner_h = self._calculate_inner_products( signal_polarizations, parameters=parameters ) @@ -774,14 +851,14 @@ def generate_phase_sample_from_marginalized_likelihood( def distance_marginalized_likelihood(self, d_inner_h, h_inner_h, *, parameters): d_inner_h_ref, h_inner_h_ref = self._setup_rho( - d_inner_h, h_inner_h, parameters=parameters) + d_inner_h, h_inner_h, parameters=parameters + ) if self.phase_marginalization: d_inner_h_ref = abs(d_inner_h_ref) else: d_inner_h_ref = d_inner_h_ref.real - return self._interp_dist_margd_loglikelihood( - d_inner_h_ref, h_inner_h_ref) + return self._interp_dist_margd_loglikelihood(d_inner_h_ref, h_inner_h_ref) def phase_marginalized_likelihood(self, d_inner_h, h_inner_h): d_inner_h = ln_i0(abs(d_inner_h)) @@ -791,28 +868,31 @@ def phase_marginalized_likelihood(self, d_inner_h, h_inner_h): else: return d_inner_h - h_inner_h / 2 - def time_marginalized_likelihood(self, d_inner_h_tc_array, h_inner_h, *, parameters): + def time_marginalized_likelihood( + self, d_inner_h_tc_array, h_inner_h, *, parameters + ): times = self._times if self.jitter_time: - times = self._times + parameters['time_jitter'] + times = self._times + parameters["time_jitter"] if not aac.is_jax_array(d_inner_h_tc_array): - _time_prior = self.priors['geocent_time'] + _time_prior = self.priors["geocent_time"] time_mask = (times >= _time_prior.minimum) & (times <= _time_prior.maximum) times = times[time_mask] if self.calibration_marginalization: d_inner_h_tc_array = d_inner_h_tc_array[:, time_mask] else: d_inner_h_tc_array = d_inner_h_tc_array[time_mask] - time_prior_array = self.priors['geocent_time'].prob(times) * self._delta_tc + time_prior_array = self.priors["geocent_time"].prob(times) * self._delta_tc if self.distance_marginalization: log_l_tc_array = self.distance_marginalized_likelihood( - d_inner_h=d_inner_h_tc_array, h_inner_h=h_inner_h, parameters=parameters) + d_inner_h=d_inner_h_tc_array, h_inner_h=h_inner_h, parameters=parameters + ) elif self.phase_marginalization: log_l_tc_array = self.phase_marginalized_likelihood( - d_inner_h=d_inner_h_tc_array, - h_inner_h=h_inner_h) + d_inner_h=d_inner_h_tc_array, h_inner_h=h_inner_h + ) elif self.calibration_marginalization: log_l_tc_array = d_inner_h_tc_array.real - h_inner_h[:, np.newaxis] / 2 else: @@ -822,8 +902,9 @@ def time_marginalized_likelihood(self, d_inner_h_tc_array, h_inner_h, *, paramet def get_calibration_log_likelihoods(self, signal_polarizations=None, *, parameters): parameters.update(self.get_sky_frame_parameters(parameters)) if signal_polarizations is None: - signal_polarizations = \ - self.waveform_generator.frequency_domain_strain(parameters) + signal_polarizations = self.waveform_generator.frequency_domain_strain( + parameters + ) total_snrs = self._CalculatedSNRs() @@ -850,14 +931,18 @@ def get_calibration_log_likelihoods(self, signal_polarizations=None, *, paramete elif self.phase_marginalization: log_l_cal_array = self.phase_marginalized_likelihood( d_inner_h=total_snrs.d_inner_h_array, - h_inner_h=total_snrs.optimal_snr_squared_array) + h_inner_h=total_snrs.optimal_snr_squared_array, + ) else: - log_l_cal_array = \ - np.real(total_snrs.d_inner_h_array - total_snrs.optimal_snr_squared_array / 2) + log_l_cal_array = np.real( + total_snrs.d_inner_h_array - total_snrs.optimal_snr_squared_array / 2 + ) return log_l_cal_array - def calibration_marginalized_likelihood(self, d_inner_h_calibration_array, h_inner_h, *, parameters): + def calibration_marginalized_likelihood( + self, d_inner_h_calibration_array, h_inner_h, *, parameters + ): if self.time_marginalization: log_l_cal_array = self.time_marginalized_likelihood( d_inner_h_tc_array=d_inner_h_calibration_array, @@ -866,26 +951,33 @@ def calibration_marginalized_likelihood(self, d_inner_h_calibration_array, h_inn ) elif self.distance_marginalization: log_l_cal_array = self.distance_marginalized_likelihood( - d_inner_h=d_inner_h_calibration_array, h_inner_h=h_inner_h, parameters=parameters) + d_inner_h=d_inner_h_calibration_array, + h_inner_h=h_inner_h, + parameters=parameters, + ) elif self.phase_marginalization: log_l_cal_array = self.phase_marginalized_likelihood( - d_inner_h=d_inner_h_calibration_array, - h_inner_h=h_inner_h) + d_inner_h=d_inner_h_calibration_array, h_inner_h=h_inner_h + ) else: log_l_cal_array = np.real(d_inner_h_calibration_array - h_inner_h / 2) return logsumexp(log_l_cal_array) - np.log(self.number_of_response_curves) def _setup_rho(self, d_inner_h, optimal_snr_squared, parameters): - optimal_snr_squared_ref = (optimal_snr_squared.real * - parameters['luminosity_distance'] ** 2 / - self._ref_dist ** 2.) - d_inner_h_ref = (d_inner_h * parameters['luminosity_distance'] / - self._ref_dist) + optimal_snr_squared_ref = ( + optimal_snr_squared.real + * parameters["luminosity_distance"] ** 2 + / self._ref_dist**2.0 + ) + d_inner_h_ref = d_inner_h * parameters["luminosity_distance"] / self._ref_dist return d_inner_h_ref, optimal_snr_squared_ref def log_likelihood(self, parameters): - return self.log_likelihood_ratio(parameters=parameters) + self.noise_log_likelihood() + return ( + self.log_likelihood_ratio(parameters=parameters) + + self.noise_log_likelihood() + ) @property def _delta_distance(self): @@ -893,35 +985,33 @@ def _delta_distance(self): @property def _dist_multiplier(self): - ''' Maximum value of ref_dist/dist_array ''' + """Maximum value of ref_dist/dist_array""" return self._ref_dist / self._distance_array[0] @property def _optimal_snr_squared_ref_array(self): - """ Optimal filter snr at fiducial distance of ref_dist Mpc """ + """Optimal filter snr at fiducial distance of ref_dist Mpc""" return np.logspace(-5, 10, self._dist_margd_loglikelihood_array.shape[0]) @property def _d_inner_h_ref_array(self): - """ Matched filter snr at fiducial distance of ref_dist Mpc """ + """Matched filter snr at fiducial distance of ref_dist Mpc""" if self.phase_marginalization: return np.logspace(-5, 10, self._dist_margd_loglikelihood_array.shape[1]) else: n_negative = self._dist_margd_loglikelihood_array.shape[1] // 2 n_positive = self._dist_margd_loglikelihood_array.shape[1] - n_negative - return np.hstack(( - -np.logspace(3, -3, n_negative), np.logspace(-3, 10, n_positive) - )) + return np.hstack( + (-np.logspace(3, -3, n_negative), np.logspace(-3, 10, n_positive)) + ) def _setup_distance_marginalization(self, lookup_table=None): if isinstance(lookup_table, str) or lookup_table is None: self.cached_lookup_table_filename = lookup_table - lookup_table = self.load_lookup_table( - self.cached_lookup_table_filename) + lookup_table = self.load_lookup_table(self.cached_lookup_table_filename) if isinstance(lookup_table, dict): if self._test_cached_lookup_table(lookup_table): - self._dist_margd_loglikelihood_array = lookup_table[ - 'lookup_table'] + self._dist_margd_loglikelihood_array = lookup_table["lookup_table"] else: self._create_lookup_table() else: @@ -936,15 +1026,14 @@ def _setup_distance_marginalization(self, lookup_table=None): @property def cached_lookup_table_filename(self): if self._lookup_table_filename is None: - self._lookup_table_filename = ( - '.distance_marginalization_lookup.npz') + self._lookup_table_filename = ".distance_marginalization_lookup.npz" return self._lookup_table_filename @cached_lookup_table_filename.setter def cached_lookup_table_filename(self, filename): if isinstance(filename, str): - if filename[-4:] != '.npz': - filename += '.npz' + if filename[-4:] != ".npz": + filename += ".npz" self._lookup_table_filename = filename def load_lookup_table(self, filename): @@ -957,59 +1046,69 @@ def load_lookup_table(self, filename): return None match, failure = self._test_cached_lookup_table(loaded_file) if match: - logger.info('Loaded distance marginalisation lookup table from ' - '{}.'.format(filename)) + logger.info( + "Loaded distance marginalisation lookup table from {}.".format( + filename + ) + ) return loaded_file else: - logger.info('Loaded distance marginalisation lookup table does ' - 'not match for {}.'.format(failure)) + logger.info( + "Loaded distance marginalisation lookup table does " + "not match for {}.".format(failure) + ) elif isinstance(filename, str): - logger.info('Distance marginalisation file {} does not ' - 'exist'.format(filename)) + logger.info( + "Distance marginalisation file {} does not exist".format(filename) + ) return None def cache_lookup_table(self): - np.savez(self.cached_lookup_table_filename, - distance_array=self._distance_array, - prior_array=self.distance_prior_array, - lookup_table=self._dist_margd_loglikelihood_array, - reference_distance=self._ref_dist, - phase_marginalization=self.phase_marginalization) + np.savez( + self.cached_lookup_table_filename, + distance_array=self._distance_array, + prior_array=self.distance_prior_array, + lookup_table=self._dist_margd_loglikelihood_array, + reference_distance=self._ref_dist, + phase_marginalization=self.phase_marginalization, + ) def _test_cached_lookup_table(self, loaded_file): pairs = dict( distance_array=self._distance_array, prior_array=self.distance_prior_array, reference_distance=self._ref_dist, - phase_marginalization=self.phase_marginalization) + phase_marginalization=self.phase_marginalization, + ) for key in pairs: if key not in loaded_file: return False, key - elif not np.allclose(np.atleast_1d(loaded_file[key]), - np.atleast_1d(pairs[key]), - rtol=1e-15): + elif not np.allclose( + np.atleast_1d(loaded_file[key]), np.atleast_1d(pairs[key]), rtol=1e-15 + ): return False, key return True, None def _create_lookup_table(self): - """ Make the lookup table """ + """Make the lookup table""" from tqdm.auto import tqdm - logger.info('Building lookup table for distance marginalisation.') + + logger.info("Building lookup table for distance marginalisation.") self._dist_margd_loglikelihood_array = np.zeros((400, 800)) scaling = self._ref_dist / self._distance_array d_inner_h_array_full = np.outer(self._d_inner_h_ref_array, scaling) - h_inner_h_array_full = np.outer(self._optimal_snr_squared_ref_array, scaling ** 2) + h_inner_h_array_full = np.outer(self._optimal_snr_squared_ref_array, scaling**2) if self.phase_marginalization: d_inner_h_array_full = ln_i0(abs(d_inner_h_array_full)) prior_term = self.distance_prior_array * self._delta_distance for ii, optimal_snr_squared_array in tqdm( - enumerate(h_inner_h_array_full), total=len(self._optimal_snr_squared_ref_array) + enumerate(h_inner_h_array_full), + total=len(self._optimal_snr_squared_ref_array), ): for jj, d_inner_h_array in enumerate(d_inner_h_array_full): self._dist_margd_loglikelihood_array[ii][jj] = logsumexp( - d_inner_h_array - optimal_snr_squared_array / 2, - b=prior_term + d_inner_h_array - optimal_snr_squared_array / 2, b=prior_term ) log_norm = logsumexp( 0 / self._distance_array, b=self.distance_prior_array * self._delta_distance @@ -1026,21 +1125,32 @@ def _setup_phase_marginalization(self, min_bound=-5, max_bound=10): def _setup_time_marginalization(self): self._delta_tc = 2 / self.waveform_generator.sampling_frequency - self._times = \ - self.interferometers.start_time + np.linspace( - 0, self.interferometers.duration, - int(self.interferometers.duration / 2 * - self.waveform_generator.sampling_frequency + 1))[1:] - self.time_prior_array = \ - self.priors['geocent_time'].prob(self._times) * self._delta_tc + self._times = ( + self.interferometers.start_time + + np.linspace( + 0, + self.interferometers.duration, + int( + self.interferometers.duration + / 2 + * self.waveform_generator.sampling_frequency + + 1 + ), + )[1:] + ) + self.time_prior_array = ( + self.priors["geocent_time"].prob(self._times) * self._delta_tc + ) def _setup_calibration_marginalization(self, calibration_lookup_table, priors=None): - self.calibration_draws, self.calibration_parameter_draws = calibration.build_calibration_lookup( - interferometers=self.interferometers, - lookup_files=calibration_lookup_table, - priors=priors, - number_of_response_curves=self.number_of_response_curves, - starting_index=self.starting_index, + self.calibration_draws, self.calibration_parameter_draws = ( + calibration.build_calibration_lookup( + interferometers=self.interferometers, + lookup_files=calibration_lookup_table, + priors=priors, + number_of_response_curves=self.number_of_response_curves, + starting_index=self.starting_index, + ) ) for name, parameters in self.calibration_parameter_draws.items(): if parameters is not None: @@ -1048,7 +1158,7 @@ def _setup_calibration_marginalization(self, calibration_lookup_table, priors=No priors[key] = DeltaFunction(0.0) self.calibration_abs_draws = dict() for name in self.calibration_draws: - self.calibration_abs_draws[name] = np.abs(self.calibration_draws[name])**2 + self.calibration_abs_draws[name] = np.abs(self.calibration_draws[name]) ** 2 @property def interferometers(self): @@ -1106,7 +1216,7 @@ def get_sky_frame_parameters(self, parameters): ======= dict: dictionary containing ra, dec, and geocent_time """ - time = parameters.get(f'{self.time_reference}_time', None) + time = parameters.get(f"{self.time_reference}_time", None) if time is None and "geocent_time" in parameters: logger.warning( f"Cannot find {self.time_reference}_time in parameters. " @@ -1115,8 +1225,11 @@ def get_sky_frame_parameters(self, parameters): if self.reference_frame is not None: try: ra, dec = zenith_azimuth_to_ra_dec( - parameters['zenith'], parameters['azimuth'], - time, self.reference_frame) + parameters["zenith"], + parameters["azimuth"], + time, + self.reference_frame, + ) except KeyError: if "ra" in parameters and "dec" in parameters: ra = parameters["ra"] @@ -1130,7 +1243,10 @@ def get_sky_frame_parameters(self, parameters): else: ra = parameters["ra"] dec = parameters["dec"] - if "geocent" not in self.time_reference and f"{self.time_reference}_time" in parameters: + if ( + "geocent" not in self.time_reference + and f"{self.time_reference}_time" in parameters + ): geocent_time = time - self.reference_ifo.time_delay_from_geocenter( ra=ra, dec=dec, time=time ) @@ -1142,11 +1258,14 @@ def get_sky_frame_parameters(self, parameters): def lal_version(self): try: from lal import git_version, __version__ + lal_version = str(__version__) logger.info("Using lal version {}".format(lal_version)) lal_git_version = str(git_version.verbose_msg).replace("\n", ";") logger.info("Using lal git version {}".format(lal_git_version)) - return "lal_version={}, lal_git_version={}".format(lal_version, lal_git_version) + return "lal_version={}, lal_git_version={}".format( + lal_version, lal_git_version + ) except (ImportError, AttributeError): return "N/A" @@ -1154,11 +1273,14 @@ def lal_version(self): def lalsimulation_version(self): try: from lalsimulation import git_version, __version__ + lalsim_version = str(__version__) logger.info("Using lalsimulation version {}".format(lalsim_version)) lalsim_git_version = str(git_version.verbose_msg).replace("\n", ";") logger.info("Using lalsimulation git version {}".format(lalsim_git_version)) - return "lalsimulation_version={}, lalsimulation_git_version={}".format(lalsim_version, lalsim_git_version) + return "lalsimulation_version={}, lalsimulation_git_version={}".format( + lalsim_version, lalsim_git_version + ) except (ImportError, AttributeError): return "N/A" @@ -1182,4 +1304,5 @@ def meta_data(self): reference_frame=self._reference_frame_str, lal_version=self.lal_version, lalsimulation_version=self.lalsimulation_version, - waveform_generator_meta_data=self.waveform_generator.meta_data) + waveform_generator_meta_data=self.waveform_generator.meta_data, + ) diff --git a/bilby/gw/likelihood/basic.py b/bilby/gw/likelihood/basic.py index 35b5994a6..abdac6b51 100644 --- a/bilby/gw/likelihood/basic.py +++ b/bilby/gw/likelihood/basic.py @@ -4,7 +4,6 @@ class BasicGravitationalWaveTransient(Likelihood): - def __init__(self, interferometers, waveform_generator): """ @@ -30,11 +29,15 @@ def __init__(self, interferometers, waveform_generator): self.waveform_generator = waveform_generator def __repr__(self): - return self.__class__.__name__ + '(interferometers={},\n\twaveform_generator={})' \ - .format(self.interferometers, self.waveform_generator) + return ( + self.__class__.__name__ + + "(interferometers={},\n\twaveform_generator={})".format( + self.interferometers, self.waveform_generator + ) + ) def noise_log_likelihood(self): - """ Calculates the real part of noise log-likelihood + """Calculates the real part of noise log-likelihood Returns ======= @@ -43,14 +46,18 @@ def noise_log_likelihood(self): """ log_l = 0 for interferometer in self.interferometers: - log_l -= 2. / self.waveform_generator.duration * ( - abs(interferometer.frequency_domain_strain) ** 2 - / interferometer.power_spectral_density_array - ).sum() + log_l -= ( + 2.0 + / self.waveform_generator.duration + * ( + abs(interferometer.frequency_domain_strain) ** 2 + / interferometer.power_spectral_density_array + ).sum() + ) return log_l def log_likelihood(self, parameters): - """ Calculates the real part of log-likelihood value + """Calculates the real part of log-likelihood value Returns ======= @@ -58,17 +65,20 @@ def log_likelihood(self, parameters): """ log_l = 0 - waveform_polarizations = \ - self.waveform_generator.frequency_domain_strain(parameters) + waveform_polarizations = self.waveform_generator.frequency_domain_strain( + parameters + ) if waveform_polarizations is None: return np.nan_to_num(-np.inf) for interferometer in self.interferometers: log_l += self.log_likelihood_interferometer( - waveform_polarizations, interferometer, parameters=parameters) + waveform_polarizations, interferometer, parameters=parameters + ) return log_l.real - def log_likelihood_interferometer(self, waveform_polarizations, - interferometer, parameters): + def log_likelihood_interferometer( + self, waveform_polarizations, interferometer, parameters + ): """ Parameters @@ -84,11 +94,14 @@ def log_likelihood_interferometer(self, waveform_polarizations, """ signal_ifo = interferometer.get_detector_response( - waveform_polarizations, parameters) + waveform_polarizations, parameters + ) residual = interferometer.frequency_domain_strain - signal_ifo - log_l = - 2. / self.waveform_generator.duration * ( - abs(residual)**2 / interferometer.power_spectral_density_array - ).sum() + log_l = ( + -2.0 + / self.waveform_generator.duration + * (abs(residual) ** 2 / interferometer.power_spectral_density_array).sum() + ) return log_l.real diff --git a/bilby/gw/likelihood/multiband.py b/bilby/gw/likelihood/multiband.py index 0fe5d8f47..12a60f98e 100644 --- a/bilby/gw/likelihood/multiband.py +++ b/bilby/gw/likelihood/multiband.py @@ -1,4 +1,3 @@ - import math import numbers @@ -7,10 +6,14 @@ from .base import GravitationalWaveTransient from ...compat.utils import array_module from ...core.utils import ( - logger, speed_of_light, solar_mass, radius_of_earth, - gravitational_constant, round_up_to_power_of_two, + logger, + speed_of_light, + solar_mass, + radius_of_earth, + gravitational_constant, + round_up_to_power_of_two, recursively_load_dict_contents_from_group, - recursively_save_dict_contents_to_group + recursively_save_dict_contents_to_group, ) from ..prior import CBCPriorDict from ..utils import ln_i0 @@ -90,20 +93,40 @@ class MBGravitationalWaveTransient(GravitationalWaveTransient): A likelihood object, able to compute the likelihood of the data given some model parameters """ + def __init__( - self, interferometers, waveform_generator, reference_chirp_mass=None, highest_mode=2, - linear_interpolation=True, accuracy_factor=5, time_offset=None, delta_f_end=None, - maximum_banding_frequency=None, minimum_banding_duration=0., weights=None, - distance_marginalization=False, phase_marginalization=False, priors=None, - time_marginalization=False, jitter_time=True, distance_marginalization_lookup_table=None, - reference_frame="sky", time_reference="geocenter" + self, + interferometers, + waveform_generator, + reference_chirp_mass=None, + highest_mode=2, + linear_interpolation=True, + accuracy_factor=5, + time_offset=None, + delta_f_end=None, + maximum_banding_frequency=None, + minimum_banding_duration=0.0, + weights=None, + distance_marginalization=False, + phase_marginalization=False, + priors=None, + time_marginalization=False, + jitter_time=True, + distance_marginalization_lookup_table=None, + reference_frame="sky", + time_reference="geocenter", ): super(MBGravitationalWaveTransient, self).__init__( - interferometers=interferometers, waveform_generator=waveform_generator, priors=priors, - distance_marginalization=distance_marginalization, phase_marginalization=phase_marginalization, + interferometers=interferometers, + waveform_generator=waveform_generator, + priors=priors, + distance_marginalization=distance_marginalization, + phase_marginalization=phase_marginalization, time_marginalization=time_marginalization, distance_marginalization_lookup_table=distance_marginalization_lookup_table, - jitter_time=jitter_time, reference_frame=reference_frame, time_reference=time_reference + jitter_time=jitter_time, + reference_frame=reference_frame, + time_reference=time_reference, ) if weights is None: self.reference_chirp_mass = reference_chirp_mass @@ -118,9 +141,10 @@ def __init__( else: if isinstance(weights, str): import h5py + logger.info(f"Loading multiband weights from {weights}.") - with h5py.File(weights, 'r') as f: - weights = recursively_load_dict_contents_from_group(f, '/') + with h5py.File(weights, "r") as f: + weights = recursively_load_dict_contents_from_group(f, "/") self.setup_multibanding_from_weights(weights) if self.time_marginalization: self._setup_time_marginalization_multiband() @@ -131,7 +155,12 @@ def reference_chirp_mass(self): @property def reference_chirp_mass_in_second(self): - return gravitational_constant * self._reference_chirp_mass * solar_mass / speed_of_light**3. + return ( + gravitational_constant + * self._reference_chirp_mass + * solar_mass + / speed_of_light**3.0 + ) @reference_chirp_mass.setter def reference_chirp_mass(self, reference_chirp_mass): @@ -174,7 +203,9 @@ def linear_interpolation(self): @linear_interpolation.setter def linear_interpolation(self, linear_interpolation): - if isinstance(linear_interpolation, bool) or isinstance(linear_interpolation, np.bool_): + if isinstance(linear_interpolation, bool) or isinstance( + linear_interpolation, np.bool_ + ): self._linear_interpolation = linear_interpolation else: raise TypeError("linear_interpolation must be a bool") @@ -215,13 +246,18 @@ def time_offset(self, time_offset): raise TypeError("time_offset must be a number") elif self.priors is not None and time_parameter in self.priors: self._time_offset = ( - self.interferometers.start_time + self.interferometers.duration - - self.priors[time_parameter].minimum + safety + self.interferometers.start_time + + self.interferometers.duration + - self.priors[time_parameter].minimum + + safety ) else: self._time_offset = 2.12 - logger.warning("time offset can not be inferred. Use the standard time offset of {} seconds.".format( - self._time_offset)) + logger.warning( + "time offset can not be inferred. Use the standard time offset of {} seconds.".format( + self._time_offset + ) + ) @property def delta_f_end(self): @@ -249,13 +285,18 @@ def delta_f_end(self, delta_f_end): raise TypeError("delta_f_end must be a number") elif self.priors is not None and time_parameter in self.priors: self._delta_f_end = 100 / ( - self.interferometers.start_time + self.interferometers.duration - - self.priors[time_parameter].maximum - safety + self.interferometers.start_time + + self.interferometers.duration + - self.priors[time_parameter].maximum + - safety ) else: - self._delta_f_end = 53. - logger.warning("delta_f_end can not be inferred. Use the standard delta_f_end of {} Hz.".format( - self._delta_f_end)) + self._delta_f_end = 53.0 + logger.warning( + "delta_f_end can not be inferred. Use the standard delta_f_end of {} Hz.".format( + self._delta_f_end + ) + ) @property def maximum_banding_frequency(self): @@ -271,7 +312,8 @@ def maximum_banding_frequency(self, maximum_banding_frequency): time-to-merger \tau(f). The user-specified frequency is used if it is lower than that frequency. """ fmax_tmp = ( - (15 / 968)**(3 / 5) * (self.highest_mode / (2 * np.pi))**(8 / 5) + (15 / 968) ** (3 / 5) + * (self.highest_mode / (2 * np.pi)) ** (8 / 5) / self.reference_chirp_mass_in_second ) if maximum_banding_frequency is not None: @@ -279,8 +321,10 @@ def maximum_banding_frequency(self, maximum_banding_frequency): if maximum_banding_frequency < fmax_tmp: fmax_tmp = maximum_banding_frequency else: - logger.warning("The input maximum_banding_frequency is too large." - "It is set to be {} Hz.".format(fmax_tmp)) + logger.warning( + "The input maximum_banding_frequency is too large." + "It is set to be {} Hz.".format(fmax_tmp) + ) else: raise TypeError("maximum_banding_frequency must be a number") self._maximum_banding_frequency = fmax_tmp @@ -335,7 +379,9 @@ def _tau(self, f): """ f_22 = 2 * f / self.highest_mode return ( - 5 / 256 * self.reference_chirp_mass_in_second + 5 + / 256 + * self.reference_chirp_mass_in_second * (np.pi * self.reference_chirp_mass_in_second * f_22) ** (-8 / 3) ) @@ -355,8 +401,11 @@ def _dtaudf(self, f): """ f_22 = 2 * f / self.highest_mode return ( - -5 / 96 * self.reference_chirp_mass_in_second - * (np.pi * self.reference_chirp_mass_in_second * f_22) ** (-8. / 3.) / f + -5 + / 96 + * self.reference_chirp_mass_in_second + * (np.pi * self.reference_chirp_mass_in_second * f_22) ** (-8.0 / 3.0) + / f ) def _find_starting_frequency(self, duration, fnow): @@ -379,25 +428,29 @@ def _find_starting_frequency(self, duration, fnow): exist. """ + def _is_above_fnext(f): """This function returns True if f > fnext""" cond1 = ( - duration - self.time_offset - self._tau(f) + duration + - self.time_offset + - self._tau(f) - self.accuracy_factor * np.sqrt(-self._dtaudf(f)) ) > 0 - cond2 = f - 1. / np.sqrt(-self._dtaudf(f)) - fnow > 0 + cond2 = f - 1.0 / np.sqrt(-self._dtaudf(f)) - fnow > 0 return cond1 and cond2 + # Bisection search for fnext fmin, fmax = fnow, self.maximum_banding_frequency if not _is_above_fnext(fmax): return None, None while fmax - fmin > 1e-2 / duration: - f = (fmin + fmax) / 2. + f = (fmin + fmax) / 2.0 if _is_above_fnext(f): fmax = f else: fmin = f - return f, 1. / np.sqrt(-self._dtaudf(f)) + return f, 1.0 / np.sqrt(-self._dtaudf(f)) def _setup_frequency_bands(self): r"""Set up frequency bands. The durations of bands geometrically decrease T, T/2. T/4, ..., where T is the @@ -409,21 +462,29 @@ def _setup_frequency_bands(self): """ self.durations = np.array([self.interferometers.duration]) - self.fb_dfb = [[self.minimum_frequency, 0.]] + self.fb_dfb = [[self.minimum_frequency, 0.0]] dnext = self.interferometers.duration / 2 while dnext > max(self.time_offset, self.minimum_banding_duration): fnow, _ = self.fb_dfb[-1] fnext, dfnext = self._find_starting_frequency(dnext, fnow) - if fnext is not None and fnext < min(self.maximum_frequency, self.maximum_banding_frequency): + if fnext is not None and fnext < min( + self.maximum_frequency, self.maximum_banding_frequency + ): self.durations = np.append(self.durations, dnext) self.fb_dfb.append([fnext, dfnext]) dnext /= 2 else: break - self.fb_dfb.append([self.maximum_frequency + self.delta_f_end, self.delta_f_end]) + self.fb_dfb.append( + [self.maximum_frequency + self.delta_f_end, self.delta_f_end] + ) self.fb_dfb = np.array(self.fb_dfb) - logger.info("The total frequency range is divided into {} bands with frequency intervals of {}.".format( - self.number_of_bands, ", ".join(["1/{} Hz".format(d) for d in self.durations]))) + logger.info( + "The total frequency range is divided into {} bands with frequency intervals of {}.".format( + self.number_of_bands, + ", ".join(["1/{} Hz".format(d) for d in self.durations]), + ) + ) def _setup_integers(self): """Set up integers needed for likelihood evaluations. This sets the following instance variables. @@ -440,10 +501,17 @@ def _setup_integers(self): dnow = self.durations[b] fnow, dfnow = self.fb_dfb[b] fnext, _ = self.fb_dfb[b + 1] - Nb = max(round_up_to_power_of_two(2. * (fnext * self.interferometers.duration + 1.)), 2**b) + Nb = max( + round_up_to_power_of_two( + 2.0 * (fnext * self.interferometers.duration + 1.0) + ), + 2**b, + ) self.Nbs = np.append(self.Nbs, Nb) self.Mbs = np.append(self.Mbs, Nb // 2**b) - self.Ks_Ke.append([math.ceil((fnow - dfnow) * dnow), math.floor(fnext * dnow)]) + self.Ks_Ke.append( + [math.ceil((fnow - dfnow) * dnow), math.floor(fnext * dnow)] + ) self.Ks_Ke = np.array(self.Ks_Ke) def _setup_waveform_frequency_points(self): @@ -463,19 +531,30 @@ def _setup_waveform_frequency_points(self): for i in range(self.number_of_bands): d = self.durations[i] Ks, Ke = self.Ks_Ke[i] - self.banded_frequency_points = np.append(self.banded_frequency_points, np.arange(Ks, Ke + 1) / d) + self.banded_frequency_points = np.append( + self.banded_frequency_points, np.arange(Ks, Ke + 1) / d + ) end_idx = start_idx + Ke - Ks self.start_end_idxs.append([start_idx, end_idx]) start_idx = end_idx + 1 self.start_end_idxs = np.array(self.start_end_idxs) - unique_frequencies, idxs = np.unique(self.banded_frequency_points, return_inverse=True) - self.waveform_generator.waveform_arguments['frequencies'] = unique_frequencies + unique_frequencies, idxs = np.unique( + self.banded_frequency_points, return_inverse=True + ) + self.waveform_generator.waveform_arguments["frequencies"] = unique_frequencies self.unique_to_original_frequencies = idxs - logger.info("The number of frequency points where waveforms are evaluated is {}.".format( - len(unique_frequencies))) - logger.info("The speed-up gain of multi-banding is {}.".format( - (self.maximum_frequency - self.minimum_frequency) * self.interferometers.duration / - len(unique_frequencies))) + logger.info( + "The number of frequency points where waveforms are evaluated is {}.".format( + len(unique_frequencies) + ) + ) + logger.info( + "The speed-up gain of multi-banding is {}.".format( + (self.maximum_frequency - self.minimum_frequency) + * self.interferometers.duration + / len(unique_frequencies) + ) + ) def _get_window_sequence(self, delta_f, start_idx, length, b): """Compute window function on frequencies with a fixed frequency interval @@ -509,56 +588,69 @@ def _get_window_sequence(self, delta_f, start_idx, length, b): ) decrease_stop = np.clip(math.ceil(fnext / delta_f) - start_idx, 0, length) - window_sequence[unity_start:decrease_start] = 1. + window_sequence[unity_start:decrease_start] = 1.0 # this if statement avoids overflow caused by vanishing dfnow if increase_start < unity_start: frequencies = (np.arange(increase_start, unity_start) + start_idx) * delta_f window_sequence[increase_start:unity_start] = ( - 1. + np.cos(np.pi * (frequencies - fnow) / dfnow) - ) / 2. + 1.0 + np.cos(np.pi * (frequencies - fnow) / dfnow) + ) / 2.0 if decrease_start < decrease_stop: - frequencies = (np.arange(decrease_start, decrease_stop) + start_idx) * delta_f + frequencies = ( + np.arange(decrease_start, decrease_stop) + start_idx + ) * delta_f window_sequence[decrease_start:decrease_stop] = ( - 1. - np.cos(np.pi * (frequencies - fnext) / dfnext) - ) / 2. + 1.0 - np.cos(np.pi * (frequencies - fnext) / dfnext) + ) / 2.0 return window_sequence def _setup_linear_coefficients(self): """Set up coefficients by which waveforms are multiplied to compute (d, h)""" - self.linear_coeffs = dict((ifo.name, np.array([])) for ifo in self.interferometers) + self.linear_coeffs = dict( + (ifo.name, np.array([])) for ifo in self.interferometers + ) N = self.Nbs[-1] for ifo in self.interferometers: logger.info("Pre-computing linear coefficients for {}".format(ifo.name)) fddata = np.zeros(N // 2 + 1, dtype=complex) - fddata[:len(ifo.frequency_domain_strain)][ifo.frequency_mask[:len(fddata)]] += np.asarray( - ifo.frequency_domain_strain[ifo.frequency_mask] / ifo.power_spectral_density_array[ifo.frequency_mask] + fddata[: len(ifo.frequency_domain_strain)][ + ifo.frequency_mask[: len(fddata)] + ] += np.asarray( + ifo.frequency_domain_strain[ifo.frequency_mask] + / ifo.power_spectral_density_array[ifo.frequency_mask] ) for b in range(self.number_of_bands): Ks, Ke = self.Ks_Ke[b] - windows = self._get_window_sequence(1. / self.durations[b], Ks, Ke - Ks + 1, b) - fddata_in_ith_band = np.copy(fddata[:int(self.Nbs[b] / 2 + 1)]) - fddata_in_ith_band[-1] = 0. # zeroing data at the Nyquist frequency - tddata = np.fft.irfft(fddata_in_ith_band)[-self.Mbs[b]:] + windows = self._get_window_sequence( + 1.0 / self.durations[b], Ks, Ke - Ks + 1, b + ) + fddata_in_ith_band = np.copy(fddata[: int(self.Nbs[b] / 2 + 1)]) + fddata_in_ith_band[-1] = 0.0 # zeroing data at the Nyquist frequency + tddata = np.fft.irfft(fddata_in_ith_band)[-self.Mbs[b] :] Ks, Ke = self.Ks_Ke[b] - fddata_in_ith_band = np.fft.rfft(tddata)[Ks:Ke + 1] + fddata_in_ith_band = np.fft.rfft(tddata)[Ks : Ke + 1] self.linear_coeffs[ifo.name] = np.append( - self.linear_coeffs[ifo.name], (4. / self.durations[b]) * windows * np.conj(fddata_in_ith_band)) + self.linear_coeffs[ifo.name], + (4.0 / self.durations[b]) * windows * np.conj(fddata_in_ith_band), + ) def _setup_quadratic_coefficients_linear_interp(self): """Set up coefficients by which the squares of waveforms are multiplied to compute (h, h) for the linear-interpolation algorithm""" logger.info("Linear-interpolation algorithm is used for (h, h).") - self.quadratic_coeffs = dict((ifo.name, np.array([])) for ifo in self.interferometers) + self.quadratic_coeffs = dict( + (ifo.name, np.array([])) for ifo in self.interferometers + ) original_duration = float(self.interferometers.duration) for b in range(self.number_of_bands): logger.info(f"Pre-computing quadratic coefficients for the {b}-th band") _start, _end = self.start_end_idxs[b] - banded_frequencies = self.banded_frequency_points[_start:_end + 1] + banded_frequencies = self.banded_frequency_points[_start : _end + 1] prefactor = 4 * self.durations[b] / original_duration # precompute window values @@ -569,19 +661,25 @@ def _setup_quadratic_coefficients_linear_interp(self): 1 / original_duration, start_idx_in_band, math.floor(_fnext * original_duration) - start_idx_in_band + 1, - b + b, ) for ifo in self.interferometers: end_idx_in_band = min( start_idx_in_band + len(window_sequence) - 1, - len(ifo.power_spectral_density_array) - 1 + len(ifo.power_spectral_density_array) - 1, + ) + _frequency_mask = np.asarray( + ifo.frequency_mask[start_idx_in_band : end_idx_in_band + 1] ) - _frequency_mask = np.asarray(ifo.frequency_mask[start_idx_in_band:end_idx_in_band + 1]) window_over_psd = np.zeros(end_idx_in_band + 1 - start_idx_in_band) - window_over_psd[_frequency_mask] = \ - 1. / ifo.power_spectral_density_array[start_idx_in_band:end_idx_in_band + 1][_frequency_mask] - window_over_psd *= window_sequence[:len(window_over_psd)] + window_over_psd[_frequency_mask] = ( + 1.0 + / ifo.power_spectral_density_array[ + start_idx_in_band : end_idx_in_band + 1 + ][_frequency_mask] + ) + window_over_psd *= window_sequence[: len(window_over_psd)] coeffs = np.zeros(len(banded_frequencies)) for k in range(len(coeffs) - 1): @@ -590,25 +688,39 @@ def _setup_quadratic_coefficients_linear_interp(self): else: start_idx_in_sum = max( start_idx_in_band, - math.ceil(original_duration * banded_frequencies[k]) + math.ceil(original_duration * banded_frequencies[k]), ) if k == len(coeffs) - 2: end_idx_in_sum = end_idx_in_band else: end_idx_in_sum = min( end_idx_in_band, - math.ceil(original_duration * banded_frequencies[k + 1]) - 1 + math.ceil(original_duration * banded_frequencies[k + 1]) + - 1, ) - frequencies_in_sum = np.arange(start_idx_in_sum, end_idx_in_sum + 1) / original_duration + frequencies_in_sum = ( + np.arange(start_idx_in_sum, end_idx_in_sum + 1) + / original_duration + ) coeffs[k] += prefactor * np.sum( - (banded_frequencies[k + 1] - frequencies_in_sum) * - window_over_psd[start_idx_in_sum - start_idx_in_band:end_idx_in_sum - start_idx_in_band + 1] + (banded_frequencies[k + 1] - frequencies_in_sum) + * window_over_psd[ + start_idx_in_sum - start_idx_in_band : end_idx_in_sum + - start_idx_in_band + + 1 + ] ) coeffs[k + 1] += prefactor * np.sum( - (frequencies_in_sum - banded_frequencies[k]) * - window_over_psd[start_idx_in_sum - start_idx_in_band:end_idx_in_sum - start_idx_in_band + 1] + (frequencies_in_sum - banded_frequencies[k]) + * window_over_psd[ + start_idx_in_sum - start_idx_in_band : end_idx_in_sum + - start_idx_in_band + + 1 + ] ) - self.quadratic_coeffs[ifo.name] = np.append(self.quadratic_coeffs[ifo.name], coeffs) + self.quadratic_coeffs[ifo.name] = np.append( + self.quadratic_coeffs[ifo.name], coeffs + ) def _setup_quadratic_coefficients_ifft_fft(self): """Set up coefficients needed for the IFFT-FFT algorithm to compute (h, h)""" @@ -617,30 +729,37 @@ def _setup_quadratic_coefficients_ifft_fft(self): # variables defined below correspond to \hat{N}^(b), \hat{T}^(b), \tilde{I}^(b)_{c, k}, h^(b)_{c, m} and # \sqrt{w^(b)(f^(b)_k)} \tilde{h}(f^(b)_k) in the paper Nhatbs = [min(2 * Mb, Nb) for Mb, Nb in zip(self.Mbs, self.Nbs)] - self.Tbhats = [self.interferometers.duration * Nbhat / Nb for Nb, Nbhat in zip(self.Nbs, Nhatbs)] + self.Tbhats = [ + self.interferometers.duration * Nbhat / Nb + for Nb, Nbhat in zip(self.Nbs, Nhatbs) + ] self.Ibcs = dict((ifo.name, []) for ifo in self.interferometers) self.hbcs = dict((ifo.name, []) for ifo in self.interferometers) self.wths = dict((ifo.name, []) for ifo in self.interferometers) for ifo in self.interferometers: logger.info("Pre-computing quadratic coefficients for {}".format(ifo.name)) full_inv_psds = np.zeros(N // 2 + 1) - full_inv_psds[:len(ifo.power_spectral_density_array)][ifo.frequency_mask[:len(full_inv_psds)]] = ( - 1 / ifo.power_spectral_density_array[ifo.frequency_mask] - ) + full_inv_psds[: len(ifo.power_spectral_density_array)][ + ifo.frequency_mask[: len(full_inv_psds)] + ] = 1 / ifo.power_spectral_density_array[ifo.frequency_mask] for b in range(self.number_of_bands): - Imb = np.fft.irfft(full_inv_psds[:self.Nbs[b] // 2 + 1]) + Imb = np.fft.irfft(full_inv_psds[: self.Nbs[b] // 2 + 1]) half_length = Nhatbs[b] // 2 - Imbc = np.append(Imb[:half_length + 1], Imb[-(Nhatbs[b] - half_length - 1):]) + Imbc = np.append( + Imb[: half_length + 1], Imb[-(Nhatbs[b] - half_length - 1) :] + ) self.Ibcs[ifo.name].append(np.fft.rfft(Imbc)) # Allocate arrays for IFFT-FFT operations self.hbcs[ifo.name].append(np.zeros(Nhatbs[b])) - self.wths[ifo.name].append(np.zeros(self.Mbs[b] // 2 + 1, dtype=complex)) + self.wths[ifo.name].append( + np.zeros(self.Mbs[b] // 2 + 1, dtype=complex) + ) # precompute windows and their squares self.windows = np.array([]) self.square_root_windows = np.array([]) for b in range(self.number_of_bands): Ks, Ke = self.Ks_Ke[b] - ws = self._get_window_sequence(1. / self.durations[b], Ks, Ke - Ks + 1, b) + ws = self._get_window_sequence(1.0 / self.durations[b], Ks, Ke - Ks + 1, b) self.windows = np.append(self.windows, ws) self.square_root_windows = np.append(self.square_root_windows, np.sqrt(ws)) @@ -648,16 +767,28 @@ def _setup_quadratic_coefficients_ifft_fft(self): def weights(self): _weights = {} for key in [ - "reference_chirp_mass", "highest_mode", "linear_interpolation", - "accuracy_factor", "time_offset", "delta_f_end", - "maximum_banding_frequency", "minimum_banding_duration", - "durations", "fb_dfb", "Nbs", "Mbs", "Ks_Ke", - "banded_frequency_points", "start_end_idxs", - "unique_to_original_frequencies", "linear_coeffs" + "reference_chirp_mass", + "highest_mode", + "linear_interpolation", + "accuracy_factor", + "time_offset", + "delta_f_end", + "maximum_banding_frequency", + "minimum_banding_duration", + "durations", + "fb_dfb", + "Nbs", + "Mbs", + "Ks_Ke", + "banded_frequency_points", + "start_end_idxs", + "unique_to_original_frequencies", + "linear_coeffs", ]: _weights[key] = getattr(self, key) - _weights["waveform_frequencies"] = \ - self.waveform_generator.waveform_arguments['frequencies'] + _weights["waveform_frequencies"] = self.waveform_generator.waveform_arguments[ + "frequencies" + ] if self.linear_interpolation: _weights["quadratic_coeffs"] = self.quadratic_coeffs else: @@ -667,7 +798,9 @@ def weights(self): _weights[key] = {} value = getattr(self, key) for ifo_name, data in value.items(): - _weights[key][ifo_name] = dict((str(b), v) for b, v in enumerate(data)) + _weights[key][ifo_name] = dict( + (str(b), v) for b, v in enumerate(data) + ) return _weights def save_weights(self, filename): @@ -680,11 +813,12 @@ def save_weights(self, filename): """ import h5py + if not filename.endswith(".hdf5"): filename += ".hdf5" logger.info(f"Saving multiband weights to {filename}") - with h5py.File(filename, 'w') as f: - recursively_save_dict_contents_to_group(f, '/', self.weights) + with h5py.File(filename, "w") as f: + recursively_save_dict_contents_to_group(f, "/", self.weights) def setup_multibanding_from_weights(self, weights): """ @@ -707,7 +841,9 @@ def setup_multibanding_from_weights(self, weights): to_set[ifo_name] = [data[str(b)] for b in range(len(data.keys()))] setattr(self, key, to_set) elif key == "waveform_frequencies": - self.waveform_generator.waveform_arguments['frequencies'] = weights["waveform_frequencies"] + self.waveform_generator.waveform_arguments["frequencies"] = weights[ + "waveform_frequencies" + ] else: setattr(self, key, value) @@ -715,17 +851,23 @@ def _setup_time_marginalization_multiband(self): N = self.Nbs[-1] // 2 self._delta_tc = self.durations[0] / N self._times = self.interferometers.start_time + np.arange(N) * self._delta_tc - self.time_prior_array = self.priors['geocent_time'].prob(self._times) * self._delta_tc + self.time_prior_array = ( + self.priors["geocent_time"].prob(self._times) * self._delta_tc + ) # allocate array which is FFTed at each likelihood evaluation self._full_d_h = np.zeros(N, dtype=complex) # idxs to convert full frequency points to banded frequency points, used for filling _full_d_h. - self._full_to_multiband = [int(f * self.durations[0]) for f in self.banded_frequency_points] + self._full_to_multiband = [ + int(f * self.durations[0]) for f in self.banded_frequency_points + ] self._beam_pattern_reference_time = ( - self.priors['geocent_time'].minimum + self.priors['geocent_time'].maximum + self.priors["geocent_time"].minimum + self.priors["geocent_time"].maximum ) / 2 self.interferometers.reference_time = self._beam_pattern_reference_time - def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array=True, parameters): + def calculate_snrs( + self, waveform_polarizations, interferometer, *, return_array=True, parameters + ): """ Compute the snrs @@ -757,35 +899,43 @@ def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array xp = array_module(strain) - d_inner_h = xp.conj(xp.dot(strain, xp.asarray(self.linear_coeffs[interferometer.name]))) + d_inner_h = xp.conj( + xp.dot(strain, xp.asarray(self.linear_coeffs[interferometer.name])) + ) if self.linear_interpolation: optimal_snr_squared = xp.vdot( - xp.abs(strain)**2, - xp.asarray(self.quadratic_coeffs[interferometer.name]) + xp.abs(strain) ** 2, + xp.asarray(self.quadratic_coeffs[interferometer.name]), ) else: - optimal_snr_squared = 0. + optimal_snr_squared = 0.0 for b in range(self.number_of_bands): Ks, Ke = self.Ks_Ke[b] start_idx, end_idx = self.start_end_idxs[b] Mb = self.Mbs[b] if b == 0: - optimal_snr_squared += (4. / self.interferometers.duration) * xp.vdot( - xp.abs(strain[start_idx:end_idx + 1])**2, - interferometer.frequency_mask[Ks:Ke + 1] * xp.asarray(self.windows[start_idx:end_idx + 1]) - / interferometer.power_spectral_density_array[Ks:Ke + 1]) + optimal_snr_squared += ( + 4.0 / self.interferometers.duration + ) * xp.vdot( + xp.abs(strain[start_idx : end_idx + 1]) ** 2, + interferometer.frequency_mask[Ks : Ke + 1] + * xp.asarray(self.windows[start_idx : end_idx + 1]) + / interferometer.power_spectral_density_array[Ks : Ke + 1], + ) else: - self.wths[interferometer.name][b][Ks:Ke + 1] = ( - xp.asarray(self.square_root_windows[start_idx:end_idx + 1]) - * strain[start_idx:end_idx + 1] + self.wths[interferometer.name][b][Ks : Ke + 1] = ( + xp.asarray(self.square_root_windows[start_idx : end_idx + 1]) + * strain[start_idx : end_idx + 1] ) self.hbcs[interferometer.name][b][-Mb:] = xp.fft.irfft( xp.asarray(self.wths[interferometer.name][b]) ) thbc = xp.fft.rfft(xp.asarray(self.hbcs[interferometer.name][b])) - optimal_snr_squared += (4. / self.Tbhats[b]) * xp.vdot( - xp.abs(thbc)**2, xp.asarray(self.Ibcs[interferometer.name][b].real)) + optimal_snr_squared += (4.0 / self.Tbhats[b]) * xp.vdot( + xp.abs(thbc) ** 2, + xp.asarray(self.Ibcs[interferometer.name][b].real), + ) complex_matched_filter_snr = d_inner_h / (optimal_snr_squared**0.5) @@ -793,8 +943,10 @@ def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array self._full_d_h[self._full_to_multiband] *= 0 for b in range(self.number_of_bands): start_idx, end_idx = self.start_end_idxs[b] - self._full_d_h[self._full_to_multiband[start_idx:end_idx + 1]] += \ - strain[start_idx:end_idx + 1] * self.linear_coeffs[interferometer.name][start_idx:end_idx + 1] + self._full_d_h[self._full_to_multiband[start_idx : end_idx + 1]] += ( + strain[start_idx : end_idx + 1] + * self.linear_coeffs[interferometer.name][start_idx : end_idx + 1] + ) d_inner_h_array = xp.fft.fft(self._full_d_h) else: d_inner_h_array = None @@ -810,11 +962,14 @@ def _rescale_signal(self, signal, new_distance): for mode in signal: signal[mode] *= self._ref_dist / new_distance - def generate_time_sample_from_marginalized_likelihood(self, signal_polarizations=None, *, parameters): + def generate_time_sample_from_marginalized_likelihood( + self, signal_polarizations=None, *, parameters + ): parameters.update(self.get_sky_frame_parameters(parameters=parameters)) if signal_polarizations is None: - signal_polarizations = \ - self.waveform_generator.frequency_domain_strain(parameters) + signal_polarizations = self.waveform_generator.frequency_domain_strain( + parameters + ) snrs = self._CalculatedSNRs() @@ -829,16 +984,17 @@ def generate_time_sample_from_marginalized_likelihood(self, signal_polarizations if self.distance_marginalization: time_log_like = self.distance_marginalized_likelihood( - d_inner_h, h_inner_h, parameters=parameters) + d_inner_h, h_inner_h, parameters=parameters + ) elif self.phase_marginalization: time_log_like = ln_i0(abs(d_inner_h)) - h_inner_h.real / 2 else: - time_log_like = (d_inner_h.real - h_inner_h.real / 2) + time_log_like = d_inner_h.real - h_inner_h.real / 2 times = self._times if self.jitter_time: times = times + parameters["time_jitter"] - time_prior_array = self.priors['geocent_time'].prob(times) + time_prior_array = self.priors["geocent_time"].prob(times) time_post = np.exp(time_log_like - max(time_log_like)) * time_prior_array time_post /= np.sum(time_post) return np.random.choice(times, p=time_post) diff --git a/bilby/gw/likelihood/relative.py b/bilby/gw/likelihood/relative.py index 72c23c958..97e3fa3c9 100644 --- a/bilby/gw/likelihood/relative.py +++ b/bilby/gw/likelihood/relative.py @@ -102,23 +102,25 @@ class RelativeBinningGravitationalWaveTransient(GravitationalWaveTransient): The relative binning likelihood does not currently support calibration marginalization. """ - def __init__(self, interferometers, - waveform_generator, - fiducial_parameters=None, - parameter_bounds=None, - maximization_kwargs=None, - update_fiducial_parameters=False, - distance_marginalization=False, - time_marginalization=False, - phase_marginalization=False, - priors=None, - distance_marginalization_lookup_table=None, - jitter_time=True, - reference_frame="sky", - time_reference="geocenter", - chi=1, - epsilon=0.5): - + def __init__( + self, + interferometers, + waveform_generator, + fiducial_parameters=None, + parameter_bounds=None, + maximization_kwargs=None, + update_fiducial_parameters=False, + distance_marginalization=False, + time_marginalization=False, + phase_marginalization=False, + priors=None, + distance_marginalization_lookup_table=None, + jitter_time=True, + reference_frame="sky", + time_reference="geocenter", + chi=1, + epsilon=0.5, + ): super(RelativeBinningGravitationalWaveTransient, self).__init__( interferometers=interferometers, waveform_generator=waveform_generator, @@ -129,7 +131,8 @@ def __init__(self, interferometers, distance_marginalization_lookup_table=distance_marginalization_lookup_table, jitter_time=jitter_time, reference_frame=reference_frame, - time_reference=time_reference) + time_reference=time_reference, + ) if fiducial_parameters is None: logger.info("Drawing fiducial parameters from prior.") @@ -159,22 +162,38 @@ def __init__(self, interferometers, if update_fiducial_parameters: # write a check to make sure prior is not None - logger.info("Using scipy optimization to find maximum likelihood parameters.") - self.parameters_to_be_updated = [key for key in priors if not isinstance( - priors[key], (DeltaFunction, Constraint, float, int))] - logger.info(f"Parameters over which likelihood is maximized: {self.parameters_to_be_updated}") + logger.info( + "Using scipy optimization to find maximum likelihood parameters." + ) + self.parameters_to_be_updated = [ + key + for key in priors + if not isinstance(priors[key], (DeltaFunction, Constraint, float, int)) + ] + logger.info( + f"Parameters over which likelihood is maximized: {self.parameters_to_be_updated}" + ) if parameter_bounds is None: logger.info("No parameter bounds were given. Using priors instead.") self.parameter_bounds = self.get_bounds_from_priors(priors) else: - self.parameter_bounds = self.get_parameter_list_from_dictionary(parameter_bounds) + self.parameter_bounds = self.get_parameter_list_from_dictionary( + parameter_bounds + ) self.fiducial_parameters = self.find_maximum_likelihood_parameters( - self.parameter_bounds, maximization_kwargs=maximization_kwargs) - logger.info(f"Fiducial likelihood: {self.log_likelihood_ratio(self.fiducial_parameters):.2f}") + self.parameter_bounds, maximization_kwargs=maximization_kwargs + ) + logger.info( + f"Fiducial likelihood: {self.log_likelihood_ratio(self.fiducial_parameters):.2f}" + ) def __repr__(self): - return self.__class__.__name__ + '(interferometers={},\n\twaveform_generator={},\n\fiducial_parameters={},' \ - .format(self.interferometers, self.waveform_generator, self.fiducial_parameters) + return ( + self.__class__.__name__ + + "(interferometers={},\n\twaveform_generator={},\n\fiducial_parameters={},".format( + self.interferometers, self.waveform_generator, self.fiducial_parameters + ) + ) def setup_bins(self): """ @@ -198,14 +217,16 @@ def setup_bins(self): & (frequency_array <= maximum_frequency) ] - d_alpha = self.chi * 2 * np.pi / np.abs( - (minimum_frequency ** gamma) * np.heaviside(-gamma, 1) - - (maximum_frequency ** gamma) * np.heaviside(gamma, 1) - ) - d_phi = np.sum( - np.sign(gamma) * d_alpha * frequency_array_useful ** gamma, - axis=0 + d_alpha = ( + self.chi + * 2 + * np.pi + / np.abs( + (minimum_frequency**gamma) * np.heaviside(-gamma, 1) + - (maximum_frequency**gamma) * np.heaviside(gamma, 1) + ) ) + d_phi = np.sum(np.sign(gamma) * d_alpha * frequency_array_useful**gamma, axis=0) d_phi_from_start = d_phi - d_phi[0] number_of_bins = int(d_phi_from_start[-1] // self.epsilon) bin_inds = list() @@ -213,7 +234,9 @@ def setup_bins(self): last_index = -1 for i in range(number_of_bins + 1): - bin_index = np.where(d_phi_from_start >= ((i / number_of_bins) * d_phi_from_start[-1]))[0][0] + bin_index = np.where( + d_phi_from_start >= ((i / number_of_bins) * d_phi_from_start[-1]) + )[0][0] if bin_index == last_index: continue bin_freq = frequency_array_useful[bin_index] @@ -230,7 +253,9 @@ def setup_bins(self): f"Set up {self.number_of_bins} bins " f"between {minimum_frequency} Hz and {maximum_frequency} Hz" ) - self.waveform_generator.waveform_arguments["frequency_bin_edges"] = self.bin_freqs + self.waveform_generator.waveform_arguments["frequency_bin_edges"] = ( + self.bin_freqs + ) self.bin_widths = self.bin_freqs[1:] - self.bin_freqs[:-1] self.bin_centers = (self.bin_freqs[1:] + self.bin_freqs[:-1]) / 2 for interferometer in self.interferometers: @@ -244,12 +269,17 @@ def set_fiducial_waveforms(self, parameters): self._set_fiducial() parameters.update(self.get_sky_frame_parameters(parameters=parameters)) self.fiducial_polarizations = self.waveform_generator.frequency_domain_strain( - parameters) + parameters + ) self._unset_fiducial() - maximum_nonzero_index = np.where(self.fiducial_polarizations["plus"] != 0j)[0][-1] + maximum_nonzero_index = np.where(self.fiducial_polarizations["plus"] != 0j)[0][ + -1 + ] logger.debug(f"Maximum Nonzero Index is {maximum_nonzero_index}") - maximum_nonzero_frequency = self.waveform_generator.frequency_array[maximum_nonzero_index] + maximum_nonzero_frequency = self.waveform_generator.frequency_array[ + maximum_nonzero_index + ] logger.debug(f"Maximum Nonzero Frequency is {maximum_nonzero_frequency}") self.maximum_frequency = maximum_nonzero_frequency @@ -258,17 +288,22 @@ def set_fiducial_waveforms(self, parameters): for interferometer in self.interferometers: logger.debug(f"Maximum Frequency is {interferometer.maximum_frequency}") - wf = interferometer.get_detector_response(self.fiducial_polarizations, parameters) + wf = interferometer.get_detector_response( + self.fiducial_polarizations, parameters + ) wf *= interferometer.frequency_array <= self.maximum_frequency self.per_detector_fiducial_waveforms[interferometer.name] = wf - def find_maximum_likelihood_parameters(self, parameter_bounds, - iterations=5, maximization_kwargs=None): + def find_maximum_likelihood_parameters( + self, parameter_bounds, iterations=5, maximization_kwargs=None + ): if maximization_kwargs is None: maximization_kwargs = dict() parameters = deepcopy(self.fiducial_parameters) self._unset_fiducial() - updated_parameters_list = self.get_parameter_list_from_dictionary(self.fiducial_parameters) + updated_parameters_list = self.get_parameter_list_from_dictionary( + self.fiducial_parameters + ) old_fiducial_ln_likelihood = self.log_likelihood_ratio(self.fiducial_parameters) logger.info(f"Fiducial ln likelihood ratio: {old_fiducial_ln_likelihood:.2f}") for it in range(iterations): @@ -280,14 +315,18 @@ def find_maximum_likelihood_parameters(self, parameter_bounds, x0=updated_parameters_list, **maximization_kwargs, ) - updated_parameters_list = output['x'] + updated_parameters_list = output["x"] updated_parameters = deepcopy(self.fiducial_parameters) - updated_parameters.update(self.get_parameter_dictionary_from_list(updated_parameters_list)) + updated_parameters.update( + self.get_parameter_dictionary_from_list(updated_parameters_list) + ) self.set_fiducial_waveforms(updated_parameters) self.setup_bins() self.compute_summary_data() new_fiducial_ln_likelihood = self.log_likelihood_ratio(updated_parameters) - logger.info(f"Fiducial ln likelihood ratio: {new_fiducial_ln_likelihood:.2f}") + logger.info( + f"Fiducial ln likelihood ratio: {new_fiducial_ln_likelihood:.2f}" + ) if new_fiducial_ln_likelihood - old_fiducial_ln_likelihood < 0.1: break old_fiducial_ln_likelihood = new_fiducial_ln_likelihood @@ -302,7 +341,9 @@ def lnlike_scipy_maximize(self, parameter_list, parameters): def get_parameter_dictionary_from_list(self, parameter_list): parameter_dictionary = dict(zip(self.parameters_to_be_updated, parameter_list)) - excluded_parameter_keys = set(self.fiducial_parameters) - set(self.parameters_to_be_updated) + excluded_parameter_keys = set(self.fiducial_parameters) - set( + self.parameters_to_be_updated + ) for key in excluded_parameter_keys: parameter_dictionary[key] = self.fiducial_parameters[key] return parameter_dictionary @@ -333,7 +374,9 @@ def compute_summary_data(self): masked_bin_inds[-1] += 1 masked_strain = interferometer.frequency_domain_strain[mask] - masked_h0 = np.asarray(self.per_detector_fiducial_waveforms[interferometer.name][mask]) + masked_h0 = np.asarray( + self.per_detector_fiducial_waveforms[interferometer.name][mask] + ) masked_psd = interferometer.power_spectral_density_array[mask] duration = interferometer.duration a0, b0, a1, b1 = np.zeros((4, self.number_of_bins), dtype=complex) @@ -355,14 +398,20 @@ def compute_summary_data(self): a0[i] = noise_weighted_inner_product(h0, strain, psd, duration) b0[i] = noise_weighted_inner_product(h0, h0, psd, duration) - a1[i] = noise_weighted_inner_product(h0, strain * delta_frequency, psd, duration) - b1[i] = noise_weighted_inner_product(h0, h0 * delta_frequency, psd, duration) + a1[i] = noise_weighted_inner_product( + h0, strain * delta_frequency, psd, duration + ) + b1[i] = noise_weighted_inner_product( + h0, h0 * delta_frequency, psd, duration + ) summary_data[interferometer.name] = (a0, a1, b0, b1) self.summary_data = summary_data - def compute_waveform_ratio_per_interferometer(self, waveform_polarizations, interferometer, parameters): + def compute_waveform_ratio_per_interferometer( + self, waveform_polarizations, interferometer, parameters + ): name = interferometer.name strain = interferometer.get_detector_response( waveform_polarizations=waveform_polarizations, @@ -382,7 +431,7 @@ def _compute_full_waveform(self, signal_polarizations, interferometer, parameter r0, r1 = self.compute_waveform_ratio_per_interferometer( waveform_polarizations=signal_polarizations, interferometer=interferometer, - parameters=parameters + parameters=parameters, ) idxs = slice(self.bin_inds[0], self.bin_inds[-1] + 1) @@ -392,10 +441,14 @@ def _compute_full_waveform(self, signal_polarizations, interferometer, parameter f = interferometer.frequency_array full_waveform_ratio = np.zeros(f.shape[0], dtype=complex) - full_waveform_ratio[idxs] = duplicated_r0 + duplicated_r1 * (f[idxs] - duplicated_fm) + full_waveform_ratio[idxs] = duplicated_r0 + duplicated_r1 * ( + f[idxs] - duplicated_fm + ) return fiducial_waveform * full_waveform_ratio - def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array=True, parameters): + def calculate_snrs( + self, waveform_polarizations, interferometer, *, return_array=True, parameters + ): r0, r1 = self.compute_waveform_ratio_per_interferometer( waveform_polarizations=waveform_polarizations, interferometer=interferometer, @@ -405,7 +458,7 @@ def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array d_inner_h = (a0 * r0.conj() + a1 * r1.conj()).sum() h_inner_h = (b0 * abs(r0) ** 2 + 2 * b1 * (r0 * r1.conj()).real).sum() optimal_snr_squared = h_inner_h - complex_matched_filter_snr = d_inner_h / (optimal_snr_squared ** 0.5) + complex_matched_filter_snr = d_inner_h / (optimal_snr_squared**0.5) if return_array and self.time_marginalization: xp = array_module(r0) @@ -414,10 +467,15 @@ def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array interferometer=interferometer, parameters=parameters, ) - d_inner_h_array = 4 / self.waveform_generator.duration * xp.fft.fft( - full_waveform[0:-1] - * interferometer.frequency_domain_strain.conj()[0:-1] - / interferometer.power_spectral_density_array[0:-1]) + d_inner_h_array = ( + 4 + / self.waveform_generator.duration + * xp.fft.fft( + full_waveform[0:-1] + * interferometer.frequency_domain_strain.conj()[0:-1] + / interferometer.power_spectral_density_array[0:-1] + ) + ) else: d_inner_h_array = None @@ -426,7 +484,7 @@ def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array d_inner_h=d_inner_h, optimal_snr_squared=optimal_snr_squared.real, complex_matched_filter_snr=complex_matched_filter_snr, - d_inner_h_array=d_inner_h_array + d_inner_h_array=d_inner_h_array, ) def _set_fiducial(self): diff --git a/bilby/gw/likelihood/roq.py b/bilby/gw/likelihood/roq.py index 60c50d65f..25d23fdf4 100644 --- a/bilby/gw/likelihood/roq.py +++ b/bilby/gw/likelihood/roq.py @@ -5,7 +5,10 @@ from .base import GravitationalWaveTransient from ...compat.utils import array_module from ...core.utils import ( - logger, create_frequency_series, speed_of_light, radius_of_earth + logger, + create_frequency_series, + speed_of_light, + radius_of_earth, ) from ..prior import CBCPriorDict from ..utils import ln_i0 @@ -122,6 +125,7 @@ class ROQGravitationalWaveTransient(GravitationalWaveTransient): - e.g., "H1": sample in the time of arrival at H1 """ + @property def roq_params(self): return self._roq_params @@ -148,28 +152,39 @@ def roq_params(self, roq_params): self._roq_params = roq_params def __init__( - self, interferometers, waveform_generator, priors, - weights=None, linear_matrix=None, quadratic_matrix=None, - roq_params=None, roq_params_check=True, roq_scale_factor=1, - distance_marginalization=False, phase_marginalization=False, - time_marginalization=False, jitter_time=True, delta_tc=None, - distance_marginalization_lookup_table=None, - reference_frame="sky", time_reference="geocenter", - parameter_conversion=None - + self, + interferometers, + waveform_generator, + priors, + weights=None, + linear_matrix=None, + quadratic_matrix=None, + roq_params=None, + roq_params_check=True, + roq_scale_factor=1, + distance_marginalization=False, + phase_marginalization=False, + time_marginalization=False, + jitter_time=True, + delta_tc=None, + distance_marginalization_lookup_table=None, + reference_frame="sky", + time_reference="geocenter", + parameter_conversion=None, ): self._delta_tc = delta_tc self._roq_params = None super(ROQGravitationalWaveTransient, self).__init__( interferometers=interferometers, - waveform_generator=waveform_generator, priors=priors, + waveform_generator=waveform_generator, + priors=priors, distance_marginalization=distance_marginalization, phase_marginalization=phase_marginalization, time_marginalization=time_marginalization, distance_marginalization_lookup_table=distance_marginalization_lookup_table, jitter_time=jitter_time, reference_frame=reference_frame, - time_reference=time_reference + time_reference=time_reference, ) self.roq_params_check = roq_params_check @@ -186,55 +201,78 @@ def __init__( elif isinstance(weights, str): self.weights = self.load_weights(weights) else: - is_hdf5_linear = isinstance(linear_matrix, str) and linear_matrix.endswith('.hdf5') - linear_matrix = self._parse_basis(linear_matrix, 'linear') - is_hdf5_quadratic = isinstance(quadratic_matrix, str) and quadratic_matrix.endswith('.hdf5') - quadratic_matrix = self._parse_basis(quadratic_matrix, 'quadratic') + is_hdf5_linear = isinstance(linear_matrix, str) and linear_matrix.endswith( + ".hdf5" + ) + linear_matrix = self._parse_basis(linear_matrix, "linear") + is_hdf5_quadratic = isinstance( + quadratic_matrix, str + ) and quadratic_matrix.endswith(".hdf5") + quadratic_matrix = self._parse_basis(quadratic_matrix, "quadratic") # retrieve roq params from a basis file if it is .hdf5 if self.roq_params is None: if is_hdf5_linear: self.roq_params = np.array( - (linear_matrix['minimum_frequency_hz'][()], - linear_matrix['maximum_frequency_hz'][()], - linear_matrix['duration_s'][()]), - dtype=[('flow', float), ('fhigh', float), ('seglen', float)] + ( + linear_matrix["minimum_frequency_hz"][()], + linear_matrix["maximum_frequency_hz"][()], + linear_matrix["duration_s"][()], + ), + dtype=[("flow", float), ("fhigh", float), ("seglen", float)], ) if is_hdf5_quadratic: if self.roq_params is None: self.roq_params = np.array( - (quadratic_matrix['minimum_frequency_hz'][()], - quadratic_matrix['maximum_frequency_hz'][()], - quadratic_matrix['duration_s'][()]), - dtype=[('flow', float), ('fhigh', float), ('seglen', float)] + ( + quadratic_matrix["minimum_frequency_hz"][()], + quadratic_matrix["maximum_frequency_hz"][()], + quadratic_matrix["duration_s"][()], + ), + dtype=[ + ("flow", float), + ("fhigh", float), + ("seglen", float), + ], ) else: - self.roq_params['flow'] = max( - self.roq_params['flow'], quadratic_matrix['minimum_frequency_hz'][()] + self.roq_params["flow"] = max( + self.roq_params["flow"], + quadratic_matrix["minimum_frequency_hz"][()], ) - self.roq_params['fhigh'] = min( - self.roq_params['fhigh'], quadratic_matrix['maximum_frequency_hz'][()] + self.roq_params["fhigh"] = min( + self.roq_params["fhigh"], + quadratic_matrix["maximum_frequency_hz"][()], ) - self.roq_params['seglen'] = min( - self.roq_params['seglen'], quadratic_matrix['duration_s'][()] + self.roq_params["seglen"] = min( + self.roq_params["seglen"], + quadratic_matrix["duration_s"][()], ) if self.roq_params is not None: for ifo in self.interferometers: self.perform_roq_params_check(ifo) self.weights = dict() - self._set_weights(linear_matrix=linear_matrix, quadratic_matrix=quadratic_matrix) + self._set_weights( + linear_matrix=linear_matrix, quadratic_matrix=quadratic_matrix + ) if is_hdf5_linear: linear_matrix.close() if is_hdf5_quadratic: quadratic_matrix.close() - self.number_of_bases_linear = len(self.weights[f'{self.interferometers[0].name}_linear']) - self.number_of_bases_quadratic = len(self.weights[f'{self.interferometers[0].name}_quadratic']) - self._cache = dict(parameters=None, basis_number_linear=None, basis_number_quadratic=None) + self.number_of_bases_linear = len( + self.weights[f"{self.interferometers[0].name}_linear"] + ) + self.number_of_bases_quadratic = len( + self.weights[f"{self.interferometers[0].name}_quadratic"] + ) + self._cache = dict( + parameters=None, basis_number_linear=None, basis_number_quadratic=None + ) self.parameter_conversion = parameter_conversion - for basis_type in ['linear', 'quadratic']: - number_of_bases = getattr(self, f'number_of_bases_{basis_type}') + for basis_type in ["linear", "quadratic"]: + number_of_bases = getattr(self, f"number_of_bases_{basis_type}") if number_of_bases > 1: self._verify_numbers_of_prior_ranges_and_frequency_nodes(basis_type) else: @@ -244,11 +282,18 @@ def __init__( self._set_unique_frequency_nodes_and_inverse() # need to fill waveform_arguments here if single basis is used, as they will never be updated. if self.number_of_bases_linear == 1 and self.number_of_bases_quadratic == 1: - frequency_nodes, linear_indices, quadratic_indices = \ + frequency_nodes, linear_indices, quadratic_indices = ( self._unique_frequency_nodes_and_inverse[0][0] - self._waveform_generator.waveform_arguments['frequency_nodes'] = frequency_nodes - self._waveform_generator.waveform_arguments['linear_indices'] = linear_indices - self._waveform_generator.waveform_arguments['quadratic_indices'] = quadratic_indices + ) + self._waveform_generator.waveform_arguments["frequency_nodes"] = ( + frequency_nodes + ) + self._waveform_generator.waveform_arguments["linear_indices"] = ( + linear_indices + ) + self._waveform_generator.waveform_arguments["quadratic_indices"] = ( + quadratic_indices + ) def _verify_numbers_of_prior_ranges_and_frequency_nodes(self, basis_type): """ @@ -260,29 +305,33 @@ def _verify_numbers_of_prior_ranges_and_frequency_nodes(self, basis_type): basis_type: str """ - number_of_bases = getattr(self, f'number_of_bases_{basis_type}') - key = f'prior_range_{basis_type}' + number_of_bases = getattr(self, f"number_of_bases_{basis_type}") + key = f"prior_range_{basis_type}" try: prior_ranges = self.weights[key] except KeyError: raise AttributeError( - f'For the use of multiple {basis_type} ROQ bases, weights should contain "{key}".') + f'For the use of multiple {basis_type} ROQ bases, weights should contain "{key}".' + ) else: for param_name in prior_ranges: if len(prior_ranges[param_name]) != number_of_bases: raise ValueError( f'The number of prior ranges for "{param_name}" does not ' - f'match the number of {basis_type} bases') - key = f'frequency_nodes_{basis_type}' + f"match the number of {basis_type} bases" + ) + key = f"frequency_nodes_{basis_type}" try: frequency_nodes = self.weights[key] except KeyError: raise AttributeError( - f'For the use of multiple {basis_type} ROQ bases, weights should contain "{key}".') + f'For the use of multiple {basis_type} ROQ bases, weights should contain "{key}".' + ) else: if len(frequency_nodes) != number_of_bases: raise ValueError( - f'The number of arrays of frequency nodes does not match the number of {basis_type} bases') + f"The number of arrays of frequency nodes does not match the number of {basis_type} bases" + ) def _verify_prior_ranges(self, basis_type): """Check if the union of prior ranges is within the ROQ basis bounds. @@ -292,7 +341,7 @@ def _verify_prior_ranges(self, basis_type): basis_type: str """ - key = f'prior_range_{basis_type}' + key = f"prior_range_{basis_type}" if key not in self.weights: return prior_ranges = self.weights[key] @@ -324,9 +373,13 @@ def _check_frequency_nodes_exist_for_single_basis(self, basis_type): basis_type: str """ - key = f'frequency_nodes_{basis_type}' - if not (key in self.weights or key in self._waveform_generator.waveform_arguments): - raise AttributeError(f'{key} should be contained in weights or waveform arguments.') + key = f"frequency_nodes_{basis_type}" + if not ( + key in self.weights or key in self._waveform_generator.waveform_arguments + ): + raise AttributeError( + f"{key} should be contained in weights or waveform arguments." + ) elif key not in self._waveform_generator.waveform_arguments: self._waveform_generator.waveform_arguments[key] = self.weights[key][0] elif key not in self.weights: @@ -340,33 +393,40 @@ def _set_unique_frequency_nodes_and_inverse(self): self._unique_frequency_nodes_and_inverse = [] for idx_linear in range(self.number_of_bases_linear): tmp = [] - frequency_nodes_linear = xp.asarray(self.weights['frequency_nodes_linear'][idx_linear]) + frequency_nodes_linear = xp.asarray( + self.weights["frequency_nodes_linear"][idx_linear] + ) size_linear = len(frequency_nodes_linear) for idx_quadratic in range(self.number_of_bases_quadratic): - frequency_nodes_quadratic = xp.asarray(self.weights['frequency_nodes_quadratic'][idx_quadratic]) + frequency_nodes_quadratic = xp.asarray( + self.weights["frequency_nodes_quadratic"][idx_quadratic] + ) frequency_nodes_unique, original_indices = xp.unique( xp.hstack((frequency_nodes_linear, frequency_nodes_quadratic)), - return_inverse=True + return_inverse=True, ) linear_indices = original_indices[:size_linear] quadratic_indices = original_indices[size_linear:] - tmp.append( - (frequency_nodes_unique, linear_indices, quadratic_indices) - ) + tmp.append((frequency_nodes_unique, linear_indices, quadratic_indices)) self._unique_frequency_nodes_and_inverse.append(tmp) def _setup_time_marginalization(self): if self._delta_tc is None: self._delta_tc = self._get_time_resolution() - tcmin = self.priors['geocent_time'].minimum - tcmax = self.priors['geocent_time'].maximum + tcmin = self.priors["geocent_time"].minimum + tcmax = self.priors["geocent_time"].maximum number_of_time_samples = int(np.ceil((tcmax - tcmin) / self._delta_tc)) # adjust delta tc so that the last time sample has an equal weight self._delta_tc = (tcmax - tcmin) / number_of_time_samples logger.info( - "delta tc for time marginalization = {} seconds.".format(self._delta_tc)) - self._times = tcmin + self._delta_tc / 2. + np.arange(number_of_time_samples) * self._delta_tc - self._beam_pattern_reference_time = (tcmin + tcmax) / 2. + "delta tc for time marginalization = {} seconds.".format(self._delta_tc) + ) + self._times = ( + tcmin + + self._delta_tc / 2.0 + + np.arange(number_of_time_samples) * self._delta_tc + ) + self._beam_pattern_reference_time = (tcmin + tcmax) / 2.0 @staticmethod def _parse_basis(basis, basis_type): @@ -385,22 +445,23 @@ def _parse_basis(basis, basis_type): basis : hdf5-like object """ - if basis_type not in ['linear', 'quadratic']: - raise ValueError(f'basis_type {basis_type} not recognized') + if basis_type not in ["linear", "quadratic"]: + raise ValueError(f"basis_type {basis_type} not recognized") if isinstance(basis, str): - logger.info(f'Loading {basis_type}_matrix from {basis}') - format = basis.split('.')[-1] - if format == 'npy': - basis = {f'basis_{basis_type}': {'0': {'basis': np.load(basis)}}} - elif format == 'hdf5': + logger.info(f"Loading {basis_type}_matrix from {basis}") + format = basis.split(".")[-1] + if format == "npy": + basis = {f"basis_{basis_type}": {"0": {"basis": np.load(basis)}}} + elif format == "hdf5": import h5py - basis = h5py.File(basis, 'r') + + basis = h5py.File(basis, "r") else: - raise IOError(f'Format {format} not recognized.') + raise IOError(f"Format {format} not recognized.") elif isinstance(basis, np.ndarray): - basis = {f'basis_{basis_type}': {'0': {'basis': basis.T}}} + basis = {f"basis_{basis_type}": {"0": {"basis": basis.T}}} else: - raise TypeError('basis needs to be str or np.ndarray') + raise TypeError("basis needs to be str or np.ndarray") return basis def _select_prior_ranges(self, prior_ranges): @@ -429,13 +490,14 @@ def _select_prior_ranges(self, prior_ranges): except KeyError: continue prior_ranges_of_this_param = prior_ranges[param_name] - in_prior_range *= \ - (prior_ranges_of_this_param[:, 1] >= prior.minimum) * \ - (prior_ranges_of_this_param[:, 0] <= prior.maximum) + in_prior_range *= (prior_ranges_of_this_param[:, 1] >= prior.minimum) * ( + prior_ranges_of_this_param[:, 0] <= prior.maximum + ) idxs_in_prior_range = np.arange(number_of_prior_ranges)[in_prior_range] - return idxs_in_prior_range, \ - dict((param_name, prior_ranges[param_name][idxs_in_prior_range]) - for param_name in param_names) + return idxs_in_prior_range, dict( + (param_name, prior_ranges[param_name][idxs_in_prior_range]) + for param_name in param_names + ) def _update_basis(self, parameters): """ @@ -452,41 +514,53 @@ def _update_basis(self, parameters): if self._cache["parameters"] == parameters: return for basis_type, number_of_bases in zip( - ['linear', 'quadratic'], [self.number_of_bases_linear, self.number_of_bases_quadratic] + ["linear", "quadratic"], + [self.number_of_bases_linear, self.number_of_bases_quadratic], ): - basis_number_key = f'basis_number_{basis_type}' + basis_number_key = f"basis_number_{basis_type}" if number_of_bases == 1: self._cache[basis_number_key] = 0 continue in_prior_range = np.ones(number_of_bases, dtype=bool) - prior_range_key = f'prior_range_{basis_type}' + prior_range_key = f"prior_range_{basis_type}" for param_name in self.weights[prior_range_key]: if param_name not in parameters: continue - in_prior_range *= \ - (self.weights[prior_range_key][param_name][:, 0] <= parameters[param_name]) * \ - (self.weights[prior_range_key][param_name][:, 1] >= parameters[param_name]) - self._cache[basis_number_key] = np.arange(number_of_bases)[in_prior_range][0] - basis_number_linear = self._cache['basis_number_linear'] - basis_number_quadratic = self._cache['basis_number_quadratic'] - frequency_nodes, linear_indices, quadratic_indices = \ - self._unique_frequency_nodes_and_inverse[basis_number_linear][basis_number_quadratic] - self._waveform_generator.waveform_arguments['frequency_nodes'] = frequency_nodes - self._waveform_generator.waveform_arguments['linear_indices'] = linear_indices - self._waveform_generator.waveform_arguments['quadratic_indices'] = quadratic_indices - self._cache['parameters'] = parameters.copy() + in_prior_range *= ( + self.weights[prior_range_key][param_name][:, 0] + <= parameters[param_name] + ) * ( + self.weights[prior_range_key][param_name][:, 1] + >= parameters[param_name] + ) + self._cache[basis_number_key] = np.arange(number_of_bases)[in_prior_range][ + 0 + ] + basis_number_linear = self._cache["basis_number_linear"] + basis_number_quadratic = self._cache["basis_number_quadratic"] + frequency_nodes, linear_indices, quadratic_indices = ( + self._unique_frequency_nodes_and_inverse[basis_number_linear][ + basis_number_quadratic + ] + ) + self._waveform_generator.waveform_arguments["frequency_nodes"] = frequency_nodes + self._waveform_generator.waveform_arguments["linear_indices"] = linear_indices + self._waveform_generator.waveform_arguments["quadratic_indices"] = ( + quadratic_indices + ) + self._cache["parameters"] = parameters.copy() @property def basis_number_linear(self): if self.number_of_bases_linear > 1: - return self._cache['basis_number_linear'] + return self._cache["basis_number_linear"] else: return 0 @property def basis_number_quadratic(self): if self.number_of_bases_quadratic > 1: - return self._cache['basis_number_quadratic'] + return self._cache["basis_number_quadratic"] else: return 0 @@ -496,9 +570,13 @@ def waveform_generator(self): @waveform_generator.setter def waveform_generator(self, waveform_generator): - self._waveform_generator = _WaveformGeneratorWrapper(waveform_generator, self._update_basis) + self._waveform_generator = _WaveformGeneratorWrapper( + waveform_generator, self._update_basis + ) - def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array=True, parameters): + def calculate_snrs( + self, waveform_polarizations, interferometer, *, return_array=True, parameters + ): """ Compute the snrs for ROQ @@ -511,52 +589,63 @@ def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array if self.time_marginalization: time_ref = self._beam_pattern_reference_time else: - time_ref = parameters['geocent_time'] + time_ref = parameters["geocent_time"] - frequency_nodes = self.waveform_generator.waveform_arguments['frequency_nodes'] - linear_indices = self.waveform_generator.waveform_arguments['linear_indices'] - quadratic_indices = self.waveform_generator.waveform_arguments['quadratic_indices'] + frequency_nodes = self.waveform_generator.waveform_arguments["frequency_nodes"] + linear_indices = self.waveform_generator.waveform_arguments["linear_indices"] + quadratic_indices = self.waveform_generator.waveform_arguments[ + "quadratic_indices" + ] h_linear = 0j h_quadratic = 0j - for mode in waveform_polarizations['linear']: + for mode in waveform_polarizations["linear"]: response = interferometer.antenna_response( - parameters['ra'], parameters['dec'], - time_ref, - parameters['psi'], - mode + parameters["ra"], parameters["dec"], time_ref, parameters["psi"], mode ) - h_linear += waveform_polarizations['linear'][mode] * response - h_quadratic += waveform_polarizations['quadratic'][mode] * response + h_linear += waveform_polarizations["linear"][mode] * response + h_quadratic += waveform_polarizations["quadratic"][mode] * response xp = array_module(h_linear) calib_factor = interferometer.calibration_model.get_calibration_factor( - xp.asarray(frequency_nodes), prefix=f'recalib_{interferometer.name}_', xp=xp, **parameters) + xp.asarray(frequency_nodes), + prefix=f"recalib_{interferometer.name}_", + xp=xp, + **parameters, + ) h_linear *= calib_factor[linear_indices] h_quadratic *= calib_factor[quadratic_indices] optimal_snr_squared = xp.vdot( - xp.abs(h_quadratic)**2, - xp.asarray(self.weights[interferometer.name + '_quadratic'][self.basis_number_quadratic]) + xp.abs(h_quadratic) ** 2, + xp.asarray( + self.weights[interferometer.name + "_quadratic"][ + self.basis_number_quadratic + ] + ), ) dt = interferometer.time_delay_from_geocenter( - parameters['ra'], parameters['dec'], time_ref) - dt_geocent = parameters['geocent_time'] - interferometer.strain_data.start_time + parameters["ra"], parameters["dec"], time_ref + ) + dt_geocent = parameters["geocent_time"] - interferometer.strain_data.start_time ifo_time = dt_geocent + dt indices, in_bounds = self._closest_time_indices( - ifo_time, xp.asarray(self.weights['time_samples'])) - indices = xp.clip(xp.asarray(indices), 0, len(self.weights['time_samples']) - 1) + ifo_time, xp.asarray(self.weights["time_samples"]) + ) + indices = xp.clip(xp.asarray(indices), 0, len(self.weights["time_samples"]) - 1) d_inner_h_tc_array = xp.einsum( - 'i,ji->j', + "i,ji->j", xp.conj(h_linear), xp.asarray( - self.weights[interferometer.name + '_linear'][self.basis_number_linear] + self.weights[interferometer.name + "_linear"][self.basis_number_linear] )[indices], ) d_inner_h = self._interp_five_samples( - xp.asarray(self.weights['time_samples'])[indices], d_inner_h_tc_array, ifo_time + xp.asarray(self.weights["time_samples"])[indices], + d_inner_h_tc_array, + ifo_time, ) with np.errstate(invalid="ignore"): @@ -569,8 +658,10 @@ def calculate_snrs(self, waveform_polarizations, interferometer, *, return_array ifo_times = self._times - interferometer.strain_data.start_time ifo_times += dt if self.jitter_time: - ifo_times += parameters['time_jitter'] - d_inner_h_array = self._calculate_d_inner_h_array(ifo_times, h_linear, interferometer.name) + ifo_times += parameters["time_jitter"] + d_inner_h_array = self._calculate_d_inner_h_array( + ifo_times, h_linear, interferometer.name + ) else: d_inner_h_array = None @@ -601,7 +692,9 @@ def _closest_time_indices(time, samples): Whether the indices are for valid times """ xp = array_module(time) - closest = xp.astype(xp.floor((time - samples[0]) / (samples[1] - samples[0])), int) + closest = xp.astype( + xp.floor((time - samples[0]) / (samples[1] - samples[0])), int + ) indices = [closest + ii for ii in [-2, -1, 0, 1, 2]] in_bounds = (indices[0] >= 0) & (indices[-1] < len(samples)) return indices, in_bounds @@ -626,12 +719,18 @@ def _interp_five_samples(time_samples, values, time): value: float The value of the function at the input time """ - r1 = (-values[0] + 8. * values[1] - 14. * values[2] + 8. * values[3] - values[4]) / 4. - r2 = values[2] - 2. * values[3] + values[4] + r1 = ( + -values[0] + + 8.0 * values[1] + - 14.0 * values[2] + + 8.0 * values[3] + - values[4] + ) / 4.0 + r2 = values[2] - 2.0 * values[3] + values[4] a = (time_samples[3] - time) / max(time_samples[1] - time_samples[0], 1e-12) - b = 1. - a - c = (a**3. - a) / 6. - d = (b**3. - b) / 6. + b = 1.0 - a + c = (a**3.0 - a) / 6.0 + d = (b**3.0 - b) / 6.0 return a * values[2] + b * values[3] + c * r1 + d * r2 def _calculate_d_inner_h_array(self, times, h_linear, ifo_name, *, xp=None): @@ -654,12 +753,16 @@ def _calculate_d_inner_h_array(self, times, h_linear, ifo_name, *, xp=None): """ if xp is None: xp = aac.array_namespace(h_linear) - roq_time_space = self.weights['time_samples'][1] - self.weights['time_samples'][0] - times_per_roq_time_space = (times - self.weights['time_samples'][0]) / roq_time_space + roq_time_space = ( + self.weights["time_samples"][1] - self.weights["time_samples"][0] + ) + times_per_roq_time_space = ( + times - self.weights["time_samples"][0] + ) / roq_time_space closest_idxs = xp.astype(xp.floor(times_per_roq_time_space), int) # Get the nearest 5 samples of d_inner_h. Calculate only the required d_inner_h values if the time # spacing is larger than 5 times the ROQ time spacing. - weights_linear = self.weights[ifo_name + '_linear'][self.basis_number_linear] + weights_linear = self.weights[ifo_name + "_linear"][self.basis_number_linear] h_linear_conj = h_linear.conj() if (times[1] - times[0]) / roq_time_space > 5: d_inner_h_m2 = weights_linear[closest_idxs - 2] @ h_linear_conj @@ -676,15 +779,21 @@ def _calculate_d_inner_h_array(self, times, h_linear, ifo_name, *, xp=None): d_inner_h_p2 = d_inner_h_at_roq_time_samples[closest_idxs + 2] # quantities required for spline interpolation b = times_per_roq_time_space - closest_idxs - a = 1. - b - c = (a**3. - a) / 6. - d = (b**3. - b) / 6. - r1 = (-d_inner_h_m2 + 8. * d_inner_h_m1 - 14. * d_inner_h_0 + 8. * d_inner_h_p1 - d_inner_h_p2) / 4. - r2 = d_inner_h_0 - 2. * d_inner_h_p1 + d_inner_h_p2 + a = 1.0 - b + c = (a**3.0 - a) / 6.0 + d = (b**3.0 - b) / 6.0 + r1 = ( + -d_inner_h_m2 + + 8.0 * d_inner_h_m1 + - 14.0 * d_inner_h_0 + + 8.0 * d_inner_h_p1 + - d_inner_h_p2 + ) / 4.0 + r2 = d_inner_h_0 - 2.0 * d_inner_h_p1 + d_inner_h_p2 return a * d_inner_h_0 + b * d_inner_h_p1 + c * r1 + d * r2 def perform_roq_params_check(self, ifo=None): - """ Perform checking that the prior and data are valid for the ROQ + """Perform checking that the prior and data are valid for the ROQ Parameters ========== @@ -696,47 +805,54 @@ def perform_roq_params_check(self, ifo=None): return else: if getattr(self, "roq_params_file", None) is not None: - msg = ("Check ROQ params {} with roq_scale_factor={}" - .format(self.roq_params_file, self.roq_scale_factor)) + msg = "Check ROQ params {} with roq_scale_factor={}".format( + self.roq_params_file, self.roq_scale_factor + ) else: - msg = ("Check ROQ params with roq_scale_factor={}" - .format(self.roq_scale_factor)) + msg = "Check ROQ params with roq_scale_factor={}".format( + self.roq_scale_factor + ) logger.info(msg) roq_params = self.roq_params - roq_minimum_frequency = roq_params['flow'] * self.roq_scale_factor - roq_maximum_frequency = roq_params['fhigh'] * self.roq_scale_factor - roq_segment_length = roq_params['seglen'] / self.roq_scale_factor + roq_minimum_frequency = roq_params["flow"] * self.roq_scale_factor + roq_maximum_frequency = roq_params["fhigh"] * self.roq_scale_factor + roq_segment_length = roq_params["seglen"] / self.roq_scale_factor try: - roq_minimum_chirp_mass = roq_params['chirpmassmin'] / self.roq_scale_factor + roq_minimum_chirp_mass = roq_params["chirpmassmin"] / self.roq_scale_factor except ValueError: roq_minimum_chirp_mass = None try: - roq_maximum_chirp_mass = roq_params['chirpmassmax'] / self.roq_scale_factor + roq_maximum_chirp_mass = roq_params["chirpmassmax"] / self.roq_scale_factor except ValueError: roq_maximum_chirp_mass = None try: - roq_minimum_component_mass = roq_params['compmin'] / self.roq_scale_factor + roq_minimum_component_mass = roq_params["compmin"] / self.roq_scale_factor except ValueError: roq_minimum_component_mass = None if float(ifo.maximum_frequency) > roq_maximum_frequency: raise BilbyROQParamsRangeError( - "Requested maximum frequency {} larger than ROQ basis fhigh {}" - .format(ifo.maximum_frequency, roq_maximum_frequency) + "Requested maximum frequency {} larger than ROQ basis fhigh {}".format( + ifo.maximum_frequency, roq_maximum_frequency + ) ) if float(ifo.minimum_frequency) < roq_minimum_frequency: raise BilbyROQParamsRangeError( - "Requested minimum frequency {} lower than ROQ basis flow {}" - .format(ifo.minimum_frequency, roq_minimum_frequency) + "Requested minimum frequency {} lower than ROQ basis flow {}".format( + ifo.minimum_frequency, roq_minimum_frequency + ) ) if float(ifo.strain_data.duration) != roq_segment_length: raise BilbyROQParamsRangeError( - "Requested duration differs from ROQ basis seglen") + "Requested duration differs from ROQ basis seglen" + ) priors = self.priors if isinstance(priors, CBCPriorDict) is False: - logger.warning("Unable to check ROQ parameter bounds: priors not understood") + logger.warning( + "Unable to check ROQ parameter bounds: priors not understood" + ) return if roq_minimum_chirp_mass is not None: @@ -744,8 +860,9 @@ def perform_roq_params_check(self, ifo=None): logger.warning("Unable to check minimum chirp mass ROQ bounds") elif priors.minimum_chirp_mass < roq_minimum_chirp_mass: raise BilbyROQParamsRangeError( - "Prior minimum chirp mass {} less than ROQ basis bound {}" - .format(priors.minimum_chirp_mass, roq_minimum_chirp_mass) + "Prior minimum chirp mass {} less than ROQ basis bound {}".format( + priors.minimum_chirp_mass, roq_minimum_chirp_mass + ) ) if roq_maximum_chirp_mass is not None: @@ -753,8 +870,9 @@ def perform_roq_params_check(self, ifo=None): logger.warning("Unable to check maximum_chirp mass ROQ bounds") elif priors.maximum_chirp_mass > roq_maximum_chirp_mass: raise BilbyROQParamsRangeError( - "Prior maximum chirp mass {} greater than ROQ basis bound {}" - .format(priors.maximum_chirp_mass, roq_maximum_chirp_mass) + "Prior maximum chirp mass {} greater than ROQ basis bound {}".format( + priors.maximum_chirp_mass, roq_maximum_chirp_mass + ) ) if roq_minimum_component_mass is not None: @@ -762,8 +880,9 @@ def perform_roq_params_check(self, ifo=None): logger.warning("Unable to check minimum component mass ROQ bounds") elif priors.minimum_component_mass < roq_minimum_component_mass: raise BilbyROQParamsRangeError( - "Prior minimum component mass {} less than ROQ basis bound {}" - .format(priors.minimum_component_mass, roq_minimum_component_mass) + "Prior minimum component mass {} less than ROQ basis bound {}".format( + priors.minimum_component_mass, roq_minimum_component_mass + ) ) def _set_weights(self, linear_matrix, quadratic_matrix): @@ -779,56 +898,90 @@ def _set_weights(self, linear_matrix, quadratic_matrix): xp = aac.array_namespace(self.interferometers.frequency_array) time_space = self._get_time_resolution() number_of_time_samples = int(self.interferometers.duration / time_space) - earth_light_crossing_time = 2 * radius_of_earth / speed_of_light + 5 * time_space + earth_light_crossing_time = ( + 2 * radius_of_earth / speed_of_light + 5 * time_space + ) start_idx = max( 0, - int(np.floor(( - self.priors['{}_time'.format(self.time_reference)].minimum - - earth_light_crossing_time - - self.interferometers.start_time - ) / time_space)) + int( + np.floor( + ( + self.priors["{}_time".format(self.time_reference)].minimum + - earth_light_crossing_time + - self.interferometers.start_time + ) + / time_space + ) + ), ) end_idx = min( number_of_time_samples - 1, - int(np.ceil(( - self.priors['{}_time'.format(self.time_reference)].maximum - + earth_light_crossing_time - - self.interferometers.start_time - ) / time_space)) + int( + np.ceil( + ( + self.priors["{}_time".format(self.time_reference)].maximum + + earth_light_crossing_time + - self.interferometers.start_time + ) + / time_space + ) + ), + ) + self.weights["time_samples"] = xp.arange(start_idx, end_idx + 1) * float( + time_space + ) + logger.info( + "Using {} ROQ time samples".format(len(self.weights["time_samples"])) ) - self.weights['time_samples'] = xp.arange(start_idx, end_idx + 1) * float(time_space) - logger.info("Using {} ROQ time samples".format(len(self.weights['time_samples']))) # select bases to be used, set prior ranges and frequency nodes if exist idxs_in_prior_range = dict() - for basis_type, matrix in zip(['linear', 'quadratic'], [linear_matrix, quadratic_matrix]): - key = f'prior_range_{basis_type}' + for basis_type, matrix in zip( + ["linear", "quadratic"], [linear_matrix, quadratic_matrix] + ): + key = f"prior_range_{basis_type}" if key in matrix: prior_ranges = {} for param_name in matrix[key]: - if 'roq_scale_power' in matrix[key][param_name].attrs: - roq_scale_factor = self.roq_scale_factor**matrix[key][param_name].attrs['roq_scale_power'] + if "roq_scale_power" in matrix[key][param_name].attrs: + roq_scale_factor = ( + self.roq_scale_factor + ** matrix[key][param_name].attrs["roq_scale_power"] + ) else: - roq_scale_factor = 1. - prior_ranges[param_name] = matrix[key][param_name][()] * roq_scale_factor - selected_idxs, selected_prior_ranges = self._select_prior_ranges(prior_ranges) + roq_scale_factor = 1.0 + prior_ranges[param_name] = ( + matrix[key][param_name][()] * roq_scale_factor + ) + selected_idxs, selected_prior_ranges = self._select_prior_ranges( + prior_ranges + ) if len(selected_idxs) == 0: - raise BilbyROQParamsRangeError(f"There are no {basis_type} ROQ bases within the prior range.") + raise BilbyROQParamsRangeError( + f"There are no {basis_type} ROQ bases within the prior range." + ) self.weights[key] = selected_prior_ranges idxs_in_prior_range[basis_type] = selected_idxs else: idxs_in_prior_range[basis_type] = [0] - if 'frequency_nodes' in matrix[f'basis_{basis_type}'][str(idxs_in_prior_range[basis_type][0])]: - self.weights[f'frequency_nodes_{basis_type}'] = [ - matrix[f'basis_{basis_type}'][str(i)]['frequency_nodes'][()] * self.roq_scale_factor - for i in idxs_in_prior_range[basis_type]] - - if 'multiband_linear' in linear_matrix: - multiband_linear = linear_matrix['multiband_linear'][()] + if ( + "frequency_nodes" + in matrix[f"basis_{basis_type}"][ + str(idxs_in_prior_range[basis_type][0]) + ] + ): + self.weights[f"frequency_nodes_{basis_type}"] = [ + matrix[f"basis_{basis_type}"][str(i)]["frequency_nodes"][()] + * self.roq_scale_factor + for i in idxs_in_prior_range[basis_type] + ] + + if "multiband_linear" in linear_matrix: + multiband_linear = linear_matrix["multiband_linear"][()] else: multiband_linear = False - if 'multiband_quadratic' in quadratic_matrix: - multiband_quadratic = quadratic_matrix['multiband_quadratic'][()] + if "multiband_quadratic" in quadratic_matrix: + multiband_quadratic = quadratic_matrix["multiband_quadratic"][()] else: multiband_quadratic = False @@ -839,45 +992,74 @@ def _set_weights(self, linear_matrix, quadratic_matrix): for ifo in self.interferometers: if self.roq_params is not None: # Get scaled ROQ quantities - roq_scaled_minimum_frequency = self.roq_params['flow'] * self.roq_scale_factor - roq_scaled_maximum_frequency = self.roq_params['fhigh'] * self.roq_scale_factor - roq_scaled_segment_length = self.roq_params['seglen'] / self.roq_scale_factor + roq_scaled_minimum_frequency = ( + self.roq_params["flow"] * self.roq_scale_factor + ) + roq_scaled_maximum_frequency = ( + self.roq_params["fhigh"] * self.roq_scale_factor + ) + roq_scaled_segment_length = ( + self.roq_params["seglen"] / self.roq_scale_factor + ) # Generate frequencies for the ROQ roq_frequencies = create_frequency_series( sampling_frequency=roq_scaled_maximum_frequency * 2, - duration=roq_scaled_segment_length) + duration=roq_scaled_segment_length, + ) roq_mask = roq_frequencies >= roq_scaled_minimum_frequency roq_frequencies = roq_frequencies[roq_mask] - overlap_frequencies, ifo_idxs_this_ifo, roq_idxs_this_ifo = np.intersect1d( - np.asarray(ifo.frequency_array[ifo.frequency_mask]), roq_frequencies, - return_indices=True) + overlap_frequencies, ifo_idxs_this_ifo, roq_idxs_this_ifo = ( + np.intersect1d( + np.asarray(ifo.frequency_array[ifo.frequency_mask]), + roq_frequencies, + return_indices=True, + ) + ) else: - overlap_frequencies = np.asarray(ifo.frequency_array[ifo.frequency_mask]) + overlap_frequencies = np.asarray( + ifo.frequency_array[ifo.frequency_mask] + ) roq_idxs_this_ifo = np.arange( - linear_matrix['basis_linear'][str(idxs_in_prior_range['linear'][0])]['basis'].shape[1], - dtype=int) + linear_matrix["basis_linear"][ + str(idxs_in_prior_range["linear"][0]) + ]["basis"].shape[1], + dtype=int, + ) ifo_idxs_this_ifo = np.arange(sum(ifo.frequency_mask)) if len(ifo_idxs_this_ifo) != len(roq_idxs_this_ifo): raise ValueError( "Mismatch between ROQ basis and frequency array for " - "{}".format(ifo.name)) + "{}".format(ifo.name) + ) logger.info( "Building ROQ weights for {} with {} frequencies between {} " "and {}.".format( - ifo.name, len(overlap_frequencies), - min(overlap_frequencies), max(overlap_frequencies))) + ifo.name, + len(overlap_frequencies), + min(overlap_frequencies), + max(overlap_frequencies), + ) + ) roq_idxs[ifo.name] = roq_idxs_this_ifo ifo_idxs[ifo.name] = ifo_idxs_this_ifo if multiband_linear: - self._set_weights_linear_multiband(linear_matrix, idxs_in_prior_range['linear']) + self._set_weights_linear_multiband( + linear_matrix, idxs_in_prior_range["linear"] + ) else: - self._set_weights_linear(linear_matrix, idxs_in_prior_range['linear'], roq_idxs, ifo_idxs) + self._set_weights_linear( + linear_matrix, idxs_in_prior_range["linear"], roq_idxs, ifo_idxs + ) if multiband_quadratic: - self._set_weights_quadratic_multiband(quadratic_matrix, idxs_in_prior_range['quadratic']) + self._set_weights_quadratic_multiband( + quadratic_matrix, idxs_in_prior_range["quadratic"] + ) else: - self._set_weights_quadratic(quadratic_matrix, idxs_in_prior_range['quadratic'], roq_idxs, ifo_idxs) + self._set_weights_quadratic( + quadratic_matrix, idxs_in_prior_range["quadratic"], roq_idxs, ifo_idxs + ) def _set_weights_linear(self, linear_matrix, basis_idxs, roq_idxs, ifo_idxs): """ @@ -899,30 +1081,38 @@ def _set_weights_linear(self, linear_matrix, basis_idxs, roq_idxs, ifo_idxs): """ for ifo in self.interferometers: - self.weights[ifo.name + '_linear'] = [] - time_space = self.weights['time_samples'][1] - self.weights['time_samples'][0] + self.weights[ifo.name + "_linear"] = [] + time_space = self.weights["time_samples"][1] - self.weights["time_samples"][0] number_of_time_samples = int(self.interferometers.duration / time_space) - start_idx = int(self.weights['time_samples'][0] / time_space) - end_idx = int(self.weights['time_samples'][-1] / time_space) + start_idx = int(self.weights["time_samples"][0] / time_space) + end_idx = int(self.weights["time_samples"][-1] / time_space) nonzero_idxs = {} data_over_psd = {} for ifo in self.interferometers: nonzero_idxs[ifo.name] = ifo_idxs[ifo.name] + int( - ifo.minimum_frequency * self.interferometers.duration) + ifo.minimum_frequency * self.interferometers.duration + ) data_over_psd[ifo.name] = ( ifo.frequency_domain_strain[ifo.frequency_mask][ifo_idxs[ifo.name]] - / ifo.power_spectral_density_array[ifo.frequency_mask][ifo_idxs[ifo.name]] + / ifo.power_spectral_density_array[ifo.frequency_mask][ + ifo_idxs[ifo.name] + ] ) xp = array_module(data_over_psd) if aac.is_numpy_namespace(xp): try: import pyfftw + ifft_input = pyfftw.empty_aligned(number_of_time_samples, dtype=complex) - ifft_output = pyfftw.empty_aligned(number_of_time_samples, dtype=complex) - ifft = pyfftw.FFTW(ifft_input, ifft_output, direction='FFTW_BACKWARD') + ifft_output = pyfftw.empty_aligned( + number_of_time_samples, dtype=complex + ) + ifft = pyfftw.FFTW(ifft_input, ifft_output, direction="FFTW_BACKWARD") except ImportError: pyfftw = None - logger.warning("You do not have pyfftw installed, falling back to numpy.fft.") + logger.warning( + "You do not have pyfftw installed, falling back to numpy.fft." + ) ifft_input = np.zeros(number_of_time_samples, dtype=complex) ifft = np.fft.ifft else: @@ -931,22 +1121,33 @@ def _set_weights_linear(self, linear_matrix, basis_idxs, roq_idxs, ifo_idxs): ifft = xp.fft.ifft for basis_idx in basis_idxs: logger.info(f"Building linear ROQ weights for the {basis_idx}-th basis.") - linear_matrix_single = xp.asarray(linear_matrix['basis_linear'][str(basis_idx)]['basis']) + linear_matrix_single = xp.asarray( + linear_matrix["basis_linear"][str(basis_idx)]["basis"] + ) basis_size = linear_matrix_single.shape[0] for ifo in self.interferometers: if pyfftw: - ifft_input[:] *= 0. + ifft_input[:] *= 0.0 else: ifft_input *= 0 - linear_weights = \ - xp.zeros((basis_size, len(self.weights['time_samples'])), dtype=complex) + linear_weights = xp.zeros( + (basis_size, len(self.weights["time_samples"])), dtype=complex + ) for i in range(basis_size): - basis_element = xp.asarray(linear_matrix_single[i][roq_idxs[ifo.name]]).conj() - ifft_input = xpx.at(ifft_input, nonzero_idxs[ifo.name]).set(data_over_psd[ifo.name] * basis_element) - linear_weights = xpx.at(linear_weights, i).set(ifft(ifft_input)[start_idx:end_idx + 1]) + basis_element = xp.asarray( + linear_matrix_single[i][roq_idxs[ifo.name]] + ).conj() + ifft_input = xpx.at(ifft_input, nonzero_idxs[ifo.name]).set( + data_over_psd[ifo.name] * basis_element + ) + linear_weights = xpx.at(linear_weights, i).set( + ifft(ifft_input)[start_idx : end_idx + 1] + ) linear_weights = linear_weights.T - linear_weights *= 4. * number_of_time_samples / float(self.interferometers.duration) - self.weights[ifo.name + '_linear'].append(linear_weights) + linear_weights *= ( + 4.0 * number_of_time_samples / float(self.interferometers.duration) + ) + self.weights[ifo.name + "_linear"].append(linear_weights) if pyfftw is not None: pyfftw.forget_wisdom() @@ -963,48 +1164,74 @@ def _set_weights_linear_multiband(self, linear_matrix, basis_idxs): """ for ifo in self.interferometers: - self.weights[ifo.name + '_linear'] = [] + self.weights[ifo.name + "_linear"] = [] xp = aac.array_namespace(self.interferometers.frequency_array) - Tbs = linear_matrix['durations_s_linear'][()] / self.roq_scale_factor - start_end_frequency_bins = linear_matrix['start_end_frequency_bins_linear'][()] - basis_dimension = np.sum(start_end_frequency_bins[:, 1] - start_end_frequency_bins[:, 0] + 1) + Tbs = linear_matrix["durations_s_linear"][()] / self.roq_scale_factor + start_end_frequency_bins = linear_matrix["start_end_frequency_bins_linear"][()] + basis_dimension = np.sum( + start_end_frequency_bins[:, 1] - start_end_frequency_bins[:, 0] + 1 + ) fhigh_basis = np.max(start_end_frequency_bins[:, 1] / Tbs) # prepare time-shifted data, which is multiplied by basis tc_shifted_data = dict() for ifo in self.interferometers: - over_whitened_frequency_data = xp.zeros(int(fhigh_basis * ifo.duration) + 1, dtype=complex) + over_whitened_frequency_data = xp.zeros( + int(fhigh_basis * ifo.duration) + 1, dtype=complex + ) over_whitened_frequency_data = xpx.at( - over_whitened_frequency_data, xp.arange(len(ifo.frequency_domain_strain))[ifo.frequency_mask] + over_whitened_frequency_data, + xp.arange(len(ifo.frequency_domain_strain))[ifo.frequency_mask], ).set( - ifo.frequency_domain_strain[ifo.frequency_mask] / ifo.power_spectral_density_array[ifo.frequency_mask] + ifo.frequency_domain_strain[ifo.frequency_mask] + / ifo.power_spectral_density_array[ifo.frequency_mask] ) over_whitened_time_data = xp.fft.irfft(over_whitened_frequency_data) - tc_shifted_data[ifo.name] = xp.zeros((basis_dimension, len(self.weights['time_samples'])), dtype=complex) + tc_shifted_data[ifo.name] = xp.zeros( + (basis_dimension, len(self.weights["time_samples"])), dtype=complex + ) start_idx_of_band = 0 for b, Tb in enumerate(Tbs): start_frequency_bin, end_frequency_bin = start_end_frequency_bins[b] fs = xp.arange(start_frequency_bin, end_frequency_bin + 1) / Tb Db = xp.fft.rfft( - over_whitened_time_data[-int(2. * fhigh_basis * Tb):] - )[start_frequency_bin:end_frequency_bin + 1] - start_idx_of_next_band = start_idx_of_band + end_frequency_bin - start_frequency_bin + 1 - this_data = xp.zeros(len(self.weights['time_samples']), dtype=complex) + over_whitened_time_data[-int(2.0 * fhigh_basis * Tb) :] + )[start_frequency_bin : end_frequency_bin + 1] + start_idx_of_next_band = ( + start_idx_of_band + end_frequency_bin - start_frequency_bin + 1 + ) + this_data = xp.zeros(len(self.weights["time_samples"]), dtype=complex) sl = slice(start_idx_of_band, start_idx_of_next_band) this_data = ( - 4. / Tb * Db[:, None] * xp.exp( - 2. * np.pi * 1j * fs[:, None] - * (xp.asarray(self.weights['time_samples'][None, :]) - ifo.duration + Tb) + 4.0 + / Tb + * Db[:, None] + * xp.exp( + 2.0 + * np.pi + * 1j + * fs[:, None] + * ( + xp.asarray(self.weights["time_samples"][None, :]) + - ifo.duration + + Tb + ) ) ) - tc_shifted_data[ifo.name] = xpx.at(tc_shifted_data[ifo.name], sl).set(this_data) + tc_shifted_data[ifo.name] = xpx.at(tc_shifted_data[ifo.name], sl).set( + this_data + ) start_idx_of_band = start_idx_of_next_band # compute inner products for basis_idx in basis_idxs: logger.info(f"Building linear ROQ weights for the {basis_idx}-th basis.") - linear_matrix_single = xp.asarray(linear_matrix['basis_linear'][str(basis_idx)]['basis'][()]) + linear_matrix_single = xp.asarray( + linear_matrix["basis_linear"][str(basis_idx)]["basis"][()] + ) for ifo in self.interferometers: - self.weights[ifo.name + '_linear'].append((linear_matrix_single.conj() @ tc_shifted_data[ifo.name]).T) + self.weights[ifo.name + "_linear"].append( + (linear_matrix_single.conj() @ tc_shifted_data[ifo.name]).T + ) def _set_weights_quadratic(self, quadratic_matrix, basis_idxs, roq_idxs, ifo_idxs): """ @@ -1025,15 +1252,25 @@ def _set_weights_quadratic(self, quadratic_matrix, basis_idxs, roq_idxs, ifo_idx """ for ifo in self.interferometers: - self.weights[ifo.name + '_quadratic'] = [] + self.weights[ifo.name + "_quadratic"] = [] xp = aac.array_namespace(self.interferometers.frequency_array) for basis_idx in basis_idxs: logger.info(f"Building quadratic ROQ weights for the {basis_idx}-th basis.") - quadratic_matrix_single = xp.asarray(quadratic_matrix['basis_quadratic'][str(basis_idx)]['basis'][()].real) + quadratic_matrix_single = xp.asarray( + quadratic_matrix["basis_quadratic"][str(basis_idx)]["basis"][()].real + ) for ifo in self.interferometers: - inv_psd = xp.asarray(1 / ifo.power_spectral_density_array[ifo.frequency_mask][ifo_idxs[ifo.name]]) - self.weights[ifo.name + '_quadratic'].append( - 4. / ifo.strain_data.duration * quadratic_matrix_single[:, roq_idxs[ifo.name]] @ inv_psd + inv_psd = xp.asarray( + 1 + / ifo.power_spectral_density_array[ifo.frequency_mask][ + ifo_idxs[ifo.name] + ] + ) + self.weights[ifo.name + "_quadratic"].append( + 4.0 + / ifo.strain_data.duration + * quadratic_matrix_single[:, roq_idxs[ifo.name]] + @ inv_psd ) del quadratic_matrix_single @@ -1050,11 +1287,15 @@ def _set_weights_quadratic_multiband(self, quadratic_matrix, basis_idxs): """ for ifo in self.interferometers: - self.weights[ifo.name + '_quadratic'] = [] + self.weights[ifo.name + "_quadratic"] = [] xp = aac.array_namespace(self.interferometers.frequency_array) - Tbs = quadratic_matrix['durations_s_quadratic'][()] / self.roq_scale_factor - start_end_frequency_bins = quadratic_matrix['start_end_frequency_bins_quadratic'][()] - basis_dimension = np.sum(start_end_frequency_bins[:, 1] - start_end_frequency_bins[:, 0] + 1) + Tbs = quadratic_matrix["durations_s_quadratic"][()] / self.roq_scale_factor + start_end_frequency_bins = quadratic_matrix[ + "start_end_frequency_bins_quadratic" + ][()] + basis_dimension = np.sum( + start_end_frequency_bins[:, 1] - start_end_frequency_bins[:, 0] + 1 + ) fhigh_basis = np.max(start_end_frequency_bins[:, 1] / Tbs) # prepare coefficients multiplied by basis multibanded_inverse_psd = dict() @@ -1062,7 +1303,7 @@ def _set_weights_quadratic_multiband(self, quadratic_matrix, basis_idxs): inverse_psd_frequency = xp.zeros(int(fhigh_basis * ifo.duration) + 1) sl = np.arange(len(ifo.power_spectral_density_array))[ifo.frequency_mask] inverse_psd_frequency = xpx.at(inverse_psd_frequency, sl).set( - 1. / xp.asarray(ifo.power_spectral_density_array[ifo.frequency_mask]) + 1.0 / xp.asarray(ifo.power_spectral_density_array[ifo.frequency_mask]) ) inverse_psd_time = xp.fft.irfft(inverse_psd_frequency) multibanded_inverse_psd[ifo.name] = xp.zeros(basis_dimension) @@ -1070,22 +1311,38 @@ def _set_weights_quadratic_multiband(self, quadratic_matrix, basis_idxs): for b, Tb in enumerate(Tbs): start_frequency_bin, end_frequency_bin = start_end_frequency_bins[b] number_of_samples_half = int(fhigh_basis * Tb) - start_idx_of_next_band = start_idx_of_band + end_frequency_bin - start_frequency_bin + 1 + start_idx_of_next_band = ( + start_idx_of_band + end_frequency_bin - start_frequency_bin + 1 + ) sl = slice(start_idx_of_band, start_idx_of_next_band) - this_data = 4. / Tb * xp.fft.rfft( - xp.concat([inverse_psd_time[:number_of_samples_half], inverse_psd_time[-number_of_samples_half:]]) - )[start_frequency_bin:end_frequency_bin + 1].real - multibanded_inverse_psd[ifo.name] = xpx.at(multibanded_inverse_psd[ifo.name], sl).set(this_data) + this_data = ( + 4.0 + / Tb + * xp.fft.rfft( + xp.concat( + [ + inverse_psd_time[:number_of_samples_half], + inverse_psd_time[-number_of_samples_half:], + ] + ) + )[start_frequency_bin : end_frequency_bin + 1].real + ) + multibanded_inverse_psd[ifo.name] = xpx.at( + multibanded_inverse_psd[ifo.name], sl + ).set(this_data) start_idx_of_band = start_idx_of_next_band # compute inner products for basis_idx in basis_idxs: logger.info(f"Building quadratic ROQ weights for the {basis_idx}-th basis.") - quadratic_matrix_single = xp.asarray(quadratic_matrix['basis_quadratic'][str(basis_idx)]['basis'][()].real) + quadratic_matrix_single = xp.asarray( + quadratic_matrix["basis_quadratic"][str(basis_idx)]["basis"][()].real + ) for ifo in self.interferometers: - self.weights[ifo.name + '_quadratic'].append( - quadratic_matrix_single @ multibanded_inverse_psd[ifo.name]) + self.weights[ifo.name + "_quadratic"].append( + quadratic_matrix_single @ multibanded_inverse_psd[ifo.name] + ) - def save_weights(self, filename, format='hdf5'): + def save_weights(self, filename, format="hdf5"): """ Save ROQ weights into a single file. Support for json format was removed in :code:`v2.7`, only hdf5 and npz are supported. @@ -1097,45 +1354,44 @@ def save_weights(self, filename, format='hdf5'): format : str The format to save the weight in, should be in :code:`hdf5, npz`. """ - if format not in ['hdf5', 'npz']: + if format not in ["hdf5", "npz"]: raise IOError(f"Format {format} not recognized.") if format not in filename: filename += "." + format logger.info(f"Saving ROQ weights to {filename}") - if format == 'npz': + if format == "npz": if self.number_of_bases_linear > 1 or self.number_of_bases_quadratic > 1: - raise ValueError(f'Format {format} not compatible with multiple bases') + raise ValueError(f"Format {format} not compatible with multiple bases") weights = dict() - weights['time_samples'] = self.weights['time_samples'] - for basis_type in ['linear', 'quadratic']: + weights["time_samples"] = self.weights["time_samples"] + for basis_type in ["linear", "quadratic"]: for ifo in self.interferometers: - key = f'{ifo.name}_{basis_type}' + key = f"{ifo.name}_{basis_type}" weights[key] = self.weights[key][0] np.savez(filename, **weights) else: import h5py - with h5py.File(filename, 'w') as f: - f.create_dataset('time_samples', - data=self.weights['time_samples']) - for basis_type in ['linear', 'quadratic']: - key = f'prior_range_{basis_type}' + + with h5py.File(filename, "w") as f: + f.create_dataset("time_samples", data=self.weights["time_samples"]) + for basis_type in ["linear", "quadratic"]: + key = f"prior_range_{basis_type}" if key in self.weights: grp = f.create_group(key) for param_name in self.weights[key]: grp.create_dataset( - param_name, data=self.weights[key][param_name]) - key = f'frequency_nodes_{basis_type}' + param_name, data=self.weights[key][param_name] + ) + key = f"frequency_nodes_{basis_type}" if key in self.weights: grp = f.create_group(key) for i in range(len(self.weights[key])): - grp.create_dataset( - str(i), data=self.weights[key][i]) + grp.create_dataset(str(i), data=self.weights[key][i]) for ifo in self.interferometers: key = f"{ifo.name}_{basis_type}" grp = f.create_group(key) for i in range(len(self.weights[key])): - grp.create_dataset( - str(i), data=self.weights[key][i]) + grp.create_dataset(str(i), data=self.weights[key][i]) def load_weights(self, filename, format=None): """ @@ -1160,7 +1416,7 @@ def load_weights(self, filename, format=None): warnings.warn( "json format for ROQ weights is deprecated, use hdf5 instead.", - DeprecationWarning + DeprecationWarning, ) elif format not in ["npz", "hdf5"]: raise IOError(f"Format {format} not recognized.") @@ -1168,31 +1424,31 @@ def load_weights(self, filename, format=None): logger.info(f"Loading ROQ weights from {filename}") if format == "npz": weights = dict(np.load(filename)) - for basis_type in ['linear', 'quadratic']: + for basis_type in ["linear", "quadratic"]: for ifo in self.interferometers: - key = f'{ifo.name}_{basis_type}' + key = f"{ifo.name}_{basis_type}" weights[key] = [weights[key]] else: weights = dict() import h5py - with h5py.File(filename, 'r') as f: - weights['time_samples'] = f['time_samples'][()] - for basis_type in ['linear', 'quadratic']: - key = f'prior_range_{basis_type}' + + with h5py.File(filename, "r") as f: + weights["time_samples"] = f["time_samples"][()] + for basis_type in ["linear", "quadratic"]: + key = f"prior_range_{basis_type}" if key in f: - idxs_in_prior_range, selected_prior_ranges = \ + idxs_in_prior_range, selected_prior_ranges = ( self._select_prior_ranges(f[key]) + ) weights[key] = selected_prior_ranges else: idxs_in_prior_range = [0] - key = f'frequency_nodes_{basis_type}' + key = f"frequency_nodes_{basis_type}" if key in f: - weights[key] = [f[key][str(i)][()] - for i in idxs_in_prior_range] + weights[key] = [f[key][str(i)][()] for i in idxs_in_prior_range] for ifo in self.interferometers: key = f"{ifo.name}_{basis_type}" - weights[key] = [f[key][str(i)][()] - for i in idxs_in_prior_range] + weights[key] = [f[key][str(i)][()] for i in idxs_in_prior_range] return weights def _get_time_resolution(self): @@ -1211,7 +1467,7 @@ def _get_time_resolution(self): Time resolution """ - def calc_fhigh(freq, psd, scaling=20.): + def calc_fhigh(freq, psd, scaling=20.0): """ Parameters @@ -1229,21 +1485,22 @@ def calc_fhigh(freq, psd, scaling=20.): The maximum frequency which must be considered """ from scipy.integrate import simpson - integrand1 = np.power(freq, -7. / 3) / psd + + integrand1 = np.power(freq, -7.0 / 3) / psd integral1 = simpson(y=integrand1, x=freq) - integrand3 = np.power(freq, 2. / 3.) / (psd * integral1) + integrand3 = np.power(freq, 2.0 / 3.0) / (psd * integral1) f_3_bar = simpson(y=integrand3, x=freq) - f_high = scaling * f_3_bar**(1 / 3) + f_high = scaling * f_3_bar ** (1 / 3) return f_high def c_f_scaling(snr): - return (np.pi**2 * snr**2 / 6)**(1 / 3) + return (np.pi**2 * snr**2 / 6) ** (1 / 3) inj_snr_sq = 0 for ifo in self.interferometers: - inj_snr_sq += max(10, ifo.meta_data.get('optimal_SNR', 30))**2 + inj_snr_sq += max(10, ifo.meta_data.get("optimal_SNR", 30)) ** 2 psd = ifo.power_spectral_density_array[ifo.frequency_mask] freq = ifo.frequency_array[ifo.frequency_mask] @@ -1257,24 +1514,29 @@ def c_f_scaling(snr): # duration / delta_t needs to be a power of 2 for IFFT number_of_time_samples = max( self.interferometers.duration / delta_t, - self.interferometers.frequency_array[-1] * self.interferometers.duration + 1) - number_of_time_samples = int(2**np.ceil(np.log2(number_of_time_samples))) + self.interferometers.frequency_array[-1] * self.interferometers.duration + + 1, + ) + number_of_time_samples = int(2 ** np.ceil(np.log2(number_of_time_samples))) delta_t = self.interferometers.duration / number_of_time_samples logger.info("ROQ time-step = {}".format(delta_t)) return delta_t def _rescale_signal(self, signal, new_distance): - for kind in ['linear', 'quadratic']: + for kind in ["linear", "quadratic"]: for mode in signal[kind]: signal[kind][mode] *= self._ref_dist / new_distance - def generate_time_sample_from_marginalized_likelihood(self, signal_polarizations=None, *, parameters): + def generate_time_sample_from_marginalized_likelihood( + self, signal_polarizations=None, *, parameters + ): from ...core.utils import random parameters.update(self.get_sky_frame_parameters(parameters=parameters)) if signal_polarizations is None: - signal_polarizations = \ - self.waveform_generator.frequency_domain_strain(parameters) + signal_polarizations = self.waveform_generator.frequency_domain_strain( + parameters + ) snrs = self._CalculatedSNRs() @@ -1289,17 +1551,20 @@ def generate_time_sample_from_marginalized_likelihood(self, signal_polarizations if self.distance_marginalization: time_log_like = self.distance_marginalized_likelihood( - d_inner_h, h_inner_h, parameters=parameters) + d_inner_h, h_inner_h, parameters=parameters + ) elif self.phase_marginalization: time_log_like = ln_i0(abs(d_inner_h)) - h_inner_h.real / 2 else: - time_log_like = (d_inner_h.real - h_inner_h.real / 2) + time_log_like = d_inner_h.real - h_inner_h.real / 2 times = self._times if self.jitter_time: times = times + parameters["time_jitter"] - time_prior_array = np.asarray(self.priors['geocent_time'].prob(times)) - time_post = np.exp(np.asarray(time_log_like - max(time_log_like))) * time_prior_array + time_prior_array = np.asarray(self.priors["geocent_time"].prob(times)) + time_post = ( + np.exp(np.asarray(time_log_like - max(time_log_like))) * time_prior_array + ) time_post /= np.sum(time_post) return random.rng.choice(times, p=time_post) diff --git a/bilby/gw/prior.py b/bilby/gw/prior.py index ce96a02dd..e80d90174 100644 --- a/bilby/gw/prior.py +++ b/bilby/gw/prior.py @@ -10,25 +10,39 @@ from ..compat.utils import xp_wrap from ..core.prior import ( - PriorDict, Uniform, Prior, DeltaFunction, Gaussian, Interped, Constraint, - conditional_prior_factory, PowerLaw, ConditionalLogUniform, - ConditionalPriorDict, ConditionalBasePrior, BaseJointPriorDist, JointPrior, + PriorDict, + Uniform, + Prior, + DeltaFunction, + Gaussian, + Interped, + Constraint, + conditional_prior_factory, + PowerLaw, + ConditionalLogUniform, + ConditionalPriorDict, + ConditionalBasePrior, + BaseJointPriorDist, + JointPrior, JointPriorDistError, ) from ..core.utils import infer_args_from_method, logger, random, interp1d from .conversion import ( convert_to_lal_binary_black_hole_parameters, - convert_to_lal_binary_neutron_star_parameters, generate_mass_parameters, - generate_tidal_parameters, fill_from_fixed_priors, + convert_to_lal_binary_neutron_star_parameters, + generate_mass_parameters, + generate_tidal_parameters, + fill_from_fixed_priors, generate_all_bbh_parameters, chirp_mass_and_mass_ratio_to_total_mass, - total_mass_and_mass_ratio_to_component_masses) + total_mass_and_mass_ratio_to_component_masses, +) from .cosmology import get_cosmology, z_at_value from .source import PARAMETER_SETS from .utils import calculate_time_to_merger -DEFAULT_PRIOR_DIR = os.path.join(os.path.dirname(__file__), 'prior_files') +DEFAULT_PRIOR_DIR = os.path.join(os.path.dirname(__file__), "prior_files") class BilbyPriorConversionError(Exception): @@ -36,7 +50,7 @@ class BilbyPriorConversionError(Exception): def convert_to_flat_in_component_mass_prior(result, fraction=0.25): - """ Converts samples with a defined prior in chirp-mass and mass-ratio to flat in component mass by resampling with + """Converts samples with a defined prior in chirp-mass and mass-ratio to flat in component mass by resampling with the posterior with weights defined as ratio in new:old prior values times the jacobian which for F(mc, q) -> G(m1, m2) is defined as J := m1^2 / mc @@ -51,14 +65,20 @@ def convert_to_flat_in_component_mass_prior(result, fraction=0.25): """ if getattr(result, "priors") is not None: - for key in ['chirp_mass', 'mass_ratio']: + for key in ["chirp_mass", "mass_ratio"]: if key not in result.priors.keys(): - BilbyPriorConversionError("{} Prior not found in result object".format(key)) + BilbyPriorConversionError( + "{} Prior not found in result object".format(key) + ) if isinstance(result.priors[key], Constraint): - BilbyPriorConversionError("{} Prior should not be a Constraint".format(key)) - for key in ['mass_1', 'mass_2']: + BilbyPriorConversionError( + "{} Prior should not be a Constraint".format(key) + ) + for key in ["mass_1", "mass_2"]: if not isinstance(result.priors[key], Constraint): - BilbyPriorConversionError("{} Prior should be a Constraint Prior".format(key)) + BilbyPriorConversionError( + "{} Prior should be a Constraint Prior".format(key) + ) else: BilbyPriorConversionError("No prior in the result: unable to convert") @@ -67,27 +87,43 @@ def convert_to_flat_in_component_mass_prior(result, fraction=0.25): old_priors = copy.copy(result.priors) posterior = result.posterior - for key in ['chirp_mass', 'mass_ratio']: - priors[key] = Constraint(priors[key].minimum, priors[key].maximum, key, latex_label=priors[key].latex_label) - for key in ['mass_1', 'mass_2']: - priors[key] = Uniform(priors[key].minimum, priors[key].maximum, key, latex_label=priors[key].latex_label, - unit=r"$M_{\odot}$") + for key in ["chirp_mass", "mass_ratio"]: + priors[key] = Constraint( + priors[key].minimum, + priors[key].maximum, + key, + latex_label=priors[key].latex_label, + ) + for key in ["mass_1", "mass_2"]: + priors[key] = Uniform( + priors[key].minimum, + priors[key].maximum, + key, + latex_label=priors[key].latex_label, + unit=r"$M_{\odot}$", + ) - weights = np.array(result.get_weights_by_new_prior(old_priors, priors, - prior_names=['chirp_mass', 'mass_ratio', 'mass_1', 'mass_2'])) + weights = np.array( + result.get_weights_by_new_prior( + old_priors, + priors, + prior_names=["chirp_mass", "mass_ratio", "mass_1", "mass_2"], + ) + ) jacobian = posterior["mass_1"] ** 2 / posterior["chirp_mass"] weights = jacobian * weights result.posterior = posterior.sample(frac=fraction, weights=weights) logger.info("Resampling posterior to flat-in-component mass") - effective_sample_size = sum(weights)**2 / sum(weights**2) + effective_sample_size = sum(weights) ** 2 / sum(weights**2) n_posterior = len(posterior) if fraction > effective_sample_size / n_posterior: logger.warning( "Sampling posterior of length {} with fraction {}, but " "effective_sample_size / len(posterior) = {}. This may produce " - "biased results" - .format(n_posterior, fraction, effective_sample_size / n_posterior) + "biased results".format( + n_posterior, fraction, effective_sample_size / n_posterior + ) ) result.posterior = posterior.sample(frac=fraction, weights=weights, replace=True) result.meta_data["reweighted_to_flat_in_component_mass"] = True @@ -102,15 +138,27 @@ class Cosmological(Interped): @property def _default_args_dict(self): from astropy import units + return dict( - redshift=dict(name='redshift', latex_label='$z$', unit=None), + redshift=dict(name="redshift", latex_label="$z$", unit=None), luminosity_distance=dict( - name='luminosity_distance', latex_label='$d_L$', unit=units.Mpc), + name="luminosity_distance", latex_label="$d_L$", unit=units.Mpc + ), comoving_distance=dict( - name='comoving_distance', latex_label='$d_C$', unit=units.Mpc)) + name="comoving_distance", latex_label="$d_C$", unit=units.Mpc + ), + ) - def __init__(self, minimum, maximum, cosmology=None, name=None, - latex_label=None, unit=None, boundary=None): + def __init__( + self, + minimum, + maximum, + cosmology=None, + name=None, + latex_label=None, + unit=None, + boundary=None, + ): """ Parameters @@ -135,34 +183,42 @@ def __init__(self, minimum, maximum, cosmology=None, name=None, The boundary condition to apply to the prior when sampling. """ from astropy import units + self.cosmology = get_cosmology(cosmology) if name not in self._default_args_dict: raise ValueError( "Name {} not recognised. Must be one of luminosity_distance, " - "comoving_distance, redshift".format(name)) + "comoving_distance, redshift".format(name) + ) self.name = name label_args = self._default_args_dict[self.name] if latex_label is not None: - label_args['latex_label'] = latex_label + label_args["latex_label"] = latex_label if unit is not None: if not isinstance(unit, units.Unit): unit = units.Unit(unit) - label_args['unit'] = unit - self.unit = label_args['unit'] + label_args["unit"] = unit + self.unit = label_args["unit"] self._minimum = dict() self._maximum = dict() self.minimum = minimum self.maximum = maximum - if name == 'redshift': + if name == "redshift": xx, yy = self._get_redshift_arrays() - elif name == 'comoving_distance': + elif name == "comoving_distance": xx, yy = self._get_comoving_distance_arrays() - elif name == 'luminosity_distance': + elif name == "luminosity_distance": xx, yy = self._get_luminosity_distance_arrays() else: - raise ValueError('Name {} not recognized.'.format(name)) - super(Cosmological, self).__init__(xx=xx, yy=yy, minimum=minimum, maximum=maximum, - boundary=boundary, **label_args) + raise ValueError("Name {} not recognized.".format(name)) + super(Cosmological, self).__init__( + xx=xx, + yy=yy, + minimum=minimum, + maximum=maximum, + boundary=boundary, + **label_args, + ) @property def minimum(self): @@ -171,7 +227,9 @@ def minimum(self): @minimum.setter def minimum(self, minimum): if (self.name in self._minimum) and (minimum < self.minimum): - self._set_limit(value=minimum, limit_dict=self._minimum, recalculate_array=True) + self._set_limit( + value=minimum, limit_dict=self._minimum, recalculate_array=True + ) else: self._set_limit(value=minimum, limit_dict=self._minimum) @@ -182,7 +240,9 @@ def maximum(self): @maximum.setter def maximum(self, maximum): if (self.name in self._maximum) and (maximum > self.maximum): - self._set_limit(value=maximum, limit_dict=self._maximum, recalculate_array=True) + self._set_limit( + value=maximum, limit_dict=self._maximum, recalculate_array=True + ) else: self._set_limit(value=maximum, limit_dict=self._maximum) @@ -201,37 +261,37 @@ def _set_limit(self, value, limit_dict, recalculate_array=False): """ cosmology = get_cosmology(self.cosmology) limit_dict[self.name] = value - if self.name == 'redshift': - limit_dict['luminosity_distance'] = \ - cosmology.luminosity_distance(value).value - limit_dict['comoving_distance'] = \ - cosmology.comoving_distance(value).value - elif self.name == 'luminosity_distance': + if self.name == "redshift": + limit_dict["luminosity_distance"] = cosmology.luminosity_distance( + value + ).value + limit_dict["comoving_distance"] = cosmology.comoving_distance(value).value + elif self.name == "luminosity_distance": if value == 0: - limit_dict['redshift'] = 0 + limit_dict["redshift"] = 0 else: - limit_dict['redshift'] = z_at_value( + limit_dict["redshift"] = z_at_value( cosmology.luminosity_distance, value * self.unit ) - limit_dict['comoving_distance'] = ( - cosmology.comoving_distance(limit_dict['redshift']).value - ) - elif self.name == 'comoving_distance': + limit_dict["comoving_distance"] = cosmology.comoving_distance( + limit_dict["redshift"] + ).value + elif self.name == "comoving_distance": if value == 0: - limit_dict['redshift'] = 0 + limit_dict["redshift"] = 0 else: - limit_dict['redshift'] = z_at_value( + limit_dict["redshift"] = z_at_value( cosmology.comoving_distance, value * self.unit ) - limit_dict['luminosity_distance'] = ( - cosmology.luminosity_distance(limit_dict['redshift']).value - ) + limit_dict["luminosity_distance"] = cosmology.luminosity_distance( + limit_dict["redshift"] + ).value if recalculate_array: - if self.name == 'redshift': + if self.name == "redshift": self.xx, self.yy = self._get_redshift_arrays() - elif self.name == 'comoving_distance': + elif self.name == "comoving_distance": self.xx, self.yy = self._get_comoving_distance_arrays() - elif self.name == 'luminosity_distance': + elif self.name == "luminosity_distance": self.xx, self.yy = self._get_luminosity_distance_arrays() try: self._update_instance() @@ -259,13 +319,13 @@ def get_corresponding_prior(self, name=None, unit=None): args_dict = {key: getattr(self, key) for key in subclass_args} self._convert_to(new=name, args_dict=args_dict) if unit is not None: - args_dict['unit'] = unit + args_dict["unit"] = unit return self.__class__(**args_dict) def _convert_to(self, new, args_dict): args_dict.update(self._default_args_dict[new]) - args_dict['minimum'] = self._minimum[args_dict['name']] - args_dict['maximum'] = self._maximum[args_dict['name']] + args_dict["minimum"] = self._minimum[args_dict["name"]] + args_dict["maximum"] = self._maximum[args_dict["name"]] def _get_comoving_distance_arrays(self): zs, p_dz = self._get_redshift_arrays() @@ -298,12 +358,13 @@ def from_repr(cls, string): def get_instantiation_dict(self): from astropy import units from .cosmology import get_available_cosmologies + available = get_available_cosmologies() instantiation_dict = super().get_instantiation_dict() if self.cosmology.name in available: - instantiation_dict['cosmology'] = self.cosmology.name + instantiation_dict["cosmology"] = self.cosmology.name if isinstance(self.unit, units.Unit): - instantiation_dict['unit'] = self.unit.to_string() + instantiation_dict["unit"] = self.unit.to_string() return instantiation_dict @@ -320,8 +381,9 @@ class UniformComovingVolume(Cosmological): """ def _get_redshift_arrays(self): - zs = np.linspace(self._minimum['redshift'] * 0.99, - self._maximum['redshift'] * 1.01, 1000) + zs = np.linspace( + self._minimum["redshift"] * 0.99, self._maximum["redshift"] * 1.01, 1000 + ) p_dz = self.cosmology.differential_comoving_volume(zs).value return zs, p_dz @@ -339,8 +401,9 @@ class UniformSourceFrame(Cosmological): """ def _get_redshift_arrays(self): - zs = np.linspace(self._minimum['redshift'] * 0.99, - self._maximum['redshift'] * 1.01, 1000) + zs = np.linspace( + self._minimum["redshift"] * 0.99, self._maximum["redshift"] * 1.01, 1000 + ) p_dz = self.cosmology.differential_comoving_volume(zs).value / (1 + zs) return zs, p_dz @@ -364,8 +427,15 @@ class UniformInComponentsChirpMass(PowerLaw): _leaves = ["_minimum", "_maximum", "alpha"] - def __init__(self, minimum, maximum, name='chirp_mass', - latex_label=r'$\mathcal{M}$', unit=None, boundary=None): + def __init__( + self, + minimum, + maximum, + name="chirp_mass", + latex_label=r"$\mathcal{M}$", + unit=None, + boundary=None, + ): """ Parameters ========== @@ -379,8 +449,14 @@ def __init__(self, minimum, maximum, name='chirp_mass', boundary: see superclass """ super(UniformInComponentsChirpMass, self).__init__( - alpha=1., minimum=minimum, maximum=maximum, - name=name, latex_label=latex_label, unit=unit, boundary=boundary) + alpha=1.0, + minimum=minimum, + maximum=maximum, + name=name, + latex_label=latex_label, + unit=unit, + boundary=boundary, + ) class UniformInComponentsMassRatio(Prior): @@ -402,8 +478,16 @@ class UniformInComponentsMassRatio(Prior): _leaves = ["norm", "icdf"] - def __init__(self, minimum, maximum, name='mass_ratio', latex_label='$q$', - unit=None, boundary=None, equal_mass=False): + def __init__( + self, + minimum, + maximum, + name="mass_ratio", + latex_label="$q$", + unit=None, + boundary=None, + equal_mass=False, + ): """ Parameters ========== @@ -424,17 +508,26 @@ def __init__(self, minimum, maximum, name='mass_ratio', latex_label='$q$', """ self.equal_mass = equal_mass super(UniformInComponentsMassRatio, self).__init__( - minimum=minimum, maximum=maximum, name=name, - latex_label=latex_label, unit=unit, boundary=boundary) + minimum=minimum, + maximum=maximum, + name=name, + latex_label=latex_label, + unit=unit, + boundary=boundary, + ) self.norm = self._integral(maximum) - self._integral(minimum) qs = np.linspace(minimum, maximum, 1000) self.icdf = interp1d( - self.cdf(qs), qs, kind='cubic', - bounds_error=False, fill_value=(minimum, maximum)) + self.cdf(qs), + qs, + kind="cubic", + bounds_error=False, + fill_value=(minimum, maximum), + ) @staticmethod def _integral(q): - return -5. * q**(-1. / 5.) * hyp2f1(-2. / 5., -1. / 5., 4. / 5., -q) + return -5.0 * q ** (-1.0 / 5.0) * hyp2f1(-2.0 / 5.0, -1.0 / 5.0, 4.0 / 5.0, -q) def cdf(self, val, *, xp=np): return (self._integral(val) - self._integral(self.minimum)) / self.norm @@ -448,7 +541,9 @@ def rescale(self, val, *, xp=None): def prob(self, val, *, xp=None): in_prior = (val >= self.minimum) & (val <= self.maximum) with np.errstate(invalid="ignore"): - prob = (1. + val)**(2. / 5.) / (val**(6. / 5.)) / self.norm * in_prior + prob = ( + (1.0 + val) ** (2.0 / 5.0) / (val ** (6.0 / 5.0)) / self.norm * in_prior + ) return prob @xp_wrap @@ -506,13 +601,14 @@ def __init__( """ self.a_prior = a_prior self.z_prior = z_prior - chi_min = min(a_prior.maximum * z_prior.minimum, - a_prior.minimum * z_prior.maximum) + chi_min = min( + a_prior.maximum * z_prior.minimum, a_prior.minimum * z_prior.maximum + ) chi_max = a_prior.maximum * z_prior.maximum if self._is_simple_aligned_prior: self.num_interp = 100_000 if num_interp is None else num_interp xx = np.linspace(chi_min, chi_max, self.num_interp) - yy = - np.log(np.abs(xx) / a_prior.maximum) / (2 * a_prior.maximum) + yy = -np.log(np.abs(xx) / a_prior.maximum) / (2 * a_prior.maximum) else: def integrand(aa, chi): @@ -526,8 +622,7 @@ def integrand(aa, chi): self.num_interp = 10_000 if num_interp is None else num_interp xx = np.linspace(chi_min, chi_max, self.num_interp) yy = [ - quad(integrand, a_prior.minimum, a_prior.maximum, chi)[0] - for chi in xx + quad(integrand, a_prior.minimum, a_prior.maximum, chi)[0] for chi in xx ] super().__init__( xx=xx, @@ -563,16 +658,27 @@ class ConditionalChiUniformSpinMagnitude(ConditionalLogUniform): p(a | \chi) &\propto \frac{1}{a} """ - def __init__(self, minimum, maximum, name, latex_label=None, unit=None, boundary=None): + def __init__( + self, minimum, maximum, name, latex_label=None, unit=None, boundary=None + ): super(ConditionalChiUniformSpinMagnitude, self).__init__( - minimum=minimum, maximum=maximum, name=name, latex_label=latex_label, unit=unit, boundary=boundary, - condition_func=self._condition_function) + minimum=minimum, + maximum=maximum, + name=name, + latex_label=latex_label, + unit=unit, + boundary=boundary, + condition_func=self._condition_function, + ) self._required_variables = [name.replace("a", "chi")] self.__class__.__name__ = "ConditionalChiUniformSpinMagnitude" self.__class__.__qualname__ = "ConditionalChiUniformSpinMagnitude" def _condition_function(self, reference_params, **kwargs): - return dict(minimum=np.abs(kwargs[self._required_variables[0]]), maximum=reference_params["maximum"]) + return dict( + minimum=np.abs(kwargs[self._required_variables[0]]), + maximum=reference_params["maximum"], + ) def __repr__(self): return Prior.__repr__(self) @@ -598,12 +704,17 @@ class ConditionalChiInPlane(ConditionalBasePrior): N^{-1} &= 2 \ln\left(\frac{a_\max}{|\chi|}\right) """ - def __init__(self, minimum, maximum, name, latex_label=None, unit=None, boundary=None): + def __init__( + self, minimum, maximum, name, latex_label=None, unit=None, boundary=None + ): super(ConditionalChiInPlane, self).__init__( - minimum=minimum, maximum=maximum, - name=name, latex_label=latex_label, - unit=unit, boundary=boundary, - condition_func=self._condition_function + minimum=minimum, + maximum=maximum, + name=name, + latex_label=latex_label, + unit=unit, + boundary=boundary, + condition_func=self._condition_function, ) self._required_variables = [name[:5]] self._reference_maximum = maximum @@ -612,14 +723,17 @@ def __init__(self, minimum, maximum, name, latex_label=None, unit=None, boundary @xp_wrap def prob(self, val, *, xp=np, **required_variables): - parameters = self.condition_func(self.reference_params.copy(), **required_variables) + parameters = self.condition_func( + self.reference_params.copy(), **required_variables + ) chi_aligned = abs(required_variables[self._required_variables[0]]) minimum = parameters.get("minimum", self.minimum) maximum = parameters.get("maximum", self.maximum) return ( - (val >= minimum) * (val <= maximum) + (val >= minimum) + * (val <= maximum) * val - / (chi_aligned ** 2 + val ** 2) + / (chi_aligned**2 + val**2) / xp.log(self._reference_maximum / chi_aligned) ) @@ -650,11 +764,13 @@ def cdf(self, val, *, xp=np, **required_variables): self.update_conditions(**required_variables) chi_aligned = abs(required_variables[self._required_variables[0]]) return xp.clip( - (val >= self.minimum) * (val <= self.maximum) + (val >= self.minimum) + * (val <= self.maximum) * xp.log(1 + (val / chi_aligned) ** 2) - / 2 / xp.log(self._reference_maximum / chi_aligned), + / 2 + / xp.log(self._reference_maximum / chi_aligned), 0, - 1 + 1, ) def rescale(self, val, *, xp=np, **required_variables): @@ -676,13 +792,16 @@ def rescale(self, val, *, xp=np, **required_variables): """ self.update_conditions(**required_variables) chi_aligned = abs(required_variables[self._required_variables[0]]) - return chi_aligned * ((self._reference_maximum / chi_aligned) ** (2 * val) - 1) ** 0.5 + return ( + chi_aligned + * ((self._reference_maximum / chi_aligned) ** (2 * val) - 1) ** 0.5 + ) def _condition_function(self, reference_params, **kwargs): with np.errstate(invalid="ignore"): maximum = ( - self._reference_maximum ** 2 - kwargs[self._required_variables[0]] ** 2 - )**0.5 + self._reference_maximum**2 - kwargs[self._required_variables[0]] ** 2 + ) ** 0.5 return dict(minimum=0, maximum=maximum) def __repr__(self): @@ -704,7 +823,9 @@ def __init__(self, minimum=-np.inf, maximum=np.inf): unphysical samples. Necessary for EoS sampling. """ - super().__init__(minimum=minimum, maximum=maximum, name=None, latex_label=None, unit=None) + super().__init__( + minimum=minimum, maximum=maximum, name=None, latex_label=None, unit=None + ) def prob(self, val, *, xp=np): """ @@ -713,7 +834,6 @@ def prob(self, val, *, xp=np): return val def ln_prob(self, val, *, xp=np): - if val: result = 0.0 elif not val: @@ -730,9 +850,9 @@ def minimum_chirp_mass(self): if "chirp_mass" in self: return self["chirp_mass"].minimum elif "mass_1" in self: - mass_1 = self['mass_1'].minimum + mass_1 = self["mass_1"].minimum if "mass_2" in self: - mass_2 = self['mass_2'].minimum + mass_2 = self["mass_2"].minimum elif "mass_ratio" in self: mass_2 = mass_1 * self["mass_ratio"].minimum if mass_1 is not None and mass_2 is not None: @@ -749,9 +869,9 @@ def maximum_chirp_mass(self): if "chirp_mass" in self: return self["chirp_mass"].maximum elif "mass_1" in self: - mass_1 = self['mass_1'].maximum + mass_1 = self["mass_1"].maximum if "mass_2" in self: - mass_2 = self['mass_2'].maximum + mass_2 = self["mass_2"].maximum elif "mass_ratio" in self: mass_2 = mass_1 * self["mass_ratio"].maximum if mass_1 is not None and mass_2 is not None: @@ -779,16 +899,18 @@ def minimum_component_mass(self): return self["mass_2"].minimum if "chirp_mass" in self and "mass_ratio" in self: total_mass = chirp_mass_and_mass_ratio_to_total_mass( - self["chirp_mass"].minimum, self["mass_ratio"].minimum) + self["chirp_mass"].minimum, self["mass_ratio"].minimum + ) _, mass_2 = total_mass_and_mass_ratio_to_component_masses( - self["mass_ratio"].minimum, total_mass) + self["mass_ratio"].minimum, total_mass + ) return mass_2 else: logger.warning("Unable to determine minimum component mass") return None def is_nonempty_intersection(self, pset): - """ Check if keys in self exist in the parameter set + """Check if keys in self exist in the parameter set Parameters ---------- @@ -809,54 +931,52 @@ def is_nonempty_intersection(self, pset): @property def spin(self): - """ Return true if priors include any spin parameters """ + """Return true if priors include any spin parameters""" return self.is_nonempty_intersection("spin") @property def precession(self): - """ Return true if priors include any precession parameters """ + """Return true if priors include any precession parameters""" return self.is_nonempty_intersection("precession_only") @property def measured_spin(self): - """ Return true if priors include any measured_spin parameters """ + """Return true if priors include any measured_spin parameters""" return self.is_nonempty_intersection("measured_spin") @property def intrinsic(self): - """ Return true if priors include any intrinsic parameters """ + """Return true if priors include any intrinsic parameters""" return self.is_nonempty_intersection("intrinsic") @property def extrinsic(self): - """ Return true if priors include any extrinsic parameters """ + """Return true if priors include any extrinsic parameters""" return self.is_nonempty_intersection("extrinsic") @property def sky(self): - """ Return true if priors include any extrinsic parameters """ + """Return true if priors include any extrinsic parameters""" return self.is_nonempty_intersection("sky") @property def distance_inclination(self): - """ Return true if priors include any extrinsic parameters """ + """Return true if priors include any extrinsic parameters""" return self.is_nonempty_intersection("distance_inclination") @property def mass(self): - """ Return true if priors include any mass parameters """ + """Return true if priors include any mass parameters""" return self.is_nonempty_intersection("mass") @property def phase(self): - """ Return true if priors include phase parameters """ + """Return true if priors include phase parameters""" return self.is_nonempty_intersection("phase") @property def _cosmological_priors(self): - return [ - key for key, prior in self.items() if isinstance(prior, Cosmological) - ] + return [key for key, prior in self.items() if isinstance(prior, Cosmological)] @property def is_cosmological(self): @@ -901,12 +1021,14 @@ def check_valid_cosmology(self, error=True, warning=False): from astropy.cosmology import cosmology_equal cosmologies = [self[key].cosmology for key in self._cosmological_priors] - if all(cosmology_equal(cosmologies[0], c, allow_equivalent=True) for c in cosmologies[1:]): + if all( + cosmology_equal(cosmologies[0], c, allow_equivalent=True) + for c in cosmologies[1:] + ): return True message = ( - "All cosmological priors must use the same cosmology. " - f"Found: {cosmologies}" + f"All cosmological priors must use the same cosmology. Found: {cosmologies}" ) if warning: logger.warning(message) @@ -914,8 +1036,10 @@ def check_valid_cosmology(self, error=True, warning=False): if error: raise ValueError(message) - def validate_prior(self, duration, minimum_frequency, N=1000, error=True, warning=False): - """ Validate the prior is suitable for use + def validate_prior( + self, duration, minimum_frequency, N=1000, error=True, warning=False + ): + """Validate the prior is suitable for use Parameters ========== @@ -936,14 +1060,16 @@ def validate_prior(self, duration, minimum_frequency, N=1000, error=True, warnin """ samples = self.sample(N) samples = generate_all_bbh_parameters(samples) - durations = np.array([ - calculate_time_to_merger( - frequency=minimum_frequency, - mass_1=mass_1, - mass_2=mass_2, - ) - for (mass_1, mass_2) in zip(samples["mass_1"], samples["mass_2"]) - ]) + durations = np.array( + [ + calculate_time_to_merger( + frequency=minimum_frequency, + mass_1=mass_1, + mass_2=mass_2, + ) + for (mass_1, mass_2) in zip(samples["mass_1"], samples["mass_2"]) + ] + ) longest_duration = max(durations) if longest_duration < duration: return True @@ -960,9 +1086,14 @@ def validate_prior(self, duration, minimum_frequency, N=1000, error=True, warnin class BBHPriorDict(CBCPriorDict): - def __init__(self, dictionary=None, filename=None, aligned_spin=False, - conversion_function=None): - """ Initialises a Prior set for Binary Black holes + def __init__( + self, + dictionary=None, + filename=None, + aligned_spin=False, + conversion_function=None, + ): + """Initialises a Prior set for Binary Black holes Parameters ========== @@ -977,17 +1108,22 @@ def __init__(self, dictionary=None, filename=None, aligned_spin=False, """ if dictionary is None and filename is None: if aligned_spin: - fname = 'aligned_spins_bbh.prior' - logger.info('Using aligned spin prior') + fname = "aligned_spins_bbh.prior" + logger.info("Using aligned spin prior") else: - fname = 'precessing_spins_bbh.prior' + fname = "precessing_spins_bbh.prior" filename = os.path.join(DEFAULT_PRIOR_DIR, fname) - logger.info('No prior given, using default BBH priors in {}.'.format(filename)) + logger.info( + "No prior given, using default BBH priors in {}.".format(filename) + ) elif filename is not None: if not os.path.isfile(filename): filename = os.path.join(DEFAULT_PRIOR_DIR, filename) - super(BBHPriorDict, self).__init__(dictionary=dictionary, filename=filename, - conversion_function=conversion_function) + super(BBHPriorDict, self).__init__( + dictionary=dictionary, + filename=filename, + conversion_function=conversion_function, + ) def default_conversion_function(self, sample): """ @@ -1035,41 +1171,66 @@ def test_redundancy(self, key, disable_logging=False): Whether the key is redundant or not """ if key in self: - logger.debug('{} already in prior'.format(key)) + logger.debug("{} already in prior".format(key)) return True - sampling_parameters = {key for key in self if not isinstance( - self[key], (DeltaFunction, Constraint))} - - mass_parameters = {'mass_1', 'mass_2', 'chirp_mass', 'total_mass', 'mass_ratio', 'symmetric_mass_ratio'} - spin_tilt_1_parameters = {'tilt_1', 'cos_tilt_1'} - spin_tilt_2_parameters = {'tilt_2', 'cos_tilt_2'} - spin_azimuth_parameters = {'phi_1', 'phi_2', 'phi_12', 'phi_jl'} - inclination_parameters = {'theta_jn', 'cos_theta_jn'} - distance_parameters = {'luminosity_distance', 'comoving_distance', 'redshift'} - - for independent_parameters, parameter_set in \ - zip([2, 2, 1, 1, 1, 1], - [mass_parameters, spin_azimuth_parameters, - spin_tilt_1_parameters, spin_tilt_2_parameters, - inclination_parameters, distance_parameters]): + sampling_parameters = { + key + for key in self + if not isinstance(self[key], (DeltaFunction, Constraint)) + } + + mass_parameters = { + "mass_1", + "mass_2", + "chirp_mass", + "total_mass", + "mass_ratio", + "symmetric_mass_ratio", + } + spin_tilt_1_parameters = {"tilt_1", "cos_tilt_1"} + spin_tilt_2_parameters = {"tilt_2", "cos_tilt_2"} + spin_azimuth_parameters = {"phi_1", "phi_2", "phi_12", "phi_jl"} + inclination_parameters = {"theta_jn", "cos_theta_jn"} + distance_parameters = {"luminosity_distance", "comoving_distance", "redshift"} + + for independent_parameters, parameter_set in zip( + [2, 2, 1, 1, 1, 1], + [ + mass_parameters, + spin_azimuth_parameters, + spin_tilt_1_parameters, + spin_tilt_2_parameters, + inclination_parameters, + distance_parameters, + ], + ): if key in parameter_set: - if len(parameter_set.intersection( - sampling_parameters)) >= independent_parameters: + if ( + len(parameter_set.intersection(sampling_parameters)) + >= independent_parameters + ): logger.disabled = disable_logging - logger.warning('{} already in prior. ' - 'This may lead to unexpected behaviour.' - .format(parameter_set.intersection(self))) + logger.warning( + "{} already in prior. " + "This may lead to unexpected behaviour.".format( + parameter_set.intersection(self) + ) + ) logger.disabled = False return True return False class BNSPriorDict(CBCPriorDict): - - def __init__(self, dictionary=None, filename=None, aligned_spin=True, - conversion_function=None): - """ Initialises a Prior set for Binary Neutron Stars + def __init__( + self, + dictionary=None, + filename=None, + aligned_spin=True, + conversion_function=None, + ): + """Initialises a Prior set for Binary Neutron Stars Parameters ========== @@ -1083,17 +1244,22 @@ def __init__(self, dictionary=None, filename=None, aligned_spin=True, BNSPriorDict.default_conversion_function """ if aligned_spin: - default_file = 'aligned_spins_bns_tides_on.prior' + default_file = "aligned_spins_bns_tides_on.prior" else: - default_file = 'precessing_spins_bns_tides_on.prior' + default_file = "precessing_spins_bns_tides_on.prior" if dictionary is None and filename is None: filename = os.path.join(DEFAULT_PRIOR_DIR, default_file) - logger.info('No prior given, using default BNS priors in {}.'.format(filename)) + logger.info( + "No prior given, using default BNS priors in {}.".format(filename) + ) elif filename is not None: if not os.path.isfile(filename): filename = os.path.join(DEFAULT_PRIOR_DIR, filename) - super(BNSPriorDict, self).__init__(dictionary=dictionary, filename=filename, - conversion_function=conversion_function) + super(BNSPriorDict, self).__init__( + dictionary=dictionary, + filename=filename, + conversion_function=conversion_function, + ) def default_conversion_function(self, sample): """ @@ -1128,7 +1294,9 @@ def default_conversion_function(self, sample): def test_redundancy(self, key, disable_logging=False): logger.disabled = disable_logging - logger.info("Performing redundancy check using BBHPriorDict(self).test_redundancy") + logger.info( + "Performing redundancy check using BBHPriorDict(self).test_redundancy" + ) bbh_redundancy = BBHPriorDict(self).test_redundancy(key) logger.disabled = False @@ -1136,19 +1304,29 @@ def test_redundancy(self, key, disable_logging=False): return True redundant = False - sampling_parameters = {key for key in self if not isinstance( - self[key], (DeltaFunction, Constraint))} + sampling_parameters = { + key + for key in self + if not isinstance(self[key], (DeltaFunction, Constraint)) + } - tidal_parameters = \ - {'lambda_1', 'lambda_2', 'lambda_tilde', 'delta_lambda_tilde'} + tidal_parameters = { + "lambda_1", + "lambda_2", + "lambda_tilde", + "delta_lambda_tilde", + } if key in tidal_parameters: if len(tidal_parameters.intersection(sampling_parameters)) > 2: redundant = True logger.disabled = disable_logging - logger.warning('{} already in prior. ' - 'This may lead to unexpected behaviour.' - .format(tidal_parameters.intersection(self))) + logger.warning( + "{} already in prior. " + "This may lead to unexpected behaviour.".format( + tidal_parameters.intersection(self) + ) + ) logger.disabled = False elif len(tidal_parameters.intersection(sampling_parameters)) == 2: redundant = True @@ -1156,44 +1334,44 @@ def test_redundancy(self, key, disable_logging=False): @property def tidal(self): - """ Return true if priors include phase parameters """ + """Return true if priors include phase parameters""" return self.is_nonempty_intersection("tidal") Prior._default_latex_labels = { - 'mass_1': '$m_1$', - 'mass_2': '$m_2$', - 'total_mass': '$M$', - 'chirp_mass': r'$\mathcal{M}$', - 'mass_ratio': '$q$', - 'symmetric_mass_ratio': r'$\eta$', - 'a_1': '$a_1$', - 'a_2': '$a_2$', - 'tilt_1': r'$\theta_1$', - 'tilt_2': r'$\theta_2$', - 'cos_tilt_1': r'$\cos\theta_1$', - 'cos_tilt_2': r'$\cos\theta_2$', - 'phi_12': r'$\Delta\phi$', - 'phi_jl': r'$\phi_{JL}$', - 'luminosity_distance': '$d_L$', - 'dec': r'$\mathrm{DEC}$', - 'ra': r'$\mathrm{RA}$', - 'iota': r'$\iota$', - 'cos_iota': r'$\cos\iota$', - 'theta_jn': r'$\theta_{JN}$', - 'cos_theta_jn': r'$\cos\theta_{JN}$', - 'psi': r'$\psi$', - 'phase': r'$\phi$', - 'geocent_time': '$t_c$', - 'time_jitter': '$t_j$', - 'lambda_1': r'$\Lambda_1$', - 'lambda_2': r'$\Lambda_2$', - 'lambda_tilde': r'$\tilde{\Lambda}$', - 'delta_lambda_tilde': r'$\delta\tilde{\Lambda}$', - 'chi_1': r'$\chi_1$', - 'chi_2': r'$\chi_2$', - 'chi_1_in_plane': r'$\chi_{1, \perp}$', - 'chi_2_in_plane': r'$\chi_{2, \perp}$', + "mass_1": "$m_1$", + "mass_2": "$m_2$", + "total_mass": "$M$", + "chirp_mass": r"$\mathcal{M}$", + "mass_ratio": "$q$", + "symmetric_mass_ratio": r"$\eta$", + "a_1": "$a_1$", + "a_2": "$a_2$", + "tilt_1": r"$\theta_1$", + "tilt_2": r"$\theta_2$", + "cos_tilt_1": r"$\cos\theta_1$", + "cos_tilt_2": r"$\cos\theta_2$", + "phi_12": r"$\Delta\phi$", + "phi_jl": r"$\phi_{JL}$", + "luminosity_distance": "$d_L$", + "dec": r"$\mathrm{DEC}$", + "ra": r"$\mathrm{RA}$", + "iota": r"$\iota$", + "cos_iota": r"$\cos\iota$", + "theta_jn": r"$\theta_{JN}$", + "cos_theta_jn": r"$\cos\theta_{JN}$", + "psi": r"$\psi$", + "phase": r"$\phi$", + "geocent_time": "$t_c$", + "time_jitter": "$t_j$", + "lambda_1": r"$\Lambda_1$", + "lambda_2": r"$\Lambda_2$", + "lambda_tilde": r"$\tilde{\Lambda}$", + "delta_lambda_tilde": r"$\delta\tilde{\Lambda}$", + "chi_1": r"$\chi_1$", + "chi_2": r"$\chi_2$", + "chi_1_in_plane": r"$\chi_{1, \perp}$", + "chi_2_in_plane": r"$\chi_{2, \perp}$", } @@ -1217,7 +1395,9 @@ def __init__(self, dictionary=None, filename=None): """ if dictionary is None and filename is not None: filename = os.path.join(DEFAULT_PRIOR_DIR, filename) - super(CalibrationPriorDict, self).__init__(dictionary=dictionary, filename=filename) + super(CalibrationPriorDict, self).__init__( + dictionary=dictionary, filename=filename + ) self.source = None def to_file(self, outdir, label): @@ -1239,9 +1419,15 @@ def to_file(self, outdir, label): outfile.write("# prior source file is {}".format(self.source)) @staticmethod - def from_envelope_file(envelope_file, minimum_frequency, - maximum_frequency, n_nodes, label, - boundary="reflective", correction_type=None): + def from_envelope_file( + envelope_file, + minimum_frequency, + maximum_frequency, + n_nodes, + label, + boundary="reflective", + correction_type=None, + ): """ Load in the calibration envelope. @@ -1290,7 +1476,10 @@ def from_envelope_file(envelope_file, minimum_frequency, This includes the frequencies of the nodes which are _not_ sampled. """ from .detector.calibration import _check_calibration_correction_type - correction_type = _check_calibration_correction_type(correction_type=correction_type) + + correction_type = _check_calibration_correction_type( + correction_type=correction_type + ) calibration_data = np.genfromtxt(envelope_file).T log_frequency_array = np.log(calibration_data[0]) @@ -1306,45 +1495,63 @@ def from_envelope_file(envelope_file, minimum_frequency, amplitude_sigma = abs(calibration_data[5] - calibration_data[3]) / 2 phase_sigma = abs(calibration_data[6] - calibration_data[4]) / 2 - log_nodes = np.linspace(np.log(minimum_frequency), - np.log(maximum_frequency), n_nodes) + log_nodes = np.linspace( + np.log(minimum_frequency), np.log(maximum_frequency), n_nodes + ) - amplitude_mean_nodes = \ - InterpolatedUnivariateSpline(log_frequency_array, amplitude_median)(log_nodes) - amplitude_sigma_nodes = \ - InterpolatedUnivariateSpline(log_frequency_array, amplitude_sigma)(log_nodes) - phase_mean_nodes = \ - InterpolatedUnivariateSpline(log_frequency_array, phase_median)(log_nodes) - phase_sigma_nodes = \ - InterpolatedUnivariateSpline(log_frequency_array, phase_sigma)(log_nodes) + amplitude_mean_nodes = InterpolatedUnivariateSpline( + log_frequency_array, amplitude_median + )(log_nodes) + amplitude_sigma_nodes = InterpolatedUnivariateSpline( + log_frequency_array, amplitude_sigma + )(log_nodes) + phase_mean_nodes = InterpolatedUnivariateSpline( + log_frequency_array, phase_median + )(log_nodes) + phase_sigma_nodes = InterpolatedUnivariateSpline( + log_frequency_array, phase_sigma + )(log_nodes) prior = CalibrationPriorDict() for ii in range(n_nodes): name = "recalib_{}_amplitude_{}".format(label, ii) latex_label = "$A^{}_{}$".format(label, ii) - prior[name] = Gaussian(mu=amplitude_mean_nodes[ii], - sigma=amplitude_sigma_nodes[ii], - name=name, latex_label=latex_label, - boundary=boundary) + prior[name] = Gaussian( + mu=amplitude_mean_nodes[ii], + sigma=amplitude_sigma_nodes[ii], + name=name, + latex_label=latex_label, + boundary=boundary, + ) for ii in range(n_nodes): name = "recalib_{}_phase_{}".format(label, ii) latex_label = r"$\phi^{}_{}$".format(label, ii) - prior[name] = Gaussian(mu=phase_mean_nodes[ii], - sigma=phase_sigma_nodes[ii], - name=name, latex_label=latex_label, - boundary=boundary) + prior[name] = Gaussian( + mu=phase_mean_nodes[ii], + sigma=phase_sigma_nodes[ii], + name=name, + latex_label=latex_label, + boundary=boundary, + ) for ii in range(n_nodes): name = "recalib_{}_frequency_{}".format(label, ii) latex_label = "$f^{}_{}$".format(label, ii) - prior[name] = DeltaFunction(peak=np.exp(log_nodes[ii]), name=name, - latex_label=latex_label) + prior[name] = DeltaFunction( + peak=np.exp(log_nodes[ii]), name=name, latex_label=latex_label + ) prior.source = os.path.abspath(envelope_file) return prior @staticmethod def constant_uncertainty_spline( - amplitude_sigma, phase_sigma, minimum_frequency, maximum_frequency, - n_nodes, label, boundary="reflective"): + amplitude_sigma, + phase_sigma, + minimum_frequency, + maximum_frequency, + n_nodes, + label, + boundary="reflective", + ): """ Make prior assuming constant in frequency calibration uncertainty. @@ -1373,8 +1580,9 @@ def constant_uncertainty_spline( Priors for the relevant parameters. This includes the frequencies of the nodes which are _not_ sampled. """ - nodes = np.logspace(np.log10(minimum_frequency), - np.log10(maximum_frequency), n_nodes) + nodes = np.logspace( + np.log10(minimum_frequency), np.log10(maximum_frequency), n_nodes + ) amplitude_mean_nodes = [0] * n_nodes amplitude_sigma_nodes = [amplitude_sigma] * n_nodes @@ -1385,22 +1593,29 @@ def constant_uncertainty_spline( for ii in range(n_nodes): name = "recalib_{}_amplitude_{}".format(label, ii) latex_label = "$A^{}_{}$".format(label, ii) - prior[name] = Gaussian(mu=amplitude_mean_nodes[ii], - sigma=amplitude_sigma_nodes[ii], - name=name, latex_label=latex_label, - boundary=boundary) + prior[name] = Gaussian( + mu=amplitude_mean_nodes[ii], + sigma=amplitude_sigma_nodes[ii], + name=name, + latex_label=latex_label, + boundary=boundary, + ) for ii in range(n_nodes): name = "recalib_{}_phase_{}".format(label, ii) latex_label = r"$\phi^{}_{}$".format(label, ii) - prior[name] = Gaussian(mu=phase_mean_nodes[ii], - sigma=phase_sigma_nodes[ii], - name=name, latex_label=latex_label, - boundary=boundary) + prior[name] = Gaussian( + mu=phase_mean_nodes[ii], + sigma=phase_sigma_nodes[ii], + name=name, + latex_label=latex_label, + boundary=boundary, + ) for ii in range(n_nodes): name = "recalib_{}_frequency_{}".format(label, ii) latex_label = "$f^{}_{}$".format(label, ii) - prior[name] = DeltaFunction(peak=nodes[ii], name=name, - latex_label=latex_label) + prior[name] = DeltaFunction( + peak=nodes[ii], name=name, latex_label=latex_label + ) return prior @@ -1435,7 +1650,7 @@ def secondary_mass_condition_function(reference_params, mass_1): dict: Updated prior limits given the provided primary mass. """ - return dict(minimum=reference_params['minimum'], maximum=mass_1) + return dict(minimum=reference_params["minimum"], maximum=mass_1) ConditionalCosmological = conditional_prior_factory(Cosmological) @@ -1467,6 +1682,7 @@ class HealPixMapPriorDist(BaseJointPriorDist): PriorDist : `bilby.gw.prior.HealPixMapPriorDist` A JointPriorDist object to store the joint prior distribution according to passed healpix map """ + def __init__(self, hp_file, names=None, bounds=None, distance=False): self.hp = self._check_imports() self.hp_file = hp_file @@ -1500,7 +1716,9 @@ def __init__(self, hp_file, names=None, bounds=None, distance=False): self.pixel_area = self.hp.nside2pixarea(self.nside) self.pixel_length = self.pixel_area ** (1 / 2.0) self.pix_xx = np.arange(self.npix) - self._all_interped = interp1d(x=self.pix_xx, y=self.prob, bounds_error=False, fill_value=0) + self._all_interped = interp1d( + x=self.pix_xx, y=self.prob, bounds_error=False, fill_value=0 + ) self.inverse_cdf = None self.distance_pdf = None self.distance_icdf = None @@ -1529,7 +1747,9 @@ def _check_imports(): try: import healpy except Exception: - raise ImportError("Must have healpy installed on this machine to use HealPixMapPrior") + raise ImportError( + "Must have healpy installed on this machine to use HealPixMapPrior" + ) return healpy @xp_wrap @@ -1561,7 +1781,9 @@ def _rescale(self, samp, *, xp=None, **kwargs): for i, val in enumerate(pix_rescale): theta, ra = self.hp.pix2ang(self.nside, int(round(val))) dec = 0.5 * np.pi - theta - sample = xpx.at(sample, i).set(xp.asarray(self.draw_from_pixel(ra, dec, int(round(val))))) + sample = xpx.at(sample, i).set( + xp.asarray(self.draw_from_pixel(ra, dec, int(round(val)))) + ) if self.distance: self.update_distance(int(round(val))) dist_samples = xpx.at(dist_samples, i).set( @@ -1590,7 +1812,7 @@ def update_distance(self, pix_idx): self.distance_pdf = lambda r: self.distnorm[pix_idx] * norm( loc=self.distmu[pix_idx], scale=self.distsigma[pix_idx] ).pdf(r) - pdfs = self.rs ** 2 * self.distance_pdf(self.rs) + pdfs = self.rs**2 * self.distance_pdf(self.rs) cdfs = np.cumsum(pdfs) / np.sum(pdfs) self.distance_icdf = interp1d(cdfs, self.rs) @@ -1644,7 +1866,9 @@ def _sample(self, size, *, random_state=None, **kwargs): if self.distance: self.update_distance(sample_pix[samp]) dist = self.draw_distance(sample_pix[samp], random_state=rng) - ra, dec = self.draw_from_pixel(ra, dec, sample_pix[samp], random_state=rng) + ra, dec = self.draw_from_pixel( + ra, dec, sample_pix[samp], random_state=rng + ) new = [ra, dec, dist] else: new = self.draw_from_pixel(ra, dec, sample_pix[samp]) @@ -1765,7 +1989,8 @@ def _ln_prob(self, samp, lnprob, outbounds, *, xp=None): if self.distance: self.update_distance(pixel) lnprob = xpx.at(lnprob, i).set( - lnprob[i] + xp.log(xp.asarray(self.distance_pdf(dist) * dist ** 2)) + lnprob[i] + + xp.log(xp.asarray(self.distance_pdf(dist) * dist**2)) ) lnprob = xp.where(xp.asarray(outbounds), -np.inf, lnprob) return lnprob @@ -1809,6 +2034,7 @@ class HealPixPrior(JointPrior): See :code:`bilby.gw.prior.HealPixMapPriorDist` for more details of how to instantiate the prior. """ + def __init__(self, dist, name=None, latex_label=None, unit=None): """ @@ -1826,5 +2052,9 @@ def __init__(self, dist, name=None, latex_label=None, unit=None): The unit of the parameter. """ if not isinstance(dist, HealPixMapPriorDist): - raise JointPriorDistError("dist object must be instance of HealPixMapPriorDist") - super(HealPixPrior, self).__init__(dist=dist, name=name, latex_label=latex_label, unit=unit) + raise JointPriorDistError( + "dist object must be instance of HealPixMapPriorDist" + ) + super(HealPixPrior, self).__init__( + dist=dist, name=name, latex_label=latex_label, unit=unit + ) diff --git a/bilby/gw/result.py b/bilby/gw/result.py index db056966f..5a652d392 100644 --- a/bilby/gw/result.py +++ b/bilby/gw/result.py @@ -6,8 +6,12 @@ from ..core.result import Result as CoreResult from ..core.utils import ( - infft, logger, check_directory_exists_and_if_not_mkdir, - latex_plot_format, safe_file_dump, safe_save_figure, + infft, + logger, + check_directory_exists_and_if_not_mkdir, + latex_plot_format, + safe_file_dump, + safe_save_figure, ) from .utils import plot_spline_pos, spline_angle_xform, asd_from_freq_series from .detector import get_empty_interferometer, Interferometer @@ -18,8 +22,8 @@ class CompactBinaryCoalescenceResult(CoreResult): Result class with additional methods and attributes specific to analyses of compact binaries. """ - def __init__(self, **kwargs): + def __init__(self, **kwargs): if "meta_data" not in kwargs: kwargs["meta_data"] = dict() if "global_meta_data" not in kwargs: @@ -29,6 +33,7 @@ def __init__(self, **kwargs): # Ensure cosmology is always stored in the meta_data if "cosmology" not in kwargs["meta_data"]["global_meta_data"]: from .cosmology import get_cosmology + kwargs["meta_data"]["global_meta_data"]["cosmology"] = get_cosmology() super(CompactBinaryCoalescenceResult, self).__init__(**kwargs) @@ -42,98 +47,101 @@ def __get_from_nested_meta_data(self, *keys): dictionary = item return item except KeyError: - raise AttributeError( - "No information stored for {}".format('/'.join(keys))) + raise AttributeError("No information stored for {}".format("/".join(keys))) @property def sampling_frequency(self): - """ Sampling frequency in Hertz""" - return self.__get_from_nested_meta_data( - 'likelihood', 'sampling_frequency') + """Sampling frequency in Hertz""" + return self.__get_from_nested_meta_data("likelihood", "sampling_frequency") @property def duration(self): - """ Duration in seconds """ - return self.__get_from_nested_meta_data( - 'likelihood', 'duration') + """Duration in seconds""" + return self.__get_from_nested_meta_data("likelihood", "duration") @property def start_time(self): - """ Start time in seconds """ - return self.__get_from_nested_meta_data( - 'likelihood', 'start_time') + """Start time in seconds""" + return self.__get_from_nested_meta_data("likelihood", "start_time") @property def time_marginalization(self): - """ Boolean for if the likelihood used time marginalization """ - return self.__get_from_nested_meta_data( - 'likelihood', 'time_marginalization') + """Boolean for if the likelihood used time marginalization""" + return self.__get_from_nested_meta_data("likelihood", "time_marginalization") @property def phase_marginalization(self): - """ Boolean for if the likelihood used phase marginalization """ - return self.__get_from_nested_meta_data( - 'likelihood', 'phase_marginalization') + """Boolean for if the likelihood used phase marginalization""" + return self.__get_from_nested_meta_data("likelihood", "phase_marginalization") @property def distance_marginalization(self): - """ Boolean for if the likelihood used distance marginalization """ + """Boolean for if the likelihood used distance marginalization""" return self.__get_from_nested_meta_data( - 'likelihood', 'distance_marginalization') + "likelihood", "distance_marginalization" + ) @property def interferometers(self): - """ List of interferometer names """ - return [name for name in self.__get_from_nested_meta_data( - 'likelihood', 'interferometers')] + """List of interferometer names""" + return [ + name + for name in self.__get_from_nested_meta_data( + "likelihood", "interferometers" + ) + ] @property def waveform_approximant(self): - """ String of the waveform approximant """ + """String of the waveform approximant""" return self.__get_from_nested_meta_data( - 'likelihood', 'waveform_arguments', 'waveform_approximant') + "likelihood", "waveform_arguments", "waveform_approximant" + ) @property def waveform_generator_class(self): - """ Waveform generator class """ + """Waveform generator class""" return self.__get_from_nested_meta_data( - 'likelihood', 'waveform_generator_class') + "likelihood", "waveform_generator_class" + ) @property def waveform_arguments(self): - """ Dict of waveform arguments """ - return self.__get_from_nested_meta_data( - 'likelihood', 'waveform_arguments') + """Dict of waveform arguments""" + return self.__get_from_nested_meta_data("likelihood", "waveform_arguments") @property def waveform_generator_meta_data(self): - """ Dict of metadata for reconstructing the waveform generator. """ + """Dict of metadata for reconstructing the waveform generator.""" return self.__get_from_nested_meta_data( - 'likelihood', 'waveform_generator_meta_data') + "likelihood", "waveform_generator_meta_data" + ) @property def reference_frequency(self): - """ Float of the reference frequency """ + """Float of the reference frequency""" return self.__get_from_nested_meta_data( - 'likelihood', 'waveform_arguments', 'reference_frequency') + "likelihood", "waveform_arguments", "reference_frequency" + ) @property def frequency_domain_source_model(self): - """ The frequency domain source model (function)""" + """The frequency domain source model (function)""" return self.__get_from_nested_meta_data( - 'likelihood', 'frequency_domain_source_model') + "likelihood", "frequency_domain_source_model" + ) @property def time_domain_source_model(self): - """ The time domain source model (function)""" + """The time domain source model (function)""" return self.__get_from_nested_meta_data( - 'likelihood', 'time_domain_source_model') + "likelihood", "time_domain_source_model" + ) @property def parameter_conversion(self): - """ The frequency domain source model (function)""" - return self.__get_from_nested_meta_data( - 'likelihood', 'parameter_conversion') + """The frequency domain source model (function)""" + return self.__get_from_nested_meta_data("likelihood", "parameter_conversion") @property def cosmology(self): @@ -144,9 +152,7 @@ def cosmology(self): .. versionadded:: 2.5.0 """ try: - return self.__get_from_nested_meta_data( - 'global_meta_data', 'cosmology' - ) + return self.__get_from_nested_meta_data("global_meta_data", "cosmology") except AttributeError as e: logger.warning( "No cosmology found in result. " @@ -156,7 +162,7 @@ def cosmology(self): return None def detector_injection_properties(self, detector): - """ Returns a dictionary of the injection properties for each detector + """Returns a dictionary of the injection properties for each detector The injection properties include the parameters injected, and information about the signal to noise ratio (SNR) given the noise @@ -175,14 +181,15 @@ def detector_injection_properties(self, detector): """ try: return self.__get_from_nested_meta_data( - 'likelihood', 'interferometers', detector) + "likelihood", "interferometers", detector + ) except AttributeError: logger.info("No injection for detector {}".format(detector)) return None @latex_plot_format - def plot_calibration_posterior(self, level=.9, format="png"): - """ Plots the calibration amplitude and phase uncertainty. + def plot_calibration_posterior(self, level=0.9, format="png"): + """Plots the calibration amplitude and phase uncertainty. Adapted from the LALInference version in bayespputils Plot is saved to {self.outdir}/{self.label}_calibration.{format} @@ -195,6 +202,7 @@ def plot_calibration_posterior(self, level=.9, format="png"): Format to save the plot, default=png, options are png/pdf """ import matplotlib.pyplot as plt + if format not in ["png", "pdf"]: raise ValueError("Format should be one of png or pdf") @@ -205,68 +213,102 @@ def plot_calibration_posterior(self, level=.9, format="png"): outdir = self.outdir parameters = posterior.keys() - ifos = np.unique([param.split('_')[1] for param in parameters if 'recalib_' in param]) + ifos = np.unique( + [param.split("_")[1] for param in parameters if "recalib_" in param] + ) if ifos.size == 0: logger.info("No calibration parameters found. Aborting calibration plot.") return for ifo in ifos: - if ifo == 'H1': - color = 'r' - elif ifo == 'L1': - color = 'g' - elif ifo == 'V1': - color = 'm' + if ifo == "H1": + color = "r" + elif ifo == "L1": + color = "g" + elif ifo == "V1": + color = "m" else: - color = 'c' + color = "c" # Assume spline control frequencies are constant - freq_params = np.sort([param for param in parameters if - 'recalib_{0}_frequency_'.format(ifo) in param]) + freq_params = np.sort( + [ + param + for param in parameters + if "recalib_{0}_frequency_".format(ifo) in param + ] + ) logfreqs = np.log([posterior[param].iloc[0] for param in freq_params]) # Amplitude calibration model plt.sca(ax1) - amp_params = np.sort([param for param in parameters if - 'recalib_{0}_amplitude_'.format(ifo) in param]) + amp_params = np.sort( + [ + param + for param in parameters + if "recalib_{0}_amplitude_".format(ifo) in param + ] + ) if len(amp_params) > 0: - amplitude = 100 * np.column_stack([posterior[param] for param in amp_params]) - plot_spline_pos(logfreqs, amplitude, color=color, level=level, - label=r"{0} (mean, {1}$\%$)".format(ifo.upper(), int(level * 100))) + amplitude = 100 * np.column_stack( + [posterior[param] for param in amp_params] + ) + plot_spline_pos( + logfreqs, + amplitude, + color=color, + level=level, + label=r"{0} (mean, {1}$\%$)".format(ifo.upper(), int(level * 100)), + ) # Phase calibration model plt.sca(ax2) - phase_params = np.sort([param for param in parameters if - 'recalib_{0}_phase_'.format(ifo) in param]) + phase_params = np.sort( + [ + param + for param in parameters + if "recalib_{0}_phase_".format(ifo) in param + ] + ) if len(phase_params) > 0: phase = np.column_stack([posterior[param] for param in phase_params]) - plot_spline_pos(logfreqs, phase, color=color, level=level, - label=r"{0} (mean, {1}$\%$)".format(ifo.upper(), int(level * 100)), - xform=spline_angle_xform) + plot_spline_pos( + logfreqs, + phase, + color=color, + level=level, + label=r"{0} (mean, {1}$\%$)".format(ifo.upper(), int(level * 100)), + xform=spline_angle_xform, + ) - ax1.tick_params(labelsize=.75 * font_size) - ax2.tick_params(labelsize=.75 * font_size) - plt.legend(loc='upper right', prop={'size': .75 * font_size}, framealpha=0.1) - ax1.set_xscale('log') - ax2.set_xscale('log') + ax1.tick_params(labelsize=0.75 * font_size) + ax2.tick_params(labelsize=0.75 * font_size) + plt.legend(loc="upper right", prop={"size": 0.75 * font_size}, framealpha=0.1) + ax1.set_xscale("log") + ax2.set_xscale("log") - ax2.set_xlabel('Frequency [Hz]', fontsize=font_size) - ax1.set_ylabel(r'Amplitude [$\%$]', fontsize=font_size) - ax2.set_ylabel('Phase [deg]', fontsize=font_size) + ax2.set_xlabel("Frequency [Hz]", fontsize=font_size) + ax1.set_ylabel(r"Amplitude [$\%$]", fontsize=font_size) + ax2.set_ylabel("Phase [deg]", fontsize=font_size) - filename = os.path.join(outdir, self.label + '_calibration.' + format) + filename = os.path.join(outdir, self.label + "_calibration." + format) fig.tight_layout() safe_save_figure( - fig=fig, filename=filename, - format=format, dpi=600, bbox_inches='tight' + fig=fig, filename=filename, format=format, dpi=600, bbox_inches="tight" ) logger.debug("Calibration figure saved to {}".format(filename)) plt.close() def plot_waveform_posterior( - self, interferometers=None, level=0.9, n_samples=None, - format='png', start_time=None, end_time=None): + self, + interferometers=None, + level=0.9, + n_samples=None, + format="png", + start_time=None, + end_time=None, + ): """ Plot the posterior for the waveform in the frequency domain and whitened time domain for all detectors. @@ -297,18 +339,29 @@ def plot_waveform_posterior( if interferometers is None: interferometers = self.interferometers elif not isinstance(interferometers, list): - raise TypeError( - 'interferometers must be a list or InterferometerList') + raise TypeError("interferometers must be a list or InterferometerList") for ifo in interferometers: self.plot_interferometer_waveform_posterior( - interferometer=ifo, level=level, n_samples=n_samples, - save=True, format=format, start_time=start_time, - end_time=end_time) + interferometer=ifo, + level=level, + n_samples=n_samples, + save=True, + format=format, + start_time=start_time, + end_time=end_time, + ) @latex_plot_format def plot_interferometer_waveform_posterior( - self, interferometer, level=0.9, n_samples=None, save=True, - format='png', start_time=None, end_time=None): + self, + interferometer, + level=0.9, + n_samples=None, + save=True, + format="png", + start_time=None, + end_time=None, + ): """ Plot the posterior for the waveform in the frequency domain and whitened time domain. @@ -364,56 +417,52 @@ def plot_interferometer_waveform_posterior( except ImportError: logger.warning( "HTML plotting requested, but plotly cannot be imported, " - "falling back to png format for waveform plot.") + "falling back to png format for waveform plot." + ) format = "png" if isinstance(interferometer, str): interferometer = get_empty_interferometer(interferometer) interferometer.set_strain_data_from_zero_noise( sampling_frequency=self.sampling_frequency, - duration=self.duration, start_time=self.start_time) + duration=self.duration, + start_time=self.start_time, + ) PLOT_DATA = False elif not isinstance(interferometer, Interferometer): - raise TypeError( - 'interferometer must be either str or Interferometer') + raise TypeError("interferometer must be either str or Interferometer") else: PLOT_DATA = True - logger.info("Generating waveform figure for {}".format( - interferometer.name)) + logger.info("Generating waveform figure for {}".format(interferometer.name)) if n_samples is None: samples = self.posterior elif n_samples > len(self.posterior): logger.debug( "Requested more waveform samples ({}) than we have " - "posterior samples ({})!".format( - n_samples, len(self.posterior) - ) + "posterior samples ({})!".format(n_samples, len(self.posterior)) ) samples = self.posterior else: samples = self.posterior.sample(n_samples, replace=False) if start_time is None: - start_time = - 0.4 + start_time = -0.4 start_time = np.mean(samples.geocent_time) + start_time if end_time is None: end_time = 0.2 end_time = np.mean(samples.geocent_time) + end_time if format == "html": - start_time = - np.inf + start_time = -np.inf end_time = np.inf - time_idxs = ( - (interferometer.time_array >= start_time) & - (interferometer.time_array <= end_time) + time_idxs = (interferometer.time_array >= start_time) & ( + interferometer.time_array <= end_time ) frequency_idxs = np.where(interferometer.frequency_mask)[0] - logger.debug("Frequency mask contains {} values".format( - len(frequency_idxs)) - ) - frequency_idxs = frequency_idxs[::max(1, len(frequency_idxs) // 4000)] - logger.debug("Downsampling frequency mask to {} values".format( - len(frequency_idxs)) + logger.debug("Frequency mask contains {} values".format(len(frequency_idxs))) + frequency_idxs = frequency_idxs[:: max(1, len(frequency_idxs) // 4000)] + logger.debug( + "Downsampling frequency mask to {} values".format(len(frequency_idxs)) ) plot_times = interferometer.time_array[time_idxs] plot_times -= interferometer.strain_data.start_time @@ -422,34 +471,36 @@ def plot_interferometer_waveform_posterior( plot_frequencies = interferometer.frequency_array[frequency_idxs] waveform_generator = self.waveform_generator_class( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, start_time=self.start_time, frequency_domain_source_model=self.frequency_domain_source_model, time_domain_source_model=self.time_domain_source_model, parameter_conversion=self.parameter_conversion, waveform_arguments=self.waveform_arguments, - **self.waveform_generator_meta_data) + **self.waveform_generator_meta_data, + ) if format == "html": fig = make_subplots( - rows=2, cols=1, + rows=2, + cols=1, row_heights=[0.5, 0.5], ) fig.update_layout( - template='plotly_white', + template="plotly_white", font=dict( family="Computer Modern", - ) + ), ) else: import matplotlib.pyplot as plt from matplotlib import rcParams + old_font_size = rcParams["font.size"] rcParams["font.size"] = 20 fig, axs = plt.subplots( - 2, 1, - gridspec_kw=dict(height_ratios=[1.5, 1]), - figsize=(16, 12.5) + 2, 1, gridspec_kw=dict(height_ratios=[1.5, 1]), figsize=(16, 12.5) ) if PLOT_DATA: @@ -459,13 +510,14 @@ def plot_interferometer_waveform_posterior( x=plot_frequencies, y=asd_from_freq_series( interferometer.frequency_domain_strain[frequency_idxs], - 1 / interferometer.strain_data.duration + 1 / interferometer.strain_data.duration, ), fill=None, - mode='lines', line_color=DATA_COLOR, + mode="lines", + line_color=DATA_COLOR, opacity=0.5, name="Data", - legendgroup='data', + legendgroup="data", ), row=1, col=1, @@ -473,12 +525,15 @@ def plot_interferometer_waveform_posterior( fig.add_trace( go.Scatter( x=plot_frequencies, - y=interferometer.amplitude_spectral_density_array[frequency_idxs], + y=interferometer.amplitude_spectral_density_array[ + frequency_idxs + ], fill=None, - mode='lines', line_color=DATA_COLOR, + mode="lines", + line_color=DATA_COLOR, opacity=0.8, name="ASD", - legendgroup='asd', + legendgroup="asd", ), row=1, col=1, @@ -488,10 +543,11 @@ def plot_interferometer_waveform_posterior( x=plot_times, y=interferometer.whitened_time_domain_strain[time_idxs], fill=None, - mode='lines', line_color=DATA_COLOR, + mode="lines", + line_color=DATA_COLOR, opacity=0.5, name="Data", - legendgroup='data', + legendgroup="data", showlegend=False, ), row=2, @@ -502,16 +558,25 @@ def plot_interferometer_waveform_posterior( plot_frequencies, asd_from_freq_series( interferometer.frequency_domain_strain[frequency_idxs], - 1 / interferometer.strain_data.duration), - color=DATA_COLOR, label='Data', alpha=0.3) + 1 / interferometer.strain_data.duration, + ), + color=DATA_COLOR, + label="Data", + alpha=0.3, + ) axs[0].loglog( plot_frequencies, interferometer.amplitude_spectral_density_array[frequency_idxs], - color=DATA_COLOR, label='ASD') + color=DATA_COLOR, + label="ASD", + ) axs[1].plot( - plot_times, interferometer.whitened_time_domain_strain[time_idxs], - color=DATA_COLOR, alpha=0.3) - logger.debug('Plotted interferometer data.') + plot_times, + interferometer.whitened_time_domain_strain[time_idxs], + color=DATA_COLOR, + alpha=0.3, + ) + logger.debug("Plotted interferometer data.") fd_waveforms = list() td_waveforms = list() @@ -520,20 +585,22 @@ def plot_interferometer_waveform_posterior( fd_waveform = interferometer.get_detector_response(wf_pols, params) fd_waveforms.append(fd_waveform[frequency_idxs]) whitened_fd_waveform = interferometer.whiten_frequency_series(fd_waveform) - td_waveform = interferometer.get_whitened_time_series_from_whitened_frequency_series( - whitened_fd_waveform - )[time_idxs] + td_waveform = ( + interferometer.get_whitened_time_series_from_whitened_frequency_series( + whitened_fd_waveform + )[time_idxs] + ) td_waveforms.append(td_waveform) fd_waveforms = asd_from_freq_series( - fd_waveforms, - 1 / interferometer.strain_data.duration) + fd_waveforms, 1 / interferometer.strain_data.duration + ) td_waveforms = np.array(td_waveforms) delta = (1 + level) / 2 upper_percentile = delta * 100 lower_percentile = (1 - delta) * 100 logger.debug( - 'Plotting posterior between the {} and {} percentiles'.format( + "Plotting posterior between the {} and {} percentiles".format( lower_percentile, upper_percentile ) ) @@ -541,38 +608,46 @@ def plot_interferometer_waveform_posterior( if format == "html": fig.add_trace( go.Scatter( - x=plot_frequencies, y=np.median(fd_waveforms, axis=0), + x=plot_frequencies, + y=np.median(fd_waveforms, axis=0), fill=None, - mode='lines', line_color=WAVEFORM_COLOR, + mode="lines", + line_color=WAVEFORM_COLOR, opacity=1, name="Median reconstructed", - legendgroup='median', + legendgroup="median", ), row=1, col=1, ) fig.add_trace( go.Scatter( - x=plot_frequencies, y=np.percentile(fd_waveforms, lower_percentile, axis=0), + x=plot_frequencies, + y=np.percentile(fd_waveforms, lower_percentile, axis=0), fill=None, - mode='lines', + mode="lines", line_color=WAVEFORM_COLOR, opacity=0.1, - name="{:.2f}% credible interval".format(upper_percentile - lower_percentile), - legendgroup='uncertainty', + name="{:.2f}% credible interval".format( + upper_percentile - lower_percentile + ), + legendgroup="uncertainty", ), row=1, col=1, ) fig.add_trace( go.Scatter( - x=plot_frequencies, y=np.percentile(fd_waveforms, upper_percentile, axis=0), - fill='tonexty', - mode='lines', + x=plot_frequencies, + y=np.percentile(fd_waveforms, upper_percentile, axis=0), + fill="tonexty", + mode="lines", line_color=WAVEFORM_COLOR, opacity=0.1, - name="{:.2f}% credible interval".format(upper_percentile - lower_percentile), - legendgroup='uncertainty', + name="{:.2f}% credible interval".format( + upper_percentile - lower_percentile + ), + legendgroup="uncertainty", showlegend=False, ), row=1, @@ -580,12 +655,14 @@ def plot_interferometer_waveform_posterior( ) fig.add_trace( go.Scatter( - x=plot_times, y=np.median(td_waveforms, axis=0), + x=plot_times, + y=np.median(td_waveforms, axis=0), fill=None, - mode='lines', line_color=WAVEFORM_COLOR, + mode="lines", + line_color=WAVEFORM_COLOR, opacity=1, name="Median reconstructed", - legendgroup='median', + legendgroup="median", showlegend=False, ), row=2, @@ -593,13 +670,16 @@ def plot_interferometer_waveform_posterior( ) fig.add_trace( go.Scatter( - x=plot_times, y=np.percentile(td_waveforms, lower_percentile, axis=0), + x=plot_times, + y=np.percentile(td_waveforms, lower_percentile, axis=0), fill=None, - mode='lines', + mode="lines", line_color=WAVEFORM_COLOR, opacity=0.1, - name="{:.2f}% credible interval".format(upper_percentile - lower_percentile), - legendgroup='uncertainty', + name="{:.2f}% credible interval".format( + upper_percentile - lower_percentile + ), + legendgroup="uncertainty", showlegend=False, ), row=2, @@ -607,13 +687,16 @@ def plot_interferometer_waveform_posterior( ) fig.add_trace( go.Scatter( - x=plot_times, y=np.percentile(td_waveforms, upper_percentile, axis=0), - fill='tonexty', - mode='lines', + x=plot_times, + y=np.percentile(td_waveforms, upper_percentile, axis=0), + fill="tonexty", + mode="lines", line_color=WAVEFORM_COLOR, opacity=0.1, - name="{:.2f}% credible interval".format(upper_percentile - lower_percentile), - legendgroup='uncertainty', + name="{:.2f}% credible interval".format( + upper_percentile - lower_percentile + ), + legendgroup="uncertainty", showlegend=False, ), row=2, @@ -623,58 +706,69 @@ def plot_interferometer_waveform_posterior( lower_limit = np.mean(fd_waveforms, axis=0)[0] / 1e3 axs[0].loglog( plot_frequencies, - np.mean(fd_waveforms, axis=0), color=WAVEFORM_COLOR, label='Mean reconstructed') + np.mean(fd_waveforms, axis=0), + color=WAVEFORM_COLOR, + label="Mean reconstructed", + ) axs[0].fill_between( plot_frequencies, np.percentile(fd_waveforms, lower_percentile, axis=0), np.percentile(fd_waveforms, upper_percentile, axis=0), color=WAVEFORM_COLOR, - label=r'{}% credible interval'.format(int(upper_percentile - lower_percentile)), - alpha=0.3) - axs[1].plot( - plot_times, np.mean(td_waveforms, axis=0), - color=WAVEFORM_COLOR) + label=r"{}% credible interval".format( + int(upper_percentile - lower_percentile) + ), + alpha=0.3, + ) + axs[1].plot(plot_times, np.mean(td_waveforms, axis=0), color=WAVEFORM_COLOR) axs[1].fill_between( - plot_times, np.percentile( - td_waveforms, lower_percentile, axis=0), + plot_times, + np.percentile(td_waveforms, lower_percentile, axis=0), np.percentile(td_waveforms, upper_percentile, axis=0), color=WAVEFORM_COLOR, - alpha=0.3) + alpha=0.3, + ) if self.injection_parameters is not None: try: hf_inj = waveform_generator.frequency_domain_strain( - self.injection_parameters) + self.injection_parameters + ) hf_inj_det = interferometer.get_detector_response( - hf_inj, self.injection_parameters) + hf_inj, self.injection_parameters + ) ht_inj_det = infft( - hf_inj_det * np.sqrt(2. / interferometer.sampling_frequency) / - interferometer.amplitude_spectral_density_array, - self.sampling_frequency)[time_idxs] + hf_inj_det + * np.sqrt(2.0 / interferometer.sampling_frequency) + / interferometer.amplitude_spectral_density_array, + self.sampling_frequency, + )[time_idxs] if format == "html": fig.add_trace( go.Scatter( x=plot_frequencies, y=asd_from_freq_series( hf_inj_det[frequency_idxs], - 1 / interferometer.strain_data.duration), + 1 / interferometer.strain_data.duration, + ), fill=None, - mode='lines', - line=dict(color=INJECTION_COLOR, dash='dot'), + mode="lines", + line=dict(color=INJECTION_COLOR, dash="dot"), name="Injection", - legendgroup='injection', + legendgroup="injection", ), row=1, col=1, ) fig.add_trace( go.Scatter( - x=plot_times, y=ht_inj_det, + x=plot_times, + y=ht_inj_det, fill=None, - mode='lines', - line=dict(color=INJECTION_COLOR, dash='dot'), + mode="lines", + line=dict(color=INJECTION_COLOR, dash="dot"), name="Injection", - legendgroup='injection', + legendgroup="injection", showlegend=False, ), row=2, @@ -685,14 +779,18 @@ def plot_interferometer_waveform_posterior( plot_frequencies, asd_from_freq_series( hf_inj_det[frequency_idxs], - 1 / interferometer.strain_data.duration), - color=INJECTION_COLOR, label='Injection', linestyle=':') + 1 / interferometer.strain_data.duration, + ), + color=INJECTION_COLOR, + label="Injection", + linestyle=":", + ) axs[1].plot( - plot_times, ht_inj_det, - color=INJECTION_COLOR, linestyle=':') - logger.debug('Plotted injection.') + plot_times, ht_inj_det, color=INJECTION_COLOR, linestyle=":" + ) + logger.debug("Plotted injection.") except IndexError as e: - logger.info('Failed to plot injection with message {}.'.format(e)) + logger.info("Failed to plot injection with message {}.".format(e)) f_domain_x_label = "$f [\\mathrm{Hz}]$" f_domain_y_label = "$\\mathrm{ASD} \\left[\\mathrm{Hz}^{-1/2}\\right]$" @@ -704,29 +802,27 @@ def plot_interferometer_waveform_posterior( fig.update_xaxes(title_text=t_domain_x_label, type="linear", row=2) fig.update_yaxes(title_text=t_domain_y_label, type="linear", row=2) else: - axs[0].set_xlim(interferometer.minimum_frequency, - interferometer.maximum_frequency) + axs[0].set_xlim( + interferometer.minimum_frequency, interferometer.maximum_frequency + ) axs[1].set_xlim(start_time, end_time) axs[0].set_ylim(lower_limit) axs[0].set_xlabel(f_domain_x_label) axs[0].set_ylabel(f_domain_y_label) axs[1].set_xlabel(t_domain_x_label) axs[1].set_ylabel(t_domain_y_label) - axs[0].legend(loc='lower left', ncol=2) + axs[0].legend(loc="lower left", ncol=2) if save: filename = os.path.join( self.outdir, - self.label + '_{}_waveform.{}'.format( - interferometer.name, format)) - if format == 'html': - plot(fig, filename=filename, include_mathjax='cdn', auto_open=False) + self.label + "_{}_waveform.{}".format(interferometer.name, format), + ) + if format == "html": + plot(fig, filename=filename, include_mathjax="cdn", auto_open=False) else: plt.tight_layout() - safe_save_figure( - fig=fig, filename=filename, - format=format, dpi=600 - ) + safe_save_figure(fig=fig, filename=filename, format=format, dpi=600) plt.close() logger.debug("Waveform figure saved to {}".format(filename)) rcParams["font.size"] = old_font_size @@ -735,11 +831,23 @@ def plot_interferometer_waveform_posterior( return fig def plot_skymap( - self, maxpts=None, trials=5, jobs=1, enable_multiresolution=True, - objid=None, instruments=None, geo=False, dpi=600, - transparent=False, colorbar=False, contour=[50, 90], - annotate=True, cmap='cylon', load_pickle=False): - """ Generate a fits file and sky map from a result + self, + maxpts=None, + trials=5, + jobs=1, + enable_multiresolution=True, + objid=None, + instruments=None, + geo=False, + dpi=600, + transparent=False, + colorbar=False, + contour=[50, 90], + annotate=True, + cmap="cylon", + load_pickle=False, + ): + """Generate a fits file and sky map from a result Code adapted from ligo.skymap.tool.ligo_skymap_from_samples and ligo.skymap.tool.plot_skymap. Note, the use of this additionally @@ -790,29 +898,31 @@ def plot_skymap( check_directory_exists_and_if_not_mkdir(self.outdir) - logger.info('Reading samples for skymap') + logger.info("Reading samples for skymap") data = self.posterior if maxpts is not None and maxpts < len(data): - logger.info('Taking random subsample of chain') + logger.info("Taking random subsample of chain") data = data.sample(maxpts) - default_obj_filename = os.path.join(self.outdir, '{}_skypost.obj'.format(self.label)) + default_obj_filename = os.path.join( + self.outdir, "{}_skypost.obj".format(self.label) + ) if load_pickle is False: try: - pts = data[['ra', 'dec', 'luminosity_distance']].values + pts = data[["ra", "dec", "luminosity_distance"]].values confidence_levels = kde.Clustered2Plus1DSkyKDE distance = True except KeyError: logger.warning("The results file does not contain luminosity_distance") - pts = data[['ra', 'dec']].values + pts = data[["ra", "dec"]].values confidence_levels = kde.Clustered2DSkyKDE distance = False - logger.info('Initialising skymap class') + logger.info("Initialising skymap class") skypost = confidence_levels(pts, trials=trials, jobs=jobs) - logger.info('Pickling skymap to {}'.format(default_obj_filename)) + logger.info("Pickling skymap to {}".format(default_obj_filename)) safe_file_dump(skypost, default_obj_filename, "pickle") else: @@ -820,38 +930,38 @@ def plot_skymap( obj_filename = load_pickle else: obj_filename = default_obj_filename - logger.info('Reading from pickle {}'.format(obj_filename)) - with open(obj_filename, 'rb') as file: + logger.info("Reading from pickle {}".format(obj_filename)) + with open(obj_filename, "rb") as file: skypost = pickle.load(file) skypost.jobs = jobs distance = isinstance(skypost, kde.Clustered2Plus1DSkyKDE) - logger.info('Making skymap') + logger.info("Making skymap") hpmap = skypost.as_healpix() if not enable_multiresolution: hpmap = bayestar.rasterize(hpmap) hpmap.meta.update(io.fits.metadata_for_version_module(version)) - hpmap.meta['creator'] = "bilby" - hpmap.meta['origin'] = 'LIGO/Virgo' - hpmap.meta['gps_creation_time'] = Time.now().gps - hpmap.meta['history'] = "" + hpmap.meta["creator"] = "bilby" + hpmap.meta["origin"] = "LIGO/Virgo" + hpmap.meta["gps_creation_time"] = Time.now().gps + hpmap.meta["history"] = "" if objid is not None: - hpmap.meta['objid'] = objid + hpmap.meta["objid"] = objid if instruments: - hpmap.meta['instruments'] = instruments + hpmap.meta["instruments"] = instruments if distance: - hpmap.meta['distmean'] = np.mean(data['luminosity_distance']) - hpmap.meta['diststd'] = np.std(data['luminosity_distance']) + hpmap.meta["distmean"] = np.mean(data["luminosity_distance"]) + hpmap.meta["diststd"] = np.std(data["luminosity_distance"]) try: - time = data['geocent_time'] - hpmap.meta['gps_time'] = time.mean() + time = data["geocent_time"] + hpmap.meta["gps_time"] = time.mean() except KeyError: - logger.warning('Cannot determine the event time from geocent_time') + logger.warning("Cannot determine the event time from geocent_time") fits_filename = os.path.join(self.outdir, "{}_skymap.fits".format(self.label)) - logger.info('Saving skymap fits-file to {}'.format(fits_filename)) + logger.info("Saving skymap fits-file to {}".format(fits_filename)) io.write_sky_map(fits_filename, hpmap, nest=True) skymap, metadata = io.fits.read_sky_map(fits_filename, nest=None) @@ -862,41 +972,50 @@ def plot_skymap( probperdeg2 = skymap / deg2perpix if geo: - obstime = Time(metadata['gps_time'], format='gps').utc.isot - ax = plt.axes(projection='geo degrees mollweide', obstime=obstime) + obstime = Time(metadata["gps_time"], format="gps").utc.isot + ax = plt.axes(projection="geo degrees mollweide", obstime=obstime) else: - ax = plt.axes(projection='astro hours mollweide') + ax = plt.axes(projection="astro hours mollweide") ax.grid() # Plot sky map. vmax = probperdeg2.max() img = ax.imshow_hpx( - (probperdeg2, 'ICRS'), nested=metadata['nest'], vmin=0., vmax=vmax, - cmap=cmap) + (probperdeg2, "ICRS"), + nested=metadata["nest"], + vmin=0.0, + vmax=vmax, + cmap=cmap, + ) # Add colorbar. if colorbar: cb = plot.colorbar(img) - cb.set_label(r'prob. per deg$^2$') + cb.set_label(r"prob. per deg$^2$") if contour is not None: confidence_levels = 100 * postprocess.find_greedy_credible_levels(skymap) contours = ax.contour_hpx( - (confidence_levels, 'ICRS'), nested=metadata['nest'], - colors='k', linewidths=0.5, levels=contour) - fmt = r'%g\%%' if rcParams['text.usetex'] else '%g%%' + (confidence_levels, "ICRS"), + nested=metadata["nest"], + colors="k", + linewidths=0.5, + levels=contour, + ) + fmt = r"%g\%%" if rcParams["text.usetex"] else "%g%%" plt.clabel(contours, fmt=fmt, fontsize=6, inline=True) # Add continents. if geo: geojson_filename = os.path.join( - os.path.dirname(plot.__file__), 'ne_simplified_coastline.json') - with open(geojson_filename, 'r') as geojson_file: - geoms = json.load(geojson_file)['geometries'] - verts = [coord for geom in geoms - for coord in zip(*geom['coordinates'])] - plt.plot(*verts, color='0.5', linewidth=0.5, - transform=ax.get_transform('world')) + os.path.dirname(plot.__file__), "ne_simplified_coastline.json" + ) + with open(geojson_filename, "r") as geojson_file: + geoms = json.load(geojson_file)["geometries"] + verts = [coord for geom in geoms for coord in zip(*geom["coordinates"])] + plt.plot( + *verts, color="0.5", linewidth=0.5, transform=ax.get_transform("world") + ) # Add a white outline to all text to make it stand out from the background. plot.outline_text(ax) @@ -904,19 +1023,19 @@ def plot_skymap( if annotate: text = [] try: - objid = metadata['objid'] + objid = metadata["objid"] except KeyError: pass else: - text.append('event ID: {}'.format(objid)) + text.append("event ID: {}".format(objid)) if contour: pp = np.round(contour).astype(int) - ii = np.round(np.searchsorted(np.sort(confidence_levels), contour) * - deg2perpix).astype(int) + ii = np.round( + np.searchsorted(np.sort(confidence_levels), contour) * deg2perpix + ).astype(int) for i, p in zip(ii, pp): - text.append( - u'{:d}% area: {:d} deg$^2$'.format(p, i)) - ax.text(1, 1, '\n'.join(text), transform=ax.transAxes, ha='right') + text.append("{:d}% area: {:d} deg$^2$".format(p, i)) + ax.text(1, 1, "\n".join(text), transform=ax.transAxes, ha="right") filename = os.path.join(self.outdir, "{}_skymap.png".format(self.label)) logger.info("Generating 2D projected skymap to {}".format(filename)) diff --git a/bilby/gw/sampler/proposal.py b/bilby/gw/sampler/proposal.py index 2ac84687e..b5f902b6c 100644 --- a/bilby/gw/sampler/proposal.py +++ b/bilby/gw/sampler/proposal.py @@ -12,10 +12,10 @@ class SkyLocationWanderJump(JumpProposal): """ def __call__(self, sample, **kwargs): - temperature = 1 / kwargs.get('inverse_temperature', 1.0) + temperature = 1 / kwargs.get("inverse_temperature", 1.0) sigma = temperature**0.5 / 2 / np.pi - sample['ra'] += random.gauss(0, sigma) - sample['dec'] += random.gauss(0, sigma) + sample["ra"] += random.gauss(0, sigma) + sample["dec"] += random.gauss(0, sigma) return super(SkyLocationWanderJump, self).__call__(sample) @@ -25,16 +25,16 @@ class CorrelatedPolarisationPhaseJump(JumpProposal): """ def __call__(self, sample, **kwargs): - alpha = sample['psi'] + sample['phase'] - beta = sample['psi'] - sample['phase'] + alpha = sample["psi"] + sample["phase"] + beta = sample["psi"] - sample["phase"] draw = random.random() if draw < 0.5: alpha = 3.0 * np.pi * random.random() else: beta = 3.0 * np.pi * random.random() - 2 * np.pi - sample['psi'] = (alpha + beta) * 0.5 - sample['phase'] = (alpha - beta) * 0.5 + sample["psi"] = (alpha + beta) * 0.5 + sample["phase"] = (alpha - beta) * 0.5 return super(CorrelatedPolarisationPhaseJump, self).__call__(sample) @@ -44,6 +44,6 @@ class PolarisationPhaseJump(JumpProposal): """ def __call__(self, sample, **kwargs): - sample['phase'] += np.pi - sample['psi'] += np.pi / 2 + sample["phase"] += np.pi + sample["psi"] += np.pi / 2 return super(PolarisationPhaseJump, self).__call__(sample) diff --git a/bilby/gw/source.py b/bilby/gw/source.py index 97ad51b51..2fad4a3ba 100644 --- a/bilby/gw/source.py +++ b/bilby/gw/source.py @@ -4,11 +4,13 @@ from ..core import utils from ..core.utils import logger from .conversion import bilby_to_lalsimulation_spins -from .utils import (lalsim_GetApproximantFromString, - lalsim_SimInspiralFD, - lalsim_SimInspiralChooseFDWaveform, - lalsim_SimInspiralChooseFDWaveformSequence, - safe_cast_mode_to_int) +from .utils import ( + lalsim_GetApproximantFromString, + lalsim_SimInspiralFD, + lalsim_SimInspiralChooseFDWaveform, + lalsim_SimInspiralChooseFDWaveformSequence, + safe_cast_mode_to_int, +) UNUSED_KWARGS_MESSAGE = """There are unused waveform kwargs. Make sure all of the waveform kwargs are correctly @@ -18,8 +20,21 @@ """ -def gwsignal_binary_black_hole(frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1, - phi_12, a_2, tilt_2, phi_jl, theta_jn, phase, **kwargs): +def gwsignal_binary_black_hole( + frequency_array, + mass_1, + mass_2, + luminosity_distance, + a_1, + tilt_1, + phi_12, + a_2, + tilt_2, + phi_jl, + theta_jn, + phase, + **kwargs, +): """ A binary black hole waveform model using GWsignal @@ -109,29 +124,35 @@ def gwsignal_binary_black_hole(frequency_array, mass_1, mass_2, luminosity_dista ) waveform_kwargs.update(kwargs) - waveform_approximant = waveform_kwargs['waveform_approximant'] + waveform_approximant = waveform_kwargs["waveform_approximant"] if waveform_approximant not in ["SEOBNRv5HM", "SEOBNRv5PHM"]: if waveform_approximant == "IMRPhenomXPHM": - logger.warning("The new waveform interface is unreviewed for this model" + - "and it is only intended for testing.") + logger.warning( + "The new waveform interface is unreviewed for this model" + + "and it is only intended for testing." + ) else: raise ValueError("The new waveform interface is unreviewed for this model.") - reference_frequency = waveform_kwargs['reference_frequency'] - minimum_frequency = waveform_kwargs['minimum_frequency'] - maximum_frequency = waveform_kwargs['maximum_frequency'] - catch_waveform_errors = waveform_kwargs['catch_waveform_errors'] - mode_array = waveform_kwargs['mode_array'] - pn_amplitude_order = waveform_kwargs['pn_amplitude_order'] + reference_frequency = waveform_kwargs["reference_frequency"] + minimum_frequency = waveform_kwargs["minimum_frequency"] + maximum_frequency = waveform_kwargs["maximum_frequency"] + catch_waveform_errors = waveform_kwargs["catch_waveform_errors"] + mode_array = waveform_kwargs["mode_array"] + pn_amplitude_order = waveform_kwargs["pn_amplitude_order"] if pn_amplitude_order != 0: # This is to mimic the behaviour in # https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c#L5542 if pn_amplitude_order == -1: if waveform_approximant in ["SpinTaylorT4", "SpinTaylorT5"]: - pn_amplitude_order = 3 # Equivalent to MAX_PRECESSING_AMP_PN_ORDER in LALSimulation + pn_amplitude_order = ( + 3 # Equivalent to MAX_PRECESSING_AMP_PN_ORDER in LALSimulation + ) else: - pn_amplitude_order = 6 # Equivalent to MAX_NONPRECESSING_AMP_PN_ORDER in LALSimulation - start_frequency = minimum_frequency * 2. / (pn_amplitude_order + 2) + pn_amplitude_order = ( + 6 # Equivalent to MAX_NONPRECESSING_AMP_PN_ORDER in LALSimulation + ) + start_frequency = minimum_frequency * 2.0 / (pn_amplitude_order + 2) else: start_frequency = minimum_frequency @@ -140,13 +161,25 @@ def gwsignal_binary_black_hole(frequency_array, mass_1, mass_2, luminosity_dista delta_frequency = frequency_array[1] - frequency_array[0] - frequency_bounds = ((frequency_array >= minimum_frequency) * - (frequency_array <= maximum_frequency)) + frequency_bounds = (frequency_array >= minimum_frequency) * ( + frequency_array <= maximum_frequency + ) - iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = bilby_to_lalsimulation_spins( - theta_jn=theta_jn, phi_jl=phi_jl, tilt_1=tilt_1, tilt_2=tilt_2, - phi_12=phi_12, a_1=a_1, a_2=a_2, mass_1=mass_1 * utils.solar_mass, mass_2=mass_2 * utils.solar_mass, - reference_frequency=reference_frequency, phase=phase) + iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = ( + bilby_to_lalsimulation_spins( + theta_jn=theta_jn, + phi_jl=phi_jl, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_12=phi_12, + a_1=a_1, + a_2=a_2, + mass_1=mass_1 * utils.solar_mass, + mass_2=mass_2 * utils.solar_mass, + reference_frequency=reference_frequency, + phase=phase, + ) + ) eccentricity = 0.0 longitude_ascending_nodes = 0.0 @@ -154,35 +187,38 @@ def gwsignal_binary_black_hole(frequency_array, mass_1, mass_2, luminosity_dista # Check if conditioning is needed condition = 0 - if wf_gen.metadata["implemented_domain"] == 'time': + if wf_gen.metadata["implemented_domain"] == "time": condition = 1 # Create dict for gwsignal generator - gwsignal_dict = {'mass1' : mass_1 * u.solMass, - 'mass2' : mass_2 * u.solMass, - 'spin1x' : spin_1x * u.dimensionless_unscaled, - 'spin1y' : spin_1y * u.dimensionless_unscaled, - 'spin1z' : spin_1z * u.dimensionless_unscaled, - 'spin2x' : spin_2x * u.dimensionless_unscaled, - 'spin2y' : spin_2y * u.dimensionless_unscaled, - 'spin2z' : spin_2z * u.dimensionless_unscaled, - 'deltaF' : delta_frequency * u.Hz, - 'f22_start' : start_frequency * u.Hz, - 'f_max': maximum_frequency * u.Hz, - 'f22_ref': reference_frequency * u.Hz, - 'phi_ref' : phase * u.rad, - 'distance' : luminosity_distance * u.Mpc, - 'inclination' : iota * u.rad, - 'eccentricity' : eccentricity * u.dimensionless_unscaled, - 'longAscNodes' : longitude_ascending_nodes * u.rad, - 'meanPerAno' : mean_per_ano * u.rad, - # 'ModeArray': mode_array, - 'condition': condition - } + gwsignal_dict = { + "mass1": mass_1 * u.solMass, + "mass2": mass_2 * u.solMass, + "spin1x": spin_1x * u.dimensionless_unscaled, + "spin1y": spin_1y * u.dimensionless_unscaled, + "spin1z": spin_1z * u.dimensionless_unscaled, + "spin2x": spin_2x * u.dimensionless_unscaled, + "spin2y": spin_2y * u.dimensionless_unscaled, + "spin2z": spin_2z * u.dimensionless_unscaled, + "deltaF": delta_frequency * u.Hz, + "f22_start": start_frequency * u.Hz, + "f_max": maximum_frequency * u.Hz, + "f22_ref": reference_frequency * u.Hz, + "phi_ref": phase * u.rad, + "distance": luminosity_distance * u.Mpc, + "inclination": iota * u.rad, + "eccentricity": eccentricity * u.dimensionless_unscaled, + "longAscNodes": longitude_ascending_nodes * u.rad, + "meanPerAno": mean_per_ano * u.rad, + # 'ModeArray': mode_array, + "condition": condition, + } if mode_array is not None: try: - mode_array = [tuple(map(safe_cast_mode_to_int, mode)) for mode in mode_array] + mode_array = [ + tuple(map(safe_cast_mode_to_int, mode)) for mode in mode_array + ] except (ValueError, TypeError) as e: raise ValueError( f"Unable to convert mode_array elements to tuples of ints. " @@ -194,17 +230,17 @@ def gwsignal_binary_black_hole(frequency_array, mass_1, mass_2, luminosity_dista extra_args = waveform_kwargs.copy() for key in [ - "waveform_approximant", - "reference_frequency", - "minimum_frequency", - "maximum_frequency", - "catch_waveform_errors", - "mode_array", - "pn_spin_order", - "pn_amplitude_order", - "pn_tidal_order", - "pn_phase_order", - "numerical_relativity_file", + "waveform_approximant", + "reference_frequency", + "minimum_frequency", + "maximum_frequency", + "catch_waveform_errors", + "mode_array", + "pn_spin_order", + "pn_amplitude_order", + "pn_tidal_order", + "pn_phase_order", + "numerical_relativity_file", ]: if key in extra_args.keys(): del extra_args[key] @@ -219,20 +255,24 @@ def gwsignal_binary_black_hole(frequency_array, mass_1, mass_2, luminosity_dista else: EDOM = ( "Internal function call failed: Input domain error" in e.args[0] - ) or "Input domain error" in e.args[ - 0 - ] + ) or "Input domain error" in e.args[0] if EDOM: - failed_parameters = dict(mass_1=mass_1, mass_2=mass_2, - spin_1=(spin_1x, spin_1y, spin_1z), - spin_2=(spin_2x, spin_2y, spin_2z), - luminosity_distance=luminosity_distance, - iota=iota, phase=phase, - eccentricity=eccentricity, - start_frequency=minimum_frequency) - logger.warning("Evaluating the waveform failed with error: {}\n".format(e) + - "The parameters were {}\n".format(failed_parameters) + - "Likelihood will be set to -inf.") + failed_parameters = dict( + mass_1=mass_1, + mass_2=mass_2, + spin_1=(spin_1x, spin_1y, spin_1z), + spin_2=(spin_2x, spin_2y, spin_2z), + luminosity_distance=luminosity_distance, + iota=iota, + phase=phase, + eccentricity=eccentricity, + start_frequency=minimum_frequency, + ) + logger.warning( + "Evaluating the waveform failed with error: {}\n".format(e) + + "The parameters were {}\n".format(failed_parameters) + + "Likelihood will be set to -inf." + ) return None else: raise @@ -244,15 +284,17 @@ def gwsignal_binary_black_hole(frequency_array, mass_1, mass_2, luminosity_dista h_cross = np.zeros_like(frequency_array, dtype=complex) if len(hplus) > len(frequency_array): - logger.debug("GWsignal waveform longer than bilby's `frequency_array`" + - "({} vs {}), ".format(len(hplus), len(frequency_array)) + - "probably because padded with zeros up to the next power of two length." + - " Truncating GWsignal array.") - h_plus = hplus[:len(h_plus)] - h_cross = hcross[:len(h_cross)] + logger.debug( + "GWsignal waveform longer than bilby's `frequency_array`" + + "({} vs {}), ".format(len(hplus), len(frequency_array)) + + "probably because padded with zeros up to the next power of two length." + + " Truncating GWsignal array." + ) + h_plus = hplus[: len(h_plus)] + h_cross = hcross[: len(h_cross)] else: - h_plus[:len(hplus)] = hplus - h_cross[:len(hcross)] = hcross + h_plus[: len(hplus)] = hplus + h_cross[: len(hcross)] = hcross h_plus *= frequency_bounds h_cross *= frequency_bounds @@ -267,9 +309,21 @@ def gwsignal_binary_black_hole(frequency_array, mass_1, mass_2, luminosity_dista def lal_binary_black_hole( - frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1, - phi_12, a_2, tilt_2, phi_jl, theta_jn, phase, **kwargs): - """ A Binary Black Hole waveform model using lalsimulation + frequency_array, + mass_1, + mass_2, + luminosity_distance, + a_1, + tilt_1, + phi_12, + a_2, + tilt_2, + phi_jl, + theta_jn, + phase, + **kwargs, +): + """A Binary Black Hole waveform model using lalsimulation Parameters ========== @@ -336,23 +390,52 @@ def lal_binary_black_hole( dict: A dictionary with the plus and cross polarisation strain modes """ waveform_kwargs = dict( - waveform_approximant='IMRPhenomPv2', reference_frequency=50.0, - minimum_frequency=20.0, maximum_frequency=frequency_array[-1], - catch_waveform_errors=False, pn_spin_order=-1, pn_tidal_order=-1, - pn_phase_order=-1, pn_amplitude_order=0) + waveform_approximant="IMRPhenomPv2", + reference_frequency=50.0, + minimum_frequency=20.0, + maximum_frequency=frequency_array[-1], + catch_waveform_errors=False, + pn_spin_order=-1, + pn_tidal_order=-1, + pn_phase_order=-1, + pn_amplitude_order=0, + ) waveform_kwargs.update(kwargs) return _base_lal_cbc_fd_waveform( - frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2, - luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase, - a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_12=phi_12, - phi_jl=phi_jl, **waveform_kwargs) + frequency_array=frequency_array, + mass_1=mass_1, + mass_2=mass_2, + luminosity_distance=luminosity_distance, + theta_jn=theta_jn, + phase=phase, + a_1=a_1, + a_2=a_2, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_12=phi_12, + phi_jl=phi_jl, + **waveform_kwargs, + ) def lal_binary_neutron_star( - frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1, - phi_12, a_2, tilt_2, phi_jl, theta_jn, phase, lambda_1, lambda_2, - **kwargs): - """ A Binary Neutron Star waveform model using lalsimulation + frequency_array, + mass_1, + mass_2, + luminosity_distance, + a_1, + tilt_1, + phi_12, + a_2, + tilt_2, + phi_jl, + theta_jn, + phase, + lambda_1, + lambda_2, + **kwargs, +): + """A Binary Neutron Star waveform model using lalsimulation Parameters ========== @@ -420,22 +503,47 @@ def lal_binary_neutron_star( dict: A dictionary with the plus and cross polarisation strain modes """ waveform_kwargs = dict( - waveform_approximant='IMRPhenomPv2_NRTidal', reference_frequency=50.0, - minimum_frequency=20.0, maximum_frequency=frequency_array[-1], - catch_waveform_errors=False, pn_spin_order=-1, pn_tidal_order=-1, - pn_phase_order=-1, pn_amplitude_order=0) + waveform_approximant="IMRPhenomPv2_NRTidal", + reference_frequency=50.0, + minimum_frequency=20.0, + maximum_frequency=frequency_array[-1], + catch_waveform_errors=False, + pn_spin_order=-1, + pn_tidal_order=-1, + pn_phase_order=-1, + pn_amplitude_order=0, + ) waveform_kwargs.update(kwargs) return _base_lal_cbc_fd_waveform( - frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2, - luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase, - a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_12=phi_12, - phi_jl=phi_jl, lambda_1=lambda_1, lambda_2=lambda_2, **waveform_kwargs) + frequency_array=frequency_array, + mass_1=mass_1, + mass_2=mass_2, + luminosity_distance=luminosity_distance, + theta_jn=theta_jn, + phase=phase, + a_1=a_1, + a_2=a_2, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_12=phi_12, + phi_jl=phi_jl, + lambda_1=lambda_1, + lambda_2=lambda_2, + **waveform_kwargs, + ) def lal_eccentric_binary_black_hole_no_spins( - frequency_array, mass_1, mass_2, eccentricity, luminosity_distance, - theta_jn, phase, **kwargs): - """ Eccentric binary black hole waveform model using lalsimulation (EccentricFD) + frequency_array, + mass_1, + mass_2, + eccentricity, + luminosity_distance, + theta_jn, + phase, + **kwargs, +): + """Eccentric binary black hole waveform model using lalsimulation (EccentricFD) Parameters ========== @@ -488,15 +596,27 @@ def lal_eccentric_binary_black_hole_no_spins( dict: A dictionary with the plus and cross polarisation strain modes """ waveform_kwargs = dict( - waveform_approximant='EccentricFD', reference_frequency=10.0, - minimum_frequency=10.0, maximum_frequency=frequency_array[-1], - catch_waveform_errors=False, pn_spin_order=-1, pn_tidal_order=-1, - pn_phase_order=-1, pn_amplitude_order=0) + waveform_approximant="EccentricFD", + reference_frequency=10.0, + minimum_frequency=10.0, + maximum_frequency=frequency_array[-1], + catch_waveform_errors=False, + pn_spin_order=-1, + pn_tidal_order=-1, + pn_phase_order=-1, + pn_amplitude_order=0, + ) waveform_kwargs.update(kwargs) return _base_lal_cbc_fd_waveform( - frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2, - luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase, - eccentricity=eccentricity, **waveform_kwargs) + frequency_array=frequency_array, + mass_1=mass_1, + mass_2=mass_2, + luminosity_distance=luminosity_distance, + theta_jn=theta_jn, + phase=phase, + eccentricity=eccentricity, + **waveform_kwargs, + ) def set_waveform_dictionary(waveform_kwargs, lambda_1=0, lambda_2=0): @@ -520,16 +640,26 @@ def set_waveform_dictionary(waveform_kwargs, lambda_1=0, lambda_2=0): """ import lal import lalsimulation as lalsim - waveform_dictionary = waveform_kwargs.pop('lal_waveform_dictionary', lal.CreateDict()) + + waveform_dictionary = waveform_kwargs.pop( + "lal_waveform_dictionary", lal.CreateDict() + ) waveform_kwargs["TidalLambda1"] = lambda_1 waveform_kwargs["TidalLambda2"] = lambda_2 - waveform_kwargs["NumRelData"] = waveform_kwargs.pop("numerical_relativity_file", None) + waveform_kwargs["NumRelData"] = waveform_kwargs.pop( + "numerical_relativity_file", None + ) for key in [ - "pn_spin_order", "pn_tidal_order", "pn_phase_order", "pn_amplitude_order" + "pn_spin_order", + "pn_tidal_order", + "pn_phase_order", + "pn_amplitude_order", ]: - waveform_kwargs[key[:2].upper() + key[3:].title().replace('_', '')] = waveform_kwargs.pop(key) + waveform_kwargs[key[:2].upper() + key[3:].title().replace("_", "")] = ( + waveform_kwargs.pop(key) + ) for key in list(waveform_kwargs.keys()).copy(): func = getattr(lalsim, f"SimInspiralWaveformParamsInsert{key}", None) @@ -545,15 +675,31 @@ def set_waveform_dictionary(waveform_kwargs, lambda_1=0, lambda_2=0): for mode in mode_array: mode = tuple(map(safe_cast_mode_to_int, mode)) lalsim.SimInspiralModeArrayActivateMode(mode_array_lal, mode[0], mode[1]) - lalsim.SimInspiralWaveformParamsInsertModeArray(waveform_dictionary, mode_array_lal) + lalsim.SimInspiralWaveformParamsInsertModeArray( + waveform_dictionary, mode_array_lal + ) return waveform_dictionary def _base_lal_cbc_fd_waveform( - frequency_array, mass_1, mass_2, luminosity_distance, theta_jn, phase, - a_1=0.0, a_2=0.0, tilt_1=0.0, tilt_2=0.0, phi_12=0.0, phi_jl=0.0, - lambda_1=0.0, lambda_2=0.0, eccentricity=0.0, **waveform_kwargs): - """ Generate a cbc waveform model using lalsimulation + frequency_array, + mass_1, + mass_2, + luminosity_distance, + theta_jn, + phase, + a_1=0.0, + a_2=0.0, + tilt_1=0.0, + tilt_2=0.0, + phi_12=0.0, + phi_jl=0.0, + lambda_1=0.0, + lambda_2=0.0, + eccentricity=0.0, + **waveform_kwargs, +): + """Generate a cbc waveform model using lalsimulation Parameters ========== @@ -596,12 +742,12 @@ def _base_lal_cbc_fd_waveform( """ import lalsimulation as lalsim - waveform_approximant = waveform_kwargs.pop('waveform_approximant') - reference_frequency = waveform_kwargs.pop('reference_frequency') - minimum_frequency = waveform_kwargs.pop('minimum_frequency') - maximum_frequency = waveform_kwargs.pop('maximum_frequency') - catch_waveform_errors = waveform_kwargs.pop('catch_waveform_errors') - pn_amplitude_order = waveform_kwargs['pn_amplitude_order'] + waveform_approximant = waveform_kwargs.pop("waveform_approximant") + reference_frequency = waveform_kwargs.pop("reference_frequency") + minimum_frequency = waveform_kwargs.pop("minimum_frequency") + maximum_frequency = waveform_kwargs.pop("maximum_frequency") + catch_waveform_errors = waveform_kwargs.pop("catch_waveform_errors") + pn_amplitude_order = waveform_kwargs["pn_amplitude_order"] waveform_dictionary = set_waveform_dictionary(waveform_kwargs, lambda_1, lambda_2) approximant = lalsim_GetApproximantFromString(waveform_approximant) @@ -615,17 +761,29 @@ def _base_lal_cbc_fd_waveform( delta_frequency = frequency_array[1] - frequency_array[0] - frequency_bounds = ((frequency_array >= minimum_frequency) * - (frequency_array <= maximum_frequency)) + frequency_bounds = (frequency_array >= minimum_frequency) * ( + frequency_array <= maximum_frequency + ) luminosity_distance = luminosity_distance * 1e6 * utils.parsec mass_1 = mass_1 * utils.solar_mass mass_2 = mass_2 * utils.solar_mass - iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = bilby_to_lalsimulation_spins( - theta_jn=theta_jn, phi_jl=phi_jl, tilt_1=tilt_1, tilt_2=tilt_2, - phi_12=phi_12, a_1=a_1, a_2=a_2, mass_1=mass_1, mass_2=mass_2, - reference_frequency=reference_frequency, phase=phase) + iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = ( + bilby_to_lalsimulation_spins( + theta_jn=theta_jn, + phi_jl=phi_jl, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_12=phi_12, + a_1=a_1, + a_2=a_2, + mass_1=mass_1, + mass_2=mass_2, + reference_frequency=reference_frequency, + phase=phase, + ) + ) longitude_ascending_nodes = 0.0 mean_per_ano = 0.0 @@ -636,27 +794,49 @@ def _base_lal_cbc_fd_waveform( wf_func = lalsim_SimInspiralFD try: hplus, hcross = wf_func( - mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, - spin_2z, luminosity_distance, iota, phase, - longitude_ascending_nodes, eccentricity, mean_per_ano, delta_frequency, - start_frequency, maximum_frequency, reference_frequency, - waveform_dictionary, approximant) + mass_1, + mass_2, + spin_1x, + spin_1y, + spin_1z, + spin_2x, + spin_2y, + spin_2z, + luminosity_distance, + iota, + phase, + longitude_ascending_nodes, + eccentricity, + mean_per_ano, + delta_frequency, + start_frequency, + maximum_frequency, + reference_frequency, + waveform_dictionary, + approximant, + ) except Exception as e: if not catch_waveform_errors: raise else: - EDOM = (e.args[0] == 'Internal function call failed: Input domain error') + EDOM = e.args[0] == "Internal function call failed: Input domain error" if EDOM: - failed_parameters = dict(mass_1=mass_1, mass_2=mass_2, - spin_1=(spin_1x, spin_1y, spin_1z), - spin_2=(spin_2x, spin_2y, spin_2z), - luminosity_distance=luminosity_distance, - iota=iota, phase=phase, - eccentricity=eccentricity, - start_frequency=start_frequency) - logger.warning("Evaluating the waveform failed with error: {}\n".format(e) + - "The parameters were {}\n".format(failed_parameters) + - "Likelihood will be set to -inf.") + failed_parameters = dict( + mass_1=mass_1, + mass_2=mass_2, + spin_1=(spin_1x, spin_1y, spin_1z), + spin_2=(spin_2x, spin_2y, spin_2z), + luminosity_distance=luminosity_distance, + iota=iota, + phase=phase, + eccentricity=eccentricity, + start_frequency=start_frequency, + ) + logger.warning( + "Evaluating the waveform failed with error: {}\n".format(e) + + "The parameters were {}\n".format(failed_parameters) + + "Likelihood will be set to -inf." + ) return None else: raise @@ -665,21 +845,25 @@ def _base_lal_cbc_fd_waveform( h_cross = np.zeros_like(frequency_array, dtype=complex) if len(hplus.data.data) > len(frequency_array): - logger.debug("LALsim waveform longer than bilby's `frequency_array`" + - "({} vs {}), ".format(len(hplus.data.data), len(frequency_array)) + - "probably because padded with zeros up to the next power of two length." + - " Truncating lalsim array.") - h_plus = hplus.data.data[:len(h_plus)] - h_cross = hcross.data.data[:len(h_cross)] + logger.debug( + "LALsim waveform longer than bilby's `frequency_array`" + + "({} vs {}), ".format(len(hplus.data.data), len(frequency_array)) + + "probably because padded with zeros up to the next power of two length." + + " Truncating lalsim array." + ) + h_plus = hplus.data.data[: len(h_plus)] + h_cross = hcross.data.data[: len(h_cross)] else: - h_plus[:len(hplus.data.data)] = hplus.data.data - h_cross[:len(hcross.data.data)] = hcross.data.data + h_plus[: len(hplus.data.data)] = hplus.data.data + h_cross[: len(hcross.data.data)] = hcross.data.data h_plus *= frequency_bounds h_cross *= frequency_bounds if wf_func == lalsim_SimInspiralFD: - dt = 1 / hplus.deltaF + (hplus.epoch.gpsSeconds + hplus.epoch.gpsNanoSeconds * 1e-9) + dt = 1 / hplus.deltaF + ( + hplus.epoch.gpsSeconds + hplus.epoch.gpsNanoSeconds * 1e-9 + ) time_shift = np.exp(-1j * 2 * np.pi * dt * frequency_array[frequency_bounds]) h_plus[frequency_bounds] *= time_shift h_cross[frequency_bounds] *= time_shift @@ -691,40 +875,111 @@ def _base_lal_cbc_fd_waveform( def binary_black_hole_roq( - frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1, - phi_12, a_2, tilt_2, phi_jl, theta_jn, phase, **waveform_arguments): + frequency_array, + mass_1, + mass_2, + luminosity_distance, + a_1, + tilt_1, + phi_12, + a_2, + tilt_2, + phi_jl, + theta_jn, + phase, + **waveform_arguments, +): waveform_kwargs = dict( - waveform_approximant='IMRPhenomPv2', reference_frequency=20.0, - catch_waveform_errors=False, pn_spin_order=-1, pn_tidal_order=-1, - pn_phase_order=-1, pn_amplitude_order=0) + waveform_approximant="IMRPhenomPv2", + reference_frequency=20.0, + catch_waveform_errors=False, + pn_spin_order=-1, + pn_tidal_order=-1, + pn_phase_order=-1, + pn_amplitude_order=0, + ) waveform_kwargs.update(waveform_arguments) return _base_roq_waveform( - frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2, - luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase, - a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_jl=phi_jl, - phi_12=phi_12, lambda_1=0.0, lambda_2=0.0, **waveform_kwargs) + frequency_array=frequency_array, + mass_1=mass_1, + mass_2=mass_2, + luminosity_distance=luminosity_distance, + theta_jn=theta_jn, + phase=phase, + a_1=a_1, + a_2=a_2, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_jl=phi_jl, + phi_12=phi_12, + lambda_1=0.0, + lambda_2=0.0, + **waveform_kwargs, + ) def binary_neutron_star_roq( - frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1, - phi_12, a_2, tilt_2, phi_jl, lambda_1, lambda_2, theta_jn, phase, - **waveform_arguments): + frequency_array, + mass_1, + mass_2, + luminosity_distance, + a_1, + tilt_1, + phi_12, + a_2, + tilt_2, + phi_jl, + lambda_1, + lambda_2, + theta_jn, + phase, + **waveform_arguments, +): waveform_kwargs = dict( - waveform_approximant='IMRPhenomD_NRTidal', reference_frequency=20.0, - catch_waveform_errors=False, pn_spin_order=-1, pn_tidal_order=-1, - pn_phase_order=-1, pn_amplitude_order=0) + waveform_approximant="IMRPhenomD_NRTidal", + reference_frequency=20.0, + catch_waveform_errors=False, + pn_spin_order=-1, + pn_tidal_order=-1, + pn_phase_order=-1, + pn_amplitude_order=0, + ) waveform_kwargs.update(waveform_arguments) return _base_roq_waveform( - frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2, - luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase, - a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_jl=phi_jl, - phi_12=phi_12, lambda_1=lambda_1, lambda_2=lambda_2, **waveform_kwargs) + frequency_array=frequency_array, + mass_1=mass_1, + mass_2=mass_2, + luminosity_distance=luminosity_distance, + theta_jn=theta_jn, + phase=phase, + a_1=a_1, + a_2=a_2, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_jl=phi_jl, + phi_12=phi_12, + lambda_1=lambda_1, + lambda_2=lambda_2, + **waveform_kwargs, + ) def lal_binary_black_hole_relative_binning( - frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1, - phi_12, a_2, tilt_2, phi_jl, theta_jn, phase, **kwargs): - """ Source model to go with RelativeBinningGravitationalWaveTransient likelihood. + frequency_array, + mass_1, + mass_2, + luminosity_distance, + a_1, + tilt_1, + phi_12, + a_2, + tilt_2, + phi_jl, + theta_jn, + phase, + **kwargs, +): + """Source model to go with RelativeBinningGravitationalWaveTransient likelihood. All parameters are the same as in the usual source models, except `fiducial` @@ -736,35 +991,79 @@ def lal_binary_black_hole_relative_binning( fiducial = kwargs.pop("fiducial", 0) waveform_kwargs = dict( - waveform_approximant='IMRPhenomPv2', reference_frequency=50.0, - minimum_frequency=20.0, maximum_frequency=frequency_array[-1], - catch_waveform_errors=False, pn_spin_order=-1, pn_tidal_order=-1, - pn_phase_order=-1, pn_amplitude_order=0) + waveform_approximant="IMRPhenomPv2", + reference_frequency=50.0, + minimum_frequency=20.0, + maximum_frequency=frequency_array[-1], + catch_waveform_errors=False, + pn_spin_order=-1, + pn_tidal_order=-1, + pn_phase_order=-1, + pn_amplitude_order=0, + ) waveform_kwargs.update(kwargs) if fiducial == 1: _ = waveform_kwargs.pop("frequency_bin_edges", None) return _base_lal_cbc_fd_waveform( - frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2, - luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase, - a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_jl=phi_jl, - phi_12=phi_12, lambda_1=0.0, lambda_2=0.0, **waveform_kwargs) + frequency_array=frequency_array, + mass_1=mass_1, + mass_2=mass_2, + luminosity_distance=luminosity_distance, + theta_jn=theta_jn, + phase=phase, + a_1=a_1, + a_2=a_2, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_jl=phi_jl, + phi_12=phi_12, + lambda_1=0.0, + lambda_2=0.0, + **waveform_kwargs, + ) else: _ = waveform_kwargs.pop("minimum_frequency", None) _ = waveform_kwargs.pop("maximum_frequency", None) waveform_kwargs["frequencies"] = waveform_kwargs.pop("frequency_bin_edges") return _base_waveform_frequency_sequence( - frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2, - luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase, - a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_jl=phi_jl, - phi_12=phi_12, lambda_1=0.0, lambda_2=0.0, **waveform_kwargs) + frequency_array=frequency_array, + mass_1=mass_1, + mass_2=mass_2, + luminosity_distance=luminosity_distance, + theta_jn=theta_jn, + phase=phase, + a_1=a_1, + a_2=a_2, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_jl=phi_jl, + phi_12=phi_12, + lambda_1=0.0, + lambda_2=0.0, + **waveform_kwargs, + ) def lal_binary_neutron_star_relative_binning( - frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1, - phi_12, a_2, tilt_2, phi_jl, lambda_1, lambda_2, theta_jn, phase, **kwargs): - """ Source model to go with RelativeBinningGravitationalWaveTransient likelihood. + frequency_array, + mass_1, + mass_2, + luminosity_distance, + a_1, + tilt_1, + phi_12, + a_2, + tilt_2, + phi_jl, + lambda_1, + lambda_2, + theta_jn, + phase, + **kwargs, +): + """Source model to go with RelativeBinningGravitationalWaveTransient likelihood. All parameters are the same as in the usual source models, except `fiducial` @@ -776,35 +1075,78 @@ def lal_binary_neutron_star_relative_binning( fiducial = kwargs.pop("fiducial", 0) waveform_kwargs = dict( - waveform_approximant='IMRPhenomPv2_NRTidal', reference_frequency=50.0, - minimum_frequency=20.0, maximum_frequency=frequency_array[-1], - catch_waveform_errors=False, pn_spin_order=-1, pn_tidal_order=-1, - pn_phase_order=-1, pn_amplitude_order=0) + waveform_approximant="IMRPhenomPv2_NRTidal", + reference_frequency=50.0, + minimum_frequency=20.0, + maximum_frequency=frequency_array[-1], + catch_waveform_errors=False, + pn_spin_order=-1, + pn_tidal_order=-1, + pn_phase_order=-1, + pn_amplitude_order=0, + ) waveform_kwargs.update(kwargs) if fiducial == 1: _ = waveform_kwargs.pop("frequency_bin_edges", None) return _base_lal_cbc_fd_waveform( - frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2, - luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase, - a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_12=phi_12, - phi_jl=phi_jl, lambda_1=lambda_1, lambda_2=lambda_2, **waveform_kwargs) + frequency_array=frequency_array, + mass_1=mass_1, + mass_2=mass_2, + luminosity_distance=luminosity_distance, + theta_jn=theta_jn, + phase=phase, + a_1=a_1, + a_2=a_2, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_12=phi_12, + phi_jl=phi_jl, + lambda_1=lambda_1, + lambda_2=lambda_2, + **waveform_kwargs, + ) else: _ = waveform_kwargs.pop("minimum_frequency", None) _ = waveform_kwargs.pop("maximum_frequency", None) waveform_kwargs["frequencies"] = waveform_kwargs.pop("frequency_bin_edges") return _base_waveform_frequency_sequence( - frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2, - luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase, - a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_jl=phi_jl, - phi_12=phi_12, lambda_1=lambda_1, lambda_2=lambda_2, **waveform_kwargs) + frequency_array=frequency_array, + mass_1=mass_1, + mass_2=mass_2, + luminosity_distance=luminosity_distance, + theta_jn=theta_jn, + phase=phase, + a_1=a_1, + a_2=a_2, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_jl=phi_jl, + phi_12=phi_12, + lambda_1=lambda_1, + lambda_2=lambda_2, + **waveform_kwargs, + ) def _base_roq_waveform( - frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1, - phi_12, a_2, tilt_2, lambda_1, lambda_2, phi_jl, theta_jn, phase, - **waveform_arguments): - """ Base source model for ROQGravitationalWaveTransient, which evaluates + frequency_array, + mass_1, + mass_2, + luminosity_distance, + a_1, + tilt_1, + phi_12, + a_2, + tilt_2, + lambda_1, + lambda_2, + phi_jl, + theta_jn, + phase, + **waveform_arguments, +): + """Base source model for ROQGravitationalWaveTransient, which evaluates waveform values at frequency nodes contained in waveform_arguments. This requires that waveform_arguments contain all of 'frequency_nodes', 'linear_indices', and 'quadratic_indices', or both 'frequency_nodes_linear' and @@ -863,46 +1205,72 @@ def _base_roq_waveform( Dict containing plus and cross modes evaluated at the linear and quadratic frequency nodes. """ - if 'frequency_nodes' not in waveform_arguments: - size_linear = len(waveform_arguments['frequency_nodes_linear']) + if "frequency_nodes" not in waveform_arguments: + size_linear = len(waveform_arguments["frequency_nodes_linear"]) frequency_nodes_combined = np.hstack( - (waveform_arguments.pop('frequency_nodes_linear'), - waveform_arguments.pop('frequency_nodes_quadratic')) + ( + waveform_arguments.pop("frequency_nodes_linear"), + waveform_arguments.pop("frequency_nodes_quadratic"), + ) ) frequency_nodes_unique, original_indices = np.unique( frequency_nodes_combined, return_inverse=True ) linear_indices = original_indices[:size_linear] quadratic_indices = original_indices[size_linear:] - waveform_arguments['frequencies'] = frequency_nodes_unique + waveform_arguments["frequencies"] = frequency_nodes_unique else: linear_indices = waveform_arguments.pop("linear_indices") quadratic_indices = waveform_arguments.pop("quadratic_indices") for key in ["frequency_nodes_linear", "frequency_nodes_quadratic"]: _ = waveform_arguments.pop(key, None) - waveform_arguments['frequencies'] = waveform_arguments.pop('frequency_nodes') + waveform_arguments["frequencies"] = waveform_arguments.pop("frequency_nodes") waveform_polarizations = _base_waveform_frequency_sequence( - frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2, - luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase, - a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_jl=phi_jl, - phi_12=phi_12, lambda_1=lambda_1, lambda_2=lambda_2, **waveform_arguments) + frequency_array=frequency_array, + mass_1=mass_1, + mass_2=mass_2, + luminosity_distance=luminosity_distance, + theta_jn=theta_jn, + phase=phase, + a_1=a_1, + a_2=a_2, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_jl=phi_jl, + phi_12=phi_12, + lambda_1=lambda_1, + lambda_2=lambda_2, + **waveform_arguments, + ) return { - 'linear': { - 'plus': waveform_polarizations['plus'][linear_indices], - 'cross': waveform_polarizations['cross'][linear_indices] + "linear": { + "plus": waveform_polarizations["plus"][linear_indices], + "cross": waveform_polarizations["cross"][linear_indices], + }, + "quadratic": { + "plus": waveform_polarizations["plus"][quadratic_indices], + "cross": waveform_polarizations["cross"][quadratic_indices], }, - 'quadratic': { - 'plus': waveform_polarizations['plus'][quadratic_indices], - 'cross': waveform_polarizations['cross'][quadratic_indices] - } } def binary_black_hole_frequency_sequence( - frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1, - phi_12, a_2, tilt_2, phi_jl, theta_jn, phase, **kwargs): - """ A Binary Black Hole waveform model using lalsimulation. This generates + frequency_array, + mass_1, + mass_2, + luminosity_distance, + a_1, + tilt_1, + phi_12, + a_2, + tilt_2, + phi_jl, + theta_jn, + phase, + **kwargs, +): + """A Binary Black Hole waveform model using lalsimulation. This generates a waveform only on specified frequency points. This is useful for likelihood requiring waveform values at a subset of all the frequency samples. For example, this is used for MBGravitationalWaveTransient. @@ -969,22 +1337,52 @@ def binary_black_hole_frequency_sequence( dict: A dictionary with the plus and cross polarisation strain modes """ waveform_kwargs = dict( - waveform_approximant='IMRPhenomPv2', reference_frequency=50.0, - catch_waveform_errors=False, pn_spin_order=-1, pn_tidal_order=-1, - pn_phase_order=-1, pn_amplitude_order=0) + waveform_approximant="IMRPhenomPv2", + reference_frequency=50.0, + catch_waveform_errors=False, + pn_spin_order=-1, + pn_tidal_order=-1, + pn_phase_order=-1, + pn_amplitude_order=0, + ) waveform_kwargs.update(kwargs) return _base_waveform_frequency_sequence( - frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2, - luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase, - a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_jl=phi_jl, - phi_12=phi_12, lambda_1=0.0, lambda_2=0.0, **waveform_kwargs) + frequency_array=frequency_array, + mass_1=mass_1, + mass_2=mass_2, + luminosity_distance=luminosity_distance, + theta_jn=theta_jn, + phase=phase, + a_1=a_1, + a_2=a_2, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_jl=phi_jl, + phi_12=phi_12, + lambda_1=0.0, + lambda_2=0.0, + **waveform_kwargs, + ) def binary_neutron_star_frequency_sequence( - frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1, - phi_12, a_2, tilt_2, phi_jl, lambda_1, lambda_2, theta_jn, phase, - **kwargs): - """ A Binary Neutron Star waveform model using lalsimulation. This generates + frequency_array, + mass_1, + mass_2, + luminosity_distance, + a_1, + tilt_1, + phi_12, + a_2, + tilt_2, + phi_jl, + lambda_1, + lambda_2, + theta_jn, + phase, + **kwargs, +): + """A Binary Neutron Star waveform model using lalsimulation. This generates a waveform only on specified frequency points. This is useful for likelihood requiring waveform values at a subset of all the frequency samples. For example, this is used for MBGravitationalWaveTransient. @@ -1055,22 +1453,52 @@ def binary_neutron_star_frequency_sequence( dict: A dictionary with the plus and cross polarisation strain modes """ waveform_kwargs = dict( - waveform_approximant='IMRPhenomPv2_NRTidal', reference_frequency=50.0, - catch_waveform_errors=False, pn_spin_order=-1, pn_tidal_order=-1, - pn_phase_order=-1, pn_amplitude_order=0) + waveform_approximant="IMRPhenomPv2_NRTidal", + reference_frequency=50.0, + catch_waveform_errors=False, + pn_spin_order=-1, + pn_tidal_order=-1, + pn_phase_order=-1, + pn_amplitude_order=0, + ) waveform_kwargs.update(kwargs) return _base_waveform_frequency_sequence( - frequency_array=frequency_array, mass_1=mass_1, mass_2=mass_2, - luminosity_distance=luminosity_distance, theta_jn=theta_jn, phase=phase, - a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_jl=phi_jl, - phi_12=phi_12, lambda_1=lambda_1, lambda_2=lambda_2, **waveform_kwargs) + frequency_array=frequency_array, + mass_1=mass_1, + mass_2=mass_2, + luminosity_distance=luminosity_distance, + theta_jn=theta_jn, + phase=phase, + a_1=a_1, + a_2=a_2, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_jl=phi_jl, + phi_12=phi_12, + lambda_1=lambda_1, + lambda_2=lambda_2, + **waveform_kwargs, + ) def _base_waveform_frequency_sequence( - frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1, - phi_12, a_2, tilt_2, lambda_1, lambda_2, phi_jl, theta_jn, phase, - **waveform_kwargs): - """ Generate a cbc waveform model on specified frequency samples + frequency_array, + mass_1, + mass_2, + luminosity_distance, + a_1, + tilt_1, + phi_12, + a_2, + tilt_2, + lambda_1, + lambda_2, + phi_jl, + theta_jn, + phase, + **waveform_kwargs, +): + """Generate a cbc waveform model on specified frequency samples Parameters ---------- @@ -1105,10 +1533,10 @@ def _base_waveform_frequency_sequence( Dict containing plus and cross modes evaluated at the linear and quadratic frequency nodes. """ - frequencies = waveform_kwargs.pop('frequencies') - reference_frequency = waveform_kwargs.pop('reference_frequency') - approximant = waveform_kwargs.pop('waveform_approximant') - catch_waveform_errors = waveform_kwargs.pop('catch_waveform_errors') + frequencies = waveform_kwargs.pop("frequencies") + reference_frequency = waveform_kwargs.pop("reference_frequency") + approximant = waveform_kwargs.pop("waveform_approximant") + catch_waveform_errors = waveform_kwargs.pop("catch_waveform_errors") waveform_dictionary = set_waveform_dictionary(waveform_kwargs, lambda_1, lambda_2) approximant = lalsim_GetApproximantFromString(approximant) @@ -1117,30 +1545,60 @@ def _base_waveform_frequency_sequence( mass_1 = mass_1 * utils.solar_mass mass_2 = mass_2 * utils.solar_mass - iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = bilby_to_lalsimulation_spins( - theta_jn=theta_jn, phi_jl=phi_jl, tilt_1=tilt_1, tilt_2=tilt_2, - phi_12=phi_12, a_1=a_1, a_2=a_2, mass_1=mass_1, mass_2=mass_2, - reference_frequency=reference_frequency, phase=phase) + iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = ( + bilby_to_lalsimulation_spins( + theta_jn=theta_jn, + phi_jl=phi_jl, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_12=phi_12, + a_1=a_1, + a_2=a_2, + mass_1=mass_1, + mass_2=mass_2, + reference_frequency=reference_frequency, + phase=phase, + ) + ) try: h_plus, h_cross = lalsim_SimInspiralChooseFDWaveformSequence( - phase, mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, - spin_2z, reference_frequency, luminosity_distance, iota, - waveform_dictionary, approximant, frequencies) + phase, + mass_1, + mass_2, + spin_1x, + spin_1y, + spin_1z, + spin_2x, + spin_2y, + spin_2z, + reference_frequency, + luminosity_distance, + iota, + waveform_dictionary, + approximant, + frequencies, + ) except Exception as e: if not catch_waveform_errors: raise else: - EDOM = (e.args[0] == 'Internal function call failed: Input domain error') + EDOM = e.args[0] == "Internal function call failed: Input domain error" if EDOM: - failed_parameters = dict(mass_1=mass_1, mass_2=mass_2, - spin_1=(spin_1x, spin_1y, spin_1z), - spin_2=(spin_2x, spin_2y, spin_2z), - luminosity_distance=luminosity_distance, - iota=iota, phase=phase) - logger.warning("Evaluating the waveform failed with error: {}\n".format(e) + - "The parameters were {}\n".format(failed_parameters) + - "Likelihood will be set to -inf.") + failed_parameters = dict( + mass_1=mass_1, + mass_2=mass_2, + spin_1=(spin_1x, spin_1y, spin_1z), + spin_2=(spin_2x, spin_2y, spin_2z), + luminosity_distance=luminosity_distance, + iota=iota, + phase=phase, + ) + logger.warning( + "Evaluating the waveform failed with error: {}\n".format(e) + + "The parameters were {}\n".format(failed_parameters) + + "Likelihood will be set to -inf." + ) return None else: raise @@ -1196,18 +1654,29 @@ def sinegaussian(frequency_array, hrss, Q, frequency, **kwargs): fm = frequency_array - frequency fp = frequency_array + frequency - negative_term = xp.exp(-fm**2 * np.pi**2 * tau**2) - positive_term = xp.exp(-fp**2 * np.pi**2 * tau**2) + negative_term = xp.exp(-(fm**2) * np.pi**2 * tau**2) + positive_term = xp.exp(-(fp**2) * np.pi**2 * tau**2) - h_plus = hrss * np.pi**0.5 * tau / 2 * ( - negative_term + positive_term - ) / (temp * (1 + xp.exp(-Q**2)))**0.5 + h_plus = ( + hrss + * np.pi**0.5 + * tau + / 2 + * (negative_term + positive_term) + / (temp * (1 + xp.exp(-(Q**2)))) ** 0.5 + ) - h_cross = -1j * hrss * np.pi**0.5 * tau / 2 * ( - negative_term - positive_term - ) / (temp * (1 - xp.exp(-Q**2)))**0.5 + h_cross = ( + -1j + * hrss + * np.pi**0.5 + * tau + / 2 + * (negative_term - positive_term) + / (temp * (1 - xp.exp(-(Q**2)))) ** 0.5 + ) - return {'plus': h_plus, 'cross': h_cross} + return {"plus": h_plus, "cross": h_cross} def supernova(frequency_array, luminosity_distance, **kwargs): @@ -1246,11 +1715,18 @@ def supernova(frequency_array, luminosity_distance, **kwargs): h_plus = scaling * (data[:, 0] + 1j * data[:, 1]) h_cross = scaling * (data[:, 2] + 1j * data[:, 3]) - return {'plus': h_plus, 'cross': h_cross} + return {"plus": h_plus, "cross": h_cross} def supernova_pca_model( - frequency_array, pc_coeff1, pc_coeff2, pc_coeff3, pc_coeff4, pc_coeff5, luminosity_distance, **kwargs + frequency_array, + pc_coeff1, + pc_coeff2, + pc_coeff3, + pc_coeff4, + pc_coeff5, + luminosity_distance, + **kwargs, ): r""" Signal model based on a five-component principal component decomposition @@ -1294,8 +1770,13 @@ def supernova_pca_model( coefficients = [pc_coeff1, pc_coeff2, pc_coeff3, pc_coeff4, pc_coeff5] strain = xp.sum( - xp.asarray([coeff * principal_components[:, ii] for ii, coeff in enumerate(coefficients)]), - axis=0 + xp.asarray( + [ + coeff * principal_components[:, ii] + for ii, coeff in enumerate(coefficients) + ] + ), + axis=0, ) # file at 10kpc @@ -1304,50 +1785,83 @@ def supernova_pca_model( h_plus = scaling * strain h_cross = scaling * strain - return {'plus': h_plus, 'cross': h_cross} + return {"plus": h_plus, "cross": h_cross} precession_only = { - "tilt_1", "tilt_2", "phi_12", "phi_jl", "chi_1_in_plane", "chi_2_in_plane", + "tilt_1", + "tilt_2", + "phi_12", + "phi_jl", + "chi_1_in_plane", + "chi_2_in_plane", } spin = { - "a_1", "a_2", "tilt_1", "tilt_2", "phi_12", "phi_jl", "chi_1", "chi_2", - "chi_1_in_plane", "chi_2_in_plane", + "a_1", + "a_2", + "tilt_1", + "tilt_2", + "phi_12", + "phi_jl", + "chi_1", + "chi_2", + "chi_1_in_plane", + "chi_2_in_plane", } mass = { - "chirp_mass", "mass_ratio", "total_mass", "mass_1", "mass_2", + "chirp_mass", + "mass_ratio", + "total_mass", + "mass_1", + "mass_2", "symmetric_mass_ratio", } -primary_spin_and_q = { - "a_1", "chi_1", "mass_ratio" -} -tidal = { - "lambda_1", "lambda_2", "lambda_tilde", "delta_lambda_tilde" -} +primary_spin_and_q = {"a_1", "chi_1", "mass_ratio"} +tidal = {"lambda_1", "lambda_2", "lambda_tilde", "delta_lambda_tilde"} phase = { - "phase", "delta_phase", + "phase", + "delta_phase", } extrinsic = { - "azimuth", "zenith", "luminosity_distance", "psi", "theta_jn", - "cos_theta_jn", "geocent_time", "time_jitter", "ra", "dec", - "H1_time", "L1_time", "V1_time", + "azimuth", + "zenith", + "luminosity_distance", + "psi", + "theta_jn", + "cos_theta_jn", + "geocent_time", + "time_jitter", + "ra", + "dec", + "H1_time", + "L1_time", + "V1_time", } sky = { - "azimuth", "zenith", "ra", "dec", + "azimuth", + "zenith", + "ra", + "dec", } distance_inclination = { - "luminosity_distance", "redshift", "theta_jn", "cos_theta_jn", -} -measured_spin = { - "chi_1", "chi_2", "a_1", "a_2", "chi_1_in_plane" + "luminosity_distance", + "redshift", + "theta_jn", + "cos_theta_jn", } +measured_spin = {"chi_1", "chi_2", "a_1", "a_2", "chi_1_in_plane"} PARAMETER_SETS = dict( - spin=spin, mass=mass, phase=phase, extrinsic=extrinsic, - tidal=tidal, primary_spin_and_q=primary_spin_and_q, + spin=spin, + mass=mass, + phase=phase, + extrinsic=extrinsic, + tidal=tidal, + primary_spin_and_q=primary_spin_and_q, intrinsic=spin.union(mass).union(phase).union(tidal), precession_only=precession_only, - sky=sky, distance_inclination=distance_inclination, + sky=sky, + distance_inclination=distance_inclination, measured_spin=measured_spin, ) diff --git a/bilby/gw/utils.py b/bilby/gw/utils.py index cf1f27ad5..b760b7a3b 100644 --- a/bilby/gw/utils.py +++ b/bilby/gw/utils.py @@ -9,9 +9,13 @@ from .geometry import zenith_azimuth_to_theta_phi from .time import greenwich_mean_sidereal_time from ..compat.utils import xp_wrap -from ..core.utils import (logger, run_commandline, - check_directory_exists_and_if_not_mkdir, - SamplesSummary, theta_phi_to_ra_dec) +from ..core.utils import ( + logger, + run_commandline, + check_directory_exists_and_if_not_mkdir, + SamplesSummary, + theta_phi_to_ra_dec, +) from ..core.utils.constants import solar_mass @@ -57,7 +61,9 @@ def psd_from_freq_series(freq_data, df, *, xp=None): @xp_wrap -def get_vertex_position_geocentric(position, _longitude=None, _elevation=None, *, xp=None): +def get_vertex_position_geocentric( + position, _longitude=None, _elevation=None, *, xp=None +): """ Calculate the position of the IFO vertex in geocentric coordinates in meters. @@ -76,23 +82,31 @@ def get_vertex_position_geocentric(position, _longitude=None, _elevation=None, * """ if _longitude is not None and _elevation is not None: - logger.warning("Syntax of get_vertex_position_geocentric changed in Bilby 3.0.0." - "Please pass latitude, longitude and elevation as the single array-like argument.") + logger.warning( + "Syntax of get_vertex_position_geocentric changed in Bilby 3.0.0." + "Please pass latitude, longitude and elevation as the single array-like argument." + ) latitude = position longitude = _longitude elevation = _elevation elif len(position) == 3: latitude, longitude, elevation = position else: - raise ValueError("Position must be an array-like of length 3 containing latitude, longitude, and elevation.") + raise ValueError( + "Position must be an array-like of length 3 containing latitude, longitude, and elevation." + ) semi_major_axis = 6378137 # for ellipsoid model of Earth, in m semi_minor_axis = 6356752.314 # in m - radius = semi_major_axis**2 * (semi_major_axis**2 * xp.cos(latitude)**2 + - semi_minor_axis**2 * xp.sin(latitude)**2)**(-0.5) + radius = semi_major_axis**2 * ( + semi_major_axis**2 * xp.cos(latitude) ** 2 + + semi_minor_axis**2 * xp.sin(latitude) ** 2 + ) ** (-0.5) x_comp = (radius + elevation) * xp.cos(latitude) * xp.cos(longitude) y_comp = (radius + elevation) * xp.cos(latitude) * xp.sin(longitude) - z_comp = ((semi_minor_axis / semi_major_axis)**2 * radius + elevation) * xp.sin(latitude) + z_comp = ((semi_minor_axis / semi_major_axis) ** 2 * radius + elevation) * xp.sin( + latitude + ) return xp.asarray([x_comp, y_comp, z_comp]) @@ -117,20 +131,29 @@ def get_vertex_position_ellipsoid(position, *, xp=None): semi_minor_axis = 6356752.314 # in m r = xp.sqrt(position[0] ** 2 + position[1] ** 2) - E = (semi_minor_axis * xp.abs(position[2]) - (semi_major_axis ** 2 - semi_minor_axis ** 2)) / (r * semi_major_axis) - F = (semi_minor_axis * xp.abs(position[2]) + (semi_major_axis ** 2 - semi_minor_axis ** 2)) / (r * semi_major_axis) + E = ( + semi_minor_axis * xp.abs(position[2]) + - (semi_major_axis**2 - semi_minor_axis**2) + ) / (r * semi_major_axis) + F = ( + semi_minor_axis * xp.abs(position[2]) + + (semi_major_axis**2 - semi_minor_axis**2) + ) / (r * semi_major_axis) P = 4 / 3 * (E * F + 1) - Q = 2 * (E ** 2 - F ** 2) - D = P ** 3 + Q ** 2 + Q = 2 * (E**2 - F**2) + D = P**3 + Q**2 v = (xp.sqrt(D) - Q) ** (1 / 3) - (xp.sqrt(D) + Q) ** (1 / 3) # Calculate solution in first quadrant and then adjust based on the sign of the original z_comp - G = 1 / 2 * (xp.sqrt(E ** 2 + v) + E) - t = xp.sqrt(G ** 2 + (F - v * G) / (2 * G - E)) - G - latitude = xp.sign(position[2]) * xp.arctan((semi_major_axis * (1 - t ** 2)) / (2 * semi_minor_axis * t)) + G = 1 / 2 * (xp.sqrt(E**2 + v) + E) + t = xp.sqrt(G**2 + (F - v * G) / (2 * G - E)) - G + latitude = xp.sign(position[2]) * xp.arctan( + (semi_major_axis * (1 - t**2)) / (2 * semi_minor_axis * t) + ) longitude = xp.arctan2(position[1], position[0]) - elevation = (r - semi_major_axis * t) * xp.cos(latitude) \ - + (position[2] - xp.sign(position[2]) * semi_minor_axis) * xp.sin(latitude) + elevation = (r - semi_major_axis * t) * xp.cos(latitude) + ( + position[2] - xp.sign(position[2]) * semi_minor_axis + ) * xp.sin(latitude) return xp.asarray([latitude, longitude, elevation]) @@ -159,7 +182,7 @@ def inner_product(aa, bb, frequency, PSD): df = frequency[1] - frequency[0] integral = integrand.sum() * df - return 4. * integral + return 4.0 * integral def noise_weighted_inner_product(aa, bb, power_spectral_density, duration): @@ -185,7 +208,9 @@ def noise_weighted_inner_product(aa, bb, power_spectral_density, duration): return 4 / duration * integrand.sum() -def matched_filter_snr(signal, frequency_domain_strain, power_spectral_density, duration): +def matched_filter_snr( + signal, frequency_domain_strain, power_spectral_density, duration +): """ Calculate the _complex_ matched filter snr of a signal. This is / optimal_snr @@ -207,11 +232,19 @@ def matched_filter_snr(signal, frequency_domain_strain, power_spectral_density, """ rho_mf = noise_weighted_inner_product( - aa=signal, bb=frequency_domain_strain, - power_spectral_density=power_spectral_density, duration=duration) - rho_mf /= optimal_snr_squared( - signal=signal, power_spectral_density=power_spectral_density, - duration=duration)**0.5 + aa=signal, + bb=frequency_domain_strain, + power_spectral_density=power_spectral_density, + duration=duration, + ) + rho_mf /= ( + optimal_snr_squared( + signal=signal, + power_spectral_density=power_spectral_density, + duration=duration, + ) + ** 0.5 + ) return rho_mf @@ -235,11 +268,21 @@ def optimal_snr_squared(signal, power_spectral_density, duration): float: The square of the optimal matched filter signal to noise ratio """ - return noise_weighted_inner_product(signal, signal, power_spectral_density, duration) + return noise_weighted_inner_product( + signal, signal, power_spectral_density, duration + ) -def overlap(signal_a, signal_b, power_spectral_density=None, delta_frequency=None, - lower_cut_off=None, upper_cut_off=None, norm_a=None, norm_b=None): +def overlap( + signal_a, + signal_b, + power_spectral_density=None, + delta_frequency=None, + lower_cut_off=None, + upper_cut_off=None, + norm_a=None, + norm_b=None, +): r""" Compute the overlap between two signals @@ -271,7 +314,9 @@ def overlap(signal_a, signal_b, power_spectral_density=None, delta_frequency=Non low_index = int(lower_cut_off / delta_frequency) up_index = int(upper_cut_off / delta_frequency) integrand = signal_a.conj() * signal_b - integrand = integrand[low_index:up_index] / power_spectral_density[low_index:up_index] + integrand = ( + integrand[low_index:up_index] / power_spectral_density[low_index:up_index] + ) integral = (4 * delta_frequency * integrand) / norm_a / norm_b return sum(integral).real @@ -333,8 +378,9 @@ def get_event_time(event): def get_open_strain_data( - name, start_time, end_time, outdir, cache=False, buffer_time=0, **kwargs): - """ A function which accesses the open strain data + name, start_time, end_time, outdir, cache=False, buffer_time=0, **kwargs +): + """A function which accesses the open strain data This uses `gwpy` to download the open data and then saves a cached copy for later use @@ -362,7 +408,8 @@ def get_open_strain_data( """ from gwpy.timeseries import TimeSeries - filename = '{}/{}_{}_{}.txt'.format(outdir, name, start_time, end_time) + + filename = "{}/{}_{}_{}.txt".format(outdir, name, start_time, end_time) if buffer_time < 0: raise ValueError("buffer_time < 0") @@ -370,26 +417,40 @@ def get_open_strain_data( end_time = end_time + buffer_time if os.path.isfile(filename) and cache: - logger.info('Using cached data from {}'.format(filename)) + logger.info("Using cached data from {}".format(filename)) strain = TimeSeries.read(filename) else: - logger.info('Fetching open data from {} to {} with buffer time {}' - .format(start_time, end_time, buffer_time)) + logger.info( + "Fetching open data from {} to {} with buffer time {}".format( + start_time, end_time, buffer_time + ) + ) try: strain = TimeSeries.fetch_open_data(name, start_time, end_time, **kwargs) - logger.info('Saving cache of data to {}'.format(filename)) + logger.info("Saving cache of data to {}".format(filename)) strain.write(filename) except Exception as e: logger.info("Unable to fetch open data, see debug for detailed info") - logger.info("Call to gwpy.timeseries.TimeSeries.fetch_open_data returned {}" - .format(e)) + logger.info( + "Call to gwpy.timeseries.TimeSeries.fetch_open_data returned {}".format( + e + ) + ) strain = None return strain -def read_frame_file(file_name, start_time, end_time, resample=None, channel=None, buffer_time=0, **kwargs): - """ A function which accesses the open strain data +def read_frame_file( + file_name, + start_time, + end_time, + resample=None, + channel=None, + buffer_time=0, + **kwargs, +): + """A function which accesses the open strain data This uses `gwpy` to download the open data and then saves a cached copy for later use @@ -416,34 +477,61 @@ def read_frame_file(file_name, start_time, end_time, resample=None, channel=None """ from gwpy.timeseries import TimeSeries + loaded = False strain = None if channel is not None: try: - strain = TimeSeries.read(source=file_name, channel=channel, start=start_time, end=end_time, **kwargs) + strain = TimeSeries.read( + source=file_name, + channel=channel, + start=start_time, + end=end_time, + **kwargs, + ) loaded = True - logger.info('Successfully loaded {}.'.format(channel)) + logger.info("Successfully loaded {}.".format(channel)) except (RuntimeError, ValueError): - logger.warning('Channel {} not found. Trying preset channel names'.format(channel)) + logger.warning( + "Channel {} not found. Trying preset channel names".format(channel) + ) if loaded is False: - ligo_channel_types = ['GDS-CALIB_STRAIN', 'DCS-CALIB_STRAIN_C01', 'DCS-CALIB_STRAIN_C02', - 'DCH-CLEAN_STRAIN_C02', 'GWOSC-16KHZ_R1_STRAIN', - 'GWOSC-4KHZ_R1_STRAIN'] - virgo_channel_types = ['Hrec_hoft_V1O2Repro2A_16384Hz', 'FAKE_h_16384Hz_4R', - 'GWOSC-16KHZ_R1_STRAIN', 'GWOSC-4KHZ_R1_STRAIN'] - channel_types = dict(H1=ligo_channel_types, L1=ligo_channel_types, V1=virgo_channel_types) + ligo_channel_types = [ + "GDS-CALIB_STRAIN", + "DCS-CALIB_STRAIN_C01", + "DCS-CALIB_STRAIN_C02", + "DCH-CLEAN_STRAIN_C02", + "GWOSC-16KHZ_R1_STRAIN", + "GWOSC-4KHZ_R1_STRAIN", + ] + virgo_channel_types = [ + "Hrec_hoft_V1O2Repro2A_16384Hz", + "FAKE_h_16384Hz_4R", + "GWOSC-16KHZ_R1_STRAIN", + "GWOSC-4KHZ_R1_STRAIN", + ] + channel_types = dict( + H1=ligo_channel_types, L1=ligo_channel_types, V1=virgo_channel_types + ) for detector in channel_types.keys(): for channel_type in channel_types[detector]: if loaded: break - channel = '{}:{}'.format(detector, channel_type) + channel = "{}:{}".format(detector, channel_type) try: - strain = TimeSeries.read(source=file_name, channel=channel, start=start_time, end=end_time, - **kwargs) + strain = TimeSeries.read( + source=file_name, + channel=channel, + start=start_time, + end=end_time, + **kwargs, + ) loaded = True - logger.info('Successfully read strain data for channel {}.'.format(channel)) + logger.info( + "Successfully read strain data for channel {}.".format(channel) + ) except (RuntimeError, ValueError): pass @@ -452,11 +540,13 @@ def read_frame_file(file_name, start_time, end_time, resample=None, channel=None strain = strain.resample(resample) return strain else: - logger.warning('No data loaded.') + logger.warning("No data loaded.") return None -def get_gracedb(gracedb, outdir, duration, calibration, detectors, query_types=None, server=None): +def get_gracedb( + gracedb, outdir, duration, calibration, detectors, query_types=None, server=None +): """ Download information about a trigger from GraceDb and create cache files for finding gravitational-wave strain data. @@ -486,21 +576,29 @@ def get_gracedb(gracedb, outdir, duration, calibration, detectors, query_types=N List of cache filenames, one per interferometer. """ candidate = gracedb_to_json(gracedb, outdir=outdir) - trigger_time = candidate['gpstime'] + trigger_time = candidate["gpstime"] gps_start_time = trigger_time - duration cache_files = [] if query_types is None: query_types = [None] * len(detectors) for i, det in enumerate(detectors): output_cache_file = gw_data_find( - det, gps_start_time, duration, calibration, - outdir=outdir, query_type=query_types[i], server=server) + det, + gps_start_time, + duration, + calibration, + outdir=outdir, + query_type=query_types[i], + server=server, + ) cache_files.append(output_cache_file) return candidate, cache_files -def gracedb_to_json(gracedb, cred=None, service_url='https://gracedb.ligo.org/api/', outdir=None): - """ Script to download a GraceDB candidate +def gracedb_to_json( + gracedb, cred=None, service_url="https://gracedb.ligo.org/api/", outdir=None +): + """Script to download a GraceDB candidate Parameters ========== @@ -515,21 +613,21 @@ def gracedb_to_json(gracedb, cred=None, service_url='https://gracedb.ligo.org/ap outdir: str, optional If given, a string identfying the location in which to store the json """ - logger.info( - 'Starting routine to download GraceDb candidate {}'.format(gracedb)) + logger.info("Starting routine to download GraceDb candidate {}".format(gracedb)) from ligo.gracedb.rest import GraceDb - logger.info('Initialise client and attempt to download') - logger.info('Fetching from {}'.format(service_url)) + logger.info("Initialise client and attempt to download") + logger.info("Fetching from {}".format(service_url)) try: client = GraceDb(cred=cred, service_url=service_url) except IOError: raise ValueError( - 'Failed to authenticate with gracedb: check your X509 ' - 'certificate is accessible and valid') + "Failed to authenticate with gracedb: check your X509 " + "certificate is accessible and valid" + ) try: candidate = client.event(gracedb) - logger.info('Successfully downloaded candidate') + logger.info("Successfully downloaded candidate") except Exception as e: raise ValueError("Unable to obtain GraceDB candidate, exception: {}".format(e)) @@ -537,17 +635,24 @@ def gracedb_to_json(gracedb, cred=None, service_url='https://gracedb.ligo.org/ap if outdir is not None: check_directory_exists_and_if_not_mkdir(outdir) - outfilepath = os.path.join(outdir, '{}.json'.format(gracedb)) - logger.info('Writing candidate to {}'.format(outfilepath)) - with open(outfilepath, 'w') as outfile: + outfilepath = os.path.join(outdir, "{}.json".format(gracedb)) + logger.info("Writing candidate to {}".format(outfilepath)) + with open(outfilepath, "w") as outfile: json.dump(json_output, outfile, indent=2) return json_output -def gw_data_find(observatory, gps_start_time, duration, calibration, - outdir='.', query_type=None, server=None): - """ Builds a gw_data_find call and process output +def gw_data_find( + observatory, + gps_start_time, + duration, + calibration, + outdir=".", + query_type=None, + server=None, +): + """Builds a gw_data_find call and process output Parameters ========== @@ -570,49 +675,52 @@ def gw_data_find(observatory, gps_start_time, duration, calibration, Path to the output cache file """ - logger.info('Building gw_data_find command line') + logger.info("Building gw_data_find command line") - observatory_lookup = dict(H1='H', L1='L', V1='V') + observatory_lookup = dict(H1="H", L1="L", V1="V") observatory_code = observatory_lookup[observatory] if query_type is None: - logger.warning('No query type provided. This may prevent data from being read.') - if observatory_code == 'V': - query_type = 'V1Online' + logger.warning("No query type provided. This may prevent data from being read.") + if observatory_code == "V": + query_type = "V1Online" else: - query_type = '{}_HOFT_C0{}'.format(observatory, calibration) + query_type = "{}_HOFT_C0{}".format(observatory, calibration) - logger.info('Using LDRDataFind query type {}'.format(query_type)) + logger.info("Using LDRDataFind query type {}".format(query_type)) - cache_file = '{}-{}_CACHE-{}-{}.lcf'.format( - observatory, query_type, gps_start_time, duration) + cache_file = "{}-{}_CACHE-{}-{}.lcf".format( + observatory, query_type, gps_start_time, duration + ) output_cache_file = os.path.join(outdir, cache_file) gps_end_time = gps_start_time + duration if server is None: server = os.environ.get("LIGO_DATAFIND_SERVER") if server is None: - logger.warning('No data_find server found, defaulting to CIT server. ' - 'To run on other clusters, pass the output of ' - '`echo $LIGO_DATAFIND_SERVER`') - server = 'ldr.ldas.cit:80' - - cl_list = ['gw_data_find'] - cl_list.append('--observatory {}'.format(observatory_code)) - cl_list.append('--gps-start-time {}'.format(int(np.floor(gps_start_time)))) - cl_list.append('--gps-end-time {}'.format(int(np.ceil(gps_end_time)))) - cl_list.append('--type {}'.format(query_type)) - cl_list.append('--output {}'.format(output_cache_file)) - cl_list.append('--server {}'.format(server)) - cl_list.append('--url-type file') - cl_list.append('--lal-cache') - cl = ' '.join(cl_list) + logger.warning( + "No data_find server found, defaulting to CIT server. " + "To run on other clusters, pass the output of " + "`echo $LIGO_DATAFIND_SERVER`" + ) + server = "ldr.ldas.cit:80" + + cl_list = ["gw_data_find"] + cl_list.append("--observatory {}".format(observatory_code)) + cl_list.append("--gps-start-time {}".format(int(np.floor(gps_start_time)))) + cl_list.append("--gps-end-time {}".format(int(np.ceil(gps_end_time)))) + cl_list.append("--type {}".format(query_type)) + cl_list.append("--output {}".format(output_cache_file)) + cl_list.append("--server {}".format(server)) + cl_list.append("--url-type file") + cl_list.append("--lal-cache") + cl = " ".join(cl_list) run_commandline(cl) return output_cache_file def convert_args_list_to_float(*args_list): - """ Converts inputs to floats, returns a list in the same order as the input""" + """Converts inputs to floats, returns a list in the same order as the input""" try: args_list = [float(arg) for arg in args_list] except ValueError: @@ -621,13 +729,33 @@ def convert_args_list_to_float(*args_list): def lalsim_SimInspiralTransformPrecessingNewInitialConditions( - theta_jn, phi_jl, tilt_1, tilt_2, phi_12, a_1, a_2, mass_1, mass_2, - reference_frequency, phase): + theta_jn, + phi_jl, + tilt_1, + tilt_2, + phi_12, + a_1, + a_2, + mass_1, + mass_2, + reference_frequency, + phase, +): from lalsimulation import SimInspiralTransformPrecessingNewInitialConditions args_list = convert_args_list_to_float( - theta_jn, phi_jl, tilt_1, tilt_2, phi_12, a_1, a_2, mass_1, mass_2, - reference_frequency, phase) + theta_jn, + phi_jl, + tilt_1, + tilt_2, + phi_12, + a_1, + a_2, + mass_1, + mass_2, + reference_frequency, + phase, + ) return SimInspiralTransformPrecessingNewInitialConditions(*args_list) @@ -635,6 +763,7 @@ def lalsim_SimInspiralTransformPrecessingNewInitialConditions( @lru_cache(maxsize=10) def lalsim_GetApproximantFromString(waveform_approximant): from lalsimulation import GetApproximantFromString + if isinstance(waveform_approximant, str): return GetApproximantFromString(waveform_approximant) else: @@ -642,11 +771,27 @@ def lalsim_GetApproximantFromString(waveform_approximant): def lalsim_SimInspiralFD( - mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, - spin_2z, luminosity_distance, iota, phase, - longitude_ascending_nodes, eccentricity, mean_per_ano, delta_frequency, - minimum_frequency, maximum_frequency, reference_frequency, - waveform_dictionary, approximant): + mass_1, + mass_2, + spin_1x, + spin_1y, + spin_1z, + spin_2x, + spin_2y, + spin_2z, + luminosity_distance, + iota, + phase, + longitude_ascending_nodes, + eccentricity, + mean_per_ano, + delta_frequency, + minimum_frequency, + maximum_frequency, + reference_frequency, + waveform_dictionary, + approximant, +): """ Safely call lalsimulation.SimInspiralFD @@ -676,10 +821,25 @@ def lalsim_SimInspiralFD( from lalsimulation import SimInspiralFD args = convert_args_list_to_float( - mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z, - luminosity_distance, iota, phase, longitude_ascending_nodes, - eccentricity, mean_per_ano, delta_frequency, minimum_frequency, - maximum_frequency, reference_frequency) + mass_1, + mass_2, + spin_1x, + spin_1y, + spin_1z, + spin_2x, + spin_2y, + spin_2z, + luminosity_distance, + iota, + phase, + longitude_ascending_nodes, + eccentricity, + mean_per_ano, + delta_frequency, + minimum_frequency, + maximum_frequency, + reference_frequency, + ) approximant = _get_lalsim_approximant(approximant) @@ -687,11 +847,27 @@ def lalsim_SimInspiralFD( def lalsim_SimInspiralChooseFDWaveform( - mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, - spin_2z, luminosity_distance, iota, phase, - longitude_ascending_nodes, eccentricity, mean_per_ano, delta_frequency, - minimum_frequency, maximum_frequency, reference_frequency, - waveform_dictionary, approximant): + mass_1, + mass_2, + spin_1x, + spin_1y, + spin_1z, + spin_2x, + spin_2y, + spin_2z, + luminosity_distance, + iota, + phase, + longitude_ascending_nodes, + eccentricity, + mean_per_ano, + delta_frequency, + minimum_frequency, + maximum_frequency, + reference_frequency, + waveform_dictionary, + approximant, +): """ Safely call lalsimulation.SimInspiralChooseFDWaveform @@ -721,10 +897,25 @@ def lalsim_SimInspiralChooseFDWaveform( from lalsimulation import SimInspiralChooseFDWaveform args = convert_args_list_to_float( - mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z, - luminosity_distance, iota, phase, longitude_ascending_nodes, - eccentricity, mean_per_ano, delta_frequency, minimum_frequency, - maximum_frequency, reference_frequency) + mass_1, + mass_2, + spin_1x, + spin_1y, + spin_1z, + spin_2x, + spin_2y, + spin_2z, + luminosity_distance, + iota, + phase, + longitude_ascending_nodes, + eccentricity, + mean_per_ano, + delta_frequency, + minimum_frequency, + maximum_frequency, + reference_frequency, + ) approximant = _get_lalsim_approximant(approximant) @@ -743,9 +934,22 @@ def _get_lalsim_approximant(approximant): def lalsim_SimInspiralChooseFDWaveformSequence( - phase, mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, - spin_2z, reference_frequency, luminosity_distance, iota, - waveform_dictionary, approximant, frequency_array): + phase, + mass_1, + mass_2, + spin_1x, + spin_1y, + spin_1z, + spin_2x, + spin_2y, + spin_2z, + reference_frequency, + luminosity_distance, + iota, + waveform_dictionary, + approximant, + frequency_array, +): """ Safely call lalsimulation.SimInspiralChooseFDWaveformSequence @@ -770,10 +974,33 @@ def lalsim_SimInspiralChooseFDWaveformSequence( from lal import REAL8Vector, CreateREAL8Vector from lalsimulation import SimInspiralChooseFDWaveformSequence - [mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z, - luminosity_distance, iota, phase, reference_frequency] = convert_args_list_to_float( - mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z, - luminosity_distance, iota, phase, reference_frequency) + [ + mass_1, + mass_2, + spin_1x, + spin_1y, + spin_1z, + spin_2x, + spin_2y, + spin_2z, + luminosity_distance, + iota, + phase, + reference_frequency, + ] = convert_args_list_to_float( + mass_1, + mass_2, + spin_1x, + spin_1y, + spin_1z, + spin_2x, + spin_2y, + spin_2z, + luminosity_distance, + iota, + phase, + reference_frequency, + ) if isinstance(approximant, int): pass @@ -788,37 +1015,49 @@ def lalsim_SimInspiralChooseFDWaveformSequence( frequency_array.data = old_frequency_array return SimInspiralChooseFDWaveformSequence( - phase, mass_1, mass_2, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, - spin_2z, reference_frequency, luminosity_distance, iota, - waveform_dictionary, approximant, frequency_array) + phase, + mass_1, + mass_2, + spin_1x, + spin_1y, + spin_1z, + spin_2x, + spin_2y, + spin_2z, + reference_frequency, + luminosity_distance, + iota, + waveform_dictionary, + approximant, + frequency_array, + ) -def lalsim_SimInspiralWaveformParamsInsertTidalLambda1( - waveform_dictionary, lambda_1): +def lalsim_SimInspiralWaveformParamsInsertTidalLambda1(waveform_dictionary, lambda_1): from lalsimulation import SimInspiralWaveformParamsInsertTidalLambda1 + try: lambda_1 = float(lambda_1) except ValueError: raise ValueError("Unable to convert lambda_1 to float") - return SimInspiralWaveformParamsInsertTidalLambda1( - waveform_dictionary, lambda_1) + return SimInspiralWaveformParamsInsertTidalLambda1(waveform_dictionary, lambda_1) -def lalsim_SimInspiralWaveformParamsInsertTidalLambda2( - waveform_dictionary, lambda_2): +def lalsim_SimInspiralWaveformParamsInsertTidalLambda2(waveform_dictionary, lambda_2): from lalsimulation import SimInspiralWaveformParamsInsertTidalLambda2 + try: lambda_2 = float(lambda_2) except ValueError: raise ValueError("Unable to convert lambda_2 to float") - return SimInspiralWaveformParamsInsertTidalLambda2( - waveform_dictionary, lambda_2) + return SimInspiralWaveformParamsInsertTidalLambda2(waveform_dictionary, lambda_2) def lalsim_SimNeutronStarEOS4ParamSDGammaCheck(g0, g1, g2, g3): from lalsimulation import SimNeutronStarEOS4ParamSDGammaCheck + try: g0 = float(g0) g1 = float(g1) @@ -834,6 +1073,7 @@ def lalsim_SimNeutronStarEOS4ParamSDGammaCheck(g0, g1, g2, g3): def lalsim_SimNeutronStarEOS4ParameterSpectralDecomposition(g0, g1, g2, g3): from lalsimulation import SimNeutronStarEOS4ParameterSpectralDecomposition + try: g0 = float(g0) g1 = float(g1) @@ -849,6 +1089,7 @@ def lalsim_SimNeutronStarEOS4ParameterSpectralDecomposition(g0, g1, g2, g3): def lalsim_SimNeutronStarEOS4ParamSDViableFamilyCheck(g0, g1, g2, g3): from lalsimulation import SimNeutronStarEOS4ParamSDViableFamilyCheck + try: g0 = float(g0) g1 = float(g1) @@ -864,6 +1105,7 @@ def lalsim_SimNeutronStarEOS4ParamSDViableFamilyCheck(g0, g1, g2, g3): def lalsim_SimNeutronStarEOS3PieceDynamicPolytrope(g0, log10p1_si, g1, log10p2_si, g2): from lalsimulation import SimNeutronStarEOS3PieceDynamicPolytrope + try: g0 = float(g0) g1 = float(g1) @@ -880,6 +1122,7 @@ def lalsim_SimNeutronStarEOS3PieceDynamicPolytrope(g0, log10p1_si, g1, log10p2_s def lalsim_SimNeutronStarEOS3PieceCausalAnalytic(v1, log10p1_si, v2, log10p2_si, v3): from lalsimulation import SimNeutronStarEOS3PieceCausalAnalytic + try: v1 = float(v1) v2 = float(v2) @@ -894,8 +1137,11 @@ def lalsim_SimNeutronStarEOS3PieceCausalAnalytic(v1, log10p1_si, v2, log10p2_si, return SimNeutronStarEOS3PieceCausalAnalytic(v1, log10p1_si, v2, log10p2_si, v3) -def lalsim_SimNeutronStarEOS3PDViableFamilyCheck(p0, log10p1_si, p1, log10p2_si, p2, causal): +def lalsim_SimNeutronStarEOS3PDViableFamilyCheck( + p0, log10p1_si, p1, log10p2_si, p2, causal +): from lalsimulation import SimNeutronStarEOS3PDViableFamilyCheck + try: p0 = float(p0) p1 = float(p1) @@ -908,7 +1154,9 @@ def lalsim_SimNeutronStarEOS3PDViableFamilyCheck(p0, log10p1_si, p1, log10p2_si, except TypeError: raise TypeError("Unable to convert EOS parameters to floats or int") - return SimNeutronStarEOS3PDViableFamilyCheck(p0, log10p1_si, p1, log10p2_si, p2, causal) + return SimNeutronStarEOS3PDViableFamilyCheck( + p0, log10p1_si, p1, log10p2_si, p2, causal + ) def lalsim_CreateSimNeutronStarFamily(eos): @@ -925,6 +1173,7 @@ def lalsim_SimNeutronStarEOSMaxPseudoEnthalpy(eos): def lalsim_SimNeutronStarEOSSpeedOfSoundGeometerized(max_pseudo_enthalpy, eos): from lalsimulation import SimNeutronStarEOSSpeedOfSoundGeometerized + try: max_pseudo_enthalpy = float(max_pseudo_enthalpy) except ValueError: @@ -949,6 +1198,7 @@ def lalsim_SimNeutronStarMaximumMass(fam): def lalsim_SimNeutronStarRadius(mass_in_SI, fam): from lalsimulation import SimNeutronStarRadius + try: mass_in_SI = float(mass_in_SI) except ValueError: @@ -961,6 +1211,7 @@ def lalsim_SimNeutronStarRadius(mass_in_SI, fam): def lalsim_SimNeutronStarLoveNumberK2(mass_in_SI, fam): from lalsimulation import SimNeutronStarLoveNumberK2 + try: mass_in_SI = float(mass_in_SI) except ValueError: @@ -992,7 +1243,9 @@ def spline_angle_xform(delta_psi, *, xp=None): return 180.0 / np.pi * xp.arctan2(xp.imag(rotation), xp.real(rotation)) -def plot_spline_pos(log_freqs, samples, nfreqs=100, level=0.9, color='k', label=None, xform=None): +def plot_spline_pos( + log_freqs, samples, nfreqs=100, level=0.9, color="k", label=None, xform=None +): """ Plot calibration posterior estimates for a spline model in log space. Adapted from the same function in lalinference.bayespputils @@ -1016,6 +1269,7 @@ def plot_spline_pos(log_freqs, samples, nfreqs=100, level=0.9, color='k', label= """ import matplotlib.pyplot as plt + freq_points = np.exp(log_freqs) freqs = np.logspace(min(log_freqs), max(log_freqs), nfreqs, base=np.exp(1)) @@ -1026,28 +1280,42 @@ def plot_spline_pos(log_freqs, samples, nfreqs=100, level=0.9, color='k', label= else: scaled_samples = xform(samples) - scaled_samples_summary = SamplesSummary(scaled_samples, average='mean') - data_summary = SamplesSummary(data, average='mean') - - plt.errorbar(freq_points, scaled_samples_summary.average, - yerr=[-scaled_samples_summary.lower_relative_credible_interval, - scaled_samples_summary.upper_relative_credible_interval], - fmt='.', color=color, lw=4, alpha=0.5, capsize=0) + scaled_samples_summary = SamplesSummary(scaled_samples, average="mean") + data_summary = SamplesSummary(data, average="mean") + + plt.errorbar( + freq_points, + scaled_samples_summary.average, + yerr=[ + -scaled_samples_summary.lower_relative_credible_interval, + scaled_samples_summary.upper_relative_credible_interval, + ], + fmt=".", + color=color, + lw=4, + alpha=0.5, + capsize=0, + ) for i, sample in enumerate(samples): temp = interp1d( - log_freqs, sample, kind="cubic", fill_value=0, - bounds_error=False)(np.log(freqs)) + log_freqs, sample, kind="cubic", fill_value=0, bounds_error=False + )(np.log(freqs)) if xform is None: data[i] = temp else: data[i] = xform(temp) plt.plot(freqs, np.mean(data, axis=0), color=color, label=label) - plt.fill_between(freqs, data_summary.lower_absolute_credible_interval, - data_summary.upper_absolute_credible_interval, - color=color, alpha=.1, linewidth=0.1) - plt.xlim(freq_points.min() - .5, freq_points.max() + 50) + plt.fill_between( + freqs, + data_summary.lower_absolute_credible_interval, + data_summary.upper_absolute_credible_interval, + color=color, + alpha=0.1, + linewidth=0.1, + ) + plt.xlim(freq_points.min() - 0.5, freq_points.max() + 50) @xp_wrap @@ -1070,7 +1338,7 @@ def ln_i0(value, *, xp=None): def calculate_time_to_merger(frequency, mass_1, mass_2, chi=0, safety=1.1): - """ Leading-order calculation of the time to merger from frequency + """Leading-order calculation of the time to merger from frequency This uses the XLALSimInspiralTaylorF2ReducedSpinChirpTime routine to estimate the time to merger. Based on the implementation in @@ -1094,12 +1362,13 @@ def calculate_time_to_merger(frequency, mass_1, mass_2, chi=0, safety=1.1): """ import lalsimulation + return safety * lalsimulation.SimInspiralTaylorF2ReducedSpinChirpTime( float(frequency), float(mass_1 * solar_mass), float(mass_2 * solar_mass), float(chi), - -1 + -1, ) diff --git a/bilby/gw/waveform_generator.py b/bilby/gw/waveform_generator.py index 3b961f34c..2460d825a 100644 --- a/bilby/gw/waveform_generator.py +++ b/bilby/gw/waveform_generator.py @@ -16,17 +16,26 @@ class WaveformGenerator(object): Waveform generators provide a unified method to call disparate source models. """ - duration = PropertyAccessor('_times_and_frequencies', 'duration') - sampling_frequency = PropertyAccessor('_times_and_frequencies', 'sampling_frequency') - start_time = PropertyAccessor('_times_and_frequencies', 'start_time') - frequency_array = PropertyAccessor('_times_and_frequencies', 'frequency_array') - time_array = PropertyAccessor('_times_and_frequencies', 'time_array') - - def __init__(self, duration=None, sampling_frequency=None, start_time=0, frequency_domain_source_model=None, - time_domain_source_model=None, parameters=None, - parameter_conversion=None, - waveform_arguments=None, use_cache=True, - ): + duration = PropertyAccessor("_times_and_frequencies", "duration") + sampling_frequency = PropertyAccessor( + "_times_and_frequencies", "sampling_frequency" + ) + start_time = PropertyAccessor("_times_and_frequencies", "start_time") + frequency_array = PropertyAccessor("_times_and_frequencies", "frequency_array") + time_array = PropertyAccessor("_times_and_frequencies", "time_array") + + def __init__( + self, + duration=None, + sampling_frequency=None, + start_time=0, + frequency_domain_source_model=None, + time_domain_source_model=None, + parameters=None, + parameter_conversion=None, + waveform_arguments=None, + use_cache=True, + ): """ The base waveform generator class. @@ -66,9 +75,11 @@ def __init__(self, duration=None, sampling_frequency=None, start_time=0, frequen with :code:`JAX`. """ - self._times_and_frequencies = CoupledTimeAndFrequencySeries(duration=duration, - sampling_frequency=sampling_frequency, - start_time=start_time) + self._times_and_frequencies = CoupledTimeAndFrequencySeries( + duration=duration, + sampling_frequency=sampling_frequency, + start_time=start_time, + ) self.frequency_domain_source_model = frequency_domain_source_model self.time_domain_source_model = time_domain_source_model self.source_parameter_keys = self._parameters_from_source_model() @@ -103,15 +114,24 @@ def __repr__(self): else: param_conv_name = utils.get_function_path(self.parameter_conversion) - return self.__class__.__name__ + '(duration={}, sampling_frequency={}, start_time={}, ' \ - 'frequency_domain_source_model={}, time_domain_source_model={}, ' \ - 'parameter_conversion={}, ' \ - 'waveform_arguments={})'\ - .format(self.duration, self.sampling_frequency, self.start_time, fdsm_name, tdsm_name, - param_conv_name, self.waveform_arguments) + return ( + self.__class__.__name__ + + "(duration={}, sampling_frequency={}, start_time={}, " + "frequency_domain_source_model={}, time_domain_source_model={}, " + "parameter_conversion={}, " + "waveform_arguments={})".format( + self.duration, + self.sampling_frequency, + self.start_time, + fdsm_name, + tdsm_name, + param_conv_name, + self.waveform_arguments, + ) + ) def frequency_domain_strain(self, parameters=None, *, xp=None): - """ Wrapper to source_model. + """Wrapper to source_model. Converts parameters with self.parameter_conversion before handing it off to the source model. Automatically refers to the time_domain_source model via NFFT if no frequency_domain_source_model is given. @@ -132,16 +152,18 @@ def frequency_domain_strain(self, parameters=None, *, xp=None): RuntimeError: If no source model is given """ - return self._calculate_strain(model=self.frequency_domain_source_model, - model_data_points=self.frequency_array, - parameters=parameters, - transformation_function=utils.nfft, - transformed_model=self.time_domain_source_model, - transformed_model_data_points=self.time_array, - xp=xp) + return self._calculate_strain( + model=self.frequency_domain_source_model, + model_data_points=self.frequency_array, + parameters=parameters, + transformation_function=utils.nfft, + transformed_model=self.time_domain_source_model, + transformed_model_data_points=self.time_array, + xp=xp, + ) def time_domain_strain(self, parameters=None, *, xp=None): - """ Wrapper to source_model. + """Wrapper to source_model. Converts parameters with self.parameter_conversion before handing it off to the source model. Automatically refers to the frequency_domain_source model via INFFT if no frequency_domain_source_model is @@ -167,34 +189,47 @@ def time_domain_strain(self, parameters=None, *, xp=None): RuntimeError: If no source model is given """ - return self._calculate_strain(model=self.time_domain_source_model, - model_data_points=self.time_array, - parameters=parameters, - transformation_function=utils.infft, - transformed_model=self.frequency_domain_source_model, - transformed_model_data_points=self.frequency_array, - xp=xp) - - def _calculate_strain(self, model, model_data_points, transformation_function, transformed_model, - transformed_model_data_points, parameters, *, xp=None): + return self._calculate_strain( + model=self.time_domain_source_model, + model_data_points=self.time_array, + parameters=parameters, + transformation_function=utils.infft, + transformed_model=self.frequency_domain_source_model, + transformed_model_data_points=self.frequency_array, + xp=xp, + ) + + def _calculate_strain( + self, + model, + model_data_points, + transformation_function, + transformed_model, + transformed_model_data_points, + parameters, + *, + xp=None, + ): if parameters is None: - parameters = self._cache.get('parameters', None) + parameters = self._cache.get("parameters", None) if parameters is None: raise ValueError("No parameters given to generate waveform.") if ( self.use_cache - and parameters == self._cache.get('parameters', None) - and self._cache['model'] == model - and self._cache['transformed_model'] == transformed_model + and parameters == self._cache.get("parameters", None) + and self._cache["model"] == model + and self._cache["transformed_model"] == transformed_model ): - return self._cache['waveform'] + return self._cache["waveform"] elif self.use_cache: - self._cache['parameters'] = parameters.copy() - self._cache['model'] = model - self._cache['transformed_model'] = transformed_model + self._cache["parameters"] = parameters.copy() + self._cache["model"] = model + self._cache["transformed_model"] = transformed_model parameters = self._format_parameters(parameters) if model is not None: - model_strain = self._strain_from_model(model_data_points, model, parameters, xp=xp) + model_strain = self._strain_from_model( + model_data_points, model, parameters, xp=xp + ) elif transformed_model is not None: model_strain = self._strain_from_transformed_model( transformed_model_data_points, @@ -206,7 +241,7 @@ def _calculate_strain(self, model, model_data_points, transformation_function, t else: raise RuntimeError("No source model given") if self.use_cache: - self._cache['waveform'] = model_strain + self._cache["waveform"] = model_strain return model_strain def _strain_from_model(self, model_data_points, model, parameters, *, xp=None): @@ -228,20 +263,25 @@ def _strain_from_transformed_model( ) if aac.is_array_api_obj(transformed_model_strain): - return transformation_function(transformed_model_strain, self.sampling_frequency) + return transformation_function( + transformed_model_strain, self.sampling_frequency + ) model_strain = dict() for key in transformed_model_strain: if transformation_function == utils.nfft: - model_strain[key], _ = \ - transformation_function(transformed_model_strain[key], self.sampling_frequency) + model_strain[key], _ = transformation_function( + transformed_model_strain[key], self.sampling_frequency + ) else: - model_strain[key] = transformation_function(transformed_model_strain[key], self.sampling_frequency) + model_strain[key] = transformation_function( + transformed_model_strain[key], self.sampling_frequency + ) return model_strain @property def parameters(self): - """ The dictionary of parameters for source model. + """The dictionary of parameters for source model. Returns ======= @@ -263,8 +303,7 @@ def _format_parameters(self, parameters): raise TypeError('"parameters" must be a dictionary.') new_parameters = parameters.copy() new_parameters, _ = self.parameter_conversion(new_parameters) - for key in self.source_parameter_keys.symmetric_difference( - new_parameters): + for key in self.source_parameter_keys.symmetric_difference(new_parameters): new_parameters.pop(key) new_parameters.update(self.waveform_arguments) return new_parameters @@ -282,8 +321,9 @@ def _parameters_from_source_model(self): elif self.time_domain_source_model is not None: model = self.time_domain_source_model else: - raise AttributeError('Either time or frequency domain source ' - 'model must be provided.') + raise AttributeError( + "Either time or frequency domain source model must be provided." + ) return set(utils.infer_parameters_from_function(model)) @property @@ -302,7 +342,8 @@ def meta_data(self): class LALCBCWaveformGenerator(WaveformGenerator): - """ A waveform generator with specific checks for LAL CBC waveforms """ + """A waveform generator with specific checks for LAL CBC waveforms""" + LAL_SIM_INSPIRAL_SPINS_FLOW = 1 def __init__(self, **kwargs): @@ -311,11 +352,22 @@ def __init__(self, **kwargs): def validate_reference_frequency(self): from lalsimulation import SimInspiralGetSpinFreqFromApproximant + waveform_approximant = self.waveform_arguments["waveform_approximant"] - waveform_approximant_number = lalsim_GetApproximantFromString(waveform_approximant) - if SimInspiralGetSpinFreqFromApproximant(waveform_approximant_number) == self.LAL_SIM_INSPIRAL_SPINS_FLOW: - if self.waveform_arguments["reference_frequency"] != self.waveform_arguments["minimum_frequency"]: - raise ValueError(f"For {waveform_approximant}, reference_frequency must equal minimum_frequency") + waveform_approximant_number = lalsim_GetApproximantFromString( + waveform_approximant + ) + if ( + SimInspiralGetSpinFreqFromApproximant(waveform_approximant_number) + == self.LAL_SIM_INSPIRAL_SPINS_FLOW + ): + if ( + self.waveform_arguments["reference_frequency"] + != self.waveform_arguments["minimum_frequency"] + ): + raise ValueError( + f"For {waveform_approximant}, reference_frequency must equal minimum_frequency" + ) class GWSignalWaveformGenerator(WaveformGenerator): @@ -368,7 +420,9 @@ def _create_generator(self, waveform_approximant=None): try: from lalsimulation.gwsignal import gwsignal_get_waveform_generator except ImportError: - raise ImportError("lalsimulation is not installed. Cannot use the GWSignal waveform generator.") + raise ImportError( + "lalsimulation is not installed. Cannot use the GWSignal waveform generator." + ) if waveform_approximant is None: waveform_approximant = self.waveform_approximant @@ -385,7 +439,9 @@ def __getstate__(self): def __setstate__(self, state): # if the waveform generator can't be pickled, will reinstantiate it if state.get("generator", None) == "": - state["generator"] = self._create_generator(waveform_approximant=state["waveform_approximant"]) + state["generator"] = self._create_generator( + waveform_approximant=state["waveform_approximant"] + ) self.__dict__.update(state) def __repr__(self): @@ -454,21 +510,23 @@ def _from_bilby_parameters(self, **parameters): pn_amplitude_order=0, ) waveform_kwargs.update(self.waveform_arguments) - reference_frequency = waveform_kwargs['reference_frequency'] - minimum_frequency = waveform_kwargs['minimum_frequency'] - maximum_frequency = waveform_kwargs['maximum_frequency'] - mode_array = waveform_kwargs['mode_array'] - pn_amplitude_order = waveform_kwargs['pn_amplitude_order'] + reference_frequency = waveform_kwargs["reference_frequency"] + minimum_frequency = waveform_kwargs["minimum_frequency"] + maximum_frequency = waveform_kwargs["maximum_frequency"] + mode_array = waveform_kwargs["mode_array"] + pn_amplitude_order = waveform_kwargs["pn_amplitude_order"] if pn_amplitude_order != 0: # This is to mimic the behaviour in # https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c#L5542 if pn_amplitude_order == -1: if self.waveform_approximant in ["SpinTaylorT4", "SpinTaylorT5"]: - pn_amplitude_order = 3 # Equivalent to MAX_PRECESSING_AMP_PN_ORDER in LALSimulation + pn_amplitude_order = ( + 3 # Equivalent to MAX_PRECESSING_AMP_PN_ORDER in LALSimulation + ) else: pn_amplitude_order = 6 # Equivalent to MAX_NONPRECESSING_AMP_PN_ORDER in LALSimulation - start_frequency = minimum_frequency * 2. / (pn_amplitude_order + 2) + start_frequency = minimum_frequency * 2.0 / (pn_amplitude_order + 2) else: start_frequency = minimum_frequency @@ -476,46 +534,50 @@ def _from_bilby_parameters(self, **parameters): params.update(parameters) parameters = params - iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = bilby_to_lalsimulation_spins( - theta_jn=parameters["theta_jn"], - phi_jl=parameters["phi_jl"], - tilt_1=parameters["tilt_1"], - tilt_2=parameters["tilt_2"], - phi_12=parameters["phi_12"], - a_1=parameters["a_1"], - a_2=parameters["a_2"], - mass_1=parameters["mass_1"] * utils.solar_mass, - mass_2=parameters["mass_2"] * utils.solar_mass, - reference_frequency=reference_frequency, - phase=parameters["phase"], + iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = ( + bilby_to_lalsimulation_spins( + theta_jn=parameters["theta_jn"], + phi_jl=parameters["phi_jl"], + tilt_1=parameters["tilt_1"], + tilt_2=parameters["tilt_2"], + phi_12=parameters["phi_12"], + a_1=parameters["a_1"], + a_2=parameters["a_2"], + mass_1=parameters["mass_1"] * utils.solar_mass, + mass_2=parameters["mass_2"] * utils.solar_mass, + reference_frequency=reference_frequency, + phase=parameters["phase"], + ) ) gwsignal_dict = { - 'mass1': parameters["mass_1"], - 'mass2': parameters["mass_2"], - 'spin1x': spin_1x, - 'spin1y': spin_1y, - 'spin1z': spin_1z, - 'spin2x': spin_2x, - 'spin2y': spin_2y, - 'spin2z': spin_2z, - 'lambda1': parameters["lambda_1"], - 'lambda2': parameters["lambda_2"], - 'deltaF': 1 / self.duration, - 'deltaT': 1 / self.sampling_frequency, - 'f22_start': start_frequency, - 'f_max': maximum_frequency, - 'f22_ref': reference_frequency, - 'phi_ref': parameters["phase"], - 'distance': parameters["luminosity_distance"] * 1e6, - 'inclination': iota, - 'eccentricity': parameters["eccentricity"], - 'meanPerAno': parameters["mean_per_ano"], - 'condition': int(self.generator.metadata["implemented_domain"] == 'time'), + "mass1": parameters["mass_1"], + "mass2": parameters["mass_2"], + "spin1x": spin_1x, + "spin1y": spin_1y, + "spin1z": spin_1z, + "spin2x": spin_2x, + "spin2y": spin_2y, + "spin2z": spin_2z, + "lambda1": parameters["lambda_1"], + "lambda2": parameters["lambda_2"], + "deltaF": 1 / self.duration, + "deltaT": 1 / self.sampling_frequency, + "f22_start": start_frequency, + "f_max": maximum_frequency, + "f22_ref": reference_frequency, + "phi_ref": parameters["phase"], + "distance": parameters["luminosity_distance"] * 1e6, + "inclination": iota, + "eccentricity": parameters["eccentricity"], + "meanPerAno": parameters["mean_per_ano"], + "condition": int(self.generator.metadata["implemented_domain"] == "time"), } # add astropy units to the parameters using the defaults from gwsignal - from lalsimulation.gwsignal.core.parameter_conventions import Cosmo_units_dictionary + from lalsimulation.gwsignal.core.parameter_conventions import ( + Cosmo_units_dictionary, + ) gwsignal_dict = { key: val << Cosmo_units_dictionary.get(key, 0) @@ -528,17 +590,17 @@ def _from_bilby_parameters(self, **parameters): extra_args = waveform_kwargs.copy() for key in [ - "waveform_approximant", - "reference_frequency", - "minimum_frequency", - "maximum_frequency", - "catch_waveform_errors", - "mode_array", - "pn_spin_order", - "pn_amplitude_order", - "pn_tidal_order", - "pn_phase_order", - "numerical_relativity_file", + "waveform_approximant", + "reference_frequency", + "minimum_frequency", + "maximum_frequency", + "catch_waveform_errors", + "mode_array", + "pn_spin_order", + "pn_amplitude_order", + "pn_tidal_order", + "pn_phase_order", + "numerical_relativity_file", ]: if key in extra_args.keys(): del extra_args[key] @@ -553,7 +615,7 @@ def frequency_domain_strain(self, parameters): GenerateFDWaveform, self._from_bilby_parameters(**parameters), self.generator, - self.waveform_arguments.get("catch_waveform_errors", False) + self.waveform_arguments.get("catch_waveform_errors", False), ) wf = self._extract_waveform(hpc, "frequency") @@ -561,17 +623,20 @@ def frequency_domain_strain(self, parameters): return None minimum_frequency = self.waveform_arguments.get("minimum_frequency", 20.0) - maximum_frequency = self.waveform_arguments.get("maximum_frequency", self.frequency_array[-1]) - frequency_bounds = ( - (self.frequency_array >= minimum_frequency) - * (self.frequency_array <= maximum_frequency) + maximum_frequency = self.waveform_arguments.get( + "maximum_frequency", self.frequency_array[-1] + ) + frequency_bounds = (self.frequency_array >= minimum_frequency) * ( + self.frequency_array <= maximum_frequency ) for key in wf: wf[key] *= frequency_bounds - if self.generator.metadata["implemented_domain"] == 'time': + if self.generator.metadata["implemented_domain"] == "time": dt = 1 / hpc.hp.df.value + hpc.hp.epoch.value - time_shift = np.exp(-1j * 2 * np.pi * dt * self.frequency_array[frequency_bounds]) + time_shift = np.exp( + -1j * 2 * np.pi * dt * self.frequency_array[frequency_bounds] + ) for key in wf: wf[key][frequency_bounds] *= time_shift @@ -584,7 +649,7 @@ def time_domain_strain(self, parameters): GenerateTDWaveform, self._from_bilby_parameters(**parameters), self.generator, - self.waveform_arguments.get("catch_waveform_errors", False) + self.waveform_arguments.get("catch_waveform_errors", False), ) return self._extract_waveform(hpc, "time") @@ -606,13 +671,14 @@ def _extract_waveform(self, hpc, kind): if len(hpc.hp) > len(array): logger.debug( f"GWsignal waveform longer than bilby's `{kind}_array`({len(hpc.hp)} " - f"vs {len(array)}). Truncating GWsignal array.") + f"vs {len(array)}). Truncating GWsignal array." + ) # set slice to force the output into a numpy array - h_plus[:] = hpc.hp[:len(h_plus)] - h_cross[:] = hpc.hc[:len(h_cross)] + h_plus[:] = hpc.hp[: len(h_plus)] + h_cross[:] = hpc.hc[: len(h_cross)] else: - h_plus[:len(hpc.hp)] = hpc.hp - h_cross[:len(hpc.hc)] = hpc.hc + h_plus[: len(hpc.hp)] = hpc.hp + h_cross[: len(hpc.hc)] = hpc.hc return dict(plus=h_plus, cross=h_cross) diff --git a/bilby/hyper/compat/pytrees.py b/bilby/hyper/compat/pytrees.py index d8199468e..98d50432b 100644 --- a/bilby/hyper/compat/pytrees.py +++ b/bilby/hyper/compat/pytrees.py @@ -14,9 +14,7 @@ def hyperparameter_likelihood_flatten(likelihood: HyperparameterLikelihood): likelihood.samples_factor, likelihood.hyper_prior, ) - aux_data += ( - likelihood.sampling_prior, - ) + aux_data += (likelihood.sampling_prior,) return children, aux_data @@ -49,6 +47,8 @@ def model_unflatten(aux_data, flat) -> Model: register_pytree_node( - HyperparameterLikelihood, hyperparameter_likelihood_flatten, hyperparameter_likelihood_unflatten + HyperparameterLikelihood, + hyperparameter_likelihood_flatten, + hyperparameter_likelihood_unflatten, ) register_pytree_node(Model, model_flatten, model_unflatten) diff --git a/bilby/hyper/likelihood.py b/bilby/hyper/likelihood.py index fa113a048..2eb4e7a76 100644 --- a/bilby/hyper/likelihood.py +++ b/bilby/hyper/likelihood.py @@ -1,4 +1,3 @@ - import logging import array_api_compat as aac @@ -11,7 +10,7 @@ class HyperparameterLikelihood(Likelihood): - """ A likelihood for inferring hyperparameter posterior distributions + """A likelihood for inferring hyperparameter posterior distributions See Eq. (34) of https://arxiv.org/abs/1809.02293 for a definition. @@ -36,16 +35,30 @@ class HyperparameterLikelihood(Likelihood): """ - def __init__(self, posteriors, hyper_prior, sampling_prior=None, - log_evidences=None, max_samples=1e100, xp=np): + def __init__( + self, + posteriors, + hyper_prior, + sampling_prior=None, + log_evidences=None, + max_samples=1e100, + xp=np, + ): if not isinstance(hyper_prior, Model): hyper_prior = Model([hyper_prior]) if sampling_prior is None: - if ('log_prior' not in posteriors[0].keys()) and ('prior' not in posteriors[0].keys()): - raise ValueError('Missing both sampling prior function and prior or log_prior ' - 'column in posterior dictionary. Must pass one or the other.') + if ("log_prior" not in posteriors[0].keys()) and ( + "prior" not in posteriors[0].keys() + ): + raise ValueError( + "Missing both sampling prior function and prior or log_prior " + "column in posterior dictionary. Must pass one or the other." + ) else: - if not (isinstance(sampling_prior, Model) or isinstance(sampling_prior, PriorDict)): + if not ( + isinstance(sampling_prior, Model) + or isinstance(sampling_prior, PriorDict) + ): sampling_prior = Model([sampling_prior]) if log_evidences is not None: self.evidence_factor = np.sum(log_evidences) @@ -63,13 +76,12 @@ def __init__(self, posteriors, hyper_prior, sampling_prior=None, self.data = self.resample_posteriors(posteriors=posteriors, xp=xp) self.n_posteriors = len(posteriors) self.samples_per_posterior = self.max_samples - self.samples_factor =\ - - self.n_posteriors * np.log(self.samples_per_posterior) + self.samples_factor = -self.n_posteriors * np.log(self.samples_per_posterior) def log_likelihood_ratio(self, parameters): probs = self.hyper_prior.prob(self.data, **parameters) xp = array_module(probs) - log_l = xp.sum(xp.log(xp.sum(probs / self.data['prior'], axis=-1))) + log_l = xp.sum(xp.log(xp.sum(probs / self.data["prior"], axis=-1))) log_l += self.samples_factor return xp.nan_to_num(log_l) @@ -77,7 +89,9 @@ def noise_log_likelihood(self): return self.evidence_factor def log_likelihood(self, parameters): - return self.noise_log_likelihood() + self.log_likelihood_ratio(parameters=parameters) + return self.noise_log_likelihood() + self.log_likelihood_ratio( + parameters=parameters + ) def resample_posteriors(self, posteriors=None, max_samples=None, xp=np): """ @@ -107,18 +121,19 @@ def resample_posteriors(self, posteriors=None, max_samples=None, xp=np): for posterior in posteriors: self.max_samples = min(len(posterior), self.max_samples) data = {key: [] for key in posteriors[0]} - if 'log_prior' in data.keys(): - data.pop('log_prior') - if 'prior' not in data.keys(): - data['prior'] = [] - logging.debug('Downsampling to {} samples per posterior.'.format( - self.max_samples)) + if "log_prior" in data.keys(): + data.pop("log_prior") + if "prior" not in data.keys(): + data["prior"] = [] + logging.debug( + "Downsampling to {} samples per posterior.".format(self.max_samples) + ) for posterior in posteriors: temp = posterior.sample(self.max_samples) if self.sampling_prior is not None: - temp['prior'] = self.sampling_prior.prob(temp, axis=0) - elif 'log_prior' in temp.keys(): - temp['prior'] = np.exp(temp['log_prior']) + temp["prior"] = self.sampling_prior.prob(temp, axis=0) + elif "log_prior" in temp.keys(): + temp["prior"] = np.exp(temp["log_prior"]) for key in data: data[key].append(temp[key]) for key in data: diff --git a/bilby/hyper/model.py b/bilby/hyper/model.py index eaa2da98c..23f40dece 100644 --- a/bilby/hyper/model.py +++ b/bilby/hyper/model.py @@ -54,10 +54,7 @@ def prob(self, data, **kwargs): probability = 1.0 for ii, function in enumerate(self.models): function_parameters = self._get_function_parameters(function, **kwargs) - if ( - self.cache - and self._cached_parameters[function] == function_parameters - ): + if self.cache and self._cached_parameters[function] == function_parameters: new_probability = self._cached_probability[function] else: new_probability = function(data, **function_parameters) diff --git a/cli_bilby/bilby_result.py b/cli_bilby/bilby_result.py index 48d546747..e78ce4598 100644 --- a/cli_bilby/bilby_result.py +++ b/cli_bilby/bilby_result.py @@ -1,4 +1,4 @@ -""" A command line interface to ease the process of batch jobs on result files +"""A command line interface to ease the process of batch jobs on result files Examples -------- @@ -21,6 +21,7 @@ individually. Note that passing extra commands in is not yet implemented. """ + import argparse import bilby @@ -181,7 +182,7 @@ def apply_lightweight(result, args): strip_keys = [ "_nested_samples", "log_likelihood_evaluations", - "log_prior_evaluations" + "log_prior_evaluations", ] for key in strip_keys: setattr(result, key, None) diff --git a/cli_bilby/plot_multiple_posteriors.py b/cli_bilby/plot_multiple_posteriors.py index 51636e48b..27a99f4f7 100644 --- a/cli_bilby/plot_multiple_posteriors.py +++ b/cli_bilby/plot_multiple_posteriors.py @@ -2,18 +2,28 @@ def setup_command_line_args(): - parser = argparse.ArgumentParser( - description="Plot corner plots from results files") - parser.add_argument("-r", "--results", nargs='+', - help="List of results files to use.") - parser.add_argument("-f", "--filename", default=None, - help="Output file name.") - parser.add_argument("-l", "--labels", nargs='+', default=None, - help="List of labels to use for each result.") - parser.add_argument("-p", "--parameters", nargs='+', default=None, - help="List of parameters.") - parser.add_argument("-e", "--evidences", action='store_true', default=False, - help="Add the evidences to the legend.") + parser = argparse.ArgumentParser(description="Plot corner plots from results files") + parser.add_argument( + "-r", "--results", nargs="+", help="List of results files to use." + ) + parser.add_argument("-f", "--filename", default=None, help="Output file name.") + parser.add_argument( + "-l", + "--labels", + nargs="+", + default=None, + help="List of labels to use for each result.", + ) + parser.add_argument( + "-p", "--parameters", nargs="+", default=None, help="List of parameters." + ) + parser.add_argument( + "-e", + "--evidences", + action="store_true", + default=False, + help="Add the evidences to the legend.", + ) args, _ = parser.parse_known_args() return args @@ -22,9 +32,12 @@ def setup_command_line_args(): def main(): args = setup_command_line_args() import bilby - results = [bilby.core.result.read_in_result(filename=r) - for r in args.results] - bilby.core.result.plot_multiple(results, filename=args.filename, - labels=args.labels, - parameters=args.parameters, - evidences=args.evidences) + + results = [bilby.core.result.read_in_result(filename=r) for r in args.results] + bilby.core.result.plot_multiple( + results, + filename=args.filename, + labels=args.labels, + parameters=args.parameters, + evidences=args.evidences, + ) diff --git a/docs/conf.py b/docs/conf.py index 961ef7370..d2c36071c 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -22,12 +22,17 @@ import subprocess import sys import bilby -sys.path.insert(0, os.path.abspath('../')) + +sys.path.insert(0, os.path.abspath("../")) def git_revision_hash(): try: - return subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode('ascii').strip() + return ( + subprocess.check_output(["git", "rev-parse", "HEAD"]) + .decode("ascii") + .strip() + ) except subprocess.CalledProcessError: return "master" @@ -58,43 +63,43 @@ def git_upstream_url(): # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. extensions = [ - 'sphinx.ext.autodoc', - 'sphinx.ext.mathjax', - 'numpydoc', - 'nbsphinx', - 'sphinx.ext.autosummary', - 'sphinx.ext.autosectionlabel', - 'sphinx_design', + "sphinx.ext.autodoc", + "sphinx.ext.mathjax", + "numpydoc", + "nbsphinx", + "sphinx.ext.autosummary", + "sphinx.ext.autosectionlabel", + "sphinx_design", "sphinx.ext.linkcode", - 'myst_parser', + "myst_parser", "sphinx_sitemap", ] autosummary_generate = True # Add any paths that contain templates here, relative to this directory. -templates_path = ['templates'] +templates_path = ["templates"] # The suffix(es) of source filenames. # You can specify multiple suffix as a list of string: # -source_suffix = ['.rst', '.md', '.txt'] +source_suffix = [".rst", ".md", ".txt"] # The master toctree document. -master_doc = 'index' +master_doc = "index" # General information about the project. -project = u'bilby' -copyright = u'2019, Greg Ashton' -author = u'Greg Ashton' +project = "bilby" +copyright = "2019, Greg Ashton" +author = "Greg Ashton" # The version info for the project you're documenting, acts as replacement for # |version| and |release|, also used in various other places throughout the # built documents. # -fullversion = bilby.__version__.split(':')[0] +fullversion = bilby.__version__.split(":")[0] # The short X.Y version. -version = '.'.join(fullversion.split('.')[:2]) +version = ".".join(fullversion.split(".")[:2]) # The full version, including alpha/beta/rc tags. release = version @@ -109,16 +114,16 @@ def git_upstream_url(): # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. # This patterns also effect to html_static_path and html_extra_path -exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store', 'requirements.txt'] +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "requirements.txt"] # The html path allows for google search console verification through # Aditya Vijaykumar's gmail ID -html_extra_path = ['robots.txt', 'google063678b5c432c237.html'] +html_extra_path = ["robots.txt", "google063678b5c432c237.html"] # The name of the Pygments (syntax highlighting) style to use. -pygments_style = 'sphinx' +pygments_style = "sphinx" # If true, `todo` and `todoList` produce output, else they produce nothing. todo_include_todos = False @@ -129,7 +134,7 @@ def git_upstream_url(): # The theme to use for HTML and HTML Help pages. See the documentation for # a list of builtin themes. # -html_theme = 'sphinx_rtd_theme' +html_theme = "sphinx_rtd_theme" # Theme options are theme-specific and customize the look and feel of a theme # further. For a list of options available for each theme, see the @@ -148,21 +153,21 @@ def git_upstream_url(): # This is required for the alabaster theme # refs: http://alabaster.readthedocs.io/en/latest/installation.html#sidebars html_sidebars = { - '**': [ - 'about.html', - 'navigation.html', - 'relations.html', # needs 'show_related': True theme option to display - 'searchbox.html', - 'donate.html', + "**": [ + "about.html", + "navigation.html", + "relations.html", # needs 'show_related': True theme option to display + "searchbox.html", + "donate.html", ] } -html_baseurl = 'https://bilby-dev.github.io/bilby/' +html_baseurl = "https://bilby-dev.github.io/bilby/" # -- Options for HTMLHelp output ------------------------------------------ # Output file base name for HTML help builder. -htmlhelp_basename = 'bilbydoc' +htmlhelp_basename = "bilbydoc" # -- Options for LaTeX output --------------------------------------------- @@ -171,15 +176,12 @@ def git_upstream_url(): # The paper size ('letterpaper' or 'a4paper'). # # 'papersize': 'letterpaper', - # The font size ('10pt', '11pt' or '12pt'). # # 'pointsize': '10pt', - # Additional stuff for the LaTeX preamble. # # 'preamble': '', - # Latex figure (float) alignment # # 'figure_align': 'htbp', @@ -189,8 +191,7 @@ def git_upstream_url(): # (source start file, target name, title, # author, documentclass [howto, manual, or own class]). latex_documents = [ - (master_doc, 'bilby.tex', u'bilby Documentation', - u'Paul Lasky', 'manual'), + (master_doc, "bilby.tex", "bilby Documentation", "Paul Lasky", "manual"), ] @@ -198,10 +199,7 @@ def git_upstream_url(): # One entry per manual page. List of tuples # (source start file, name, description, authors, manual section). -man_pages = [ - (master_doc, 'bilby', u'bilby Documentation', - [author], 1) -] +man_pages = [(master_doc, "bilby", "bilby Documentation", [author], 1)] # -- Options for Texinfo output ------------------------------------------- @@ -210,9 +208,15 @@ def git_upstream_url(): # (source start file, target name, title, author, # dir menu entry, description, category) texinfo_documents = [ - (master_doc, 'bilby', u'bilby Documentation', - author, 'bilby', 'One line description of project.', - 'Miscellaneous'), + ( + master_doc, + "bilby", + "bilby Documentation", + author, + "bilby", + "One line description of project.", + "Miscellaneous", + ), ] numpydoc_show_class_members = False @@ -225,9 +229,9 @@ def linkcode_resolve(domain, info): """ Adapted from https://github.com/aaugustin/websockets/blob/8e1628a14e0dd2ca98871c7500484b5d42d16b67/docs/conf.py """ - if domain != 'py': + if domain != "py": return None - if not info['module']: + if not info["module"]: return None try: diff --git a/examples/core_examples/alternative_samplers/linear_regression_bilby_mcmc.py b/examples/core_examples/alternative_samplers/linear_regression_bilby_mcmc.py index 50fb5d769..0fb197996 100644 --- a/examples/core_examples/alternative_samplers/linear_regression_bilby_mcmc.py +++ b/examples/core_examples/alternative_samplers/linear_regression_bilby_mcmc.py @@ -5,6 +5,7 @@ data with background Gaussian noise """ + import bilby import matplotlib.pyplot as plt import numpy as np diff --git a/examples/core_examples/alternative_samplers/linear_regression_pymc.py b/examples/core_examples/alternative_samplers/linear_regression_pymc.py index d5f6e514c..87972c907 100644 --- a/examples/core_examples/alternative_samplers/linear_regression_pymc.py +++ b/examples/core_examples/alternative_samplers/linear_regression_pymc.py @@ -5,6 +5,7 @@ data with background Gaussian noise """ + import bilby import matplotlib.pyplot as plt import numpy as np diff --git a/examples/core_examples/gaussian_example.py b/examples/core_examples/gaussian_example.py index f8b80ec37..0e8859ece 100644 --- a/examples/core_examples/gaussian_example.py +++ b/examples/core_examples/gaussian_example.py @@ -3,6 +3,7 @@ An example of how to use bilby to perform parameter estimation for non-gravitational wave data consisting of a Gaussian with a mean and variance """ + import bilby import numpy as np from bilby.core.utils import random diff --git a/examples/core_examples/grid_example.py b/examples/core_examples/grid_example.py index f2afa36c8..877029461 100644 --- a/examples/core_examples/grid_example.py +++ b/examples/core_examples/grid_example.py @@ -5,6 +5,7 @@ data with background Gaussian noise """ + import bilby import matplotlib.pyplot as plt import numpy as np diff --git a/examples/core_examples/hyper_parameter_example.py b/examples/core_examples/hyper_parameter_example.py index 5b388bf1d..dc42a5b07 100644 --- a/examples/core_examples/hyper_parameter_example.py +++ b/examples/core_examples/hyper_parameter_example.py @@ -2,6 +2,7 @@ """ An example of how to use bilby to perform parameter estimation for hyper params """ + import bilby import matplotlib.pyplot as plt import numpy as np diff --git a/examples/core_examples/linear_regression.py b/examples/core_examples/linear_regression.py index 045fcbca7..09a2fb025 100644 --- a/examples/core_examples/linear_regression.py +++ b/examples/core_examples/linear_regression.py @@ -5,6 +5,7 @@ data with background Gaussian noise """ + import bilby import matplotlib.pyplot as plt import numpy as np diff --git a/examples/core_examples/linear_regression_grid.py b/examples/core_examples/linear_regression_grid.py index 24030f4c3..712c42942 100644 --- a/examples/core_examples/linear_regression_grid.py +++ b/examples/core_examples/linear_regression_grid.py @@ -5,6 +5,7 @@ This will compare the output of using a stochastic sampling method to evaluating the posterior on a grid. """ + import bilby import matplotlib.pyplot as plt import numpy as np diff --git a/examples/core_examples/linear_regression_unknown_noise.py b/examples/core_examples/linear_regression_unknown_noise.py index fe869e8cb..fa5c50b83 100644 --- a/examples/core_examples/linear_regression_unknown_noise.py +++ b/examples/core_examples/linear_regression_unknown_noise.py @@ -5,6 +5,7 @@ data with background Gaussian noise with unknown variance. """ + import bilby import matplotlib.pyplot as plt import numpy as np diff --git a/examples/core_examples/linear_regression_with_Fisher.py b/examples/core_examples/linear_regression_with_Fisher.py index ebf4188af..d7f6da2b7 100644 --- a/examples/core_examples/linear_regression_with_Fisher.py +++ b/examples/core_examples/linear_regression_with_Fisher.py @@ -6,6 +6,7 @@ estimated using the Fisher Information Matrix approximation. """ + import copy import bilby diff --git a/examples/core_examples/logo/sample_logo.py b/examples/core_examples/logo/sample_logo.py index ed817b7a9..3f407c412 100644 --- a/examples/core_examples/logo/sample_logo.py +++ b/examples/core_examples/logo/sample_logo.py @@ -1,4 +1,5 @@ -""" Script used to generate the samples for the bilby logo """ +"""Script used to generate the samples for the bilby logo""" + import bilby import numpy as np import scipy.interpolate as si diff --git a/examples/core_examples/occam_factor_example.py b/examples/core_examples/occam_factor_example.py index e00ffe6a6..d5549f78f 100644 --- a/examples/core_examples/occam_factor_example.py +++ b/examples/core_examples/occam_factor_example.py @@ -30,6 +30,7 @@ improved by increasing this to say 500 or 1000. """ + import bilby import matplotlib.pyplot as plt import numpy as np @@ -94,8 +95,7 @@ def polynomial(self, x, parameters): def log_likelihood(self, parameters): res = self.y - self.polynomial(self.x, parameters) return -0.5 * ( - np.sum((res / self.sigma) ** 2) - + self.N * np.log(2 * np.pi * self.sigma**2) + np.sum((res / self.sigma) ** 2) + self.N * np.log(2 * np.pi * self.sigma**2) ) diff --git a/examples/core_examples/radioactive_decay.py b/examples/core_examples/radioactive_decay.py index 0e669979b..5993bcfd4 100644 --- a/examples/core_examples/radioactive_decay.py +++ b/examples/core_examples/radioactive_decay.py @@ -4,6 +4,7 @@ non-gravitational wave data. In this case, fitting the half-life and initial radionuclide number for Polonium 214. """ + import bilby import matplotlib.pyplot as plt import numpy as np diff --git a/examples/gw_examples/data_examples/GW150914.py b/examples/gw_examples/data_examples/GW150914.py index f1aff756e..dd915d0f5 100755 --- a/examples/gw_examples/data_examples/GW150914.py +++ b/examples/gw_examples/data_examples/GW150914.py @@ -9,6 +9,7 @@ [1] https://gwpy.github.io/docs/stable/timeseries/remote-access.html """ + import bilby from gwpy.timeseries import TimeSeries diff --git a/examples/gw_examples/data_examples/GW190425.py b/examples/gw_examples/data_examples/GW190425.py index efe9da71a..2adc103c1 100644 --- a/examples/gw_examples/data_examples/GW190425.py +++ b/examples/gw_examples/data_examples/GW190425.py @@ -9,6 +9,7 @@ [1] https://gwpy.github.io/docs/stable/timeseries/remote-access.html """ + import bilby from gwpy.timeseries import TimeSeries diff --git a/examples/gw_examples/injection_examples/binary_neutron_star_example.py b/examples/gw_examples/injection_examples/binary_neutron_star_example.py index 45cd39037..0a69bcf63 100644 --- a/examples/gw_examples/injection_examples/binary_neutron_star_example.py +++ b/examples/gw_examples/injection_examples/binary_neutron_star_example.py @@ -8,7 +8,6 @@ tidal deformabilities """ - import bilby from bilby.core.utils.random import seed diff --git a/examples/gw_examples/injection_examples/bns_eos_example.py b/examples/gw_examples/injection_examples/bns_eos_example.py index 2a0820555..ac90af9c0 100644 --- a/examples/gw_examples/injection_examples/bns_eos_example.py +++ b/examples/gw_examples/injection_examples/bns_eos_example.py @@ -7,7 +7,6 @@ WARNING: The code is extremely slow. """ - import bilby from bilby.core.utils.random import seed from bilby.gw.eos import EOSFamily, TabularEOS diff --git a/examples/gw_examples/injection_examples/bns_polytrope_eos_example.py b/examples/gw_examples/injection_examples/bns_polytrope_eos_example.py index 845f196ce..c86617c2a 100644 --- a/examples/gw_examples/injection_examples/bns_polytrope_eos_example.py +++ b/examples/gw_examples/injection_examples/bns_polytrope_eos_example.py @@ -14,7 +14,6 @@ LALSimNeutronStarEOSDynamicPolytrope.c. """ - import bilby from bilby.core.utils.random import seed diff --git a/examples/gw_examples/injection_examples/bns_spectral_pca_eos_example.py b/examples/gw_examples/injection_examples/bns_spectral_pca_eos_example.py index 9c9e95464..2f0b21472 100644 --- a/examples/gw_examples/injection_examples/bns_spectral_pca_eos_example.py +++ b/examples/gw_examples/injection_examples/bns_spectral_pca_eos_example.py @@ -10,7 +10,6 @@ in the appendix of https://arxiv.org/pdf/2001.01747.pdf. """ - import bilby from bilby.core.utils.random import seed diff --git a/examples/gw_examples/injection_examples/change_sampled_parameters.py b/examples/gw_examples/injection_examples/change_sampled_parameters.py index c5ad7f618..5395bcff1 100644 --- a/examples/gw_examples/injection_examples/change_sampled_parameters.py +++ b/examples/gw_examples/injection_examples/change_sampled_parameters.py @@ -7,6 +7,7 @@ mass ratio and redshift. The cosmology is according to the Planck 2015 data release. """ + import bilby import numpy as np from bilby.core.utils.random import seed diff --git a/examples/gw_examples/injection_examples/create_your_own_source_model.py b/examples/gw_examples/injection_examples/create_your_own_source_model.py index fab1a650c..974e6471a 100644 --- a/examples/gw_examples/injection_examples/create_your_own_source_model.py +++ b/examples/gw_examples/injection_examples/create_your_own_source_model.py @@ -2,6 +2,7 @@ """ A script to demonstrate how to use your own source model """ + import bilby import numpy as np from bilby.core.utils.random import seed diff --git a/examples/gw_examples/injection_examples/custom_proposal_example.py b/examples/gw_examples/injection_examples/custom_proposal_example.py index 49bdcceed..a2ec37ccc 100644 --- a/examples/gw_examples/injection_examples/custom_proposal_example.py +++ b/examples/gw_examples/injection_examples/custom_proposal_example.py @@ -7,6 +7,7 @@ Due to how cpnest creates parallel processes, the multiprocessing start method needs to be set on some operating systems. """ + import multiprocessing multiprocessing.set_start_method("fork") # noqa diff --git a/examples/gw_examples/injection_examples/fake_sampler_example.py b/examples/gw_examples/injection_examples/fake_sampler_example.py index 6563a657d..24459e9f8 100755 --- a/examples/gw_examples/injection_examples/fake_sampler_example.py +++ b/examples/gw_examples/injection_examples/fake_sampler_example.py @@ -119,9 +119,9 @@ def main(): ) # update the waveform generator to use our higher-order mode waveform - likelihood.waveform_generator.waveform_arguments[ - "waveform_approximant" - ] = "IMRPhenomXHM" + likelihood.waveform_generator.waveform_arguments["waveform_approximant"] = ( + "IMRPhenomXHM" + ) # call the FakeSampler to compute the new likelihoods new_result = bilby.run_sampler( diff --git a/examples/gw_examples/injection_examples/jax_fast_tutorial.py b/examples/gw_examples/injection_examples/jax_fast_tutorial.py index bfb7f1b08..073f26873 100644 --- a/examples/gw_examples/injection_examples/jax_fast_tutorial.py +++ b/examples/gw_examples/injection_examples/jax_fast_tutorial.py @@ -9,6 +9,7 @@ We optionally use ripple waveforms and a JIT-compiled likelihood. """ + import os # Set OMP_NUM_THREADS to stop lalsimulation taking over my computer diff --git a/examples/gw_examples/injection_examples/marginalized_likelihood.py b/examples/gw_examples/injection_examples/marginalized_likelihood.py index 21494c97a..13bdcbf5c 100644 --- a/examples/gw_examples/injection_examples/marginalized_likelihood.py +++ b/examples/gw_examples/injection_examples/marginalized_likelihood.py @@ -6,6 +6,7 @@ We also demonstrate how the posterior distribution for the marginalised parameter can be recovered in post-processing. """ + import bilby from bilby.core.utils.random import seed diff --git a/examples/gw_examples/injection_examples/non_tensor.py b/examples/gw_examples/injection_examples/non_tensor.py index 8425a8ebb..3e862969f 100644 --- a/examples/gw_examples/injection_examples/non_tensor.py +++ b/examples/gw_examples/injection_examples/non_tensor.py @@ -6,6 +6,7 @@ We adapt the sine-Gaussian burst model to include vector polarizations with an unknown contribution from the vector modes. """ + import bilby import numpy as np from bilby.core.utils.random import seed diff --git a/examples/gw_examples/injection_examples/plot_skymap.py b/examples/gw_examples/injection_examples/plot_skymap.py index 2a9d3c516..ea1c51dbb 100644 --- a/examples/gw_examples/injection_examples/plot_skymap.py +++ b/examples/gw_examples/injection_examples/plot_skymap.py @@ -3,6 +3,7 @@ Example script which produces posterior samples of ra and dec and generates a skymap """ + import bilby from bilby.core.utils.random import seed diff --git a/examples/gw_examples/injection_examples/plot_time_domain_data.py b/examples/gw_examples/injection_examples/plot_time_domain_data.py index d1c29247a..0fbb2c792 100644 --- a/examples/gw_examples/injection_examples/plot_time_domain_data.py +++ b/examples/gw_examples/injection_examples/plot_time_domain_data.py @@ -3,6 +3,7 @@ This example demonstrates how to simulate some data, add an injected signal and plot the data. """ + from bilby.core.utils.random import seed from bilby.gw.detector import get_empty_interferometer from bilby.gw.source import lal_binary_black_hole diff --git a/examples/gw_examples/injection_examples/relative_binning.py b/examples/gw_examples/injection_examples/relative_binning.py index 12f5c7ca1..fbcf4708e 100644 --- a/examples/gw_examples/injection_examples/relative_binning.py +++ b/examples/gw_examples/injection_examples/relative_binning.py @@ -7,6 +7,7 @@ and distance using a uniform in comoving volume prior on luminosity distance between luminosity distances of 100Mpc and 5Gpc, the cosmology is Planck15. """ + from copy import deepcopy import bilby @@ -172,7 +173,7 @@ ) weights = np.exp(weights) print( - f"Reweighting efficiency is {np.mean(weights)**2 / np.mean(weights**2) * 100:.2f}%" + f"Reweighting efficiency is {np.mean(weights) ** 2 / np.mean(weights**2) * 100:.2f}%" ) print(f"Binned vs unbinned log Bayes factor {np.log(np.mean(weights)):.2f}") diff --git a/examples/gw_examples/injection_examples/reproduce_mpa1_eos.py b/examples/gw_examples/injection_examples/reproduce_mpa1_eos.py index 6d2ee4fc2..97f14b78a 100644 --- a/examples/gw_examples/injection_examples/reproduce_mpa1_eos.py +++ b/examples/gw_examples/injection_examples/reproduce_mpa1_eos.py @@ -2,6 +2,7 @@ """ A script to demonstrate how to create and plot EoS's with Bilby """ + from bilby.gw import eos # In this script we're going to use Bilby to plot the MPA1 EoS from tabulated data diff --git a/examples/gw_examples/injection_examples/sine_gaussian_example.py b/examples/gw_examples/injection_examples/sine_gaussian_example.py index a6aa425ce..9010523fe 100644 --- a/examples/gw_examples/injection_examples/sine_gaussian_example.py +++ b/examples/gw_examples/injection_examples/sine_gaussian_example.py @@ -3,6 +3,7 @@ Tutorial to demonstrate running parameter estimation on a sine gaussian injected signal. """ + import bilby from bilby.core.utils.random import seed diff --git a/examples/gw_examples/injection_examples/standard_15d_cbc_tutorial.py b/examples/gw_examples/injection_examples/standard_15d_cbc_tutorial.py index 922d65143..f7de1eb11 100644 --- a/examples/gw_examples/injection_examples/standard_15d_cbc_tutorial.py +++ b/examples/gw_examples/injection_examples/standard_15d_cbc_tutorial.py @@ -6,6 +6,7 @@ This will take many hours to run. """ + import bilby from bilby.core.utils import random diff --git a/examples/gw_examples/supernova_example/supernova_example.py b/examples/gw_examples/supernova_example/supernova_example.py index 910979ce6..125479449 100644 --- a/examples/gw_examples/supernova_example/supernova_example.py +++ b/examples/gw_examples/supernova_example/supernova_example.py @@ -10,6 +10,7 @@ conda install -c conda-forge pymultinest """ + import bilby import numpy as np from bilby.core.utils.random import seed diff --git a/examples/tutorials/compare_samplers.ipynb b/examples/tutorials/compare_samplers.ipynb index 0ddab8627..5af956b4e 100644 --- a/examples/tutorials/compare_samplers.ipynb +++ b/examples/tutorials/compare_samplers.ipynb @@ -179,7 +179,7 @@ " resume=False,\n", " clean=True,\n", " verbose=False,\n", - " **samplers[sampler]\n", + " **samplers[sampler],\n", " )\n", " results[sampler] = result" ] diff --git a/examples/tutorials/conditional_priors.ipynb b/examples/tutorials/conditional_priors.ipynb index 0838e5da9..5f6282b8e 100644 --- a/examples/tutorials/conditional_priors.ipynb +++ b/examples/tutorials/conditional_priors.ipynb @@ -22,8 +22,12 @@ "import numpy as np\n", "import pandas as pd\n", "from bilby.core.prior import (\n", - " Prior, PriorDict, ConditionalPriorDict,\n", - " Uniform, ConditionalUniform, Constraint, \n", + " Prior,\n", + " PriorDict,\n", + " ConditionalPriorDict,\n", + " Uniform,\n", + " ConditionalUniform,\n", + " Constraint,\n", ")\n", "from corner import corner\n", "from scipy.stats import semicircular\n", @@ -63,8 +67,9 @@ "outputs": [], "source": [ "class SemiCircular(Prior):\n", - "\n", - " def __init__(self, radius=1, center=0, name=None, latex_label=None, unit=None, boundary=None):\n", + " def __init__(\n", + " self, radius=1, center=0, name=None, latex_label=None, unit=None, boundary=None\n", + " ):\n", " super(SemiCircular, self).__init__(\n", " minimum=center - radius,\n", " maximum=center + radius,\n", @@ -91,7 +96,7 @@ "\n", "\n", "def conditional_func_y(reference_parameters, x):\n", - " condition = np.sqrt(reference_parameters[\"maximum\"]-x**2)\n", + " condition = np.sqrt(reference_parameters[\"maximum\"] - x**2)\n", " return dict(minimum=-condition, maximum=condition)" ] }, @@ -124,31 +129,31 @@ "\n", "def convert_to_radial(parameters):\n", " p = parameters.copy()\n", - " p['r'] = p['x']**2 + p['y']**2\n", + " p[\"r\"] = p[\"x\"] ** 2 + p[\"y\"] ** 2\n", " return p\n", "\n", + "\n", "def sample_circle_with_constraint():\n", " d = PriorDict(\n", - " dictionary=dict(\n", - " x=Uniform(-1, 1),\n", - " y=Uniform(-1, 1),\n", - " r=Constraint(0, 1),\n", - " ),\n", - " conversion_function=convert_to_radial\n", - " )\n", + " dictionary=dict(\n", + " x=Uniform(-1, 1),\n", + " y=Uniform(-1, 1),\n", + " r=Constraint(0, 1),\n", + " ),\n", + " conversion_function=convert_to_radial,\n", + " )\n", " return pd.DataFrame(d.sample(N))\n", "\n", "\n", "def sample_circle_with_conditional():\n", " d = ConditionalPriorDict(\n", - " dictionary=dict(\n", - " x=SemiCircular(),\n", - " y=ConditionalUniform(\n", - " condition_func=conditional_func_y, \n", - " minimum=-1, maximum=1\n", - " )\n", - " )\n", + " dictionary=dict(\n", + " x=SemiCircular(),\n", + " y=ConditionalUniform(\n", + " condition_func=conditional_func_y, minimum=-1, maximum=1\n", + " ),\n", " )\n", + " )\n", " return pd.DataFrame(d.sample(N))\n", "\n", "\n", @@ -179,22 +184,34 @@ "source": [ "class BoundedUniform(ConditionalUniform):\n", " \"\"\"Conditional Uniform prior where prior sample < previous prior sample\n", - " \n", + "\n", " This is ensured by fixing the maximum bound to be the previous prior sample value.\n", " \"\"\"\n", - " def __init__(self, idx: int, minimum, maximum, name=None, latex_label=None,\n", - " unit=None, boundary=None):\n", + "\n", + " def __init__(\n", + " self,\n", + " idx: int,\n", + " minimum,\n", + " maximum,\n", + " name=None,\n", + " latex_label=None,\n", + " unit=None,\n", + " boundary=None,\n", + " ):\n", " super(BoundedUniform, self).__init__(\n", - " minimum=minimum, maximum=maximum, name=name, \n", - " latex_label=latex_label, unit=unit,\n", - " boundary=boundary, condition_func=self.bounds_condition\n", + " minimum=minimum,\n", + " maximum=maximum,\n", + " name=name,\n", + " latex_label=latex_label,\n", + " unit=unit,\n", + " boundary=boundary,\n", + " condition_func=self.bounds_condition,\n", " )\n", " self.idx = idx\n", " self.previous_name = f\"{name[:-1]}{self.idx - 1}\"\n", - " self._required_variables = [self.previous_name] \n", + " self._required_variables = [self.previous_name]\n", " # this is used in prior.sample(... **required_variables)\n", "\n", - "\n", " def bounds_condition(self, reference_params, **required_variables):\n", " previous_sample = required_variables[self.previous_name]\n", " return dict(maximum=previous_sample)\n", @@ -203,11 +220,13 @@ "def make_uniform_conditonal_priordict(n_priors=3):\n", " priors = ConditionalPriorDict()\n", " for i in range(n_priors):\n", - " if i==0:\n", + " if i == 0:\n", " priors[f\"uni{i}\"] = Uniform(minimum=0, maximum=1, name=f\"uni{i}\")\n", " else:\n", - " priors[f\"uni{i}\"] = BoundedUniform(idx=i, minimum=0, maximum=1, name=f\"uni{i}\")\n", - " return priors\n" + " priors[f\"uni{i}\"] = BoundedUniform(\n", + " idx=i, minimum=0, maximum=1, name=f\"uni{i}\"\n", + " )\n", + " return priors" ] }, { @@ -217,7 +236,12 @@ "outputs": [], "source": [ "samples = pd.DataFrame(make_uniform_conditonal_priordict(3).sample(10000))\n", - "fig = corner(samples.values, **CORNER_KWARGS, color=\"tab:blue\", labels=[f\"$A_{ii}$\" for ii in range(3)])\n", + "fig = corner(\n", + " samples.values,\n", + " **CORNER_KWARGS,\n", + " color=\"tab:blue\",\n", + " labels=[f\"$A_{ii}$\" for ii in range(3)],\n", + ")\n", "plt.show()\n", "plt.close()" ] diff --git a/examples/tutorials/fitting_with_x_and_y_errors.ipynb b/examples/tutorials/fitting_with_x_and_y_errors.ipynb index 38818b666..a679347bd 100644 --- a/examples/tutorials/fitting_with_x_and_y_errors.ipynb +++ b/examples/tutorials/fitting_with_x_and_y_errors.ipynb @@ -20,7 +20,7 @@ "import numpy as np\n", "from bilby.core.utils import random\n", "\n", - "#sets seed of bilby's generator \"rng\" to \"123\"\n", + "# sets seed of bilby's generator \"rng\" to \"123\"\n", "random.seed(123)\n", "\n", "%matplotlib inline" @@ -98,7 +98,15 @@ " m=bilby.core.prior.Uniform(0, 30, \"m\"), c=bilby.core.prior.Uniform(0, 30, \"c\")\n", ")\n", "\n", - "sampler_kwargs = dict(priors=priors, sampler=\"bilby_mcmc\", nsamples=1000, printdt=5, outdir=\"outdir\", verbose=False, clean=True)" + "sampler_kwargs = dict(\n", + " priors=priors,\n", + " sampler=\"bilby_mcmc\",\n", + " nsamples=1000,\n", + " printdt=5,\n", + " outdir=\"outdir\",\n", + " verbose=False,\n", + " clean=True,\n", + ")" ] }, { diff --git a/scripts/check_subject_format.py b/scripts/check_subject_format.py index 49a40e2e1..b8d45edd9 100644 --- a/scripts/check_subject_format.py +++ b/scripts/check_subject_format.py @@ -20,7 +20,9 @@ def main(): args = parser.parse_args() if args.github_event: - title = json.loads(args.github_event.read_text(encoding="utf-8"))["pull_request"]["title"] + title = json.loads(args.github_event.read_text(encoding="utf-8"))[ + "pull_request" + ]["title"] else: message = args.commit_msg_file.read_text(encoding="utf-8") title = message.partition("\n")[0] diff --git a/test/bilby_mcmc/test_chain.py b/test/bilby_mcmc/test_chain.py index befc5b1d8..6308a12ad 100644 --- a/test/bilby_mcmc/test_chain.py +++ b/test/bilby_mcmc/test_chain.py @@ -22,12 +22,14 @@ def tearDown(self): shutil.rmtree(self.outdir) def create_random_sample(self): - return Sample({ - "a": np.random.normal(0, 1), - "b": np.random.normal(0, 1), - LOGLKEY: np.random.normal(0, 1), - LOGPKEY: -1 - }) + return Sample( + { + "a": np.random.normal(0, 1), + "b": np.random.normal(0, 1), + LOGLKEY: np.random.normal(0, 1), + LOGPKEY: -1, + } + ) def create_chain(self, n=1000): chain = Chain(initial_sample=self.initial_sample) @@ -43,7 +45,7 @@ def test_append(self): chain = Chain(initial_sample=self.initial_sample) chain.append(self.create_random_sample()) self.assertEqual(chain.position, 1) - self.assertEqual(len(chain.get_1d_array('a')), 2) + self.assertEqual(len(chain.get_1d_array("a")), 2) def test_append_within_init_space(self): chain = Chain(initial_sample=self.initial_sample) @@ -54,7 +56,7 @@ def test_append_within_init_space(self): self.assertEqual(chain.position, N) # N samples + 1 initial position - self.assertEqual(len(chain.get_1d_array('a')), N + 1) + self.assertEqual(len(chain.get_1d_array("a")), N + 1) def test_append_with_extending(self): block_length = 100 @@ -153,7 +155,7 @@ def test_plot(self): chain.plot(outdir=self.outdir, label="test") self.assertTrue(os.path.exists(f"{self.outdir}/test_checkpoint_trace.png")) priors = dict( - a=bilby.core.prior.Uniform(-10, 10, latex_label='a'), + a=bilby.core.prior.Uniform(-10, 10, latex_label="a"), b=bilby.core.prior.Uniform(-10, 10), ) chain.thin_by_nact = 0.5 @@ -180,18 +182,18 @@ def test_dict_access(self): def test_list_access(self): s = Sample(self.sample_dict) slist = s.list - self.assertEqual(slist, [self.sample_dict['a'], self.sample_dict['b']]) + self.assertEqual(slist, [self.sample_dict["a"], self.sample_dict["b"]]) def test_setitem(self): s = Sample(self.sample_dict) # Set existing parameter - s['a'] = 100 - self.assertEqual(s['a'], 100) + s["a"] = 100 + self.assertEqual(s["a"], 100) # Add parameter - s['c'] = 100 - self.assertEqual(s['c'], 100) + s["c"] = 100 + self.assertEqual(s["c"], 100) def test_parameter_only_dict(self): s = Sample(self.sample_dict) @@ -201,9 +203,9 @@ def test_update(self): sample_dict = dict(a=1, b=2) curr = Sample(sample_dict) prop = curr.copy() - prop['a'] = 200 - self.assertEqual(prop['a'], 200) - self.assertEqual(curr['a'], 1) + prop["a"] = 200 + self.assertEqual(prop["a"], 200) + self.assertEqual(curr["a"], 1) class TestACT(unittest.TestCase): diff --git a/test/bilby_mcmc/test_proposals.py b/test/bilby_mcmc/test_proposals.py index cd36f94fa..f55d4c3fa 100644 --- a/test/bilby_mcmc/test_proposals.py +++ b/test/bilby_mcmc/test_proposals.py @@ -15,7 +15,8 @@ class GivenProposal(proposals.BaseProposal): - """ A simple proposal class used for testing """ + """A simple proposal class used for testing""" + def __init__(self, priors, weight=1, subset=None, sigma=0.01): super(GivenProposal, self).__init__(priors, weight, subset) @@ -26,10 +27,14 @@ def propose(self, chain): class TestBaseProposals(unittest.TestCase): def create_priors(self, ndim=2, boundary=None): - priors = bilby.core.prior.PriorDict({ - f'x{i}': bilby.core.prior.Uniform(-10, 10, name=f'x{i}', boundary=boundary) - for i in range(ndim) - }) + priors = bilby.core.prior.PriorDict( + { + f"x{i}": bilby.core.prior.Uniform( + -10, 10, name=f"x{i}", boundary=boundary + ) + for i in range(ndim) + } + ) priors["fixedA"] = bilby.core.prior.DeltaFunction(1) priors["infinite_support"] = bilby.core.prior.Normal(0, 1) priors["half_infinite_support"] = bilby.core.prior.HalfNormal(1) @@ -130,7 +135,7 @@ def get_simple_proposals(self): def proposal_check(self, prop, ndim=2, N=100): chain = self.create_chain(ndim=ndim) - if getattr(prop, 'needs_likelihood_and_priors', False): + if getattr(prop, "needs_likelihood_and_priors", False): return print(f"Testing {prop.__class__.__name__}") diff --git a/test/bilby_mcmc/test_sampler.py b/test/bilby_mcmc/test_sampler.py index 61e3638a0..2a664c2cf 100644 --- a/test/bilby_mcmc/test_sampler.py +++ b/test/bilby_mcmc/test_sampler.py @@ -25,6 +25,7 @@ def setUp(self): def model(time, m, c): return time * m + c + injection_parameters = dict(m=0.5, c=0.2) sampling_frequency = 10 time_duration = 10 @@ -37,11 +38,11 @@ def model(time, m, c): # From hereon, the syntax is exactly equivalent to other bilby examples # We make a prior priors = dict() - priors['m'] = bilby.core.prior.Uniform(0, 5, 'm') - priors['c'] = bilby.core.prior.Uniform(-2, 2, 'c') + priors["m"] = bilby.core.prior.Uniform(0, 5, "m") + priors["c"] = bilby.core.prior.Uniform(-2, 2, "c") priors = bilby.core.prior.PriorDict(priors) - search_parameter_keys = ['m', 'c'] + search_parameter_keys = ["m", "c"] use_ratio = False bilby.core.sampler.base_sampler.initialize_global_variables( @@ -64,7 +65,7 @@ def test_None_proposal_cycle(self): beta=1, Tindex=0, Eindex=0, - use_ratio=False + use_ratio=False, ) def test_default_proposal_cycle(self): @@ -74,7 +75,7 @@ def test_default_proposal_cycle(self): beta=1, Tindex=0, Eindex=0, - use_ratio=False + use_ratio=False, ) nsteps = 0 @@ -89,9 +90,7 @@ def test_default_proposal_cycle(self): def test_get_expected_outputs(): label = "par0" outdir = os.path.join("some", "bilby_pipe", "dir") - filenames, directories = Bilby_MCMC.get_expected_outputs( - outdir=outdir, label=label - ) + filenames, directories = Bilby_MCMC.get_expected_outputs(outdir=outdir, label=label) assert len(filenames) == 1 assert len(directories) == 0 assert os.path.join(outdir, f"{label}_resume.pickle") in filenames diff --git a/test/check_author_list.py b/test/check_author_list.py index 95d5be1ca..adaaac12b 100644 --- a/test/check_author_list.py +++ b/test/check_author_list.py @@ -1,4 +1,4 @@ -""" A script to verify that the .AUTHOR.md file is up to date """ +"""A script to verify that the .AUTHOR.md file is up to date""" import re import subprocess @@ -9,22 +9,25 @@ AUTHORS_list = " ".join([line for line in f]).lower() -lines = subprocess.check_output(["git", "shortlog", "HEAD", "-sn"]).decode("utf-8").split("\n") +lines = ( + subprocess.check_output(["git", "shortlog", "HEAD", "-sn"]) + .decode("utf-8") + .split("\n") +) if len(lines) == 0: raise Exception("No authors to check against") def remove_accents(raw_text): - - raw_text = re.sub(u"[àáâãäå]", 'a', raw_text) - raw_text = re.sub(u"[èéêë]", 'e', raw_text) - raw_text = re.sub(u"[ìíîï]", 'i', raw_text) - raw_text = re.sub(u"[òóôõö]", 'o', raw_text) - raw_text = re.sub(u"[ùúûü]", 'u', raw_text) - raw_text = re.sub(u"[ýÿ]", 'y', raw_text) - raw_text = re.sub(u"[ß]", 'ss', raw_text) - raw_text = re.sub(u"[ñ]", 'n', raw_text) + raw_text = re.sub("[àáâãäå]", "a", raw_text) + raw_text = re.sub("[èéêë]", "e", raw_text) + raw_text = re.sub("[ìíîï]", "i", raw_text) + raw_text = re.sub("[òóôõö]", "o", raw_text) + raw_text = re.sub("[ùúûü]", "u", raw_text) + raw_text = re.sub("[ýÿ]", "y", raw_text) + raw_text = re.sub("[ß]", "ss", raw_text) + raw_text = re.sub("[ñ]", "n", raw_text) return raw_text @@ -32,7 +35,7 @@ def remove_accents(raw_text): fail_test = False for line in lines: line = line.replace(".", " ") - line = re.sub('([A-Z][a-z]+)', r' \1', re.sub('([A-Z]+)', r' \1', line)) + line = re.sub("([A-Z][a-z]+)", r" \1", re.sub("([A-Z]+)", r" \1", line)) line = remove_accents(line) for element in line.split()[1:]: element = element.lower() @@ -41,4 +44,6 @@ def remove_accents(raw_text): fail_test += True if fail_test: - raise Exception("Author check list failed.. have you added your name to the .AUTHOR file?") + raise Exception( + "Author check list failed.. have you added your name to the .AUTHOR file?" + ) diff --git a/test/check_subject_format_test.py b/test/check_subject_format_test.py index 81641019c..c4b1d24b9 100644 --- a/test/check_subject_format_test.py +++ b/test/check_subject_format_test.py @@ -16,14 +16,17 @@ def run_validator(self, title, github_event): with tempfile.TemporaryDirectory() as directory: source = Path(directory) / "message" if github_event: - source.write_text(json.dumps({"pull_request": {"title": title}}), encoding="utf-8") + source.write_text( + json.dumps({"pull_request": {"title": title}}), encoding="utf-8" + ) option = "--github-event" else: source.write_text(title, encoding="utf-8") option = "--commit-msg-file" return subprocess.run( [sys.executable, str(SCRIPT), option, str(source)], - capture_output=True, text=True, + capture_output=True, + text=True, ) def test_valid_titles(self): @@ -31,11 +34,14 @@ def test_valid_titles(self): f"{acronym}: Valid summary" for acronym in "API BENCH BLD BUG CI DEP DEV DOC ENH MAINT REV STY TST REL".split() ] + [ - 'MAINT/TST: Update tests', 'DOC: Explain `café` and $(echo example)', - 'BUG(gw): Fix sampling', 'MAINT/TST(core.prior): Update tests', - 'DOC(gw/likelihood): Update documentation', 'DEV(pre-commit): Update hooks', - 'MAINT/CI: Enforce commit message and PR titles', - 'CI(tests): Update test workflow', + "MAINT/TST: Update tests", + "DOC: Explain `café` and $(echo example)", + "BUG(gw): Fix sampling", + "MAINT/TST(core.prior): Update tests", + "DOC(gw/likelihood): Update documentation", + "DEV(pre-commit): Update hooks", + "MAINT/CI: Enforce commit message and PR titles", + "CI(tests): Update test workflow", ] for title in titles: for github_event in (True, False): @@ -44,27 +50,43 @@ def test_valid_titles(self): self.assertEqual(result.returncode, 0, result.stderr) def test_invalid_titles(self): - for title in ('', 'Fix sampling', 'bug: Fix sampling', 'FEAT: Add sampling', - 'BUG:Fix sampling', 'BUG: ', 'BUG: ', ' BUG: Fix sampling', - 'BUG/INVALID: Fix sampling', 'BUG/: Fix sampling', - 'BUG(): Fix sampling', 'BUG( ): Fix sampling', - 'BUG(core prior): Fix sampling', 'BUG(core\tprior): Fix sampling', - 'BUG((gw)): Fix sampling', 'BUG(gw: Fix sampling', - 'BUG(gw)): Fix sampling', 'BUG (gw): Fix sampling', - 'BUG(gw): ', 'BUG(gw)/TST: Fix sampling'): + for title in ( + "", + "Fix sampling", + "bug: Fix sampling", + "FEAT: Add sampling", + "BUG:Fix sampling", + "BUG: ", + "BUG: ", + " BUG: Fix sampling", + "BUG/INVALID: Fix sampling", + "BUG/: Fix sampling", + "BUG(): Fix sampling", + "BUG( ): Fix sampling", + "BUG(core prior): Fix sampling", + "BUG(core\tprior): Fix sampling", + "BUG((gw)): Fix sampling", + "BUG(gw: Fix sampling", + "BUG(gw)): Fix sampling", + "BUG (gw): Fix sampling", + "BUG(gw): ", + "BUG(gw)/TST: Fix sampling", + ): for github_event in (True, False): with self.subTest(title=title, github_event=github_event): result = self.run_validator(title, github_event) self.assertEqual(result.returncode, 1) self.assertIn("Bilby subject format", result.stderr) - self.assertEqual(result.stderr.startswith("::error::"), github_event) + self.assertEqual( + result.stderr.startswith("::error::"), github_event + ) def test_commit_body_is_allowed(self): - result = self.run_validator('BUG: Fix sampling\n\nExplain the fix.\n', False) + result = self.run_validator("BUG: Fix sampling\n\nExplain the fix.\n", False) self.assertEqual(result.returncode, 0, result.stderr) def test_multiline_pr_title_is_rejected(self): - result = self.run_validator('BUG: Fix sampling\nExtra line', True) + result = self.run_validator("BUG: Fix sampling\nExtra line", True) self.assertEqual(result.returncode, 1) diff --git a/test/conftest.py b/test/conftest.py index 0c511d1ec..a42398f59 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -8,7 +8,10 @@ def pytest_addoption(parser): parser.addoption( - "--skip-roqs", action="store_true", default=False, help="Skip all tests that require ROQs" + "--skip-roqs", + action="store_true", + default=False, + help="Skip all tests that require ROQs", ) parser.addoption( "--array-backend", @@ -19,7 +22,9 @@ def pytest_addoption(parser): def pytest_configure(config): config.addinivalue_line("markers", "requires_roqs: mark a test that requires ROQs") - config.addinivalue_line("markers", "array_backend: mark that a test uses all array backends") + config.addinivalue_line( + "markers", "array_backend: mark that a test uses all array backends" + ) def pytest_collection_modifyitems(config, items): @@ -53,6 +58,7 @@ def _xp(request): xp = jax.numpy case "torch": import torch + # torch starts a lot of threads, so disable this on the first import # to avoid segfaults try: @@ -65,7 +71,6 @@ def _xp(request): xp = torch case _: try: - xp = importlib.import_module(backend) except ImportError: raise ValueError(f"Unknown backend for testing: {backend}") @@ -80,9 +85,11 @@ def _rng(xp): return resolve_random_state(12345) elif aac.is_jax_namespace(xp): import jax.random + return resolve_random_state(jax.random.key(12345)) elif aac.is_torch_namespace(xp): import torch + return resolve_random_state(torch.Tensor([12345])) else: raise ValueError(f"Unknown array namespace {xp} for RNG") diff --git a/test/core/grid_test.py b/test/core/grid_test.py index 781077f34..70f4b7fcc 100644 --- a/test/core/grid_test.py +++ b/test/core/grid_test.py @@ -23,7 +23,9 @@ def dim(self): return len(self.cov[0]) def log_likelihood(self, parameters): - return -parameters["x0"]**2 / 2 - parameters["x1"]**2 / 2 - np.log(2 * np.pi) + return ( + -(parameters["x0"] ** 2) / 2 - parameters["x1"] ** 2 / 2 - np.log(2 * np.pi) + ) @pytest.mark.array_backend @@ -115,33 +117,64 @@ def test_parameter_names(self): def test_no_marginalization(self): # test arrays are the same if no parameters are given to marginalize # over - self.assertTrue(np.array_equal( - self.grid.ln_likelihood, - self.grid.marginalize_ln_likelihood(not_parameters=self.grid.parameter_names))) + self.assertTrue( + np.array_equal( + self.grid.ln_likelihood, + self.grid.marginalize_ln_likelihood( + not_parameters=self.grid.parameter_names + ), + ) + ) def test_marginalization_shapes(self): self.assertEqual(0, len(self.grid.marginalize_ln_likelihood().shape)) - marg1 = self.grid.marginalize_ln_likelihood(parameters=self.grid.parameter_names[0]) - marg2 = self.grid.marginalize_ln_likelihood(parameters=self.grid.parameter_names[1]) + marg1 = self.grid.marginalize_ln_likelihood( + parameters=self.grid.parameter_names[0] + ) + marg2 = self.grid.marginalize_ln_likelihood( + parameters=self.grid.parameter_names[1] + ) self.assertTupleEqual((self.grid_size,), marg1.shape) self.assertTupleEqual((self.grid_size,), marg2.shape) - self.assertTupleEqual((self.grid_size, self.grid_size), self.grid.ln_likelihood.shape) - self.assertTupleEqual((self.grid_size, self.grid_size), self.grid.ln_posterior.shape) + self.assertTupleEqual( + (self.grid_size, self.grid_size), self.grid.ln_likelihood.shape + ) + self.assertTupleEqual( + (self.grid_size, self.grid_size), self.grid.ln_posterior.shape + ) def test_marginalization_opposite(self): - self.assertTrue(np.array_equal( - self.grid.marginalize_ln_likelihood(parameters=self.grid.parameter_names[0]), - self.grid.marginalize_ln_likelihood(not_parameters=self.grid.parameter_names[1]))) + self.assertTrue( + np.array_equal( + self.grid.marginalize_ln_likelihood( + parameters=self.grid.parameter_names[0] + ), + self.grid.marginalize_ln_likelihood( + not_parameters=self.grid.parameter_names[1] + ), + ) + ) - self.assertTrue(np.array_equal( - self.grid.marginalize_ln_likelihood(parameters=self.grid.parameter_names[1]), - self.grid.marginalize_ln_likelihood(not_parameters=self.grid.parameter_names[0]))) + self.assertTrue( + np.array_equal( + self.grid.marginalize_ln_likelihood( + parameters=self.grid.parameter_names[1] + ), + self.grid.marginalize_ln_likelihood( + not_parameters=self.grid.parameter_names[0] + ), + ) + ) def test_max_marginalized_likelihood(self): # marginalised likelihoods should have max values of 1 (as they are not # properly normalised) - self.assertEqual(1.0, self.grid.marginalize_likelihood(self.grid.parameter_names[0]).max()) - self.assertEqual(1.0, self.grid.marginalize_likelihood(self.grid.parameter_names[1]).max()) + self.assertEqual( + 1.0, self.grid.marginalize_likelihood(self.grid.parameter_names[0]).max() + ) + self.assertEqual( + 1.0, self.grid.marginalize_likelihood(self.grid.parameter_names[1]).max() + ) def test_ln_evidence(self): ln_z = self.grid.ln_evidence @@ -161,9 +194,17 @@ def test_fail_grid_size(self): ) def test_mesh_grid(self): - self.assertTupleEqual((self.grid_size, self.grid_size), self.grid.mesh_grid[0].shape) - self.assertEqual(self.priors[self.grid.parameter_names[0]].minimum, self.grid.mesh_grid[0][0, 0]) - self.assertEqual(self.priors[self.grid.parameter_names[1]].maximum, self.grid.mesh_grid[0][-1, -1]) + self.assertTupleEqual( + (self.grid_size, self.grid_size), self.grid.mesh_grid[0].shape + ) + self.assertEqual( + self.priors[self.grid.parameter_names[0]].minimum, + self.grid.mesh_grid[0][0, 0], + ) + self.assertEqual( + self.priors[self.grid.parameter_names[1]].maximum, + self.grid.mesh_grid[0][-1, -1], + ) def test_grid_integer_points(self): n_points = [10, 20] @@ -177,8 +218,12 @@ def test_grid_integer_points(self): ) self.assertTupleEqual(tuple(n_points), grid.mesh_grid[0].shape) - self.assertEqual(grid.mesh_grid[0][0, 0], self.priors[self.grid.parameter_names[0]].minimum) - self.assertEqual(grid.mesh_grid[0][-1, -1], self.priors[self.grid.parameter_names[1]].maximum) + self.assertEqual( + grid.mesh_grid[0][0, 0], self.priors[self.grid.parameter_names[0]].minimum + ) + self.assertEqual( + grid.mesh_grid[0][-1, -1], self.priors[self.grid.parameter_names[1]].maximum + ) def test_grid_dict_points(self): n_points = {"x0": 15, "x1": 18} @@ -191,8 +236,12 @@ def test_grid_dict_points(self): xp=self.xp, ) self.assertTupleEqual((n_points["x0"], n_points["x1"]), grid.mesh_grid[0].shape) - self.assertEqual(grid.mesh_grid[0][0, 0], self.priors[self.grid.parameter_names[0]].minimum) - self.assertEqual(grid.mesh_grid[0][-1, -1], self.priors[self.grid.parameter_names[1]].maximum) + self.assertEqual( + grid.mesh_grid[0][0, 0], self.priors[self.grid.parameter_names[0]].minimum + ) + self.assertEqual( + grid.mesh_grid[0][-1, -1], self.priors[self.grid.parameter_names[1]].maximum + ) def test_grid_from_array(self): x0s = np.linspace(self.priors["x0"].minimum, self.priors["x0"].maximum, 13) @@ -209,8 +258,12 @@ def test_grid_from_array(self): ) self.assertTupleEqual((len(x0s), len(x1s)), grid.mesh_grid[0].shape) - self.assertEqual(grid.mesh_grid[0][0, 0], self.priors[self.grid.parameter_names[0]].minimum) - self.assertEqual(grid.mesh_grid[0][-1, -1], self.priors[self.grid.parameter_names[1]].maximum) + self.assertEqual( + grid.mesh_grid[0][0, 0], self.priors[self.grid.parameter_names[0]].minimum + ) + self.assertEqual( + grid.mesh_grid[0][-1, -1], self.priors[self.grid.parameter_names[1]].maximum + ) self.assertTrue(np.array_equal(grid.sample_points["x0"], x0s)) self.assertTrue(np.array_equal(grid.sample_points["x1"], x1s)) @@ -224,7 +277,11 @@ def test_save_and_load_from_filename(self): self.assertEqual(new_grid.n_dims, self.grid.n_dims) self.assertTrue(np.array_equal(new_grid.mesh_grid[0], self.grid.mesh_grid[0])) for par in new_grid.parameter_names: - self.assertTrue(np.array_equal(new_grid.sample_points[par], self.grid.sample_points[par])) + self.assertTrue( + np.array_equal( + new_grid.sample_points[par], self.grid.sample_points[par] + ) + ) self.assertEqual(new_grid.ln_evidence, self.grid.ln_evidence) self.assertTrue(np.array_equal(new_grid.ln_likelihood, self.grid.ln_likelihood)) self.assertTrue(np.array_equal(new_grid.ln_posterior, self.grid.ln_posterior)) @@ -237,8 +294,10 @@ def test_save_and_load_from_outdir_label(self): self.assertEqual(self.grid.n_dims, new_grid.n_dims) self.assertTrue(np.array_equal(new_grid.mesh_grid[0], self.grid.mesh_grid[0])) for par in new_grid.parameter_names: - self.assertTrue(np.array_equal( - new_grid.sample_points[par], self.grid.sample_points[par]) + self.assertTrue( + np.array_equal( + new_grid.sample_points[par], self.grid.sample_points[par] + ) ) self.assertEqual(self.grid.ln_evidence, new_grid.ln_evidence) self.assertTrue(np.array_equal(self.grid.ln_likelihood, new_grid.ln_likelihood)) @@ -254,7 +313,11 @@ def test_save_and_load_gzip(self): self.assertEqual(self.grid.n_dims, new_grid.n_dims) self.assertTrue(np.array_equal(self.grid.mesh_grid[0], new_grid.mesh_grid[0])) for par in new_grid.parameter_names: - self.assertTrue(np.array_equal(self.grid.sample_points[par], new_grid.sample_points[par])) + self.assertTrue( + np.array_equal( + self.grid.sample_points[par], new_grid.sample_points[par] + ) + ) self.assertEqual(self.grid.ln_evidence, new_grid.ln_evidence) self.assertTrue(np.array_equal(self.grid.ln_likelihood, new_grid.ln_likelihood)) self.assertTrue(np.array_equal(self.grid.ln_posterior, new_grid.ln_posterior)) diff --git a/test/core/likelihood_test.py b/test/core/likelihood_test.py index 851961baa..b11bc9614 100644 --- a/test/core/likelihood_test.py +++ b/test/core/likelihood_test.py @@ -176,7 +176,9 @@ def test_model_parameters(self): parameter1=self.parameter1_value, parameter2=self.parameter2_value ) full_parameters = dict( - parameter1=self.parameter1_value, parameter2=self.parameter2_value, sigma=sigma + parameter1=self.parameter1_value, + parameter2=self.parameter2_value, + sigma=sigma, ) self.assertDictEqual( expected_model_parameters, @@ -197,7 +199,9 @@ def setUp(self): self.N = 100 self.sigma = 0.1 self.x = self.xp.linspace(0, 1, self.N) - self.y = 2 * self.x + 1 + self.xp.asarray(np.random.normal(0, self.sigma, self.N)) + self.y = ( + 2 * self.x + 1 + self.xp.asarray(np.random.normal(0, self.sigma, self.N)) + ) def test_function(x, m, c): return m * x + c @@ -271,7 +275,9 @@ def setUp(self): self.nu = self.N - 2 self.sigma = 1 self.x = self.xp.linspace(0, 1, self.N) - self.y = 2 * self.x + 1 + self.xp.asarray(np.random.normal(0, self.sigma, self.N)) + self.y = ( + 2 * self.x + 1 + self.xp.asarray(np.random.normal(0, self.sigma, self.N)) + ) def test_function(x, m, c): return m * x + c @@ -403,14 +409,22 @@ def test_neg_rate_array(self): likelihood.log_likelihood(parameters) def test_init_y(self): - self.assertEqual(aac.get_namespace(self.y), aac.get_namespace(self.poisson_likelihood.y)) - np.testing.assert_array_equal(np.asarray(self.y), np.asarray(self.poisson_likelihood.y)) + self.assertEqual( + aac.get_namespace(self.y), aac.get_namespace(self.poisson_likelihood.y) + ) + np.testing.assert_array_equal( + np.asarray(self.y), np.asarray(self.poisson_likelihood.y) + ) def test_set_y_to_array(self): new_y = self.xp.arange(0, 50, step=2) self.poisson_likelihood.y = new_y - self.assertEqual(aac.get_namespace(new_y), aac.get_namespace(self.poisson_likelihood.y)) - np.testing.assert_array_equal(np.asarray(new_y), np.asarray(self.poisson_likelihood.y)) + self.assertEqual( + aac.get_namespace(new_y), aac.get_namespace(self.poisson_likelihood.y) + ) + np.testing.assert_array_equal( + np.asarray(new_y), np.asarray(self.poisson_likelihood.y) + ) def test_set_y_to_positive_int(self): new_y = 5 @@ -450,7 +464,7 @@ def test_log_likelihood_zero_func_return_element(self): self.assertEqual(-np.inf, poisson_likelihood.log_likelihood(dict())) def test_log_likelihood_dummy(self): - """ Merely tests if it goes into the right if else bracket """ + """Merely tests if it goes into the right if else bracket""" poisson_likelihood = PoissonLikelihood( x=self.x, y=self.y, func=lambda x: self.xp.linspace(1, 100, self.N) ) @@ -553,7 +567,7 @@ def test_set_y_to_nd_array_with_negative_element(self): self.exponential_likelihood.y = self.xp.asarray([4.3, -1.2, 4]) def test_log_likelihood_default(self): - """ Merely tests that it ends up at the right place in the code """ + """Merely tests that it ends up at the right place in the code""" exponential_likelihood = ExponentialLikelihood( x=self.x, y=self.y, func=lambda x: self.xp.asarray([4.2]) ) @@ -724,7 +738,9 @@ def test_log_likelihood_correctly_sums(self): + self.second_likelihood.log_likelihood(self.second_parmaeters) + self.third_likelihood.log_likelihood(self.third_parmaeters) ) - self.assertEqual(expected, self.joint_likelihood.log_likelihood(self.parameters)) + self.assertEqual( + expected, self.joint_likelihood.log_likelihood(self.parameters) + ) def test_log_noise_likelihood(self): self.first_likelihood.noise_log_likelihood = mock.MagicMock(return_value=1) @@ -762,7 +778,6 @@ def test_jitted_likelihood(self): class TestGPLikelihood(unittest.TestCase): - def setUp(self) -> None: self.t = [1, 2, 3] self.y = [4, 5, 6] @@ -772,10 +787,18 @@ def setUp(self) -> None: self.gp_mock = mock.MagicMock() self.gp_mock.compute = mock.MagicMock() self.parameter_dict = dict(a=1, b=2) - self.gp_mock.get_parameter_dict = mock.MagicMock(return_value=dict(self.parameter_dict)) + self.gp_mock.get_parameter_dict = mock.MagicMock( + return_value=dict(self.parameter_dict) + ) self.gp_class = mock.MagicMock(return_value=self.gp_mock) self.celerite_likelihood = bilby.core.likelihood._GPLikelihood( - kernel=self.kernel, mean_model=self.mean_model, t=self.t, y=self.y, yerr=self.yerr, gp_class=self.gp_class) + kernel=self.kernel, + mean_model=self.mean_model, + t=self.t, + y=self.y, + yerr=self.yerr, + gp_class=self.gp_class, + ) def tearDown(self) -> None: del self.t @@ -805,39 +828,44 @@ def test_gp_class(self): def test_gp_instantiation(self): self.celerite_likelihood.GPClass.assert_called_once_with( - kernel=self.kernel, mean=self.mean_model, fit_mean=True, fit_white_noise=True) + kernel=self.kernel, + mean=self.mean_model, + fit_mean=True, + fit_white_noise=True, + ) def test_gp_mock(self): self.celerite_likelihood.gp.compute.assert_called_once_with( - self.celerite_likelihood.t, yerr=self.celerite_likelihood.yerr) + self.celerite_likelihood.t, yerr=self.celerite_likelihood.yerr + ) def test_set_parameters_no_exceptions(self): self.celerite_likelihood.gp.set_parameter = mock.MagicMock() self.celerite_likelihood.mean_model.set_parameter = mock.MagicMock() expected_a = 5 self.celerite_likelihood.set_parameters(dict(a=expected_a)) - self.celerite_likelihood.gp.set_parameter.assert_called_once_with(name="a", value=5) + self.celerite_likelihood.gp.set_parameter.assert_called_once_with( + name="a", value=5 + ) class TestFunctionMeanModel(unittest.TestCase): - def test_function_to_celerite_mean_model(self): def func(x, a, b, c): - return a * x ** 2 + b * x + c + return a * x**2 + b * x + c mean_model = bilby.core.likelihood.function_to_celerite_mean_model(func=func) self.assertListEqual(["a", "b", "c"], list(mean_model.parameter_names)) def test_function_to_george_mean_model(self): def func(x, a, b, c): - return a * x ** 2 + b * x + c + return a * x**2 + b * x + c mean_model = bilby.core.likelihood.function_to_celerite_mean_model(func=func) self.assertListEqual(["a", "b", "c"], list(mean_model.parameter_names)) class TestCeleriteLikelihoodEvaluation(unittest.TestCase): - def setUp(self) -> None: import celerite @@ -847,13 +875,23 @@ def func(x, a): self.t = [0, 1, 2] self.y = [0, 1, 2] self.yerr = [0.4, 0.5, 0.6] - self.parameters = {"kernel:log_S0": 0, "kernel:log_Q": 0, "kernel:log_omega0": 0, "mean:a": 1} + self.parameters = { + "kernel:log_S0": 0, + "kernel:log_Q": 0, + "kernel:log_omega0": 0, + "mean:a": 1, + } self.kernel = celerite.terms.SHOTerm(log_S0=0, log_Q=0, log_omega0=0) self.MeanModel = bilby.likelihood.function_to_celerite_mean_model(func=func) self.mean_model = self.MeanModel(a=1) self.celerite_likelihood = bilby.core.likelihood.CeleriteLikelihood( - kernel=self.kernel, mean_model=self.mean_model, t=self.t, y=self.y, yerr=self.yerr) + kernel=self.kernel, + mean_model=self.mean_model, + t=self.t, + y=self.y, + yerr=self.yerr, + ) def tearDown(self) -> None: del self.t @@ -871,14 +909,21 @@ def test_log_l_evalutation(self): self.assertAlmostEqual(expected, log_l, places=5) def test_set_parameters(self): - combined_params = {"kernel:log_S0": 2, "kernel:log_Q": 2, "kernel:log_omega0": 2, "mean:a": 2} + combined_params = { + "kernel:log_S0": 2, + "kernel:log_Q": 2, + "kernel:log_omega0": 2, + "mean:a": 2, + } self.celerite_likelihood.set_parameters(combined_params) - registered = {key: self.celerite_likelihood.gp.get_parameter(key) for key in combined_params} + registered = { + key: self.celerite_likelihood.gp.get_parameter(key) + for key in combined_params + } self.assertDictEqual(combined_params, registered) class TestGeorgeLikelihoodEvaluation(unittest.TestCase): - def setUp(self) -> None: import george @@ -894,7 +939,12 @@ def func(x, a): self.MeanModel = bilby.likelihood.function_to_celerite_mean_model(func=func) self.mean_model = self.MeanModel(a=1) self.george_likelihood = bilby.core.likelihood.GeorgeLikelihood( - kernel=self.kernel, mean_model=self.mean_model, t=self.t, y=self.y, yerr=self.yerr) + kernel=self.kernel, + mean_model=self.mean_model, + t=self.t, + y=self.y, + yerr=self.yerr, + ) def tearDown(self) -> None: pass @@ -905,9 +955,15 @@ def test_likelihood_value(self): self.assertAlmostEqual(expected, log_l, places=5) def test_set_parameters(self): - combined_params = {"kernel:k1:log_constant": 2, "kernel:k2:metric:log_M_0_0": 2, "mean:a": 2} + combined_params = { + "kernel:k1:log_constant": 2, + "kernel:k2:metric:log_M_0_0": 2, + "mean:a": 2, + } self.george_likelihood.set_parameters(combined_params) - registered = {key: self.george_likelihood.gp.get_parameter(key) for key in combined_params} + registered = { + key: self.george_likelihood.gp.get_parameter(key) for key in combined_params + } self.assertDictEqual(combined_params, registered) diff --git a/test/core/prior/analytical_test.py b/test/core/prior/analytical_test.py index a8a0730bd..1ebf88c8f 100644 --- a/test/core/prior/analytical_test.py +++ b/test/core/prior/analytical_test.py @@ -37,7 +37,9 @@ def test_array_sample(self): self.assertAlmostEqual(zeros / N, 1 / nvalues, places=int(np.log10(np.sqrt(N)))) self.assertAlmostEqual(ones / N, 1 / nvalues, places=int(np.log10(np.sqrt(N)))) self.assertAlmostEqual(twos / N, 1 / nvalues, places=int(np.log10(np.sqrt(N)))) - self.assertAlmostEqual(threes / N, 1 / nvalues, places=int(np.log10(np.sqrt(N)))) + self.assertAlmostEqual( + threes / N, 1 / nvalues, places=int(np.log10(np.sqrt(N))) + ) def test_single_probability(self): N = 3 @@ -67,7 +69,9 @@ def test_array_probability(self): N = 3 values = [1.1, 2.2, 300.0] discrete_value_prior = bilby.core.prior.DiscreteValues(values) - probs = discrete_value_prior.prob(self.xp.asarray([1.1, 2.2, 2.2, 300.0, 200.0])) + probs = discrete_value_prior.prob( + self.xp.asarray([1.1, 2.2, 2.2, 300.0, 200.0]) + ) self.assertEqual(aac.get_namespace(probs), self.xp) np.testing.assert_array_equal(np.asarray(probs), np.array([1 / N] * 4 + [0])) @@ -88,9 +92,13 @@ def test_array_lnprobability(self): N = 3 values = [1.1, 2.2, 300.0] discrete_value_prior = bilby.core.prior.DiscreteValues(values) - ln_probs = discrete_value_prior.ln_prob(self.xp.asarray([1.1, 2.2, 2.2, 300, 150])) + ln_probs = discrete_value_prior.ln_prob( + self.xp.asarray([1.1, 2.2, 2.2, 300, 150]) + ) self.assertEqual(aac.get_namespace(ln_probs), self.xp) - np.testing.assert_array_equal(np.asarray(ln_probs), np.array([-np.log(N)] * 4 + [-np.inf])) + np.testing.assert_array_equal( + np.asarray(ln_probs), np.array([-np.log(N)] * 4 + [-np.inf]) + ) @pytest.mark.array_backend @@ -140,9 +148,9 @@ def test_array_probability(self): probs = categorical_prior.prob(self.xp.asarray([0, 1, 1, 2, 3])) self.assertEqual(aac.get_namespace(probs), self.xp) - self.assertTrue(np.all( - np.asarray(probs) == np.array([1 / N, 1 / N, 1 / N, 1 / N, 0]) - )) + self.assertTrue( + np.all(np.asarray(probs) == np.array([1 / N, 1 / N, 1 / N, 1 / N, 0])) + ) def test_single_lnprobability(self): N = 3 @@ -161,9 +169,9 @@ def test_array_lnprobability(self): categorical_prior = bilby.core.prior.Categorical(N) ln_prob = categorical_prior.ln_prob(self.xp.asarray([0, 1, 1, 2, 3])) self.assertEqual(aac.get_namespace(ln_prob), self.xp) - self.assertTrue(np.all( - np.asarray(ln_prob) == np.array([-np.log(N)] * 4 + [-np.inf]) - )) + self.assertTrue( + np.all(np.asarray(ln_prob) == np.array([-np.log(N)] * 4 + [-np.inf])) + ) @pytest.mark.array_backend @@ -196,8 +204,12 @@ def test_array_sample(self): for i in categorical_prior.values: case = np.sum(s == i) cases += case - self.assertAlmostEqual(case / N, categorical_prior.prob(i), places=int(np.log10(np.sqrt(N)))) - self.assertAlmostEqual(case / N, weights[i] / np.sum(weights), places=int(np.log10(np.sqrt(N)))) + self.assertAlmostEqual( + case / N, categorical_prior.prob(i), places=int(np.log10(np.sqrt(N))) + ) + self.assertAlmostEqual( + case / N, weights[i] / np.sum(weights), places=int(np.log10(np.sqrt(N))) + ) self.assertEqual(cases, N) def test_single_probability(self): @@ -259,9 +271,7 @@ def test_cdf(self): categorical_prior = bilby.core.prior.WeightedCategorical(N, weights=weights) sample = categorical_prior.sample(size=10, random_state=self.rng) original = self.xp.asarray(sample) - new = self.xp.asarray(categorical_prior.rescale( - categorical_prior.cdf(sample) - )) + new = self.xp.asarray(categorical_prior.rescale(categorical_prior.cdf(sample))) np.testing.assert_array_equal(original, new) self.assertEqual(type(new), type(original)) diff --git a/test/core/prior/base_test.py b/test/core/prior/base_test.py index d2133eaca..14014e135 100644 --- a/test/core/prior/base_test.py +++ b/test/core/prior/base_test.py @@ -66,9 +66,7 @@ def test_repr_round_trips_with_numpy_scalar_parameters(self): (np.float32(0.0), np.float32(1.0)), (np.int64(0), np.int64(10)), ]: - prior = bilby.core.prior.Uniform( - minimum=minimum, maximum=maximum, name="x" - ) + prior = bilby.core.prior.Uniform(minimum=minimum, maximum=maximum, name="x") self.assertNotIn("np.", repr(prior)) reconstructed = bilby.core.prior.PriorDict({"x": repr(prior)})["x"] self.assertEqual(prior, reconstructed) @@ -171,19 +169,28 @@ def setUp(self): ) def conversion_func(parameters): - return dict(a=parameters["a"], b=parameters["b"], c=parameters["a"] / parameters["b"]) - self.priors = bilby.core.prior.PriorDict(self.priors, conversion_function=conversion_func) + return dict( + a=parameters["a"], + b=parameters["b"], + c=parameters["a"] / parameters["b"], + ) + + self.priors = bilby.core.prior.PriorDict( + self.priors, conversion_function=conversion_func + ) def test_prob_integrate_to_one(self): keys = ["a", "b", "c"] n_samples = 1000000 - samples = self.priors.sample_subset(keys=keys, size=n_samples, random_state=self.rng) + samples = self.priors.sample_subset( + keys=keys, size=n_samples, random_state=self.rng + ) prob = self.priors.prob(samples, axis=0) self.assertEqual(aac.get_namespace(prob), self.xp) prob = np.asarray(prob) dm1 = self.priors["a"].maximum - self.priors["a"].minimum dm2 = self.priors["b"].maximum - self.priors["b"].minimum - prior_volume = (dm1 * dm2) + prior_volume = dm1 * dm2 n_accepted = np.sum(prob) integral = prior_volume * n_accepted / n_samples # binomial random distribution diff --git a/test/core/prior/conditional_test.py b/test/core/prior/conditional_test.py index e5c967de2..86daf2cba 100644 --- a/test/core/prior/conditional_test.py +++ b/test/core/prior/conditional_test.py @@ -147,7 +147,9 @@ def test_rescale_ln_prob_update_conditions(self): def test_cdf_calls_update_conditions(self): self.prior = bilby.core.prior.ConditionalUniform( - condition_func=self.condition_func, minimum=self.minimum, maximum=self.maximum + condition_func=self.condition_func, + minimum=self.minimum, + maximum=self.maximum, ) with mock.patch.object(self.prior, "update_conditions") as m: self.prior.cdf( @@ -218,7 +220,9 @@ def condition_func_3(reference_parameters, var_1, var_2): var_2=self.xp.asarray(0.5), var_3=self.xp.asarray(0.4), ) - self.test_value = 1 / np.prod([self.test_sample[f"var_{ii}"] for ii in range(3)]) + self.test_value = 1 / np.prod( + [self.test_sample[f"var_{ii}"] for ii in range(3)] + ) for key, value in dict( var_0=self.prior_0, var_1=self.prior_1, @@ -279,7 +283,10 @@ def test_prob_illegal_conditions(self): self.conditional_priors.prob(sample=self.test_sample, xp=self.xp) def test_ln_prob(self): - self.assertEqual(np.log(self.test_value), self.conditional_priors.ln_prob(sample=self.test_sample)) + self.assertEqual( + np.log(self.test_value), + self.conditional_priors.ln_prob(sample=self.test_sample), + ) def test_ln_prob_illegal_conditions(self): del self.conditional_priors["var_0"] @@ -317,7 +324,7 @@ def condition_func(reference_params, a): priors["b"] = bilby.core.prior.ConditionalPowerLaw( condition_func=condition_func, minimum=1, maximum=2, alpha=-2 ) - samples = (priors.sample(2)) + samples = priors.sample(2) for key in priors: # verify the least recently sampled includes works for scalar # only condition functions @@ -352,7 +359,7 @@ def test_rescale_with_joint_prior(self): # set multivariate Gaussian distribution names = ["mvgvar_0", "mvgvar_1"] mu = [[0.79, -0.83]] - cov = [[[0.03, 0.], [0., 0.04]]] + cov = [[[0.03, 0.0], [0.0, 0.04]]] mvg = bilby.core.prior.MultivariateGaussianDist(names, mus=mu, covs=cov) priordict = bilby.core.prior.ConditionalPriorDict( @@ -390,9 +397,9 @@ def test_cdf(self): sample = self.conditional_priors.sample() np.testing.assert_array_equal( self.conditional_priors.rescale( - sample.keys(), - self.conditional_priors.cdf(sample=sample).values() - ), np.array(list(sample.values())) + sample.keys(), self.conditional_priors.cdf(sample=sample).values() + ), + np.array(list(sample.values())), ) def test_rescale_illegal_conditions(self): @@ -467,7 +474,6 @@ def _tp_conditional_uniform(ref_params, period): @pytest.mark.array_backend @pytest.mark.usefixtures("xp_class") class TestDirichletPrior(unittest.TestCase): - def setUp(self): self.priors = bilby.core.prior.DirichletPriorDict(5) diff --git a/test/core/prior/dict_test.py b/test/core/prior/dict_test.py index c27947838..9dc73d1ee 100644 --- a/test/core/prior/dict_test.py +++ b/test/core/prior/dict_test.py @@ -16,7 +16,6 @@ def __init__(self, names, mus, covs, weights): class FakeJointPriorDist(bilby.core.prior.BaseJointPriorDist): - def __init__(self, names, bounds=None): super().__init__(names=names, bounds=bounds) @@ -27,9 +26,7 @@ def __init__(self, names, bounds=None): @pytest.mark.array_backend @pytest.mark.usefixtures("xp_class") class TestPriorDict(unittest.TestCase): - def setUp(self): - self.first_prior = bilby.core.prior.Uniform( name="a", minimum=0, maximum=1, unit="kg", boundary=None ) @@ -109,12 +106,8 @@ def test_read_from_file(self): latex_label="$q$", unit=None, ), - a_1=bilby.core.prior.Uniform( - name="a_1", minimum=0, maximum=0.99 - ), - a_2=bilby.core.prior.Uniform( - name="a_2", minimum=0, maximum=0.99 - ), + a_1=bilby.core.prior.Uniform(name="a_1", minimum=0, maximum=0.99), + a_2=bilby.core.prior.Uniform(name="a_2", minimum=0, maximum=0.99), tilt_1=bilby.core.prior.Sine(name="tilt_1"), tilt_2=bilby.core.prior.Sine(name="tilt_2"), phi_12=bilby.core.prior.Uniform( @@ -145,17 +138,36 @@ def test_read_from_file(self): self.assertDictEqual(expected, self.prior_set_from_file) fake_dist = FakeJointPriorDist(names=["testAfake", "testBfake"]) - testAfake = bilby.core.prior.JointPrior(dist=fake_dist, name="testAfake", unit="unit") - testBfake = bilby.core.prior.JointPrior(dist=fake_dist, name="testBfake", unit="unit") - base_dist = bilby.core.prior.BaseJointPriorDist(names=["testAbase", "testBbase"]) - testAbase = bilby.core.prior.JointPrior(dist=base_dist, name="testAbase", unit="unit") - testBbase = bilby.core.prior.JointPrior(dist=base_dist, name="testBbase", unit="unit") - expected_joint = dict(testAfake=testAfake, testBfake=testBfake, testAbase=testAbase, testBbase=testBbase) + testAfake = bilby.core.prior.JointPrior( + dist=fake_dist, name="testAfake", unit="unit" + ) + testBfake = bilby.core.prior.JointPrior( + dist=fake_dist, name="testBfake", unit="unit" + ) + base_dist = bilby.core.prior.BaseJointPriorDist( + names=["testAbase", "testBbase"] + ) + testAbase = bilby.core.prior.JointPrior( + dist=base_dist, name="testAbase", unit="unit" + ) + testBbase = bilby.core.prior.JointPrior( + dist=base_dist, name="testBbase", unit="unit" + ) + expected_joint = dict( + testAfake=testAfake, + testBfake=testBfake, + testAbase=testAbase, + testBbase=testBbase, + ) self.assertDictEqual(expected_joint, self.joint_prior_from_file) - self.assertTrue(id(self.joint_prior_from_file["testAfake"].dist) - == id(self.joint_prior_from_file["testBfake"].dist)) - self.assertTrue(id(self.joint_prior_from_file["testAbase"].dist) - == id(self.joint_prior_from_file["testBbase"].dist)) + self.assertTrue( + id(self.joint_prior_from_file["testAfake"].dist) + == id(self.joint_prior_from_file["testBfake"].dist) + ) + self.assertTrue( + id(self.joint_prior_from_file["testAbase"].dist) + == id(self.joint_prior_from_file["testBbase"].dist) + ) def test_to_file(self): """ @@ -231,10 +243,14 @@ def test_prior_set_from_dict_but_using_a_string(self): unit=None, ), a_1=bilby.core.prior.Uniform( - name="a_1", minimum=0, maximum=0.99, + name="a_1", + minimum=0, + maximum=0.99, ), a_2=bilby.core.prior.Uniform( - name="a_2", minimum=0, maximum=0.99, + name="a_2", + minimum=0, + maximum=0.99, ), tilt_1=bilby.core.prior.Sine(name="tilt_1"), tilt_2=bilby.core.prior.Sine(name="tilt_2"), @@ -275,7 +291,8 @@ def test_dict_argument_is_not_string_or_dict(self): def test_sample_subset_correct_size(self): size = 7 samples = self.prior_set_from_dict.sample_subset( - keys=self.prior_set_from_dict.keys(), size=size, + keys=self.prior_set_from_dict.keys(), + size=size, random_state=self.rng, ) self.assertEqual(len(self.prior_set_from_dict), len(samples)) @@ -302,7 +319,6 @@ def test_sample_subset_constrained_as_array(self): self.assertTrue(out.shape == (len(keys), size)) def test_sample_subset_constrained(self): - def conversion_function(parameters): converted_parameters = parameters.copy() converted_parameters["delta_mass"] = ( @@ -366,7 +382,9 @@ def test_sample_returns_correct_type(self): self.assertEqual(aac.get_namespace(samples[key]), self.xp) def test_prob(self): - samples = self.prior_set_from_dict.sample_subset(keys=["mass", "speed"], random_state=self.rng) + samples = self.prior_set_from_dict.sample_subset( + keys=["mass", "speed"], random_state=self.rng + ) expected = self.first_prior.prob(samples["mass"]) * self.second_prior.prob( samples["speed"] ) @@ -374,7 +392,9 @@ def test_prob(self): self.assertEqual(aac.get_namespace(expected), self.xp) def test_ln_prob(self): - samples = self.prior_set_from_dict.sample_subset(keys=["mass", "speed"], random_state=self.rng) + samples = self.prior_set_from_dict.sample_subset( + keys=["mass", "speed"], random_state=self.rng + ) expected = self.first_prior.ln_prob( samples["mass"] ) + self.second_prior.ln_prob(samples["speed"]) @@ -405,10 +425,11 @@ def test_cdf(self): """ sample = self.prior_set_from_dict.sample(random_state=self.rng) original = self.xp.asarray(list(sample.values())) - new = self.xp.asarray(self.prior_set_from_dict.rescale( - sample.keys(), - self.prior_set_from_dict.cdf(sample=sample).values() - )) + new = self.xp.asarray( + self.prior_set_from_dict.rescale( + sample.keys(), self.prior_set_from_dict.cdf(sample=sample).values() + ) + ) self.assertLess(max(abs(original - new)), 1e-10) self.assertEqual(aac.get_namespace(new), self.xp) @@ -496,7 +517,9 @@ def setUp(self): ), pp=bilby.core.prior.HalfGaussian(name="test", unit="unit", sigma=1), qq=bilby.core.prior.HalfNormal(name="test", unit="unit", sigma=1), - rr=bilby.core.prior.LogGaussian(name="test", unit="unit", mu=0, sigma=1), + rr=bilby.core.prior.LogGaussian( + name="test", unit="unit", mu=0, sigma=1 + ), ss=bilby.core.prior.LogNormal(name="test", unit="unit", mu=0, sigma=1), tt=bilby.core.prior.Exponential(name="test", unit="unit", mu=1), uu=bilby.core.prior.StudentT( @@ -554,7 +577,7 @@ def setUp(self): ) def test_read_write_to_json(self): - """ Interped prior is removed as there is numerical error in the recovered prior.""" + """Interped prior is removed as there is numerical error in the recovered prior.""" self.priors.to_json(outdir="prior_files", label="json_test") new_priors = bilby.core.prior.PriorDict.from_json( filename="prior_files/json_test_prior.json" @@ -565,7 +588,9 @@ def test_read_write_to_json(self): self.assertLess(max(abs(old_interped.xx - new_interped.xx)), 1e-15) self.assertLess(max(abs(old_interped.yy - new_interped.yy)), 1e-15) self.assertTrue(id(new_priors["testa"].dist) == id(new_priors["testb"].dist)) - self.assertTrue(id(new_priors["testAfake"].dist) == id(new_priors["testBfake"].dist)) + self.assertTrue( + id(new_priors["testAfake"].dist) == id(new_priors["testBfake"].dist) + ) class TestLoadPrior(unittest.TestCase): @@ -622,10 +647,13 @@ def test_constrained_uniform_ln_prob(xp, positive_uniform_prior, normalized, val values = xp.asarray(values) expected = xp.where(values > 0, 0.0 if normalized else -np.log(2), -xp.inf) with patch.object( - positive_uniform_prior, "normalize_constraint_factor", + positive_uniform_prior, + "normalize_constraint_factor", wraps=positive_uniform_prior.normalize_constraint_factor, ) as normalization: - actual = positive_uniform_prior.ln_prob({"x": values}, axis=0, normalized=normalized) + actual = positive_uniform_prior.ln_prob( + {"x": values}, axis=0, normalized=normalized + ) assert normalization.call_count == int(normalized) assert xp.asarray(actual).shape == values.shape assert xp.allclose(actual, expected) @@ -663,7 +691,9 @@ def test_check_ln_prob_jax_jit_and_grad(xp, positive_uniform_prior): @jax.jit def log_probability(ln_density): return positive_uniform_prior.check_ln_prob( - {"x": values}, ln_density, xp=xp, + {"x": values}, + ln_density, + xp=xp, ) ln_density = xp.full_like(values, -np.log(2)) @@ -683,7 +713,10 @@ def log_probability(ln_density): @pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("values", [0.5, [-0.5, 0.5]]) def test_check_ln_prob_torch_preserves_tensor( - positive_uniform_prior, dtype, device, values, + positive_uniform_prior, + dtype, + device, + values, ): torch = pytest.importorskip("torch") if device == "cuda" and not torch.cuda.is_available(): @@ -692,10 +725,14 @@ def test_check_ln_prob_torch_preserves_tensor( ln_density = torch.full_like(values, -np.log(2), requires_grad=True) xp = aac.get_namespace(ln_density) actual = positive_uniform_prior.check_ln_prob( - {"x": values}, ln_density, xp=xp, + {"x": values}, + ln_density, + xp=xp, ) expected = torch.where( - values > 0, torch.zeros_like(values), -torch.inf, + values > 0, + torch.zeros_like(values), + -torch.inf, ) torch.testing.assert_close(actual, expected) assert actual.dtype == ln_density.dtype @@ -774,11 +811,10 @@ def test_without_available_default_priors_no_prior_is_set(self): class TestLoadPriorWithCosmologicalParameters(unittest.TestCase): - def test_load(self): prior_file = os.path.join( os.path.dirname(os.path.realpath(__file__)), - "prior_files/prior_with_cosmo_params.prior" + "prior_files/prior_with_cosmo_params.prior", ) prior_dict = bilby.gw.prior.BBHPriorDict(filename=prior_file) cosmology = prior_dict["luminosity_distance"].cosmology diff --git a/test/core/prior/joint_test.py b/test/core/prior/joint_test.py index c99373b00..24371c2e0 100644 --- a/test/core/prior/joint_test.py +++ b/test/core/prior/joint_test.py @@ -49,10 +49,8 @@ def _test_mvg_ln_prob_diff_expected(self, mvg, mus, sigmas, corrcoefs): for i in np.ndindex(4, 4, 4): ln_prob_mus_sigmas_d_i = mvg.ln_prob(mus + sigmas * (d @ i)) diff_ln_prob = ln_prob_mus - ln_prob_mus_sigmas_d_i - diff_ln_prob_expected = 0.5 * np.sum(np.array(i)**2) - self.assertTrue( - np.allclose(diff_ln_prob, diff_ln_prob_expected) - ) + diff_ln_prob_expected = 0.5 * np.sum(np.array(i) ** 2) + self.assertTrue(np.allclose(diff_ln_prob, diff_ln_prob_expected)) def test_mvg_unit_scales(self): # test using order-unity standard deviations and correlations @@ -60,7 +58,7 @@ def test_mvg_unit_scales(self): corrcoefs = np.identity(3) mus = np.array([3, 1, 2]) mvg = bilby.core.prior.MultivariateGaussianDist( - names=['a', 'b', 'c'], + names=["a", "b", "c"], mus=mus, sigmas=sigmas, corrcoefs=corrcoefs, @@ -74,11 +72,13 @@ def test_mvg_cw_scales(self): # parameters of a continuous wave signal T = 365.25 * 86400 rho = 10 - sigmas = np.array([ - 5 * np.sqrt(3) / (2 * np.pi * T * rho), - 6 * np.sqrt(5) / (np.pi * T**2 * rho), - 60 * np.sqrt(7) / (np.pi * T**3 * rho) - ]) + sigmas = np.array( + [ + 5 * np.sqrt(3) / (2 * np.pi * T * rho), + 6 * np.sqrt(5) / (np.pi * T**2 * rho), + 60 * np.sqrt(7) / (np.pi * T**3 * rho), + ] + ) corrcoefs = np.identity(3) corrcoefs[0, 2] = corrcoefs[2, 0] = -np.sqrt(21) / 5 diff --git a/test/core/prior/prior_test.py b/test/core/prior/prior_test.py index e3cd94dba..51455b8d2 100644 --- a/test/core/prior/prior_test.py +++ b/test/core/prior/prior_test.py @@ -21,9 +21,7 @@ os.path.dirname(os.path.realpath(__file__)), "prior_files/GW150914_testing_skymap.fits", ) -hp_dist = bilby.gw.prior.HealPixMapPriorDist( - hp_map_file, names=["testra", "testdec"] -) +hp_dist = bilby.gw.prior.HealPixMapPriorDist(hp_map_file, names=["testra", "testdec"]) hp_3d_dist = bilby.gw.prior.HealPixMapPriorDist( hp_map_file, names=["testra", "testdec", "testdistance"], distance=True ) @@ -49,7 +47,11 @@ def setUp(self): # need to reset this for the repr test to get equality correct hp_dist.requested_parameters = {"testra": None, "testdec": None} - hp_3d_dist.requested_parameters = {"testra": None, "testdec": None, "testdistance": None} + hp_3d_dist.requested_parameters = { + "testra": None, + "testdec": None, + "testdistance": None, + } def condition_func(reference_params, test_param): return reference_params.copy() @@ -88,7 +90,9 @@ def condition_func(reference_params, test_param): maximum=5, ), bilby.core.prior.FromFile( - file_name=os.path.join(os.path.dirname(__file__), "prior_files/fromfile.txt"), + file_name=os.path.join( + os.path.dirname(__file__), "prior_files/fromfile.txt" + ), minimum=0, maximum=1, unit="unit", @@ -263,13 +267,11 @@ def condition_func(reference_params, test_param): ] if aac.is_torch_namespace(self.xp): self.priors = [ - p for p in self.priors - if not isinstance(p, bilby.core.prior.Interped) + p for p in self.priors if not isinstance(p, bilby.core.prior.Interped) ] elif aac.is_jax_namespace(self.xp): self.priors = [ - p for p in self.priors - if not isinstance(p, bilby.core.prior.StudentT) + p for p in self.priors if not isinstance(p, bilby.core.prior.StudentT) ] def tearDown(self): @@ -296,7 +298,9 @@ def evaluate_prior(prior_, val): sample = jax.numpy.asarray(alt_prior.sample(3)) evaluate_prior(alt_prior, sample) new_cache_size = evaluate_prior._cache_size() - message = f"Cache size increased by more than 1 for {prior.__class__.__name__}" + message = ( + f"Cache size increased by more than 1 for {prior.__class__.__name__}" + ) assert new_cache_size <= cache_size + 1, message def test_grad(self): @@ -309,15 +313,18 @@ def test_grad(self): for prior in self.priors: # gradients don't work yet for some priors - if isinstance(prior, ( - bilby.core.prior.Beta, - bilby.core.prior.ConditionalPowerLaw, - bilby.core.prior.Gamma, - bilby.core.prior.Interped, - bilby.core.prior.JointPrior, - bilby.core.prior.Triangular, - bilby.core.prior.WeightedDiscreteValues, - )): + if isinstance( + prior, + ( + bilby.core.prior.Beta, + bilby.core.prior.ConditionalPowerLaw, + bilby.core.prior.Gamma, + bilby.core.prior.Interped, + bilby.core.prior.JointPrior, + bilby.core.prior.Triangular, + bilby.core.prior.WeightedDiscreteValues, + ), + ): continue sample = jnp.asarray(prior.sample()) for method in [prior.prob, prior.ln_prob, prior.cdf]: @@ -342,15 +349,21 @@ def test_minimum_rescaling(self): # the edge of the prior is extremely suppressed for these priors # and so the rescale function doesn't quite return the lower bound continue - if isinstance(prior, bilby.gw.prior.HealPixPrior) and aac.is_torch_namespace(self.xp): + if isinstance( + prior, bilby.gw.prior.HealPixPrior + ) and aac.is_torch_namespace(self.xp): # HealPix rescaling requires interpolation continue with self.subTest(prior=prior): if bilby.core.prior.JointPrior in prior.__class__.__mro__: minimum_sample = prior.rescale(self.xp.asarray(0)) if prior.dist.filled_rescale(): - self.assertAlmostEqual(np.asarray(minimum_sample[0]), prior.minimum) - self.assertAlmostEqual(np.asarray(minimum_sample[1]), prior.minimum) + self.assertAlmostEqual( + np.asarray(minimum_sample[0]), prior.minimum + ) + self.assertAlmostEqual( + np.asarray(minimum_sample[1]), prior.minimum + ) else: minimum_sample = prior.rescale(self.xp.asarray(0)) self.assertAlmostEqual(np.asarray(minimum_sample), prior.minimum) @@ -358,15 +371,21 @@ def test_minimum_rescaling(self): def test_maximum_rescaling(self): """Test the the rescaling works as expected.""" for prior in self.priors: - if isinstance(prior, bilby.gw.prior.HealPixPrior) and aac.is_torch_namespace(self.xp): + if isinstance( + prior, bilby.gw.prior.HealPixPrior + ) and aac.is_torch_namespace(self.xp): # HealPix rescaling requires interpolation continue with self.subTest(prior=prior): if bilby.core.prior.JointPrior in prior.__class__.__mro__: maximum_sample = prior.rescale(self.xp.asarray(0)) if prior.dist.filled_rescale(): - self.assertAlmostEqual(np.asarray(maximum_sample[0]), prior.maximum) - self.assertAlmostEqual(np.asarray(maximum_sample[1]), prior.maximum) + self.assertAlmostEqual( + np.asarray(maximum_sample[0]), prior.maximum + ) + self.assertAlmostEqual( + np.asarray(maximum_sample[1]), prior.maximum + ) elif isinstance(prior, bilby.gw.prior.AlignedSpin): maximum_sample = prior.rescale(self.xp.asarray(1)) self.assertGreater(np.asarray(maximum_sample), 0.997) @@ -380,7 +399,9 @@ def test_many_sample_rescaling(self): if isinstance(prior, bilby.core.prior.analytical.SymmetricLogUniform): # SymmetricLogUniform has support down to -maximum continue - if isinstance(prior, bilby.gw.prior.HealPixPrior) and aac.is_torch_namespace(self.xp): + if isinstance( + prior, bilby.gw.prior.HealPixPrior + ) and aac.is_torch_namespace(self.xp): # HealPix rescaling requires interpolation continue many_samples = prior.rescale(self.xp.asarray(np.random.uniform(0, 1, 1000))) @@ -389,7 +410,10 @@ def test_many_sample_rescaling(self): continue with self.subTest(prior=prior): self.assertTrue( - all((many_samples >= prior.minimum) & (many_samples <= prior.maximum)) + all( + (many_samples >= prior.minimum) + & (many_samples <= prior.maximum) + ) ) self._validate_return_type(many_samples) @@ -527,7 +551,9 @@ def test_cdf_is_inverse_of_rescaling(self): def test_cdf_one_above_domain(self): for prior in self.priors: - if isinstance(prior, bilby.gw.prior.HealPixPrior) and aac.is_torch_namespace(self.xp): + if isinstance( + prior, bilby.gw.prior.HealPixPrior + ) and aac.is_torch_namespace(self.xp): # HealPix rescaling requires interpolation continue if prior.maximum == np.inf: @@ -543,7 +569,9 @@ def test_cdf_zero_below_domain(self): if isinstance(prior, bilby.core.prior.analytical.SymmetricLogUniform): # SymmetricLogUniform has support down to -maximum continue - if isinstance(prior, bilby.gw.prior.HealPixPrior) and aac.is_torch_namespace(self.xp): + if isinstance( + prior, bilby.gw.prior.HealPixPrior + ) and aac.is_torch_namespace(self.xp): # HealPix rescaling requires interpolation continue if ( @@ -578,10 +606,10 @@ def test_studentt_fail(self): def test_beta_fail(self): with self.assertRaises(ValueError): - bilby.core.prior.Beta(name="test", unit="unit", alpha=-2.0, beta=2.0), + (bilby.core.prior.Beta(name="test", unit="unit", alpha=-2.0, beta=2.0),) with self.assertRaises(ValueError): - bilby.core.prior.Beta(name="test", unit="unit", alpha=2.0, beta=-2.0), + (bilby.core.prior.Beta(name="test", unit="unit", alpha=2.0, beta=-2.0),) def test_multivariate_gaussian_fail(self): with self.assertRaises(ValueError): @@ -727,7 +755,9 @@ def test_probability_surrounding_domain(self): # SymmetricLogUniform has support down to -maximum continue with np.errstate(invalid="ignore"): - surround_domain = self.xp.linspace(prior.minimum - 1, prior.maximum + 1, 1000) + surround_domain = self.xp.linspace( + prior.minimum - 1, prior.maximum + 1, 1000 + ) indomain = (surround_domain >= prior.minimum) | ( surround_domain <= prior.maximum ) @@ -889,10 +919,18 @@ def test_accuracy(self): ) if isinstance(prior, (testTuple)): with self.subTest(prior=prior): - np.testing.assert_almost_equal(prior.prob(self.xp.asarray(domain)), scipy_prob) - np.testing.assert_almost_equal(prior.ln_prob(self.xp.asarray(domain)), scipy_lnprob) - np.testing.assert_almost_equal(prior.cdf(self.xp.asarray(domain)), scipy_cdf) - if isinstance(prior, bilby.core.prior.StudentT) and "jax" in str(self.xp): + np.testing.assert_almost_equal( + prior.prob(self.xp.asarray(domain)), scipy_prob + ) + np.testing.assert_almost_equal( + prior.ln_prob(self.xp.asarray(domain)), scipy_lnprob + ) + np.testing.assert_almost_equal( + prior.cdf(self.xp.asarray(domain)), scipy_cdf + ) + if isinstance(prior, bilby.core.prior.StudentT) and "jax" in str( + self.xp + ): # JAX implementation of StudentT prior rescale is not accurate enough continue np.testing.assert_almost_equal( @@ -986,7 +1024,9 @@ def test_set_maximum_setting(self): continue with self.subTest(prior=prior): prior.maximum = (prior.maximum + prior.minimum) / 2 - self.assertTrue(max(prior.sample(10000, random_state=self.rng)) < prior.maximum) + self.assertTrue( + max(prior.sample(10000, random_state=self.rng)) < prior.maximum + ) def test_set_minimum_setting(self): for prior in self.priors: @@ -1013,7 +1053,9 @@ def test_set_minimum_setting(self): continue with self.subTest(prior=prior): prior.minimum = (prior.maximum + prior.minimum) / 2 - self.assertTrue(min(prior.sample(10000, random_state=self.rng)) > prior.minimum) + self.assertTrue( + min(prior.sample(10000, random_state=self.rng)) > prior.minimum + ) if __name__ == "__main__": diff --git a/test/core/prior/slabspike_test.py b/test/core/prior/slabspike_test.py index 7c5716b8a..5d27c2368 100644 --- a/test/core/prior/slabspike_test.py +++ b/test/core/prior/slabspike_test.py @@ -6,13 +6,27 @@ import bilby from bilby.core.prior.slabspike import SlabSpikePrior -from bilby.core.prior.analytical import Uniform, PowerLaw, LogUniform, TruncatedGaussian, \ - Beta, Gaussian, Cosine, Sine, HalfGaussian, LogNormal, Exponential, StudentT, Logistic, \ - Cauchy, Gamma, ChiSquared +from bilby.core.prior.analytical import ( + Uniform, + PowerLaw, + LogUniform, + TruncatedGaussian, + Beta, + Gaussian, + Cosine, + Sine, + HalfGaussian, + LogNormal, + Exponential, + StudentT, + Logistic, + Cauchy, + Gamma, + ChiSquared, +) class TestSlabSpikePrior(unittest.TestCase): - def setUp(self): self.minimum = 0 self.maximum = 1 @@ -20,7 +34,10 @@ def setUp(self): self.spike_height = 0.3 self.slab = bilby.core.prior.Prior(minimum=self.minimum, maximum=self.maximum) self.prior = SlabSpikePrior( - slab=self.slab, spike_location=self.spike_loc, spike_height=self.spike_height) + slab=self.slab, + spike_location=self.spike_loc, + spike_height=self.spike_height, + ) def tearDown(self): del self.minimum @@ -66,7 +83,6 @@ def test_set_spike_height_domain_edge(self): @pytest.mark.array_backend @pytest.mark.usefixtures("xp_class") class TestSlabSpikeClasses(unittest.TestCase): - def setUp(self): self.minimum = self.xp.asarray(0.4) self.maximum = self.xp.asarray(2.4) @@ -77,28 +93,41 @@ def setUp(self): Uniform(minimum=self.minimum, maximum=self.maximum), PowerLaw(minimum=self.minimum, maximum=self.maximum, alpha=2), LogUniform(minimum=self.minimum, maximum=self.maximum), - TruncatedGaussian(minimum=self.minimum, maximum=self.maximum, mu=0, sigma=1), + TruncatedGaussian( + minimum=self.minimum, maximum=self.maximum, mu=0, sigma=1 + ), Beta(minimum=self.minimum, maximum=self.maximum, alpha=1, beta=1), Gaussian(mu=0, sigma=1), - Cosine(minimum=self.xp.asarray(-np.pi / 2), maximum=self.xp.asarray(np.pi / 2)), + Cosine( + minimum=self.xp.asarray(-np.pi / 2), maximum=self.xp.asarray(np.pi / 2) + ), Sine(minimum=self.xp.asarray(0), maximum=self.xp.asarray(np.pi)), HalfGaussian(sigma=1), LogNormal(mu=1, sigma=2), Exponential(mu=2), Logistic(mu=2, scale=1), Cauchy(alpha=1, beta=2), - Gamma(k=1, theta=1.), + Gamma(k=1, theta=1.0), ChiSquared(nu=2), ] if not aac.is_jax_namespace(self.xp): - StudentT(df=2), - self.slab_spikes = [SlabSpikePrior(slab, spike_height=self.spike_height, spike_location=self.spike_loc) - for slab in self.slabs] - self.test_nodes_finite_support = self.xp.linspace(self.minimum, self.maximum, 1000) + (StudentT(df=2),) + self.slab_spikes = [ + SlabSpikePrior( + slab, spike_height=self.spike_height, spike_location=self.spike_loc + ) + for slab in self.slabs + ] + self.test_nodes_finite_support = self.xp.linspace( + self.minimum, self.maximum, 1000 + ) self.test_nodes_infinite_support = self.xp.linspace(-10, 10, 1000) - self.test_nodes = [self.test_nodes_finite_support - if np.isinf(slab.minimum) or np.isinf(slab.maximum) - else self.test_nodes_finite_support for slab in self.slabs] + self.test_nodes = [ + self.test_nodes_finite_support + if np.isinf(slab.minimum) or np.isinf(slab.maximum) + else self.test_nodes_finite_support + for slab in self.slabs + ] def tearDown(self): del self.minimum @@ -110,7 +139,9 @@ def tearDown(self): del self.test_nodes_infinite_support def test_prob_on_slab(self): - for slab, slab_spike, test_nodes in zip(self.slabs, self.slab_spikes, self.test_nodes): + for slab, slab_spike, test_nodes in zip( + self.slabs, self.slab_spikes, self.test_nodes + ): expected = slab.prob(test_nodes) * slab_spike.slab_fraction actual = slab_spike.prob(test_nodes) self.assertTrue(np.allclose(expected, actual, rtol=1e-5)) @@ -121,7 +152,9 @@ def test_prob_on_spike(self): self.assertEqual(np.inf, slab_spike.prob(self.spike_loc)) def test_ln_prob_on_slab(self): - for slab, slab_spike, test_nodes in zip(self.slabs, self.slab_spikes, self.test_nodes): + for slab, slab_spike, test_nodes in zip( + self.slabs, self.slab_spikes, self.test_nodes + ): expected = slab.ln_prob(test_nodes) + np.log(slab_spike.slab_fraction) actual = slab_spike.ln_prob(test_nodes) self.assertTrue(np.array_equal(expected, actual)) @@ -135,12 +168,16 @@ def test_ln_prob_on_spike(self): def test_inverse_cdf_below_spike_with_spike_at_minimum(self): for slab in self.slabs: - slab_spike = SlabSpikePrior(slab=slab, spike_height=0.4, spike_location=slab.minimum) + slab_spike = SlabSpikePrior( + slab=slab, spike_height=0.4, spike_location=slab.minimum + ) self.assertEqual(0, slab_spike.inverse_cdf_below_spike) def test_inverse_cdf_below_spike_with_spike_at_maximum(self): for slab in self.slabs: - slab_spike = SlabSpikePrior(slab=slab, spike_height=0.4, spike_location=slab.maximum) + slab_spike = SlabSpikePrior( + slab=slab, spike_height=0.4, spike_location=slab.maximum + ) expected = 1 - slab_spike.spike_height actual = slab_spike.inverse_cdf_below_spike self.assertEqual(expected, actual) @@ -149,7 +186,9 @@ def test_inverse_cdf_below_spike_arbitrary_position(self): pass def test_cdf_below_spike(self): - for slab, slab_spike, test_nodes in zip(self.slabs, self.slab_spikes, self.test_nodes): + for slab, slab_spike, test_nodes in zip( + self.slabs, self.slab_spikes, self.test_nodes + ): test_nodes = test_nodes[np.where(test_nodes < self.spike_loc)] expected = slab.cdf(test_nodes) * slab_spike.slab_fraction actual = slab_spike.cdf(test_nodes) @@ -164,9 +203,13 @@ def test_cdf_at_spike(self): self.assertEqual(aac.get_namespace(actual), self.xp) def test_cdf_above_spike(self): - for slab, slab_spike, test_nodes in zip(self.slabs, self.slab_spikes, self.test_nodes): + for slab, slab_spike, test_nodes in zip( + self.slabs, self.slab_spikes, self.test_nodes + ): test_nodes = test_nodes[np.where(test_nodes > self.spike_loc)] - expected = slab.cdf(test_nodes) * slab_spike.slab_fraction + self.spike_height + expected = ( + slab.cdf(test_nodes) * slab_spike.slab_fraction + self.spike_height + ) actual = slab_spike.cdf(test_nodes) np.testing.assert_allclose(expected, actual, rtol=1e-12) self.assertEqual(aac.get_namespace(actual), self.xp) @@ -185,7 +228,9 @@ def test_cdf_at_maximum(self): def test_rescale_no_spike(self): for slab in self.slabs: - slab_spike = SlabSpikePrior(slab=slab, spike_height=0, spike_location=slab.minimum) + slab_spike = SlabSpikePrior( + slab=slab, spike_height=0, spike_location=slab.minimum + ) vals = self.xp.linspace(0, 1, 1000) expected = slab.rescale(vals) actual = slab_spike.rescale(vals) @@ -207,18 +252,24 @@ def test_rescale_at_spike(self): for slab, slab_spike in zip(self.slabs, self.slab_spikes): vals = self.xp.linspace( slab_spike.inverse_cdf_below_spike, - slab_spike.inverse_cdf_below_spike + slab_spike.spike_height, 1000 + slab_spike.inverse_cdf_below_spike + slab_spike.spike_height, + 1000, + ) + expected = self.xp.ones(len(vals)) * slab.rescale( + vals[0] / slab_spike.slab_fraction ) - expected = self.xp.ones(len(vals)) * slab.rescale(vals[0] / slab_spike.slab_fraction) actual = slab_spike.rescale(vals) self.assertTrue(np.allclose(expected, actual, rtol=1e-5)) self.assertEqual(aac.get_namespace(actual), self.xp) def test_rescale_above_spike(self): for slab, slab_spike in zip(self.slabs, self.slab_spikes): - vals = self.xp.linspace(slab_spike.inverse_cdf_below_spike + self.spike_height, 1, 1000) + vals = self.xp.linspace( + slab_spike.inverse_cdf_below_spike + self.spike_height, 1, 1000 + ) expected = self.xp.ones(len(vals)) * slab.rescale( - (vals - self.spike_height) / slab_spike.slab_fraction) + (vals - self.spike_height) / slab_spike.slab_fraction + ) actual = slab_spike.rescale(vals) self.assertTrue(np.allclose(expected, actual, rtol=1e-5)) self.assertEqual(aac.get_namespace(actual), self.xp) diff --git a/test/core/result_test.py b/test/core/result_test.py index 94a2f5d5e..c7d2b8cb4 100644 --- a/test/core/result_test.py +++ b/test/core/result_test.py @@ -16,7 +16,6 @@ @pytest.mark.array_backend @pytest.mark.usefixtures("xp_class") class TestJson(unittest.TestCase): - def setUp(self): self.encoder = bilby.core.utils.BilbyJsonEncoder self.decoder = bilby.core.utils.decode_bilby_json @@ -56,7 +55,6 @@ def test_dataframe_encoding(self): class TestResult(unittest.TestCase): - @pytest.fixture(autouse=True) def init_outdir(self, tmp_path): # Use pytest's tmp_path fixture to create a temporary directory @@ -143,9 +141,7 @@ def test_result_file_name_pickle(self): ) def test_fail_save_and_load_missing_inputs(self): - with self.assertRaises( - ValueError, msg="No information given to load file" - ): + with self.assertRaises(ValueError, msg="No information given to load file"): bilby.core.result.read_in_result() def test_fail_save_and_load_no_extension(self): @@ -239,21 +235,24 @@ def test_unset_posterior(self): _ = self.result.posterior def test_save_and_load_json(self): - self._save_and_load_test(extension='json') + self._save_and_load_test(extension="json") def test_save_and_load_json_gzip(self): - self._save_and_load_test(extension='json', gzip=True) + self._save_and_load_test(extension="json", gzip=True) def test_save_and_load_pkl(self): - self._save_and_load_test(extension='pkl') + self._save_and_load_test(extension="pkl") def test_save_and_load_hdf5(self): - self._save_and_load_test(extension='hdf5') + self._save_and_load_test(extension="hdf5") def _save_and_load_test(self, extension, gzip=False): self.result.save_to_file(extension=extension, gzip=gzip) loaded_result = bilby.core.result.read_in_result( - outdir=self.result.outdir, label=self.result.label, extension=extension, gzip=gzip + outdir=self.result.outdir, + label=self.result.label, + extension=extension, + gzip=gzip, ) self.assertTrue( np.array_equal( @@ -284,26 +283,34 @@ def _save_and_load_test(self, extension, gzip=False): self.assertEqual(self.result.sampling_time, loaded_result.sampling_time) def test_save_and_dont_overwrite_json(self): - self._save_and_dont_overwrite_test(extension='json') + self._save_and_dont_overwrite_test(extension="json") def test_save_and_dont_overwrite_pkl(self): - self._save_and_dont_overwrite_test(extension='pkl') + self._save_and_dont_overwrite_test(extension="pkl") def test_save_and_dont_overwrite_hdf5(self): - self._save_and_dont_overwrite_test(extension='hdf5') + self._save_and_dont_overwrite_test(extension="hdf5") def _save_and_dont_overwrite_test(self, extension): self.result.save_to_file(overwrite=False, extension=extension) self.result.save_to_file(overwrite=False, extension=extension) - self.assertTrue(os.path.isfile(f"{self.result.outdir}/{self.result.label}_result.{extension}.old")) + self.assertTrue( + os.path.isfile( + f"{self.result.outdir}/{self.result.label}_result.{extension}.old" + ) + ) def _save_with_outdir_and_filename(self, filename, outdir, template): - self.result.save_to_file(filename=filename, outdir=outdir, extension="json", gzip=False) + self.result.save_to_file( + filename=filename, outdir=outdir, extension="json", gzip=False + ) self.assertTrue(os.path.isfile(template)) def test_save_with_outdir_and_filename_different_outdir(self): self._save_with_outdir_and_filename( - f"{self.other_outdir}/result", self.other_outdir2, f"{self.other_outdir2}/result" + f"{self.other_outdir}/result", + self.other_outdir2, + f"{self.other_outdir2}/result", ) def test_save_with_outdir_and_filename_same_outdir(self): @@ -323,27 +330,35 @@ def test_save_with_filename_only(self): def test_save_with_outdir_no_filename(self): self._save_with_outdir_and_filename( - None, self.other_outdir, os.path.join(self.other_outdir, f"{self.result.label}_result.json") + None, + self.other_outdir, + os.path.join(self.other_outdir, f"{self.result.label}_result.json"), ) def test_save_no_filename_or_outdir(self): self._save_with_outdir_and_filename( - None, None, os.path.join(self.result.outdir, f"{self.result.label}_result.json") + None, + None, + os.path.join(self.result.outdir, f"{self.result.label}_result.json"), ) def test_save_and_overwrite_json(self): - self._save_and_overwrite_test(extension='json') + self._save_and_overwrite_test(extension="json") def test_save_and_overwrite_pkl(self): - self._save_and_overwrite_test(extension='pkl') + self._save_and_overwrite_test(extension="pkl") def test_save_and_overwrite_hdf5(self): - self._save_and_overwrite_test(extension='hdf5') + self._save_and_overwrite_test(extension="hdf5") def _save_and_overwrite_test(self, extension): self.result.save_to_file(overwrite=True, extension=extension) self.result.save_to_file(overwrite=True, extension=extension) - self.assertFalse(os.path.isfile(f"{self.result.outdir}/{self.result.label}_result.{extension}.old")) + self.assertFalse( + os.path.isfile( + f"{self.result.outdir}/{self.result.label}_result.{extension}.old" + ) + ) def test_save_samples(self): self.result.save_posterior_samples() @@ -530,18 +545,24 @@ def test_to_arviz(self): self.assertTrue("x" in az.posterior and "y" in az.posterior) for var in ["x", "y"]: - self.assertTrue(np.array_equal(az.posterior[var].values.squeeze(), - self.result.posterior[var].values)) + self.assertTrue( + np.array_equal( + az.posterior[var].values.squeeze(), + self.result.posterior[var].values, + ) + ) self.assertTrue(len(np.squeeze(az.prior[var])) == Nprior) - self.assertTrue(np.array_equal(az.log_likelihood["log_likelihood"].values.squeeze(), - log_likelihood)) - self.assertTrue( - az.posterior.attrs["inference_library"] == "bilby: {}".format( - self.result.sampler + np.array_equal( + az.log_likelihood["log_likelihood"].values.squeeze(), log_likelihood ) ) + + self.assertTrue( + az.posterior.attrs["inference_library"] + == "bilby: {}".format(self.result.sampler) + ) self.assertTrue( az.posterior.attrs["inference_library_version"] == bilby.utils.get_version_information() @@ -551,12 +572,14 @@ def test_to_arviz(self): del az self.result.posterior["log_likelihood"] = log_likelihood az = self.result.to_arviz() - self.assertTrue(np.array_equal(az.log_likelihood["log_likelihood"].values.squeeze(), - log_likelihood)) + self.assertTrue( + np.array_equal( + az.log_likelihood["log_likelihood"].values.squeeze(), log_likelihood + ) + ) @patch("builtins.__import__") def test_to_arviz_not_installed(self, mock_import): - def import_side_effect(name, *args): if name == "arviz": raise ImportError @@ -569,17 +592,16 @@ def import_side_effect(name, *args): self.assertEqual( str(excinfo.exception), - "ArviZ is not installed, so cannot convert to InferenceData." + "ArviZ is not installed, so cannot convert to InferenceData.", ) def test_result_caching(self): - class SimpleLikelihood(bilby.Likelihood): def __init__(self): super().__init__() def log_likelihood(self, parameters): - return -parameters["x"]**2 + return -(parameters["x"] ** 2) likelihood = SimpleLikelihood() priors = dict(x=bilby.core.prior.Uniform(-5, 5, "x")) @@ -592,14 +614,14 @@ class NotAResult(bilby.core.result.Result): result = bilby.run_sampler( likelihood, priors, - sampler='bilby_mcmc', + sampler="bilby_mcmc", outdir=self.outdir, nsamples=10, L1steps=1, proposal_cycle="default_noGMnoKD", printdt=1, check_point_plot=False, - result_class=NotAResult + result_class=NotAResult, ) assert isinstance(result, NotAResult) @@ -609,7 +631,7 @@ class NotAResult(bilby.core.result.Result): cached_result = bilby.run_sampler( likelihood, priors, - sampler='bilby_mcmc', + sampler="bilby_mcmc", outdir=self.outdir, nsamples=10, L1steps=1, @@ -645,9 +667,7 @@ def setUp(self): np.random.seed(7) lal_dict = lal.CreateDict() lal.DictInsertREAL8Value(lal_dict, "test_value", 1.23) - priors = bilby.prior.PriorDict( - dict(x=bilby.prior.Uniform(0, 1, "x")) - ) + priors = bilby.prior.PriorDict(dict(x=bilby.prior.Uniform(0, 1, "x"))) result = bilby.core.result.Result( label="label", outdir=self.outdir, @@ -888,17 +908,18 @@ def test_sanity_check_labels(self): class TestPPPlots(unittest.TestCase): - @pytest.fixture(autouse=True) def init_outdir(self, tmp_path): # Use pytest's tmp_path fixture to create a temporary directory self.outdir = str(tmp_path / "test_pp_plots") def setUp(self): - priors = bilby.core.prior.PriorDict(dict( - a=bilby.core.prior.Uniform(0, 1, latex_label="$a$"), - b=bilby.core.prior.Uniform(0, 1, latex_label="$b$"), - )) + priors = bilby.core.prior.PriorDict( + dict( + a=bilby.core.prior.Uniform(0, 1, latex_label="$a$"), + b=bilby.core.prior.Uniform(0, 1, latex_label="$b$"), + ) + ) self.results = [ bilby.core.result.Result( label=str(ii), @@ -931,6 +952,7 @@ def __init__(self, mean=0, sigma=1): A very simple Gaussian likelihood for testing """ from scipy.stats import norm + super().__init__() self.mean = mean self.sigma = sigma @@ -941,11 +963,12 @@ def log_likelihood(self, parameters): class TestReweight(unittest.TestCase): - def setUp(self): - self.priors = bilby.core.prior.PriorDict(dict( - mu=bilby.core.prior.TruncatedNormal(0, 1, minimum=-5, maximum=5), - )) + self.priors = bilby.core.prior.PriorDict( + dict( + mu=bilby.core.prior.TruncatedNormal(0, 1, minimum=-5, maximum=5), + ) + ) self.result = bilby.core.result.Result( search_parameter_keys=list(self.priors.keys()), priors=self.priors, @@ -992,11 +1015,11 @@ def test_reweight_different_likelihood_weights_correct(self): should be close to the original evidence within statistical error. """ from scipy.stats import norm + new, weights, _, _, _, _ = self._run_reweighting(sigma=0.5) - expected_weights = ( - norm(0, 0.5).pdf(self.result.posterior["mu"]) - / norm(0, 1).pdf(self.result.posterior["mu"]) - ) + expected_weights = norm(0, 0.5).pdf(self.result.posterior["mu"]) / norm( + 0, 1 + ).pdf(self.result.posterior["mu"]) self.assertLess(min(abs(weights - expected_weights)), 1e-10) self.assertLess(abs(new.log_evidence - self.result.log_evidence), 0.05) self.assertNotEqual(new.log_evidence, self.result.log_evidence) @@ -1005,7 +1028,6 @@ def test_reweight_different_likelihood_weights_correct(self): @pytest.mark.array_backend @pytest.mark.usefixtures("xp_class") class TestResultSaveAndRead(unittest.TestCase): - @pytest.fixture(autouse=True) def init_outdir(self, tmp_path): # Use pytest's tmp_path fixture to create a temporary directory @@ -1050,9 +1072,15 @@ def setUp(self): result.log_noise_evidence = 13 self.result = result - @parameterized.parameterized.expand([ - ".h5", ".hdf5", ".json", ".pkl", ".pickle", - ]) + @parameterized.parameterized.expand( + [ + ".h5", + ".hdf5", + ".json", + ".pkl", + ".pickle", + ] + ) def test_save_and_read_filename_with_extension_and_extension_none(self, ext): # Should use the extension from filename filename = os.path.join(self.result.outdir, f"custom_name.{ext}") @@ -1061,17 +1089,19 @@ def test_save_and_read_filename_with_extension_and_extension_none(self, ext): bilby.core.result.read_in_result(filename=filename) os.remove(filename) - @parameterized.parameterized.expand([ - ("json",), - ("pkl",), - ("pickle",), - (True,), - ]) + @parameterized.parameterized.expand( + [ + ("json",), + ("pkl",), + ("pickle",), + (True,), + ] + ) def test_save_and_read_filename_with_extension_and_extension(self, extension): """Test all the extensions that are support when the filename is provided""" filename = os.path.join(self.result.outdir, "custom_name.hdf5") expected = filename - with self.assertLogs(logger, level='WARNING') as cm: + with self.assertLogs(logger, level="WARNING") as cm: self.result.save_to_file(filename=filename, extension=extension) self.assertIn("does not match the provided extension", cm.output[0]) self.assertTrue(os.path.isfile(expected)) @@ -1099,14 +1129,16 @@ def test_save_to_file_defaults_to_pickle_with_incorrect_extension(self): bilby.core.result.read_in_result(filename=expected) os.remove(expected) - @parameterized.parameterized.expand([ - ("json", "hdf5"), - ("json", "pkl"), - ("hdf5", "json"), - ("pkl", "json"), - ("json", "pkl"), - ("hdf5", "pkl"), - ]) + @parameterized.parameterized.expand( + [ + ("json", "hdf5"), + ("json", "pkl"), + ("hdf5", "json"), + ("pkl", "json"), + ("json", "pkl"), + ("hdf5", "pkl"), + ] + ) def test_save_and_read_incorrect_extension(self, save_extension, read_extension): """Test that an incorrect extension raises a somewhat helpful error""" filename = os.path.join(self.result.outdir, "my_result") @@ -1114,7 +1146,9 @@ def test_save_and_read_incorrect_extension(self, save_extension, read_extension) with self.assertRaises( (FileLoadError, IOError), msg=f"Failed to read in file {filename}" ): - bilby.core.result.read_in_result(filename=filename, extension=read_extension) + bilby.core.result.read_in_result( + filename=filename, extension=read_extension + ) os.remove(filename) diff --git a/test/core/sampler/base_sampler_test.py b/test/core/sampler/base_sampler_test.py index 7ef8757f3..31fffc478 100644 --- a/test/core/sampler/base_sampler_test.py +++ b/test/core/sampler/base_sampler_test.py @@ -10,7 +10,9 @@ import bilby from bilby.core import prior -loaded_samplers = {k: v.load() for k, v in bilby.core.sampler.IMPLEMENTED_SAMPLERS.items()} +loaded_samplers = { + k: v.load() for k, v in bilby.core.sampler.IMPLEMENTED_SAMPLERS.items() +} class MockLikelihood(bilby.core.likelihood.Likelihood): @@ -45,7 +47,7 @@ def setUp(self, soft_init=False): outdir=test_directory, use_ratio=False, skip_import_verification=True, - soft_init=soft_init + soft_init=soft_init, ) def tearDown(self): @@ -116,8 +118,11 @@ def test_get_random_draw(self): def test_initial_points_with_random_state(self): global_rng = bilby.core.utils.random.rng global_state = copy.deepcopy(global_rng.bit_generator.state) - with patch.object(self.sampler, "check_draw", return_value=True), patch.object( - self.sampler, "prior_transform", side_effect=lambda unit: unit + with ( + patch.object(self.sampler, "check_draw", return_value=True), + patch.object( + self.sampler, "prior_transform", side_effect=lambda unit: unit + ), ): points, _, _ = self.sampler.get_initial_points_from_prior( 5, random_state=np.random.default_rng(1234) @@ -128,9 +133,13 @@ def test_initial_points_with_random_state(self): def test_initial_points_default_random_state(self): rng = np.random.default_rng(1234) - with patch.object(bilby.core.utils.random.Generator, "rng", rng), patch.object( - self.sampler, "check_draw", return_value=True - ), patch.object(self.sampler, "prior_transform", side_effect=lambda unit: unit): + with ( + patch.object(bilby.core.utils.random.Generator, "rng", rng), + patch.object(self.sampler, "check_draw", return_value=True), + patch.object( + self.sampler, "prior_transform", side_effect=lambda unit: unit + ), + ): points, _, _ = self.sampler.get_initial_points_from_prior(5) expected = np.random.default_rng(1234).uniform(0, 1, (5, self.sampler.ndim)) np.testing.assert_array_equal(points, expected) @@ -252,9 +261,7 @@ def test_pool_creates_properly_no_pool(self, sampler_name): @parameterized.expand(samplers) def test_pool_creates_properly_pool(self, sampler): - sampler = loaded_samplers[sampler]( - self.likelihood, self.priors, npool=2 - ) + sampler = loaded_samplers[sampler](self.likelihood, self.priors, npool=2) sampler._setup_pool() if hasattr(sampler, "setup_sampler"): sampler.setup_sampler() diff --git a/test/core/sampler/dynesty_test.py b/test/core/sampler/dynesty_test.py index 5e305781c..9291ded5f 100644 --- a/test/core/sampler/dynesty_test.py +++ b/test/core/sampler/dynesty_test.py @@ -258,12 +258,14 @@ def make_trace_plot(*args, **kwargs): all(not plt.fignum_exists(figure.number) for figure in created_figures) ) - @parameterized.parameterized.expand(( - ("unif", "single"), - ("unif", "multi"), - ("rslice", "single"), - ("rslice", "multi"), - )) + @parameterized.parameterized.expand( + ( + ("unif", "single"), + ("unif", "multi"), + ("rslice", "single"), + ("rslice", "multi"), + ) + ) def test_dynesty_native_methods_initialize(self, sample, bound): """ Make sure that we can initialize the native dynesty samplers. @@ -301,7 +303,6 @@ def test_get_expected_outputs(): class ProposalsTest(unittest.TestCase): - def test_boundaries(self): inputs = np.array([0.1, 1.1, -1.3]) expected = np.array([0.1, 0.1, 0.7]) @@ -449,7 +450,6 @@ def test_converges_to_correct_value(self): class TestReproducibility(unittest.TestCase): - @staticmethod def model(x, m, c): return m * x + c diff --git a/test/core/sampler/general_sampler_tests.py b/test/core/sampler/general_sampler_tests.py index 38700c632..e57d8b03b 100644 --- a/test/core/sampler/general_sampler_tests.py +++ b/test/core/sampler/general_sampler_tests.py @@ -26,7 +26,6 @@ def test_get_sampler_class(): def test_get_sampler_class_not_implemented(): """Assert an error is raised if the sampler is not recognized""" with pytest.raises( - ValueError, - match=r"Sampler not_a_valid_sampler not yet implemented" + ValueError, match=r"Sampler not_a_valid_sampler not yet implemented" ): get_sampler_class("not_a_valid_sampler") diff --git a/test/core/sampler/implemented_samplers_test.py b/test/core/sampler/implemented_samplers_test.py index c34398185..bd46591b9 100644 --- a/test/core/sampler/implemented_samplers_test.py +++ b/test/core/sampler/implemented_samplers_test.py @@ -20,10 +20,12 @@ def test_allowed_keys(): def test_values(): # Values and keys should have the same lengths - assert len(list(IMPLEMENTED_SAMPLERS.values())) \ - == len(list(IMPLEMENTED_SAMPLERS.keys())) - assert len(list(IMPLEMENTED_SAMPLERS.values())) \ - == len(list(IMPLEMENTED_SAMPLERS._samplers.values())) + assert len(list(IMPLEMENTED_SAMPLERS.values())) == len( + list(IMPLEMENTED_SAMPLERS.keys()) + ) + assert len(list(IMPLEMENTED_SAMPLERS.values())) == len( + list(IMPLEMENTED_SAMPLERS._samplers.values()) + ) def test_items(): diff --git a/test/core/sampler/proposal_test.py b/test/core/sampler/proposal_test.py index 1b2331af0..5eefec6ba 100644 --- a/test/core/sampler/proposal_test.py +++ b/test/core/sampler/proposal_test.py @@ -148,11 +148,13 @@ def test_set_step_size(self): def test_jump_proposal_call(self): sample = proposal.Sample(dict(reflective=0.0, periodic=0.0, default=0.0)) new_sample = self.jump_proposal(sample) - expected = proposal.Sample(dict( - reflective=0.5942057242396577, - periodic=-0.02692301311556511, - default=-0.7450848662857457, - )) + expected = proposal.Sample( + dict( + reflective=0.5942057242396577, + periodic=-0.02692301311556511, + default=-0.7450848662857457, + ) + ) self.assertDictEqual(expected, new_sample) @@ -169,7 +171,9 @@ def setUp(self): ) bilby.core.utils.random.seed(5) self.jump_proposal = proposal.EnsembleWalk( - random_number_generator=bilby.core.utils.random.rng.uniform, n_points=4, priors=self.priors + random_number_generator=bilby.core.utils.random.rng.uniform, + n_points=4, + priors=self.priors, ) self.coordinates = [proposal.Sample(self.priors.sample()) for _ in range(10)] @@ -185,7 +189,10 @@ def test_set_n_points(self): self.assertEqual(3, self.jump_proposal.n_points) def test_random_number_generator_init(self): - self.assertEqual(bilby.core.utils.random.rng.uniform, self.jump_proposal.random_number_generator) + self.assertEqual( + bilby.core.utils.random.rng.uniform, + self.jump_proposal.random_number_generator, + ) def test_get_center_of_mass(self): samples = [ @@ -200,11 +207,13 @@ def test_get_center_of_mass(self): def test_jump_proposal_call(self): sample = proposal.Sample(dict(periodic=0.1, reflective=0.1, default=0.1)) new_sample = self.jump_proposal(sample, coordinates=self.coordinates) - expected = proposal.Sample(dict( - periodic=0.437075089594473, - reflective=-0.18027731528487945, - default=-0.17570046901727415, - )) + expected = proposal.Sample( + dict( + periodic=0.437075089594473, + reflective=-0.18027731528487945, + default=-0.17570046901727415, + ) + ) for key, value in new_sample.items(): self.assertAlmostEqual(expected[key], value) @@ -236,22 +245,20 @@ def test_set_get_scale(self): self.assertEqual(5.0, self.jump_proposal.scale) def test_jump_proposal_call(self): - sample = proposal.Sample( - dict(periodic=0.1, reflective=0.1, default=0.1) - ) + sample = proposal.Sample(dict(periodic=0.1, reflective=0.1, default=0.1)) new_sample = self.jump_proposal(sample, coordinates=self.coordinates) - expected = proposal.Sample(dict( - periodic=0.5790181653312239, - reflective=-0.028378746842481914, - default=-0.23534241783479043, - )) + expected = proposal.Sample( + dict( + periodic=0.5790181653312239, + reflective=-0.028378746842481914, + default=-0.23534241783479043, + ) + ) for key, value in new_sample.items(): self.assertAlmostEqual(expected[key], value) def test_log_j_after_call(self): - sample = proposal.Sample( - dict(periodic=0.2, reflective=0.2, default=0.2) - ) + sample = proposal.Sample(dict(periodic=0.2, reflective=0.2, default=0.2)) self.jump_proposal(sample=sample, coordinates=self.coordinates) self.assertAlmostEqual(-3.2879289432183088, self.jump_proposal.log_j, 10) @@ -289,14 +296,14 @@ def test_set_get_sigma(self): self.assertEqual(2, self.jump_proposal.sigma) def test_jump_proposal_call(self): - sample = proposal.Sample( - dict(periodic=0.1, reflective=0.1, default=0.1) + sample = proposal.Sample(dict(periodic=0.1, reflective=0.1, default=0.1)) + expected = proposal.Sample( + dict( + periodic=0.09440864471444077, + reflective=0.567962015300636, + default=0.0657296821780595, + ) ) - expected = proposal.Sample(dict( - periodic=0.09440864471444077, - reflective=0.567962015300636, - default=0.0657296821780595, - )) new_sample = self.jump_proposal(sample, coordinates=self.coordinates) for key, value in new_sample.items(): self.assertAlmostEqual(expected[key], value) diff --git a/test/core/series_test.py b/test/core/series_test.py index c2b8dccdb..3a534647a 100644 --- a/test/core/series_test.py +++ b/test/core/series_test.py @@ -95,10 +95,16 @@ def test_time_array_setter(self): self.series.time_array = new_time_array self.assertTrue(np.array_equal(new_time_array, self.series.time_array)) self.assertAlmostEqual( - np.asarray(new_sampling_frequency), np.asarray(self.series.sampling_frequency), places=1 + np.asarray(new_sampling_frequency), + np.asarray(self.series.sampling_frequency), + places=1, + ) + self.assertAlmostEqual( + np.asarray(new_duration), np.asarray(self.series.duration), places=1 + ) + self.assertAlmostEqual( + np.asarray(new_start_time), np.asarray(self.series.start_time), places=1 ) - self.assertAlmostEqual(np.asarray(new_duration), np.asarray(self.series.duration), places=1) - self.assertAlmostEqual(np.asarray(new_start_time), np.asarray(self.series.start_time), places=1) def test_time_array_without_sampling_frequency(self): self.series.sampling_frequency = None diff --git a/test/core/utils_test.py b/test/core/utils_test.py index aee0c5437..dd5859ab8 100644 --- a/test/core/utils_test.py +++ b/test/core/utils_test.py @@ -121,7 +121,7 @@ def test_builtin_function(self): "errors", "newline", "closefd", - "opener" + "opener", ] actual = utils.infer_parameters_from_function(open) self.assertListEqual(expected, actual) @@ -216,9 +216,9 @@ def test_get_sampling_frequency_from_frequency_array(self): self.assertEqual(self.sampling_frequency, new_sampling_freq) def test_get_sampling_frequency_from_frequency_array_unequally_sampled(self): - self.frequency_array = xpx.at( - self.frequency_array, -1 - ).set(self.frequency_array[-1] + 0.0001) + self.frequency_array = xpx.at(self.frequency_array, -1).set( + self.frequency_array[-1] + 0.0001 + ) with self.assertRaises(ValueError): _, _ = utils.get_sampling_frequency_and_duration_from_frequency_array( self.frequency_array @@ -390,7 +390,9 @@ def setUp(self): self.xx = np.linspace(0, 1, 10) self.yy = np.linspace(0, 1, 10) self.zz = np.random.random((10, 10)) - self.interpolant = bilby.core.utils.BoundedRectBivariateSpline(self.xx, self.yy, self.zz) + self.interpolant = bilby.core.utils.BoundedRectBivariateSpline( + self.xx, self.yy, self.zz + ) def tearDown(self): pass @@ -405,8 +407,8 @@ def test_returns_none_for_floats_outside_range(self): def test_returns_float_for_float_and_array(self): input_array = self.xp.asarray(np.random.random(10)) self.assertEqual(aac.get_namespace(self.interpolant(input_array, 0.5)), self.xp) - self.assertEqual(aac.get_namespace( - self.interpolant(input_array, input_array)), self.xp + self.assertEqual( + aac.get_namespace(self.interpolant(input_array, input_array)), self.xp ) self.assertEqual(aac.get_namespace(self.interpolant(0.5, input_array)), self.xp) @@ -435,7 +437,7 @@ def setUp(self): self.lnfunc1 = self.xp.log(self.x) self.func1int = (self.x[-1] ** 2 - self.x[0] ** 2) / 2 with np.errstate(divide="ignore"): - self.lnfunc2 = self.xp.log(self.x ** 2) + self.lnfunc2 = self.xp.log(self.x**2) self.func2int = (self.x[-1] ** 3 - self.x[0] ** 3) / 3 self.irregularx = self.xp.asarray( @@ -457,7 +459,7 @@ def setUp(self): ) with np.errstate(divide="ignore"): self.lnfunc1irregular = self.xp.log(self.irregularx) - self.lnfunc2irregular = self.xp.log(self.irregularx ** 2) + self.lnfunc2irregular = self.xp.log(self.irregularx**2) self.irregulardxs = self.xp.diff(self.irregularx) def test_incorrect_step_type(self): @@ -473,23 +475,30 @@ def test_integral_func1(self): res2 = utils.logtrapzexp(self.lnfunc1, self.dxs) self.assertTrue(np.abs(res1 - res2) < 1e-12) - self.assertTrue(np.abs((self.xp.exp(res1) - self.func1int) / self.func1int) < 1e-12) + self.assertTrue( + np.abs((self.xp.exp(res1) - self.func1int) / self.func1int) < 1e-12 + ) def test_integral_func2(self): res = utils.logtrapzexp(self.lnfunc2, self.dxs) - self.assertTrue(np.abs((self.xp.exp(res) - self.func2int) / self.func2int) < 1e-4) + self.assertTrue( + np.abs((self.xp.exp(res) - self.func2int) / self.func2int) < 1e-4 + ) def test_integral_func1_irregular_steps(self): res = utils.logtrapzexp(self.lnfunc1irregular, self.irregulardxs) - self.assertTrue(np.abs((self.xp.exp(res) - self.func1int) / self.func1int) < 1e-12) + self.assertTrue( + np.abs((self.xp.exp(res) - self.func1int) / self.func1int) < 1e-12 + ) def test_integral_func2_irregular_steps(self): res = utils.logtrapzexp(self.lnfunc2irregular, self.irregulardxs) - self.assertTrue(np.abs((self.xp.exp(res) - self.func2int) / self.func2int) < 1e-2) + self.assertTrue( + np.abs((self.xp.exp(res) - self.func2int) / self.func2int) < 1e-2 + ) class TestSavingNumpyRandomGenerator(unittest.TestCase): - @pytest.fixture(autouse=True) def init_outdir(self, tmp_path): # Use pytest's tmp_path fixture to create a temporary directory @@ -505,9 +514,7 @@ def setUp(self): def test_hdf5(self): with h5py.File(self.outdir / "test.h5", "w") as f: - bilby.core.utils.recursively_save_dict_contents_to_group( - f, "/", self.data - ) + bilby.core.utils.recursively_save_dict_contents_to_group(f, "/", self.data) a = self.data["rng"].random() with h5py.File(self.outdir / "test.h5", "r") as f: @@ -517,12 +524,12 @@ def test_hdf5(self): self.assertEqual(a, b) def test_json(self): - with open(self.outdir / "test.json", 'w') as file: + with open(self.outdir / "test.json", "w") as file: json.dump(self.data, file, indent=2, cls=bilby.core.utils.BilbyJsonEncoder) a = self.data["rng"].random() - with open(self.outdir / "test.json", 'r') as file: + with open(self.outdir / "test.json", "r") as file: data = json.load(file, object_hook=bilby.core.utils.decode_bilby_json) b = data["rng"].random() @@ -533,14 +540,10 @@ def test_hdf5_seed_sequence(self): data = {"seed_sequence": seed_sequence} with h5py.File(self.outdir / "test_seed_sequence.h5", "w") as f: - bilby.core.utils.recursively_save_dict_contents_to_group( - f, "/", data - ) + bilby.core.utils.recursively_save_dict_contents_to_group(f, "/", data) with h5py.File(self.outdir / "test_seed_sequence.h5", "r") as f: - loaded = bilby.core.utils.recursively_load_dict_contents_from_group( - f, "/" - ) + loaded = bilby.core.utils.recursively_load_dict_contents_from_group(f, "/") self.assertIsInstance(loaded["seed_sequence"], np.random.SeedSequence) self.assertEqual(loaded["seed_sequence"].state, seed_sequence.state) @@ -553,26 +556,23 @@ def test_json_seed_sequence(self): json.dump(data, file, indent=2, cls=bilby.core.utils.BilbyJsonEncoder) with open(self.outdir / "test_seed_sequence.json", "r") as file: - loaded = json.load( - file, object_hook=bilby.core.utils.decode_bilby_json - ) + loaded = json.load(file, object_hook=bilby.core.utils.decode_bilby_json) self.assertIsInstance(loaded["seed_sequence"], np.random.SeedSequence) self.assertEqual(loaded["seed_sequence"].state, seed_sequence.state) def test_pickle(self): - with open(self.outdir / "test.pkl", 'wb') as file: + with open(self.outdir / "test.pkl", "wb") as file: dill.dump(self.data, file) a = self.data["rng"].random() - with open(self.outdir / "test.pkl", 'rb') as file: + with open(self.outdir / "test.pkl", "rb") as file: data = dill.load(file) b = data["rng"].random() self.assertEqual(a, b) class TestSavingLalDict(unittest.TestCase): - @pytest.fixture(autouse=True) def init_outdir(self, tmp_path): # Use pytest's tmp_path fixture to create a temporary directory @@ -598,9 +598,7 @@ def _assert_lal_dicts_equal(self, a, b): def test_hdf5(self): with h5py.File(self.outdir / "test.h5", "w") as f: - bilby.core.utils.recursively_save_dict_contents_to_group( - f, "/", self.data - ) + bilby.core.utils.recursively_save_dict_contents_to_group(f, "/", self.data) with h5py.File(self.outdir / "test.h5", "r") as f: data = bilby.core.utils.recursively_load_dict_contents_from_group(f, "/") @@ -608,10 +606,10 @@ def test_hdf5(self): self._assert_lal_dicts_equal(self.data["lal_dict"], data["lal_dict"]) def test_json(self): - with open(self.outdir / "test.json", 'w') as file: + with open(self.outdir / "test.json", "w") as file: json.dump(self.data, file, indent=2, cls=bilby.core.utils.BilbyJsonEncoder) - with open(self.outdir / "test.json", 'r') as file: + with open(self.outdir / "test.json", "r") as file: data = json.load(file, object_hook=bilby.core.utils.decode_bilby_json) self._assert_lal_dicts_equal(self.data["lal_dict"], data["lal_dict"]) @@ -628,7 +626,6 @@ def test_decode_lal_dict(self): class TestGlobalMetaData(unittest.TestCase): - @pytest.fixture(autouse=True) def set_caplog(self, caplog): self._caplog = caplog @@ -656,7 +653,9 @@ def test_set_rng(self): def test_set_cosmology(self): bilby.gw.cosmology.set_cosmology("Planck15_LAL") - self.assertTrue(global_meta_data["cosmology"] is bilby.gw.cosmology.COSMOLOGY[0]) + self.assertTrue( + global_meta_data["cosmology"] is bilby.gw.cosmology.COSMOLOGY[0] + ) def test_update(self): bilby.core.utils.meta_data.logger.propagate = True @@ -673,7 +672,6 @@ def test_init(self): class TestRandomUtils(unittest.TestCase): - def setUp(self): # Ensure a clean import of the random module if "bilby.core.utils.random" in sys.modules: @@ -692,7 +690,9 @@ def test_no_warning_when_accessed_via_module(self): random.seed(42) for warning in w: - self.assertNotIn("Detected that `rng` was likely imported directly", str(warning.message)) + self.assertNotIn( + "Detected that `rng` was likely imported directly", str(warning.message) + ) def test_warning_when_imported_directly(self): from bilby.core.utils.random import rng @@ -704,7 +704,9 @@ def test_warning_when_imported_directly(self): sys.modules["fake_module"] = fake_module try: - with self.assertWarnsRegex(RuntimeWarning, "Detected that `rng` was likely imported directly"): + with self.assertWarnsRegex( + RuntimeWarning, "Detected that `rng` was likely imported directly" + ): random.seed(123) finally: # always clean up the fake module even if the test fails diff --git a/test/gw/conversion_test.py b/test/gw/conversion_test.py index ae4cc0777..12db91fa2 100644 --- a/test/gw/conversion_test.py +++ b/test/gw/conversion_test.py @@ -17,8 +17,8 @@ def setUp(self): self.mass_2 = self.xp.asarray(1.3) self.mass_ratio = self.xp.asarray(13 / 14) self.total_mass = self.xp.asarray(2.7) - self.chirp_mass = (self.mass_1 * self.mass_2) ** 0.6 / self.total_mass ** 0.2 - self.symmetric_mass_ratio = (self.mass_1 * self.mass_2) / self.total_mass ** 2 + self.chirp_mass = (self.mass_1 * self.mass_2) ** 0.6 / self.total_mass**0.2 + self.symmetric_mass_ratio = (self.mass_1 * self.mass_2) / self.total_mass**2 self.cos_angle = self.xp.asarray(-1.0) self.angle = self.xp.pi self.lambda_1 = self.xp.asarray(300.0) @@ -27,17 +27,13 @@ def setUp(self): 8 / 13 * ( - ( - 1 - + 7 * self.symmetric_mass_ratio - - 31 * self.symmetric_mass_ratio ** 2 - ) + (1 + 7 * self.symmetric_mass_ratio - 31 * self.symmetric_mass_ratio**2) * (self.lambda_1 + self.lambda_2) + (1 - 4 * self.symmetric_mass_ratio) ** 0.5 * ( 1 + 9 * self.symmetric_mass_ratio - - 11 * self.symmetric_mass_ratio ** 2 + - 11 * self.symmetric_mass_ratio**2 ) * (self.lambda_1 - self.lambda_2) ) @@ -50,14 +46,14 @@ def setUp(self): * ( 1 - 13272 / 1319 * self.symmetric_mass_ratio - + 8944 / 1319 * self.symmetric_mass_ratio ** 2 + + 8944 / 1319 * self.symmetric_mass_ratio**2 ) * (self.lambda_1 + self.lambda_2) + ( 1 - 15910 / 1319 * self.symmetric_mass_ratio - + 32850 / 1319 * self.symmetric_mass_ratio ** 2 - + 3380 / 1319 * self.symmetric_mass_ratio ** 3 + + 32850 / 1319 * self.symmetric_mass_ratio**2 + + 3380 / 1319 * self.symmetric_mass_ratio**3 ) * (self.lambda_1 - self.lambda_2) ) @@ -96,10 +92,14 @@ def test_symmetric_mass_ratio_to_mass_ratio(self): self.assertEqual(aac.get_namespace(mass_ratio), self.xp) def test_chirp_mass_and_total_mass_to_symmetric_mass_ratio(self): - symmetric_mass_ratio = conversion.chirp_mass_and_total_mass_to_symmetric_mass_ratio( - self.chirp_mass, self.total_mass + symmetric_mass_ratio = ( + conversion.chirp_mass_and_total_mass_to_symmetric_mass_ratio( + self.chirp_mass, self.total_mass + ) + ) + self.assertAlmostEqual( + float(self.symmetric_mass_ratio), float(symmetric_mass_ratio) ) - self.assertAlmostEqual(float(self.symmetric_mass_ratio), float(symmetric_mass_ratio)) self.assertEqual(aac.get_namespace(symmetric_mass_ratio), self.xp) def test_chirp_mass_and_mass_ratio_to_total_mass(self): @@ -110,9 +110,9 @@ def test_chirp_mass_and_mass_ratio_to_total_mass(self): self.assertEqual(aac.get_namespace(total_mass), self.xp) def test_chirp_mass_and_mass_ratio_to_component_masses(self): - mass_1, mass_2 = \ - conversion.chirp_mass_and_mass_ratio_to_component_masses( - self.chirp_mass, self.mass_ratio) + mass_1, mass_2 = conversion.chirp_mass_and_mass_ratio_to_component_masses( + self.chirp_mass, self.mass_ratio + ) self.assertAlmostEqual(self.mass_1, mass_1) self.assertAlmostEqual(self.mass_2, mass_2) self.assertEqual(aac.get_namespace(mass_1), self.xp) @@ -207,9 +207,11 @@ def test_identity_conversion(self): lambda_1=self.lambda_1, lambda_2=self.lambda_2, lambda_tilde=self.lambda_tilde, - delta_lambda_tilde=self.delta_lambda_tilde + delta_lambda_tilde=self.delta_lambda_tilde, + ) + identity_samples, blank_list = conversion.identity_map_conversion( + original_samples ) - identity_samples, blank_list = conversion.identity_map_conversion(original_samples) assert blank_list == [] for key, val in identity_samples.items(): assert val == self.__dict__[key] @@ -224,10 +226,10 @@ def setUp(self): self.mass_parameters["mass_ratio"] = conversion.component_masses_to_mass_ratio( **self.component_mass_pars ) - self.mass_parameters[ - "symmetric_mass_ratio" - ] = conversion.component_masses_to_symmetric_mass_ratio( - **self.component_mass_pars + self.mass_parameters["symmetric_mass_ratio"] = ( + conversion.component_masses_to_symmetric_mass_ratio( + **self.component_mass_pars + ) ) self.mass_parameters["chirp_mass"] = conversion.component_masses_to_chirp_mass( **self.component_mass_pars @@ -239,13 +241,13 @@ def setUp(self): self.all_component_pars = self.component_tidal_parameters.copy() self.all_component_pars.update(self.component_mass_pars) self.tidal_parameters = self.component_tidal_parameters.copy() - self.tidal_parameters[ - "lambda_tilde" - ] = conversion.lambda_1_lambda_2_to_lambda_tilde(**self.all_component_pars) - self.tidal_parameters[ - "delta_lambda_tilde" - ] = conversion.lambda_1_lambda_2_to_delta_lambda_tilde( - **self.all_component_pars + self.tidal_parameters["lambda_tilde"] = ( + conversion.lambda_1_lambda_2_to_lambda_tilde(**self.all_component_pars) + ) + self.tidal_parameters["delta_lambda_tilde"] = ( + conversion.lambda_1_lambda_2_to_delta_lambda_tilde( + **self.all_component_pars + ) ) def tearDown(self): @@ -371,7 +373,14 @@ def test_bbh_zero_aligned_spin_to_spherical_with_magnitude(self): key: self.parameters[key] for key in ["a_1", "a_2", "tilt_1", "tilt_2", "phi_12", "phi_jl"] }, - dict(a_1=a_1, a_2=a_2, tilt_1=tilt_1, tilt_2=tilt_2, phi_jl=phi_jl, phi_12=phi_12), + dict( + a_1=a_1, + a_2=a_2, + tilt_1=tilt_1, + tilt_2=tilt_2, + phi_jl=phi_jl, + phi_12=phi_12, + ), ) def test_bbh_zero_aligned_spin_to_spherical_without_magnitude(self): @@ -493,9 +502,9 @@ def setUp(self): "lambda_tilde", "delta_lambda_tilde", ] - self.data_frame = pd.DataFrame({ - key: [value] * 100 for key, value in self.parameters.items() - }) + self.data_frame = pd.DataFrame( + {key: [value] * 100 for key, value in self.parameters.items()} + ) def test_generate_all_bbh_parameters(self): self._generate( @@ -519,7 +528,9 @@ def test_generate_bbh_parameters_with_likelihood(self): priors = bilby.gw.prior.BBHPriorDict() priors["geocent_time"] = bilby.core.prior.Uniform(0.4, 0.6) ifos = bilby.gw.detector.InterferometerList(["H1", "L1"]) - ifos.set_strain_data_from_power_spectral_densities(duration=1, sampling_frequency=256) + ifos.set_strain_data_from_power_spectral_densities( + duration=1, sampling_frequency=256 + ) wfg = bilby.gw.waveform_generator.WaveformGenerator( frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole ) @@ -561,15 +572,20 @@ def test_generate_bbh_parameters_with_likelihood(self): converted[f"{ifo.name}_log_likelihood"], converted_with_time_jitter[f"{ifo.name}_log_likelihood"], ) - self.assertNotEqual(converted["mass_1"], initial_likelihood_parameters["mass_1"]) + self.assertNotEqual( + converted["mass_1"], initial_likelihood_parameters["mass_1"] + ) def test_identity_generation_no_likelihood(self): - test_fixed_prior = bilby.core.prior.PriorDict({ - "test_param_a": bilby.core.prior.DeltaFunction(0, name="test_param_a"), - "test_param_b": bilby.core.prior.DeltaFunction(1, name="test_param_b") - } + test_fixed_prior = bilby.core.prior.PriorDict( + { + "test_param_a": bilby.core.prior.DeltaFunction(0, name="test_param_a"), + "test_param_b": bilby.core.prior.DeltaFunction(1, name="test_param_b"), + } + ) + output_sample = conversion.identity_map_generation( + self.parameters, priors=test_fixed_prior ) - output_sample = conversion.identity_map_generation(self.parameters, priors=test_fixed_prior) assert output_sample.pop("test_param_a") == 0 assert output_sample.pop("test_param_b") == 1 for key, val in self.parameters.items(): @@ -583,7 +599,9 @@ def test_identity_generation_with_likelihood(self): # Note we do *not* switch to azimuth/zenith, because the identity generation function # is not intended to be capable of that conversion ifos = bilby.gw.detector.InterferometerList(["H1"]) - ifos.set_strain_data_from_power_spectral_densities(duration=1, sampling_frequency=256) + ifos.set_strain_data_from_power_spectral_densities( + duration=1, sampling_frequency=256 + ) wfg = bilby.gw.waveform_generator.WaveformGenerator( frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole ) @@ -595,7 +613,9 @@ def test_identity_generation_with_likelihood(self): time_marginalization=True, reference_frame="sky", ) - output_sample = conversion.identity_map_generation(self.parameters, priors=priors, likelihood=likelihood) + output_sample = conversion.identity_map_generation( + self.parameters, priors=priors, likelihood=likelihood + ) extra_expected = [ "phase", "geocent_time", @@ -636,60 +656,72 @@ def test_comoving_luminosity_with_cosmology(self): @pytest.mark.usefixtures("xp_class") class TestGenerateMassParameters(unittest.TestCase): def setUp(self): - self.expected_values = {'mass_1': self.xp.asarray(2.0), - 'mass_2': self.xp.asarray(1.0), - 'chirp_mass': self.xp.asarray(1.2167286837864113), - 'total_mass': self.xp.asarray(3.0), - 'mass_1_source': self.xp.asarray(4.0), - 'mass_2_source': self.xp.asarray(2.0), - 'chirp_mass_source': self.xp.asarray(2.433457367572823), - 'total_mass_source': self.xp.asarray(6), - 'symmetric_mass_ratio': self.xp.asarray(0.2222222222222222), - 'mass_ratio': self.xp.asarray(0.5)} + self.expected_values = { + "mass_1": self.xp.asarray(2.0), + "mass_2": self.xp.asarray(1.0), + "chirp_mass": self.xp.asarray(1.2167286837864113), + "total_mass": self.xp.asarray(3.0), + "mass_1_source": self.xp.asarray(4.0), + "mass_2_source": self.xp.asarray(2.0), + "chirp_mass_source": self.xp.asarray(2.433457367572823), + "total_mass_source": self.xp.asarray(6), + "symmetric_mass_ratio": self.xp.asarray(0.2222222222222222), + "mass_ratio": self.xp.asarray(0.5), + } def helper_generation_from_keys(self, keys, expected_values, source=False): # Explicitly test the helper generate_component_masses - local_test_vars = \ - {key: expected_values[key] for key in keys} - local_test_vars_with_component_masses = \ - conversion.generate_component_masses(local_test_vars, source=source) + local_test_vars = {key: expected_values[key] for key in keys} + local_test_vars_with_component_masses = conversion.generate_component_masses( + local_test_vars, source=source + ) if source: - self.assertTrue("mass_1_source" in local_test_vars_with_component_masses.keys()) - self.assertTrue("mass_2_source" in local_test_vars_with_component_masses.keys()) + self.assertTrue( + "mass_1_source" in local_test_vars_with_component_masses.keys() + ) + self.assertTrue( + "mass_2_source" in local_test_vars_with_component_masses.keys() + ) else: self.assertTrue("mass_1" in local_test_vars_with_component_masses.keys()) self.assertTrue("mass_2" in local_test_vars_with_component_masses.keys()) for key in local_test_vars_with_component_masses.keys(): self.assertAlmostEqual( np.asarray(local_test_vars_with_component_masses[key]), - np.asarray(self.expected_values[key])) + np.asarray(self.expected_values[key]), + ) # Test the function more generally - local_all_mass_parameters = \ - conversion.generate_mass_parameters(local_test_vars, source=source) + local_all_mass_parameters = conversion.generate_mass_parameters( + local_test_vars, source=source + ) if source: self.assertEqual( set(local_all_mass_parameters.keys()), - set(["mass_1_source", - "mass_2_source", - "chirp_mass_source", - "total_mass_source", - "symmetric_mass_ratio", - "mass_ratio", - ] - ) + set( + [ + "mass_1_source", + "mass_2_source", + "chirp_mass_source", + "total_mass_source", + "symmetric_mass_ratio", + "mass_ratio", + ] + ), ) else: self.assertEqual( set(local_all_mass_parameters.keys()), - set(["mass_1", - "mass_2", - "chirp_mass", - "total_mass", - "symmetric_mass_ratio", - "mass_ratio", - ] - ) + set( + [ + "mass_1", + "mass_2", + "chirp_mass", + "total_mass", + "symmetric_mass_ratio", + "mass_ratio", + ] + ), ) for key in local_all_mass_parameters.keys(): self.assertAlmostEqual( @@ -702,149 +734,168 @@ def helper_generation_from_keys(self, keys, expected_values, source=False): ) def test_from_mass_1_and_mass_2(self): - self.helper_generation_from_keys(["mass_1", "mass_2"], - self.expected_values) + self.helper_generation_from_keys(["mass_1", "mass_2"], self.expected_values) def test_from_mass_1_and_mass_ratio(self): - self.helper_generation_from_keys(["mass_1", "mass_ratio"], - self.expected_values) + self.helper_generation_from_keys(["mass_1", "mass_ratio"], self.expected_values) def test_from_mass_2_and_mass_ratio(self): - self.helper_generation_from_keys(["mass_2", "mass_ratio"], - self.expected_values) + self.helper_generation_from_keys(["mass_2", "mass_ratio"], self.expected_values) def test_from_mass_1_and_total_mass(self): - self.helper_generation_from_keys(["mass_2", "total_mass"], - self.expected_values) + self.helper_generation_from_keys(["mass_2", "total_mass"], self.expected_values) def test_from_chirp_mass_and_mass_ratio(self): - self.helper_generation_from_keys(["chirp_mass", "mass_ratio"], - self.expected_values) + self.helper_generation_from_keys( + ["chirp_mass", "mass_ratio"], self.expected_values + ) def test_from_chirp_mass_and_symmetric_mass_ratio(self): - self.helper_generation_from_keys(["chirp_mass", "symmetric_mass_ratio"], - self.expected_values) + self.helper_generation_from_keys( + ["chirp_mass", "symmetric_mass_ratio"], self.expected_values + ) def test_from_chirp_mass_and_symmetric_mass_1(self): - self.helper_generation_from_keys(["chirp_mass", "mass_1"], - self.expected_values) + self.helper_generation_from_keys(["chirp_mass", "mass_1"], self.expected_values) def test_from_chirp_mass_and_symmetric_mass_2(self): - self.helper_generation_from_keys(["chirp_mass", "mass_2"], - self.expected_values) + self.helper_generation_from_keys(["chirp_mass", "mass_2"], self.expected_values) def test_from_mass_1_source_and_mass_2_source(self): - self.helper_generation_from_keys(["mass_1_source", "mass_2_source"], - self.expected_values, source=True) + self.helper_generation_from_keys( + ["mass_1_source", "mass_2_source"], self.expected_values, source=True + ) def test_from_mass_1_source_and_mass_ratio(self): - self.helper_generation_from_keys(["mass_1_source", "mass_ratio"], - self.expected_values, source=True) + self.helper_generation_from_keys( + ["mass_1_source", "mass_ratio"], self.expected_values, source=True + ) def test_from_mass_2_source_and_mass_ratio(self): - self.helper_generation_from_keys(["mass_2_source", "mass_ratio"], - self.expected_values, source=True) + self.helper_generation_from_keys( + ["mass_2_source", "mass_ratio"], self.expected_values, source=True + ) def test_from_mass_1_source_and_total_mass(self): - self.helper_generation_from_keys(["mass_2_source", "total_mass_source"], - self.expected_values, source=True) + self.helper_generation_from_keys( + ["mass_2_source", "total_mass_source"], self.expected_values, source=True + ) def test_from_chirp_mass_source_and_mass_ratio(self): - self.helper_generation_from_keys(["chirp_mass_source", "mass_ratio"], - self.expected_values, source=True) + self.helper_generation_from_keys( + ["chirp_mass_source", "mass_ratio"], self.expected_values, source=True + ) def test_from_chirp_mass_source_and_symmetric_mass_ratio(self): - self.helper_generation_from_keys(["chirp_mass_source", "symmetric_mass_ratio"], - self.expected_values, source=True) + self.helper_generation_from_keys( + ["chirp_mass_source", "symmetric_mass_ratio"], + self.expected_values, + source=True, + ) def test_from_chirp_mass_source_and_symmetric_mass_1(self): - self.helper_generation_from_keys(["chirp_mass_source", "mass_1_source"], - self.expected_values, source=True) + self.helper_generation_from_keys( + ["chirp_mass_source", "mass_1_source"], self.expected_values, source=True + ) def test_from_chirp_mass_source_and_symmetric_mass_2(self): - self.helper_generation_from_keys(["chirp_mass_source", "mass_2_source"], - self.expected_values, source=True) + self.helper_generation_from_keys( + ["chirp_mass_source", "mass_2_source"], self.expected_values, source=True + ) @pytest.mark.array_backend @pytest.mark.usefixtures("xp_class") class TestEquationOfStateConversions(unittest.TestCase): - ''' + """ Class to test equation of state conversions. The test points were generated from a simulation independent of bilby using the original lalsimulation calls. Specific cases tested are described within each function. - ''' + """ + def setUp(self): - self.mass_1_source_spectral = self.xp.asarray([ - 4.922542724434885, - 4.350626907771598, - 4.206155335439082, - 1.7822696459661311, - 1.3091740103047926 - ]) - self.mass_2_source_spectral = self.xp.asarray([ - 3.459974694590303, - 1.2276461777181447, - 3.7287707089639976, - 0.3724016563531846, - 1.055042934805801 - ]) - self.spectral_pca_gamma_0 = self.xp.asarray([ - 0.7074873121348357, - 0.05855931126849878, - 0.7795329261793462, - 1.467907561566463, - 2.9066488405635624 - ]) - self.spectral_pca_gamma_1 = self.xp.asarray([ - -0.29807111670823816, - 2.027708558522935, - -1.4415775226512115, - -0.7104870098896858, - -0.4913817181089619 - ]) - self.spectral_pca_gamma_2 = self.xp.asarray([ - 0.25625095371021156, - -0.19574096643220049, - -0.2710238103460012, - 0.22815820981582358, - -0.1543413205016374 - ]) - self.spectral_pca_gamma_3 = self.xp.asarray([ - -0.04030365100175101, - 0.05698030777919032, - -0.045595911403040264, - -0.023480394227900117, - -0.07114492992285618 - ]) + self.mass_1_source_spectral = self.xp.asarray( + [ + 4.922542724434885, + 4.350626907771598, + 4.206155335439082, + 1.7822696459661311, + 1.3091740103047926, + ] + ) + self.mass_2_source_spectral = self.xp.asarray( + [ + 3.459974694590303, + 1.2276461777181447, + 3.7287707089639976, + 0.3724016563531846, + 1.055042934805801, + ] + ) + self.spectral_pca_gamma_0 = self.xp.asarray( + [ + 0.7074873121348357, + 0.05855931126849878, + 0.7795329261793462, + 1.467907561566463, + 2.9066488405635624, + ] + ) + self.spectral_pca_gamma_1 = self.xp.asarray( + [ + -0.29807111670823816, + 2.027708558522935, + -1.4415775226512115, + -0.7104870098896858, + -0.4913817181089619, + ] + ) + self.spectral_pca_gamma_2 = self.xp.asarray( + [ + 0.25625095371021156, + -0.19574096643220049, + -0.2710238103460012, + 0.22815820981582358, + -0.1543413205016374, + ] + ) + self.spectral_pca_gamma_3 = self.xp.asarray( + [ + -0.04030365100175101, + 0.05698030777919032, + -0.045595911403040264, + -0.023480394227900117, + -0.07114492992285618, + ] + ) self.spectral_gamma_0 = [ 1.1259406796075457, 0.3191335618787259, 1.3651245109783452, 1.3540140238735314, - 1.4551949842961993 + 1.4551949842961993, ] self.spectral_gamma_1 = [ 0.26791504475282835, 0.3930374252139248, 0.11438399886108475, 0.14181113477953, - -0.11989033256620368 + -0.11989033256620368, ] self.spectral_gamma_2 = [ -0.06810849354463173, -0.038250139296677754, -0.0801540229444505, -0.05230330841791625, - -0.005197303281460286 + -0.005197303281460286, ] self.spectral_gamma_3 = [ 0.002848121360389597, 0.000872447754855139, 0.005528747386660879, 0.0024325946344566484, - 0.00043890906202786106 + 0.00043890906202786106, ] self.mass_1_source_polytrope = [ 2.2466565877822573, @@ -852,7 +903,7 @@ def setUp(self): 4.123897187899834, 2.014160764697004, 1.414796714032148, - 2.0919349759766614 + 2.0919349759766614, ] self.mass_2_source_polytrope = [ 0.36696047254774256, @@ -860,7 +911,7 @@ def setUp(self): 1.650477659961306, 1.310399737462001, 0.5470843356210495, - 1.2311162283818198 + 1.2311162283818198, ] self.polytrope_log10_pressure_1 = [ 34.05849276958394, @@ -868,7 +919,7 @@ def setUp(self): 33.07579629429792, 33.93412833210738, 34.24096323517809, - 35.293288373856534 + 35.293288373856534, ] self.polytrope_log10_pressure_2 = [ 33.82891829901602, @@ -876,7 +927,7 @@ def setUp(self): 34.940095188881976, 34.72710820593933, 35.42780071717415, - 35.648689969687915 + 35.648689969687915, ] self.polytrope_gamma_0 = [ 2.359580734009537, @@ -884,7 +935,7 @@ def setUp(self): 4.784129809424835, 1.4900432021657437, 1.0037220431922798, - 4.183994058757201 + 4.183994058757201, ] self.polytrope_gamma_1 = [ 1.9497583698697314, @@ -892,7 +943,7 @@ def setUp(self): 2.8228335336587826, 4.032519623275465, 1.10894361284508, - 3.168076721819637 + 3.168076721819637, ] self.polytrope_gamma_2 = [ 4.6001755196585385, @@ -900,63 +951,79 @@ def setUp(self): 4.429607300132092, 1.8176338276795763, 2.9938859949129797, - 1.300271383168368 + 1.300271383168368, ] - self.lambda_1_spectral = [0., 0., 0., 0., 1275.7253186286332] - self.lambda_2_spectral = [0., 0., 0., 0., 4504.897675043909] - self.lambda_1_polytrope = [0., 0., 0., 0., 0., 234.66424898184766] - self.lambda_2_polytrope = [0., 0., 0., 0., 0., 3710.931378294547] + self.lambda_1_spectral = [0.0, 0.0, 0.0, 0.0, 1275.7253186286332] + self.lambda_2_spectral = [0.0, 0.0, 0.0, 0.0, 4504.897675043909] + self.lambda_1_polytrope = [0.0, 0.0, 0.0, 0.0, 0.0, 234.66424898184766] + self.lambda_2_polytrope = [0.0, 0.0, 0.0, 0.0, 0.0, 3710.931378294547] self.eos_check_spectral = [0, 0, 0, 0, 1] self.eos_check_polytrope = [0, 0, 0, 0, 0, 1] def test_spectral_pca_to_spectral(self): for i in range(len(self.mass_1_source_spectral)): - spectral_gamma_0, spectral_gamma_1, spectral_gamma_2, spectral_gamma_3 = \ + spectral_gamma_0, spectral_gamma_1, spectral_gamma_2, spectral_gamma_3 = ( conversion.spectral_pca_to_spectral( self.spectral_pca_gamma_0[i], self.spectral_pca_gamma_1[i], self.spectral_pca_gamma_2[i], - self.spectral_pca_gamma_3[i] + self.spectral_pca_gamma_3[i], ) - self.assertAlmostEqual(float(spectral_gamma_0), self.spectral_gamma_0[i], places=5) - self.assertAlmostEqual(float(spectral_gamma_1), self.spectral_gamma_1[i], places=5) - self.assertAlmostEqual(float(spectral_gamma_2), self.spectral_gamma_2[i], places=5) - self.assertAlmostEqual(float(spectral_gamma_3), self.spectral_gamma_3[i], places=5) - for val in [spectral_gamma_0, spectral_gamma_1, spectral_gamma_2, spectral_gamma_3]: + ) + self.assertAlmostEqual( + float(spectral_gamma_0), self.spectral_gamma_0[i], places=5 + ) + self.assertAlmostEqual( + float(spectral_gamma_1), self.spectral_gamma_1[i], places=5 + ) + self.assertAlmostEqual( + float(spectral_gamma_2), self.spectral_gamma_2[i], places=5 + ) + self.assertAlmostEqual( + float(spectral_gamma_3), self.spectral_gamma_3[i], places=5 + ) + for val in [ + spectral_gamma_0, + spectral_gamma_1, + spectral_gamma_2, + spectral_gamma_3, + ]: self.assertEqual(aac.get_namespace(val), self.xp) def test_spectral_params_to_lambda_1_lambda_2(self): - ''' + """ The points cover 5 test cases: - Fail SimNeutronStarEOS4ParamSDGammaCheck() - Fail max_speed_of_sound_ <=1.1 - Fail mass_1_source <= max_mass - Fail mass_2_source >= min_mass - Passes all and produces accurate lambda_1, lambda_2, eos_check values - ''' + """ for i in range(len(self.mass_1_source_spectral)): - spectral_gamma_0, spectral_gamma_1, spectral_gamma_2, spectral_gamma_3 = \ + spectral_gamma_0, spectral_gamma_1, spectral_gamma_2, spectral_gamma_3 = ( conversion.spectral_pca_to_spectral( self.spectral_pca_gamma_0[i], self.spectral_pca_gamma_1[i], self.spectral_pca_gamma_2[i], - self.spectral_pca_gamma_3[i] + self.spectral_pca_gamma_3[i], ) - lambda_1, lambda_2, eos_check = \ + ) + lambda_1, lambda_2, eos_check = ( conversion.spectral_params_to_lambda_1_lambda_2( spectral_gamma_0, spectral_gamma_1, spectral_gamma_2, spectral_gamma_3, self.mass_1_source_spectral[i], - self.mass_2_source_spectral[i] + self.mass_2_source_spectral[i], ) + ) self.assertAlmostEqual(self.lambda_1_spectral[i], float(lambda_1), places=0) self.assertAlmostEqual(self.lambda_2_spectral[i], float(lambda_2), places=0) self.assertAlmostEqual(self.eos_check_spectral[i], eos_check) def test_polytrope_or_causal_params_to_lambda_1_lambda_2_causal(self): - ''' + """ The points cover 6 test cases: - Fail log10_pressure1 >= log10_pressure2 - Fail SimNeutronStarEOS3PDViableFamilyCheck() @@ -964,9 +1031,9 @@ def test_polytrope_or_causal_params_to_lambda_1_lambda_2_causal(self): - Fail mass_1_source <= max_mass - Fail mass_2_source >= min_mass - Passes all and produces accurate lambda_1, lambda_2, eos_check values - ''' + """ for i in range(len(self.mass_1_source_polytrope)): - lambda_1, lambda_2, eos_check = \ + lambda_1, lambda_2, eos_check = ( conversion.polytrope_or_causal_params_to_lambda_1_lambda_2( self.polytrope_gamma_0[i], self.polytrope_log10_pressure_1[i], @@ -975,8 +1042,9 @@ def test_polytrope_or_causal_params_to_lambda_1_lambda_2_causal(self): self.polytrope_gamma_2[i], self.mass_1_source_polytrope[i], self.mass_2_source_polytrope[i], - 0 + 0, ) + ) self.assertAlmostEqual(self.lambda_1_polytrope[i], lambda_1, places=2) self.assertAlmostEqual(self.lambda_2_polytrope[i], lambda_2, places=1) self.assertAlmostEqual(self.eos_check_polytrope[i], eos_check) diff --git a/test/gw/cosmology_test.py b/test/gw/cosmology_test.py index dacbff49c..a3f831bdb 100644 --- a/test/gw/cosmology_test.py +++ b/test/gw/cosmology_test.py @@ -77,7 +77,6 @@ def test_getting_cosmology_non_standard_default(self): class TestPlanck15LALCosmology(unittest.TestCase): - def setUp(self): pass @@ -88,12 +87,7 @@ def test_redshift_to_luminosity_distance(self): # Change to use CreateDefaultCosmologicalParameters once it is # available in a release omega = lal.CreateCosmologicalParameters( - h=0.679, - om=0.3065, - ol=1 - 0.3065, - w0=-1.0, - w1=0.0, - w2=0.0 + h=0.679, om=0.3065, ol=1 - 0.3065, w0=-1.0, w1=0.0, w2=0.0 ) dist_lal = lal.LuminosityDistance(omega=omega, z=z) # Results are the same to within 9 decimal places diff --git a/test/gw/detector/calibration_test.py b/test/gw/detector/calibration_test.py index d06967b16..aae6885d4 100644 --- a/test/gw/detector/calibration_test.py +++ b/test/gw/detector/calibration_test.py @@ -130,7 +130,9 @@ def test_generate_draws(self): draws, parameters = calibration._generate_calibration_draws( self.ifo, self.priors, self.number_of_draws ) - self.assertEqual(draws.shape, (self.number_of_draws, sum(self.ifo.frequency_mask))) + self.assertEqual( + draws.shape, (self.number_of_draws, sum(self.ifo.frequency_mask)) + ) self.assertListEqual(list(self.priors.keys()), list(parameters.keys())) @parameterized.expand([("template",), ("data",), (None,)]) @@ -165,14 +167,16 @@ def test_build_calibration_lookup(self): for ifo in ifos: ifo.minimum_frequency = 20 ifo.maximum_frequency = 1024 - priors.update(prior.CalibrationPriorDict.constant_uncertainty_spline( - amplitude_sigma=0.1, - phase_sigma=0.1, - minimum_frequency=ifo.minimum_frequency, - maximum_frequency=ifo.maximum_frequency, - n_nodes=10, - label=ifo.name, - )) + priors.update( + prior.CalibrationPriorDict.constant_uncertainty_spline( + amplitude_sigma=0.1, + phase_sigma=0.1, + minimum_frequency=ifo.minimum_frequency, + maximum_frequency=ifo.maximum_frequency, + n_nodes=10, + label=ifo.name, + ) + ) ifo.calibration_model = calibration.CubicSpline( prefix=f"recalib_{ifo.name}_", n_points=10, diff --git a/test/gw/detector/geometry_test.py b/test/gw/detector/geometry_test.py index a7c27bb59..ea2016db1 100644 --- a/test/gw/detector/geometry_test.py +++ b/test/gw/detector/geometry_test.py @@ -154,37 +154,49 @@ def test_y_with_latitude_update(self): def test_detector_tensor_with_x_azimuth_update(self): original = self.geometry.detector_tensor self.geometry.xarm_azimuth += 1 - self.assertNotEqual(np.max(np.asarray(abs(self.geometry.detector_tensor - original))), 0) + self.assertNotEqual( + np.max(np.asarray(abs(self.geometry.detector_tensor - original))), 0 + ) self.assertEqual(aac.get_namespace(self.geometry.detector_tensor), self.xp) def test_detector_tensor_with_y_azimuth_update(self): original = self.geometry.detector_tensor self.geometry.yarm_azimuth += 1 - self.assertNotEqual(np.max(np.asarray(abs(self.geometry.detector_tensor - original))), 0) + self.assertNotEqual( + np.max(np.asarray(abs(self.geometry.detector_tensor - original))), 0 + ) self.assertEqual(aac.get_namespace(self.geometry.detector_tensor), self.xp) def test_detector_tensor_with_x_tilt_update(self): original = self.geometry.detector_tensor self.geometry.xarm_tilt += 1 - self.assertNotEqual(np.max(np.asarray(abs(self.geometry.detector_tensor - original))), 0) + self.assertNotEqual( + np.max(np.asarray(abs(self.geometry.detector_tensor - original))), 0 + ) self.assertEqual(aac.get_namespace(self.geometry.detector_tensor), self.xp) def test_detector_tensor_with_y_tilt_update(self): original = self.geometry.detector_tensor self.geometry.yarm_tilt += 1 - self.assertNotEqual(np.max(np.asarray(abs(self.geometry.detector_tensor - original))), 0) + self.assertNotEqual( + np.max(np.asarray(abs(self.geometry.detector_tensor - original))), 0 + ) self.assertEqual(aac.get_namespace(self.geometry.detector_tensor), self.xp) def test_detector_tensor_with_longitude_update(self): original = self.geometry.detector_tensor self.geometry.longitude += 1 - self.assertNotEqual(np.max(np.asarray(abs(self.geometry.detector_tensor - original))), 0) + self.assertNotEqual( + np.max(np.asarray(abs(self.geometry.detector_tensor - original))), 0 + ) self.assertEqual(aac.get_namespace(self.geometry.detector_tensor), self.xp) def test_detector_tensor_with_latitude_update(self): original = self.geometry.detector_tensor self.geometry.latitude += 1 - self.assertNotEqual(np.max(np.asarray(abs(self.geometry.detector_tensor - original))), 0) + self.assertNotEqual( + np.max(np.asarray(abs(self.geometry.detector_tensor - original))), 0 + ) self.assertEqual(aac.get_namespace(self.geometry.detector_tensor), self.xp) def test_unit_vector_along_arm_default(self): diff --git a/test/gw/detector/injection_test.py b/test/gw/detector/injection_test.py index 95d6a3da0..5cf1f2fbe 100644 --- a/test/gw/detector/injection_test.py +++ b/test/gw/detector/injection_test.py @@ -52,7 +52,7 @@ def test_injection_into_timeseries_matches_ifo_injections(): mismatch = 1 - ( sum(whitened_data_1 * whitened_data_2) - / sum(whitened_data_1**2)**0.5 - / sum(whitened_data_2**2)**0.5 + / sum(whitened_data_1**2) ** 0.5 + / sum(whitened_data_2**2) ** 0.5 ) assert mismatch < 3e-3 diff --git a/test/gw/detector/interferometer_test.py b/test/gw/detector/interferometer_test.py index 01da695d1..47ea5cbca 100644 --- a/test/gw/detector/interferometer_test.py +++ b/test/gw/detector/interferometer_test.py @@ -60,8 +60,7 @@ def setUp(self): lambda _: self.wg_polarizations ) self.parameters = dict( - ra=0.0, dec=0.0, geocent_time=0.0, psi=0.0, - mass_1=100, mass_2=100 + ra=0.0, dec=0.0, geocent_time=0.0, psi=0.0, mass_1=100, mass_2=100 ) bilby.core.utils.check_directory_exists_and_if_not_mkdir(self.outdir) @@ -326,13 +325,24 @@ def test_template_template_inner_product(self): signal_1 = np.ones_like(self.ifo.power_spectral_density_array) signal_2 = np.ones_like(self.ifo.power_spectral_density_array) * 2 signal_1_optimal = self.ifo.optimal_snr_squared(signal=signal_1) - signal_1_optimal_by_template_template = self.ifo.template_template_inner_product( - signal_1=signal_1, - signal_2=signal_1 + signal_1_optimal_by_template_template = ( + self.ifo.template_template_inner_product( + signal_1=signal_1, signal_2=signal_1 + ) + ) + self.assertAlmostEqual( + signal_1_optimal, + signal_1_optimal_by_template_template, + delta=signal_1_optimal / 1e14, + ) + signal_1_signal_2_inner_product = self.ifo.template_template_inner_product( + signal_1=signal_1, signal_2=signal_2 + ) + self.assertAlmostEqual( + signal_1_optimal * 2, + signal_1_signal_2_inner_product, + delta=signal_1_optimal / 1e14, ) - self.assertAlmostEqual(signal_1_optimal, signal_1_optimal_by_template_template, delta=signal_1_optimal / 1e14) - signal_1_signal_2_inner_product = self.ifo.template_template_inner_product(signal_1=signal_1, signal_2=signal_2) - self.assertAlmostEqual(signal_1_optimal * 2, signal_1_signal_2_inner_product, delta=signal_1_optimal / 1e14) def test_repr(self): expected = ( @@ -363,6 +373,7 @@ def test_to_and_from_pkl_loading(self): def test_to_and_from_pkl_wrong_class(self): import dill + with open("./outdir/psd.pkl", "wb") as ff: dill.dump(self.ifo.power_spectral_density, ff) filename = self.ifo._filename_from_outdir_label_extension( @@ -380,7 +391,7 @@ def test_psd_not_impacted_by_window_factor(monkeypatch): ifo.power_spectral_density.get_power_spectral_density_array( frequency_array=ifo.strain_data.frequency_array ), - ifo.power_spectral_density_array + ifo.power_spectral_density_array, ) @@ -526,9 +537,16 @@ def test_eq_false_different_ifo_strain_data(self): class TestInterferometerAntennaPatternAgainstLAL(unittest.TestCase): def setUp(self): self.name = "name" - self.ifo_names = ['H1', 'L1', 'V1', 'K1', 'G1', 'ET'] - self.lal_prefixes = {'H1': 'H1', 'L1': 'L1', 'V1': 'V1', 'K1': 'K1', 'G1': 'G1', 'ET': 'E1'} - self.polarizations = ['plus', 'cross', 'breathing', 'longitudinal', 'x', 'y'] + self.ifo_names = ["H1", "L1", "V1", "K1", "G1", "ET"] + self.lal_prefixes = { + "H1": "H1", + "L1": "L1", + "V1": "V1", + "K1": "K1", + "G1": "G1", + "ET": "E1", + } + self.polarizations = ["plus", "cross", "breathing", "longitudinal", "x", "y"] self.ifos = bilby.gw.detector.InterferometerList(self.ifo_names) self.gpstime = 1305303144 self.trial = 100 @@ -548,19 +566,25 @@ def test_antenna_pattern_vs_lal(self): f_lal = np.zeros((self.trial, 6)) for n, ifo_name in enumerate(self.ifo_names): - response = lalsimulation.DetectorPrefixToLALDetector(self.lal_prefixes[ifo_name]).response + response = lalsimulation.DetectorPrefixToLALDetector( + self.lal_prefixes[ifo_name] + ).response ifo = self.ifos[n] for i in range(self.trial): - ra = 2. * np.pi * np.random.uniform() - dec = np.pi * np.random.uniform() - np.pi / 2. + ra = 2.0 * np.pi * np.random.uniform() + dec = np.pi * np.random.uniform() - np.pi / 2.0 psi = np.pi * np.random.uniform() - f_lal[i] = lal.ComputeDetAMResponseExtraModes(response, ra, dec, psi, gmst) + f_lal[i] = lal.ComputeDetAMResponseExtraModes( + response, ra, dec, psi, gmst + ) for m, pol in enumerate(self.polarizations): - f_bilby[i, m] = ifo.antenna_response(ra, dec, self.gpstime, psi, pol) + f_bilby[i, m] = ifo.antenna_response( + ra, dec, self.gpstime, psi, pol + ) std = np.std(f_bilby - f_lal, axis=0) for m, pol in enumerate(self.polarizations): - with self.subTest(':'.join((ifo_name, pol))): + with self.subTest(":".join((ifo_name, pol))): self.assertAlmostEqual(std[m], 0.0, places=7) def test_time_delay_vs_lal(self): @@ -571,12 +595,11 @@ def test_time_delay_vs_lal(self): det = lal.cached_detector_by_prefix[self.lal_prefixes[ifo_name]] for i in range(self.trial): gpstime = np.random.uniform(1205303144, 1405303144) - ra = 2. * np.pi * np.random.uniform() - dec = np.pi * np.random.uniform() - np.pi / 2. - delays[i] = ( - lal.TimeDelayFromEarthCenter(det.location, ra, dec, gpstime) - - ifo.time_delay_from_geocenter(ra, dec, gpstime) - ) + ra = 2.0 * np.pi * np.random.uniform() + dec = np.pi * np.random.uniform() - np.pi / 2.0 + delays[i] = lal.TimeDelayFromEarthCenter( + det.location, ra, dec, gpstime + ) - ifo.time_delay_from_geocenter(ra, dec, gpstime) std = max(abs(delays)) with self.subTest(ifo_name): @@ -588,33 +611,33 @@ class TestInterferometerWhitenedStrain(unittest.TestCase): def setUp(self): self.duration = 64 self.sampling_frequency = 4096 - self.ifo = bilby.gw.detector.get_empty_interferometer('H1') + self.ifo = bilby.gw.detector.get_empty_interferometer("H1") self.ifo.set_strain_data_from_power_spectral_density( - sampling_frequency=self.sampling_frequency, duration=self.duration) + sampling_frequency=self.sampling_frequency, duration=self.duration + ) self.waveform_generator = bilby.gw.waveform_generator.WaveformGenerator( duration=self.duration, sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole, - waveform_arguments={ - "waveform_approximant": "IMRPhenomXP" - }) + waveform_arguments={"waveform_approximant": "IMRPhenomXP"}, + ) self.parameters = { - 'mass_1': 10, - 'mass_2': 10, - 'a_1': 0, - 'a_2': 0, - 'tilt_1': 0, - 'tilt_2': 0, - 'phi_12': 0, - 'phi_jl': 0, - 'theta_jn': 0, - 'luminosity_distance': 40, - 'phase': 0, - 'ra': 0, - 'dec': 0, - 'geocent_time': 62, - 'psi': 0 + "mass_1": 10, + "mass_2": 10, + "a_1": 0, + "a_2": 0, + "tilt_1": 0, + "tilt_2": 0, + "phi_12": 0, + "phi_jl": 0, + "theta_jn": 0, + "luminosity_distance": 40, + "phase": 0, + "ra": 0, + "dec": 0, + "geocent_time": 62, + "psi": 0, } def tearDown(self): @@ -645,33 +668,43 @@ def test_time_domain_whitened_strain(self): def test_frequency_domain_noise_and_signal_whitening(self): # Inject some (loud) signal - self.ifo.inject_signal(waveform_generator=self.waveform_generator, parameters=self.parameters) + self.ifo.inject_signal( + waveform_generator=self.waveform_generator, parameters=self.parameters + ) # Make the template separately - waveform_polarizations = self.waveform_generator.frequency_domain_strain(parameters=self.parameters) - signal_ifo = self.ifo.get_detector_response( - waveform_polarizations=waveform_polarizations, + waveform_polarizations = self.waveform_generator.frequency_domain_strain( parameters=self.parameters ) + signal_ifo = self.ifo.get_detector_response( + waveform_polarizations=waveform_polarizations, parameters=self.parameters + ) # Whiten the template whitened_signal_ifo = self.ifo.whiten_frequency_series(signal_ifo) mask = self.ifo.frequency_mask - white = self.ifo.whitened_frequency_domain_strain[mask] - whitened_signal_ifo[mask] + white = ( + self.ifo.whitened_frequency_domain_strain[mask] - whitened_signal_ifo[mask] + ) self._check_frequency_series_whiteness(white) def test_time_domain_noise_and_signal_whitening(self): # Inject some (loud) signal - self.ifo.inject_signal(waveform_generator=self.waveform_generator, parameters=self.parameters) + self.ifo.inject_signal( + waveform_generator=self.waveform_generator, parameters=self.parameters + ) # Make the template separately - waveform_polarizations = self.waveform_generator.frequency_domain_strain(parameters=self.parameters) - signal_ifo = self.ifo.get_detector_response( - waveform_polarizations=waveform_polarizations, + waveform_polarizations = self.waveform_generator.frequency_domain_strain( parameters=self.parameters ) + signal_ifo = self.ifo.get_detector_response( + waveform_polarizations=waveform_polarizations, parameters=self.parameters + ) # Whiten the template in FD whitened_signal_ifo_fd = self.ifo.whiten_frequency_series(signal_ifo) # Get whitened template in TD - whitened_signal_ifo_td = self.ifo.get_whitened_time_series_from_whitened_frequency_series( - whitened_signal_ifo_fd + whitened_signal_ifo_td = ( + self.ifo.get_whitened_time_series_from_whitened_frequency_series( + whitened_signal_ifo_fd + ) ) whitened_td = self.ifo.whitened_time_domain_strain - whitened_signal_ifo_td self._check_time_series_whiteness(whitened_td) diff --git a/test/gw/detector/networks_test.py b/test/gw/detector/networks_test.py index a3ef2286e..45ff6b59a 100644 --- a/test/gw/detector/networks_test.py +++ b/test/gw/detector/networks_test.py @@ -415,7 +415,9 @@ def c(a): pair_a = a(lat_1 - lat_0, long_1 - long_0, lat_0, lat_1) pair_c = c(pair_a) distance = bilby.core.utils.radius_of_earth * pair_c - self.assertAlmostEqual(float(distance) / 1000, float(pair[0].length), delta=1) + self.assertAlmostEqual( + float(distance) / 1000, float(pair[0].length), delta=1 + ) def test_individual_positions_exact(self): """ @@ -426,7 +428,9 @@ def test_individual_positions_exact(self): for pair in combinations(self.triangular_ifo, 2): vertex_0 = pair[0].vertex_position_geocentric() vertex_1 = pair[1].vertex_position_geocentric() - self.assertAlmostEqual(float(pair[0].length), float(pair[1].length), delta=1e-6) + self.assertAlmostEqual( + float(pair[0].length), float(pair[1].length), delta=1e-6 + ) distance = self.xp.linalg.vector_norm(vertex_0 - vertex_1) / 1000 self.assertAlmostEqual(float(distance), float(pair[0].length), delta=1e-6) @@ -441,12 +445,20 @@ def test_arm_endpoints_coincide_with_neighbouring_vertices(self): for ii, ifo in enumerate(self.triangular_ifo): vertex = ifo.vertex_position_geocentric() arm_length = ifo.length * 1000 - next_vertex = self.triangular_ifo[(ii + 1) % n_ifos].vertex_position_geocentric() - previous_vertex = self.triangular_ifo[(ii - 1) % n_ifos].vertex_position_geocentric() + next_vertex = self.triangular_ifo[ + (ii + 1) % n_ifos + ].vertex_position_geocentric() + previous_vertex = self.triangular_ifo[ + (ii - 1) % n_ifos + ].vertex_position_geocentric() x_arm_end = vertex + arm_length * ifo.geometry.x y_arm_end = vertex + arm_length * ifo.geometry.y - self.assertLess(float(self.xp.linalg.vector_norm(x_arm_end - next_vertex)), 1e-6) - self.assertLess(float(self.xp.linalg.vector_norm(y_arm_end - previous_vertex)), 1e-6) + self.assertLess( + float(self.xp.linalg.vector_norm(x_arm_end - next_vertex)), 1e-6 + ) + self.assertLess( + float(self.xp.linalg.vector_norm(y_arm_end - previous_vertex)), 1e-6 + ) def test_detector_tensors_sum_to_zero(self): """ diff --git a/test/gw/detector/psd_test.py b/test/gw/detector/psd_test.py index 1f718c6ae..1cdbec70d 100644 --- a/test/gw/detector/psd_test.py +++ b/test/gw/detector/psd_test.py @@ -74,8 +74,10 @@ def test_from_amplitude_spectral_density_array(self): self.assertTrue(np.array_equal(self.asd_array, actual.asd_array)) def test_from_power_spectral_density_array(self): - actual = bilby.gw.detector.PowerSpectralDensity.from_power_spectral_density_array( - frequency_array=self.frequency_array, psd_array=self.psd_array + actual = ( + bilby.gw.detector.PowerSpectralDensity.from_power_spectral_density_array( + frequency_array=self.frequency_array, psd_array=self.psd_array + ) ) self.assertTrue(np.array_equal(self.psd_array, actual.psd_array)) self.assertTrue(np.array_equal(self.asd_array, actual.asd_array)) @@ -171,8 +173,10 @@ def test_power_spectral_density_interpolated_from_psd_file(self): self.assertAlmostEqual(expected, psd.power_spectral_density_interpolated(2)) def test_from_amplitude_spectral_density_file(self): - psd = bilby.gw.detector.PowerSpectralDensity.from_amplitude_spectral_density_file( - asd_file=self.asd_file + psd = ( + bilby.gw.detector.PowerSpectralDensity.from_amplitude_spectral_density_file( + asd_file=self.asd_file + ) ) self.assertEqual(self.asd_file, psd.asd_file) self.assertTrue(np.allclose(self.psd_array, psd.psd_array, atol=1e-60)) @@ -247,11 +251,15 @@ def setUp(self): self.psd_from_file_2 = bilby.gw.detector.PowerSpectralDensity.from_aligo() self.frequency_array = np.linspace(1, 100) self.psd_array = np.linspace(1, 100) - self.psd_from_array_1 = bilby.gw.detector.PowerSpectralDensity.from_power_spectral_density_array( - frequency_array=self.frequency_array, psd_array=self.psd_array + self.psd_from_array_1 = ( + bilby.gw.detector.PowerSpectralDensity.from_power_spectral_density_array( + frequency_array=self.frequency_array, psd_array=self.psd_array + ) ) - self.psd_from_array_2 = bilby.gw.detector.PowerSpectralDensity.from_power_spectral_density_array( - frequency_array=self.frequency_array, psd_array=self.psd_array + self.psd_from_array_2 = ( + bilby.gw.detector.PowerSpectralDensity.from_power_spectral_density_array( + frequency_array=self.frequency_array, psd_array=self.psd_array + ) ) def tearDown(self): diff --git a/test/gw/detector/strain_data_test.py b/test/gw/detector/strain_data_test.py index 0f82a40a2..ed65285f8 100644 --- a/test/gw/detector/strain_data_test.py +++ b/test/gw/detector/strain_data_test.py @@ -36,10 +36,11 @@ def test_frequency_mask(self): def test_frequency_mask_2(self): strain_data = bilby.gw.detector.InterferometerStrainData( - minimum_frequency=20, maximum_frequency=512) + minimum_frequency=20, maximum_frequency=512 + ) strain_data.set_from_time_domain_strain( time_domain_strain=np.random.normal(0, 1, 4096), - time_array=np.arange(0, 4, 4 / 4096) + time_array=np.arange(0, 4, 4 / 4096), ) # Test from init @@ -56,10 +57,11 @@ def test_frequency_mask_2(self): def test_notches_frequency_mask(self): strain_data = bilby.gw.detector.InterferometerStrainData( - minimum_frequency=20, maximum_frequency=512, notch_list=[(100, 101)]) + minimum_frequency=20, maximum_frequency=512, notch_list=[(100, 101)] + ) strain_data.set_from_time_domain_strain( time_domain_strain=np.random.normal(0, 1, 4096), - time_array=np.arange(0, 4, 4 / 4096) + time_array=np.arange(0, 4, 4 / 4096), ) # Test from init @@ -173,7 +175,7 @@ def test_time_domain_window_no_roll_off_no_alpha(self): len(self.ifosd._time_domain_strain), alpha=self.ifosd.alpha ) self.assertEqual(expected_window, self.ifosd.time_domain_window()) - self.assertEqual(np.mean(expected_window ** 2), self.ifosd.window_factor) + self.assertEqual(np.mean(expected_window**2), self.ifosd.window_factor) def test_time_domain_window_sets_roll_off_directly(self): self.ifosd._time_domain_strain = np.array([3]) @@ -364,7 +366,9 @@ def setUp(self): self.maximum_frequency = 1024 def test_init(self): - notch = bilby.gw.detector.strain_data.Notch(self.minimum_frequency, self.maximum_frequency) + notch = bilby.gw.detector.strain_data.Notch( + self.minimum_frequency, self.maximum_frequency + ) self.assertEqual(notch.minimum_frequency, self.minimum_frequency) self.assertEqual(notch.maximum_frequency, self.maximum_frequency) @@ -384,7 +388,9 @@ def test_init_fail(self): bilby.gw.detector.strain_data.Notch(30, 20) def test_idxs(self): - notch = bilby.gw.detector.strain_data.Notch(self.minimum_frequency, self.maximum_frequency) + notch = bilby.gw.detector.strain_data.Notch( + self.minimum_frequency, self.maximum_frequency + ) freqs = np.linspace(0, 2048, 100) idxs = notch.get_idxs(freqs) self.assertEqual(len(idxs), len(freqs)) @@ -394,7 +400,6 @@ def test_idxs(self): class TestNotchList(unittest.TestCase): - def test_init_single(self): notch_list_of_tuples = [(32, 34)] notch_list = bilby.gw.detector.strain_data.NotchList(notch_list_of_tuples) diff --git a/test/gw/eos/eos_test.py b/test/gw/eos/eos_test.py index 8f23b6a95..7b7e78907 100644 --- a/test/gw/eos/eos_test.py +++ b/test/gw/eos/eos_test.py @@ -8,11 +8,12 @@ KNOWN_TOV_RESULT = 1 KNOWN_EOS_PRIOR_RESULTS = 1 ENERGY_FROM_PRESSURE_RESULTS = 0.04047517810698063 -EOS_FROM_TABLE = EOSFamily(TabularEOS('AP4', True)) -EOS_FROM_SPRECTRAL_DECOMPOSITION = EOSFamily(SpectralDecompositionEOS(gammas=[0.8651, 0.1548, -0.0151, -0.0002], - p0=1.5e33, - e0=2.02e14, - xmax=7.04)) +EOS_FROM_TABLE = EOSFamily(TabularEOS("AP4", True)) +EOS_FROM_SPRECTRAL_DECOMPOSITION = EOSFamily( + SpectralDecompositionEOS( + gammas=[0.8651, 0.1548, -0.0151, -0.0002], p0=1.5e33, e0=2.02e14, xmax=7.04 + ) +) MASSES_TO_TEST = numpy.linspace(0.9, 2.0, 12) PRESSURE = 1e-11 ENTHALPY = 1.2 @@ -20,49 +21,80 @@ # Setup for MPA1 Comparison PRESSURE_0 = 5.3716e32 -ENERGY_DENSITY_0 = 1.1555e35 / (utils.speed_of_light * 100) ** 2. +ENERGY_DENSITY_0 = 1.1555e35 / (utils.speed_of_light * 100) ** 2.0 XMAX = 12.3081 GAMMAS = [1.0215, 0.1653, -0.0235, -0.0004] -LALSIM_MPA1 = lalsim.SimNeutronStarEOS4ParameterSpectralDecomposition(GAMMAS[0], GAMMAS[1], GAMMAS[2], GAMMAS[3]) +LALSIM_MPA1 = lalsim.SimNeutronStarEOS4ParameterSpectralDecomposition( + GAMMAS[0], GAMMAS[1], GAMMAS[2], GAMMAS[3] +) BILBY_MPA1 = SpectralDecompositionEOS(GAMMAS, PRESSURE_0, ENERGY_DENSITY_0, XMAX) MPA1_PRESSURES = BILBY_MPA1.e_pdat.T[0] -LALSIM_MPA1_ENERGY_DENSITY = [lalsim.SimNeutronStarEOSEnergyDensityOfPressureGeometerized(PRESSURE, LALSIM_MPA1) - for PRESSURE in MPA1_PRESSURES[:-1]] +LALSIM_MPA1_ENERGY_DENSITY = [ + lalsim.SimNeutronStarEOSEnergyDensityOfPressureGeometerized(PRESSURE, LALSIM_MPA1) + for PRESSURE in MPA1_PRESSURES[:-1] +] BILBY_MPA1_ENERGY_DENSITY = BILBY_MPA1.e_pdat.T[1][:-1] class TestEOSFamily(unittest.TestCase): def test_spectral_decomposition_energy_from_pressure(self): - self.assertAlmostEqual(EOS_FROM_TABLE.eos.energy_from_pressure(PRESSURE), - 3.2736497985232014e-10) + self.assertAlmostEqual( + EOS_FROM_TABLE.eos.energy_from_pressure(PRESSURE), 3.2736497985232014e-10 + ) - self.assertAlmostEqual(EOS_FROM_SPRECTRAL_DECOMPOSITION.eos.energy_from_pressure(PRESSURE), - 3.270622527256167e-10) + self.assertAlmostEqual( + EOS_FROM_SPRECTRAL_DECOMPOSITION.eos.energy_from_pressure(PRESSURE), + 3.270622527256167e-10, + ) def test_spectral_decomposition_pressure_from_pseudo_enthalpy(self): - self.assertAlmostEqual(EOS_FROM_TABLE.eos.pressure_from_pseudo_enthalpy(ENTHALPY), - 2.7338376042831513e-09) + self.assertAlmostEqual( + EOS_FROM_TABLE.eos.pressure_from_pseudo_enthalpy(ENTHALPY), + 2.7338376042831513e-09, + ) - self.assertAlmostEqual(EOS_FROM_SPRECTRAL_DECOMPOSITION.eos.pressure_from_pseudo_enthalpy(ENTHALPY), - 2.754018499535077e-09) + self.assertAlmostEqual( + EOS_FROM_SPRECTRAL_DECOMPOSITION.eos.pressure_from_pseudo_enthalpy( + ENTHALPY + ), + 2.754018499535077e-09, + ) def test_spectral_decomposition_energy_density_from_pseudo_enthalpy(self): - self.assertAlmostEqual(EOS_FROM_TABLE.eos.energy_density_from_pseudo_enthalpy(ENTHALPY), - 2.9486942467903607e-09) + self.assertAlmostEqual( + EOS_FROM_TABLE.eos.energy_density_from_pseudo_enthalpy(ENTHALPY), + 2.9486942467903607e-09, + ) - self.assertAlmostEqual(EOS_FROM_SPRECTRAL_DECOMPOSITION.eos.energy_density_from_pseudo_enthalpy(ENTHALPY), - 3.0402598495601078e-09) + self.assertAlmostEqual( + EOS_FROM_SPRECTRAL_DECOMPOSITION.eos.energy_density_from_pseudo_enthalpy( + ENTHALPY + ), + 3.0402598495601078e-09, + ) def test_spectral_decomposition_pseudo_enthalpy_from_energy_density(self): - self.assertAlmostEqual(EOS_FROM_TABLE.eos.pseudo_enthalpy_from_energy_density(ENERGY_DENSITY), - 0.024415755781136812) - print(EOS_FROM_SPRECTRAL_DECOMPOSITION.eos.pseudo_enthalpy_from_energy_density(ENERGY_DENSITY)) - self.assertAlmostEqual(EOS_FROM_SPRECTRAL_DECOMPOSITION.eos.pseudo_enthalpy_from_energy_density(ENERGY_DENSITY), - 0.02420629785967365) + self.assertAlmostEqual( + EOS_FROM_TABLE.eos.pseudo_enthalpy_from_energy_density(ENERGY_DENSITY), + 0.024415755781136812, + ) + print( + EOS_FROM_SPRECTRAL_DECOMPOSITION.eos.pseudo_enthalpy_from_energy_density( + ENERGY_DENSITY + ) + ) + self.assertAlmostEqual( + EOS_FROM_SPRECTRAL_DECOMPOSITION.eos.pseudo_enthalpy_from_energy_density( + ENERGY_DENSITY + ), + 0.02420629785967365, + ) class TestBilbyLALSimComparison(unittest.TestCase): def test_spectral_decomposition_MPA1(self): - numpy.testing.assert_allclose(LALSIM_MPA1_ENERGY_DENSITY, BILBY_MPA1_ENERGY_DENSITY, rtol=1e6) + numpy.testing.assert_allclose( + LALSIM_MPA1_ENERGY_DENSITY, BILBY_MPA1_ENERGY_DENSITY, rtol=1e6 + ) diff --git a/test/gw/likelihood/base_test.py b/test/gw/likelihood/base_test.py index ff0e7c130..05fe08cd0 100644 --- a/test/gw/likelihood/base_test.py +++ b/test/gw/likelihood/base_test.py @@ -67,9 +67,7 @@ def test_noise_log_likelihood(self): """Test noise log likelihood matches precomputed value""" nll = self.likelihood.noise_log_likelihood() self.assertEqual(aac.get_namespace(nll), self.xp) - self.assertAlmostEqual( - -4014.1787704539474, float(nll), 3 - ) + self.assertAlmostEqual(-4014.1787704539474, float(nll), 3) def test_log_likelihood(self): """Test log likelihood matches precomputed value""" @@ -81,7 +79,8 @@ def test_log_likelihood_ratio(self): """Test log likelihood ratio returns the correct value""" llr = self.likelihood.log_likelihood_ratio(self.parameters) self.assertAlmostEqual( - float(self.likelihood.log_likelihood(self.parameters)) - float(self.likelihood.noise_log_likelihood()), + float(self.likelihood.log_likelihood(self.parameters)) + - float(self.likelihood.noise_log_likelihood()), float(llr), 3, ) @@ -91,7 +90,10 @@ def test_likelihood_zero_when_waveform_is_none(self): """Test log likelihood returns np.nan_to_num(-np.inf) when the waveform is None""" self.likelihood.waveform_generator.frequency_domain_strain = lambda x: None - self.assertEqual(self.likelihood.log_likelihood_ratio(self.parameters), np.nan_to_num(-np.inf)) + self.assertEqual( + self.likelihood.log_likelihood_ratio(self.parameters), + np.nan_to_num(-np.inf), + ) def test_repr(self): expected = ( @@ -118,7 +120,9 @@ def test_interferometers_setting_list(self): self.assertListEqual( bilby.gw.detector.InterferometerList(ifos), self.likelihood.interferometers ) - self.assertIsInstance(self.likelihood.interferometers, bilby.gw.detector.InterferometerList) + self.assertIsInstance( + self.likelihood.interferometers, bilby.gw.detector.InterferometerList + ) def test_interferometers_setting_interferometer_list(self): ifos = bilby.gw.detector.InterferometerList( @@ -131,7 +135,9 @@ def test_interferometers_setting_interferometer_list(self): self.assertListEqual( bilby.gw.detector.InterferometerList(ifos), self.likelihood.interferometers ) - self.assertIsInstance(self.likelihood.interferometers, bilby.gw.detector.InterferometerList) + self.assertIsInstance( + self.likelihood.interferometers, bilby.gw.detector.InterferometerList + ) def test_meta_data(self): expected = dict( @@ -161,7 +167,7 @@ def test_reference_frame_agrees_with_default(self): interferometers=self.interferometers, waveform_generator=self.waveform_generator, priors=self.prior.copy(), - reference_frame="H1L1" + reference_frame="H1L1", ) parameters = self.parameters.copy() del parameters["ra"], parameters["dec"] @@ -177,7 +183,7 @@ def test_reference_frame_agrees_with_default(self): self.assertEqual(aac.get_namespace(parameters["dec"]), self.xp) self.assertEqual( new_likelihood.log_likelihood_ratio(parameters), - self.likelihood.log_likelihood_ratio(parameters) + self.likelihood.log_likelihood_ratio(parameters), ) def test_time_reference_agrees_with_default(self): @@ -185,13 +191,13 @@ def test_time_reference_agrees_with_default(self): interferometers=self.interferometers, waveform_generator=self.waveform_generator, priors=self.prior.copy(), - time_reference="H1" + time_reference="H1", ) ifo = bilby.gw.detector.get_empty_interferometer("H1") time_delay = ifo.time_delay_from_geocenter( ra=self.parameters["ra"], dec=self.parameters["dec"], - time=self.parameters["geocent_time"] + time=self.parameters["geocent_time"], ) parameters = self.parameters.copy() parameters["H1_time"] = parameters["geocent_time"] + time_delay diff --git a/test/gw/likelihood/basic_test.py b/test/gw/likelihood/basic_test.py index 516ecbc32..3e2c77b88 100644 --- a/test/gw/likelihood/basic_test.py +++ b/test/gw/likelihood/basic_test.py @@ -56,9 +56,7 @@ def tearDown(self): def test_noise_log_likelihood(self): """Test noise log likelihood matches precomputed value""" nll = self.likelihood.noise_log_likelihood() - self.assertAlmostEqual( - -4014.1787704539474, float(nll), 3 - ) + self.assertAlmostEqual(-4014.1787704539474, float(nll), 3) self.assertEqual(aac.get_namespace(nll), self.xp) def test_log_likelihood(self): @@ -71,7 +69,8 @@ def test_log_likelihood_ratio(self): """Test log likelihood ratio returns the correct value""" llr = self.likelihood.log_likelihood_ratio(self.parameters) self.assertAlmostEqual( - self.likelihood.log_likelihood(self.parameters) - self.likelihood.noise_log_likelihood(), + self.likelihood.log_likelihood(self.parameters) + - self.likelihood.noise_log_likelihood(), llr, 3, ) @@ -81,7 +80,10 @@ def test_likelihood_zero_when_waveform_is_none(self): """Test log likelihood returns np.nan_to_num(-np.inf) when the waveform is None""" self.likelihood.waveform_generator.frequency_domain_strain = lambda x: None - self.assertEqual(self.likelihood.log_likelihood_ratio(self.parameters), np.nan_to_num(-np.inf)) + self.assertEqual( + self.likelihood.log_likelihood_ratio(self.parameters), + np.nan_to_num(-np.inf), + ) def test_repr(self): expected = "BasicGravitationalWaveTransient(interferometers={},\n\twaveform_generator={})".format( diff --git a/test/gw/likelihood/jit_test.py b/test/gw/likelihood/jit_test.py index 1fc06e255..bbcfe0685 100644 --- a/test/gw/likelihood/jit_test.py +++ b/test/gw/likelihood/jit_test.py @@ -12,7 +12,7 @@ from bilby.gw.waveform_generator import WaveformGenerator try: - from bilby.compat import pytrees as _ # noqa + from bilby.compat import pytrees as _ # noqa from bilby.gw.compat import pytrees as _ # noqa except ImportError: pass @@ -75,10 +75,12 @@ def likelihood(xp, **marginalizations): parameter_conversion=null_convert, use_cache=False, ) - priors = PriorDict(dict( - phase=Uniform(0, 2 * np.pi), - geocent_time=Uniform(0, 4), - )) + priors = PriorDict( + dict( + phase=Uniform(0, 2 * np.pi), + geocent_time=Uniform(0, 4), + ) + ) likelihood = GravitationalWaveTransient( interferometers=interferometers, @@ -103,7 +105,7 @@ def parameters(xp): ) params.update({f"recalib_H1_amplitude_{i}": 0.0 for i in range(10)}) params.update({f"recalib_H1_phase_{i}": 0.0 for i in range(10)}) - params.update({f"recalib_H1_frequency_{i}": 20 * 10**(i / 10) for i in range(10)}) + params.update({f"recalib_H1_frequency_{i}": 20 * 10 ** (i / 10) for i in range(10)}) return params @@ -125,4 +127,8 @@ def test_jitted_likelihood_with_time_marginalization(xp, parameters): @pytest.mark.array_backend def test_jitted_likelihood_with_phase_time_marginalization(xp, parameters): - _evaluate_with_jit(likelihood(xp, phase_marginalization=True, time_marginalization=True), parameters, xp) + _evaluate_with_jit( + likelihood(xp, phase_marginalization=True, time_marginalization=True), + parameters, + xp, + ) diff --git a/test/gw/likelihood/marginalization_test.py b/test/gw/likelihood/marginalization_test.py index b5da6ba16..d3aca94d9 100644 --- a/test/gw/likelihood/marginalization_test.py +++ b/test/gw/likelihood/marginalization_test.py @@ -167,6 +167,7 @@ class TestMarginalizations(unittest.TestCase): For time, this is strongly dependent on the specific time grid used. The `time_jitter` parameter makes this a weaker dependence during sampling. """ + _parameters = product( ["regular", "roq", "relbin", "multiband"], ["luminosity_distance", "geocent_time", "phase"], @@ -219,7 +220,7 @@ def setUp(self): reference_frequency=20.0, minimum_frequency=20.0, waveform_approximant="IMRPhenomPv2", - ) + ), ) self.interferometers.inject_signal( parameters=self.parameters, waveform_generator=self.waveform_generator @@ -228,7 +229,7 @@ def setUp(self): self.priors = bilby.gw.prior.BBHPriorDict() self.priors["geocent_time"] = bilby.prior.Uniform( minimum=self.parameters["geocent_time"] - 0.1, - maximum=self.parameters["geocent_time"] + 0.1 + maximum=self.parameters["geocent_time"] + 0.1, ) self.relbin_waveform_generator = bilby.gw.waveform_generator.WaveformGenerator( @@ -240,7 +241,7 @@ def setUp(self): reference_frequency=20.0, minimum_frequency=20.0, waveform_approximant="IMRPhenomPv2", - ) + ), ) self.multiband_waveform_generator = bilby.gw.WaveformGenerator( @@ -251,7 +252,7 @@ def setUp(self): waveform_arguments=dict( reference_frequency=20.0, waveform_approximant="IMRPhenomPv2", - ) + ), ) @property @@ -287,8 +288,10 @@ def roq_waveform_generator(self): reference_frequency=20.0, waveform_approximant="IMRPhenomPv2", frequency_nodes_linear=np.load(f"{self.roq_dir}/fnodes_linear.npy"), - frequency_nodes_quadratic=np.load(f"{self.roq_dir}/fnodes_quadratic.npy"), - ) + frequency_nodes_quadratic=np.load( + f"{self.roq_dir}/fnodes_quadratic.npy" + ), + ), ) def tearDown(self): @@ -302,11 +305,22 @@ def tearDown(self): @classmethod def tearDownClass(cls): """remove lookup tables so that they are not used accidentally in subsequent tests""" - for filename in [cls.lookup_phase, cls.lookup_no_phase, cls.path_to_roq_weights]: + for filename in [ + cls.lookup_phase, + cls.lookup_no_phase, + cls.path_to_roq_weights, + ]: if os.path.exists(filename): os.remove(filename) - def likelihood_kwargs(self, kind, time_marginalization, phase_marginalization, distance_marginalization, priors): + def likelihood_kwargs( + self, + kind, + time_marginalization, + phase_marginalization, + distance_marginalization, + priors, + ): if priors is None: priors = deepcopy(self.priors) if distance_marginalization and phase_marginalization: @@ -325,11 +339,13 @@ def likelihood_kwargs(self, kind, time_marginalization, phase_marginalization, d priors=priors, ) if kind == "roq": - kwargs.update(dict( - linear_matrix=self.roq_linear_matrix_file, - quadratic_matrix=self.roq_quadratic_matrix_file, - waveform_generator=self.roq_waveform_generator, - )) + kwargs.update( + dict( + linear_matrix=self.roq_linear_matrix_file, + quadratic_matrix=self.roq_quadratic_matrix_file, + waveform_generator=self.roq_waveform_generator, + ) + ) if os.path.exists(self.__class__.path_to_roq_weights): kwargs["weights"] = self.__class__.path_to_roq_weights elif kind == "relbin": @@ -338,9 +354,8 @@ def likelihood_kwargs(self, kind, time_marginalization, phase_marginalization, d elif kind == "multiband": kwargs["waveform_generator"] = self.multiband_waveform_generator kwargs["reference_chirp_mass"] = ( - (self.parameters["mass_1"] * self.parameters["mass_2"])**0.6 / - (self.parameters["mass_1"] + self.parameters["mass_2"])**0.2 - ) + self.parameters["mass_1"] * self.parameters["mass_2"] + ) ** 0.6 / (self.parameters["mass_1"] + self.parameters["mass_2"]) ** 0.2 return kwargs def get_likelihood( @@ -349,10 +364,14 @@ def get_likelihood( time_marginalization=False, phase_marginalization=False, distance_marginalization=False, - priors=None + priors=None, ): kwargs = self.likelihood_kwargs( - kind, time_marginalization, phase_marginalization, distance_marginalization, priors + kind, + time_marginalization, + phase_marginalization, + distance_marginalization, + priors, ) if kind == "regular": cls_ = bilby.gw.likelihood.GravitationalWaveTransient @@ -403,11 +422,11 @@ def _template(self, marginalized, non_marginalized, key, prior=None, values=None @parameterized.expand( _parameters, name_func=lambda func, num, param: ( - f"{func.__name__}_{num}__{param.args[0]}_{param.args[1]}_" + "_".join([ - ["D", "T", "P"][ii] for ii, val - in enumerate(param.args[-3:]) if val - ]) - ) + f"{func.__name__}_{num}__{param.args[0]}_{param.args[1]}_" + + "_".join( + [["D", "T", "P"][ii] for ii, val in enumerate(param.args[-3:]) if val] + ) + ), ) def test_marginalisation(self, kind, key, distance, time, phase): if all([distance, time, phase]): @@ -449,13 +468,16 @@ def test_time_marginalisation_full_segment(self, kind): ) @parameterized.expand( - itertools.product(["regular", "roq", "relbin", "multiband"], *itertools.repeat([True, False], 3)), + itertools.product( + ["regular", "roq", "relbin", "multiband"], + *itertools.repeat([True, False], 3), + ), name_func=lambda func, num, param: ( - f"{func.__name__}_{num}__{param.args[0]}_" + "_".join([ - ["D", "P", "T"][ii] for ii, val - in enumerate(param.args[1:]) if val - ]) - ) + f"{func.__name__}_{num}__{param.args[0]}_" + + "_".join( + [["D", "P", "T"][ii] for ii, val in enumerate(param.args[1:]) if val] + ) + ), ) def test_marginalization_reconstruction(self, kind, distance, phase, time): marginalizations = dict( @@ -476,11 +498,12 @@ def test_marginalization_reconstruction(self, kind, distance, phase, time): ) output = like.generate_posterior_sample_from_marginalized_likelihood(params) for key in marginalizations: - self.assertFalse(marginalizations[key] and reference_values[key] == output[key]) + self.assertFalse( + marginalizations[key] and reference_values[key] == output[key] + ) class CalibrationMarginalization(unittest.TestCase): - def setUp(self): self.ifos = bilby.gw.detector.InterferometerList(["H1", "L1"]) self.ifos.set_strain_data_from_power_spectral_densities( @@ -492,7 +515,9 @@ def setUp(self): maximum_frequency=512, n_points=5, ) - self.ifos[1].calibration_model = calibration.Precomputed.constant_uncertainty_spline( + self.ifos[ + 1 + ].calibration_model = calibration.Precomputed.constant_uncertainty_spline( amplitude_sigma=0.1, phase_sigma=0.1, frequency_array=self.ifos[1].frequency_array[self.ifos[1].frequency_mask], @@ -502,14 +527,16 @@ def setUp(self): ) self.priors = bilby.gw.prior.BBHPriorDict() self.priors["geocent_time"] = bilby.core.prior.Uniform(0, 4) - self.priors.update(bilby.gw.prior.CalibrationPriorDict.constant_uncertainty_spline( - amplitude_sigma=0.1, - phase_sigma=0.1, - minimum_frequency=20, - maximum_frequency=512, - n_nodes=5, - label="H1", - )) + self.priors.update( + bilby.gw.prior.CalibrationPriorDict.constant_uncertainty_spline( + amplitude_sigma=0.1, + phase_sigma=0.1, + minimum_frequency=20, + maximum_frequency=512, + n_nodes=5, + label="H1", + ) + ) self.wfg = bilby.gw.waveform_generator.WaveformGenerator( frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole, parameter_conversion=bilby.gw.conversion.convert_to_lal_binary_black_hole_parameters, diff --git a/test/gw/likelihood/multiband_test.py b/test/gw/likelihood/multiband_test.py index 33954d8bc..46a208d3f 100644 --- a/test/gw/likelihood/multiband_test.py +++ b/test/gw/likelihood/multiband_test.py @@ -34,7 +34,7 @@ def setUp(self): phase=1.3, geocent_time=self.xp.asarray(1187008882), ra=self.xp.asarray(1.3), - dec=self.xp.asarray(-1.2) + dec=self.xp.asarray(-1.2), ) # Network SNR is ~50 # torch needs sky parameters to be tensors @@ -42,8 +42,9 @@ def setUp(self): bilby.core.utils.random.seed(70817) rng = bilby.core.utils.random.rng self.ifos.set_strain_data_from_power_spectral_densities( - sampling_frequency=self.sampling_frequency, duration=self.duration, - start_time=self.test_parameters['geocent_time'] - self.duration + 2. + sampling_frequency=self.sampling_frequency, + duration=self.duration, + start_time=self.test_parameters["geocent_time"] - self.duration + 2.0, ) for ifo in self.ifos: ifo.minimum_frequency = self.fmin @@ -56,16 +57,18 @@ def setUp(self): prefix=f"recalib_{ifo.name}_", minimum_frequency=ifo.minimum_frequency, maximum_frequency=ifo.maximum_frequency, - n_points=spline_calibration_nodes + n_points=spline_calibration_nodes, ) for i in range(spline_calibration_nodes): self.test_parameters[f"recalib_{ifo.name}_amplitude_{i}"] = 0 self.test_parameters[f"recalib_{ifo.name}_phase_{i}"] = 0 # Calibration errors of 5% in amplitude and 5 degrees in phase - self.calibration_parameters[f"recalib_{ifo.name}_amplitude_{i}"] = \ + self.calibration_parameters[f"recalib_{ifo.name}_amplitude_{i}"] = ( rng.normal(loc=0, scale=0.05) - self.calibration_parameters[f"recalib_{ifo.name}_phase_{i}"] = \ + ) + self.calibration_parameters[f"recalib_{ifo.name}_phase_{i}"] = ( rng.normal(loc=0, scale=5 * np.pi / 180) + ) self.priors = bilby.gw.prior.BBHPriorDict() self.priors.pop("mass_1") @@ -73,55 +76,64 @@ def setUp(self): self.priors["chirp_mass"] = bilby.core.prior.Uniform(5.5, 6.5) self.priors["mass_ratio"] = bilby.core.prior.Uniform(0.125, 1) self.priors["geocent_time"] = bilby.core.prior.Uniform( - self.test_parameters['geocent_time'] - 0.1, - self.test_parameters['geocent_time'] + 0.1) + self.test_parameters["geocent_time"] - 0.1, + self.test_parameters["geocent_time"] + 0.1, + ) def tearDown(self): - del ( - self.ifos, - self.priors - ) + del (self.ifos, self.priors) - @parameterized.expand([ - ("IMRPhenomD", True, 2, False, 1.5e-2), - ("IMRPhenomD", True, 2, True, 1.5e-2), - ("IMRPhenomD", False, 2, False, 5e-3), - ("IMRPhenomD", False, 2, True, 6e-3), - ("IMRPhenomHM", False, 4, False, 8e-4), - ("IMRPhenomHM", False, 4, True, 1e-3) - ]) + @parameterized.expand( + [ + ("IMRPhenomD", True, 2, False, 1.5e-2), + ("IMRPhenomD", True, 2, True, 1.5e-2), + ("IMRPhenomD", False, 2, False, 5e-3), + ("IMRPhenomD", False, 2, True, 6e-3), + ("IMRPhenomHM", False, 4, False, 8e-4), + ("IMRPhenomHM", False, 4, True, 1e-3), + ] + ) def test_matches_original_likelihood( - self, waveform_approximant, linear_interpolation, highest_mode, add_cal_errors, tolerance + self, + waveform_approximant, + linear_interpolation, + highest_mode, + add_cal_errors, + tolerance, ): """ Check if multi-band likelihood values match original likelihood values """ wfg = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant=waveform_approximant - ) + ), ) wfg = BackendWaveformGenerator(wfg, self.xp) self.ifos.inject_signal(parameters=self.test_parameters, waveform_generator=wfg) wfg_mb = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.binary_black_hole_frequency_sequence, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant=waveform_approximant - ) + ), ) wfg_mb = BackendWaveformGenerator(wfg_mb, self.xp) likelihood = bilby.gw.likelihood.GravitationalWaveTransient( interferometers=self.ifos, waveform_generator=wfg ) likelihood_mb = bilby.gw.likelihood.MBGravitationalWaveTransient( - interferometers=self.ifos, waveform_generator=wfg_mb, - reference_chirp_mass=self.test_parameters['chirp_mass'], - priors=self.priors.copy(), linear_interpolation=linear_interpolation, - highest_mode=highest_mode + interferometers=self.ifos, + waveform_generator=wfg_mb, + reference_chirp_mass=self.test_parameters["chirp_mass"], + priors=self.priors.copy(), + linear_interpolation=linear_interpolation, + highest_mode=highest_mode, ) parameters = deepcopy(self.test_parameters) if add_cal_errors: @@ -137,35 +149,41 @@ def test_large_accuracy_factor(self): """ waveform_approximant = "IMRPhenomD" wfg = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant=waveform_approximant - ) + ), ) wfg = BackendWaveformGenerator(wfg, self.xp) self.ifos.inject_signal(parameters=self.test_parameters, waveform_generator=wfg) wfg_mb = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.binary_black_hole_frequency_sequence, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant=waveform_approximant - ) + ), ) wfg_mb = BackendWaveformGenerator(wfg_mb, self.xp) likelihood = bilby.gw.likelihood.GravitationalWaveTransient( interferometers=self.ifos, waveform_generator=wfg ) likelihood_mb = bilby.gw.likelihood.MBGravitationalWaveTransient( - interferometers=self.ifos, waveform_generator=wfg_mb, - reference_chirp_mass=self.test_parameters['chirp_mass'], - priors=self.priors.copy(), accuracy_factor=5 + interferometers=self.ifos, + waveform_generator=wfg_mb, + reference_chirp_mass=self.test_parameters["chirp_mass"], + priors=self.priors.copy(), + accuracy_factor=5, ) likelihood_mb_more_accurate = bilby.gw.likelihood.MBGravitationalWaveTransient( - interferometers=self.ifos, waveform_generator=wfg_mb, - reference_chirp_mass=self.test_parameters['chirp_mass'], - priors=self.priors.copy(), accuracy_factor=50 + interferometers=self.ifos, + waveform_generator=wfg_mb, + reference_chirp_mass=self.test_parameters["chirp_mass"], + priors=self.priors.copy(), + accuracy_factor=50, ) self.assertLess( abs( @@ -175,7 +193,8 @@ def test_large_accuracy_factor(self): abs( likelihood.log_likelihood_ratio(self.test_parameters) - likelihood_mb.log_likelihood_ratio(self.test_parameters) - ) / 2 + ) + / 2, ) def test_reference_chirp_mass_from_prior(self): @@ -183,33 +202,39 @@ def test_reference_chirp_mass_from_prior(self): Check if reference chirp mass is automatically determined from prior if no number has been passed """ wfg_mb = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.binary_black_hole_frequency_sequence, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant="IMRPhenomD" - ) + ), ) likelihood1 = bilby.gw.likelihood.MBGravitationalWaveTransient( - interferometers=self.ifos, waveform_generator=wfg_mb, + interferometers=self.ifos, + waveform_generator=wfg_mb, reference_chirp_mass=self.priors["chirp_mass"].minimum, - priors=self.priors.copy() + priors=self.priors.copy(), ) likelihood2 = bilby.gw.likelihood.MBGravitationalWaveTransient( - interferometers=self.ifos, waveform_generator=wfg_mb, - priors=self.priors.copy() + interferometers=self.ifos, + waveform_generator=wfg_mb, + priors=self.priors.copy(), + ) + self.assertAlmostEqual( + likelihood1.reference_chirp_mass, likelihood2.reference_chirp_mass ) - self.assertAlmostEqual(likelihood1.reference_chirp_mass, likelihood2.reference_chirp_mass) def test_no_reference_chirp_mass(self): """ Check if an error is raised if either reference_chirp_mass or priors is not specified. """ wfg_mb = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.binary_black_hole_frequency_sequence, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant="IMRPhenomD" - ) + ), ) with self.assertRaises(TypeError): bilby.gw.likelihood.MBGravitationalWaveTransient( @@ -221,11 +246,12 @@ def test_cannot_determine_reference_chirp_mass(self): Check if an error is raised if priors does not contain necessary information to determine reference chirp mass """ wfg_mb = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.binary_black_hole_frequency_sequence, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant="IMRPhenomD" - ) + ), ) for key in ["chirp_mass", "mass_1", "mass_2"]: if key in self.priors: @@ -235,7 +261,7 @@ def test_cannot_determine_reference_chirp_mass(self): interferometers=self.ifos, waveform_generator=wfg_mb, priors=self.priors ) - @parameterized.expand([(True, ), (False, )]) + @parameterized.expand([(True,), (False,)]) def test_inout_weights(self, linear_interpolation): """ Check if multiband weights can be saved as a file, and a likelihood object constructed from the weights file @@ -243,28 +269,29 @@ def test_inout_weights(self, linear_interpolation): """ waveform_approximant = "IMRPhenomD" wfg = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant=waveform_approximant - ) + ), ) wfg = BackendWaveformGenerator(wfg, self.xp) - self.ifos.inject_signal( - parameters=self.test_parameters, waveform_generator=wfg - ) + self.ifos.inject_signal(parameters=self.test_parameters, waveform_generator=wfg) wfg_mb = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.binary_black_hole_frequency_sequence, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant=waveform_approximant - ) + ), ) wfg_mb = BackendWaveformGenerator(wfg_mb, self.xp) likelihood_mb = bilby.gw.likelihood.MBGravitationalWaveTransient( - interferometers=self.ifos, waveform_generator=wfg_mb, - reference_chirp_mass=self.test_parameters['chirp_mass'], + interferometers=self.ifos, + waveform_generator=wfg_mb, + reference_chirp_mass=self.test_parameters["chirp_mass"], linear_interpolation=linear_interpolation, ) llr = likelihood_mb.log_likelihood_ratio(self.test_parameters) @@ -278,51 +305,60 @@ def test_inout_weights(self, linear_interpolation): # reset waveform generator to check if likelihood recovered from the weights file properly adds banded # frequency points to waveform arguments wfg_mb = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.binary_black_hole_frequency_sequence, waveform_arguments=dict( - reference_frequency=self.fmin, waveform_approximant=waveform_approximant - ) + reference_frequency=self.fmin, + waveform_approximant=waveform_approximant, + ), ) wfg_mb = BackendWaveformGenerator(wfg_mb, self.xp) - likelihood_mb_from_weights = bilby.gw.likelihood.MBGravitationalWaveTransient( - interferometers=self.ifos, waveform_generator=wfg_mb, weights=filepath + likelihood_mb_from_weights = ( + bilby.gw.likelihood.MBGravitationalWaveTransient( + interferometers=self.ifos, + waveform_generator=wfg_mb, + weights=filepath, + ) ) - llr_from_weights = likelihood_mb_from_weights.log_likelihood_ratio(self.test_parameters) + llr_from_weights = likelihood_mb_from_weights.log_likelihood_ratio( + self.test_parameters + ) self.assertEqual(aac.get_namespace(llr_from_weights), self.xp) self.assertAlmostEqual(llr, llr_from_weights) - @parameterized.expand([(True, ), (False, )]) + @parameterized.expand([(True,), (False,)]) def test_from_dict_weights(self, linear_interpolation): """ Check if a likelihood object constructed from dictionary-like weights produce the same likelihood value """ waveform_approximant = "IMRPhenomD" wfg = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant=waveform_approximant - ) + ), ) wfg = BackendWaveformGenerator(wfg, self.xp) - self.ifos.inject_signal( - parameters=self.test_parameters, waveform_generator=wfg - ) + self.ifos.inject_signal(parameters=self.test_parameters, waveform_generator=wfg) wfg_mb = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.binary_black_hole_frequency_sequence, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant=waveform_approximant - ) + ), ) wfg_mb = BackendWaveformGenerator(wfg_mb, self.xp) likelihood_mb = bilby.gw.likelihood.MBGravitationalWaveTransient( - interferometers=self.ifos, waveform_generator=wfg_mb, - reference_chirp_mass=self.test_parameters['chirp_mass'], + interferometers=self.ifos, + waveform_generator=wfg_mb, + reference_chirp_mass=self.test_parameters["chirp_mass"], linear_interpolation=linear_interpolation, ) llr = likelihood_mb.log_likelihood_ratio(self.test_parameters) @@ -331,29 +367,39 @@ def test_from_dict_weights(self, linear_interpolation): # reset waveform generator to check if likelihood recovered from the weights properly adds banded # frequency points to waveform arguments wfg_mb = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.binary_black_hole_frequency_sequence, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant=waveform_approximant - ) + ), ) wfg_mb = BackendWaveformGenerator(wfg_mb, self.xp) weights = likelihood_mb.weights likelihood_mb_from_weights = bilby.gw.likelihood.MBGravitationalWaveTransient( interferometers=self.ifos, waveform_generator=wfg_mb, weights=weights ) - llr_from_weights = likelihood_mb_from_weights.log_likelihood_ratio(self.test_parameters) + llr_from_weights = likelihood_mb_from_weights.log_likelihood_ratio( + self.test_parameters + ) self.assertEqual(aac.get_namespace(llr_from_weights), self.xp) self.assertAlmostEqual(llr, llr_from_weights) - @parameterized.expand([ - ("IMRPhenomD", True, 2, False, 1e-2), - ("IMRPhenomD", True, 2, True, 1e-2), - ("IMRPhenomHM", False, 4, False, 5e-3), - ]) + @parameterized.expand( + [ + ("IMRPhenomD", True, 2, False, 1e-2), + ("IMRPhenomD", True, 2, True, 1e-2), + ("IMRPhenomHM", False, 4, False, 5e-3), + ] + ) def test_matches_original_likelihood_low_maximum_frequency( - self, waveform_approximant, linear_interpolation, highest_mode, add_cal_errors, tolerance + self, + waveform_approximant, + linear_interpolation, + highest_mode, + add_cal_errors, + tolerance, ): """ Test for maximum frequency < sampling frequency / 2 @@ -362,38 +408,41 @@ def test_matches_original_likelihood_low_maximum_frequency( ifo.maximum_frequency = self.sampling_frequency / 8 wfg = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant=waveform_approximant - ) + ), ) wfg = BackendWaveformGenerator(wfg, self.xp) self.ifos.inject_signal(parameters=self.test_parameters, waveform_generator=wfg) wfg_mb = bilby.gw.WaveformGenerator( - duration=self.duration, sampling_frequency=self.sampling_frequency, + duration=self.duration, + sampling_frequency=self.sampling_frequency, frequency_domain_source_model=bilby.gw.source.binary_black_hole_frequency_sequence, waveform_arguments=dict( reference_frequency=self.fmin, waveform_approximant=waveform_approximant - ) + ), ) wfg_mb = BackendWaveformGenerator(wfg_mb, self.xp) likelihood = bilby.gw.likelihood.GravitationalWaveTransient( interferometers=self.ifos, waveform_generator=wfg ) likelihood_mb = bilby.gw.likelihood.MBGravitationalWaveTransient( - interferometers=self.ifos, waveform_generator=wfg_mb, - reference_chirp_mass=self.test_parameters['chirp_mass'], - priors=self.priors.copy(), linear_interpolation=linear_interpolation, - highest_mode=highest_mode + interferometers=self.ifos, + waveform_generator=wfg_mb, + reference_chirp_mass=self.test_parameters["chirp_mass"], + priors=self.priors.copy(), + linear_interpolation=linear_interpolation, + highest_mode=highest_mode, ) parameters = deepcopy(self.test_parameters) if add_cal_errors: parameters.update(self.calibration_parameters) llrmb = likelihood_mb.log_likelihood_ratio(parameters) self.assertLess( - abs(likelihood.log_likelihood_ratio(parameters) - llrmb), - tolerance + abs(likelihood.log_likelihood_ratio(parameters) - llrmb), tolerance ) self.assertEqual(aac.get_namespace(llrmb), self.xp) diff --git a/test/gw/likelihood/relative_binning_test.py b/test/gw/likelihood/relative_binning_test.py index 9d2e72f79..607377293 100644 --- a/test/gw/likelihood/relative_binning_test.py +++ b/test/gw/likelihood/relative_binning_test.py @@ -15,8 +15,10 @@ def setUp(self): sampling_frequency = 8192 chirp_mass = 13 mass_ratio = 0.5 - mass_1, mass_2 = bilby.gw.conversion.chirp_mass_and_mass_ratio_to_component_masses( - chirp_mass=chirp_mass, mass_ratio=mass_ratio + mass_1, mass_2 = ( + bilby.gw.conversion.chirp_mass_and_mass_ratio_to_component_masses( + chirp_mass=chirp_mass, mass_ratio=mass_ratio + ) ) self.test_parameters = dict( chirp_mass=13, @@ -37,14 +39,18 @@ def setUp(self): fiducial=0, ) self.fiducial_parameters = self.test_parameters.copy() - del self.fiducial_parameters["chirp_mass"], self.fiducial_parameters["mass_ratio"] + del ( + self.fiducial_parameters["chirp_mass"], + self.fiducial_parameters["mass_ratio"], + ) self.fiducial_parameters["mass_1"] = mass_1 self.fiducial_parameters["mass_2"] = mass_2 ifos = bilby.gw.detector.InterferometerList(["H1", "L1", "V1"]) ifos.set_strain_data_from_power_spectral_densities( - sampling_frequency=sampling_frequency, duration=duration, - start_time=self.test_parameters['geocent_time'] - duration + 2., + sampling_frequency=sampling_frequency, + duration=duration, + start_time=self.test_parameters["geocent_time"] - duration + 2.0, random_state=self.rng, ) for ifo in ifos: @@ -57,16 +63,18 @@ def setUp(self): prefix=f"recalib_{ifo.name}_", minimum_frequency=ifo.minimum_frequency, maximum_frequency=ifo.maximum_frequency, - n_points=spline_calibration_nodes + n_points=spline_calibration_nodes, ) for i in range(spline_calibration_nodes): self.test_parameters[f"recalib_{ifo.name}_amplitude_{i}"] = 0 self.test_parameters[f"recalib_{ifo.name}_phase_{i}"] = 0 # Calibration errors of 5% in amplitude and 5 degrees in phase - self.calibration_parameters[f"recalib_{ifo.name}_amplitude_{i}"] = \ + self.calibration_parameters[f"recalib_{ifo.name}_amplitude_{i}"] = ( self.rng.normal(loc=0, scale=0.05) - self.calibration_parameters[f"recalib_{ifo.name}_phase_{i}"] = \ + ) + self.calibration_parameters[f"recalib_{ifo.name}_phase_{i}"] = ( self.rng.normal(loc=0, scale=5 * np.pi / 180) + ) priors = bilby.gw.prior.BBHPriorDict() priors.pop("mass_1") @@ -79,16 +87,24 @@ def setUp(self): approximant = "IMRPhenomXP" non_bin_wfg = bilby.gw.WaveformGenerator( - duration=duration, sampling_frequency=sampling_frequency, + duration=duration, + sampling_frequency=sampling_frequency, frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole, waveform_arguments=dict( - reference_frequency=fmin, minimum_frequency=fmin, waveform_approximant=approximant) + reference_frequency=fmin, + minimum_frequency=fmin, + waveform_approximant=approximant, + ), ) bin_wfg = bilby.gw.waveform_generator.WaveformGenerator( - duration=duration, sampling_frequency=sampling_frequency, + duration=duration, + sampling_frequency=sampling_frequency, frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole_relative_binning, waveform_arguments=dict( - reference_frequency=fmin, waveform_approximant=approximant, minimum_frequency=fmin) + reference_frequency=fmin, + waveform_approximant=approximant, + minimum_frequency=fmin, + ), ) ifos.inject_signal( parameters=self.test_parameters, @@ -98,11 +114,13 @@ def setUp(self): self.ifos = ifos self.non_bin = bilby.gw.likelihood.GravitationalWaveTransient( - interferometers=ifos, waveform_generator=deepcopy(non_bin_wfg), - priors=priors.copy() + interferometers=ifos, + waveform_generator=deepcopy(non_bin_wfg), + priors=priors.copy(), ) self.binned = bilby.gw.likelihood.RelativeBinningGravitationalWaveTransient( - interferometers=ifos, waveform_generator=deepcopy(bin_wfg), + interferometers=ifos, + waveform_generator=deepcopy(bin_wfg), fiducial_parameters=self.fiducial_parameters, priors=priors.copy(), epsilon=0.05, @@ -126,10 +144,10 @@ def test_matches_non_binned_many(self): self.assertLess( abs(regular_ln_l - binned_ln_l) / abs(self.reference_ln_l - regular_ln_l), - 0.1 + 0.1, ) - @parameterized.expand([(False, ), (True, )]) + @parameterized.expand([(False,), (True,)]) @pytest.mark.flaky(reruns=3, only_rerun=["AssertionError"]) def test_matches_non_binned(self, add_cal_errors): parameters = deepcopy(self.test_parameters) @@ -145,12 +163,24 @@ def test_optimization_gives_good_match(self): fiducial_parameters["chirp_mass"] *= 0.99 priors = self.priors.copy() for key in [ - "ra", "dec", "geocent_time", "phase", "psi", "theta_jn", "luminosity_distance", - "a_1", "a_2", "tilt_1", "tilt_2", "phi_12", "phi_jl", + "ra", + "dec", + "geocent_time", + "phase", + "psi", + "theta_jn", + "luminosity_distance", + "a_1", + "a_2", + "tilt_1", + "tilt_2", + "phi_12", + "phi_jl", ]: priors[key] = self.test_parameters[key] binned = bilby.gw.likelihood.RelativeBinningGravitationalWaveTransient( - interferometers=self.ifos, waveform_generator=deepcopy(self.bin_wfg), + interferometers=self.ifos, + waveform_generator=deepcopy(self.bin_wfg), priors=priors, fiducial_parameters=fiducial_parameters, epsilon=0.05, @@ -166,7 +196,8 @@ def test_very_small_epsilon_returns_good_value(self): test that we avoid this. """ binned = bilby.gw.likelihood.RelativeBinningGravitationalWaveTransient( - interferometers=self.ifos, waveform_generator=deepcopy(self.bin_wfg), + interferometers=self.ifos, + waveform_generator=deepcopy(self.bin_wfg), fiducial_parameters=self.fiducial_parameters, priors=self.priors.copy(), epsilon=0.001, @@ -208,7 +239,7 @@ def test_likelihood_when_waveform_extends_beyond_maximum_frequency(self): ifos.set_strain_data_from_zero_noise( sampling_frequency=sampling_frequency, duration=duration, - start_time=test_parameters['geocent_time'] - duration + 2. + start_time=test_parameters["geocent_time"] - duration + 2.0, ) for ifo in ifos: ifo.minimum_frequency = fmin @@ -227,8 +258,8 @@ def test_likelihood_when_waveform_extends_beyond_maximum_frequency(self): waveform_arguments=dict( reference_frequency=fmin, minimum_frequency=fmin, - waveform_approximant=approximant - ) + waveform_approximant=approximant, + ), ) bin_wfg = bilby.gw.waveform_generator.WaveformGenerator( duration=duration, @@ -237,8 +268,8 @@ def test_likelihood_when_waveform_extends_beyond_maximum_frequency(self): waveform_arguments=dict( reference_frequency=fmin, waveform_approximant=approximant, - minimum_frequency=fmin - ) + minimum_frequency=fmin, + ), ) ifos.inject_signal( parameters=test_parameters, @@ -249,7 +280,7 @@ def test_likelihood_when_waveform_extends_beyond_maximum_frequency(self): non_bin = bilby.gw.likelihood.GravitationalWaveTransient( interferometers=ifos, waveform_generator=deepcopy(non_bin_wfg), - priors=priors.copy() + priors=priors.copy(), ) binned = bilby.gw.likelihood.RelativeBinningGravitationalWaveTransient( interferometers=ifos, diff --git a/test/gw/likelihood/roq_test.py b/test/gw/likelihood/roq_test.py index b470e5cef..dd1030759 100644 --- a/test/gw/likelihood/roq_test.py +++ b/test/gw/likelihood/roq_test.py @@ -18,7 +18,6 @@ class ROQBasisMixin: - @property def roq_dir(self): trial_roq_paths = [ @@ -197,7 +196,6 @@ def test_create_roq_weights_with_params(self): self.assertEqual(aac.get_namespace(roq_llr), self.xp) def test_create_roq_weights_frequency_mismatch_works_with_params(self): - self.ifos[0].maximum_frequency = self.ifos[0].maximum_frequency / 2 bilby.gw.likelihood.ROQGravitationalWaveTransient( interferometers=self.ifos, @@ -262,9 +260,9 @@ def test_create_roq_weights_fails_with_min_component_mass_outside_bounds(self): def test_create_roq_weights_fails_with_max_frequency(self): ifos = bilby.gw.detector.InterferometerList(["H1"]) ifos.set_strain_data_from_power_spectral_densities( - sampling_frequency=2 ** 14, duration=4 + sampling_frequency=2**14, duration=4 ) - ifos[0].maximum_frequency = 2 ** 13 + ifos[0].maximum_frequency = 2**13 with self.assertRaises(BilbyROQParamsRangeError): bilby.gw.likelihood.ROQGravitationalWaveTransient( interferometers=ifos, @@ -381,6 +379,7 @@ class TestROQLikelihoodHDF5(unittest.TestCase, ROQBasisMixin): respectively, and 2 quadratic bases constructed over 8Msun= self.priors["chirp_mass"].minimum) * - (mc_ranges_linear[:, 0] <= self.priors["chirp_mass"].maximum) + (mc_ranges_linear[:, 1] >= self.priors["chirp_mass"].minimum) + * (mc_ranges_linear[:, 0] <= self.priors["chirp_mass"].maximum) ) number_of_bases_quadratic = np.sum( - (mc_ranges_quadratic[:, 1] >= self.priors["chirp_mass"].minimum) * - (mc_ranges_quadratic[:, 0] <= self.priors["chirp_mass"].maximum) + (mc_ranges_quadratic[:, 1] >= self.priors["chirp_mass"].minimum) + * (mc_ranges_quadratic[:, 0] <= self.priors["chirp_mass"].maximum) ) self.assertEqual(likelihood.number_of_bases_linear, number_of_bases_linear) - self.assertEqual(likelihood.number_of_bases_quadratic, number_of_bases_quadratic) - self.assertEqual(len(likelihood.weights['frequency_nodes_linear']), number_of_bases_linear) - self.assertEqual(len(likelihood.weights['frequency_nodes_quadratic']), number_of_bases_quadratic) + self.assertEqual( + likelihood.number_of_bases_quadratic, number_of_bases_quadratic + ) + self.assertEqual( + len(likelihood.weights["frequency_nodes_linear"]), number_of_bases_linear + ) + self.assertEqual( + len(likelihood.weights["frequency_nodes_quadratic"]), + number_of_bases_quadratic, + ) for ifo in interferometers: - self.assertEqual(len(likelihood.weights['{}_linear'.format(ifo.name)]), number_of_bases_linear) - self.assertEqual(len(likelihood.weights['{}_quadratic'.format(ifo.name)]), number_of_bases_quadratic) + self.assertEqual( + len(likelihood.weights["{}_linear".format(ifo.name)]), + number_of_bases_linear, + ) + self.assertEqual( + len(likelihood.weights["{}_quadratic".format(ifo.name)]), + number_of_bases_quadratic, + ) @parameterized.expand( product( @@ -563,10 +589,12 @@ def test_number_of_loaded_bases(self, basis_linear, basis_quadratic, mc_range, r [_path_to_basis, _path_to_basis_mb], [(8, 9), (8, 14)], [1, 2], - [False, True] + [False, True], ) ) - def test_likelihood_accuracy(self, basis_linear, basis_quadratic, mc_range, roq_scale_factor, add_cal_errors): + def test_likelihood_accuracy( + self, basis_linear, basis_quadratic, mc_range, roq_scale_factor, add_cal_errors + ): "Compare with log likelihood ratios computed by the non-ROQ likelihood" # The maximum error of log likelihood ratio. It is set to be larger for roq_scale_factor=1 as the injected SNR # is higher. @@ -578,21 +606,47 @@ def test_likelihood_accuracy(self, basis_linear, basis_quadratic, mc_range, roq_ raise self.assertLess_likelihood_errors( - basis_linear, basis_quadratic, mc_range, roq_scale_factor, add_cal_errors, max_llr_error + basis_linear, + basis_quadratic, + mc_range, + roq_scale_factor, + add_cal_errors, + max_llr_error, ) - @parameterized.expand([(_path_to_basis_mb, 100, 1024), (_path_to_basis_mb, 20, 200), (_path_to_basis_mb, 100, 200)]) - def test_likelihood_accuracy_narrower_frequency_range(self, basis, minimum_frequency, maximum_frequency): + @parameterized.expand( + [ + (_path_to_basis_mb, 100, 1024), + (_path_to_basis_mb, 20, 200), + (_path_to_basis_mb, 100, 200), + ] + ) + def test_likelihood_accuracy_narrower_frequency_range( + self, basis, minimum_frequency, maximum_frequency + ): """Compare with log likelihood ratios computed by the non-ROQ likelihood in the case where analyzed frequency range is narrower than the basis frequency range""" self.assertLess_likelihood_errors( - basis, basis, (8, 9), 1, False, 1.5e-1, - minimum_frequency=minimum_frequency, maximum_frequency=maximum_frequency + basis, + basis, + (8, 9), + 1, + False, + 1.5e-1, + minimum_frequency=minimum_frequency, + maximum_frequency=maximum_frequency, ) def assertLess_likelihood_errors( - self, basis_linear, basis_quadratic, mc_range, roq_scale_factor, add_cal_errors, max_llr_error, - minimum_frequency=None, maximum_frequency=None + self, + basis_linear, + basis_quadratic, + mc_range, + roq_scale_factor, + add_cal_errors, + max_llr_error, + minimum_frequency=None, + maximum_frequency=None, ): self.minimum_frequency *= roq_scale_factor self.sampling_frequency *= roq_scale_factor @@ -617,7 +671,7 @@ def assertLess_likelihood_errors( interferometers.set_strain_data_from_zero_noise( sampling_frequency=self.sampling_frequency, duration=self.duration, - start_time=self.injection_parameters["geocent_time"] - self.duration + 1 + start_time=self.injection_parameters["geocent_time"] - self.duration + 1, ) if add_cal_errors: @@ -630,14 +684,16 @@ def assertLess_likelihood_errors( prefix=prefix, minimum_frequency=ifo.minimum_frequency, maximum_frequency=ifo.maximum_frequency, - n_points=spline_calibration_nodes + n_points=spline_calibration_nodes, ) for i in range(spline_calibration_nodes): # 5% in amplitude, 5deg in phase - self.injection_parameters[f"{prefix}amplitude_{i}"] = \ - rng.normal(loc=0, scale=0.05) - self.injection_parameters[f"{prefix}phase_{i}"] = \ - rng.normal(loc=0, scale=5 * np.pi / 180) + self.injection_parameters[f"{prefix}amplitude_{i}"] = rng.normal( + loc=0, scale=0.05 + ) + self.injection_parameters[f"{prefix}phase_{i}"] = rng.normal( + loc=0, scale=5 * np.pi / 180 + ) waveform_generator = bilby.gw.WaveformGenerator( duration=self.duration, @@ -645,16 +701,18 @@ def assertLess_likelihood_errors( frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole, waveform_arguments=dict( reference_frequency=self.reference_frequency, - waveform_approximant=self.waveform_approximant - ) + waveform_approximant=self.waveform_approximant, + ), ) waveform_generator = BackendWaveformGenerator(waveform_generator, self.xp) - interferometers.inject_signal(waveform_generator=waveform_generator, parameters=self.injection_parameters) + interferometers.inject_signal( + waveform_generator=waveform_generator, parameters=self.injection_parameters + ) likelihood = bilby.gw.GravitationalWaveTransient( interferometers=interferometers, waveform_generator=waveform_generator, - priors=self.priors + priors=self.priors, ) search_waveform_generator = bilby.gw.waveform_generator.WaveformGenerator( @@ -663,19 +721,23 @@ def assertLess_likelihood_errors( frequency_domain_source_model=bilby.gw.source.binary_black_hole_roq, waveform_arguments=dict( reference_frequency=self.reference_frequency, - waveform_approximant=self.waveform_approximant - ) + waveform_approximant=self.waveform_approximant, + ), + ) + search_waveform_generator = BackendWaveformGenerator( + search_waveform_generator, self.xp ) - search_waveform_generator = BackendWaveformGenerator(search_waveform_generator, self.xp) likelihood_roq = bilby.gw.likelihood.ROQGravitationalWaveTransient( interferometers=interferometers, priors=self.priors, waveform_generator=search_waveform_generator, linear_matrix=f"{self.roq_dir}/{basis_linear}", quadratic_matrix=f"{self.roq_dir}/{basis_quadratic}", - roq_scale_factor=roq_scale_factor + roq_scale_factor=roq_scale_factor, ) - for mc in np.linspace(self.priors["chirp_mass"].minimum, self.priors["chirp_mass"].maximum, 11): + for mc in np.linspace( + self.priors["chirp_mass"].minimum, self.priors["chirp_mass"].maximum, 11 + ): parameters = self.injection_parameters.copy() parameters["chirp_mass"] = mc llr = likelihood.log_likelihood_ratio(parameters) @@ -683,17 +745,21 @@ def assertLess_likelihood_errors( self.assertLess(np.abs(llr - llr_roq), max_llr_error) self.assertEqual(aac.get_namespace(llr_roq), self.xp) - @parameterized.expand([ - ("log_likelihood_ratio", ), - ("compute_per_detector_log_likelihood", ), - ("generate_posterior_sample_from_marginalized_likelihood", ) - ]) + @parameterized.expand( + [ + ("log_likelihood_ratio",), + ("compute_per_detector_log_likelihood",), + ("generate_posterior_sample_from_marginalized_likelihood",), + ] + ) def test_method_switches_basis(self, method): basis_linear = self._path_to_basis basis_quadratic = self._path_to_basis interferometers = bilby.gw.detector.InterferometerList(["H1", "L1"]) interferometers.set_strain_data_from_power_spectral_densities( - sampling_frequency=self.sampling_frequency, duration=self.duration, start_time=1.2 - self.duration + 1 + sampling_frequency=self.sampling_frequency, + duration=self.duration, + start_time=1.2 - self.duration + 1, ) for ifo in interferometers: ifo.minimum_frequency = 20 @@ -703,8 +769,8 @@ def test_method_switches_basis(self, method): frequency_domain_source_model=bilby.gw.source.binary_black_hole_roq, waveform_arguments=dict( reference_frequency=self.reference_frequency, - waveform_approximant=self.waveform_approximant - ) + waveform_approximant=self.waveform_approximant, + ), ) priors = deepcopy(self.priors) priors["chirp_mass"].minimum = 8 @@ -715,7 +781,7 @@ def test_method_switches_basis(self, method): priors=priors, waveform_generator=waveform_generator, linear_matrix=f"{self.roq_dir}/{basis_linear}", - quadratic_matrix=f"{self.roq_dir}/{basis_quadratic}" + quadratic_matrix=f"{self.roq_dir}/{basis_quadratic}", ) parameters = self.injection_parameters.copy() for mc in np.linspace(8, 14, 11): @@ -732,10 +798,15 @@ class TestCreateROQLikelihood(unittest.TestCase, ROQBasisMixin): respectively, and 2 quadratic bases constructed over 8Msun= self.minimum_frequency + self.full_frequencies_to_sequence = ( + self.frequency_array >= self.minimum_frequency + ) self.frequencies = self.frequency_array[self.full_frequencies_to_sequence] self.waveform_kwargs = dict( waveform_approximant="IMRPhenomHM", @@ -396,14 +404,16 @@ def test_valid_parameters(self): bilby.gw.source.binary_black_hole_frequency_sequence( self.frequency_array, frequencies=self.frequencies, **self.parameters ), - dict + dict, ) def test_waveform_error_catching(self): self.bad_parameters.update(self.waveform_kwargs) self.assertIsNone( bilby.gw.source.binary_black_hole_frequency_sequence( - self.frequency_array, frequencies=self.frequencies, **self.bad_parameters + self.frequency_array, + frequencies=self.frequencies, + **self.bad_parameters, ) ) @@ -413,7 +423,9 @@ def test_waveform_error_raising(self): raise_error_parameters["catch_waveform_errors"] = False with self.assertRaises(Exception): bilby.gw.source.binary_black_hole_frequency_sequence( - self.frequency_array, frequencies=self.frequencies, **raise_error_parameters + self.frequency_array, + frequencies=self.frequencies, + **raise_error_parameters, ) def test_match_LalBBH(self): @@ -422,18 +434,23 @@ def test_match_LalBBH(self): self.frequency_array, frequencies=self.frequencies, **self.parameters ) lalbbh = bilby.gw.source.lal_binary_black_hole( - self.frequency_array, minimum_frequency=self.minimum_frequency, **self.parameters + self.frequency_array, + minimum_frequency=self.minimum_frequency, + **self.parameters, ) self.assertEqual(freqseq.keys(), lalbbh.keys()) for mode in freqseq: - diff = np.sum(np.abs(freqseq[mode] - lalbbh[mode][self.full_frequencies_to_sequence])**2.) - norm = np.sum(np.abs(freqseq[mode])**2.) + diff = np.sum( + np.abs(freqseq[mode] - lalbbh[mode][self.full_frequencies_to_sequence]) + ** 2.0 + ) + norm = np.sum(np.abs(freqseq[mode]) ** 2.0) self.assertLess(diff / norm, 1e-15) def test_match_LalBBH_specify_modes(self): parameters = copy(self.parameters) parameters.update(self.waveform_kwargs) - parameters['mode_array'] = [[2, 2]] + parameters["mode_array"] = [[2, 2]] freqseq = bilby.gw.source.binary_black_hole_frequency_sequence( self.frequency_array, frequencies=self.frequencies, **parameters ) @@ -442,16 +459,19 @@ def test_match_LalBBH_specify_modes(self): ) self.assertEqual(freqseq.keys(), lalbbh.keys()) for mode in freqseq: - diff = np.sum(np.abs(freqseq[mode] - lalbbh[mode][self.full_frequencies_to_sequence])**2.) - norm = np.sum(np.abs(freqseq[mode])**2.) + diff = np.sum( + np.abs(freqseq[mode] - lalbbh[mode][self.full_frequencies_to_sequence]) + ** 2.0 + ) + norm = np.sum(np.abs(freqseq[mode]) ** 2.0) self.assertLess(diff / norm, 1e-15) def test_match_LalBBH_nonGR(self): parameters = copy(self.parameters) parameters.update(self.waveform_kwargs) wf_dict = lal.CreateDict() - lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi0(wf_dict, 1.) - parameters['lal_waveform_dictionary'] = wf_dict + lalsimulation.SimInspiralWaveformParamsInsertNonGRDChi0(wf_dict, 1.0) + parameters["lal_waveform_dictionary"] = wf_dict freqseq = bilby.gw.source.binary_black_hole_frequency_sequence( self.frequency_array, frequencies=self.frequencies, **parameters ) @@ -460,8 +480,11 @@ def test_match_LalBBH_nonGR(self): ) self.assertEqual(freqseq.keys(), lalbbh.keys()) for mode in freqseq: - diff = np.sum(np.abs(freqseq[mode] - lalbbh[mode][self.full_frequencies_to_sequence])**2.) - norm = np.sum(np.abs(freqseq[mode])**2.) + diff = np.sum( + np.abs(freqseq[mode] - lalbbh[mode][self.full_frequencies_to_sequence]) + ** 2.0 + ) + norm = np.sum(np.abs(freqseq[mode]) ** 2.0) self.assertLess(diff / norm, 1e-15) @@ -480,11 +503,13 @@ def setUp(self): theta_jn=1.7, phase=0.0, lambda_1=1000.0, - lambda_2=1000.0 + lambda_2=1000.0, ) self.minimum_frequency = 50.0 self.frequency_array = bilby.core.utils.create_frequency_series(2048, 16) - self.full_frequencies_to_sequence = self.frequency_array >= self.minimum_frequency + self.full_frequencies_to_sequence = ( + self.frequency_array >= self.minimum_frequency + ) self.frequencies = self.frequency_array[self.full_frequencies_to_sequence] self.waveform_kwargs = dict( waveform_approximant="IMRPhenomPv2_NRTidal", @@ -503,7 +528,7 @@ def test_with_valid_parameters(self): bilby.gw.source.binary_neutron_star_frequency_sequence( self.frequency_array, frequencies=self.frequencies, **self.parameters ), - dict + dict, ) def test_fails_without_tidal_parameters(self): @@ -521,12 +546,17 @@ def test_match_LalBNS(self): self.frequency_array, frequencies=self.frequencies, **self.parameters ) lalbns = bilby.gw.source.lal_binary_neutron_star( - self.frequency_array, minimum_frequency=self.minimum_frequency, **self.parameters + self.frequency_array, + minimum_frequency=self.minimum_frequency, + **self.parameters, ) self.assertEqual(freqseq.keys(), lalbns.keys()) for mode in freqseq: - diff = np.sum(np.abs(freqseq[mode] - lalbns[mode][self.full_frequencies_to_sequence])**2.) - norm = np.sum(np.abs(freqseq[mode])**2.) + diff = np.sum( + np.abs(freqseq[mode] - lalbns[mode][self.full_frequencies_to_sequence]) + ** 2.0 + ) + norm = np.sum(np.abs(freqseq[mode]) ** 2.0) self.assertLess(diff / norm, 1e-5) @@ -558,7 +588,7 @@ def setUp(self): minimum_frequency=20.0, catch_waveform_errors=True, fiducial=False, - frequency_bin_edges=np.arange(20, 1500, 50) + frequency_bin_edges=np.arange(20, 1500, 50), ) self.frequency_array = bilby.core.utils.create_frequency_series(2048, 4) self.bad_parameters = copy(self.parameters) diff --git a/test/gw/utils_test.py b/test/gw/utils_test.py index bb683533c..63612b25c 100644 --- a/test/gw/utils_test.py +++ b/test/gw/utils_test.py @@ -37,7 +37,7 @@ def test_asd_from_freq_series(self): self.assertEqual(aac.get_namespace(asd), self.xp) asd = np.asarray(asd) freq_data = np.asarray(freq_data) - self.assertTrue(np.all(asd == freq_data * 2 * df ** 0.5)) + self.assertTrue(np.all(asd == freq_data * 2 * df**0.5)) def test_psd_from_freq_series(self): freq_data = self.xp.asarray([1, 2, 3]) @@ -46,7 +46,7 @@ def test_psd_from_freq_series(self): self.assertEqual(aac.get_namespace(psd), self.xp) psd = np.asarray(psd) freq_data = np.asarray(freq_data) - self.assertTrue(np.all(psd == (freq_data * 2 * df ** 0.5) ** 2)) + self.assertTrue(np.all(psd == (freq_data * 2 * df**0.5) ** 2)) def test_inner_product(self): aa = self.xp.asarray([1, 2, 3]) @@ -112,6 +112,7 @@ def test_overlap(self): @pytest.mark.skip(reason="GWOSC unstable: avoiding this test") def test_get_event_time(self): from urllib3.exceptions import NewConnectionError + events = [ "GW150914", "GW170104", @@ -283,13 +284,22 @@ def test_get_vertex_position_geocentric_ellipsoid_conversion(self): Test that the conversion between geocentric and ellipsoid coordinates is consistent for each quadrant, above and below the ellipsoid. """ - for lat_sign, long_sign, height_sign in itertools.product([-1, 1], [-1, 1], [-1, 1]): - latitude_true = lat_sign * self.xp.deg2rad(self.xp.asarray(46 + 27. / 60 + 18.528 / 3600)) - longitude_true = long_sign * self.xp.deg2rad(self.xp.asarray(119 + 24. / 60 + 27.5657 / 3600)) + for lat_sign, long_sign, height_sign in itertools.product( + [-1, 1], [-1, 1], [-1, 1] + ): + latitude_true = lat_sign * self.xp.deg2rad( + self.xp.asarray(46 + 27.0 / 60 + 18.528 / 3600) + ) + longitude_true = long_sign * self.xp.deg2rad( + self.xp.asarray(119 + 24.0 / 60 + 27.5657 / 3600) + ) elevation_true = height_sign * 142.554 vertex = gwutils.get_vertex_position_geocentric( - self.xp.asarray([latitude_true, longitude_true, elevation_true])) - latitude, longitude, elevation = gwutils.get_vertex_position_ellipsoid(vertex) + self.xp.asarray([latitude_true, longitude_true, elevation_true]) + ) + latitude, longitude, elevation = gwutils.get_vertex_position_ellipsoid( + vertex + ) self.assertAlmostEqual(float(latitude), float(latitude_true), 5) self.assertAlmostEqual(float(longitude), float(longitude_true), 5) self.assertAlmostEqual(float(elevation), float(elevation_true), 5) @@ -318,7 +328,6 @@ def test_safe_cast_mode_to_int(self): @pytest.mark.array_backend @pytest.mark.usefixtures("xp_class") class TestSkyFrameConversion(unittest.TestCase): - def setUp(self) -> None: self.priors = bilby.core.prior.PriorDict() self.priors["ra"] = bilby.core.prior.Uniform(0, 2 * np.pi) @@ -365,6 +374,7 @@ def test_conversion_gives_correct_prior(self) -> None: @pytest.mark.array_backend def test_ln_i0_mathces_scipy(xp): from scipy.special import i0 + values = xp.linspace(-10, 10, 101) assert max(abs(gwutils.ln_i0(values) - xp.log(i0(values)))) < 1e-10 diff --git a/test/gw/waveform_generator_test.py b/test/gw/waveform_generator_test.py index b6fd4f8a9..dc1c87015 100644 --- a/test/gw/waveform_generator_test.py +++ b/test/gw/waveform_generator_test.py @@ -77,9 +77,15 @@ def test_repr(self): self.waveform_generator.duration, self.waveform_generator.sampling_frequency, self.waveform_generator.start_time, - bilby.core.utils.get_function_path(self.waveform_generator.frequency_domain_source_model), - bilby.core.utils.get_function_path(self.waveform_generator.time_domain_source_model), - bilby.core.utils.get_function_path(bilby.gw.conversion.convert_to_lal_binary_black_hole_parameters), + bilby.core.utils.get_function_path( + self.waveform_generator.frequency_domain_source_model + ), + bilby.core.utils.get_function_path( + self.waveform_generator.time_domain_source_model + ), + bilby.core.utils.get_function_path( + bilby.gw.conversion.convert_to_lal_binary_black_hole_parameters + ), self.waveform_generator.waveform_arguments, ) ) @@ -96,9 +102,15 @@ def test_repr_with_time_domain_source_model(self): self.waveform_generator.duration, self.waveform_generator.sampling_frequency, self.waveform_generator.start_time, - bilby.core.utils.get_function_path(self.waveform_generator.frequency_domain_source_model), - bilby.core.utils.get_function_path(self.waveform_generator.time_domain_source_model), - bilby.core.utils.get_function_path(bilby.gw.conversion.convert_to_lal_binary_black_hole_parameters), + bilby.core.utils.get_function_path( + self.waveform_generator.frequency_domain_source_model + ), + bilby.core.utils.get_function_path( + self.waveform_generator.time_domain_source_model + ), + bilby.core.utils.get_function_path( + bilby.gw.conversion.convert_to_lal_binary_black_hole_parameters + ), self.waveform_generator.waveform_arguments, ) ) @@ -116,8 +128,12 @@ def conversion_func(): self.waveform_generator.duration, self.waveform_generator.sampling_frequency, self.waveform_generator.start_time, - bilby.core.utils.get_function_path(self.waveform_generator.frequency_domain_source_model), - bilby.core.utils.get_function_path(self.waveform_generator.time_domain_source_model), + bilby.core.utils.get_function_path( + self.waveform_generator.frequency_domain_source_model + ), + bilby.core.utils.get_function_path( + self.waveform_generator.time_domain_source_model + ), bilby.core.utils.get_function_path(conversion_func), self.waveform_generator.waveform_arguments, ) @@ -130,7 +146,9 @@ def test_duration(self): def test_sampling_frequency(self): self.assertEqual(self.waveform_generator.sampling_frequency, 4096) - self.assertEqual(aac.get_namespace(self.waveform_generator.sampling_frequency), self.xp) + self.assertEqual( + aac.get_namespace(self.waveform_generator.sampling_frequency), self.xp + ) def test_source_model(self): self.assertEqual( @@ -139,10 +157,14 @@ def test_source_model(self): ) def test_frequency_array_type(self): - self.assertEqual(aac.array_namespace(self.waveform_generator.frequency_array), self.xp) + self.assertEqual( + aac.array_namespace(self.waveform_generator.frequency_array), self.xp + ) def test_time_array_type(self): - self.assertEqual(aac.array_namespace(self.waveform_generator.time_array), self.xp) + self.assertEqual( + aac.array_namespace(self.waveform_generator.time_array), self.xp + ) def test_source_model_parameters(self): formatted_parameters = self.waveform_generator._format_parameters( @@ -420,42 +442,50 @@ def test_caching_without_parameters(self): def test_frequency_domain_caching_and_using_time_domain_strain_without_parameters( self, ): - self.assertFalse(_test_caching_different_domain( - self.waveform_generator.frequency_domain_strain, - self.waveform_generator.time_domain_strain, - self.simulation_parameters, - None, - )) + self.assertFalse( + _test_caching_different_domain( + self.waveform_generator.frequency_domain_strain, + self.waveform_generator.time_domain_strain, + self.simulation_parameters, + None, + ) + ) def test_frequency_domain_caching_and_using_time_domain_strain_with_parameters( self, ): - self.assertFalse(_test_caching_different_domain( - self.waveform_generator.frequency_domain_strain, - self.waveform_generator.time_domain_strain, - self.simulation_parameters, - self.simulation_parameters, - )) + self.assertFalse( + _test_caching_different_domain( + self.waveform_generator.frequency_domain_strain, + self.waveform_generator.time_domain_strain, + self.simulation_parameters, + self.simulation_parameters, + ) + ) def test_time_domain_caching_and_using_frequency_domain_strain_without_parameters( self, ): - self.assertFalse(_test_caching_different_domain( - self.waveform_generator.time_domain_strain, - self.waveform_generator.frequency_domain_strain, - self.simulation_parameters, - None, - )) + self.assertFalse( + _test_caching_different_domain( + self.waveform_generator.time_domain_strain, + self.waveform_generator.frequency_domain_strain, + self.simulation_parameters, + None, + ) + ) def test_time_domain_caching_and_using_frequency_domain_strain_with_parameters( self, ): - self.assertFalse(_test_caching_different_domain( - self.waveform_generator.time_domain_strain, - self.waveform_generator.frequency_domain_strain, - self.simulation_parameters, - self.simulation_parameters, - )) + self.assertFalse( + _test_caching_different_domain( + self.waveform_generator.time_domain_strain, + self.waveform_generator.frequency_domain_strain, + self.simulation_parameters, + self.simulation_parameters, + ) + ) def test_frequency_domain_caching_changing_model(self): original_waveform = self.waveform_generator.frequency_domain_strain( @@ -660,42 +690,50 @@ def test_caching_without_parameters(self): def test_frequency_domain_caching_and_using_time_domain_strain_without_parameters( self, ): - self.assertFalse(_test_caching_different_domain( - self.waveform_generator.frequency_domain_strain, - self.waveform_generator.time_domain_strain, - self.simulation_parameters, - None, - )) + self.assertFalse( + _test_caching_different_domain( + self.waveform_generator.frequency_domain_strain, + self.waveform_generator.time_domain_strain, + self.simulation_parameters, + None, + ) + ) def test_frequency_domain_caching_and_using_time_domain_strain_with_parameters( self, ): - self.assertFalse(_test_caching_different_domain( - self.waveform_generator.frequency_domain_strain, - self.waveform_generator.time_domain_strain, - self.simulation_parameters, - self.simulation_parameters, - )) + self.assertFalse( + _test_caching_different_domain( + self.waveform_generator.frequency_domain_strain, + self.waveform_generator.time_domain_strain, + self.simulation_parameters, + self.simulation_parameters, + ) + ) def test_time_domain_caching_and_using_frequency_domain_strain_without_parameters( self, ): - self.assertFalse(_test_caching_different_domain( - self.waveform_generator.time_domain_strain, - self.waveform_generator.frequency_domain_strain, - self.simulation_parameters, - None, - )) + self.assertFalse( + _test_caching_different_domain( + self.waveform_generator.time_domain_strain, + self.waveform_generator.frequency_domain_strain, + self.simulation_parameters, + None, + ) + ) def test_time_domain_caching_and_using_frequency_domain_strain_with_parameters( self, ): - self.assertFalse(_test_caching_different_domain( - self.waveform_generator.time_domain_strain, - self.waveform_generator.frequency_domain_strain, - self.simulation_parameters, - self.simulation_parameters, - )) + self.assertFalse( + _test_caching_different_domain( + self.waveform_generator.time_domain_strain, + self.waveform_generator.frequency_domain_strain, + self.simulation_parameters, + self.simulation_parameters, + ) + ) def _test_caching_different_domain(func1, func2, params1, params2): @@ -708,12 +746,13 @@ def _test_caching_different_domain(func1, func2, params1, params2): class TestGWSignalGenerator(unittest.TestCase): - def get_wfgen(self, **kwargs): default_kwargs = dict( duration=4, sampling_frequency=256, - waveform_arguments=dict(waveform_approximant=kwargs.pop("waveform_approximant", "IMRPhenomXP")), + waveform_arguments=dict( + waveform_approximant=kwargs.pop("waveform_approximant", "IMRPhenomXP") + ), ) default_kwargs.update(kwargs) return bilby.gw.waveform_generator.GWSignalWaveformGenerator(**default_kwargs) @@ -790,7 +829,9 @@ def test_warns_for_disabled_eccentric_parameters(self): eccentricity=0.1, mean_per_ano=0.2, ) - with mock.patch.object(bilby.gw.waveform_generator.logger, "warning") as warning: + with mock.patch.object( + bilby.gw.waveform_generator.logger, "warning" + ) as warning: wfg._format_parameters(parameters) warning.assert_called_once_with( "The following parameters are set but will be ignored since " @@ -808,7 +849,9 @@ def test_warns_for_disabled_tides(self): lambda_1=1000.0, lambda_2=1000.0, ) - with mock.patch.object(bilby.gw.waveform_generator.logger, "warning") as warning: + with mock.patch.object( + bilby.gw.waveform_generator.logger, "warning" + ) as warning: wfg._format_parameters(parameters) warning.assert_called_once_with( "The following parameters are set but will be ignored since " @@ -824,13 +867,17 @@ def test_does_not_warn_for_unrelated_parameters(self): luminosity_distance=1.0, phase=0.0, ) - with mock.patch.object(bilby.gw.waveform_generator.logger, "warning") as warning: + with mock.patch.object( + bilby.gw.waveform_generator.logger, "warning" + ) as warning: wfg._format_parameters(parameters) warning.assert_not_called() def test_eccentric_parameters_work(self): wfg = self.get_wfgen( - eccentric=True, spinning=False, waveform_approximant="EccentricFD", + eccentric=True, + spinning=False, + waveform_approximant="EccentricFD", ) parameters_1 = dict( mass_1=10, @@ -856,7 +903,9 @@ def test_eccentric_parameters_work(self): def test_tidal_parameters_work(self): wfg = self.get_wfgen( - tidal=True, spinning=False, waveform_approximant="IMRPhenomD_NRTidalv2", + tidal=True, + spinning=False, + waveform_approximant="IMRPhenomD_NRTidalv2", sampling_frequency=16384, ) parameters_1 = dict( @@ -974,7 +1023,9 @@ def test_raises_error_with_extra_kwargs(self): def test_caught_error_return_none(self): wfg = self.get_wfgen( spinning=False, - waveform_arguments=dict(waveform_approximant="IMRPhenomD", catch_waveform_errors=True), + waveform_arguments=dict( + waveform_approximant="IMRPhenomD", catch_waveform_errors=True + ), ) parameters = dict( mass_1=10.0, diff --git a/test/hyper/hyper_pe_test.py b/test/hyper/hyper_pe_test.py index 83b3c0a51..f5e348dd0 100644 --- a/test/hyper/hyper_pe_test.py +++ b/test/hyper/hyper_pe_test.py @@ -16,7 +16,6 @@ def __call__(self, a, b, c): class _ToyClassVariableNames: - variable_names = ["a", "b", "c"] def __call__(self, **kwargs): @@ -57,11 +56,13 @@ def test_evidence_factor_without_evidences(self): ) self.assertTrue(np.isnan(like.evidence_factor)) - @parameterized.expand([ - ("func", _toy_function), - ("class_no_names", _ToyClassNoVariableNames()), - ("class_with_names", _ToyClassVariableNames()), - ]) + @parameterized.expand( + [ + ("func", _toy_function), + ("class_no_names", _ToyClassNoVariableNames()), + ("class_with_names", _ToyClassVariableNames()), + ] + ) def test_get_function_parameters(self, _, model): expected = dict(a=1, b=2, c=3) model = hyp.model.Model([model]) diff --git a/test/import_test.py b/test/import_test.py index 5dd56cc56..155c82874 100644 --- a/test/import_test.py +++ b/test/import_test.py @@ -10,8 +10,15 @@ unique_packages = set(sys.modules) unwanted = { - "lal", "lalsimulation", "matplotlib", - "h5py", "dill", "tqdm", "tables", "deepdish", "corner", + "lal", + "lalsimulation", + "matplotlib", + "h5py", + "dill", + "tqdm", + "tables", + "deepdish", + "corner", } for filename in ["mcmc_requirements.txt", "optional_requirements.txt"]: diff --git a/test/integration/example_test.py b/test/integration/example_test.py index 266001fc1..960aa1643 100644 --- a/test/integration/example_test.py +++ b/test/integration/example_test.py @@ -69,7 +69,7 @@ def tearDownClass(cls): @parameterized.parameterized.expand(core_args) def test_core_examples(self, name, fname): - """ Loop over examples to check they run """ + """Loop over examples to check they run""" bilby.core.utils.command_line_args.bilby_test_mode = False ignore = ["15d_gaussian"] if any([item in fname for item in ignore]): @@ -78,7 +78,7 @@ def test_core_examples(self, name, fname): @parameterized.parameterized.expand(gw_args) def test_gw_examples(self, name, fname): - """ Loop over examples to check they run """ + """Loop over examples to check they run""" bilby.core.utils.command_line_args.bilby_test_mode = True _execute_file(name, fname) diff --git a/test/integration/noise_realisation_test.py b/test/integration/noise_realisation_test.py index 6bada5f58..7f1da34a4 100644 --- a/test/integration/noise_realisation_test.py +++ b/test/integration/noise_realisation_test.py @@ -39,7 +39,7 @@ def test_noise_normalisation(self): # generate some toy-model signal for matched filtering SNR testing n_avg = 1000 snr = np.zeros(n_avg) - mu = np.exp(-((time_array - duration / 2.0) ** 2) / (2.0 * 0.1 ** 2)) * np.sin( + mu = np.exp(-((time_array - duration / 2.0) ** 2) / (2.0 * 0.1**2)) * np.sin( 2 * np.pi * 100 * time_array ) muf, frequency_array = bilby.core.utils.nfft(mu, sampling_frequency) diff --git a/test/integration/other_test.py b/test/integration/other_test.py index 7f49825f9..5ef0478ac 100644 --- a/test/integration/other_test.py +++ b/test/integration/other_test.py @@ -36,7 +36,7 @@ def tearDownClass(self): """ def test_make_standard_data(self): - " Load in the saved standard data and compare with new data " + "Load in the saved standard data and compare with new data" # Load in the saved standard data frequencies_saved, hf_real_saved, hf_imag_saved = np.loadtxt( diff --git a/test/integration/sample_from_the_prior_test.py b/test/integration/sample_from_the_prior_test.py index 7fce8b162..e475b7d4d 100644 --- a/test/integration/sample_from_the_prior_test.py +++ b/test/integration/sample_from_the_prior_test.py @@ -93,7 +93,7 @@ def test_fifteen_dimensional_cbc(self): nact=10, outdir=self.outdir, label=label, - save=False + save=False, ) pvalues = [ ks_2samp_wrapper( diff --git a/test/integration/sampler_run_test.py b/test/integration/sampler_run_test.py index fd515037a..d2f059807 100644 --- a/test/integration/sampler_run_test.py +++ b/test/integration/sampler_run_test.py @@ -31,10 +31,7 @@ emcee=dict(iterations=1000, nwalkers=10), ) -sampler_imports = dict( - bilby_mcmc="bilby", - dynamic_dynesty="dynesty" -) +sampler_imports = dict(bilby_mcmc="bilby", dynamic_dynesty="dynesty") # Only load the native samplers loaded_samplers = { @@ -99,7 +96,9 @@ def test_run_sampler_single(self, sampler): @parameterized.expand(_sampler_kwargs.keys()) def test_run_sampler_schwimmbad(self, sampler): pool = SerialPool() - initialize_global_variables(self.likelihood, self.priors, ["m", "c"], True, {"m": 0, "c": 0}) + initialize_global_variables( + self.likelihood, self.priors, ["m", "c"], True, {"m": 0, "c": 0} + ) self._run_sampler(sampler, pool_size=1, pool=pool) @parameterized.expand(_sampler_kwargs.keys()) @@ -129,7 +128,9 @@ def test_interrupt_sampler_single(self, sampler): @parameterized.expand(_sampler_kwargs.keys()) def test_interrupt_sampler_schwimmbad(self, sampler): pool = SerialPool() - initialize_global_variables(self.likelihood, self.priors, ["m", "c"], True, {"m": 0, "c": 0}) + initialize_global_variables( + self.likelihood, self.priors, ["m", "c"], True, {"m": 0, "c": 0} + ) self._run_with_signal_handling(sampler, pool_size=1, pool=pool) @parameterized.expand(_sampler_kwargs.keys()) @@ -157,7 +158,9 @@ def trigger_signal(): with self.assertRaises((SystemExit, KeyboardInterrupt)): try: while True: - self._run_sampler(sampler=sampler, pool_size=pool_size, exit_code=5, **kwargs) + self._run_sampler( + sampler=sampler, pool_size=pool_size, exit_code=5, **kwargs + ) except SystemExit as error: self.assertEqual(error.code, 5) raise diff --git a/test/integration/test_waveforms.py b/test/integration/test_waveforms.py index c564c1052..39af32103 100644 --- a/test/integration/test_waveforms.py +++ b/test/integration/test_waveforms.py @@ -67,7 +67,6 @@ def test_TaylorF2(self): self.run_for_approximant(waveform_approximant, source="bns") def run_for_approximant(self, waveform_approximant, source): - if source == "bbh": injection_parameters = self.BBH_precessing_injection_parameters frequency_domain_source_model = bilby.gw.source.lal_binary_black_hole @@ -75,7 +74,9 @@ def run_for_approximant(self, waveform_approximant, source): injection_parameters = self.BNS_precessing_injection_parameters frequency_domain_source_model = bilby.gw.source.lal_binary_neutron_star else: - raise ValueError("Source can only be 'bbh' or 'bns', but was '{}'".format(source)) + raise ValueError( + "Source can only be 'bbh' or 'bns', but was '{}'".format(source) + ) # create a waveform generator for bilby duration = 4.0 @@ -144,7 +145,7 @@ def run_for_approximant(self, waveform_approximant, source): (waveform_generator.frequency_array)[-1], lambda_1, lambda_2, - **waveform_arguments + **waveform_arguments, ) h_plus = get_lalsim_waveform["plus"] @@ -247,7 +248,7 @@ def lalsim_FD_waveform( fmax, lambda_1=None, lambda_2=None, - **kwarg + **kwarg, ): mass1 = m1 * lal.MSUN_SI mass2 = m2 * lal.MSUN_SI