Python Forum

I have written the following code in Jupyter notebook to animate the colours of the graph nodes.

%matplotlib notebook

import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation
from functools import partial

def update(i,G,pos,t,ax):
    #fig.clf()    
    cmap = matplotlib.cm.get_cmap('Blues')
    color_map = []
    for node in G:
        color_map.append(cmap(t[i]))
        
    nx.draw_networkx_nodes(G,pos,node_color=color_map,ax=ax)

t = np.linspace(0,1,10)
G = nx.petersen_graph()
pos = nx.spring_layout(G) # positions for all nodes   

fig = plt.figure()
ax = fig.add_axes([0,0,1,1])
#fig.add_axes(rect,label=label1)
plt.axis('off')

upd = partial(update,G=G,pos=pos,t=t,ax=ax)
nx.draw_networkx_edges(G,pos,width=1.0,alpha=0.5,ax=ax) #nodelist=[0,1,2,3], #node_color='r', #node_size=500,
    
ani = matplotlib.animation.FuncAnimation(fig,upd,repeat=True,interval=100,frames=len(t),repeat_delay=10)
plt.show()

Although this works in my case, I would like to be able to execute the whole program inside a function like so

%matplotlib notebook
    
    import networkx as nx
    import matplotlib.pyplot as plt
    import numpy as np
    import matplotlib.animation
    from functools import partial
    
    def update(i,G,pos,t,ax):
        #fig.clf()    
        cmap = matplotlib.cm.get_cmap('Blues')
        color_map = []
        for node in G:
            color_map.append(cmap(t[i]))
            
        nx.draw_networkx_nodes(G,pos,node_color=color_map,ax=ax)
    
    def anim():
        t = np.linspace(0,1,10)
        G = nx.petersen_graph()
        pos = nx.spring_layout(G) # positions for all nodes   
    
        fig = plt.figure()
        ax = fig.add_axes([0,0,1,1])
        #fig.add_axes(rect,label=label1)
        plt.axis('off')
    
        upd = partial(update,G=G,pos=pos,t=t,ax=ax)
        nx.draw_networkx_edges(G,pos,width=1.0,alpha=0.5,ax=ax) #nodelist=[0,1,2,3], #node_color='r', #node_size=500,
        
        ani = matplotlib.animation.FuncAnimation(fig,upd,repeat=True,interval=100,frames=len(t),repeat_delay=10)
        plt.show()
        
    anim()

This one, however, does not work. Could someone tell me why?

adamG