From 657d79c8c76a39a053791c8010daa4678576110c Mon Sep 17 00:00:00 2001 From: Evan Spearman Date: Sat, 22 Jul 2023 20:24:29 -0500 Subject: [PATCH] chore: Operation types now have nodes instead of individual operations --- __init__.py | 17 +- src/comfymath/control.py | 105 +------- src/comfymath/convert.py | 85 +++++-- src/comfymath/float.py | 303 +++++++++++------------ src/comfymath/int.py | 240 +++++++++---------- src/comfymath/number.py | 240 +++++-------------- src/comfymath/vec.py | 501 +++++++++++++++++++++++++++++++++++++++ src/comfymath/vec2.py | 259 -------------------- src/comfymath/vec3.py | 260 -------------------- src/comfymath/vec4.py | 261 -------------------- 10 files changed, 892 insertions(+), 1379 deletions(-) create mode 100644 src/comfymath/vec.py delete mode 100644 src/comfymath/vec2.py delete mode 100644 src/comfymath/vec3.py delete mode 100644 src/comfymath/vec4.py diff --git a/__init__.py b/__init__.py index 264be62..0ef2a2b 100644 --- a/__init__.py +++ b/__init__.py @@ -2,9 +2,7 @@ from .src.comfymath.convert import NODE_CLASS_MAPPINGS as convert_NCM from .src.comfymath.int import NODE_CLASS_MAPPINGS as int_NCM from .src.comfymath.float import NODE_CLASS_MAPPINGS as float_NCM from .src.comfymath.number import NODE_CLASS_MAPPINGS as number_NCM -from .src.comfymath.vec2 import NODE_CLASS_MAPPINGS as vec2_NCM -from .src.comfymath.vec3 import NODE_CLASS_MAPPINGS as vec3_NCM -from .src.comfymath.vec4 import NODE_CLASS_MAPPINGS as vec4_NCM +from .src.comfymath.vec import NODE_CLASS_MAPPINGS as vec_NCM from .src.comfymath.control import NODE_CLASS_MAPPINGS as control_NCM @@ -13,10 +11,15 @@ NODE_CLASS_MAPPINGS = { **int_NCM, **float_NCM, **number_NCM, - **vec2_NCM, - **vec3_NCM, - **vec4_NCM, + **vec_NCM, **control_NCM, } -NODE_DISPLAY_NAME_MAPPINGS = {key: key for key in NODE_CLASS_MAPPINGS} + +def remove_cm_prefix(node_mapping: str) -> str: + if node_mapping.startswith("CM_"): + return node_mapping[3:] + return node_mapping + + +NODE_DISPLAY_NAME_MAPPINGS = {key: remove_cm_prefix(key) for key in NODE_CLASS_MAPPINGS} diff --git a/src/comfymath/control.py b/src/comfymath/control.py index 95079ab..7927db0 100644 --- a/src/comfymath/control.py +++ b/src/comfymath/control.py @@ -1,104 +1,3 @@ -from typing import Any, Mapping, Callable +from typing import Any, Mapping - -def _pascalcase(input_str: str) -> str: - return input_str.lower().title().replace("_", "") - - -def _get_input_types_method(input_type: str) -> Callable[[], Mapping[str, Any]]: - input_type_tuple: tuple = (input_type,) - if input_type == "INT": - input_type_tuple = ("INT", {"default": 0}) - if input_type == "FLOAT": - input_type_tuple = ("FLOAT", {"default": 0.0}) - if input_type == "STRING": - input_type_tuple = ("STRING", {"default": ""}) - return lambda: { - "required": { - "condition": ("INT", {"default": 0}), - "then": input_type_tuple, - "else_": input_type_tuple, - } - } - - -def _op_if(self, condition: int, then: Any, else_: Any) -> Any: - if condition != 0: - return (then,) - else: - return (else_,) - - -def _make_if_node_class(input_type: str, category: str) -> type: - name = f"{_pascalcase(input_type)}If" - - class_dict = { - "INPUT_TYPES": _get_input_types_method(input_type), - "RETURN_TYPES": (input_type,), - "FUNCTION": "op", - "CATEGORY": category, - "op": _op_if, - } - - return type(name, (), class_dict) - - -_TYPES = { - "INT": "base", - "FLOAT": "base", - "STRING": "base", - "CONDITIONING": "base", - "IMAGE": "base", - "LATENT": "base", - "MODEL": "base", - "CLIP": "base", - "VAE": "base", - "CLIP_VISION": "base", - "CONTROL_NET": "base", - "CLIP_VISION_OUTPUT": "base", - "STYLE_MODEL": "base", - "GLIGEN": "base", - "MASK": "base", - # Comfy Math - "VEC2": "math", - "VEC3": "math", - "VEC4": "math", - # Was Node Suite - "CROP_DATA": "was", - "LIST": "was", - "MIDAS_MODEL": "was", - "SEED": "was", - "DICT": "was", - "NUMBER": "was", - "BLIP_MODEL": "was", - "CLIPSEG_MODEL": "was", - "LANG_SAM_MODEL": "was", - "SAM_MODEL": "was", - "SAM_PARAMETERS": "was", - "IMAGE_BOUNDS": "was", - "UPSCALE_MODEL": "was", - # Impact Pack - "ONNX_DETECTOR": "impact", - "BBOX_DETECTOR": "impact", - "SEGM_DETECTOR": "impact", - "SEGS": "impact", - "KSAMPLER": "impact", - "KSAMPLER_ADVANCED": "impact", - "REGIONAL_PROMPTS": "impact", - "DETAILER_PIPE": "impact", - "PK_HOOK": "impact", - "UPSCALER": "impact", -} - - -IF_CLASS_MAPPINGS = { - f"{_pascalcase(type_name)}If": _make_if_node_class( - type_name, f"math/control/if/{category}" - ) - for type_name, category in _TYPES.items() -} - - -NODE_CLASS_MAPPINGS = { - **IF_CLASS_MAPPINGS, -} +NODE_CLASS_MAPPINGS: Mapping[str, Any] = {} diff --git a/src/comfymath/convert.py b/src/comfymath/convert.py index b6ebbf4..e9ca3d7 100644 --- a/src/comfymath/convert.py +++ b/src/comfymath/convert.py @@ -1,8 +1,6 @@ from typing import Any, Mapping -from .vec2 import Vec2, VEC2_ZERO -from .vec3 import Vec3, VEC3_ZERO -from .vec4 import Vec4, VEC4_ZERO +from .vec import Vec2, VEC2_ZERO, Vec3, VEC3_ZERO, Vec4, VEC4_ZERO from .number import number @@ -44,6 +42,7 @@ class IntToNumber: def op(self, a: int) -> tuple[number]: return (a,) + class NumberToInt: @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: @@ -61,7 +60,7 @@ class FloatToNumber: @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: return {"required": {"a": ("FLOAT", {"default": 0.0})}} - + RETURN_TYPES = ("NUMBER",) FUNCTION = "op" CATEGORY = "math/conversion" @@ -69,11 +68,12 @@ class FloatToNumber: def op(self, a: float) -> tuple[number]: return (a,) + class NumberToFloat: @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: return {"required": {"a": ("NUMBER", {"default": 0.0})}} - + RETURN_TYPES = ("FLOAT",) FUNCTION = "op" CATEGORY = "math/conversion" @@ -100,6 +100,23 @@ class ComposeVec2: return ((x, y),) +class FillVec2: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "a": ("FLOAT", {"default": 0.0}), + } + } + + RETURN_TYPES = ("VEC2",) + FUNCTION = "op" + CATEGORY = "math/conversion" + + def op(self, a: float) -> tuple[Vec2]: + return ((a, a),) + + class BreakoutVec2: @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: @@ -132,6 +149,23 @@ class ComposeVec3: return ((x, y, z),) +class FillVec3: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "a": ("FLOAT", {"default": 0.0}), + } + } + + RETURN_TYPES = ("VEC3",) + FUNCTION = "op" + CATEGORY = "math/conversion" + + def op(self, a: float) -> tuple[Vec3]: + return ((a, a, a),) + + class BreakoutVec3: @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: @@ -165,6 +199,23 @@ class ComposeVec4: return ((x, y, z, w),) +class FillVec4: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "a": ("FLOAT", {"default": 0.0}), + } + } + + RETURN_TYPES = ("VEC4",) + FUNCTION = "op" + CATEGORY = "math/conversion" + + def op(self, a: float) -> tuple[Vec4]: + return ((a, a, a, a),) + + class BreakoutVec4: @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: @@ -179,16 +230,16 @@ class BreakoutVec4: NODE_CLASS_MAPPINGS = { - "FloatToInt": FloatToInt, - "IntToFloat": IntToFloat, - "IntToNumber": IntToNumber, - "NumberToInt": NumberToInt, - "FloatToNumber": FloatToNumber, - "NumberToFloat": NumberToFloat, - "ComposeVec2": ComposeVec2, - "ComposeVec3": ComposeVec3, - "ComposeVec4": ComposeVec4, - "BreakoutVec2": BreakoutVec2, - "BreakoutVec3": BreakoutVec3, - "BreakoutVec4": BreakoutVec4, + "CM_FloatToInt": FloatToInt, + "CM_IntToFloat": IntToFloat, + "CM_IntToNumber": IntToNumber, + "CM_NumberToInt": NumberToInt, + "CM_FloatToNumber": FloatToNumber, + "CM_NumberToFloat": NumberToFloat, + "CM_ComposeVec2": ComposeVec2, + "CM_ComposeVec3": ComposeVec3, + "CM_ComposeVec4": ComposeVec4, + "CM_BreakoutVec2": BreakoutVec2, + "CM_BreakoutVec3": BreakoutVec3, + "CM_BreakoutVec4": BreakoutVec4, } diff --git a/src/comfymath/float.py b/src/comfymath/float.py index 5828aa0..0e74b33 100644 --- a/src/comfymath/float.py +++ b/src/comfymath/float.py @@ -1,188 +1,159 @@ import math -from dataclasses import dataclass -from typing import Callable +from typing import Any, Callable, Mapping + +DEFAULT_FLOAT = ("FLOAT", {"default": 0.0}) + +FLOAT_UNARY_OPERATIONS: Mapping[str, Callable[[float], float]] = { + "Neg": lambda a: -a, + "Inc": lambda a: a + 1, + "Dec": lambda a: a - 1, + "Abs": lambda a: abs(a), + "Sqr": lambda a: a * a, + "Cube": lambda a: a * a * a, + "Sqrt": lambda a: math.sqrt(a), + "Exp": lambda a: math.exp(a), + "Ln": lambda a: math.log(a), + "Log10": lambda a: math.log10(a), + "Log2": lambda a: math.log2(a), + "Sin": lambda a: math.sin(a), + "Cos": lambda a: math.cos(a), + "Tan": lambda a: math.tan(a), + "Asin": lambda a: math.asin(a), + "Acos": lambda a: math.acos(a), + "Atan": lambda a: math.atan(a), + "Sinh": lambda a: math.sinh(a), + "Cosh": lambda a: math.cosh(a), + "Tanh": lambda a: math.tanh(a), + "Asinh": lambda a: math.asinh(a), + "Acosh": lambda a: math.acosh(a), + "Atanh": lambda a: math.atanh(a), + "Round": lambda a: round(a), + "Floor": lambda a: math.floor(a), + "Ceil": lambda a: math.ceil(a), + "Trunc": lambda a: math.trunc(a), + "Erf": lambda a: math.erf(a), + "Erfc": lambda a: math.erfc(a), + "Gamma": lambda a: math.gamma(a), + "Radians": lambda a: math.radians(a), + "Degrees": lambda a: math.degrees(a), +} + +FLOAT_UNARY_CONDITIONS: Mapping[str, Callable[[float], bool]] = { + "IsZero": lambda a: a == 0.0, + "IsPositive": lambda a: a > 0.0, + "IsNegative": lambda a: a < 0.0, + "IsNonZero": lambda a: a != 0.0, + "IsPositiveInfinity": lambda a: math.isinf(a) and a > 0.0, + "IsNegativeInfinity": lambda a: math.isinf(a) and a < 0.0, + "IsNaN": lambda a: math.isnan(a), + "IsFinite": lambda a: math.isfinite(a), + "IsInfinite": lambda a: math.isinf(a), + "IsEven": lambda a: a % 2 == 0.0, + "IsOdd": lambda a: a % 2 != 0.0, +} + +FLOAT_BINARY_OPERATIONS: Mapping[str, Callable[[float, float], float]] = { + "Add": lambda a, b: a + b, + "Sub": lambda a, b: a - b, + "Mul": lambda a, b: a * b, + "Div": lambda a, b: a / b, + "Mod": lambda a, b: a % b, + "Pow": lambda a, b: a**b, + "FloorDiv": lambda a, b: a // b, + "Max": lambda a, b: max(a, b), + "Min": lambda a, b: min(a, b), + "Log": lambda a, b: math.log(a, b), + "Atan2": lambda a, b: math.atan2(a, b), +} + +FLOAT_BINARY_CONDITIONS: Mapping[str, Callable[[float, float], bool]] = { + "Eq": lambda a, b: a == b, + "Neq": lambda a, b: a != b, + "Gt": lambda a, b: a > b, + "Gte": lambda a, b: a >= b, + "Lt": lambda a, b: a < b, + "Lte": lambda a, b: a <= b, +} -@dataclass class FloatUnaryOperation: - name: str - function: Callable[[float], float] + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(FLOAT_UNARY_OPERATIONS.keys()),), + "a": DEFAULT_FLOAT, + } + } + + RETURN_TYPES = ("FLOAT",) + FUNCTION = "op" + CATEGORY = "math/float" + + def op(self, op: str, a: float) -> tuple[float]: + return (FLOAT_UNARY_OPERATIONS[op](a),) -@dataclass class FloatUnaryCondition: - name: str - function: Callable[[float], bool] + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(FLOAT_UNARY_CONDITIONS.keys()),), + "a": DEFAULT_FLOAT, + } + } + + RETURN_TYPES = ("BOOL",) + FUNCTION = "op" + CATEGORY = "math/float" + + def op(self, op: str, a: float) -> tuple[bool]: + return (FLOAT_UNARY_CONDITIONS[op](a),) -@dataclass class FloatBinaryOperation: - name: str - function: Callable[[float, float], float] + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(FLOAT_BINARY_OPERATIONS.keys()),), + "a": DEFAULT_FLOAT, + "b": DEFAULT_FLOAT, + } + } + + RETURN_TYPES = ("FLOAT",) + FUNCTION = "op" + CATEGORY = "math/float" + + def op(self, op: str, a: float, b: float) -> tuple[float]: + return (FLOAT_BINARY_OPERATIONS[op](a, b),) -@dataclass class FloatBinaryCondition: - name: str - function: Callable[[float, float], bool] - - -FLOAT_UNARY_OPERATIONS = [ - FloatUnaryOperation("Neg", lambda a: -a), - FloatUnaryOperation("Inc", lambda a: a + 1), - FloatUnaryOperation("Dec", lambda a: a - 1), - FloatUnaryOperation("Abs", lambda a: abs(a)), - FloatUnaryOperation("Sqr", lambda a: a * a), - FloatUnaryOperation("Cube", lambda a: a * a * a), - FloatUnaryOperation("Sqrt", lambda a: math.sqrt(a)), - FloatUnaryOperation("Exp", lambda a: math.exp(a)), - FloatUnaryOperation("Ln", lambda a: math.log(a)), - FloatUnaryOperation("Log10", lambda a: math.log10(a)), - FloatUnaryOperation("Log2", lambda a: math.log2(a)), - FloatUnaryOperation("Sin", lambda a: math.sin(a)), - FloatUnaryOperation("Cos", lambda a: math.cos(a)), - FloatUnaryOperation("Tan", lambda a: math.tan(a)), - FloatUnaryOperation("Asin", lambda a: math.asin(a)), - FloatUnaryOperation("Acos", lambda a: math.acos(a)), - FloatUnaryOperation("Atan", lambda a: math.atan(a)), - FloatUnaryOperation("Sinh", lambda a: math.sinh(a)), - FloatUnaryOperation("Cosh", lambda a: math.cosh(a)), - FloatUnaryOperation("Tanh", lambda a: math.tanh(a)), - FloatUnaryOperation("Asinh", lambda a: math.asinh(a)), - FloatUnaryOperation("Acosh", lambda a: math.acosh(a)), - FloatUnaryOperation("Atanh", lambda a: math.atanh(a)), - FloatUnaryOperation("Round", lambda a: round(a)), - FloatUnaryOperation("Floor", lambda a: math.floor(a)), - FloatUnaryOperation("Ceil", lambda a: math.ceil(a)), - FloatUnaryOperation("Trunc", lambda a: math.trunc(a)), - FloatUnaryOperation("Erf", lambda a: math.erf(a)), - FloatUnaryOperation("Erfc", lambda a: math.erfc(a)), - FloatUnaryOperation("Gamma", lambda a: math.gamma(a)), - FloatUnaryOperation("Radians", lambda a: math.radians(a)), - FloatUnaryOperation("Degrees", lambda a: math.degrees(a)), -] - -FLOAT_UNARY_CONDITIONS = [ - FloatUnaryCondition("IsZero", lambda a: a == 0.0), - FloatUnaryCondition("IsPositive", lambda a: a > 0.0), - FloatUnaryCondition("IsNegative", lambda a: a < 0.0), - FloatUnaryCondition("IsNonZero", lambda a: a != 0.0), - FloatUnaryCondition("IsPositiveInfinity", lambda a: math.isinf(a) and a > 0.0), - FloatUnaryCondition("IsNegativeInfinity", lambda a: math.isinf(a) and a < 0.0), - FloatUnaryCondition("IsNaN", lambda a: math.isnan(a)), - FloatUnaryCondition("IsFinite", lambda a: math.isfinite(a)), - FloatUnaryCondition("IsInfinite", lambda a: math.isinf(a)), - FloatUnaryCondition("IsEven", lambda a: a % 2 == 0.0), - FloatUnaryCondition("IsOdd", lambda a: a % 2 != 0.0), -] - -FLOAT_BINARY_OPERATIONS = [ - FloatBinaryOperation("Add", lambda a, b: a + b), - FloatBinaryOperation("Sub", lambda a, b: a - b), - FloatBinaryOperation("Mul", lambda a, b: a * b), - FloatBinaryOperation("Div", lambda a, b: a / b), - FloatBinaryOperation("Mod", lambda a, b: a % b), - FloatBinaryOperation("Pow", lambda a, b: a**b), - FloatBinaryOperation("FloorDiv", lambda a, b: a // b), - FloatBinaryOperation("Max", lambda a, b: max(a, b)), - FloatBinaryOperation("Min", lambda a, b: min(a, b)), - FloatBinaryOperation("Log", lambda a, b: math.log(a, b)), - FloatBinaryOperation("Atan2", lambda a, b: math.atan2(a, b)), -] - -FLOAT_BINARY_CONDITIONS = [ - FloatBinaryCondition("Eq", lambda a, b: a == b), - FloatBinaryCondition("Neq", lambda a, b: a != b), - FloatBinaryCondition("Gt", lambda a, b: a > b), - FloatBinaryCondition("Gte", lambda a, b: a >= b), - FloatBinaryCondition("Lt", lambda a, b: a < b), - FloatBinaryCondition("Lte", lambda a, b: a <= b), -] - - -def _get_float_unary_op_node_class(op: FloatUnaryOperation) -> type: - name = f"Float{op.name}" - class_dict = { - "INPUT_TYPES": lambda: {"required": {"a": ("FLOAT", {"default": 0.0})}}, - "RETURN_TYPES": ("FLOAT",), - "FUNCTION": "op", - "CATEGORY": "math/float", - "op": op.function, - } - return type(name, (), class_dict) - - -def _get_float_unary_cond_node_class(op: FloatUnaryCondition) -> type: - name = f"Float{op.name}" - class_dict = { - "INPUT_TYPES": lambda: {"required": {"a": ("FLOAT", {"default": 0.0})}}, - "RETURN_TYPES": ("INT",), - "FUNCTION": "op", - "CATEGORY": "math/float", - "op": lambda a: int(op.function(a)), - } - return type(name, (), class_dict) - - -def _get_float_binary_op_node_class(op: FloatBinaryOperation) -> type: - name = f"Float{op.name}" - class_dict = { - "INPUT_TYPES": lambda: { + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { "required": { - "a": ("FLOAT", {"default": 0.0}), - "b": ("FLOAT", {"default": 0.0}), + "op": (list(FLOAT_BINARY_CONDITIONS.keys()),), + "a": DEFAULT_FLOAT, + "b": DEFAULT_FLOAT, } - }, - "RETURN_TYPES": ("FLOAT",), - "FUNCTION": "op", - "CATEGORY": "math/float", - "op": op.function, - } - return type(name, (), class_dict) + } + RETURN_TYPES = ("BOOL",) + FUNCTION = "op" + CATEGORY = "math/float" -def _get_float_binary_cond_node_class(op: FloatBinaryCondition) -> type: - name = f"Float{op.name}" - class_dict = { - "INPUT_TYPES": lambda: { - "required": { - "a": ("FLOAT", {"default": 0.0}), - "b": ("FLOAT", {"default": 0.0}), - } - }, - "RETURN_TYPES": ("INT",), - "FUNCTION": "op", - "CATEGORY": "math/float", - "op": lambda a, b: int(op.function(a, b)), - } - return type(name, (), class_dict) + def op(self, op: str, a: float, b: float) -> tuple[bool]: + return (FLOAT_BINARY_CONDITIONS[op](a, b),) -FLOAT_UNARY_OPERATION_CLASS_MAPPINGS = { - f"Float{op.name}": _get_float_unary_op_node_class(op) - for op in FLOAT_UNARY_OPERATIONS -} - -FLOAT_UNARY_CONDITION_CLASS_MAPPINGS = { - f"Float{op.name}": _get_float_unary_cond_node_class(op) - for op in FLOAT_UNARY_CONDITIONS -} - - -FLOAT_BINARY_OPERATION_CLASS_MAPPINGS = { - f"Float{op.name}": _get_float_binary_op_node_class(op) - for op in FLOAT_BINARY_OPERATIONS -} - -FLOAT_BINARY_CONDITION_CLASS_MAPPINGS = { - f"Float{op.name}": _get_float_binary_cond_node_class(op) - for op in FLOAT_BINARY_CONDITIONS -} - NODE_CLASS_MAPPINGS = { - **FLOAT_UNARY_OPERATION_CLASS_MAPPINGS, - **FLOAT_UNARY_CONDITION_CLASS_MAPPINGS, - **FLOAT_BINARY_OPERATION_CLASS_MAPPINGS, - **FLOAT_BINARY_CONDITION_CLASS_MAPPINGS, + "CM_FloatUnaryOperation": FloatUnaryOperation, + "CM_FloatUnaryCondition": FloatUnaryCondition, + "CM_FloatBinaryOperation": FloatBinaryOperation, + "CM_FloatBinaryCondition": FloatBinaryCondition, } diff --git a/src/comfymath/int.py b/src/comfymath/int.py index ebe91a7..576fc6a 100644 --- a/src/comfymath/int.py +++ b/src/comfymath/int.py @@ -1,151 +1,129 @@ import math -from dataclasses import dataclass -from typing import Callable, TypeAlias +from typing import Any, Callable, Mapping + +DEFAULT_INT = ("INT", {"default": 0}) + +INT_UNARY_OPERATIONS: Mapping[str, Callable[[int], int]] = { + "Abs": lambda a: abs(a), + "Neg": lambda a: -a, + "Inc": lambda a: a + 1, + "Dec": lambda a: a - 1, + "Sqr": lambda a: a * a, + "Cube": lambda a: a * a * a, + "Not": lambda a: ~a, + "Factorial": lambda a: math.factorial(a), +} + +INT_UNARY_CONDITIONS: Mapping[str, Callable[[int], bool]] = { + "IsZero": lambda a: a == 0, + "IsNonZero": lambda a: a != 0, + "IsPositive": lambda a: a > 0, + "IsNegative": lambda a: a < 0, + "IsEven": lambda a: a % 2 == 0, + "IsOdd": lambda a: a % 2 == 1, +} + +INT_BINARY_OPERATIONS: Mapping[str, Callable[[int, int], int]] = { + "Add": lambda a, b: a + b, + "Sub": lambda a, b: a - b, + "Mul": lambda a, b: a * b, + "Div": lambda a, b: a // b, + "Mod": lambda a, b: a % b, + "Pow": lambda a, b: a**b, + "And": lambda a, b: a & b, + "Nand": lambda a, b: ~a & b, + "Or": lambda a, b: a | b, + "Nor": lambda a, b: ~a & b, + "Xor": lambda a, b: a ^ b, + "Xnor": lambda a, b: ~a ^ b, + "Shl": lambda a, b: a << b, + "Shr": lambda a, b: a >> b, + "Max": lambda a, b: max(a, b), + "Min": lambda a, b: min(a, b), +} + +INT_BINARY_CONDITIONS: Mapping[str, Callable[[int, int], bool]] = { + "Eq": lambda a, b: a == b, + "Neq": lambda a, b: a != b, + "Gt": lambda a, b: a > b, + "Lt": lambda a, b: a < b, + "Geq": lambda a, b: a >= b, + "Leq": lambda a, b: a <= b, +} -@dataclass class IntUnaryOperation: - name: str - function: Callable[[int], int] + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": {"op": (list(INT_UNARY_OPERATIONS.keys()),), "a": DEFAULT_INT} + } + + RETURN_TYPES = ("INT",) + FUNCTION = "op" + CATEGORY = "math/int" + + def op(self, op: str, a: int) -> tuple[int]: + return (INT_UNARY_OPERATIONS[op](a),) -@dataclass class IntUnaryCondition: - name: str - function: Callable[[int], bool] + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": {"op": (list(INT_UNARY_CONDITIONS.keys()),), "a": DEFAULT_INT} + } + + RETURN_TYPES = ("BOOL",) + FUNCTION = "op" + CATEGORY = "math/int" + + def op(self, op: str, a: int) -> tuple[bool]: + return (INT_UNARY_CONDITIONS[op](a),) -@dataclass class IntBinaryOperation: - name: str - function: Callable[[int, int], int] + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(INT_BINARY_OPERATIONS.keys()),), + "a": DEFAULT_INT, + "b": DEFAULT_INT, + } + } + + RETURN_TYPES = ("INT",) + FUNCTION = "op" + CATEGORY = "math/int" + + def op(self, op: str, a: int, b: int) -> tuple[int]: + return (INT_BINARY_OPERATIONS[op](a, b),) -@dataclass class IntBinaryCondition: - name: str - function: Callable[[int, int], bool] + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(INT_BINARY_CONDITIONS.keys()),), + "a": DEFAULT_INT, + "b": DEFAULT_INT, + } + } + RETURN_TYPES = ("BOOL",) + FUNCTION = "op" + CATEGORY = "math/int" -INT_UNARY_OPERATIONS = [ - IntUnaryOperation("Abs", lambda a: abs(a)), - IntUnaryOperation("Neg", lambda a: -a), - IntUnaryOperation("Inc", lambda a: a + 1), - IntUnaryOperation("Dec", lambda a: a - 1), - IntUnaryOperation("Sqr", lambda a: a * a), - IntUnaryOperation("Cube", lambda a: a * a * a), - IntUnaryOperation("Not", lambda a: ~a), - IntUnaryOperation("Factorial", lambda a: math.factorial(a)), -] + def op(self, op: str, a: int, b: int) -> tuple[bool]: + return (INT_BINARY_CONDITIONS[op](a, b),) -INT_UNARY_CONDITIONS = [ - IntUnaryCondition("IsZero", lambda a: a == 0), - IntUnaryCondition("IsNonZero", lambda a: a != 0), - IntUnaryCondition("IsPositive", lambda a: a > 0), - IntUnaryCondition("IsNegative", lambda a: a < 0), - IntUnaryCondition("IsEven", lambda a: a % 2 == 0), - IntUnaryCondition("IsOdd", lambda a: a % 2 == 1), -] - -INT_BINARY_OPERATIONS = [ - IntBinaryOperation("Add", lambda a, b: a + b), - IntBinaryOperation("Sub", lambda a, b: a - b), - IntBinaryOperation("Mul", lambda a, b: a * b), - IntBinaryOperation("Div", lambda a, b: a // b), - IntBinaryOperation("Mod", lambda a, b: a % b), - IntBinaryOperation("Pow", lambda a, b: a**b), - IntBinaryOperation("And", lambda a, b: a & b), - IntBinaryOperation("Nand", lambda a, b: ~a & b), - IntBinaryOperation("Or", lambda a, b: a | b), - IntBinaryOperation("Nor", lambda a, b: ~a & b), - IntBinaryOperation("Xor", lambda a, b: a ^ b), - IntBinaryOperation("Xnor", lambda a, b: ~a ^ b), - IntBinaryOperation("Shl", lambda a, b: a << b), - IntBinaryOperation("Shr", lambda a, b: a >> b), - IntBinaryOperation("Max", lambda a, b: max(a, b)), - IntBinaryOperation("Min", lambda a, b: min(a, b)), -] - -INT_BINARY_CONDITIONS = [ - IntBinaryCondition("Eq", lambda a, b: a == b), - IntBinaryCondition("Neq", lambda a, b: a != b), - IntBinaryCondition("Gt", lambda a, b: a > b), - IntBinaryCondition("Lt", lambda a, b: a < b), - IntBinaryCondition("Geq", lambda a, b: a >= b), - IntBinaryCondition("Leq", lambda a, b: a <= b), -] - - -def _get_int_unary_op_node_class(op: IntUnaryOperation) -> type: - name = f"Int{op.name}" - class_dict = { - "INPUT_TYPES": lambda: {"required": {"a": ("INT", {"default": 0})}}, - "RETURN_TYPES": ("INT",), - "FUNCTION": "op", - "CATEGORY": "math/int", - "op": lambda a: op.function(a), - } - return type(name, (), class_dict) - -def _get_int_unary_cond_node_class(op: IntUnaryCondition) -> type: - name = f"Int{op.name}" - class_dict = { - "INPUT_TYPES": lambda: {"required": {"a": ("INT", {"default": 0})}}, - "RETURN_TYPES": ("INT",), - "FUNCTION": "op", - "CATEGORY": "math/int", - "op": lambda a: int(op.function(a)), - } - return type(name, (), class_dict) - - -def _get_int_binary_op_node_class(op: IntBinaryOperation) -> type: - name = f"Int{op.name}" - class_dict = { - "INPUT_TYPES": lambda: { - "required": {"a": ("INT", {"default": 0}), "b": ("INT", {"default": 0})} - }, - "RETURN_TYPES": ("INT",), - "FUNCTION": "op", - "CATEGORY": "math/int", - "op": lambda a, b: op.function(a, b), - } - return type(name, (), class_dict) - -def _get_int_binary_cond_node_class(op: IntBinaryCondition) -> type: - name = f"Int{op.name}" - class_dict = { - "INPUT_TYPES": lambda: { - "required": {"a": ("INT", {"default": 0}), "b": ("INT", {"default": 0})} - }, - "RETURN_TYPES": ("INT",), - "FUNCTION": "op", - "CATEGORY": "math/int", - "op": lambda a, b: int(op.function(a, b)), - } - return type(name, (), class_dict) - - -FLOAT_UNARY_OPERATION_CLASS_MAPPINGS = { - f"Int{op.name}": _get_int_unary_op_node_class(op) for op in INT_UNARY_OPERATIONS -} - -FLOAT_UNARY_CONDITION_CLASS_MAPPINGS = { - f"Int{op.name}": _get_int_unary_cond_node_class(op) for op in INT_UNARY_CONDITIONS -} - -FLOAT_BINARY_OPERATION_CLASS_MAPPINGS = { - f"Int{op.name}": _get_int_binary_op_node_class(op) for op in INT_BINARY_OPERATIONS -} - -FLOAT_BINARY_CONDITION_CLASS_MAPPINGS = { - f"Int{op.name}": _get_int_binary_cond_node_class(op) for op in INT_BINARY_CONDITIONS -} NODE_CLASS_MAPPINGS = { - **FLOAT_UNARY_OPERATION_CLASS_MAPPINGS, - **FLOAT_UNARY_CONDITION_CLASS_MAPPINGS, - **FLOAT_BINARY_OPERATION_CLASS_MAPPINGS, - **FLOAT_BINARY_CONDITION_CLASS_MAPPINGS, + "CM_IntUnaryOperation": IntUnaryOperation, + "CM_IntUnaryCondition": IntUnaryCondition, + "CM_IntBinaryOperation": IntBinaryOperation, + "CM_IntBinaryCondition": IntBinaryCondition, } diff --git a/src/comfymath/number.py b/src/comfymath/number.py index cfed571..aa7cf29 100644 --- a/src/comfymath/number.py +++ b/src/comfymath/number.py @@ -1,205 +1,95 @@ from dataclasses import dataclass -from typing import Callable, TypeAlias, Sequence +from typing import Any, Callable, Mapping, TypeAlias -from .int import ( - INT_UNARY_OPERATIONS, - INT_UNARY_CONDITIONS, - INT_BINARY_OPERATIONS, - INT_BINARY_CONDITIONS, -) from .float import ( FLOAT_UNARY_OPERATIONS, FLOAT_UNARY_CONDITIONS, FLOAT_BINARY_OPERATIONS, FLOAT_BINARY_CONDITIONS, - FloatUnaryOperation, - FloatUnaryCondition, - FloatBinaryOperation, - FloatBinaryCondition, ) +DEFAULT_NUMBER = ("NUMBER", {"default": 0.0}) + number: TypeAlias = int | float -@dataclass class NumberUnaryOperation: - name: str - function: Callable[[number], number] + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(FLOAT_UNARY_OPERATIONS.keys()),), + "a": DEFAULT_NUMBER, + } + } + + RETURN_TYPES = ("NUMBER",) + FUNCTION = "op" + CATEGORY = "math/number" + + def op(self, op: str, a: number) -> tuple[float]: + return (FLOAT_UNARY_OPERATIONS[op](float(a)),) -@dataclass class NumberUnaryCondition: - name: str - function: Callable[[number], bool] + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(FLOAT_UNARY_CONDITIONS.keys()),), + "a": DEFAULT_NUMBER, + } + } + + RETURN_TYPES = ("BOOL",) + FUNCTION = "op" + CATEGORY = "math/number" + + def op(self, op: str, a: number) -> tuple[bool]: + return (FLOAT_UNARY_CONDITIONS[op](float(a)),) -@dataclass class NumberBinaryOperation: - name: str - function: Callable[[number, number], number] + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(FLOAT_BINARY_OPERATIONS.keys()),), + "a": DEFAULT_NUMBER, + "b": DEFAULT_NUMBER, + } + } + + RETURN_TYPES = ("NUMBER",) + FUNCTION = "op" + CATEGORY = "math/number" + + def op(self, op: str, a: number, b: number) -> tuple[float]: + return (FLOAT_BINARY_OPERATIONS[op](float(a), float(b)),) -@dataclass class NumberBinaryCondition: - name: str - function: Callable[[number, number], bool] - - -def _float_unary_operation_to_num_unary_operation( - op: FloatUnaryOperation, -) -> NumberUnaryOperation: - return NumberUnaryOperation(op.name, op.function) - - -def _float_unary_condition_to_num_unary_condition( - op: FloatUnaryCondition, -) -> NumberUnaryCondition: - return NumberUnaryCondition(op.name, op.function) - - -def _float_binary_operation_to_num_binary_operation( - op: FloatBinaryOperation, -) -> NumberBinaryOperation: - return NumberBinaryOperation(op.name, op.function) - - -def _float_binary_condition_to_num_binary_condition( - op: FloatBinaryCondition, -) -> NumberBinaryCondition: - return NumberBinaryCondition(op.name, op.function) - - -def _combine_unary_operations() -> Sequence[NumberUnaryOperation]: - float_unary_op_names = {op.name for op in FLOAT_UNARY_OPERATIONS} - int_unary_op_names = {op.name for op in INT_UNARY_OPERATIONS} - num_unary_op_names = float_unary_op_names & int_unary_op_names - return [ - _float_unary_operation_to_num_unary_operation(op) - for op in FLOAT_UNARY_OPERATIONS - if op.name in num_unary_op_names - ] - - -def _combine_unary_conditions() -> Sequence[NumberUnaryCondition]: - float_unary_cond_names = {op.name for op in FLOAT_UNARY_CONDITIONS} - int_unary_cond_names = {op.name for op in INT_UNARY_CONDITIONS} - num_unary_cond_names = float_unary_cond_names & int_unary_cond_names - return [ - _float_unary_condition_to_num_unary_condition(op) - for op in FLOAT_UNARY_CONDITIONS - if op.name in num_unary_cond_names - ] - - -def _combine_binary_operations() -> Sequence[NumberBinaryOperation]: - float_binary_op_names = {op.name for op in FLOAT_BINARY_OPERATIONS} - int_binary_op_names = {op.name for op in INT_BINARY_OPERATIONS} - num_binary_op_names = float_binary_op_names & int_binary_op_names - return [ - _float_binary_operation_to_num_binary_operation(op) - for op in FLOAT_BINARY_OPERATIONS - if op.name in num_binary_op_names - ] - - -def _combine_binary_conditions() -> Sequence[NumberBinaryCondition]: - float_binary_cond_names = {op.name for op in FLOAT_BINARY_CONDITIONS} - int_binary_cond_names = {op.name for op in INT_BINARY_CONDITIONS} - num_binary_cond_names = float_binary_cond_names & int_binary_cond_names - return [ - _float_binary_condition_to_num_binary_condition(op) - for op in FLOAT_BINARY_CONDITIONS - if op.name in num_binary_cond_names - ] - - -def _get_number_unary_op_node_class(op: NumberUnaryOperation) -> type: - name = f"Number{op.name}" - class_dict = { - "INPUT_TYPES": lambda: {"required": {"a": ("NUMBER", {"default": 0.0})}}, - "RETURN_TYPES": ("NUMBER",), - "FUNCTION": "op", - "CATEGORY": "math/number", - "op": op.function, - } - return type(name, (), class_dict) - - -def _get_number_unary_cond_node_class(op: NumberUnaryCondition) -> type: - name = f"Number{op.name}" - class_dict = { - "INPUT_TYPES": lambda: {"required": {"a": ("NUMBER", {"default": 0.0})}}, - "RETURN_TYPES": ("INT",), - "FUNCTION": "op", - "CATEGORY": "math/number", - "op": lambda a: int(op.function(a)), - } - return type(name, (), class_dict) - - -def _get_number_binary_op_node_class(op: NumberBinaryOperation) -> type: - name = f"Number{op.name}" - class_dict = { - "INPUT_TYPES": lambda: { + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { "required": { - "a": ("NUMBER", {"default": 0.0}), - "b": ("NUMBER", {"default": 0.0}), + "op": (list(FLOAT_BINARY_CONDITIONS.keys()),), + "a": DEFAULT_NUMBER, + "b": DEFAULT_NUMBER, } - }, - "RETURN_TYPES": ("NUMBER",), - "FUNCTION": "op", - "CATEGORY": "math/number", - "op": op.function, - } - return type(name, (), class_dict) + } + RETURN_TYPES = ("BOOL",) + FUNCTION = "op" + CATEGORY = "math/float" -def _get_number_binary_cond_node_class(op: NumberBinaryCondition) -> type: - name = f"Number{op.name}" - class_dict = { - "INPUT_TYPES": lambda: { - "required": { - "a": ("NUMBER", {"default": 0.0}), - "b": ("NUMBER", {"default": 0.0}), - } - }, - "RETURN_TYPES": ("NUMBER",), - "FUNCTION": "op", - "CATEGORY": "math/number", - "op": lambda a, b: int(op.function(a, b)), - } - return type(name, (), class_dict) + def op(self, op: str, a: number, b: number) -> tuple[bool]: + return (FLOAT_BINARY_CONDITIONS[op](float(a), float(b)),) -NUMBER_UNARY_OPERATIONS = _combine_unary_operations() -NUMBER_UNARY_CONDITIONS = _combine_unary_conditions() -NUMBER_BINARY_OPERATIONS = _combine_binary_operations() -NUMBER_BINARY_CONDITIONS = _combine_binary_conditions() - -NUMBER_UNARY_OPERATION_CLASS_MAPPINGS = { - f"Number{op.name}": _get_number_unary_op_node_class(op) - for op in NUMBER_UNARY_OPERATIONS -} - -NUMBER_UNARY_CONDITION_CLASS_MAPPINGS = { - f"Number{op.name}": _get_number_unary_cond_node_class(op) - for op in NUMBER_UNARY_CONDITIONS -} - - -NUMBER_BINARY_OPERATION_CLASS_MAPPINGS = { - f"Number{op.name}": _get_number_binary_op_node_class(op) - for op in NUMBER_BINARY_OPERATIONS -} - -NUMBER_BINARY_CONDITION_CLASS_MAPPINGS = { - f"Number{op.name}": _get_number_binary_cond_node_class(op) - for op in NUMBER_BINARY_CONDITIONS -} - NODE_CLASS_MAPPINGS = { - **NUMBER_UNARY_OPERATION_CLASS_MAPPINGS, - **NUMBER_UNARY_CONDITION_CLASS_MAPPINGS, - **NUMBER_BINARY_OPERATION_CLASS_MAPPINGS, - **NUMBER_BINARY_CONDITION_CLASS_MAPPINGS, + "CM_NumberUnaryOperation": NumberUnaryOperation, + "CM_NumberUnaryCondition": NumberUnaryCondition, + "CM_NumberBinaryOperation": NumberBinaryOperation, + "CM_NumberBinaryCondition": NumberBinaryCondition, } diff --git a/src/comfymath/vec.py b/src/comfymath/vec.py new file mode 100644 index 0000000..d45a797 --- /dev/null +++ b/src/comfymath/vec.py @@ -0,0 +1,501 @@ +import numpy + +from typing import Any, Callable, Mapping, TypeAlias + +Vec2: TypeAlias = tuple[float, float] +VEC2_ZERO = (0.0, 0.0) +DEFAULT_VEC2 = ("VEC2", {"default": VEC2_ZERO}) + +Vec3: TypeAlias = tuple[float, float, float] +VEC3_ZERO = (0.0, 0.0, 0.0) +DEFAULT_VEC3 = ("VEC3", {"default": VEC3_ZERO}) + +Vec4: TypeAlias = tuple[float, float, float, float] +VEC4_ZERO = (0.0, 0.0, 0.0, 0.0) +DEFAULT_VEC4 = ("VEC4", {"default": VEC4_ZERO}) + +VEC_UNARY_OPERATIONS: Mapping[str, Callable[[numpy.ndarray], numpy.ndarray]] = { + "Neg": lambda a: -a, + "Normalize": lambda a: a / numpy.linalg.norm(a), +} + +VEC_TO_SCALAR_UNARY_OPERATION: Mapping[str, Callable[[numpy.ndarray], float]] = { + "Norm": lambda a: numpy.linalg.norm(a).astype(float), +} + +VEC_UNARY_CONDITIONS: Mapping[str, Callable[[numpy.ndarray], bool]] = { + "IsZero": lambda a: not numpy.any(a).astype(bool), + "IsNotZero": lambda a: numpy.any(a).astype(bool), + "IsNormalized": lambda a: numpy.allclose(a, a / numpy.linalg.norm(a)), + "IsNotNormalized": lambda a: not numpy.allclose(a, a / numpy.linalg.norm(a)), +} + +VEC_BINARY_OPERATIONS: Mapping[ + str, Callable[[numpy.ndarray, numpy.ndarray], numpy.ndarray] +] = { + "Add": lambda a, b: a + b, + "Sub": lambda a, b: a - b, + "Cross": lambda a, b: numpy.cross(a, b), +} + +VEC_TO_SCALAR_BINARY_OPERATION: Mapping[ + str, Callable[[numpy.ndarray, numpy.ndarray], float] +] = { + "Dot": lambda a, b: numpy.dot(a, b), + "Distance": lambda a, b: numpy.linalg.norm(a - b).astype(float), +} + +VEC_BINARY_CONDITIONS: Mapping[str, Callable[[numpy.ndarray, numpy.ndarray], bool]] = { + "Eq": lambda a, b: numpy.allclose(a, b), + "Neq": lambda a, b: not numpy.allclose(a, b), +} + +VEC_SCALAR_OPERATION: Mapping[str, Callable[[numpy.ndarray, float], numpy.ndarray]] = { + "Mul": lambda a, b: a * b, + "Div": lambda a, b: a / b, +} + + +def _vec2_from_numpy(a: numpy.ndarray) -> Vec2: + return ( + float(a[0]), + float(a[1]), + ) + + +def _vec3_from_numpy(a: numpy.ndarray) -> Vec3: + return ( + float(a[0]), + float(a[1]), + float(a[2]), + ) + + +def _vec4_from_numpy(a: numpy.ndarray) -> Vec4: + return ( + float(a[0]), + float(a[1]), + float(a[2]), + float(a[3]), + ) + + +class Vec2UnaryOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_UNARY_OPERATIONS.keys()),), + "a": DEFAULT_VEC2, + } + } + + RETURN_TYPES = ("VEC2",) + FUNCTION = "op" + CATEGORY = "math/vec2" + + def op(self, op: str, a: Vec2) -> tuple[Vec2]: + return (_vec2_from_numpy(VEC_UNARY_OPERATIONS[op](numpy.array(a))),) + + +class Vec2ToScalarUnaryOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_TO_SCALAR_UNARY_OPERATION.keys()),), + "a": DEFAULT_VEC2, + } + } + + RETURN_TYPES = ("FLOAT",) + FUNCTION = "op" + CATEGORY = "math/vec2" + + def op(self, op: str, a: Vec2) -> tuple[float]: + return (VEC_TO_SCALAR_UNARY_OPERATION[op](numpy.array(a)),) + + +class Vec2UnaryCondition: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_UNARY_CONDITIONS.keys()),), + "a": DEFAULT_VEC2, + } + } + + RETURN_TYPES = ("BOOL",) + FUNCTION = "op" + CATEGORY = "math/vec2" + + def op(self, op: str, a: Vec2) -> tuple[bool]: + return (VEC_UNARY_CONDITIONS[op](numpy.array(a)),) + + +class Vec2BinaryOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_BINARY_OPERATIONS.keys()),), + "a": DEFAULT_VEC2, + "b": DEFAULT_VEC2, + } + } + + RETURN_TYPES = ("VEC2",) + FUNCTION = "op" + CATEGORY = "math/vec2" + + def op(self, op: str, a: Vec2, b: Vec2) -> tuple[Vec2]: + return ( + _vec2_from_numpy(VEC_BINARY_OPERATIONS[op](numpy.array(a), numpy.array(b))), + ) + + +class Vec2ToScalarBinaryOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_TO_SCALAR_BINARY_OPERATION.keys()),), + "a": DEFAULT_VEC2, + "b": DEFAULT_VEC2, + } + } + + RETURN_TYPES = ("FLOAT",) + FUNCTION = "op" + CATEGORY = "math/vec2" + + def op(self, op: str, a: Vec2, b: Vec2) -> tuple[float]: + return (VEC_TO_SCALAR_BINARY_OPERATION[op](numpy.array(a), numpy.array(b)),) + + +class Vec2BinaryCondition: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_BINARY_CONDITIONS.keys()),), + "a": DEFAULT_VEC2, + "b": DEFAULT_VEC2, + } + } + + RETURN_TYPES = ("BOOL",) + FUNCTION = "op" + CATEGORY = "math/vec2" + + def op(self, op: str, a: Vec2, b: Vec2) -> tuple[bool]: + return (VEC_BINARY_CONDITIONS[op](numpy.array(a), numpy.array(b)),) + + +class Vec2ScalarOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_SCALAR_OPERATION.keys()),), + "a": DEFAULT_VEC2, + "b": ("FLOAT",), + } + } + + RETURN_TYPES = ("VEC2",) + FUNCTION = "op" + CATEGORY = "math/vec2" + + def op(self, op: str, a: Vec2, b: float) -> tuple[Vec2]: + return (_vec2_from_numpy(VEC_SCALAR_OPERATION[op](numpy.array(a), b)),) + + +class Vec3UnaryOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_UNARY_OPERATIONS.keys()),), + "a": DEFAULT_VEC3, + } + } + + RETURN_TYPES = ("VEC3",) + FUNCTION = "op" + CATEGORY = "math/vec3" + + def op(self, op: str, a: Vec3) -> tuple[Vec3]: + return (_vec3_from_numpy(VEC_UNARY_OPERATIONS[op](numpy.array(a))),) + + +class Vec3ToScalarUnaryOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_TO_SCALAR_UNARY_OPERATION.keys()),), + "a": DEFAULT_VEC3, + } + } + + RETURN_TYPES = ("FLOAT",) + FUNCTION = "op" + CATEGORY = "math/vec3" + + def op(self, op: str, a: Vec3) -> tuple[float]: + return (VEC_TO_SCALAR_UNARY_OPERATION[op](numpy.array(a)),) + + +class Vec3UnaryCondition: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_UNARY_CONDITIONS.keys()),), + "a": DEFAULT_VEC3, + } + } + + RETURN_TYPES = ("BOOL",) + FUNCTION = "op" + CATEGORY = "math/vec3" + + def op(self, op: str, a: Vec3) -> tuple[bool]: + return (VEC_UNARY_CONDITIONS[op](numpy.array(a)),) + + +class Vec3BinaryOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_BINARY_OPERATIONS.keys()),), + "a": DEFAULT_VEC3, + "b": DEFAULT_VEC3, + } + } + + RETURN_TYPES = ("VEC3",) + FUNCTION = "op" + CATEGORY = "math/vec3" + + def op(self, op: str, a: Vec3, b: Vec3) -> tuple[Vec3]: + return ( + _vec3_from_numpy(VEC_BINARY_OPERATIONS[op](numpy.array(a), numpy.array(b))), + ) + + +class Vec3ToScalarBinaryOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_TO_SCALAR_BINARY_OPERATION.keys()),), + "a": DEFAULT_VEC3, + "b": DEFAULT_VEC3, + } + } + + RETURN_TYPES = ("FLOAT",) + FUNCTION = "op" + CATEGORY = "math/vec3" + + def op(self, op: str, a: Vec3, b: Vec3) -> tuple[float]: + return (VEC_TO_SCALAR_BINARY_OPERATION[op](numpy.array(a), numpy.array(b)),) + + +class Vec3BinaryCondition: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_BINARY_CONDITIONS.keys()),), + "a": DEFAULT_VEC3, + "b": DEFAULT_VEC3, + } + } + + RETURN_TYPES = ("BOOL",) + FUNCTION = "op" + CATEGORY = "math/vec3" + + def op(self, op: str, a: Vec3, b: Vec3) -> tuple[bool]: + return (VEC_BINARY_CONDITIONS[op](numpy.array(a), numpy.array(b)),) + + +class Vec3ScalarOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_SCALAR_OPERATION.keys()),), + "a": DEFAULT_VEC3, + "b": ("FLOAT",), + } + } + + RETURN_TYPES = ("VEC3",) + FUNCTION = "op" + CATEGORY = "math/vec3" + + def op(self, op: str, a: Vec3, b: float) -> tuple[Vec3]: + return (_vec3_from_numpy(VEC_SCALAR_OPERATION[op](numpy.array(a), b)),) + + +class Vec4UnaryOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_UNARY_OPERATIONS.keys()),), + "a": DEFAULT_VEC4, + } + } + + RETURN_TYPES = ("VEC4",) + FUNCTION = "op" + CATEGORY = "math/vec4" + + def op(self, op: str, a: Vec4) -> tuple[Vec4]: + return (_vec4_from_numpy(VEC_UNARY_OPERATIONS[op](numpy.array(a))),) + + +class Vec4ToScalarUnaryOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_TO_SCALAR_UNARY_OPERATION.keys()),), + "a": DEFAULT_VEC4, + } + } + + RETURN_TYPES = ("FLOAT",) + FUNCTION = "op" + CATEGORY = "math/vec4" + + def op(self, op: str, a: Vec4) -> tuple[float]: + return (VEC_TO_SCALAR_UNARY_OPERATION[op](numpy.array(a)),) + + +class Vec4UnaryCondition: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_UNARY_CONDITIONS.keys()),), + "a": DEFAULT_VEC4, + } + } + + RETURN_TYPES = ("BOOL",) + FUNCTION = "op" + CATEGORY = "math/vec4" + + def op(self, op: str, a: Vec4) -> tuple[bool]: + return (VEC_UNARY_CONDITIONS[op](numpy.array(a)),) + + +class Vec4BinaryOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_BINARY_OPERATIONS.keys()),), + "a": DEFAULT_VEC4, + "b": DEFAULT_VEC4, + } + } + + RETURN_TYPES = ("VEC4",) + FUNCTION = "op" + CATEGORY = "math/vec4" + + def op(self, op: str, a: Vec4, b: Vec4) -> tuple[Vec4]: + return ( + _vec4_from_numpy(VEC_BINARY_OPERATIONS[op](numpy.array(a), numpy.array(b))), + ) + + +class Vec4ToScalarBinaryOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_TO_SCALAR_BINARY_OPERATION.keys()),), + "a": DEFAULT_VEC4, + "b": DEFAULT_VEC4, + } + } + + RETURN_TYPES = ("FLOAT",) + FUNCTION = "op" + CATEGORY = "math/vec4" + + def op(self, op: str, a: Vec4, b: Vec4) -> tuple[float]: + return (VEC_TO_SCALAR_BINARY_OPERATION[op](numpy.array(a), numpy.array(b)),) + + +class Vec4BinaryCondition: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_BINARY_CONDITIONS.keys()),), + "a": DEFAULT_VEC4, + "b": DEFAULT_VEC4, + } + } + + RETURN_TYPES = ("BOOL",) + FUNCTION = "op" + CATEGORY = "math/vec4" + + def op(self, op: str, a: Vec4, b: Vec4) -> tuple[bool]: + return (VEC_BINARY_CONDITIONS[op](numpy.array(a), numpy.array(b)),) + + +class Vec4ScalarOperation: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "op": (list(VEC_SCALAR_OPERATION.keys()),), + "a": DEFAULT_VEC4, + "b": ("FLOAT",), + } + } + + RETURN_TYPES = ("VEC4",) + FUNCTION = "op" + CATEGORY = "math/vec4" + + def op(self, op: str, a: Vec4, b: float) -> tuple[Vec4]: + return (_vec4_from_numpy(VEC_SCALAR_OPERATION[op](numpy.array(a), b)),) + + +NODE_CLASS_MAPPINGS = { + "CM_Vec2UnaryOperation": Vec2UnaryOperation, + "CM_Vec2UnaryCondition": Vec2UnaryCondition, + "CM_Vec2ToScalarUnaryOperation": Vec2ToScalarUnaryOperation, + "CM_Vec2BinaryOperation": Vec2BinaryOperation, + "CM_Vec2BinaryCondition": Vec2BinaryCondition, + "CM_Vec2ToScalarBinaryOperation": Vec2ToScalarBinaryOperation, + "CM_Vec2ScalarOperation": Vec2ScalarOperation, + "CM_Vec3UnaryOperation": Vec3UnaryOperation, + "CM_Vec3UnaryCondition": Vec3UnaryCondition, + "CM_Vec3ToScalarUnaryOperation": Vec3ToScalarUnaryOperation, + "CM_Vec3BinaryOperation": Vec3BinaryOperation, + "CM_Vec3BinaryCondition": Vec3BinaryCondition, + "CM_Vec3ToScalarBinaryOperation": Vec3ToScalarBinaryOperation, + "CM_Vec3ScalarOperation": Vec3ScalarOperation, + "CM_Vec4UnaryOperation": Vec4UnaryOperation, + "CM_Vec4UnaryCondition": Vec4UnaryCondition, + "CM_Vec4ToScalarUnaryOperation": Vec4ToScalarUnaryOperation, + "CM_Vec4BinaryOperation": Vec4BinaryOperation, + "CM_Vec4BinaryCondition": Vec4BinaryCondition, + "CM_Vec4ToScalarBinaryOperation": Vec4ToScalarBinaryOperation, + "CM_Vec4ScalarOperation": Vec4ScalarOperation, +} diff --git a/src/comfymath/vec2.py b/src/comfymath/vec2.py deleted file mode 100644 index 149b2ee..0000000 --- a/src/comfymath/vec2.py +++ /dev/null @@ -1,259 +0,0 @@ -import numpy - -from abc import ABC, abstractmethod -from typing import Any, Mapping, TypeAlias - -Vec2: TypeAlias = tuple[float, float] -VEC2_ZERO = (0.0, 0.0) - - -def _vec2_from_numpy(a: numpy.ndarray) -> Vec2: - return ( - float(a[0]), - float(a[1]), - ) - - -class Vec2UnaryOperator(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return {"required": {"a": ("VEC2", {"default": VEC2_ZERO})}} - - RETURN_TYPES = ("VEC2",) - FUNCTION = "op" - - def op(self, a: Vec2) -> tuple[Vec2]: - return (_vec2_from_numpy(self.op_numpy(numpy.array(a))),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray) -> numpy.ndarray: - pass - - -class Vec2BinaryOperator(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { - "required": { - "a": ("VEC2", {"default": VEC2_ZERO}), - "b": ("VEC2", {"default": VEC2_ZERO}), - } - } - - RETURN_TYPES = ("VEC2",) - FUNCTION = "op" - - def op(self, a: Vec2, b: Vec2) -> tuple[Vec2]: - return (_vec2_from_numpy(self.op_numpy(numpy.array(a), numpy.array(b))),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: - pass - - -class Vec2UnaryQuery(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return {"required": {"a": ("VEC2", {"default": VEC2_ZERO})}} - - RETURN_TYPES = ("INT",) - FUNCTION = "op" - - def op(self, a: Vec2) -> tuple[int]: - return (self.op_numpy(numpy.array(a)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray) -> int: - pass - - -class Vec2BinaryQuery(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { - "required": { - "a": ("VEC2", {"default": VEC2_ZERO}), - "b": ("VEC2", {"default": VEC2_ZERO}), - } - } - - RETURN_TYPES = ("INT",) - FUNCTION = "op" - - def op(self, a: Vec2, b: Vec2) -> tuple[int]: - return (self.op_numpy(numpy.array(a), numpy.array(b)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> int: - pass - - -class Vec2ToScalarUnary(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return {"required": {"a": ("VEC2", {"default": VEC2_ZERO})}} - - RETURN_TYPES = ("FLOAT",) - FUNCTION = "op" - - def op(self, a: Vec2) -> tuple[float]: - return (self.op_numpy(numpy.array(a)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray) -> float: - pass - - -class Vec2ToScalarBinary(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { - "required": { - "a": ("VEC2", {"default": VEC2_ZERO}), - "b": ("VEC2", {"default": VEC2_ZERO}), - } - } - - RETURN_TYPES = ("FLOAT",) - FUNCTION = "op" - - def op(self, a: Vec2, b: Vec2) -> tuple[float]: - return (self.op_numpy(numpy.array(a), numpy.array(b)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> float: - pass - - -class Vec2ScalarOperation(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { - "required": { - "a": ("VEC2", {"default": VEC2_ZERO}), - "b": ("FLOAT", {"default": 0.0}), - } - } - - RETURN_TYPES = ("VEC2",) - FUNCTION = "op" - - def op(self, a: Vec2, b: float) -> tuple[Vec2]: - return (_vec2_from_numpy(self.op_numpy(numpy.array(a), b)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray, b: float) -> numpy.ndarray: - pass - - -class Vec2Add(Vec2BinaryOperator): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: - return a + b - - CATEGORY = "math/vec2" - - -class Vec2Sub(Vec2BinaryOperator): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: - return a - b - - CATEGORY = "math/vec2" - - -class Vec2Dot(Vec2ToScalarBinary): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> float: - return a.dot(b).astype(float) - - CATEGORY = "math/vec2" - - -class Vec2Cross(Vec2BinaryOperator): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: - return numpy.cross(a, b) - - CATEGORY = "math/vec2" - - -class Vec2Eq(Vec2BinaryQuery): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> int: - return int(a == b) - - CATEGORY = "math/vec2" - - -class Vec2Ne(Vec2BinaryQuery): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> int: - return int(a != b) - - CATEGORY = "math/vec2" - - -class Vec2Neg(Vec2UnaryOperator): - def op_numpy(self, a: numpy.ndarray) -> numpy.ndarray: - return -a - - CATEGORY = "math/vec2" - - -class Vec2Norm(Vec2ToScalarUnary): - def op_numpy(self, a: numpy.ndarray) -> float: - return numpy.linalg.norm(a).astype(float) - - CATEGORY = "math/vec2" - - -class Vec2Normalize(Vec2UnaryOperator): - def op_numpy(self, a: numpy.ndarray) -> numpy.ndarray: - return a / numpy.linalg.norm(a) - - CATEGORY = "math/vec2" - - -class Vec2Distance(Vec2ToScalarBinary): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> float: - return numpy.linalg.norm(a - b).astype(float) - - CATEGORY = "math/vec2" - - -class Vec2ScalarMul(Vec2ScalarOperation): - def op_numpy(self, a: numpy.ndarray, b: float) -> numpy.ndarray: - return a * b - - CATEGORY = "math/vec2" - - -NODE_CLASS_MAPPINGS = { - "Vec2Add": Vec2Add, - "Vec2Sub": Vec2Sub, - "Vec2Dot": Vec2Dot, - "Vec2Cross": Vec2Cross, - "Vec2Eq": Vec2Eq, - "Vec2Ne": Vec2Ne, - "Vec2Neg": Vec2Neg, - "Vec2Norm": Vec2Norm, - "Vec2Normalize": Vec2Normalize, - "Vec2Distance": Vec2Distance, - "Vec2ScalarMul": Vec2ScalarMul, -} diff --git a/src/comfymath/vec3.py b/src/comfymath/vec3.py deleted file mode 100644 index 434fbc9..0000000 --- a/src/comfymath/vec3.py +++ /dev/null @@ -1,260 +0,0 @@ -import numpy - -from abc import ABC, abstractmethod -from typing import Any, Mapping, TypeAlias - -Vec3: TypeAlias = tuple[float, float, float] -VEC3_ZERO = (0.0, 0.0, 0.0) - - -def _vec3_from_numpy(a: numpy.ndarray) -> Vec3: - return ( - float(a[0]), - float(a[1]), - float(a[2]), - ) - - -class Vec3UnaryOperator(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return {"required": {"a": ("VEC3", {"default": VEC3_ZERO})}} - - RETURN_TYPES = ("VEC3",) - FUNCTION = "op" - - def op(self, a: Vec3) -> tuple[Vec3]: - return (_vec3_from_numpy(self.op_numpy(numpy.array(a))),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray) -> numpy.ndarray: - pass - - -class Vec3BinaryOperator(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { - "required": { - "a": ("VEC3", {"default": VEC3_ZERO}), - "b": ("VEC3", {"default": VEC3_ZERO}), - } - } - - RETURN_TYPES = ("VEC3",) - FUNCTION = "op" - - def op(self, a: Vec3, b: Vec3) -> tuple[Vec3]: - return (_vec3_from_numpy(self.op_numpy(numpy.array(a), numpy.array(b))),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: - pass - - -class Vec3UnaryQuery(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return {"required": {"a": ("VEC3", {"default": VEC3_ZERO})}} - - RETURN_TYPES = ("INT",) - FUNCTION = "op" - - def op(self, a: Vec3) -> tuple[int]: - return (self.op_numpy(numpy.array(a)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray) -> int: - pass - - -class Vec3BinaryQuery(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { - "required": { - "a": ("VEC3", {"default": VEC3_ZERO}), - "b": ("VEC3", {"default": VEC3_ZERO}), - } - } - - RETURN_TYPES = ("INT",) - FUNCTION = "op" - - def op(self, a: Vec3, b: Vec3) -> tuple[int]: - return (self.op_numpy(numpy.array(a), numpy.array(b)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> int: - pass - - -class Vec3ToScalarUnary(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return {"required": {"a": ("VEC2", {"default": VEC3_ZERO})}} - - RETURN_TYPES = ("FLOAT",) - FUNCTION = "op" - - def op(self, a: Vec3) -> tuple[float]: - return (self.op_numpy(numpy.array(a)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray) -> float: - pass - - -class Vec3ToScalarBinary(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { - "required": { - "a": ("VEC2", {"default": VEC3_ZERO}), - "b": ("VEC2", {"default": VEC3_ZERO}), - } - } - - RETURN_TYPES = ("FLOAT",) - FUNCTION = "op" - - def op(self, a: Vec3, b: Vec3) -> tuple[float]: - return (self.op_numpy(numpy.array(a), numpy.array(b)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> float: - pass - - -class Vec3ScalarOperation(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { - "required": { - "a": ("VEC3", {"default": VEC3_ZERO}), - "b": ("FLOAT", {"default": 0.0}), - } - } - - RETURN_TYPES = ("VEC3",) - FUNCTION = "op" - - def op(self, a: Vec3, b: float) -> tuple[Vec3]: - return (_vec3_from_numpy(self.op_numpy(numpy.array(a), b)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray, b: float) -> numpy.ndarray: - pass - - -class Vec3Add(Vec3BinaryOperator): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: - return a + b - - CATEGORY = "math/vec3" - - -class Vec3Sub(Vec3BinaryOperator): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: - return a - b - - CATEGORY = "math/vec3" - - -class Vec3Dot(Vec3ToScalarBinary): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> float: - return a.dot(b).astype(float) - - CATEGORY = "math/vec3" - - -class Vec3Cross(Vec3BinaryOperator): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: - return numpy.cross(a, b) - - CATEGORY = "math/vec3" - - -class Vec3Eq(Vec3BinaryQuery): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> int: - return int(a == b) - - CATEGORY = "math/vec3" - - -class Vec3Ne(Vec3BinaryQuery): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> int: - return int(a != b) - - CATEGORY = "math/vec3" - - -class Vec3Neg(Vec3UnaryOperator): - def op_numpy(self, a: numpy.ndarray) -> numpy.ndarray: - return -a - - CATEGORY = "math/vec3" - - -class Vec3Norm(Vec3ToScalarUnary): - def op_numpy(self, a: numpy.ndarray) -> float: - return numpy.linalg.norm(a).astype(float) - - CATEGORY = "math/vec3" - - -class Vec3Normalize(Vec3UnaryOperator): - def op_numpy(self, a: numpy.ndarray) -> numpy.ndarray: - return a / numpy.linalg.norm(a) - - CATEGORY = "math/vec3" - - -class Vec3Distance(Vec3ToScalarBinary): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> float: - return numpy.linalg.norm(a - b).astype(float) - - CATEGORY = "math/vec3" - - -class Vec3ScalarMul(Vec3ScalarOperation): - def op_numpy(self, a: numpy.ndarray, b: float) -> numpy.ndarray: - return a * b - - CATEGORY = "math/vec3" - - -NODE_CLASS_MAPPINGS = { - "Vec3Add": Vec3Add, - "Vec3Sub": Vec3Sub, - "Vec3Dot": Vec3Dot, - "Vec3Cross": Vec3Cross, - "Vec3Eq": Vec3Eq, - "Vec3Ne": Vec3Ne, - "Vec3Neg": Vec3Neg, - "Vec3Norm": Vec3Norm, - "Vec3Normalize": Vec3Normalize, - "Vec3Distance": Vec3Distance, - "Vec3ScalarMul": Vec3ScalarMul, -} diff --git a/src/comfymath/vec4.py b/src/comfymath/vec4.py deleted file mode 100644 index 22de5dd..0000000 --- a/src/comfymath/vec4.py +++ /dev/null @@ -1,261 +0,0 @@ -import numpy - -from abc import ABC, abstractmethod -from typing import Any, Mapping, TypeAlias - -Vec4: TypeAlias = tuple[float, float, float, float] -VEC4_ZERO = (0.0, 0.0, 0.0, 0.0) - - -def _vec4_from_numpy(a: numpy.ndarray) -> Vec4: - return ( - float(a[0]), - float(a[1]), - float(a[2]), - float(a[3]), - ) - - -class Vec4UnaryOperator(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return {"required": {"a": ("VEC4", {"default": VEC4_ZERO})}} - - RETURN_TYPES = ("VEC4",) - FUNCTION = "op" - - def op(self, a: Vec4) -> tuple[Vec4]: - return (_vec4_from_numpy(self.op_numpy(numpy.array(a))),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray) -> numpy.ndarray: - pass - - -class Vec4BinaryOperator(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { - "required": { - "a": ("VEC4", {"default": VEC4_ZERO}), - "b": ("VEC4", {"default": VEC4_ZERO}), - } - } - - RETURN_TYPES = ("VEC4",) - FUNCTION = "op" - - def op(self, a: Vec4, b: Vec4) -> tuple[Vec4]: - return (_vec4_from_numpy(self.op_numpy(numpy.array(a), numpy.array(b))),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: - pass - - -class Vec4UnaryQuery(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return {"required": {"a": ("VEC4", {"default": VEC4_ZERO})}} - - RETURN_TYPES = ("INT",) - FUNCTION = "op" - - def op(self, a: Vec4) -> tuple[int]: - return (self.op_numpy(numpy.array(a)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray) -> int: - pass - - -class Vec4BinaryQuery(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { - "required": { - "a": ("VEC4", {"default": VEC4_ZERO}), - "b": ("VEC4", {"default": VEC4_ZERO}), - } - } - - RETURN_TYPES = ("INT",) - FUNCTION = "op" - - def op(self, a: Vec4, b: Vec4) -> tuple[int]: - return (self.op_numpy(numpy.array(a), numpy.array(b)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> int: - pass - - -class Vec4ToScalarUnary(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return {"required": {"a": ("VEC4", {"default": VEC4_ZERO})}} - - RETURN_TYPES = ("FLOAT",) - FUNCTION = "op" - - def op(self, a: Vec4) -> tuple[float]: - return (self.op_numpy(numpy.array(a)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray) -> float: - pass - - -class Vec4ToScalarBinary(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { - "required": { - "a": ("VEC4", {"default": VEC4_ZERO}), - "b": ("VEC4", {"default": VEC4_ZERO}), - } - } - - RETURN_TYPES = ("FLOAT",) - FUNCTION = "op" - - def op(self, a: Vec4, b: Vec4) -> tuple[float]: - return (self.op_numpy(numpy.array(a), numpy.array(b)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> float: - pass - - -class Vec4ScalarOperation(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { - "required": { - "a": ("VEC4", {"default": VEC4_ZERO}), - "b": ("FLOAT", {"default": 0.0}), - } - } - - RETURN_TYPES = ("VEC4",) - FUNCTION = "op" - - def op(self, a: Vec4, b: float) -> tuple[Vec4]: - return (_vec4_from_numpy(self.op_numpy(numpy.array(a), b)),) - - @abstractmethod - def op_numpy(self, a: numpy.ndarray, b: float) -> numpy.ndarray: - pass - - -class Vec4Add(Vec4BinaryOperator): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: - return a + b - - CATEGORY = "math/vec4" - - -class Vec4Sub(Vec4BinaryOperator): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: - return a - b - - CATEGORY = "math/vec4" - - -class Vec4Dot(Vec4ToScalarBinary): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> float: - return a.dot(b).astype(float) - - CATEGORY = "math/vec4" - - -class Vec4Cross(Vec4BinaryOperator): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: - return numpy.cross(a, b) - - CATEGORY = "math/vec4" - - -class Vec4Eq(Vec4BinaryQuery): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> int: - return int(a == b) - - CATEGORY = "math/vec4" - - -class Vec4Ne(Vec4BinaryQuery): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> int: - return int(a != b) - - CATEGORY = "math/vec4" - - -class Vec4Neg(Vec4UnaryOperator): - def op_numpy(self, a: numpy.ndarray) -> numpy.ndarray: - return -a - - CATEGORY = "math/vec4" - - -class Vec4Norm(Vec4ToScalarUnary): - def op_numpy(self, a: numpy.ndarray) -> float: - return numpy.linalg.norm(a).astype(float) - - CATEGORY = "math/vec4" - - -class Vec4Normalize(Vec4UnaryOperator): - def op_numpy(self, a: numpy.ndarray) -> numpy.ndarray: - return a / numpy.linalg.norm(a) - - CATEGORY = "math/vec4" - - -class Vec4Distance(Vec4ToScalarBinary): - def op_numpy(self, a: numpy.ndarray, b: numpy.ndarray) -> float: - return numpy.linalg.norm(a - b).astype(float) - - CATEGORY = "math/vec4" - - -class Vec4ScalarMul(Vec4ScalarOperation): - def op_numpy(self, a: numpy.ndarray, b: float) -> numpy.ndarray: - return a * b - - CATEGORY = "math/vec4" - - -NODE_CLASS_MAPPINGS = { - "Vec4Add": Vec4Add, - "Vec4Sub": Vec4Sub, - "Vec4Dot": Vec4Dot, - "Vec4Cross": Vec4Cross, - "Vec4Eq": Vec4Eq, - "Vec4Ne": Vec4Ne, - "Vec4Neg": Vec4Neg, - "Vec4Norm": Vec4Norm, - "Vec4Normalize": Vec4Normalize, - "Vec4Distance": Vec4Distance, - "Vec4ScalarMul": Vec4ScalarMul, -}