Add Jones Pupil Analysis (#432)

* feat: Add JonesPupil analysis class

- Added `JonesPupil` class to `optiland/analysis/jones_pupil.py`.
- Exposed `JonesPupil` in `optiland/analysis/__init__.py`.
- Added `JonesPupil` to `optiland/analysis/jones_pupil.py`.
- Added unit tests in `tests/analysis/test_jones_pupil.py`.
- Jones pupil analysis generates a grid of rays and computes the Jones matrix elements (Jxx, Jxy, Jyx, Jyy) by projecting the 3x3 global polarization matrix onto a local basis defined by the ray direction.
- Visualization displays the real and imaginary parts of the Jones matrix elements.

* Fix formatting and linting issues in JonesPupil analysis.

- Removed unused `Literal` import.
- Sorted imports in `TYPE_CHECKING` block.
- Combined nested `if` statements for clarity.
- Removed unused variables: `num_fields`, `num_wavelengths`, `x_in`, `y_in`.
- Resolved long line length warning.
- Applied `ruff` auto-fixes.

* feat: Add Jones Pupil analysis module, example notebook, and documentation.

* fix: update tests after Jones pupil update

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: Kramer <[email protected]>

Author: google-labs-jules[bot]

README.md

@@ -136,7 +136,7 @@ Optiland is continually evolving to provide new functionalities for optical desi
 - [ ] **Multi-Path Sequential Ray Tracing**
 - [x] **Multiple Configurations (Zoom Lenses)**
 - [ ] **Thin Film Design and Optimization** 
-- [ ] **Jones Pupils**
+- [x] **Jones Pupils**
 - [ ] **Additional Freeforms (Superconic, etc.)**
 - [x] **Image Simulation**
 - [ ] **Interferogram Analysis**

docs/api/analysis/optiland.analysis.jones_pupil.rst

@@ -0,0 +1,11 @@
+optiland.analysis.jones\_pupil
+==============================
+
+.. automodule:: optiland.analysis.jones_pupil
+
+   
+   .. rubric:: Classes
+
+   .. autosummary::
+   
+      JonesPupil

docs/api/api_analysis.rst

@@ -18,6 +18,7 @@ encircled energy, field curvature, distortion, etc.
    analysis.intensity
    analysis.image_simulation
    analysis.irradiance
+   analysis.jones_pupil
    analysis.pupil_aberration
    analysis.ray_fan
    analysis.rms_vs_field

docs/functionalities.rst

@@ -24,6 +24,8 @@ Analysis Tools
   Perform precise ray-based evaluations for both idealized and physically realistic systems.
 - **Polarization Ray Tracing**:
   Model vectorial light propagation, including polarization effects and birefringent materials.
+- **Jones Pupil Analysis**:
+  Visualize spatially resolved Jones matrix elements at the exit pupil to assess polarization properties.
 - **Comprehensive Optical Analysis**:
   Generate spot diagrams, ray aberration fans, OPD maps, distortion plots, and more.
 - **Wavefront Analysis**:

docs/gallery/analysis.rst

@@ -16,6 +16,7 @@ This section contains examples of typical analysis tasks that can be performed w
     analysis/y_ybar
     analysis/rms_spot_size_vs_field
     analysis/rms_wavefront_error_vs_field
+    analysis/jones_pupil
     analysis/pupil_aberration
     analysis/irradiance
     analysis/through_focus_spot_diagram

docs/gallery/analysis/jones_pupil.ipynb

@@ -14,3 +14,4 @@
 from .intensity import RadiantIntensity
 from .through_focus_mtf import ThroughFocusMTF
 from .through_focus_spot_diagram import ThroughFocusSpotDiagram
+from .jones_pupil import JonesPupil

optiland/analysis/__init__.py

@@ -0,0 +1,228 @@
+"""Jones Pupil Analysis
+
+This module provides a Jones pupil analysis for optical systems.
+
+Kramer Harrison, 2025
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+import optiland.backend as be
+from optiland.analysis.base import BaseAnalysis
+from optiland.rays import PolarizationState
+
+if TYPE_CHECKING:
+    from matplotlib.axes import Axes
+    from matplotlib.figure import Figure
+
+    from optiland.optic import Optic
+
+
+class JonesPupil(BaseAnalysis):
+    """Generates and plots Jones pupil maps.
+
+    This class computes the spatially resolved Jones matrix at the exit pupil
+    (or image plane) as a function of normalized pupil coordinates. It visualizes
+    the real and imaginary parts of the Jones matrix elements (Jxx, Jxy, Jyx, Jyy).
+
+    Attributes:
+        optic: Instance of the optic object to be assessed.
+    Attributes:
+        optic: Instance of the optic object to be assessed.
+        field: Field at which data is generated (Hx, Hy).
+        wavelengths: Wavelengths at which data is generated.
+        grid_size: The side length of the square grid of rays (NxN).
+        data: Contains Jones matrix data in a list, ordered by wavelength.
+    """
+
+    def __init__(
+        self,
+        optic: Optic,
+        field: tuple[float, float] = (0, 0),
+        wavelengths: str | list = "all",
+        grid_size: int = 65,
+    ):
+        """Initializes the JonesPupil analysis.
+
+        Args:
+            optic: An instance of the optic object to be assessed.
+            field: The normalized field coordinates (Hx, Hy) at which to
+                generate data. Defaults to (0, 0).
+            wavelengths: Wavelengths at which to generate data. If 'all', all
+                defined wavelengths are used. Defaults to "all".
+            grid_size: The number of points along one dimension of the pupil grid.
+                Defaults to 65.
+        """
+        self.field = field
+        self.grid_size = grid_size
+        super().__init__(optic, wavelengths)
+
+    def view(
+        self,
+        fig_to_plot_on: Figure | None = None,
+        figsize: tuple[float, float] = (16, 8),
+    ) -> tuple[Figure, list[Axes]]:
+        """Displays the Jones pupil plots.
+
+        Args:
+            fig_to_plot_on: An existing Matplotlib figure to plot on. If None,
+                a new figure is created. Defaults to None.
+            figsize: The figure size for the output window. Defaults to (16, 8).
+
+        Returns:
+            A tuple containing the Matplotlib figure and a list of its axes.
+        """
+        # Select primary wavelength index
+        wl_idx = 0
+        if self.optic.primary_wavelength in self.wavelengths:
+            wl_idx = self.wavelengths.index(self.optic.primary_wavelength)
+
+        data_fw = self.data[wl_idx]
+
+        if fig_to_plot_on:
+            fig = fig_to_plot_on
+            fig.clear()
+        else:
+            fig = plt.figure(figsize=figsize)
+
+        # 2 rows (Real, Imag), 4 columns (Jxx, Jxy, Jyx, Jyy)
+        axs = fig.subplots(2, 4, sharex=True, sharey=True)
+
+        # Elements to plot
+        elements = [
+            ("Jxx", data_fw["J"][:, 0, 0]),
+            ("Jxy", data_fw["J"][:, 0, 1]),
+            ("Jyx", data_fw["J"][:, 1, 0]),
+            ("Jyy", data_fw["J"][:, 1, 1]),
+        ]
+
+        px = be.to_numpy(data_fw["Px"]).reshape(self.grid_size, self.grid_size)
+        py = be.to_numpy(data_fw["Py"]).reshape(self.grid_size, self.grid_size)
+        mask = px**2 + py**2 <= 1.0
+
+        for col, (name, values) in enumerate(elements):
+            val_np = be.to_numpy(values).reshape(self.grid_size, self.grid_size)
+            val_np[~mask] = np.nan
+
+            # Real part
+            ax_real = axs[0, col]
+            im_real = ax_real.pcolormesh(
+                px, py, np.real(val_np), shading="nearest", cmap="viridis"
+            )
+            ax_real.set_title(f"Re({name})")
+            ax_real.set_aspect("equal")
+            fig.colorbar(im_real, ax=ax_real, fraction=0.046, pad=0.04)
+
+            # Imag part
+            ax_imag = axs[1, col]
+            im_imag = ax_imag.pcolormesh(
+                px, py, np.imag(val_np), shading="nearest", cmap="viridis"
+            )
+            ax_imag.set_title(f"Im({name})")
+            ax_imag.set_aspect("equal")
+            fig.colorbar(im_imag, ax=ax_imag, fraction=0.046, pad=0.04)
+
+        # Labels
+        for ax in axs[:, 0]:
+            ax.set_ylabel("Py")
+        for ax in axs[-1, :]:
+            ax.set_xlabel("Px")
+
+        field_val = self.field
+        wl_val = self.wavelengths[wl_idx]
+        fig.suptitle(f"Jones Pupil - Field: {field_val}, Wavelength: {wl_val:.4f} µm")
+        fig.tight_layout()
+
+        return fig, fig.get_axes()
+
+    def _generate_data(self):
+        """Generates Jones matrix data for all fields and wavelengths."""
+        # Generate pupil grid
+        x = be.linspace(-1.0, 1.0, self.grid_size)
+        y = be.linspace(-1.0, 1.0, self.grid_size)
+        Px_grid, Py_grid = be.meshgrid(x, y)
+        Px = Px_grid.flatten()
+        Py = Py_grid.flatten()
+
+        data = []
+        Hx, Hy = self.field
+        for wl in self.wavelengths:
+            data.append(self._generate_single_data(Hx, Hy, Px, Py, wl))
+
+        return data
+
+    def _generate_single_data(self, Hx, Hy, Px, Py, wavelength):
+        """Generates data for a single field and wavelength configuration."""
+
+        # Handle polarization state
+        original_pol = self.optic.polarization
+        if original_pol == "ignore":
+            # Temporarily enable polarization to get PolarizedRays
+            self.optic.set_polarization(PolarizationState())
+
+        try:
+            rays = self.optic.trace_generic(
+                Hx=Hx, Hy=Hy, Px=Px, Py=Py, wavelength=wavelength
+            )
+        finally:
+            if original_pol == "ignore":
+                self.optic.set_polarization("ignore")
+
+        if not hasattr(rays, "p"):
+            # Fallback if rays are not polarized (should not happen w/ check above)
+            raise RuntimeError("Ray tracing did not return polarized rays.")
+
+        # Ray direction vectors (normalized)
+        k = be.stack([rays.L, rays.M, rays.N], axis=1)
+        # Normalize k (should be already, but to be safe)
+        k_norm = be.linalg.norm(k, axis=1)
+        k = k / be.unsqueeze_last(k_norm)
+
+        # Construct local basis vectors (Standard Polar Projection / Dipole-like)
+        # v ~ Y-axis: perpendicular to k and X=[1,0,0]
+        x_axis = be.array([1.0, 0.0, 0.0])
+        # Broadcast x_axis to match k shape
+        x_axis = be.broadcast_to(x_axis, k.shape)
+
+        v = be.cross(k, x_axis)
+        v_norm = be.linalg.norm(v, axis=1)
+
+        # Avoid division by zero
+        v = v / be.unsqueeze_last(v_norm + 1e-15)
+
+        # u ~ X-axis: perpendicular to v and k
+        u = be.cross(v, k)
+        u_norm = be.linalg.norm(u, axis=1)
+        # Avoid division by zero
+        u = u / be.unsqueeze_last(u_norm + 1e-15)
+
+        # Project global P onto local basis (u, v)
+        # Jxx = u . (P . x_in)
+        # Jxy = u . (P . y_in)
+        # Jyx = v . (P . x_in)
+        # Jyy = v . (P . y_in)
+
+        # p has shape (N, 3, 3)
+        # P . x_in is simply the first column of p
+        # P . y_in is simply the second column of p
+
+        P_x_in = rays.p[:, :, 0]  # Shape (N, 3)
+        P_y_in = rays.p[:, :, 1]  # Shape (N, 3)
+
+        # Dot products
+        Jxx = be.sum(u * P_x_in, axis=1)
+        Jxy = be.sum(u * P_y_in, axis=1)
+        Jyx = be.sum(v * P_x_in, axis=1)
+        Jyy = be.sum(v * P_y_in, axis=1)
+
+        # Stack into (N, 2, 2)
+        row1 = be.stack([Jxx, Jxy], axis=1)
+        row2 = be.stack([Jyx, Jyy], axis=1)
+        J = be.stack([row1, row2], axis=1)
+
+        return {"Px": Px, "Py": Py, "J": J}

optiland/analysis/jones_pupil.py

@@ -60,6 +60,8 @@ def max_y_field(self):
     @property
     def max_field(self):
         """float: Maximum radial field value."""
+        if not self.fields:
+            return 0.0
         return be.max(be.sqrt(self.x_fields**2 + self.y_fields**2))
 
     @property

optiland/fields/field_group.py

@@ -0,0 +1,93 @@
+
+import pytest
+import numpy as np
+import matplotlib.pyplot as plt
+import optiland.backend as be
+from optiland.samples.objectives import CookeTriplet
+from optiland.analysis.jones_pupil import JonesPupil
+from optiland.rays import PolarizationState
+
+
+def test_jones_pupil_initialization(set_test_backend):
+    optic = CookeTriplet()
+    optic.set_polarization("ignore") # Default state for test
+    jp = JonesPupil(optic)
+    assert jp.optic == optic
+    assert jp.grid_size == 65
+    assert jp.field == (0, 0)
+    assert jp.wavelengths is not None
+
+def test_jones_pupil_generate_data(set_test_backend):
+    optic = CookeTriplet()
+    optic.set_polarization("ignore") # Default state for test
+    jp = JonesPupil(optic, grid_size=5)
+    data = jp.data
+
+    # Check structure: list of wavelengths -> dict
+    assert isinstance(data, list)
+    assert len(data) == len(jp.wavelengths)
+    
+    # Check data for first wavelength
+    single_data = data[0]
+    assert isinstance(single_data, dict)
+    assert "Px" in single_data
+    assert "Py" in single_data
+    assert "J" in single_data
+
+    # Check shapes
+    num_rays = 5 * 5
+    assert single_data["Px"].shape == (num_rays,)
+    assert single_data["J"].shape == (num_rays, 2, 2)
+
+    J = single_data["J"]
+    Jxx = J[:, 0, 0]
+    Jxy = J[:, 0, 1]
+
+    # Center ray (index 12 for 5x5 grid)
+    center_idx = num_rays // 2
+
+    # Use backend agnostic checks
+    val_xx = be.to_numpy(Jxx[center_idx])
+    val_xy = be.to_numpy(Jxy[center_idx])
+
+    assert np.abs(val_xx) > 0.5
+    assert np.abs(val_xy) < 0.1
+
+def test_jones_pupil_polarization_handling(set_test_backend):
+    optic = CookeTriplet()
+    optic.set_polarization("ignore") # Default state for test
+    # Ensure it works even if optic polarization is 'ignore'
+    optic.set_polarization("ignore")
+    jp = JonesPupil(optic, grid_size=3)
+    # Trace happens in __init__ / first access to data
+
+    # Ensure optic state is restored
+    assert optic.polarization == "ignore"
+
+    # Data should be valid
+    assert jp.data is not None
+
+def test_jones_pupil_view(set_test_backend):
+    optic = CookeTriplet()
+    optic.set_polarization("ignore") # Default state for test
+    jp = JonesPupil(optic, grid_size=5)
+    fig, axs = jp.view()
+
+    assert fig is not None
+    # 2 rows, 4 columns = 8 axes + 8 colorbars = 16 axes
+    assert len(axs) == 16
+
+    # Clean up
+    plt.close(fig)
+
+def test_jones_pupil_view_custom_field(set_test_backend):
+    optic = CookeTriplet()
+    optic.set_polarization("ignore") # Default state for test
+    # Test with off-axis field
+    jp = JonesPupil(optic, field=(0, 1.0), grid_size=5)
+    data = jp.data
+    assert len(data) == len(jp.wavelengths)
+    
+    # Just verify valid execution
+    single_data = data[0]
+    assert "J" in single_data