Rewrite a double loop in a nicer and maybe shorter way

Question

I am wondering if the following code can be written in a somewhat nicer way. Basically, I want to calculate z = f(x, y) for a (x, y) meshgrid.

a = linspace(0, xr, 100)                                                                  
b = linspace(0, yr, 100)                                                                  

for i in xrange(100):
   for j in xrange(100):
      z[i][j] = f(a[i],b[j])

Answer 1

Yeah. Your code as presented in the question is nice.

Do not ever think that few lines is "nice" or "cool". What counts is clarity, readability and maintainability. Other people should be able to understand your code (and you should understand it in 12 months, when you need to find a bug).

Many programmers, especially young ones, think that "clever" solutions are desirable. They are not. And that's what is so nice with the python community. We are much less afflicted by that mistake than others.

Answer 2

you could do something like

z = [[f(item_a, item_b) for item_b in b] for item_a in a]

Answer 3

You could use itertools' product:

[f(i,j) for i,j in product( a, b )]

and if you really want to shrink those 5 lines into 1 then:

[f(i,j) for i,j in product( linspace(0,xr,100), linspace(0,yr,100)]

To take it even further if you want a function of xr and yr where you can also preset the ranges of 0 and 100 to something else:

def ranged_linspace( _start, _end, _function ):
    def output_z( xr, yr ):
        return [_function( i, j ) for i,j in product( linspace( _start, xr, _end ), linspace( _start, yr, _end ) )]
    return output_z