Basics
Install and set up:
(v1.6) pkg> add TensorCastjulia> using TensorCast
julia> V = [10, 20, 30];
julia> M = Array(reshape(1:12, 3, 4))
3×4 Matrix{Int64}:
1 4 7 10
2 5 8 11
3 6 9 12Broadcasting
Here's a simple use of @cast, broadcasting addition over a matrix, a vector, and a scalar. Using := makes a new array, after which the left and right will be equal for all possible values of the indices:
julia> @cast A[j,i] := M[i,j] + 10 * V[i]
4×3 Matrix{Int64}:
101 202 303
104 205 306
107 208 309
110 211 312
julia> all(A[j,i] == M[i,j] + 10 * V[i] for i in 1:3, j in 1:4)
trueTo make this happen, first M is transposed, and then the axis of V is re-oriented to lie in the second dimension. The macro @pretty prints out what @cast produces:
julia> @pretty @cast A[j,i] := M[i,j] + 10 * V[i]
begin
local panda = transmute(M, (2, 1))
local bat = transmute(V, (nothing, 1))
A = @__dot__(panda + 10bat)
endThe function transmute is a generalised version of permutedims. It transposes M, and then places V's first dimension to lie along the second dimension – also a transpose, here, but it allows for things like transmute(M, (2,nothing,1)). This is from TransmuteDims.jl.
Of course @__dot__ is the full name of the broadcasting macro @., which simply produces A = pelican .+ 10 .* termite. And the animal names (in place of generated symbols) are from the indispensible MacroTools.jl.
Slicing
Here are two ways to slice M into its columns, both of these produce S = sliceview(M, (:, *)) which is simply collect(eachcol(M)):
julia> @cast S[j][i] := M[i,j]
4-element Vector{SubArray{Int64, 1, Matrix{Int64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}:
[1, 2, 3]
[4, 5, 6]
[7, 8, 9]
[10, 11, 12]
julia> all(S[j][i] == M[i,j] for i in 1:3, j in 1:4)
true
julia> @cast S2[j] := M[:,j];
julia> S == S2
trueWe can glue slices back together with the same notation, M2[i,j] := S[j][i]. Combining this with slicing from M[:,j] is a convenient way to perform mapslices operations:
julia> @cast C[i,j] := cumsum(M[:,j])[i]
3×4 lazystack(::Vector{Vector{Int64}}) with eltype Int64:
1 4 7 10
3 9 15 21
6 15 24 33
julia> C == mapslices(cumsum, M, dims=1)
trueThis is both faster and more general than mapslices, as it does not have to infer what shape the function produces:
julia> using BenchmarkTools
julia> f(M) = @cast _[i,j] := cumsum(M[:,j])[i];
julia> M10 = rand(10,1000);
julia> @btime mapslices(cumsum, $M10, dims=1);
630.725 μs (7508 allocations: 400.28 KiB)
julia> @btime f($M10);
64.056 μs (2006 allocations: 297.25 KiB)It's also more readable, in less trivial examples:
julia> using LinearAlgebra, Random; Random.seed!(42);
julia> T = rand(3,3,5);
julia> @cast E[i,n] := eigen(T[:,:,n]).values[i];
julia> E == dropdims(mapslices(x -> eigen(x).values, T, dims=(1,2)), dims=2)
trueFixed indices
Sometimes it's useful to insert a trivial index into the output – for instance to match the output of mapslices(x -> eigen(... just above:
julia> @cast E3[i,_,n] := eigen(T[:,:,n]).values[i];
julia> size(E3)
(3, 1, 5)Using _ on the right will demand that size(E3,2) == 1; but you can also fix indices to other values. If these are variables not integers, they must be interpolated M[$row,j] to distinguish them from index names:
julia> col = 3;
julia> @cast O[i,j] := (M[i,1], M[j,$col])
3×3 Matrix{Tuple{Int64, Int64}}:
(1, 7) (1, 8) (1, 9)
(2, 7) (2, 8) (2, 9)
(3, 7) (3, 8) (3, 9)Reshaping
Sometimes it's useful to combine two (or more) indices into one, which may be written either (i,j) or i⊗j:
julia> vec(M) == @cast _[(i,j)] := M[i,j]
trueThe next-simplest version of this is precisely what the built-in function kron does:
julia> W = 1 ./ [10,20,30,40];
julia> @cast K[_, i⊗j] := V[i] * W[j]
1×12 Matrix{Float64}:
1.0 2.0 3.0 0.5 1.0 1.5 0.333333 0.666667 1.0 0.25 0.5 0.75
julia> @cast K[1, (i,j)] := V[i] * W[j] # identical!
1×12 Matrix{Float64}:
1.0 2.0 3.0 0.5 1.0 1.5 0.333333 0.666667 1.0 0.25 0.5 0.75
julia> K == kron(W, V)'
true
julia> all(K[i + 3(j-1)] == V[i] * W[j] for i in 1:3, j in 1:4)
true
julia> Base.kron(A::Array{T,3}, X::Array{T′,3}) where {T,T′} = # extend kron to 3-tensors
@cast _[x⊗a, y⊗b, z⊗c] := A[a,b,c] * X[x,y,z]If an array on the right has a combined index, then it may be ambiguous how to divide up its range. You can resolve this by providing explicit ranges, after the main expression:
julia> @cast A[i,j] := collect(1:12)[i⊗j] i in 1:2
2×6 Matrix{Int64}:
1 3 5 7 9 11
2 4 6 8 10 12
julia> @cast A[i,j] := collect(1:12)[i⊗j] (i ∈ 1:4, j ∈ 1:3)
4×3 Matrix{Int64}:
1 5 9
2 6 10
3 7 11
4 8 12Writing into a given array with = instead of := will also remove ambiguities, as size(A) is known:
julia> @cast A[i,j] = 10 * collect(1:12)[i⊗j];Repeating
If the right hand side is independent of an index, then the same result is repeated. The range of the index must still be known:
julia> @cast R[r,(n,c)] := M[r,c]^2 (n in 1:3)
3×12 Matrix{Int64}:
1 1 1 16 16 16 49 49 49 100 100 100
4 4 4 25 25 25 64 64 64 121 121 121
9 9 9 36 36 36 81 81 81 144 144 144
julia> R == repeat(M .^ 2, inner=(1,3))
true
julia> @cast similar(R)[r,(c,n)] = M[r,c] # repeat(M, outer=(1,3)), uses size(R)
3×12 Matrix{Int64}:
1 4 7 10 1 4 7 10 1 4 7 10
2 5 8 11 2 5 8 11 2 5 8 11
3 6 9 12 3 6 9 12 3 6 9 12Glue & reshape
As a less trivial example of these combined indices, here is one way to combine a list of matrices into one large grid. Notice that x varies faster than i, and so on – the linear order of index x⊗i agrees with that of two indices x,i.
julia> list = [fill(k,2,2) for k in 1:8];
julia> @cast mat[x⊗i, y⊗j] |= list[i⊗j][x,y] i in 1:2
4×8 Matrix{Int64}:
1 1 3 3 5 5 7 7
1 1 3 3 5 5 7 7
2 2 4 4 6 6 8 8
2 2 4 4 6 6 8 8
julia> vec(mat) == @cast _[xi⊗yj] := mat[xi, yj]
true
julia> mat == hvcat((4,4), transpose(reshape(list,2,4))...)
trueAlternatively, this reshapes each matrix to a vector, and makes them columns of the output:
julia> @cast colwise[x⊗y,i] := (list[i][x,y])^2
4×8 Matrix{Int64}:
1 4 9 16 25 36 49 64
1 4 9 16 25 36 49 64
1 4 9 16 25 36 49 64
1 4 9 16 25 36 49 64
julia> colwise == reduce(hcat, vec.(list)) .^ 2
trueIndex values
Mostly the indices appearing in @cast expressions are just notation, to indicate what permutation / reshape is required. But if an index appears outside of square brackets, this is understood as a value, implemented by broadcasting over a range (appropriately permuted):
julia> @cast _[i,j] := M[i,j]^2 * (i >= j)
3×4 Matrix{Int64}:
1 0 0 0
4 25 0 0
9 36 81 0
julia> ans == M .^2 .* (axes(M,1) .>= transpose(axes(M,2))) # what this generates
true
julia> using OffsetArrays
julia> @cast _[r,c] := r^2 + c^2 (r in -1:1, c in -7:7)
3×15 OffsetArray(::Matrix{Int64}, -1:1, -7:7) with eltype Int64 with indices -1:1×-7:7:
50 37 26 17 10 5 2 1 2 5 10 17 26 37 50
49 36 25 16 9 4 1 0 1 4 9 16 25 36 49
50 37 26 17 10 5 2 1 2 5 10 17 26 37 50Writing $i will interpolate the variable i, distinct from the index i:
julia> i, k = 10, 100;
julia> @cast ones(3)[i] = i + $i + k
3-element Vector{Float64}:
111.0
112.0
113.0Reverse & shuffle
A minus in front of an index will reverse that direction, and a tilde will shuffle it. Both create views, which you may explicitly collect using |=:
julia> @cast M2[i,j] := M[i,-j]
3×4 view(::Matrix{Int64}, :, 4:-1:1) with eltype Int64:
10 7 4 1
11 8 5 2
12 9 6 3
julia> all(M2[i,j] == M[i, end+begin-j] for i in 1:3, j in 1:4)
true
julia> using Random; Random.seed!(42);
julia> @cast M3[i,j] |= M[i,~j]
3×4 Matrix{Int64}:
1 10 7 4
2 11 8 5
3 12 9 6Note that the minus is a slight deviation from the rule that left equals right for all indices, it should really be M[i, end+1-j].
Primes '
Acting on indices, A[i'] is normalised to A[i′] unicode \prime (which looks identical in some fonts). Acting on elements, C[i,j]' means adjoint.(C), elementwise complex conjugate, equivalent to (C')[j,i]. If the elements are matrices, as in C[:,:,k]', then adjoint is conjugate transpose.
julia> @cast C[i,i'] := (1:4)[i⊗i′] + im (i ∈ 1:2, i′ ∈ 1:2)
2×2 Matrix{Complex{Int64}}:
1+1im 3+1im
2+1im 4+1im
julia> @cast _[i,j] := C[i,j]'
2×2 Matrix{Complex{Int64}}:
1-1im 3-1im
2-1im 4-1im
julia> C' == @cast _[j,i] := C[i,j]'
true
julia> @cast cubes[1,k′] := (k')^3 (k' in 1:5) # k' outside square brackets
1×5 Matrix{Int64}:
1 8 27 64 125Splats
You can use a slice of an array as the arguments of a function, or a struct:
julia> struct Tri{T} x::T; y::T; z::T end;
julia> @cast triples[k] := Tri(M[:,k]...)
4-element Vector{Tri{Int64}}:
Tri{Int64}(1, 2, 3)
Tri{Int64}(4, 5, 6)
Tri{Int64}(7, 8, 9)
Tri{Int64}(10, 11, 12)
julia> ans == Base.splat(Tri).(eachcol(M))
trueThis one can also be done as reinterpret(reshape, Tri{Int64}, M).
Arrays of functions
Besides arrays of numbers (and arrays of arrays) you can also broadcast an array of functions, which is done by calling Core._apply(f, xs...) = f(xs...):
julia> funs = [identity, sqrt];
julia> @cast applied[i,j] := funs[i](V[j])
2×3 Matrix{Real}:
10 20 30
3.16228 4.47214 5.47723
julia> @pretty @cast applied[i,j] := funs[i](V[j])
begin
@boundscheck funs isa Tuple || (ndims(funs) == 1 || throw(ArgumentError("expected a vector or tuple funs[i]")))
@boundscheck V isa Tuple || (ndims(V) == 1 || throw(ArgumentError("expected a vector or tuple V[j]")))
local octopus = transmute(V, Val((nothing, 1)))
applied = @__dot__(_apply(funs, octopus))
endRepeated indices
The only situation in which repeated indices are allowed is when they either extract the diagonal of a matrix, or create a diagonal matrix:
julia> @cast D[i] |= C[i,i]
2-element Vector{Complex{Int64}}:
1 + 1im
4 + 1im
julia> D2 = @cast _[i,i] := V[i]
3×3 Diagonal{Int64, Vector{Int64}}:
10 ⋅ ⋅
⋅ 20 ⋅
⋅ ⋅ 30All indices appearing on the right must also appear on the left. There is no implicit sum over repeated indices on different tensors. To sum over things, you need @reduce or @matmul, described on the next page.