绘制3D网格图，如何访问网格图数组中的点？

Question

我正在尝试使用numpy和meshgrid在matplotlib中绘制一个图表。

我想在数组的中心区域填充'1'，而其他地方填充'0'。我试着设置这个数组，但for循环似乎从来没有进入（在for循环内的打印语句没有输出）。有什么建议吗？

import numpy as np
import pylab as py
from scipy import *
from numpy.fft import fft
import mpl_toolkits.mplot3d.axes3d as p3

def fft2dplot(
    type = 'rect', aperature = 16, method = 'transfer',
    wavelength = 1, distance = 1000000):
    dict = {
        'rect' : 'Rectangle', 'circ' : 'Circle', 
        'transfer' : 'Tranfer Function', 'integral' : 'Integral'}
    #Scale is not correct
    scale = aperature/(distance*wavelength) #in mm
    #Range for aperature, 
    x = y = arange(-aperature*8,aperature*8, 1)
    X,Y = np.meshgrid(x,y)
    print len(X)
    X = Y = Z = X*0

    #These following statements never enter (type == rect passes, but for loop don't)
    if type == 'rect':
        for u in x[-aperature/2:aperature/2]:
                for w in y[-aperature/2:aperature/2]:
                    Z[u,w] = 1
                    print "I'm here"

    fig = py.figure()
    ax = p3.Axes3D(fig)
    #ax.contour3D(X,Y,Z)
    ax.plot_wireframe(X, Y, Z, rstride=1, cstride=1)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    py.show()

Answer 1

这是我认为你想要的代码。你需要避免使用循环，而且你不能那样进行索引（请看下面的评论）。

import numpy as np
import pylab as py
from scipy import *  # it is better practice to do: from scipy import module1, module2, etc...
from numpy.fft import fft
import mpl_toolkits.mplot3d.axes3d as p3

def fft2dplot(
    type = 'rect', aperature = 16, method = 'transfer',
    wavelength = 1, distance = 1000000):
    # here you are overwriting the built-in 'dict' suggest using a different variable name
    dict = {
        'rect' : 'Rectangle', 'circ' : 'Circle', 
        'transfer' : 'Tranfer Function', 'integral' : 'Integral'}
    #Scale is not correct
    scale = aperature/(distance*wavelength) #in mm
    #Range for aperature, 
    x = y = arange(-aperature*8,aperature*8, 1)
    X,Y = np.meshgrid(x,y)
    print len(X)
    # you were losing references to the meshgrid arrays X and Y here.
    # I've provided an alternate method of getting an array of zeros
    # Z = np.zeros(X.shape) works as well
    #X = Y = Z = X*0
    Z = X.copy()
    Z[:,:] = 0


    # Indexing does not work like this.  Indices are unsigned integers 0 and above
    # or a boolean array the same shape as the indexed array.
    #These following statements never enter (type == rect passes, but for loop don't)
##    if type == 'rect':
##        for u in x[-aperature/2:aperature/2]:
##                for w in y[-aperature/2:aperature/2]:
##                    Z[u,w] = 1
##                    print "I'm here"

    # what you need is to define your indices based on your conditions.  
    # Here we create a boolean array of indices that indicate where we meet your
    # conditions.
    # Then we use that array to assign those elements of the Z array to the value 1
    idxs = (X>-aperature/2)&(X<aperature/2)&(Y>-aperature/2)&(Y<aperature/2)
    Z[idxs] = 1

    fig = py.figure()
    ax = p3.Axes3D(fig)
    #ax.contour3D(X,Y,Z)
    ax.plot_wireframe(X, Y, Z, rstride=1, cstride=1)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    py.show()