Python Plotly chart update with two dropdowns

I am trying to build a plotly scatterplot in Jupyter Lab to be able to see dependencies between various columns in a DataFrame. I want to have two dropdown menus (corresponding to the X and Y axes), in each of which a full list of the DF columns will be available. When I select a column in any of the menus, the data on the appropriate axis should be replaced by the column I selected (so, if I select the same column for X and Y, I would expect a straight line).

Below is my current implementation with a sample DataFrame:

# Creating the DataFrame    
temp = pd.DataFrame(np.random.randint(0, 1000, (100, 10)))
col_list = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J']
temp.columns = col_list

# Init figure with the A column on both axes by default
fig = go.Figure()

default_col = 0
fig.add_trace(
    go.Scatter(
        x=temp[col_list[default_col]].values,
        y=temp[col_list[default_col]].values,
        name="Metric correlation",
        mode="markers"
    ),
)
fig.update_xaxes(title_text=col_list[default_col])
fig.update_yaxes(title_text=col_list[default_col])

col_list = temp.columns

# Building options for each of the lists
btns_x = [
    dict(
        label=c,
        method="update",
        args=[
            {"x": temp[c].fillna(0).values,
             'xaxis': {'title': c}
             }],
    ) for c in col_list]


btns_y = [
    dict(
        label=c,
        method="update",
        args=[
            {"y": temp[c].fillna(0).values,
             'yaxis': {'title': c}
             }],
    ) for c in col_list]

# Adding the lists to the figure
fig.update_layout(
    updatemenus=[
        dict(
            buttons=btns_x,
            # method="update",
            direction="down",
            pad={"r": 10, "t": 10},
            showactive=True,
            x=0.1,
            xanchor="left",
            y=1.1,
            yanchor="top"
        ),
        dict(
            buttons=btns_y,
            # method="update",
            direction="down",
            pad={"r": 10, "t": 10},
            showactive=True,
            x=0.1,
            xanchor="right",
            y=1.1,
            yanchor="top"
        ),
    ]
)
fig.update_layout(width=1000, height=1000)
fig.show()

The figure draws correctly initially:

Still, there are a few problems:

1. When I change values in a dropdown, it only works once. On the next tries nothing happens
2. If I first change the value on one dropdown and then on the other, all data disappears from the graph (see screenshot below)
3. The axes labels are not being updated

Answer

It’s just about being systematic around the list comprehensions. Below fully works, allows selection of any column and updates appropriate axis title.

import pandas as pd
import numpy as np
import plotly.express as px

temp = pd.DataFrame(np.random.randint(0, 1000, (100, 10)))
col_list = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"]
temp.columns = col_list

fig = px.scatter(temp, x="A", y="B")

fig.update_layout(
    updatemenus=[
        {
            "buttons": [
                {
                    "label": c,
                    "method": "update",
                    "args": [
                        {axis: [temp[c]]},
                        {f"{axis}axis": {"title": {"text": c}}},
                    ],
                }
                for c in temp.columns
            ],
            "x": 0 if axis == "x" else 0.1,
            "y": 1.2,
        }
        for axis in "xy"
    ]
)