galaxy.py

import taichi as ti

ti.init(arch=ti.gpu)


def main():
    n = 800
    pixels = ti.field(dtype=ti.f32, shape=(n, n))
    star_positions = ti.Vector.field(2, dtype=ti.i32, shape=(1000))

    star_offsets = ti.Vector.field(2, dtype=ti.i32, shape=())
    rotation_angle = ti.field(dtype=ti.f32, shape=())

    @ti.kernel
    def initialize():
        for i, j in pixels:
            pixels[i, j] = 0.0

        # Initialize star positions
        for k in range(1000):
            i, j = int(ti.random() * n), int(ti.random() * n)
            star_positions[k] = ti.Vector([i, j])

    @ti.func
    def rotate(position, angle):
        x = position[0] - n / 2
        y = position[1] - n / 2

        new_x = x * ti.cos(angle) - y * ti.sin(angle)
        new_y = x * ti.sin(angle) + y * ti.cos(angle)

        new_position = ti.Vector([new_x + n / 2, new_y + n / 2])
        return new_position

    @ti.kernel
    def paint():
        for i, j in pixels:
            pixels[i, j] = 0.0

        for k in range(1000):
            i, j = star_positions[k]
            pixels[i, j] = 1.0
            new_position = rotate(star_positions[k] + star_offsets[None], rotation_angle[None])
            if 0 <= new_position[0] < n and 0 <= new_position[1] < n:
                pixels[int(new_position[0]), int(new_position[1])] = 1.0

    gui = ti.GUI("Stars", res=(n, n))

    initialize()
    star_offsets[None] = ti.Vector([0, 0])
    rotation_angle[None] = 0

    while gui.running:
        for e in gui.get_events(ti.GUI.PRESS):
            if e.key == ti.GUI.ESCAPE:
                gui.running = False
            elif e.key == ti.GUI.UP:
                star_offsets[None][0] -= 1
            elif e.key == ti.GUI.DOWN:
                star_offsets[None][0] += 1
            elif e.key == ti.GUI.LEFT:
                star_offsets[None][1] -= 1
            elif e.key == ti.GUI.RIGHT:
                star_offsets[None][1] += 1
            elif e.key == 'a':
                rotation_angle[None] -= 0.02
            elif e.key == 'd':
                rotation_angle[None] += 0.02

        paint()
        gui.set_image(pixels)
        gui.show()


if __name__ == "__main__":
    main()