misc_nodes.py 83 KB

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