procfunc 0.30.0__py3-none-any.whl

procfunc/tracer/trace.py

@@ -0,0 +1,222 @@
+"""
+torch.fx / jax-like function-compute-graph tracing tool, but specially designed for procedural generation functions
+
+
+This tool was heavily inspired by torch.fx.symbolic_trace https://docs.pytorch.org/docs/2.6/fx.html
+"""
+
+import builtins
+import inspect
+import logging
+import math
+import random
+from types import ModuleType
+from typing import Any, Callable
+
+import numpy as np
+import numpy.linalg
+
+import procfunc as pf
+from procfunc import compute_graph as cg
+from procfunc.util import pytree
+
+from .patch import (
+    Patcher,
+    PatchFunctionTarget,
+    TraceLevel,
+)
+from .proxy import RngProxy
+
+logger = logging.getLogger(__name__)
+
+_autowrap_modules: list[tuple[ModuleType, bool, TraceLevel]] = []
+
+
+def autowrap_module(
+    module: ModuleType,
+    allow_exec: bool = False,
+    trace_level: TraceLevel = TraceLevel.PRIMITIVES,
+):
+    _autowrap_modules.append((module, allow_exec, trace_level))
+
+
+autowrap_module(math, allow_exec=True)
+autowrap_module(np, allow_exec=True)
+autowrap_module(numpy.linalg, allow_exec=True)
+
+
+_banned_modules: list[ModuleType] = [
+    np.random,
+    random,
+]
+
+
+def add_banned_module(module: ModuleType):
+    _banned_modules.append(module)
+
+
+_patch_function_targets: list[PatchFunctionTarget] = []
+
+
+def add_wrap_target(target: PatchFunctionTarget):
+    _patch_function_targets.append(target)
+
+
+WRAP_BUILTINS = ["min", "max", "abs", "round", "sum"]
+
+for _builtin_name in WRAP_BUILTINS:
+    add_wrap_target(
+        PatchFunctionTarget(
+            frame=builtins.__dict__,
+            name=_builtin_name,
+            trace_level=TraceLevel.PRIMITIVES,
+            normalize=False,
+            allow_exec=True,
+            source_name="builtins",
+        )
+    )
+
+WRAP_CONSTRUCTORS = ["Vector", "Color", "Euler"]
+
+for _name in WRAP_CONSTRUCTORS:
+    add_wrap_target(
+        PatchFunctionTarget(
+            frame=pf.__dict__,
+            name=_name,
+            trace_level=TraceLevel.PRIMITIVES,
+            normalize=False,
+            allow_exec=True,
+            source_name="mathutils",
+        )
+    )
+
+
+_search_scopes: list[dict] = []
+
+
+def add_search_scope(module: ModuleType):
+    """
+    Causes an intermediate module to be searched to discover any existing targets that need to be patched
+
+    e.g. for procfunc.nodes.to_mesh_object, the original to_mesh_object is already a target, but we need to
+    add_search_scope on the `nodes` module so that that module's references to to_mesh_object get wrapped.
+    """
+
+    _search_scopes.append(module.__dict__)
+
+
+def _map_args(
+    func: Callable,
+    **inputs: Any,
+) -> dict[str, cg.Proxy | Any]:
+    signature = inspect.signature(func)
+
+    res = {}
+
+    for name, param in signature.parameters.items():
+        if name in inputs:
+            val = inputs[name]
+            if isinstance(val, np.random.Generator):
+                node = cg.ConstantNode(value=val)
+                res[name] = RngProxy(node, val, dirty=False)
+            elif isinstance(val, cg.ConstantNode) and isinstance(
+                val.value, np.random.Generator
+            ):
+                res[name] = RngProxy(val, val.value, dirty=False)
+            else:
+                res[name] = val
+            continue
+
+        if param.default is not param.empty:
+            res[name] = param.default
+            continue
+
+        node = cg.InputPlaceholderNode(
+            name=name, default_value=None, metadata={"varname": name}
+        )
+        res[name] = cg.Proxy(node)
+
+    return res
+
+
+def trace(
+    func: Callable,
+    trace_level: TraceLevel = TraceLevel.GENERATORS,
+    name: str | None = None,
+    **inputs: Any,
+):
+    """
+    Turn a python function into a graph datastructure.
+
+    Using this datastructure is (usually) equivelent to executing the function.
+
+    Args:
+        func: The function to trace
+        trace_level: Granularity of the graph. Functions at this level become leaves;
+            finer functions are traced through. choice() peeks through all options when
+            trace_level >= RANDOM_CONTROL, or resolves to the chosen branch when finer.
+    """
+
+    logger.debug(f"Tracing {func} {id(func)=} with {inputs=} {trace_level=}")
+
+    if name is None:
+        assert hasattr(func, "__name__")
+        assert isinstance(func.__name__, str)
+        name = func.__name__
+
+    proxy_args = _map_args(func, **inputs)
+
+    if pf.context.globals.current_trace_level is not None:
+        # TODO we can lift this restriction fairly(?) easily by having a global patcher & saving/restoring this state
+        # rather than setting to false at end of function. see fx.trace for an example
+        raise RuntimeError(
+            f"Can't trace {name}, tracing is already in progress for another function. "
+            "Nested tracing is not yet supported - contact the developers to request"
+        )
+
+    pf.context.globals.current_trace_level = trace_level.value
+
+    patcher = Patcher(
+        trace_level=trace_level,
+        autopatch_wrap_modules=_autowrap_modules,
+        autopatch_remove_modules=_banned_modules,
+        search_scopes=_search_scopes,
+        patch_functions=_patch_function_targets,
+    )
+
+    try:
+        func = patcher.apply_preexecute_patches(func, trace_level)
+        logger.debug(f"Executing {func.__name__} {id(func)=}")
+        func_result = func(**proxy_args)
+    finally:
+        pf.context.globals.current_trace_level = None
+        patcher.unpatch_all()
+
+    metadata = {
+        "operations": [
+            (trace, {"func": func, "trace_level": trace_level}),
+        ],
+    }
+
+    def extract_node(v):
+        if isinstance(v, cg.Proxy):
+            return v.node
+        return cg.ConstantNode(value=v)
+
+    outputs = pytree.PyTree(func_result)
+    outputs = outputs.map(extract_node)
+
+    input_nodes = {
+        k: v.node
+        for k, v in proxy_args.items()
+        if isinstance(v, cg.Proxy) and isinstance(v.node, cg.InputPlaceholderNode)
+    }
+
+    compgraph = cg.ComputeGraph(
+        inputs=pytree.PyTree(input_nodes),
+        outputs=outputs,
+        name=name,
+        metadata=metadata,
+    )
+
+    return compgraph

procfunc/transforms/__init__.py

@@ -0,0 +1,49 @@
+from .cleanup import (
+    coerce_shaders_to_materialresult,
+    eliminate_duplicate_result_types,
+    eliminate_duplicate_subgraphs,
+    extract_shader_vectors_as_inputs,
+    fill_graph_defaults_with_call_node,
+    remove_v1_name_from_graph,
+    replace_ids,
+)
+from .convert import (
+    colors_to_hsv_definition,
+)
+from .distribution import (
+    distribution_to_mode,
+    outlier_distribution,
+)
+from .extract_materials import (
+    extract_materials_from_graph,
+    extract_materials_from_graphs,
+)
+from .infer_distribution import (
+    infer_distribution_hypercube,
+    infer_nodegroup_distributions,
+)
+from .parameters import (
+    extract_parameter_distributions,
+)
+from .util import (
+    map_graph_list,
+    map_subgraphs,
+)
+
+__all__ = [
+    "coerce_shaders_to_materialresult",
+    "eliminate_duplicate_result_types",
+    "eliminate_duplicate_subgraphs",
+    "extract_shader_vectors_as_inputs",
+    "fill_graph_defaults_with_call_node",
+    "remove_v1_name_from_graph",
+    "replace_ids",
+    "colors_to_hsv_definition",
+    "distribution_to_mode",
+    "outlier_distribution",
+    "infer_distribution_hypercube",
+    "infer_nodegroup_distributions",
+    "extract_parameter_distributions",
+    "map_graph_list",
+    "map_subgraphs",
+]

procfunc/transforms/cleanup.py

@@ -0,0 +1,214 @@
+import logging
+from collections import defaultdict
+from typing import Any, Callable
+
+from procfunc import compute_graph as cg
+from procfunc import types as t
+from procfunc.nodes import types as nt
+from procfunc.nodes.shader import coord, geometry
+from procfunc.util import pytree
+
+logger = logging.getLogger(__name__)
+
+
+def remove_v1_name_from_graph(
+    _call_node: cg.Node, graph: cg.ComputeGraph
+) -> cg.ComputeGraph:
+    if graph.name.startswith("nodegroup_"):
+        graph.name = graph.name.replace("nodegroup_", "")
+    if graph.name.startswith("shader_"):
+        graph.name = graph.name.replace("shader_", "")
+    return graph
+
+
+def eliminate_duplicate_subgraphs(
+    graphs: list[cg.ComputeGraph],
+) -> list[cg.ComputeGraph]:
+    unique: list[cg.ComputeGraph] = []
+    removed: list[cg.ComputeGraph] = []
+    # maps id(duplicate_subgraph) -> canonical subgraph to replace it with
+    replacements: dict[int, cg.ComputeGraph] = {}
+
+    for topgraph in graphs:
+        subgraphs = reversed(
+            list(cg.traverse_nested_graphs(topgraph, yield_call_nodes=True))
+        )
+        for _call_node, subgraph in subgraphs:
+            match = next((g for g in unique if cg.graph_nodes_equal(subgraph, g)), None)
+            if match is not None:
+                if len(subgraph.name) < len(match.name):
+                    match.name = subgraph.name
+                replacements[id(subgraph)] = match
+                removed.append(subgraph)
+            else:
+                unique.append(subgraph)
+
+    # second pass: update ALL call nodes (in all nested subgraphs) that reference a replaced subgraph
+    for topgraph in graphs:
+        for subgraph in cg.traverse_nested_graphs(topgraph):
+            for node in cg.traverse_depth_first(subgraph):
+                if (
+                    isinstance(node, cg.SubgraphCallNode)
+                    and id(node.subgraph) in replacements
+                ):
+                    node.subgraph = replacements[id(node.subgraph)]
+
+    logger.debug(f"Eliminated duplicated subgraphs {[g.name for g in removed]}")
+
+    return graphs
+
+
+def eliminate_duplicate_result_types(
+    graphs: list[cg.ComputeGraph],
+    uses_threshold: int = 1,
+) -> list[cg.ComputeGraph]:
+    rettype_uses: dict[type, list[cg.ComputeGraph]] = defaultdict(list)
+
+    for graph in graphs:
+        for subgraph in cg.traverse_nested_graphs(graph):
+            result_type = subgraph.outputs.toplevel_type()
+            if result_type is None or not pytree.is_type_namedtuple(result_type):
+                continue
+
+            for rt in rettype_uses.keys():
+                if list(rt._fields) == list(result_type._fields):
+                    result_type = rt
+                    break
+
+            rettype_uses[result_type].append(subgraph)
+
+    for rettype, uses in rettype_uses.items():
+        if len(uses) <= uses_threshold:
+            continue
+        first_rettype = uses[0].outputs.toplevel_type()
+        for subgraph in uses[1:]:
+            subgraph.outputs.spec.container = first_rettype
+
+    return graphs
+
+
+def fill_graph_defaults_with_call_node(
+    call_node: cg.SubgraphCallNode,
+    graph: cg.ComputeGraph,
+) -> cg.ComputeGraph:
+    if call_node is None:
+        return graph
+
+    if any(
+        isinstance(arg.default_value, float) and arg.default_value != 0.0
+        for arg in graph.inputs.values()
+    ):
+        logger.debug(
+            f"Skipping {graph.name} because it has nondefault existing default args"
+        )
+        return graph
+
+    for name, inpnode in graph.inputs.items():
+        fillval = call_node.kwargs.get(name, None)
+        if fillval is not None and not isinstance(fillval, cg.Node):
+            inpnode.kwargs["default_value"] = fillval
+
+    return graph
+
+
+def coerce_shaders_to_materialresult(
+    _call_node: cg.Node, subgraph: cg.ComputeGraph
+) -> cg.ComputeGraph:
+    if subgraph.outputs.toplevel_type() is t.Material:
+        return subgraph
+    outputs = subgraph.outputs.dict()
+    surface = outputs.get("surface") or outputs.get("bsdf")
+    if surface is None:
+        return subgraph
+    shader_outputs = {
+        "surface": surface,
+        "displacement": outputs.get("displacement"),
+        "volume": outputs.get("volume"),
+    }
+    if len(outputs) > len(shader_outputs):
+        logger.warning(
+            f"{coerce_shaders_to_materialresult.__name__} skipping due to extra outputs: {outputs.keys()}"
+        )
+        return subgraph
+    logger.debug(
+        f"{coerce_shaders_to_materialresult.__name__} converted {subgraph.name} output"
+    )
+    subgraph.outputs = pytree.PyTree(t.Material(**shader_outputs))
+
+    return subgraph
+
+
+def replace_ids(
+    graph: cg.ComputeGraph,
+    ids: set[int],
+    val: Any,
+):
+    """
+    Pull out hardcoded arguments to be inputs to the graph instead
+
+    Args:
+        graph: The graph to extract constants from
+        extract_mask: A mask of which args to extract. The key is a tuple of the parent node id and the arg name.
+    """
+
+    assert isinstance(graph, cg.ComputeGraph)
+
+    for name, parent, child in cg.traverse_depth_first(
+        graph, yield_consts=True, yield_name=True, yield_parent=True
+    ):
+        if id(child) not in ids:
+            continue
+        if isinstance(name, int):
+            args = list(parent.args)
+            args[name] = val
+            parent.args = tuple(args)
+        else:
+            parent.kwargs[name] = val
+
+    return graph
+
+
+def extract_as_input(
+    graph: cg.ComputeGraph,
+    nodes: set[int],
+    name: str,
+    arg_type: type,
+):
+    inp = cg.InputPlaceholderNode(
+        default_value=None, metadata={"known_value_type": arg_type, "varname": name}
+    )
+
+    inputs = graph.inputs.obj()
+    assert isinstance(inputs, dict), inputs
+    inputs[name] = inp
+    graph.inputs = pytree.PyTree(inputs)
+
+    return replace_ids(graph, nodes, inp)
+
+
+def extract_shader_vectors_as_inputs(
+    graph: cg.ComputeGraph,
+    extract_funcs: list[Callable[..., Any]] | None = None,
+):
+    """
+    Pull out shader vectors as inputs to the graph instead
+    """
+
+    if extract_funcs is None:
+        extract_funcs = [coord, geometry]
+
+    def _is_vector_target(node: cg.FunctionCallNode) -> bool:
+        return isinstance(node, cg.FunctionCallNode) and node.func in extract_funcs
+
+    vector_nodes = set(
+        id(node)
+        for node in cg.traverse_depth_first(graph)
+        if _is_vector_target(node)
+        or (isinstance(node, cg.GetAttributeNode) and _is_vector_target(node.args[0]))
+    )
+
+    if len(vector_nodes) == 0:
+        return graph
+
+    extract_as_input(graph, vector_nodes, "vector", nt.ProcNode[t.Vector])
+    return graph

procfunc/transforms/convert.py

@@ -0,0 +1,20 @@
+import procfunc as pf
+from procfunc import compute_graph as cg
+
+
+def colors_to_hsv_definition(graph: cg.ComputeGraph) -> cg.ComputeGraph:
+    for node in cg.traverse_depth_first(graph):
+        for i, arg in enumerate(node.args):
+            if isinstance(arg, pf.Color):
+                hsv = tuple(round(x, 4) for x in arg.hsv)
+                node.args[i] = cg.FunctionCallNode(
+                    pf.color.hsv_to_rgba, args=(), kwargs={"hsv": hsv}
+                )
+        for key, arg in node.kwargs.items():
+            if isinstance(arg, pf.Color):
+                hsv = tuple(round(x, 4) for x in arg.hsv)
+                node.kwargs[key] = cg.FunctionCallNode(
+                    pf.color.hsv_to_rgba, args=(), kwargs={"hsv": hsv}
+                )
+
+    return graph

procfunc/transforms/distribution.py

@@ -0,0 +1,191 @@
+import enum
+import logging
+from typing import Any, Callable
+
+import numpy as np
+
+import procfunc as pf
+from procfunc import compute_graph as cg
+from procfunc.compute_graph import transform_compute_graph
+from procfunc.random import random_distrib_funcs
+
+logger = logging.getLogger(__name__)
+
+
+class NumpyRandomDistrib(enum.Enum):
+    # from https://numpy.org/doc/2.2/reference/random/generator.html#distributions
+    BETA = "beta"
+    BINOMIAL = "binomial"
+    CHISQUARE = "chisquare"
+    EXPONENTIAL = "exponential"
+    F = "f"
+    GAMMA = "gamma"
+    GAUSSIAN = "gaussian"
+    GEOMETRIC = "geometric"
+    GUMBEL = "gumbel"
+    HYPERGEOMETRIC = "hypergeometric"
+    LAPLACE = "laplace"
+    LOGISTIC = "logistic"
+    LOGNORMAL = "lognormal"
+    LOGSERIES = "logseries"
+    MULTINOMIAL = "multinomial"
+    MULTIVARIATE_NORMAL = "multivariate_normal"
+    NEGATIVE_BINOMIAL = "negative_binomial"
+    NORMAL = "normal"
+    NONCENTRAL_CHISQUARE = "noncentral_chisquare"
+    NONCENTRAL_F = "noncentral_f"
+    PARETO = "pareto"
+    POISSON = "poisson"
+    POWER = "power"
+    RAYLEIGH = "rayleigh"
+    SHUFFLE = "shuffle"
+    STANDARD_CAUCHY = "standard_cauchy"
+    STANDARD_EXPONENTIAL = "standard_exponential"
+    STANDARD_GAMMA = "standard_gamma"
+    STANDARD_NORMAL = "standard_normal"
+    STANDARD_LOGISTIC = "standard_logistic"
+    STANDARD_LAPLACE = "standard_laplace"
+    STANDARD_PARETO = "standard_pareto"
+    STANDARD_T = "standard_t"
+    TRIANGULAR = "triangular"
+    UNIFORM = "uniform"
+    VONMISES = "vonmises"
+    WALD = "wald"
+    WEIBULL = "weibull"
+    ZIPF = "zipf"
+    INTEGER = "integers"
+
+
+FUNCNAME_TO_DISTRIB = {v.value: v for v in NumpyRandomDistrib.__members__.values()}
+
+
+def as_distribution(
+    node: cg.Node,
+) -> NumpyRandomDistrib | Callable[[np.random.Generator, ...], Any] | None:
+    match node:
+        case cg.FunctionCallNode(func=x) if x in random_distrib_funcs:
+            return x
+        case cg.MethodCallNode(method_name=method, args=(arg_0,)) if (
+            method in FUNCNAME_TO_DISTRIB
+            and isinstance(arg_0, cg.Node)
+            and arg_0.metadata.get("known_value_type") is np.random.Generator
+        ):
+            return FUNCNAME_TO_DISTRIB[method]
+        case _:
+            return None
+
+
+def distribution_to_mode(
+    compute_graph: cg.ComputeGraph,
+    graph_name: str | None = None,
+) -> cg.ComputeGraph:
+    """
+    Transform a generator to use the mode of the distribution instead of random samples
+    """
+
+    def map_to_mode(node: cg.Node) -> cg.Node | float:
+        match as_distribution(node):
+            case NumpyRandomDistrib.UNIFORM:
+                assert len(node.args) == 1 or len(node.kwargs) == 0, (
+                    "implementation may be errored for mix of args and kwargs"
+                )
+                low = node.kwargs.get("low", node.args[1])
+                high = node.kwargs.get("high", node.args[2])
+                if isinstance(low, cg.Node) or isinstance(high, cg.Node):
+                    logger.warning(
+                        f"Uniform mode not implemented for {node=} with non-constant {low=} {high=}"
+                    )
+                    return node
+                return (low + high) / 2
+            case NumpyRandomDistrib.NORMAL:
+                assert len(node.args) == 1 or len(node.kwargs) == 0, (
+                    "implementation may be errored for mix of args and kwargs"
+                )
+                mean = node.kwargs.get("mean", node.args[1])
+                _std = node.kwargs.get("std", node.args[2])
+                if isinstance(mean, cg.Node):
+                    logger.warning(
+                        f"Normal mode not implemented for {node=} with non-constant {mean=}"
+                    )
+                    return node
+                return mean
+            case _:
+                return node
+
+    return transform_compute_graph(
+        compute_graph,
+        map_to_mode,
+        graph_name=graph_name or compute_graph.name + "_mode",
+    )
+
+
+def map_to_outlier(
+    node: cg.Node, pct: float = 0.05, normal_clip_std: float = 3.0
+) -> cg.Node:
+    varname = node.metadata.get("varname", None)
+    varname = (varname + "_outlier") if varname else None
+
+    match as_distribution(node):
+        case NumpyRandomDistrib.UNIFORM:
+            assert len(node.args) == 1 or len(node.kwargs) == 0, (
+                "cant handle arg/kwarg mix"
+            )
+            low = node.kwargs.get("low", node.args[1])
+            high = node.kwargs.get("high", node.args[2])
+            if isinstance(low, cg.Node) or isinstance(high, cg.Node):
+                logger.warning(
+                    f"outlier not implemented for {node=} with {min=} {max=}"
+                )
+                return node
+            rng = node.args[0]
+            assert isinstance(rng, cg.Node), f"got {node.args[0]=}"
+            return cg.FunctionCallNode(
+                func=pf.random.uniform_tails,
+                args=(rng,),
+                kwargs=dict(low=low, high=high, tail_pct=pct),
+                varname=varname,
+            )
+        case NumpyRandomDistrib.NORMAL:
+            assert len(node.args) == 1 or len(node.kwargs) == 0, (
+                "cant handle arg/kwarg mix"
+            )
+            mean = node.kwargs.get("mean", node.args[1])
+            std = node.kwargs.get("std", node.args[2])
+            if isinstance(mean, cg.Node):
+                logger.warning(f"outlier not implemented for {node=} with {mean=}")
+                return node
+            rng = node.args[0]
+            return cg.FunctionCallNode(
+                func=pf.random.uniform_tails,
+                args=(rng,),
+                kwargs=dict(
+                    tail_pct=pct,
+                    low=mean - normal_clip_std * std,
+                    high=mean + normal_clip_std * std,
+                ),
+                varname=varname,
+            )
+        case func if func in pf.random.random_distrib_funcs:
+            logger.warning(
+                f"{outlier_distribution.__name__} not implemented for {func}"
+            )
+            return node
+        case _:
+            return node
+
+
+def outlier_distribution(
+    compute_graph: cg.ComputeGraph,
+    pct: float = 0.05,
+    graph_name: str | None = None,
+    normal_clip_std: float = 3.0,
+) -> cg.ComputeGraph:
+    """
+    Transform a generator to generate outliers with a given probability
+    """
+
+    return transform_compute_graph(
+        compute_graph,
+        lambda node: map_to_outlier(node, pct, normal_clip_std),
+        graph_name=graph_name or compute_graph.name + "_outlier",
+    )