Перепишите двойную петлю лучше и, возможно, короче

Question

Мне интересно, может ли следующий код быть написан несколько лучше.По сути, я хочу вычислить z = f(x, y) для (x, y) сетки.

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

Да. Ваш код, представленный в вопросе, хорош.

Никогда не думайте, что несколько строк - это "хорошо" или "круто". То, что имеет значение, является ясностью, удобочитаемостью и ремонтопригодностью. Другие люди должны понимать ваш код (и вы должны понимать его через 12 месяцев, когда вам нужно будет найти ошибку).

Многие программисты, особенно молодые, считают, что «умные» решения желательны. Они не. И это то, что так приятно с сообществом питонов. Мы гораздо меньше страдаем от этой ошибки, чем другие.

Answer 2

вы могли бы сделать что-то вроде

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

Answer 3

Вы можете использовать itertools 'product:

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

и если вы действительно хотите сократить эти 5 строк до 1, то:

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

Чтобы продвинуться еще дальше, если вам нужна функция xr и yr, где вы также можете установить диапазоны от 0 и 100 до чего-то еще:

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

Answer 4

Если вы установите все сразу, вы можете использовать понимание списка;

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

Если вам нужно использовать z, который уже есть, вы не сможете этого сделать, и ваш код будет примерно самым лучшим, что вы получите.

Дополнение: Я не знаю, с чем это lingrid делает, но если он генерирует список из 100 элементов, используйте понимание списка aaronasterling; нет смысла создавать дополнительный итератор, если вам это не нужно.

Answer 5

Ваш linspace на самом деле выглядит так, как будто это может быть np.linspace.Если это так, вы можете работать с массивами numpy без явной итерации:

z = f(x[:, np.newaxis], y)

Например:

>>> import numpy as np
>>> x = np.linspace(0, 9, 10)
>>> y = np.linspace(0, 90, 10)
>>> x[:, np.newaxis] + y  # or f(x[:, np.newaxis], y)
array([[  0.,  10.,  20.,  30.,  40.,  50.,  60.,  70.,  80.,  90.],
       [  1.,  11.,  21.,  31.,  41.,  51.,  61.,  71.,  81.,  91.],
       [  2.,  12.,  22.,  32.,  42.,  52.,  62.,  72.,  82.,  92.],
       [  3.,  13.,  23.,  33.,  43.,  53.,  63.,  73.,  83.,  93.],
       [  4.,  14.,  24.,  34.,  44.,  54.,  64.,  74.,  84.,  94.],
       [  5.,  15.,  25.,  35.,  45.,  55.,  65.,  75.,  85.,  95.],
       [  6.,  16.,  26.,  36.,  46.,  56.,  66.,  76.,  86.,  96.],
       [  7.,  17.,  27.,  37.,  47.,  57.,  67.,  77.,  87.,  97.],
       [  8.,  18.,  28.,  38.,  48.,  58.,  68.,  78.,  88.,  98.],
       [  9.,  19.,  29.,  39.,  49.,  59.,  69.,  79.,  89.,  99.]])

Но вы также можете использовать np.ogrid вместо двух linspace:

импортировать numpy как np

>>> x, y = np.ogrid[0:10, 0:100:10]
>>> x + y  # or f(x, y)
array([[ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 1, 11, 21, 31, 41, 51, 61, 71, 81, 91],
       [ 2, 12, 22, 32, 42, 52, 62, 72, 82, 92],
       [ 3, 13, 23, 33, 43, 53, 63, 73, 83, 93],
       [ 4, 14, 24, 34, 44, 54, 64, 74, 84, 94],
       [ 5, 15, 25, 35, 45, 55, 65, 75, 85, 95],
       [ 6, 16, 26, 36, 46, 56, 66, 76, 86, 96],
       [ 7, 17, 27, 37, 47, 57, 67, 77, 87, 97],
       [ 8, 18, 28, 38, 48, 58, 68, 78, 88, 98],
       [ 9, 19, 29, 39, 49, 59, 69, 79, 89, 99]])

Это в некоторой степени зависит от того, кем вы являетесь f.Если он содержит такие функции, как math.sin, вам необходимо заменить их на numpy.sin.

Если речь не идет о numpy, тогда вам следует придерживаться либо вашего варианта, либо, необязательно, использования enumerate при цикле:

for idx1, ai in enumerate(a):
   for idx2, bj in enumerate(b):
      z[idx1][idx2] = f(ai, bj)

Это имеет то преимущество, что вам не нужно жестко кодировать ваш range (или xrange) или использовать len(a) в качестве ввода.Но в целом, если нет большой разницы в производительности ¹ , используйте метод, который вы и другие, используя свой код, поймете легко.

---

¹ Если вашa и b равны numpy.array с, тогда будет существенная разница в производительности, поскольку numpy может обрабатывать массивы намного быстрее, если не требуется преобразований list <-> numpy.array.

Answer 6

Я думаю, что это код одной строки, который вы ищете

z =  [[a+b for b in linspace(0,yr,100)] for a in linspace(0,xr,100)]

Answer 7

Это показывает общий результат. a превращается в список длиной 6, а b - длиной 4. В результате получается список из 6 списков, и каждый вложенный список имеет длину 4 элемента.

>>> def f(x,y):
...     return x+y
... 
>>> a, b = list(range(0, 12, 2)), list(range(0, 12, 3))
>>> print len(a), len(b)
6 4
>>> result = [[f(aa, bb) for bb in b] for aa in a]
>>> print result
[[0, 3, 6, 9], [2, 5, 8, 11], [4, 7, 10, 13], [6, 9, 12, 15], [8, 11, 14, 17], [10, 13, 16, 19]]