Plot 4D graph in python2.7

Question

I want to plot red, blue and green colors on the three axis and an array which stores the value corresoding to each combination of color in python2.7....when i run my program either becomes unresponsive for 24 hours or it gives me memory error. Here is my code:

import pylab
import math
from itertools import product
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np

N=[]
p=np.zeros((256,256,256))
S=[]
fig=plt.figure()
ax=fig.gca(projection='3d')
X=np.arange(0,256,1) #for one of the features either red, blue or green
Y=np.arange(0,256,1)
X,Y = np.meshgrid(X,Y)
R=np.sqrt(X**2 + Y**2)
Z=R/np.sqrt(2)
N=p.flatten();
N=(p[i,j,k] for k in Z)            
surf=ax.plot_surface(X,Y,Z, rstride=1, cstride=1,
                    facecolors=cm.jet(N),
                    linewidth=0, antialiased=False, shade=False)
plt.show()

Please help. I have read the previous posts, and have used them, still I am getting memory error. Here p is a containing values of combinations of red, green and blue. For simplicity I have initialized it to zero...it is giving the following error..colset.append(fcolors[rs][cs]) IndexError: index out of bounds

Answer 1

First, your program is slow because you're doing a lot of unnecessary work building N. You're building a 70 MB list a few bytes at a time (256*256*256=16,777,216 appends!). A better (faster, memory efficient) way to build p is to use numpy's array broadcasting, and then reuse p to make N:

import numpy as np
a = np.arange(256)
p = a[:,np.newaxis,np.newaxis] * a[np.newaxis,:,np.newaxis] * a[np.newaxis,np.newaxis,:] 
N = p.flatten()

Second and more importantly, you're not using plot_surface() correctly. According to the docs, X, Y and Z should be 2D arrays. X and Y lay down a 2D grid and Z provides the "height" for each point on that 2D grid. If you want to manually set the facecolor, it should also be a 2D array. You should look at the example in the docs for a working example.

EDIT:

I'm not sure what your plot is intended to look like, so lets walk through the MPL demo.

Make the necessary imports and create an axis object (yours does this correctly):

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import matplotlib.pyplot as plt
import numpy as np

fig = plt.figure()
ax = fig.gca(projection='3d')

Next, make an X/Y grid and corresponding Z. In your program, X, Y and Z are 1D. They describe a line in 3D space, not a surface.

X = np.arange(-5, 5, 0.25)
Y = np.arange(-5, 5, 0.25)
X, Y = np.meshgrid(X, Y)  # <-- returns a 2D grid from initial 1D arrays
R = np.sqrt(X**2 + Y**2)
Z = np.sin(R)

Lets first plot the simplest thing possible. No colors, default anti-aliasing, lines, etc.

surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1)
plt.show()

Now add a colors. Note that the color comes from the Z component.

surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.jet)
plt.show()

Now manually control the colors (MPL inspiration).

colortuple = ('y', 'k')  # only use two colors: yellow and black
xlen, ylen = X.shape  # get length of 
colors = np.empty(X.shape, dtype=str)  # make a 2D array of strings
for i in range(xlen):
    for j in range(ylen):
        index = (i + j) % 2  # alternating 0's and 1's
        colors[i,j] = colortuple[index]
surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, 
        facecolors=colors)

If you want to color based on some other metric, you can create your own colormap. There are many answered questions on how to do that.

Edit 2:

Colors can also be specified as RGB sequences. For something like your red on X, green on Y description you could do this:

xlen, ylen = X.shape
colors = np.zeros((xlen,ylen,3))
jspan = np.linspace(0., 1., ylen)
ispan = np.linspace(0., 1., xlen)
for i in range(xlen):
    colors[i,:,0] = jspan
for j in range(ylen):
    colors[:,j,1] = ispan

surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, facecolors=colors,)