ball_in_field_simpler.py

"""
All the equations below assumes that the agent is a ball with radius R, and the potential force
is being produced only by L2 distance between current agents location and target point
"""

import taichi as ti
import numpy as np


ti.init(arch=ti.cpu, default_fp=ti.f32)

constants = {
    "radius": 0.05,
    "g": 9.8,
    "f": 0.001,
    "ro": 1000,
}
constants["volume"] = 4 * np.pi * (constants["radius"] ** 3) / 3
constants["mass"] = constants["volume"] * constants["ro"]

sim_steps = 4000


@ti.func
def sim_step(dt,):
    """Makes one step of the simulation

    Args:
        dt (float): time intrval

    Returns:
        list: [description]
    """

    acceleration = (2 * (target_x - x)) / constants["mass"]

    ti.atomic_add(v, acceleration * dt)
    ti.atomic_add(x, v * dt)


@ti.kernel
def run_simulation():
    t0 = 0
    for t1 in range(0, 20, sim_steps):
        dt = t1 - t0
        sim_step(dt=dt,)
        t0 = t1


if __name__ == "__main__":
    target_x = ti.Vector(2, dt=ti.f32, shape=1)
    x = ti.Vector(2, dt=ti.f32, shape=1)
    v = ti.Vector(2, dt=ti.f32, shape=1)

    x[0] = [2.0, 0.5]
    target_x[0] = [0.0, 0.0]
    v[0] = [0.1, 0.5]

    run_simulation()