# Licensed under a 3-clause BSD style license - see LICENSE.rst
import numpy as np
import astropy.units as u
import matplotlib.pyplot as plt
from gammapy.maps import Map
from gammapy.modeling.models import PowerLawSpectralModel
__all__ = ["PSFKernel"]
[docs]class PSFKernel:
"""PSF kernel for `~gammapy.maps.Map`.
This is a container class to store a PSF kernel
that can be used to convolve `~gammapy.maps.WcsNDMap` objects.
It is usually computed from an `~gammapy.irf.PSFMap`.
Parameters
----------
psf_kernel_map : `~gammapy.maps.Map`
PSF kernel stored in a Map
Examples
--------
::
from gammapy.maps import Map, WcsGeom, MapAxis
from gammapy.irf import PSFMap
from astropy import units as u
# Define energy axis
energy_axis_true = MapAxis.from_energy_bounds(
"0.1 TeV", "10 TeV", nbin=4, name="energy_true
)
# Create WcsGeom and map
geom = WcsGeom.create(binsz=0.02 * u.deg, width=2.0 * u.deg, axes=[energy_axis_true])
some_map = Map.from_geom(geom)
# Fill map at three positions
some_map.fill_by_coord([[0, 0, 0], [0, 0, 0], [0.3, 1, 3]])
psf = PSFMap.from_gauss(
energy_axis_true=energy_axis_true, sigma=[0.3, 0.2, 0.1] * u.deg
)
kernel = psf.get_psf_kernel(geom=geom, max_radius=1*u.deg)
# Do the convolution
some_map_convolved = some_map.convolve(kernel)
some_map_convolved.plot_grid();
"""
def __init__(self, psf_kernel_map, normalize=True):
self._psf_kernel_map = psf_kernel_map
if normalize:
self.normalize()
[docs] def normalize(self):
"""Force normalisation of the kernel"""
data = self.psf_kernel_map.data
if self.psf_kernel_map.geom.is_image:
axis = (0, 1)
else:
axis = (1, 2)
with np.errstate(divide="ignore", invalid="ignore"):
data = np.nan_to_num(data / data.sum(axis=axis, keepdims=True))
self.psf_kernel_map.data = data
@property
def data(self):
"""Access the PSFKernel numpy array"""
return self._psf_kernel_map.data
@property
def psf_kernel_map(self):
"""The map object holding the kernel (`~gammapy.maps.Map`)"""
return self._psf_kernel_map
[docs] @classmethod
def read(cls, *args, **kwargs):
"""Read kernel Map from file."""
psf_kernel_map = Map.read(*args, **kwargs)
return cls(psf_kernel_map)
[docs] @classmethod
def from_spatial_model(cls, model, geom, max_radius=None, factor=4):
"""Create PSF kernel from spatial model
Parameters
----------
geom : `~gammapy.maps.WcsGeom`
Map geometry
model : `~gammapy.modeling.models.SpatiaModel`
Gaussian width.
max_radius : `~astropy.coordinates.Angle`
Desired kernel map size.
factor : int
Oversample factor to compute the PSF
Returns
-------
kernel : `~gammapy.irf.PSFKernel`
the kernel Map with reduced geometry according to the max_radius
"""
if max_radius is None:
max_radius = model.evaluation_radius
geom = geom.to_odd_npix(max_radius=max_radius)
model.position = geom.center_skydir
geom = geom.upsample(factor=factor)
map = model.integrate_geom(geom, oversampling_factor=1)
return cls(psf_kernel_map=map.downsample(factor=factor))
[docs] @classmethod
def from_gauss(cls, geom, sigma, max_radius=None, factor=4):
"""Create Gaussian PSF.
This is used for testing and examples.
The map geometry parameters (pixel size, energy bins) are taken from ``geom``.
The Gaussian width ``sigma`` is a scalar,
TODO : support array input if it should vary along the energy axis.
Parameters
----------
geom : `~gammapy.maps.WcsGeom`
Map geometry
sigma : `~astropy.coordinates.Angle`
Gaussian width.
max_radius : `~astropy.coordinates.Angle`
Desired kernel map size.
factor : int
Oversample factor to compute the PSF
Returns
-------
kernel : `~gammapy.irf.PSFKernel`
the kernel Map with reduced geometry according to the max_radius
"""
from gammapy.modeling.models import GaussianSpatialModel
gauss = GaussianSpatialModel(sigma=sigma)
return cls.from_spatial_model(
model=gauss, geom=geom, max_radius=max_radius, factor=factor
)
[docs] def write(self, *args, **kwargs):
"""Write the Map object which contains the PSF kernel to file."""
self.psf_kernel_map.write(*args, **kwargs)
[docs] def to_image(self, spectrum=None, exposure=None, keepdims=True):
"""Transform 3D PSFKernel into a 2D PSFKernel.
Parameters
----------
spectrum : `~gammapy.modeling.models.SpectralModel`
Spectral model to compute the weights.
Default is power-law with spectral index of 2.
exposure : `~astropy.units.Quantity` or `~numpy.ndarray`
1D array containing exposure in each true energy bin.
It must have the same size as the PSFKernel energy axis.
Default is uniform exposure over energy.
keepdims : bool
If true, the resulting PSFKernel will keep an energy axis with one bin.
Default is True.
Returns
-------
weighted_kernel : `~gammapy.irf.PSFKernel`
the weighted kernel summed over energy
"""
map = self.psf_kernel_map
if spectrum is None:
spectrum = PowerLawSpectralModel(index=2.0)
if exposure is None:
exposure = np.ones(map.geom.axes[0].center.shape)
exposure = u.Quantity(exposure)
if exposure.shape != map.geom.axes[0].center.shape:
raise ValueError("Incorrect exposure_array shape")
# Compute weights vector
for name in map.geom.axes.names:
if "energy" in name:
energy_name = name
energy_edges = map.geom.axes[energy_name].edges
weights = spectrum.integral(
energy_min=energy_edges[:-1], energy_max=energy_edges[1:], intervals=True
)
weights *= exposure
weights /= weights.sum()
spectrum_weighted_kernel = map.copy()
spectrum_weighted_kernel.quantity *= weights[:, np.newaxis, np.newaxis]
return self.__class__(spectrum_weighted_kernel.sum_over_axes(keepdims=keepdims))
[docs] def slice_by_idx(self, slices):
"""Slice by idx"""
kernel = self.psf_kernel_map.slice_by_idx(slices=slices)
return self.__class__(psf_kernel_map=kernel)
[docs] def plot_kernel(self, ax=None, energy=None, add_cbar=False, **kwargs):
"""Plot PSF kernel.
Parameters
----------
energy : `~astropy.units.Quantity`, optional
If None, the PSF kernel is summed over the energy axis. Otherwise, the kernel
corresponding to the energy bin including the given energy is shown.
ax : `~matplotlib.axes.Axes`, optional
Axis
kwargs: dict
Keyword arguments passed to `~matplotlib.axes.Axes.imshow`.
Returns
-------
ax : `~matplotlib.axes.Axes`
Axis
"""
ax = plt.gca() if ax is None else ax
if energy:
kernel_map = self.psf_kernel_map
energy_center = kernel_map.geom.axes["energy_true"].center.to(energy.unit)
idx = np.argmin(np.abs(energy_center.value - energy.value))
kernel_map.get_image_by_idx([idx]).plot(ax=ax, add_cbar=add_cbar, **kwargs)
else:
self.to_image().psf_kernel_map.plot(ax=ax, add_cbar=add_cbar, **kwargs)
return ax
[docs] def peek(self, figsize=(15, 5)):
"""Quick-look summary plots.
Parameters
----------
figsize : tuple
Size of the figure.
"""
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=figsize)
axes[0].set_title("Energy-integrated PSF kernel")
self.plot_kernel(ax=axes[0], add_cbar=True)
axes[1].set_title("PSF kernel at 1 TeV")
self.plot_kernel(ax=axes[1], add_cbar=True, energy=1 * u.TeV)
plt.tight_layout()