Motion/InverseKinematics.py

##############################
#
# based on http://theorangeduck.com/page/deep-learning-framework-character-motion-synthesis-and-editing
#
##############################

import numpy as np
import scipy.linalg as linalg

try:
    from . import Animation
    from . import AnimationStructure
    from .Quaternions import Quaternions
except:
    import Animation
    import AnimationStructure
    from Quaternions import Quaternions


class BasicInverseKinematics:
    """
    Basic Inverse Kinematics Solver
    
    This is an extremely simple full body IK
    solver.
    
    It works given the following conditions:
    
        * All joint targets must be specified
        * All joint targets must be in reach
        * All joint targets must not differ 
          extremely from the starting pose
        * No bone length constraints can be violated
        * The root translation and rotation are 
          set to good initial values
    
    It works under the observation that if the
    _directions_ the joints are pointing toward
    match the _directions_ of the vectors between
    the target joints then the pose should match
    that of the target pose.
    
    Therefore it iterates over joints rotating
    each joint such that the vectors between it
    and it's children match that of the target
    positions.
    
    Parameters
    ----------
    
    animation : Animation
        animation input
        
    positions : (F, J, 3) ndarray
        target positions for each frame F
        and each joint J
    
    iterations : int
        Optional number of iterations.
        If the above conditions are met
        1 iteration should be enough,
        therefore the default is 1
        
    silent : bool
        Optional if to suppress output
        defaults to False
    """
    
    def __init__(self, animation, positions, iterations=1, silent=True):
    
        self.animation = animation
        self.positions = positions
        self.iterations = iterations
        self.silent = silent
        
    def __call__(self):
        
        children = AnimationStructure.children_list(self.animation.parents)
        
        for i in range(self.iterations):
        
            for j in AnimationStructure.joints(self.animation.parents):
                
                c = np.array(children[j])
                if len(c) == 0: continue

                anim_transforms = Animation.transforms_global(self.animation)
                anim_positions = anim_transforms[:,:,:3,3] # look at the translation vector inside a rotation matrix
                anim_rotations = Quaternions.from_transforms(anim_transforms)
                # self.animation.rotations[1,0]*(self.animation.offsets[1] + self.animation.rotations[1,1]*self.animation.offsets[2]) ==
                #         anim_positions[1, 2]

                jdirs = anim_positions[:,c] - anim_positions[:,np.newaxis,j]  #  limb vectors given by animation
                ddirs = self.positions[:,c] - self.positions[:,np.newaxis,j]  #  limb vectors given by input positions (target)

                if len(c) > 1 and (jdirs==0).all() and (ddirs==0).all(): # there was an expansion of high degree vertices (expand_topology())
                    # self.animation.rotations[:, j] remains untouched
                    continue


                jsums = np.sqrt(np.sum(jdirs**2.0, axis=-1)) + 1e-20
                dsums = np.sqrt(np.sum(ddirs**2.0, axis=-1)) + 1e-20
                
                jdirs = jdirs / jsums[:,:,np.newaxis]
                ddirs = ddirs / dsums[:,:,np.newaxis]
                
                angles = np.arccos(np.sum(jdirs * ddirs, axis=2).clip(-1, 1))

                axises = np.cross(jdirs, ddirs)
                # if jdirs.shape[1] == 1: # for a single child reconstruction should be exact
                #     assert np.allclose(Quaternions.from_angle_axis(angles, np.cross(jdirs, ddirs)) * jdirs, ddirs)
                axises = -anim_rotations[:,j,np.newaxis] * axises

                # find out which of the given bones are not of zero length
                # zero length bones produces a meaningless rotation
                jdirs_positive = (jsums > 1e-4)
                ddirs_positive = (dsums > 1e-4)
                dirs_positive = jdirs_positive[0]

                rotations = Quaternions.from_angle_axis(angles, axises)

                if rotations.shape[1] == 1:
                    averages = rotations[:,0]
                else:
                    averages = Quaternions.exp(rotations[:,dirs_positive].log().mean(axis=-2))

                self.animation.rotations[:,j] = self.animation.rotations[:,j] * averages
            
            if not self.silent:
                anim_positions = Animation.positions_global(self.animation)
                error = np.mean(np.sum((anim_positions - self.positions)**2.0, axis=-1)**0.5) # axis=-1)
                print('[BasicInverseKinematics] Iteration %i Error: %f' % (i+1, error))
        
        return self.animation


class JacobianInverseKinematics:
    """
    Jacobian Based Full Body IK Solver
    
    This is a full body IK solver which
    uses the dampened least squares inverse
    jacobian method.
    
    It should remain fairly stable and effective
    even for joint positions which are out of
    reach and it can also take any number of targets
    to treat as end effectors.
    
    Parameters
    ----------
    
    animation : Animation
        animation to solve inverse problem on

    targets : {int : (F, 3) ndarray}
        Dictionary of target positions for each
        frame F, mapping joint index to 
        a target position
    
    references : (F, 3)
        Optional list of J joint position
        references for which the result
        should bias toward
        
    iterations : int
        Optional number of iterations to
        compute. More iterations results in
        better accuracy but takes longer to 
        compute. Default is 10.
        
    recalculate : bool
        Optional if to recalcuate jacobian
        each iteration. Gives better accuracy
        but slower to compute. Defaults to True
        
    damping : float
        Optional damping constant. Higher
        damping increases stability but
        requires more iterations to converge.
        Defaults to 5.0
        
    secondary : float
        Force, or bias toward secondary target.
        Defaults to 0.25
        
    silent : bool
        Optional if to suppress output
        defaults to False
    """
    
    def __init__(self, animation, targets,
        references=None, iterations=10,
        recalculate=True, damping=2.0,
        secondary=0.25, translate=False,
        silent=False, weights=None,
        weights_translate=None):
        
        self.animation = animation
        self.targets = targets
        self.references = references
        
        self.iterations  = iterations
        self.recalculate = recalculate
        self.damping   = damping
        self.secondary   = secondary
        self.translate   = translate
        self.silent      = silent
        self.weights     = weights
        self.weights_translate = weights_translate
        
    def cross(self, a, b):
        o = np.empty(b.shape)
        o[...,0] = a[...,1]*b[...,2] - a[...,2]*b[...,1]
        o[...,1] = a[...,2]*b[...,0] - a[...,0]*b[...,2]
        o[...,2] = a[...,0]*b[...,1] - a[...,1]*b[...,0]
        return o
        
    def jacobian(self, x, fp, fr, ts, dsc, tdsc):
        
        """ Find parent rotations """
        prs = fr[:,self.animation.parents]
        prs[:,0] = Quaternions.id((1))
        
        """ Find global positions of target joints """
        tps = fp[:,np.array(list(ts.keys()))]
        
        """ Get partial rotations """
        qys = Quaternions.from_angle_axis(x[:,1:prs.shape[1]*3:3], np.array([[[0,1,0]]]))
        qzs = Quaternions.from_angle_axis(x[:,2:prs.shape[1]*3:3], np.array([[[0,0,1]]]))
        
        """ Find axis of rotations """
        es = np.empty((len(x),fr.shape[1]*3, 3))
        es[:,0::3] = ((prs * qzs) * qys) * np.array([[[1,0,0]]])
        es[:,1::3] = ((prs * qzs) * np.array([[[0,1,0]]]))
        es[:,2::3] = ((prs * np.array([[[0,0,1]]])))
        
        """ Construct Jacobian """
        j = fp.repeat(3, axis=1)
        j = dsc[np.newaxis,:,:,np.newaxis] * (tps[:,np.newaxis,:] - j[:,:,np.newaxis])
        j = self.cross(es[:,:,np.newaxis,:], j)
        j = np.swapaxes(j.reshape((len(x), fr.shape[1]*3, len(ts)*3)), 1, 2)
                
        if self.translate:
            
            es = np.empty((len(x),fr.shape[1]*3, 3))
            es[:,0::3] = prs * np.array([[[1,0,0]]])
            es[:,1::3] = prs * np.array([[[0,1,0]]])
            es[:,2::3] = prs * np.array([[[0,0,1]]])
            
            jt = tdsc[np.newaxis,:,:,np.newaxis] * es[:,:,np.newaxis,:].repeat(tps.shape[1], axis=2)
            jt = np.swapaxes(jt.reshape((len(x), fr.shape[1]*3, len(ts)*3)), 1, 2)
            
            j = np.concatenate([j, jt], axis=-1)
        
        return j
        
    #@profile(immediate=True)
    def __call__(self, descendants=None, gamma=1.0):
        
        self.descendants = descendants
        
        """ Calculate Masses """
        if self.weights is None:
            self.weights = np.ones(self.animation.shape[1])
            
        if self.weights_translate is None:
            self.weights_translate = np.ones(self.animation.shape[1])
        
        """ Calculate Descendants """
        if self.descendants is None:
            self.descendants = AnimationStructure.descendants_mask(self.animation.parents)
        
        self.tdescendants = np.eye(self.animation.shape[1]) + self.descendants
        
        self.first_descendants = self.descendants[:,np.array(list(self.targets.keys()))].repeat(3, axis=0).astype(int)
        self.first_tdescendants = self.tdescendants[:,np.array(list(self.targets.keys()))].repeat(3, axis=0).astype(int)
        
        """ Calculate End Effectors """
        self.endeff = np.array(list(self.targets.values()))
        self.endeff = np.swapaxes(self.endeff, 0, 1) 
        
        if not self.references is None:
            self.second_descendants = self.descendants.repeat(3, axis=0).astype(int)
            self.second_tdescendants = self.tdescendants.repeat(3, axis=0).astype(int)
            self.second_targets = dict([(i, self.references[:,i]) for i in xrange(self.references.shape[1])])
        
        nf = len(self.animation)
        nj = self.animation.shape[1]
        
        if not self.silent:
            gp = Animation.positions_global(self.animation)
            gp = gp[:,np.array(list(self.targets.keys()))]            
            error = np.mean(np.sqrt(np.sum((self.endeff - gp)**2.0, axis=2)))
            print('[JacobianInverseKinematics] Start | Error: %f' % error)
        
        for i in range(self.iterations):

            """ Get Global Rotations & Positions """
            gt = Animation.transforms_global(self.animation)
            gp = gt[:,:,:,3]
            gp = gp[:,:,:3] / gp[:,:,3,np.newaxis]
            gr = Quaternions.from_transforms(gt)
            
            x = self.animation.rotations.euler().reshape(nf, -1)
            w = self.weights.repeat(3)
            
            if self.translate:
                x = np.hstack([x, self.animation.positions.reshape(nf, -1)])
                w = np.hstack([w, self.weights_translate.repeat(3)])
            
            """ Generate Jacobian """
            if self.recalculate or i == 0:
                j = self.jacobian(x, gp, gr, self.targets, self.first_descendants, self.first_tdescendants)
            
            """ Update Variables """            
            l = self.damping * (1.0 / (w + 0.001))
            d = (l*l) * np.eye(x.shape[1])
            e = gamma * (self.endeff.reshape(nf,-1) - gp[:,np.array(list(self.targets.keys()))].reshape(nf, -1))
            
            x += np.array(list(map(lambda jf, ef:
                linalg.lu_solve(linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef)), j, e)))
            
            """ Generate Secondary Jacobian """
            if self.references is not None:
                
                ns = np.array(list(map(lambda jf:
                    np.eye(x.shape[1]) - linalg.solve(jf.T.dot(jf) + d, jf.T.dot(jf)), j)))
                    
                if self.recalculate or i == 0:
                    j2 = self.jacobian(x, gp, gr, self.second_targets, self.second_descendants, self.second_tdescendants)
                        
                e2 = self.secondary * (self.references.reshape(nf, -1) - gp.reshape(nf, -1))
                
                x += np.array(list(map(lambda nsf, j2f, e2f:
                    nsf.dot(linalg.lu_solve(linalg.lu_factor(j2f.T.dot(j2f) + d), j2f.T.dot(e2f))), ns, j2, e2)))

            """ Set Back Rotations / Translations """
            self.animation.rotations = Quaternions.from_euler(
                x[:,:nj*3].reshape((nf, nj, 3)), order='xyz', world=True)
                
            if self.translate:
                self.animation.positions = x[:,nj*3:].reshape((nf,nj, 3))
                
            """ Generate Error """
            
            if not self.silent:
                gp = Animation.positions_global(self.animation)
                gp = gp[:,np.array(list(self.targets.keys()))]
                error = np.mean(np.sum((self.endeff - gp)**2.0, axis=2)**0.5)
                print('[JacobianInverseKinematics] Iteration %i | Error: %f' % (i+1, error))
            

class BasicJacobianIK:
    """
    Same interface as BasicInverseKinematics
    but uses the Jacobian IK Solver Instead
    """
    
    def __init__(self, animation, positions, iterations=10, silent=True, **kw):
       
        targets = dict([(i, positions[:,i]) for i in range(positions.shape[1])])
        self.ik = JacobianInverseKinematics(animation, targets, iterations=iterations, silent=silent, **kw)
       
    def __call__(self, **kw):
        return self.ik(**kw)
        

class ICP:
    
    
    def __init__(self, 
        anim, rest, weights, mesh, goal,
        find_closest=True, damping=10,
        iterations=10, silent=True,
        translate=True, recalculate=True,
        weights_translate=None):
        
        self.animation = anim
        self.rest = rest
        self.vweights = weights
        self.mesh = mesh
        self.goal = goal
        self.find_closest = find_closest
        self.iterations = iterations
        self.silent = silent
        self.translate = translate
        self.damping   = damping
        self.weights = None
        self.weights_translate = weights_translate
        self.recalculate = recalculate
        
    def cross(self, a, b):
        o = np.empty(b.shape)
        o[...,0] = a[...,1]*b[...,2] - a[...,2]*b[...,1]
        o[...,1] = a[...,2]*b[...,0] - a[...,0]*b[...,2]
        o[...,2] = a[...,0]*b[...,1] - a[...,1]*b[...,0]
        return o
        
    def jacobian(self, x, fp, fr, goal, weights, des_r, des_t):
        
        """ Find parent rotations """
        prs = fr[:,self.animation.parents]
        prs[:,0] = Quaternions.id((1))
        
        """ Get partial rotations """
        qys = Quaternions.from_angle_axis(x[:,1:prs.shape[1]*3:3], np.array([[[0,1,0]]]))
        qzs = Quaternions.from_angle_axis(x[:,2:prs.shape[1]*3:3], np.array([[[0,0,1]]]))
        
        """ Find axis of rotations """
        es = np.empty((len(x),fr.shape[1]*3, 3))
        es[:,0::3] = ((prs * qzs) * qys) * np.array([[[1,0,0]]])
        es[:,1::3] = ((prs * qzs) * np.array([[[0,1,0]]]))
        es[:,2::3] = ((prs * np.array([[[0,0,1]]])))
        
        """ Construct Jacobian """
        j = fp.repeat(3, axis=1)
        j = des_r[np.newaxis,:,:,:,np.newaxis] * (goal[:,np.newaxis,:,np.newaxis] - j[:,:,np.newaxis,np.newaxis])
        j = np.sum(j * weights[np.newaxis,np.newaxis,:,:,np.newaxis], 3)
        j = self.cross(es[:,:,np.newaxis,:], j)
        j = np.swapaxes(j.reshape((len(x), fr.shape[1]*3, goal.shape[1]*3)), 1, 2)
        
        if self.translate:
            
            es = np.empty((len(x),fr.shape[1]*3, 3))
            es[:,0::3] = prs * np.array([[[1,0,0]]])
            es[:,1::3] = prs * np.array([[[0,1,0]]])
            es[:,2::3] = prs * np.array([[[0,0,1]]])
            
            jt = des_t[np.newaxis,:,:,:,np.newaxis] * es[:,:,np.newaxis,np.newaxis,:].repeat(goal.shape[1], axis=2)
            jt = np.sum(jt * weights[np.newaxis,np.newaxis,:,:,np.newaxis], 3)
            jt = np.swapaxes(jt.reshape((len(x), fr.shape[1]*3, goal.shape[1]*3)), 1, 2)
            
            j = np.concatenate([j, jt], axis=-1)
        
        return j
        
    #@profile(immediate=True)
    def __call__(self, descendants=None, maxjoints=4, gamma=1.0, transpose=False):
        
        """ Calculate Masses """
        if self.weights is None:
            self.weights = np.ones(self.animation.shape[1])
            
        if self.weights_translate is None:
            self.weights_translate = np.ones(self.animation.shape[1]) 
        
        nf = len(self.animation)
        nj = self.animation.shape[1]
        nv = self.goal.shape[1]
        
        weightids = np.argsort(-self.vweights, axis=1)[:,:maxjoints]
        weightvls = np.array(list(map(lambda w, i: w[i], self.vweights, weightids)))
        weightvls = weightvls / weightvls.sum(axis=1)[...,np.newaxis]
        
        if descendants is None:
            self.descendants = AnimationStructure.descendants_mask(self.animation.parents)
        else:
            self.descendants = descendants
        
        des_r = np.eye(nj) + self.descendants
        des_r = des_r[:,weightids].repeat(3, axis=0)

        des_t = np.eye(nj) + self.descendants
        des_t = des_t[:,weightids].repeat(3, axis=0)
        
        if not self.silent:
            curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints)
            error = np.mean(np.sqrt(np.sum((curr - self.goal)**2.0, axis=-1)))
            print('[ICP] Start | Error: %f' % error)
        
        for i in range(self.iterations):
            
            """ Get Global Rotations & Positions """
            gt = Animation.transforms_global(self.animation)
            gp = gt[:,:,:,3]
            gp = gp[:,:,:3] / gp[:,:,3,np.newaxis]
            gr = Quaternions.from_transforms(gt)
            
            x = self.animation.rotations.euler().reshape(nf, -1)
            w = self.weights.repeat(3)
            
            if self.translate:
                x = np.hstack([x, self.animation.positions.reshape(nf, -1)])
                w = np.hstack([w, self.weights_translate.repeat(3)])
            
            """ Get Current State """
            curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints)

            """ Find Cloest Points """
            if self.find_closest:
                mapping = np.argmin(
                    (curr[:,:,np.newaxis] - 
                     self.goal[:,np.newaxis,:])**2.0, axis=2)
                e = gamma * (np.array(list(map(lambda g, m: g[m], self.goal, mapping))) - curr).reshape(nf, -1)
            else:
                e = gamma * (self.goal - curr).reshape(nf, -1)
            
            """ Generate Jacobian """
            if self.recalculate or i == 0:
                j = self.jacobian(x, gp, gr, self.goal, weightvls, des_r, des_t)
            
            """ Update Variables """            
            l = self.damping * (1.0 / (w + 1e-10))
            d = (l*l) * np.eye(x.shape[1])
            
            if transpose:
                x += np.array(list(map(lambda jf, ef: jf.T.dot(ef), j, e)))
            else:
                x += np.array(list(map(lambda jf, ef:
                    linalg.lu_solve(linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef)), j, e)))
            
            """ Set Back Rotations / Translations """
            self.animation.rotations = Quaternions.from_euler(
                x[:,:nj*3].reshape((nf, nj, 3)), order='xyz', world=True)
                
            if self.translate:
                self.animation.positions = x[:,nj*3:].reshape((nf, nj, 3))
        
            if not self.silent:
                curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh)
                error = np.mean(np.sqrt(np.sum((curr - self.goal)**2.0, axis=-1)))
                print('[ICP] Iteration %i | Error: %f' % (i+1, error))
                

def animation_from_positions(positions, parents, offsets=None, ik_iterations=1):

    root_idx = np.where(parents == -1)[0][0]

    if offsets is None:
        orig_offsets = Animation.offsets_from_positions(positions, parents)

        # compute offsets over all joints except for root
        idx_no_root = np.delete(np.arange(positions.shape[1]), root_idx)
        orig_offsets_no_root = orig_offsets[:, idx_no_root]

        # prevent a zero offset (can happen in first iterations of generated motion). such offset is ill posed and
        # results in ambigious angles when computing inverse kinematics
        orig_offsets_no_root[orig_offsets_no_root == 0] = 1e-6 * np.random.randn()

        # when synthesising motion, bone length is not always fixed among frames. We calculate mean bone length and
        # use it for the offset
        bone_lens = np.linalg.norm(orig_offsets_no_root, axis=2)[:, :, np.newaxis]
        normed_offsets_no_root = np.divide(orig_offsets_no_root, bone_lens, out=np.zeros_like(orig_offsets_no_root), where=bone_lens!=0)
        offsets_no_root = normed_offsets_no_root[0] * bone_lens.mean(axis=0)

        offsets = orig_offsets[0]
        offsets[idx_no_root] = offsets_no_root

    anim, sorted_order, parents = Animation.animation_from_offsets(offsets, parents, shape=positions.shape)
    positions = positions[:, sorted_order]
    sorted_root_idx = np.where(sorted_order == root_idx)[0][0]

    anim.positions[:, sorted_root_idx] = positions[:, 0] # keep root positions. important for IK

    # apply IK
    ik = BasicInverseKinematics(anim, positions, silent=True, iterations=ik_iterations)
    ik()
    return anim, sorted_order, parents