bilby-dev / bilby

A unified framework for stochastic sampling packages and gravitational-wave inference in Python.
https://bilby-dev.github.io/bilby/
MIT License
68 stars 75 forks source link

[BUG] relative binning - correctly set last bin index #48

Closed noahewolfe closed 2 weeks ago

noahewolfe commented 1 month ago

If the last bin index is already chosen to be the last point in the frequency array, we shouldn't increment by 1, else we have an invalid index.

I observed this with very long injections that have power even at the maximum frequency, and thus, the bin selection will already include the last frequency point in the bin indices.

ColmTalbot commented 1 month ago

Hi @noahewolfe, can you provide a failing example, ideally in a way we can add to unit tests, or at least so that people can come back and reproduce this later?

adivijaykumar commented 1 month ago

Thanks @noahewolfe; I thought I had taken care of these issues in https://git.ligo.org/lscsoft/bilby/-/merge_requests/1310, but clearly not. As Colm said, an example and unit test would be very helpful.

noahewolfe commented 1 month ago

Can do! I wrote up something, copying in large part the existing relative binning unit test (this one) to construct a unit test, which should also provide an example.

I'll leave the exact code example at the very bottom, as it's a bit long right now. This isn't exactly what I was working with when I first reported this bug, but, it reproduces the same error, namely:

======================================================================
ERROR: test_matches_non_binned_many (__main__. TestRelativeBinningLikelihoodSubSolarMassBBH.test_matches_non_binned_many)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/noah.wolfe/ligo-ptas/scripts/long_relative_binning_unittest.py", line 90, in setUp
    self.binned = RelativeBinningGravitationalWaveTransient(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/noah.wolfe/.conda/envs/ligo-ptas/lib/python3.11/site-packages/bilby/gw/likelihood/relative.py", line 148, in __init__
    self.compute_summary_data()
  File "/home/noah.wolfe/.conda/envs/ligo-ptas/lib/python3.11/site-packages/bilby/gw/likelihood/relative.py", line 336, in compute_summary_data
    stop = masked_frequency_array[end_idx]
           ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^
IndexError: index 32449 is out of bounds for axis 0 with size 32449

When I include the fix in this pull request, the tests both pass.

Here's my go at a unit test; if this kind of thing looks good, I can update the pull request with this as a new file under test/gw/likelihood directory! (Admittedly, probably only the test_very_small_epsilon_returns_good_value test is nessecary for this, and test_matches_non_binned_many doesn't do a whole lot since I've chosen a narrow prior-- open to any additional thoughts y'all have on what this test should look like)

import unittest
from copy import deepcopy

import bilby
from bilby.gw.source import (
    lal_binary_black_hole,
    lal_binary_black_hole_relative_binning
)
from bilby.gw.likelihood import (
    GravitationalWaveTransient,
    RelativeBinningGravitationalWaveTransient
)
import numpy as np

class TestRelativeBinningLikelihoodSubSolarMassBBH(unittest.TestCase):
    def setUp(self):
        duration = 16
        fmin = 20
        sampling_frequency = 8192
        self.test_parameters = dict(
            chirp_mass=0.435,
            mass_ratio=1.0,
            a_1=0.3,
            a_2=0.4,
            tilt_1=1.0,
            tilt_2=0.2,
            phi_12=1.0,
            phi_jl=2.0,
            luminosity_distance=40,
            theta_jn=0.4,
            psi=0.659,
            phase=1.3,
            geocent_time=1187008882,
            ra=1.3,
            dec=-1.2,
        )
        self.fiducial_parameters = self.test_parameters.copy()

        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.
        )
        for ifo in ifos:
            ifo.minimum_frequency = fmin

        priors = bilby.gw.prior.BBHPriorDict()
        priors["chirp_mass"] = bilby.core.prior.Uniform(0.434, 0.436)
        priors["mass_ratio"] = bilby.core.prior.Uniform(0.125, 1)
        priors["geocent_time"] = self.test_parameters["geocent_time"]
        priors["fiducial"] = 0
        self.priors = priors

        approximant = "IMRPhenomXP"
        non_bin_wfg = bilby.gw.WaveformGenerator(
            duration=duration, sampling_frequency=sampling_frequency,
            frequency_domain_source_model=lal_binary_black_hole,
            waveform_arguments=dict(
                reference_frequency=fmin,
                minimum_frequency=fmin,
                waveform_approximant=approximant
            )
        )
        bin_wfg = bilby.gw.waveform_generator.WaveformGenerator(
            duration=duration,
            sampling_frequency=sampling_frequency,
            frequency_domain_source_model=lal_binary_black_hole_relative_binning,
            waveform_arguments=dict(
                reference_frequency=fmin,
                waveform_approximant=approximant,
                minimum_frequency=fmin
            )
        )
        ifos.inject_signal(
            parameters=self.test_parameters,
            waveform_generator=non_bin_wfg,
            raise_error=False,
        )
        self.ifos = ifos

        self.non_bin = GravitationalWaveTransient(
            interferometers=ifos,
            waveform_generator=deepcopy(non_bin_wfg),
            priors=priors.copy()
        )
        self.binned = RelativeBinningGravitationalWaveTransient(
            interferometers=ifos,
            waveform_generator=deepcopy(bin_wfg),
            fiducial_parameters=self.fiducial_parameters,
            priors=priors.copy(),
            epsilon=0.05,
        )
        self.non_bin.parameters.update(self.test_parameters)
        self.reference_ln_l = self.non_bin.log_likelihood_ratio()
        self.bin_wfg = bin_wfg
        self.priors = priors

    def tearDown(self):
        del (
            self.non_bin,
            self.binned,
        )

    def test_matches_non_binned_many(self):
        for _ in range(100):
            parameters = self.priors.sample()
            self.non_bin.parameters.update(parameters)
            self.binned.parameters.update(parameters)
            regular_ln_l = self.non_bin.log_likelihood_ratio()
            binned_ln_l = self.binned.log_likelihood_ratio()
            self.assertLess(
                abs(regular_ln_l - binned_ln_l)
                / abs(self.reference_ln_l - regular_ln_l),
                0.1
            )

    def test_very_small_epsilon_returns_good_value(self):
        """
        If the frequency bins cover less than one bin, the likeilhood is nan,
        test that we avoid this.
        """
        binned = RelativeBinningGravitationalWaveTransient(
            interferometers=self.ifos,
            waveform_generator=deepcopy(self.bin_wfg),
            fiducial_parameters=self.fiducial_parameters,
            priors=self.priors.copy(),
            epsilon=0.001,
        )
        binned.parameters.update(self.test_parameters)
        self.assertFalse(np.isnan(binned.log_likelihood_ratio()))

if __name__ == "__main__":
    unittest.main()
adivijaykumar commented 1 month ago

Thanks @noahewolfe! I think I misunderstood the problem earlier. I think this fix is great. For the unit test, I don't think we need an extra test class---I'd just simply replace the test parameters here with one for a subsolar mass system.

adivijaykumar commented 1 month ago

Also, I think the prior width you have in your example is okay, since at those masses an SSM candidate would have a chirp mass posterior that fits within that prior (and hence log-likelihood would vary quite a bit). In fact, my fear was the prior was too wide, but that doesn't seem to be the case at least with the tests you ran.

ColmTalbot commented 1 month ago

Thanks @noahewolfe! I think I misunderstood the problem earlier. I think this fix is great. For the unit test, I don't think we need an extra test class---I'd just simply replace the test parameters here with one for a subsolar mass system.

I would vote for adding another test case using the same code over replacing the existing one. The existing one clearly probes different behaviour!

adivijaykumar commented 1 month ago

I kinda thought if the parameters are changed to be consistent with subsolar mass, it would be a superset of things we were testing for currently. But maybe I am missing something.

ColmTalbot commented 1 month ago

We're presumably currently testing a case where

masked_bin_inds[-1] >= len(masked_frequency_array) - 1

I think this corresponds to the waveform not extending as far as the maximum frequency.

The change here is for the waveform extending beyond the maximum frequency.

I think both of these cases are worth including in the tests, maybe not the full set of tests, but at least verifying that the binning scheme can be set up and a likelihood can be evaluated.

ColmTalbot commented 4 weeks ago

@noahewolfe is there anything you need to get moving with the tests?

noahewolfe commented 2 weeks ago

Hi Colm/all, sorry for the delay! I'm working on this now. I'll incorporate the relevant test as an additional method of the TestRelativeBinningLikelihood class.