Selectedpoint highlights data points for each category instead of from the dataset as a whole

Dash Python
2

UPDATE:
I’ve worked out a way to do it, but I’d be interested to know if there’s a better (maybe built-in) way.

Solution: update the figure dictionary directly by mapping the group indices to the absolute indices as per the following updated callback. Note that in this solution, the DataFrame must be sorted by it’s grouping: i.e. all Setosa together, then all Versicolour, then, all Virginica (although ordering of the groups shouldn’t matter).

@app.callback(
    Output('iris-scatter', 'figure'),
    Input('iris-scatter', 'figure'),
    Input('iris-dropdown', 'value')
    )
def update_plot(fig_old, selection):
    """
    Take in the point(s) to be selected from the dropdown and highlight them on the graph
    If no points are selected, then highlight all points
      ARGS:
            fig_old -- the current figure (dict) to be updated
            selection -- the absolute indices of the points to be highlighted

    RETURN:
            the fig with updated selection
    
    """

    # select all points if selection not specified/empty
    if (selection is None) or (len(selection) == 0):
        points = [i for i in range(len(df))]        # all points
        fig = go.Figure(fig_old)
        fig.update_traces(selectedpoints=points)

    # select *only* the indented points
    else:        
        # get the dict for each group for the `fig_old` dict
        groups = [i for i in fig_old['data']]
        k = None
        
        # loop through the groups (species)
        for gi, group in enumerate(groups):
            xi = [i for i in range(len(group['x']))]                    # indices of points within the group (use x-coords)
            k = xi if k is None else [i + k[-1] + 1 for i in xi]        # add indices from last group to current xi

            # choose the group indices corresponding to the absolute indices
            group_points = [xi[i] for i, p in enumerate(k) if p in selection]   

            # update `selectedpoints` in the group dict
            fig_old['data'][gi]['selectedpoints'] = group_points

        # read as figure for further changes
        fig = go.Figure(fig_old)

    fig.update_layout(uirevision="Don't Change")

    return fig

Cheers!