コード例 #1
0
ファイル: test_spectrum.py プロジェクト: AtreyeeS/gammapy
class TestSpectrumOnOff:
    """ Test ON OFF SpectrumDataset"""
    def setup(self):
        etrue = np.logspace(-1, 1, 10) * u.TeV
        self.e_true = MapAxis.from_energy_edges(etrue, name="energy_true")
        ereco = np.logspace(-1, 1, 5) * u.TeV
        elo = ereco[:-1]
        ehi = ereco[1:]
        self.e_reco = MapAxis.from_energy_edges(ereco, name="energy")

        start = u.Quantity([0], "s")
        stop = u.Quantity([1000], "s")
        time_ref = Time("2010-01-01 00:00:00.0")
        self.gti = GTI.create(start, stop, time_ref)
        self.livetime = self.gti.time_sum

        self.on_region = make_region("icrs;circle(0.,1.,0.1)")
        off_region = make_region("icrs;box(0.,1.,0.1, 0.2,30)")
        self.off_region = off_region.union(
            make_region("icrs;box(-1.,-1.,0.1, 0.2,150)"))
        self.wcs = WcsGeom.create(npix=300, binsz=0.01, frame="icrs").wcs

        self.aeff = RegionNDMap.create(region=self.on_region,
                                       wcs=self.wcs,
                                       axes=[self.e_true],
                                       unit="cm2")
        self.aeff.data += 1

        data = np.ones(elo.shape)
        data[-1] = 0  # to test stats calculation with empty bins

        axis = MapAxis.from_edges(ereco, name="energy", interp="log")
        self.on_counts = RegionNDMap.create(
            region=self.on_region,
            wcs=self.wcs,
            axes=[axis],
            meta={"EXPOSURE": self.livetime.to_value("s")},
        )
        self.on_counts.data += 1
        self.on_counts.data[-1] = 0

        self.off_counts = RegionNDMap.create(region=self.off_region,
                                             wcs=self.wcs,
                                             axes=[axis])
        self.off_counts.data += 10

        acceptance = RegionNDMap.from_geom(self.on_counts.geom)
        acceptance.data += 1

        data = np.ones(elo.shape)
        data[-1] = 0

        acceptance_off = RegionNDMap.from_geom(self.off_counts.geom)
        acceptance_off.data += 10

        self.edisp = EDispKernelMap.from_diagonal_response(
            self.e_reco, self.e_true, self.on_counts.geom.to_image())

        exposure = self.aeff * self.livetime
        exposure.meta["livetime"] = self.livetime

        mask_safe = RegionNDMap.from_geom(self.on_counts.geom, dtype=bool)
        mask_safe.data += True

        self.dataset = SpectrumDatasetOnOff(
            counts=self.on_counts,
            counts_off=self.off_counts,
            exposure=exposure,
            edisp=self.edisp,
            acceptance=acceptance,
            acceptance_off=acceptance_off,
            name="test",
            gti=self.gti,
            mask_safe=mask_safe,
        )

    def test_spectrum_dataset_on_off_create(self):
        e_reco = MapAxis.from_edges(u.Quantity([0.1, 1, 10.0], "TeV"),
                                    name="energy")
        e_true = MapAxis.from_edges(u.Quantity([0.05, 0.5, 5, 20.0], "TeV"),
                                    name="energy_true")
        geom = RegionGeom(region=None, axes=[e_reco])
        empty_dataset = SpectrumDatasetOnOff.create(geom=geom,
                                                    energy_axis_true=e_true)

        assert empty_dataset.counts.data.sum() == 0
        assert empty_dataset.data_shape[0] == 2
        assert empty_dataset.counts_off.data.sum() == 0
        assert empty_dataset.counts_off.geom.axes[0].nbin == 2
        assert_allclose(empty_dataset.acceptance_off, 0)
        assert_allclose(empty_dataset.acceptance, 0)
        assert empty_dataset.acceptance.data.shape[0] == 2
        assert empty_dataset.acceptance_off.data.shape[0] == 2
        assert empty_dataset.gti.time_sum.value == 0
        assert len(empty_dataset.gti.table) == 0
        assert empty_dataset.energy_range[0] is None

    def test_create_stack(self):
        geom = RegionGeom(region=None, axes=[self.e_reco])

        stacked = SpectrumDatasetOnOff.create(geom=geom,
                                              energy_axis_true=self.e_true)
        stacked.mask_safe.data += True

        stacked.stack(self.dataset)
        assert_allclose(stacked.energy_range.value,
                        self.dataset.energy_range.value)

    def test_alpha(self):
        assert self.dataset.alpha.data.shape == (4, 1, 1)
        assert_allclose(self.dataset.alpha.data, 0.1)

    def test_npred_no_edisp(self):
        const = 1 * u.Unit("cm-2 s-1 TeV-1")
        model = SkyModel(spectral_model=ConstantSpectralModel(const=const))
        livetime = 1 * u.s

        aeff = RegionNDMap.create(
            region=self.on_region,
            unit="cm2",
            axes=[self.e_reco.copy(name="energy_true")],
        )

        aeff.data += 1
        dataset = SpectrumDatasetOnOff(
            counts=self.on_counts,
            counts_off=self.off_counts,
            exposure=aeff * livetime,
            models=model,
        )
        energy = aeff.geom.axes[0].edges
        expected = aeff.data[0] * (energy[-1] - energy[0]) * const * livetime

        assert_allclose(dataset.npred_signal().data.sum(), expected.value)

    def test_to_spectrum_dataset(self):
        ds = self.dataset.to_spectrum_dataset()

        assert isinstance(ds, SpectrumDataset)
        assert_allclose(ds.background.data.sum(), 4)

    @requires_dependency("matplotlib")
    def test_peek(self):
        dataset = self.dataset.copy()
        dataset.models = SkyModel(spectral_model=PowerLawSpectralModel())

        with mpl_plot_check():
            dataset.peek()

    @requires_dependency("matplotlib")
    def test_plot_fit(self):
        dataset = self.dataset.copy()
        dataset.models = SkyModel(spectral_model=PowerLawSpectralModel())

        with mpl_plot_check():
            dataset.plot_fit()

    @requires_dependency("matplotlib")
    def test_plot_off_regions(self):
        from gammapy.visualization import plot_spectrum_datasets_off_regions

        with mpl_plot_check():
            plot_spectrum_datasets_off_regions([self.dataset])

    def test_to_from_ogip_files(self, tmp_path):
        dataset = self.dataset.copy(name="test")
        dataset.write(tmp_path / "test.fits")
        newdataset = SpectrumDatasetOnOff.read(tmp_path / "test.fits")

        expected_regions = compound_region_to_list(self.off_counts.geom.region)
        regions = compound_region_to_list(newdataset.counts_off.geom.region)

        assert newdataset.counts.meta["RESPFILE"] == "test_rmf.fits"
        assert newdataset.counts.meta["BACKFILE"] == "test_bkg.fits"
        assert newdataset.counts.meta["ANCRFILE"] == "test_arf.fits"

        assert_allclose(self.on_counts.data, newdataset.counts.data)
        assert_allclose(self.off_counts.data, newdataset.counts_off.data)
        assert_allclose(self.edisp.edisp_map.data,
                        newdataset.edisp.edisp_map.data)
        assert_time_allclose(newdataset.gti.time_start, dataset.gti.time_start)

        assert len(regions) == len(expected_regions)
        assert regions[0].center.is_equivalent_frame(
            expected_regions[0].center)
        assert_allclose(regions[1].angle, expected_regions[1].angle)

    def test_to_from_ogip_files_no_mask(self, tmp_path):
        dataset = self.dataset.copy(name="test")
        dataset.mask_safe = None
        dataset.write(tmp_path / "test.fits")
        newdataset = SpectrumDatasetOnOff.read(tmp_path / "test.fits")

        assert_allclose(newdataset.mask_safe.data, True)

    def test_to_from_ogip_files_zip(self, tmp_path):
        dataset = self.dataset.copy(name="test")
        dataset.write(tmp_path / "test.fits.gz")
        newdataset = SpectrumDatasetOnOff.read(tmp_path / "test.fits.gz")

        assert newdataset.counts.meta["RESPFILE"] == "test_rmf.fits.gz"
        assert newdataset.counts.meta["BACKFILE"] == "test_bkg.fits.gz"
        assert newdataset.counts.meta["ANCRFILE"] == "test_arf.fits.gz"

    def test_to_from_ogip_files_no_edisp(self, tmp_path):

        mask_safe = RegionNDMap.from_geom(self.on_counts.geom, dtype=bool)
        mask_safe.data += True

        acceptance = RegionNDMap.from_geom(self.on_counts.geom, data=1.0)

        exposure = self.aeff * self.livetime
        exposure.meta["livetime"] = self.livetime

        dataset = SpectrumDatasetOnOff(
            counts=self.on_counts,
            exposure=exposure,
            mask_safe=mask_safe,
            acceptance=acceptance,
            name="test",
        )
        dataset.write(tmp_path / "pha_obstest.fits")
        newdataset = SpectrumDatasetOnOff.read(tmp_path / "pha_obstest.fits")

        assert_allclose(self.on_counts.data, newdataset.counts.data)
        assert newdataset.counts_off is None
        assert newdataset.edisp is None
        assert newdataset.gti is None

    def test_energy_mask(self):
        mask = self.dataset.counts.geom.energy_mask(energy_min=0.3 * u.TeV,
                                                    energy_max=6 * u.TeV)
        desired = [False, True, True, False]
        assert_allclose(mask.data[:, 0, 0], desired)

        mask = self.dataset.counts.geom.energy_mask(energy_max=6 * u.TeV)
        desired = [True, True, True, False]
        assert_allclose(mask.data[:, 0, 0], desired)

        mask = self.dataset.counts.geom.energy_mask(energy_min=1 * u.TeV)
        desired = [False, False, True, True]
        assert_allclose(mask.data[:, 0, 0], desired)

    def test_str(self):
        model = SkyModel(spectral_model=PowerLawSpectralModel())
        dataset = SpectrumDatasetOnOff(
            counts=self.on_counts,
            counts_off=self.off_counts,
            models=model,
            exposure=self.aeff * self.livetime,
            edisp=self.edisp,
            acceptance=RegionNDMap.from_geom(geom=self.on_counts.geom, data=1),
            acceptance_off=RegionNDMap.from_geom(geom=self.off_counts.geom,
                                                 data=10),
        )
        assert "SpectrumDatasetOnOff" in str(dataset)
        assert "wstat" in str(dataset)

    def test_fake(self):
        """Test the fake dataset"""
        source_model = SkyModel(spectral_model=PowerLawSpectralModel())
        dataset = SpectrumDatasetOnOff(
            name="test",
            counts=self.on_counts,
            counts_off=self.off_counts,
            models=source_model,
            exposure=self.aeff * self.livetime,
            edisp=self.edisp,
            acceptance=RegionNDMap.from_geom(geom=self.on_counts.geom, data=1),
            acceptance_off=RegionNDMap.from_geom(geom=self.off_counts.geom,
                                                 data=10),
        )
        real_dataset = dataset.copy()

        background = RegionNDMap.from_geom(dataset.counts.geom)
        background.data += 1
        dataset.fake(npred_background=background, random_state=314)

        assert real_dataset.counts.data.shape == dataset.counts.data.shape
        assert real_dataset.counts_off.data.shape == dataset.counts_off.data.shape
        assert dataset.counts_off.data.sum() == 39
        assert dataset.counts.data.sum() == 5

    def test_info_dict(self):
        info_dict = self.dataset.info_dict()

        assert_allclose(info_dict["counts"], 3)
        assert_allclose(info_dict["counts_off"], 40)
        assert_allclose(info_dict["acceptance"], 4)
        assert_allclose(info_dict["acceptance_off"], 40)

        assert_allclose(info_dict["alpha"], 0.1)
        assert_allclose(info_dict["excess"], -1, rtol=1e-2)
        assert_allclose(info_dict["ontime"].value, 1e3)
        assert_allclose(info_dict["sqrt_ts"], -0.501005, rtol=1e-2)

        assert info_dict["name"] == "test"

    def test_resample_energy_axis(self):
        axis = MapAxis.from_edges([0.1, 1, 10] * u.TeV,
                                  name="energy",
                                  interp="log")
        grouped = self.dataset.resample_energy_axis(energy_axis=axis)

        assert grouped.counts.data.shape == (2, 1, 1)
        # exposure should be untouched
        assert_allclose(grouped.exposure.data, 1000)
        assert_allclose(np.squeeze(grouped.counts), [2, 1])
        assert_allclose(np.squeeze(grouped.counts_off), [20, 20])
        assert grouped.edisp.edisp_map.data.shape == (9, 2, 1, 1)
        assert_allclose(np.squeeze(grouped.acceptance), [2, 2])
        assert_allclose(np.squeeze(grouped.acceptance_off), [20, 20])

    def test_to_image(self):
        grouped = self.dataset.to_image()

        assert grouped.counts.data.shape == (1, 1, 1)
        # exposure should be untouched
        assert_allclose(grouped.exposure.data, 1000)
        assert_allclose(np.squeeze(grouped.counts), 3)
        assert_allclose(np.squeeze(grouped.counts_off), 40)
        assert grouped.edisp.edisp_map.data.shape == (9, 1, 1, 1)
        assert_allclose(np.squeeze(grouped.acceptance), 4)
        assert_allclose(np.squeeze(grouped.acceptance_off), 40)
コード例 #2
0
ファイル: test_spectrum.py プロジェクト: gaia-verna/gammapy
class TestSpectrumOnOff:
    """ Test ON OFF SpectrumDataset"""

    def setup(self):
        etrue = np.logspace(-1, 1, 10) * u.TeV
        self.e_true = etrue
        ereco = np.logspace(-1, 1, 5) * u.TeV
        elo = ereco[:-1]
        ehi = ereco[1:]
        self.e_reco = ereco
        self.aeff = EffectiveAreaTable(etrue[:-1], etrue[1:], np.ones(9) * u.cm ** 2)
        self.edisp = EDispKernel.from_diagonal_response(etrue, ereco)

        start = u.Quantity([0], "s")
        stop = u.Quantity([1000], "s")
        time_ref = Time("2010-01-01 00:00:00.0")
        self.gti = GTI.create(start, stop, time_ref)
        self.livetime = self.gti.time_sum

        self.on_region = make_region("icrs;circle(0.,1.,0.1)")
        off_region = make_region("icrs;box(0.,1.,0.1, 0.2,30)")
        self.off_region = off_region.union(
            make_region("icrs;box(-1.,-1.,0.1, 0.2,150)")
        )
        self.wcs = WcsGeom.create(npix=300, binsz=0.01, frame="icrs").wcs

        data = np.ones(elo.shape)
        data[-1] = 0  # to test stats calculation with empty bins

        axis = MapAxis.from_edges(ereco, name="energy", interp="log")
        self.on_counts = RegionNDMap.create(
            region=self.on_region, wcs=self.wcs, axes=[axis]
        )
        self.on_counts.data += 1
        self.on_counts.data[-1] = 0

        self.off_counts = RegionNDMap.create(
            region=self.off_region, wcs=self.wcs, axes=[axis]
        )
        self.off_counts.data += 10

        acceptance = RegionNDMap.from_geom(self.on_counts.geom)
        acceptance.data += 1

        data = np.ones(elo.shape)
        data[-1] = 0

        acceptance_off = RegionNDMap.from_geom(self.off_counts.geom)
        acceptance_off.data += 10

        self.dataset = SpectrumDatasetOnOff(
            counts=self.on_counts,
            counts_off=self.off_counts,
            aeff=self.aeff,
            edisp=self.edisp,
            livetime=self.livetime,
            acceptance=acceptance,
            acceptance_off=acceptance_off,
            name="test",
            gti=self.gti,
        )

    def test_spectrumdatasetonoff_create(self):
        e_reco = u.Quantity([0.1, 1, 10.0], "TeV")
        e_true = u.Quantity([0.05, 0.5, 5, 20.0], "TeV")
        empty_dataset = SpectrumDatasetOnOff.create(e_reco, e_true)

        assert empty_dataset.counts.data.sum() == 0
        assert empty_dataset.data_shape[0] == 2
        assert empty_dataset.counts_off.data.sum() == 0
        assert empty_dataset.counts_off.geom.axes[0].nbin == 2
        assert_allclose(empty_dataset.acceptance_off, 1)
        assert_allclose(empty_dataset.acceptance, 1)
        assert empty_dataset.acceptance.data.shape[0] == 2
        assert empty_dataset.acceptance_off.data.shape[0] == 2
        assert empty_dataset.livetime.value == 0
        assert len(empty_dataset.gti.table) == 0
        assert empty_dataset.energy_range[0] is None

    def test_create_stack(self):
        stacked = SpectrumDatasetOnOff.create(self.e_reco, self.e_true)
        stacked.stack(self.dataset)
        assert_allclose(stacked.energy_range.value, self.dataset.energy_range.value)

    def test_alpha(self):
        assert self.dataset.alpha.data.shape == (4, 1, 1)
        assert_allclose(self.dataset.alpha.data, 0.1)

    def test_npred_no_edisp(self):
        const = 1 * u.Unit("cm-2 s-1 TeV-1")
        model = SkyModel(spectral_model=ConstantSpectralModel(const=const))
        livetime = 1 * u.s

        e_reco = self.on_counts.geom.axes[0].edges
        aeff = EffectiveAreaTable(e_reco[:-1], e_reco[1:], np.ones(4) * u.cm ** 2)
        dataset = SpectrumDatasetOnOff(
            counts=self.on_counts,
            counts_off=self.off_counts,
            aeff=aeff,
            models=model,
            livetime=livetime,
        )

        energy = aeff.energy.edges
        expected = aeff.data.data[0] * (energy[-1] - energy[0]) * const * livetime

        assert_allclose(dataset.npred_sig().data.sum(), expected.value)

    @requires_dependency("matplotlib")
    def test_peek(self):
        dataset = self.dataset.copy()
        dataset.models = SkyModel(spectral_model=PowerLawSpectralModel())

        with mpl_plot_check():
            dataset.peek()

    @requires_dependency("matplotlib")
    def test_plot_fit(self):
        dataset = self.dataset.copy()
        dataset.models = SkyModel(spectral_model=PowerLawSpectralModel())

        with mpl_plot_check():
            dataset.plot_fit()

    def test_to_from_ogip_files(self, tmp_path):
        dataset = self.dataset.copy(name="test")
        dataset.to_ogip_files(outdir=tmp_path)
        newdataset = SpectrumDatasetOnOff.from_ogip_files(tmp_path / "pha_obstest.fits")

        expected_regions = compound_region_to_list(self.off_counts.geom.region)
        regions = compound_region_to_list(newdataset.counts_off.geom.region)

        assert_allclose(self.on_counts.data, newdataset.counts.data)
        assert_allclose(self.off_counts.data, newdataset.counts_off.data)
        assert_allclose(self.edisp.pdf_matrix, newdataset.edisp.pdf_matrix)
        assert_time_allclose(newdataset.gti.time_start, dataset.gti.time_start)

        assert len(regions) == len(expected_regions)
        assert regions[0].center.is_equivalent_frame(expected_regions[0].center)
        assert_allclose(regions[1].angle, expected_regions[1].angle)

    def test_to_from_ogip_files_no_edisp(self, tmp_path):

        mask_safe = RegionNDMap.from_geom(self.on_counts.geom, dtype=bool)
        mask_safe.data += True

        dataset = SpectrumDatasetOnOff(
            counts=self.on_counts,
            aeff=self.aeff,
            livetime=self.livetime,
            mask_safe=mask_safe,
            acceptance=1,
            name="test",
        )
        dataset.to_ogip_files(outdir=tmp_path)
        newdataset = SpectrumDatasetOnOff.from_ogip_files(tmp_path / "pha_obstest.fits")

        assert_allclose(self.on_counts.data, newdataset.counts.data)
        assert newdataset.counts_off is None
        assert newdataset.edisp is None
        assert newdataset.gti is None

    def test_energy_mask(self):
        mask = self.dataset.counts.geom.energy_mask(emin=0.3 * u.TeV, emax=6 * u.TeV)
        desired = [False, True, True, False]
        assert_allclose(mask[:, 0, 0], desired)

        mask = self.dataset.counts.geom.energy_mask(emax=6 * u.TeV)
        desired = [True, True, True, False]
        assert_allclose(mask[:, 0, 0], desired)

        mask = self.dataset.counts.geom.energy_mask(emin=1 * u.TeV)
        desired = [False, False, True, True]
        assert_allclose(mask[:, 0, 0], desired)

    def test_str(self):
        model = SkyModel(spectral_model=PowerLawSpectralModel())
        dataset = SpectrumDatasetOnOff(
            counts=self.on_counts,
            counts_off=self.off_counts,
            models=model,
            aeff=self.aeff,
            livetime=self.livetime,
            edisp=self.edisp,
            acceptance=1,
            acceptance_off=10,
        )
        assert "SpectrumDatasetOnOff" in str(dataset)
        assert "wstat" in str(dataset)

    def test_fake(self):
        """Test the fake dataset"""
        source_model = SkyModel(spectral_model=PowerLawSpectralModel())
        dataset = SpectrumDatasetOnOff(
            counts=self.on_counts,
            counts_off=self.off_counts,
            models=source_model,
            aeff=self.aeff,
            livetime=self.livetime,
            edisp=self.edisp,
            acceptance=1,
            acceptance_off=10,
        )
        real_dataset = dataset.copy()

        background = RegionNDMap.from_geom(dataset.counts.geom)
        background.data += 1
        dataset.fake(background_model=background, random_state=314)

        assert real_dataset.counts.data.shape == dataset.counts.data.shape
        assert real_dataset.counts_off.data.shape == dataset.counts_off.data.shape
        assert dataset.counts_off.data.sum() == 39
        assert dataset.counts.data.sum() == 5

    def test_info_dict(self):
        info_dict = self.dataset.info_dict()

        assert_allclose(info_dict["n_on"], 3)
        assert_allclose(info_dict["n_off"], 40)
        assert_allclose(info_dict["a_on"], 1)
        assert_allclose(info_dict["a_off"], 10)

        assert_allclose(info_dict["alpha"], 0.1)
        assert_allclose(info_dict["excess"], -1)
        assert_allclose(info_dict["livetime"].value, 1e3)

        assert info_dict["name"] == "test"