[Suggestion] Modify `mi.traverse()` to allow "pinning" scene parameter names to instance ID

[leroyvn]

Summary

This is a suggestion to contribute to solving the issue of scene parameters appearing under hard-to-predict names after scene tree traversal (see #508 for context).

The problem

When running this script:

import mitsuba as mi

mi_scene = mi.load_dict(
    {
        "type": "scene",
        "some_bsdf": {"type": "diffuse"},
        "rectangle": {
            "type": "rectangle",
            "bsdf": {
                "type": "ref",
                "id": "some_bsdf",
            },
        },
        "disk": {
            "type": "disk",
            "bsdf": {
                "type": "ref",
                "id": "some_bsdf",
            },
        },
    }
)
print(mi.traverse(mi_scene))

we get the following output:

SceneParameters[
  ----------------------------------------------------------------------------------
  Name                           Flags    Type  Parent
  ----------------------------------------------------------------------------------
  disk.bsdf.reflectance.value    ∂        float UniformSpectrum
  disk.to_world                  ∂, D     Transform4f Disk
  rectangle.to_world             ∂, D     Transform4f Rectangle
]

The problem here is that node names are determined by the scene tree structure, which depends on the order in which objects are processed during scene loading. This, from my understanding, depends on the alphabetical order of the keys in the scene dictionary. In this example, the BSDF appears as a child of "disk", and will appear as a child of "rectangle" if "disk" is renamed "zzz".

This behaviour makes it complicated to infer scene parameter names when assembling scenes from many scene dictionary fragments (typically when building a scene with a generator like we have in Eradiate). I provided a more confusing example in discussion #508.

Proposal

I believe a way to improve the predictibility of node names would be to offer to users the possibility to override node names with the underlying instance's ID. Typically, it seems reasonable in the aforementioned example to expect that the reflectance of some_bsdf can be found as some_bsdf.reflectance.value.

To do so, I suggest two things:

• Add to the traversal logic some node name override triggered upon detection of a non-empty ID string. The node name is simply replaced by the ID, and traversal continues as before. Child nodes appear following the original naming hierarchy, starting from the overridden node:

  SceneParameters[
  ----------------------------------------------------------------------------------
  Name                           Flags    Type  Parent
  ----------------------------------------------------------------------------------
  disk.to_world                  ∂, D     Transform4f Disk
  rectangle.to_world             ∂, D     Transform4f Rectangle
  some_bsdf.reflectance.value    ∂        float UniformSpectrum
  ]

• Make this behaviour optional with an additional parameter that accepts a boolean:

  def traverse(node: mi.Object, name_id_override: bool = False) -> SceneParameters:
    ...

This also results in a more intuitive behaviour when declaring BSDFs, phase functions, etc. as top-level objects in the scene dictionary and referencing them later on.

I experimented with this idea in my project, with the added possibility to restrict node name override using regular expressions passed to name_id_override.

Does such a modification look like a good idea to you?

[leroyvn]

Update: I've been experimenting further and it turns out that this sometimes doesn't work when triggering scene updates. The issue is that upon calling SceneParameters.set_dirty(), reverse scene traversal on an ID-aliased entry will eventually raise because parent node name is inferred from the current node's path. For instance, in the previous example, some_bsdf, which has a depth of 1 (it is a child of one of the shapes) has a "level-0" name. I patched SceneParameters and traverse() to track node name aliases, so that hierarchy climbing can resume from the original path once the top level of the aliased branch has been reached.

This brings up a secondary question: If a node has multiple parents, this is not tracked. This means that all parents except one (that is hard to predict) have to be set as dirty manually. Should this be a source of concern?

[njroussel]

This brings up a secondary question: If a node has multiple parents, this is not tracked. This means that all parents except one (that is hard to predict) have to be set as dirty manually. Should this be a source of concern?

That is correct. I never realized this. It can be concerning, I think for our applications we just never ran into situations where a child update actually required some complex update in the parent.

Overall I think the issue you pointed out here and the proposal are on the right track, with some aspects that still need to be fleshed out. I'll note this done - we have a few other ideas for larger changes involving scene descriptions and structure, we should consider how traverse interacts with those.

[leroyvn]

In case that would be useful, this is what I currently have:

import drjit as dr
import mitsuba as mi
from mitsuba.python.util import SceneParameters as _MitsubaSceneParameters

class SceneParameters(_MitsubaSceneParameters):
    def __init__(self, properties=None, hierarchy=None, aliases=None):
        super().__init__(properties, hierarchy)
        self.aliases = aliases if aliases is not None else {}

    def set_dirty(self, key: str):
        # Inherit docstring

        value, _, node, flags = self.properties[key]

        is_nondifferentiable = flags & mi.ParamFlags.NonDifferentiable.value
        if is_nondifferentiable and dr.grad_enabled(value):
            mi.Log(
                mi.LogLevel.Warn,
                f"Parameter '{key}' is marked as non-differentiable but has "
                "gradients enabled, unexpected results may occur!",
            )

        node_key = key  # Key of current node
        while node is not None:
            parent, depth = self.hierarchy[node]

            name = node_key
            if parent is not None:
                if "." not in name and depth > 0:
                    # We've hit the top level from an ID-aliased node:
                    # Resolve the alias to finish climbing the hierarchy
                    node_key = self.aliases[name]
                node_key, name = node_key.rsplit(".", 1)

            self.nodes_to_update.setdefault((depth, node), set())
            self.nodes_to_update[(depth, node)].add(name)

            node = parent

        return self.properties[key]

def mi_traverse(
    obj: mi.Object, name_id_override: str | list[str] | bool | None = None
) -> mi.SceneParameters:
    """
    Traverse a node of the Mitsuba scene graph and return scene parameters as
    a mutable mapping.

    Parameters
    ----------
    obj : mitsuba.Object
        Mitsuba scene graph node to be traversed.

    name_id_override : str or list of str, optional
        If set, this argument will be used to select nodes in the scene tree
        whose names will be "pinned" to their ID. Passed values are used as
        regular expressions, with all that it implies regarding ID string
        matching. If this parameter is set to ``True``, a regex that matches
        anything is used.

    Returns
    -------
    SceneParameters
    """

    if name_id_override is None or name_id_override is False:
        name_id_override = []

    if name_id_override is True:
        name_id_override = [r".*"]

    if type(name_id_override) is not list:
        name_id_override = [name_id_override]

    import re

    regexps = [re.compile(k).match for k in name_id_override]

    class SceneTraversal(mi.TraversalCallback):
        def __init__(
            self,
            node,
            parent=None,
            properties=None,
            hierarchy=None,
            prefixes=None,
            name=None,
            depth=0,
            flags=+mi.ParamFlags.Differentiable,
            aliases=None,
        ):
            mi.TraversalCallback.__init__(self)
            self.properties = dict() if properties is None else properties
            self.hierarchy = dict() if hierarchy is None else hierarchy
            self.prefixes = set() if prefixes is None else prefixes
            self.aliases = dict() if aliases is None else aliases

            node_id = node.id()
            if name_id_override and node_id:
                for r in regexps:
                    if r(node_id):
                        if node_id != name:
                            self.aliases[node_id] = name
                        name = node_id
                        break

            if name is not None:
                ctr, name_len = 1, len(name)
                while name in self.prefixes:
                    name = f"{name[:name_len]}_{ctr}"
                    ctr += 1
                self.prefixes.add(name)

            self.name = name
            self.node = node
            self.depth = depth
            self.hierarchy[node] = (parent, depth)
            self.flags = flags

        def put_parameter(self, name, ptr, flags, cpptype=None):
            name = name if self.name is None else self.name + "." + name

            flags = self.flags | flags
            # Non-differentiable parameters shouldn't be flagged as discontinuous
            if (flags & mi.ParamFlags.NonDifferentiable) != 0:
                flags = flags & ~mi.ParamFlags.Discontinuous

            self.properties[name] = (ptr, cpptype, self.node, self.flags | flags)

        def put_object(self, name, node, flags):
            if node is None or node in self.hierarchy:
                return

            cb = SceneTraversal(
                node=node,
                parent=self.node,
                properties=self.properties,
                hierarchy=self.hierarchy,
                prefixes=self.prefixes,
                name=name if self.name is None else f"{self.name}.{name}",
                depth=self.depth + 1,
                flags=self.flags | flags,
                aliases=self.aliases,
            )
            node.traverse(cb)

    cb = SceneTraversal(obj)
    obj.traverse(cb)

    return SceneParameters(cb.properties, cb.hierarchy, cb.aliases)