Animating data sorting in python

Question

I was trying to develop a program that visualises the sorting of a list in python as a scatter plot, and I do not know where to start. After googling, I came up with the following code that does the same thing but on a bar chart:

# import all the modules
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import matplotlib as mp
import numpy as np
import random
  
# set the style of the graph
plt.style.use('fivethirtyeight')
  
# input the size of the array (list here)
# and shuffle the elements to create
# a random list
n = int(input("enter array size\n"))
a = [i for i in range(1, n+1)]
random.shuffle(a)
  
# insertion sort
  
  
def insertionsort(a):
    for j in range(1, len(a)):
        key = a[j]
        i = j-1
  
        while(i >= 0 and a[i] > key):
            a[i+1] = a[i]
            i -= 1
  
            # yield the current position
            # of elements in a
            yield a
        a[i+1] = key
        yield a
  
  
# generator object returned by the function
generator = insertionsort(a)
  
# to set the colors of the bars.
data_normalizer = mp.colors.Normalize()
color_map = mp.colors.LinearSegmentedColormap(
    "my_map",
    {
        "red": [(0, 1.0, 1.0),
                (1.0, .5, .5)],
        "green": [(0, 0.5, 0.5),
                  (1.0, 0, 0)],
        "blue": [(0, 0.50, 0.5),
                 (1.0, 0, 0)]
    }
)
  
  
fig, ax = plt.subplots()
  
# the bar container
rects = ax.bar(range(len(a)), a, align="edge",
               color=color_map(data_normalizer(range(n))))
  
# setting the view limit of x and y axes
ax.set_xlim(0, len(a))
ax.set_ylim(0, int(1.1*len(a)))
  
# the text to be shown on the upper left
# indicating the number of iterations
# transform indicates the position with
# relevance to the axes coordinates.
text = ax.text(0.01, 0.95, "", transform=ax.transAxes)
iteration = [0]
  
# function to be called repeatedly to animate
  
  
def animate(A, rects, iteration):
  
    # setting the size of each bar equal
    # to the value of the elements
    for rect, val in zip(rects, A):
        rect.set_height(val)
  
    iteration[0] += 1
    text.set_text("iterations : {}".format(iteration[0]))
  
  
anim = FuncAnimation(fig, func=animate,
                     fargs=(rects, iteration), frames=generator, interval=50,
                     repeat=False)
  
plt.show()

Is there anyone that could give me advices on how to convert this into a scatter plot? (Especially for the def animate part, how should I update the values of the scatter plot?) I have tried reading documentations but it does not seem to give me the help that I need, I would really appreciate it if someone could point me in the correct direction.

Answer 1

One way to do this is to simply remove the line that defines your bar plot rects = ax.bar(range(len(a)), a, align="edge",color=color_map(data_normalizer(range(n)))) and instead declare a scatter plot dots=ax.scatter(range(len(a)),a,color=color_map(data_normalizer(range(n)))). As you noticed, updating the scatter plot is a little different than updating the bar plots. The approach I took was to use set_offsets to update the points on the scatter plot (the equivalent for the bar plot is set_height). You can find more details about this and more ways to update your scatter plot and create scatter plot animations by following the link r-beginners provided in the comments.

Overall the code looks like that:

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import matplotlib as mp
import numpy as np
import random

plt.style.use('fivethirtyeight')

n = int(input("enter array size\n"))
a = [i for i in range(1, n+1)]
random.shuffle(a)
  
def insertionsort(a):
    for j in range(1, len(a)):
        key = a[j]
        i = j-1
  
        while(i >= 0 and a[i] > key):
            a[i+1] = a[i]
            i -= 1
            yield a
        a[i+1] = key
        yield a
  
generator = insertionsort(a)
data_normalizer = mp.colors.Normalize()
color_map = mp.colors.LinearSegmentedColormap(
    "my_map",
    {
        "red": [(0, 1.0, 1.0),
                (1.0, .5, .5)],
        "green": [(0, 0.5, 0.5),
                  (1.0, 0, 0)],
        "blue": [(0, 0.50, 0.5),
                 (1.0, 0, 0)]
    }
)
  
fig, ax = plt.subplots()
  
dots=ax.scatter(range(len(a)),a,color=color_map(data_normalizer(range(n))))
ax.set_xlim(0, len(a))
ax.set_ylim(0, int(1.1*len(a)))
  
text = ax.text(0.01, 0.95, "", transform=ax.transAxes)
iteration = [0]

def animate(A,dots,iteration):

    up_points=np.array([np.arange(len(a)),np.array(A)]).T
    dots.set_offsets(up_points)
    iteration[0] += 1
    text.set_text("iterations : {}".format(iteration[0]))
    
anim=FuncAnimation(fig,func=animate,fargs=(dots,iteration,),frames=generator,interval=30,repeat=True)
plt.show()

And the output with 25 points at 30 fps gives: