Updated another-dev branch

PSF was added to build_aeff.py and some tests for it

src/km3irf/build_aeff.py

@@ -7,7 +7,7 @@ import pandas as pd
 import uproot as ur
 
 from km3io import OfflineReader
-from .irf_tools import aeff_2D
+from .irf_tools import aeff_2D, psf_3D
 
 # import matplotlib.pyplot as plt
 # from matplotlib.colors import LogNorm
@@ -17,6 +17,7 @@ import astropy.units as u
 from astropy.io import fits
 
 from scipy.stats import binned_statistic
+from scipy.ndimage import gaussian_filter1d, gaussian_filter
 
 
 # from collections import defaultdict
@@ -105,7 +106,6 @@ class DataContainer:
         theta_binC = np.arccos(0.5 * (cos_theta_binE[:-1] + cos_theta_binE[1:]))
 
         # Fill histograms for effective area
-        # !!! Check number of events uproot vs pandas
         aeff_all = (
             aeff_2D(
                 e_bins=energy_binE,
@@ -124,6 +124,98 @@ class DataContainer:
 
         return None
 
+    def build_psf(
+        self,
+        df_pass,
+        cos_theta_binE=np.linspace(1, -1, 7),
+        energy_binE=np.logspace(2, 8, 25),
+        rad_binE=np.concatenate(
+            (
+                np.linspace(0, 1, 21),
+                np.linspace(1, 5, 41)[1:],
+                np.linspace(5, 30, 51)[1:],
+                [180.0],
+            )
+        ),
+        norm=False,
+        smooth=True,
+        smooth_norm=True,
+        output="psf.fits",
+    ):
+        """
+        Build Point Spread Function 3D .fits
+
+        df_pass: incoming data frame
+
+        cos_theta_binE: numpy array of linear bins for cos of zenith angle theta
+
+        energy_binE: log numpy array of enegy bins
+
+        rad_binE: numpy array oflinear radial bins
+        (20 bins for 0-1 deg, 40 bins for 1-5 deg, 50 bins for 5-30 deg, + 1 final bin up to 180 deg)
+
+        output: name of generated PSF file with extension .fits
+
+        """
+        theta_binE = np.arccos(cos_theta_binE)
+        # Bin centers
+        energy_binC = np.sqrt(energy_binE[:-1] * energy_binE[1:])
+        theta_binC = np.arccos(0.5 * (cos_theta_binE[:-1] + cos_theta_binE[1:]))
+        rad_binC = 0.5 * (rad_binE[1:] + rad_binE[:-1])
+
+        # Fill histogram for PSF
+        psf = psf_3D(
+            e_bins=energy_binE,
+            r_bins=rad_binE,
+            t_bins=theta_binE,
+            dataset=df_pass,
+            weights=1,
+        )
+
+        # compute dP/dOmega
+        sizes_rad_bins = np.diff(rad_binE**2)
+        norma = psf.sum(axis=0, keepdims=True)
+        psf /= sizes_rad_bins[:, None, None] * (np.pi / 180) ** 2 * np.pi
+
+        # Normalization for PSF
+        if norm:
+            psf = np.nan_to_num(psf / norma)
+
+        # Smearing
+        if smooth and not norm:
+            s1 = gaussian_filter1d(psf, 0.5, axis=0, mode="nearest")
+            s2 = gaussian_filter1d(psf, 2, axis=0, mode="nearest")
+            s3 = gaussian_filter1d(psf, 4, axis=0, mode="nearest")
+            s4 = gaussian_filter1d(psf, 6, axis=0, mode="constant")
+            psf = np.concatenate(
+                (s1[:10], s2[10:20], s3[20:60], s4[60:-1], [psf[-1]]), axis=0
+            )
+            # smooth edges between the different ranges
+            psf[10:-1] = gaussian_filter1d(psf[10:-1], 1, axis=0, mode="nearest")
+            if smooth_norm:
+                norm_psf_sm = (
+                    psf * sizes_rad_bins[:, None, None] * (np.pi / 180) ** 2 * np.pi
+                ).sum(axis=0, keepdims=True)
+                psf = np.nan_to_num(psf / norm_psf_sm)
+        elif smooth and norm:
+            raise Exception("smooth and norm cannot be True at the same time")
+
+        new_psf_file = WritePSF(
+            energy_binC,
+            energy_binE,
+            theta_binC,
+            theta_binE,
+            rad_binC,
+            rad_binE,
+            psf_T=psf,
+        )
+        new_psf_file.to_fits(file_name=output)
+
+        return None
+
+    def build_edisp():
+        pass
+
 
 def unpack_data(no_bdt, uproot_file):
     """
@@ -246,3 +338,89 @@ class WriteAeff:
         aeff_fits.writeto(file_name, overwrite=True)
 
         return print(f"file {file_name} is written successfully!")
+
+
+# Class for writing PSF to fits files
+class WritePSF:
+    def __init__(
+        self,
+        energy_binC,
+        energy_binE,
+        theta_binC,
+        theta_binE,
+        rad_binC,
+        rad_binE,
+        psf_T,
+    ):
+        self.col1 = fits.Column(
+            name="ENERG_LO",
+            format="{}E".format(len(energy_binC)),
+            unit="GeV",
+            array=[energy_binE[:-1]],
+        )
+        self.col2 = fits.Column(
+            name="ENERG_HI",
+            format="{}E".format(len(energy_binC)),
+            unit="GeV",
+            array=[energy_binE[1:]],
+        )
+        self.col3 = fits.Column(
+            name="THETA_LO",
+            format="{}E".format(len(theta_binC)),
+            unit="rad",
+            array=[theta_binE[:-1]],
+        )
+        self.col4 = fits.Column(
+            name="THETA_HI",
+            format="{}E".format(len(theta_binC)),
+            unit="rad",
+            array=[theta_binE[1:]],
+        )
+        self.col5 = fits.Column(
+            name="RAD_LO",
+            format="{}E".format(len(rad_binC)),
+            unit="deg",
+            array=[rad_binE[:-1]],
+        )
+        self.col6 = fits.Column(
+            name="RAD_HI",
+            format="{}E".format(len(rad_binC)),
+            unit="deg",
+            array=[rad_binE[1:]],
+        )
+        self.col7 = fits.Column(
+            name="RPSF",
+            format="{}D".format(len(energy_binC) * len(theta_binC) * len(rad_binC)),
+            dim="({},{},{})".format(len(energy_binC), len(theta_binC), len(rad_binC)),
+            unit="sr-1",
+            array=[psf_T],
+        )
+
+    def to_fits(self, file_name):
+        cols = fits.ColDefs(
+            [
+                self.col1,
+                self.col2,
+                self.col3,
+                self.col4,
+                self.col5,
+                self.col6,
+                self.col7,
+            ]
+        )
+        hdu = fits.PrimaryHDU()
+        hdu2 = fits.BinTableHDU.from_columns(cols)
+        hdu2.header["EXTNAME"] = "PSF_2D_TABLE"
+        hdu2.header[
+            "HDUDOC"
+        ] = "https://github.com/open-gamma-ray-astro/gamma-astro-data-formats"
+        hdu2.header["HDUVERS"] = "0.2"
+        hdu2.header["HDUCLASS"] = "GADF"
+        hdu2.header["HDUCLAS1"] = "RESPONSE"
+        hdu2.header["HDUCLAS2"] = "RPSF"
+        hdu2.header["HDUCLAS3"] = "FULL-ENCLOSURE"
+        hdu2.header["HDUCLAS4"] = "PSF_TABLE"
+        psf_fits = fits.HDUList([hdu, hdu2])
+        psf_fits.writeto(file_name, overwrite=True)
+
+        return print(f"file {file_name} is written successfully!")