Does Mitsuba support multiple 2D sampling simultaneously?

[ciy405x]

So, I have a list of 2D PDFs. Let's consider this as a tensor of shape (N, H, W).
From what I see, mi.Hierarchical2D0 and others seem to support single PDF sampling (i.e., for a PDF tensor of shape (H, W)).
Also, by passing (M, 2) mi.Point2f as a sample argument to the sample(self, sample, param=[], active=True) method, I can get a tensor of shape (M, 2).
I'm wondering if this function can be expanded to perform sampling over a 'batch' of PDFs, such that I can get a tensor of shape (N, M, 2) or (N * M, 2) mi.Point2f.
Like sampling with ShapePtr, do Distributions also have a Ptr?

[njroussel]

Hi @ciy405x

This is possible! Take a look at this discussions: mitsuba-renderer/drjit#150
In particular, you'll want to have a look at dr.switch (documentation).

Here's an example with DiscreteDistribution2D. It's a bit nonsensical, but this should give you an idea on how to actually use dr.switch.

import mitsuba as mi
import drjit as dr
import numpy as np

mi.set_variant('cuda_ad_rgb')

def build_distr(i: float):
    return mi.DiscreteDistribution2D(np.array(
        [[1, 2, 3],
         [i, 1, 3]])
    )

distributions = []
for i in range(10):
    distributions.append(build_distr(i))

def sampler_for_distr(i: int) :
    def sample_func(sample, active):
        return distributions[i].sample(sample, active)

    return sample_func

samplers = [sampler_for_distr(i) for i in range(10)]

samples = mi.Point2f(dr.linspace(mi.Float, 0, 1, 10), dr.linspace(mi.Float, 0, 1, 10))
out = dr.switch(dr.arange(mi.UInt, 10), samplers, samples, True)

out[0] # Point2f for position (width == 10)
out[1] # PMF value at position (width == 10)
out[2] # Samples ready for re-use (width == 10)

[Mephisto405]

This looks absolutely amazing method, but I cannot find it in the official documentation. Will it be updated soon? Anyway, I will follow your kind snippet :) Thx, a lot!!

[njroussel]

Indeed, I've added it now (https://drjit.readthedocs.io/en/latest/reference.html#drjit.switch). Thank you!