Add configurable maximum depth for hidden geometry sampling

mitransient/integrators/transientnlospath.py

@@ -151,6 +151,13 @@ class TransientNLOSPath(TransientADIntegrator):
          If True, ray directions are sampled using the Hidden Geometry Sampling technique.
          See [Royo2022] for more information about Hidden Geometry Sampling (default: false)
 
+     * - nlos_hidden_geometry_sampling_max_depth
+       - |int|
+       - Limits Hidden Geometry Sampling to path vertices with depth smaller than
+         this value. A value of 1 therefore applies HGS only at the first path
+         vertex (depth 0). A value of -1 disables the depth limit and preserves
+         the previous behavior. (default: -1)
+
      * - nlos_hidden_geometry_sampling_do_rroulette
        - |bool|
        - Only relevant when `nlos_hidden_geometry_sampling` is True.
@@ -235,6 +242,12 @@ def __init__(self, props: mi.Properties):
             'nlos_laser_sampling', False)
         self.hg_sampling: bool = props.get(
             'nlos_hidden_geometry_sampling', False)
+        self.hg_sampling_max_depth: int = props.get(
+            'nlos_hidden_geometry_sampling_max_depth', -1)
+        if self.hg_sampling_max_depth == 0 or self.hg_sampling_max_depth < -1:
+            raise RuntimeError(
+                'nlos_hidden_geometry_sampling_max_depth must be -1 '
+                'or a positive integer')
         self.hg_sampling_do_rroulette = (
             props.get('nlos_hidden_geometry_sampling_do_rroulette', False)
             and
@@ -248,6 +261,17 @@ def __init__(self, props: mi.Properties):
         self.account_first_and_last_bounces: bool = props.get(
             'account_first_and_last_bounces', False)
 
+    def _apply_hg_sampling_depth_limit(
+            self, depth: mi.UInt32, do_hg_sample: mi.Bool,
+            pdf_bsdf_method: mi.Float) -> Tuple[mi.Bool, mi.Float]:
+        if self.hg_sampling_max_depth != -1:
+            hg_sampling_active = depth < self.hg_sampling_max_depth
+            do_hg_sample &= hg_sampling_active
+            pdf_bsdf_method = dr.select(
+                hg_sampling_active, pdf_bsdf_method, mi.Float(1.0))
+
+        return do_hg_sample, pdf_bsdf_method
+
     def prepare(self, scene: mi.Scene, sensor: mi.Sensor, seed: mi.UInt32, spp: int, aovs: List):
         # prepare laser sampling
         if isinstance(sensor, int):
@@ -808,6 +832,10 @@ def sample(self,
                 do_hg_sample = mi.Bool(self.hg_sampling)
                 pdf_bsdf_method = mi.Float(1.0)
 
+            do_hg_sample, pdf_bsdf_method = \
+                self._apply_hg_sampling_depth_limit(
+                    depth, do_hg_sample, pdf_bsdf_method)
+
             active_hg = active_next & do_hg_sample
             bsdf_sample_hg, bsdf_weight_hg = self.hidden_geometry_sample(
                 scene, sampler, bsdf,

tests/integration/test_nlos_hgs_max_depth.py

@@ -0,0 +1,138 @@
+import drjit as dr
+import mitsuba as mi
+import pytest
+
+mi.set_variant('llvm_ad_rgb')
+
+import mitransient  # noqa: F401
+
+
+def _integrator(hg_max_depth=None):
+    config = {
+        'type': 'transient_nlos_path',
+        'nlos_hidden_geometry_sampling': True,
+    }
+
+    if hg_max_depth is not None:
+        config['nlos_hidden_geometry_sampling_max_depth'] = hg_max_depth
+
+    return mi.load_dict(config)
+
+
+def _apply(integrator, depth, do_hg_sample=True, method_pdf=0.5):
+    do_hg_sample, method_pdf = (
+        integrator._apply_hg_sampling_depth_limit(
+            mi.UInt32(depth),
+            mi.Bool(do_hg_sample),
+            mi.Float(method_pdf),
+        )
+    )
+
+    dr.eval(do_hg_sample, method_pdf)
+
+    return bool(do_hg_sample[0]), float(method_pdf[0])
+
+
+def test_hg_sampling_max_depth_defaults_to_unlimited():
+    integrator = _integrator()
+
+    assert integrator.hg_sampling_max_depth == -1
+
+    # The historical behavior must be preserved at arbitrary depth.
+    do_hg_sample, method_pdf = _apply(
+        integrator,
+        depth=10,
+        do_hg_sample=True,
+        method_pdf=0.5,
+    )
+
+    assert do_hg_sample
+    assert method_pdf == pytest.approx(0.5)
+
+
+def test_hg_sampling_max_depth_one_only_allows_depth_zero():
+    integrator = _integrator(1)
+
+    do_hg_sample, method_pdf = _apply(
+        integrator,
+        depth=0,
+        do_hg_sample=True,
+        method_pdf=0.5,
+    )
+
+    assert do_hg_sample
+    assert method_pdf == pytest.approx(0.5)
+
+    do_hg_sample, method_pdf = _apply(
+        integrator,
+        depth=1,
+        do_hg_sample=True,
+        method_pdf=0.5,
+    )
+
+    assert not do_hg_sample
+
+    # Once HGS is disallowed, BSDF sampling is no longer selected with
+    # Russian Roulette and therefore has method probability 1.
+    assert method_pdf == pytest.approx(1.0)
+
+
+def test_hg_sampling_max_depth_n_allows_depths_below_n():
+    integrator = _integrator(3)
+
+    for depth in (0, 1, 2):
+        do_hg_sample, method_pdf = _apply(
+            integrator,
+            depth=depth,
+            do_hg_sample=True,
+            method_pdf=0.5,
+        )
+
+        assert do_hg_sample
+        assert method_pdf == pytest.approx(0.5)
+
+    do_hg_sample, method_pdf = _apply(
+        integrator,
+        depth=3,
+        do_hg_sample=True,
+        method_pdf=0.5,
+    )
+
+    assert not do_hg_sample
+    assert method_pdf == pytest.approx(1.0)
+
+
+def test_hg_sampling_depth_limit_preserves_bsdf_choice():
+    integrator = _integrator(2)
+
+    # At an allowed depth, a BSDF choice made by the existing
+    # HGS/BSDF Russian Roulette remains a BSDF choice with p=0.5.
+    do_hg_sample, method_pdf = _apply(
+        integrator,
+        depth=1,
+        do_hg_sample=False,
+        method_pdf=0.5,
+    )
+
+    assert not do_hg_sample
+    assert method_pdf == pytest.approx(0.5)
+
+    # Beyond the HGS depth limit, BSDF sampling becomes the only method.
+    do_hg_sample, method_pdf = _apply(
+        integrator,
+        depth=2,
+        do_hg_sample=False,
+        method_pdf=0.5,
+    )
+
+    assert not do_hg_sample
+    assert method_pdf == pytest.approx(1.0)
+
+
+@pytest.mark.parametrize("value", [0, -2])
+def test_hg_sampling_max_depth_rejects_invalid_values(value):
+    with pytest.raises(
+        RuntimeError,
+        match="must be -1 or a positive integer",
+    ):
+        _integrator(value)