Apply function to arrays of arbitrary dimensions

New to Julia
,
1

I want to write a function which will do some calculations with vectors and rotation matrices in 3D space (so 3 element, single dimension vectors and 3x3 element matrices of two dimensions). What is the best way to apply this function to collections of vectors/matrices? I would ideally be able to apply such a function to a single vector and matrix (3 and 3x3 elements) as well as collections of higher dimension, like 3xN vectors with 3x3xN matrices or 3 vectors with 3x3xNxM matrices or 3xN vectors with 3x3xNxM matrices, or even a set of 3xN vectors, 3xM vectors and 3x3xNxM matrices. So really to inputs of arbitrary dimensions.

I can think of several approaches:

  • Just write loops wherever the function is used and call it on single vectors/matrices. This is what I do in the below example.
  • Create custom types for vector arrays of size 3 and matrices of size 3x3 and then use broadcasting.
    • I’m not sure how broadcasting would work for some mixed dimensions like the last part of the example, can Julia broadcast “along” a single dimension?
  • In fact, if such broadcasting “along” a single dimension is possible, than perhaps I can just use broadcasting without custom types.
    • Perhaps this can work but only for arrays of a certain layout of dimensions, I didn’t find anything in the documentation.
  • Use multiple dispatch for each set of dimensions I use, where each function is just a loop which calls the base function.
    • This seems unwieldy, especially once there are more than two inputs.

Probably there are multiple correct answers, some factors I think are important are performance, maintainability, following Julia convention, portability (can I still use these functions with CUDA.jl for instance).

I am new to Julia and coming from a Python/C background, so I apologize if there are some obvious things I’m missing. And perhaps I am overthinking things and the loops that I already have are best!

Thanks for any insight!

Here’s an example (I’ve used 2D vectors/matrices just to be simpler):

function rotate_x(angle)
    rotation = [
        [cos(angle) sin(angle)]
        [-sin(angle) cos(angle)]
    ]
end


function rotate_offset(rotation, offset, vector)
    rotation * vector + offset
end


function main()
    # Rotate and offset a single vector
    # Make the rotation matrix
    rotation = rotate_x(pi/6)
    # Make the input vector
    vector = [0.0, 1.0]
    # Make the offset vector
    offset = [0.3, 0.0]
    # Compute the result
    result = rotate_offset(rotation, offset, vector)
    println(rotation)
    println(vector)
    println(offset)
    println(result)
    println()

    # Rotate and offset an single dimension of vectors
    # Make the 99 rotation matrices
    rotation = zeros(2, 2, 99)
    for i in 1:size(rotation, 3)
        rotation[:, :, i] = rotate_x(pi/6)
    end
    # Make the 99 input vectors
    vector = zeros(2, 99)
    vector[2, :] .= 1.0
    # Make the 99 offset vectors
    offset = zeros(2, 99)
    kffset[1, :] .= 0.3
    # Compute the result
    result = zeros(2, 99)
    for i in 1:size(vector, 2)
        result[:, i] = rotate_offset(
            rotation[:, :, i], offset[:, i], vector[:, i])
    end
    println(rotation[:, :, 44])
    println(vector[:, 44])
    println(offset[:, 44])
    println(result[:, 44])
    println()

    # Apply a rotation and offset along some more arbitrary dimensions
    # Make the 99x55 shaped rotation matrices
    rotation = zeros(2, 2, 99, 55)
    for j in 1:size(rotation, 4)
        for i in 1:size(rotation, 3)
            rotation[:, :, i, j] = rotate_x(pi/6)
        end
    end
    # Make the 99 input vectors
    vector = zeros(2, 99)
    vector[2, :] .= 1.0
    # Make the 55 offset vectors
    offset = zeros(2, 55)
    offset[1, :] .= 0.3
    # Compute the result
    result = zeros(2, 99, 55)
    for j in 1:size(rotation, 4)
        for i in 1:size(rotation, 3)
            result[:, i, j] = rotate_offset(
                rotation[:, :, i, j], offset[:, j], vector[:, i])
        end
    end
    println(rotation[:, :, 44, 33])
    println(vector[:, 44])
    println(offset[:, 33])
    println(result[:, 44, 33])
    println()
end

main()
2

I think you could benefit from using StaticArrays for your vectors / matrices. The advantage of StaticArrays is that you get your vectors / matrices as single objects, rather than making tensors where your first or first and second dimensions mean vector / matrix. The memory layout is effectively the same and you can even switch back and forth using reinterpret. Here I’ll show you an example with the GeometryBasics package just because it uses geometry primitives based on StaticArrays and I had it available. Here I just do some random computations on differently sized collections of vectors / matrices (pretend they’re rotation matrices):

using GeometryBasics
using LinearAlgebra

vectors = rand(Vec3{Float64}, 100)
matrices = rand(Mat3{Float64}, 100, 3)
vectors_2 = rand(Vec3{Float64}, 100, 4, 5)

rotated_vectors = matrices .* vectors
cross_products = cross.(vectors, vectors_2)
5 Likes
3

Adding to the previous answer, StaticArrays is indeed what you probably want. For me it was a revelation when I was explained this:

julia> using BenchmarkTools

julia> using StaticArrays

julia> x = [ SVector{3,Float64}(1,1,1) for i in 1:10 ];

julia> function f!(x)
         for i in 1:length(x)
           x[i] = 2*x[i] + 1
         end
         nothing
       end
f! (generic function with 1 method)

julia> @btime f!(x)
  363.464 ns (0 allocations: 0 bytes)

julia>

The vectors are three dimensional, but you can operate on collections of them as if they were numbers, without allocating anything. This makes the code so much cleaner when working with this kind of data. And in addition that turns out to be very performant (and with a structure adaptable to loop vectorization).

5 Likes
4

Thanks both for the responses! This seems like a good solution, like I was thinking in my second suggestion but without reinventing the wheel, and with some performance benefits as well! I might want to use such code in CUDA in the future, but I think that is a separate discussion.

I do still wonder if there are some advanced broadcasting techniques for arrays of partially matching dimensions, like the third part of my example above. Is this something that exists and people actually use? I realize as well that this type of operation may be cleaner and more understandable just using loops, which is partly why Julia is great!

P.S. Can I mark both answers as a solution? Does it make a difference?