Source code for gammapy.spectrum.simulation

# Licensed under a 3 - clause BSD style license - see LICENSE.rst
from __future__ import absolute_import, division, print_function, unicode_literals
from collections import OrderedDict
import logging
import astropy.units as u
from ..utils.random import get_random_state
from ..utils.energy import EnergyBounds
from .utils import CountsPredictor
from .core import PHACountsSpectrum
from .observation import SpectrumObservation, SpectrumObservationList

__all__ = [
    'SpectrumSimulation'
]

log = logging.getLogger(__name__)


[docs]class SpectrumSimulation(object): """Simulate `~gammapy.spectrum.SpectrumObservation`. For a usage example see :gp-extra-notebook:`spectrum_simulation` Parameters ---------- livetime : `~astropy.units.Quantity` Livetime source_model : `~gammapy.spectrum.models.SpectralModel` Source model aeff : `~gammapy.irf.EffectiveAreaTable` Effective Area edisp : `~gammapy.irf.EnergyDispersion`, optional Energy Dispersion background_model : `~gammapy.spectrum.models.SpectralModel`, optional Background model alpha : float, optional Exposure ratio between source and background """ def __init__(self, livetime, source_model, aeff, edisp=None, background_model=None, alpha=None): self.livetime = livetime self.source_model = source_model self.aeff = aeff self.edisp = edisp self.background_model = background_model self.alpha = alpha self.on_vector = None self.off_vector = None self.obs = None self.result = SpectrumObservationList() @property def npred_source(self): """Predicted source `~gammapy.spectrum.CountsSpectrum`. Calls :func:`gammapy.spectrum.utils.CountsPredictor`. """ predictor = CountsPredictor(livetime=self.livetime, aeff=self.aeff, edisp=self.edisp, model=self.source_model) predictor.run() return predictor.npred @property def npred_background(self): """Predicted background (`~gammapy.spectrum.CountsSpectrum`). Calls :func:`gammapy.spectrum.utils.CountsPredictor`. """ predictor = CountsPredictor(livetime=self.livetime, aeff=self.aeff, edisp=self.edisp, model=self.background_model) predictor.run() return predictor.npred @property def e_reco(self): """Reconstructed energy binning.""" if self.edisp is not None: temp = self.edisp.e_reco.bins else: temp = self.aeff.energy.bins return EnergyBounds(temp)
[docs] def run(self, seed): """Simulate `~gammapy.spectrum.SpectrumObservationList`. The seeds will be set as observation ID. Previously produced results will be overwritten. Parameters ---------- seed : array of ints Random number generator seeds """ self.reset() n_obs = len(seed) log.info("Simulating {} observations".format(n_obs)) for counter, current_seed in enumerate(seed): progress = ((counter + 1) / n_obs) * 100 if progress % 10 == 0: log.info("Progress : {} %".format(progress)) self.simulate_obs(seed=current_seed, obs_id=current_seed) self.result.append(self.obs)
[docs] def reset(self): """Clear all results.""" self.result = SpectrumObservationList() self.obs = None self.on_vector = None self.off_vector = None
[docs] def simulate_obs(self, obs_id, seed='random-seed'): """Simulate one `~gammapy.spectrum.SpectrumObservation`. The result is stored as ``obs`` attribute Parameters ---------- obs_id : int Observation identifier seed : {int, 'random-seed', 'global-rng', `~numpy.random.RandomState`} see :func:~`gammapy.utils.random.get_random_state` """ random_state = get_random_state(seed) self.simulate_source_counts(random_state) if self.background_model is not None: self.simulate_background_counts(random_state) obs = SpectrumObservation(on_vector=self.on_vector, off_vector=self.off_vector, aeff=self.aeff, edisp=self.edisp) obs.obs_id = obs_id self.obs = obs
[docs] def simulate_source_counts(self, rand): """Simulate source `~gammapy.spectrum.PHACountsSpectrum`. Source counts are added to the on vector. Parameters ---------- rand: `~numpy.random.RandomState` random state """ on_counts = rand.poisson(self.npred_source.data.data.value) on_vector = PHACountsSpectrum(energy_lo=self.e_reco.lower_bounds, energy_hi=self.e_reco.upper_bounds, data=on_counts, backscal=1, meta=self._get_meta()) on_vector.livetime = self.livetime self.on_vector = on_vector
[docs] def simulate_background_counts(self, rand): """Simulate background `~gammapy.spectrum.PHACountsSpectrum`. Background counts are added to the on vector. Furthermore background counts divided by alpha are added to the off vector. TODO: At the moment the source IRFs are used. Make it possible to pass dedicated background IRFs. Parameters ---------- rand: `~numpy.random.RandomState` random state """ bkg_counts = rand.poisson(self.npred_background.data.data.value) off_counts = rand.poisson(self.npred_background.data.data.value / self.alpha) # Add background to on_vector self.on_vector.data.data += bkg_counts * u.ct # Create off vector off_vector = PHACountsSpectrum(energy_lo=self.e_reco.lower_bounds, energy_hi=self.e_reco.upper_bounds, data=off_counts, backscal=1. / self.alpha, is_bkg=True, meta=self._get_meta()) off_vector.livetime = self.livetime self.off_vector = off_vector
def _get_meta(self): """Meta info added to simulated counts spectra.""" meta = OrderedDict() meta['CREATOR'] = self.__class__.__name__ return meta