用Taichi（python）实现星际穿越特效

import taichi as ti
import math
ti.init(arch=ti.gpu)

#
fps=60.0
dt=1.0/fps

#首先设置一个窗口
res = (1920,1280)
pixels = ti.Vector.field(n=3,dtype=float,shape=res)#设置像素的参数

#world space varibles
grid_size = 100
z_near,z_far = 200,3200
N=(res[0]//grid_size,res[1]//grid_size,(z_far-z_near)//grid_size)
pos=ti.Vector.field(n=3,dtype=float,shape=N)
color=ti.Vector.field(n=3,dtype=float,shape=N)
vel=ti.Vector.field(n=3,dtype=float,shape=())

#camera related varibles
fov=120
tan_half_fov=math.tan((fov/360)*math.pi)

@ti.func
def rand3():
    return ti.Vector([ti.random(),ti.random(),ti.random()])

@ti.func
def spawn(I):
    pos[I]=(I+rand3())*grid_size
    color[I]=rand3()+ti.Vector([1.0,2.0,1.0])
    
#init the stars
@ti.kernel
def init():
    for i,j,k in pos:
        I=ti.Vector([i,j,k])
        spawn(I)

#step
@ti.kernel
def step():
    for I in ti.grouped(pos):
        pos[I] += vel[None] * dt
        if pos[I].z < z_near:
            pos[I].z+=z_far-z_near

#projections
@ti.func
def project(pos_3d):
    center=ti.Vector(res)/2
    w=tan_half_fov*pos_3d.z
    rel_pos=(ti.Vector([pos_3d.x,pos_3d.y])-center)/w
    screen_space_pos=rel_pos*res[1]+center
    return screen_space_pos

#画
@ti.func
def draw_particle(c,rad,color):
    x_range=(ti.max(int(c.x-rad-1),0)),ti.min(int(c.x+rad+1),res[0])#使得遍历元素少了
    y_range=(ti.max(int(c.y-rad-1),0)),ti.min(int(c.y+rad+1),res[1])
    for I in ti.grouped(ti.ndrange(x_range,y_range)):
        dist=(I-c).norm()
        if dist<rad:
            pixels[I]=color

@ti.func
def draw_line(p1,p2,rad,color):
    x_range=(ti.max(int(ti.min(p1.x,p2.x)-rad-1),0),ti.min(int(ti.max(p1.x,p2.x)+rad+1),res[0]))
    y_range=(ti.max(int(ti.min(p1.y,p2.y)-rad-1),0),ti.min(int(ti.max(p1.y,p2.y)+rad+1),res[1]))
    for I in ti.grouped(ti.ndrange(x_range,y_range)):
        dist=ti.min((I-p1).norm(),(I-p2).norm())
        p1I=I-p1
        p1p2_normalized=(p2-p1).normalized()
        p=p1+p1I.dot(p1p2_normalized)*p1p2_normalized
        dist2=(I-p).norm()
        if 0< p1I.dot(p1p2_normalized)<(p2-p1).norm():
            dist=dist2
        if dist<rad:
            pixels[I]=color

@ti.kernel
def paint():
    for i ,j in pixels:
        pixels[i,j]=ti.Vector([0.0,0.0,0.0])
    
    for I in ti.grouped(pos):
        p_current=pos[I]
        p_previous=pos[I]-vel[None]*dt
        rad =2.0*(1.0-pos[I].z/z_far)**2
        col=color[I]*(1.0-pos[I].z/z_far)**0.5
        col *=1.0-(2*((pos[I].x/res[0]-0.5)**2+(pos[I].y/res[1]-0.5)**2))**0.2
       # draw_particle(project(pos[I]),rad,color[I])
        draw_line(project(p_current),project(p_previous),rad,col)

#main
init()
gui = ti.GUI("happy 1025",res,fast_gui = True)
t=0.0
while gui.running:
    t+=dt
    vel[None].z=-1000*(1+math.cos(t*math.pi))
    step()
    paint()
    gui.set_image(pixels)
    gui.show()