Branched flow with taichi

branch.py

import math

import numpy as np
import random
import matplotlib.pyplot as plt

medium_size = 1024
x = np.linspace(-1., 1., medium_size)
y = np.linspace(-1., 1., medium_size)
x, y = np.meshgrid(x, y, indexing='ij')

u = np.random.randn(medium_size, medium_size)
v = np.random.randn(medium_size, medium_size)
z = (u.astype(np.complex64) + 1.j * v.astype(np.complex64)) / (0.001 + (x*x+y*y).astype(np.complex64))
n = np.abs(np.fft.ifft2(z)).astype(np.float32)
n_np = 1 + 3.2 * n

print("Made field")

import taichi as ti
import taichi_glsl as ts

ti.init(arch=ti.gpu)

resolution = 1024

counts = ti.field(dtype=float, shape=(resolution, resolution))
pixels = ti.field(dtype=float, shape=(resolution, resolution))

n = ti.field(dtype=float, shape=(resolution, resolution))
force_x = ti.field(dtype=float, shape=(resolution, resolution))
force_y = ti.field(dtype=float, shape=(resolution, resolution))

num_particles = 500000
p = ti.field(dtype=ti.f32, shape=(num_particles, 2))
q = ti.field(dtype=ti.f32, shape=(num_particles, 2))

dt = 0.0005
dx = 2. / resolution
dy = 2. / resolution
idx = resolution / 2
idy = resolution / 2
offset_x = resolution / 2
offset_y = resolution / 2

@ti.kernel
def compute_force():
    for i, j in n:
        force_x[i, j] = 0.5 * idx * (n[i + 1, j] - n[i - 1, j])
        force_y[i, j] = 0.5 * idy * (n[i, j + 1] - n[i, j - 1])


@ti.kernel
def init_state():
    for i in range(num_particles):
        p[i, 0] = ti.cos(2 * math.pi * i / num_particles)
        p[i, 1] = ti.sin(2 * math.pi * i / num_particles)
        q[i, 0] = 0
        q[i, 1] = 0

@ti.kernel
def update():
    for i in range(num_particles):
        xy = ts.vec2(idx * q[i, 0] + offset_x, idy * q[i, 1] + offset_y)
        p[i, 0] -= dt * ts.bilerp(force_x, xy)
        p[i, 1] -= dt * ts.bilerp(force_y, xy)
        q[i, 0] += dt * p[i, 0]
        q[i, 1] += dt * p[i, 1]

@ti.kernel
def clear():
    for i, j in pixels:
        counts[i, j] = 0

@ti.kernel
def paint1():
    for i in range(num_particles):
        x_coord = ti.cast(idx * q[i, 0] + offset_x, ti.i32)
        y_coord = ti.cast(idy * q[i, 1] + offset_y, ti.i32)
        counts[x_coord, y_coord] += 0.0025

@ti.kernel
def paint2():
    for i, j in pixels:
        pixels[i, j] = 0.25 * ti.log(1 + counts[i, j])

gui = ti.GUI("Branched flow", res=(resolution, resolution))

n.from_numpy(n_np)
init_state()
compute_force()
clear()

for i in range(1000000):
    for j in range(5):
        update()
        paint1()
    paint2()
    gui.set_image(pixels)
    gui.show()

Note line 80 is an atomic write that sucks up much of the performance. Need to find a better way.