Implement NaimaModel wrapper class

gammapy/spectrum/models.py

@@ -25,6 +25,7 @@
     "TableModel",
     "AbsorbedSpectralModel",
     "Absorption",
+    "NaimaModel",
 ]
 
 
@@ -1464,3 +1465,128 @@ def evaluate(self, energy, **kwargs):
         flux = self.spectral_model.evaluate(energy=energy, **kwargs)
         absorption = self.absorption.evaluate(energy=energy, parameter=parameter)
         return flux * absorption
+
+
+class NaimaModel(SpectralModel):
+    r"""A wrapper for Naima models
+
+    This class provides an interface with the models defined in the `~naima.models` module.
+    The model accepts as a positional argument a `Naima <https://naima.readthedocs.io/en/latest/>`_
+    radiative model instance, used to compute the non-thermal emission from populations of
+    relativistic electrons or protons due to interactions with the ISM or with radiation and magnetic fields.
+
+    One of the advantages provided by this class consists in the possibility of performing a maximum
+    likelihood spectral fit of the model's parameters directly on observations, as opposed to the MCMC
+    `fit to flux points <https://naima.readthedocs.io/en/latest/mcmc.html>`_ featured in
+    Naima. All the parameters defining the parent population of charged particles are stored as
+    `~gammapy.utils.modeling.Parameter` and left free by default. In case that the radiative model is `
+    ~naima.radiative.Synchrotron`, the magnetic field strength may also be fitted. Parameters can be
+    freezed/unfreezed before the fit, and maximum/minimum values can be set to limit the parameters space to
+    the physically interesting region.
+
+    Parameters
+    ----------
+    radiative_model : `~naima.models.BaseRadiative`
+        An instance of a radiative model defined in `~naima.models`
+    distance : `~astropy.units.Quantity`, optional
+        Distance to the source. If set to 0, the intrinsic differential
+        luminosity will be returned. Default is 1 kpc
+    seed : str or list of str, optional
+        Seed photon field(s) to be considered for the `radiative_model` flux computation,
+        in case of a `~naima.models.InverseCompton` model. It can be a subset of the
+        `seed_photon_fields` list defining the `radiative_model`. Default is the whole list
+        of photon fields
+
+    Examples
+    --------
+    Create and plot a spectral model that convolves an `ExponentialCutoffPowerLaw` electron distribution
+    with an `InverseCompton` radiative model, in the presence of multiple seed photon fields.
+
+    .. plot::
+        :include-source:
+
+        import naima
+        from gammapy.spectrum.models import NaimaModel
+        import astropy.units as u
+        import matplotlib.pyplot as plt
+
+
+        particle_distribution = naima.models.ExponentialCutoffPowerLaw(1e30 / u.eV, 10 * u.TeV, 3.0, 30 * u.TeV)
+        radiative_model = naima.radiative.InverseCompton(
+            particle_distribution,
+            seed_photon_fields=[
+                "CMB",
+                ["FIR", 26.5 * u.K, 0.415 * u.eV / u.cm ** 3],
+            ],
+            Eemin=100 * u.GeV,
+        )
+
+        model = NaimaModel(radiative_model, distance=1.5 * u.kpc)
+
+        opts = {
+            "energy_range" : [10 * u.GeV, 80 * u.TeV],
+            "energy_power" : 2,
+            "flux_unit" : "erg-1 cm-2 s-1",
+        }
+
+        # Plot the total inverse Compton emission
+        model.plot(label='IC (total)', **opts)
+
+        # Plot the separate contributions from each seed photon field
+        for seed, ls in zip(['CMB','FIR'], ['-','--']):
+            model = NaimaModel(radiative_model, seed=seed, distance=1.5 * u.kpc)
+            model.plot(label="IC ({})".format(seed), ls=ls, color="gray", **opts)
+
+        plt.legend(loc='best')
+        plt.show()
+    """
+    #TODO: prevent users from setting new attributes after init
+
+    def __init__(self, radiative_model, distance=1.0 * u.kpc, seed=None):
+        import naima
+        self.radiative_model = radiative_model
+        self._particle_distribution = self.radiative_model.particle_distribution
+        self.distance = Parameter("distance", distance, frozen=True)
+        self.seed = seed
+
+        # This ensures the support of naima.models.TableModel
+        if isinstance(self._particle_distribution, naima.models.TableModel):
+            param_names = ["amplitude"]
+        else:
+            param_names = self._particle_distribution.param_names
+
+        parameters = []
+        for name in param_names:
+            value = getattr(self._particle_distribution, name)
+            setattr(self, name, Parameter(name, value))
+            parameters.append(getattr(self, name))
+
+        # In case of a synchrotron radiative model, append B to the fittable parameters
+        if 'B' in self.radiative_model.param_names:
+            B = getattr(self.radiative_model, "B")
+            setattr(self, "B", Parameter("B", B))
+            parameters.append(getattr(self, "B"))
+
+        super().__init__(parameters)
+
+
+    def evaluate(self, energy, **kwargs):
+        """Evaluate the model"""
+        for name, value in kwargs.items():
+            setattr(self._particle_distribution, name, value)
+
+        distance = self.distance.quantity
+
+        # Flattening the input energy list and later reshaping the flux list
+        # prevents some radiative models from displaying broadcasting problems.
+        if self.seed == None:
+            dnde = self.radiative_model.flux(energy.flatten(), distance=distance)
+        else:
+            dnde = self.radiative_model.flux(
+                energy.flatten(), seed=self.seed, distance=distance
+            )
+
+        dnde = dnde.reshape(energy.shape)
+
+        unit = 1 / (energy.unit * u.cm ** 2 * u.s)
+        return dnde.to(unit)

gammapy/spectrum/tests/test_models.py

@@ -1,6 +1,7 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 import pytest
 import numpy as np
+import naima
 import astropy.units as u
 from ...utils.energy import EnergyBounds
 from ...utils.testing import assert_quantity_allclose
@@ -15,6 +16,7 @@
     AbsorbedSpectralModel,
     Absorption,
     ConstantModel,
+    NaimaModel,
 )
 
 
@@ -192,8 +194,58 @@ def table_model():
 except ImportError:
     pass
 
+# Add Naima models
+particle_distributions = [
+    naima.models.PowerLaw(amplitude=2e33 / u.eV, e_0=10 * u.TeV, alpha=2.5),
+    naima.models.ExponentialCutoffBrokenPowerLaw(
+        amplitude=2e33 / u.eV,
+        e_0=10 * u.TeV,
+        alpha_1=2.5,
+        alpha_2=2.7,
+        e_break=900 * u.GeV,
+        e_cutoff=10 * u.TeV,
+    ),
+    naima.models.LogParabola(amplitude=2e33 / u.eV, e_0=10 * u.TeV, alpha=1.3, beta=0.5),
+]
+radiative_models = [
+    naima.radiative.PionDecay(particle_distributions[0], nh=1 * u.cm ** -3),
+    naima.radiative.InverseCompton(particle_distributions[1], seed_photon_fields=["CMB"]),
+    naima.radiative.Synchrotron(particle_distributions[2], B=2 * u.G),
+]
+
+TEST_MODELS.append(
+    dict(
+        name="naima1",
+        model=NaimaModel(radiative_models[0]),
+        val_at_2TeV=9.725347355450884e-14 * u.Unit("cm-2 s-1 TeV-1"),
+        integral_1_10TeV=3.530537143620737e-13 * u.Unit("cm-2 s-1"),
+        eflux_1_10TeV=7.643559573105779e-13 * u.Unit("TeV cm-2 s-1"),
+    )
+)
+
+TEST_MODELS.append(
+    dict(
+        name="naima2",
+        model=NaimaModel(radiative_models[1]),
+        val_at_2TeV=4.347836316893546e-12 * u.Unit("cm-2 s-1 TeV-1"),
+        integral_1_10TeV=1.5958109911918303e-11 * u.Unit("cm-2 s-1"),
+        eflux_1_10TeV=2.851281562480875e-11 * u.Unit("TeV cm-2 s-1"),
+    )
+)
+
+TEST_MODELS.append(
+    dict(
+        name="naima3",
+        model=NaimaModel(radiative_models[2]),
+        val_at_2TeV=1.0565840392550432e-24 * u.Unit("cm-2 s-1 TeV-1"),
+        integral_1_10TeV=4.4491861907713736e-13 * u.Unit("cm-2 s-1"),
+        eflux_1_10TeV=4.594120986691428e-13 * u.Unit("TeV cm-2 s-1"),
+    )
+)
+
 
 @requires_dependency("uncertainties")
+@requires_dependency("naima")
 @pytest.mark.parametrize("spectrum", TEST_MODELS, ids=[_["name"] for _ in TEST_MODELS])
 def test_models(spectrum):
     model = spectrum["model"]