Custom Python plugin

Overview

In Mitsuba it easy to add custom code for BSDFs, integrators, emitters, sensors, and more. This tutorial will show you how to create custom plugins in Python and register them for use.

To illustrate this, we are going to implement a new tinted dielectric BSDF, that behaves much like a regular dielectric BSDF but adds a colorful tint to the reflections at grazing angle. We will then register this new BSDF and use it to render a simple scene.

🚀 You will learn how to:

  • Create a custom BSDF plugin in Python

  • Register a custom plugin to the system

  • Use a registered custom plugin

Setup

Custom Mitsuba plugins written in Python work best with Just-In-Time (JIT) variants. This is because it would be pretty inefficient to execute Python BSDF code for millions of scalar light paths. JIT variants, on the other hand, only execute a few calls to those methods on arrays containing millions of entries at once, mitigating the overhead coming from the the Python layer. In this example we will therefore stick with the llvm_ad_rgb variant.

[1]:
import drjit as dr
import mitsuba as mi

mi.set_variant('llvm_ad_rgb')

Implementation

As mentioned in the tutorial overview, we are going to implement a tinted dielectric BSDF in this tutorial. This code is very similar to the actual C++ implementation of the dielectric BSDF, so we will not look at it in great detail.

First, our BSDF Python class MyBSDF needs to inherit from BSDF. This allows us to override the constructor method as well as sample(), eval() and pdf().

The constructor takes a Properties object as an argument, which can be used to read parameters defined in the XML scene description or passed in the load_dict() dictionary. Here we read the index of refraction ratio eta as well as the tint color tint from props. In the constructor, we also properly set the BSDF members used in other methods like m_flags and m_components.

Similarly to the regular dielectric BSDF, the eval() and pdf() methods of our custom BSDF should always return 0.0 and never be called as it is a degenerate BSDF described by a Dirac delta distribution.

Regarding the sample() method, apart from the computation of the tinted reflection value value_r, the rest of code should be identical to the C++ implementation of dielectric.

Note that it is also possible to override the to_string() method which is called in any printing/logging routine.

Finally, we override the implementation of the traverse() and parameters_changed() methods to expose the tint parameter via the traverse mechanism. This will allow us to edit this parameter after the BSDF is instanciated.

[2]:
class MyBSDF(mi.BSDF):
    def __init__(self, props):
        mi.BSDF.__init__(self, props)

        # Read 'eta' and 'tint' properties from `props`
        self.eta = 1.33
        if props.has_property('eta'):
            self.eta = props['eta']

        self.tint = mi.Color3f(props['tint'])

        # Set the BSDF flags
        reflection_flags   = mi.BSDFFlags.DeltaReflection   | mi.BSDFFlags.FrontSide | mi.BSDFFlags.BackSide
        transmission_flags = mi.BSDFFlags.DeltaTransmission | mi.BSDFFlags.FrontSide | mi.BSDFFlags.BackSide
        self.m_components  = [reflection_flags, transmission_flags]
        self.m_flags = reflection_flags | transmission_flags

    def sample(self, ctx, si, sample1, sample2, active):
        # Compute Fresnel terms
        cos_theta_i = mi.Frame3f.cos_theta(si.wi)
        r_i, cos_theta_t, eta_it, eta_ti = mi.fresnel(cos_theta_i, self.eta)
        t_i = dr.maximum(1.0 - r_i, 0.0)

        # Pick between reflection and transmission
        selected_r = (sample1 <= r_i) & active

        # Fill up the BSDFSample struct
        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, tint based on the incident angle (more tint at grazing angle)
        value_r = dr.lerp(mi.Color3f(self.tint), mi.Color3f(1.0), dr.clip(cos_theta_i, 0.0, 1.0))

        # For transmission, radiance must be scaled to account for the solid angle compression
        value_t = mi.Color3f(1.0) * dr.square(eta_ti)

        value = dr.select(selected_r, value_r, value_t)

        return (bs, value)

    def eval(self, ctx, si, wo, active):
        return 0.0

    def pdf(self, ctx, si, wo, active):
        return 0.0

    def eval_pdf(self, ctx, si, wo, active):
        return 0.0, 0.0

    def traverse(self, callback):
        callback.put_parameter('tint', self.tint, mi.ParamFlags.Differentiable)

    def parameters_changed(self, keys):
        print("🏝️ there is nothing to do here 🏝️")

    def to_string(self):
        return ('MyBSDF[\n'
                '    eta=%s,\n'
                '    tint=%s,\n'
                ']' % (self.eta, self.tint))

Plugin registration

There’s only one more thing to do before we can use our custom BSDF in scenes. We need to register it in the system so it can be used. This can be done by calling the register_bsdf() function and specifying the name to be used to instantiate this plugin. The function takes a constructor lambda function as the second parameter.

📑 Note

Similar functions exist for other types of plugins, e.g.

  • register_integrator

  • register_emitter

  • register_sensor

  • register_film

  • register_mesh

  • register_texture

  • register_volume

  • register_phasefunction

  • register_medium

  • register_sampler

[3]:
mi.register_bsdf("mybsdf", lambda props: MyBSDF(props))

Plugin instantiation

You can now use this plugin like you would with any other BSDF plugin and set the appropriate properties of the BSDF expected in its constructor in the XML or dict representation.

[4]:
my_bsdf = mi.load_dict({
    'type' : 'mybsdf',
    'tint' : [0.2, 0.9, 0.2],
    'eta' : 1.33
})

my_bsdf
[4]:
MyBSDF[
    eta=1.33,
    tint=[[0.2, 0.9, 0.2]],
]

Rendering

Finally, let’s use our custom BSDF in an actual scene and render it to see how our tinted BSDF looks like.

[5]:
scene = mi.load_dict({
    'type': 'scene',
    'integrator': {
        'type': 'path'
    },
    'light': {
        'type': 'constant',
        'radiance': 0.99,
    },
    'sphere' : {
        'type': 'sphere',
        'bsdf': my_bsdf
    },
    'sensor': {
        'type': 'perspective',
        'to_world': mi.ScalarTransform4f().look_at(origin=[0, -5, 5],
                                                 target=[0, 0, 0],
                                                 up=[0, 0, 1]),
    }
})

image = mi.render(scene)

mi.Bitmap(image).convert(srgb_gamma=True)
[5]:

Edit parameters

As expected, it is possible to access our custom BSDF’s parameters using the traverse mechanism.

[6]:
params = mi.traverse(scene)
params
[6]:
SceneParameters[
  ----------------------------------------------------------------------------------------
  Name                                 Flags    Type           Parent
  ----------------------------------------------------------------------------------------
  light.sampling_weight                         float          ConstantBackgroundEmitter
  light.radiance.value                 ∂        Float          UniformSpectrum
  sensor.near_clip                              float          PerspectiveCamera
  sensor.far_clip                               float          PerspectiveCamera
  sensor.shutter_open                           float          PerspectiveCamera
  sensor.shutter_open_time                      float          PerspectiveCamera
  sensor.film.size                              ScalarVector2u HDRFilm
  sensor.film.crop_size                         ScalarVector2u HDRFilm
  sensor.film.crop_offset                       ScalarPoint2u  HDRFilm
  sensor.x_fov                         ∂, D     Float          PerspectiveCamera
  sensor.principal_point_offset_x      ∂, D     Float          PerspectiveCamera
  sensor.principal_point_offset_y      ∂, D     Float          PerspectiveCamera
  sensor.to_world                      ∂, D     Transform4f    PerspectiveCamera
  sphere.bsdf.tint                     ∂        Color3f        BSDF
  sphere.silhouette_sampling_weight             float          Sphere
  sphere.to_world                      ∂, D     Transform4f    Sphere
]

We can then update the tint value:

[7]:
key = 'sphere.bsdf.tint'
params[key] = mi.Color3f(0.9, 0.2, 0.2)
params.update();
🏝️ there is nothing to do here 🏝️

When re-rendering this scene, we now see that the new tint is indeed used!

[8]:
image = mi.render(scene)

mi.Bitmap(image).convert(srgb_gamma=True)
[8]:

See also