link_nodes.py 41 KB

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