diff --git a/.gitignore b/.gitignore
index 13da1a2f..09970e0e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -33,3 +33,4 @@ __pycache__
results/.ipynb_checkpoints
results/estimator_comparison_simple
results/estimator_comparison_optimization
+results/estimator_comparison_attached_reparam
diff --git a/results/Estimator comparison (attached reparameterization).ipynb b/results/Estimator comparison (attached reparameterization).ipynb
new file mode 100644
index 00000000..07ed9ef8
--- /dev/null
+++ b/results/Estimator comparison (attached reparameterization).ipynb
@@ -0,0 +1,168 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "53581a91",
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2022-07-28T11:30:15.423479Z",
+ "start_time": "2022-07-28T11:30:15.091516Z"
+ }
+ },
+ "outputs": [],
+ "source": [
+ "import mitsuba as mi\n",
+ "mi.set_variant(\"llvm_ad_rgb\")\n",
+ "import drjit as dr\n",
+ "\n",
+ "import os\n",
+ "base_dir = 'estimator_comparison_attached_reparam'\n",
+ "if not os.path.exists(base_dir):\n",
+ " os.makedirs(base_dir)\n",
+ " \n",
+ "import numpy as np\n",
+ "import matplotlib.pyplot as plt\n",
+ "import cmap_diff\n",
+ "%config InlineBackend.figure_formats = ['svg']\n",
+ "%matplotlib inline\n",
+ " \n",
+ "mi.Thread.thread().logger().set_log_level(mi.LogLevel.Warn)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "5737c51e",
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2022-07-28T11:31:29.924896Z",
+ "start_time": "2022-07-28T11:30:15.431353Z"
+ }
+ },
+ "outputs": [],
+ "source": [
+ "scene_names = [\n",
+ " ('attached_disc_roughness', 'plane.bsdf.alpha.data'),\n",
+ " ('attached_disc_normalmap', 'plane.bsdf.normalmap.data'),\n",
+ "]\n",
+ "\n",
+ "methods = [\n",
+ " ('Detached BSDF sampling', {'method': 'bs_detached'}),\n",
+ " ('Attached BSDF sampling (naive)', {'method': 'bs_attached'}),\n",
+ " ('Attached BSDF sampling (reparam)', {'method': 'bs_attached_reparam',\n",
+ " 'reparam_kappa': 1e6, 'reparam_rays': 48}),\n",
+ "]\n",
+ "\n",
+ "for scene_name, param_key in scene_names:\n",
+ " print(\"*\", scene_name)\n",
+ " scene_dir = \"{}/{}\".format(base_dir, scene_name)\n",
+ " if not os.path.exists(scene_dir):\n",
+ " os.makedirs(scene_dir)\n",
+ " \n",
+ " \n",
+ " scene = mi.load_file('scenes/{}.xml'.format(scene_name), res=256)\n",
+ " params = mi.traverse(scene)\n",
+ " params.keep([param_key])\n",
+ " params[param_key] = dr.maximum(0.1, params[param_key])\n",
+ " \n",
+ " # Primal rendering\n",
+ " integrator = mi.load_dict({'type': 'estimator_comparison', 'method': 'primal_mis', 'hide_emitters': True})\n",
+ " image = mi.render(scene, params, integrator=integrator, seed=0, spp=32)\n",
+ " mi.util.convert_to_bitmap(image, uint8_srgb=False).write('{}/primal.exr'.format(scene_dir))\n",
+ "\n",
+ " # Compare various methods ...\n",
+ " for method_name, method_dict in methods:\n",
+ " print(\" \", method_name)\n",
+ " integrator_dict = {'type': 'estimator_comparison', 'hide_emitters': True}\n",
+ " for k, v in method_dict.items():\n",
+ " integrator_dict[k] = v\n",
+ " integrator = mi.load_dict(integrator_dict)\n",
+ "\n",
+ " # Differentiable rendering ...\n",
+ " dr.enable_grad(params[param_key])\n",
+ " image = mi.render(scene, params, integrator=integrator, seed=0, spp=256, antithetic_pass=False)\n",
+ " # ... and propagate back to input parameters\n",
+ " dr.backward(image)\n",
+ " param_grad = dr.grad(params[param_key])\n",
+ " dr.set_grad(params[param_key], 0.0)\n",
+ " dr.disable_grad(params[param_key])\n",
+ " \n",
+ " mi.util.convert_to_bitmap(param_grad, uint8_srgb=False).write('{}/{}.exr'.format(scene_dir, method_dict['method']))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "4cf0ff0c",
+ "metadata": {
+ "ExecuteTime": {
+ "end_time": "2022-07-28T11:27:06.540158Z",
+ "start_time": "2022-07-28T11:27:05.821131Z"
+ }
+ },
+ "outputs": [],
+ "source": [
+ "scales = [0.5, 2]\n",
+ "\n",
+ "for scene_idx, scene_name in enumerate([k for k, _ in scene_names]):\n",
+ " scene_dir = \"{}/{}\".format(base_dir, scene_name)\n",
+ " \n",
+ " fig, axes = plt.subplots(ncols=4, figsize=(10,3))\n",
+ " for ax in axes:\n",
+ " ax.set_xticks([]); ax.set_yticks([])\n",
+ " \n",
+ " image_primal = np.array(mi.Bitmap(\"{}/primal.exr\".format(scene_dir)))\n",
+ " image_primal = np.clip(image_primal**(1/2.2), 0.0, 1.0) # Crude gamma correction\n",
+ " axes[0].imshow(image_primal)\n",
+ " axes[0].set_title(\"Primal\")\n",
+ " \n",
+ " for method_idx, method_name in enumerate([v['method'] for k, v in methods]):\n",
+ " vminmax = scales[scene_idx]\n",
+ " image_grad = np.array(mi.Bitmap(\"{}/{}.exr\".format(scene_dir, method_name)))\n",
+ " data_ = image_grad[:,:,0] if len(image_grad.shape) == 3 else image_grad\n",
+ " im = axes[method_idx+1].imshow(data_, cmap='diff', vmin=-vminmax, vmax=+vminmax)\n",
+ " \n",
+ " axes[1].set_title(\"Detached\")\n",
+ " axes[2].set_title(\"Attached (naïve)\")\n",
+ " axes[3].set_title(\"Attached (reparam.)\")\n",
+ " plt.suptitle(scene_name, weight='bold', size=14)\n",
+ " \n",
+ " outname = '{}/comparison.jpg'.format(scene_dir)\n",
+ " plt.savefig(outname, dpi=300, pad_inches=0.1, bbox_inches='tight')\n",
+ " \n",
+ " plt.tight_layout()\n",
+ " plt.show()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "c70f0559",
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.9.12"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/results/scenes/attached_disc_normalmap.xml b/results/scenes/attached_disc_normalmap.xml
new file mode 100644
index 00000000..6793c637
--- /dev/null
+++ b/results/scenes/attached_disc_normalmap.xml
@@ -0,0 +1,57 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/results/scenes/attached_disc_roughness.xml b/results/scenes/attached_disc_roughness.xml
new file mode 100644
index 00000000..86f2875d
--- /dev/null
+++ b/results/scenes/attached_disc_roughness.xml
@@ -0,0 +1,67 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/results/scenes/textures/normalmap_lily.png b/results/scenes/textures/normalmap_lily.png
new file mode 100644
index 00000000..dfa36993
Binary files /dev/null and b/results/scenes/textures/normalmap_lily.png differ
diff --git a/results/scenes/textures/textured_roughness.jpg b/results/scenes/textures/textured_roughness.jpg
new file mode 100644
index 00000000..28726a14
Binary files /dev/null and b/results/scenes/textures/textured_roughness.jpg differ
diff --git a/src/python/python/ad/__init__.py b/src/python/python/ad/__init__.py
index 4ace1ccf..601cce0a 100644
--- a/src/python/python/ad/__init__.py
+++ b/src/python/python/ad/__init__.py
@@ -1,3 +1,4 @@
from .integrators import *
from .optimizers import *
from .reparam import *
+from .attached_reparam import *
diff --git a/src/python/python/ad/attached_reparam.py b/src/python/python/ad/attached_reparam.py
new file mode 100644
index 00000000..3cc9aa18
--- /dev/null
+++ b/src/python/python/ad/attached_reparam.py
@@ -0,0 +1,279 @@
+from __future__ import annotations as __annotations__ # Delayed parsing of type annotations
+
+import drjit as dr
+import mitsuba as mi
+
+import typing
+if typing.TYPE_CHECKING:
+ from typing import Tuple
+
+class _AttachedReparameterizeOp(dr.CustomOp):
+ """
+ Dr.Jit custom operation that reparameterizes rays generated via
+ attached sampling strategies. Based on the paper
+
+ "Monte Carlo Estimators for Differential Light Transport"
+ (Proceedings of SIGGRAPH'21) by Tizan Zeltner, Sébastien Speierer,
+ Iliyan Georgiev, and Wenzel Jakob.
+
+ This is needed to to avoid bias caused by the discontinuous visibility
+ function in gradient-based optimization involving attached samples.
+ """
+ def eval(self, scene, rng, wo, pdf, si,
+ num_rays, kappa, exponent, antithetic, naive, active):
+ # Stash all of this information for the forward/backward passes
+ self.scene = scene
+ self.rng = rng
+ self.wo = wo
+ self.pdf = pdf
+ self.si = si
+ self.num_rays = num_rays
+ self.kappa = kappa
+ self.exponent = exponent
+ self.antithetic = antithetic
+ self.naive = naive
+ self.active = active
+
+ # The reparameterization is simply the identity in primal mode
+ return self.wo, dr.full(mi.Float, 1, dr.width(wo))
+
+
+ def scale_factor(self):
+ """
+ Smooth scale factor that is used to slow down sample movement close
+ to geometric discontinuities.
+ """
+
+ # The computation is completely independent of any differentiated scene
+ # parameters, i.e. everything here is detached.
+ with dr.suspend_grad():
+ # Initialize some accumulators
+ Z = mi.Float(0.0)
+ dZ = mi.Vector3f(0.0)
+ B = mi.Float(0.0)
+ dB_lhs = mi.Vector3f(0.0)
+ it = mi.UInt32(0)
+ rng = self.rng
+ si = self.si
+ ray_d = si.to_world(self.wo)
+ ray = si.spawn_ray(ray_d)
+ ray_frame = mi.Frame3f(ray_d)
+
+ loop = mi.Loop(name="reparameterize_attached_direction(): scale factor",
+ state=lambda: (it, Z, dZ, B, dB_lhs, rng.state))
+
+ # Unroll the entire loop in wavefront mode
+ # loop.set_uniform(True) # TODO can we turn this back on? (see self.active in loop condiction)
+ loop.set_max_iterations(self.num_rays)
+ loop.set_eval_stride(self.num_rays)
+
+ while loop(self.active & (it < self.num_rays)):
+ rng_state_backup = rng.state
+ sample = mi.Point2f(rng.next_float32(),
+ rng.next_float32())
+
+ if self.antithetic:
+ repeat = dr.eq(it & 1, 0)
+ rng.state[repeat] = rng_state_backup
+ else:
+ repeat = mi.Bool(False)
+
+ # Sample an auxiliary ray from a von Mises Fisher distribution
+ omega_local = mi.warp.square_to_von_mises_fisher(sample, self.kappa)
+
+ # Antithetic sampling (optional)
+ omega_local.x[repeat] = -omega_local.x
+ omega_local.y[repeat] = -omega_local.y
+
+ aux_ray = mi.Ray3f(
+ o = ray.o,
+ d = ray_frame.to_world(omega_local),
+ time = ray.time,
+ wavelengths = ray.wavelengths)
+
+ # Compute an intersection that includes the boundary test
+ aux_si = self.scene.ray_intersect(aux_ray,
+ ray_flags=mi.RayFlags.All | mi.RayFlags.BoundaryTest,
+ coherent=False)
+
+ aux_hit = aux_si.is_valid()
+
+ # Standard boundary test evaluation
+ b_test = dr.select(aux_hit, aux_si.boundary_test, 1.0)
+ # The horizon of the local hemisphere also counts as an edge here,
+ # as e.g. normal maps introduce a discontinuity there.
+ horizon = dr.abs(dr.dot(si.n, aux_ray.d))**2
+ aux_B = b_test * horizon
+
+ # Augmented boundary test. This quantity will be smoothed as a
+ # result of this convolution process.
+ B_direct = dr.power(1.0 - aux_B, 3.0)
+
+ # Inverse of vMF density without normalization constant
+ inv_vmf_density = dr.rcp(dr.fma(sample.y, dr.exp(-2 * self.kappa), 1 - sample.y))
+
+ # Compute harmonic weight, being wary of division by near-zero values
+ w_denom = inv_vmf_density + aux_B - 1
+ w_denom_rcp = dr.select(w_denom > 1e-4, dr.rcp(w_denom), 0.0)
+ w = dr.power(w_denom_rcp, self.exponent) * inv_vmf_density
+
+ # Analytic weight gradient w.r.t. `ray.d` (detaching inv_vmf_density gradient)
+ tmp1 = inv_vmf_density * w * w_denom_rcp * self.kappa * self.exponent
+ tmp2 = ray_frame.to_world(mi.Vector3f(omega_local.x, omega_local.y, 0))
+ d_w_omega = dr.clamp(tmp1, -1e10, 1e10) * tmp2
+
+ Z += w
+ dZ += d_w_omega
+ B += w * B_direct
+ dB_lhs += d_w_omega * B_direct
+ it += 1
+
+ inv_Z = dr.rcp(dr.maximum(Z, 1e-8))
+ B = B * inv_Z
+ dB = (dB_lhs - dZ * B) * inv_Z
+
+ # Ignore inactive lanes
+ B = dr.select(self.active, B, 0.0)
+ dB = dr.select(self.active, dB, 0.0)
+
+ return B, dB
+
+
+ def forward(self):
+ """
+ Propagate the gradients in the forward direction to 'ray.d' and the
+ jacobian determinant 'det'. From a warp field point of view, the
+ derivative of 'ray.d' is the warp field direction at 'ray', and
+ the derivative of 'det' is the divergence of the warp field at 'ray'.
+ """
+
+ # Setup inputs and their gradients
+ wo, pdf = mi.Vector3f(self.wo), mi.Float(self.pdf)
+ dr.enable_grad([wo, pdf])
+ dr.set_grad(wo, self.grad_in('wo'))
+ dr.set_grad(pdf, self.grad_in('pdf'))
+
+ # Naïve reparameterization: simply cancel out all movement
+ S = wo
+ det_S = dr.detach(pdf) / pdf
+
+ if not self.naive:
+ # Only scale down movement close to geometric discontinuities
+ B, dB = self.scale_factor()
+ S_tmp = B * S
+ det_S_tmp = dr.dot(dB, self.si.to_world(S)) + B * det_S
+ S = S_tmp
+ det_S = det_S_tmp
+
+ # Apply reparameterization to direction
+ wo_R = dr.normalize(wo - S + dr.detach(S))
+ det_R = mi.Float(1.0) - det_S + dr.detach(det_S)
+
+ # Set output gradients
+ dr.forward_to(wo_R, det_R)
+ self.set_grad_out((dr.select(self.active, dr.grad(wo_R), 0.0),
+ dr.select(self.active, dr.grad(det_R), 0.0)))
+
+
+ def backward(self):
+ # Setup inputs/outputs and their gradients
+ grad_dir, grad_det = self.grad_out()
+ grad_dir = dr.select(self.active, grad_dir, 0.0)
+ grad_det = dr.select(self.active, grad_det, 0.0)
+
+ wo = mi.Vector3f(self.wo)
+ pdf = mi.Float(self.pdf)
+ dr.enable_grad([wo, pdf])
+
+ # Naïve attached parameterization: simply cancel out all movement
+ S = wo
+ det_S = dr.detach(pdf) / pdf
+
+ if not self.naive:
+ # Only scale down movement close to geometric discontinuities
+ B, dB = self.scale_factor()
+
+ S_tmp = B * S
+ det_S_tmp = dr.dot(dB, self.si.to_world(S)) + B * det_S
+ S = S_tmp; det_S = det_S_tmp
+ del S_tmp, det_S_tmp, B, dB
+
+ # Apply reparameterization to direction
+ wo_R = dr.normalize(wo - S + dr.detach(S))
+ det_R = mi.Float(1.0) - det_S + dr.detach(det_S)
+
+ # Set input gradients
+ dr.set_grad(wo_R, grad_dir)
+ dr.set_grad(det_R, grad_det)
+ dr.enqueue(dr.ADMode.Backward, wo_R, det_R)
+ dr.traverse(mi.Float, dr.ADMode.Backward)
+ self.set_grad_in('wo', dr.grad(wo))
+ self.set_grad_in('pdf', dr.grad(pdf))
+
+
+ def name(self):
+ return "reparameterize_ray()"
+
+
+def reparameterize_attached_direction(scene: mitsuba.Scene,
+ rng: mitsuba.PCG32,
+ wo: mitsuba.Vector3f,
+ pdf: mitsuba.Float,
+ si: mitsuba.SurfaceInteraction3f,
+ num_rays: int=4,
+ kappa: float=1e5,
+ exponent: float=3.0,
+ antithetic: bool=False,
+ naive: bool=False,
+ active: mitsuba.Bool = True
+) -> Tuple[mitsuba.Vector3f, mitsuba.Float]:
+ """
+ Reparameterize given ray by "attaching" the derivatives of its direction to
+ moving geometry in the scene.
+
+ Parameter ``scene`` (``mitsuba.Scene``):
+ Scene containing all shapes.
+
+ Parameter ``rng`` (``mitsuba.PCG32``):
+ Random number generator used to sample auxiliary ray directions.
+
+ Parameter ``params`` (``mitsuba.SceneParameters``):
+ Scene parameters
+
+ Parameter ``wo`` (``mitsuba.Vector3f``):
+ (Attached) sampled direction
+
+ Parameter ``pdf`` (``mitsuba.Float``):
+ (Attached) sampling PDF
+
+ Parameter ``si`` (``mitsuba.SurfaceInteraction3f``):
+ Current surface interaction
+
+ Parameter ``num_rays`` (``int``):
+ Number of auxiliary rays to trace when performing the convolution.
+
+ Parameter ``kappa`` (``float``):
+ Kappa parameter of the von Mises Fisher distribution used to sample the
+ auxiliary rays.
+
+ Parameter ``exponent`` (``float``):
+ Exponent used in the computation of the harmonic weights
+
+ Parameter ``antithetic`` (``bool``):
+ Should antithetic sampling be enabled to improve convergence?
+ (Default: False)
+
+ Parameter ``naive`` (``bool``):
+ Should the naïve version of the reparameterization be used that falls
+ back to detached sampling, i.e. cancelling all sample movement?
+
+ Parameter ``active`` (``mitsuba.Bool``):
+ Boolean array specifying the active lanes
+
+ Returns → (mitsuba.Vector3f, mitsuba.Float):
+ Returns the reparameterized ray direction and the Jacobian
+ determinant of the change of variables.
+ """
+
+ return dr.custom(_AttachedReparameterizeOp, scene, rng, wo, pdf, si,
+ num_rays, kappa, exponent, antithetic, naive, active)
diff --git a/src/python/python/ad/integrators/common.py b/src/python/python/ad/integrators/common.py
index ece5134a..134949d1 100644
--- a/src/python/python/ad/integrators/common.py
+++ b/src/python/python/ad/integrators/common.py
@@ -88,6 +88,7 @@ def render(self: mi.SamplingIntegrator,
δL=None,
state_in=None,
reparam=None,
+ attached_reparam=None,
antithetic_pass=antithetic_pass,
active=mi.Bool(True)
)
@@ -151,6 +152,17 @@ def render_forward(self: mi.SamplingIntegrator,
else:
reparam = None
+ # Potentially also create a wrapper for attached reparameterization
+ if hasattr(self, 'method') and 'attached_reparam' in self.method and hasattr(self, 'attached_reparam'):
+ attached_reparam = _AttachedReparamWrapper(
+ scene=scene,
+ attached_reparam=self.attached_reparam,
+ wavefront_size=sampler.wavefront_size(),
+ seed=seed
+ )
+ else:
+ attached_reparam = None
+
# Generate a set of rays starting at the sensor, keep track of
# derivatives wrt. sample positions ('pos') if there are any
ray, weight, pos, det = self.sample_rays(scene, sensor,
@@ -163,6 +175,7 @@ def render_forward(self: mi.SamplingIntegrator,
sampler=sampler,
ray=ray,
reparam=reparam,
+ attached_reparam=attached_reparam,
antithetic_pass=antithetic_pass,
active=mi.Bool(True)
)
@@ -234,6 +247,17 @@ def render_backward(self: mi.SamplingIntegrator,
else:
reparam = None
+ # Potentially also create a wrapper for attached reparameterization
+ if hasattr(self, 'method') and 'attached_reparam' in self.method and hasattr(self, 'attached_reparam'):
+ attached_reparam = _AttachedReparamWrapper(
+ scene=scene,
+ attached_reparam=self.attached_reparam,
+ wavefront_size=sampler.wavefront_size(),
+ seed=seed
+ )
+ else:
+ attached_reparam = None
+
# Generate a set of rays starting at the sensor, keep track of
# derivatives wrt. sample positions ('pos') if there are any
ray, weight, pos, det = self.sample_rays(scene, sensor,
@@ -246,6 +270,7 @@ def render_backward(self: mi.SamplingIntegrator,
sampler=sampler,
ray=ray,
reparam=reparam,
+ attached_reparam=attached_reparam,
antithetic_pass=antithetic_pass,
active=mi.Bool(True)
)
@@ -673,6 +698,17 @@ def render_forward(self: mi.SamplingIntegrator,
else:
reparam = None
+ # Potentially also create a wrapper for attached reparameterization
+ if hasattr(self, 'method') and 'attached_reparam' in self.method and hasattr(self, 'attached_reparam'):
+ attached_reparam = _AttachedReparamWrapper(
+ scene=scene,
+ attached_reparam=self.attached_reparam,
+ wavefront_size=sampler.wavefront_size(),
+ seed=seed
+ )
+ else:
+ attached_reparam = None
+
# Generate a set of rays starting at the sensor, keep track of
# derivatives wrt. sample positions ('pos') if there are any
ray, weight, pos, det = self.sample_rays(scene, sensor,
@@ -688,6 +724,7 @@ def render_forward(self: mi.SamplingIntegrator,
δL=None,
state_in=None,
reparam=None,
+ attached_reparam=None,
antithetic_pass=antithetic_pass,
active=mi.Bool(True)
)
@@ -761,6 +798,7 @@ def render_forward(self: mi.SamplingIntegrator,
δL=None,
state_in=state_out,
reparam=reparam,
+ attached_reparam=attached_reparam,
antithetic_pass=antithetic_pass,
active=mi.Bool(True)
)
@@ -899,6 +937,17 @@ def render_backward(self: mi.SamplingIntegrator,
else:
reparam = None
+ # Potentially also create a wrapper for attached reparameterization
+ if hasattr(self, 'method') and 'attached_reparam' in self.method and hasattr(self, 'attached_reparam'):
+ attached_reparam = _AttachedReparamWrapper(
+ scene=scene,
+ attached_reparam=self.attached_reparam,
+ wavefront_size=sampler.wavefront_size(),
+ seed=seed
+ )
+ else:
+ attached_reparam = None
+
# Generate a set of rays starting at the sensor, keep track of
# derivatives wrt. sample positions ('pos') if there are any
ray, weight, pos, det = self.sample_rays(scene, sensor,
@@ -914,6 +963,7 @@ def render_backward(self: mi.SamplingIntegrator,
δL=None,
state_in=None,
reparam=None,
+ attached_reparam=None,
antithetic_pass=antithetic_pass,
active=mi.Bool(True)
)
@@ -977,6 +1027,7 @@ def render_backward(self: mi.SamplingIntegrator,
δL=δL,
state_in=state_out,
reparam=reparam,
+ attached_reparam=attached_reparam,
antithetic_pass=antithetic_pass,
active=mi.Bool(True)
)
@@ -1256,6 +1307,60 @@ def __call__(self,
depth, active)
+class _AttachedReparamWrapper:
+ """
+ This class is an implementation detail of ``ADIntegrator``, which performs
+ necessary initialization steps and subsequently wraps an attached
+ reparameterization technique. It serves the following important purposes:
+
+ 1. Ensuring the availability of uncorrelated random variates.
+ 2. Exposing the underlying RNG state to recorded loops.
+ """
+
+ # AttachedReparamWrapper instances can be provided as dr.Loop state
+ # variables. For this to work we must declare relevant fields
+ DRJIT_STRUCT = { 'rng' : mi.PCG32 }
+
+ def __init__(self,
+ scene : mi.Scene,
+ attached_reparam: Callable[
+ [mi.Scene, mi.PCG32,
+ mi.Vector3f,
+ mi.Float,
+ mi.SurfaceInteraction3f,
+ mi.Bool],
+ Tuple[mi.Vector3f, mi.Float]],
+ wavefront_size : int,
+ seed : int):
+
+ self.scene = scene
+ self.attached_reparam = attached_reparam
+
+ # Create a uniform random number generator that won't show any
+ # correlation with the main sampler. PCG32Sampler.seed() uses
+ # the same logic except for the XOR with -1
+
+ idx = dr.arange(mi.UInt32, wavefront_size)
+ tmp = dr.opaque(mi.UInt32, 0xffffffff ^ seed)
+ v0, v1 = mi.sample_tea_32(tmp, idx)
+ self.rng = mi.PCG32(initstate=v0, initseq=v1)
+
+ def __call__(self,
+ wo: mi.Vector3f,
+ pdf: mi.Float,
+ si: mi.SurfaceInteraction3f,
+ active: Union[mi.Bool, bool] = True
+ ) -> Tuple[mi.Vector3f, mi.Float]:
+ """
+ This function takes a directional sample and a boolean active mask as
+ input and returns the reparameterized ray direction and the Jacobian
+ determinant of the change of variables.
+ The sample includes both the sampled direction (`wo`) and its sampling
+ PDF (`pdf`) generated and is assume dto be the result of an attached
+ sampling technique.
+ """
+ return self.attached_reparam(self.scene, self.rng, wo, pdf, si, active)
+
# ---------------------------------------------------------------------------
# Helper functions used by various differentiable integrators
# ---------------------------------------------------------------------------
diff --git a/src/python/python/ad/integrators/estimator_comparison.py b/src/python/python/ad/integrators/estimator_comparison.py
index d453b54a..31b62fc3 100644
--- a/src/python/python/ad/integrators/estimator_comparison.py
+++ b/src/python/python/ad/integrators/estimator_comparison.py
@@ -62,11 +62,58 @@ def __init__(self, props):
# - "mis_detached_detached_diff"
# Detached MIS weights, detached diff. BSDF sampling, detached emitter sampling
#
+ # - "bs_attached_reparam"
+ # Attached reparameterized BSDF sampling
+ #
+ # - "mis_attached_attached_reparam"
+ # Attached MIS weights, attached reparameterized BSDF sampling, detached emitter sampling
+ #
self.method = props.string('method', 'none')
# Hide directly visible emitters? (Only relevant in primal modes.)
self.hide_emitters = props.get('hide_emitters', False)
+ # Specifies the number of auxiliary rays used to evaluate the
+ # reparameterization
+ self.reparam_rays = props.get('reparam_rays', 16)
+
+ # Specifies the von Mises Fisher distribution parameter for sampling
+ # auxiliary rays in Bangaru et al.'s [2020] parameterization
+ self.reparam_kappa = props.get('reparam_kappa', 1e5)
+
+ # Harmonic weight exponent in Bangaru et al.'s [2020] parameterization
+ self.reparam_exp = props.get('reparam_exp', 3.0)
+
+ # Enable antithetic sampling in the reparameterization?
+ self.reparam_antithetic = props.get('reparam_antithetic', False)
+
+ # Unroll the loop tracing auxiliary rays in the reparameterization?
+ self.reparam_unroll = props.get('reparam_unroll', False)
+
+ # Use naïve version of the reparameterization that falls back to
+ # detached sampling, i.e. cancelling all sample movement?
+ self.reparam_naive = props.get('reparam_naive', False)
+
+
+ def attached_reparam(self,
+ scene: mitsuba.render.Scene,
+ rng: mitsuba.core.PCG32,
+ wo: mitsuba.core.Vector3f,
+ pdf: mitsuba.core.Float,
+ si: mitsuba.render.SurfaceInteraction3f,
+ active: mitsuba.core.Bool):
+ """
+ Helper function for attached reparameterization
+ """
+
+ return mi.ad.reparameterize_attached_direction(scene, rng, wo, pdf, si,
+ num_rays=self.reparam_rays,
+ kappa=self.reparam_kappa,
+ exponent=self.reparam_exp,
+ antithetic=self.reparam_antithetic,
+ naive=self.reparam_naive,
+ active=active)
+
def sample(self,
mode: dr.ADMode,
scene: mi.Scene,
@@ -75,6 +122,9 @@ def sample(self,
δL: Optional[mi.Spectrum],
state_in: Optional[mi.Spectrum],
antithetic_pass: Optional[bool],
+ attached_reparam: Optional[
+ Callable[[mi.Vector3f, mi.Float, mi.Bool],
+ Tuple[mi.Vector3f, mi.Float]]],
active: mi.Bool,
**kwargs # Absorbs unused arguments
) -> Tuple[mi.Spectrum,
@@ -528,6 +578,81 @@ def mis_weight(pdf_a, pdf_b):
L += mis * bsdf_val / bs.pdf * emitter_val
+ elif self.method == 'bs_attached_reparam':
+ # ------------ Attached BSDF sampling -------------------------------------------
+ # ... but using the attached+reparameterized BSDF sampling strategy (Section 5.2)
+ # -------------------------------------------------------------------------------
+
+ with dr.resume_grad(when=not primal):
+ # Use a BSDF sampling strategy, everything is attached here.
+ bs, _ = bsdf.sample(
+ ctx, si, sampler.next_1d(active), sampler.next_2d(active), active
+ )
+ active &= dr.neq(bs.pdf, 0.0)
+
+ # Reparameterize
+ wo_R, det_R = attached_reparam(bs.wo, bs.pdf, si, active)
+
+ # Re-evaluate BSDF with reparam-ed direction
+ bsdf_value = bsdf.eval(ctx, si, wo_R, active)
+
+ ray = si.spawn_ray(si.to_world(wo_R))
+ si_bsdf = scene.ray_intersect(ray, active)
+
+ # Evaluate emitter, also being differentiated due to diff. `wo`
+ delta = mi.has_flag(bs.sampled_type, mi.BSDFFlags.Delta)
+ ds = mi.DirectionSample3f(scene, si_bsdf, si)
+ active_b = active & dr.neq(ds.emitter, None) & ~delta
+ emitter_val = ds.emitter.eval(si_bsdf, active_b)
+
+ L = bsdf_value / bs.pdf * emitter_val * det_R
+
+ elif self.method == 'mis_attached_attached_reparam':
+ # ------------ MIS with attached weights and attached strategies ----------------
+ # ... but using the attached+reparameterized BSDF sampling strategy (Section 5.2)
+ # -------------------------------------------------------------------------------
+
+ with dr.resume_grad(when=not primal):
+ # First, use an emitter sampling strategy
+ active_e = active & mi.has_flag(bsdf.flags(), mi.BSDFFlags.Smooth)
+ ds, emitter_weight = scene.sample_emitter_direction(
+ si, sampler.next_2d(active_e), True, active_e
+ )
+ active_e &= dr.neq(ds.pdf, 0.0)
+ wo = si.to_local(ds.d)
+
+ # Evaluate BSDF
+ bsdf_val, bsdf_pdf = bsdf.eval_pdf(ctx, si, wo, active_e)
+ mis = dr.select(ds.delta, 1.0, mis_weight(ds.pdf, bsdf_pdf))
+
+ L += mis * bsdf_val * emitter_weight
+
+ # Second, use a BSDF sampling strategy
+ bs, _ = bsdf.sample(
+ ctx, si, sampler.next_1d(active), sampler.next_2d(active), active
+ )
+ active &= dr.neq(bs.pdf, 0.0)
+
+ # Reparameterize
+ wo_R, det_R = attached_reparam(bs.wo, bs.pdf, si, active)
+
+ # Re-evaluate BSDF with reparam-ed direction
+ bsdf_value = bsdf.eval(ctx, si, wo_R, active)
+ bsdf_pdf = bsdf.eval(ctx, si, wo_R, active)
+
+ ray = si.spawn_ray(si.to_world(wo_R))
+ si_bsdf = scene.ray_intersect(ray, active)
+
+ # Evaluate emitter, also being differentiated due to diff. `wo`
+ delta = mi.has_flag(bs.sampled_type, mi.BSDFFlags.Delta)
+ ds = mi.DirectionSample3f(scene, si_bsdf, si)
+ active_b = active & dr.neq(ds.emitter, None) & ~delta
+ emitter_pdf = scene.pdf_emitter_direction(si, ds, active_b)
+ mis = dr.select(active_b, mis_weight(bsdf_pdf, emitter_pdf), 1.0)
+ emitter_val = ds.emitter.eval(si_bsdf, active_b)
+
+ L += mis * bsdf_value / bs.pdf * emitter_val * det_R
+
if not primal:
with dr.resume_grad():
# Propagate derivatives from/to 'L' based on 'mode'