environment map optimization support

docs/examples/10_inverse_rendering/invert_bunny.py

@@ -0,0 +1,68 @@
+# Simple inverse rendering example: render a cornell box reference image, then
+# then replace one of the scene parameters and try to recover it using
+# differentiable rendering and gradient-based optimization.
+
+import enoki as ek
+import mitsuba
+mitsuba.set_variant('gpu_autodiff_rgb')
+
+from mitsuba.core import Float
+from mitsuba.core.xml import load_file
+from mitsuba.python.util import traverse
+from mitsuba.python.autodiff import render, write_bitmap, Adam
+import time
+
+# Load example scene
+scene = load_file('bunny.xml')
+
+# Find differentiable scene parameters
+params = traverse(scene)
+
+# Make a backup copy
+param_res = params['my_envmap.resolution']
+param_ref = Float(params['my_envmap.data'])
+
+# Discard all parameters except for one we want to differentiate
+params.keep(['my_envmap.data'])
+
+# Render a reference image (no derivatives used yet)
+image_ref = render(scene, spp=1)
+crop_size = scene.sensors()[0].film().crop_size()
+write_bitmap('out_ref.png', image_ref, crop_size)
+
+# Change to a uniform white lighting environment
+params['my_envmap.data'] = ek.full(Float, 1.0, len(params['my_envmap.data']))
+params.update()
+
+# Construct an Adam optimizer that will adjust the parameters 'params'
+#  opt = Adam(params, lr=.02)
+opt = Adam(params, lr=.02)
+print(opt)
+
+time_a = time.time()
+
+iterations = 100
+for it in range(iterations):
+    # Perform a differentiable rendering of the scene
+    image = render(scene, optimizer=opt, unbiased=True, spp=1)
+    write_bitmap('out_%03i.png' % it, image, crop_size)
+    write_bitmap('envmap_%03i.png' % it, params['my_envmap.data'],
+                 (param_res[1], param_res[0]))
+
+    # Objective: MSE between 'image' and 'image_ref'
+    ob_val = ek.hsum(ek.sqr(image - image_ref)) / len(image)
+
+    # Back-propagate errors to input parameters
+    ek.backward(ob_val)
+
+    # Optimizer: take a gradient step
+    opt.step()
+
+    # Compare iterate against ground-truth value
+    err_ref = ek.hsum(ek.sqr(param_ref - params['my_envmap.data']))
+    print('Iteration %03i: error=%g' % (it, err_ref[0]), end='\r')
+
+time_b = time.time()
+
+print()
+print('%f ms per iteration' % (((time_b - time_a) * 1000) / iterations))

include/mitsuba/render/bsdf.h

@@ -432,7 +432,7 @@ class MTS_EXPORT_RENDER BSDF : public Object {
     }
 
     /// Return a string identifier
-    std::string id() const override { return m_id; }
+    std::string id() const override;
 
     /// Return a human-readable representation of the BSDF
     std::string to_string() const override = 0;

include/mitsuba/render/endpoint.h

@@ -216,6 +216,9 @@ class MTS_EXPORT_RENDER Endpoint : public Object {
      */
     virtual void set_scene(const Scene *scene);
 
+    /// Return a string identifier
+    std::string id() const override;
+
     //! @}
     // =============================================================
 
@@ -233,6 +236,7 @@ class MTS_EXPORT_RENDER Endpoint : public Object {
     Shape *m_shape = nullptr;
     bool m_needs_sample_2 = true;
     bool m_needs_sample_3 = true;
+    std::string m_id;
 };
 
 MTS_EXTERN_CLASS_RENDER(Endpoint)

src/emitters/envmap.cpp

@@ -109,7 +109,7 @@ class EnvironmentMapEmitter final : public Emitter<Float, Spectrum> {
                 }
 
                 *lum_ptr++ = lum * sin_theta;
-                store_unaligned(ptr, coeff);
+                store(ptr, coeff);
                 ptr += 4;
             }
         }
@@ -124,7 +124,37 @@ class EnvironmentMapEmitter final : public Emitter<Float, Spectrum> {
     }
 
     void parameters_changed() override {
-        // TODO update warp for better importance sampling when data has changed
+        m_data.managed();
+
+        std::unique_ptr<ScalarFloat[]> luminance(new ScalarFloat[hprod(m_resolution)]);
+
+        ScalarFloat *ptr     = (ScalarFloat *) m_data.data(),
+                    *lum_ptr = (ScalarFloat *) luminance.get();
+
+
+        for (size_t y = 0; y < m_resolution.y(); ++y) {
+            ScalarFloat sin_theta =
+                std::sin(y / ScalarFloat(m_resolution.y() - 1) * math::Pi<ScalarFloat>);
+
+            for (size_t x = 0; x < m_resolution.x(); ++x) {
+                ScalarVector4f coeff = load<ScalarVector4f>(ptr);
+                ScalarFloat lum;
+
+                if constexpr (is_monochromatic_v<Spectrum>) {
+                    lum = coeff.x();
+                } else if constexpr (is_rgb_v<Spectrum>) {
+                    lum = mitsuba::luminance(ScalarColor3f(head<3>(coeff)));
+                } else {
+                    static_assert(is_spectral_v<Spectrum>);
+                    lum = srgb_model_mean(head<3>(coeff)) * coeff.w();
+                }
+
+                *lum_ptr++ = lum * sin_theta;
+                ptr += 4;
+            }
+        }
+
+        m_warp = Warp(luminance.get(), m_resolution);
     }
 
     void set_scene(const Scene *scene) override {

src/librender/bsdf.cpp

@@ -13,6 +13,8 @@ MTS_VARIANT Spectrum BSDF<Float, Spectrum>::eval_null_transmission(
     return 0.f;
 }
 
+MTS_VARIANT std::string BSDF<Float, Spectrum>::id() const { return m_id; }
+
 template <typename Index>
 std::string type_mask_to_string(Index type_mask) {
     std::ostringstream oss;

src/librender/endpoint.cpp

@@ -5,7 +5,7 @@
 
 NAMESPACE_BEGIN(mitsuba)
 
-MTS_VARIANT Endpoint<Float, Spectrum>::Endpoint(const Properties &props) {
+MTS_VARIANT Endpoint<Float, Spectrum>::Endpoint(const Properties &props) : m_id(props.id()) {
     m_world_transform = props.animated_transform("to_world", ScalarTransform4f()).get();
 }
 
@@ -48,6 +48,8 @@ MTS_VARIANT Spectrum Endpoint<Float, Spectrum>::eval(const SurfaceInteraction3f
     NotImplementedError("eval");
 }
 
+MTS_VARIANT std::string Endpoint<Float, Spectrum>::id() const { return m_id; }
+
 MTS_IMPLEMENT_CLASS_VARIANT(Endpoint, Object)
 MTS_INSTANTIATE_CLASS(Endpoint)
 NAMESPACE_END(mitsuba)

src/python/python/autodiff.py

@@ -91,7 +91,7 @@ def _render_helper(scene, spp=None, sensor_index=0):
     return values / (weight + 1e-8)
 
 
-def write_bitmap(filename, data, resolution):
+def write_bitmap(filename, data, resolution, write_async=True):
     """
     Write the linearized RGB image in `data` to a PNG/EXR/.. file with
     resolution `resolution`.
@@ -110,7 +110,12 @@ def write_bitmap(filename, data, resolution):
        filename.endswith('.jpeg'):
         bitmap = bitmap.convert(Bitmap.PixelFormat.RGB,
                                 Struct.Type.UInt8, True)
-    bitmap.write_async(filename, quality=0 if filename.endswith('png') else -1)
+    quality = 0 if filename.endswith('png') else -1
+
+    if write_async:
+        bitmap.write_async(filename, quality=quality)
+    else:
+        bitmap.write(filename, quality=quality)
 
 
 def render(scene,