From 63a75169059f683a31d340fd96a80c17f9645a5e Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 9 Oct 2023 19:05:08 -0400 Subject: [PATCH] Refactor simulation API to use `simulation_zone` decorator --- api/node_mapper.py | 29 +++++----- api/static/simulation.py | 53 +++++++++++++------ api/types.py | 11 ++++ book/src/api/advanced-scripting/simulation.md | 30 +++++------ 4 files changed, 76 insertions(+), 47 deletions(-) diff --git a/api/node_mapper.py b/api/node_mapper.py index 54321bc..0fd4dbd 100644 --- a/api/node_mapper.py +++ b/api/node_mapper.py @@ -14,6 +14,21 @@ class OutputsList(dict): __setattr__ = dict.__setitem__ __delattr__ = dict.__delitem__ +def set_or_create_link(x, node_input): + if issubclass(type(x), Type): + State.current_node_tree.links.new(x._socket, node_input) + else: + def link_constant(): + constant = Type(value=x) + State.current_node_tree.links.new(constant._socket, node_input) + if node_input.hide_value: + link_constant() + else: + try: + node_input.default_value = x + except: + link_constant() + def build_node(node_type): def build(_primary_arg=None, **kwargs): for k, v in kwargs.copy().items(): @@ -48,20 +63,6 @@ def build(_primary_arg=None, **kwargs): if node_input2.name.lower().replace(' ', '_') == argname and node_input2.type == node_input.type: all_with_name.append(node_input2) if argname in kwargs: - def set_or_create_link(x, node_input): - if issubclass(type(x), Type): - State.current_node_tree.links.new(x._socket, node_input) - else: - def link_constant(): - constant = Type(value=x) - State.current_node_tree.links.new(constant._socket, node_input) - if node_input.hide_value: - link_constant() - else: - try: - node_input.default_value = x - except: - link_constant() value = kwargs[argname] if isinstance(value, enum.Enum): value = value.value diff --git a/api/static/simulation.py b/api/static/simulation.py index 9805449..ccb7d8d 100644 --- a/api/static/simulation.py +++ b/api/static/simulation.py @@ -1,25 +1,46 @@ +import bpy import inspect import typing -class SimulationInput: - class DeltaTime: pass - class ElapsedTime: pass - -def simulation(block: typing.Callable[typing.Any, 'Geometry']): +def simulation_zone(block: typing.Callable): """ Create a simulation input/output block. - > Only available in the `geometry-node-simulation` branch of Blender 3.5. + > Only available in Blender 3.6+. """ - def wrapped(geometry: 'Geometry', *args, **kwargs): - from geometry_script import simulation_input, simulation_output - simulation_in = simulation_input(geometry=geometry) + def wrapped(*args, **kwargs): + from geometry_script.api.node_mapper import OutputsList, set_or_create_link + from geometry_script.api.state import State + from geometry_script.api.types import Type, socket_class_to_data_type + signature = inspect.signature(block) - for key, value in signature.parameters.items(): - match value.annotation: - case SimulationInput.DeltaTime: - kwargs[key] = simulation_in.delta_time - case SimulationInput.ElapsedTime: - kwargs[key] = simulation_in.elapsed_time - return simulation_output(geometry=block(simulation_in.geometry, *args, **kwargs)).geometry + + # setup zone + simulation_in = State.current_node_tree.nodes.new(bpy.types.GeometryNodeSimulationInput.__name__) + simulation_out = State.current_node_tree.nodes.new(bpy.types.GeometryNodeSimulationOutput.__name__) + simulation_in.pair_with_output(simulation_out) + + # clear state items + for item in simulation_out.state_items: + simulation_out.state_items.remove(item) + + # create state items from block signature + state_items = {} + for param in [*signature.parameters.values()][1:]: + state_items[param.name] = (param.annotation, param.default, None, None) + for i, arg in enumerate(state_items.items()): + simulation_out.state_items.new(socket_class_to_data_type(arg[1][0].socket_type), arg[0].replace('_', ' ').title()) + set_or_create_link(kwargs[arg[0]] if arg[0] in kwargs else args[i], simulation_in.inputs[i]) + + step = block(*[Type(o) for o in simulation_in.outputs[:-1]]) + + if isinstance(step, Type): + step = (step,) + for i, result in enumerate(step): + State.current_node_tree.links.new(result._socket, simulation_out.inputs[i]) + + if len(simulation_out.outputs[:-1]) == 1: + return Type(simulation_out.outputs[0]) + else: + return OutputsList({o.name.lower().replace(' ', '_'): Type(o) for o in simulation_out.outputs[:-1]}) return wrapped \ No newline at end of file diff --git a/api/types.py b/api/types.py index 620e17a..2b89168 100644 --- a/api/types.py +++ b/api/types.py @@ -20,6 +20,15 @@ def socket_type_to_data_type(socket_type): case _: return socket_type +def socket_class_to_data_type(socket_class_name): + match socket_class_name: + case 'NodeSocketGeometry': + return 'GEOMETRY' + case 'NodeSocketFloat': + return 'FLOAT' + case _: + return socket_class_name + # The base class all exposed socket types conform to. class _TypeMeta(type): def __getitem__(self, args): @@ -217,6 +226,8 @@ def transfer(self, attribute, **kwargs): return self.transfer_attribute(data_type=data_type, attribute=attribute, **kwargs) def __getitem__(self, subscript): + if self._socket.type == 'VECTOR' and isinstance(subscript, int): + return self._get_xyz_component(subscript) if isinstance(subscript, tuple): accessor = subscript[0] args = subscript[1:] diff --git a/book/src/api/advanced-scripting/simulation.md b/book/src/api/advanced-scripting/simulation.md index d471439..3c31614 100644 --- a/book/src/api/advanced-scripting/simulation.md +++ b/book/src/api/advanced-scripting/simulation.md @@ -1,26 +1,22 @@ # Simulation -> This API is subject to change as future builds of Blender with simulation nodes are released. - -The `geometry-nodes-simulation` branch of Blender 3.5 includes support for "simulation nodes". +Blender 3.6 includes simulation nodes. Using a *Simulation Input* and *Simulation Output* node, you can create effects that change over time. -As a convenience, the `@simulation` decorator is provided to make simulation node blocks easier to create. +As a convenience, the `@simulation_zone` decorator is provided to make simulation node blocks easier to create. ```python -@simulation -def move_over_time( - geometry: Geometry, # the first input must be `Geometry` - speed: Float, - dt: SimulationInput.DeltaTime, # Automatically passes the delta time on any argument annotated with `SimulationInput.DeltaTime`. - elapsed: SimulationInput.ElapsedTime, # Automatically passes the elapsed time -) -> Geometry: - return geometry.set_position( - offset=combine_xyz(x=speed) - ) -``` +from geometry_script import * -Every frame the argument `geometry` will be set to the geometry from the previous frame. This allows the offset to accumulate over time. +@tree +def test_sim(geometry: Geometry): + @simulation_zone + def my_sim(delta_time, geometry: Geometry, value: Float): + return (geometry, value) + return my_sim(geometry, 0.26).value +``` -The `SimulationInput.DeltaTime`/`SimulationInput.ElapsedTime` types mark arguments that should be given the outputs from the *Simulation Input* node. \ No newline at end of file +The first argument should always be `delta_time`. Any other arguments must also be returned as a tuple with their modified values. +Each frame, the result from the previous frame is passed into the zone's inputs. +The initial call to `my_sim` in `test_sim` provides the initial values for the simulation. \ No newline at end of file