I want to plot a scalar density as a function of two variables x and y, which can potentially be scaled logarithmically. I essentially run simulations for each pair of x and y and want to report the data using a nice colormap. However, I run into the problem that I cannot make imshow scale the data correctly. While pcolormesh works reliably, it produces files that are orders of magnitudes larger and often cannot be displayed without artifacts (like thin white lines between data points).

Here’s some code to reproduce the problem:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.image import NonUniformImage

# calculate axis positions
x = np.geomspace(1, 100, 5)
dx = np.sqrt(x[1] / x[0])  # half the difference between points in logspace
y = np.linspace(0, 1, 3)
dy = (y[1] - y[0]) / 2  # half the difference between points in linspace
extent = (x[0] / dx, x[-1] * dx, y[0] - dy, y[-1] + dy)

# get some random image data to plot
z = np.random.uniform(size=(len(x), len(y)))
# create figure axes
fig, ax = plt.subplots(ncols=3, figsize=(12, 3))

# use imshow to plot array
ax[0].imshow(z.T, origin="lower", aspect="auto", extent=extent)
ax[0].set_xscale("log")
ax[0].set_title("imshow")

# use NonUniformImage to plot array
im = NonUniformImage(ax[1], extent=extent)
im.set_data(x, y, z.T)
ax[1].add_image(im)
ax[1].set_xscale("log")
ax[1].set_title("NonUniformImage")

# use pcolormesh to plot array
x2 = np.geomspace(*extent[:2], 6)
y2 = np.linspace(*extent[2:], 4)
ax[2].pcolormesh(x2, y2, z.T)
ax[2].set_title("pcolormesh")

# set axis scales
for i in range(3):
    ax[i].set_xlim(*extent[:2])
    ax[i].set_ylim(*extent[2:])
    ax[i].set_xscale("log")

plt.show()

Running this example results in the following picture

Clearly, imshow is distorting the image, presumably because it assumes that the image contains data on a linearly scaled axis.
The second panel shows my attempt at using NonUniformImage, which gets things completely wrong for some reason.
The third panel shows what I want to see, albeit with using pcolormesh, which has the severe drawbacks I mentioned above.

Essentially, I just want to show a “normal” image with rectangular pixels of equal size on a log-scaled axis. I think this should be possible, but I was not able to achieve this. Any help would be much appreciated!

Note that this older answer does not work properly since it simply adds an axes with logarithmic ticks, so the user cannot reliably interact with the result (e.g., to change the ticks afterwards).