3D matplotlib: color based on x-axis position

Question

Dear Stackoverflow Users,

I use 3D matplotlib to create 3D envelopes. So far I have managed to achieve almost what I want, but there is the last detail that I would like to solve: I would like the envelope to be colored in accordance with the values ​​of the x axis, and not in accordance with the values ​​of the z axis. I admit that I copied parts of the code to get a graph, without understanding each line in detail, there are several lines that remain cryptic to me. Each line that I don’t understand is marked with the comment “Here is a line I don’t understand,” so if one of you suspects that I need a modification, I don’t understand, they know this, and this can help solve the problem. Here is the working code:

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

This generates the following figure:

How can I use the same colors, but match them with x values ​​instead of z values, since they appear automatically?

Thanks in advance! 🙂

Yours faithfully!

Answer 1

Setting the color of the trisurf graph to another than its Z value is not possible, because unfortunately plot_trisurf ignores the facecolors argument .

However, using the usual surface_plot allows surface_plot to provide an array of facecolors colors.

 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( (XX.min())/float((XX.min()).max()) ) ax2.plot_surface(X,Y,Z ,facecolors=colors, linewidth=0, shade=False ) ax.set_xlabel("x") ax2.set_xlabel("x") plt.show()