Fix some Microfacet bugs

ray-tracing-material/Cargo.toml

@@ -4,4 +4,5 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+rand_distr = "0.4.3"
 ray-tracing-core = { path = "../ray-tracing-core" }

ray-tracing-material/src/microfacet.rs

@@ -1,3 +1,4 @@
+use rand_distr::{Distribution, Normal};
 use ray_tracing_core::{material::Material, prelude::*};
 use std::fmt::Debug;
 
@@ -13,7 +14,7 @@ impl<D: MicrofacetDistribution + Debug> Microfacet<D> {
     }
 }
 
-fn fresnel_real(cos_theta_in: Float, nu1: Float, nu2: Float) -> Float {
+pub fn fresnel_real(cos_theta_in: Float, nu1: Float, nu2: Float) -> Float {
     let nu_rel = nu1 / nu2;
     let cos_theta_tr = Float::sqrt(1.0 - nu_rel * nu_rel * (1.0 - cos_theta_in * cos_theta_in));
 
@@ -34,12 +35,28 @@ impl<R: Rng, D: MicrofacetDistribution + Debug + Sync> Material<R> for Microface
 
             let g = self.dist.g1(w_in, w_h) * self.dist.g1(w_out, w_h);
 
-            self.color * fresnel_real(Dir3::dot(w_in, w_h), 1.0, 1.3) * g * self.dist.d(w_h)
+            self.color * fresnel_real(Dir3::dot(w_in, w_h), 1.0, 1.5) * g * self.dist.d(w_h)
                 / (4.0 * w_in.y() * w_out.y()).max(0.0)
         } else {
             Color::black()
         }
     }
+
+    fn sample(&self, w_in: Dir3, rng: &mut R) -> ray_tracing_core::material::SampleResult {
+        let w_h = self.dist.sample_d(rng);
+
+        let w_out = (2.0 * (Dir3::dot(w_in, w_h) * w_h) - w_in).normalize();
+
+        if w_out.y() > 0.0 {
+            let g = self.dist.g1(w_in, w_h) * self.dist.g1(w_out, w_h);
+            ray_tracing_core::material::SampleResult::new(
+                w_out,
+                self.color * fresnel_real(Dir3::dot(w_in, w_h), 1.0, 1.5) * g,
+            )
+        } else {
+            ray_tracing_core::material::SampleResult::new(w_out, Color::black())
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -57,25 +74,43 @@ impl MicrofacetDistribution for BeckmannDistribution {
     fn g1(&self, w: Dir3, w_h: Dir3) -> Float {
         if w.y() > 0.0 && Dir3::dot(w, w_h) > 0.0 {
             let cos_theta = w.y();
-            let a = self.alpha * (Float::sqrt(1.0 - cos_theta * cos_theta) / cos_theta);
-            (3.535 * a + 2.181 * a * a) / (1.0 + 2.276 * a + 2.577 * a * a)
+            let a = 1.0 / (self.alpha * (Float::sqrt(1.0 - cos_theta * cos_theta) / cos_theta));
+            if a < 1.6 {
+                (3.535 * a + 2.181 * a * a) / (1.0 + 2.276 * a + 2.577 * a * a)
+            } else {
+                1.0
+            }
         } else {
-            1.0
+            0.0
         }
     }
 
     fn d(&self, w_h: Dir3) -> Float {
         if w_h.y() > 0.0 {
             let cos_theta_squared = w_h.y() * w_h.y();
-            1.0 / (FloatConsts::PI
-                * self.alpha
-                * self.alpha
-                * cos_theta_squared
-                * cos_theta_squared)
+            let tan_theta_squared =
+                w_h.x() * w_h.x() / cos_theta_squared + w_h.z() * w_h.z() / cos_theta_squared;
+            Float::exp(-tan_theta_squared / (self.alpha * self.alpha))
+                / (FloatConsts::PI
+                    * self.alpha
+                    * self.alpha
+                    * cos_theta_squared
+                    * cos_theta_squared)
         } else {
             0.0
         }
     }
+
+    fn sample_d<R: Rng>(&self, rng: &mut R) -> Dir3 {
+        let n = Normal::new(0.0, self.alpha).unwrap();
+
+        let mx = n.sample(rng);
+        let mz = n.sample(rng);
+
+        let y = Float::sqrt(1.0 / (mx * mx + 1.0 + mz * mz));
+
+        Dir3::new(mx * y, y, mz * y).normalize()
+    }
 }
 
 pub trait MicrofacetDistribution {
@@ -83,5 +118,5 @@ pub trait MicrofacetDistribution {
 
     fn d(&self, w_h: Dir3) -> Float;
 
-    // fn sample_d(&self) -> Dir3;
+    fn sample_d<R: Rng>(&self, rng: &mut R) -> Dir3;
 }