How to update interactive figure in loop with JupyterLab

Question

I am trying to update an interactive matplotlib figure while in a loop using JupyterLab. I am able to do this if I create the figure in a different cell from the loop, but I would prefer to create the figure and run the loop in the same cell.

Simple Code Example:

import matplotlib.pyplot as plt
import time

%matplotlib widget

fig = plt.figure()

for i in range(5):
    x = list(range(i+2))
    xx = [x**2 for x in x]
    plt.clf()
    plt.plot(x, xx)
    fig.canvas.draw()
    
    time.sleep(1)

If fig = plt.figure() is in the same cell as the loop the figure is not updated until the loop finishes:

If I create the figure in a different cell I get the dynamic update, but I would like to be able to create the figure in the same cell if possible so the output is below the loop:

I have tried several answers in other questions (here, here, and here) however, they do not seem to work with interactive figures in JupyterLab. I am using the jupyter/scipy-notebook docker image as my environment, so I believe everything is set up correctly.

Is there any way to get the dynamic update in the same cell as the figure is created?

Answer 1

If you don't want to use asyncio, you can use display(..., display_id=True) to obtain a handle and use .update() on it:

import matplotlib.pyplot as plt
import time
%matplotlib widget
fig = plt.figure()

hfig = display(fig, display_id=True)


def update():
    for i in range(5):
        print(i)
        x = list(range(i + 2))
        xx = [x**2 for x in x]
        plt.clf()
        plt.plot(x, xx)
        fig.canvas.draw()
        hfig.update(fig)
        time.sleep(1)

update()

plt.close(fig)

PS: Make sure to close your figure!

Answer 2

You can use asyncio, taking advantage of the IPython event loop:

import matplotlib.pyplot as plt
import asyncio
%matplotlib widget
fig = plt.figure()


async def update():
    for i in range(5):
        print(i)
        x = list(range(i + 2))
        xx = [x**2 for x in x]
        plt.clf()
        plt.plot(x, xx)
        fig.canvas.draw()
        await asyncio.sleep(1)


loop = asyncio.get_event_loop()
loop.create_task(update());