upvector calculater for Autodesk Maya

CalculateUpVector.py

# -*- coding: utf-8 -*-
"""
FK から アップベクタの位置を逆算
参考：
    http://www.comtec.daikin.co.jp/DC/UsersNotes/Ritaro/tutorial/mobu_02/#1
"""
__author__ = 'yamahigashi'
__version__ = '1.0.0'

# ==================================================================================

maya_useNewAPI = True
TYPE_ID = 0x00076564

import sys
from math import cos, sin

import maya.api.OpenMaya as api
from maya.api.OpenMaya import MVector
from maya.api.OpenMaya import MTransformationMatrix


FULL_EXTENSION_TOLERANCE = api.MFloatVector.kTolerance * 10  # 何回か計算してるので誤差増やす


# ==================================================================================
#  helper functions
def _clamp(n, min_n, max_n):
    return max(min(max_n, n), min_n)
# ==================================================================================


class CalculateUpVector(api.MPxNode):

    type_name = 'calculateUpVector'
    type_id = api.MTypeId(TYPE_ID)

    @classmethod
    def initialize(cls):

        fnMatrix = api.MFnMatrixAttribute()
        fnUnit = api.MFnUnitAttribute()
        fnNumeric = api.MFnNumericAttribute()
        # fnEnum = api.MFnEnumAttribute()

        # ---------------------------- inputs ------------------------------- #
        cls.aRootParentMatrix = fnMatrix.create('rootParentMatrix', 'rpm')
        cls.aRootTranslateX = fnUnit.create('rootTranslateX', 'rtx', api.MFnUnitAttribute.kDistance, 0.)
        cls.aRootTranslateY = fnUnit.create('rootTranslateY', 'rty', api.MFnUnitAttribute.kDistance, 0.)
        cls.aRootTranslateZ = fnUnit.create('rootTranslateZ', 'rtz', api.MFnUnitAttribute.kDistance, 0.)
        cls.aRootTranslate = fnNumeric.create('rootTranslate', 'rt', cls.aRootTranslateX, cls.aRootTranslateY, cls.aRootTranslateZ)

        cls.aMiddleParentMatrix = fnMatrix.create('middleParentMatrix', 'mpm')
        cls.aMiddleTranslateX = fnUnit.create('middleTranslateX', 'mtx', api.MFnUnitAttribute.kDistance, 0.)
        cls.aMiddleTranslateY = fnUnit.create('middleTranslateY', 'mty', api.MFnUnitAttribute.kDistance, 0.)
        cls.aMiddleTranslateZ = fnUnit.create('middleTranslateZ', 'mtz', api.MFnUnitAttribute.kDistance, 0.)
        cls.aMiddleTranslate = fnNumeric.create('middleTranslate', 'mt', cls.aMiddleTranslateX, cls.aMiddleTranslateY, cls.aMiddleTranslateZ)

        cls.aEffectorParentMatrix = fnMatrix.create('effectorParentMatrix', 'epm')
        cls.aEffectorTranslateX = fnUnit.create('effectorTranslateX', 'etx', api.MFnUnitAttribute.kDistance, 0.)
        cls.aEffectorTranslateY = fnUnit.create('effectorTranslateY', 'ety', api.MFnUnitAttribute.kDistance, 0.)
        cls.aEffectorTranslateZ = fnUnit.create('effectorTranslateZ', 'etz', api.MFnUnitAttribute.kDistance, 0.)
        cls.aEffectorTranslate = fnNumeric.create('effectorTranslate', 'et', cls.aEffectorTranslateX, cls.aEffectorTranslateY, cls.aEffectorTranslateZ)

        cls.aUpDirX = fnUnit.create('upDirectionX', 'udx', api.MFnUnitAttribute.kDistance, 0.)
        cls.aUpDirY = fnUnit.create('upDirectionY', 'udy', api.MFnUnitAttribute.kDistance, 0.)
        cls.aUpDirZ = fnUnit.create('upDirectionZ', 'udz', api.MFnUnitAttribute.kDistance, 1.)
        cls.aUpDir = fnNumeric.create('upDirection', 'ud', cls.aUpDirX, cls.aUpDirY, cls.aUpDirZ)

        cls.aOutTranslateX = fnUnit.create('outTranslateX', 'tx', api.MFnUnitAttribute.kDistance, 0.)
        cls.aOutTranslateY = fnUnit.create('outTranslateY', 'ty', api.MFnUnitAttribute.kDistance, 0.)
        cls.aOutTranslateZ = fnUnit.create('outTranslateZ', 'tz', api.MFnUnitAttribute.kDistance, 0.)
        cls.aOutTranslate = fnNumeric.create('outTranslate', 't', cls.aOutTranslateX, cls.aOutTranslateY, cls.aOutTranslateZ)

        '''
        cls.aUpAxis = fnEnum.create('upAxis', 'upa')
        fnEnum.addField('+x', 0)
        fnEnum.addField('+y', 1)
        fnEnum.addField('+z', 2)
        fnEnum.addField('-x', 3)
        fnEnum.addField('-y', 4)
        fnEnum.addField('-z', 5)
        '''

        cls.aFlip = fnNumeric.create('flip', 'flip', api.MFnNumericData.kBoolean)

        # ---------------------------- output ------------------------------- #
        cls.aRootParentMatrix = fnMatrix.create('rootParentMatrix', 'rpm')
        cls.aUpVectorX = fnUnit.create('upVectorX', 'ux', api.MFnUnitAttribute.kDistance, 0.)
        cls.aUpVectorY = fnUnit.create('upVectorY', 'uy', api.MFnUnitAttribute.kDistance, 0.)
        cls.aUpVectorZ = fnUnit.create('upVectorZ', 'uz', api.MFnUnitAttribute.kDistance, 1.)
        cls.aUpVector = fnNumeric.create('upVector', 'upv', cls.aUpVectorX, cls.aUpVectorY, cls.aUpVectorZ)

        # ------------------------------------------------------------------- #
        cls.addAttribute(cls.aRootParentMatrix)
        cls.addAttribute(cls.aMiddleParentMatrix)
        cls.addAttribute(cls.aEffectorParentMatrix)
        cls.addAttribute(cls.aRootTranslate)
        cls.addAttribute(cls.aMiddleTranslate)
        cls.addAttribute(cls.aEffectorTranslate)
        cls.addAttribute(cls.aUpVector)
        cls.addAttribute(cls.aFlip)
        cls.addAttribute(cls.aUpDir)
        # cls.addAttribute(cls.aUpAxis)
        cls.addAttribute(cls.aOutTranslate)

        # ------------------------------------------------------------------- #
        cls.attributeAffects(cls.aRootParentMatrix, cls.aOutTranslate)
        cls.attributeAffects(cls.aMiddleParentMatrix, cls.aOutTranslate)
        cls.attributeAffects(cls.aEffectorParentMatrix, cls.aOutTranslate)
        cls.attributeAffects(cls.aRootTranslate, cls.aOutTranslate)
        cls.attributeAffects(cls.aMiddleTranslate, cls.aOutTranslate)
        cls.attributeAffects(cls.aEffectorTranslate, cls.aOutTranslate)
        cls.attributeAffects(cls.aUpVector, cls.aOutTranslate)
        cls.attributeAffects(cls.aUpDir, cls.aOutTranslate)
        # cls.attributeAffects(cls.aUpAxis, cls.aOutTranslate)
        cls.attributeAffects(cls.aFlip, cls.aOutTranslate)

    def compute(self, plug, block):
        '''
        if plug.isChild:
            plug = plug.parent()
        if plug != self.aOutTranslate:
            return
        '''

        rootPos = MVector(block.inputValue(self.aRootTranslate).asDouble3())
        midPos = MVector(block.inputValue(self.aMiddleTranslate).asDouble3())
        effPos = MVector(block.inputValue(self.aEffectorTranslate).asDouble3())
        flip = block.inputValue(self.aFlip).asBool()
        # upAxis = block.inputValue(self.aUpAxis).asInt()

        # calculate wold positions
        rootParentMatrix = MTransformationMatrix(block.inputValue(self.aRootParentMatrix).asMatrix())
        middleParentMatrix = MTransformationMatrix(block.inputValue(self.aMiddleParentMatrix).asMatrix())
        effectorParentMatrix = MTransformationMatrix(block.inputValue(self.aEffectorParentMatrix).asMatrix())

        rootWorldPos = rootParentMatrix.translateBy(rootPos, api.MSpace.kObject).translation(api.MSpace.kWorld)
        midWorldPos = middleParentMatrix.translateBy(midPos, api.MSpace.kObject).translation(api.MSpace.kWorld)
        effWorldPos = effectorParentMatrix.translateBy(effPos, api.MSpace.kObject).translation(api.MSpace.kWorld)

        # calculate UpVector
        res1 = self.computeWithCos(block, rootWorldPos, midWorldPos, effWorldPos, flip)
        # block.outputValue(self.aOutTranslate).set3Double(*res1)

        res2 = self.computeWithSin(block, rootWorldPos, midWorldPos, effWorldPos, flip)
        block.outputValue(self.aOutTranslate).set3Double(*((res1 + res2) / 2))

    def computeWithSin(self, block, rootWorldPos, midWorldPos, effWorldPos, flip):

        modifier1 = 10
        rootToMid = midWorldPos - rootWorldPos
        rootToEff = effWorldPos - rootWorldPos
        upperLength = rootToMid.length()
        halfRootToEff = (rootToEff / 2)

        angle = rootToMid.angle(rootToEff)
        upVFactor = _clamp(
            upperLength * sin(angle) * modifier1,
            upperLength * 1.3,
            upperLength * 2.2)

        if flip:
            upVFactor *= -1

        tmp = rootToMid - halfRootToEff
        if tmp.length() < FULL_EXTENSION_TOLERANCE:
            norm = self.getUpNormal(block)
        else:
            norm = tmp.normal()

        upV = norm * upVFactor
        res = rootWorldPos + halfRootToEff + upV
        return res

    def computeWithCos(self, block, rootWorldPos, midWorldPos, effWorldPos, flip):
        #
        rootToMid = midWorldPos - rootWorldPos
        rootToEff = effWorldPos - rootWorldPos
        upperLength = rootToMid.length()

        angle = rootToMid.angle(rootToEff)
        midLen = upperLength * cos(angle)

        # upper
        normRootToEff = rootToEff.normal()
        upper = normRootToEff * midLen + rootWorldPos

        if flip:
            tmpUpper = upper - midWorldPos
        else:
            tmpUpper = midWorldPos - upper

        # switch if root to effector is in full extension
        if tmpUpper.length() < FULL_EXTENSION_TOLERANCE:
            norm = self.getUpNormal(block)

        else:
            norm = tmpUpper.normal()

        res = norm * rootToMid.length() + upper
        return res

    def getUpDir(self, block):
        return MVector(block.inputValue(self.aUpDir).asDouble3()).normal()

    def getUpNormal(self, block):
        n = self.getUpDir(block)
        middleParentMatrix = MTransformationMatrix(block.inputValue(self.aMiddleParentMatrix).asMatrix())
        return middleParentMatrix.translateBy(n, api.MSpace.kObject).translation(api.MSpace.kWorld)


# ------------------------------------------------------------------------------
def _registerNode(plugin, cls):
    try:
        plugin.registerNode(cls.type_name, cls.type_id, lambda: cls(), cls.initialize)

    except RuntimeError:
        sys.stderr.write('Failed to register node: ' + cls.type_name)
        raise


def _deregisterNode(plugin, cls):
    try:
        plugin.deregisterNode(cls.type_id)

    except RuntimeError:
        sys.stderr.write('Failed to deregister node: ' + cls.type_name)
        raise


def initializePlugin(mobj):
    plugin = api.MFnPlugin(mobj, __author__, __version__, 'Any')
    _registerNode(plugin, CalculateUpVector)


def uninitializePlugin(mobj):
    plugin = api.MFnPlugin(mobj)
    _deregisterNode(plugin, CalculateUpVector)


''' test
spaceLocator -n joint1;
spaceLocator -n joint2;
spaceLocator -n joint3;
spaceLocator -n upv;

nodeEdCreateNodeCommand "calculateUpVector";

connectAttr joint1.parentMatrix calculateUpVector1.rpm;
connectAttr joint1.translate calculateUpVector1.rt;

connectAttr joint2.parentMatrix calculateUpVector1.mpm;
connectAttr joint2.translate calculateUpVector1.mt;

connectAttr joint3.parentMatrix calculateUpVector1.epm;
connectAttr joint3.translate calculateUpVector1.et;

connectAttr -f calculateUpVector1.outTranslate upv.translate;
'''