在 plotly 中创建具有固定单元格宽度的分组热图

Question

I would like to plot several heatmaps on one plotly figure, where the columns of all heatmaps should have the same width (within and across heatmaps). My first idea was to use make_subplots together with go.Heatmap but with this approach the width of each cell changes accordingly to how many columns each subplot has. I need it the other way around. I would like to set a fixed cell width that needs to be automatically determined using the entire width of the whole figure and the overall number of columns. Another approach would be to use px.imshow where everything is plotted as one big heatmap, and so all columns have the same width. But with this approach, I am not able to plot the titles for each group of columns. This thread is related to this one but I need a solution in Python not R.

Code to create example data:

import numpy as np
from numpy.random import default_rng
import pandas as pd
from plotly.subplots import make_subplots
import plotly.graph_objects as go

rng = default_rng(42)

# create example dataframe
n_variates = 2
variables = ['apple','pear', 'orange','mandarin','grapefruit']
sources = ['pome','pome','citrus','citrus','citrus']
variates = np.repeat(list(range(n_variates)),len(variables))
values = rng.random(n_variates * len(variables))
df = pd.DataFrame({'variable':np.tile(variables,n_variates),
                   'source':np.tile(sources,n_variates),
                   'variate':variates,
                   'value':values})

First approach: Using make_subplots and go.Heatmap (Can create titles but cells have different widths):

# use subplots and go.Heatmap
z_list = []
source_names = []
for source_name,source in df.groupby('source'):
    z = source.pivot(columns=['source','variable'],index='variate',values='value')
    z_list.append(z)
    source_names.append(source_name)

fig = make_subplots(rows=1,cols=len(z_list),
                    subplot_titles=source_names,
                    horizontal_spacing=0,
                    shared_yaxes=True)
for col,z in enumerate(z_list):
    fig.add_trace(go.Heatmap(z=z,x=z.columns.get_level_values(1),coloraxis = "coloraxis"),1,col+1)
    
fig.update_yaxes(dtick=1)
fig.show()

Second approach: Using px.imshow (Cells have the same width but group titles are missing):

# use px.imshow()
import plotly.express as px
z = df.pivot(columns=['source','variable'],index='variate',values='value')
source_labels,variable_labels = z.columns.get_level_values(0),z.columns.get_level_values(1)
variate_labels = z.index.tolist()
fig = px.imshow(z.values,x=variable_labels,y=variate_labels)
fig.update_xaxes(type='category')
fig.update_yaxes(type='category')
fig.show()

Answer 1

After much trial and error on my part, I think it would be easiest to add a title to the second example with the annotation feature.

import plotly.express as px

z = df.pivot(columns=['source','variable'], index='variate', values='value')
source_labels,variable_labels = z.columns.get_level_values(0), z.columns.get_level_values(1)
variate_labels = z.index.tolist()
fig = px.imshow(z.values, x=variable_labels, y=variate_labels, text_auto='.2f')

fig.add_annotation(x=0.3, y=1.1, xref='paper', yref='paper', text=df.source.unique()[0], showarrow=False)
fig.add_annotation(x=0.62, y=1.1, xref='paper', yref='paper', text=df.source.unique()[1], showarrow=False)

fig.update_xaxes(type='category')
fig.update_yaxes(type='category')
fig.show()

If you use the x-axis with the category variable criteria as it is, you get the following But the title is not the center of the group.

fig.add_annotation(x=variable_labels[0], y=1.1, xref='x', yref='paper', text=df.source.unique()[0], showarrow=False)
fig.add_annotation(x=variable_labels[3], y=1.1, xref='x', yref='paper', text=df.source.unique()[1], showarrow=False)