misc_nodes.py 62 KB

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