How to update graph in a loop?

Question

I created a little model to show a disease could spread. I succeded in showing a different graph for each iteration but I would like to plot a single graph that gets updated at each iteration and shows how the particles move from one iteration to the other.

This is where i call the data i want to plot:

def plotter(population):
for people in population:
    if people.status==0:
        plt.scatter(people.positionx,people.positiony,c='b')
    else:
        if people.healthstatus==0:
            plt.scatter(people.positionx,people.positiony,c='g')
        if people.healthstatus==1:
            plt.scatter(people.positionx,people.positiony,c='y')
        if people.healthstatus==2:
            plt.scatter(people.positionx,people.positiony,c='r')

this is the main where iterate the model

def main(iterations,populationsize):
    popde=generator(populationsize)
    population=popde[0]
    dead=popde[1]
    plt.ion()

    for numit in range(iterations):
        population=movement(population)
        popde2=infection(population,populationsize,dead)
        population=popde2[0]
        dead=popde2[1]
        populationsize=popde2[2]
        plotter(population)
        plt.pause(0.1)
        plt.draw()

The code works perfectly fine, it's just a style issue

I tried looking for other solutions on the web but I couldn't the one that fits my problem. Thanks in advance to all those who will help!

Answer 1

I assume that population is a list and that the attributes positionx and positiony of the object people are scalar numbers (float or int). For future questions consider posting a minimal working example. The following code is based on this answer to a similar question.

It is not necessary to create an individual scatter plot for each person. The code below collects all x-, y-positions and colors in one timestep and plots them in a single scatter plot object. This scatter plot is created before the loop (saved to the variable sc) and then passed to the plotter function, which updates the plot properties (positions & colors) each iteration.

import matplotlib.pyplot as plt

def main(iterations,populationsize):
    popde=generator(populationsize)
    population=popde[0]
    dead=popde[1]
    plt.ion()

    sc = plt.scatter([], [])  # initialise empty scatter plot

    for numit in range(iterations):
        population=movement(population)
        popde2=infection(population,populationsize,dead)
        population=popde2[0]
        dead=popde2[1]
        populationsize=popde2[2]
        plotter(population, sc)  # pass scatter object to plotter function
        plt.draw()  # update the current figure
        plt.pause(0.1)


def plotter(population, sc):
    # create list of all x- and y-positions
    posx = [people.positionx for people in population]
    posy = [people.positiony for people in population]

    # create list of colors
    color = ['#88888']*len(population)  # initialise all colors as grey
    for ip, people in enumerate(population):
        if people.status == 0:
            color[ip] = 'b'
        elif people.healthstatus == 0:
            color[ip] = 'g'
        elif people.healthstatus == 1:
            color[ip] = 'y'
        elif people.healthstatus == 2:
            color[ip] = 'r'
        # if none of these criteria match, the marker remains grey

    sc.set_offsets(np.c_[posx, posy])  # update marker positions
    sc.set_color(color)  # update marker color

As a side note, I would recommend to use the "object-oriented coding style" of matplotlib when writing scripts or functions. The above code always plots in the currently active figure, so if you have previously created another figure or the active figure changes while the code is running you might suddenly be plotting in the wrong figure.