Is it possible to manually set the rgb colors of a heatmap?

In Matplotlib, the set_facecolors on a QuadMesh (created via pcolormesh) allows to send an array of rgb(a) values to directly change the colors of the mesh.

Is it possible to do the same with Plotly’s Heatmap?
I can only find ways to create custom colormaps, or set the z values, but no way to directly set the rgb values of the pixels.

In my case, I would like to use this to represent masked values of a numpy array as grayscale, and unmasked using the viridis colormap, on a single Heatmap.

Example:

import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm

a = np.random.random((10, 15))
mask = a < 0.5
b = np.ma.masked_where(mask, a)

cmap = cm.get_cmap('viridis')
cmap_masked = cm.get_cmap('gray')

colors = cmap(a)
colors_masked = cmap_masked(a)
colors[mask] = colors_masked[mask]

fig, ax = plt.subplots(1, 2, figsize=(8, 3))

m1 = ax[0].pcolormesh(b)
m2 = ax[1].pcolormesh(a)
m2.set_array(None)  # necessary for set_facecolors to work
m2.set_facecolors(colors.reshape(150, 4))

ax[0].set_title('Simple numpy masked array')
ax[1].set_title('With masked values as grayscale')

I know I could achieve something similar by using two Heatmaps on top of each other, but I have a mechanism where a common colormapper is used to send colors to matplotlib meshes and also point clouds with Pythreejs.
We would like to support more than one backend for the 2d plots, and it would be nice if all could use the same mechanism.

I also know that it is possible to set the colors if I use Image instead of Heatmap, but we need to use the latter because not all pixels have the same size.

Many thanks for any help!

Maybe this helps:

Update:

I can achieve something close to what I need using a transparent Heatmap and a background layout image.
Here is the code:


import numpy as np
from matplotlib import cm
import plotly.graph_objects as go
from PIL import Image

N = 15
M = 10

x = np.arange(N + 1)
y = np.arange(M + 1)
a = np.random.random((10, 15))
mask = a < 0.5
b = np.ma.masked_where(mask, a)

cmap = cm.get_cmap('viridis')
cmap_masked = cm.get_cmap('gray')

colors = cmap(a)
colors_masked = cmap_masked(a)
colors[mask] = colors_masked[mask]

img = Image.fromarray(np.flipud(np.uint8(colors*255)))

fig = go.FigureWidget(layout={
    'width': 600,
    'height': 400,
    'margin': {
        'l': 0,
        'r': 0,
        't': 0,
        'b': 0
    }
})

# Add invisible heatmap trace.
# This trace is added to help the autoresize logic work and adds a colorbar next to our image
fig.add_trace(
    go.Heatmap(x=x, y=y, z=a, opacity=0, colorscale='viridis')
)

height = fig.layout.height - fig.layout.margin.b - fig.layout.margin.t
width = fig.layout.width - fig.layout.margin.l - fig.layout.margin.r

# Add background image
fig.update_layout(
    images=[go.layout.Image(
        x=0,
        sizex=15,
        y=10,
        sizey=10,
        xref="x",
        yref="y",
        opacity=1.0,
        layer="above",
        sizing="stretch",
        source=img.resize((width, height), 0))]
)

fig

which yields basically exactly what I was looking for:
Screenshot at 2022-12-20 17-16-35

Two caveats:

  1. This does not work for cases where the pixels do not all have the same size
  2. The background image gets messed up when I change one of the axes to a logarithmic scale.
    In the images below, I changed the x axis to a log scale.
    The version with the background image (Left) has the pixels with incorrect sizes (as it they are still on the linear scale?).
    A version with just a heatmap (Right) behaves correctly with log scale.

I was not successful in looking for a parameter on the background image that would fix this.