Question

I’m trying to plot a graph in matplotlib using numpy and meshgrid.

I want to make a region in the center of the array filled with ‘ones’ and elsewhere zeros. I tried setting up the array, but the for loop doesn’t seem to ever enter (print statements don’t print inside for loops). Any pointers?

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

Here is the code I think you want. You want to avoid loops and you can’t index that way (see the comments below).

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