Differentiate shadow and intersection in ray logs

pybits/visualize_rays.py

@@ -11,7 +11,8 @@
         ("px", np.uint16),
         ("py", np.uint16),
         ("bounces", np.uint8),
-        ("inf", np.uint8),
+        ("shadow", np.uint8),
+        ("intersect", np.uint8),
         ("ox", np.float32),
         ("oy", np.float32),
         ("oz", np.float32),
@@ -24,34 +25,52 @@
 
 def read(path):
     rays = pd.DataFrame(np.fromfile(path, RAY_DTYPE))
-    rays["tx"] = rays.ox + rays.dx + rays.inf * rays.dx * 10.0
-    rays["ty"] = rays.oy + rays.dy + rays.inf * rays.dy * 10.0
-    rays["tz"] = rays.oz + rays.dz + rays.inf * rays.dz * 10.0
+    rays["tx"] = rays.ox + rays.dx
+    rays["ty"] = rays.oy + rays.dy
+    rays["tz"] = rays.oz + rays.dz
     return rays
 
 
 def visualize_as_one_entity(rays):
-    colors = plt.colormaps["viridis"](np.linspace(0.0, 1.0, len(rays)))
+    def plot(name, rays, color):
+        segments = rays[["ox", "oy", "oz", "tx", "ty", "tz"]].to_numpy()
+        segments = segments.reshape((len(rays), 2, 3))
+        rerun.log(
+            f"world/rays/{name}",
+            rerun.LineStrips3D(segments, radii=0.001, colors=color),
+            static=True,
+        )
 
-    segments = rays[["ox", "oy", "oz", "tx", "ty", "tz"]].to_numpy()
-    segments = segments.reshape((len(rays), 2, 3))
-    rerun.log(
-        "world/rays",
-        rerun.LineStrips3D(segments, radii=0.001, colors=colors),
-        static=True,
-    )
+    rays_environment = rays[(rays.shadow == 0) & (rays.intersect == 0)]
+    rays_surface = rays[(rays.shadow == 0) & (rays.intersect == 1)]
+    rays_light = rays[(rays.shadow == 1) & (rays.intersect == 0)]
+    rays_shadow = rays[(rays.shadow == 1) & (rays.intersect == 1)]
+    plot("environment", rays_environment, [255, 0, 0])
+    plot("surface", rays_surface, [0, 255, 0])
+    plot("light", rays_light, [255, 255, 0])
+    plot("shadow", rays_shadow, [0, 255, 255])
 
 
 def visualize_grouped_per_path(rays):
     colormap = plt.colormaps["viridis"].resampled(rays.bounces.max())
     rays["color"] = rays.bounces / rays.bounces.max()
     paths = rays.groupby(["iteration", "px", "py"])
 
-    for (iteration, px, py), path in paths:
+    def plot(name, path):
         segments = path[["ox", "oy", "oz", "tx", "ty", "tz"]].values
         segments = segments.reshape((len(path), 2, 3))
         rerun.log(
-            f"world/rays/ray_{iteration}_{px}x{py}",
+            f"world/rays/{name}",
             rerun.LineStrips3D(segments, radii=0.001, colors=colormap(path.color)),
             static=True,
         )
+
+    for (iteration, px, py), path in paths:
+        path_environment = path[(path.shadow == 0) & (path.intersect == 0)]
+        path_surface = path[(path.shadow == 0) & (path.intersect == 1)]
+        path_light = path[(path.shadow == 1) & (path.intersect == 0)]
+        path_shadow = path[(path.shadow == 1) & (path.intersect == 1)]
+        plot(f"{iteration}_{px}x{py}/environment", path_environment)
+        plot(f"{iteration}_{px}x{py}/surface", path_surface)
+        plot(f"{iteration}_{px}x{py}/light", path_light)
+        plot(f"{iteration}_{px}x{py}/shadow", path_shadow)

recreate-venv.sh

@@ -3,4 +3,4 @@
 rm -rf ./.venv/
 
 uv venv .venv
-uv pip install numpy pandas matplotlib rerun-sdk==0.17 ruff-lsp pyright
+uv pip install numpy pandas matplotlib rerun-sdk==0.18.2 ruff-lsp pyright

tracing/src/pathtracer.rs

@@ -29,18 +29,22 @@ impl Pathtracer {
             let intersection = self
                 .kdtree
                 .intersect(&self.scene.geometries, &ray, 0.0..=f32::MAX);
+            ray_logger
+                .log_ray(
+                    &intersection
+                        .as_ref()
+                        .map_or(ray.clone(), |isect| ray.extended(isect.intersection.t)),
+                    bounce,
+                    intersection.is_some(),
+                )
+                .unwrap();
             if intersection.is_none() {
-                ray_logger.log_infinite(&ray, bounce).unwrap();
                 return accumulated_radiance + accumulated_transport * self.scene.environment;
             }
             let intersection = intersection.unwrap();
             let intersection_index = intersection.index;
             let intersection = intersection.intersection;
 
-            ray_logger
-                .log_finite(&ray.extended(intersection.t), bounce)
-                .unwrap();
-
             let wi = -ray.direction;
             let n = self.scene.get_normal(intersection_index, &intersection);
             let uv = self.scene.get_texcoord(intersection_index, &intersection);
@@ -64,6 +68,9 @@ impl Pathtracer {
                     let intersection =
                         self.kdtree
                             .intersect(&self.scene.geometries, &shadow_ray, 0.0..=1.0);
+                    ray_logger
+                        .log_shadow(&shadow_ray, bounce, intersection.is_some())
+                        .unwrap();
                     if intersection.is_some() {
                         return Vec3::ZERO;
                     }

tracing/src/raylogger.rs

@@ -38,7 +38,8 @@ impl RayLoggerWriter {
     fn write(
         &mut self,
         ray: &Ray,
-        infinite: bool,
+        shadow: u8,
+        intersect: u8,
         iteration: u16,
         pixel_x: u16,
         pixel_y: u16,
@@ -49,18 +50,19 @@ impl RayLoggerWriter {
         }
 
         if let Self::File(buf) = self {
-            let mut bytes = [0u8; 32];
+            let mut bytes = [0u8; 33];
             bytes[0..2].copy_from_slice(&iteration.to_le_bytes());
             bytes[2..4].copy_from_slice(&pixel_x.to_le_bytes());
             bytes[4..6].copy_from_slice(&pixel_y.to_le_bytes());
             bytes[6..7].copy_from_slice(&bounces.to_le_bytes());
-            bytes[7..8].copy_from_slice(&(infinite as u8).to_le_bytes());
-            bytes[8..12].copy_from_slice(&ray.origin.x.to_le_bytes());
-            bytes[12..16].copy_from_slice(&ray.origin.y.to_le_bytes());
-            bytes[16..20].copy_from_slice(&ray.origin.z.to_le_bytes());
-            bytes[20..24].copy_from_slice(&ray.direction.x.to_le_bytes());
-            bytes[24..28].copy_from_slice(&ray.direction.y.to_le_bytes());
-            bytes[28..32].copy_from_slice(&ray.direction.z.to_le_bytes());
+            bytes[7..8].copy_from_slice(&shadow.to_le_bytes());
+            bytes[8..9].copy_from_slice(&intersect.to_le_bytes());
+            bytes[9..13].copy_from_slice(&ray.origin.x.to_le_bytes());
+            bytes[13..17].copy_from_slice(&ray.origin.y.to_le_bytes());
+            bytes[17..21].copy_from_slice(&ray.origin.z.to_le_bytes());
+            bytes[21..25].copy_from_slice(&ray.direction.x.to_le_bytes());
+            bytes[25..29].copy_from_slice(&ray.direction.y.to_le_bytes());
+            bytes[29..33].copy_from_slice(&ray.direction.z.to_le_bytes());
             buf.write_all(&bytes)
         } else {
             Ok(())
@@ -92,21 +94,25 @@ pub struct RayLoggerWithIterationAndPixel<'a> {
 }
 
 impl RayLoggerWithIterationAndPixel<'_> {
-    pub fn log_infinite(&mut self, ray: &Ray, bounces: u8) -> Result<(), Error> {
+    pub fn log_ray(&mut self, ray: &Ray, bounces: u8, intersect: bool) -> Result<(), Error> {
+        let shadow = false;
         self.writer.write(
             ray,
-            true,
+            shadow as u8,
+            intersect as u8,
             self.iteration,
             self.pixel_x,
             self.pixel_y,
             bounces,
         )
     }
 
-    pub fn log_finite(&mut self, ray: &Ray, bounces: u8) -> Result<(), Error> {
+    pub fn log_shadow(&mut self, ray: &Ray, bounces: u8, intersect: bool) -> Result<(), Error> {
+        let shadow = true;
         self.writer.write(
             ray,
-            false,
+            shadow as u8,
+            intersect as u8,
             self.iteration,
             self.pixel_x,
             self.pixel_y,