diff --git a/src/km3irf/build_aeff.py b/src/km3irf/build_aeff.py
index ee20c1c9b69309fb5fe17eb83492dc93887f20dd..8083e37061f679c24211e13b0262447004a160b8 100644
--- a/src/km3irf/build_aeff.py
+++ b/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,
@@ -125,6 +125,100 @@ 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):
"""
retrieve information from data and pack it to pandas DataFrame
@@ -246,3 +340,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!")