Skip to content

Vectors and matrices

Particles, colours, directions and local deformation all have several components. Tack lets you express their arithmetic as small values instead of maintaining one scalar name per component. A value lives inside an iteration; a vector or matrix array stores one such value per element across iterations and launches.

Start with positions and velocities

import numpy as np
import tack

tack.set_default_device("cpu")
pos = tack.empty((3,), tack.vec2(tack.f32))
vel = tack.empty((3,), tack.vec2(tack.f32))
tack.copyto(pos, np.array([[0, 0], [1, 0], [0, 1]], dtype=np.float32))
vel.fill([1.0, 0.0])

@tack.kernel
def advance(pos, vel, dt):
    for i in range(pos.shape[0]):
        pos[i] += dt * vel[i]

advance(pos, vel, 0.5)
print(pos.numpy())    # [[0.5, 0.0], [1.5, 0.0], [0.5, 1.0]]

Each iteration loads a two-component velocity and updates its own position. dt broadcasts to both components. Multiplying two vectors with * multiplies corresponding components; use dot or @ for a dot product.

A local assignment p = pos[i] loads a value. Changing p.x changes that local; write pos[i] = p to store it back. Writing pos[i].x targets the array directly. Vectors have a fixed width for each compiled specialization. There are no runtime-length vector values or swizzle constructors such as v.rgb/tack.vec4(v, 1.0); construct the components explicitly.

Vectors

Choose a vector element dtype through vec2, vec3, vec4, or vec(n). The component dtype defaults to f32. The ordinary array constructors allocate its storage, and vector operations are scalarized at the IR level:

v = tack.empty((n,), tack.vec3(tack.f32))

@tack.kernel
def normalize_vectors(v, n):
    for i in range(n):
        vec = v[i]                   # loads 3 components
        length = sqrt(vec[0]**2 + vec[1]**2 + vec[2]**2)
        v[i] = vec / length          # stores 3 components

Arithmetic works on whole vectors, component by component, and a scalar operand is applied to every component. So do the math builtins, the casts and conditional expressions:

@tack.kernel
def step(pos, vel, grid, out, dt, n):
    for i in range(n):
        v = vel[i]
        v = min(max(v, -10.0), 10.0)        # clamp each component
        cell = int(floor(pos[i] / 0.25))    # a vector of i32 cell indices
        v = v if pos[i].y > 0.0 else -v     # one condition, whole vectors
        speed = v.norm()                    # methods are listed below
        vel[i] = v
        pos[i] += v * dt                    # augmented store to an array element
        out[i] = grid[cell] * speed         # grid[cell] is grid[cell[0], cell[1], cell[2]]

In a kernel the element-type names build vector values: tack.vec3(x, y, z) from three scalars, tack.vec3(s) with every component s, and tack.vec3(v) from a 3-vector (vec2 and vec4 likewise). A list of scalars is also a vector, of any width, so it can stand wherever a vector is expected: pos[i] = [x, y], v = [0.0, 0.0], w.dot([1.0, 0.0]), p - [1.0, 0.0].

The vector methods are norm(), norm_sqr(), dot(w), cross(w) (3-vectors), normalized(), and the reductions over components sum(), min() and max(). norm(eps) is sqrt(norm_sqr() + eps), a length that is never zero. normalized(eps) divides by norm() + eps, for a vector that may be zero. min and max as functions take two or more values, as in Python: min(a, b, c).

Comparing vectors compares component by component and gives a mask, a vector of 0/1. any and all reduce a mask to one truth value, and tack.select(mask, a, b) picks per component:

@tack.kernel
def confine(pos, vel, alive, n):
    for i in range(n):
        p = pos[i]
        inside = -1.0 <= p <= 1.0                  # a mask, every link evaluated
        if not all(inside):
            alive[i] = 0
        vel[i] = tack.select(inside, vel[i], -vel[i])    # reflect the components outside
        pos[i] = tack.select(p > 1.0, 1.0, p)             # a scalar arm fills every component
        hits = (p > 0.9) and (vel[i] > 0.0)              # and/or act per component on masks
        if any(hits):
            alive[i] += hits.sum()

A mask is an ordinary integer vector, so mask.sum() counts. if and while take a scalar condition: a mask there is rejected with the hint to reduce it with any or all.

A device function can return several values, vectors among them, to be unpacked at the call:

@tack.func
def closest_hit(origin, direction):
    ...
    return distance, normal, color        # a scalar and two vectors

distance, normal, color = closest_hit(o, d)

A vector array's shape excludes its component dimension. A vector array of shape (n,) and dtype tack.vec3(tack.f32) downloads as (n, 3) with numpy(), and tack.copyto takes values of that shape, or values that broadcast to it (one (3,) vector, say), not the flat storage shape. size is n; scalar_size is 3 * n. scalar_dtype gives the component type. Copies and reshapes preserve the element type; reshaping changes only the array dimensions, without copying storage.

Two vectors in one operation must have the same number of components. An assignment evaluates its whole right side before it stores anything, so v = v.cross(w) and pos[i] = pos[i].cross(axis[i]) read the old components. A tuple assignment does the same and then assigns its targets from left to right, and its targets may be array elements: pos[i], vel[i] = p, v, or a[i], b[i] = b[i], a[i] to swap two.

A store to an array element must match the array: a vector of the array's width, or a scalar, which sets every component (vel[i] = 0.0). A vector of another width, or a vector stored into an array of scalars, is rejected.

Components are scalars. vec[0] and vec.x read one (x, y, z, w name the first four); vec[1] = x, vec.y = x and vec[1] += x write one. The same forms work on an array element without naming the vector first: v[i][2] and v[i].z read a component, v[i][2] = x and v[i].z += x store one. The index may be a runtime value (vec[k] for a loop variable k), which lowers to a chain of selects rather than a branch; a typed or integer tack.constant read that way is instead one load from a constant array. A runtime index must be in [0, n): outside it, on either side, a read gives the last component and a write does nothing, since a kernel cannot raise. Only a literal index counts from the end (vec[-1]), and a literal out of range is rejected at lowering.

Matrices

Matrices are small and fixed-size, up to 4×4. Like vectors they are scalarized: a matrix is its entries, and every operation expands to scalar arithmetic at lowering.

F = tack.empty((n,), tack.mat2(tack.f32))    # a 2x2 matrix per particle
C = tack.empty((n,), tack.mat2(tack.f32))

@tack.kernel
def update(F, C, x, dt, n):
    for p in range(n):
        F[p] = (tack.identity(2) + dt * C[p]) @ F[p]   # '@' multiplies
        J = F[p].determinant()
        stress = (F[p] - F[p].inverse().transpose()) * J      # '*' is entry by entry
        x[p] += stress @ tack.vec2(0.0, -1.0) * dt        # matrix @ vector
  • Building one. tack.mat2([a, b], [c, d]) from rows of scalars, tack.mat2(u, v) from vectors as rows, tack.mat2(a, b, c, d) from four scalars in row-major order (mat3, mat4 likewise); tack.mat(row, ...) for any shape up to 4×4, such as three rows of two; tack.identity(n) for a literal n, as np.identity; or u.outer_product(v).
  • Products. A @ B for matrices; A @ v and v @ A take a vector as a column and as a row and give a vector; u @ v is the dot product. Shapes must agree. +, -, *, /, the math builtins and scalar operands work entry by entry, as for vectors.
  • Methods. transpose(); for square matrices trace(); for 2×2 and 3×3 determinant() and inverse(). inverse() divides by the determinant without checking it. norm() and sum() run over all entries.
  • Entries. A[i, j] reads one and A[i, j] = x or A[i, j] += x writes one, on a matrix variable or directly on an array element (F[p][0, 1]). The indices may be runtime values, which must then be in range.
  • Arrays. A matrix array stores each matrix in row-major order. tack.copyto takes (*shape, n, m) and numpy() returns it. F[p] = A, F[p] += A, F[p] *= s and tack.atomic_add(F, p, A) work as they do for vector arrays.
  • Device functions take and return matrices, alone or among several values.

A matrix and a vector with the same number of entries are different things: storing one where the other belongs, or adding them, is rejected.

The statistics in tack.algorithms (dot, norm, var, ...) accept a vector array and reduce over all its components in storage order.

Matrices in a real computation

In MPM, each particle stores a 2×2 affine velocity matrix. A matrix-vector product predicts a velocity contribution at a nearby grid node; an outer product reconstructs the affine matrix during interpolation. The example shows why * and @ mean different operations, and why a matrix array is clearer than a four-component vector used as a matrix by hand.

The dye tutorial reuses one interpolation function on RGB vectors. N-body accumulates three-component directions. Use those examples after the small position update above.

The same names serve both sides: on the host tack.vec3() and tack.mat2() are element dtypes that allocate arrays of vectors and matrices, and in a kernel tack.vec3(x, y, z) and tack.mat2(...) build values.