Animated rotating graph with fixed position of axis in matplotlib

Question

I have been trying to create a rotating 3D graph like this which sequentially adds datapoints to the plot. The code is working as intended and the plot looks as I expected. However, my question is whether there is an easier or "cleaner" way to keep the plot rotating while the z-axis stays fixed. My workaround here was to remove the plot's axes and redraw an artificial axis with ticks and labels via single lines and then undo the motion of the viewpoint by changing the position of my artificial axis in each frame. This seems to be overly complicated, so I am looking forward to your suggestions. Thank you!

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt   
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Line3DCollection
import matplotlib.animation as animation

from IPython.display import HTML

from math import floor, ceil
%matplotlib inline


### === Data preparation
flattened_df = np.random.rand(240)

data = []
for i, element in enumerate(flattened_df, 1):
    #print(i, element)
    if element=="***":
        continue
    z = flattened_df[i-1]
    month = (i-1) % 12
    angle = np.deg2rad(month*30)
    x = np.cos(angle)
    y = np.sin(angle)
    data += [[x,y,z]]
    
data = np.array(data).astype("float") #some of input csv values are strings instead of floats

z = np.array(data[:].T[2]) 

### === 3D LineCollection
#segments for the 3D LineCollection
segments = [[[x0, y0, z0], [x1, y1, z1]] for x0, y0, z0, x1, y1, z1 in 
         zip(data[:-1].T[0], data[:-1].T[1], data[:-1].T[2], 
             data[1:].T[0], data[1:].T[1], data[1:].T[2])]

# Create the 3D-line collection object
# color according to y value
#lc = Line3DCollection(segments, cmap=plt.get_cmap('summer'))

# color according to order of appearance (=time)
colormap = plt.get_cmap("summer")
colors = [colormap(k) for k in np.arange(0, 1, 1/len(data))]
#print(len(colors), len(data), colors)
lc = Line3DCollection(segments, colors = colors)
    
lc.set_array(z) 
lc.set_linewidth(2)

### ==== plot
# Create a figure and a 3D Axes
fig = plt.figure()
ax = Axes3D(fig)
ax.set_xlim3d(-1.2, 1.2)
ax.set_ylim3d(-1.2, 1.2)

lowest_tick = floor(2 * data[:].T[2].min())/2 # lowest z value rounded to next lower .5
highest_tick = ceil(2 * data[:].T[2].max())/2 # highest z value rounded to next highest .5

z_tick_values = np.arange(lowest_tick, highest_tick+0.1, 0.5) #+0.1 to include last point
ax.set_zlim3d(lowest_tick-0.5, highest_tick+0.5)

# Hide grid lines
ax.grid(False)

# Hide axes
plt.axis("off")

months = ["Jan","Feb","Mar","Apr","May","Jun","Jul","Aug","Sep","Oct","Nov","Dec"]
for i, month in enumerate(months):
    angle = np.deg2rad(i*30) #30*12 = 360
    x = np.cos(angle)
    y = np.sin(angle)
    ax.text(x, y, lowest_tick-0.5, month, ha="center")

ax.add_collection3d(lc)

z_axis, = ax.plot3D([-1,-1], [0,0], [lowest_tick-0.5, highest_tick+0.5], c="b")
z_ticks = []
z_tick_labels = []

for n in z_tick_values: 
    tick, = ax.plot3D([-1,-1],[0,-1],[n, n], c="b")
    z_ticks.append(tick)
    text = ax.text(-1, 0, n, n, zdir=None, ha="right", va="center")
    z_tick_labels.append(text)
      
# initialization function: plot the background of each frame
def init():
    lc.set_segments([])
    z_axis.set_data([], [])
    return_tuple = (lc, z_axis)
    for tick in z_ticks:
        tick.set_data([], [])
        return_tuple += (tick,)
    for label in z_tick_labels:
        label.set_position(([], []))
        return_tuple += (label,)
    return return_tuple


def animate(i):
    lc.set_segments(segments[:(i//5)])
    lc.set_color(colors[:(i//5)])
    
    angle = np.deg2rad(i-90)
    x_ax = 1.5*np.cos(angle)
    y_ax = 1.5*np.sin(angle)
    
    z_axis.set_data(np.array([x_ax,x_ax]),np.array([y_ax,y_ax]))
    z_axis.set_3d_properties(np.array([lowest_tick-0.5, highest_tick+0.5]))
    
    for n, tick in enumerate(z_ticks):
        tick.set_data(np.array([x_ax, 1.05 * x_ax]), np.array([y_ax, 1.05 * y_ax]))
        tick.set_3d_properties(np.array([z_tick_values[n], z_tick_values[n]]))
    for n, label in enumerate(z_tick_labels):
        label.set_position((1.15 * x_ax, 1.15 * y_ax))
        label.set_3d_properties(z_tick_values[n], zdir=None)
    ax.view_init(elev=10., azim=i)
    return fig,

# call the animator.
anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=len(data)*5, interval=20, blit=True)
anim.save("rotating_graph.mp4")
HTML(anim.to_html5_video())

Result:

Rotating 3D Graph

Animated rotating graph with fixed position of axis in matplotlib

0 Answers