link_containers.py 38 KB

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