Source code for eradiate.kernel._render

from __future__ import annotations

import logging
import typing as t
import warnings

import attrs
import mitsuba as mi
from import tqdm

from ._kernel_dict import UpdateMapTemplate
from .._config import ProgressLevel, config
from ..attrs import documented, parse_docs
from ..contexts import KernelContext
from ..rng import SeedState, root_seed_state

logger = logging.getLogger(__name__)

# ------------------------------------------------------------------------------
#                         Parameter lookup strategies
# ------------------------------------------------------------------------------

[docs]@parse_docs @attrs.frozen class TypeIdLookupStrategy: """ This parameter ID lookup strategy searches for a Mitsuba type and object ID match. Instances are callables which take, as argument, the current node during a Mitsuba scene tree traversal and, optionally, its path in the Mitsuba scene tree. If the lookup succeeds, the full parameter path is returned. """ node_type: type = documented( attrs.field(validator=attrs.validators.instance_of(type)), doc="Type of the node which will be looked up.", type="type", ) node_id: str = documented( attrs.field(validator=attrs.validators.instance_of(str)), doc="ID of the node which will be looked up.", type="str", ) parameter_relpath: str = documented( attrs.field(validator=attrs.validators.instance_of(str)), doc="Parameter path relative to its parent object.", type="str", ) def __call__(self, node, node_path: str | None = None) -> str | None: if isinstance(node, self.node_type) and == self.node_id: prefix = f"{node_path}." if node_path is not None else "" return f"{prefix}{self.parameter_relpath}" else: return None
# ------------------------------------------------------------------------------ # Mitsuba scene traversal # ------------------------------------------------------------------------------
[docs]@parse_docs @attrs.define class MitsubaObjectWrapper: """ This container aggregates a Mitsuba object, its associated parameters and a set of updaters that can be used to modify the scene parameters. See Also -------- :func:`mi_traverse` """ obj: mi.Object = documented( attrs.field(repr=lambda x: "Scene[...]" if isinstance(x, mi.Scene) else str(x)), doc="Mitsuba object.", type="mitsuba.Object", ) parameters: mi.SceneParameters | None = documented( attrs.field( default=None, repr=lambda x: "SceneParameters[...]" if x.__class__.__name__ == "SceneParameters" else str(None), ), doc="Mitsuba scene parameter map.", type="mitsuba.SceneParameters", init_type="mitsuba.SceneParameters, optional", default="None", ) umap_template: UpdateMapTemplate | None = documented( attrs.field( default=None, repr=lambda x: "UpdateMapTemplate[...]" if isinstance(x, UpdateMapTemplate) else str(x), ), doc="An update map template, which can be rendered and used to update " "Mitsuba scene parameters depending on context information.", type=".UpdateMapTemplate", init_type=".UpdateMapTemplate, optional", default="None", )
[docs] def drop_parameters(self) -> None: """ Reduce the size of the scene parameter table :attr:`.parameters` by only keeping elements whose keys are listed in the parameter update map template :attr:`.umap_template`. For parameters associated with a lookup protocol, the looked up parameter ID is checked and used. """ if self.umap_template is not None: keys = [] for name, param in self.umap_template.items(): if param.lookup_strategy is not None: if param.parameter_id is not None: keys.append(param.parameter_id) else: warnings.warn( f"Parameter '{name}' has a lookup strategy but the " "associated parameter ID is undefined; was a " "parameter lookup performed during the Mitsuba " "scene traversal?" ) else: keys.append(name) self.parameters.keep(keys)
[docs]def mi_traverse( obj: mi.Object, umap_template: UpdateMapTemplate | None = None, ) -> MitsubaObjectWrapper: """ Traverse a node of the Mitsuba scene graph and return a container holding the Mitsuba scene, its parameter map and an updated parameter update map. Parameters ---------- obj : mitsuba.Object Mitsuba scene graph node to be traversed. umap_template : .UpdateMapTemplate, optional An additional update map template which is to be updated during traversal. This is used to perform parameter lookup during traversal. Returns ------- MitsubaObjectWrapper A container holding the traversed object, the corresponding parameter map and the parameter update map (if any). Notes ----- This is a reimplementation of the :func:`mitsuba.traverse` function. """ umap_template = ( UpdateMapTemplate( if umap_template is not None else UpdateMapTemplate() ) lookups = { k: v for k, v in umap_template.items() if v.parameter_id is None and v.lookup_strategy is not None } class SceneTraversal(mi.TraversalCallback): def __init__( self, node, parent=None, properties=None, hierarchy=None, prefixes=None, name=None, depth=0, flags=+mi.ParamFlags.Differentiable, ): mi.TraversalCallback.__init__(self) = 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 if name is not None: ctr, name_len = 1, len(name) while name in self.prefixes: name = "%s_%i" % (name[:name_len], ctr) ctr += 1 self.prefixes.add(name) = name self.node = node self.depth = depth self.hierarchy[node] = (parent, depth) self.flags = flags # Try and recover a parameter ID from this node for name, uparam in list(lookups.items()): lookup_result = uparam.lookup_strategy(self.node, if lookup_result is not None: uparam.parameter_id = lookup_result del lookups[ name ] # Remove successful lookups to accelerate future searches def put_parameter(self, name, ptr, flags, cpptype=None): name = name if is None else + "." + 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[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,, hierarchy=self.hierarchy, prefixes=self.prefixes, name=name if is None else f"{}.{name}", depth=self.depth + 1, flags=self.flags | flags, ) node.traverse(cb) cb = SceneTraversal(obj) obj.traverse(cb) # Check if there are unsuccessful lookups if lookups: warnings.warn( "There were unsuccessful Mitsuba scene parameter lookups: " f"{list(lookups.keys())}" ) return MitsubaObjectWrapper( obj=obj, parameters=mi.SceneParameters(, cb.hierarchy), umap_template=umap_template, )
# ------------------------------------------------------------------------------ # Mitsuba scene render # ------------------------------------------------------------------------------
[docs]def mi_render( mi_scene: MitsubaObjectWrapper, ctxs: list[KernelContext], sensors: None | int | list[int] = None, spp: int = 0, seed_state: SeedState | None = None, ) -> dict[t.Any, mi.Bitmap]: """ Render a Mitsuba scene multiple times given specified contexts and sensor indices. Parameters ---------- mi_scene : .MitsubaObjectWrapper Mitsuba scene to render. ctxs : list of :class:`.KernelContext` List of contexts used to generate the parameter update table at each iteration. sensors : int or list of int, optional Sensor indices to render. If ``None`` (default), all sensors are rendered. spp : int, optional, default: 0 Number of samples per pixel. If set to 0 (default), the value set in the original scene definition takes precedence. seed_state : .SeedState, optional Seed state used to generate seeds to initialize Mitsuba's RNG at each run. If unset, Eradiate's root seed state is used. Returns ------- dict A nested dictionary mapping context and sensor indices to rendered bitmaps. Notes ----- This function wraps sequential calls to :func:`mitsuba.render`. """ if seed_state is None: logger.debug("Using default RNG seed generator") seed_state = root_seed_state results = {} # Loop on contexts with tqdm( initial=0, total=len(ctxs), unit_scale=1.0, leave=True, bar_format="{desc}{n:g}/{total:g}|{bar}| {elapsed}, ETA={remaining}", disable=(config.progress < ProgressLevel.SPECTRAL_LOOP) or len(ctxs) <= 1, ) as pbar: for ctx in ctxs: pbar.set_description( f"Eradiate [{ctx.index_formatted}]", refresh=True, ) logger.debug("Updating Mitsuba scene parameters") mi_scene.parameters.update(mi_scene.umap_template.render(ctx)) if sensors is None: mi_sensors = [ (i, sensor) for i, sensor in enumerate(mi_scene.obj.sensors()) ] else: if isinstance(sensors, int): sensors = [sensors] mi_sensors = [(i, mi_scene.obj.sensors()[i]) for i in sensors] # Loop on sensors for i_sensor, mi_sensor in mi_sensors: # Render sensor seed = int( logger.debug( 'Running Mitsuba for sensor "%s" with seed value %s',, seed, ) mi.render(mi_scene.obj, sensor=i_sensor, seed=seed, spp=spp) # Store result in a new Bitmap object siah = if siah not in results: results[siah] = {} results[siah][] = mi.Bitmap( pbar.update() return results