How can I make a density plot with log-scaled axes in matplotlib?

Question

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. I also vaguely remember that the example shown in the first column used to work a few years back, but apparently this is not the case anymore. Any help with solving this 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).

Answer 1

(Updated answer)

If I understand correctly, currently imshow with log scale axes stretches a linear image. NonUniformImage doesn't seem to do such transformation. So, maybe the following works for you:

Using NonUniformImage with linear

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_title("imshow")

# use NonUniformImage to plot array
im = NonUniformImage(ax[1], extent=extent, interpolation='nearest')
# calculate linear "cell centers" for the given extent
x_lin = np.linspace(*extent[:2], 11)[1::2]
im.set_data(x_lin, 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()

(Original answer)

Using imshow with linear axes and adding a dummy log scale ticks

Here is a solution using imshow with the given extent, without setting log scale. Create a second x-axis via ax.twiny() with log scale, and only show the x ticks of that new x-axis.

import matplotlib.pyplot as plt
import numpy as np

# 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(figsize=(6, 3))

# use imshow to plot the array
ax.imshow(z.T, origin="lower", aspect="auto", extent=extent)
ax.set_title("imshow with twiny")

# create a twin ax
ax2 = ax.twiny()
# set the twin ax to log
ax2.set_xscale("log")
# copy the limits of the original ax
ax2.set_xlim(ax.get_xlim())
# by default, twiny()'s ticks are put at the top, place them at the bottom
ax2.tick_params(axis='x', which='both', labelbottom=True, bottom=True, labeltop=False, top=False)
# hide the original ticks
ax.tick_params(axis='x', which='both', labelbottom=False, bottom=False)
plt.show()