Improved integration for common Astromodels Functions, and custom function cleanup

cosipy/response/ExtendedSourceResponse.py

@@ -1,8 +1,11 @@
-from histpy import Histogram
 import numpy as np
+
 import astropy.units as u
-import gc
-from astromodels.functions.function import Function1D, FunctionMeta, ModelAssertionViolation, Function2D, Function3D
+
+from astromodels import Function3D
+
+from histpy import Histogram
+
 from .functions import get_integrated_extended_model
 from .functions_3d import get_integrated_extended_model_3d
 
@@ -29,26 +32,34 @@ def __init__(self, *args, **kwargs):
         """
         Initialize an ExtendedSourceResponse object.
         """
-        # Not to track the under/overflow bins
+
         kwargs['track_overflow'] = False
+        kwargs['sparse'] = False
 
         super().__init__(*args, **kwargs)
 
-        if not list(self.axes.labels) == ['NuLambda', 'Ei', 'Em', 'Phi', 'PsiChi']:
+        self.post_init()
+
+    def post_init(self):
+        """
+        Do init operations specific to our subclass
+        """
+        if not tuple(self.axes.labels) == ('NuLambda', 'Ei', 'Em', 'Phi', 'PsiChi'):
             # 'NuLambda' should be 'lb' if it is in the gal. coordinates?
             raise ValueError(f"The input axes {self.axes.labels} is not supported by ExtendedSourceResponse class.")
 
+        # unit required for get_expectation() input
+        self._exp_unit = (u.s * u.cm**2 * u.sr)**(-1)
+
     @classmethod
-    def open(cls, filename, name='hist'):
+    def _open(cls, name='hist'):
         """
-        Load data from an HDF5 file.
+        Load response from an HDF5 group.
 
         Parameters
         ----------
-        filename : str
-            The path to the HDF5 file.
         name : str, optional
-            The name of the histogram group in the HDF5 file (default is 'hist').
+            The name of the histogram group (default is 'hist').
 
         Returns
         -------
@@ -60,11 +71,16 @@ def open(cls, filename, name='hist'):
         ValueError
             If the shape of the contents does not match the axes.
         """
-        resp = super().open(filename, name)
+
+        resp = super()._open(name)
+
+        resp.track_overflow(False)
 
         if resp.is_sparse:
             resp = resp.to_dense()
-
+
+        resp.post_init()
+
         return resp
 
     def get_expectation(self, allsky_image_model):
@@ -73,28 +89,24 @@ def get_expectation(self, allsky_image_model):
 
         Parameters
         ----------
-        allsky_image_model : Histogram 
+        allsky_image_model : Histogram
             The all-sky image model to use for calculation.
 
         Returns
         -------
         Histogram
             A histogram representing the calculated expectation.
         """
-        if self.axes[0].label == allsky_image_model.axes[0].label \
-            and self.axes[1].label == allsky_image_model.axes[1].label \
-            and np.all(self.axes[0].edges == allsky_image_model.axes[0].edges) \
-            and np.all(self.axes[1].edges == allsky_image_model.axes[1].edges) \
-            and allsky_image_model.unit == u.Unit('1/(s*cm*cm*sr)'):
-
-            contents = np.tensordot(allsky_image_model.contents, self.contents, axes=([0,1], [0,1]))
-            contents *= self.axes[0].pixarea()
-
-            return Histogram(edges=self.axes[2:], contents=contents, copy_contents=False)
-
-        else:
+
+        if allsky_image_model.axes[:2] != self.axes[:2] or \
+           allsky_image_model.unit != self._exp_unit:
             raise ValueError(f"The input allskymodel mismatches with the extended source response.")
 
+        contents = np.tensordot(allsky_image_model.contents, self.contents, axes=((0,1), (0,1)))
+        contents *= self.axes[0].pixarea()
+
+        return Histogram(edges=self.axes[2:], contents=contents, copy_contents=False)
+
     def get_expectation_from_astromodel(self, source):
         """
         Calculate expectation from an astromodels extended source model.
@@ -115,13 +127,14 @@ def get_expectation_from_astromodel(self, source):
             A histogram representing the calculated expectation based on the
             provided extended source model.
         """
-
-        if isinstance(source.spatial_shape, Function3D):
 
-            allsky_image_model = get_integrated_extended_model_3d(source, image_axis = self.axes[0], energy_axis = self.axes[1])
-
+        if isinstance(source.spatial_shape, Function3D):
+            allsky_image_model = get_integrated_extended_model_3d(source,
+                                                                  image_axis = self.axes[0],
+                                                                  energy_axis = self.axes[1])
         else:
+            allsky_image_model = get_integrated_extended_model(source,
+                                                               image_axis = self.axes[0],
+                                                               energy_axis = self.axes[1])
 
-            allsky_image_model = get_integrated_extended_model(source, image_axis = self.axes[0], energy_axis = self.axes[1]) 
-
         return self.get_expectation(allsky_image_model)

cosipy/response/PointSourceResponse.py

@@ -100,7 +100,7 @@ def get_expectation(self, spectrum, polarization=None, flux=None):
         if flux is None:
             energy_axis = self.photon_energy_axis
             flux = get_integrated_spectral_model(spectrum, energy_axis)
-
+            
         expectation = np.tensordot(contents, flux.contents, axes=(0, 0))
 
         # if self is sparse, expectation will be a SparseArray with
@@ -110,6 +110,6 @@ def get_expectation(self, spectrum, polarization=None, flux=None):
                          copy_contents = False)
 
         if not hist.unit == u.dimensionless_unscaled:
-            raise RuntimeError("Expectation should be dimensionless, but has units of " + str(hist.unit) + ".")
-
+            raise RuntimeError(f"Expectation should be dimensionless, but has units of {(hist.unit)}.")
+        
         return hist

cosipy/response/functions.py

@@ -1,8 +1,7 @@
 import numpy as np
+
 import astropy.units as u
-from astropy.units import Quantity
 from astropy.coordinates import Galactic
-from scipy import integrate
 
 from histpy import Histogram
 
@@ -23,8 +22,8 @@
 )
 
 def get_integrated_spectral_model(spectrum, energy_axis):
-    """
-    Get the photon fluxes integrated over given energy bins with an input astropy spectral model
+    """Get the photon fluxes integrated over given energy bins with an
+    input astropy spectral model
 
     Parameters
     ----------
@@ -45,8 +44,38 @@ def get_integrated_spectral_model(spectrum, energy_axis):
 
     Notes
     -----
-    This function determines the unit of the spectrum, performs the integration
-    over each energy bin, and returns the result as a Histogram object.
+    This function determines the unit of the spectrum, performs the
+    integration over each energy bin, and returns the result as a
+    Histogram object.
+
+    """
+
+    from cosipy.response.integrals import get_integral_values
+
+    spectrum_unit = get_spectrum_unit(spectrum)
+
+    flux_values = get_integral_values(spectrum, energy_axis.edges.value)
+
+    flux = Histogram(energy_axis,
+                     contents = flux_values,
+                     unit = spectrum_unit * energy_axis.unit,
+                     copy_contents = False)
+
+    return flux
+
+
+def get_spectrum_unit(spectrum):
+    """
+    Get the unit of the spectral model.
+
+    Parameters
+    ----------
+    spectrum : astromodels.functions
+        One-dimensional spectral function from astromodels.
+
+    Returns:
+    astropy.unit for the spectrum
+
     """
 
     from cosipy.threeml import Band_Eflux
@@ -77,33 +106,18 @@ def get_integrated_spectral_model(spectrum, energy_axis):
                 except:
                     raise RuntimeError("Spectrum not yet supported because units are unknown.")
 
-    if isinstance(spectrum, DiracDelta):
-        flux = [spectrum.value.value
-                if spectrum.zero_point.value >= lo_lim and spectrum.zero_point.value <= hi_lim
-                else 0
-                for lo_lim, hi_lim in zip(energy_axis.lower_bounds.value,
-                                          energy_axis.upper_bounds.value)]
+    return spectrum_unit
 
-    else:
-        flux = [integrate.quad(spectrum, lo_lim, hi_lim)[0]
-                for lo_lim, hi_lim in zip(energy_axis.lower_bounds.value,
-                                          energy_axis.upper_bounds.value)]
-
-    flux = Histogram(energy_axis,
-                     contents = flux,
-                     unit = spectrum_unit * energy_axis.unit)
-
-    return flux
 
 def get_integrated_extended_model(extendedmodel, image_axis, energy_axis):
-    """
-    Calculate the integrated flux map for an extended source model.
+    """Calculate the integrated flux map for an extended source model.
 
     Parameters
     ----------
     extendedmodel : astromodels.ExtendedSource
-        An astromodels extended source model object. This model represents
-        the spatial and spectral distribution of an extended astronomical source.
+        An astromodels extended source model object. This model
+        represents the spatial and spectral distribution of an
+        extended astronomical source.
     image_axis : histpy.HealpixAxis
         Spatial axis for the image.
     energy_axis : histpy.Axis
@@ -116,21 +130,26 @@ def get_integrated_extended_model(extendedmodel, image_axis, energy_axis):
 
     Notes
     -----
-    This function first integrates the spectral model over the energy bins,
-    then combines it with the spatial distribution to create a 2D flux map.
+    This function first integrates the spectral model over the energy
+    bins, then combines it with the spatial distribution to create a
+    2D flux map.
+
     """
 
     if not isinstance(image_axis.coordsys, Galactic):
         raise ValueError
 
-    integrated_flux = get_integrated_spectral_model(spectrum = extendedmodel.spectrum.main.shape, energy_axis = energy_axis)
+    integrated_flux = \
+        get_integrated_spectral_model(spectrum = extendedmodel.spectrum.main.shape,
+                                      energy_axis = energy_axis)
+
+    l, b = image_axis.pix2ang(np.arange(image_axis.npix), lonlat=True)
+    normalized_map = extendedmodel.spatial_shape(l, b) / u.sr
 
-    npix = image_axis.npix
-    coords = image_axis.pix2skycoord(np.arange(npix))
-    normalized_map = extendedmodel.spatial_shape(coords.l.deg, coords.b.deg) / u.sr
+    flux = np.tensordot(normalized_map, integrated_flux.contents, axes = 0)
 
     flux_map = Histogram((image_axis, energy_axis),
-                         contents = np.tensordot(normalized_map, integrated_flux.contents, axes = 0),
+                         contents = flux,
                          copy_contents = False)
 
     return flux_map

cosipy/response/functions_3d.py

@@ -1,18 +1,18 @@
 import numpy as np
-import astropy.units as u
-from astropy.units import Quantity
-from astropy.coordinates import Galactic
+
 from scipy import integrate
 from scipy.interpolate import interp1d
-from histpy import Histogram
 
-import sys
+import astropy.units as u
+from astropy.coordinates import Galactic
+
+from histpy import Histogram
 
 import logging
 logger = logging.getLogger(__name__)
 
 def get_integrated_extended_model_3d(extendedmodel, image_axis, energy_axis):
-    
+
     """
     Calculate the integrated flux map for an extended source model.
 
@@ -31,48 +31,51 @@ def get_integrated_extended_model_3d(extendedmodel, image_axis, energy_axis):
     flux_map : histpy.Histogram
         2D histogram representing the integrated flux map.
     """
-
+
+    from cosipy.threeml.custom_functions import GalpropHealpixModel
+
     if not isinstance(image_axis.coordsys, Galactic):
         raise ValueError
 
-    npix = image_axis.npix
-    coords = image_axis.pix2skycoord(np.arange(npix))
+    l, b = image_axis.pix2ang(np.arange(image_axis.npix), lonlat=True)
 
-    # GalpropHealpixModel: 
-    if type(extendedmodel.spatial_shape).__name__ == "GalpropHealpixModel":
-
-        # The norm is updated internally by 3ML for each likelihood call:
+    if isinstance(extendedmodel.spatial_shape, GalpropHealpixModel):
+
+        # The norm is updated internally by 3ML for each likelihood call
         norm = extendedmodel.spatial_shape.K.value
-       
-        # Make sure the dummy spectral parameter is fixed:
+
+        # Make sure the dummy spectral parameter is fixed
         extendedmodel.spectrum.main.Constant.k.free = False
 
         # First get differential intensity from GALPROP model (ph/cm2/s/sr/MeV),
         # and then integrate over energy bins. We attach the integrated flux
-        # to the extended model instance so that it only needs to be calculated once. 
+        # to the extended model instance so that it only needs to be calculated once.
         if not isinstance(extendedmodel.spatial_shape._result, np.ndarray):
-        
-            intensity = (1/norm)*extendedmodel.spatial_shape.evaluate(coords.l.deg,coords.b.deg,energy_axis.edges.to(u.MeV),norm)
-        
-            # Integrate over energy bins for each sky position:
+
+            intensity = (1/norm)*extendedmodel.spatial_shape.evaluate(l, b, energy_axis.edges.to(u.MeV), norm)
+
+            # Integrate over energy bins for each sky position
             extendedmodel.spatial_shape.intg_flux = np.zeros((intensity.shape[0], intensity.shape[1]-1))
-
-            # Convert units outside loop to optimize speed:
-            energy = energy_axis.edges.to(u.MeV)
-            bin_low = energy_axis.lower_bounds.to(u.MeV)/u.MeV # unitless for integration
-            bin_high = energy_axis.upper_bounds.to(u.MeV)/u.MeV # unitless for integration
-
-            # Integrate spectrum over energy bins for each spatial pixel:
+
+            # Convert units outside loop to optimize speed
+            energy_edges = energy_axis.edges.to_value(u.MeV)
+
+            # Integrate spectrum over energy bins for each spatial pixel
             logger.info("Integrating intensity over energy bins...")
             for j in range(len(intensity)):
-
-                interp_func = interp1d(energy, intensity[j], bounds_error=False, fill_value='extrapolate')
-
-                extendedmodel.spatial_shape.intg_flux[j] = Quantity([integrate.quad(interp_func, lo_lim, hi_lim)[0]
-                         for lo_lim,hi_lim
-                         in zip(bin_low, bin_high)])
-
-        flux_map = Histogram([image_axis, energy_axis], \
-                contents = norm*extendedmodel.spatial_shape.intg_flux*((u.s * u.cm**2 * u.sr) ** (-1)))
+
+                interp_func = interp1d(energy_edges, intensity[j], bounds_error=False, fill_value='extrapolate')
+
+                extendedmodel.spatial_shape.intg_flux[j] = \
+                    np.array([integrate.quad(interp_func, lo_lim, hi_lim)[0]
+                              for lo_lim, hi_lim
+                              in zip(energy_edges[:-1], energy_edges[1:])])
+
+        flux = norm * extendedmodel.spatial_shape.intg_flux
+
+        flux_map = Histogram((image_axis, energy_axis), \
+                             contents = flux, \
+                             unit = (u.s * u.cm**2 * u.sr) ** (-1),
+                             copy_contents = False)
 
     return flux_map