Heatmap in Julia

Viewed 84

I am trying to create a heat map. I have the following code. Is there a simple way to create a heatmap, s.t. lst1 is the x-axis, lst2 is the y-axis and lst ist the intensity in the graph?

lst1 = []
lst2 = []
lst3 = []
for i in range(0,3.5,step = 0.5)
    for j in range(0,4,step = 0.5)
        println(i,j)
        a = f(parameter,i,j)
        push!(lst1,i)
        push!(lst2,j)
        push!(lst3,a)
        print("($i , $j): $a %")
    end
end
2 Answers

A somewhat shorter way of doing this, utilizing broadcasts, might be:

lst1 = 0:0.5:3.5
lst2 = 0:0.5:4

lst3 = f.(Ref(parameter), lst1, lst2')

lst1 and lst2 are constructed using the colon operator but are equivalent to the range call you showed.

lst3 is constructed using the Julia broadcast operator. Here, we wrap parameter in a Ref (think of it as a zero dimensional array or a pointer) to indicate that it should not be expanded during the broadcast. We pass lst1 as is, and its form mimics a column vector. We then pass the transpose of lst2 (obtained by the ' operator) which makes it have the form of a row vector.

These two differing dimensions cause the broadcast to create a Matrix with the first axis being lst1 and the second axis being lst2. For clarity you can examine the output of tuple.(lst1, lst2'), which will show you essentially the values passed into the function f.

In the end, lst3 will be a Matrix of elements of the type which f returns.

To actually plot this you should consider using the Plots.jl or Makie.jl packages.

Sorry I can't provide code examples right now, posting this on mobile. Will reformat later.

You may do it like this with Plots.jl

using Plots

f = (x, y) -> x^2 + y^2

x = range(-1, 1; length=10)
y = range(-2, 2; length=40)
z = fill(NaN, size(y, 1), size(x, 1))

for i in eachindex(x), j in eachindex(y)
    z[j, i] = f(x[i], y[j])
end
heatmap(x, y, z)

enter image description here

Plots.heatmap waits for rectangular matrix of z-values, which I preallocated before the cycle. In this matrix, x-direction corresponds to the columns and y-direction to the rows.

z as matrix      z when plotted

z11 z12 z13      z31 z32 z33
z21 z22 z23      z21 z22 z23
z31 z32 z33      z11 z12 z13

Take a look at

z = [1 2 3; 4 5 6; 7 8 9]
heatmap(z)  # heatmap(1:size(z, 2), 1:size(z, 1), z)

enter image description here

Things to improve

  • replace fill with uninitialised matrix constructor
  • construct x-y pairs, e.g. by Iterators.product instead of nested loops
  • use broadcasting
Related