How to create annotated heatmaps in subplots?

@jvd23

Here is the big surprise related to updating annotations, introduced, probably, in Plotly 4.0.0:
If the initial annotations in a fig.layout consist in list, annots, of length L, then when we are updating it
with a list, newannots, of length n>=L,
the first L annotations from newannots update successively the elements in annots, and the following n-L are new annotations that extend annots.

Example:


layout = go.Layout(
            annotations=[
                go.layout.Annotation(text="Initial 1"),
                go.layout.Annotation(text="Initial 2"),
            ]
        )
print(f'Initial annotations:\n\n {layout}')
layout.update(
            annotations=[
                go.layout.Annotation(width=10),
                go.layout.Annotation(width=20),
                go.layout.Annotation(width=30),
                go.layout.Annotation(width=40),
                go.layout.Annotation(width=50),
            ]
        )

print(f'updated annotations:\n\n{layout}')

Hence a more performant working code for the example with two subplots is like this:

import numpy as np
import plotly.graph_objs as go
import plotly.figure_factory as ff
from plotly.subplots import make_subplots
import string
#Define data for heatmap
N=5
x = np.array([10*k for k in range(N)])
y = np.linspace(0, 2, N) 
z1 = np.random.randint(5,15, (N,N))
z2 = np.random.randint(10,27, (N,N))
mytext = np.array(list(string.ascii_uppercase))[:25].reshape(N,N)


fig1 = ff.create_annotated_heatmap(z1, x.tolist(), y.tolist(),  colorscale='matter')
fig2 = ff.create_annotated_heatmap(z2, x.tolist(), y.tolist(), annotation_text=mytext, colorscale='Viridis')

fig = make_subplots(subplot_titles=('A', 'B'),
    rows=1, cols=2,
    horizontal_spacing=0.075,)

print(f'Annotations defined by make-subplots:\n\n{fig.layout.annotations}\n\n')


#update font size in subplot titles; This is the tricks that just overwrites  the font size,
# with the same value, and leads to the right final annotations:
newfont= [go.layout.Annotation(font_size=16)]*2
print(f'A superflous, but necessary update :\n\n{newfont}')

fig.add_trace(fig1.data[0], 1, 1)
fig.add_trace(fig2.data[0], 1, 2)

annot1 = list(fig1.layout.annotations)
annot2 = list(fig2.layout.annotations)
for k  in range(len(annot2)):
    annot2[k]['xref'] = 'x2'
    annot2[k]['yref'] = 'y2'

annot1.extend(annot2)

newfont.extend(annot1)
fig.update_layout(annotations=newfont)  
fig.update_layout(width=700, height=400)  

But if you define n subplots, with n >=3, it is cumbersome to extend succesively, annot1 = fig1.layout.annotations,
with fig.layout.annotations from the heatmap 2, 3, 4, ā€¦ definition.

To avoid lines of code like these ones:

annot1.extend(annot2)
annot1.extend(annot3)
.
.
.

we define a recursive function to extend an empty list with the sublists from a lists of lists:

def recursive_extend (mylist, nr):
    #mylist is a list of lists
    # initial nr =len(mylist)
    result = []
    
    if nr> 1:
        result.extend(mylist[nr-1])
        result.extend(recursive_extend( mylist, nr-1))
    else: 
        result.extend(mylist[nr-1])
    
    return result

Let us illustrate how is defined a figure consisting in 4 subplots of annotated heatmaps:

N=5
x = np.array([10*k for k in range(N)])
y = np.linspace(0, 2, N) 
z1 = np.random.randint(5,15, (N,N))
z2 = np.random.randint(10,27, (N,N))
mytext = np.array(list(string.ascii_uppercase))[:25].reshape(N,N)


fig1 = ff.create_annotated_heatmap(z1, x.tolist(), y.tolist(),  colorscale='matter')
fig2 = ff.create_annotated_heatmap(z2, x.tolist(), y.tolist(), annotation_text=mytext, colorscale='Viridis')
fig3 = ff.create_annotated_heatmap(z1, x.tolist(), y.tolist(),  colorscale='deep')
fig4 = ff.create_annotated_heatmap(z2, x.tolist(), y.tolist(), annotation_text=mytext, colorscale='Plasma')

fig = make_subplots(subplot_titles=('A', 'B', 'C', 'D'),
    rows=2, cols=2,
    horizontal_spacing=0.075, vertical_spacing=0.11)


fig.add_trace(fig1.data[0], 1, 1)
fig.add_trace(fig2.data[0], 1, 2)
fig.add_trace(fig3.data[0], 2, 1)
fig.add_trace(fig4.data[0], 2, 2)


newfont = [go.layout.Annotation(font_size=16)]*len(fig.layout.annotations)
#print(f'A superflous update that forces the right updates of the annotations:\n\n{newfont}')
fig_annots = [newfont, fig1.layout.annotations, fig2.layout.annotations, fig3.layout.annotations, fig4.layout.annotations]

for j in range(2, len(fig_annots)):
    for k  in range(len(fig_annots[j])):
        fig_annots[j][k]['xref'] = f'x{j}'
        fig_annots[j][k]['yref'] = f'y{j}'
    
new_annotations = recursive_extend (fig_annots[::-1], len(fig_annots))# Note that fig_annots is reverted to ensure that
                                                                      # the elements of newfonts are the first 4 annotations  
fig.update_layout(annotations=new_annotations)

Note: When we are concatenating a relative big number of lists it is recomended to use extend, not +, because + is less performant.