[Lang] ti.cross now supports 2D vectors

docs/scalar_tensor.rst

@@ -92,7 +92,7 @@ You can access an element of the Taichi tensor by an index or indices.
     This sets the element value at index ``2`` of 1D tensor ``b`` to ``5``:
     ::
 
-        a[2] = 2
+        b[2] = 5
 
     .. note ::
 
@@ -158,6 +158,7 @@ Meta data
     :return: (SNode) the parent of ``a``'s containing SNode
 
     ::
+
         x = ti.var(ti.i32)
         y = ti.var(ti.i32)
         blk1 = ti.root.dense(ti.ij, (6, 5))

docs/vector.rst

@@ -138,16 +138,24 @@ Methods
 
 .. function:: ti.cross(a, b)
 
-    :parameter a: (Vector, 3 component)
-    :parameter b: (Vector, 3 component)
-    :return: (Vector, 3D) the cross product of ``a`` and ``b``
+    :parameter a: (Vector, 2, or 3 component)
+    :parameter b: (Vector with the same component number as a)
+    :return: (Vector, 3D) the cross product of ``a`` and ``b``, if the component number is 3
+    :return: (Scalar) the last component of cross product of ``Vector([a,0])`` and ``Vector([b,0])``, if the component number is 2
 
     We use a right-handed coordinate system. E.g.,
     ::
 
         a = ti.Vector([1, 2, 3])
         b = ti.Vector([4, 5, 6])
-        c = ti.cross(a, b) # [2*6 - 5*3, 4*3 - 1*6, 1*5 - 4*2]
+        c = ti.cross(a, b)
+        # [2*6 - 5*3, 4*3 - 1*6, 1*5 - 4*2] = [-3, 6, -3]
+
+        a2 = ti.Vector([1, 2])
+        b2 = ti.Vector([4, 5])
+        c2 = ti.cross(a, b)
+        # [2*0 - 5*0, 4*0 - 1*0, 1*5 - 4*2] = [0, 0, -3]
+        # here it only calculates and returns the last component -3
 
 
 .. function:: ti.outer_product(a, b)

python/taichi/lang/matrix.py

@@ -541,10 +541,12 @@ def __ti_repr__(self):
             yield '['
 
         for j in range(self.m):
-            if j: yield ', '
+            if j:
+                yield ', '
             yield '['
             for i in range(self.n):
-                if i: yield ', '
+                if i:
+                    yield ', '
                 yield self(i, j)
             yield ']'
 
@@ -596,13 +598,20 @@ def dot(self, other):
 
     @staticmethod
     def cross(a, b):
-        assert a.m == 1 and a.n == 3
-        assert b.m == 1 and b.n == 3
-        return Vector([
-            a(1) * b(2) - a(2) * b(1),
-            a(2) * b(0) - a(0) * b(2),
-            a(0) * b(1) - a(1) * b(0),
-        ])
+        if a.n == 3 and a.m == 1 and b.n == 3 and b.m == 1:
+            return Matrix([
+                a(1) * b(2) - a(2) * b(1),
+                a(2) * b(0) - a(0) * b(2),
+                a(0) * b(1) - a(1) * b(0),
+            ])
+
+        elif a.n == 2 and a.m == 1 and b.n == 2 and b.m == 1:
+            return a(0) * b(1) - a(1) * b(0)
+
+        else:
+            raise Exception(
+                "cross product is only supported between pairs of 3D or those of 2D vectors"
+            )
 
     @staticmethod
     def outer_product(a, b):

tests/python/test_linalg.py

@@ -3,6 +3,106 @@
 from taichi import approx
 
 
+@ti.all_archs
+def test_basic_utils():
+    a = ti.Vector(3, dt=ti.f32)
+    b = ti.Vector(3, dt=ti.f32)
+    abT = ti.Matrix(3, 3, dt=ti.f32)
+    aNormalized = ti.Vector(3, dt=ti.f32)
+
+    normA = ti.var(ti.f32)
+    normSqrA = ti.var(ti.f32)
+
+    @ti.layout
+    def place():
+        ti.root.place(a, b, abT, aNormalized, normA, normSqrA)
+
+    @ti.kernel
+    def init():
+        a[None] = ti.Vector([1.0, 2.0, 3.0])
+        b[None] = ti.Vector([4.0, 5.0, 6.0])
+        abT[None] = ti.Matrix.outer_product(a, b)
+
+        normA[None] = a.norm()
+        normSqrA[None] = a.norm_sqr()
+
+        aNormalized[None] = ti.Matrix.normalized(a)
+
+    init()
+
+    for i in range(3):
+        for j in range(3):
+            assert abT[None][i, j] == a[None][i] * b[None][j]
+
+    sqrt14 = np.sqrt(14.0)
+    invSqrt14 = 1.0 / sqrt14
+    assert normSqrA[None] == 14.0
+    assert normA[None] == approx(sqrt14)
+    assert aNormalized[None][0] == approx(1.0 * invSqrt14)
+    assert aNormalized[None][1] == approx(2.0 * invSqrt14)
+    assert aNormalized[None][2] == approx(3.0 * invSqrt14)
+
+
+@ti.all_archs
+def test_cross():
+    a = ti.Vector(3, dt=ti.f32)
+    b = ti.Vector(3, dt=ti.f32)
+    c = ti.Vector(3, dt=ti.f32)
+
+    a2 = ti.Vector(2, dt=ti.f32)
+    b2 = ti.Vector(2, dt=ti.f32)
+    c2 = ti.var(dt=ti.f32)
+
+    @ti.layout
+    def place():
+        ti.root.place(a, b, c, a2, b2, c2)
+
+    @ti.kernel
+    def init():
+        a[None] = ti.Vector([1.0, 2.0, 3.0])
+        b[None] = ti.Vector([4.0, 5.0, 6.0])
+        c[None] = ti.Matrix.cross(a, b)
+
+        a2[None] = ti.Vector([1.0, 2.0])
+        b2[None] = ti.Vector([4.0, 5.0])
+        c2[None] = ti.Matrix.cross(a2, b2)
+
+    init()
+    assert c[None][0] == -3.0
+    assert c[None][1] == 6.0
+    assert c[None][2] == -3.0
+    assert c2[None] == -3.0
+
+
+@ti.all_archs
+def test_dot():
+    a = ti.Vector(3, dt=ti.f32)
+    b = ti.Vector(3, dt=ti.f32)
+    c = ti.var(dt=ti.f32)
+
+    a2 = ti.Vector(2, dt=ti.f32)
+    b2 = ti.Vector(2, dt=ti.f32)
+    c2 = ti.var(dt=ti.f32)
+
+    @ti.layout
+    def place():
+        ti.root.place(a, b, c, a2, b2, c2)
+
+    @ti.kernel
+    def init():
+        a[None] = ti.Vector([1.0, 2.0, 3.0])
+        b[None] = ti.Vector([4.0, 5.0, 6.0])
+        c[None] = ti.Matrix.dot(a, b)
+
+        a2[None] = ti.Vector([1.0, 2.0])
+        b2[None] = ti.Vector([4.0, 5.0])
+        c2[None] = ti.Matrix.dot(a2, b2)
+
+    init()
+    assert c[None] == 32.0
+    assert c2[None] == 14.0
+
+
 @ti.all_archs
 def test_transpose():
     dim = 3