Add simulate_dataset() convenience function

gammapy/cube/simulate.py

@@ -0,0 +1,82 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+"""Simulate observations"""
+import astropy.units as u
+from ..cube import MapDataset, PSFKernel
+from ..cube import make_map_exposure_true_energy, make_map_background_irf
+from ..maps import WcsNDMap
+from ..cube.models import BackgroundModel
+from ..utils.random import get_random_state
+
+__all__ = ["simulate_dataset"]
+
+
+def simulate_dataset(
+    skymodel, geom, pointing, irfs, livetime=1 * u.h, offset=0 * u.deg, max_radius=0.8 * u.deg, random_state='random-seed'
+):
+
+    """Simulate a 3D dataset
+
+    Simulate a source defined with a sky model for a given pointing,
+    geometry and irfs for a given exposure time.
+    This will return a dataset object which includes the counts cube,
+    the exposure cube, the psf cube, the background model and the sky model.
+
+    Parameters
+    ----------
+    skymodel : `~gammapy.cube.models.SkyModel`
+        Background model map
+    geom : `~gammapy.maps.wcs.WcsGeom`
+        Geometry object for the observation
+    pointing : `~astropy.coordinates.SkyCoord`
+        Pointing position
+    irfs : dict
+        Irfs used for simulating the observation
+    livetime : `~astropy.units.quantity.Quantity`
+        Livetime exposure of the simulated observation
+    offset : `~astropy.units.quantity.Quantity`
+        Offset from the center of the pointing position.
+        This is used for the PSF and Edisp estimation
+    max_radius : `~astropy.coordinates.Angle`
+        The maximum radius of the PSF kernel.
+    random_state: {int, 'random-seed', 'global-rng', `~numpy.random.RandomState`}
+        Defines random number generator initialisation.
+
+    Returns
+    ---------
+    dataset : `~gammapy.cube.fit.MapDataset`
+        A dataset of the simulated observation.
+    """
+
+    background = make_map_background_irf(
+        pointing=pointing, ontime=livetime, bkg=irfs["bkg"], geom=geom
+    )
+
+    background_model = BackgroundModel(background)
+
+    psf = irfs["psf"].to_energy_dependent_table_psf(theta=offset)
+    psf_kernel = PSFKernel.from_table_psf(psf, geom, max_radius=max_radius)
+
+    exposure = make_map_exposure_true_energy(
+        pointing=pointing, livetime=livetime, aeff=irfs["aeff"], geom=geom
+    )
+
+    if "edisp" in irfs:
+        energy = geom.axes[0].edges * geom.axes[0].unit
+        edisp = irfs["edisp"].to_energy_dispersion(offset, e_reco=energy, e_true=energy)
+    else:
+        edisp = None
+
+    dataset = MapDataset(
+        model=skymodel,
+        exposure=exposure,
+        background_model=background_model,
+        psf=psf_kernel,
+        edisp=edisp,
+    )
+
+    npred_map = dataset.npred()
+    rng = get_random_state(random_state)
+    counts = rng.poisson(npred_map.data)
+    dataset.counts = WcsNDMap(geom, counts)
+
+    return dataset

gammapy/cube/tests/test_simulate.py

@@ -0,0 +1,54 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+import numpy as np
+from numpy.testing import assert_allclose
+from astropy.coordinates import SkyCoord
+import astropy.units as u
+from ...irf import load_cta_irfs
+from ...maps import WcsGeom, MapAxis
+from ...spectrum.models import PowerLaw
+from ...image.models import SkyGaussian
+from ...cube.models import SkyModel
+from ...cube import MapDataset
+from ..simulate import simulate_dataset
+
+
+def test_simulate():
+    filename = (
+        "$GAMMAPY_DATA/cta-1dc/caldb/data/cta/1dc/bcf/South_z20_50h/irf_file.fits"
+    )
+    irfs = load_cta_irfs(filename)
+
+    # Define sky model to simulate the data
+    spatial_model = SkyGaussian(lon_0="0 deg", lat_0="0 deg", sigma="0.2 deg")
+    spectral_model = PowerLaw(
+        index=2, amplitude="1e-11 cm-2 s-1 TeV-1", reference="1 TeV"
+    )
+    sky_model_simu = SkyModel(
+        spatial_model=spatial_model, spectral_model=spectral_model
+    )
+
+    # Define map geometry
+    axis = MapAxis.from_edges(
+        np.logspace(-1, 1.0, 20), unit="TeV", name="energy", interp="log"
+    )
+    geom = WcsGeom.create(
+        skydir=(0, 0), binsz=0.025, width=(1, 1), coordsys="GAL", axes=[axis]
+    )
+
+    # Define some observation parameters
+    pointing = SkyCoord(0 * u.deg, 0 * u.deg, frame="galactic")
+
+    dataset = simulate_dataset(
+        sky_model_simu, geom, pointing, irfs, livetime=10 * u.h, random_state=42
+    )
+
+    assert isinstance(dataset, MapDataset)
+    assert isinstance(dataset.model, SkyModel)
+
+    assert dataset.counts.data.dtype is np.dtype("int")
+    assert_allclose(dataset.counts.data[5, 20, 20],5)
+    assert_allclose(dataset.exposure.data[5, 20, 20],16122681639.856329)
+    assert_allclose(dataset.background_model.map.data[5, 20, 20],0.9765544250896915)
+    assert_allclose(dataset.psf.data[5, 32, 32],0.044402823)
+    assert_allclose(dataset.edisp.data.data[10,10],0.6623215756621856)
+