Some doubts about types

For function arguments, you almost always want to be more generic. Indeed, you can leave off the type declarations completely and it is perfectly fine — Julia will type-specialize your function at compile time based on the type of arguments that you pass. But it can be nice to add type declarations to indicate the expected types, for clarity, and also to allow you to define different methods for different types.

For example, a very generic definition of your function would be something like:

function propagator_freespace_labframe(R::AbstractMatrix{<:Real}, kn::Real, AlphaBlocks::AbstractArray{<:Complex,3})
    N = size(R,2)
    T = float(promote_type(eltype(R), typeof(kn), real(eltype(AlphaBlocks))))
    A = Matrix{Complex{T}}(I, 3*N, 3*N)
    I3 = Matrix{Complex{T}}(I, 3, 3)

    rk_to_rj = Vector{T}(undef,3)
    rjkhat = Vector{T}(undef,3)
    rjkrjk = Array{T}(undef,(3, 3))
    Ajk = Array{Complex{T}}(undef,(3, 3))

This way, you accept any type of argument, and the type of the temporary arrays uses a floating-point type T that is chosen based on the types that are passed as arguments.

On a separate note, note that you are allocating tons of copies of your arrays with slicing expressions like AlphaBlocks[:,:,jj], and your vectorized expressions allocate zillions of temporary arrays. You need to be cautious about writing Matlab-style code in Julia, which does not take full advantage of the language. See also the More Dots: Syntactic Loop Fusion in Julia blog post.

In this case, it looks like your vectorized operations in your inner loops are operating on 3-component coordinate vectors. For this application, it will be much faster and more elegant to use StaticArrays, in which the 3 dimensions are known at compile time and all of the 3-element loops are inlined.

That is, I’m guessing that AlphaBlocks should be a Vector{SMatrix{3,3,Complex{T}}}, R should be a Vector{SVector{3,T}}, A should be a Matrix{SMatrix{3,3,Complex{T}}}, and so on. Then you would do things like rk_to_rj = R[jj] - R[kk] and A[jj,kk] = Ajk * AlphaBlocks[kk].

Note also that explicitly forming I3 is a waste — the whole point of the I object is that things like rjkrjk - I are computed efficiently by only subtracting the diagonal elements, with I implicitly “expanding” to a matrix of the correct size.

8 Likes