[FIXED] NX Graph is placed inside Matplotlib Button

Issue

I'm facing problems trying to separate a simple generated Graph from the Buttons in the plot.
Basically, the idea is that the nodes in the graph should change colour whenever either of the Buttons are pressed. I'm unable to separate the graph from the "Next" Button. And it seems that the graphs inside the button are superseding whenever the Buttons are pressed.

Any ideas how I can separate the graph and prevent the superposition of the graphs upon Button press?
Here's an example code:

import matplotlib.pyplot as plt
import networkx as nx
from matplotlib.widgets import Button
from itertools import product

# Create the 2D graph and define the range
N = 6
G = nx.Graph()

for i in range(1, (N*N)+1):
    G.add_node(i)

def update_node_colors(colors):
    node_pos = list(product(range(0, N), repeat=2))
    pos = dict( (i + 1 + (N - 1 - j) * N, (i,j)) for i,j in node_pos )
    nx.draw(G, pos, with_labels=True, node_color=colors)
    plt.draw()

def prev_button(event):
    colors = ["yellow" for _ in G.nodes()]
    update_node_colors(colors)

def next_button(event):
    colors = ["red" for _ in G.nodes()]
    update_node_colors(colors)

fig, ax = plt.subplots(1, 1, figsize=(7,6))

# Give position to nodes
node_pos = list(product(range(0, N), repeat=2))
pos = dict( (i + 1 + (N - 1 - j) * N, (i,j)) for i,j in node_pos )

nx.draw(G, pos, with_labels=True)

plt.subplots_adjust(bottom=0.2)

ax_prev = plt.axes([0.85, 0.15, 0.15, 0.07])
button_prev = Button(ax_prev, "Prev", color = "green", hovercolor = "blue")
button_prev.on_clicked(prev_button)

ax_next = plt.axes([0.85, 0.05, 0.15, 0.07])
button_next = Button(ax_next, "Next", color = "orange", hovercolor = "red")
button_next.on_clicked(next_button)

#plt.axis('off')
plt.show()

Solution

You are creating an figure and an axis with:fig, ax = plt.subplots(1, 1, figsize=(7,6)) Later you create 2 more axis instances (ax_prev and ax_next) via plt.axes(). When you do this you program switches focus on the last created axis instance. In this case this is ax_next. This means that nx.draw() will now use this axis to draw in.

To switch the focus back onto the main instance ax you can use plt.sca(ax). Then nx.draw() will use the slected main axis ax.

The final code could look like this for example:

def update_node_colors(colors):
    plt.sca(ax)
    node_pos = list(product(range(0, N), repeat=2))
    pos = dict( (i + 1 + (N - 1 - j) * N, (i,j)) for i,j in node_pos )
    nx.draw(G, pos, with_labels=True, node_color=colors)
    plt.draw()

Answered By - tetris programming