Rotating cube

pygame-rotating-cube.py

import pygame
import numpy as np

# Import pygame.locals for easier access to key coordinates
from pygame.locals import (
    K_UP,
    K_DOWN,
    K_LEFT,
    K_RIGHT,
    K_ESCAPE,
    KEYDOWN,
    KEYUP,
    QUIT,
)

pygame.init()

# Set up the drawing window
SCR_WIDTH = 500
SCR_HEIGHT = 500
screen = pygame.display.set_mode([SCR_WIDTH, SCR_HEIGHT])

# Get pixels array
pxls = pygame.PixelArray(screen)

# Run until the user asks to quit
running = True

# Globals
ax = 0                          # X-axis rotarion angle
ay = 0                          # Y-axis rotation angle
D = 0.003                       # Angle increment constant

# Make array with cube edge coordinates
CUBE_SIZE = 100
cube = []
square = [(-1, -1), (1, -1), (1, 1), (-1, 1)]

for k in range(4):
    x0 = CUBE_SIZE/2 * square[k][0]
    y0 = CUBE_SIZE/2 * square[k][1]
    x1 = CUBE_SIZE/2 * square[(k+1)%4][0]
    y1 = CUBE_SIZE/2 * square[(k+1)%4][1]
    z0 = -CUBE_SIZE/2
    z1 = CUBE_SIZE/2
    # x0 = x0 + 55
    # x1 = x1 + 55
    cube.append({'x0':x0, 'y0':y0, 'z0':z0, 'x1':x1, 'y1':y1, 'z1':z0})
    cube.append({'x0':x0, 'y0':y0, 'z0':z1, 'x1':x1, 'y1':y1, 'z1':z1})
    cube.append({'x0':x0, 'y0':y0, 'z0':z0, 'x1':x0, 'y1':y0, 'z1':z1})

# Calculate projection coordinates
# x - on the axis perpendicular to the view
# z - on the axis of view
# f - focal distance
def projection(x, z, f):
    return x * f / z

# Calculate affine transform of 2d coordinate
def affine2d(coords, matrix):
    res = np.matmul([coords[0], coords[1], 1.0], matrix)
    return (res[0]/res[2], res[1]/res[2])

# Rotate in 2d
def rotate2d(c, a):
    sa = np.sin(a)
    ca = np.cos(a)
    rotator = np.array([[ca, -sa, 0], [sa, ca, 0], [0, 0, 1]])
    return affine2d(c, rotator)

# Rotate in 3d (trivial)
def rotate3d(c, ax, ay, az):
    x = c[0]
    y = c[1]
    z = c[2]
    (y, z) = rotate2d((y,z), ax)
    (x, z) = rotate2d((x,z), ay)
    (y, x) = rotate2d((y,x), az)
    return (x, y, z)

# Rotate of list of 3d lines
def rotate3d_poly(p, ax, ay, az):
    pout = []
    for l in p:
        (x0, y0, z0) = rotate3d((l['x0'], l['y0'], l['z0']), ax, ay, az)
        (x1, y1, z1) = rotate3d((l['x1'], l['y1'], l['z1']), ax, ay, az)
        pout.append({'x0':x0, 'y0':y0, 'z0':z0, 'x1':x1, 'y1':y1, 'z1':z1})

    return pout

# Draw list of lines on canvas
def draw3d_poly(canvas, p, f, x, y, z):
    for l in p:
        x0 = projection(l['x0'], l['z0'] + z, f) + x
        y0 = projection(l['y0'], l['z0'] + z, f) + y
        x1 = projection(l['x1'], l['z1'] + z, f) + x
        y1 = projection(l['y1'], l['z1'] + z, f) + y
        pygame.draw.line(canvas, (0,0,0), (x0, y0), (x1, y1), 2)
        # pygame.draw.aaline(canvas, (0,0,0), (x0, y0), (x1, y1))

# PyGame main loop
while running:
    for event in pygame.event.get():
        # Rotate if key is down
        if event.type == KEYDOWN:
            if event.key == K_UP:    ax = D
            if event.key == K_DOWN:  ax = -D
            if event.key == K_LEFT:  ay = D
            if event.key == K_RIGHT: ay = -D
            if event.key == K_ESCAPE:
                running = False

        # Stop rotation if key is up
        if event.type == KEYUP:
            ax = 0
            ay = 0

        # Did the user click the window close button?
        if event.type == QUIT:
            running = False

    # Rotate cube
    cube = rotate3d_poly(cube, ax, ay, 0)

    # Clear screen
    screen.fill((255, 255, 255))

    # pxls[x][y] = (0, 0, 0)
    # pygame.surfarray.blit_array(screen, pxls)

    # Draw rotated cube
    draw3d_poly(screen, cube, 250, SCR_WIDTH/2, SCR_HEIGHT/2, 200)

    # Flip the display
    pygame.display.flip()

# Quit if not running
pygame.quit()