Add support for per-emitter sampling weight (#721)

This adds support for a per emitter `sampling_weight`, similar to Mitsuba 0.6. This allows specifying a non-uniform sampling for direct illumination. If no weights are specified, we revert back to the already existing uniform sampling and skip any extra complexity due to the non-uniform case.
mitsuba-renderer · May 9, 2023 · 9a5f4c0 · 9a5f4c0
1 parent 96b219d
commit 9a5f4c0
Show file tree

Hide file tree

Showing 11 changed files with 195 additions and 6 deletions.
diff --git a/include/mitsuba/python/docstr.h b/include/mitsuba/python/docstr.h
@@ -2345,6 +2345,8 @@ static const char *__doc_mitsuba_Emitter_class = R"doc()doc";
 
 static const char *__doc_mitsuba_Emitter_flags = R"doc(Flags for all components combined.)doc";
 
+static const char *__doc_mitsuba_Emitter_sampling_weight = R"doc(The emitter's sampling weight.)doc";
+
 static const char *__doc_mitsuba_Emitter_is_environment = R"doc(Is this an environment map light emitter?)doc";
 
 static const char *__doc_mitsuba_Emitter_m_flags = R"doc(Combined flags for all properties of this emitter.)doc";

diff --git a/include/mitsuba/render/emitter.h b/include/mitsuba/render/emitter.h
@@ -53,16 +53,30 @@ template <typename Float, typename Spectrum>
 class MI_EXPORT_LIB Emitter : public Endpoint<Float, Spectrum> {
 public:
     MI_IMPORT_BASE(Endpoint, m_shape)
+    MI_IMPORT_TYPES()
 
     /// Is this an environment map light emitter?
     bool is_environment() const {
         return has_flag(m_flags, EmitterFlags::Infinite) &&
                !has_flag(m_flags, EmitterFlags::Delta);
     }
 
+    /// The emitter's sampling weight.
+    ScalarFloat sampling_weight() const { return m_sampling_weight; }
+
     /// Flags for all components combined.
     uint32_t flags(dr::mask_t<Float> /*active*/ = true) const { return m_flags; }
 
+    void traverse(TraversalCallback *callback) override;
+
+    void parameters_changed(const std::vector<std::string> &keys = {}) override;
+
+    /// Return whether the emitter parameters have changed
+    bool dirty() const { return m_dirty; }
+
+    /// Modify the emitter's "dirty" flag
+    void set_dirty(bool dirty) { m_dirty = dirty; }
+
     DRJIT_VCALL_REGISTER(Float, mitsuba::Emitter)
 
     MI_DECLARE_CLASS()
@@ -74,6 +88,12 @@ class MI_EXPORT_LIB Emitter : public Endpoint<Float, Spectrum> {
 protected:
     /// Combined flags for all properties of this emitter.
     uint32_t m_flags;
+
+    /// Sampling weight
+    ScalarFloat m_sampling_weight;
+
+    /// True if the emitters's parameters have changed
+    bool m_dirty = false;
 };
 
 MI_EXTERN_CLASS(Emitter)
@@ -96,6 +116,7 @@ DRJIT_VCALL_TEMPLATE_BEGIN(mitsuba::Emitter)
     DRJIT_VCALL_GETTER(flags, uint32_t)
     DRJIT_VCALL_GETTER(shape, const typename Class::Shape *)
     DRJIT_VCALL_GETTER(medium, const typename Class::Medium *)
+    DRJIT_VCALL_GETTER(sampling_weight, float)
 DRJIT_VCALL_TEMPLATE_END(mitsuba::Emitter)
 
 //! @}

diff --git a/include/mitsuba/render/scene.h b/include/mitsuba/render/scene.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <mitsuba/core/distr_1d.h>
 #include <mitsuba/core/spectrum.h>
 #include <mitsuba/render/emitter.h>
 #include <mitsuba/render/shapegroup.h>
@@ -554,6 +555,9 @@ class MI_EXPORT_LIB Scene : public Object {
 
     using ShapeKDTree = mitsuba::ShapeKDTree<Float, Spectrum>;
 
+    /// Updates the discrete distribution used to select an emitter
+    void update_emitter_sampling_distribution();
+
 protected:
     /// Acceleration data structure (IAS) (type depends on implementation)
     void *m_accel = nullptr;
@@ -572,6 +576,7 @@ class MI_EXPORT_LIB Scene : public Object {
     ref<Integrator> m_integrator;
     ref<Emitter> m_environment;
     ScalarFloat m_emitter_pmf;
+    std::unique_ptr<DiscreteDistribution<Float>> m_emitter_distr = nullptr;
 
     bool m_shapes_grad_enabled;
 };

diff --git a/src/emitters/area.cpp b/src/emitters/area.cpp
@@ -76,6 +76,7 @@ class AreaLight final : public Emitter<Float, Spectrum> {
     }
 
     void traverse(TraversalCallback *callback) override {
+        Base::traverse(callback);
         callback->put_object("radiance", m_radiance.get(), +ParamFlags::Differentiable);
     }
 

diff --git a/src/emitters/constant.cpp b/src/emitters/constant.cpp
@@ -68,6 +68,7 @@ class ConstantBackgroundEmitter final : public Emitter<Float, Spectrum> {
     }
 
     void traverse(TraversalCallback *callback) override {
+        Base::traverse(callback);
         callback->put_object("radiance", m_radiance.get(), +ParamFlags::Differentiable);
     }
 

diff --git a/src/emitters/directionalarea.cpp b/src/emitters/directionalarea.cpp
@@ -75,6 +75,7 @@ class DirectionalArea final : public Emitter<Float, Spectrum> {
     }
 
     void traverse(TraversalCallback *callback) override {
+        Base::traverse(callback);
         callback->put_object("radiance", m_radiance.get(), +ParamFlags::Differentiable);
     }
 

diff --git a/src/emitters/point.cpp b/src/emitters/point.cpp
@@ -86,6 +86,7 @@ class PointLight final : public Emitter<Float, Spectrum> {
     }
 
     void traverse(TraversalCallback *callback) override {
+        Base::traverse(callback);
         callback->put_parameter("position", (Point3f &) m_position.value(), +ParamFlags::NonDifferentiable);
         callback->put_object("intensity", m_intensity.get(), +ParamFlags::Differentiable);
     }
@@ -95,6 +96,7 @@ class PointLight final : public Emitter<Float, Spectrum> {
             m_position = m_position.value(); // update scalar part as well
             dr::make_opaque(m_position);
         }
+        Base::parameters_changed(keys);
     }
 
     std::pair<Ray3f, Spectrum> sample_ray(Float time, Float wavelength_sample,

diff --git a/src/render/emitter.cpp b/src/render/emitter.cpp
@@ -1,13 +1,27 @@
+#include <mitsuba/core/properties.h>
 #include <mitsuba/core/spectrum.h>
 #include <mitsuba/render/emitter.h>
 #include <mitsuba/render/endpoint.h>
 
 NAMESPACE_BEGIN(mitsuba)
 
 MI_VARIANT Emitter<Float, Spectrum>::Emitter(const Properties &props)
-    : Base(props) {}
+    : Base(props) {
+        m_sampling_weight = props.get<ScalarFloat>("sampling_weight", 1.0f);
+    }
 MI_VARIANT Emitter<Float, Spectrum>::~Emitter() { }
 
+MI_VARIANT
+void Emitter<Float, Spectrum>::traverse(TraversalCallback *callback) {
+    callback->put_parameter("sampling_weight", m_sampling_weight, +ParamFlags::NonDifferentiable);
+}
+
+MI_VARIANT
+void Emitter<Float, Spectrum>::parameters_changed(const std::vector<std::string> &keys) {
+    set_dirty(true);
+    Base::parameters_changed(keys);
+}
+
 MI_IMPLEMENT_CLASS_VARIANT(Emitter, Endpoint, "emitter")
 MI_INSTANTIATE_CLASS(Emitter)
 NAMESPACE_END(mitsuba)
diff --git a/src/render/python/emitter_v.cpp b/src/render/python/emitter_v.cpp
@@ -84,6 +84,7 @@ MI_PY_EXPORT(Emitter) {
     MI_PY_TRAMPOLINE_CLASS(PyEmitter, Emitter, Endpoint)
         .def(py::init<const Properties&>())
         .def_method(Emitter, is_environment)
+        .def_method(Emitter, sampling_weight)
         .def_method(Emitter, flags, "active"_a = true)
         .def_readwrite("m_needs_sample_2", &PyEmitter::m_needs_sample_2)
         .def_readwrite("m_needs_sample_3", &PyEmitter::m_needs_sample_3)
@@ -151,6 +152,7 @@ MI_PY_EXPORT(Emitter) {
                 D(Endpoint, sample_wavelengths))
         .def("flags", [](EmitterPtr ptr) { return ptr->flags(); }, D(Emitter, flags))
         .def("shape", [](EmitterPtr ptr) { return ptr->shape(); }, D(Endpoint, shape))
+        .def("sampling_weight", [](EmitterPtr ptr) { return ptr->sampling_weight(); }, D(Emitter, sampling_weight))
         .def("is_environment",
              [](EmitterPtr ptr) { return ptr->is_environment(); },
              D(Emitter, is_environment));

diff --git a/src/render/scene.cpp b/src/render/scene.cpp
@@ -84,11 +84,37 @@ MI_VARIANT Scene<Float, Spectrum>::Scene(const Properties &props) {
     m_emitters_dr = dr::load<DynamicBuffer<EmitterPtr>>(
         m_emitters.data(), m_emitters.size());
 
-    m_emitter_pmf = m_emitters.empty() ? 0.f : (1.f / m_emitters.size());
+    update_emitter_sampling_distribution();
 
     m_shapes_grad_enabled = false;
 }
 
+MI_VARIANT
+void Scene<Float, Spectrum>::update_emitter_sampling_distribution() {
+    // Check if we need to use non-uniform emitter sampling.
+    bool non_uniform_sampling = false;
+    for (auto &e : m_emitters) {
+        if (e->sampling_weight() != ScalarFloat(1.0)) {
+            non_uniform_sampling = true;
+            break;
+        }
+    }
+    size_t n_emitters = m_emitters.size();
+    if (non_uniform_sampling) {
+        std::unique_ptr<ScalarFloat[]> sample_weights(new ScalarFloat[n_emitters]);
+        for (size_t i = 0; i < n_emitters; ++i)
+            sample_weights[i] = m_emitters[i]->sampling_weight();
+        m_emitter_distr = std::make_unique<DiscreteDistribution<Float>>(
+            sample_weights.get(), n_emitters);
+    } else {
+        // By default use uniform sampling with constant PMF
+        m_emitter_pmf = m_emitters.empty() ? 0.f : (1.f / n_emitters);
+    }
+    // Clear emitter's dirty flag
+    for (auto &e : m_emitters)
+        e->set_dirty(false);
+}
+
 MI_VARIANT Scene<Float, Spectrum>::~Scene() {
     if constexpr (dr::is_cuda_v<Float>)
         accel_release_gpu();
@@ -169,6 +195,11 @@ Scene<Float, Spectrum>::sample_emitter(Float index_sample, Mask active) const {
             return { UInt32(-1), 0.f, index_sample };
     }
 
+    if (m_emitter_distr != nullptr) {
+        auto [index, reused_sample, pmf] = m_emitter_distr->sample_reuse_pmf(index_sample);
+        return {index, dr::rcp(pmf), reused_sample};
+    }
+
     uint32_t emitter_count = (uint32_t) m_emitters.size();
     ScalarFloat emitter_count_f = (ScalarFloat) emitter_count;
     Float index_sample_scaled = index_sample * emitter_count_f;
@@ -178,9 +209,12 @@ Scene<Float, Spectrum>::sample_emitter(Float index_sample, Mask active) const {
     return { index, emitter_count_f, index_sample_scaled - Float(index) };
 }
 
-MI_VARIANT Float Scene<Float, Spectrum>::pdf_emitter(UInt32 /*index*/,
-                                                      Mask /*active*/) const {
-    return m_emitter_pmf;
+MI_VARIANT Float Scene<Float, Spectrum>::pdf_emitter(UInt32 index,
+                                                      Mask active) const {
+    if (m_emitter_distr == nullptr)
+        return m_emitter_pmf;
+    else
+        return m_emitter_distr->eval_pmf_normalized(index, active);
 }
 
 MI_VARIANT std::tuple<typename Scene<Float, Spectrum>::Ray3f, Spectrum,
@@ -283,7 +317,12 @@ Scene<Float, Spectrum>::pdf_emitter_direction(const Interaction3f &ref,
                                               const DirectionSample3f &ds,
                                               Mask active) const {
     MI_MASK_ARGUMENT(active);
-    return ds.emitter->pdf_direction(ref, ds, active) * m_emitter_pmf;
+    Float emitter_pmf;
+    if (m_emitter_distr == nullptr)
+        emitter_pmf = m_emitter_pmf;
+    else
+        emitter_pmf = ds.emitter->sampling_weight() * m_emitter_distr->normalization();
+    return ds.emitter->pdf_direction(ref, ds, active) * emitter_pmf;
 }
 
 MI_VARIANT Spectrum Scene<Float, Spectrum>::eval_emitter_direction(
@@ -334,6 +373,15 @@ MI_VARIANT void Scene<Float, Spectrum>::parameters_changed(const std::vector<std
         if (m_shapes_grad_enabled)
             break;
     }
+
+    // Check if emitters were modified and we potentially need to update
+    // the emitter sampling distribution.
+    for (auto &e : m_emitters) {
+        if (e->dirty()) {
+            update_emitter_sampling_distribution();
+            break;
+        }
+    }
 }
 
 MI_VARIANT std::string Scene<Float, Spectrum>::to_string() const {

diff --git a/src/render/tests/test_scene.py b/src/render/tests/test_scene.py
@@ -154,3 +154,95 @@ def render():
 
     import drjit as dr
     dr.eval(pi)
+
+
+def test05_test_uniform_emitter_pdf(variants_all_backends_once):
+    scene = mi.load_dict({
+        'type': 'scene',
+        'shape': {'type': 'sphere', 'emitter': {'type':'area'}},
+        'emitter_0': {'type':'point', 'sampling_weight': 1.0},
+        'emitter_1': {'type':'constant', 'sampling_weight': 1.0},
+    })
+    assert dr.allclose(scene.pdf_emitter(0), 1.0 / 3.0)
+    assert dr.allclose(scene.pdf_emitter(1), 1.0 / 3.0)
+    assert dr.allclose(scene.pdf_emitter(2), 1.0 / 3.0)
+
+
+def test06_test_nonuniform_emitter_pdf(variants_all_backends_once):
+    import numpy as np
+
+    weights = [1.3, 3.8, 0.0]
+    scene = mi.load_dict({
+        'type': 'scene',
+        'shape': {'type': 'sphere', 'emitter': {'type':'area', 'sampling_weight': weights[0]}},
+        'emitter_0': {'type':'point', 'sampling_weight': weights[1]},
+        'emitter_1': {'type':'constant', 'sampling_weight': weights[2]},
+    })
+    weights = [emitter.sampling_weight() for emitter in scene.emitters()]
+    pdf = np.array(weights) / np.sum(weights)
+    assert dr.allclose(scene.pdf_emitter(0), pdf[0])
+    assert dr.allclose(scene.pdf_emitter(1), pdf[1])
+    assert dr.allclose(scene.pdf_emitter(2), pdf[2])
+
+
+def test07_test_uniform_emitter_sampling(variants_all_backends_once):
+    scene = mi.load_dict({
+        'type': 'scene',
+        'shape': {'type': 'sphere', 'emitter': {'type':'area', 'sampling_weight': 1.0}},
+        'emitter_0': {'type':'point', 'sampling_weight': 1.0},
+        'emitter_1': {'type':'constant', 'sampling_weight': 1.0},
+    })
+    sample = 0.6
+    index, weight, reused_sample = scene.sample_emitter(sample)
+    assert dr.allclose(index, 1)
+    assert dr.allclose(weight, 3.0)
+    assert dr.allclose(reused_sample, sample * 3 - 1)
+
+
+def test08_test_nonuniform_emitter_sampling(variants_all_backends_once):
+    sample = 0.75
+    weights = [1.3, 3.8, 0.0]
+    scene = mi.load_dict({
+        'type': 'scene',
+        'shape': {'type': 'sphere', 'emitter': {'type':'area', 'sampling_weight': weights[0]}},
+        'emitter_0': {'type':'point', 'sampling_weight': weights[1]},
+        'emitter_1': {'type':'constant', 'sampling_weight': weights[2]},
+    })
+    index, weight, reused_sample = scene.sample_emitter(sample)
+    distr = mi.DiscreteDistribution([emitter.sampling_weight() for emitter in scene.emitters()])
+    ref_index, ref_reused_sample, ref_pmf = distr.sample_reuse_pmf(sample)
+    assert dr.allclose(index, ref_index)
+    assert dr.allclose(weight, 1.0 / ref_pmf)
+    assert dr.allclose(reused_sample, ref_reused_sample)
+
+
+def test09_test_emitter_sampling_weight_update(variants_all_backends_once):
+    import numpy as np
+
+    weights = [2.0, 1.0, 0.5]
+    scene = mi.load_dict({
+        'type': 'scene',
+        'emitter_0': {'type':'point', 'sampling_weight': weights[0]},
+        'emitter_1': {'type':'constant', 'sampling_weight': weights[1]},
+        'emitter_2': {'type':'directional', 'sampling_weight': weights[2]},
+    })
+
+    params = mi.traverse(scene)
+    params['emitter_0.sampling_weight'] = 0.8
+    params['emitter_1.sampling_weight'] = 0.05
+    params['emitter_2.sampling_weight'] = 1.2
+    params.update()
+
+    sample = 0.75
+    weights = [emitter.sampling_weight() for emitter in scene.emitters()]
+    distr = mi.DiscreteDistribution(weights)
+    index, weight, reused_sample = scene.sample_emitter(sample)
+    ref_index, ref_reused_sample, ref_pmf = distr.sample_reuse_pmf(sample)
+    assert dr.allclose(index, ref_index)
+    assert dr.allclose(weight, 1.0 / ref_pmf)
+    assert dr.allclose(reused_sample, ref_reused_sample)
+
+    pdf = np.array(weights) / np.sum(weights)
+    assert dr.allclose(scene.pdf_emitter(0), pdf[0])
+    assert dr.allclose(scene.pdf_emitter(1), pdf[1])
+    assert dr.allclose(scene.pdf_emitter(2), pdf[2])