Search code examples
pythonpython-3.xmathvectorinverse-kinematics

Rotate and translate a vector towards a target


I'm trying to learn and understand Inverse Kinematics by making it in PyGame. I've made a Bone class that has members a which is a PyGame Vector2 an angle and length. Finally it has a property b which is the calculated based on the previous 3 members.

In code it's like this:

import math

from pygame import Vector2

class Bone():
    def __init__(self, x: float, y: float, length: float) -> None:
        self.a      = Vector2(x, y)
        self.angle  = 0
        self.length = length

    @property
    def b(self):
        return Vector2(self.a.x + self.length * math.cos(self.angle), 
                       self.a.y + self.length * math.sin(self.angle))

    def rotate_and_translate(self, target: Vector2):
        dir = target - self.a
        self.angle = math.atan2(dir.y, dir.x)

        #???
        self.a.update(target.x + dir.x - self.length * np.cos(self.angle),target.y + dir.y - self.length * np.sin(self.angle))
        #???

        print(f'dir: {dir}, dir_mag: {dir.magnitude()}, target + dir: {target + dir}, a: {self.a}')

    def rotate(self, target: Vector2):
        """
        a(x1,y1)     b(x2, y2)
          x-----------x
           \ alpha
            \  
             \ target - a
              \
               \
                x(target)                
        """
        dir = target - self.a
        self.angle = math.atan2(dir.y, dir.x)

The main program looks like so:

import sys, pygame

from bone import Bone
from pygame.math import Vector2

pygame.init()

shape  = width, height = 1500, 850
screen = pygame.display.set_mode(shape)

def draw(bones: list[Bone], target: Vector2):
    black = 0, 0, 0
    white = 255, 255, 255

    screen.fill(white)

    for bone in bones:
        pygame.draw.aaline(screen, black, bone.a, bone.b)

    pygame.draw.circle(screen, black, (int(target[0]), int(target[1])), 2)
    pygame.display.flip()

def IK(bones: list[Bone], target: Vector2):
    i = len(bones) - 2

    bones[-1].rotate_and_translate(target)

    while i >= 0:
        if i != 0:
            bones[i].rotate_and_translate(bones[i + 1].a)
        else:
            bones[i].rotate(bones[i + 1].a)
        i -= 1

    return bones

def main():
    bones = []
    root  = Bone(width / 2, height / 2, 100)

    bones.append(root)

    #for i in range(1, 1):
    #    bones.append(Bone(bones[i - 1].b.x, bones[i - 1].b.y, 100))

    while 1:
        target = Vector2(pygame.mouse.get_pos())

        bones = IK(bones, target)

        for event in pygame.event.get():
            if event.type == pygame.QUIT: sys.exit()
    
        draw(bones, target)

 if __name__ == "__main__":
     main()

To summarize, I instantiate the bones (currently just 1) and am trying to get it to both rotate and translate so that its b property is on the target. As I understand it, every bone except the root can rotate_and_translate(), while the root can only rotate().

Currently only the rotation is working as intended and I'm a bit stumped on how to rotate and translate properly.

The closest I got to the answer was the bone following the point, but the angle was fixed and it would never rotate.

This is all the code there is for now. As always any and all advice is appreciated.

EDIT: I added the self.a.update() line into rotate_and_translate() which yields behaviour somewhat close to what I want but it's x and y values constantly flip between a set of 2 values at every point...


Solution

  • OK, so in my infinite silliness I was using the correct formula the wrong way.

    Basically, the way I want the bone to behave, b would equate to the target. and b is calculated like so:

       Vector2(self.a.x + self.length * np.cos(self.angle), 
               self.a.y + self.length * np.sin(self.angle))
    

    Which reliably yields a point that is a fixed distance away from a at all times.

    So I just assumed that it'd make sense to use the same formula to calculate the new a aswell and I was right. The problem was that I foolishly thought id have to factor in the coordinates of dir when calculating a new a, which was a mistake since dir is the difference between target and a which then yields the angle through arctan2. All I needed to do was remove the + dir.x and + dir.y from the formula and it worked like a charm:

       self.a.update(target.x - self.length * np.cos(self.angle), 
                     target.y - self.length * np.sin(self.angle))