text: # Licensed under a 3 - clause BSD style license - see LICENSE.rst
  from collections import OrderedDict
  import logging
  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__)
  
  
  class SpectrumSimulation:
      """Simulate `~gammapy.spectrum.SpectrumObservation`.
  
      For a usage example see :gp-notebook:`spectrum_simulation`
  
      Parameters
      ----------
      livetime : `~astropy.units.Quantity`
          Livetime
      source_model : `~gammapy.spectrum.models.SpectralModel`
          Source model
      aeff : `~gammapy.irf.EffectiveAreaTable`, optional
          Effective Area
      edisp : `~gammapy.irf.EnergyDispersion`, optional
          Energy Dispersion
      e_true : `~astropy.units.Quantity`, optional
          Desired energy axis of the prediced counts vector if no IRFs are given
      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=None,
          edisp=None,
          e_true=None,
          background_model=None,
          alpha=None,
      ):
          self.livetime = livetime
          self.source_model = source_model
          self.aeff = aeff
          self.edisp = edisp
          self.e_true = e_true
          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,
              e_true=self.e_true,
              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,
              e_true=self.e_true,
              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:
              if self.aeff is not None:
                  temp = self.aeff.energy.bins
              else:
                  temp = self.e_true
          return EnergyBounds(temp)
  
      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)
  
      def reset(self):
          """Clear all results."""
          self.result = SpectrumObservationList()
          self.obs = None
          self.on_vector = None
          self.off_vector = None
  
      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
  
      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
  
      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
  
          # 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.0 / self.alpha,
              is_bkg=True,
              meta=self._get_meta(),
          )
          off_vector.livetime = self.livetime
          self.off_vector = off_vector
  
      def _get_meta(self):
          return OrderedDict([("CREATOR", self.__class__.__name__)])
  