Fix/image space non air material (#76)

* Add final lens component to optical system if present

* Update path length calculations to include wavelength parameter

* chore: Add inverse transformation for global coordinate system in image surface

* feat: Enhance SurfaceFactory to support image surface creation and material configuration

* test: Add unit tests for image surface generation in SurfaceFactory

* fix: Update ImageSurface constructor to correctly assign material_post parameter

* refactor: Simplify surface creation logic in SurfaceFactory by removing image surface handling

* fix: Correct refractive index assignment in Wavefront class to use material_post parameter

* refactor: Remove obsolete image surface generation tests from TestSurfaceFactory

* fix: Update ImageSurface tests to correctly handle material_post parameter

* refactor: Remove redundant exit pupil location calculation in Paraxial class

* fix: Update path length and OPD image calculations in Wavefront tests to include wavelength parameter

* chore: Add TODO for tests on non-air object and image spaces in Paraxial tests

Author: HarrisonKramer

optiland/paraxial.py

@@ -172,17 +172,10 @@ def XPL(self):
             float: exit pupil location relative to the image surface
         """
         stop_index = self.surfaces.stop_index
-        num_surfaces = len(self.surfaces.surfaces)
-        if stop_index == num_surfaces-2:
-            positions = self.optic.surface_group.positions
-            loc_relative = positions[-2] - positions[-1]
-            return loc_relative[0]
-
         z_start = self.surfaces.positions[stop_index]
         wavelength = self.optic.primary_wavelength
         y, u = self._trace_generic(0.0, 0.1, z_start, wavelength,
                                    skip=stop_index+1)
-
         loc_relative = -y[-1] / u[-1]
         return loc_relative[0]
 

optiland/surfaces/image_surface.py

@@ -25,11 +25,11 @@ class ImageSurface(Surface):
     """
 
     def __init__(self, geometry: BaseGeometry, material_pre: BaseMaterial,
-                 aperture: BaseAperture = None):
+                 material_post: BaseMaterial, aperture: BaseAperture = None):
         super().__init__(
             geometry=geometry,
             material_pre=material_pre,
-            material_post=material_pre,
+            material_post=material_post,
             is_stop=False,
             aperture=aperture
         )
@@ -51,8 +51,13 @@ def _trace_paraxial(self, rays: ParaxialRays):
         t = -rays.z
         rays.propagate(t)
 
+        # inverse transform coordinate system
+        self.geometry.globalize(rays)
+
         self._record(rays)
 
+        return rays
+
     def _interact(self, rays):
         """
         Interacts rays with the surface.

optiland/surfaces/surface_factory.py

@@ -44,6 +44,7 @@ def __init__(self, surface_group):
         self._surface_group = surface_group
         self.last_thickness = 0
         self._use_absolute_cs = False
+        self._last_surface = None  # placeholder for last surface object
 
     def create_surface(self, surface_type, index, is_stop, material,
                        **kwargs):
@@ -318,7 +319,7 @@ def _configure_zernike_geometry(self, cs, **kwargs):
                                              norm_radius)
 
         return geometry
-    
+
     def _configure_material(self, index, material):
         """
         Configures the material for a surface based on the given index and
@@ -335,8 +336,16 @@ def _configure_material(self, index, material):
             tuple: A tuple containing the material before and after the
                 surface.
         """
+        # object surface
         if index == 0:
             material_pre = None
+
+        # image surface
+        elif (index == self._surface_group.num_surfaces
+              and self._last_surface is not None):
+            material_pre = self._last_surface.material_post
+
+        # intermediate surface
         else:
             previous_surface = self._surface_group.surfaces[index-1]
             material_pre = previous_surface.material_post

optiland/visualization/system.py

@@ -103,6 +103,10 @@ def _identify_components(self):
 
                     lens_surfaces = []
 
+            # add final lens, if any
+            if lens_surfaces:
+                self._add_component('lens', lens_surfaces)
+
     def _add_component(self, component_name, *args):
         """
         Adds a component to the list of components.

optiland/wavefront.py

@@ -86,7 +86,7 @@ def _generate_data(self, fields, wavelengths):
 
                 # Reference sphere center and radius
                 xc, yc, zc, R = self._get_reference_sphere(pupil_z)
-                opd_ref = self._get_path_length(xc, yc, zc, R)
+                opd_ref = self._get_path_length(xc, yc, zc, R, wavelength)
                 opd_ref = self._correct_tilt(field, opd_ref, x=0, y=0)
 
                 field_data.append(self._generate_field_data(field, wavelength,
@@ -116,7 +116,7 @@ def _generate_field_data(self, field, wavelength, opd_ref, xc, yc, zc, R):
         # trace distribution through pupil
         self.optic.trace(*field, wavelength, None, self.distribution)
         intensity = self.optic.surface_group.intensity[-1, :]
-        opd = self._get_path_length(xc, yc, zc, R)
+        opd = self._get_path_length(xc, yc, zc, R, wavelength)
         opd = self._correct_tilt(field, opd)
         return (opd_ref - opd) / (wavelength * 1e-3), intensity
 
@@ -159,7 +159,7 @@ def _get_reference_sphere(self, pupil_z):
 
         return xc, yc, zc, R
 
-    def _get_path_length(self, xc, yc, zc, r):
+    def _get_path_length(self, xc, yc, zc, r, wavelength):
         """
         Calculates the optical path difference.
 
@@ -168,12 +168,13 @@ def _get_path_length(self, xc, yc, zc, r):
             yc (float): The y-coordinate of the reference sphere center.
             zc (float): The z-coordinate of the reference sphere center.
             r (float): The radius of the reference sphere.
+            wavelength (float): The wavelength of the light.
 
         Returns:
             float: The optical path difference.
         """
         opd = self.optic.surface_group.opd[-1, :]
-        return opd - self._opd_image_to_xp(xc, yc, zc, r)
+        return opd - self._opd_image_to_xp(xc, yc, zc, r, wavelength)
 
     def _correct_tilt(self, field, opd, x=None, y=None):
         """
@@ -203,7 +204,7 @@ def _correct_tilt(self, field, opd, x=None, y=None):
                                (1 - y) * np.sin(np.radians(y_tilt)) * EPD / 2)
         return opd - tilt_correction
 
-    def _opd_image_to_xp(self, xc, yc, zc, R):
+    def _opd_image_to_xp(self, xc, yc, zc, R, wavelength):
         """
         Finds propagation distance from image plane to reference sphere.
 
@@ -212,6 +213,7 @@ def _opd_image_to_xp(self, xc, yc, zc, R):
             yc (float): The y-coordinate of the reference sphere center.
             zc (float): The z-coordinate of the reference sphere center.
             R (float): The radius of the reference sphere.
+            wavelength (float): The wavelength of the light.
 
         Returns:
             float: Propagation distance from image plane to reference sphere.
@@ -232,7 +234,10 @@ def _opd_image_to_xp(self, xc, yc, zc, R):
         d = b ** 2 - 4 * a * c
         t = (-b - np.sqrt(d)) / (2 * a)
         t[t < 0] = (-b[t < 0] + np.sqrt(d[t < 0])) / (2 * a[t < 0])
-        return t
+
+        # refractive index in image space
+        n = self.optic.image_surface.material_post.n(wavelength)
+        return n * t
 
 
 class OPDFan(Wavefront):

tests/test_image_surface.py

@@ -12,24 +12,26 @@ def setup_image_surface():
     cs = CoordinateSystem()
     geometry = StandardGeometry(cs, 100)
     material_pre = IdealMaterial(1, 0)
+    material_post = IdealMaterial(1, 0)
     image_surface = ImageSurface(
         geometry=geometry,
         material_pre=material_pre,
+        material_post=material_post,
         aperture=None
     )
-    return image_surface, geometry, material_pre
+    return image_surface, geometry, material_pre, material_post
 
 
 def test_initialization(setup_image_surface):
-    image_surface, geometry, material_pre = setup_image_surface
+    image_surface, geometry, material_pre, material_post = setup_image_surface
     assert image_surface.geometry == geometry
     assert image_surface.material_pre == material_pre
-    assert image_surface.material_post == material_pre
+    assert image_surface.material_post == material_post
     assert not image_surface.is_stop
 
 
 def test_trace_paraxial(setup_image_surface):
-    image_surface, _, _ = setup_image_surface
+    image_surface, _, _, _ = setup_image_surface
     y = np.array([1])
     u = np.array([0])
     z = np.array([-10])
@@ -39,7 +41,7 @@ def test_trace_paraxial(setup_image_surface):
 
 
 def test_interact(setup_image_surface):
-    image_surface, _, _ = setup_image_surface
+    image_surface, _, _, _ = setup_image_surface
     x = np.random.rand(10)
     rays = RealRays(x, x, x, x, x, x, x, x)
     modified_rays = image_surface._interact(rays)

tests/test_paraxial.py

@@ -23,6 +23,7 @@
     PetzvalLens,
     Telephoto
 )
+# TODO: add tests for non-air object and image spaces
 
 
 def get_optic_data():

tests/test_wavefront.py

@@ -67,7 +67,7 @@ def test_get_path_length(self):
         w = wavefront.Wavefront(optic)
         w._trace_chief_ray((0, 0), 0.55)
         xc, yc, zc, R = w._get_reference_sphere(pupil_z=100)
-        path_length = w._get_path_length(xc, yc, zc, R)
+        path_length = w._get_path_length(xc, yc, zc, R, 0.55)
         assert np.allclose(path_length, np.array([34.84418309]))
 
     def test_correct_tilt(self):
@@ -86,7 +86,7 @@ def test_opd_image_to_xp(self):
         w = wavefront.Wavefront(optic)
         w._trace_chief_ray((0, 0), 0.55)
         xc, yc, zc, R = w._get_reference_sphere(pupil_z=100)
-        t = w._opd_image_to_xp(xc, yc, zc, R)
+        t = w._opd_image_to_xp(xc, yc, zc, R, 0.55)
         assert np.allclose(t, np.array([39.454938]))