[Example] add quadtree.py

[archibate]

Related issue = #814

Learning bitmasked.

[Click here for the format server]

[KLozes]

This is a neat example! I went ahead and improved it a little. Not sure how I can I can push to this branch so here are my changes.

First, and improvement to the parent() function in expr.py. Now we can input an integer to grab parents of parents
expr.py

    def parent(self, n=1):
        import taichi as ti
        p = self.ptr.snode()
        for i in range(n):
            p = p.parent
        return Expr(ti.core.global_var_expr_from_snode(p))

Second, in quadtree.py, we can now specifty the number of levels with an integer, instead of making them manually
quadtree.py

import taichi as ti
import numpy as np

ti.init(arch=ti.x64)

RES = 1024
K = 2
R = 7
N = K ** R

Broot = ti.root
B = ti.root
for r in range(R):
    B = B.bitmasked(ti.ij, (K, K))

qt = ti.var(ti.f32)
B.place(qt)

img = ti.Vector(3, dt=ti.f32, shape=(RES, RES))

@ti.kernel
def action(p: ti.ext_arr()):
    a = ti.cast(p[0] * N, ti.i32)
    b = ti.cast(p[1] * N, ti.i32)
    qt[a, b] = 1

@ti.func
def draw_rect(b, i, j, s, k, dx, dy):
    x = i // s
    y = j // s
    a = 0
    if dx and i % k == 0 or dy and j % k == 0:
        a += ti.is_active(b, [x, y])
        a += ti.is_active(b, [x - dx, y - dy])
    return a

@ti.kernel
def paint():
    for i, j in img:
        for k in ti.static(range(3)):
            img[i, j][k] *= 0.85
    for i, j in img:
        s = RES // N
        for r in ti.static(range(R)):
            k = RES // K ** (R-r)
            ia = draw_rect(qt.parent(r+1), i, j, s, k, 1, 0)
            ja = draw_rect(qt.parent(r+1), i, j, s, k, 0, 1)
            img[i, j][0] += (ia + ja) * ((R-r) / R) ** 2


def vec2_npf32(m):
    return np.array([m[0], m[1]], dtype=np.float32)

gui = ti.GUI('Quadtree', (RES, RES))
while not gui.get_event(ti.GUI.PRESS):
    Broot.deactivate_all()
    pos = gui.get_cursor_pos()
    action(vec2_npf32(pos))
    paint()
    gui.set_image(img.to_numpy(as_vector=True))
    gui.show()

[archibate]

we can now specifty the number of levels with an integer, instead of making them manually

Cool! This is a useful feature, now we can easily tweak R and obtain different results!

[KLozes]

Another thing. Is it possible to turn off the cursor when on the gui? That way the use should be able to see the subdivision better.

[archibate]

Yes, we can! But it's wise to do this in a separate PR since it involves many C-side works: https://zgserver.com/linux-x11.html

[KLozes]

Ultimately, I think we should try to make a general purpose Barnes-hut or FMM solver, like we have for MGPCG. That way users can apply it to solve lots of nbody problems, like astrophysical, molecular dynamics, vortex particles, etc.

[yuanming-hu]

Looks good to me! Thank you, guys!

(Sorry about my delayed review...)