Plotly heatmaps in subplots: how to set aspect ratio?

Question

There is a similar question, but it only asks this for a regular figure. The respective answers do not work with subplots:

https://stackoverflow.com/a/55239488/2215205
https://stackoverflow.com/a/71579181/2215205

1. px.imshow(df, aspect='equal') is ignored entirely with fig.add_trace().
2. Manipulating the figure layout in that respect is only applied to the first subplot, not all of them:

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go

## Example data
df = pd.DataFrame([[0,0,0,0,0],
                   [0,0,1,0,0],
                   [0,1,2,1,0],
                   [0,0,1,0,0],
                   [0,0,0,0,0],
                  ])

## Subplots w/ go.Heatmap()
layout = dict(
    # yaxis=...,
    # xaxis=...,
    # width=plot_width,
    # height=plot_height,
    # autosize=False,
    yaxis_scaleanchor="x" ,  # <---- only works in 1st subplot :(
    )

fig = go.Figure(layout=layout).set_subplots(rows=1, cols=2, subplot_titles=['#1', '#2'])

fig.add_trace(
    go.Heatmap(
        z = df,
        coloraxis = "coloraxis",  # use a common coloraxis for all subplots
    ), 
    1,1)
fig.add_trace(
    go.Heatmap(
        z = df*2,
        coloraxis = "coloraxis",  # use a common coloraxis for all subplots
    ),
    1,2)

## Update y-axes
fig.update_yaxes(
    title_text="yaxis #1", 
    row=1, col=1, 
    scaleanchor='x',  # <---- only works in 1st subplot :(  ==> works in Colab !!
    ) 
fig.update_yaxes(
    title_text="yaxis #2",
    row=1, col=2,
    scaleanchor='x',  # <---- only works in 1st subplot :( ==> works in Colab !!
    )

## Update layout
# fig['layout']['yaxis']['scaleanchor']='x' # <--- updates only first subplot :(

fig.update_layout(
    title='2D Heatmaps',
    autosize=False,
    # yaxis_scaleanchor="x",  # <--- updates only first subplot :(
    yaxis = dict(scaleanchor = 'x'),   # <--- updates only first subplot :(
    # coloraxis_colorscale='Viridis',
    )

# fig.show()

Any help is deeply appreciated.

Please find a Jupyter-Notebook here with my efforts: https://colab.research.google.com/drive/13NKNmgAXudXp62UlmXKDhVFgpFe1YeQN?usp=sharing

Answer 1

Note that each subplot has its own x axis / y axis (subplot1: xaxis/yaxis, subplot2: xaxis2/yaxis2, etc.) and those axes are defined in the layout. The second x axis, xaxis2, can be referenced with the string 'x2', same for the 2nd y axis : 'y2' refers to yaxis2, and so on.

fig.update_layout(
    title='2D Heatmaps',
    autosize=False,
    yaxis_scaleanchor="x",      # <- updates 1st subplot :)
    yaxis2_scaleanchor="x2"     # <- updates 2nd subplot :)
)