CODEWITHSUNDEEP
pythonnumpyfor-loopnumpy-ndarray
yankeefan11
yankeefan11

Reputation: 485

Python - For loops along grid diagonals

I have a grid which is x=np.linspace(-1,1,n); y=x. I perform some calculation for each x and y with a for loop 

for a,b in enumerate(x): 
   for c,d in enumrate(y): 
     functionstuff(x,y)

To save time, I only need to do the computations on the diamond within the diagonals defined between the points (-1,0)-(0,1)-(1,0)-(0,-1). So I am trying to figure out how to best do this. So for a given x, I want to go from +(1-|x|) to -(1-|x|).

So how would I go about this with adjusting the step size between the smaller ys (so that I am not doing x = -1, y = 0 n times).

So I can make an numpy.array object to loop over:

for ii,x1 in enumerate(x1):

    y = np.linspace(1-np.abs(x1),-(1-np.abs(x1)),2*round((1-np.abs(x1))/dy) + 1)     
    for jj,y1 in enumerate(y):
       f[ii][jj] = DoStuff(x1,y1)
       xplot[ii][jj] = x1
       yplot[ii][jj] = y1

When I want to look at plot_surface(xplot,yplot,f), I just get a line along a diagonal (my guess being that I had to define xplot and yplot as a zero array before looping over)

Upvotes: 2

Views: 106

Answers (1)

jyalim
jyalim

Reputation: 3419

In numpy it helps to take advantage of broadcasting when you can. In this case, you can set up a two-dimensional grid, X and Y and compute on it rather than looping over the elements of x and y:

import numpy as np
x = np.linspace(-1,1,N)
y = x.copy()
X,Y = np.meshgrid(x,y,indexing='ij')
Z = functionstuff(X,Y)

You'll find this much more computationally efficient.

If you want to restrict your domain, you can do so with masking:

MASK = (Y<=1+X) & (Y<=1-X) & (Y>=X-1) & (Y>=-X-1)
Z    = np.zeros(X.shape)
Z[MASK] = functionstuff(X[MASK],Y[MASK])

If functionstuff = lambda x,y: exp(-(x**2+y**2)/2), then no masking results in this first image, and masking results in the second:

No Masking Masking

Upvotes: 1

Related Questions