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()
```