Heatmap with irregularly spaced grids

I’m trying to create a figure with two heatmaps where

  1. the x and y variables are irregularly spaced,
  2. the scale of the colors are the same in both heatmaps, and
  3. there is only one colorbar.

In Python, I would do

import numpy as np
import matplotlib.pyplot as plt

# Generate non-linear grid
n = 20
x = np.zeros(n)

x[0] = 0

for i in range(1, n):
    x[i] = x[i - 1] + (1 - x[i - 1]) / (n - i) ** 1.25

# Generate z variable
z = np.zeros((2, n, n))
for i in range(n):
    for j in range(n):
        r = x[i]**2 + x[j]**2
        z[0, i, j] = np.sin(10 * r) / (1 + r)

z[1] = z[0]

# Plot z variables against non-linear grids
x1, x2 = np.meshgrid(x, x, indexing='ij')
plt.subplot(211)
plt.pcolormesh(x1, x2, z[0], shading='auto')
plt.subplot(212)
plt.pcolormesh(x1, x2, z[1], shading='auto')

# Insert one colormap
plt.subplots_adjust(bottom=0.1, right=0.8, top=0.9)
cax = plt.axes([0.85, 0.1, 0.075, 0.8])
plt.colorbar(cax=cax)

# Synchronize scale of colormap
zmin = np.minimum(z.min(), z2.min())
zmax = np.maximum(z.max(), z2.max())
for im in plt.gca().get_images():
    im.set_clim(vmin=zmin, vmax=zmax)

plt.show()

with result
Figure_1

But in Julia I can’t figure out what to do. What I have thus far is

using Plots

# Generate non-linear grid
n = 20
x = zeros(n)
x[1] = 0
for i in 2:n
    x[i] = x[i - 1] + (1 - x[i - 1]) / (n - i + 1)^1.25
end

# Generate z variables
z = zeros(n, n)
for i in 1:n
    for j in 1:n
        r = x[i]^2 + x[j]^2
        z[i, j] = sin(10*r) / (1 + r)
    end
end
z2 = 0.5 .* z

# Plot z variables against the INDICES of the non-linear grids
plot(heatmap(z), heatmap(z2), layout=(2, 1))

Hi there! Check this out :slight_smile:

2 Likes

Great! That takes care of item 1. As for item 3, this probably fixes it. That leaves item 2 :nerd_face:

I would do it all in Makie. Note that the discourse thread you linked to uses Plots, which is a different plotting package. If you choose to use Makie, look at the color bar docs, you have a lot of control over placement, scale etc.

Makie has a slight learning curve, but it is well worth it :slight_smile:

2 Likes

lmk if you want me to implement it in Makie!

1 Like

+1 for Makie. There was a recent post here regarding your second point, including some discussion about making the API for it even simpler: One colorbar for multiple axes - #6 by ederag

1 Like

One way of achieving this using Plots.jl is shown here and adapted below.

# 1 - INPUT DATA: generate non-linear grid and z variables
n = 20
x = zeros(n)
for i in 2:n
    x[i] = x[i-1] + (1 - x[i-1]) / (n-i+1)^1.25
end
r = @. x^2 + x'^2
z = @. sin(10*r) / (1 + r)
z2 =  0.5z

# 2 - PLOT DATA: multiple heatmaps with single colorbar
using Plots; gr()
clims = extrema([z; z2])
p1 = heatmap(x, x, z, ratio=1, clims=clims, cb=false, lims=extrema(x), tick_dir=:out)
p2 = heatmap(x, x, z2, ratio=1, clims=clims, cb=false, lims=extrema(x), tick_dir=:out)
p3 = scatter([0], [0], clims=clims, zcolor=clims, grid=false,
    xlims=(1,1.1), framestyle=:none, label="", cb_title="legend")
l = @layout[grid(1,2) a{0.05w}]
plot(p1, p2, p3, layout=l, size=(1000,400))

4 Likes

Thanks for the offer! If you could, it would be great!

That was great! With a few minor adjustments it results in almost exactly what I’m looking for :slight_smile: There is only the minor detail that the tick labels on the color bar protrude into the image’s right side… :thinking:

# 2 - PLOT DATA: multiple heatmaps with single colorbar
using Plots;  # No gr()
clims = extrema([z; z2])
p1 = heatmap(x, x, z, clims=clims, cb=false, lims=extrema(x), tick_dir=:out, c=:viridis)  # Add c=:viridis and remove ratio=1
p2 = heatmap(x, x, z2, clims=clims, cb=false, lims=extrema(x), tick_dir=:out, c=:viridis)  # Add c=:viridis and remove ratio=1
p3 = scatter([0], [0], clims=clims, zcolor=clims, grid=false,
    xlims=(1,1.1), framestyle=:none, label="", c=:viridis)  # Add c=:viridis and remove cb_title="legend"
l = @layout[grid(2,1) a{0.05w}]  # grid(1,2) --> grid(2,1)
plot(p1, p2, p3, layout=l)

plot_30

Use Measures - see for example here.

2 Likes