Source code for eradiate.kernel._kernel_dict

"""
This module contains the infrastructure used to generate Mitsuba scene
dictionaries and scene parameter maps.
"""

from __future__ import annotations

import enum
import typing as t
from collections import UserDict

import attrs
import mitsuba as mi

from ..attrs import documented, parse_docs
from ..contexts import KernelContext
from ..util.misc import nest


[docs] @parse_docs @attrs.define class InitParameter: """ This class declares an Eradiate parameter in a Mitsuba scene dictionary. It holds an evaluation protocol for this parameter depending on context information. """ #: Sentinel value indicating that a parameter is not used UNUSED: t.ClassVar[object] = object() evaluator: t.Callable = documented( attrs.field(validator=attrs.validators.is_callable()), doc="A callable that returns the value of the parameter for a given " "context, with signature ``f(ctx: KernelContext) -> Any``.", type="callable", ) def __call__(self, ctx: KernelContext) -> t.Any: return self.evaluator(ctx)
[docs] @parse_docs @attrs.define class UpdateParameter: """ This class declares an Eradiate parameter in a Mitsuba scene parameter update map. It holds an evaluation protocol depending on context information. See Also -------- :class:`.KernelContext`, :class:`.TypeIdLookupStrategy` """ #: Sentinel value indicating that a parameter is not used UNUSED: t.ClassVar[object] = object()
[docs] class Flags(enum.Flag): """ Update parameter flags. """ NONE = 0 SPECTRAL = enum.auto() #: Varies during the spectral loop GEOMETRIC = enum.auto() #: Triggers a scene rebuild ALL = SPECTRAL | GEOMETRIC
evaluator: t.Callable = documented( attrs.field(validator=attrs.validators.is_callable()), doc="A callable that returns the value of the parameter for a given " "context, with signature ``f(ctx: KernelContext) -> Any``.", type="callable", ) flags: Flags = documented( attrs.field(default=Flags.ALL), doc="Flags specifying parameter attributes. By default, the declared " "parameter will pass all filters.", type=".Flags", default=".Flags.ALL", ) lookup_strategy: None | (t.Callable[[mi.Object, str], str | None]) = documented( attrs.field(default=None), doc="A callable that searches a Mitsuba scene tree node for a desired " "parameter ID: with signature " "``f(node: mi.Object, node_relpath: str) -> Optional[str]``.", type="callable or None", init_type="callable, optional", default="None", ) parameter_id: str | None = documented( attrs.field(default=None), doc="The full ID of the Mitsuba scene parameter to update.", type="str or None", init_type="str, optional", ) def __call__(self, ctx: KernelContext) -> t.Any: return self.evaluator(ctx)
[docs] @attrs.define(slots=False) class KernelDictTemplate(UserDict): """ A dict-like structure which defines the structure of an instantiable Mitsuba scene dictionary. Entries are indexed by dot-separated paths which can then be expanded to a nested dictionary using the :meth:`.render` method. Each entry can be either a hard-coded value which can be directly interpreted by the :func:`mitsuba.load_dict` function, or an :class:`.InitParameter` object which must be rendered before the template can be instantiated. """ data: dict[str, InitParameter] = attrs.field(factory=dict)
[docs] def render( self, ctx: KernelContext, nested: bool = True, drop: bool = True ) -> dict: """ Render the template as a nested dictionary using a parameter map to fill in empty fields. Parameters ---------- ctx : :class:`.KernelContext` A kernel dictionary context. nested : bool, optional If ``True``, the returned dictionary will be nested and suitable for instantiation by Mitsuba; otherwise, the returned dictionary will be flat. drop : bool, optional If ``True``, drop unused parameters. Parameters may be unused either because they were filtered out by the flags or because context information implied it. Returns ------- dict """ result = {} for k, v in list(self.items()): value = v(ctx) if isinstance(v, InitParameter) else v if (value is InitParameter.UNUSED) and drop: continue else: result[k] = value return nest(result, sep=".") if nested else result
[docs] @attrs.define(slots=False) class UpdateMapTemplate(UserDict): """ A dict-like structure which contains the structure of a Mitsuba scene parameter update map. Entries are indexed by dot-separated paths which can then be expanded to a nested dictionary using the :meth:`.render` method. """ data: dict[str, UpdateParameter] = attrs.field(factory=dict)
[docs] def remove(self, keys: str | list[str]) -> None: """ Remove all parameters matching the given regular expression. Parameters ---------- keys : str or list of str Regular expressions matching the parameters to remove. Notes ----- This method mutates the parameter map. """ if not isinstance(keys, list): keys = [keys] import re regexps = [re.compile(k).match for k in keys] keys = [k for k in self.keys() if any(r(k) for r in regexps)] for key in keys: del self.data[key]
[docs] def keep(self, keys: str | list[str]) -> None: """ Keep only parameters matching the given regular expression. Parameters ---------- keys : str or list of str Regular expressions matching the parameters to keep. Notes ----- This method mutates the parameter map. """ if not isinstance(keys, list): keys = [keys] import re regexps = [re.compile(k).match for k in keys] keys = [k for k in self.keys() if any(r(k) for r in regexps)] result = {k: self.data[k] for k in keys} self.data = result
[docs] def render( self, ctx: KernelContext, flags: UpdateParameter.Flags = UpdateParameter.Flags.ALL, drop: bool = False, ) -> dict: """ Evaluate the parameter map for a set of arguments. Parameters ---------- ctx : :class:`.KernelContext` A kernel dictionary context. flags : :class:`.ParamFlags` Parameter flags. Only parameters with at least one of the specified will pass the filter. drop : bool, optional If ``True``, drop unused parameters. Parameters may be unused either because they were filtered out by the flags or because context information implied it. Returns ------- dict Raises ------ ValueError If a value is not a :class:`.UpdateParameter`. ValueError If ``drop`` is ``False`` and the rendered parameter map contains an unused parameter. """ unused = [] result = {} for k in list( self.keys() ): # Ensures correct iteration even if the loop mutates the mapping v = self[k] if isinstance(v, UpdateParameter): key = k if v.parameter_id is None else v.parameter_id if v.flags & flags: result[key] = v(ctx) else: unused.append(k) if not drop: result[key] = UpdateParameter.UNUSED # Check for leftover empty values if not drop and unused: raise ValueError(f"Unevaluated parameters: {unused}") return result