constraints.py

from dataclasses import dataclass
import numpy

import pygame
from pygame import Vector2


WINDOW_WIDTH = 500
WINDOW_HEIGHT = 400
MAX_FPS = 60


pygame.init()
window = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
pygame.display.set_caption("Constraint-Based Physics Playground")
clock = pygame.Clock()
is_running = True

#font = pygame.Font("FiraCode-Regular.ttf", 14)
font = pygame.Font(None, 20)


@dataclass
class MassPoint:
    position: Vector2
    velocity: Vector2
    mass: float

particle = MassPoint(Vector2(250, 250), Vector2(), 5.0)
target = MassPoint(Vector2(250, 250), Vector2(), 5.0)

baumgarte = 0.004
iterations = 1
dt = 1 / 60


while is_running:
    clock.tick(MAX_FPS)
    fps = clock.get_fps()

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            is_running = False

    mouse = Vector2(*pygame.mouse.get_pos())

    keys = pygame.key.get_pressed()
    if keys[pygame.K_UP]: baumgarte += 0.0001
    if keys[pygame.K_DOWN]: baumgarte -= 0.0001

    target.position = mouse.copy()


    # Position constraint
    # C(p) = p - a = 0
    c = particle.position - target.position
    c = numpy.array([[c.x, c.y]]).transpose()

    # Baumgarte stabilization
    # We feed some of position error into velocity constraint
    bias = (-baumgarte / dt) * c
    
    # Velocity matrix
    # (Transpose to make it a column matrix)
    v = numpy.array([[particle.velocity.x, particle.velocity.y]]).transpose()

    # Jacobian matrix
    j = numpy.array([[1, 0], [0, 1]])
    jt = j.transpose()

    # Mass matrix
    m = numpy.array([[particle.mass, 0], [0, particle.mass]])
    mi = numpy.linalg.inv(m)

    # Now we can write the velocity constraint terms in the general form Ax=b
    # https://dyn4j.org/2010/07/equality-constraints/
    #
    # A -> J M^-1 Jt  (This is also called effective mass?)
    # x -> λ          (Lambda is impulse magnitude, also called lagrange multiplier)
    # b -> -JV
    A = j @ mi @ jt
    b = -j @ v + bias

    # Initial guess for lambda
    x = numpy.zeros_like(b)

    # Gauss-Seidel iterative solver
    # Algorithm from https://stackoverflow.com/a/66764437
    for _ in range(iterations):
        x0 = x.copy()

        for i in range(A.shape[0]):
            x[i] = (b[i] - numpy.dot(A[i, :i], x[:i]) - numpy.dot(A[i, (i + 1):], x0[(i + 1):])) / A[i, i]

        # With the newly solved lambda apply impulse and  update the velocity
        impulse = mi @ jt @ x
        v += impulse
        particle.velocity.x = v[0, 0]
        particle.velocity.y = v[1, 0]

        error = numpy.linalg.norm(x0 - x)

        # Okay I'm not sure about this part, if I don't update b
        # everything goes chaos
        b = -j @ v + bias

    # Apply damping
    particle.velocity *= 0.98

    # Integrate velocities
    particle.position += particle.velocity


    window.fill((255, 255, 255))

    pygame.draw.circle(window, (255, 80, 0), particle.position, 12, 2)
    pygame.draw.circle(window, (70, 255, 150), target.position, 5)

    lines = (
        f"FPS: {round(fps, 2)}",
        "",
        "C(p) = p - a = 0",
        f"Iterations: {iterations}",
        f"Baumgarte:  {round(baumgarte, 5)}",
        f"Δt:         {round(dt, 3)}",
        f"Error:      {round(error, 6)}"
    )

    y_gap = 17
    for i, line in enumerate(lines):
        if not line: continue
        window.blit(font.render(line, True, (0, 0, 0)), (5, 5 + i * y_gap))

    pygame.display.flip()

pygame.quit()