optical data: replace material by index

Author: victorpoughon

src/torchlensmaker/elements/light_sources.py

@@ -92,25 +92,25 @@ def forward(self, inputs: OpticalData) -> OpticalData:
             rays_object=cat_optional(inputs.rays_object, rays_object),
             var_base=var_base,
             var_object=var_object,
-            material=self.material,
         )
 
         # now sample wavelength
-
-        if non_chromatic.rays_wavelength is not None:
-            raise RuntimeError("Rays already have wavelength data")
-
         if "wavelength" not in non_chromatic.sampling:
             raise RuntimeError("Missing 'wavelength' key in sampling configuration")
 
         chromatic_space = non_chromatic.sampling["wavelength"].sample1d(
             self.wavelength_lower, self.wavelength_upper, non_chromatic.dtype
         )
 
-        return cartesian_wavelength(non_chromatic, chromatic_space).replace(
+        chromatic = cartesian_wavelength(non_chromatic, chromatic_space).replace(
             var_wavelength=chromatic_space
         )
 
+        # index of refraction
+        rays_index = self.material.refractive_index(chromatic.rays_wavelength)
+
+        return chromatic.replace(rays_index=rays_index)
+
 
 class RaySource(LightSourceBase):
     """

src/torchlensmaker/elements/optical_surfaces.py

@@ -237,30 +237,18 @@ def forward(self, data: OpticalData) -> tuple[OpticalData, Tensor]:
         # Zero rays special case
         # (needs to happen after self.collision_surface is called to enable rendering of it)
         if data.P.numel() == 0:
-            return data.replace(
-                material=self.material, dfk=new_dfk, ifk=new_ifk
-            ), torch.full((data.P.shape[0],), True)
-
-        # Compute indices of refraction (scalars for non dispersive materials,
-        # tensors for dispersive materials)
-        if (
-            data.rays_wavelength is None
-        ):  # TODO this should be empty but not None when number of rays is zero
-            if not isinstance(data.material, NonDispersiveMaterial) or not isinstance(
-                self.material, NonDispersiveMaterial
-            ):
-                raise RuntimeError(
-                    f"Cannot compute refraction with dispersive material "
-                    f"because optical data has no wavelength variable "
-                    f"(got materials {data.material} and {self.material})"
-                )
-
-            n1 = torch.as_tensor(data.material.n)
-            n2 = torch.as_tensor(self.material.n)
+            return data.replace(dfk=new_dfk, ifk=new_ifk), torch.full(
+                (data.P.shape[0],), True
+            )
 
-        else:
-            n1 = data.material.refractive_index(data.rays_wavelength)
-            n2 = self.material.refractive_index(data.rays_wavelength)
+        # Compute indices of refraction
+        n1 = data.rays_index
+        n2 = self.material.refractive_index(data.rays_wavelength)
+        assert n1.shape == n2.shape == (data.P.shape[0],), (
+            n1.shape,
+            n2.shape,
+            data.P.shape,
+        )
 
         # Snell's law happens here
         # Compute refraction on the full frame rays (including non-colliding
@@ -282,6 +270,7 @@ def forward(self, data: OpticalData) -> tuple[OpticalData, Tensor]:
             new_rays_wavelength = filter_optional_tensor(
                 data.rays_wavelength, both_valid
             )
+            new_rays_index = filter_optional_tensor(n2, both_valid)
         else:
             # keep tir rays
             new_P = collision_points[valid_collision]
@@ -291,14 +280,15 @@ def forward(self, data: OpticalData) -> tuple[OpticalData, Tensor]:
             new_rays_wavelength = filter_optional_tensor(
                 data.rays_wavelength, valid_collision
             )
+            new_rays_index = filter_optional_tensor(n2, valid_collision)
 
         return data.replace(
             P=new_P,
             V=new_V,
             rays_base=new_rays_base,
             rays_object=new_rays_object,
             rays_wavelength=new_rays_wavelength,
-            material=self.material,
+            rays_index=new_rays_index,
             dfk=new_dfk,
             ifk=new_ifk,
         ), valid_refraction
@@ -328,6 +318,9 @@ def forward(self, data: OpticalData) -> OpticalData:
             rays_wavelength=filter_optional_tensor(
                 data.rays_wavelength, valid_collision
             ),
+            rays_index=filter_optional_tensor(
+                data.rays_index, valid_collision
+            ),
             dfk=new_dfk,
             ifk=new_ifk,  # correct but useless cause Aperture is only circular plane currently
         )

src/torchlensmaker/optical_data.py

@@ -50,7 +50,12 @@ class OpticalData:
     # Light rays in parametric form: P + tV
     P: Float[torch.Tensor, "N D"]
     V: Float[torch.Tensor, "N D"]
+
+    # Light rays wavelength in nm
     rays_wavelength: Float[torch.Tensor, " N"]
+    
+    # Light rays index of refraction
+    rays_index: Float[torch.Tensor, " N"]
 
     # Rays variables
     # Tensors of shape (N, 2|3) or None
@@ -65,9 +70,6 @@ class OpticalData:
     var_object: Optional[torch.Tensor]
     var_wavelength: Optional[torch.Tensor]
 
-    # Material model for this batch of rays
-    material: MaterialModel
-
     # Loss accumulator
     # Tensor of dim 0
     loss: torch.Tensor
@@ -124,13 +126,13 @@ def default_input(
         ifk=ifk,
         P=torch.empty((0, dim), dtype=dtype),
         V=torch.empty((0, dim), dtype=dtype),
+        rays_wavelength=torch.empty((0,), dtype=dtype),
+        rays_index=torch.empty((0,), dtype=dtype),
         rays_base=None,
         rays_object=None,
         rays_image=None,
-        rays_wavelength=None,
         var_base=None,
         var_object=None,
         var_wavelength=None,
-        material=get_material_model("vacuum"),
         loss=torch.tensor(0.0, dtype=dtype),
     )