How to properly implement Heatmap efficiently?

I am currently trying to visualise a solution to a set of differential equations, where u is my solution, and though my code works the running of the heatmap to visualise it takes at least 200 seconds each time. It should be noted here that u is a vector of length 2^10 and m is 200,000. I understand that this is a large problem but I just don’t see why copying values across and plotting is taking so long.

begin

U = zeros(m,n)

for i = 1:m
    U[i, :] = u[:,i]
end

Plots.heatmap(U)

end

This code is not runnable but I am hoping that someone can still explain if I am doing something inefficient here. If it needs to be run, then I can copy all the code as unfortunately there is no other way to make it self contained.

Thank you for the help.

I should note that I am using Pluto

Four things:

  • Every u[:, i] is a slice, which creates an allocated vector. You can try using view(u, :, i).
  • Writing to U row-wise is inefficient, perhaps you can reorganize U to receive data columnwise.
  • Make sure your code is in a function, and that this is not happening in global scope.
  • Finally, are you just trying to create a transpose? Can you just write Plots.heatmap(u') instead?
1 Like

I agree with everything DNF said, but I would like to also point out that you are (if I am not mistaken) trying to plot a heatmap of size 200\,000 \times 1024. That’s a lot of points, and I am afraid that most of the time is spent actually plotting (and not creating U).

You might consider switching from plotting a heatmap to plotting an image. Makie.jl has a function for that (image). I am not sure how to do it in Plots, but it seems that you’d have to convert U to a matrix of RGB/Gray (a type from ColorTypes.jl) and just use plot.

3 Likes

You might also want to try outside pluto.

Thank you for the help, I ended up using the simple transpose (don’t know how I missed that) and also reduced the number of points by a factor of 10. This seems to have done the job, although the time taken still varies sometimes.

I will also try the image method.

What do you mean by in a function?

If you run this code

for i = 1:m
    U[i, :] = u[:,i]
end

in global scope, and not inside a function, it will be slow. If it isn’t already in a function, you should write

function mytranspose!(U, u)
    for i = 1:m
        U[i, :] = u[:,i]  # better to use view here
    end
    return U
end

and then call the function in your script.

I strongly recommend reading Performance Tips · The Julia Language
Any question you have on Julia performance should be addressed by first consulting that page, before looking elsewhere.

I tried the following out of curiosity:

julia> using Plots

julia> const U = rand(1000, 1000);

julia> @time plot(heatmap(U)); # Second run
  0.015866 seconds (897 allocations: 15.411 MiB)

julia> @time plot(Gray.(U)); # Second run
  0.009538 seconds (456 allocations: 7.696 MiB)

You can see that plot is significantly faster than heatmap already for this size of U. I suspect that the difference will be even more substantial for bigger sizes.

Note that the results are different though:

heatmap

plot

So you’d have to add a colorbar etc. Moreover, plot rasterizes the image (which allows for faster plotting), while heatmap produces vector graphics.

Some minor suggestions (that should be covered by the performance tips):

# You are overwriting everything, no reason to first place zeros
#U = zeros(m,n)
U = Matrix{eltype(u)}(undef, m, n) 

# By construction, i should be in-bounds. 
# Also, you may wish to iterate over an axes
#for i = 1:m
@inbounds for i in axes(U, 1)
    # Avoid allocating an intermediate array when slicing
    # Also, use .= for assigning values
    U[i, :] .= view(u, i)
    #U[i, :] = u[:,i]
end