sample_from_the_prior_test.py 3.03 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83
from __future__ import absolute_import
import shutil
import os
import logging

import unittest
import numpy as np
import bilby
from scipy.stats import ks_2samp, kstest


class Test(unittest.TestCase):
    outdir = 'outdir_for_tests'

    @classmethod
    def setUpClass(self):
        if os.path.isdir(self.outdir):
            try:
                shutil.rmtree(self.outdir)
            except OSError:
                logging.warning(
                    "{} not removed prior to tests".format(self.outdir))

    @classmethod
    def tearDownClass(self):
        if os.path.isdir(self.outdir):
            try:
                shutil.rmtree(self.outdir)
            except OSError:
                logging.warning(
                    "{} not removed prior to tests".format(self.outdir))

    def test_fifteen_dimensional_cbc(self):
        duration = 4.
        sampling_frequency = 2048.
        label = 'full_15_parameters'
        np.random.seed(88170235)

        waveform_arguments = dict(waveform_approximant='IMRPhenomPv2',
                                  reference_frequency=50., minimum_frequency=20.)
        waveform_generator = bilby.gw.WaveformGenerator(
            duration=duration, sampling_frequency=sampling_frequency,
            frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole,
            parameter_conversion=bilby.gw.conversion.convert_to_lal_binary_black_hole_parameters,
            waveform_arguments=waveform_arguments)

        ifos = bilby.gw.detector.InterferometerList(['H1', 'L1'])
        ifos.set_strain_data_from_power_spectral_densities(
            sampling_frequency=sampling_frequency, duration=duration,
            start_time=0)

        priors = bilby.gw.prior.BBHPriorDict()
        priors.pop('mass_1')
        priors.pop('mass_2')
        priors['chirp_mass'] = bilby.prior.Uniform(
            name='chirp_mass', latex_label='$M$', minimum=10.0, maximum=100.0,
            unit='$M_{\\odot}$')
        priors['mass_ratio'] = bilby.prior.Uniform(
            name='mass_ratio', latex_label='$q$', minimum=0.5, maximum=1.0)
        priors['geocent_time'] = bilby.core.prior.Uniform(
            minimum=-0.1, maximum=0.1)

        likelihood = bilby.gw.GravitationalWaveTransient(
            interferometers=ifos, waveform_generator=waveform_generator,
            priors=priors, distance_marginalization=False,
            phase_marginalization=False, time_marginalization=False)

        likelihood = bilby.core.likelihood.ZeroLikelihood(likelihood)

        result = bilby.run_sampler(
            likelihood=likelihood, priors=priors, sampler='dynesty',
            npoints=1000, walks=100, outdir=self.outdir, label=label)
        pvalues = [ks_2samp(result.priors[key].sample(10000),
                            result.posterior[key].values).pvalue
                   for key in priors.keys()]
        print("P values per parameter")
        for key, p in zip(priors.keys(), pvalues):
            print(key, p)
        self.assertGreater(kstest(pvalues, "uniform").pvalue, 0.01)


if __name__ == '__main__':
    unittest.main()