How can I fix legend entries not matching with my matplotlib line chart on a streamlit app?

Question

I want to display a plot in a streamlit app. After selecting the spec in a sidebar widget, I display the selected graphs in a line chart. The problem I have is with the legend. When I select or unselect a spec, the graph is updated correctly but the legend is not.

So for example I have Specs 1 though 5 displayed correctly, but according to the legend Spec 1 is Spec 4, which is not true...

I tried filtering for different things and locating the legend entry at different locations, but that did nothing. Below is my current code.

# Import libraries
import pandas as pd
import streamlit as st
import matplotlib.pyplot as plt

#----------------------- Stress - Strain-----------------------
# Load data from source file
df = pd.read_excel(
    io = 'MRE.xlsx',
    engine='openpyxl',
    sheet_name='Sheet1',
    skiprows=0,
    usecols='A:E',
    nrows=11168,
)

#SIDEBAR -> add things that can be filtered for
st.sidebar.header('Enter filter for Stress and Strain display')

spec = st.sidebar.multiselect(
    'Select spec:',
    options=df['Spec'].unique(),
    default=df['Spec'].unique()
)

df_selection = df.query(
    'Spec == @spec'
)

# Plot data and add legend, title, ...
fig = plt.figure()
ax = fig.add_subplot()
df_selection.groupby("Spec").plot(x="Average strain (%)",
                                  y="Average stress (MPa)",
                                  legend=False,
                                  ax=ax,
                                  linewidth = 0.65,
                                  figsize=(10,6))
plt.xlabel('Strain (%/100)')
plt.ylabel('Stress (MPa)')
plt.title('Stress-Strain diagram')
plt.grid(visible=True)

# Shrink current axis's height by 10% on the bottom
# box = ax.get_position()
# ax.set_position([box.x0, box.y0 + box.height * 0.1,
#                   box.width, box.height * 0.9])

# # Put a legend below current axis
plt.legend(df_selection['Spec'].unique(),loc='upper center', bbox_to_anchor=(0.5, -0.15),
          ncol=5)
plt.show()
st.pyplot(fig)

How the legend entry should look like vs. how it actually looks

Sorry, I don't know how to share the data for the df, so I'll share the github link.

Answer 1

Found a solution that seems to work, at least at the moment. I separated the data, so each spec has their own column. And then I used plt.plot on the new variables instead of the group_by. Seems to work for now.

t = pd.DataFrame();
u = pd.DataFrame();

    values = range(0,len(spec));

for i in values:
    t[i] = new_df_strain.iloc[:,i];
    u[i] = new_df_stress.iloc[:,i];

t2 = t.set_axis(axis = 1, labels = spec, inplace=False);
u2 = u.set_axis(axis = 1, labels = spec, inplace=False);

# Plot data and add legend, title, ...

fig = plt.figure();
plt.rcParams['font.size'] = 35;
ax = fig.add_subplot();
plt.plot(t2,
     u2,
     linewidth = 1.5,);

plt.xlabel('Strain (%)');
plt.ylabel('Stress (MPa)');