misc_nodes.py 71 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825
  1. from .node_container_common import *
  2. from .base_definitions import MantisNode, NodeSocket, FLOAT_EPSILON
  3. from .xForm_containers import xFormArmature, xFormBone
  4. from .misc_nodes_socket_templates import *
  5. from math import pi, tau
  6. def TellClasses():
  7. return [
  8. # utility
  9. InputFloat,
  10. InputIntNode,
  11. InputVector,
  12. InputBoolean,
  13. InputBooleanThreeTuple,
  14. InputRotationOrder,
  15. InputTransformSpace,
  16. InputString,
  17. InputMatrix,
  18. InputExistingGeometryObject,
  19. InputExistingGeometryData,
  20. UtilityGeometryOfXForm,
  21. UtilityNameOfXForm,
  22. UtilityPointFromCurve,
  23. UtilityMatrixFromCurve,
  24. UtilityMatricesFromCurve,
  25. UtilityNumberOfCurveSegments,
  26. UtilityMatrixFromCurveSegment,
  27. UtilityGetCurvePoint,
  28. UtilityGetNearestFactorOnCurve,
  29. UtilityKDChoosePoint,
  30. UtilityKDChooseXForm,
  31. UtilityMetaRig,
  32. UtilityBoneProperties,
  33. UtilityDriverVariable,
  34. UtilityDriver,
  35. UtilityFCurve,
  36. UtilityKeyframe,
  37. UtilitySwitch,
  38. UtilityCombineThreeBool,
  39. UtilityCombineVector,
  40. UtilitySeparateVector,
  41. UtilityCatStrings,
  42. UtilityGetBoneLength,
  43. UtilityPointFromBoneMatrix,
  44. UtilitySetBoneLength,
  45. UtilityMatrixSetLocation,
  46. UtilityMatrixGetLocation,
  47. UtilityMatrixFromXForm,
  48. UtilityAxesFromMatrix,
  49. UtilityBoneMatrixHeadTailFlip,
  50. UtilityMatrixTransform,
  51. UtilityMatrixInvert,
  52. UtilityMatrixCompose,
  53. UtilityMatrixAlignRoll,
  54. UtilityTransformationMatrix,
  55. UtilityIntToString,
  56. UtilityArrayGet,
  57. UtilityArrayLength,
  58. UtilitySetBoneMatrixTail,
  59. # Control flow switches
  60. UtilityCompare,
  61. UtilityChoose,
  62. # useful NoOp:
  63. UtilityPrint,
  64. ]
  65. def matrix_from_head_tail(head, tail, normal=None):
  66. from mathutils import Vector, Matrix
  67. if normal is None:
  68. rotation = Vector((0,1,0)).rotation_difference((tail-head).normalized()).to_matrix()
  69. m= Matrix.LocRotScale(head, rotation, None)
  70. else: # construct a basis matrix
  71. m = Matrix.Identity(3)
  72. axis = (tail-head).normalized()
  73. conormal = axis.cross(normal)
  74. m[0]=conormal
  75. m[1]=axis
  76. m[2]=normal
  77. m = m.transposed().to_4x4()
  78. m.translation=head.copy()
  79. m[3][3]=(tail-head).length
  80. return m
  81. def get_mesh_from_curve(curve_name : str, execution_id : str, bContext, ribbon=True):
  82. from bpy import data
  83. curve = data.objects.get(curve_name)
  84. assert curve.type == 'CURVE', f"ERROR: object is not a curve: {curve_name}"
  85. from .utilities import mesh_from_curve
  86. curve_type='ribbon' if ribbon else 'wire'
  87. m_name = curve_name+'.'+str(hash(curve_name+'.'+curve_type+'.'+execution_id))
  88. if not (m := data.meshes.get(m_name)):
  89. m = mesh_from_curve(curve, bContext, ribbon)
  90. m.name = m_name
  91. return m
  92. def cleanup_curve(curve_name : str, execution_id : str) -> None:
  93. import bpy
  94. curve = bpy.data.objects.get(curve_name)
  95. m_name = curve_name+'.'+str(hash(curve.name+'.'+ execution_id))
  96. if (mesh := bpy.data.meshes.get(m_name)):
  97. bpy.data.meshes.remove(mesh)
  98. def kd_find(node, points, ref_pt, num_points):
  99. if num_points == 0:
  100. raise RuntimeError(f"Cannot find 0 points for {node}")
  101. from mathutils import kdtree
  102. kd = kdtree.KDTree(len(points))
  103. for i, pt in enumerate(points):
  104. try:
  105. kd.insert(pt, i)
  106. except (TypeError, ValueError) as e:
  107. prRed(f"Cannot get point from for {node}")
  108. raise e
  109. kd.balance()
  110. try:
  111. if num_points == 1: # make it a list to keep it consistent with
  112. result = [kd.find(ref_pt)] # find_n which returns a list
  113. else:
  114. result = kd.find_n(ref_pt, num_points)
  115. # the result of kd.find has some other stuff we don't care about
  116. except (TypeError, ValueError) as e:
  117. prRed(f"Reference Point {ref_pt} invalid for {node}")
  118. raise e
  119. return result
  120. def array_choose_relink(node, indices, array_input, output, ):
  121. """
  122. Used to choose the correct link to send out of an array-choose node.
  123. """
  124. keep_links = []
  125. for index in indices:
  126. l = node.inputs[array_input].links[index]
  127. keep_links.append(l)
  128. for link in node.outputs[output].links:
  129. to_node = link.to_node
  130. for l in keep_links:
  131. l.from_node.outputs[l.from_socket].connect(to_node, link.to_socket)
  132. link.die()
  133. def array_choose_data(node, data, output):
  134. """
  135. Used to choose the correct data to send out of an array-choose node.
  136. """
  137. # We need to make new outputs and link from each one based on the data in the array...
  138. node.outputs.init_sockets([output+"."+str(i).zfill(4) for i in range(len(data)) ])
  139. for i, data_item in enumerate(data):
  140. node.parameters[output+"."+str(i).zfill(4)] = data_item
  141. for link in node.outputs[output].links:
  142. to_node = link.to_node
  143. for i in range(len(data)):
  144. # Make a link from the new output.
  145. node.outputs[output+"."+str(i).zfill(4)].connect(to_node, link.to_socket)
  146. link.die()
  147. def zero_radius_error_message(node, curve):
  148. return f"ERROR: cannot get matrix from zero-radius curve point "\
  149. "in curve object: {curve.name} for node: {node}. "\
  150. "This is a limitation of Mantis (For now). Please inspect the curve and ensure "\
  151. "that each curve point has a radius greater than 0. Sometimes, this error is " \
  152. "caused by drivers. "
  153. #*#-------------------------------#++#-------------------------------#*#
  154. # B A S E C L A S S E S
  155. #*#-------------------------------#++#-------------------------------#*#
  156. class SimpleInputNode(MantisNode):
  157. def __init__(self, signature, base_tree, socket_templates=[]):
  158. super().__init__(signature, base_tree, socket_templates)
  159. self.node_type = 'UTILITY'
  160. self.prepared, self.executed = True, True
  161. #*#-------------------------------#++#-------------------------------#*#
  162. # U T I L I T Y N O D E S
  163. #*#-------------------------------#++#-------------------------------#*#
  164. class InputFloat(SimpleInputNode):
  165. '''A node representing float input'''
  166. def __init__(self, signature, base_tree):
  167. super().__init__(signature, base_tree)
  168. outputs = ["Float Input"]
  169. self.outputs.init_sockets(outputs)
  170. self.init_parameters()
  171. class InputIntNode(SimpleInputNode):
  172. '''A node representing integer input'''
  173. def __init__(self, signature, base_tree):
  174. super().__init__(signature, base_tree)
  175. outputs = ["Integer"]
  176. self.outputs.init_sockets(outputs)
  177. self.init_parameters()
  178. class InputVector(SimpleInputNode):
  179. '''A node representing vector input'''
  180. def __init__(self, signature, base_tree):
  181. super().__init__(signature, base_tree)
  182. outputs = [""]
  183. self.outputs.init_sockets(outputs)
  184. self.init_parameters()
  185. class InputBoolean(SimpleInputNode):
  186. '''A node representing boolean input'''
  187. def __init__(self, signature, base_tree):
  188. super().__init__(signature, base_tree)
  189. outputs = [""]
  190. self.outputs.init_sockets(outputs)
  191. self.init_parameters()
  192. class InputBooleanThreeTuple(SimpleInputNode):
  193. '''A node representing a tuple of three booleans'''
  194. def __init__(self, signature, base_tree):
  195. super().__init__(signature, base_tree)
  196. outputs = [""]
  197. self.outputs.init_sockets(outputs)
  198. self.init_parameters()
  199. class InputRotationOrder(SimpleInputNode):
  200. '''A node representing string input for rotation order'''
  201. def __init__(self, signature, base_tree):
  202. super().__init__(signature, base_tree)
  203. outputs = [""]
  204. self.outputs.init_sockets(outputs)
  205. self.init_parameters()
  206. class InputTransformSpace(SimpleInputNode):
  207. '''A node representing string input for transform space'''
  208. def __init__(self, signature, base_tree):
  209. super().__init__(signature, base_tree)
  210. outputs = [""]
  211. self.outputs.init_sockets(outputs)
  212. self.init_parameters()
  213. def evaluate_input(self, input_name):
  214. return self.parameters[""]
  215. class InputString(SimpleInputNode):
  216. '''A node representing string input'''
  217. def __init__(self, signature, base_tree):
  218. super().__init__(signature, base_tree)
  219. outputs = [""]
  220. self.outputs.init_sockets(outputs)
  221. self.init_parameters()
  222. class InputMatrix(SimpleInputNode):
  223. '''A node representing axis-angle quaternion input'''
  224. def __init__(self, signature, base_tree):
  225. super().__init__(signature, base_tree)
  226. outputs = ["Matrix",]
  227. self.outputs.init_sockets(outputs)
  228. self.init_parameters()
  229. class UtilityMatrixFromCurve(MantisNode):
  230. '''Get a matrix from a curve'''
  231. def __init__(self, signature, base_tree):
  232. super().__init__(signature, base_tree, MatrixFromCurveSockets)
  233. self.init_parameters()
  234. self.node_type = "UTILITY"
  235. def bPrepare(self, bContext = None,):
  236. from mathutils import Matrix
  237. import bpy
  238. mat = Matrix.Identity(4)
  239. curve_name = self.evaluate_input("Curve")
  240. curve = bpy.data.objects.get(curve_name)
  241. if not curve:
  242. prRed(f"WARN: No curve found for {self}. Using an identity matrix instead.")
  243. mat[3][3] = 1.0
  244. elif curve.type != "CURVE":
  245. prRed(f"WARN: Object {curve.name} is not a curve. Using an identity matrix instead.")
  246. mat[3][3] = 1.0
  247. else:
  248. if bContext is None: bContext = bpy.context # is this wise?
  249. m = get_mesh_from_curve(curve.name, self.base_tree.execution_id, bContext)
  250. from .utilities import data_from_ribbon_mesh
  251. #
  252. num_divisions = self.evaluate_input("Total Divisions")
  253. if num_divisions <= 0:
  254. raise GraphError("The number of divisions in the curve must be 1 or greater.")
  255. m_index = self.evaluate_input("Matrix Index")
  256. if m_index >= num_divisions:
  257. prRed(m_index, num_divisions)
  258. raise GraphError(f"{self} tried to get a matrix-index greater the total number of divisions."
  259. "The matrix index starts at 0. You're probably off by +1.")
  260. spline_index = self.evaluate_input("Spline Index")
  261. if spline_index > len(curve.data.splines)-1:
  262. raise GraphError(f"{self} is attempting to read from a spline in {curve.name} that does not exist."
  263. " Try and reduce the value of Spline Index.")
  264. splines_factors = [ [] for i in range (spline_index)]
  265. factors = [1/num_divisions*m_index, 1/num_divisions*(m_index+1)]
  266. splines_factors.append(factors)
  267. data = data_from_ribbon_mesh(m, splines_factors, curve.matrix_world)
  268. if data[spline_index][1][0] < FLOAT_EPSILON: # radius is None:
  269. raise RuntimeError(zero_radius_error_message(self, curve))
  270. head=data[spline_index][0][0]
  271. tail= data[spline_index][0][1]
  272. axis = (tail-head).normalized()
  273. if axis.length_squared < FLOAT_EPSILON:
  274. raise RuntimeError(f"Failed to read the curve {curve.name}.")
  275. normal=data[spline_index][2][0]
  276. # make sure the normal is perpendicular to the tail
  277. from .utilities import make_perpendicular
  278. normal = make_perpendicular(axis, normal)
  279. mat = matrix_from_head_tail(head, tail, normal)
  280. # this is in world space... let's just convert it back
  281. mat.translation = head - curve.location
  282. # TODO HACK TODO
  283. # all the nodes should work in world-space, and it should be the responsibility
  284. # of the xForm node to convert!
  285. self.parameters["Matrix"] = mat
  286. self.prepared = True
  287. self.executed = True
  288. def bFinalize(self, bContext=None):
  289. cleanup_curve(self.evaluate_input("Curve"), self.base_tree.execution_id)
  290. class UtilityPointFromCurve(MantisNode):
  291. '''Get a point from a curve'''
  292. def __init__(self, signature, base_tree):
  293. super().__init__(signature, base_tree, PointFromCurveSockets)
  294. self.init_parameters()
  295. self.node_type = "UTILITY"
  296. def bPrepare(self, bContext = None,):
  297. import bpy
  298. curve = bpy.data.objects.get(self.evaluate_input("Curve"))
  299. if not curve:
  300. raise RuntimeError(f"No curve found for {self}.")
  301. elif curve.type != "CURVE":
  302. raise GraphError(f"ERROR: Object {curve.name} is not a curve.")
  303. else:
  304. if bContext is None: bContext = bpy.context # is this wise?
  305. m = get_mesh_from_curve(curve.name, self.base_tree.execution_id, bContext)
  306. from .utilities import data_from_ribbon_mesh
  307. #
  308. num_divisions = 1
  309. spline_index = self.evaluate_input("Spline Index")
  310. splines_factors = [ [] for i in range (spline_index)]
  311. factors = [self.evaluate_input("Factor")]
  312. splines_factors.append(factors)
  313. data = data_from_ribbon_mesh(m, splines_factors, curve.matrix_world)
  314. p = data[spline_index][0][0] - curve.location
  315. self.parameters["Point"] = p
  316. self.prepared, self.executed = True, True
  317. def bFinalize(self, bContext=None):
  318. cleanup_curve(self.evaluate_input("Curve"), self.base_tree.execution_id)
  319. class UtilityMatricesFromCurve(MantisNode):
  320. '''Get matrices from a curve'''
  321. def __init__(self, signature, base_tree):
  322. super().__init__(signature, base_tree, MatricesFromCurveSockets)
  323. self.init_parameters()
  324. self.node_type = "UTILITY"
  325. def bPrepare(self, bContext = None,):
  326. import time
  327. # start_time = time.time()
  328. #
  329. from mathutils import Matrix
  330. import bpy
  331. m = Matrix.Identity(4)
  332. curve_name = self.evaluate_input("Curve")
  333. curve = bpy.data.objects.get(curve_name)
  334. if not curve:
  335. prRed(f"WARN: No curve found for {self}. Using an identity matrix instead.")
  336. m[3][3] = 1.0
  337. elif curve.type != "CURVE":
  338. prRed(f"WARN: Object {curve.name} is not a curve. Using an identity matrix instead.")
  339. m[3][3] = 1.0
  340. else:
  341. if bContext is None: bContext = bpy.context # is this wise?
  342. mesh = get_mesh_from_curve(curve.name, self.base_tree.execution_id, bContext)
  343. from .utilities import data_from_ribbon_mesh
  344. num_divisions = self.evaluate_input("Total Divisions")
  345. spline_index = self.evaluate_input("Spline Index")
  346. splines_factors = [ [] for i in range (spline_index)]
  347. factors = [0.0] + [(1/num_divisions*(i+1)) for i in range(num_divisions)]
  348. splines_factors.append(factors)
  349. data = data_from_ribbon_mesh(mesh, splines_factors, curve.matrix_world)
  350. # [spline_index][points,tangents,normals][datapoint_index]
  351. from .utilities import make_perpendicular
  352. matrices=[]
  353. for i in range(num_divisions):
  354. if data[spline_index][1][i] < FLOAT_EPSILON: # radius is None:
  355. raise RuntimeError(zero_radius_error_message(self, curve))
  356. m = matrix_from_head_tail (
  357. data[spline_index][0][i], data[spline_index][0][i+1],
  358. make_perpendicular((data[spline_index][0][i+1]-data[spline_index][0][i]).normalized(), data[spline_index][2][i]),)
  359. m.translation = data[spline_index][0][i] - curve.location
  360. matrices.append(m)
  361. for link in self.outputs["Matrices"].links:
  362. for i, m in enumerate(matrices):
  363. name = "Matrix"+str(i).zfill(4)
  364. if not (out := self.outputs.get(name)): # reuse them if there are multiple links.
  365. out = self.outputs[name] = NodeSocket(name = name, node=self)
  366. c = out.connect(link.to_node, link.to_socket)
  367. # prOrange(c)
  368. self.parameters[name] = m
  369. # print (mesh)
  370. link.die()
  371. self.prepared = True
  372. self.executed = True
  373. # prGreen(f"Matrices from curves took {time.time() - start_time} seconds.")
  374. def bFinalize(self, bContext=None):
  375. import bpy
  376. curve_name = self.evaluate_input("Curve")
  377. curve = bpy.data.objects.get(curve_name)
  378. m_name = curve.name+'.'+self.base_tree.execution_id
  379. if (mesh := bpy.data.meshes.get(m_name)):
  380. prGreen(f"Freeing mesh data {m_name}...")
  381. bpy.data.meshes.remove(mesh)
  382. class UtilityNumberOfCurveSegments(MantisNode):
  383. def __init__(self, signature, base_tree):
  384. super().__init__(signature, base_tree)
  385. inputs = [
  386. "Curve" ,
  387. "Spline Index" ,
  388. ]
  389. outputs = [
  390. "Number of Segments" ,
  391. ]
  392. self.inputs.init_sockets(inputs)
  393. self.outputs.init_sockets(outputs)
  394. self.init_parameters()
  395. self.node_type = "UTILITY"
  396. def bPrepare(self, bContext = None,):
  397. import bpy
  398. curve_name = self.evaluate_input("Curve")
  399. curve = bpy.data.objects.get(curve_name)
  400. spline = curve.data.splines[self.evaluate_input("Spline Index")]
  401. if spline.type == "BEZIER":
  402. self.parameters["Number of Segments"] = len(spline.bezier_points)-1
  403. else:
  404. self.parameters["Number of Segments"] = len(spline.points)-1
  405. self.prepared = True
  406. self.executed = True
  407. class UtilityMatrixFromCurveSegment(MantisNode):
  408. def __init__(self, signature, base_tree):
  409. super().__init__(signature, base_tree, MatrixFromCurveSegmentSockets)
  410. self.init_parameters()
  411. self.node_type = "UTILITY"
  412. def bPrepare(self, bContext = None,):
  413. import bpy
  414. curve = bpy.data.objects.get(self.evaluate_input("Curve"))
  415. if not curve:
  416. raise RuntimeError(f"No curve found for {self}.")
  417. elif curve.type != "CURVE":
  418. raise GraphError(f"ERROR: Object {curve.name} is not a curve.")
  419. else:
  420. if bContext is None: bContext = bpy.context # is this wise?
  421. m = get_mesh_from_curve(curve.name, self.base_tree.execution_id, bContext)
  422. from .utilities import data_from_ribbon_mesh
  423. # this section is dumb, but it is because the data_from_ribbon_mesh
  424. # function is designed to pull data from many splines at once (for optimization)
  425. # so we have to give it empty splines for each one we skip.
  426. # TODO: Refactor this to make it so I can select spline index
  427. spline_index = self.evaluate_input("Spline Index")
  428. spline = curve.data.splines[spline_index]
  429. splines_factors = [ [] for i in range (spline_index)]
  430. factors = [0.0]
  431. points = spline.bezier_points if spline.type == 'BEZIER' else spline.points
  432. total_length=0.0
  433. for i in range(len(points)-1):
  434. total_length+= (seg_length := (points[i+1].co - points[i].co).length)
  435. factors.append(seg_length)
  436. prev_length = 0.0
  437. for i in range(len(factors)):
  438. factors[i] = prev_length+factors[i]/total_length
  439. prev_length=factors[i]
  440. # Why does this happen? Floating point error?
  441. if factors[i]>1.0: factors[i] = 1.0
  442. splines_factors.append(factors)
  443. #
  444. data = data_from_ribbon_mesh(m, splines_factors, curve.matrix_world)
  445. segment_index = self.evaluate_input("Segment Index")
  446. if data[spline_index][1][segment_index] < FLOAT_EPSILON: # radius is None:
  447. raise RuntimeError(zero_radius_error_message(self, curve))
  448. head=data[spline_index][0][segment_index]
  449. tail= data[spline_index][0][segment_index+1]
  450. axis = (tail-head).normalized()
  451. normal=data[spline_index][2][segment_index]
  452. # make sure the normal is perpendicular to the tail
  453. from .utilities import make_perpendicular
  454. normal = make_perpendicular(axis, normal)
  455. m = matrix_from_head_tail(head, tail, normal)
  456. m.translation = head - curve.location
  457. self.parameters["Matrix"] = m
  458. self.prepared, self.executed = True, True
  459. def bFinalize(self, bContext=None):
  460. cleanup_curve(self.evaluate_input("Curve"), self.base_tree.execution_id)
  461. class UtilityGetCurvePoint(MantisNode):
  462. def __init__(self, signature, base_tree):
  463. super().__init__(signature, base_tree, GetCurvePointSockets)
  464. self.init_parameters()
  465. self.node_type = "UTILITY"
  466. def bPrepare(self, bContext=None):
  467. import bpy
  468. my_curve=self.evaluate_input("Curve")
  469. if my_curve is None:
  470. raise RuntimeError(f"Error with curve name for {self}, {my_curve}")
  471. curve = bpy.data.objects.get(my_curve)
  472. if not curve:
  473. raise RuntimeError(f"No curve found for {self}.")
  474. elif curve.type != "CURVE":
  475. raise GraphError(f"ERROR: Object {curve.name} is not a curve.")
  476. spline = curve.data.splines[self.evaluate_input("Spline Index")]
  477. if spline.type == 'BEZIER':
  478. bez_pt = spline.bezier_points[self.evaluate_input("Index")]
  479. self.parameters["Point"]=bez_pt.co
  480. self.parameters["Left Handle"]=bez_pt.handle_left
  481. self.parameters["Right Handle"]=bez_pt.handle_right
  482. else:
  483. pt = spline.points[self.evaluate_input("Index")]
  484. self.parameters["Point"]=(pt.co[0], pt.co[1], pt.co[2])
  485. self.prepared, self.executed = True, True
  486. class UtilityGetNearestFactorOnCurve(MantisNode):
  487. def __init__(self, signature, base_tree):
  488. super().__init__(signature, base_tree, GetNearestFactorOnCurveSockets)
  489. self.init_parameters()
  490. self.node_type = "UTILITY"
  491. def bPrepare(self, bContext = None,):
  492. import bpy
  493. curve = bpy.data.objects.get(self.evaluate_input("Curve"))
  494. if not curve:
  495. raise RuntimeError(f"No curve found for {self}.")
  496. elif curve.type != "CURVE":
  497. raise GraphError(f"ERROR: Object {curve.name} is not a curve.")
  498. else:
  499. if bContext is None: bContext = bpy.context # is this wise?
  500. m = get_mesh_from_curve(curve.name,
  501. self.base_tree.execution_id,
  502. bContext, ribbon=False)
  503. # this is confusing but I am not re-writing these old functions
  504. from .utilities import FindNearestPointOnWireMesh as nearest_point
  505. spline_index = self.evaluate_input("Spline Index")
  506. ref_pt = self.evaluate_input("Reference Point")
  507. splines_points = [ [] for i in range (spline_index)]
  508. splines_points.append([ref_pt])
  509. pt = nearest_point(m, splines_points)[spline_index][0]
  510. self.parameters["Factor"] = pt
  511. self.prepared, self.executed = True, True
  512. class UtilityKDChoosePoint(MantisNode):
  513. def __init__(self, signature, base_tree):
  514. super().__init__(signature, base_tree)
  515. inputs = [
  516. "Reference Point" ,
  517. "Points" ,
  518. "Number to Find" ,
  519. ]
  520. outputs = [
  521. "Result Point" ,
  522. "Result Index" ,
  523. "Result Distance" ,
  524. ]
  525. self.inputs.init_sockets(inputs)
  526. self.outputs.init_sockets(outputs)
  527. self.init_parameters()
  528. self.node_type = "UTILITY"
  529. def bPrepare(self, bContext = None,):
  530. from mathutils import Vector
  531. points= []
  532. ref_point = self.evaluate_input('Reference Point')
  533. num_points = self.evaluate_input('Number to Find')
  534. for i, l in enumerate(self.inputs['Points'].links):
  535. pt = self.evaluate_input('Points', i)
  536. points.append(pt)
  537. if not isinstance(pt, Vector):
  538. prRed(f"Cannot get point from {l.from_node} for {self}")
  539. assert ref_point is not None, wrapRed(f"Reference Point {ref_point} is invalid in node {self}")
  540. result = kd_find(self, points, ref_point, num_points)
  541. indices = [ found_point[1] for found_point in result ]
  542. distances = [ found_point[2] for found_point in result ]
  543. array_choose_relink(self, indices, "Points", "Result Point")
  544. array_choose_data(self, indices, "Result Index")
  545. array_choose_data(self, distances, "Result Distance")
  546. self.prepared, self.executed = True, True
  547. class UtilityKDChooseXForm(MantisNode):
  548. def __init__(self, signature, base_tree):
  549. super().__init__(signature, base_tree)
  550. inputs = [
  551. "Reference Point" ,
  552. "xForm Nodes" ,
  553. "Get Point Head/Tail" ,
  554. "Number to Find" ,
  555. ]
  556. outputs = [
  557. "Result xForm" ,
  558. "Result Index" ,
  559. "Result Distance" ,
  560. ]
  561. self.inputs.init_sockets(inputs)
  562. self.outputs.init_sockets(outputs)
  563. self.init_parameters()
  564. self.node_type = "UTILITY"
  565. def bPrepare(self, bContext = None,):
  566. if len(self.hierarchy_dependencies)==0 and len(self.hierarchy_connections)==0 and \
  567. len(self.connections)==0 and len(self.dependencies)==0:
  568. self.prepared, self.executed = True, True
  569. return #Either it is already done or it doesn't matter.
  570. from mathutils import Vector
  571. points= []
  572. ref_point = self.evaluate_input('Reference Point')
  573. num_points = self.evaluate_input('Number to Find')
  574. for i, l in enumerate(self.inputs['xForm Nodes'].links):
  575. matrix = l.from_node.evaluate_input('Matrix')
  576. if matrix is None:
  577. raise GraphError(f"Cannot get point from {l.from_node} for {self}. Does it have a matrix?")
  578. pt = matrix.translation
  579. if head_tail := self.evaluate_input("Get Point Head/Tail"):
  580. # get the Y-axis of the basis, assume it is normalized
  581. y_axis = Vector((matrix[0][1],matrix[1][1], matrix[2][1]))
  582. pt = pt.lerp(pt+y_axis*matrix[3][3], head_tail)
  583. points.append(pt)
  584. if not isinstance(pt, Vector):
  585. prRed(f"Cannot get point from {l.from_node} for {self}")
  586. assert ref_point is not None, wrapRed(f"Reference Point {ref_point} is invalid in node {self}")
  587. result = kd_find(self, points, ref_point, num_points)
  588. indices = [ found_point[1] for found_point in result ]
  589. distances = [ found_point[2] for found_point in result ]
  590. array_choose_relink(self, indices, "xForm Nodes", "Result xForm")
  591. array_choose_data(self, indices, "Result Index")
  592. array_choose_data(self, distances, "Result Distance")
  593. self.prepared, self.executed = True, True
  594. class UtilityMetaRig(MantisNode):
  595. '''A node representing an armature object'''
  596. def __init__(self, signature, base_tree):
  597. super().__init__(signature, base_tree)
  598. inputs = [
  599. "Meta-Armature" ,
  600. "Meta-Bone" ,
  601. ]
  602. outputs = [
  603. "Matrix" ,
  604. ]
  605. self.inputs.init_sockets(inputs)
  606. self.outputs.init_sockets(outputs)
  607. self.init_parameters()
  608. self.node_type = "UTILITY"
  609. def bPrepare(self, bContext = None,):
  610. #kinda clumsy, whatever
  611. import bpy
  612. from mathutils import Matrix
  613. m = Matrix.Identity(4)
  614. meta_rig = self.evaluate_input("Meta-Armature")
  615. meta_bone = self.evaluate_input("Meta-Bone")
  616. if meta_rig:
  617. if ( armOb := bpy.data.objects.get(meta_rig) ):
  618. m = armOb.matrix_world
  619. if ( b := armOb.data.bones.get(meta_bone)):
  620. # calculate the correct object-space matrix
  621. m = Matrix.Identity(3)
  622. bones = [] # from the last ancestor, mult the matrices until we get to b
  623. while (b): bones.append(b); b = b.parent
  624. while (bones): b = bones.pop(); m = m @ b.matrix
  625. m = Matrix.Translation(b.head_local) @ m.to_4x4()
  626. #
  627. m[3][3] = b.length # this is where I arbitrarily decided to store length
  628. # else:
  629. # prRed("no bone for MetaRig node ", self)
  630. else:
  631. raise RuntimeError(wrapRed(f"No meta-rig input for MetaRig node {self}"))
  632. self.parameters["Matrix"] = m
  633. self.prepared = True
  634. self.executed = True
  635. class UtilityBoneProperties(SimpleInputNode):
  636. '''A node representing a bone's gettable properties'''
  637. def __init__(self, signature, base_tree):
  638. super().__init__(signature, base_tree)
  639. outputs = [
  640. "matrix" ,
  641. "matrix_local" ,
  642. "matrix_basis" ,
  643. "head" ,
  644. "tail" ,
  645. "length" ,
  646. "rotation" ,
  647. "location" ,
  648. "scale" ,
  649. ]
  650. self.outputs.init_sockets(outputs)
  651. self.init_parameters()
  652. def fill_parameters(self, prototype=None):
  653. return
  654. # TODO this should probably be moved to Links
  655. class UtilityDriverVariable(MantisNode):
  656. '''A node representing an armature object'''
  657. def __init__(self, signature, base_tree):
  658. super().__init__(signature, base_tree)
  659. inputs = [
  660. "Variable Type" ,
  661. "Property" ,
  662. "Property Index" ,
  663. "Evaluation Space",
  664. "Rotation Mode" ,
  665. "xForm 1" ,
  666. "xForm 2" ,
  667. ]
  668. outputs = [
  669. "Driver Variable",
  670. ]
  671. self.inputs.init_sockets(inputs)
  672. self.outputs.init_sockets(outputs)
  673. self.init_parameters()
  674. self.node_type = "DRIVER" # MUST be run in Pose mode
  675. self.prepared = True
  676. def reset_execution(self):
  677. super().reset_execution()
  678. # clear this to ensure there are no stale reference pointers
  679. self.parameters["Driver Variable"] = None
  680. self.prepared=True
  681. def evaluate_input(self, input_name):
  682. if input_name == 'Property':
  683. if self.inputs.get('Property'):
  684. if self.inputs['Property'].is_linked:
  685. # get the name instead...
  686. trace = trace_single_line(self, input_name)
  687. return trace[1].name # the name of the socket
  688. return self.parameters["Property"]
  689. return super().evaluate_input(input_name)
  690. def GetxForm(self, index=1):
  691. trace = trace_single_line(self, "xForm 1" if index == 1 else "xForm 2")
  692. for node in trace[0]:
  693. if (node.__class__ in [xFormArmature, xFormBone]):
  694. return node #this will fetch the first one, that's good!
  695. return None
  696. def bExecute(self, bContext = None,):
  697. prepare_parameters(self)
  698. #prPurple ("Executing Driver Variable Node")
  699. xF1 = self.GetxForm()
  700. xF2 = self.GetxForm(index=2)
  701. # kinda clumsy
  702. xForm1, xForm2 = None, None
  703. if xF1 : xForm1 = xF1.bGetObject()
  704. if xF2 : xForm2 = xF2.bGetObject()
  705. v_type = self.evaluate_input("Variable Type")
  706. i = self.evaluate_input("Property Index"); dVarChannel = ""
  707. if (i >= 0): #negative values will use the vector property.
  708. if self.evaluate_input("Property") == 'location':
  709. if i == 0: dVarChannel = "LOC_X"
  710. elif i == 1: dVarChannel = "LOC_Y"
  711. elif i == 2: dVarChannel = "LOC_Z"
  712. else: raise RuntimeError("Invalid property index for %s" % self)
  713. if self.evaluate_input("Property") == 'rotation':
  714. if i == 0: dVarChannel = "ROT_X"
  715. elif i == 1: dVarChannel = "ROT_Y"
  716. elif i == 2: dVarChannel = "ROT_Z"
  717. elif i == 3: dVarChannel = "ROT_W"
  718. else: raise RuntimeError("Invalid property index for %s" % self)
  719. if self.evaluate_input("Property") == 'scale':
  720. if i == 0: dVarChannel = "SCALE_X"
  721. elif i == 1: dVarChannel = "SCALE_Y"
  722. elif i == 2: dVarChannel = "SCALE_Z"
  723. elif i == 3: dVarChannel = "SCALE_AVG"
  724. else: raise RuntimeError("Invalid property index for %s" % self)
  725. if self.evaluate_input("Property") == 'scale_average':
  726. dVarChannel = "SCALE_AVG"
  727. if dVarChannel: v_type = "TRANSFORMS"
  728. my_var = {
  729. "owner" : xForm1, # will be filled in by Driver
  730. "prop" : self.evaluate_input("Property"), # will be filled in by Driver
  731. "type" : v_type,
  732. "space" : self.evaluate_input("Evaluation Space"),
  733. "rotation_mode" : self.evaluate_input("Rotation Mode"),
  734. "xForm 1" : xForm1,#self.GetxForm(index = 1),
  735. "xForm 2" : xForm2,#self.GetxForm(index = 2),
  736. "channel" : dVarChannel,}
  737. self.parameters["Driver Variable"] = my_var
  738. self.executed = True
  739. class UtilityKeyframe(MantisNode):
  740. '''A node representing a keyframe for a F-Curve'''
  741. def __init__(self, signature, base_tree):
  742. super().__init__(signature, base_tree)
  743. inputs = [
  744. "Frame" ,
  745. "Value" ,
  746. ]
  747. outputs = [
  748. "Keyframe" ,
  749. ]
  750. additional_parameters = {"Keyframe":{}}
  751. self.inputs.init_sockets(inputs)
  752. self.outputs.init_sockets(outputs)
  753. self.init_parameters( additional_parameters=additional_parameters)
  754. self.node_type = "DRIVER" # MUST be run in Pose mode
  755. setup_custom_props(self)
  756. def bPrepare(self, bContext = None,):
  757. key = self.parameters["Keyframe"]
  758. from mathutils import Vector
  759. key["co"]= Vector( (self.evaluate_input("Frame"), self.evaluate_input("Value"),))
  760. key["type"]="GENERATED"
  761. key["interpolation"] = "LINEAR"
  762. # eventually this will have the right data, TODO
  763. # self.parameters["Keyframe"] = key
  764. self.prepared = True
  765. self.executed = True
  766. class UtilityFCurve(MantisNode):
  767. '''A node representing an armature object'''
  768. def __init__(self, signature, base_tree):
  769. super().__init__(signature, base_tree)
  770. inputs = [
  771. "Extrapolation Mode",
  772. ]
  773. outputs = [
  774. "fCurve",
  775. ]
  776. self.inputs.init_sockets(inputs)
  777. self.outputs.init_sockets(outputs)
  778. self.init_parameters()
  779. self.node_type = "UTILITY"
  780. setup_custom_props(self)
  781. self.prepared = True
  782. def reset_execution(self):
  783. super().reset_execution()
  784. self.prepared=True
  785. def evaluate_input(self, input_name):
  786. return super().evaluate_input(input_name)
  787. def bExecute(self, bContext = None,):
  788. prepare_parameters(self)
  789. from .utilities import get_node_prototype
  790. np = get_node_prototype(self.signature, self.base_tree)
  791. extrap_mode = self.evaluate_input("Extrapolation Mode")
  792. keys = [] # ugly but whatever
  793. #['amplitude', 'back', 'bl_rna', 'co', 'co_ui', 'easing', 'handle_left', 'handle_left_type', 'handle_right', 'handle_right_type',
  794. # 'interpolation', 'period', 'rna_type', 'select_control_point', 'select_left_handle', 'select_right_handle', 'type']
  795. for k in self.inputs.keys():
  796. if k == 'Extrapolation Mode' : continue
  797. # print (self.inputs[k])
  798. if (key := self.evaluate_input(k)) is None:
  799. prOrange(f"WARN: No keyframe connected to {self}:{k}. Skipping Link.")
  800. else:
  801. keys.append(key)
  802. if len(keys) <1:
  803. prOrange(f"WARN: no keys in fCurve {self}.")
  804. keys.append(extrap_mode)
  805. self.parameters["fCurve"] = keys
  806. self.executed = True
  807. #TODO make the fCurve data a data class instead of a dict
  808. class UtilityDriver(MantisNode):
  809. '''A node representing an armature object'''
  810. def __init__(self, signature, base_tree):
  811. super().__init__(signature, base_tree)
  812. inputs = [
  813. "Driver Type" ,
  814. "Expression" ,
  815. "fCurve" ,
  816. ]
  817. outputs = [
  818. "Driver",
  819. ]
  820. from .drivers import MantisDriver
  821. additional_parameters = {
  822. "Driver":MantisDriver(),
  823. }
  824. self.inputs.init_sockets(inputs)
  825. self.outputs.init_sockets(outputs)
  826. self.init_parameters(additional_parameters=additional_parameters)
  827. self.node_type = "DRIVER" # MUST be run in Pose mode
  828. setup_custom_props(self)
  829. self.prepared = True
  830. def reset_execution(self):
  831. super().reset_execution()
  832. from .drivers import MantisDriver
  833. self.parameters["Driver"]=MantisDriver()
  834. self.prepared=True
  835. def bExecute(self, bContext = None,):
  836. prepare_parameters(self)
  837. from .drivers import MantisDriver
  838. #prPurple("Executing Driver Node")
  839. my_vars = []
  840. keys = self.evaluate_input("fCurve")
  841. if keys is None or len(keys) <2:
  842. prWhite(f"INFO: no fCurve connected to {self}; using default fCurve.")
  843. from mathutils import Vector
  844. keys = [
  845. {"co":Vector( (0, 0,)), "type":"GENERATED", "interpolation":"LINEAR" },
  846. {"co":Vector( (1, 1,)), "type":"GENERATED", "interpolation":"LINEAR" },
  847. "CONSTANT",]
  848. for inp in list(self.inputs.keys() )[3:]:
  849. if (new_var := self.evaluate_input(inp)):
  850. new_var["name"] = inp
  851. my_vars.append(new_var)
  852. else:
  853. raise RuntimeError(f"Failed to initialize Driver variable for {self}")
  854. my_driver ={ "owner" : None,
  855. "prop" : None, # will be filled out in the node that uses the driver
  856. "expression" : self.evaluate_input("Expression"),
  857. "ind" : -1, # same here
  858. "type" : self.evaluate_input("Driver Type"),
  859. "vars" : my_vars,
  860. "keys" : keys[:-1],
  861. "extrapolation" : keys[-1] }
  862. my_driver = MantisDriver(my_driver)
  863. self.parameters["Driver"].update(my_driver)
  864. print("Initializing driver %s " % (wrapPurple(self.__repr__())) )
  865. self.executed = True
  866. class UtilitySwitch(MantisNode):
  867. '''A node representing an armature object'''
  868. def __init__(self, signature, base_tree):
  869. super().__init__(signature, base_tree)
  870. inputs = {
  871. "Parameter" ,
  872. "Parameter Index" ,
  873. "Invert Switch" ,
  874. }
  875. outputs = [
  876. "Driver",
  877. ]
  878. from .drivers import MantisDriver
  879. additional_parameters = {
  880. "Driver":MantisDriver(),
  881. }
  882. self.inputs.init_sockets(inputs)
  883. self.outputs.init_sockets(outputs)
  884. self.init_parameters(additional_parameters=additional_parameters)
  885. self.node_type = "DRIVER" # MUST be run in Pose mode
  886. self.prepared = True
  887. def evaluate_input(self, input_name):
  888. if input_name == 'Parameter':
  889. if self.inputs['Parameter'].is_connected:
  890. trace = trace_single_line(self, input_name)
  891. return trace[1].name # the name of the socket
  892. return self.parameters["Parameter"]
  893. return super().evaluate_input(input_name)
  894. def GetxForm(self,):
  895. trace = trace_single_line(self, "Parameter" )
  896. for node in trace[0]:
  897. if (node.__class__ in [xFormArmature, xFormBone]):
  898. return node #this will fetch the first one, that's good!
  899. return None
  900. def reset_execution(self):
  901. super().reset_execution()
  902. from .drivers import MantisDriver
  903. self.parameters["Driver"]=MantisDriver()
  904. self.prepared=True
  905. def bExecute(self, bContext = None,):
  906. #prepare_parameters(self)
  907. #prPurple ("Executing Switch Node")
  908. xForm = self.GetxForm()
  909. if xForm : xForm = xForm.bGetObject()
  910. if not xForm:
  911. raise RuntimeError("Could not evaluate xForm for %s" % self)
  912. from .drivers import MantisDriver
  913. my_driver ={ "owner" : None,
  914. "prop" : None, # will be filled out in the node that uses the driver
  915. "ind" : -1, # same here
  916. "type" : "SCRIPTED",
  917. "vars" : [ { "owner" : xForm,
  918. "prop" : self.evaluate_input("Parameter"),
  919. "name" : "a",
  920. "type" : "SINGLE_PROP", } ],
  921. "keys" : [ { "co":(0,0),
  922. "interpolation": "LINEAR",
  923. "type":"KEYFRAME",}, #display type
  924. { "co":(1,1),
  925. "interpolation": "LINEAR",
  926. "type":"KEYFRAME",},],
  927. "extrapolation": 'CONSTANT', }
  928. my_driver ["expression"] = "a"
  929. my_driver = MantisDriver(my_driver)
  930. # this makes it so I can check for type later!
  931. if self.evaluate_input("Invert Switch") == True:
  932. my_driver ["expression"] = "1 - a"
  933. # this way, regardless of what order things are handled, the
  934. # driver is sent to the next node.
  935. # In the case of some drivers, the parameter may be sent out
  936. # before it's filled in (because there is a circular dependency)
  937. # I want to support this behaviour because Blender supports it.
  938. # We do not make a copy. We update the driver, so that
  939. # the same instance is filled out.
  940. self.parameters["Driver"].update(my_driver)
  941. print("Initializing driver %s " % (wrapPurple(self.__repr__())) )
  942. self.executed = True
  943. class UtilityCombineThreeBool(MantisNode):
  944. '''A node for combining three booleans into a boolean three-tuple'''
  945. def __init__(self, signature, base_tree):
  946. super().__init__(signature, base_tree)
  947. inputs = [
  948. "X" ,
  949. "Y" ,
  950. "Z" ,
  951. ]
  952. outputs = [
  953. "Three-Bool",
  954. ]
  955. self.inputs.init_sockets(inputs)
  956. self.outputs.init_sockets(outputs)
  957. self.init_parameters()
  958. self.node_type = "UTILITY"
  959. def reset_execution(self): # need to make sure any references are deleted
  960. super().reset_execution() # so we prepare the node again to reset them
  961. if self.parameters["Three-Bool"] is not None:
  962. for param in self.parameters["Three-Bool"].values():
  963. if isinstance(param, dict):
  964. self.prepared=False; break
  965. def bPrepare(self, bContext = None,):
  966. self.parameters["Three-Bool"] = (
  967. self.evaluate_input("X"),
  968. self.evaluate_input("Y"),
  969. self.evaluate_input("Z"), )
  970. self.prepared = True
  971. self.executed = True
  972. # Note this is a copy of the above. This needs to be de-duplicated.
  973. class UtilityCombineVector(MantisNode):
  974. '''A node for combining three floats into a vector'''
  975. def __init__(self, signature, base_tree):
  976. super().__init__(signature, base_tree)
  977. super().__init__(signature, base_tree)
  978. inputs = [
  979. "X" ,
  980. "Y" ,
  981. "Z" ,
  982. ]
  983. outputs = [
  984. "Vector",
  985. ]
  986. self.inputs.init_sockets(inputs)
  987. self.outputs.init_sockets(outputs)
  988. self.init_parameters()
  989. self.node_type = "UTILITY"
  990. def reset_execution(self): # need to make sure any references are deleted
  991. super().reset_execution() # so we prepare the node again to reset them
  992. if self.parameters["Vector"] is not None:
  993. for param in self.parameters["Vector"]:
  994. if isinstance(param, dict):
  995. self.prepared=False; break
  996. def bPrepare(self, bContext = None,):
  997. #prPurple("Executing CombineVector Node")
  998. prepare_parameters(self)
  999. self.parameters["Vector"] = (
  1000. self.evaluate_input("X"),
  1001. self.evaluate_input("Y"),
  1002. self.evaluate_input("Z"), )
  1003. self.prepared, self.executed = True, True
  1004. class UtilitySeparateVector(MantisNode):
  1005. '''A node for separating a vector into three floats'''
  1006. def __init__(self, signature, base_tree):
  1007. super().__init__(signature, base_tree)
  1008. inputs = [
  1009. "Vector"
  1010. ]
  1011. outputs = [
  1012. "X" ,
  1013. "Y" ,
  1014. "Z" ,
  1015. ]
  1016. self.inputs.init_sockets(inputs)
  1017. self.outputs.init_sockets(outputs)
  1018. self.init_parameters()
  1019. self.node_type = "UTILITY"
  1020. def bPrepare(self, bContext = None,):
  1021. self.parameters["X"] = self.evaluate_input("Vector")[0]
  1022. self.parameters["Y"] = self.evaluate_input("Vector")[1]
  1023. self.parameters["Z"] = self.evaluate_input("Vector")[2]
  1024. self.prepared, self.executed = True, True
  1025. class UtilityCatStrings(MantisNode):
  1026. '''A node representing an armature object'''
  1027. def __init__(self, signature, base_tree):
  1028. super().__init__(signature, base_tree)
  1029. inputs = [
  1030. "String_1" ,
  1031. "String_2" ,
  1032. ]
  1033. outputs = [
  1034. "OutputString" ,
  1035. ]
  1036. self.inputs.init_sockets(inputs)
  1037. self.outputs.init_sockets(outputs)
  1038. self.init_parameters()
  1039. self.node_type = "UTILITY"
  1040. def bPrepare(self, bContext = None,):
  1041. self.parameters["OutputString"] = self.evaluate_input("String_1")+self.evaluate_input("String_2")
  1042. self.prepared, self.executed = True, True
  1043. # TODO move this to the Xform file
  1044. class InputExistingGeometryObject(MantisNode):
  1045. '''A node representing an existing object'''
  1046. def __init__(self, signature, base_tree):
  1047. super().__init__(signature, base_tree)
  1048. inputs = [
  1049. "Name" ,
  1050. ]
  1051. outputs = [
  1052. "Object" ,
  1053. ]
  1054. self.inputs.init_sockets(inputs)
  1055. self.outputs.init_sockets(outputs)
  1056. self.init_parameters()
  1057. self.node_type = "XFORM"
  1058. def reset_execution(self):
  1059. super().reset_execution()
  1060. self.prepared=False
  1061. def bPrepare(self, bContext=None):
  1062. from bpy import data
  1063. ob = None
  1064. if name := self.evaluate_input("Name"):
  1065. ob= data.objects.get( name )
  1066. if ob is None and name:
  1067. prRed(f"No object found with name {name} in {self}")
  1068. self.bObject=ob
  1069. self.prepared, self.executed = True, True
  1070. def bGetObject(self, mode=''):
  1071. return self.bObject
  1072. class InputExistingGeometryData(MantisNode):
  1073. '''A node representing existing object data'''
  1074. def __init__(self, signature, base_tree):
  1075. super().__init__(signature, base_tree)
  1076. inputs = [
  1077. "Name" ,
  1078. ]
  1079. outputs = [
  1080. "Geometry" ,
  1081. ]
  1082. self.inputs.init_sockets(inputs)
  1083. self.outputs.init_sockets(outputs)
  1084. self.init_parameters()
  1085. self.node_type = "UTILITY"
  1086. self.prepared = True; self.executed = True
  1087. def reset_execution(self):
  1088. super().reset_execution()
  1089. self.prepared, self.executed = True, True
  1090. # the mode argument is only for interface consistency
  1091. def bGetObject(self, mode=''):
  1092. from bpy import data
  1093. # first try Curve, then try Mesh
  1094. bObject = data.curves.get(self.evaluate_input("Name"))
  1095. if not bObject:
  1096. bObject = data.meshes.get(self.evaluate_input("Name"))
  1097. if bObject is None:
  1098. raise RuntimeError(f"Could not find a mesh or curve datablock named \"{self.evaluate_input('Name')}\" for node {self}")
  1099. return bObject
  1100. class UtilityGeometryOfXForm(MantisNode):
  1101. '''A node representing existing object data'''
  1102. def __init__(self, signature, base_tree):
  1103. super().__init__(signature, base_tree)
  1104. inputs = [
  1105. "xForm" ,
  1106. ]
  1107. outputs = [
  1108. "Geometry" ,
  1109. ]
  1110. self.inputs.init_sockets(inputs)
  1111. self.outputs.init_sockets(outputs)
  1112. self.init_parameters()
  1113. self.node_type = "UTILITY"
  1114. self.prepared = True
  1115. self.executed = True
  1116. def reset_execution(self):
  1117. super().reset_execution()
  1118. self.prepared, self.executed = True, True
  1119. # mode for interface consistency
  1120. def bGetObject(self, mode=''):
  1121. if not (self.inputs.get('xForm') and self.inputs['xForm'].links):
  1122. prOrange(f"WARN: Cannot retrieve data from {self}, there is no xForm node connected.")
  1123. return None
  1124. xf = self.inputs["xForm"].links[0].from_node
  1125. if xf.node_type == 'XFORM':
  1126. xf_ob = xf.bGetObject()
  1127. if (xf_ob is not None) and xf_ob.type in ['MESH', 'CURVE']:
  1128. return xf_ob.data
  1129. prOrange(f"WARN: Cannot retrieve data from {self}, the connected xForm is not a mesh or curve.")
  1130. return None
  1131. class UtilityNameOfXForm(MantisNode):
  1132. '''A node representing existing object data'''
  1133. def __init__(self, signature, base_tree):
  1134. super().__init__(signature, base_tree)
  1135. inputs = [
  1136. "xForm" ,
  1137. ]
  1138. outputs = [
  1139. "Name" ,
  1140. ]
  1141. self.inputs.init_sockets(inputs)
  1142. self.outputs.init_sockets(outputs)
  1143. self.init_parameters()
  1144. self.node_type = "UTILITY"
  1145. # mode for interface consistency
  1146. def bPrepare(self, bContext = None,):
  1147. if not (self.inputs.get('xForm') and self.inputs['xForm'].links):
  1148. raise RuntimeError( f"WARN: Cannot retrieve data from {self},"
  1149. " there is no xForm node connected.")
  1150. xf = self.inputs["xForm"].links[0].from_node
  1151. self.parameters["Name"] = xf.evaluate_input('Name')
  1152. self.prepared, self.executed = True, True
  1153. class UtilityGetBoneLength(MantisNode):
  1154. '''A node to get the length of a bone matrix'''
  1155. def __init__(self, signature, base_tree):
  1156. super().__init__(signature, base_tree)
  1157. inputs = [
  1158. "Bone Matrix" ,
  1159. ]
  1160. outputs = [
  1161. "Bone Length" ,
  1162. ]
  1163. self.inputs.init_sockets(inputs)
  1164. self.outputs.init_sockets(outputs)
  1165. self.init_parameters()
  1166. self.node_type = "UTILITY"
  1167. def bPrepare(self, bContext = None,):
  1168. if (l := self.evaluate_input("Bone Matrix")) is not None:
  1169. self.parameters["Bone Length"] = l[3][3]
  1170. else:
  1171. other = self.inputs["Bone Matrix"].links[0].from_node
  1172. raise RuntimeError(f"Cannot get matrix for {self} from {other}")
  1173. self.prepared, self.executed = True, True
  1174. class UtilityPointFromBoneMatrix(MantisNode):
  1175. '''A node representing an armature object'''
  1176. def __init__(self, signature, base_tree):
  1177. super().__init__(signature, base_tree)
  1178. inputs = [
  1179. "Bone Matrix" ,
  1180. "Head/Tail" ,
  1181. ]
  1182. outputs = [
  1183. "Point" ,
  1184. ]
  1185. self.inputs.init_sockets(inputs)
  1186. self.outputs.init_sockets(outputs)
  1187. self.init_parameters()
  1188. self.node_type = "UTILITY"
  1189. # TODO: find out why this is sometimes not ready at bPrepare phase
  1190. def bPrepare(self, bContext = None,):
  1191. from mathutils import Vector
  1192. matrix = self.evaluate_input("Bone Matrix")
  1193. head, rotation, _scale = matrix.copy().decompose()
  1194. tail = head.copy() + (rotation @ Vector((0,1,0)))*matrix[3][3]
  1195. self.parameters["Point"] = head.lerp(tail, self.evaluate_input("Head/Tail"))
  1196. self.prepared, self.executed = True, True
  1197. class UtilitySetBoneLength(MantisNode):
  1198. '''Sets the length of a Bone's matrix'''
  1199. def __init__(self, signature, base_tree):
  1200. super().__init__(signature, base_tree)
  1201. inputs = [
  1202. "Bone Matrix" ,
  1203. "Length" ,
  1204. ]
  1205. outputs = [
  1206. "Bone Matrix" ,
  1207. ]
  1208. self.inputs.init_sockets(inputs)
  1209. self.outputs.init_sockets(outputs)
  1210. self.init_parameters()
  1211. self.node_type = "UTILITY"
  1212. def bPrepare(self, bContext = None,):
  1213. from mathutils import Vector
  1214. if matrix := self.evaluate_input("Bone Matrix"):
  1215. matrix = matrix.copy()
  1216. # print (self.inputs["Length"].links)
  1217. matrix[3][3] = self.evaluate_input("Length")
  1218. self.parameters["Length"] = self.evaluate_input("Length")
  1219. self.parameters["Bone Matrix"] = matrix
  1220. else:
  1221. raise RuntimeError(f"Cannot get matrix for {self}")
  1222. self.prepared, self.executed = True, True
  1223. class UtilityMatrixSetLocation(MantisNode):
  1224. '''Sets the location of a matrix'''
  1225. def __init__(self, signature, base_tree):
  1226. super().__init__(signature, base_tree)
  1227. inputs = [
  1228. "Matrix" ,
  1229. "Location" ,
  1230. ]
  1231. outputs = [
  1232. "Matrix" ,
  1233. ]
  1234. self.inputs.init_sockets(inputs)
  1235. self.outputs.init_sockets(outputs)
  1236. self.init_parameters()
  1237. self.node_type = "UTILITY"
  1238. def bPrepare(self, bContext = None,):
  1239. from mathutils import Vector
  1240. if matrix := self.evaluate_input("Matrix"):
  1241. matrix = matrix.copy()
  1242. # print (self.inputs["Length"].links)
  1243. loc = self.evaluate_input("Location")
  1244. matrix[0][3] = loc[0]; matrix[1][3] = loc[1]; matrix[2][3] = loc[2]
  1245. self.parameters["Matrix"] = matrix
  1246. self.prepared, self.executed = True, True
  1247. class UtilityMatrixGetLocation(MantisNode):
  1248. '''Gets the location of a matrix'''
  1249. def __init__(self, signature, base_tree):
  1250. super().__init__(signature, base_tree)
  1251. inputs = [
  1252. "Matrix" ,
  1253. ]
  1254. outputs = [
  1255. "Location" ,
  1256. ]
  1257. self.inputs.init_sockets(inputs)
  1258. self.outputs.init_sockets(outputs)
  1259. self.init_parameters()
  1260. self.node_type = "UTILITY"
  1261. def bPrepare(self, bContext = None,):
  1262. from mathutils import Vector
  1263. if matrix := self.evaluate_input("Matrix"):
  1264. self.parameters["Location"] = matrix.to_translation()
  1265. self.prepared = True; self.executed = True
  1266. class UtilityMatrixFromXForm(MantisNode):
  1267. """Returns the matrix of the given xForm node."""
  1268. def __init__(self, signature, base_tree):
  1269. super().__init__(signature, base_tree)
  1270. inputs = [
  1271. "xForm" ,
  1272. ]
  1273. outputs = [
  1274. "Matrix" ,
  1275. ]
  1276. self.node_type = "UTILITY"
  1277. self.inputs.init_sockets(inputs)
  1278. self.outputs.init_sockets(outputs)
  1279. self.init_parameters()
  1280. def GetxForm(self):
  1281. trace = trace_single_line(self, "xForm")
  1282. for node in trace[0]:
  1283. if (node.node_type == 'XFORM'):
  1284. return node
  1285. raise GraphError("%s is not connected to an xForm" % self)
  1286. def bPrepare(self, bContext = None,):
  1287. from mathutils import Vector, Matrix
  1288. self.parameters["Matrix"] = Matrix.Identity(4)
  1289. if matrix := self.GetxForm().parameters.get("Matrix"):
  1290. self.parameters["Matrix"] = matrix.copy()
  1291. elif hasattr(self.GetxForm().bObject, "matrix"):
  1292. self.parameters["Matrix"] = self.GetxForm().bObject.matrix.copy()
  1293. elif hasattr(self.GetxForm().bObject, "matrix_world"):
  1294. self.parameters["Matrix"] = self.GetxForm().bObject.matrix_world.copy()
  1295. else:
  1296. prRed(f"Could not find matrix for {self} - check if the referenced object exists.")
  1297. self.prepared = True; self.executed = True
  1298. class UtilityAxesFromMatrix(MantisNode):
  1299. """Returns the axes of the given matrix."""
  1300. def __init__(self, signature, base_tree):
  1301. super().__init__(signature, base_tree)
  1302. inputs = [
  1303. "Matrix" ,
  1304. ]
  1305. outputs = [
  1306. "X Axis" ,
  1307. "Y Axis" ,
  1308. "Z Axis" ,
  1309. ]
  1310. self.inputs.init_sockets(inputs)
  1311. self.outputs.init_sockets(outputs)
  1312. self.init_parameters()
  1313. self.node_type = "UTILITY"
  1314. def bPrepare(self, bContext = None,):
  1315. from mathutils import Vector
  1316. if matrix := self.evaluate_input("Matrix"):
  1317. matrix= matrix.copy().to_3x3(); matrix.transpose()
  1318. self.parameters['X Axis'] = matrix[0]
  1319. self.parameters['Y Axis'] = matrix[1]
  1320. self.parameters['Z Axis'] = matrix[2]
  1321. self.prepared = True; self.executed = True
  1322. class UtilityBoneMatrixHeadTailFlip(MantisNode):
  1323. def __init__(self, signature, base_tree):
  1324. super().__init__(signature, base_tree)
  1325. inputs = [
  1326. "Bone Matrix" ,
  1327. ]
  1328. outputs = [
  1329. "Bone Matrix" ,
  1330. ]
  1331. self.inputs.init_sockets(inputs)
  1332. self.outputs.init_sockets(outputs)
  1333. self.init_parameters()
  1334. self.node_type = "UTILITY"
  1335. def bPrepare(self, bContext = None,):
  1336. from mathutils import Vector, Matrix, Quaternion
  1337. from bpy.types import Bone
  1338. if matrix := self.evaluate_input("Bone Matrix"):
  1339. axis, roll = Bone.AxisRollFromMatrix(matrix.to_3x3())
  1340. new_mat = Bone.MatrixFromAxisRoll(-1*axis, roll)
  1341. length = matrix[3][3]
  1342. new_mat.resize_4x4() # last column contains
  1343. new_mat[0][3] = matrix[0][3] + axis[0]*length # x location
  1344. new_mat[1][3] = matrix[1][3] + axis[1]*length # y location
  1345. new_mat[2][3] = matrix[2][3] + axis[2]*length # z location
  1346. new_mat[3][3] = length # length
  1347. self.parameters["Bone Matrix"] = new_mat
  1348. self.prepared, self.executed = True, True
  1349. class UtilityMatrixTransform(MantisNode):
  1350. def __init__(self, signature, base_tree):
  1351. super().__init__(signature, base_tree)
  1352. inputs = [
  1353. "Matrix 1" ,
  1354. "Matrix 2" ,
  1355. ]
  1356. outputs = [
  1357. "Out Matrix" ,
  1358. ]
  1359. self.inputs.init_sockets(inputs)
  1360. self.outputs.init_sockets(outputs)
  1361. self.init_parameters()
  1362. self.node_type = "UTILITY"
  1363. def bPrepare(self, bContext = None,):
  1364. from mathutils import Vector
  1365. mat1 = self.evaluate_input("Matrix 1"); mat2 = self.evaluate_input("Matrix 2")
  1366. if mat1 and mat2:
  1367. mat1copy = mat1.copy()
  1368. self.parameters["Out Matrix"] = mat2 @ mat1copy
  1369. self.parameters["Out Matrix"].translation = mat1copy.to_translation()+ mat2.to_translation()
  1370. else:
  1371. raise RuntimeError(wrapRed(f"Node {self} did not receive all matrix inputs..."
  1372. " found input 1? {mat1 is not None}, 2? {mat2 is not None}"))
  1373. self.prepared, self.executed = True, True
  1374. class UtilityMatrixInvert(MantisNode):
  1375. def __init__(self, signature, base_tree):
  1376. super().__init__(signature, base_tree, MatrixInvertSockets)
  1377. self.init_parameters()
  1378. self.node_type = "UTILITY"
  1379. def bPrepare(self, bContext = None,):
  1380. from mathutils import Vector
  1381. mat1 = self.evaluate_input("Matrix 1")
  1382. if mat1:
  1383. mat1copy = mat1.copy()
  1384. try:
  1385. self.parameters["Matrix"] = mat1copy.inverted()
  1386. except ValueError as e:
  1387. prRed(f"ERROR: {self}: The matrix cannot be inverted."); prOrange(mat1)
  1388. raise e
  1389. else:
  1390. raise RuntimeError(wrapRed(f"Node {self} did not receive all matrix inputs..."
  1391. " found input 1? {mat1 is not None}"))
  1392. self.prepared, self.executed = True, True
  1393. class UtilityMatrixCompose(MantisNode):
  1394. def __init__(self, signature, base_tree):
  1395. super().__init__(signature, base_tree, MatrixComposeSockets)
  1396. self.init_parameters()
  1397. self.node_type = "UTILITY"
  1398. def bPrepare(self, bContext = None,):
  1399. from mathutils import Matrix
  1400. matrix= Matrix.Identity(3)
  1401. matrix[0] = self.evaluate_input('X Basis Vector')
  1402. matrix[1] = self.evaluate_input('Y Basis Vector')
  1403. matrix[2] = self.evaluate_input('Z Basis Vector')
  1404. matrix.transpose(); matrix=matrix.to_4x4()
  1405. matrix.translation = self.evaluate_input('Translation')
  1406. self.parameters['Matrix']=matrix
  1407. self.prepared = True; self.executed = True
  1408. class UtilityMatrixAlignRoll(MantisNode):
  1409. def __init__(self, signature, base_tree):
  1410. super().__init__(signature, base_tree, MatrixAlignRollSockets)
  1411. self.init_parameters()
  1412. self.node_type = "UTILITY"
  1413. def bPrepare(self, bContext = None,):
  1414. from mathutils import Vector, Matrix
  1415. align_axis = Vector(self.evaluate_input('Alignment Vector'))
  1416. # why do I have to construct a vector here?
  1417. # why is the socket returning a bpy_prop_array ?
  1418. print(align_axis)
  1419. if align_axis.length_squared==0:
  1420. raise RuntimeError(f"WARN: cannot align matrix in {self}"
  1421. " because the alignment vector is zero.")
  1422. input=self.evaluate_input('Matrix').copy()
  1423. y_axis= input.to_3x3().transposed()[1]
  1424. from .utilities import project_point_to_plane
  1425. projected=project_point_to_plane(
  1426. align_axis.normalized(), Vector((0,0,0)), y_axis).normalized()
  1427. # now that we have the projected vector, transform the points from
  1428. # the plane of the y_axis to flat space and get the signed angle
  1429. from math import atan2
  1430. try:
  1431. flattened = (input.to_3x3().inverted() @ projected)
  1432. except ValueError:
  1433. raise ValueError(f"Cannot align the matrix in {self} because it is degenerate.")
  1434. rotation = Matrix.Rotation(atan2(flattened.x, flattened.z), 4, y_axis)
  1435. matrix = rotation @ input.copy()
  1436. matrix.translation=input.translation
  1437. matrix[3][3] = input[3][3]
  1438. self.parameters['Matrix'] = matrix
  1439. self.prepared = True; self.executed = True
  1440. # NOTE: I tried other ways of setting the matrix, including composing
  1441. # it directly from the Y axis, the normalized projection of the align
  1442. # axis, and their cross-product. That only nearly worked.
  1443. # this calculation should not work better, but it does. Why?
  1444. class UtilityTransformationMatrix(MantisNode):
  1445. def __init__(self, signature, base_tree):
  1446. super().__init__(signature, base_tree)
  1447. inputs = [
  1448. "Operation" ,
  1449. "Vector" ,
  1450. "W" ,
  1451. ]
  1452. outputs = [
  1453. "Matrix" ,
  1454. ]
  1455. self.inputs.init_sockets(inputs)
  1456. self.outputs.init_sockets(outputs)
  1457. self.init_parameters()
  1458. self.node_type = "UTILITY"
  1459. def bPrepare(self, bContext = None,):
  1460. from mathutils import Matrix, Vector
  1461. if (operation := self.evaluate_input("Operation")) == 'ROTATE_AXIS_ANGLE':
  1462. # this can, will, and should fail if the axis is 0,0,0
  1463. self.parameters["Matrix"] = rotMat = Matrix.Rotation(self.evaluate_input("W"), 4, Vector(self.evaluate_input("Vector")).normalized())
  1464. elif (operation := self.evaluate_input("Operation")) == 'TRANSLATE':
  1465. m = Matrix.Identity(4)
  1466. if axis := self.evaluate_input("Vector"):
  1467. m[0][3]=axis[0];m[1][3]=axis[1];m[2][3]=axis[2]
  1468. self.parameters['Matrix'] = m
  1469. elif (operation := self.evaluate_input("Operation")) == 'SCALE':
  1470. self.parameters["Matrix"] = Matrix.Scale(self.evaluate_input("W"), 4, Vector(self.evaluate_input("Vector")).normalized())
  1471. else:
  1472. raise NotImplementedError(self.evaluate_input("Operation").__repr__()+ " Operation not yet implemented.")
  1473. self.prepared = True; self.executed = True
  1474. class UtilityIntToString(MantisNode):
  1475. def __init__(self, signature, base_tree):
  1476. super().__init__(signature, base_tree)
  1477. inputs = [
  1478. "Number" ,
  1479. "Zero Padding" ,
  1480. ]
  1481. outputs = [
  1482. "String" ,
  1483. ]
  1484. self.inputs.init_sockets(inputs)
  1485. self.outputs.init_sockets(outputs)
  1486. self.init_parameters()
  1487. self.node_type = "UTILITY"
  1488. def bPrepare(self, bContext = None,):
  1489. number = self.evaluate_input("Number")
  1490. zeroes = self.evaluate_input("Zero Padding")
  1491. # I'm casting to int because I want to support any number, even though the node asks for int.
  1492. self.parameters["String"] = str(int(number)).zfill(int(zeroes))
  1493. self.prepared = True; self.executed = True
  1494. class UtilityArrayGet(MantisNode):
  1495. def __init__(self, signature, base_tree):
  1496. super().__init__(signature, base_tree)
  1497. inputs = [
  1498. "Index" ,
  1499. "OoB Behaviour" ,
  1500. "Array" ,
  1501. ]
  1502. outputs = [
  1503. "Output" ,
  1504. ]
  1505. self.inputs.init_sockets(inputs)
  1506. self.outputs.init_sockets(outputs)
  1507. self.init_parameters()
  1508. self.node_type = "UTILITY"
  1509. def bPrepare(self, bContext = None,):
  1510. if len(self.hierarchy_dependencies)==0 and len(self.hierarchy_connections)==0 and \
  1511. len(self.connections)==0 and len(self.dependencies)==0:
  1512. self.prepared, self.executed = True, True
  1513. return #Either it is already done or it doesn't matter.
  1514. elif self.prepared == False:
  1515. # sort the array entries
  1516. for inp in self.inputs.values():
  1517. inp.links.sort(key=lambda a : -a.multi_input_sort_id)
  1518. oob = self.evaluate_input("OoB Behaviour")
  1519. index = self.evaluate_input("Index")
  1520. from .utilities import cap, wrap
  1521. # we must assume that the array has sent the correct number of links
  1522. if oob == 'WRAP':
  1523. index = index % len(self.inputs['Array'].links)
  1524. if oob == 'HOLD':
  1525. index = cap(index, len(self.inputs['Array'].links)-1)
  1526. array_choose_relink(self, [index], "Array", "Output")
  1527. self.prepared, self.executed = True, True
  1528. class UtilityArrayLength(MantisNode):
  1529. def __init__(self, signature, base_tree):
  1530. super().__init__(signature, base_tree)
  1531. inputs = [
  1532. "Array" ,
  1533. ]
  1534. outputs = [
  1535. "Length" ,
  1536. ]
  1537. self.inputs.init_sockets(inputs)
  1538. self.outputs.init_sockets(outputs)
  1539. self.init_parameters()
  1540. self.node_type = "UTILITY"
  1541. def bPrepare(self, bContext = None,):
  1542. self.parameters["Length"] = len(self.inputs["Array"].links)
  1543. self.prepared, self.executed = True, True
  1544. class UtilitySetBoneMatrixTail(MantisNode):
  1545. def __init__(self, signature, base_tree):
  1546. super().__init__(signature, base_tree)
  1547. inputs = {
  1548. "Matrix" ,
  1549. "Tail Location" ,
  1550. }
  1551. outputs = [
  1552. "Result" ,
  1553. ]
  1554. self.inputs.init_sockets(inputs)
  1555. self.outputs.init_sockets(outputs)
  1556. self.init_parameters()
  1557. self.node_type = "UTILITY"
  1558. def bPrepare(self, bContext = None,):
  1559. from mathutils import Matrix
  1560. matrix = self.evaluate_input("Matrix")
  1561. if matrix is None: matrix = Matrix.Identity(4)
  1562. #just do this for now lol
  1563. self.parameters["Result"] = matrix_from_head_tail(matrix.translation, self.evaluate_input("Tail Location"))
  1564. self.prepared = True; self.executed = True
  1565. class UtilityPrint(MantisNode):
  1566. def __init__(self, signature, base_tree):
  1567. super().__init__(signature, base_tree)
  1568. inputs = [
  1569. "Input" ,
  1570. ]
  1571. self.inputs.init_sockets(inputs)
  1572. self.init_parameters()
  1573. self.node_type = "UTILITY"
  1574. def bPrepare(self, bContext = None,):
  1575. if my_input := self.evaluate_input("Input"):
  1576. print("Preparation phase: ", wrapWhite(self), wrapGreen(my_input))
  1577. self.prepared = True
  1578. def bExecute(self, bContext = None,):
  1579. if my_input := self.evaluate_input("Input"):
  1580. print("Execution phase: ", wrapWhite(self), wrapGreen(my_input))
  1581. self.executed = True
  1582. class UtilityCompare(MantisNode):
  1583. def __init__(self, signature, base_tree):
  1584. super().__init__(signature, base_tree, CompareSockets)
  1585. self.init_parameters()
  1586. self.node_type = "UTILITY"
  1587. def bPrepare(self, bContext = None,):
  1588. operation=self.evaluate_input("Comparison")
  1589. a = self.evaluate_input("A")
  1590. b = self.evaluate_input("B")
  1591. if isinstance(a, str) and isinstance(b, str) and \
  1592. operation not in ['EQUAL', 'NOT_EQUAL']:
  1593. raise GraphError("Strings do not have numerical value to"
  1594. " compute greater than or less than.")
  1595. match operation:
  1596. case "EQUAL":
  1597. self.parameters["Result"] = a == b
  1598. case "NOT_EQUAL":
  1599. self.parameters["Result"] = a != b
  1600. case "GREATER_THAN":
  1601. self.parameters["Result"] = a > b
  1602. case "GREATER_THAN_EQUAL":
  1603. self.parameters["Result"] = a >= b
  1604. case "LESS_THAN":
  1605. self.parameters["Result"] = a < b
  1606. case "LESS_THAN_EQUAL":
  1607. self.parameters["Result"] = a <= b
  1608. self.prepared = True; self.executed = True
  1609. class UtilityChoose(MantisNode):
  1610. def __init__(self, signature, base_tree):
  1611. super().__init__(signature, base_tree)
  1612. inputs = [
  1613. "Condition" ,
  1614. "A" ,
  1615. "B" ,
  1616. ]
  1617. outputs = [
  1618. "Result" ,
  1619. ]
  1620. self.inputs.init_sockets(inputs)
  1621. self.outputs.init_sockets(outputs)
  1622. self.init_parameters()
  1623. self.node_type = "UTILITY"
  1624. def reset_execution(self):
  1625. prepared=self.prepared
  1626. super().reset_execution()
  1627. # prevent this node from attempting to prepare again.
  1628. self.prepared, self.executed = prepared, prepared
  1629. def bPrepare(self, bContext = None,):
  1630. if self.outputs['Result'].links: # otherwise this doesn't matter as it is not connected.
  1631. prGreen(f"Executing Choose Node {self}")
  1632. condition = self.evaluate_input("Condition")
  1633. if self.evaluate_input('A') is not None and self.evaluate_input('B') is not None:
  1634. self.parameters['Result'] = self.evaluate_input('B') if condition else self.evaluate_input('A')
  1635. elif self.evaluate_input('A') is None and self.evaluate_input('B') is None:
  1636. if condition: link = self.inputs['B'].links[0]
  1637. else: link = self.inputs['A'].links[0]
  1638. from_node = link.from_node; from_socket = link.from_socket
  1639. for link in self.outputs['Result'].links:
  1640. from_node.outputs[from_socket].connect(link.to_node, link.to_socket)
  1641. link.die()
  1642. self.flush_links()
  1643. # attempting to init the connections seems more error prone than leaving them be.
  1644. else:
  1645. raise GraphError(f"Choose Node {self} has incorrect types.")
  1646. self.prepared = True; self.executed = True