Merge pull request #2278 from JouvinLea/refactor_crab_spectrum

Refactor class CrabSpectrum in a function

gammapy/spectrum/crab.py

@@ -3,7 +3,7 @@
 from astropy import units as u
 from .models import PowerLaw, LogParabola, ExponentialCutoffPowerLaw, SpectralModel
 
-__all__ = ["CrabSpectrum"]
+__all__ = ["create_crab_spectral_model"]
 
 # HESS publication: 2006A&A...457..899A
 hess_pl = {
@@ -63,8 +63,8 @@ def evaluate(energy):
         return flux / energy ** 2
 
 
-class CrabSpectrum:
-    """Crab nebula spectral model.
+def create_crab_spectral_model(reference="meyer"):
+    """Create the Crab nebula spectral model depending of the reference given.
 
     The Crab nebula is often used as a standard candle in gamma-ray astronomy.
     Fluxes and sensitivities are often quoted relative to the Crab spectrum.
@@ -86,18 +86,18 @@ class CrabSpectrum:
     Let's first import what we need::
 
         import astropy.units as u
-        from gammapy.spectrum import CrabSpectrum
+        from gammapy.spectrum import create_crab_spectral_model
         from gammapy.spectrum.models import PowerLaw
 
     Plot the 'hess_ecpl' reference Crab spectrum between 1 TeV and 100 TeV::
 
-        crab_hess_ecpl = CrabSpectrum('hess_ecpl')
-        crab_hess_ecpl.model.plot([1, 100] * u.TeV)
+        crab_hess_ecpl = create_crab_spectral_model('hess_ecpl')
+        crab_hess_ecpl.plot([1, 100] * u.TeV)
 
     Use a reference crab spectrum as unit to measure a differential flux (at 10 TeV)::
 
         >>> pwl = PowerLaw(index=2.3, amplitude=1e-12 * u.Unit('1 / (cm2 s TeV)'), reference=1 * u.TeV)
-        >>> crab = CrabSpectrum('hess_pl').model
+        >>> crab = create_crab_spectral_model('hess_pl')
         >>> energy = 10 * u.TeV
         >>> dnde_cu = (pwl(energy) / crab(energy)).to('%')
         >>> print(dnde_cu)
@@ -112,26 +112,28 @@ class CrabSpectrum:
         3.5350582166 %
     """
 
-    references = ["meyer", "hegra", "hess_pl", "hess_ecpl", "magic_lp", "magic_ecpl"]
-    """Available references (see class docstring)."""
-
-    def __init__(self, reference="meyer"):
-
-        if reference == "meyer":
-            model = MeyerCrabModel()
-        elif reference == "hegra":
-            model = PowerLaw(**hegra)
-        elif reference == "hess_pl":
-            model = PowerLaw(**hess_pl)
-        elif reference == "hess_ecpl":
-            model = ExponentialCutoffPowerLaw(**hess_ecpl)
-        elif reference == "magic_lp":
-            model = LogParabola(**magic_lp)
-        elif reference == "magic_ecpl":
-            model = ExponentialCutoffPowerLaw(**magic_ecpl)
-        else:
-            fmt = "Invalid reference: {!r}. Choices: {!r}"
-            raise ValueError(fmt.format(reference, self.references))
-
-        self.model = model
-        self.reference = reference
+    if reference == "meyer":
+        model = MeyerCrabModel()
+    elif reference == "hegra":
+        model = PowerLaw(**hegra)
+    elif reference == "hess_pl":
+        model = PowerLaw(**hess_pl)
+    elif reference == "hess_ecpl":
+        model = ExponentialCutoffPowerLaw(**hess_ecpl)
+    elif reference == "magic_lp":
+        model = LogParabola(**magic_lp)
+    elif reference == "magic_ecpl":
+        model = ExponentialCutoffPowerLaw(**magic_ecpl)
+    else:
+        fmt = "Invalid reference: {!r}. Choices: {!r}"
+        references = [
+            "meyer",
+            "hegra",
+            "hess_pl",
+            "hess_ecpl",
+            "magic_lp",
+            "magic_ecpl",
+        ]
+        raise ValueError(fmt.format(reference, references))
+
+    return model

gammapy/spectrum/tests/test_crab.py

@@ -2,7 +2,7 @@
 import pytest
 import astropy.units as u
 from ...utils.testing import assert_quantity_allclose
-from ...spectrum import CrabSpectrum
+from ...spectrum import create_crab_spectral_model
 
 CRAB_SPECTRA = [
     dict(
@@ -49,18 +49,18 @@ def test_crab_spectrum(spec):
     energy = 2 * u.TeV
     emin, emax = [1, 1e3] * u.TeV
 
-    crab_spectrum = CrabSpectrum(spec["name"])
+    crab_spectrum = create_crab_spectral_model(reference=spec["name"])
 
-    dnde = crab_spectrum.model(energy)
+    dnde = crab_spectrum(energy)
     assert_quantity_allclose(dnde, spec["dnde"])
 
-    flux = crab_spectrum.model.integral(emin, emax)
+    flux = crab_spectrum.integral(emin, emax)
     assert_quantity_allclose(flux, spec["flux"])
 
-    index = crab_spectrum.model.spectral_index(energy)
+    index = crab_spectrum.spectral_index(energy)
     assert_quantity_allclose(index, spec["index"], rtol=1e-5)
 
 
 def test_invalid_format():
     with pytest.raises(ValueError):
-        CrabSpectrum("spam")
+        create_crab_spectral_model("spam")

tutorials/hess.ipynb

@@ -42,7 +42,7 @@
     "from gammapy.cube import MapMaker, MapDataset, PSFKernel, MapMakerRing\n",
     "from gammapy.cube.models import SkyModel, BackgroundModel\n",
     "from gammapy.spectrum.models import PowerLaw\n",
-    "from gammapy.spectrum import CrabSpectrum\n",
+    "from gammapy.spectrum import create_crab_spectral_model\n",
     "from gammapy.image.models import SkyPointSource\n",
     "from gammapy.detect import compute_lima_on_off_image\n",
     "from gammapy.scripts import SpectrumAnalysisIACT\n",
@@ -334,7 +334,7 @@
    "outputs": [],
    "source": [
     "plt.figure(figsize=(10, 8))\n",
-    "crab_ref = CrabSpectrum(\"hess_pl\").model\n",
+    "crab_ref = create_crab_spectral_model(\"hess_pl\")\n",
     "\n",
     "dataset_fp = analysis.spectrum_result\n",
     "\n",

tutorials/light_curve.ipynb

@@ -218,9 +218,9 @@
    "outputs": [],
    "source": [
     "# Let's compare to the expected flux of this source\n",
-    "from gammapy.spectrum import CrabSpectrum\n",
+    "from gammapy.spectrum import create_crab_spectral_model\n",
     "\n",
-    "crab_spec = CrabSpectrum().model\n",
+    "crab_spec = create_crab_spectral_model()\n",
     "crab_flux = crab_spec.integral(*energy_range).to(\"cm-2 s-1\")\n",
     "crab_flux"
    ]

tutorials/spectrum_analysis.ipynb

@@ -101,7 +101,7 @@
     "    SpectrumExtraction,\n",
     "    SpectrumDatasetOnOff,\n",
     "    SpectrumDatasetOnOffStacker,\n",
-    "    CrabSpectrum,\n",
+    "    create_crab_spectral_model,\n",
     "    FluxPointsEstimator,\n",
     "    FluxPointsDataset,\n",
     ")"
@@ -619,7 +619,7 @@
     "model_best_joint.plot(**plot_kwargs, label=\"Joint analysis result\", ls=\"--\")\n",
     "model_best_joint.plot_error(**plot_kwargs)\n",
     "\n",
-    "CrabSpectrum().model.plot(**plot_kwargs, label=\"Crab reference\")\n",
+    "create_crab_spectral_model().plot(**plot_kwargs, label=\"Crab reference\")\n",
     "plt.legend()"
    ]
   },