Feat: add dielectric to path tracer

Author: RDambrosio016

examples/cuda/cpu/path_tracer/src/main.rs

@@ -7,7 +7,7 @@ pub mod viewer;
 
 use common::Camera;
 use path_tracer_gpu::{
-    material::{DiffuseMaterial, MaterialKind, MetallicMaterial},
+    material::{DielectricMaterial, DiffuseMaterial, MaterialKind, MetallicMaterial},
     scene::Scene,
     sphere::Sphere,
     Object,
@@ -37,13 +37,19 @@ fn main() -> Result<(), Box<dyn Error>> {
         }),
         MaterialKind::Metallic(MetallicMaterial {
             color: Vec3::new(1.0, 0.7, 0.7),
-            roughness: 0.02,
+            roughness: 0.2,
+        }),
+        MaterialKind::Dielectric(DielectricMaterial {
+            ior: 1.0,
+            color: Vec3::new(1.0, 1.0, 1.0),
         }),
     ];
 
     let objects = vec![
         Object::Sphere(Sphere::new(Vec3::new(0.0, 0.0, -1.0), 0.5, 0)),
         Object::Sphere(Sphere::new(Vec3::new(1.1, 0.2, -0.7), 0.2, 2)),
+        Object::Sphere(Sphere::new(Vec3::new(-0.5, 0.4, -0.9), 0.3, 3)),
+        Object::Sphere(Sphere::new(Vec3::new(0.7, 0.0, -0.2), 0.2, 1)),
         // Object::Sphere(Sphere::new(Vec3::new(0.0, -200.5, -1.0), 200.0, 1)),
     ];
     let cpu_scene = Scene {

examples/cuda/gpu/path_tracer_gpu/src/material.rs

@@ -1,11 +1,7 @@
-use crate::{
-    hittable::HitRecord,
-    math::{random_in_unit_sphere, reflect},
-    Ray, Vec3,
-};
+use crate::{hittable::HitRecord, math::*, Ray, Vec3};
 use cust_core::DeviceCopy;
 use enum_dispatch::enum_dispatch;
-use gpu_rand::DefaultRand;
+use gpu_rand::{DefaultRand, GpuRand};
 
 #[enum_dispatch]
 pub trait Material {
@@ -18,6 +14,7 @@ pub trait Material {
 pub enum MaterialKind {
     Diffuse(DiffuseMaterial),
     Metallic(MetallicMaterial),
+    Dielectric(DielectricMaterial),
 }
 
 #[derive(Clone, Copy, PartialEq, DeviceCopy)]
@@ -63,3 +60,47 @@ impl Material for MetallicMaterial {
         }
     }
 }
+
+#[derive(Clone, Copy, PartialEq, DeviceCopy)]
+pub struct DielectricMaterial {
+    pub ior: f32,
+    pub color: Vec3,
+}
+
+impl Material for DielectricMaterial {
+    fn scatter(&self, incoming: Ray, hit: HitRecord, rng: &mut DefaultRand) -> (Vec3, Option<Ray>) {
+        let outward_norm;
+        let ni_over_nt: f32;
+        let cos: f32;
+
+        if incoming.dir.dot(hit.normal) > 0.0 {
+            outward_norm = -hit.normal;
+            ni_over_nt = self.ior;
+            cos = self.ior * incoming.dir.dot(hit.normal) / incoming.dir.magnitude();
+        } else {
+            outward_norm = hit.normal;
+            ni_over_nt = 1.0 / self.ior;
+            cos = -incoming.dir.dot(hit.normal) / incoming.dir.magnitude();
+        }
+
+        if let Some(refracted) = refract(incoming.dir, outward_norm, ni_over_nt) {
+            if rng.normal_f32() > schlick(cos, self.ior) {
+                return (
+                    self.color,
+                    Some(Ray {
+                        origin: hit.point,
+                        dir: refracted,
+                    }),
+                );
+            }
+        }
+
+        (
+            self.color,
+            Some(Ray {
+                origin: hit.point,
+                dir: reflect(incoming.dir.normalized(), hit.normal),
+            }),
+        )
+    }
+}

examples/cuda/gpu/path_tracer_gpu/src/math.rs

@@ -35,9 +35,19 @@ pub fn reflect(v: Vec3, n: Vec3) -> Vec3 {
     v - 2.0 * v.dot(n) * n
 }
 
-pub fn refract(uv: Vec3, n: Vec3, etai_over_etat: f32) -> Vec3 {
-    let cos_t = (-uv).dot(n).min(1.0);
-    let r_out_perp = etai_over_etat * (uv + cos_t * n);
-    let r_out_parallel = -((1.0 - r_out_perp.magnitude_squared()).abs()).sqrt() * n;
-    r_out_perp + r_out_parallel
+pub fn refract(v: Vec3, n: Vec3, ni_over_nt: f32) -> Option<Vec3> {
+    let uv = v.normalized();
+    let dt = uv.dot(n);
+    let discriminant = 1.0 - ni_over_nt * ni_over_nt * (1.0 - dt * dt);
+    if discriminant > 0.0 {
+        Some(ni_over_nt * (uv - n * dt) - discriminant.sqrt() * n)
+    } else {
+        None
+    }
+}
+
+pub fn schlick(cos: f32, ref_idx: f32) -> f32 {
+    let r0 = (1.0 - ref_idx) / (1.0 + ref_idx);
+    let r0sq = r0 * r0;
+    r0sq + (1.0 - r0sq) * (1.0 - cos).powf(5.0)
 }

examples/cuda/gpu/path_tracer_gpu/src/optix.rs

@@ -1,13 +1,14 @@
 use crate::{
     hittable::{HitRecord, Hittable},
+    material::Material,
     render::{self, generate_ray},
     scene::{Scene, MAX_BOUNCES},
     sphere::Sphere,
     Ray, Vec2, Vec3, Viewport,
 };
 use cuda_std::kernel;
 use cust_core::DeviceCopy;
-use gpu_rand::DefaultRand;
+use gpu_rand::{DefaultRand, GpuRand};
 use optix_device::{
     closesthit, get_launch_index,
     glam::Vec3Swizzles,

examples/cuda/gpu/path_tracer_gpu/src/render.rs

@@ -1,8 +1,10 @@
 use crate::*;
 
+const BACKGROUND_BLUE_MULTIPLIER: f32 = 0.7;
+
 pub fn color(ray: Ray) -> Vec3 {
     let unit = ray.dir.normalized();
-    let t = 0.5 * (unit.y + 1.0);
+    let t = BACKGROUND_BLUE_MULTIPLIER * (unit.y + 1.0);
     (1.0 - t) * Vec3::one() + t * Vec3::new(0.5, 0.7, 1.0)
 }
 

examples/cuda/gpu/path_tracer_gpu/src/scene.rs

@@ -58,10 +58,7 @@ impl Scene<'_> {
                     return Vec3::zero();
                 }
             } else {
-                let unit = cur_ray.dir.normalized();
-                let t = 0.5 * (unit.y + 1.0);
-                let c = (1.0 - t) * Vec3::one() + t * Vec3::new(0.5, 0.7, 1.0);
-                return attenuation * c;
+                return attenuation * render::color(cur_ray);
             }
         }
         Vec3::zero()