Create 3D Plot (not surface, scatter), where colour depends on z values

Question

I want to create and save a number of sequential plots so I can then make an mp4 movie out of those plots. I want the color of the plot to depend on z (the value of the third axis):

The code I am using:

import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator
import numpy as np

file_dir1 = r"C:\Users\files\final_files\B_6_sec\_read.csv" 


specs23 = pd.read_csv(file_dir1, sep=',')

choose_file   = specs23          # Choose file betwenn specs21, specs22,...

quant         = 0               #  Choose between 0,1,...,according to the following list

column        = ['$\rho$', '$V_{x}$', '$V_{y}$', '$V_{z}$','$B_{x}$', '$B_{y}$','$B_{z}$','$Temperature$']

choose_column = choose_file[column[quant]] 
                               
resolution    = 1024                                       # Specify resolution of grid 

t_steps       = int(len(specs23)/resolution)               # Specify number of timesteps




fig, ax = plt.subplots(subplot_kw={"projection": "3d"},figsize=(15,10))

# Make data.
X = np.arange(0, resolution, 1)
Y = np.arange(0, int(len(specs23)/resolution),1)
X, Y = np.meshgrid(X, Y)

Z = choose_file[column[quant]].values


new_z = np.zeros((t_steps,resolution))   # Selected quantity as a function of x,t
    

###  Plot figure ###


for i in range(0,int(len(choose_file)/resolution)):
    zs = choose_column[i*resolution:resolution*(i+1)].values
    new_z[i] = zs
        

for i in range(len(X)):
    ax.plot(X[i], Y[i], new_z[i]) #%// color binded to "z" values


ax.zaxis.set_major_locator(LinearLocator(10))
# A StrMethodFormatter is used automatically
ax.zaxis.set_major_formatter('{x:.02f}')


plt.show()

What I am getting looks like this:

I would like to look it like this:

I have created the second plot using the LineCollection module. The problem is that it prints all the lines at once not allowing me to save each separately to create a movie.

You can find the dataframe I am using to create the figure here:

https://www.dropbox.com/s/idbeuhyxqfy9xvw/_read.csv?dl=0

Answer 1

The poster wants two things

1. lines with colors depending on z-values
2. animation of the lines over time

In order to achieve(1) one needs to cut up each line in separate segments and assign a color to each segment; in order to obtain a colorbar, we need to create a scalarmappable object that knows about the outer limits of the colors.

For achieving 2, one needs to either (a) save each frame of the animation and combine it after storing all the frames, or (b) leverage the animation module in matplotlib. I have used the latter in the example below and achieved the following:

from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt, numpy as np
from mpl_toolkits.mplot3d.art3d import Line3DCollection

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

# generate data
x = np.linspace(-5, 5, 500)
y = np.linspace(-5, 5, 500)
z = np.exp(-(x - 2)**2)

# uggly
segs = np.array([[(x1,y2), (x2, y2), (z1, z2)] for x1, x2, y1, y2, z1, z2 in zip(x[:-1], x[1:], y[:-1], y[1:], z[:-1], z[1:])])
segs = np.moveaxis(segs, 1, 2)

# setup segments

# get bounds
bounds_min = segs.reshape(-1, 3).min(0)
bounds_max = segs.reshape(-1, 3).max(0)

# setup colorbar stuff
# get bounds of colors
norm = plt.cm.colors.Normalize(bounds_min[2], bounds_max[2])
cmap = plt.cm.plasma
# setup scalar mappable for colorbar
sm   = plt.cm.ScalarMappable(norm, plt.cm.plasma)

# get average of segment
avg = segs.mean(1)[..., -1]
# get colors
colors = cmap(norm(avg))
# generate colors
lc = Line3DCollection(segs, norm = norm, cmap = cmap, colors = colors)
ax.add_collection(lc)

def update(idx):
    segs[..., -1] = np.roll(segs[..., -1], idx)
    lc.set_offsets(segs)
    return lc

ax.set_xlim(bounds_min[0], bounds_max[0])
ax.set_ylim(bounds_min[1], bounds_max[1])
ax.set_zlim(bounds_min[2], bounds_max[2])
fig.colorbar(sm)

from matplotlib import animation
frames = np.linspace(0, 30, 10, 0).astype(int)
ani = animation.FuncAnimation(fig, update, frames = frames)
ani.save("./test_roll.gif", savefig_kwargs = dict(transparent = False))

fig.show()