Python stingray.Powerspectrum类代码示例

 def rebin_several(self, df):
     """
     TODO: Not sure how to write tests for the rebin method!
     """
     ps = Powerspectrum(lc=self.lc, norm="Leahy")
     bin_ps = ps.rebin(df)
     assert np.isclose(bin_ps.freq[0], bin_ps.df, atol=1e-4, rtol=1e-4)

 def test_classical_significances_trial_correction(self):
     ps = Powerspectrum(lc=self.lc, norm="leahy")
     # change the powers so that just one exceeds the threshold
     ps.power = np.zeros_like(ps.power) + 2.0
     index = 1
     ps.power[index] = 10.0
     threshold = 0.01
     pval = ps.classical_significances(threshold=threshold,
                                       trial_correction=True)
     assert np.size(pval) == 0

    def test_fractional_rms_in_leahy_norm(self):
        """
        fractional rms should only be *approximately* equal the standard
        deviation divided by the mean of the light curve. Therefore, we allow
        for a larger tolerance in np.isclose()
        """
        ps = Powerspectrum(lc=self.lc, norm="Leahy")
        rms_ps, rms_err = ps.compute_rms(min_freq=ps.freq[0],
                                         max_freq=ps.freq[-1])

        rms_lc = np.std(self.lc.counts) / np.mean(self.lc.counts)
        assert np.isclose(rms_ps, rms_lc, atol=0.01)

    def test_rebin_makes_right_attributes(self):
        ps = Powerspectrum(lc=self.lc, norm="Leahy")
        # replace powers
        ps.ps = np.ones_like(ps.ps) * 2.0
        rebin_factor = 2.0
        bin_ps = ps.rebin(rebin_factor*ps.df)

        assert bin_ps.freq is not None
        assert bin_ps.ps is not None
        assert bin_ps.df == rebin_factor * 1.0 / self.lc.tseg
        assert bin_ps.norm.lower() == "leahy"
        assert bin_ps.m == 2
        assert bin_ps.n == self.lc.time.shape[0]
        assert bin_ps.nphots == np.sum(self.lc.counts)

    def test_pvals_is_numpy_array(self):
        ps = Powerspectrum(lc=self.lc, norm="leahy")
        # change the powers so that just one exceeds the threshold
        ps.power = np.zeros_like(ps.power) + 2.0

        index = 1
        ps.power[index] = 10.0

        threshold = 1.0

        pval = ps.classical_significances(threshold=threshold,
                                          trial_correction=True)

        assert isinstance(pval, np.ndarray)
        assert pval.shape[0] == 2

    def test_classical_significances_threshold(self):
        ps = Powerspectrum(lc=self.lc, norm="leahy")

        # change the powers so that just one exceeds the threshold
        ps.ps = np.zeros(ps.ps.shape[0])+2.0

        index = 1
        ps.ps[index] = 10.0

        threshold = 0.01

        pval = ps.classical_significances(threshold=threshold,
                                          trial_correction=False)
        assert pval[0, 0] < threshold
        assert pval[1, 0] == index

    def test_fractional_rms_in_frac_norm_is_consistent(self):
        time = np.arange(0, 100, 1) + 0.5

        poisson_counts = np.random.poisson(100.0,
                                           size=time.shape[0])

        lc = Lightcurve(time, counts=poisson_counts, dt=1,
                            gti=[[0, 100]])
        ps = Powerspectrum(lc=lc, norm="leahy")
        rms_ps_l, rms_err_l = ps.compute_rms(min_freq=ps.freq[1],
                                         max_freq=ps.freq[-1], white_noise_offset=0)

        ps = Powerspectrum(lc=lc, norm="frac")
        rms_ps, rms_err = ps.compute_rms(min_freq=ps.freq[1],
                                         max_freq=ps.freq[-1], white_noise_offset=0)
        assert np.allclose(rms_ps, rms_ps_l, atol=0.01)
        assert np.allclose(rms_err, rms_err_l, atol=0.01)

    def test_compute_highest_outlier_works(self):

        mp_ind = 5
        max_power = 1000.0

        ps = Powerspectrum()
        ps.freq = np.arange(10)
        ps.power = np.ones_like(ps.freq)
        ps.power[mp_ind] = max_power
        ps.m = 1
        ps.df = ps.freq[1]-ps.freq[0]
        ps.norm = "leahy"

        model = models.Const1D()
        p_amplitude = lambda amplitude: \
            scipy.stats.norm(loc=1.0, scale=1.0).pdf(
                amplitude)

        priors = {"amplitude": p_amplitude}

        lpost = PSDPosterior(ps.freq, ps.power, model, 1)
        lpost.logprior = set_logprior(lpost, priors)

        pe = PSDParEst(ps)

        res = pe.fit(lpost, [1.0])

        res.mfit = np.ones_like(ps.freq)

        max_y, max_x, max_ind = pe._compute_highest_outlier(lpost, res)

        assert np.isclose(max_y[0], 2*max_power)
        assert np.isclose(max_x[0], ps.freq[mp_ind])
        assert max_ind == mp_ind

    def setup_class(cls):
        m = 1
        nfreq = 100000
        freq = np.arange(nfreq)
        noise = np.random.exponential(size=nfreq)
        power = noise * 2.0

        ps = Powerspectrum()
        ps.freq = freq
        ps.power = power
        ps.m = m
        ps.df = freq[1] - freq[0]
        ps.norm = "leahy"

        cls.ps = ps
        cls.a_mean, cls.a_var = 2.0, 1.0

        cls.model = models.Const1D()

        p_amplitude = lambda amplitude: \
            scipy.stats.norm(loc=cls.a_mean, scale=cls.a_var).pdf(amplitude)

        cls.priors = {"amplitude": p_amplitude}
        cls.lpost = PSDPosterior(cls.ps.freq, cls.ps.power, cls.model,
                                 m=cls.ps.m)
        cls.lpost.logprior = set_logprior(cls.lpost, cls.priors)

    def test_calibrate_highest_outlier_works_with_mvn(self):
        m = 1
        nfreq = 10000
        seed = 100
        freq = np.linspace(1, 10, nfreq)
        rng = np.random.RandomState(seed)
        noise = rng.exponential(size=nfreq)
        model = models.Const1D()
        model.amplitude = 2.0
        p = model(freq)
        power = noise * p

        ps = Powerspectrum()
        ps.freq = freq
        ps.power = power
        ps.m = m
        ps.df = freq[1] - freq[0]
        ps.norm = "leahy"

        nsim = 10

        loglike = PSDLogLikelihood(ps.freq, ps.power, model, m=1)

        pe = PSDParEst(ps)

        pval = pe.calibrate_highest_outlier(loglike, [2.0], sample=None,
                                            max_post=False, seed=seed,
                                            nsim=nsim)

        assert pval > 0.001

    def test_simulate_highest_outlier_works(self):
        m = 1
        nfreq = 100000
        seed = 100
        freq = np.linspace(1, 10, nfreq)
        rng = np.random.RandomState(seed)
        noise = rng.exponential(size=nfreq)
        model = models.Const1D()
        model.amplitude = 2.0
        p = model(freq)
        power = noise * p

        ps = Powerspectrum()
        ps.freq = freq
        ps.power = power
        ps.m = m
        ps.df = freq[1] - freq[0]
        ps.norm = "leahy"

        nsim = 10

        loglike = PSDLogLikelihood(ps.freq, ps.power, model, m=1)

        s_all = np.atleast_2d(np.ones(nsim) * 2.0).T

        pe = PSDParEst(ps)

        res = pe.fit(loglike, [2.0], neg=True)

        maxpow_sim = pe.simulate_highest_outlier(s_all, loglike, [2.0],
                                                 max_post=False, seed=seed)

        assert maxpow_sim.shape[0] == nsim
        assert np.all(maxpow_sim > 20.00) and np.all(maxpow_sim < 31.0)

    def setup_class(cls):

        cls.m = 10
        nfreq = 1000000
        freq = np.arange(nfreq)
        noise = scipy.stats.chi2(2.*cls.m).rvs(size=nfreq)/np.float(cls.m)
        power = noise

        ps = Powerspectrum()
        ps.freq = freq
        ps.power = power
        ps.m = cls.m
        ps.df = freq[1]-freq[0]
        ps.norm = "leahy"


        cls.ps = ps
        cls.a_mean, cls.a_var = 2.0, 1.0

        cls.model = models.Const1D()

        p_amplitude = lambda amplitude: \
            scipy.stats.norm(loc=cls.a_mean, scale=cls.a_var).pdf(amplitude)

        cls.priors = {"amplitude":p_amplitude}

    def test_calibrate_lrt_works_with_mvn(self):

        m = 1
        nfreq = 10000
        freq = np.linspace(1, 10, nfreq)
        rng = np.random.RandomState(100)
        noise = rng.exponential(size=nfreq)
        model = models.Const1D()
        model.amplitude = 2.0
        p = model(freq)
        power = noise * p

        ps = Powerspectrum()
        ps.freq = freq
        ps.power = power
        ps.m = m
        ps.df = freq[1] - freq[0]
        ps.norm = "leahy"

        loglike = PSDLogLikelihood(ps.freq, ps.power, model, m=1)

        model2 = models.PowerLaw1D() + models.Const1D()
        model2.x_0_0.fixed = True
        loglike2 = PSDLogLikelihood(ps.freq, ps.power, model2, 1)

        pe = PSDParEst(ps)

        pval = pe.calibrate_lrt(loglike, [2.0], loglike2,
                                [2.0, 1.0, 2.0], sample=None,
                                max_post=False, nsim=10,
                                seed=100)

        assert pval > 0.001

    def test_generate_data_produces_correct_distribution(self):
        model = models.Const1D()

        model.amplitude = 2.0

        p = model(self.ps.freq)

        seed = 100
        rng = np.random.RandomState(seed)

        noise = rng.exponential(size=len(p))
        power = noise*p

        ps = Powerspectrum()
        ps.freq = self.ps.freq
        ps.power = power
        ps.m = 1
        ps.df = self.ps.freq[1]-self.ps.freq[0]
        ps.norm = "leahy"

        lpost = PSDLogLikelihood(ps.freq, ps.power, model, m=1)

        pe = PSDParEst(ps)

        rng2 = np.random.RandomState(seed)
        sim_data = pe._generate_data(lpost, [2.0], rng2)

        assert np.allclose(ps.power, sim_data.power)

    def test_calibrate_lrt_works_with_sampling(self):
        m = 1
        nfreq = 10000
        freq = np.linspace(1, 10, nfreq)
        rng = np.random.RandomState(100)
        noise = rng.exponential(size=nfreq)
        model = models.Const1D()
        model.amplitude = 2.0
        p = model(freq)
        power = noise * p

        ps = Powerspectrum()
        ps.freq = freq
        ps.power = power
        ps.m = m
        ps.df = freq[1] - freq[0]
        ps.norm = "leahy"

        lpost = PSDPosterior(ps.freq, ps.power, model, m=1)

        p_amplitude_1 = lambda amplitude: \
            scipy.stats.norm(loc=2.0, scale=1.0).pdf(amplitude)

        p_alpha_0 = lambda alpha: \
            scipy.stats.uniform(0.0, 5.0).pdf(alpha)

        p_amplitude_0 = lambda amplitude: \
            scipy.stats.norm(loc=self.a2_mean, scale=self.a2_var).pdf(
                amplitude)


        priors = {"amplitude": p_amplitude_1}

        priors2 = {"amplitude_1": p_amplitude_1,
                      "amplitude_0": p_amplitude_0,
                      "alpha_0": p_alpha_0}


        lpost.logprior = set_logprior(lpost, priors)

        model2 = models.PowerLaw1D() + models.Const1D()
        model2.x_0_0.fixed = True
        lpost2 = PSDPosterior(ps.freq, ps.power, model2, 1)
        lpost2.logprior = set_logprior(lpost2, priors2)

        pe = PSDParEst(ps)

        pval = pe.calibrate_lrt(lpost, [2.0], lpost2,
                                [2.0, 1.0, 2.0], sample=None,
                                max_post=True, nsim=10, nwalkers=100,
                                burnin=100, niter=20,
                                seed=100)

        assert pval > 0.001