Search code examples
pythonjupyter-notebookipythongoogle-colaboratory

How to wait for the user to click a point in a figure in IPython notebook?


I took the following steps to setup an IPython backend in Google Colab notebook:

!pip install ipympl
from google.colab import output
output.enable_custom_widget_manager()  

Then I log the (x,y) location of the user's click on a figure:

%matplotlib ipympl
import matplotlib
import matplotlib.pyplot as plt

fig, ax = plt.subplots()

def onclick(event):
    ix, iy = event.xdata, event.ydata
    print(ix, iy)

cid = fig.canvas.mpl_connect('button_press_event', onclick)

I need the code to wait here until the user selects at least one data point. However, the code will run if I have another command in the same notebook, for example:

print ('done!')

will run without waiting for the user to pick a data point. I tried using this before:

plt.waitforbuttonpress()
print('done!')

However, the compiler gets stuck at plt.waitforbuttonpress() and it doesn't let the user click the figure. Thanks for your help


Solution

  • You could, instead of putting the code you want to run after a button_press_event, after the event listener, you could instead put it in the onclick function. Something like this:

    %matplotlib ipympl
    import matplotlib
    import matplotlib.pyplot as plt
    
    fig, ax = plt.subplots()
    
    def onclick(event):
        ix, iy = event.xdata, event.ydata
        print(ix, iy)
        print('done!')
    
    cid = fig.canvas.mpl_connect('button_press_event', onclick)
    

    Or:

    %matplotlib ipympl
    import matplotlib
    import matplotlib.pyplot as plt
    
    fig, ax = plt.subplots()
    
    def after_onclick():
        print('done!')
    def onclick(event):
        ix, iy = event.xdata, event.ydata
        print(ix, iy)
        after_onclick()
    
    cid = fig.canvas.mpl_connect('button_press_event', onclick)
    

    The reason for this is because mpl_connect works rather oddly (see this question). Instead of waiting for the event to be registered and pausing the code execution, it will run the entire file but keep the event listener open.

    The problem with this way is that it will run the code every time the event is registered.

    If that is not what you want and you only want it to run once try this:

    %matplotlib ipympl
    import matplotlib
    import matplotlib.pyplot as plt
    
    fig, ax = plt.subplots()
    
    def after_onclick():
        print('done')
        fig.canvas.mpl_disconnect(cid)
    def onclick(event):
        ix, iy = event.xdata, event.ydata
        print(ix, iy)
        after_onclick()
    
    cid = fig.canvas.mpl_connect('button_press_event', onclick)
    

    This version adds fig.canvas.mpl_disconnect(cid) so that the code will only run once.