From 471204223f79e4910bbd6742955472cd601ded34 Mon Sep 17 00:00:00 2001 From: Lahargoue Date: Tue, 18 Aug 2026 16:38:14 +0200 Subject: [PATCH] Make hidden geometry sampling order-independent --- mitransient/integrators/transientnlospath.py | 26 +++- tests/integration/test_hgs_order.py | 156 +++++++++++++++++++ 2 files changed, 180 insertions(+), 2 deletions(-) create mode 100644 tests/integration/test_hgs_order.py diff --git a/mitransient/integrators/transientnlospath.py b/mitransient/integrators/transientnlospath.py index de7c103..9f7582e 100644 --- a/mitransient/integrators/transientnlospath.py +++ b/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 diff --git a/tests/integration/test_hgs_order.py b/tests/integration/test_hgs_order.py new file mode 100644 index 0000000..49f8259 --- /dev/null +++ b/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)}" + )