tests: add tests for fileio, image sim, & variables

Author: HarrisonKramer

tests/analysis/test_image_simulation.py

@@ -0,0 +1,73 @@
+
+import pytest
+import numpy as np
+import optiland.backend as be
+from optiland.analysis.image_simulation import ImageSimulationEngine, DistortionWarper, PSFBasisGenerator
+from optiland.samples.objectives import CookeTriplet
+
+class TestImageSimulation:
+    @pytest.fixture
+    def optic(self):
+        return CookeTriplet()
+
+    @pytest.fixture
+    def source_image(self):
+        # Create a small dummy image (green square in black background)
+        img = np.zeros((32, 32, 3), dtype=np.float32)
+        img[10:22, 10:22, 1] = 1.0
+        return img
+
+    def test_engine_init(self, optic, source_image):
+        engine = ImageSimulationEngine(optic, source_image)
+        assert engine.source_image.shape == (3, 32, 32) # Transposed to (C, H, W)
+        assert engine.config is not None
+
+    def test_engine_run(self, optic, source_image):
+        # Use very loose config for speed
+        config = {
+            "psf_grid_shape": (3, 3),
+            "psf_size": 32,
+            "num_rays": 32,
+            "n_components": 2,
+            "oversample": 1,
+            "wavelengths": [0.55] # Mono
+        }
+        engine = ImageSimulationEngine(optic, source_image, config=config)
+        result = engine.run()
+        
+        # Let's check typical RGB case
+        config["wavelengths"] = [0.65, 0.55, 0.45]
+        engine = ImageSimulationEngine(optic, source_image, config=config)
+        result = engine.run()
+        
+        assert result.shape == (32, 32, 3)
+        assert be.max(result) > 0 # Should have some signal
+
+    def test_distortion_warper(self, optic):
+        warper = DistortionWarper(optic)
+        H, W = 32, 32
+        dist_map = warper.generate_distortion_map(wavelength=0.55, image_shape=(H, W))
+        assert dist_map.shape == (1, H, W, 2)
+        
+        # Warp a dummy image
+        img = be.ones((1, H, W))
+        warped = warper.warp_image(img, dist_map)
+        assert warped.shape == (1, H, W)
+
+    def test_psf_basis_generator(self, optic):
+        gen = PSFBasisGenerator(
+            optic, 
+            wavelength=0.55, 
+            grid_shape=(3, 3), 
+            num_rays=32, 
+            psf_grid_size=32
+        )
+        eigen_psfs, coeffs, mean_psf = gen.generate_basis(n_components=2)
+        
+        assert eigen_psfs.shape == (2, 32, 32)
+        assert coeffs.shape == (2, 3, 3)
+        assert mean_psf.shape == (32, 32)
+        
+        resized_coeffs = gen.resize_coefficient_map(coeffs, (64, 64))
+        assert resized_coeffs.shape == (2, 64, 64)
+

tests/test_fileio.py

@@ -12,9 +12,11 @@
     ZemaxFileSourceHandler,
     load_zemax_file,
 )
+from optiland.fileio.converters import ZemaxToOpticConverter
 from optiland.materials import Material
 from optiland.optic import Optic
 from optiland.samples.objectives import HeliarLens
+from optiland.geometries import ToroidalGeometry
 import optiland.backend as be
 from .utils import assert_allclose
 
@@ -288,3 +290,148 @@ def test_remove_surface_after_load(set_test_backend, tmp_path):
     expected_positions_after_object = be.array([0.0, 10.0, 30.0])
 
     assert_allclose(positions[1:], expected_positions_after_object)
+
+
+class TestZemaxToOpticConverterExtended:
+    def test_configure_aperture_floating_stop_no_diameter(self):
+        zemax_data = {
+            "surfaces": {
+                0: {"type": "standard", "is_stop": True, "radius": 0.0, "conic": 0.0, "thickness": 0.0, "material": "Air"},
+            },
+            "aperture": {"floating_stop": True},
+            "fields": {"type": "angle", "x": [0], "y": [0]},
+            "wavelengths": {"primary_index": 0, "data": [0.55]},
+        }
+        converter = ZemaxToOpticConverter(zemax_data)
+        converter.optic = Optic()
+        converter._configure_surfaces()
+        
+        with pytest.raises(ValueError, match="Floating stop aperture specified but no stop diameter found"):
+            converter._configure_aperture()
+
+    def test_configure_aperture_no_valid_type(self):
+        zemax_data = {
+            "surfaces": {},
+            "aperture": {"floating_stop": False},
+            "fields": {"type": "angle", "x": [0], "y": [0]},
+            "wavelengths": {"primary_index": 0, "data": [0.55]},
+        }
+        converter = ZemaxToOpticConverter(zemax_data)
+        converter.optic = Optic()
+        
+        with pytest.raises(ValueError, match="No valid aperture type found"):
+            converter._configure_aperture()
+
+    def test_configure_surface_coefficients_unsupported_type(self):
+        zemax_data = {
+            "surfaces": {
+                0: {"type": "unsupported_surface_type", "radius": 10},
+            },
+            "aperture": {"EPD": 10},
+            "fields": {"type": "angle", "x": [0], "y": [0]},
+            "wavelengths": {"primary_index": 0, "data": [0.55]},
+        }
+        converter = ZemaxToOpticConverter(zemax_data)
+        
+        with pytest.raises(ValueError, match="Unsupported Zemax surface type"):
+            converter._configure_surface_coefficients({"type": "unsupported_surface_type"})
+
+    def test_configure_fields_vignette_warning(self, capsys):
+        zemax_data = {
+            "surfaces": {},
+            "aperture": {"EPD": 10},
+            "fields": {
+                "type": "angle",
+                "x": [0],
+                "y": [0],
+                "vignette_decenter_x": [0.1],
+                "vignette_decenter_y": [0.0]
+            },
+            "wavelengths": {"primary_index": 0, "data": [0.55]},
+        }
+        converter = ZemaxToOpticConverter(zemax_data)
+        converter.optic = Optic()
+        converter._configure_fields()
+        
+        captured = capsys.readouterr()
+        assert "Warning: Vignette decentering is not supported." in captured.out
+
+    def test_configure_surfaces_coordinate_break(self):
+        zemax_data = {
+            "surfaces": {
+                0: {
+                    "type": "coordinate_break",
+                    "param_0": 1.0, # dx
+                    "param_1": 2.0, # dy
+                    "thickness": 5.0, # dz (thickness)
+                    "param_2": 10.0, # rx deg
+                    "param_3": 20.0, # ry deg
+                    "param_4": 30.0, # rz deg
+                    "conic": 0.0,
+                },
+                1: {
+                    "type": "standard",
+                    "radius": 100.0,
+                    "thickness": 10.0,
+                    "conic": 0.0,
+                    "material": "N-BK7"
+                }
+            },
+            "aperture": {"EPD": 10},
+            "fields": {"type": "angle", "x": [0], "y": [0]},
+            "wavelengths": {"primary_index": 0, "data": [0.55]},
+        }
+        converter = ZemaxToOpticConverter(zemax_data)
+        optic = converter.convert()
+        
+        surf = optic.surface_group.surfaces[0]
+        assert surf.geometry.radius == 100.0
+        
+        cs = surf.geometry.cs
+        assert cs.x != 0 or cs.y != 0 or cs.z != 0 or cs.rx != 0 or cs.ry != 0 or cs.rz != 0
+
+    def test_configure_surfaces_toroidal(self):
+        zemax_data = {
+            "surfaces": {
+                0: {
+                    "type": "toroidal",
+                    "radius": 50.0,       # radius_y
+                    "param_1": 60.0,      # radius_x
+                    "param_2": 0.1,       # coeff start
+                    "thickness": 5.0,
+                    "conic": 0.0,
+                    "material": "Air"
+                }
+            },
+            "aperture": {"EPD": 10},
+            "fields": {"type": "angle", "x": [0], "y": [0]},
+            "wavelengths": {"primary_index": 0, "data": [0.55]},
+        }
+        converter = ZemaxToOpticConverter(zemax_data)
+        optic = converter.convert()
+        
+        surf = optic.surface_group.surfaces[0]
+        assert isinstance(surf.geometry, ToroidalGeometry)
+        assert surf.geometry.R_yz == 50.0
+        assert surf.geometry.R_rot == 60.0
+
+    def test_configure_surfaces_infinity_thickness(self):
+        zemax_data = {
+            "surfaces": {
+                0: {
+                    "type": "standard",
+                    "radius": be.inf,
+                    "thickness": be.inf,
+                    "conic": 0.0,
+                    "material": "Air"
+                }
+            },
+            "aperture": {"EPD": 10},
+            "fields": {"type": "angle", "x": [0], "y": [0]},
+            "wavelengths": {"primary_index": 0, "data": [0.55]},
+        }
+        converter = ZemaxToOpticConverter(zemax_data)
+        optic = converter.convert()
+        
+        surf = optic.surface_group.surfaces[0]
+        assert be.isinf(surf.thickness)

tests/test_variable.py

@@ -1,7 +1,9 @@
 import optiland.backend as be
 import pytest
 from unittest.mock import patch
+import numpy as np
 
+from optiland.optic import Optic
 from optiland.coordinate_system import CoordinateSystem
 from optiland.geometries import (
     ChebyshevPolynomialGeometry,
@@ -17,6 +19,13 @@
 from optiland.materials.abbe import AbbeMaterial
 from optiland.optimization.variable.material import MaterialVariable
 from optiland.optimization.scaling.identity import IdentityScaler
+from optiland.optimization.variable import (
+    DecenterVariable,
+    TiltVariable,
+    AsphereCoeffVariable,
+    ZernikeCoeffVariable,
+    Variable
+)
 from .utils import assert_allclose
 
 
@@ -703,4 +712,99 @@ def test_delitem(self, set_test_backend):
         var_manager = variable.VariableManager()
         var_manager.add(optic, "radius", surface_number=1)
         del var_manager[0]
-        assert len(var_manager) == 0
+        assert len(var_manager) == 0
+
+
+class TestVariablesExtended:
+    @pytest.fixture
+    def optic(self):
+        optic = Optic()
+        optic.add_surface(index=0, surface_type="plane")
+        optic.add_surface(index=1, surface_type="standard", radius=100)
+        optic.add_surface(index=2, surface_type="even_asphere", radius=50, coefficients=[0.0, 0.1])
+        optic.add_surface(index=3, surface_type="zernike", radius=200, coefficients=[0.0]*5)
+        return optic
+
+    def test_decenter_variable(self, optic):
+        # Surface 1 is standard
+        var = DecenterVariable(optic, surface_number=1, axis="x")
+        assert var.get_value() == 0.0
+        
+        var.update_value(0.5)
+        assert optic.surface_group.surfaces[1].geometry.cs.x == 0.5
+        assert var.get_value() == 0.5
+        
+        # Test y and z
+        var_y = DecenterVariable(optic, surface_number=1, axis="y")
+        var_y.update_value(-0.2)
+        assert optic.surface_group.surfaces[1].geometry.cs.y == -0.2
+        
+        var_z = DecenterVariable(optic, surface_number=1, axis="z")
+        var_z.update_value(1.0)
+        assert optic.surface_group.surfaces[1].geometry.cs.z == 1.0
+
+    def test_decenter_variable_invalid_axis(self, optic):
+        with pytest.raises(ValueError, match='Invalid axis "r"'):
+            DecenterVariable(optic, surface_number=1, axis="r")
+            
+        var = DecenterVariable(optic, surface_number=1, axis="x")
+        var.axis = "r"
+        with pytest.raises(ValueError, match='Invalid axis "r"'):
+            var.get_value()
+        with pytest.raises(ValueError, match='Invalid axis "r"'):
+            var.update_value(1.0)
+
+    def test_tilt_variable(self, optic):
+        var = TiltVariable(optic, surface_number=1, axis="x")
+        var.update_value(0.1) # radians
+        assert optic.surface_group.surfaces[1].geometry.cs.rx == 0.1
+        
+        var_y = TiltVariable(optic, surface_number=1, axis="y")
+        var_y.update_value(0.2)
+        assert optic.surface_group.surfaces[1].geometry.cs.ry == 0.2
+        
+
+    def test_tilt_variable_invalid_axis(self, optic):
+        with pytest.raises(ValueError, match='Invalid axis "r"'):
+            TiltVariable(optic, surface_number=1, axis="r")
+
+    def test_asphere_coeff_variable_get_padding(self, optic):
+        # Surface 2 is even_asphere with 2 coeffs
+        var = AsphereCoeffVariable(optic, surface_number=2, coeff_number=5)
+        # Should return 0 and pad
+        val = var.get_value()
+        assert val == 0.0
+        # Check padding happened
+        assert len(optic.surface_group.surfaces[2].geometry.coefficients) >= 6
+
+    def test_asphere_coeff_variable_update(self, optic):
+        var = AsphereCoeffVariable(optic, surface_number=2, coeff_number=1)
+        var.update_value(0.5)
+        assert optic.surface_group.surfaces[2].geometry.coefficients[1] == 0.5
+
+    def test_zernike_coeff_variable_padding(self, optic):
+        # Surface 3 is zernike with 5 coeffs (indices 0-4)
+        var = ZernikeCoeffVariable(optic, surface_number=3, coeff_index=10)
+        
+        # Test update padding
+        var.update_value(1.5)
+        coeffs = optic.surface_group.surfaces[3].geometry.coefficients
+        assert len(coeffs) >= 11
+        assert coeffs[10] == 1.5
+        
+        # Test get padding
+        var2 = ZernikeCoeffVariable(optic, surface_number=3, coeff_index=15)
+        val = var2.get_value()
+        assert val == 0.0
+        assert len(optic.surface_group.surfaces[3].geometry.coefficients) >= 16
+
+    def test_variable_wrapper(self, optic):
+        # Test generic Variable wrapper
+        v = Variable(optic, type_name="decenter", surface_number=1, axis="x")
+        v.update(1.0)
+        assert v.value == 1.0
+        assert optic.surface_group.surfaces[1].geometry.cs.x == 1.0
+
+    def test_variable_wrapper_invalid_type(self, optic):
+        with pytest.raises(ValueError, match='Invalid variable type "invalid"'):
+            Variable(optic, type_name="invalid")