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]:
