3D reskin plot

reskin_3d_plot.py

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

class SubplotAnimation(TimedAnimation):
    def __init__(self, num_steps=50):
        fig = plt.figure()
        self.ax1 = fig.add_subplot(2, 1, 1) 
        self.ax2 = fig.add_subplot(2, 1, 2, projection='3d')
        self.ax2.set(xlim3d=(-1, 1), xlabel='X')
        self.ax2.set(ylim3d=(-1, 1), ylabel='Y')
        self.ax2.set(zlim3d=(0, 1), zlabel='Z')

        # TODO Dummy image data, replace with real data
        def rand_im(ax, i, num_steps):
            im = np.zeros((100, 100, 3), dtype=int)
            im[:, :, 0] = i*(255/num_steps)
            im[:, :, 1] = 0
            im[:, :, 2] = 0
            return im
        self.ims = [rand_im(self.ax1, i, num_steps) for i, _ in enumerate(range(num_steps))]

        # TODO Dummy arrow data, replace with real data
        self.xs = [
            [np.sin(x) for x in np.linspace(0, np.pi, num_steps)],
            [np.sin(x) for x in np.linspace(0, np.pi, num_steps)],
            [np.sin(x) for x in np.linspace(0, np.pi, num_steps)],
            [np.sin(x) for x in np.linspace(0, np.pi, num_steps)],
            [np.sin(x) for x in np.linspace(0, np.pi, num_steps)]
        ]
        self.ys = [
            [np.cos(y) for y in np.linspace(0, np.pi, num_steps)],
            [np.cos(y) for y in np.linspace(0, np.pi, num_steps)],
            [np.cos(y) for y in np.linspace(0, np.pi, num_steps)],
            [np.cos(y) for y in np.linspace(0, np.pi, num_steps)],
            [np.cos(y) for y in np.linspace(0, np.pi, num_steps)]
        ]
        self.zs = [
            [z for z in np.linspace(0.1, 1, num_steps)],
            [z for z in np.linspace(0.1, 1, num_steps)],
            [z for z in np.linspace(0.1, 1, num_steps)],
            [z for z in np.linspace(0.1, 1, num_steps)],
            [z for z in np.linspace(0.1, 1, num_steps)]
        ]

        self.t = np.linspace(0, num_steps, num_steps-1)
        self.data= np.array([list(zip(self.xs[i], self.ys[i], self.zs[i])) for i in range(len(self.zs))])
        TimedAnimation.__init__(self, fig, interval=50, blit=True, repeat_delay=1000)
    
    def _draw_frame(self, framedata):
        i = framedata
        self.ax1.imshow(self.ims[i], animated=True)
        self.quiver.remove()
        # center top left right bottom
        self.quiver = self.ax2.quiver([0.0, 0.0, -0.5, 0.5, 0.0], [0.0, 0.5, 0.0, 0.0, -0.5], [0, 0, 0, 0, 0], 
            self.data[:, i, 0], self.data[:, i, 1], self.data[:, i, 2])

    def new_frame_seq(self):
        return iter(range(self.t.size))

    def _init_draw(self):
        self.quiver = self.ax2.quiver([], [], [], [], [], [])

ani = SubplotAnimation()
ani.save('test.gif')
# plt.show()