Fix NEE in volume integrators

The integrators would not properly handle emitter sampling when the
medium's boundary shape was smooth. It would continue to account for the
medium's extinction even though the ray had escaped it.

src/integrators/volpath.cpp

@@ -324,10 +324,11 @@ class VolumetricPathIntegrator : public MonteCarloIntegrator<Float, Spectrum> {
 
 
     /// Samples an emitter in the scene and evaluates its attenuated contribution
+    template <typename Interaction>
     std::tuple<Spectrum, DirectionSample3f>
-    sample_emitter(const Interaction3f &ref_interaction, const Scene *scene,
-                   Sampler *sampler, MediumPtr medium, UInt32 channel,
-                   Mask active) const {
+    sample_emitter(const Interaction &ref_interaction, const Scene *scene,
+                   Sampler *sampler, MediumPtr medium,
+                   UInt32 channel, Mask active) const {
         Spectrum transmittance(1.0f);
 
         auto [ds, emitter_val] = scene->sample_emitter_direction(ref_interaction, sampler->next_2d(active), false, active);
@@ -340,6 +341,11 @@ class VolumetricPathIntegrator : public MonteCarloIntegrator<Float, Spectrum> {
 
         Ray3f ray = ref_interaction.spawn_ray(ds.d);
 
+        // Potentially escaping the medium if this is the current medium's boundary
+        if constexpr (std::is_convertible_v<Interaction, SurfaceInteraction3f>)
+            dr::masked(medium, ref_interaction.is_medium_transition()) =
+                ref_interaction.target_medium(ray.d);
+
         Float total_dist = 0.f;
         SurfaceInteraction3f si = dr::zeros<SurfaceInteraction3f>();
         Mask needs_intersection = true;

src/integrators/volpathmis.cpp

@@ -373,8 +373,9 @@ class VolpathMisIntegratorImpl final : public MonteCarloIntegrator<Float, Spectr
         return { result, valid_ray };
     }
 
+    template <typename Interaction>
     std::tuple<WeightMatrix, WeightMatrix, Spectrum, DirectionSample3f>
-    sample_emitter(const Interaction3f &ref_interaction, const Scene *scene,
+    sample_emitter(const Interaction &ref_interaction, const Scene *scene,
                    Sampler *sampler, MediumPtr medium,
                    const WeightMatrix &p_over_f, UInt32 channel,
                    Mask active) const {
@@ -392,6 +393,11 @@ class VolpathMisIntegratorImpl final : public MonteCarloIntegrator<Float, Spectr
 
         Ray3f ray = ref_interaction.spawn_ray(ds.d);
 
+        // Potentially escaping the medium if this is the current medium's boundary
+        if constexpr (std::is_convertible_v<Interaction, SurfaceInteraction3f>)
+            dr::masked(medium, ref_interaction.is_medium_transition()) =
+                ref_interaction.target_medium(ray.d);
+
         Float total_dist = 0.f;
         SurfaceInteraction3f si = dr::zeros<SurfaceInteraction3f>();
 

src/python/python/ad/integrators/prbvolpath.py

@@ -250,13 +250,14 @@ def sample(self,
                     sample_emitters = mei.medium.use_emitter_sampling()
                     specular_chain &= ~act_medium_scatter
                     specular_chain |= act_medium_scatter & ~sample_emitters
+
                     active_e_medium = act_medium_scatter & sample_emitters
                     active_e = active_e_surface | active_e_medium
-                    ref_interaction = dr.zeros(mi.Interaction3f)
-                    ref_interaction[act_medium_scatter] = mei
-                    ref_interaction[active_surface] = si
+
                     nee_sampler = sampler if is_primal else sampler.clone()
-                    emitted, ds = self.sample_emitter(ref_interaction, scene, sampler, medium, channel, active_e, mode=dr.ADMode.Primal)
+                    emitted, ds = self.sample_emitter(mei, si, active_e_medium, active_e_surface,
+                        scene, nee_sampler, medium, channel, active_e, mode=dr.ADMode.Primal)
+
                     # Query the BSDF for that emitter-sampled direction
                     bsdf_val, bsdf_pdf = bsdf.eval_pdf(ctx, si, si.to_local(ds.d), active_e_surface)
                     phase_val = phase.eval(phase_ctx, mei, ds.d, active_e_medium)
@@ -267,8 +268,10 @@ def sample(self,
                     L[active_e] += dr.detach(contrib if is_primal else -contrib)
 
                     if not is_primal:
-                        self.sample_emitter(ref_interaction, scene, nee_sampler,
-                            medium, channel, active_e, adj_emitted=contrib, δL=δL, mode=mode)
+                        self.sample_emitter(mei, si, active_e_medium, active_e_surface,
+                            scene, nee_sampler, medium, channel, active_e, adj_emitted=contrib,
+                            δL=δL, mode=mode)
+
                         if dr.grad_enabled(nee_weight) or dr.grad_enabled(emitted):
                             dr.backward(δL * contrib)
 
@@ -309,20 +312,26 @@ def sample(self,
 
         return L if is_primal else δL, valid_ray, L
 
-    def sample_emitter(self, ref_interaction, scene, sampler, medium, channel,
+    def sample_emitter(self, mei, si, active_medium, active_surface, scene, sampler, medium, channel,
                        active, adj_emitted=None, δL=None, mode=None):
 
         is_primal = mode == dr.ADMode.Primal
 
         active = mi.Bool(active)
-        medium = dr.select(active, medium, dr.zeros(mi.MediumPtr))
+
+        ref_interaction = dr.zeros(mi.Interaction3f)
+        ref_interaction[active_medium] = mei
+        ref_interaction[active_surface] = si
 
         ds, emitter_val = scene.sample_emitter_direction(ref_interaction, sampler.next_2d(active), False, active)
         ds = dr.detach(ds)
         invalid = dr.eq(ds.pdf, 0.0)
         emitter_val[invalid] = 0.0
         active &= ~invalid
 
+        medium = dr.select(active, medium, dr.zeros(mi.MediumPtr))
+        medium[(active_surface & si.is_medium_transition())] = si.target_medium(ds.d)
+
         ray = ref_interaction.spawn_ray(ds.d)
         total_dist = mi.Float(0.0)
         si = dr.zeros(mi.SurfaceInteraction3f)