network_animation does not behave as intended in Jupyter Notebooks

[mi-les]

Summary
I had difficulty using anim = wntr.graphics.network_animation(wn) and getting the desired behavior in a Jupyter Notebook (+VS Code), i.e. something where I can replay the network results and click through the timesteps.

The best way to do this would be to use from IPython.display import HTML and then HTML(anim.to_jshtml()). However, the example code below would lead to problems: instead of the desired behavior as demonstrated here, I ended up with a lot of single frames plotted below one another, and an empty interactive video frame.

I figured that this is because the function plot_network, which is called by network_animation, plots each frame with plt.show(block=False).

I'd suggest having an extra argument in plot_network, for example show_plot, set to True by default, but changed to False when called within the network_animation function. With this solution, also demonstrated below, everything works as expected.

Example

import wntr
from IPython.display import HTML

wn = wntr.network.WaterNetworkModel("Net3.inp")
wntr.graphics.plot_network(wn, title=wn.name)

wn.options.hydraulic.demand_model = 'PDD'
wn.options.time.duration = 5*3600
sim = wntr.sim.WNTRSimulator(wn)
results = sim.run_sim()

pressure = results.node['pressure']
threshold = 21.09 # 30 psi
pressure_above_threshold = wntr.metrics.query(pressure, np.greater, threshold)

expected_demand = wntr.metrics.expected_demand(wn)
demand = results.node['demand']
wsa = wntr.metrics.water_service_availability(expected_demand, demand)

anim = wntr.graphics.network_animation(wn, node_attribute=wsa)
HTML(anim.to_jshtml())

Environment
Provide information on your computing environment.

• Operating system: Windows
• Python version: 3.9.13
• WNTR version: 1.0.0
• Jupyter Notebook run in VS Code

Additional context
suggestion:

def plot_network(wn,  node_attribute=None,  ... show_plot=True):
    ...
    if show_plot:
        plt.show()

    return ax


def network_animation(wn, ...):
    ...
    ax = plot_network(wn, ..., show_plot=False)

    def update(n):
        ...

        fig.clf()
        ax = plt.gca()

        ax = plot_network(wn,  ..., show_plot=False)

        return ax

    anim = animation.FuncAnimation(fig, update, interval=50, frames=len(index), blit=False, repeat=repeat)

    return anim

I'll be happy to create a PR for this if needed.

[kaklise]

Thanks for posting this issue. I don't believe I've ever created a network animation in jupyter notebooks, so this is great to know. plt.show(block=False) should have been added to all our graphics functions, but it looks like this is not the case (block=False allows the tests to continue running without closing the plot), see #265. There are other ways we could handle this if needed.

For animation, it sounds like we don't want to show the plots to avoid multiple plots in a jupyter notebook. Feel free to post a PR and we can test this out for the various run environments (command line and testing, jupyter notebooks, spyder, etc.).