Refactor simulation API to use simulation_zone decorator

Author: carson-katri

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

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

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:]

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.
+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.