Code optimization for vpt

bxdf/brdf.py

@@ -292,7 +292,7 @@ def sample_new_rays(self, incid: vec3, normal: vec3):
         return ret_dir, ret_spec, pdf
 
     @ti.func
-    def get_pdf(self, outdir: vec3, normal: vec3, incid: vec3, medium):
+    def get_pdf(self, outdir: vec3, normal: vec3, incid: vec3):
         """ 
             Solid angle PDF for a specific incident direction - BRDF sampling
             Some PDF has nothing to do with backward incid (from eye to the surface), like diffusive 

renderer/vanilla_renderer.py

@@ -12,7 +12,7 @@
 from emitters.abtract_source import LightSource
 
 from scene.obj_desc import ObjDescriptor
-from sampler.general_sampling import mis_weight
+from sampler.general_sampling import balance_heuristic
 
 @ti.data_oriented
 class Renderer(PathTracer):
@@ -72,8 +72,8 @@ def render(self):
                     if ti.static(self.use_mis):
                         mis_w = 1.0
                         if not emitter.is_delta:
-                            bsdf_pdf = self.get_pdf(obj_id, light_dir, normal, ray_d, self.world.medium)
-                            mis_w    = mis_weight(light_pdf, bsdf_pdf)
+                            bsdf_pdf = self.get_pdf(obj_id, light_dir, normal, ray_d)
+                            mis_w    = balance_heuristic(light_pdf, bsdf_pdf)
                         direct_int  += direct_spec * shadow_int * mis_w / emitter_pdf
                     else:
                         direct_int += direct_spec * shadow_int / emitter_pdf
@@ -98,7 +98,7 @@ def render(self):
                         emitter_pdf = 0.0
                         if hit_light >= 0 and self.is_delta(obj_id) == 0:
                             emitter_pdf = self.src_field[hit_light].solid_angle_pdf(ray_d, normal, min_depth)
-                        emission_weight = mis_weight(ray_pdf, emitter_pdf)
+                        emission_weight = balance_heuristic(ray_pdf, emitter_pdf)
 
             self.color[i, j] += ti.select(ti.math.isnan(color), 0., color)
             self.pixels[i, j] = self.color[i, j] / self.cnt[None]

renderer/vpt.py

@@ -14,9 +14,7 @@
 from tracer.path_tracer import PathTracer
 from emitters.abtract_source import LightSource
 
-from bxdf.medium import Medium
 from scene.obj_desc import ObjDescriptor
-from sampler.general_sampling import mis_weight
 
 """
     VPT todo: 

sampler/general_sampling.py

@@ -10,7 +10,7 @@
 from taichi.math import vec3
 
 __all__ = ['cosine_hemisphere', 'uniform_hemisphere', 'sample_triangle', 
-            'mis_weight', 'mod_phong_hemisphere', 'frensel_hemisphere', 'random_rgb']
+            'balance_heuristic', 'mod_phong_hemisphere', 'frensel_hemisphere', 'random_rgb']
 
 pi_inv = 1. / tm.pi
 
@@ -93,7 +93,7 @@ def sample_triangle(dv1: vec3, dv2: vec3):
     return triangle_pt
 
 @ti.func
-def mis_weight(pdf_a: float, pdf_b: float):
+def balance_heuristic(pdf_a: float, pdf_b: float):
     """ Balanced heuristic function for MIS weight computation """
     return ti.select(pdf_a > 1e-7, pdf_a / (pdf_a + pdf_b), 0.)
 

tracer/path_tracer.py

@@ -128,12 +128,12 @@ def eval(self, idx: int, incid: vec3, out: vec3, normal: vec3, is_mi: int, in_fr
         return ret_spec
 
     @ti.func
-    def get_pdf(self, idx: int, outdir: vec3, normal: vec3, incid: vec3, medium):
+    def get_pdf(self, idx: int, outdir: vec3, normal: vec3, incid: vec3):
         pdf = 0.
         if ti.is_active(self.brdf_nodes, idx):      # active means the object is attached to BRDF
-            pdf = self.brdf_field[idx].get_pdf(outdir, normal, incid, medium)
+            pdf = self.brdf_field[idx].get_pdf(outdir, normal, incid)
         else:
-            pdf = self.bsdf_field[idx].get_pdf(outdir, normal, incid, medium)
+            pdf = self.bsdf_field[idx].get_pdf(outdir, normal, incid, self.world.medium)
         return pdf
 
     @ti.func
@@ -147,6 +147,7 @@ def is_delta(self, idx: int):
 
     @ti.func
     def is_scattering(self, idx: int):           # check if the object with index idx is a scattering medium
+        # FIXME: if sigma_t is too small, set the scattering medium to det-refract
         is_scattering = False
         if not ti.is_active(self.brdf_nodes, idx):
             is_scattering = self.bsdf_field[idx].medium.is_scattering()

tracer/tracer_base.py

@@ -108,25 +108,24 @@ def aabb_test(self, aabb_idx, ray: vec3, ray_o: vec3):
         t2 = ti.max(t_min, t_max)
         t_near  = ti.max(ti.max(t1.x, t1.y), t1.z)
         t_far   = ti.min(ti.min(t2.x, t2.y), t2.z)
-        return t_near < t_far
+        return t_near < t_far, t_near
 
     @ti.func
     def ray_intersect(self, ray, start_p, min_depth = -1.0):
         """
             Intersection function and mesh organization can be reused
+            TODO: BVH can further accelerate this method
         """
         obj_id = -1
         tri_id = -1
-        if min_depth > 0.0:
-            min_depth -= 5e-4
-        else:
-            min_depth = 1e7
+        min_depth = ti.select(min_depth > 0.0, min_depth - 5e-4, 1e7)
         for aabb_idx in range(self.num_objects):
-            if self.aabb_test(aabb_idx, ray, start_p) == False: continue
+            aabb_intersect, t_near = self.aabb_test(aabb_idx, ray, start_p)
+            if aabb_intersect == False: continue
+            if t_near > min_depth: continue
             tri_num = self.mesh_cnt[aabb_idx]
             if tri_num:
                 for mesh_idx in range(tri_num):
-                    normal = self.normals[aabb_idx, mesh_idx]   # back-face culling removed
                     # Sadly, Taichi does not support slicing. I think this restrict the use cases of Matrix field
                     p1 = self.meshes[aabb_idx, mesh_idx, 0]
                     vec1 = self.precom_vec[aabb_idx, mesh_idx, 0]
@@ -147,10 +146,8 @@ def ray_intersect(self, ray, start_p, min_depth = -1.0):
                 c2ray_norm = center_norm2 - proj_norm ** 2  # center to ray distance ** 2
                 if c2ray_norm >= radius2: continue
                 ray_t = proj_norm
-                if center_norm2 > radius2 + 5e-4:
-                    ray_t -= ti.sqrt(radius2 - c2ray_norm)
-                else:
-                    ray_t += ti.sqrt(radius2 - c2ray_norm)
+                ray_cut = ti.sqrt(radius2 - c2ray_norm)
+                ray_t += ti.select(center_norm2 > radius2 + 5e-4, -ray_cut, ray_cut)
                 if ray_t > 5e-4 and ray_t < min_depth:
                     min_depth = ray_t
                     obj_id = aabb_idx
@@ -172,13 +169,12 @@ def does_intersect(self, ray, start_p, depth = -1.0) -> bool:
             Taichi does not support compile-time branching. Actually it does, but not flexible, for e.g \\
             C++ supports compile-time branching via template parameter, but Taichi can not "pass" compile-time constants
         """
-        if depth > 0.0:
-            depth -= 5e-4
-        else:
-            depth = 1e7
+        depth = ti.select(depth > 0.0, depth - 5e-4, 1e7)
         flag = False
         for aabb_idx in range(self.num_objects):
-            if self.aabb_test(aabb_idx, ray, start_p) == False: continue
+            aabb_intersect, t_near =  self.aabb_test(aabb_idx, ray, start_p)
+            if aabb_intersect == False: continue
+            if t_near > depth: continue
             tri_num = self.mesh_cnt[aabb_idx]
             if tri_num:
                 for mesh_idx in range(tri_num):