How to plot animated scatterplot using seaborn / matplotlib changing colour of fixed or moving marker over time

Question

I am trying to plot the readings of several air pollution sensors over time. In the first case the location (lat, long) of each sensor is fixed on the scatterplot but at different times of the day, or times of the year the colour will change depending upon the level of pollution. (In a more advanced case the same situation but with mobile sensors so coordinates will change and colours will change over time). I am running into problems using 'set_data' in the initialisation and animation functions

Most of the examples I have found online on SO and matplotlib documentation related to plotting animated line graphs rather than scatterplots. The link identified as a duplicate covers the same topic but found the first option complex and hard to adjust to my own fairly simple needs. The second solution given gave me NameError: name 'xrange' is not defined This has proved challenging as seaborn and scatterplots seem to have different code structure to lines. Hence I have asked this question.

My original goal was to use seaborn scatterplot with x and y fixed but with hue changing in each frame depending on the pollution level (see code below) but I seem to get a problem when setting initialisation function

AttributeError: 'AxesSubplot' object has no attribute 'set_data'

I have subsequently attempted to use matplotlib scatter using x, y and facecolors as variables but with the same problem

import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import random
%matplotlib notebook

# Suppose 3 fixed sensors, each with a result every day for 5 days
days = sorted(list(range(5))*3)
channel = (list(range(3)))*5
long = (random.sample(range(10, 20), 3))*5
lat = (random.sample(range(25, 35), 3))*5
colours = ['green', 'yellow', 'orange', 'red', 'brown']
colour = random.choices(colours, k=15)

# create dataframe
data = pd.DataFrame(list(zip(days, channel, long, lat, colour)), columns = ['day', 'sensor', 'x', 'y', 'colour'] )

# Set up the plot to be animated
fig, ax = plt.subplots(figsize=(10,10))
plt.xlim(10, 20)
plt.xlabel('Longitude',fontsize=20)
plt.ylim(25, 35)
plt.ylabel('Latitude',fontsize=20)
plt.title('Daily changes in pollution levels',fontsize=20)
p = sns.scatterplot([], [], hue= [], markers='o',s=500, ax=ax)
for i, txt in enumerate(data.sensor):
    ax.annotate(txt, xy=(data.x[i], data.y[i]), textcoords='offset points', xytext=(10,10), fontsize=20, weight='bold')

# initialization function 
def init(): 
    # creating an empty plot/frame 
    p.set_data([], [], []) 
    return p

# animation function 
def animate(i): 
    # x, y, hue values to be plotted 
    for j in range(0,3):
        x = data.loc[(data.day ==i) & (data.sensor ==j), 'x']
        y = data.loc[(data.day ==i) & (data.sensor ==j), 'y'] 
        hue = data.loc[(data.day ==i) & (data.sensor ==j), 'colour']  
    # set/update the x and y axes data 
        p.set_data(x, y, hue)  
    # return plot object 
    return p

# call the animator     
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=15, interval=20, blit=True) 

plt.show()

# save the animation as mp4 video file 
anim.save('sensors.mp4', writer = 'ffmpeg', fps = 5)

I suspect there are other errors but the biggest obstacle is the error message AttributeError: 'AxesSubplot' object has no attribute 'set_data' in reference to the initiation function when I attempt to save animation but I cannot find an alternative way of doing this.

Very grateful for advice on this or any other obvious errors

Answer 1

Here you go:

import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import random
%matplotlib notebook

# Suppose 3 fixed sensors, each with a result every day for 5 days
days = sorted(list(range(5))*3)
channel = (list(range(3)))*5
long = (random.sample(range(10, 20), 3))*5
lat = (random.sample(range(25, 35), 3))*5
colours = ['green', 'yellow', 'orange', 'red', 'brown']
colour = random.choices(colours, k=15)

# create dataframe
data = pd.DataFrame(list(zip(days, channel, long, lat, colour)), columns = ['day', 'sensor', 'x', 'y', 'colour'] )
# print (data)## use this to better understand what is being plotted

# Set up the plot to be animated
# I'm using ax. here, but you could also just use plt.
fig, ax = plt.subplots(figsize=(5,5))
ax.set_xlim(10, 20)
ax.set_xlabel('Longitude',fontsize=20)
ax.set_ylim(25, 35)
ax.set_ylabel('Latitude',fontsize=20)
ax.set_title('Daily changes in pollution levels',fontsize=20)

for i, txt in enumerate(data.sensor):
    ax.annotate(txt, xy=(data.x[i], data.y[i]), textcoords='offset points', xytext=(10,10), fontsize=20, weight='bold')

# for convenience: define a function which prepares the data
def get_data(day=0,sensor_id=0):
    x = data.loc[(data.day ==day) & (data.sensor ==sensor_id), 'x']
    y = data.loc[(data.day ==day) & (data.sensor ==sensor_id), 'y'] 
    col = data.loc[(data.day ==day) & (data.sensor ==sensor_id), 'colour']  
    return x,y,col

# initialization function 
def init(): 
    # plot the first day (day=0) here:
    for j in range(3):
        x,y,col=get_data(day=0,sensor_id=j)
        scat = ax.scatter(x,y,c=col, s=100)
    return scat

# animation function 
def animate(i): 
    for j in range(0,3):        
        x,y,col=get_data(day=i,sensor_id=j)
        # print(i,col)## use this to understand "where" we are
        scat = ax.scatter(x,y,c=col, s=100)

    # return plot object 
    return scat

# call the animator     
# you are iterating over day=i, so you only have 5 frames here
# also you cant use blit=True here
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=5, interval=200,)

plt.show()