How to create an animated scatterplot in Python and safe as .mp4?

Question

I have an array A in python which contains n x n x n matrices, on each position having either the value 1 or -1.

np.random.choice([1, -1], size=(n, n, n))

So say A contains 20 of these matrices. I would like to create a 3d-scatterplot visualizing these matrices, showing a blue dot if the value of the matrix at a position is 1, and a red dot if it is -1.

Each entry of A should be shown for a fixed amount of time, like 0.25 seconds, before the next is shown. Since doing this at runtime for large n is very computationally expensive i would like to render it as .mp4.

To get a better understanding, this is how i would like it to look like: click (source: https://github.com/ddatsko/3d-ising-model)

Thanks for the help!

(Context: I want to visualize the optimization of a 3D Ising-Model)

I tried it like this, and it works but is painfully slow for bigger than single digit n's.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter

plt.rcParams['animation.ffmpeg_path']='path-to-ffmpeg.exe'

metadata = dict(title='Random-Boolean-Matrix', artist='Me')
writer = FFMpegWriter(fps=2, metadata=metadata)

fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))

matrix_size = 3

# init A (array of 3d matrices)
A = []
for _ in range(10):
    A.append(np.random.choice([1, -1], size=(matrix_size, matrix_size, matrix_size)))

plt.xlim(0, matrix_size)
plt.ylim(0, matrix_size)

def plot_matrix(matrix, ax):
    ax.clear()
    ax.set_zlim(0, matrix_size)
    for i in range(matrix_size):
        for j in range(matrix_size):
            for k in range(matrix_size):
                if matrix[i, j, k] == 1:
                    ax.scatter(i, j, k, c='b', marker='o')
                else:
                    ax.scatter(i, j, k, c='r', marker='o')

with writer.saving(fig, "matrix_animation.mp4", 100):
    for i in range(10):
        plot_matrix(A[i], ax)
        writer.grab_frame()

Answer 1

Numpy has a function called np.argwhere(...) that returns the indices of all non-zero elements in an array. This is useful for your example because the numpy function can also be used to find elements of an array that satisfy some condition i.e. elements equal to -1 or 1. That eliminates the need for a triple for loop and speeds up your code a lot! Here's how it works...

indPos = [] # list of indices for the positive elements
indNeg = [] # list of indices for the negative elements
for p in range(len(A)):
    # appends the list of indices for each frame
    # when A[p]==__ evaluates to true, that element is nonzero,
    # so the index for that element is returned and appended to the list
    indPos.append(np.argwhere(A[p]==1))
    indNeg.append(np.argwhere(A[p]==-1))

Using FuncAnimation from matplotlib creates an animated object that can be displayed in a notebook or written to a file. The final output looks like this for matrix_size = 8...

Code

%matplotlib notebook
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter
from matplotlib import animation
from matplotlib.animation import FuncAnimation

plt.style.use('dark_background')
fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))
ax.view_init(25,135)
ax.set_axis_off()
ax.xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
ax.yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
ax.zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
ax.xaxis._axinfo["grid"]['color'] =  (1,1,1,0)
ax.yaxis._axinfo["grid"]['color'] =  (1,1,1,0)
ax.zaxis._axinfo["grid"]['color'] =  (1,1,1,0)
        
matrix_size = 8

# init A (array of 3d matrices)
A = []
for _ in range(30):
    A.append(np.random.choice([1, -1], size=(matrix_size, matrix_size, matrix_size)))

plt.xlim(0, matrix_size)
plt.ylim(0, matrix_size)

indPos = []
indNeg = []
for p in range(len(A)):
    indPos.append(np.argwhere(A[p]==1))
    indNeg.append(np.argwhere(A[p]==-1))

pos = ax.scatter(indPos[0][:,0],indPos[0][:,1],indPos[0][:,2], c='b', marker='o')
neg = ax.scatter(indNeg[0][:,0],indNeg[0][:,1],indNeg[0][:,2], c='r', marker='o')
def animate(i, indPos, indNeg):
    pos._offsets3d = (indPos[i][:,0],indPos[i][:,1],indPos[i][:,2])
    neg._offsets3d = (indNeg[i][:,0],indNeg[i][:,1],indNeg[i][:,2])
    return pos, neg


anim = FuncAnimation(fig, animate, frames=len(A), interval=200, fargs=(indPos, indNeg), blit=True, repeat=False)
writervideo = animation.FFMpegWriter(fps =5, bitrate =80)
anim.save("matrix_animation.mp4")