Create 3D Plot (not surface, scatter), where colour depends on z values
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:
The poster wants two things
- lines with colors depending on z-values
- 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()