Make hidden geometry sampling independent of scene shape ordering

mitransient/integrators/transientnlospath.py

@@ -274,8 +274,27 @@ def prepare(self, scene: mi.Scene, sensor: mi.Sensor, seed: mi.UInt32, spp: int,
 
             # self.hidden_geometries_distribution = mi.DiscreteDistribution(pdf_hidden_area)
             # ------------------------------------------------------------------------------------
+            scene_shapes = scene.shapes()
+            shape_ids = [shape.id() for shape in scene_shapes]
+
+            # Mitsuba does not guarantee that scene.shapes() preserves a stable
+            # ordering after scene optimization. When possible, use shape IDs to
+            # define a canonical ordering for the HGS discrete distribution.
+            #
+            # Keep the original ordering as a fallback for scenes without unique,
+            # non-empty IDs to preserve the existing behavior.
+            if (all(shape_ids)
+                    and len(set(shape_ids)) == len(shape_ids)):
+                shape_indices = sorted(
+                    range(len(scene_shapes)),
+                    key=lambda i: shape_ids[i]
+                )
+            else:
+                shape_indices = list(range(len(scene_shapes)))
+
             surface_areas = []
-            for shape in scene.shapes():
+            for shape_index in shape_indices:
+                shape = scene_shapes[shape_index]
                 surface_areas.append(
                     0.0 if (shape == sensor.get_shape()
                             and not self.hg_sampling_includes_relay_wall) else shape.surface_area()[0]
@@ -288,6 +307,7 @@ def prepare(self, scene: mi.Scene, sensor: mi.Sensor, seed: mi.UInt32, spp: int,
                 raise AssertionError('Hidden geometry sampling is activated, '
                                      'but the hidden geometry in the scene has zero surface area?')
 
+            self.hidden_geometry_shape_indices = mi.UInt(shape_indices)
             self.hidden_geometries_distribution = mi.DiscreteDistribution(
                 surface_areas)
 
@@ -422,8 +442,10 @@ def _sample_hidden_geometry_position(
                 sample2.x, active)
         sample2.x = new_sample
 
+        shape_index = dr.gather(
+            mi.UInt, self.hidden_geometry_shape_indices, index, active)
         shape: mi.ShapePtr = dr.gather(
-            mi.ShapePtr, scene.shapes_dr(), index, active)
+            mi.ShapePtr, scene.shapes_dr(), shape_index, active)
         ps = shape.sample_position(ref.time, sample2, active)
         ps.pdf *= shape_pdf
 

tests/integration/test_hgs_order.py

@@ -0,0 +1,156 @@
+import pytest
+import drjit as dr
+import mitsuba as mi
+
+mi.set_variant("llvm_ad_rgb")
+
+import mitransient  # noqa: F401
+
+
+def integrator():
+    return {
+        "type": "transient_nlos_path",
+        "max_depth": -1,
+        "nlos_laser_sampling": False,
+        "nlos_hidden_geometry_sampling": True,
+        "nlos_hidden_geometry_sampling_includes_relay_wall": False,
+        "temporal_filter": "box",
+    }
+
+
+def sensor():
+    return {
+        "type": "nlos_capture_meter",
+        "sampler": {
+            "type": "independent",
+            "sample_count": 1,
+            "seed": 0,
+        },
+        "sensor_origin": [-0.5, 0.0, 0.25],
+        "film": {
+            "type": "transient_hdr_film",
+            "width": 1,
+            "height": 1,
+            "temporal_bins": 8,
+            "bin_width_opl": 0.1,
+            "start_opl": 0.0,
+            "rfilter": {"type": "box"},
+        },
+    }
+
+
+def relay_wall():
+    return {
+        "type": "rectangle",
+        "bsdf": {
+            "type": "diffuse",
+            "reflectance": {
+                "type": "rgb",
+                "value": [1.0, 1.0, 1.0],
+            },
+        },
+        "sensor": sensor(),
+    }
+
+
+def hidden_shape(center_x, scale):
+    T = mi.ScalarTransform4f
+    return {
+        "type": "rectangle",
+        "to_world": (
+            T.translate([center_x, 0.0, 1.0])
+            @ T.scale([scale, scale, scale])
+        ),
+        "bsdf": {
+            "type": "diffuse",
+            "reflectance": {
+                "type": "rgb",
+                "value": [0.5, 0.5, 0.5],
+            },
+        },
+    }
+
+
+def emitter():
+    T = mi.ScalarTransform4f
+    return {
+        "type": "projector",
+        "to_world": T.translate([-0.5, 0.0, 0.25]),
+        "irradiance": {
+            "type": "rgb",
+            "value": [1.0, 1.0, 1.0],
+        },
+        "fov": 20.0,
+    }
+
+
+def make_scene(order):
+    scene_dict = {
+        "type": "scene",
+        "integrator": integrator(),
+        "relay_wall": relay_wall(),
+        "laser": emitter(),
+    }
+
+    objects = {
+        "shape_a": hidden_shape(-2.0, 0.5),
+        "shape_b": hidden_shape(+2.0, 1.0),
+    }
+
+    for name in order:
+        scene_dict[name] = objects[name]
+
+    return mi.load_dict(
+        scene_dict,
+        parallel=False,
+        optimize=False,
+    )
+
+
+def sample_hidden_geometry_position(order):
+    scene = make_scene(order)
+    integrator = scene.integrator()
+
+    integrator.prepare(
+        scene=scene,
+        sensor=0,
+        seed=mi.UInt32(0),
+        spp=1,
+        aovs=[],
+    )
+
+    ref = dr.zeros(mi.Interaction3f)
+    sample2 = mi.Point2f(
+        mi.Float(0.10),
+        mi.Float(0.50),
+    )
+
+    ps = integrator._sample_hidden_geometry_position(
+        ref,
+        scene,
+        sample2,
+        mi.Bool(True),
+    )
+
+    dr.eval(ps.p, ps.pdf)
+
+    p = (
+        float(ps.p.x[0]),
+        float(ps.p.y[0]),
+        float(ps.p.z[0]),
+    )
+    pdf = float(ps.pdf[0])
+
+    return p, pdf
+
+
+def test_hidden_geometry_sampling_is_independent_of_scene_shape_order():
+    p_ab, pdf_ab = sample_hidden_geometry_position(["shape_a", "shape_b"])
+    p_ba, pdf_ba = sample_hidden_geometry_position(["shape_b", "shape_a"])
+
+    assert pdf_ab == pytest.approx(pdf_ba)
+    assert p_ab == pytest.approx(p_ba), (
+        "Hidden geometry sampling depends on scene.shapes() ordering. "
+        f"sample(order=['shape_a','shape_b'])={list(p_ab)}, "
+        f"sample(order=['shape_b','shape_a'])={list(p_ba)}"
+    )