link_containers.py 39 KB

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