execute.py

@@ -72,6 +72,22 @@ def _on_depsgraph_update_post(
         _skip_next_autorun = False
         return
 
+    run_node_tree(scene, depsgraph)
+
+
+@bpy.app.handlers.persistent  # type: ignore
+def _on_frame_changed_post(
+    scene: bpy.types.Scene, depsgraph: bpy.types.Depsgraph
+) -> None:
+    global _skip_next_autorun
+    _skip_next_autorun = True
+
+    run_node_tree(scene, depsgraph)
+
+
+def run_node_tree(
+    scene: bpy.types.Scene, depsgraph: bpy.types.Depsgraph
+) -> None:
     for tree in bpy.data.node_groups:
         if tree.bl_idname != "PowerShipNodeTree":
             continue
@@ -114,10 +130,12 @@ def register() -> None:
     )
 
     bpy.app.handlers.depsgraph_update_post.append(_on_depsgraph_update_post)
+    bpy.app.handlers.frame_change_post.append(_on_frame_changed_post)
 
 
 def unregister() -> None:
     bpy.app.handlers.depsgraph_update_post.remove(_on_depsgraph_update_post)
+    bpy.app.handlers.frame_change_post.remove(_on_frame_changed_post)
 
     del bpy.types.Scene.powership_mode
 

nodes.py

@@ -4,7 +4,7 @@ import functools
 from collections import deque
 from copy import copy
 from math import degrees, radians
-from typing import TypeVar, Callable, Optional, Iterable
+from typing import TypeVar, Callable, Optional, Iterable, Union, Any, Tuple
 
 import bpy
 import nodeitems_utils
@@ -28,7 +28,7 @@ class PowerShipNodeTree(bpy.types.NodeTree):
 
         :return: whether the event ran to completion.
         """
-        event_to_node_types: dict[str, type[AbstractPowerShipNode]] = {
+        event_to_node_types: dict[str, type[AbstractPowerShipEventNode]] = {
             "FORWARD": ForwardSolveNode,
             "BACKWARD": BackwardSolveNode,
         }
@@ -67,7 +67,10 @@ class PowerShipNodeTree(bpy.types.NodeTree):
 
     def _prepare_nodes(self) -> None:
         for node in self.nodes:
-            node.reset_run()
+            if isinstance(node, AbstractPowerShipNode):
+                node.reset_run()
+            else:
+                raise TypeError(f"Node type unknown {type(node)}")
 
     def _run_from_node(
         self, depsgraph: bpy.types.Depsgraph, start_node: "AbstractPowerShipNode"
@@ -114,10 +117,16 @@ class NodeSocketExecute(bpy.types.NodeSocket):
     bl_label = "Execute"
     link_limit = 1
 
-    def draw(self, context, layout, node, text):
+    def draw(
+        self,
+        context: bpy.types.Context,
+        layout: bpy.types.UILayout,
+        node: bpy.types.Node,
+        text: Union[str, Any],
+    ) -> None:
         layout.label(text=text)
 
-    def draw_color(self, context, node):
+    def draw_color(self, context: bpy.types.Context, node: bpy.types.Node) -> Tuple[float, float, float, float]:
         return (1.0, 1.0, 1.0, 1.0)
 
 
@@ -385,7 +394,7 @@ class ForwardSolveNode(AbstractPowerShipEventNode):
     bl_label = "Forward Solve"
     bl_icon = "PLAY"
 
-    def init(self, context):
+    def init(self, context: bpy.types.Context) -> None:
         self.add_execution_socket_output()
 
 
@@ -396,7 +405,7 @@ class BackwardSolveNode(AbstractPowerShipEventNode):
     bl_label = "Backward Solve"
     bl_icon = "PLAY"
 
-    def init(self, context):
+    def init(self, context: bpy.types.Context) -> None:
         self.add_execution_socket_output()
 
 
@@ -421,7 +430,7 @@ class GetBoneNode(AbstractPowerShipNode):
         super().draw_buttons(context, layout)
         layout.prop(self, "head_tail")
 
-    def init(self, context):
+    def init(self, context: bpy.types.Context) -> None:
         self.inputs.new("NodeSocketObject", "Armature")
         self.inputs.new("NodeSocketString", "Bone")
 
@@ -541,7 +550,7 @@ class SetControlNode(AbstractPowerShipNode):
         super().draw_buttons(context, layout)
         layout.prop(self, "space")
 
-    def init(self, context):
+    def init(self, context: bpy.types.Context) -> None:
         self.add_execution_sockets()
 
         self.add_optional_input_socket("NodeSocketVector", "Location")
@@ -589,6 +598,299 @@ class SetControlNode(AbstractPowerShipNode):
                     control_obj.scale = control_scale
 
 
+class ToVector(AbstractPowerShipNode):
+    bl_idname = "ToVector"
+    bl_label = "To Vector"
+    bl_icon = "EMPTY_ARROWS"
+
+    def init(self, context: bpy.types.Context) -> None:
+        self.inputs.new("NodeSocketFloat", "X")
+        self.inputs.new("NodeSocketFloat", "Y")
+        self.inputs.new("NodeSocketFloat", "Z")
+        self.outputs.new("NodeSocketVector", "Vector")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        x = self.inputs["X"].default_value
+        y = self.inputs["Y"].default_value
+        z = self.inputs["Z"].default_value
+        self.outputs["Vector"].default_value = Vector((x, y, z))
+
+
+class SplitVector(AbstractPowerShipNode):
+    bl_idname = "SplitVector"
+    bl_label = "Split Vector"
+    bl_icon = "EMPTY_ARROWS"
+
+    def init(self, context: bpy.types.Context) -> None:
+        self.inputs.new("NodeSocketVector", "Vector")
+        self.outputs.new("NodeSocketFloat", "X")
+        self.outputs.new("NodeSocketFloat", "Y")
+        self.outputs.new("NodeSocketFloat", "Z")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        v = self.inputs["Vector"].default_value
+        self.outputs["X"].default_value = v.x
+        self.outputs["Y"].default_value = v.y
+        self.outputs["Z"].default_value = v.z
+
+
+class Distance(AbstractPowerShipNode):
+    """Calculates distance between two points."""
+
+    bl_idname = "Distance"
+    bl_label = "Distance"
+    bl_icon = "EMPTY_ARROWS"
+
+    def init(self, context: bpy.types.Context) -> None:
+        self.inputs.new("NodeSocketVector", "U")
+        self.inputs.new("NodeSocketVector", "V")
+        self.outputs.new("NodeSocketFloat", "Float")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        u = self.inputs["U"].default_value
+        v = self.inputs["V"].default_value
+        self.outputs["Float"].default_value = (u - v).length
+
+
+class NormalFromPoints(AbstractPowerShipNode):
+    """Calculates normal from three points (plane)."""
+
+    bl_idname = "NormalFromPoints"
+    bl_label = "Normal from Points"
+    bl_icon = "EMPTY_ARROWS"
+
+    def init(self, context: bpy.types.Context) -> None:
+        self.inputs.new("NodeSocketVector", "U")
+        self.inputs.new("NodeSocketVector", "V")
+        self.inputs.new("NodeSocketVector", "W")
+        self.outputs.new("NodeSocketVector", "Result")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        u = self.inputs["U"].default_value
+        v = self.inputs["V"].default_value
+        w = self.inputs["W"].default_value
+        a = v - u
+        b = w - u
+        normal = a.cross(b).normalized()
+        self.outputs["Result"].default_value = normal
+
+
+_enum_up_axis_items = (
+    ("X", "X", ""),
+    ("Y", "Y", ""),
+    ("Z", "Z", ""),
+)
+
+_enum_track_axis_items = (
+    ("X", "X", ""),
+    ("Y", "Y", ""),
+    ("Z", "Z", ""),
+    ("-X", "-X", ""),
+    ("-Y", "-Y", ""),
+    ("-Z", "-Z", ""),
+)
+
+
+class RotateTowards(AbstractPowerShipNode):
+    """Calculate the rotation from a vector with a track and up axis."""
+
+    bl_idname = "RotateTowards"
+    bl_label = "Rotate Towards"
+    bl_icon = "EMPTY_ARROWS"
+
+    track: bpy.props.EnumProperty(  # type: ignore
+        name="Track", items=_enum_track_axis_items, default="Z"
+    )
+
+    up: bpy.props.EnumProperty(  # type: ignore
+        name="Up", items=_enum_up_axis_items, default="Y"
+    )
+
+    def draw_buttons(
+        self, context: bpy.types.Context, layout: bpy.types.UILayout
+    ) -> None:
+        super().draw_buttons(context, layout)
+        layout.prop(self, "up")
+        layout.prop(self, "track")
+
+    def init(self, context: bpy.types.Context) -> None:
+        self.inputs.new("NodeSocketVector", "Vector")
+        self.inputs.new("NodeSocketVector", "RotateTo")
+        self.outputs.new("NodeSocketQuaternion", "Rotation")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        origin = self.inputs["Vector"].default_value
+        destination = self.inputs["RotateTo"].default_value
+        vec = Vector((destination - origin))
+        rot = vec.to_track_quat(self.track, self.up)
+        self.outputs["Rotation"].default_value = rot.normalized()
+
+
+class OffsetRotation(AbstractPowerShipNode):
+    """Offset a rotation."""
+
+    bl_idname = "OffsetRotation"
+    bl_label = "Offset Rotation"
+    bl_icon = "EMPTY_ARROWS"
+
+    def init(self, context: bpy.types.Context) -> None:
+        self.inputs.new("NodeSocketQuaternion", "Base")
+        self.inputs.new("NodeSocketQuaternion", "Offset")
+        self.outputs.new("NodeSocketQuaternion", "Rotation")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        base = self.inputs["Base"].default_value
+        offset = self.inputs["Offset"].default_value
+        self.outputs["Rotation"].default_value = base @ offset
+
+
+class MapRange(AbstractPowerShipNode):
+    bl_idname = "MapRange"
+    bl_label = "Map Range"
+    bl_icon = "EMPTY_ARROWS"
+
+    def init(self, context: bpy.types.Context) -> None:
+        self.inputs.new("NodeSocketFloat", "Value")
+        self.inputs.new("NodeSocketFloat", "From Min")
+        self.inputs.new("NodeSocketFloat", "From Max")
+        self.inputs.new("NodeSocketFloat", "To Min")
+        self.inputs.new("NodeSocketFloat", "To Max")
+        self.outputs.new("NodeSocketFloat", "Result")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        val = self.inputs["Value"].default_value
+        fmin = self.inputs["From Min"].default_value
+        fmax = self.inputs["From Max"].default_value
+        tmin = self.inputs["To Min"].default_value
+        tmax = self.inputs["To Max"].default_value
+
+        factor = (tmax - tmin) / (fmax - fmin)
+        offset = tmin - factor * fmin
+        self.outputs["Result"].default_value = factor * val + offset
+
+
+class AngleFromVectors(AbstractPowerShipNode):
+    """Calculate the angle between two vectors. Output in degrees."""
+
+    bl_idname = "AngleFromVectors"
+    bl_label = "Angle From Vectors"
+    bl_icon = "EMPTY_ARROWS"
+
+    def init(self, context: bpy.types.Context) -> None:
+        self.inputs.new("NodeSocketVector", "U")
+        self.inputs.new("NodeSocketVector", "V")
+        self.outputs.new("NodeSocketFloat", "Angle")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        u = self.inputs["U"].default_value
+        v = self.inputs["V"].default_value
+
+        angle = 0
+        if not (u.length_squared == 0 or v.length_squared == 0):
+            angle = u.angle(v)
+
+        self.outputs["Angle"].default_value = degrees(angle)
+
+
+_enum_vector_math_operations = [
+    ("ADD", "Add", ""),
+    ("SUBSTRACT", "Substract", ""),
+    ("MULTIPLY", "Mutliply", ""),
+    ("DIVIDE", "Divide", ""),
+    ("CROSS", "Cross", ""),
+]
+
+
+class VectorMath(AbstractPowerShipNode):
+    bl_idname = "VectorMath"
+    bl_label = "Vector Math"
+    bl_icon = "EMPTY_ARROWS"
+
+    operation: bpy.props.EnumProperty(  # type: ignore
+        name="Operation",
+        items=_enum_vector_math_operations,
+    )
+
+    def draw_buttons(
+        self, context: bpy.types.Context, layout: bpy.types.UILayout
+    ) -> None:
+        super().draw_buttons(context, layout)
+        layout.prop(self, "operation")
+
+    def init(self, context: bpy.types.Context) -> None:
+        self.inputs.new("NodeSocketVector", "U")
+        self.inputs.new("NodeSocketVector", "V")
+        self.outputs.new("NodeSocketVector", "Result")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        u = self.inputs["U"].default_value
+        v = self.inputs["V"].default_value
+
+        match self.operation:
+            case "ADD":
+                res = u + v
+            case "MULTIPLY":
+                res = u * v
+            case "SUBSTRACT":
+                res = u - v
+            case "DIVIDE":
+                res = Vector((x / y if y != 0.0 else 0.0 for x, y in zip(u, v)))
+            case "CROSS":
+                res = u.cross(v)
+            case _:
+                raise ValueError(f"Vector math operation not found: {self.operation}\n")
+
+        self.outputs["Result"].default_value = res
+
+
+_enum_math_operations = [
+    ("ADD", "Add", ""),
+    ("SUBSTRACT", "Substract", ""),
+    ("MULTIPLY", "Mutliply", ""),
+    ("DIVIDE", "Divide", ""),
+]
+
+
+class Math(AbstractPowerShipNode):
+    bl_idname = "Math"
+    bl_label = "Math"
+    bl_icon = "EMPTY_ARROWS"
+
+    operation: bpy.props.EnumProperty(  # type: ignore
+        name="Operation",
+        items=_enum_math_operations,
+    )
+
+    def draw_buttons(
+        self, context: bpy.types.Context, layout: bpy.types.UILayout
+    ) -> None:
+        super().draw_buttons(context, layout)
+        layout.prop(self, "operation")
+
+    def init(self, context: bpy.types.Context) -> None:
+        self.inputs.new("NodeSocketFloat", "U")
+        self.inputs.new("NodeSocketFloat", "V")
+        self.outputs.new("NodeSocketFloat", "Result")
+
+    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
+        u = self.inputs["U"].default_value
+        v = self.inputs["V"].default_value
+
+        match self.operation:
+            case "ADD":
+                res = u + v
+            case "MULTIPLY":
+                res = u * v
+            case "SUBSTRACT":
+                res = u - v
+            case "DIVIDE":
+                res = u / v if v != 0 else 0
+            case _:
+                raise ValueError(f"Math operation not found: {self.operation}\n")
+
+        self.outputs["Result"].default_value = res
+
+
 class SetCursorNode(AbstractAlwaysExecuteNode):
     """Sets the location and/or rotation of the 3D cursor"""
 
@@ -596,7 +898,7 @@ class SetCursorNode(AbstractAlwaysExecuteNode):
     bl_label = "Set Cursor"
     bl_icon = "CURSOR"
 
-    def init(self, context):
+    def init(self, context: bpy.types.Context) -> None:
         self.add_execution_sockets()
         self.add_optional_input_socket("NodeSocketVector", "Location")
         self.add_optional_input_socket("NodeSocketQuaternion", "Rotation")
@@ -620,7 +922,7 @@ class SetBoneNode(AbstractPowerShipNode):
     bl_label = "Set Bone"
     bl_icon = "BONE_DATA"
 
-    def init(self, context):
+    def init(self, context: bpy.types.Context) -> None:
         self.add_execution_sockets()
 
         self.add_optional_input_socket("NodeSocketVector", "Location")
@@ -660,14 +962,24 @@ class SetBoneNode(AbstractPowerShipNode):
         if control_scale is not None:
             scale = control_scale
 
+        # TODO: Fix jittering bone scale which happens
+        # esp. when multiple bones are parented to the rotated bone
+        # rounding helps but does not entirely fix the issue.
+        scale = [round(x, 4) for x in scale]
         v_nil = Vector((0, 0, 0))
         bone_rest_rot_scale = bone.bone.matrix_local.copy()
-        bone_rest_rot_scale.translation = v_nil
 
-        mat_rot_scale = Matrix.LocRotScale(v_nil, rot, scale) @ bone_rest_rot_scale
-
-        mat_loc = Matrix.Translation(loc)
-        bone_mat_world = mat_loc @ mat_rot_scale
+        match self.space:
+            case "WORLD":
+                bone_mat_world = Matrix.LocRotScale(loc, rot, scale)
+                loc, rot, scale = bone_mat_world.decompose()
+            case "CHANNELS":
+                bone_rest_rot_scale.translation = v_nil
+                mat_rot_scale = (
+                    Matrix.LocRotScale(v_nil, rot, scale) @ bone_rest_rot_scale
+                )
+                mat_loc = Matrix.Translation(loc)
+                bone_mat_world = mat_loc @ mat_rot_scale
 
         bone.matrix = arm_matrix.inverted() @ bone_mat_world
 
@@ -680,7 +992,7 @@ class TwoBoneIKNode(AbstractPowerShipNode):
     # Set to True to remove the cleanup and keep the IK constraint and temporary empties.
     _debug = False
 
-    def init(self, context):
+    def init(self, context: bpy.types.Context) -> None:
         self.add_execution_sockets()
 
         self.inputs.new("NodeSocketObject", "Armature")
@@ -875,63 +1187,6 @@ class ClampNode(AbstractPowerShipNode):
         self.outputs["Result"].default_value = clamped
 
 
-class AbstractTwoValueMathNode(AbstractPowerShipNode):
-    def init(self, context: bpy.types.Context) -> None:
-        self.inputs.new("NodeSocketFloat", "A")
-        self.inputs.new("NodeSocketFloat", "B")
-        self.outputs.new("NodeSocketFloat", "Result")
-
-    def execute(self, depsgraph: bpy.types.Depsgraph) -> None:
-        a = self._get_input_value("A", float) or 0.0
-        b = self._get_input_value("B", float) or 0.0
-        self.outputs["Result"].default_value = self.calculate(a, b)
-
-    def calculate(self, a: float, b: float) -> float:
-        return 0
-
-
-class AddNode(AbstractTwoValueMathNode):
-    """Add two values"""
-
-    bl_idname = "AddNode"
-    bl_label = "Add"
-
-    def calculate(self, a: float, b: float) -> float:
-        return a + b
-
-
-class SubtractNode(AbstractTwoValueMathNode):
-    """Subtract two values"""
-
-    bl_idname = "SubtractNode"
-    bl_label = "Subtract"
-
-    def calculate(self, a: float, b: float) -> float:
-        return a - b
-
-
-class MultiplyNode(AbstractTwoValueMathNode):
-    """Multiply two values"""
-
-    bl_idname = "MultiplyNode"
-    bl_label = "Multiply"
-
-    def calculate(self, a: float, b: float) -> float:
-        return a * b
-
-
-class DivideNode(AbstractTwoValueMathNode):
-    """Divide two values; division by zero results in NaN"""
-
-    bl_idname = "DivideNode"
-    bl_label = "Divide"
-
-    def calculate(self, a: float, b: float) -> float:
-        if b == 0:
-            return float("nan")
-        return a / b
-
-
 def _on_num_sockets_change(self: "SequenceNode", context: bpy.types.Context) -> None:
     self.recreate(context)
 
@@ -955,7 +1210,7 @@ class SequenceNode(AbstractPowerShipNode):
         super().draw_buttons(context, layout)
         layout.prop(self, "num_socks")
 
-    def init(self, context):
+    def init(self, context: bpy.types.Context) -> None:
         self.add_execution_socket_input()
         for index in range(self.num_socks):
             self.outputs.new(
@@ -966,8 +1221,8 @@ class SequenceNode(AbstractPowerShipNode):
 
 class PowerShipNodeCategory(nodeitems_utils.NodeCategory):
     @classmethod
-    def poll(cls, context):
-        return context.space_data.tree_type == "PowerShipNodeTree"
+    def poll(cls, context: bpy.types.Context) -> bool:
+        return context.space_data.tree_type == "PowerShipNodeTree"  # type: ignore
 
 
 class PowerShip_OT_rebuild_node(bpy.types.Operator):
@@ -1030,10 +1285,16 @@ node_categories = [
         "MATH",
         "Math",
         items=[
-            nodeitems_utils.NodeItem("AddNode"),
-            nodeitems_utils.NodeItem("SubtractNode"),
-            nodeitems_utils.NodeItem("MultiplyNode"),
-            nodeitems_utils.NodeItem("DivideNode"),
+            nodeitems_utils.NodeItem("RotateTowards"),
+            nodeitems_utils.NodeItem("OffsetRotation"),
+            nodeitems_utils.NodeItem("AngleFromVectors"),
+            nodeitems_utils.NodeItem("NormalFromPoints"),
+            nodeitems_utils.NodeItem("SplitVector"),
+            nodeitems_utils.NodeItem("ToVector"),
+            nodeitems_utils.NodeItem("Distance"),
+            nodeitems_utils.NodeItem("MapRange"),
+            nodeitems_utils.NodeItem("VectorMath"),
+            nodeitems_utils.NodeItem("Math"),
             nodeitems_utils.NodeItem("ClampNode"),
             nodeitems_utils.NodeItem("ToEulerNode"),
             nodeitems_utils.NodeItem("FromEulerNode"),
@@ -1065,13 +1326,20 @@ classes = (
     TwoBoneIKNode,
     SetCursorNode,
     SequenceNode,
+    # Math Nodes
+    RotateTowards,
+    AngleFromVectors,
+    OffsetRotation,
+    NormalFromPoints,
+    ToVector,
+    SplitVector,
+    Distance,
+    MapRange,
+    VectorMath,
+    Math,
     ToEulerNode,
     FromEulerNode,
     ClampNode,
-    AddNode,
-    SubtractNode,
-    MultiplyNode,
-    DivideNode,
     # Operators:
     PowerShip_OT_rebuild_node,
 )