时间戳

首页 > 解决方案 > 创建 3D 绘图（不是表面，散点图），其中颜色取决于 z 值

python - 创建 3D 绘图（不是表面，散点图），其中颜色取决于 z 值

问题描述

我想创建并保存一些连续的情节，这样我就可以从这些情节中制作一部 mp4 电影。我希望绘图的颜色取决于 z（第三轴的值）：

我正在使用的代码：

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()

我得到的看起来像这样：

在此处输入图像描述

我想看起来像这样：

在此处输入图像描述

我使用 LineCollection 模块创建了第二个图。问题是它一次打印所有行，不允许我单独保存每行以创建电影。

您可以在此处找到我用来创建图形的数据框：

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

标签： pythonmatplotlib3dfigure

解决方案

海报想要两件事

  1. 颜色取决于 z 值的线条
  2. 线条随时间的动画

为了实现（1）需要将每条线切割成单独的段并为每个段分配颜色；为了获得颜色条，我们需要创建一个可缩放的对象，该对象知道颜色的外部限制。

为了实现 2，需要（a）保存动画的每一帧并在存储所有帧后将其组合，或者（b）利用 matplotlib 中的动画模块。我在下面的示例中使用了后者并实现了以下目标：

在此处输入图像描述

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()

推荐阅读