[Example] add quadtree.py

examples/quadtree.py

@@ -0,0 +1,61 @@
+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()

python/taichi/lang/expr.py

@@ -352,10 +352,12 @@ def __ti_float__(self):
         import taichi as ti
         return ti.cast(self, ti.get_runtime().default_fp)
 
-    def parent(self):
+    def parent(self, n=1):
         import taichi as ti
-        return Expr(ti.core.global_var_expr_from_snode(
-            self.ptr.snode().parent))
+        p = self.ptr.snode()
+        for i in range(n):
+            p = p.parent
+        return Expr(ti.core.global_var_expr_from_snode(p))
 
     def snode(self):
         from .snode import SNode

python/taichi/lang/impl.py

@@ -20,6 +20,8 @@ def expr_init(rhs):
             return rhs
         elif isinstance(rhs, ti.ndrange):
             return rhs
+        elif isinstance(rhs, ti.SNode):
+            return rhs
         elif hasattr(rhs, '_data_oriented'):
             return rhs
         else:

python/taichi/misc/gui.py

@@ -74,7 +74,7 @@ def set_image(self, img):
                          2] == self.res, "Image resolution does not match GUI resolution"
         self.core.set_img(np.ascontiguousarray(img).ctypes.data)
 
-    def circle(self, pos, color, radius=1):
+    def circle(self, pos, color=0xFFFFFF, radius=1):
         self.canvas.circle_single(pos[0], pos[1], color, radius)
 
     def circles(self, pos, color=0xFFFFFF, radius=1):
@@ -119,10 +119,20 @@ def circles(self, pos, color=0xFFFFFF, radius=1):
     def triangle(self, a, b, c, color=0xFFFFFF):
         self.canvas.triangle_single(a[0], a[1], b[0], b[1], c[0], c[1], color)
 
-    def line(self, begin, end, radius, color):
+    def line(self, begin, end, radius=1, color=0xFFFFFF):
         self.canvas.path_single(begin[0], begin[1], end[0], end[1], color,
                                 radius)
 
+    def rect(self, topleft, bottomright, radius=1, color=0xFFFFFF):
+        a = topleft[0], topleft[1]
+        b = bottomright[0], topleft[1]
+        c = bottomright[0], bottomright[1]
+        d = topleft[0], bottomright[1]
+        self.line(a, b, radius, color)
+        self.line(b, c, radius, color)
+        self.line(c, d, radius, color)
+        self.line(d, a, radius, color)
+
     def show(self, file=None):
         self.core.update()
         if file: