link_nodes.py 40 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901
  1. from .node_common import *
  2. from bpy.types import Bone, NodeTree
  3. from .base_definitions import MantisNode, GraphError, FLOAT_EPSILON
  4. from .link_socket_templates import *
  5. def TellClasses():
  6. return [
  7. # special
  8. LinkInherit,
  9. # copy
  10. LinkCopyLocation,
  11. LinkCopyRotation,
  12. LinkCopyScale,
  13. LinkCopyTransforms,
  14. LinkTransformation,
  15. # limit
  16. LinkLimitLocation,
  17. LinkLimitRotation,
  18. LinkLimitScale,
  19. LinkLimitDistance,
  20. # tracking
  21. LinkStretchTo,
  22. LinkDampedTrack,
  23. LinkLockedTrack,
  24. LinkTrackTo,
  25. #misc
  26. LinkInheritConstraint,
  27. LinkArmature,
  28. # IK
  29. LinkInverseKinematics,
  30. LinkSplineIK,
  31. # stuff that snaps or limits a bone
  32. LinkFloor,
  33. LinkShrinkWrap,
  34. # Drivers
  35. LinkDrivenParameter,
  36. ]
  37. # set the name if it is available, otherwise just use the constraint's nice name
  38. set_constraint_name = lambda mantis_node : mantis_node.evaluate_input("Name") if mantis_node.evaluate_input("Name") else mantis_node.__class__.__name__
  39. class MantisLinkNode(MantisNode):
  40. def __init__(self, signature : tuple,
  41. base_tree : NodeTree,
  42. socket_templates : list[SockTemplate]=[]):
  43. super().__init__(signature, base_tree, socket_templates)
  44. self.node_type = 'LINK'
  45. self.prepared = True; self.bObject=[]
  46. def bTransformPass(self, bContext=None):
  47. parent_xForm_info = get_parent_xForm_info(self, 'Input Relationship')
  48. self.parameters['Output Relationship'] = parent_xForm_info
  49. def evaluate_input(self, input_name, index=0):
  50. # should catch 'Target', 'Pole Target' and ArmatureConstraint targets, too
  51. if ('Target' in input_name) and input_name not in ["Target Space", "Use Target Z"]:
  52. socket = self.inputs.get(input_name)
  53. if socket.is_linked:
  54. return socket.links[0].from_node
  55. return None
  56. else:
  57. return super().evaluate_input(input_name)
  58. def gen_property_socket_map(self) -> dict:
  59. props_sockets = super().gen_property_socket_map()
  60. if (os := self.inputs.get("Owner Space")) and os.is_connected and os.links[0].from_node.node_type == 'XFORM':
  61. del props_sockets['owner_space']
  62. if (ts := self.inputs.get("Target Space")) and ts.is_connected and ts.links[0].from_node.node_type == 'XFORM':
  63. del props_sockets['target_space']
  64. return props_sockets
  65. def set_custom_space(self):
  66. for c in self.bObject:
  67. if (os := self.inputs.get("Owner Space")) and os.is_connected and os.links[0].from_node.node_type == 'XFORM':
  68. c.owner_space='CUSTOM'
  69. xf = self.inputs["Owner Space"].links[0].from_node.bGetObject(mode="OBJECT")
  70. if isinstance(xf, Bone):
  71. c.space_object=self.inputs["Owner Space"].links[0].from_node.bGetParentArmature(); c.space_subtarget=xf.name
  72. else:
  73. c.space_object=xf
  74. if ts := self.inputs.get("Target_Space") and ts.is_connected and ts.links[0].from_node.node_type == 'XFORM':
  75. c.target_space='CUSTOM'
  76. xf = self.inputs["Target_Space Space"].links[0].from_node.bGetObject(mode="OBJECT")
  77. if isinstance(xf, Bone):
  78. c.space_object=self.inputs["Target_Space Space"].links[0].from_node.bGetParentArmature(); c.space_subtarget=xf.name
  79. else:
  80. c.space_object=xf
  81. def GetxForm(mantis_node, output_name="Output Relationship"):
  82. break_condition= lambda node : node.node_type=='XFORM'
  83. xforms = trace_line_up_branching(mantis_node, output_name, break_condition)
  84. return_me=[]
  85. for xf in xforms:
  86. if xf.node_type != 'XFORM':
  87. continue
  88. if xf in return_me:
  89. continue
  90. return_me.append(xf)
  91. return return_me
  92. def reset_execution(self):
  93. super().reset_execution()
  94. self.prepared = True; self.bObject = []
  95. def bFinalize(self, bContext=None):
  96. finish_drivers(self)
  97. #*#-------------------------------#++#-------------------------------#*#
  98. # L I N K N O D E S
  99. #*#-------------------------------#++#-------------------------------#*#
  100. class LinkInherit(MantisLinkNode):
  101. '''A node representing inheritance'''
  102. def __init__(self, signature, base_tree):
  103. super().__init__(signature, base_tree, LinkInheritSockets)
  104. self.init_parameters()
  105. self.set_traverse([('Parent', 'Inheritance')])
  106. def bTransformPass(self, bContext=None):
  107. parent_xForm_info = get_parent_xForm_info(self, 'Parent')
  108. self.parameters['Inheritance'] = parent_xForm_info
  109. self.executed = True
  110. def GetxForm(self):
  111. # I think this is only run in display update.
  112. trace = trace_single_line_up(self, "Inheritance")
  113. for node in trace[0]:
  114. if (node.node_type == 'XFORM'):
  115. return node
  116. raise GraphError("%s is not connected to a downstream xForm" % self)
  117. class LinkCopyLocation(MantisLinkNode):
  118. '''A node representing Copy Location'''
  119. def __init__(self, signature : tuple,
  120. base_tree : NodeTree,):
  121. super().__init__(signature, base_tree, LinkCopyLocationSockets)
  122. additional_parameters = { "Name":None }
  123. self.init_parameters(additional_parameters=additional_parameters)
  124. self.set_traverse([("Input Relationship", "Output Relationship")])
  125. def bRelationshipPass(self, context):
  126. prepare_parameters(self)
  127. for xf in self.GetxForm():
  128. c = xf.bGetObject().constraints.new('COPY_LOCATION')
  129. self.get_target_and_subtarget(c)
  130. print(wrapGreen("Creating ")+wrapWhite("Copy Location")+
  131. wrapGreen(" Constraint for bone: ") +
  132. wrapOrange(xf.bGetObject().name))
  133. if constraint_name := self.evaluate_input("Name"):
  134. c.name = constraint_name
  135. self.bObject.append(c)
  136. self.set_custom_space()
  137. props_sockets = self.gen_property_socket_map()
  138. evaluate_sockets(self, c, props_sockets)
  139. self.executed = True
  140. class LinkCopyRotation(MantisLinkNode):
  141. '''A node representing Copy Rotation'''
  142. def __init__(self, signature, base_tree):
  143. super().__init__(signature, base_tree, LinkCopyRotationSockets)
  144. additional_parameters = { "Name":None }
  145. self.init_parameters(additional_parameters=additional_parameters)
  146. self.set_traverse([("Input Relationship", "Output Relationship")])
  147. def bRelationshipPass(self, context):
  148. prepare_parameters(self)
  149. for xf in self.GetxForm():
  150. c = xf.bGetObject().constraints.new('COPY_ROTATION')
  151. self.get_target_and_subtarget(c)
  152. print(wrapGreen("Creating ")+wrapWhite("Copy Rotation")+
  153. wrapGreen(" Constraint for bone: ") +
  154. wrapOrange(xf.bGetObject().name))
  155. rotation_order = self.evaluate_input("RotationOrder")
  156. if ((rotation_order == 'QUATERNION') or (rotation_order == 'AXIS_ANGLE')):
  157. c.euler_order = 'AUTO'
  158. else:
  159. try:
  160. c.euler_order = rotation_order
  161. except TypeError: # it's a driver or incorrect
  162. c.euler_order = 'AUTO'
  163. if constraint_name := self.evaluate_input("Name"):
  164. c.name = constraint_name
  165. self.bObject.append(c)
  166. self.set_custom_space()
  167. props_sockets = self.gen_property_socket_map()
  168. evaluate_sockets(self, c, props_sockets)
  169. self.executed = True
  170. class LinkCopyScale(MantisLinkNode):
  171. '''A node representing Copy Scale'''
  172. def __init__(self, signature, base_tree):
  173. super().__init__(signature, base_tree, LinkCopyScaleSockets)
  174. additional_parameters = { "Name":None }
  175. self.init_parameters(additional_parameters=additional_parameters)
  176. self.set_traverse([("Input Relationship", "Output Relationship")])
  177. def bRelationshipPass(self, context):
  178. prepare_parameters(self)
  179. for xf in self.GetxForm():
  180. c = xf.bGetObject().constraints.new('COPY_SCALE')
  181. self.get_target_and_subtarget(c)
  182. print(wrapGreen("Creating ")+wrapWhite("Copy Scale")+
  183. wrapGreen(" Constraint for bone: ") +
  184. wrapOrange(xf.bGetObject().name))
  185. if constraint_name := self.evaluate_input("Name"):
  186. c.name = constraint_name
  187. self.bObject.append(c)
  188. if self.inputs["Owner Space"].is_connected and self.inputs["Owner Space"].links[0].from_node.node_type == 'XFORM':
  189. c.owner_space='CUSTOM'
  190. xf = self.inputs["Owner Space"].links[0].from_node.bGetObject(mode="OBJECT")
  191. if isinstance(xf, Bone):
  192. c.space_object=self.inputs["Owner Space"].links[0].from_node.bGetParentArmature(); c.space_subtarget=xf.name
  193. else:
  194. c.space_object=xf
  195. if self.inputs["Target Space"].is_connected and self.inputs["Target Space"].links[0].from_node.node_type == 'XFORM':
  196. c.target_space='CUSTOM'
  197. xf = self.inputs["Target Space"].links[0].from_node.bGetObject(mode="OBJECT")
  198. if isinstance(xf, Bone):
  199. c.space_object=self.inputs["Owner Space"].links[0].from_node.bGetParentArmature(); c.space_subtarget=xf.name
  200. else:
  201. c.space_object=xf
  202. props_sockets = self.gen_property_socket_map()
  203. evaluate_sockets(self, c, props_sockets)
  204. self.executed = True
  205. class LinkCopyTransforms(MantisLinkNode):
  206. '''A node representing Copy Transfoms'''
  207. def __init__(self, signature, base_tree):
  208. super().__init__(signature, base_tree, LinkCopyTransformsSockets)
  209. additional_parameters = { "Name":None }
  210. self.init_parameters(additional_parameters=additional_parameters)
  211. self.set_traverse([("Input Relationship", "Output Relationship")])
  212. def bRelationshipPass(self, context):
  213. prepare_parameters(self)
  214. for xf in self.GetxForm():
  215. c = xf.bGetObject().constraints.new('COPY_TRANSFORMS')
  216. self.get_target_and_subtarget(c)
  217. print(wrapGreen("Creating ")+wrapWhite("Copy Transforms")+
  218. wrapGreen(" Constraint for bone: ") +
  219. wrapOrange(xf.bGetObject().name))
  220. if constraint_name := self.evaluate_input("Name"):
  221. c.name = constraint_name
  222. self.bObject.append(c)
  223. self.set_custom_space()
  224. props_sockets = self.gen_property_socket_map()
  225. evaluate_sockets(self, c, props_sockets)
  226. self.executed = True
  227. class LinkTransformation(MantisLinkNode):
  228. '''A node representing Copy Transfoms'''
  229. def __init__(self, signature, base_tree):
  230. super().__init__(signature, base_tree, LinkTransformationSockets)
  231. self.init_parameters(additional_parameters={"Name":None })
  232. self.set_traverse([("Input Relationship", "Output Relationship")])
  233. def ui_modify_socket(self, ui_socket, socket_name=None):
  234. from_suffix, to_suffix = '', ''
  235. if self.evaluate_input("Map From") == 'ROTATION': from_suffix='_rot'
  236. elif self.evaluate_input("Map From") == 'SCALE': from_suffix='_scale'
  237. if self.evaluate_input("Map To") == 'ROTATION': to_suffix='_rot'
  238. elif self.evaluate_input("Map To") == 'SCALE': to_suffix='_scale'
  239. if ('To' in ui_socket.name or 'From' in ui_socket.name) and (from_suffix or to_suffix):
  240. for s_temp in self.socket_templates:
  241. if s_temp.name == ui_socket.name: break
  242. if 'from' in s_temp.blender_property:
  243. socket_name=s_temp.blender_property+from_suffix
  244. else:
  245. socket_name=s_temp.blender_property+to_suffix
  246. return self.update_socket_value(socket_name, ui_socket.default_value)
  247. return super().ui_modify_socket(ui_socket, socket_name)
  248. def bRelationshipPass(self, context):
  249. prepare_parameters(self)
  250. for xf in self.GetxForm():
  251. c = xf.bGetObject().constraints.new('TRANSFORM')
  252. self.get_target_and_subtarget(c)
  253. print(wrapGreen("Creating ")+wrapWhite("Transformation")+
  254. wrapGreen(" Constraint for bone: ") +
  255. wrapOrange(xf.bGetObject().name))
  256. if constraint_name := self.evaluate_input("Name"):
  257. c.name = constraint_name
  258. self.bObject.append(c)
  259. self.set_custom_space()
  260. props_sockets = self.gen_property_socket_map()
  261. # we have to fix the blender-property for scale/rotation
  262. # because Blender stores these separately.
  263. # I do not care that this code is ugly.
  264. from_suffix, to_replace = '', ''
  265. if self.evaluate_input("Map From") == 'ROTATION':
  266. from_suffix='_rot'
  267. elif self.evaluate_input("Map From") == 'SCALE':
  268. from_suffix='_scale'
  269. if self.evaluate_input("Map To") == 'ROTATION':
  270. to_replace='_rot'
  271. elif self.evaluate_input("Map To") == 'SCALE':
  272. to_replace='_scale'
  273. if from_suffix:
  274. for axis in ['x', 'y', 'z']:
  275. stub='from_min_'+axis
  276. props_sockets[stub+from_suffix]=props_sockets[stub]
  277. del props_sockets[stub]
  278. stub='from_max_'+axis
  279. props_sockets[stub+from_suffix]=props_sockets[stub]
  280. del props_sockets[stub]
  281. if to_replace:
  282. for axis in ['x', 'y', 'z']:
  283. stub='to_min_'+axis
  284. props_sockets[stub+to_replace]=props_sockets[stub]
  285. del props_sockets[stub]
  286. stub='to_max_'+axis
  287. props_sockets[stub+to_replace]=props_sockets[stub]
  288. del props_sockets[stub]
  289. evaluate_sockets(self, c, props_sockets)
  290. self.executed = True
  291. class LinkLimitLocation(MantisLinkNode):
  292. def __init__(self, signature, base_tree):
  293. super().__init__(signature, base_tree, LinkLimitLocationScaleSockets)
  294. self.init_parameters(additional_parameters={ "Name":None })
  295. self.set_traverse([("Input Relationship", "Output Relationship")])
  296. def bRelationshipPass(self, context):
  297. prepare_parameters(self)
  298. for xf in self.GetxForm():
  299. c = xf.bGetObject().constraints.new('LIMIT_LOCATION')
  300. print(wrapGreen("Creating ")+wrapWhite("Limit Location")+
  301. wrapGreen(" Constraint for bone: ") +
  302. wrapOrange(xf.bGetObject().name))
  303. if constraint_name := self.evaluate_input("Name"):
  304. c.name = constraint_name
  305. self.bObject.append(c)
  306. self.set_custom_space()
  307. props_sockets = self.gen_property_socket_map()
  308. evaluate_sockets(self, c, props_sockets)
  309. self.executed = True
  310. class LinkLimitRotation(MantisLinkNode):
  311. def __init__(self, signature, base_tree):
  312. super().__init__(signature, base_tree, LinkLimitRotationSockets)
  313. self.init_parameters(additional_parameters={ "Name":None })
  314. self.set_traverse([("Input Relationship", "Output Relationship")])
  315. def bRelationshipPass(self, context):
  316. prepare_parameters(self)
  317. for xf in self.GetxForm():
  318. c = xf.bGetObject().constraints.new('LIMIT_ROTATION')
  319. print(wrapGreen("Creating ")+wrapWhite("Limit Rotation")+
  320. wrapGreen(" Constraint for bone: ") +
  321. wrapOrange(xf.bGetObject().name))
  322. if constraint_name := self.evaluate_input("Name"):
  323. c.name = constraint_name
  324. self.bObject.append(c)
  325. self.set_custom_space()
  326. props_sockets = self.gen_property_socket_map()
  327. evaluate_sockets(self, c, props_sockets)
  328. self.executed = True
  329. class LinkLimitScale(MantisLinkNode):
  330. def __init__(self, signature, base_tree):
  331. super().__init__(signature, base_tree, LinkLimitLocationScaleSockets)
  332. self.init_parameters(additional_parameters={ "Name":None })
  333. self.set_traverse([("Input Relationship", "Output Relationship")])
  334. def bRelationshipPass(self, context):
  335. prepare_parameters(self)
  336. for xf in self.GetxForm():
  337. c = xf.bGetObject().constraints.new('LIMIT_SCALE')
  338. print(wrapGreen("Creating ")+wrapWhite("Limit Scale")+
  339. wrapGreen(" Constraint for bone: ") +
  340. wrapOrange(xf.bGetObject().name))
  341. if constraint_name := self.evaluate_input("Name"):
  342. c.name = constraint_name
  343. self.bObject.append(c)
  344. self.set_custom_space()
  345. props_sockets = self.gen_property_socket_map()
  346. evaluate_sockets(self, c, props_sockets)
  347. self.executed = True
  348. class LinkLimitDistance(MantisLinkNode):
  349. def __init__(self, signature, base_tree):
  350. super().__init__(signature, base_tree, LinkLimitDistanceSockets)
  351. self.init_parameters(additional_parameters={ "Name":None })
  352. self.set_traverse([("Input Relationship", "Output Relationship")])
  353. def bRelationshipPass(self, context):
  354. prepare_parameters(self)
  355. for xf in self.GetxForm():
  356. print(wrapGreen("Creating ")+wrapWhite("Limit Distance")+
  357. wrapGreen(" Constraint for bone: ") +
  358. wrapOrange(xf.bGetObject().name))
  359. c = xf.bGetObject().constraints.new('LIMIT_DISTANCE')
  360. self.get_target_and_subtarget(c)
  361. if constraint_name := self.evaluate_input("Name"):
  362. c.name = constraint_name
  363. self.bObject.append(c)
  364. self.set_custom_space()
  365. props_sockets = self.gen_property_socket_map()
  366. evaluate_sockets(self, c, props_sockets)
  367. self.executed = True
  368. # Tracking
  369. class LinkStretchTo(MantisLinkNode):
  370. def __init__(self, signature, base_tree):
  371. super().__init__(signature, base_tree, LinkStretchToSockets)
  372. self.init_parameters(additional_parameters={ "Name":None })
  373. self.set_traverse([("Input Relationship", "Output Relationship")])
  374. def bRelationshipPass(self, context):
  375. prepare_parameters(self)
  376. for xf in self.GetxForm():
  377. print(wrapGreen("Creating ")+wrapWhite("Stretch-To")+
  378. wrapGreen(" Constraint for bone: ") +
  379. wrapOrange(xf.bGetObject().name))
  380. c = xf.bGetObject().constraints.new('STRETCH_TO')
  381. self.get_target_and_subtarget(c)
  382. if constraint_name := self.evaluate_input("Name"):
  383. c.name = constraint_name
  384. self.bObject.append(c)
  385. props_sockets = self.gen_property_socket_map()
  386. evaluate_sockets(self, c, props_sockets)
  387. if (self.evaluate_input("Original Length") == 0):
  388. # this is meant to be set automatically.
  389. c.rest_length = xf.bGetObject().bone.length
  390. self.executed = True
  391. class LinkDampedTrack(MantisLinkNode):
  392. def __init__(self, signature, base_tree):
  393. super().__init__(signature, base_tree, LinkDampedTrackSockets)
  394. self.init_parameters(additional_parameters={ "Name":None })
  395. self.set_traverse([("Input Relationship", "Output Relationship")])
  396. def bRelationshipPass(self, context):
  397. prepare_parameters(self)
  398. for xf in self.GetxForm():
  399. print(wrapGreen("Creating ")+wrapWhite("Damped Track")+
  400. wrapGreen(" Constraint for bone: ") +
  401. wrapOrange(xf.bGetObject().name))
  402. c = xf.bGetObject().constraints.new('DAMPED_TRACK')
  403. self.get_target_and_subtarget(c)
  404. if constraint_name := self.evaluate_input("Name"):
  405. c.name = constraint_name
  406. self.bObject.append(c)
  407. props_sockets = self.gen_property_socket_map()
  408. evaluate_sockets(self, c, props_sockets)
  409. self.executed = True
  410. class LinkLockedTrack(MantisLinkNode):
  411. def __init__(self, signature, base_tree):
  412. super().__init__(signature, base_tree,LinkLockedTrackSockets)
  413. self.init_parameters(additional_parameters={"Name":None })
  414. self.set_traverse([("Input Relationship", "Output Relationship")])
  415. def bRelationshipPass(self, context):
  416. prepare_parameters(self)
  417. for xf in self.GetxForm():
  418. print(wrapGreen("Creating ")+wrapWhite("Locked Track")+
  419. wrapGreen(" Constraint for bone: ") +
  420. wrapOrange(xf.bGetObject().name))
  421. c = xf.bGetObject().constraints.new('LOCKED_TRACK')
  422. self.get_target_and_subtarget(c)
  423. if constraint_name := self.evaluate_input("Name"):
  424. c.name = constraint_name
  425. self.bObject.append(c)
  426. props_sockets = self.gen_property_socket_map()
  427. evaluate_sockets(self, c, props_sockets)
  428. self.executed = True
  429. class LinkTrackTo(MantisLinkNode):
  430. def __init__(self, signature, base_tree):
  431. super().__init__(signature, base_tree, LinkTrackToSockets)
  432. self.init_parameters(additional_parameters={"Name":None })
  433. self.set_traverse([("Input Relationship", "Output Relationship")])
  434. def bRelationshipPass(self, context):
  435. prepare_parameters(self)
  436. for xf in self.GetxForm():
  437. print(wrapGreen("Creating ")+wrapWhite("Track-To")+
  438. wrapGreen(" Constraint for bone: ") +
  439. wrapOrange(xf.bGetObject().name))
  440. c = xf.bGetObject().constraints.new('TRACK_TO')
  441. self.get_target_and_subtarget(c)
  442. if constraint_name := self.evaluate_input("Name"):
  443. c.name = constraint_name
  444. self.bObject.append(c)
  445. props_sockets = self.gen_property_socket_map()
  446. evaluate_sockets(self, c, props_sockets)
  447. self.executed = True
  448. class LinkInheritConstraint(MantisLinkNode):
  449. def __init__(self, signature, base_tree):
  450. super().__init__(signature, base_tree, LinkInheritConstraintSockets)
  451. self.init_parameters(additional_parameters={"Name":None })
  452. self.set_traverse([("Input Relationship", "Output Relationship")])
  453. def bRelationshipPass(self, context):
  454. prepare_parameters(self)
  455. for xf in self.GetxForm():
  456. print(wrapGreen("Creating ")+wrapWhite("Child-Of")+
  457. wrapGreen(" Constraint for bone: ") +
  458. wrapOrange(xf.bGetObject().name))
  459. c = xf.bGetObject().constraints.new('CHILD_OF')
  460. self.get_target_and_subtarget(c)
  461. if constraint_name := self.evaluate_input("Name"):
  462. c.name = constraint_name
  463. self.bObject.append(c)
  464. props_sockets = self.gen_property_socket_map()
  465. evaluate_sockets(self, c, props_sockets)
  466. c.set_inverse_pending
  467. self.executed = True
  468. class LinkInverseKinematics(MantisLinkNode):
  469. def __init__(self, signature, base_tree):
  470. super().__init__(signature, base_tree, LinkInverseKinematicsSockets)
  471. self.init_parameters(additional_parameters={"Name":None })
  472. self.set_traverse([("Input Relationship", "Output Relationship")])
  473. def get_base_ik_bone(self, ik_bone):
  474. chain_length : int = (self.evaluate_input("Chain Length"))
  475. if not isinstance(chain_length, (int, float)):
  476. raise GraphError(f"Chain Length must be an integer number in {self}::Chain Length")
  477. if chain_length == 0:
  478. chain_length = int("inf")
  479. base_ik_bone = ik_bone; i=1
  480. while (i<chain_length) and (base_ik_bone.parent):
  481. base_ik_bone=base_ik_bone.parent; i+=1
  482. return base_ik_bone
  483. # We need to do the calculation in a "full circle", meaning the pole_angle
  484. # can go over pi or less than -pi - but the actuall constraint value must
  485. # be clamped in that range.
  486. # so we simply wrap the value.
  487. # not very efficient but it's OK
  488. def set_pole_angle(self, constraint, angle: float) -> None:
  489. from math import pi
  490. from .utilities import wrap
  491. constraint.pole_angle = wrap(-pi, pi, angle)
  492. def calc_pole_angle_pre(self, c, ik_bone):
  493. """
  494. This function gets us most of the way to a correct IK pole angle. Unfortunately,
  495. due to the unpredictable nature of the iterative IK calculation, I can't figure
  496. out an exact solution. So we do a bisect search in calc_pole_angle_post().
  497. """
  498. # TODO: instead of these checks, convert all to armature local space. But this is tedious.
  499. if not c.target:
  500. raise GraphError(f"IK Constraint {self} must have target.")
  501. elif c.target.type != "ARMATURE":
  502. raise NotImplementedError(f"Currently, IK Constraint Target for {self} must be a bone within the same armature.")
  503. if c.pole_target.type != "ARMATURE":
  504. raise NotImplementedError(f"Currently, IK Constraint Pole Target for {self} must be a bone within the same armature.")
  505. ik_handle = c.target.pose.bones[c.subtarget]
  506. if ik_handle.id_data != ik_bone.id_data:
  507. raise NotImplementedError(f"Currently, IK Constraint Target for {self} must be a bone within the same armature.")
  508. ik_pole = c.pole_target.pose.bones[c.pole_subtarget]
  509. if ik_pole.id_data != ik_bone.id_data:
  510. raise NotImplementedError(f"Currently,IK Constraint Pole Target for {self} must be a bone within the same armature.")
  511. base_ik_bone = self.get_base_ik_bone(ik_bone)
  512. start_effector = base_ik_bone.bone.head_local
  513. end_effector = ik_handle.bone.head_local
  514. pole_location = ik_pole.bone.head_local
  515. # this is the X-Axis of the bone's rest-pose, added to its bone
  516. knee_location = base_ik_bone.bone.matrix_local.col[0].xyz+start_effector
  517. ik_axis = (end_effector-start_effector).normalized()
  518. from .utilities import project_point_to_plane
  519. pole_planar_projection = project_point_to_plane(pole_location, start_effector, ik_axis)
  520. # this planar projection is necessary because the IK axis is different than the base_bone's y axis
  521. planar_projection = project_point_to_plane(knee_location, start_effector, ik_axis)
  522. knee_direction =(planar_projection - start_effector).normalized()
  523. pole_direction =(pole_planar_projection - start_effector).normalized()
  524. return knee_direction.angle(pole_direction)
  525. def calc_pole_angle_post(self, c, ik_bone, context):
  526. """
  527. This function should give us a completely accurate result for IK.
  528. """
  529. from time import time
  530. start_time=time()
  531. def signed_angle(vector_u, vector_v, normal):
  532. # it seems that this fails if the vectors are exactly aligned under certain circumstances.
  533. angle = vector_u.angle(vector_v, 0.0) # So we use a fallback of 0
  534. # Normal specifies orientation
  535. if angle != 0 and vector_u.cross(vector_v).angle(normal) < 1:
  536. angle = -angle
  537. return angle
  538. # we have already checked for valid data.
  539. ik_handle = c.target.pose.bones[c.subtarget]
  540. base_ik_bone = self.get_base_ik_bone(ik_bone)
  541. start_effector = base_ik_bone.bone.head_local
  542. angle = c.pole_angle
  543. dg = context.view_layer.depsgraph
  544. dg.update()
  545. ik_axis = (ik_handle.bone.head_local-start_effector).normalized()
  546. center_point = start_effector +(ik_axis*base_ik_bone.bone.length)
  547. knee_direction = base_ik_bone.bone.tail_local - center_point
  548. current_knee_direction = base_ik_bone.tail-center_point
  549. error=signed_angle(current_knee_direction, knee_direction, ik_axis)
  550. if error == 0:
  551. prGreen("No Fine-tuning needed."); return
  552. # Flip it if needed
  553. dot_before=current_knee_direction.dot(knee_direction)
  554. if dot_before < 0 and angle!=0: # then it is not aligned and we should check the inverse
  555. angle = -angle; c.pole_angle=angle
  556. dg.update()
  557. current_knee_direction = base_ik_bone.tail-center_point
  558. dot_after=current_knee_direction.dot(knee_direction)
  559. if dot_after < dot_before: # they are somehow less aligned
  560. prPurple("Mantis has gone down an unexpected code path. Please report this as a bug.")
  561. angle = -angle; self.set_pole_angle(c, angle)
  562. dg.update()
  563. # now we can do a bisect search to find the best value.
  564. error_threshhold = FLOAT_EPSILON
  565. max_iterations=600
  566. error=signed_angle(current_knee_direction, knee_direction, ik_axis)
  567. if error == 0:
  568. prGreen("No Fine-tuning needed."); return
  569. angle+=error
  570. alt_angle = angle+(error*-2) # should be very near the center when flipped here
  571. # we still need to bisect search because the relationship of pole_angle <==> error is somewhat unpredictable
  572. upper_bounds = alt_angle if alt_angle > angle else angle
  573. lower_bounds = alt_angle if alt_angle < angle else angle
  574. i, error_identical = 0, 0
  575. while ( True ):
  576. if (i>=max_iterations):
  577. prOrange(f"IK Pole Angle Set reached max iterations of {i-error_identical} in {time()-start_time} seconds")
  578. break
  579. if (abs(error)<error_threshhold) or (upper_bounds<=lower_bounds) or (error_identical > 3):
  580. prPurple(f"IK Pole Angle Set converged after {i-error_identical} iterations with error={error} in {time()-start_time} seconds")
  581. break
  582. # get the center-point betweeen the bounds
  583. try_angle = lower_bounds + (upper_bounds-lower_bounds)/2
  584. self.set_pole_angle(c, try_angle); dg.update()
  585. prev_error = error
  586. error = signed_angle((base_ik_bone.tail-center_point), knee_direction, ik_axis)
  587. error_identical+= int(error == prev_error)
  588. if error>0: upper_bounds=try_angle
  589. if error<0: lower_bounds=try_angle
  590. i+=1
  591. def bRelationshipPass(self, context):
  592. prepare_parameters(self)
  593. for xf in self.GetxForm():
  594. print(wrapGreen("Creating ")+wrapOrange("Inverse Kinematics")+
  595. wrapGreen(" Constraint for bone: ") +
  596. wrapOrange(xf.bGetObject().name))
  597. ik_bone = xf.bGetObject()
  598. c = xf.bGetObject().constraints.new('IK')
  599. self.get_target_and_subtarget(c)
  600. self.get_target_and_subtarget(c, input_name = 'Pole Target')
  601. if constraint_name := self.evaluate_input("Name"):
  602. c.name = constraint_name
  603. self.bObject.append(c)
  604. c.chain_count = 1 # so that, if there are errors, this doesn't print
  605. # a whole bunch of circular dependency crap from having infinite chain length
  606. if (c.pole_target):
  607. self.set_pole_angle(c, self.calc_pole_angle_pre(c, ik_bone))
  608. props_sockets = self.gen_property_socket_map()
  609. evaluate_sockets(self, c, props_sockets)
  610. c.use_location = self.evaluate_input("Position") > 0
  611. c.use_rotation = self.evaluate_input("Rotation") > 0
  612. self.executed = True
  613. def bFinalize(self, bContext = None):
  614. # adding a test here
  615. if bContext:
  616. for i, constraint in enumerate(self.bObject):
  617. ik_bone = self.GetxForm()[i].bGetObject(mode='POSE')
  618. if constraint.pole_target:
  619. prWhite(f"Fine-tuning IK Pole Angle for {self}")
  620. # make sure to enable it first
  621. enabled_before = constraint.mute
  622. constraint.mute = False
  623. self.calc_pole_angle_post(constraint, ik_bone, bContext)
  624. constraint.mute = enabled_before
  625. super().bFinalize(bContext)
  626. def ik_report_error(pb, context, do_print=False):
  627. dg = context.view_layer.depsgraph
  628. dg.update()
  629. loc1, rot_quaternion1, scl1 = pb.matrix.decompose()
  630. loc2, rot_quaternion2, scl2 = pb.bone.matrix_local.decompose()
  631. location_error=(loc1-loc2).length
  632. rotation_error = rot_quaternion1.rotation_difference(rot_quaternion2).angle
  633. scale_error = (scl1-scl2).length
  634. if location_error < FLOAT_EPSILON: location_error = 0
  635. if abs(rotation_error) < FLOAT_EPSILON: rotation_error = 0
  636. if scale_error < FLOAT_EPSILON: scale_error = 0
  637. if do_print:
  638. print (f"IK Location Error: {location_error}")
  639. print (f"IK Rotation Error: {rotation_error}")
  640. print (f"IK Scale Error : {scale_error}")
  641. return (location_error, rotation_error, scale_error)
  642. # This is kinda a weird design decision?
  643. class LinkDrivenParameter(MantisLinkNode):
  644. '''A node representing an armature object'''
  645. def __init__(self, signature, base_tree):
  646. super().__init__(signature, base_tree, LinkDrivenParameterSockets)
  647. self.init_parameters(additional_parameters={ "Name":None })
  648. self.set_traverse([("Input Relationship", "Output Relationship")])
  649. def bRelationshipPass(self, bContext = None,):
  650. prepare_parameters(self)
  651. prGreen("Executing Driven Parameter node")
  652. prop = self.evaluate_input("Parameter")
  653. index = self.evaluate_input("Index")
  654. value = self.evaluate_input("Value")
  655. for xf in self.GetxForm():
  656. ob = xf.bGetObject(mode="POSE")
  657. # IMPORTANT: this node only works on pose bone attributes.
  658. self.bObject.append(ob)
  659. length=1
  660. if hasattr(ob, prop):
  661. try:
  662. length = len(getattr(ob, prop))
  663. except TypeError:
  664. pass
  665. except AttributeError:
  666. pass
  667. else:
  668. raise AttributeError(f"Cannot Set value {prop} on object because it does not exist.")
  669. def_value = 0.0
  670. if length>1:
  671. def_value=[0.0]*length
  672. self.parameters["Value"] = tuple( 0.0 if i != index else value for i in range(length))
  673. props_sockets = {
  674. prop: ("Value", def_value)
  675. }
  676. evaluate_sockets(self, ob, props_sockets)
  677. self.executed = True
  678. def bFinalize(self, bContext = None):
  679. driver = self.evaluate_input("Value")
  680. try:
  681. for i, val in enumerate(self.parameters["Value"]):
  682. from .drivers import MantisDriver
  683. if isinstance(val, MantisDriver):
  684. driver["ind"] = i
  685. val = driver
  686. except AttributeError:
  687. self.parameters["Value"] = driver
  688. except TypeError:
  689. self.parameters["Value"] = driver
  690. super().bFinalize(bContext)
  691. class LinkArmature(MantisLinkNode):
  692. '''A node representing an armature object'''
  693. def __init__(self, signature, base_tree,):
  694. super().__init__(signature, base_tree, LinkArmatureSockets)
  695. self.init_parameters(additional_parameters={"Name":None })
  696. self.set_traverse([("Input Relationship", "Output Relationship")])
  697. setup_custom_props(self) # <-- this takes care of the runtime-added sockets
  698. def bRelationshipPass(self, bContext = None,):
  699. prepare_parameters(self)
  700. for xf in self.GetxForm():
  701. print(wrapGreen("Creating ")+wrapOrange("Armature")+
  702. wrapGreen(" Constraint for bone: ") +
  703. wrapOrange(xf.bGetObject().name))
  704. c = xf.bGetObject().constraints.new('ARMATURE')
  705. if constraint_name := self.evaluate_input("Name"):
  706. c.name = constraint_name
  707. self.bObject.append(c)
  708. # get number of targets
  709. num_targets = len( list(self.inputs.values())[6:] )//2
  710. props_sockets = self.gen_property_socket_map()
  711. targets_weights = {}
  712. for i in range(num_targets):
  713. target = c.targets.new()
  714. target_input_name = list(self.inputs.keys())[i*2+6 ]
  715. weight_input_name = list(self.inputs.keys())[i*2+6+1]
  716. self.get_target_and_subtarget(target, target_input_name)
  717. weight_value=self.evaluate_input(weight_input_name)
  718. if not isinstance(weight_value, float):
  719. weight_value=0
  720. targets_weights[i]=weight_value
  721. props_sockets["targets[%d].weight" % i] = (weight_input_name, 0)
  722. # targets_weights.append({"weight":(weight_input_name, 0)})
  723. evaluate_sockets(self, c, props_sockets)
  724. for target, value in targets_weights.items():
  725. c.targets[target].weight=value
  726. self.executed = True
  727. class LinkSplineIK(MantisLinkNode):
  728. '''A node representing an armature object'''
  729. def __init__(self, signature, base_tree):
  730. super().__init__(signature, base_tree, LinkSplineIKSockets)
  731. self.init_parameters(additional_parameters={"Name":None })
  732. self.set_traverse([("Input Relationship", "Output Relationship")])
  733. def bRelationshipPass(self, bContext = None,):
  734. prepare_parameters(self)
  735. if not self.inputs['Target'].is_linked:
  736. raise GraphError(f"ERROR: {self} is not connected to a target curve.")
  737. for xf in self.GetxForm():
  738. print(wrapGreen("Creating ")+wrapOrange("Spline-IK")+
  739. wrapGreen(" Constraint for bone: ") +
  740. wrapOrange(xf.bGetObject().name))
  741. c = xf.bGetObject().constraints.new('SPLINE_IK')
  742. # set the spline - we need to get the right one
  743. spline_index = self.evaluate_input("Spline Index")
  744. from .utilities import get_extracted_spline_object
  745. proto_curve = self.inputs['Target'].links[0].from_node.bGetObject()
  746. curve = get_extracted_spline_object(proto_curve, spline_index, self.mContext)
  747. # link it to the view layer
  748. if (curve.name not in bContext.view_layer.active_layer_collection.collection.objects):
  749. bContext.view_layer.active_layer_collection.collection.objects.link(curve)
  750. c.target=curve
  751. if constraint_name := self.evaluate_input("Name"):
  752. c.name = constraint_name
  753. self.bObject.append(c)
  754. props_sockets = self.gen_property_socket_map()
  755. evaluate_sockets(self, c, props_sockets)
  756. self.executed = True
  757. class LinkFloor(MantisLinkNode):
  758. '''A node representing an armature object'''
  759. def __init__(self, signature, base_tree,):
  760. super().__init__(signature, base_tree, LinkFloorSockets)
  761. self.init_parameters(additional_parameters={"Name":None })
  762. self.set_traverse([("Input Relationship", "Output Relationship")])
  763. def bRelationshipPass(self, bContext = None,):
  764. prepare_parameters(self)
  765. for xf in self.GetxForm():
  766. print(wrapGreen("Creating ")+wrapOrange("Floor")+
  767. wrapGreen(" Constraint for bone: ") +
  768. wrapOrange(xf.bGetObject().name))
  769. c = xf.bGetObject().constraints.new('FLOOR')
  770. self.get_target_and_subtarget(c)
  771. if constraint_name := self.evaluate_input("Name"):
  772. c.name = constraint_name
  773. self.bObject.append(c)
  774. self.set_custom_space()
  775. props_sockets = self.gen_property_socket_map()
  776. evaluate_sockets(self, c, props_sockets)
  777. self.executed = True
  778. class LinkShrinkWrap(MantisLinkNode):
  779. '''A node representing a shrinkwrap relationship.'''
  780. def __init__(self, signature, base_tree,):
  781. super().__init__(signature, base_tree, LinkShrinkWrapSockets)
  782. self.init_parameters(additional_parameters={"Name":None })
  783. self.set_traverse([("Input Relationship", "Output Relationship")])
  784. def bRelationshipPass(self, bContext = None,):
  785. prepare_parameters(self)
  786. for xf in self.GetxForm():
  787. print(wrapGreen("Creating ")+wrapOrange("Shrinkwrap")+
  788. wrapGreen(" Constraint for bone: ") +
  789. wrapOrange(xf.bGetObject().name))
  790. c = xf.bGetObject().constraints.new('SHRINKWRAP')
  791. self.get_target_and_subtarget(c)
  792. if constraint_name := self.evaluate_input("Name"):
  793. c.name = constraint_name
  794. props_sockets = self.gen_property_socket_map()
  795. print (props_sockets.keys())
  796. self.bObject.append(c)
  797. # self.set_custom_space() # this needs to be overridden for me TODO
  798. evaluate_sockets(self, c, props_sockets)
  799. self.executed = True