Initial plot not clearing when using dropdown menu to choose new data

Question

I am teaching myself how to use the ipywidgets package by doing analysis on the total covid deaths for each country and want to make an interactive plot with a dropdown menu for Country choice.

My code works almost perfectly. The only thing is that when I chose a new country, the initial plot does not clear and I am left with 2 plots, the correct plot from choice, and the initial plot (as seen in the image below).

Can someone please point me in the right direction as to how I can remove the initial plot?

Here is the code that I am using:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import ipywidgets as widgets
import seaborn as sns

# Get dataset
data_url = "https://raw.githubusercontent.com/owid/covid-19-data/master/public/data/cases_deaths/total_deaths_per_million.csv"
df = pd.read_csv(data_url, index_col=0, parse_dates=[0], engine='python')

# na values = 0
df.fillna(0, inplace=True)
df.head()

# add year-week column
df['Year_Week'] = df.index.to_period('W').strftime('%Y-%U')

# keep only last day of week and change to datetime type
df = df.groupby(df['Year_Week']).last('1D')
df.index = pd.to_datetime(df.index + '-0', format='%Y-%U-%w')

# drop columns that aren't a country
df_country = df.drop(['World', 
                    'Africa', 
                    'Asia', 
                    'Europe', 
                    'European Union', 
                    'High income', 
                    'Low income', 
                    'Lower middle income', 
                    'North America', 
                    'South America', 
                    'Upper middle income'], 
                    axis=1)

# create function to update plot based on selected country
def update_plot(country):
    ax.clear()  # clear existing plot
    ax.plot(df.index, df_country[country])  # plot selected country
    
    # set x-axis tick locations and labels
    xticks = pd.date_range(start=df_country.index[0].strftime('%Y-01-01'), end=df_country.index[-1], freq='AS')
    xticklabels = [x.strftime('%Y') for x in xticks]
    ax.set_xticks(xticks)
    ax.set_xticklabels(xticklabels)
    ax.set_title(f"Total deaths per million ({country})")  # update plot title
    ax.set_xlabel("Date")
    ax.set_ylabel("Deaths per million")
    fig.canvas.draw()  # redraw canvas

# create drop-down menu with country names as options
country_dropdown = widgets.Dropdown(
    options=df_country.columns, 
    value=df_country.columns[0], 
    description='Country'
)

# create plot
fig, ax = plt.subplots()
update_plot(country_dropdown.value) # initial plot

# set up widget interaction
output = widgets.Output()
display(country_dropdown, output)

def on_change(change):
    if change['type'] == 'change' and change['name'] == 'value':
        with output:
            output.clear_output()
            update_plot(change['new'])
            display(fig)

country_dropdown.observe(on_change)

Answer 1

If you're working in a notebook, make your plot interactive by adding %matplotlib widget at the beginning of the cell.

Then, you don't need to use Output, simply update your fig and re-draw it in the event callback function.

%matplotlib widget

# [...]

# create plot
fig, ax = plt.subplots()
update_plot(country_dropdown.value) # initial plot

# set up widget interaction
display(country_dropdown)

def on_change(change):
    if change['type'] == 'change' and change['name'] == 'value':
        update_plot(change['new'])
        plt.draw()
        
country_dropdown.observe(on_change)

Execute the cell

Change the dropdown value. The figure is updated in the cell output without the initial one