Updating Heatmaps using extendData

How can I use extendData to update a heatmap?

I have a heatmap that I’d like to update efficiently, but I can’t figure out how to use extendData to make it happen.

Here’s my minimal example code, which starts by drawing the first 3 columns a diagonal heatmap. It should add more columns to the right as the button is pressed.

I’ve gotten this working using partial property updates and Patch(), so you can see my desired behavior. However, my actual use case has a very large heatmap, so redrawing the whole thing using Patch() is undesirable. I’d rather add columns as they are requested, with maxPoints removing the leftmost columns as needed to keep the figure a reasonable size.

I’m using dash 3.0.4.

"""Simple MRE for issues using extendData and heatmaps"""

import numpy as np
import plotly.graph_objects as go
from dash import Dash, Input, Output, Patch, State, callback, dcc, html

# Sample data for the heatmap
STARTING_COLS = 3
TOTAL_COLS = 10
DIAG_ARR = np.diag(np.ones(TOTAL_COLS)*100).astype(np.int8)


def main():
    """Initialize the Dash app and layout."""
    app = Dash(__name__)

    app.layout = html.Div([
        dcc.Graph(figure=create_initial_heatmap(), id='heatmap'),
        html.Button('Extend to the Right', id='update-button')
    ])

    app.run(debug=True)


def create_initial_heatmap() -> go.Figure:
    """Create the initial heatmap figure with starting columns."""
    return go.Figure(go.Heatmap(
        z=DIAG_ARR[:, :STARTING_COLS],
        x0=0,
        dx=1,
        y0=0,
        dy=1,
        zmin=0,
        zmax=100
    ))


@callback(
    Output('heatmap', 'figure'),
    Input('update-button', 'n_clicks'),
    State('heatmap', 'figure'),
    prevent_initial_call=True
)
def update_heatmap_via_patch(n_clicks: int | None,
                             prev_fig: dict) -> Patch:
    """WORKING CALLBACK: update the heatmap data via partial property update"""

    print(f'Number of clicks: {n_clicks}.')
    print(f'Previous figure data: {prev_fig["data"]}')

    heatmap_patch = Patch()
    heatmap_patch.data[0].z = DIAG_ARR[:, :STARTING_COLS + n_clicks]
    return heatmap_patch


# @callback(
#     Output('heatmap', 'extendData'),
#     Input('update-button', 'n_clicks'),
#     State('heatmap', 'figure'),
#     prevent_initial_call=True
# )
# def update_heatmap_extend_data(n_clicks: int | None,
#                                prev_fig: dict) -> dict:
#     """NOT WORKING CALLBACK: update the heatmap data via extendData"""

#     print(f'Number of clicks: {n_clicks}.')
#     print(f'Previous figure data: {prev_fig["data"]}')
#     if n_clicks is None or n_clicks > TOTAL_COLS - STARTING_COLS - 1:
#         return {}

#     col_to_add = min(STARTING_COLS + n_clicks, TOTAL_COLS - 1)

#     data_to_add = DIAG_ARR[:, col_to_add].tolist()
#     print(f'Adding column {col_to_add}: {data_to_add}')

#     # Create new data to extend
#     new_data = {
#         'updateData': {'z': data_to_add},
#         'traceIndices': [0],
#         'maxPoints': TOTAL_COLS}

#     return new_data


if __name__ == '__main__':
    main()

This is similar to extendData for HeatMap and Real-Time Heatmap Streaming with dcc.Interval and extendData, but those threads seem abandoned.

I have a hack that works for adding rows using extendData. It’s definitely a hack, so I’d appreciate a better solution if possible.

TL; DR

In the code below, you must hit “Update via Patch” first. After that, you can “Extend Rows” and behavior works as expected. If you click “Extend Rows” first, you’ll need to refresh the page.

This prints updates into the command-line to help explain.

Advantages

Examining the Dev Tools callback graph shows:

  • The patch-button callback requires 51 ms of compute and 311 bytes of download from server to client.
  • The extend-button callback requires (on average) 1 ms of compute and 115 bytes of download from server to client

These benefits would be more pronounced when working with a larger heatmap.
Additionally, the large upload data transfer on the callbacks is due to the use of State for printing out the figure data. If this were removed, upload transfers would decrease substantially.

Code

"""Simple MRE for issues using extendData and heatmaps"""

import numpy as np
import plotly.graph_objects as go
from dash import Dash, Input, Output, Patch, State, callback, dcc, html
from dash.exceptions import PreventUpdate

# Sample data for the heatmap
STARTING_ROWS = 3
TOTAL_ROWS = 20
MAX_ROWS = 10
DIAG_ARR = np.diag(np.ones(TOTAL_ROWS)*100).astype(np.int8)


def main() -> None:
    """Initialize the Dash app and layout."""
    app = Dash(__name__)

    app.layout = html.Div([
        dcc.Graph(figure=create_initial_heatmap(), id='heatmap'),
        html.Button('Update via Patch', id='patch-button'),
        html.Button('Extend Rows', id='extend-button')
    ])

    app.run(debug=True)


def create_initial_heatmap() -> go.Figure:
    """Create the initial heatmap figure with starting columns."""
    return go.Figure(go.Heatmap(
        z=DIAG_ARR[:STARTING_ROWS, :],
        x0=0,
        dx=1,
        y0=0,
        dy=1,
        zmin=0,
        zmax=100
    ))


@callback(
    Output('heatmap', 'figure'),
    Input('patch-button', 'n_clicks'),
    State('heatmap', 'figure'),
    prevent_initial_call=True
)
def update_heatmap_via_patch(n_clicks: int | None,
                             prev_fig: dict) -> Patch:
    """Update the heatmap data using a Patch object. Reformats figure data"""

    print(f'Number of clicks: {n_clicks}.')
    print(f'Previous figure data: {prev_fig["data"]}')

    heatmap_patch = Patch()
    heatmap_patch.data[0].z = DIAG_ARR[:STARTING_ROWS, :]
    return heatmap_patch


@callback(
    Output('heatmap', 'extendData'),
    Input('extend-button', 'n_clicks'),
    State('heatmap', 'figure'),
    prevent_initial_call=True
)
def update_heatmap_extend_data(n_clicks: int | None,
                               prev_fig: dict) -> tuple[dict, list[int], int]:
    """Update heatmap via extendData. Must be reformatted using Patch first"""

    print(f'Number of clicks: {n_clicks}.')
    print(f'Previous figure data: {prev_fig["data"]}')
    if n_clicks is None or n_clicks > TOTAL_ROWS - STARTING_ROWS:
        raise PreventUpdate('No more rows to add, max reached.')

    row_to_add = min(STARTING_ROWS + n_clicks - 1, TOTAL_ROWS - 1)

    data_to_add = DIAG_ARR[row_to_add:row_to_add + 1, :].tolist()
    print(f'Adding row {row_to_add}: {data_to_add}')

    # Create new data to extend
    return {'z': [data_to_add]}, [0], MAX_ROWS


if __name__ == '__main__':
    main()

Explanation

The trick is that the initial heatmap z data is of the following form, which is challenging to update with extendData.
'z': {'dtype': 'i1', 'bdata': 'ZAAAAAAAAAAAAAAAAAAAAAAAAAAAZAAAAAAAAAAAAAAAAAAAAAAAAAAAZAAAAAAAAAAAAAAAAAAAAAAA', 'shape': '3, 20', '_inputArray': [{'0': 100, '1': 0, '2': 0, '3': 0, '4': 0, '5': 0, '6': 0, '7': 0, '8': 0, '9': 0, '10': 0, '11': 0, '12': 0, '13': 0, '14': 0, '15': 0, '16': 0, '17': 0, '18': 0, '19': 0}, {'0': 0, '1': 100, '2': 0, '3': 0, '4': 0, '5': 0, '6': 0, '7': 0, '8': 0, '9': 0, '10': 0, '11': 0, '12': 0, '13': 0, '14': 0, '15': 0, '16': 0, '17': 0, '18': 0, '19': 0}, {'0': 0, '1': 0, '2': 100, '3': 0, '4': 0, '5': 0, '6': 0, '7': 0, '8': 0, '9': 0, '10': 0, '11': 0, '12': 0, '13': 0, '14': 0, '15': 0, '16': 0, '17': 0, '18': 0, '19': 0}]}

So I first use Patch() to reformat the z data into a nested list:

'z': [[100, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 100, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 100, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]

Once Patch() has reformatted it, I can leverage extendData to add new rows, and maxPoints works properly.

Currently, though, I can only add rows this way. Adding columns is more challenging, because it would require editing the inner lists within the nested list.