[lang] Support matrix initialization with a list of vectors

docs/linalg.rst

@@ -17,4 +17,145 @@ Matrices
 - ``R, S = ti.polar_decompose(A, ti.f32)``
 - ``U, sigma, V = ti.svd(A, ti.f32)`` (Note that ``sigma`` is a ``3x3`` diagonal matrix)
 
-TODO: doc here better like Vector.
+TODO: doc here better like Vector. WIP
+
+A matrix in Taichi can have two forms:
+
+  - as a temporary local variable. An ``n by m`` matrix consists of ``n * m`` scalar values.
+  - as a an element of a global tensor. In this case, the tensor is an ND array of ``n by m`` matrices.
+
+Declaration
+-----------
+
+As global tensors of matrices
+++++++++++++++++++++++++++++
+
+.. function:: ti.Matrix(n, m, dt=type, shape=shape)
+
+    :parameter n: (scalar) the number of rows in the matrix
+    :parameter m: (scalar) the number of columns in the matrix
+    :parameter type: (DataType) data type of the components
+    :parameter shape: (scalar or tuple) shape the tensor of vectors, see :ref:`tensor`
+
+    For example, this creates a 5x4 tensor of 3x3 matrices:
+    ::
+
+        # Python-scope
+        a = ti.Matrix(3, 3, dt=ti.f32, shape=(5, 4))
+
+.. note::
+
+    In Python-scope, ``ti.var`` declares :ref:`scalar_tensor`, while ``ti.Matrix`` declares tensors of matrices.
+
+
+As a temporary local variable
++++++++++++++++++++++++++++++
+
+.. function:: ti.Matrix([x, y, ...])
+
+    :parameter x: (scalar) the first component of the vector
+    :parameter y: (scalar) the second component of the vector
+
+    For example, this creates a 3x1 matrix with components (2, 3, 4):
+    ::
+
+        # Taichi-scope
+        a = ti.Matrix([2, 3, 4])
+
+.. note::
+
+    this is equivalent to ti.Vector([x, y, ...])
+
+
+.. function:: ti.Matrix([[x, y, ...], [z, w, ...], ...])
+
+    :parameter x: (scalar) the first component of the first column
+    :parameter y: (scalar) the second component of the first column
+    :parameter z: (scalar) the first component of the second column
+    :parameter w: (scalar) the second component of the second column
+
+    For example, this creates a 2x2 matrix with components (2, 3) in the first column and (4, 5) in the second column:
+    ::
+
+        # Taichi-scope
+        a = ti.Matrix([[2, 3], [4, 5])
+
+
+.. function:: ti.Matrix(rows=[v0, v1, v2, ...])
+
+    :parameter v0: (vector) vector of elements forming first row (or column)
+    :parameter v1: (vector) vector of elements forming second row (or column)
+    :parameter v2: (vector) vector of elements forming third row (or column)
+
+    For example, this creates a 3x3 matrix by concactinating vectors into rows (or columns):
+    ::
+
+        # Taichi-scope
+        v0 = ti.Vector([1.0, 2.0, 3.0])
+        v1 = ti.Vector([4.0, 5.0, 6.0])
+        v2 = ti.Vector([7.0, 8.0, 9.0])
+
+        # to specify data in rows
+        a = ti.Matrix(rows=[v0, v1, v2])
+
+        # to specify data in columns instead
+        a = ti.Matrix(cols=[v0, v1, v2])
+
+        # lists can be used instead of vectors
+        a = ti.Matrix(rows=[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])
+
+
+Accessing components
+--------------------
+
+As global tensors of vectors
+++++++++++++++++++++++++++++
+.. attribute:: a[p, q, ...][i, j]
+
+    :parameter a: (tensor of matrices) the tensor of matrices
+    :parameter p: (scalar) index of the first tensor dimension
+    :parameter q: (scalar) index of the second tensor dimension
+    :parameter i: (scalar) row index of the matrix
+    :parameter j: (scalar) column index of the matrix
+
+    This extracts the first component of vector ``a[6, 3]``:
+    ::
+
+        x = a[6, 3][0]
+
+        # or
+        vec = a[6, 3]
+        x = vec[0]
+
+.. note::
+
+    **Always** use two pair of square brackets to access scalar elements from tensors of matrices.
+
+     - The indices in the first pair of brackets locate the matrix inside the tensor of matrices;
+     - The indices in the second pair of brackets locate the scalar element inside the matrix.
+
+    For 0-D tensors of matrices, indices in the first pair of brackets should be ``[None]``.
+
+
+
+As a temporary local variable
++++++++++++++++++++++++++++++
+
+.. attribute:: a[i, j]
+
+    :parameter a: (Matrix) the matrix
+    :parameter i: (scalar) row index of the matrix
+    :parameter j: (scalar) column index of the matrix
+
+    For example, this extracts the element in row 0 column 1 of matrix ``a``:
+    ::
+
+        x = a[0, 1]
+
+    This sets the element in row 1 column 3 of ``a`` to 4:
+    ::
+
+        a[1, 3] = 4
+
+Methods
+-------

docs/vector.rst

@@ -5,9 +5,8 @@ Vectors
 
 A vector in Taichi can have two forms:
 
-  - as a temporary local variable, an ``n``-D vector consists of ``n`` scalar values.
-  - as a global tensor, where each tensor element is a vector. In this case, an ``n``-D vector consists of ``n`` global tensors of scalars.
-    The tensors of scalars, when treated together, can be considered to be **a global tensor of vectors**.
+  - as a temporary local variable. An ``n`` component vector consists of ``n`` scalar values.
+  - as an element of a global tensor. In this case, the tensor is an ND array of ``n`` component vectors
 
 Declaration
 -----------
@@ -21,7 +20,7 @@ As global tensors of vectors
     :parameter type: (DataType) data type of the components
     :parameter shape: (scalar or tuple) shape the tensor of vectors, see :ref:`tensor`
 
-    For example, this creates a 5x4 tensor of 3D vectors:
+    For example, this creates a 5x4 tensor of 3 component vectors:
     ::
 
         # Python-scope
@@ -57,7 +56,7 @@ As global tensors of vectors
     :parameter a: (Vector) the vector
     :parameter p: (scalar) index of the first tensor dimension
     :parameter q: (scalar) index of the second tensor dimension
-    :parameter i: (scalar) index of the vector dimension
+    :parameter i: (scalar) index of the vector component
 
     This extracts the first component of vector ``a[6, 3]``:
     ::
@@ -72,8 +71,8 @@ As global tensors of vectors
 
     **Always** use two pair of square brackets to access scalar elements from tensors of vectors.
 
-     - The indices in the first pair of brackets locate the vector index inside the tensor of vectors;
-     - The indices in the second pair of brackets locate the scalar element index inside the vector.
+     - The indices in the first pair of brackets locate the vector inside the tensor of vectors;
+     - The indices in the second pair of brackets locate the scalar element inside the vector.
 
     For 0-D tensors of vectors, indices in the first pair of brackets should be ``[None]``.
 
@@ -136,8 +135,8 @@ Methods
 
 .. function:: ti.cross(a, b)
 
-    :parameter a: (Vector, 3D)
-    :parameter b: (Vector, 3D)
+    :parameter a: (Vector, 3 component)
+    :parameter b: (Vector, 3 component)
     :return: (Vector, 3D) the cross product of ``a`` and ``b``
 
     We use right-handed coordinate system, E.g.,
@@ -163,7 +162,7 @@ Methods
         # c = [[1*4, 1*5, 1*6], [2*4, 2*5, 2*6], [3*4, 3*5, 3*6]]
 
 .. note::
-    This is not the same as `ti.cross`. ``a`` and ``b`` do not have to be 3D vectors.
+    This is not the same as `ti.cross`. ``a`` and ``b`` do not have to be 3 component vectors.
 
 
 .. function:: a.cast(dt)

python/taichi/lang/matrix.py

@@ -19,16 +19,36 @@ class Matrix:
     is_taichi_class = True
 
     def __init__(self,
-                 n,
+                 n=1,
                  m=1,
                  dt=None,
                  empty=False,
                  shape=None,
                  layout=None,
                  needs_grad=False,
-                 keep_raw=False):
+                 keep_raw=False,
+                 rows=None,
+                 cols=None):
         self.grad = None
-        if isinstance(n, list):
+        if rows is not None or cols is not None:
+            if cols is not None:
+                rows = cols
+            self.n = len(rows)
+            if isinstance(rows[0], Matrix):
+                self.m = rows[0].n
+                self.entries = [row(i) for row in rows for i in range(row.n)]
+            elif isinstance(rows[0], list):
+                self.m = len(rows[0])
+                self.entries = [x for row in rows for x in row]
+            else:
+                raise Exception(
+                    'rows/cols must be list of lists or list of vectors')
+            if cols is not None:
+                t = self.transposed(self)
+                self.n = t.n
+                self.m = t.m
+                self.entries = t.entries
+        elif isinstance(n, list):
             if n == []:
                 mat = []
             elif not isinstance(n[0], list):
@@ -40,6 +60,9 @@ def __init__(self,
                         mat = [list([expr.Expr(x)]) for x in n]
                 else:
                     mat = [[x] for x in n]
+            elif isinstance(n[0], Matrix):
+                raise Exception(
+                    'cols/rows required when using list of vectors')
             else:
                 mat = n
             self.n = len(mat)

taichi/transforms/ir_printer.cpp

@@ -203,9 +203,8 @@ class IRPrinter : public IRVisitor {
   }
 
   void visit(WhileControlStmt *stmt) override {
-    print("{} : while control {}, {}",
-        stmt->name(), stmt->mask ? stmt->mask->name() : "nullptr",
-        stmt->cond->name());
+    print("{} : while control {}, {}", stmt->name(),
+          stmt->mask ? stmt->mask->name() : "nullptr", stmt->cond->name());
   }
 
   void visit(ContinueStmt *stmt) override {

tests/python/test_linalg.py

@@ -143,3 +143,33 @@ def fill():
     for i in range(3):
         for j in range(3):
             assert a[i][j] == int(i == j)
+
+
+@ti.all_archs
+def test_init_matrix_from_vectors():
+    m1 = ti.Matrix(3, 3, dt=ti.f32, shape=(3))
+    m2 = ti.Matrix(3, 3, dt=ti.f32, shape=(3))
+    m3 = ti.Matrix(3, 3, dt=ti.f32, shape=(3))
+    m4 = ti.Matrix(3, 3, dt=ti.f32, shape=(3))
+
+    @ti.kernel
+    def fill():
+        for i in range(3):
+            a = ti.Vector([1.0, 4.0, 7.0])
+            b = ti.Vector([2.0, 5.0, 8.0])
+            c = ti.Vector([3.0, 6.0, 9.0])
+            m1[i] = ti.Matrix(rows=[a, b, c])
+            m2[i] = ti.Matrix(cols=[a, b, c])
+            m3[i] = ti.Matrix(
+                rows=[[1.0, 4.0, 7.0], [2.0, 5.0, 8.0], [3.0, 6.0, 9.0]])
+            m4[i] = ti.Matrix(
+                cols=[[1.0, 4.0, 7.0], [2.0, 5.0, 8.0], [3.0, 6.0, 9.0]])
+
+    fill()
+
+    for j in range(3):
+        for i in range(3):
+            assert m1[0][i, j] == int(i + 3 * j + 1)
+            assert m2[0][j, i] == int(i + 3 * j + 1)
+            assert m3[0][i, j] == int(i + 3 * j + 1)
+            assert m4[0][j, i] == int(i + 3 * j + 1)