3D matplotlib:颜色取决于x轴位置

时间:2017-03-17 08:00:21

标签: python matplotlib colors 3d

亲爱的Stackoverflow用户,

我正在使用3D matplotlib来生成3D信封。到目前为止,我几乎得到了我想要的东西,但是我想要解决的最后一个细节:我希望信封根据x轴值而不是根据z轴值着色。 我承认我复制了部分代码以获取图形而不详细了解每一行,有几行对我来说仍然含糊不清。我不理解的每一行都标有注释“这里我不理解”,所以如果你们中的一个人怀疑我需要的修改是一行,我不明白,他们知道它可能帮助解决问题。这是工作代码:

# ----- System libraries and plot parameters-----

import argparse
import re
import glob, os, sys
import subprocess
import math
import copy
import hashlib
import scipy
from scipy import optimize
import time 
from decimal import *
import matplotlib.pyplot as plt
import matplotlib.pylab  as pylab
import matplotlib.colors as colors
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.ticker import MaxNLocator
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
from numpy.random import randn, shuffle
from scipy import linspace, meshgrid, arange, empty, concatenate, newaxis, shape
import numpy as np
import numpy
from mpl_toolkits.axes_grid1 import make_axes_locatable

params = {'legend.fontsize' :  70,
          'figure.figsize'  :  (80, 30),
          'axes.labelsize'  :  70,
          'axes.titlesize'  :  70,
          'xtick.labelsize' :  70,
          'ytick.labelsize' :  70}
pylab.rcParams.update(params)
FFMPEG_BIN = "C:\Users\User\Desktop\ffmpeg-20170125-2080bc3-win64-static\bin\ffmpeg.exe"

parser = argparse.ArgumentParser(description='utility to print 3D sigma profiles', formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('--name', type=str, help='name of prf and pot files without the extension, example for tempjob1.prf: --name="tempjob1"', default=["all"])    
args = parser.parse_args()

#parse sigma profile
name    = args.name + ".prf"
with open(name) as f:
    sig_prof_set = f.read().splitlines()
sigma       = list()
profile     = list()
sigma_set   = list()
profile_set = list()
dieze       = 0
for l in sig_prof_set:
    if dieze < 2: #the first dummy compound should not be taken into account and once we reach the second compound, it is the first layer so we start the filling
        if "#" in l:
            dieze += 1 
        pass
    else:
        if "#" in l:
            if dieze > 1: #each time we reach a dieze, we store the sigma profile gathered into the sigma profile set and empty the list for the next
                sigma_set.append(sigma)
                profile_set.append(profile)
                sigma       = list()
                profile     = list()
            dieze += 1 #the first dummy compound should not be taken into account       
        else:
            splitted = l.split()
            sigma.append(splitted[0])
            profile.append(splitted[1])

#display 3D plot
fig = plt.figure()

#convert data to numpy arrays
sigma_set     = numpy.array(sigma_set)
profile_set   = numpy.array(profile_set)
potential_set = numpy.array(potential_set)

#shape data for graphs
layer             = numpy.array(range(len(sigma_set)))
layer_flatten     = list()
sigma_flatten     = list()
profile_flatten   = list()
potential_flatten = list()

#X is sigma, Y is layer number, Z is profile or potential
for i in layer:
    for j in range(len(sigma_set[0])):
        layer_flatten.append(layer[i])
        sigma_flatten.append(float(sigma_set[i][j]))
        profile_flatten.append(float(profile_set[i][j]))
        potential_flatten.append(float(potential_set[i][j]))

#assign graph data        
X  = numpy.array(sigma_flatten)
Y  = numpy.array(layer_flatten)
Z1 = numpy.array(profile_flatten)
Z2 = numpy.array(potential_flatten)

#actually make 3D plot
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d') #Here line I don't understand

surf = ax.plot_trisurf(X, Y, Z1, cmap=cm.jet, linewidth=0)
fig.colorbar(surf)

#set title of graph and axes
title = ax.set_title("Z-dependent sigma-profile")
title.set_y(1.01)                          #Here line I don't understand
ax.xaxis.set_major_locator(MaxNLocator(5)) #Here line I don't understand
ax.yaxis.set_major_locator(MaxNLocator(6)) #Here line I don't understand
ax.zaxis.set_major_locator(MaxNLocator(5)) #Here line I don't understand
ax.set_xlabel('sigma (e/A^2)')
ax.set_ylabel('layer')
ax.set_zlabel('p(sigma)')
ax.xaxis.labelpad = 100
ax.yaxis.labelpad = 70
ax.zaxis.labelpad = 70

fig.tight_layout()                         #Here line I don't understand

#save the figure
fig.savefig('3D_sig_prf{}.png'.format(args.name))

这会生成下图: the 3D plot colored according to z values

我如何使用相同的颜色,但是将它们与x值而不是z值相关联,因为它们似乎是自动的?

提前致谢!

祝你好运!

1 个答案:

答案 0 :(得分:2)

由于遗憾的是trisurf ignores the facecolors参数,因此无法将plot_trisurf绘图的颜色设置为Z值以外的颜色。

然而,使用普通surface_plot可以为facecolors提供一系列颜色。

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

X,Y = np.meshgrid(np.arange(10), np.arange(10))
Z = np.sin(X) + np.sin(Y)
x = X.flatten()
y = Y.flatten()
z = Z.flatten()

fig = plt.figure(figsize=(9,3.2))
plt.subplots_adjust(0,0.07,1,1,0,0)
ax = fig.add_subplot(121, projection='3d')
ax2 = fig.add_subplot(122, projection='3d')
ax.set_title("trisurf with color acc. to z")
ax2.set_title("surface with color acc. to x")

ax.plot_trisurf(x,y,z ,  cmap="magma")

colors =plt.cm.magma( (X-X.min())/float((X-X.min()).max()) )
ax2.plot_surface(X,Y,Z ,facecolors=colors, linewidth=0, shade=False )

ax.set_xlabel("x")
ax2.set_xlabel("x")
plt.show()

enter image description here