refactor: clean up Jones functionalities (#507)

Author: HarrisonKramer

optiland/jones.py

@@ -117,38 +117,21 @@ def calculate_matrix(
         return jones_matrix
 
 
-class JonesPolarizerH(BaseJones):
-    """Class representing the Jones matrix for a horizontal polarizer."""
-
-    def calculate_matrix(
-        self,
-        rays: RealRays,
-        reflect: bool = False,
-        aoi: be.ndarray = None,
-    ):
-        """Calculate the Jones matrix for the given rays.
-
-        Args:
-            rays (RealRays): Object representing the rays.
-            reflect (bool, optional): Indicates whether the rays are reflected
-                or not. Defaults to False.
-            aoi (be.ndarray, optional): Array representing the angle of
-                incidence. Defaults to None.
-
-        Returns:
-            be.ndarray: The calculated Jones matrix.
-
-        """
-        jones_matrix = be.to_complex(be.zeros((be.size(rays.x), 3, 3)))
-        jones_matrix[:, 0, 0] = 1
-        jones_matrix[:, 1, 1] = 0
-        jones_matrix[:, 2, 2] = 1
+class ConstantJones(BaseJones):
+    """Base class for constant Jones matrices.
 
-        return jones_matrix
+    Args:
+        j00 (complex): The (0, 0) element of the Jones matrix.
+        j01 (complex): The (0, 1) element of the Jones matrix.
+        j10 (complex): The (1, 0) element of the Jones matrix.
+        j11 (complex): The (1, 1) element of the Jones matrix.
+    """
 
-
-class JonesPolarizerV(BaseJones):
-    """Class representing the Jones matrix for a vertical polarizer."""
+    def __init__(self, j00: complex, j01: complex, j10: complex, j11: complex):
+        self.j00 = j00
+        self.j01 = j01
+        self.j10 = j10
+        self.j11 = j11
 
     def calculate_matrix(
         self,
@@ -170,139 +153,55 @@ def calculate_matrix(
 
         """
         jones_matrix = be.to_complex(be.zeros((be.size(rays.x), 3, 3)))
-        jones_matrix[:, 0, 0] = 0
-        jones_matrix[:, 1, 1] = 1
+        jones_matrix[:, 0, 0] = self.j00
+        jones_matrix[:, 0, 1] = self.j01
+        jones_matrix[:, 1, 0] = self.j10
+        jones_matrix[:, 1, 1] = self.j11
         jones_matrix[:, 2, 2] = 1
 
         return jones_matrix
 
 
-class JonesPolarizerL45(BaseJones):
-    """Class representing the Jones matrix for a linear polarizer at 45 degrees."""
-
-    def calculate_matrix(
-        self,
-        rays: RealRays,
-        reflect: bool = False,
-        aoi: be.ndarray = None,
-    ):
-        """Calculate the Jones matrix for the given rays.
-
-        Args:
-            rays (RealRays): Object representing the rays.
-            reflect (bool, optional): Indicates whether the rays are reflected
-                or not. Defaults to False.
-            aoi (be.ndarray, optional): Array representing the angle of
-                incidence. Defaults to None.
+class JonesPolarizerH(ConstantJones):
+    """Class representing the Jones matrix for a horizontal polarizer."""
 
-        Returns:
-            be.ndarray: The calculated Jones matrix.
+    def __init__(self):
+        super().__init__(1, 0, 0, 0)
 
-        """
-        jones_matrix = be.to_complex(be.zeros((be.size(rays.x), 3, 3)))
-        jones_matrix[:, 0, 0] = 0.5
-        jones_matrix[:, 0, 1] = 0.5
-        jones_matrix[:, 1, 0] = 0.5
-        jones_matrix[:, 1, 1] = 0.5
-        jones_matrix[:, 2, 2] = 1
 
-        return jones_matrix
+class JonesPolarizerV(ConstantJones):
+    """Class representing the Jones matrix for a vertical polarizer."""
 
+    def __init__(self):
+        super().__init__(0, 0, 0, 1)
 
-class JonesPolarizerL135(BaseJones):
-    """Class representing the Jones matrix for a linear polarizer at 135 degrees."""
 
-    def calculate_matrix(
-        self,
-        rays: RealRays,
-        reflect: bool = False,
-        aoi: be.ndarray = None,
-    ):
-        """Calculate the Jones matrix for the given rays.
+class JonesPolarizerL45(ConstantJones):
+    """Class representing the Jones matrix for a linear polarizer at 45 degrees."""
 
-        Args:
-            rays (RealRays): Object representing the rays.
-            reflect (bool, optional): Indicates whether the rays are reflected
-                or not. Defaults to False.
-            aoi (be.ndarray, optional): Array representing the angle of
-                incidence. Defaults to None.
+    def __init__(self):
+        super().__init__(0.5, 0.5, 0.5, 0.5)
 
-        Returns:
-            be.ndarray: The calculated Jones matrix.
 
-        """
-        jones_matrix = be.to_complex(be.zeros((be.size(rays.x), 3, 3)))
-        jones_matrix[:, 0, 0] = 0.5
-        jones_matrix[:, 0, 1] = -0.5
-        jones_matrix[:, 1, 0] = -0.5
-        jones_matrix[:, 1, 1] = 0.5
-        jones_matrix[:, 2, 2] = 1
+class JonesPolarizerL135(ConstantJones):
+    """Class representing the Jones matrix for a linear polarizer at 135 degrees."""
 
-        return jones_matrix
+    def __init__(self):
+        super().__init__(0.5, -0.5, -0.5, 0.5)
 
 
-class JonesPolarizerRCP(BaseJones):
+class JonesPolarizerRCP(ConstantJones):
     """Class representing the Jones matrix for a right circular polarizer."""
 
-    def calculate_matrix(
-        self,
-        rays: RealRays,
-        reflect: bool = False,
-        aoi: be.ndarray = None,
-    ):
-        """Calculate the Jones matrix for the given rays.
-
-        Args:
-            rays (RealRays): Object representing the rays.
-            reflect (bool, optional): Indicates whether the rays are reflected
-                or not. Defaults to False.
-            aoi (be.ndarray, optional): Array representing the angle of
-                incidence. Defaults to None.
-
-        Returns:
-            be.ndarray: The calculated Jones matrix.
-
-        """
-        jones_matrix = be.to_complex(be.zeros((be.size(rays.x), 3, 3)))
-        jones_matrix[:, 0, 0] = 0.5
-        jones_matrix[:, 0, 1] = 1j * 0.5
-        jones_matrix[:, 1, 0] = -1j * 0.5
-        jones_matrix[:, 1, 1] = 0.5
-        jones_matrix[:, 2, 2] = 1
-
-        return jones_matrix
+    def __init__(self):
+        super().__init__(0.5, 1j * 0.5, -1j * 0.5, 0.5)
 
 
-class JonesPolarizerLCP(BaseJones):
+class JonesPolarizerLCP(ConstantJones):
     """Class representing the Jones matrix for a left circular polarizer."""
 
-    def calculate_matrix(
-        self,
-        rays: RealRays,
-        reflect: bool = False,
-        aoi: be.ndarray = None,
-    ):
-        """Calculate the Jones matrix for the given rays.
-
-        Args:
-            rays (RealRays): Object representing the rays.
-            reflect (bool, optional): Indicates whether the rays are reflected
-                or not. Defaults to False.
-            aoi (be.ndarray, optional): Array representing the angle of
-                incidence. Defaults to None.
-
-        Returns:
-            be.ndarray: The calculated Jones matrix.
-
-        """
-        jones_matrix = be.to_complex(be.zeros((be.size(rays.x), 3, 3)))
-        jones_matrix[:, 0, 0] = 0.5
-        jones_matrix[:, 0, 1] = -1j * 0.5
-        jones_matrix[:, 1, 0] = 1j * 0.5
-        jones_matrix[:, 1, 1] = 0.5
-        jones_matrix[:, 2, 2] = 1
-
-        return jones_matrix
+    def __init__(self):
+        super().__init__(0.5, -1j * 0.5, 1j * 0.5, 0.5)
 
 
 class JonesLinearDiattenuator(BaseJones):

optiland/rays/polarization_state.py

@@ -101,44 +101,26 @@ def create_polarization(pol_type: str):
     """
     if pol_type == "unpolarized":
         return PolarizationState(is_polarized=False)
-    if pol_type == "H":
-        Ex = 1
-        Ey = 0
-        phase_x = 0
-        phase_y = 0
-    elif pol_type == "V":
-        Ex = 0
-        Ey = 1
-        phase_x = 0
-        phase_y = 0
-    elif pol_type == "L+45":
-        Ex = 1
-        Ey = 1
-        phase_x = 0
-        phase_y = 0
-    elif pol_type == "L-45":
-        Ex = 1
-        Ey = -1
-        phase_x = 0
-        phase_y = 0
-    elif pol_type == "RCP":
-        Ex = be.sqrt(2) / 2
-        Ey = be.sqrt(2) / 2
-        phase_x = 0
-        phase_y = -be.pi / 2
-    elif pol_type == "LCP":
-        Ex = be.sqrt(2) / 2
-        Ey = be.sqrt(2) / 2
-        phase_x = 0
-        phase_y = be.pi / 2
-    else:
-        raise ValueError(
-            "Invalid polarization type. Must be H, V, L+45, L-45, RCP or LCP.",
+
+    pol_map = {
+        "H": (1.0, 0.0, 0.0, 0.0),
+        "V": (0.0, 1.0, 0.0, 0.0),
+        "L+45": (1.0, 1.0, 0.0, 0.0),
+        "L-45": (1.0, -1.0, 0.0, 0.0),
+        "RCP": (be.sqrt(2) / 2, be.sqrt(2) / 2, 0.0, -be.pi / 2),
+        "LCP": (be.sqrt(2) / 2, be.sqrt(2) / 2, 0.0, be.pi / 2),
+    }
+
+    if pol_type in pol_map:
+        Ex, Ey, phase_x, phase_y = pol_map[pol_type]
+        return PolarizationState(
+            is_polarized=True,
+            Ex=Ex,
+            Ey=Ey,
+            phase_x=phase_x,
+            phase_y=phase_y,
         )
-    return PolarizationState(
-        is_polarized=True,
-        Ex=Ex,
-        Ey=Ey,
-        phase_x=phase_x,
-        phase_y=phase_y,
+
+    raise ValueError(
+        "Invalid polarization type. Must be H, V, L+45, L-45, RCP or LCP.",
     )