Using ShapePtr in python or is there any way to sample points on multiple shapes simultaneously?

[Mephisto405]

Hello,

Given a scene, I want to sample many points uniformly on all the shapes in the scene.
For instance, the following snippets do the work:

scene = mi.load_file("data/cornell-box/scene_v3.xml")
sampler = mi.PCG32(size=1)
time = 0.5
batch_size = 32

for shape in scene.shapes():
    m_area.append(shape.surface_area())
m_area = np.array(m_area)[:,0]
if len(m_area):
    m_area /= m_area.sum()

    batch_indices = np.random.choice(range(len(m_area)), size=batch_size, p=m_area)
    for i in batch_indices:
        pos_sample = mi.Point2f(sampler.next_float32(), sampler.next_float32())
        ps = m_shapes[i].sample_position(time, pos_sample)
        # for further applications
        si = mi.SurfaceInteraction3f(ps, dr.zeros(mi.Color0f))
        si.wi = ...

But I guess this code is not optimized, and I feel like there may be some Mitsuba-nic (analogous to pythonic) way to vectorize the for i in batch_indices part.
I figure out that ShapePtr->sample_position in Mitsuba C does the job, but I cannot find similar stuff in the Python version.
So how to use ShapePtr in python, or is there any way to sample points on multiple shapes simultaneously?

[njroussel]

Hi !

Indeed, Scene.shapes() returns a Python list 🙁 . What you're looking for is Scene.shapes_dr() which returns a ShapePtr.

[Mephisto405]

Oh yeah, but I found that the scene.shapes_dr().sample_position() returns AttributeError: 'mitsuba.llvm_ad_rgb.ShapePtr' object has no attribute 'sample_position'.

[njroussel]

Oh! I forgot about this.
It was a missing Pyhton binding, we've added it in the latest release: https://github.com/mitsuba-renderer/mitsuba3/releases/tag/v3.2.0.

If you have your own custom Mitsuba build you can cherry-pick this commit: bdce950
If you're using the pre-built PyPI packages, you can upgrade it to v3.2.0.

[Mephisto405]

Kkkk just released a day ago? What a lovely timely release! You save my days literally (for neural network training) 😀😃. Thx for letting me know.

[Mephisto405]

So shameful to ask every line while programming, but I ask a question again.
The question is: it's quite unclear how to put sample input of type mitsuba.Point2f in the `ShapePtr.sample_position' method.

# scene = mi.load_file("data/cornell-box/scene_v3.xml")
batch_rng = mi.PCG32(size=32*len(scene.shapes()))
pos_sample = mi.Point2f(batch_rng.next_float32(), batch_rng.next_float32())
scene.shapes_dr().sample_position(0.5, pos_sample)
# RuntimeError: jit_var_and(r11763, r11761): operands have incompatible sizes! (sizes: 256, 8)

I feel like I should make some kind of (2, 32, 8) (or (32, 8, 2)??) shape random tensor and put it into the sample_position method, but I can't find a way to do that with mi.Point2f.

[njroussel]

Here's a snippet where I'm reproducing what you're trying to achieve I believe:

import mitsuba as mi
import drjit as dr

mi.set_variant('cuda_ad_rgb')
scene = mi.load_dict(mi.cornell_box())
nb_shapes = len(scene.shapes())

N = 32
batch_rng = mi.PCG32(size=N * nb_shapes)

indices = dr.arange(mi.UInt, 0, nb_shapes)
indices = dr.repeat(indices, N)

shapes = dr.gather(mi.ShapePtr, scene.shapes_dr(), indices)
pos_sample = mi.Point2f(batch_rng.next_float32(), batch_rng.next_float32())
positions = shapes.sample_position(0.5, pos_sample)
print(f"{positions.p.numpy().shape}") # (32 * nb_shapes, 3)

The key here is that every operand is expected to have the same width. In your code, shapes_dr() is a 8-wide ShapePtr whereas pos_sample is a 256-wide Point2f. Mitsuba/Drjit operations will typically not increase the number of dimensions of your operations by doing some "cartesian product"-like combination of all possible operands.
Here the dr.gather operation acts like a upsampling step to get a 256-wide ShapePtr.