nd_cube.py

import pygame as pg
import numpy as np
import math as m
from numba import njit
import time


# no longer used
def draw_points(obj: np.ndarray, win: pg.Surface, n: int):
    """Draws the points defining the cube.

    :param obj: the object to draw
    :param win: the surface to draw to
    :param n: the number of dimensions
    """
    offset = np.array((450, 325))
    for vector in obj:
        pos = vector*3**n*10 + offset
        pos = [pos[0], pos[1]]
        pg.draw.circle(win, (100, 200, 100), pos, radius=2)


@njit
def prepare_drawing(obj: np.ndarray, n: int, upscale=False) -> np.ndarray:
    """prepares for drawing of object

    :param obj: the object to draw
    :param n: the number of dimensions
    :param upscale: If true: scales up the cube
    :return: the object
    """
    offset = np.array((450, 325), dtype=np.float32)
    scale = 7 ** n * 2
    if not upscale:
        scale /= 3

    for index, vector in enumerate(obj):
        obj[index] = vector*scale + offset

    return obj


@njit(fastmath=True)
def find_reference_points(obj: np.ndarray, lines: np.ndarray) -> np.ndarray:
    ref_array = np.empty((len(lines), 2, len(obj[0])), dtype=np.uint16)

    for vector in lines:
        point1 = obj[vector[0]]
        point2 = obj[vector[1]]
        ref_array[vector[0], 0] = point1 + 0.3*(point2 - point1)
        ref_array[vector[1], 1] = point1 + 0.7*(point2 - point1)

    return ref_array


@njit
def draw_np_old(w: int, h: int, array: np.ndarray, lines: np.ndarray) -> np.ndarray:
    surface = np.zeros((w, h), dtype=np.uint8)
    color = 100

    for vector in lines:
        x0, y0 = int(array[vector[0], 0]), int(array[vector[0], 1])
        x1, y1 = int(array[vector[1], 0]), int(array[vector[1], 1])
        dx = abs(x1 - x0)
        if x0 < x1:
            sx = 1
        else:
            sx = -1
        dy = -abs(y1 - y0)
        if y0 < y1:
            sy = 1
        else:
            sy = -1
        error = dx + dy

        while True:
            surface[int(x0), int(y0)] = color
            if x0 == x1 and y0 == y1:
                break
            e2 = 2*error
            if e2 >= dy:
                if x0 == x1:
                    break
                error += dy
                x0 += sx
            if e2 <= dx:
                if y0 == y1:
                    break
                error += dx
                y0 += sy

    smooth_surface = np.zeros((w, h, 3), dtype=np.uint8)
    for x in range(1, w-1):
        for y in range(1, h-1):
            if surface[x, y] == color:
                smooth_surface[x, y, 1] = color

            else:
                value = 0
                for x_ in range(x-1, x+2):
                    for y_ in range(y-1, y+2):
                        value += surface[x_, y_]

                smooth_surface[x, y, 1] = value//8

    return smooth_surface


@njit
def draw_np(w: int, h: int, array: np.ndarray, lines: np.ndarray) -> np.ndarray:
    surface = np.zeros((w, h), dtype=np.uint8)
    color = np.uint8(100)
    x_max = np.uint8(0)
    x_min = np.uint8(w)
    y_max = np.uint8(0)
    y_min = np.uint8(h)

    for vector in lines:
        x0, y0 = int(array[vector[0], 0]), int(array[vector[0], 1])
        x1, y1 = int(array[vector[1], 0]), int(array[vector[1], 1])
        x_max = max(x_max, x0, x1)
        x_min = min(x_min, x0, x1)
        y_max = max(y_max, y0, y1)
        y_min = min(y_min, y0, y1)
        dx = abs(x1 - x0)
        if x0 < x1:
            sx = 1
        else:
            sx = -1
        dy = -abs(y1 - y0)
        if y0 < y1:
            sy = 1
        else:
            sy = -1
        error = dx + dy

        while True:
            surface[int(x0), int(y0)] = color
            if x0 == x1 and y0 == y1:
                break
            e2 = 2*error
            if e2 >= dy:
                if x0 == x1:
                    break
                error += dy
                x0 += sx
            if e2 <= dx:
                if y0 == y1:
                    break
                error += dx
                y0 += sy

    smooth_surface = np.zeros((w, h, 3), dtype=np.uint8)
    for x in range(x_min - 1, x_max + 1):
        for y in range(y_min - 1, y_max + 1):
            if surface[x, y] == color:
                smooth_surface[x, y, 1] = color

            else:
                value = 0
                for x_ in range(x-1, x+2):
                    for y_ in range(y-1, y+2):
                        value += surface[x_, y_]

                smooth_surface[x, y, 1] = value//8

    return smooth_surface


def draw_array(surface: np.ndarray, win: pg.Surface):
    pg.surfarray.blit_array(win, surface)


def draw(obj: np.ndarray, lines: np.ndarray, win: pg.Surface, ref_array: np.ndarray | None = None):
    """Draws the lines defining the cube

    :param obj: the cube to draw
    :param lines: the lines in the cube
    :param win: the surface to draw to
    :param ref_array: used to skip drawing lines on top of other lines
    """
    if len(lines) < 20000 or ref_array is None:
        for vector in lines:
            pg.draw.aaline(win, (100, 200, 100), obj[vector[0]], obj[vector[1]])

    else:
        for vector in lines:
            pos1 = win.get_at(ref_array[vector[0], 0])
            pos2 = win.get_at(ref_array[vector[1], 1])
            if pos1[0] == 0 and pos2[0] == 0:
                pg.draw.aaline(win, (100, 200, 100), obj[vector[0]], obj[vector[1]])


def rotate_obj_nd_old(obj: np.ndarray, angle: float, axis1: int, axis2: int) -> np.ndarray:
    """Rotates the obj the given angle in the plane defined by axis1 and axis2

    :param obj: the object to rotate
    :param angle: the angle to rotate by
    :param axis1: the first dimension to rotate about
    :param axis2: the second dimension to rotate about
    :return: the rotated object
    """
    n = len(obj[0])
    axis1 -= 1
    axis2 -= 1
    rot_matrix = np.identity(n, dtype=np.float32)

    rot_matrix[axis1, axis1] = m.cos(angle)
    rot_matrix[axis1, axis2] = -m.sin(angle)
    rot_matrix[axis2, axis2] = m.cos(angle)
    rot_matrix[axis2, axis1] = m.sin(angle)

    for index, vector in enumerate(obj):
        # obj[index] = np.matmul(rot_matrix, vector)
        obj[index] = rot_matrix @ vector

    return obj


@njit(fastmath=True)
def rotate_obj_nd(obj: np.ndarray, angle: float, axis1: int, axis2: int) -> np.ndarray:
    """Rotates the obj the given angle in the plane defined by axis1 and axis2

    :param obj: the object to rotate
    :param angle: the angle to rotate by
    :param axis1: the first dimension to rotate about
    :param axis2: the second dimension to rotate about
    :return: the rotated object
    """
    axis1 -= 1
    axis2 -= 1
    cos = np.float32(m.cos(angle))
    sin = np.float32(m.sin(angle))

    for vector in obj:
        vector[axis1], vector[axis2] = cos*vector[axis1] - sin*vector[axis2], sin*vector[axis1] + cos*vector[axis2]

    return obj


def cube_column(array: np.ndarray, n: int):
    """Generates a column for the array of points

    :param array: the array of points
    :param n: the number of dimensions
    """
    direction = [1, -1]
    count = 0
    while array[-1, -n] == 0:
        for c1 in direction:
            for _ in range(int(2**n/2)):
                array[count, -n] = c1
                count += 1


def cube_nd(n: int) -> np.ndarray:
    """Generates a n-dimensional cube

    :param n: number of dimensions
    :return: array of points
    """
    array = np.zeros((2**n, n), dtype=np.float32)
    for index in range(n):
        cube_column(array, index+1)

    return array


@njit
def generate_lines(array: np.ndarray, n: int) -> np.ndarray:
    """Generates the lines confining the cube

    :param array: the array of points
    :param n: the number of dimensions
    :return: array of lines
    """
    lines = np.empty((10 * len(array), 2), dtype=np.uint32)
    count_lines = 0
    num_points = 2**n
    for index1, vector1 in enumerate(array):
        for index2 in range(n):
            new_vector = vector1.copy()
            if new_vector[index2] == -1:
                new_vector[index2] = 1
            else:
                new_vector[index2] = -1

            start = 0
            end = num_points - 1
            for index4 in range(n):
                if new_vector[index4] == 1:
                    end = int((end+start)/2)
                else:
                    start = int((end+start)/2)

            if index1 < end:
                lines[count_lines] = np.array((index1, end))
                count_lines += 1

    return lines[0:count_lines]


@njit
def perspective_projection(array: np.ndarray) -> np.ndarray:
    """projects the array to a lower dimension with perspective

    :param array: the array of points to project
    :return: array
    """
    for index, vector in enumerate(array):
        array[index] = vector/(vector[-1]-7)

    return array[:len(array)+1, :-1]


def generate_cube(n: int) -> tuple[np.ndarray, np.ndarray]:
    """Generates points and lines that defines a n-dimensional cube

    :param n: number of dimensions
    :return: points, lines"""

    array = cube_nd(n)
    lines = generate_lines(array, n)

    return array, lines


def main():
    """main function"""
    print("Use the arrow keys to rotate the hypercube")
    while True:
        try:
            n = int(input("Input number of dimensions: "))
            break

        except ValueError:
            pass

    start = time.time()
    array, lines = generate_cube(n)
    end = time.time()
    print(f"Time to compute: {int((end - start)*1000)} ms\n")

    print(f"Number of vertices: {len(array)}\n"
          f"Number of edges: {len(lines)}, Total length of edges: {len(lines) * 2} u.l.\n")

    angle = 0.0
    manual_rot1 = 0.0
    manual_rot2 = 0.0

    pg.init()
    w, h = 900, 650
    win = pg.display.set_mode((w, h))

    clock = pg.time.Clock()
    run = True

    while run:
        clock.tick(60)
        if int(angle*100) % 60 == 1:
            # pass
            print(f"fps: {clock.get_fps()}")
        for event in pg.event.get():
            if event.type == pg.QUIT:
                run = False

        pressed = pg.key.get_pressed()
        if pressed[pg.K_UP]:
            manual_rot1 += 0.03

        if pressed[pg.K_DOWN]:
            manual_rot1 -= 0.03

        if pressed[pg.K_LEFT]:
            manual_rot2 -= 0.03

        if pressed[pg.K_RIGHT]:
            manual_rot2 += 0.03

        angle += 0.008

        # t0 = time.time()
        pro_array = array.copy()
        # t1 = time.time()
        for index in range(len(pro_array[0]) - 2):
            if len(pro_array[0]) == 3:

                rotate_obj_nd(pro_array, manual_rot1, 2, 3)

                rotate_obj_nd(pro_array, manual_rot2, 1, 2)

                rotate_obj_nd(pro_array, angle, 1, 3)

            elif len(pro_array[0]) == 4:
                rotate_obj_nd(pro_array, angle, 3, 4)
            elif len(pro_array[0]) == 6:
                rotate_obj_nd(pro_array, angle, 5, 6)
            elif len(pro_array[0]) == 8:
                rotate_obj_nd(pro_array, angle, 5, 8)
            elif len(pro_array[0]) == 10:
                rotate_obj_nd(pro_array, angle, 4, 8)

            pro_array = perspective_projection(pro_array)

        # t2 = time.time()
        # print(f"rotation time: {(t2-t1)*1000}")

        # win.fill((0, 0, 0))

        # t3 = time.time()
        pro_array = prepare_drawing(pro_array, len(array[0]), upscale=True)
        # ref_array = find_reference_points(pro_array, lines)
        # t4 = time.time()
        # print(f"prepare drawing: {(t4-t3)*1000}")
        # draw(pro_array, lines, win)
        surface = draw_np(w, h, pro_array, lines)
        pg.surfarray.blit_array(win, surface)
        # t5 = time.time()
        # print(f"drawing time: {(t5-t4)*1000}")
        pg.display.update()

    pg.quit()


if __name__ == "__main__":
    main()