Merge pull request #88 from thejevans/documentation_and_tools

Documentation and tools

mla/threeml/IceCubeLike.py

@@ -32,7 +32,7 @@
 r"""This IceCube plugin is currently under develop by Kwok Lung Fan"""
 
 
-class NeutrinoPointSource(PointSource):
+class NeutrinoPointSource(PointSource,Node):
     """
     Class for NeutrinoPointSource. It is inherited from astromodels PointSource class.
     """

mla/threeml/__init__.py

@@ -3,3 +3,4 @@
 from . import spectral
 from . import data_handlers
 from . import sob_terms
+from . import IceCubeLike

mla/threeml/data_handlers.py

@@ -20,7 +20,13 @@
 
 @dataclasses.dataclass
 class ThreeMLDataHandler(data_handlers.NuSourcesDataHandler):
-    """Docstring"""
+    """
+    Inheritance class from NuSourcesDataHandler.
+    For time independent 3ML analysis.
+
+    Additional init argument:
+        injection_spectrum: spectral.BaseSpectrum
+    """
 
     injection_spectrum: spectral.BaseSpectrum
     _injection_spectrum: spectral.BaseSpectrum = dataclasses.field(
@@ -35,42 +41,71 @@ def __post_init__(self) -> None:
         )
 
     def build_signal_energy_histogram(
-        self, spectrum: spectral.BaseSpectrum, bins: np.ndarray
+        self, spectrum: spectral.BaseSpectrum, bins: np.ndarray, scale: float
     ) -> np.ndarray:
-        """Docstring"""
+        """
+        Building the signal energy histogram.
+        Only used when using MC instead of IRF to build signal energy histogram.
+
+        Args:
+            spectrum: signal spectrum
+            bins: 2d bins in sindec and logE
+        """
         return np.histogram2d(
             self.reduced_reco_sim["sindec"],
             self.reduced_reco_sim["logE"],
             bins=bins,
             weights=self.reduced_reco_sim["ow"]
-            * spectrum(self.reduced_reco_sim["trueE"]),
+            * spectrum(self.reduced_reco_sim["trueE"] * scale),
             density=True,
         )[0]
 
     def cut_reconstructed_sim(self, dec: float, sampling_width: float) -> np.ndarray:
-        """Docstring"""
+        """
+        Cutting the MC based on reconstructed dec.
+        Only use when using MC instead of IRF to build signal energy histogram.
+
+        Args:
+            dec: declination of the source
+            sampling_width: size of the sampling band in reconstruction dec.
+        """
         dec_dist = np.abs(dec - self._full_sim["dec"])
         close = dec_dist < sampling_width
         return self._full_sim[close].copy()
 
     @property
     def reduced_reco_sim(self) -> np.ndarray:
-        """Docstring"""
+        """
+        Return the reduced sim based on reconstructed dec.
+        This is the return output of cut_reconstructed_sim.
+        """
         return self._reduced_reco_sim
 
     @reduced_reco_sim.setter
     def reduced_reco_sim(self, reduced_reco_sim: np.ndarray) -> None:
-        """Docstring"""
+        """
+        setting the reduced sim based on reconstructed dec directly.
+
+        Args:
+            reduced_reco_sim: reduced sim based on reconstructed dec
+        """
         self._reduced_reco_sim = reduced_reco_sim.copy()
 
     @property
     def injection_spectrum(self) -> spectral.BaseSpectrum:
-        """Docstring"""
+        """
+        Getting the injection spectrum
+        """
         return self._injection_spectrum
 
     @injection_spectrum.setter
     def injection_spectrum(self, inject_spectrum: spectral.BaseSpectrum) -> None:
-        """Docstring"""
+        """
+        Setting the injection spectrum
+
+        Args:
+            inject_spectrum: spectrum used for injection
+        """
         if isinstance(inject_spectrum, property):
             # initial value not specified, use default
             inject_spectrum = ThreeMLDataHandler._injection_spectrum

mla/threeml/sob_terms.py

@@ -35,14 +35,21 @@
 
 @dataclasses.dataclass
 class ThreeMLPSEnergyTerm(sob_terms.SoBTerm):
-    """Docstring"""
+    """
+    Energy term for 3ML. Constructs only from 3ML energy term factory
+    """
 
     _energysobhist: np.ndarray
     _sin_dec_idx: np.ndarray
     _log_energy_idx: np.ndarray
 
     def update_sob_hist(self, factory: sob_terms.SoBTermFactory) -> None:
-        """Docstring"""
+        """
+        Updating the signal-over-background energy histogram.
+
+        Args:
+            factory: energy term factory
+        """
         self._energysobhist = factory.cal_sob_map()
 
     @property
@@ -70,7 +77,15 @@ class ThreeMLBaseEnergyTermFactory(sob_terms.SoBTermFactory):
 
 @dataclasses.dataclass
 class ThreeMLPSEnergyTermFactory(ThreeMLBaseEnergyTermFactory):
-    """Docstring"""
+    """
+    This is the class for using MC directly to build the Energy terms.
+    We sugguest using the IRF for Energy term factory due to speed.
+
+    Args:
+        data_handler: 3ML data handler
+        source: 3ML source object
+        spectrum: signal spectrum
+    """
 
     data_handler: data_handlers.ThreeMLDataHandler
     source: sources.PointSource
@@ -85,6 +100,7 @@ class ThreeMLPSEnergyTermFactory(ThreeMLBaseEnergyTermFactory):
     _bins: np.ndarray = dataclasses.field(init=False, repr=False)
     _ow_hist: np.ndarray = dataclasses.field(init=False, repr=False)
     _ow_ebin: np.ndarray = dataclasses.field(init=False, repr=False)
+    _unit_scale: float = dataclasses.field(init=False, repr=False)
 
     def __post_init__(self) -> None:
         """Docstring"""
@@ -102,6 +118,7 @@ def __post_init__(self) -> None:
             self.source.location[1],
             self.data_handler.config["reco_sampling_width"],
         )
+        self._unit_scale = self.config["Energy_convesion(ToGeV)"]
         self._bins = np.array([self._sin_dec_bins, self._log_energy_bins])
         self._init_bg_sob_map()
         self._build_ow_hist()
@@ -175,7 +192,7 @@ def cal_sob_map(self) -> np.ndarray:
             log(energy) for a given gamma.
         """
         sig_h = self.data_handler.build_signal_energy_histogram(
-            self.spectrum, self._bins
+            self.spectrum, self._bins, self._unit_scale
         )
         bin_centers = self._log_energy_bins[:-1] + np.diff(self._log_energy_bins) / 2
         # Normalize histogram by dec band
@@ -259,6 +276,7 @@ class ThreeMLPSIRFEnergyTermFactory(ThreeMLPSEnergyTermFactory):
     _irf: np.ndarray = dataclasses.field(init=False, repr=False)
     _sindec_bounds: np.ndarray = dataclasses.field(init=False, repr=False)
     _ntrueebin: int = dataclasses.field(init=False, repr=False)
+    _unit_scale: float = dataclasses.field(init=False, repr=False)
 
     def __post_init__(self) -> None:
         """Docstring"""
@@ -291,6 +309,7 @@ def __post_init__(self) -> None:
         self._sindec_bounds = np.array([lower_sindec_index, uppper_sindec_index])
         self._bins = np.array([self._sin_dec_bins, self._log_energy_bins])
         self._truelogebin = self.config["list_truelogebin"]
+        self._unit_scale = self.config["Energy_convesion(ToGeV)"]
         self._init_bg_sob_map()
         self._build_ow_hist()
         self._init_irf()
@@ -373,7 +392,7 @@ def _init_irf(self) -> None:
     def build_sig_h(self, spectrum: spectral.BaseSpectrum) -> np.ndarray:
         """Docstring"""
         sig = np.zeros(self._bg_sob.shape)
-        flux = spectrum(self._trueebin)
+        flux = spectrum(self._trueebin * self._unit_scale)  # converting unit
         sig[self._sindec_bounds[0]:self._sindec_bounds[1], :] = np.dot(
             self._irf[self._sindec_bounds[0]:self._sindec_bounds[1], :, :], flux
         )
@@ -462,4 +481,5 @@ def generate_config(cls):
         config["list_truelogebin"] = np.arange(
             2, 9.01 + 0.01, 0.01
         )
+        config["Energy_convesion(ToGeV)"] = 1e6  # GeV to keV
         return config