wip: started implementation of mitsuba2-transient-nlos in mitsuba3

I've left a to-do list in TODO(diego) comments at the start of files

mitransient/__init__.py

@@ -1,4 +1,26 @@
+"""
+TODO(diego):
+- Import mitsuba, if it's not available then raise an exception
+    telling the user to source the setpath.sh of the mitsuba installation
+    or to install mitsuba using pip.
+- Ensure that mitransient has autocompletion with the same attributes as mitsuba
+    Needs testing: set the __all__ or __dict__ variables to equal mitsuba's
+    __all__ or __dict__ variables, but make sure that our own functions
+    overwrite Mitsuba's (e.g. our sensor class)
+- Check other files' __init__.py files to see if their contents can be moved
+    to the same function
+"""
+
+
+try:
+    import mitsuba
+except ModuleNotFoundError:
+    raise Exception(
+        'The mitsuba installation could not be found. '
+        'Please install it using pip or source the setpath.sh file of your Mitsuba installation.')
+
 from .integrators import *
 from .render import *
+from .sensors import *
 
-from .utils import *
+from .utils import show_video, speed_of_light

mitransient/integrators/common.py

@@ -1,3 +1,9 @@
+"""
+TODO(diego): Figure out how to handle large numbers of samples per pixel
+so that the progress bar updates more frequently
+(similar to mitsuba2-transient-nlos's multiple passes)
+"""
+
 # Delayed parsing of type annotations
 from __future__ import annotations as __annotations__
 
@@ -54,7 +60,7 @@ def prepare_transient(self, scene, sensor):
         Prepare the integrator to perform a transient simulation
         '''
         import numpy as np
-        from mitransient.render import TransientBlock
+        from mitransient.render.transient_block import TransientBlock
 
         if isinstance(sensor, int):
             sensor = scene.sensors()[sensor]

mitransient/integrators/transientpathnlos.py

@@ -0,0 +1,287 @@
+"""
+TODO(diego): modify as per mitsuba2-transient-nlos
+"""
+
+from __future__ import annotations  # Delayed parsing of type annotations
+
+import drjit as dr
+import mitsuba as mi
+
+from mitransient.integrators.common import TransientRBIntegrator, mis_weight
+
+
+class TransientPath(TransientRBIntegrator):
+    r"""
+    .. _integrator-prb:
+
+    Path Replay Backpropagation (:monosp:`prb`)
+    -------------------------------------------
+
+    .. pluginparameters::
+
+     * - max_depth
+       - |int|
+       - Specifies the longest path depth in the generated output image (where -1
+         corresponds to :math:`\infty`). A value of 1 will only render directly
+         visible light sources. 2 will lead to single-bounce (direct-only)
+         illumination, and so on. (Default: 6)
+
+     * - rr_depth
+       - |int|
+       - Specifies the path depth, at which the implementation will begin to use
+         the *russian roulette* path termination criterion. For example, if set to
+         1, then path generation many randomly cease after encountering directly
+         visible surfaces. (Default: 5)
+
+    This plugin implements a basic Path Replay Backpropagation (PRB) integrator
+    with the following properties:
+
+    - Emitter sampling (a.k.a. next event estimation).
+
+    - Russian Roulette stopping criterion.
+
+    - No reparameterization. This means that the integrator cannot be used for
+      shape optimization (it will return incorrect/biased gradients for
+      geometric parameters like vertex positions.)
+
+    - Detached sampling. This means that the properties of ideal specular
+      objects (e.g., the IOR of a glass vase) cannot be optimized.
+
+    See ``prb_basic.py`` for an even more reduced implementation that removes
+    the first two features.
+
+    See the papers :cite:`Vicini2021` and :cite:`Zeltner2021MonteCarlo`
+    for details on PRB, attached/detached sampling, and reparameterizations.
+
+    .. tabs::
+
+        .. code-tab:: python
+
+            'type': 'prb',
+            'max_depth': 8
+    """
+
+    def sample(self,
+               mode: dr.ADMode,
+               scene: mi.Scene,
+               sampler: mi.Sampler,
+               ray: mi.Ray3f,
+               δL: Optional[mi.Spectrum],
+               state_in: Optional[mi.Spectrum],
+               active: mi.Bool,
+               add_transient,
+               **kwargs  # Absorbs unused arguments
+               ) -> Tuple[mi.Spectrum,
+                          mi.Bool, mi.Spectrum]:
+        """
+        See ``TransientADIntegrator.sample()`` for a description of this interface and
+        the role of the various parameters and return values.
+        """
+
+        # Rendering a primal image? (vs performing forward/reverse-mode AD)
+        primal = mode == dr.ADMode.Primal
+
+        # Standard BSDF evaluation context for path tracing
+        bsdf_ctx = mi.BSDFContext()
+
+        # --------------------- Configure loop state ----------------------
+
+        # Copy input arguments to avoid mutating the caller's state
+        ray = mi.Ray3f(dr.detach(ray))
+        depth = mi.UInt32(0)                          # Depth of current vertex
+        L = mi.Spectrum(0 if primal else state_in)    # Radiance accumulator
+        # Differential/adjoint radiance
+        δL = mi.Spectrum(δL if δL is not None else 0)
+        β = mi.Spectrum(1)                            # Path throughput weight
+        η = mi.Float(1)                               # Index of refraction
+        active = mi.Bool(active)                      # Active SIMD lanes
+        distance = mi.Float(0)                        # Distance of the path
+
+        # Variables caching information from the previous bounce
+        prev_si = dr.zeros(mi.SurfaceInteraction3f)
+        prev_bsdf_pdf = mi.Float(1.0)
+        prev_bsdf_delta = mi.Bool(True)
+
+        if self.camera_unwarp:
+            si = scene.ray_intersect(mi.Ray3f(ray),
+                                     ray_flags=mi.RayFlags.All,
+                                     coherent=mi.Mask(True))
+
+            distance[si.is_valid()] = -si.t
+
+        # Record the following loop in its entirety
+        loop = mi.Loop(name="Path Replay Backpropagation (%s)" % mode.name,
+                       state=lambda: (sampler, ray, depth, L, δL, β, η, active,
+                                      prev_si, prev_bsdf_pdf, prev_bsdf_delta,
+                                      distance))
+
+        # Specify the max. number of loop iterations (this can help avoid
+        # costly synchronization when when wavefront-style loops are generated)
+        loop.set_max_iterations(self.max_depth)
+
+        while loop(active):
+            # Compute a surface interaction that tracks derivatives arising
+            # from differentiable shape parameters (position, normals, etc.)
+            # In primal mode, this is just an ordinary ray tracing operation.
+
+            with dr.resume_grad(when=not primal):
+                si = scene.ray_intersect(ray,
+                                         ray_flags=mi.RayFlags.All,
+                                         coherent=dr.eq(depth, 0))
+
+            # Update distance
+            distance += dr.select(active, si.t, 0.0) * η
+
+            # Get the BSDF, potentially computes texture-space differentials
+            bsdf = si.bsdf(ray)
+
+            # ---------------------- Direct emission ----------------------
+
+            # Compute MIS weight for emitter sample from previous bounce
+            ds = mi.DirectionSample3f(scene, si=si, ref=prev_si)
+
+            mis = mis_weight(
+                prev_bsdf_pdf,
+                scene.pdf_emitter_direction(prev_si, ds, ~prev_bsdf_delta)
+            )
+
+            with dr.resume_grad(when=not primal):
+                Le = β * mis * ds.emitter.eval(si)
+
+            # Add transient contribution
+            add_transient(Le, distance, ray.wavelengths, active)
+
+            # ---------------------- Emitter sampling ----------------------
+
+            # Should we continue tracing to reach one more vertex?
+            active_next = (depth + 1 < self.max_depth) & si.is_valid()
+
+            # Is emitter sampling even possible on the current vertex?
+            active_em = active_next & mi.has_flag(
+                bsdf.flags(), mi.BSDFFlags.Smooth)
+
+            # If so, randomly sample an emitter without derivative tracking.
+            ds, em_weight = scene.sample_emitter_direction(
+                si, sampler.next_2d(), True, active_em)
+            active_em &= dr.neq(ds.pdf, 0.0)
+
+            with dr.resume_grad(when=not primal):
+                if not primal:
+                    # Given the detached emitter sample, *recompute* its
+                    # contribution with AD to enable light source optimization
+                    ds.d = dr.normalize(ds.p - si.p)
+                    em_val = scene.eval_emitter_direction(si, ds, active_em)
+                    em_weight = dr.select(
+                        dr.neq(ds.pdf, 0), em_val / ds.pdf, 0)
+                    dr.disable_grad(ds.d)
+
+                # Evaluate BSDF * cos(theta) differentiably
+                wo = si.to_local(ds.d)
+                bsdf_value_em, bsdf_pdf_em = bsdf.eval_pdf(
+                    bsdf_ctx, si, wo, active_em)
+                mis_em = dr.select(
+                    ds.delta, 1, mis_weight(ds.pdf, bsdf_pdf_em))
+                Lr_dir = β * mis_em * bsdf_value_em * em_weight
+
+            # Add contribution direct emitter sampling
+            add_transient(Lr_dir, distance + ds.dist *
+                          η, ray.wavelengths, active)
+
+            # ------------------ Detached BSDF sampling -------------------
+
+            bsdf_sample, bsdf_weight = bsdf.sample(bsdf_ctx, si,
+                                                   sampler.next_1d(),
+                                                   sampler.next_2d(),
+                                                   active_next)
+
+            # ---- Update loop variables based on current interaction -----
+
+            L = (L + Le + Lr_dir) if primal else (L - Le - Lr_dir)
+            ray = si.spawn_ray(si.to_world(bsdf_sample.wo))
+            η *= bsdf_sample.eta
+            β *= bsdf_weight
+
+            # Information about the current vertex needed by the next iteration
+
+            prev_si = dr.detach(si, True)
+            prev_bsdf_pdf = bsdf_sample.pdf
+            prev_bsdf_delta = mi.has_flag(
+                bsdf_sample.sampled_type, mi.BSDFFlags.Delta)
+
+            # -------------------- Stopping criterion ---------------------
+
+            # Don't run another iteration if the throughput has reached zero
+            β_max = dr.max(β)
+            active_next &= dr.neq(β_max, 0)
+
+            # Russian roulette stopping probability (must cancel out ior^2
+            # to obtain unitless throughput, enforces a minimum probability)
+            rr_prob = dr.minimum(β_max * η**2, .95)
+
+            # Apply only further along the path since, this introduces variance
+            rr_active = depth >= self.rr_depth
+            β[rr_active] *= dr.rcp(rr_prob)
+            rr_continue = sampler.next_1d() < rr_prob
+            active_next &= ~rr_active | rr_continue
+
+            # ------------------ Differential phase only ------------------
+
+            if not primal:
+                with dr.resume_grad():
+                    # 'L' stores the indirectly reflected radiance at the
+                    # current vertex but does not track parameter derivatives.
+                    # The following addresses this by canceling the detached
+                    # BSDF value and replacing it with an equivalent term that
+                    # has derivative tracking enabled. (nit picking: the
+                    # direct/indirect terminology isn't 100% accurate here,
+                    # since there may be a direct component that is weighted
+                    # via multiple importance sampling)
+
+                    # Recompute 'wo' to propagate derivatives to cosine term
+                    wo = si.to_local(ray.d)
+
+                    # Re-evaluate BSDF * cos(theta) differentiably
+                    bsdf_val = bsdf.eval(bsdf_ctx, si, wo, active_next)
+
+                    # Detached version of the above term and inverse
+                    bsdf_val_det = bsdf_weight * bsdf_sample.pdf
+                    inv_bsdf_val_det = dr.select(dr.neq(bsdf_val_det, 0),
+                                                 dr.rcp(bsdf_val_det), 0)
+
+                    # Differentiable version of the reflected indirect
+                    # radiance. Minor optional tweak: indicate that the primal
+                    # value of the second term is always 1.
+                    Lr_ind = L * \
+                        dr.replace_grad(1, inv_bsdf_val_det * bsdf_val)
+
+                    # Differentiable Monte Carlo estimate of all contributions
+                    Lo = Le + Lr_dir + Lr_ind
+
+                    if dr.flag(dr.JitFlag.VCallRecord) and not dr.grad_enabled(Lo):
+                        raise Exception(
+                            "The contribution computed by the differential "
+                            "rendering phase is not attached to the AD graph! "
+                            "Raising an exception since this is usually "
+                            "indicative of a bug (for example, you may have "
+                            "forgotten to call dr.enable_grad(..) on one of "
+                            "the scene parameters, or you may be trying to "
+                            "optimize a parameter that does not generate "
+                            "derivatives in detached PRB.)")
+
+                    # Propagate derivatives from/to 'Lo' based on 'mode'
+                    if mode == dr.ADMode.Backward:
+                        dr.backward_from(δL * Lo)
+                    else:
+                        δL += dr.forward_to(Lo)
+
+            depth[si.is_valid()] += 1
+            active = active_next
+
+        return (
+            L if primal else δL,  # Radiance/differential radiance
+            dr.neq(depth, 0),    # Ray validity flag for alpha blending
+            L                    # State for the differential phase
+        )
+
+
+mi.register_integrator("transient_path", lambda props: TransientPath(props))

mitransient/render/__init__.py

@@ -1 +1,31 @@
-from .transient_block import *
+# Import/re-import all files in this folder
+
+import os
+import importlib
+import glob
+
+# Make sure mitsuba.python.util is imported before the integrators
+import mitsuba.util
+
+do_reload = 'common' in globals()
+
+if mitsuba.variant() is not None and not mitsuba.variant().startswith('scalar'):
+    # Make sure `common.py` is reloaded before the integrators
+    if do_reload:
+        importlib.reload(globals()['common'])
+
+    for f in glob.glob(os.path.join(os.path.dirname(__file__), "*.py")):
+        if not os.path.isfile(f) or f.endswith('__init__.py'):
+            continue
+
+        name = os.path.basename(f)[:-3]
+        if do_reload and not name.startswith('common'):
+            importlib.reload(globals()[name])
+        else:
+            print(f'Importing {name}')
+            importlib.import_module(f'.{name}', package=__name__)
+
+        del name
+    del f
+
+del os, glob, importlib, do_reload