[Example] Add comet.py and fix to_numpy() on sparse matrix fields (ta…

…ichi-dev#1583)

* tmp save

* [Example] Add comet.py and fix to_numpy() on sparse matrix fields

* finalize

* [skip ci] enforce code format

* [skip ci] ti.no_activate for better performance

* [skip ci] enforce code format

Co-authored-by: Taichi Gardener <[email protected]>

examples/comet.py

@@ -0,0 +1,86 @@
+import taichi as ti
+import taichi_glsl as tl
+ti.init(arch=[ti.cuda, ti.metal])
+
+dim = 3
+N = 1024 * 8
+dt = 2e-4
+steps = 7
+sun = ti.Vector([0.5, 0.5] if dim == 2 else [0.5, 0.5, 0.0])
+gravity = 0.5
+pressure = 0.3
+pps_scale = 0.4
+colorinit = 0.3
+colordeg = 1.6
+initvel = 0.07
+res = 640
+
+inv_m = ti.field(ti.f32)
+color = ti.field(ti.f32)
+x = ti.Vector.field(dim, ti.f32)
+v = ti.Vector.field(dim, ti.f32)
+ti.root.bitmasked(ti.i, N).place(x, v, inv_m, color)
+xcnt = ti.field(ti.i32, ())
+img = ti.field(ti.f32, (res, res))
+
+
+@ti.kernel
+def substep():
+    ti.no_activate(x)
+    for i in x:
+        r = x[i] - sun
+        ir = r / r.norm(1e-3)**3
+        acci = pressure * ir * inv_m[i]
+        acci += -gravity * ir
+        v[i] += acci * dt
+
+        x[i] += v[i] * dt
+        color[i] *= ti.exp(-dt * colordeg)
+
+        if not all(-0.1 <= x[i] <= 1.1):
+            ti.deactivate(x.snode().parent(), [i])
+
+
+@ti.kernel
+def generate():
+    r = x[0] - sun
+    ir = 1 / r.norm(1e-3)**2
+    pps = int(pps_scale * ir)
+    for _ in range(pps):
+        r = x[0]
+        if ti.static(dim == 3):
+            r = tl.randUnit3D()
+        else:
+            r = tl.randUnit2D()
+        xi = ti.atomic_add(xcnt[None], 1) % (N - 1) + 1
+        x[xi] = x[0]
+        v[xi] = r * initvel + v[0]
+        inv_m[xi] = 0.5 + ti.random()
+        color[xi] = colorinit
+
+
+@ti.kernel
+def render():
+    for p in ti.grouped(img):
+        img[p] = 1e-6 / (p / res - ti.Vector([sun.x, sun.y])).norm(1e-4)**3
+    for i in x:
+        p = int(ti.Vector([x[i].x, x[i].y]) * res)
+        img[p] += color[i]
+
+
+inv_m[0] = 0
+x[0].x = +0.5
+x[0].y = -0.01
+v[0].x = +0.6
+v[0].y = +0.4
+color[0] = 1
+
+gui = ti.GUI('Comet', res)
+while gui.running:
+    gui.running = not gui.get_event(gui.ESCAPE)
+    generate()
+    for s in range(steps):
+        substep()
+    render()
+    gui.set_image(img)
+    gui.show()

python/taichi/lang/__init__.py

@@ -211,7 +211,7 @@ def init(arch=None,
 
 def no_activate(*args):
     for v in args:
-        taichi_lang_core.no_activate(v.ptr)
+        taichi_lang_core.no_activate(v.snode().ptr)
 
 
 def cache_shared(*args):

python/taichi/lang/matrix.py

@@ -582,6 +582,9 @@ def dtype(self):
     def data_type(self):
         return self.dtype
 
+    def snode(self):
+        return self.loop_range().snode()
+
     def make_grad(self):
         ret = self.empty_copy()
         for i in range(len(ret.entries)):
@@ -673,7 +676,7 @@ def to_numpy(self, keep_dims=False, as_vector=None):
         if not self.is_global():
             return np.array(self.entries).reshape(shape_ext)
 
-        ret = np.empty(self.shape + shape_ext, dtype=to_numpy_type(self.dtype))
+        ret = np.zeros(self.shape + shape_ext, dtype=to_numpy_type(self.dtype))
         from .meta import matrix_to_ext_arr
         matrix_to_ext_arr(self, ret, as_vector)
         import taichi as ti