diff --git a/examples/quadtree.py b/examples/quadtree.py new file mode 100644 index 0000000000000..16aa398fa6218 --- /dev/null +++ b/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() diff --git a/python/taichi/lang/expr.py b/python/taichi/lang/expr.py index 84934be242460..5952522cc7405 100644 --- a/python/taichi/lang/expr.py +++ b/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 diff --git a/python/taichi/lang/impl.py b/python/taichi/lang/impl.py index aed3e76b429d6..fc3c472f03a64 100644 --- a/python/taichi/lang/impl.py +++ b/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: diff --git a/python/taichi/misc/gui.py b/python/taichi/misc/gui.py index ee2d17895b824..88ed22b4115e1 100644 --- a/python/taichi/misc/gui.py +++ b/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: