Particle Motion Simulation with Taichi

About

The Taichi physics library is a tool designed to efficiently perform high-performance numerical simulations, enabling concise descriptions and fast execution of physics-based simulations.
Taichi leverages Python’s simple syntax while utilizing JIT (Just-In-Time) compilation to harness the power of GPUs and CPUs, achieving high-speed simulations. It is capable of efficiently implementing physics simulations such as fluids, rigid bodies, cloth, and particle systems.
As far as I know, it was originally developed by a Chinese researcher at MIT. It seems that Taichi is used behind Genesis simulator.
For learning purposes, I tested it with a simple sample.

Code

import os
import glob
import argparse

import matplotlib.pyplot as plt
import imageio
import taichi as ti


def main(args: argparse.Namespace):
    
    ti.init(arch=ti.gpu) 

    # parameters
    particles_n = args.particles_n 
    dt = args.time_diff
    gravity = ti.Vector([0, -9.8])
    iterations_total = args.iterations_total
    output_dir = args.output_dir
    video_name = args.video_name
    
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)


    # The state of particles
    positions = ti.Vector.field(2, dtype=ti.f32, shape=particles_n)  # position
    velocities = ti.Vector.field(2, dtype=ti.f32, shape=particles_n)  # velocity

    # The size of boundary
    boundary = ti.Vector([1.0, 1.0])

    # Initilize Kernel
    @ti.kernel
    def initialize():
        for i in range(particles_n):
            # Random Position
            positions[i] = ti.Vector([ti.random(), ti.random()])
            # Random Velocity
            velocities[i] = ti.Vector([ti.random() - 0.5, ti.random() - 0.5]) * 2

    # Kernel : update particle
    @ti.kernel
    def update():
        for i in range(particles_n):
            # update velocity with gravity
            velocities[i] += dt * gravity
            # update positions based on velocity
            positions[i] += dt * velocities[i]

            # Bound on boundary
            for d in ti.static(range(2)):  # on X, Y Axis
                if positions[i][d] < 0 or positions[i][d] > boundary[d]:
                    velocities[i][d] *= -0.8  # Reverse velocity and attenuate
                    positions[i][d] = max(0, min(positions[i][d], boundary[d]))

    # 
    # 
    # 
    initialize()

    # Main Loop
    for step in range(iterations_total):
        # Update Particle
        update()

        # Convert position info with numpy
        positions_np = positions.to_numpy()

        # Plot images with matplot
        plt.figure(figsize=(5, 5))
        plt.scatter(positions_np[:, 0], positions_np[:, 1], s=2, c='blue')
        plt.xlim(0, 1)
        plt.ylim(0, 1)
        plt.title(f"Step: {step}")
        plt.xlabel("x")
        plt.ylabel("y")
        plt.savefig(f"{output_dir}/frame_{step:04d}.png")
        plt.close()

    # Create Video file
    with imageio.get_writer(
        f"{output_dir}/{video_name}", fps=30
    ) as writer:
        for step in range(iterations_total):
            image_path = f"{output_dir}/frame_{step:04d}.png"
            writer.append_data(imageio.imread(image_path))

    png_files = glob.glob(os.path.join(output_dir, "*.png"))
    for file_path in png_files:
        try:
            os.remove(file_path)
        except Exception as e:
            print(f"Failed Removing: {file_path}, Error: {e}")
    print(f"Simulation video saved as {video_name}")


if __name__ == '__main__':
    
    parser = argparse.ArgumentParser(
        description='Particle Motion'
    )
    parser.add_argument(
        '--particles_n', '-PN', type=int,
        default=1000, help=''
    )
    parser.add_argument(
        '--time_diff', '-TD', type=float,
        default=0.01, help=''
    )
    parser.add_argument(
        '--iterations_total', '-IT', type=int,
        default=300, help=''
    )
    parser.add_argument(
        '--output_dir', '-OD', type=str,
        default='./frames', help=''
    )
    parser.add_argument(
        '--video_name', '-VN', type=str,
        default='particle_simulation.mp4', help=''
    )
    args = parser.parse_args()
    main(args)

Result