Add HeightProfile with wavelength-dependent material support

gustavu92 · gustavu92 · commit 69580d5a1c94 · 2026-02-15T01:51:48.000+01:00
diff --git a/optiland/interactions/phase_interaction_model.py b/optiland/interactions/phase_interaction_model.py
@@ -65,8 +65,8 @@ def interact_real_rays(self, rays: RealRays) -> RealRays:
         k_iz = n1 * k0 * n_i
 
         # 3. Get phase and ambient gradient (grad(f))
-        phase_val = self.phase_profile.get_phase(x, y)
-        phi_x, phi_y, phi_z = self.phase_profile.get_gradient(x, y)
+        phase_val = self.phase_profile.get_phase(x, y, rays.w)
+        phi_x, phi_y, phi_z = self.phase_profile.get_gradient(x, y, rays.w)
         grad_f_x = phi_x
         grad_f_y = phi_y
         grad_f_z = phi_z
diff --git a/optiland/phase/__init__.py b/optiland/phase/__init__.py
@@ -3,13 +3,15 @@
 from .base import BasePhaseProfile
 from .constant import ConstantPhaseProfile
 from .grid import GridPhaseProfile
+from .height_profile import HeightProfile
 from .linear_grating import LinearGratingPhaseProfile
 from .radial import RadialPhaseProfile
 
 __all__ = [
     "BasePhaseProfile",
     "ConstantPhaseProfile",
     "GridPhaseProfile",
+    "HeightProfile",
     "LinearGratingPhaseProfile",
     "RadialPhaseProfile",
 ]
diff --git a/optiland/phase/base.py b/optiland/phase/base.py
@@ -37,7 +37,7 @@ def efficiency(self) -> float:
         return 1.0
 
     @abc.abstractmethod
-    def get_phase(self, x: be.Array, y: be.Array) -> be.Array:
+    def get_phase(self, x: be.Array, y: be.Array, wavelength: be.Array) -> be.Array:
         """Calculates the phase added by the profile at coordinates (x, y).
 
         Args:
@@ -51,7 +51,7 @@ def get_phase(self, x: be.Array, y: be.Array) -> be.Array:
 
     @abc.abstractmethod
     def get_gradient(
-        self, x: be.Array, y: be.Array
+        self, x: be.Array, y: be.Array, wavelength: be.Array
     ) -> tuple[be.Array, be.Array, be.Array]:
         """Calculates the gradient of the phase at coordinates (x, y).
 
@@ -66,7 +66,7 @@ def get_gradient(
         raise NotImplementedError
 
     @abc.abstractmethod
-    def get_paraxial_gradient(self, y: be.Array) -> be.Array:
+    def get_paraxial_gradient(self, y: be.Array, wavelength: be.Array) -> be.Array:
         """Calculates the paraxial phase gradient at y-coordinate.
 
         This is the gradient d_phi/dy evaluated at x=0.
diff --git a/optiland/phase/constant.py b/optiland/phase/constant.py
@@ -24,7 +24,7 @@ class ConstantPhaseProfile(BasePhaseProfile):
     def __init__(self, phase: float = 0.0):
         self.phase = phase
 
-    def get_phase(self, x: be.Array, y: be.Array) -> be.Array:
+    def get_phase(self, x: be.Array, y: be.Array, wavelength: be.Array = None) -> be.Array:
         """Calculates the phase added by the profile at coordinates (x, y).
 
         Args:
@@ -37,7 +37,7 @@ def get_phase(self, x: be.Array, y: be.Array) -> be.Array:
         return be.full_like(x, self.phase)
 
     def get_gradient(
-        self, x: be.Array, y: be.Array
+        self, x: be.Array, y: be.Array, wavelength: be.Array = None
     ) -> tuple[be.Array, be.Array, be.Array]:
         """Calculates the gradient of the phase at coordinates (x, y).
 
@@ -51,7 +51,7 @@ def get_gradient(
         """
         return be.zeros_like(x), be.zeros_like(y), be.zeros_like(x)
 
-    def get_paraxial_gradient(self, y: be.Array) -> be.Array:
+    def get_paraxial_gradient(self, y: be.Array, wavelength: be.Array = None) -> be.Array:
         """Calculates the paraxial phase gradient at y-coordinate.
 
         This is the gradient d_phi/dy evaluated at x=0.
diff --git a/optiland/phase/grid.py b/optiland/phase/grid.py
@@ -56,7 +56,7 @@ def __init__(
             self.y_coords, self.x_coords, self.phase_grid
         )
 
-    def get_phase(self, x: be.Array, y: be.Array) -> be.Array:
+    def get_phase(self, x: be.Array, y: be.Array, wavelength: be.Array = None) -> be.Array:
         """Calculates the phase added by the profile at coordinates (x, y).
 
         Args:
@@ -69,7 +69,7 @@ def get_phase(self, x: be.Array, y: be.Array) -> be.Array:
         return self._spline.ev(be.to_numpy(y), be.to_numpy(x))
 
     def get_gradient(
-        self, x: be.Array, y: be.Array
+        self, x: be.Array, y: be.Array, wavelength: be.Array = None
     ) -> tuple[be.Array, be.Array, be.Array]:
         """Calculates the gradient of the phase at coordinates (x, y).
 
@@ -88,7 +88,7 @@ def get_gradient(
         d_phi_dz = be.zeros_like(x)
         return d_phi_dx, d_phi_dy, d_phi_dz
 
-    def get_paraxial_gradient(self, y: be.Array) -> be.Array:
+    def get_paraxial_gradient(self, y: be.Array, wavelength: be.Array = None) -> be.Array:
         """Calculates the paraxial phase gradient at y-coordinate.
 
         This is the gradient d_phi/dy evaluated at x=0.
diff --git a/optiland/phase/height_profile.py b/optiland/phase/height_profile.py
@@ -0,0 +1,119 @@
+"""
+Provides a phase profile based on a height map and dispersive material.
+"""
+
+from __future__ import annotations
+
+from optiland import backend as be
+from optiland.materials.base import BaseMaterial
+from optiland.phase.base import BasePhaseProfile
+
+try:
+    from scipy.interpolate import RectBivariateSpline
+except ImportError:
+    RectBivariateSpline = None
+
+
+class HeightProfile(BasePhaseProfile):
+    """A phase profile defined by a height map and a dispersive material.
+
+    The phase is calculated as:
+        phi(x, y, λ) = (2π / λ) * (n_material(λ) - 1) * h(x, y)
+
+    Assumes air as the reference medium.
+
+    Args:
+        x_coords (be.Array): X-coordinates of the height map grid.
+        y_coords (be.Array): Y-coordinates of the height map grid.
+        height_map (be.Array): Height values at grid points
+            with shape (len(y_coords), len(x_coords)).
+        material: Material providing wavelength-dependent refractive index n(λ).
+    """
+
+    phase_type = "height_profile"
+
+    def __init__(
+        self,
+        x_coords: be.Array,
+        y_coords: be.Array,
+        height_map: be.Array,
+        material: BaseMaterial,
+    ):
+        if RectBivariateSpline is None:
+            raise ImportError(
+                "scipy is required for HeightProfile. Install with: pip install scipy"
+            )
+
+        self.x_coords = be.to_numpy(x_coords)
+        self.y_coords = be.to_numpy(y_coords)
+        self.height_map = be.to_numpy(height_map)
+        self.material = material
+
+        self._spline = RectBivariateSpline(
+            self.y_coords,
+            self.x_coords,
+            self.height_map,
+        )
+
+    def _interpolate_height(self, x: be.Array, y: be.Array) -> be.Array:
+        return self._spline.ev(be.to_numpy(y), be.to_numpy(x))
+
+    def _interpolate_gradient(
+        self, x: be.Array, y: be.Array
+    ) -> tuple[be.Array, be.Array]:
+        x_np = be.to_numpy(x)
+        y_np = be.to_numpy(y)
+        dh_dx = self._spline.ev(y_np, x_np, dy=1)
+        dh_dy = self._spline.ev(y_np, x_np, dx=1)
+        return dh_dx, dh_dy
+
+    def get_phase(
+        self,
+        x: be.Array,
+        y: be.Array,
+        wavelength: be.Array,
+    ) -> be.Array:
+        h = self._interpolate_height(x, y)
+        n = self.material.n(wavelength)
+        return 2 * be.pi / (wavelength * 1e-3) * (n - 1.0) * h
+
+    def get_gradient(
+        self,
+        x: be.Array,
+        y: be.Array,
+        wavelength: be.Array,
+    ) -> tuple[be.Array, be.Array, be.Array]:
+        dh_dx, dh_dy = self._interpolate_gradient(x, y)
+        n = self.material.n(wavelength)
+        factor = 2 * be.pi / (wavelength * 1e-3) * (n - 1.0)
+        return factor * dh_dx, factor * dh_dy, be.zeros_like(x)
+
+    def get_paraxial_gradient(
+        self,
+        y: be.Array,
+        wavelength: be.Array,
+    ) -> be.Array:
+        dh_dy = self._spline.ev(
+            be.to_numpy(y),
+            be.zeros_like(y),
+            dx=1,
+        )
+        n = self.material.n(wavelength)
+        return 2 * be.pi / (wavelength * 1e-3) * (n - 1.0) * dh_dy
+
+    def to_dict(self) -> dict:
+        data = super().to_dict()
+        data["x_coords"] = self.x_coords.tolist()
+        data["y_coords"] = self.y_coords.tolist()
+        data["height_map"] = self.height_map.tolist()
+        data["material"] = getattr(self.material, "name", str(self.material))
+        return data
+
+    @classmethod
+    def from_dict(cls, data: dict) -> HeightProfile:
+        return cls(
+            x_coords=be.array(data["x_coords"]),
+            y_coords=be.array(data["y_coords"]),
+            height_map=be.array(data["height_map"]),
+            material=data["material"],
+        )
diff --git a/optiland/phase/linear_grating.py b/optiland/phase/linear_grating.py
@@ -57,7 +57,7 @@ def __init__(
     def efficiency(self) -> float:
         return self._efficiency
 
-    def get_phase(self, x: be.Array, y: be.Array) -> be.Array:
+    def get_phase(self, x: be.Array, y: be.Array, wavelength: be.Array = None) -> be.Array:
         """Calculates the phase added by the profile at coordinates (x, y).
 
         Args:
@@ -70,7 +70,7 @@ def get_phase(self, x: be.Array, y: be.Array) -> be.Array:
         return self._K_x * x + self._K_y * y
 
     def get_gradient(
-        self, x: be.Array, y: be.Array
+        self, x: be.Array, y: be.Array, wavelength: be.Array = None
     ) -> tuple[be.Array, be.Array, be.Array]:
         """Calculates the gradient of the phase at coordinates (x, y).
 
diff --git a/optiland/phase/radial.py b/optiland/phase/radial.py
@@ -23,7 +23,7 @@ class RadialPhaseProfile(BasePhaseProfile):
     def __init__(self, coefficients: list[float]):
         self.coefficients = coefficients
 
-    def get_phase(self, x: be.Array, y: be.Array) -> be.Array:
+    def get_phase(self, x: be.Array, y: be.Array, wavelength: be.Array = None) -> be.Array:
         """Calculates the phase added by the profile at coordinates (x, y).
 
         Args:
@@ -41,7 +41,7 @@ def get_phase(self, x: be.Array, y: be.Array) -> be.Array:
         return phase
 
     def get_gradient(
-        self, x: be.Array, y: be.Array
+        self, x: be.Array, y: be.Array, wavelength: be.Array = None
     ) -> tuple[be.Array, be.Array, be.Array]:
         """Calculates the gradient of the phase at coordinates (x, y).
 
@@ -74,7 +74,7 @@ def get_gradient(
 
         return d_phi_dx, d_phi_dy, d_phi_dz
 
-    def get_paraxial_gradient(self, y: be.Array) -> be.Array:
+    def get_paraxial_gradient(self, y: be.Array, wavelength: be.Array = None) -> be.Array:
         """Calculates the paraxial phase gradient at y-coordinate.
 
         This is the gradient d_phi/dy evaluated at x=0.
diff --git a/tests/test_height_profile_phase.py b/tests/test_height_profile_phase.py
@@ -0,0 +1,54 @@
+import pytest
+import numpy as np
+from optiland import backend as be
+from optiland.phase.height_profile import HeightProfile
+from optiland.materials.ideal import IdealMaterial
+from .utils import assert_allclose
+
+pytest.importorskip("scipy")
+
+@pytest.fixture
+def height_data():
+    x = be.linspace(-1, 1, 5)
+    y = be.linspace(-1, 1, 4)
+    height_map = be.array([[i + j for i in x] for j in y])
+    material = IdealMaterial(n=1.5)
+    return x, y, height_map, material
+
+def test_height_profile_init(height_data):
+    x, y, height_map, material = height_data
+    profile = HeightProfile(x, y, height_map, material)
+    assert profile.x_coords.shape[0] == len(x)
+    assert profile.y_coords.shape[0] == len(y)
+    assert profile.height_map.shape == (len(y), len(x))
+    assert profile.material is material
+
+def test_height_profile_get_phase(height_data):
+    x, y, height_map, material = height_data
+    profile = HeightProfile(x, y, height_map, material)
+    phase = profile.get_phase(be.array([0.0]), be.array([0.0]), be.array([1.0]))
+    assert phase.shape == (1,)
+    assert isinstance(phase.item(), float)
+
+def test_height_profile_get_gradient(height_data):
+    x, y, height_map, material = height_data
+    profile = HeightProfile(x, y, height_map, material)
+    grad_x, grad_y, grad_z = profile.get_gradient(be.array([0.0]), be.array([0.0]), be.array([1.0]))
+    assert grad_x.shape == grad_y.shape == grad_z.shape
+    assert_allclose(grad_z, be.zeros_like(grad_z))
+
+def test_height_profile_get_paraxial_gradient(height_data):
+    x, y, height_map, material = height_data
+    profile = HeightProfile(x, y, height_map, material)
+    paraxial = profile.get_paraxial_gradient(be.array([0.0, 0.5]), be.array([1.0]))
+    assert paraxial.shape[0] == 2
+
+def test_height_profile_to_from_dict(height_data):
+    x, y, height_map, material = height_data
+    profile = HeightProfile(x, y, height_map, material)
+    data = profile.to_dict()
+    new_profile = HeightProfile.from_dict(data)
+    assert isinstance(new_profile, HeightProfile)
+    assert_allclose(new_profile.x_coords, x)
+    assert_allclose(new_profile.y_coords, y)
+    assert_allclose(new_profile.height_map, height_map)
