Custom Mitsuba BRDF (Python plugin)

[sapo17]

Summary

First, thanks for the amazing work!

I am trying to implement a custom Diffuse BRDF using my custom PyTorch sampling function. I am getting unexpected results. I would appreciate it if you could tell me what I am missing. See description for details.

System configuration

System information:

OS: Windows-10
CPU: Intel64 Family 6 Model 183 Stepping 1, GenuineIntel
GPU: NVIDIA GeForce RTX 4080
Python: 3.11.8 | packaged by Anaconda, Inc. | (main, Feb 26 2024, 21:34:05) [MSC v.1916 64 bit (AMD64)]
NVidia driver: 551.86
CUDA: 12.1.66
LLVM: 18.1.1

Dr.Jit: 0.4.4
Mitsuba: 3.5.0
Is custom build? False
Compiled with: MSVC 19.38.33133.0
Variants:
scalar_rgb
scalar_spectral
cuda_ad_rgb
llvm_ad_rgb

Installed Mitsuba with:

   # conda venv.yml file
  - pip
  - pip:
    - mitsuba

Installed PyTorch with:

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

Description

My custom Diffuse BRDF uses its own sampling method that is supposed to transform sample2 into a vector.

Normally, my sampling method is a bit more complicated and involved. But, I managed to reproduce it with the following:

import mitsuba as mi
import drjit as dr
import torch


class CustomDiffuseBSDF(mi.BSDF):
    def __init__(self: mi.BSDF, props: mi.Properties) -> None:
        mi.BSDF.__init__(self, props)

        self.m_albedo = props.get("reflectance", [0.5])

        reflection_flags = (
            mi.BSDFFlags.DiffuseReflection | mi.BSDFFlags.FrontSide
        )
        self.m_components = [reflection_flags]
        self.m_flags = reflection_flags

    def sample(
        self: mi.BSDF,
        ctx: mi.BSDFContext,
        si: mi.SurfaceInteraction3f,
        sample1: float,
        sample2: mi.Point2f,
        active: bool = True,
    ):

        cos_theta_i = mi.Frame3f.cos_theta(si.wi)

        # fill up the BSDFSample
        bs = mi.BSDFSample3f()

        active &= cos_theta_i > 0.0
        """ NOTE: The original code is like this:
            if (unlikely(dr::none_or<false>(active) ||
                     !ctx.is_enabled(BSDFFlags::DiffuseReflection)))
            return { bs, 0.f };

            # Somehow cannot combined `active` with the below expression
        """
        if not ctx.is_enabled(mi.BSDFFlags.DiffuseReflection):
            return (bs, 0.0)

        # test = mi.Point2f([1, 1])
        sample2_flat = dr.ravel(sample2)
        wo_torch = self.square_to_cosine_hemisphere(
            mi.TensorXf(sample2_flat, shape=(1, 2))
        )
        wo_flat = dr.ravel(wo_torch)
        bs.wo = dr.unravel(mi.Vector3f, wo_flat)

        # bs.wo = mi.warp.square_to_cosine_hemisphere(sample2)
        bs.pdf = mi.warp.square_to_cosine_hemisphere_pdf(bs.wo)
        bs.eta = 1.0
        bs.sampled_type = mi.UInt32(+mi.BSDFFlags.DiffuseReflection)
        bs.sampled_component = 0

        value = self.m_albedo.eval(si, active)

        return (
            bs,
            dr.select(active & (bs.pdf > 0.0), mi.depolarizer(value), 0.0),
        )

    def eval(
        self: mi.BSDF,
        ctx: mi.BSDFContext,
        si: mi.SurfaceInteraction3f,
        wo: mi.Vector3f,
        active: bool = True,
    ) -> mi.Color3f:
        if not ctx.is_enabled(mi.BSDFFlags.DiffuseReflection):
            return 0.0

        cos_theta_i = mi.Frame3f.cos_theta(si.wi)
        cos_theta_o = mi.Frame3f.cos_theta(wo)

        active &= (cos_theta_i > 0.0) & (cos_theta_o > 0.0)

        value = self.m_albedo.eval(si, active) * dr.inv_pi * cos_theta_o

        return mi.depolarizer(value) & active

    def pdf(
        self: mi.BSDF,
        ctx: mi.BSDFContext,
        si: mi.SurfaceInteraction3f,
        wo: mi.Vector3f,
        active: bool = True,
    ) -> float:
        if not ctx.is_enabled(mi.BSDFFlags.DiffuseReflection):
            return 0.0

        cos_theta_i = mi.Frame3f.cos_theta(si.wi)
        cos_theta_o = mi.Frame3f.cos_theta(wo)

        pdf = mi.warp.square_to_cosine_hemisphere_pdf(wo)

        return dr.select(
            dr.and_(cos_theta_i > 0.0, cos_theta_o > 0.0), pdf, 0.0
        )

    def eval_pdf(
        self: mi.BSDF,
        ctx: mi.BSDFContext,
        si: mi.SurfaceInteraction3f,
        wo: mi.Vector3f,
        active: bool = True,
    ):
        if not ctx.is_enabled(mi.BSDFFlags.DiffuseReflection):
            return 0.0, 0.0

        cos_theta_i = mi.Frame3f.cos_theta(si.wi)
        cos_theta_o = mi.Frame3f.cos_theta(wo)

        active &= (cos_theta_i > 0.0) & (cos_theta_o > 0.0)

        value = self.m_albedo.eval(si, active) * dr.inv_pi * cos_theta_o
        pdf = mi.warp.square_to_cosine_hemisphere_pdf(wo)

        return (mi.depolarizer(value) & active, dr.select(active, pdf, 0.0))

    def traverse(self: mi.Object, callback) -> None:
        callback.put_parameter(
            "reflectance", self.m_albedo, mi.ParamFlags.Differentiable
        )

    def parameters_changed(self: mi.Object, keys=...) -> None:
        print("There is nothing to do here")

    def to_string(self):
        return (
            "CustomDiffuseBSDF[\n" "    reflectance=%s,\n" "]" % (self.m_albedo)
        )

    def square_to_uniform_disk(self, sample: torch.Tensor):
        radius = torch.sqrt(sample[:, 0])
        theta = 2 * torch.pi * sample[:, 1]
        x = radius * torch.cos(theta)
        y = radius * torch.sin(theta)
        return torch.stack((x, y), dim=1)

    @dr.wrap_ad(source="drjit", target="torch")
    def square_to_cosine_hemisphere(self, sample):
        print(sample)
        p = self.square_to_uniform_disk(sample)
        z = torch.clamp(1.0 - torch.norm(p, p=2) ** 2, min=0.0)
        z = z.view(1, 1)
        return torch.cat((p, z), dim=1)

As you can see a custom square_to_cosine_hemisphere function is being used for sampling.

When I render the scene like this:

import mitsuba as mi
mi.set_variant("cuda_ad_rgb")
mi.set_log_level(mi.LogLevel.Warn)
import custom_diffuse_brdf

mi.register_bsdf(
    "CustomDiffuseBSDF",
    lambda props: custom_diffuse_brdf.CustomDiffuseBSDF(props),
)

scene = mi.load_file("../scenes/matpreview/scene_custom_diffuse.xml")

image = mi.render(scene, spp=64)

plt.axis("off")
plt.imshow(image ** (1.0 / 2.2)); # approximate sRGB tonemapping

the result looks like this:

However, the expected result is this (i.e., just changing sampling to bs.wo = mi.warp.square_to_cosine_hemisphere(sample2)):

The scene is defined like this:

<scene version="3.0.0">
    <default name="spp" value="64"/>
    <default name="width" value="683"/>
    <default name="height" value="512"/>
    <default name="max_depth" value="-1"/>

    <integrator type="path">
        <integer name="max_depth" value="$max_depth"/>
    </integrator>

    <sensor type="perspective" id="camera">
        <string name="fov_axis" value="smaller" />
        <float name="focus_distance" value="6.0" />
        <float name="fov" value="28.8415" />
        <transform name="to_world">
            <lookat target="3.04072, -2.85176, 2.80939"
                    origin="3.69558, -3.46243, 3.25463"
                    up="-0.317366, 0.312466, 0.895346" />
        </transform>

        <sampler type="independent">
            <integer name="sample_count" value="$spp" />
        </sampler>

        <film type="hdrfilm" id="film">
            <integer name="width" value="$width" />
            <integer name="height" value="$height" />
            <string name="pixel_format" value="rgb" />
            <rfilter type="gaussian" />
        </film>
    </sensor>

    <emitter type="envmap" id="emitter-envmap">
        <string name="filename" value="envmap.exr" />
        <transform name="to_world">
            <rotate y="1" angle="-180" />
            <matrix value="-0.224951 -0.000001 -0.974370 0.000000 -0.974370 0.000000 0.224951 0.000000 0.000000 1.000000 -0.000001 8.870000 0.000000 0.000000 0.000000 1.000000" />
        </transform>
        <float name="scale" value="3" />
    </emitter>

    <bsdf type="diffuse" id="bsdf-diffuse">
        <rgb name="reflectance" value="0.18 0.18 0.18" />
    </bsdf>

    <texture type="checkerboard" id="texture-checkerboard">
        <rgb name="color0" value="0.4" />
        <rgb name="color1" value="0.2" />
        <transform name="to_uv">
            <scale x="8.000000" y="8.000000" />
        </transform>
    </texture>

    <bsdf type="diffuse" id="bsdf-plane">
        <ref name="reflectance" id="texture-checkerboard" />
    </bsdf>

    <bsdf type="CustomDiffuseBSDF" id="bsdf-matpreview">
        <rgb name="reflectance" value="0.9 0.18 0.18" />
    </bsdf>

    <shape type="serialized" id="shape-plane">
        <string name="filename" value="matpreview.serialized" />
        <integer name="shape_index" value="0" />
        <transform name="to_world">
            <rotate z="1" angle="-4.3" />
            <matrix value="3.38818 -4.06354 0 -1.74958 4.06354 3.38818 0 1.43683 0 0 5.29076 -0.0120714 0 0 0 1" />
        </transform>
        <ref name="bsdf" id="bsdf-plane" />
    </shape>

    <shape type="serialized" id="shape-matpreview-interior">
        <string name="filename" value="matpreview.serialized" />
        <integer name="shape_index" value="1" />
        <transform name="to_world">
            <matrix value="1 0 0 0 0 1 0 0 0 0 1 0.0252155 0 0 0 1" />
        </transform>
        <ref name="bsdf" id="bsdf-diffuse" />
    </shape>

    <shape type="serialized" id="shape-matpreview-exterior">
        <string name="filename" value="matpreview.serialized" />
        <integer name="shape_index" value="2" />
        <transform name="to_world">
            <matrix value="0.614046 0.614047 0 -1.78814e-07 -0.614047 0.614046 0 2.08616e-07 0 0 0.868393 1.02569 0 0 0 1" />
            <translate z="0.01" />
        </transform>

        <ref name="bsdf" id="bsdf-matpreview" />
    </shape>

</scene>

When I use the custom sampling method and print passed original sample l get an output like this:

tensor([[-1.0000e-06,  0.0000e+00]], device='cuda:0', requires_grad=True)
tensor([[1., 0.]], device='cuda:0', requires_grad=True)

More strangely, when I replace sample2 using arbitrary constant values like this:

test = mi.Point2f([1, 1])
sample2_flat = dr.ravel(test)
wo_torch = self.square_to_cosine_hemisphere(
    mi.TensorXf(sample2_flat, shape=(1, 2))
)
wo_flat = dr.ravel(wo_torch)
bs.wo = dr.unravel(mi.Vector3f, wo_flat)

I get identical rendering result just as in the first image. The same thing happens if I just set bs.wo = mi.Vector3f(0, 0, 0).

Could you please help me with what I am doing wrong?

Steps to reproduce

The above-mentioned scene and code should help you to reproduce my results. Please let me know if you need further information :)

[njroussel]

Hi @sapo17

I'm surprised this ran without any error message 🤯
What you are doing isn't technically possible in megakernel/recorded mode, you need to run your code in wavefront mode with:

dr.set_flag(dr.JitFlag.VCallRecord, False)
dr.set_flag(dr.JitFlag.LoopRecord, False)

There are a couple of relevant discussion threads about this already. Here's one of them: #812
Hopefully this fixes the rendering too.

[sapo17]

Thank you @njroussel this fixed my issue :-) Also thanks for the reference, I must have missed it.

FYI Unfortunately I haven't got any errors or warnings (assuming jupyter notebook did not suppressed it somehow)