misc_nodes.py 75 KB

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