mantis/math_nodes.py

from .utilities.node_common import *
from .base_definitions import MantisNode
from .math_nodes_socket_templates import *

def TellClasses():
    return [
            MathStaticInt,
            MathStaticFloat,
            MathStaticVector,
           ]

def math_operation(operation, a, b):
    match operation:
        case "ADD":
            return a+b
        case "SUBTRACT":
            return a-b
        case "MULTIPLY":
            return a*b
        case "DIVIDE":
            return a/b
        case "FLOOR_DIVIDE":
            return a//b
        case "MODULUS":
            return a%b
        case "POWER":
            return a**b
        case "ABSOLUTE":
            return abs(a)
        case "MAXIMUM":
            return max(a, b)
        case "MINIMUM":
            return min(a, b)
        case "GREATER THAN":
            return float(a > b)
        case "LESS THAN":
            return float(a < b)
        case "ARCTAN2":
            from math import atan2
            return atan2(a,b)
        case "FLOOR":
            from math import floor
            return floor(a)
        case "CEIL":
            from math import ceil
            return ceil(a)
        case "ROUND":
            return round(a)

#*#-------------------------------#++#-------------------------------#*#
# M A T H  N O D E S
#*#-------------------------------#++#-------------------------------#*#
class MathStaticInt(MantisNode):
    '''A node representing an integer math operation'''
    __slots__ = ()
    node_type = "UTILITY"
    def __init__(self, signature, base_tree):
        super().__init__(signature, base_tree, StaticIntSockets)
        self.init_parameters()

    def bPrepare(self, bContext = None,):
        a = self.evaluate_input("Int A"); b = self.evaluate_input("Int B")
        result = math_operation(self.evaluate_input("Operation"), a, b)
        self.parameters["Result Int"] = int(result)
        self.prepared, self.executed = True, True

class MathStaticFloat(MantisNode):
    '''A node representing a float math operation'''
    __slots__ = ()
    node_type = "UTILITY"
    def __init__(self, signature, base_tree):
        super().__init__(signature, base_tree, StaticFloatSockets)
        self.init_parameters()

    def bPrepare(self, bContext = None,):
        a = self.evaluate_input("Float A"); b = self.evaluate_input("Float B")
        result = math_operation(self.evaluate_input("Operation"), a, b)
        self.parameters["Result Float"] = result
        self.prepared, self.executed = True, True

class MathStaticVector(MantisNode):
    '''A node representing a vector math operation'''
    __slots__ = ()
    node_type = "UTILITY"
    def __init__(self, signature, base_tree):
        super().__init__(signature, base_tree, StaticVectorSockets)
        self.init_parameters()

    def bPrepare(self, bContext = None,):
        from mathutils import Vector
        a = Vector(self.evaluate_input("Vector A")).copy()
        b = Vector(self.evaluate_input("Vector B")).copy()
        s = self.evaluate_input("Scalar A")
        if hasattr(s, '__iter__'):
            average = lambda iterable : sum(iterable)/len(iterable)
            s = average(s)
        f_result = float("NaN")
        v_result = None
        if self.evaluate_input("Operation") == "ADD":
            v_result = a+b
        if self.evaluate_input("Operation") == "SUBTRACT":
            v_result = a-b
        if self.evaluate_input("Operation") == "MULTIPLY":
            v_result = a*b
        if self.evaluate_input("Operation") == "DIVIDE":
            v_result = Vector((0,0,0)) # channel-wise division
            for i, (axis_a, axis_b) in enumerate(zip(a,b)):
                v_result[i] = axis_a/axis_b
        if self.evaluate_input("Operation") == "POWER":
            v_result = a**b
        # since these are unary, we need to make a copy lest we create spooky effects elsewhere.
        a = a.copy()
        if self.evaluate_input("Operation") == "SCALE":
            v_result = a.normalized() * s
        if self.evaluate_input("Operation") == "LENGTH":
            f_result =  a.magnitude
        if self.evaluate_input("Operation") == "CROSS":
            v_result =  a.cross(b)
        if self.evaluate_input("Operation") == "DOT":
            f_result =  a.dot(b)
        if self.evaluate_input("Operation") == "NORMALIZE":
            v_result =  a.normalized()
        if self.evaluate_input("Operation") == "LINEAR_INTERP":
            v_result =  a.lerp(b, s).copy()
        self.parameters["Result Float"] = f_result
        # if v_result:
        self.parameters["Result Vector"] = v_result
        self.prepared = True
        self.executed = True