3D matplotlib: color depending on x axis position

Question

Dear Stackoverflow users,

I'm using 3D matplotlib to generate 3D envelopes. So far I got success in getting almost what I want, but there is a last detail I would like to solve: I would like the envelope to be colored according to x axis values and not according to z axis values. I admit I copied parts of the code to get the graph without understanding each line in detail, there are a few lines that remain cryptic to me. Each line I don't understand is marked by a comment "Here line I don't understand", so that if one of you suspect that the modification I need is in a line I don't understand, they know it and it may 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 associate them to x values instead of z values as they seem to be automatically?

Thanks in advance! 🙂

Best regards!

Answer 1

Setting the color of a trisurf plot to something other than its Z values is not possible, since unfortunately plot_trisurf ignores the facecolors argument.

However using a normal surface_plot makes it possible to supply an array of colors to 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()