gammasimtools 0.6.1__py3-none-any.whl → 0.8.2__py3-none-any.whl

simtools/{psf_analysis.py → ray_tracing/psf_analysis.py}

@@ -1,12 +1,15 @@
 """
 Module to analyse psf images (e.g. results from ray tracing simulations).
-Main functionalities are: computing centroids, psf containers etc.
 
-Author: Raul R Prado
+Main functionalities are: computing centroids, psf containers etc.
 
 """
 
+import gzip
 import logging
+import shlex
+import shutil
+import subprocess
 from math import fabs, pi, sqrt
 from pathlib import Path
 
@@ -29,18 +32,29 @@ class PSFImage:
 
     Parameters
     ----------
-    focal_length: float, optional
+    focal_length: float
         Focal length of the system in cm. If not given, PSF can only be computed in cm.
-    total_scattered_area: float, optional
+    total_scattered_area: float
         Scatter area of all photons in cm^2. If not given, effective area cannot be computed.
+    containment_fraction: float
+        Containment fraction for PSF calculation.
+    simtel_path: str
+        Path to sim_telarray installation.
     """
 
-    def __init__(self, focal_length=None, total_scattered_area=None):
-        """
-        Initialize PSFImage class.
-        """
-
+    __PSF_RADIUS = "Radius [cm]"
+    __PSF_CUMULATIVE = "Cumulative PSF"
+
+    def __init__(
+        self,
+        focal_length=None,
+        total_scattered_area=None,
+        containment_fraction=None,
+        simtel_path=None,
+    ):
+        """Initialize PSFImage class."""
         self._logger = logging.getLogger(__name__)
+        self.simtel_path = simtel_path
 
         self._total_photons = None
         self._number_of_detected_photons = None
@@ -52,11 +66,65 @@ class PSFImage:
         self.centroid_y = None
         self._total_area = total_scattered_area
         self._stored_psf = {}
-        if focal_length is not None:
-            self._cm_to_deg = 180.0 / pi / focal_length
-            self._has_focal_length = True
+        try:
+            self._cm_to_deg = 180.0 / pi / focal_length if focal_length is not None else None
+        except ZeroDivisionError:
+            self._cm_to_deg = None
+            self._logger.warning("Focal length is zero; no conversion from cm to deg possible.")
+        self._containment_fraction = containment_fraction
+
+    def process_photon_list(self, photon_file, use_rx):
+        """
+        Read and process a photon list file generated by sim_telarray.
+
+        Parameters
+        ----------
+        photons_file: str
+            Name of sim_telarray file with photon list.
+        use_rx: bool
+            Use the RX method for analysis.
+        """
+        if use_rx:
+            self._process_simtel_file_using_rx(photon_file)
         else:
-            self._has_focal_length = False
+            self.read_photon_list_from_simtel_file(photon_file)
+
+    def _process_simtel_file_using_rx(self, photon_file):
+        """
+        Process a simtel file with photon lists using the RX method.
+
+        Parameters
+        ----------
+        photons_file: str
+            Name of sim_telarray file with photon list.
+        """
+        try:
+            rx_output = subprocess.Popen(  # pylint: disable=consider-using-with
+                shlex.split(
+                    f"{self.simtel_path}/sim_telarray/bin/rx -f {self._containment_fraction:.2f} -v"
+                ),
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+            )
+            with gzip.open(photon_file, "rb") as _stdin:
+                with rx_output.stdin:
+                    shutil.copyfileobj(_stdin, rx_output.stdin)
+                    try:
+                        rx_output = rx_output.communicate()[0].splitlines()[-1:][0].split()
+                    except IndexError as e:
+                        raise IndexError(f"Unexpected output format from rx: {rx_output}") from e
+        except FileNotFoundError as e:
+            raise FileNotFoundError(f"Photon list file not found: {photon_file}") from e
+
+        try:
+            self.set_psf(2 * float(rx_output[0]), fraction=self._containment_fraction, unit="cm")
+            self.centroid_x = float(rx_output[1])
+            self.centroid_y = float(rx_output[2])
+            self._effective_area = float(rx_output[5])
+        except IndexError as e:
+            raise IndexError(f"Unexpected output format from rx: {rx_output}") from e
+        except ValueError as e:
+            raise ValueError(f"Invalid output format from rx: {rx_output}") from e
 
     def read_photon_list_from_simtel_file(self, photons_file):
         """
@@ -76,8 +144,6 @@ class PSFImage:
         self._logger.info(f"Reading sim_telarray file {photons_file}")
         self._total_photons = 0
         if Path(photons_file).suffix == ".gz":
-            import gzip  # pylint: disable=import-outside-toplevel
-
             file_open_function = gzip.open
         else:
             file_open_function = open
@@ -138,9 +204,6 @@ class PSFImage:
                     f" {self._total_area} != {total_area_in_file}"
                     " - Keeping the original value"
                 )
-            else:
-                # Do nothing - Keep the original value of _total_area
-                pass
         elif b"#" in line or len(words) == 0:
             # Skipping comments
             pass
@@ -149,9 +212,14 @@ class PSFImage:
             self.photon_pos_x.append(float(words[2]))
             self.photon_pos_y.append(float(words[3]))
 
-    def get_effective_area(self):
+    def get_effective_area(self, tel_transmission=1.0):
         """
-        Return effective area pre calculated
+        Return effective area pre calculated.
+
+        Parameters
+        ----------
+        telescope_transmission : float
+            Telescope transmission parameter.
 
         Returns
         -------
@@ -160,14 +228,14 @@ class PSFImage:
             scattering area was not set).
         """
         if "_effective_area" in self.__dict__ and self._effective_area is not None:
-            return self._effective_area
+            return self._effective_area * tel_transmission
 
         self._logger.error("Effective Area could not be calculated")
         return None
 
     def set_effective_area(self, value):
         """
-        Set effective area
+        Set effective area.
 
         Parameters
         ----------
@@ -179,7 +247,7 @@ class PSFImage:
 
     def get_psf(self, fraction=0.8, unit="cm"):
         """
-        Return PSF
+        Return PSF.
 
         Parameters
         ----------
@@ -194,7 +262,7 @@ class PSFImage:
             Containing diameter for a certain intensity fraction (PSF).
 
         """
-        if unit == "deg" and not self._has_focal_length:
+        if unit == "deg" and self._cm_to_deg is None:
             self._logger.error("PSF cannot be computed in deg because focal length is not set")
             return None
         if fraction not in self._stored_psf:
@@ -215,7 +283,7 @@ class PSFImage:
         unit: str
             'cm' or 'deg'. 'deg' will not work if focal length was not set.
         """
-        if unit == "deg" and not self._has_focal_length:
+        if unit == "deg" and self._cm_to_deg is None:
             self._logger.error("PSF cannot be set in deg because focal length is not set")
             return
         unit_factor = 1 if unit == "cm" else 1.0 / self._cm_to_deg
@@ -235,6 +303,7 @@ class PSFImage:
     def _find_psf(self, fraction):
         """
         Try to find PSF by a smart algorithm first.
+
         If it fails, _find_radius_by_scanning is called and do it by brute force.
 
         Parameters
@@ -262,7 +331,7 @@ class PSFImage:
         scale = 0.5 * sqrt(current_radius * current_radius / start_number)
         delta_number = start_number - target_number
         n_iter = 0
-        max_iter = 100
+        max_iter = 1000
         tolerance = self._number_of_detected_photons / 1000.0
         found_radius = False
         while not found_radius and n_iter < max_iter:
@@ -276,7 +345,6 @@ class PSFImage:
             found_radius = fabs(delta_number) < tolerance
 
         if found_radius:
-            # Diameter = 2 * radius
             return 2 * current_radius
 
         self._logger.warning("Could not find PSF efficiently - trying by scanning")
@@ -302,8 +370,9 @@ class PSFImage:
 
         def scan(dr, rad_min, rad_max):
             """
-            Scan the image from rad_min to rad_max in steps of dr until
-            it finds target_number photons inside.
+            Scan the image from rad_min to rad_max until it finds target_number photons inside.
+
+            Scanning is done in steps of dr.
 
             Returns
             -------
@@ -321,7 +390,7 @@ class PSFImage:
             found_radius = False
             while not found_radius:
                 s0, s1 = self._sum_photons_in_radius(r0), self._sum_photons_in_radius(r1)
-                if s0 < target_number < s1:
+                if s0 < target_number <= s1:
                     found_radius = True
                     break
                 if r1 > rad_max:
@@ -342,9 +411,7 @@ class PSFImage:
         return radius
 
     def _sum_photons_in_radius(self, radius):
-        """
-        Return the number of photons inside a certain radius.
-        """
+        """Return the number of photons inside a certain radius."""
         return np.searchsorted(self.photon_r, radius)
 
     def get_image_data(self, centralized=True):
@@ -367,9 +434,12 @@ class PSFImage:
             x_pos_data = np.array(self.photon_pos_x)
             y_pos_data = np.array(self.photon_pos_y)
         d_type = {"names": ("X", "Y"), "formats": ("f8", "f8")}
-        return np.core.records.fromarrays(np.c_[x_pos_data, y_pos_data].T, dtype=d_type)
+        result = np.recarray((len(x_pos_data),), dtype=d_type)
+        result.X = x_pos_data
+        result.Y = y_pos_data
+        return result
 
-    def plot_image(self, centralized=True, **kwargs):
+    def plot_image(self, centralized=True, file_name=None, **kwargs):
         """
         Plot 2D image as histogram (in cm).
 
@@ -395,11 +465,12 @@ class PSFImage:
         kwargs_for_psf = collect_kwargs("psf", kwargs)
 
         ax = plt.gca()
-        # Image histogram
+        ax.set_xlabel("X Position (cm)")
+        ax.set_ylabel("Y Position (cm)")
         ax.hist2d(data["X"], data["Y"], **kwargs_for_image)
         ax.set_aspect("equal", "datalim")
 
-        # PSF circle
+        # PSF circle (80%)
         center = (0, 0) if centralized else (self.centroid_x, self.centroid_y)
         circle = plt.Circle(center, self.get_psf(0.8) / 2, **kwargs_for_psf)
         ax.add_artist(circle)
@@ -407,6 +478,10 @@ class PSFImage:
         ax.axhline(0, color="k", linestyle="--", zorder=3, linewidth=0.5)
         ax.axvline(0, color="k", linestyle="--", zorder=3, linewidth=0.5)
 
+        if file_name is not None:
+            plt.savefig(file_name)
+            plt.close()
+
     def get_cumulative_data(self, radius=None):
         """
         Provide cumulative data (intensity vs radius).
@@ -420,9 +495,8 @@ class PSFImage:
         -------
         (radius, intensity)
         """
-
         if radius is not None:
-            radius_all = radius.to(u.cm).value
+            radius_all = radius.to(u.cm).value if isinstance(radius, u.Quantity) else radius
         else:
             radius_all = list(np.linspace(0, 1.6 * self.get_psf(0.8), 30))
 
@@ -430,20 +504,32 @@ class PSFImage:
         for rad in radius_all:
             intensity.append(self._sum_photons_in_radius(rad) / self._number_of_detected_photons)
         d_type = {
-            "names": ("Radius [cm]", "Cumulative PSF"),
+            "names": (self.__PSF_RADIUS, self.__PSF_CUMULATIVE),
             "formats": ("f8", "f8"),
         }
-        return np.core.records.fromarrays(np.c_[radius_all, intensity].T, dtype=d_type)
+        result = np.recarray((len(radius_all),), dtype=d_type)
+        result[self.__PSF_RADIUS] = radius_all
+        result[self.__PSF_CUMULATIVE] = intensity
 
-    def plot_cumulative(self, **kwargs):
+        return result
+
+    def plot_cumulative(self, file_name=None, d80=None, **kwargs):
         """Plot cumulative data (intensity vs radius).
 
         Parameters
         ----------
         **kwargs:
-            image_* for the histogram plot and psf_* for the psf circle."""
+            image_* for the histogram plot and psf_* for the psf circle.
+        """
         data = self.get_cumulative_data()
-        plt.plot(data["Radius [cm]"], data["Cumulative PSF"], **kwargs)
-
-
-# end of PSFImage
+        ax = plt.gca()
+        plt.tight_layout(pad=1.5)
+        ax.set_xlabel("Radius (cm)")
+        ax.set_ylabel("Contained light %")
+        plt.plot(data[self.__PSF_RADIUS], data[self.__PSF_CUMULATIVE], **kwargs)
+        plt.axvline(x=self.get_psf(0.8) / 2, color="b", linestyle="--", linewidth=1)
+        if d80 is not None:
+            plt.axvline(x=d80 / 2.0, color="r", linestyle="--", linewidth=1)
+        if file_name is not None:
+            plt.savefig(file_name)
+            plt.close()