Wavelength-dependent BSDF

[cvejarano]

Hi! First of all, thanks for Mitsuba, I am trying to use it at work to bridge the gap between realistic 3D rendering and advanced optical simulation and so far it seems very promising.

I can program in Python, but I am completely new to JIT compilation / DrJit.

The problem

I need to simulate a surface characterized by a wavelength-dependent reflectance: for a given incident ray, it either reflects or transmits it with a certain probability, and this probability depends on the wavelength.

What I have done

I started from Mitsuba v3.5.2.

I understand that each ray has associated by default 4 different wavelenghts, so I compiled a custom spectral variant with only 1 wavelength per ray (similar to #1309):

    "llvm_spectral_1f": {
        "float": "dr::LLVMArray<float>",
        "spectrum": "Spectrum<Float, 1>"
    },

Starting from this example: https://mitsuba.readthedocs.io/en/stable/src/others/custom_plugin.html, this is the sample() function of a simple BSDF which reflects everything above 600nm and refracts everything below:

    def sample(self, ctx, si, sample1, sample2, active):
        
        # Compute Fresnel terms (requited for refraction).
        # WRT the original example, we ignore the reflection coeff.
        cos_theta_i = mi.Frame3f.cos_theta(si.wi)
        _, cos_theta_t, eta_it, eta_ti = mi.fresnel(cos_theta_i, self.eta)

        # Reflection/transmission depends on the wavelength.
        ray_wl = si.wavelengths[0]
        reflectance = dr.select(ray_wl > 600, 1.0, 0.0)   #  <---- HERE -----------------------------------
        r_i = dr.minimum(reflectance, 1.0)
        t_i = dr.maximum(1.0 - r_i, 0.0)

        # Choose between reflection and transmission, based on reflectance.
        selected_r = (sample1 <= r_i) & active

        # Fill up the BSDFSample struct. See:
        # https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.BSDF.sample
        bs = mi.BSDFSample3f()
        bs.pdf = dr.select(selected_r, r_i, t_i)
        bs.sampled_component = dr.select(selected_r, mi.UInt32(0), mi.UInt32(1))
        bs.sampled_type = dr.select(
            selected_r,
            mi.UInt32(+mi.BSDFFlags.DeltaReflection),
            mi.UInt32(+mi.BSDFFlags.DeltaTransmission),
        )
        bs.wo = dr.select(
            selected_r, mi.reflect(si.wi), mi.refract(si.wi, cos_theta_t, eta_ti)
        )
        bs.eta = dr.select(selected_r, 1.0, eta_it)

        # For reflection, don't modify the output. For transmission, radiance must
        # be scaled to account for the solid angle compression.
        value_r = mi.Spectrum(1.0)
        value_t = mi.Spectrum(1.0) * dr.sqr(eta_ti)
        value = dr.select(selected_r, value_r, value_t)

        return (bs, value)

And this seems to work (apart from the fact that the result is a bit noisy, despite using 4096 samples, but I guess is due to reducing the number of wavelengths per ray):

What I want now is to be able to select the reflectance from a LUT, instead of using a simple condition, so supposing I have loaded two arrays with wavelengths and the corresponding reflectances, it would be something like this (naïvely, with numpy):

        ...
        ray_wl = si.wavelengths[0]
        reflectance = np.interp(ray_wl, self.reflectance_wavelengths, self.reflectance_vals)
        ...

Except that this doesn't work. With a scalar variant I get this error:

Traceback (most recent call last):
  File "src/mitsuba/mitsuba_sample_cubes.py", line 33, in <module>
    image = mi.render(scene, spp=4096)
  File "/xxx/mitsuba3/build/python/mitsuba/python/util.py", line 522, in render
    return dr.custom(_RenderOp, scene, sensor, params, integrator,
  File "/xxx/mitsuba3/build/python/drjit/router.py", line 5776, in custom
    output = inst.eval(**{ k: _dr.detach(v) for k, v in kwargs.items() })
  File "/xxx/mitsuba3/build/python/mitsuba/python/util.py", line 377, in eval
    return self.integrator.render(
RuntimeError: When called outside a bound function, py::cast() cannot do Python -> C++ conversions which require the creation of temporary values

And with the LLVM variant:

jitc_var_schedule(r3260): placeholder variables are used to record computation symbolically
and cannot be scheduled for evaluation! This error message could appear for
the following reasons:

1. You are using DrJit's loop or virtual function call recording feature
   and tried to perform an operation that is not permitted in this restricted
   execution mode. Please see the documentation of recorded loops/virtual
   function calls to learn about these restrictions.

2. You are accessing a variable that was modified as part of a recorded
   loop and forgot to specify it as a loop variable. Please see the
   drjit::Loop documentation for details.
Traceback (most recent call last):
  File "src/mitsuba/mitsuba_sample_cubes.py", line 33, in <module>
    image = mi.render(scene, spp=4096)
  File "/xxx/mitsuba3/build/python/mitsuba/python/util.py", line 522, in render
    return dr.custom(_RenderOp, scene, sensor, params, integrator,
  File "/xxx/mitsuba3/build/python/drjit/router.py", line 5776, in custom
    output = inst.eval(**{ k: _dr.detach(v) for k, v in kwargs.items() })
  File "/xxx/mitsuba3/build/python/mitsuba/python/util.py", line 377, in eval
    return self.integrator.render(
  File "/xxx/src/mitsuba/custom_bsdf.py", line 120, in sample
    reflectance = np.interp(ray_wl, self.reflectance_wavelengths, self.reflectance_vals)
  File "<__array_function__ internals>", line 5, in interp
  File "/usr/lib/python3/dist-packages/numpy/lib/function_base.py", line 1412, in interp
    return interp_func(x, xp, fp, left, right)
ValueError: Invalid __array_interface__ value, must be a dict

So...

What would be the right way to do this with DrJit ?

[njroussel]

Hi @cvejarano

You'll want to use dr.binary_search() . By using that function you'll be able to figure out the indices of the two wavelength values in self.reflectance_wavelengths and self.reflectance_vals that you should be interpolating between (use dr.gather to get them). You can then use dr.lerp() to actually do the interpolation (careful the interpolant needs to be contained in [0, 1]).

[merlinND]

Note that binary search requires the input array / search range to be sorted.

[cvejarano]

@njroussel Thanks for your answer. I did an even simpler test with a nearest-neighbor interpolation: just rounding the wavelength and trying to index the reflectance array with that as an index. Even that didn't work, so I kept on debugging and at the end, the reason was that I was not converting the reflectance array to a DrJit type!
So, rookie mistake; maybe obvious in hindsight. Anyway, now it works. Thanks for your help!