[lang] [test] Improve code coverage in SNode

python/taichi/lang/expr.py

@@ -67,9 +67,6 @@ def __getitem__(self, key):
     def loop_range(self):
         return self
 
-    def serialize(self):
-        return self.ptr.serialize()
-
     @python_scope
     def initialize_accessor(self):
         if self.getter:

python/taichi/lang/matrix.py

@@ -36,9 +36,7 @@ def __init__(self,
             if rows is not None and cols is not None:
                 raise Exception("cannot specify both rows and columns")
             self.dt = dt
-            import taichi as ti
-            mat = ti.Matrix.cols(cols) if cols is not None else ti.Matrix.rows(
-                rows)
+            mat = Matrix.cols(cols) if cols is not None else Matrix.rows(rows)
             self.n = mat.n
             self.m = mat.m
             self.entries = mat.entries
@@ -362,8 +360,7 @@ def at(self, index):
         return ret
 
     def empty_copy(self):
-        import taichi as ti
-        return ti.Matrix.empty(self.n, self.m)
+        return Matrix.empty(self.n, self.m)
 
     def zeros_copy(self):
         return Matrix(self.n, self.m)
@@ -698,15 +695,13 @@ def __repr__(self):
     @taichi_scope
     def zero(dt, n, m=1):
         import taichi as ti
-        return ti.Matrix([[ti.cast(0, dt) for _ in range(m)]
-                          for _ in range(n)])
+        return Matrix([[ti.cast(0, dt) for _ in range(m)] for _ in range(n)])
 
     @staticmethod
     @taichi_scope
-    def one(dt, n):
+    def one(dt, n, m=1):
         import taichi as ti
-        return ti.Matrix([[ti.cast(1, dt) for _ in range(n)]
-                          for _ in range(n)])
+        return Matrix([[ti.cast(1, dt) for _ in range(m)] for _ in range(n)])
 
     @staticmethod
     @taichi_scope
@@ -715,31 +710,29 @@ def unit(n, i, dt=None):
         if dt is None:
             dt = ti.get_runtime().default_ip
         assert 0 <= i < n
-        return ti.Matrix([ti.cast(int(j == i), dt) for j in range(n)])
+        return Matrix([ti.cast(int(j == i), dt) for j in range(n)])
 
     @staticmethod
     @taichi_scope
     def identity(dt, n):
         import taichi as ti
-        return ti.Matrix([[ti.cast(int(i == j), dt) for j in range(n)]
-                          for i in range(n)])
+        return Matrix([[ti.cast(int(i == j), dt) for j in range(n)]
+                       for i in range(n)])
 
     @staticmethod
     @taichi_scope
     def rotation2d(alpha):
         import taichi as ti
-        return ti.Matrix([[ti.cos(alpha), -ti.sin(alpha)],
-                          [ti.sin(alpha), ti.cos(alpha)]])
+        return Matrix([[ti.cos(alpha), -ti.sin(alpha)],
+                       [ti.sin(alpha), ti.cos(alpha)]])
 
     @staticmethod
-    def var(n, m, dt, **kwargs):
-        import taichi as ti
-        return ti.Matrix(n=n, m=m, dt=dt, **kwargs)
+    def var(n, m, dt, shape=None, offset=None, **kwargs):
+        return Matrix(n=n, m=m, dt=dt, shape=shape, offset=offset, **kwargs)
 
     @staticmethod
     def rows(rows):
-        import taichi as ti
-        mat = ti.Matrix()
+        mat = Matrix()
         mat.n = len(rows)
         if isinstance(rows[0], Matrix):
             for row in rows:
@@ -763,14 +756,11 @@ def rows(rows):
 
     @staticmethod
     def cols(cols):
-        import taichi as ti
-        return ti.Matrix.rows(cols).transpose()
+        return Matrix.rows(cols).transpose()
 
     @staticmethod
     def empty(n, m):
-        import taichi as ti
-        mat = ti.Matrix([[None] * m for _ in range(n)])
-        return mat
+        return Matrix([[None] * m for _ in range(n)])
 
     def __hash__(self):
         # TODO: refactor KernelTemplateMapper
@@ -816,10 +806,11 @@ def Vector(n, dt=None, shape=None, offset=None, **kwargs):
     return Matrix(n, 1, dt=dt, shape=shape, offset=offset, **kwargs)
 
 
+Vector.var = Vector
 Vector.zero = Matrix.zero
 Vector.one = Matrix.one
 Vector.dot = Matrix.dot
 Vector.cross = Matrix.cross
 Vector.outer_product = Matrix.outer_product
 Vector.unit = Matrix.unit
-Vector.normalized = Matrix.normalized
+Vector.normalized = Matrix.normalized

tests/python/test_linalg.py

@@ -239,11 +239,19 @@ def test_mat_inverse():
 
 
 @ti.all_archs
-def test_unit_vectors():
-    a = ti.Vector(3, dt=ti.i32, shape=3)
+def test_matrix_factories():
+    import math
+
+    a = ti.Vector.var(3, dt=ti.i32, shape=3)
+    b = ti.Matrix.var(2, 2, dt=ti.f32, shape=2)
+    c = ti.Matrix.var(2, 3, dt=ti.f32, shape=2)
 
     @ti.kernel
     def fill():
+        b[0] = ti.Matrix.identity(ti.f32, 2)
+        b[1] = ti.Matrix.rotation2d(math.pi / 3)
+        c[0] = ti.Matrix.zero(ti.f32, 2, 3)
+        c[1] = ti.Matrix.one(ti.f32, 2, 3)
         for i in ti.static(range(3)):
             a[i] = ti.Vector.unit(3, i)
 
@@ -253,6 +261,13 @@ def fill():
         for j in range(3):
             assert a[i][j] == int(i == j)
 
+    sqrt3o2 = math.sqrt(3) / 2
+    assert b.at(0).to_numpy() == approx(np.array([[1.0, 0.0], [0.0, 1.0]]))
+    assert b.at(1).to_numpy() == approx(
+        np.array([[0.5, -sqrt3o2], [sqrt3o2, 0.5]]))
+    assert c.at(0).to_numpy() == approx(np.zeros((2, 3)))
+    assert c.at(1).to_numpy() == approx(np.ones((2, 3)))
+
 
 # TODO: move codes below to test_matrix.py:
 

tests/python/test_tensor_reflection.py

@@ -2,7 +2,7 @@
 
 
 @ti.all_archs
-def test_POT1():
+def test_POT():
     val = ti.var(ti.i32)
 
     n = 4
@@ -11,22 +11,32 @@ def test_POT1():
 
     ti.root.dense(ti.i, n).dense(ti.j, m).dense(ti.k, p).place(val)
 
-    assert val.dim() == 3
     assert val.shape() == (n, m, p)
+    assert val.dim() == 3
+    assert val.data_type() == ti.f32
 
 
 @ti.all_archs
-def test_POT2():
+def test_non_POT():
     val = ti.var(ti.i32)
 
     n = 3
     m = 7
     p = 11
 
-    ti.root.dense(ti.i, n).dense(ti.j, m).dense(ti.k, p).place(val)
+    blk1 = ti.root.dense(ti.i, n)
+    blk2 = blk1.dense(ti.j, m)
+    blk3 = blk2.dense(ti.k, p)
+    blk3.place(val)
 
-    assert val.dim() == 3
     assert val.shape() == (n, m, p)
+    assert val.dim() == 3
+    assert val.data_type() == ti.f32
+    assert val.snode() == blk3
+    assert val.snode().parent() == blk2
+    assert val.parent(0) == blk3
+    assert val.parent() == blk2
+    assert val.parent(2) == blk1
 
 
 @ti.all_archs
@@ -37,7 +47,39 @@ def test_unordered():
     m = 7
     p = 11
 
-    ti.root.dense(ti.k, n).dense(ti.i, m).dense(ti.j, p).place(val)
+    blk1 = ti.root.dense(ti.k, n)
+    blk2 = blk1.dense(ti.i, m)
+    blk3 = blk2.dense(ti.j, p)
+    blk3.place(val)
+
+    assert val.data_type() == ti.f32
+    assert val.shape() == (n, m, p)
+    assert val.dim() == 3
+    assert val.snode() == blk3
+    assert val.snode().parent() == blk2
+    assert val.parent(0) == blk3
+    assert val.parent() == blk2
+    assert val.parent(2) == blk1
+
+
+@ti.all_archs
+def test_unordered_matrix():
+    val = ti.Matrix(3, 2, ti.i32)
+
+    n = 3
+    m = 7
+    p = 11
+
+    blk1 = ti.root.dense(ti.k, n)
+    blk2 = blk1.dense(ti.i, m)
+    blk3 = blk2.dense(ti.j, p)
+    blk3.place(val)
 
     assert val.dim() == 3
     assert val.shape() == (n, m, p)
+    assert val.data_type() == ti.f32
+    assert val.snode() == blk3
+    assert val.snode().parent() == blk2
+    assert val.parent(0) == blk3
+    assert val.parent() == blk2
+    assert val.parent(2) == blk1