how to update a matplotlib heatmap plot without creating a new window - python-3.x

Question

I have matrix class that inherits from list. This class can display itself as a matplotlib heatmap representation of the matrix.
I'm trying to have the class written such that when I change values in the matrix, I can call the matrix's method plot() and it'll update the plot to reflect the matrix changes in the heatmap.
However, every time I run the method plot(), it creates a new heatmap in a new window instead of updating the existing plot. How could I get it simply to update the existing plot?
In the code below, there are three main parts: the main function shows how an instance of the matrix class is created, plotted and updated; the matrix class is basically a list object, with some minor functionality (including plotting) bolted on; the function plotList() is the function the matrix class calls in order to generate the plot object initially.
import time
import random
import matplotlib.pyplot as plt
plt.ion()
import numpy as np
def main():
print("plot 2 x 2 matrix and display it changing in a loop")
matrix = Matrix(
numberOfColumns = 2,
numberOfRows = 2,
randomise = True
)
# Plot the matrix.
matrix.plot()
# Change the matrix, redrawing it after each change.
for row in range(len(matrix)):
for column in range(len(matrix[row])):
input("Press Enter to continue.")
matrix[row][column] = 10
matrix.plot()
input("Press Enter to terminate.")
matrix.closePlot()
class Matrix(list):
def __init__(
self,
*args,
numberOfColumns = 3,
numberOfRows = 3,
element = 0.0,
randomise = False,
randomiseLimitLower = -0.2,
randomiseLimitUpper = 0.2
):
# list initialisation
super().__init__(self, *args)
self.numberOfColumns = numberOfColumns
self.numberOfRows = numberOfRows
self.element = element
self.randomise = randomise
self.randomiseLimitLower = randomiseLimitLower
self.randomiseLimitUpper = randomiseLimitUpper
# fill with default element
for column in range(self.numberOfColumns):
self.append([element] * self.numberOfRows)
# fill with pseudorandom elements
if self.randomise:
random.seed()
for row in range(self.numberOfRows):
for column in range(self.numberOfColumns):
self[row][column] = random.uniform(
self.randomiseLimitUpper,
self.randomiseLimitLower
)
# plot
self._plot = plotList(
list = self,
mode = "return"
)
# for display or redraw plot behaviour
self._plotShown = False
def plot(self):
# display or redraw plot
self._plot.draw()
if self._plotShown:
#self._plot = plotList(
# list = self,
# mode = "return"
# )
array = np.array(self)
fig, ax = plt.subplots()
heatmap = ax.pcolor(array, cmap = plt.cm.Blues)
self._plot.draw()
else:
self._plot.show()
self._plotShown = True
def closePlot(self):
self._plot.close()
def plotList(
list = list,
mode = "plot" # plot/return
):
# convert list to NumPy array
array = np.array(list)
# create axis labels
labelsColumn = []
labelsRow = []
for rowNumber in range(0, len(list)):
labelsRow.append(rowNumber + 1)
for columnNumber in range(0, len(list[rowNumber])):
labelsColumn.append(columnNumber)
fig, ax = plt.subplots()
heatmap = ax.pcolor(array, cmap = plt.cm.Blues)
# display plot or return plot object
if mode == "plot":
plt.show()
elif mode == "return":
return(plt)
else:
Exception
if __name__ == '__main__':
main()
I'm using Python 3 in Ubuntu.

Answer 1

The method plot(self) creates a new figure in the line fig, ax = plt.subplots(). To use an existing figure you can give your figure a number or name when it's first created in plotList():
fig = plt.figure('matrix figure')
ax = fig.add_subplot(111)
then use
plt.figure('matrix figure')
ax = gca() # gets current axes
to make that the active figure and axes. Alternately, you might want to the figure and axis created in plotList and pass them to plot.