Merge pull request #3752 from facero/plot_interactive

Human readable energy units string formatting for plot_interactive & plot_grid

gammapy/maps/core.py

@@ -7,6 +7,7 @@
 from astropy import units as u
 from astropy.io import fits
 from gammapy.utils.scripts import make_path
+from gammapy.utils.units import energy_unit_format
 from .axes import MapAxis
 from .coord import MapCoord
 from .geom import pix_tuple_to_idx
@@ -1057,10 +1058,15 @@ def plot_grid(self, figsize=None, ncols=3, **kwargs):
             image_wcs.plot(ax=ax, **kwargs)
 
             if axis.node_type == "center":
-                info = f"{axis.center[idx]:.1f}"
+                if axis.name == "energy" or axis.name == "energy_true":
+                    info =  energy_unit_format(axis.center[idx])
+                else:
+                    info = f"{axis.center[idx]:.1f}"
             else:
-                info = f"{axis.edges[idx]:.1f} - {axis.edges[idx + 1]:.1f} "
-
+                if axis.name == "energy" or axis.name == "energy_true":
+                    info = f"{energy_unit_format(axis.edges[idx])} - {energy_unit_format(axis.edges[idx+1])}"
+                else:
+                    info = f"{axis.edges[idx]:.1f} - {axis.edges[idx + 1]:.1f} "
             ax.set_title(f"{axis.name.capitalize()} " + info)
             lon, lat = ax.coords[0], ax.coords[1]
             lon.set_ticks_position("b")
@@ -1121,7 +1127,16 @@ def plot_interactive(self, rc_params=None, **kwargs):
         interact_kwargs = {}
 
         for axis in self.geom.axes:
-            options = axis.as_plot_labels
+            if axis.node_type == "center":
+                if axis.name == "energy" or axis.name == "energy_true":
+                    options =  energy_unit_format(axis.center)
+                else:
+                    options = axis.as_plot_labels
+            elif axis.name == "energy" or axis.name == "energy_true":
+                E = energy_unit_format(axis.edges)
+                options = [f"{E[i]} - {E[i+1]}" for i in range(len(E)-1)]
+            else:
+                options = axis.as_plot_labels
             interact_kwargs[axis.name] = SelectionSlider(
                 options=options,
                 description=f"Select {axis.name}:",

gammapy/utils/tests/test_units.py

@@ -1,8 +1,30 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
-from gammapy.utils.units import standardise_unit
-
+import pytest
+from gammapy.utils.units import standardise_unit, energy_unit_format
+import astropy.units as u
+import numpy as np
+from gammapy.maps import MapAxis
 
 def test_standardise_unit():
     assert standardise_unit("ph cm-2 s-1") == "cm-2 s-1"
     assert standardise_unit("ct cm-2 s-1") == "cm-2 s-1"
     assert standardise_unit("cm-2 s-1") == "cm-2 s-1"
+
+
+axis = MapAxis.from_nodes([1e-1, 200, 3.5e3, 4.6e4], name="energy", unit="GeV")
+values = [
+    (1530 *u.eV, "1.53 keV"),
+    (1530 *u.keV, "1.53 MeV"),
+    (1530e4 *u.keV, "15.3 GeV"),
+    (1530 *u.GeV, "1.53 TeV"),
+    (1530.5e8 *u.keV, "153 TeV"),
+    (1530.5 *u.TeV, "1.53 PeV"),
+    (np.array([1e3,3.5e6,400.4e12,1512.5e12])*u.eV, ('1.00 keV', '3.50 MeV', '400 TeV','1.51 PeV')),
+    ([1.54e2*u.GeV, 4300*u.keV, 300.6e12*u.eV], ('154 GeV', '4.30 MeV', '301 TeV')),
+    (axis.center, ('100 MeV', '200 GeV', '3.50 TeV', '46.0 TeV')),
+    ([u.Quantity(x) for x in axis.as_plot_labels ] ,('100 MeV', '200 GeV', '3.50 TeV', '46.0 TeV'))
+]
+
+@pytest.mark.parametrize("q, expect", values)
+def test_energy_unit_format(q, expect):
+    assert energy_unit_format(q) == expect

gammapy/utils/units.py

@@ -2,6 +2,8 @@
 """Units and Quantity related helper functions"""
 import logging
 import astropy.units as u
+import numpy as np
+from math import floor
 
 __all__ = ["standardise_unit", "unit_from_fits_image_hdu"]
 
@@ -64,3 +66,34 @@ def unit_from_fits_image_hdu(header):
         unit = ""
 
     return standardise_unit(unit)
+
+def energy_unit_format(E):
+    """
+    Format energy quantities to a string representation that is more comfortable to read
+    by switching to the most relevant unit (keV, MeV, GeV, TeV) and changing the float precision.
+
+    Parameters
+    ----------    
+    E: `~astropy.units.Quantity`
+        Quantity or list of quantities
+
+    Returns
+    -------
+    str : str
+        Returns a string or tuple of strings with energy unit formatted        
+    """    
+    try:
+         iter(E)
+    except TypeError:
+         pass
+    else:
+         return tuple(map(energy_unit_format, E))
+
+    i = floor(np.log10(E.to_value(u.eV)) / 3) # a new unit every 3 decades
+    unit = (u.eV, u.keV, u.MeV, u.GeV, u.TeV, u.PeV)[i] if i < 5 else u.PeV        
+
+    v = E.to_value(unit)
+    i=floor(np.log10(v))
+    prec = (2,1,0)[i] if i < 3 else 0
+
+    return f"{v:0.{prec}f} {unit}"