Suggestion: add support for VisPy as a second "preview backend"

[jeertmans]

Note

Edit: as maintaining a second preview backend would require a lot of extra work from the Sionna team, I developed a minimal, but fully compatible, separate package, sionna-vispy, that can be imported to use VisPy instead of pythreejs viewer.

---

Hi everyone!

This is a late follow-up of #219. After playing a bit myself with the different plotting libraries, I found VisPy
to be a great solution for previewing large and complex scenes,
while offerings a nice high-level API.

Currently, it looks like the only supported previewer is PyThreeJS. While it produces a beautiful and interactive canvas, the main issue (to me) is that it does not work outside interactive environments (like Notebooks).

This constraint makes it hard to:

1. run Sionna related projects as command-line scripts;
2. and embed Sionna's visualization inside other applications.

VisPy has the following nice advantages:

• it works on most macOS, Ubuntu, and Linux, and Jupyter Notebooks;
• it supports many backends, so you are not constrained to one;
• it provides both interactive plots (even in Jupyter Notebooks);
• the syntax seems relatively close to that of PyThreeJS (see example below);
• it uses OpenGL, so this is quite performant;
• and you can easily embed it inside more complex applications, e.g., with Qt.

But also some disadvantages:

• this is one more dependency;
• you need to install the backend manually (or constrain the user to one), and it is usually quite large (a few hundreds Mo);
• it requires a relatively recent GPU and does not produce JS plots (so it is only interactive while the kernel is alive);

Further questions

If you think this could be a great fit for this project, I'd be happy to help to implement
a VisPy version of scene.preview.

However, a few questions should be answered first:

• Should we modify InteractiveDisplay to support multiple backends, or should we create another, separate, method for plotting with VisPy? In my opinion, the first solution allows reusing a lot of code, and makes it more transparent to the user.
• Should we add VisPy as a direct dependency, or as an "extra"?
• The same question for any VisPy backend: should we impose one, or let the user choose?
• The same question, but for the jupyter-rfb backend, that is only used inside notebooks?

Example

The following example code shows how to preview the scene from the Ray Tracing tutorial.
The code for plot_scene is a slightly modified version of

sionna/sionna/rt/previewer.py

Lines 201 to 244 in 2cb12fd

	def plot_scene(self):
	"""
	Plots the meshes that make the scene.
	"""
	shapes = self._scene.mi_scene.shapes()
	n = len(shapes)
	if n <= 0:
	return
	palette = None
	si = dr.zeros(mi.SurfaceInteraction3f)
	si.wi = mi.Vector3f(0, 0, 1)
	# Shapes (e.g. buildings)
	vertices, faces, albedos = [], [], []
	f_offset = 0
	for i, s in enumerate(shapes):
	null_transmission = s.bsdf().eval_null_transmission(si).numpy()
	if np.min(null_transmission) > 0.99:
	# The BSDF for this shape was probably set to `null`, do not
	# include it in the scene preview.
	continue
	n_vertices = s.vertex_count()
	v = s.vertex_position(dr.arange(mi.UInt32, n_vertices))
	vertices.append(v.numpy())
	f = s.face_indices(dr.arange(mi.UInt32, s.face_count()))
	faces.append(f.numpy() + f_offset)
	f_offset += n_vertices
	albedo = s.bsdf().eval_diffuse_reflectance(si).numpy()
	if not np.any(albedo > 0.):
	if palette is None:
	palette = matplotlib.cm.get_cmap('Pastel1_r')
	albedo[:] = palette((i % palette.N + 0.5) / palette.N)[:3]
	albedos.append(np.tile(albedo, (n_vertices, 1)))
	# Plot all objects as a single PyThreeJS mesh, which is must faster
	# than creating individual mesh objects in large scenes.
	self._plot_mesh(np.concatenate(vertices, axis=0),
	np.concatenate(faces, axis=0),
	persist=True, # The scene geometry is persistent
	colors=np.concatenate(albedos, axis=0))

where only the call to self._plot_mesh is replaced with VisPy code.

You can find the full code below:

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

from sionna.rt import load_scene
from vispy import app
from vispy.scene import SceneCanvas
from vispy.scene.visuals import Mesh


def plot_scene(scene, view):
    # This code is copied
    shapes = scene.mi_scene.shapes()
    n = len(shapes)
    if n <= 0:
        return

    palette = None
    si = dr.zeros(mi.SurfaceInteraction3f)
    si.wi = mi.Vector3f(0, 0, 1)

    # Shapes (e.g. buildings)
    vertices, faces, albedos = [], [], []
    f_offset = 0
    for i, s in enumerate(shapes):
        null_transmission = s.bsdf().eval_null_transmission(si).numpy()
        if np.min(null_transmission) > 0.99:
            # The BSDF for this shape was probably set to `null`, do not
            # include it in the scene preview.
            continue

        n_vertices = s.vertex_count()
        v = s.vertex_position(dr.arange(mi.UInt32, n_vertices))
        vertices.append(v.numpy())
        f = s.face_indices(dr.arange(mi.UInt32, s.face_count()))
        faces.append(f.numpy() + f_offset)
        f_offset += n_vertices

        albedo = s.bsdf().eval_diffuse_reflectance(si).numpy()
        if not np.any(albedo > 0.):
            if palette is None:
                palette = matplotlib.cm.get_cmap('Pastel1_r')
            albedo[:] = palette((i % palette.N + 0.5) / palette.N)[:3]

        albedos.append(np.tile(albedo, (n_vertices, 1)))

    # This code is new
    view.add(
        Mesh(
            vertices=np.concatenate(vertices, axis=0),
            faces=np.concatenate(faces, axis=0),
            vertex_colors=np.concatenate(albedos, axis=0),
        )
    )


scene = load_scene(sionna.rt.scene.munich)

#  scene.preview()  # This fails outside IPython

canvas = SceneCanvas(keys="interactive", bgcolor="white")
view = canvas.central_widget.add_view()
view.camera = "turntable"
view.camera.depth_value = 1e3 

plot_scene(scene, view)

view.camera.set_range()  # Center scene

canvas.show()
app.run()

[jeertmans]

Hello @jhoydis, as you suggested in our discussion last week, I wrote a small Python library, sionna-vispy, that implements InteractiveDisplay using VisPy instead of pythreejs. I'm sharing this here in case anyone is interested.

The library is quite minimal, but should correctly support all the features from the current previewer, and provides a nice context-based mechanism to patch the pythreejs previewer with the VisPy one.

[merlinND]

Thanks @jeertmans for offering this non-intrusive and useful solution!

[jeertmans]

You're welcome @merlinND! I hope other people will find it useful, and maybe find some bugs that I can later fix :-)

[sunskyday]

very cool, It`s very useful, thanks very much

[jeertmans]

Thanks! :-)

Feel free to give any feedback on the GitHub repository itself