taichi_particles.py

import taichi as ti
import numpy as np
import math

@ti.data_oriented
class Particles():
    def __init__(self,
                 x=1024,
                 y=1024,
                 n=128):
        self._x = x
        self._y = y
        self._n = n
        self._pos = ti.Vector.field(2, dtype=ti.f32, shape=(self._n))
        self._vel = ti.Vector.field(2, dtype=ti.f32, shape=(self._n))
        self.px_rgb = ti.Vector.field(3, dtype=ti.f32, shape=(self._x, self._y))
        self.randomise()

    @ti.kernel
    def randomise(self):
        for b in range(self._n):
            self._pos[b] = ti.Vector([ti.random(ti.f32)*self._x, ti.random(ti.f32)*self._y])
            self._vel[b] = ti.Vector([ti.random(ti.f32)-0.5, ti.random(ti.f32)-0.5])

    @ti.kernel
    def move(self):
        for b in range(self._n):
            self._pos[b] = self._pos[b] + 0.5 * self._vel[b]
            x = self._pos[b][0]
            y = self._pos[b][1]
            if   (x > self._x): self._pos[b][0] = 1
            elif (y > self._y): self._pos[b][1] = 1
            elif (x < 0):       self._pos[b][0] = self._x-1
            elif (y < 0):       self._pos[b][1] = self._y-1

    @ti.kernel
    def clear(self):
        for i,j in ti.ndrange((0, self._x),(0, self._y)):
            self.px_rgb[i, j] = ti.Vector([0.0,0.0,0.0])

    @ti.kernel
    def render(self):
        for b in range(self._n):
            xi = ti.cast(self._pos[b][0], ti.i32) - 2.0
            xj = ti.cast(self._pos[b][0], ti.i32) + 2.0
            yi = ti.cast(self._pos[b][1], ti.i32) - 2.0
            yj = ti.cast(self._pos[b][1], ti.i32) + 2.0
            for x in range(xi, xj):
                for y in range(yi, yj):
                    self.px_rgb[x, y] = ti.Vector([1.0, 1.0, 1.0])

    def process(self):
        self.move()
        self.clear()
        self.render()
        return self.px_rgb

def main():
    ti.init()
    x = 512
    y = 512
    n = 32
    particles = Particles(x, y, n)
    window = ti.ui.Window("Particles", (x, y))
    canvas = window.get_canvas()

    while window.running:
        canvas.set_image(particles.process())
        window.show()

if __name__ == '__main__':
    main()