diff --git a/docs/scalar_tensor.rst b/docs/scalar_tensor.rst index 94ae57da44e70..46b3c34e35d39 100644 --- a/docs/scalar_tensor.rst +++ b/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)) diff --git a/docs/vector.rst b/docs/vector.rst index 534007dfcfbf1..4e4303c5d5a06 100644 --- a/docs/vector.rst +++ b/docs/vector.rst @@ -138,16 +138,22 @@ 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 components) + :parameter b: (Vector of the same size as a) + :return: (scalar (for 2D inputs), or 3D Vector (for 3D inputs)) the cross product of ``a`` and ``b`` 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) + # c = [2*6 - 5*3, 4*3 - 1*6, 1*5 - 4*2] = [-3, 6, -3] + + p = ti.Vector([1, 2]) + q = ti.Vector([4, 5]) + r = ti.cross(a, b) + # r = 1*5 - 4*2 = -3 .. function:: ti.outer_product(a, b) diff --git a/python/taichi/lang/matrix.py b/python/taichi/lang/matrix.py index 7caa934c1d5e3..3a5bf03c6a319 100644 --- a/python/taichi/lang/matrix.py +++ 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 2D/3D vectors" + ) @staticmethod def outer_product(a, b): diff --git a/tests/python/test_linalg.py b/tests/python/test_linalg.py index 94bbdc7152eab..1c52168524e88 100644 --- a/tests/python/test_linalg.py +++ 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