Added airy disc visualization to spot_diagram.py

optiland/analysis/spot_diagram.py

@@ -7,6 +7,7 @@
 from copy import deepcopy
 import numpy as np
 import matplotlib.pyplot as plt
+import matplotlib.patches as patches
 
 plt.rcParams.update({'font.size': 12, 'font.family': 'cambria'})
 
@@ -65,12 +66,13 @@
         self.data = self._generate_data(self.fields, self.wavelengths,
                                         num_rings, distribution)
 
-    def view(self, figsize=(12, 4)):
+    def view(self, figsize=(12, 4), add_airy_disk = False):
         """View the spot diagram
 
         Args:
             figsize (tuple): the figure size of the output window.
                 Default is (12, 4).
+            add_airy_disk (bool): Airy disc visualization controller.
 
         Returns:
             None
@@ -88,10 +90,25 @@
         geometric_size = self.geometric_spot_radius()
         axis_lim = np.max(geometric_size)
 
+        wavelength = self.optic.wavelengths.primary_wavelength.value
+        centroids = self.centroid()
+        chief_ray_centers = self.generate_chief_rays_centers(wavelength=wavelength)
+        airy_rad_x, airy_rad_y = self.airy_disc_x_y(wavelength=wavelength)
+
         # plot wavelengths for each field
         for k, field_data in enumerate(data):
+            # Calculate the real centroid difference for the current field
+            real_centroid_x = chief_ray_centers[k][0] - centroids[k][0]
+            real_centroid_y = chief_ray_centers[k][1] - centroids[k][1]
             self._plot_field(axs[k], field_data, self.fields[k],
-                             axis_lim, self.wavelengths)
+                             axis_lim, self.wavelengths,
+                             add_airy_disk=add_airy_disk,
+                             field_index=k,
+                             airy_rad_x=airy_rad_x,
+                             airy_rad_y=airy_rad_y,
+                             real_centroid_x=real_centroid_x, 
+                             real_centroid_y=real_centroid_y
+)
 
         # remove empty axes
         for k in range(N, num_rows * 3):
@@ -101,6 +118,151 @@
         plt.tight_layout()
         plt.show()
 
+    def angle_from_cosine(self, a, b):
+        """Calculate the angle (in radians) between two vectors given their cosine values.
+
+        Returns:
+            theta_rad (float): angle in radians.
+        """
+        # Compute the angle using arccos
+        a = a / np.linalg.norm(a)
+        b = b / np.linalg.norm(b)
+        theta = np.arccos(np.clip(np.dot(a, b), -1, 1)) 
+        #theta_deg = np.degrees(theta_rad)  # Convert to degrees
+
+        return theta
+
+    def f_number(self, n, theta):
+        """Calculates the F#
+
+        Returns:
+            N_w (float): Physical F number.
+        """
+        N_w = 1 / (2 * n * np.sin(theta))
+        return N_w
+
+    def airy_radius(self, n_w, wavelength):
+        """Calculates the airy radius
+        Returns:
+            r (float): Airy radius.
+        """
+        r = 1.22 * n_w * wavelength
+        return r
+
+    def generate_marginal_rays(self, H_x, H_y, wavelength):
+        """Generates marginal rays at the stop at each pupil max, Px = (-1, +1), Py = (-1, +1)
+
+        Returns:
+            ray_tuple (tuple): Contains marginal each rays generated by trace_generic method, 
+            in north (Py=1), south (Py=-1), east (Px=1), west (Px=-1) directions.
+        """
+        ray_north = self.optic.trace_generic(Hx=H_x, Hy=H_y, Px=0,  Py=1, wavelength=wavelength)
+        ray_south = self.optic.trace_generic(Hx=H_x, Hy=H_y, Px=0,  Py=-1, wavelength=wavelength)
+        ray_east  = self.optic.trace_generic(Hx=H_x, Hy=H_y, Px=1,  Py=0, wavelength=wavelength)
+        ray_west  = self.optic.trace_generic(Hx=H_x, Hy=H_y, Px=-1, Py=0, wavelength=wavelength)
+
+        ray_tuple = ray_north, ray_south, ray_east, ray_west
+
+        return ray_tuple
+
+    def generate_marginal_rays_cosines(self, H_x, H_y, wavelength):
+        """Generates directional cosines for each marginal ray. Calculates one field at a time (one height at a time).
+
+        Returns:
+            cosines_tuple (tuple): Contains directional cosines of marginal each rays. 
+        """
+        ray_north, ray_south, ray_east, ray_west = self.generate_marginal_rays(H_x, H_y, wavelength)
+
+        north_cosines = np.array([ray_north.L, ray_north.M, ray_north.N]).ravel()
+        south_cosines = np.array([ray_south.L, ray_south.M, ray_south.N]).ravel()
+        east_cosines  = np.array([ray_east.L,  ray_east.M,  ray_east.N]).ravel()
+        west_cosines  = np.array([ray_west.L,  ray_west.M,  ray_west.N]).ravel()
+
+        cosines_tuple = north_cosines, south_cosines, east_cosines, west_cosines
+        return cosines_tuple
+
+    def generate_chief_rays_cosines(self, wavelength):
+        """Generates directional cosines for all chief rays of each field.
+
+        Returns:
+            chief_ray_cosines_list (ndarray): Mx3 numpy array, containing M number of directional cosines for all axes (x, y ,z).
+            [
+            field 1 -> (ray_chief1.L, ray_chief1.M, ray_chief1.N),
+            field 2 -> (ray_chief1.L, ray_chief1.M, ray_chief1.N),
+                .               .
+                .               .
+                .               .
+            field M -> (ray_chiefM.L, ray_chiefM.M, ray_chiefM.N)
+            ]
+        """
+
+
+        coords = self.fields
+        chief_ray_cosines_list = []
+        for H_x, H_y in coords:
+            # Always pass the field values—even for 'angle' type.
+            ray_chief = self.optic.trace_generic(Hx=H_x, Hy=H_y, Px=0, Py=0, wavelength=wavelength)
+            chief_ray_cosines_list.append(np.array([ray_chief.L, ray_chief.M, ray_chief.N]).ravel())
+        chief_ray_cosines_list = np.array(chief_ray_cosines_list)
+        return chief_ray_cosines_list
+
+    def generate_chief_rays_centers(self, wavelength):
+        """Generates the position of each chief ray. It is used to find centers of the airy discs in the function _plot_field. 
+
+        Returns:
+            chief_ray_centers (ndarray): Contains the x, y coordinates of chief ray centers for each field.
+        """
+        coords = self.fields
+        chief_ray_centers = []
+        for H_x, H_y in coords:
+            # Always pass the field values—even for 'angle' type.
+            ray_chief = self.optic.trace_generic(Hx=H_x, Hy=H_y, Px=0, Py=0, wavelength=wavelength)
+            x, y = ray_chief.x, ray_chief.y
+            chief_ray_centers.append([x, y])
+
+        chief_ray_centers = np.array(chief_ray_centers) 
+        return chief_ray_centers
+
+    def airy_disc_x_y(self, wavelength):
+        """Generates the airy radius for each x-y axes, for each field.
+
+        Returns:
+            airy_rad_tuple (tuple): Contains x axis centers (ndarray), and y axis centers (ndarray).
+        """
+
+        chief_ray_cosines_list = self.generate_chief_rays_cosines(wavelength)
+
+        airy_rad_x_list = []
+        airy_rad_y_list = []
+
+        for idx, (H_x, H_y) in enumerate(self.fields):
+            # Get marginal rays for the current field
+            north_cos, south_cos, east_cos, west_cos = self.generate_marginal_rays_cosines(H_x, H_y, wavelength)
+
+            chief_cos = chief_ray_cosines_list[idx]
+
+            # Compute relative angles along x and y axes
+            rel_angle_north = self.angle_from_cosine(chief_cos, north_cos)
+            rel_angle_south = self.angle_from_cosine(chief_cos, south_cos)
+            rel_angle_east  = self.angle_from_cosine(chief_cos, east_cos)
+            rel_angle_west  = self.angle_from_cosine(chief_cos, west_cos)
+
+            avg_angle_x = (rel_angle_north + rel_angle_south) / 2
+            avg_angle_y = (rel_angle_east  + rel_angle_west)  / 2
+
+            N_w_x = self.f_number(n=1, theta=avg_angle_x)
+            N_w_y = self.f_number(n=1, theta=avg_angle_y)
+
+            airy_rad_x = self.airy_radius(N_w_x, wavelength)
+            airy_rad_y = self.airy_radius(N_w_y, wavelength)
+
+            airy_rad_x_list.append(airy_rad_x * 1e-3) # * 1e-3 to convert normalized spots
+            airy_rad_y_list.append(airy_rad_y * 1e-3) #
+
+        airy_rad_tuple = (airy_rad_x_list, airy_rad_y_list)
+
+        return airy_rad_tuple
+
     def centroid(self):
         """Centroid of each spot
 
@@ -219,7 +381,11 @@
         return [x, y, intensity]
 
     def _plot_field(self, ax, field_data, field, axis_lim,
-                    wavelengths, buffer=1.05):
+                    wavelengths, buffer=1.05, 
+                    add_airy_disk=False, field_index=None,
+                    airy_rad_x=None, airy_rad_y=None,
+                    real_centroid_x=None, real_centroid_y=None
+):
         """
         Plot the field data on the given axis.
 
@@ -244,10 +410,39 @@
             ax.scatter(x[mask], y[mask], s=10,
                        label=f'{wavelengths[k]:.4f} µm',
                        marker=markers[k % 3], alpha=0.7)
-            ax.axis('square')
-            ax.set_xlabel('X (mm)')
-            ax.set_ylabel('Y (mm)')
-            ax.set_xlim((-axis_lim*buffer, axis_lim*buffer))
-            ax.set_ylim((-axis_lim*buffer, axis_lim*buffer))
+
+        if add_airy_disk and field_index is not None:
+            # Draw ellipse ONLY for the current field_index
+            ellipse = patches.Ellipse(
+                (real_centroid_x , real_centroid_y),
+                #(0, 0),
+                width=2*airy_rad_y[field_index],    # diameter, not radius
+                height=2*airy_rad_x[field_index],
+                linestyle="--",
+                edgecolor="black",
+                fill=False,
+                linewidth=2
+            )
+            ax.add_patch(ellipse)
+
+            offset = abs(max(real_centroid_y, real_centroid_x))
+
+            # Find the maximum extent among the geometric spot radius and the airy disk radii.
+            max_airy_x = max(airy_rad_x)
+            max_airy_y = max(airy_rad_y)
+            max_extent = max(axis_lim, max_airy_x, max_airy_y)
+
+            # Apply a buffer to ensure the data fits nicely in the plot.
+            adjusted_axis_lim = (offset + max_extent) * buffer
+        else:
+            # Without the airy disk, just use the geometric spot radius with a buffer.
+            adjusted_axis_lim = axis_lim * buffer
+
+
+        ax.axis('square')
+        ax.set_xlabel('X (mm)')
+        ax.set_ylabel('Y (mm)')
+        ax.set_xlim((-adjusted_axis_lim, adjusted_axis_lim))
+        ax.set_ylim((-adjusted_axis_lim, adjusted_axis_lim))
         ax.set_title(f'Hx: {field[0]:.3f}, Hy: {field[1]:.3f}')
-        ax.grid(alpha=0.25)
+        ax.grid(alpha=0.25)