misc_nodes.py 65 KB

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