Python替换嵌套的for循环。

huangapple
go评论93阅读模式
英文:

Python replace nested for loop

问题

以下是已翻译的代码部分：

  1. import itertools
  2. import numpy as np
  4. t = np.random.rand(3,3,3) 
  6. def foo(T):    
  8.     res = np.zeros((3,3,3))
  9.     
  10.     for a in range(3):
  11.         for b in range(3):
  12.             for c in range(3):
  13.                 idx = [a,b,c]
  14.                 combinations = list(itertools.permutations(idx, len(idx)))
  15.                 idx_arrays = tuple(np.array(idx) for idx in zip(*combinations))     
  16.                 res[a, b, c] = (1.0/len(combinations))*np.sum(T[idx_arrays])
  17.         
  18.         
  19.     return res
  21. sol = foo(t)

要使代码更快，您想要替换外部的for循环。是否可以实现这一点？理想情况下，这应该在不使用numpy的einsum方法的情况下完成。

英文:

I have the following code:

  1. import itertools
  2. import numpy as np
  4. t = np.random.rand(3,3,3) 
  6. def foo(T):    
  8.     res = np.zeros((3,3,3))
  9.     
  10.     for a in range(3):
  11.         for b in range(3):
  12.             for c in range(3):
  13.                 idx = [a,b,c]
  14.                 combinations = list(itertools.permutations(idx, len(idx)))
  15.                 idx_arrays = tuple(np.array(idx) for idx in zip(*combinations))     
  16.                 res[a, b, c] = (1.0/len(combinations))*np.sum(T[idx_arrays])
  17.         
  18.         
  19.     return res
  21. sol = foo(t)

The code above should be the equivalent to doing:

  1. for a in range(3):
  2.         for b in range(3):
  3.             for c in range(3):
  4.                 res2[a,b,c] = (1.0/6)*(
  5.                                         t[a,b,c]
  6.                                         + t[a,c,b]
  7.                                         + t[b,a,c]
  8.                                         + t[b,c,a]
  9.                                         + t[c, a, b]
  10.                                         + t[c,b,a]
  11.                                         )

To make the code faster, I would like to replace the outer for loops. Can this be done? Ideally this should be achieved without using einsum method of numpy.

答案1

得分: 1

以下是翻译好的代码部分：

  1. def foo(t):
  2.     res = np.zeros_like(t).astype(t.dtype)
  3.     A, B, C = t.shape
  4.     for a in range(A):
  5.         for b in range(a, B):
  6.             for c in range(b, C):
  7.                 r = (1.0 / 6) * (
  8.                     t[a, b, c]
  9.                     + t[a, c, b]
  10.                     + t[b, a, c]
  11.                     + t[b, c, a]
  12.                     + t[c, a, b]
  13.                     + t[c, b, a]
  14.                 )
  15.                 res[a, b, c] = r
  16.                 res[a, c, b] = r
  17.                 res[b, a, c] = r
  18.                 res[b, c, a] = r
  19.                 res[c, a, b] = r
  20.                 res[c, b, a] = r
  21.     return res
英文:

You can start improving by not doing duplicated computations:

  1. def foo(t):
  2.     res = np.zeros_like(t).astype(t.dtype)
  3.     A, B, C = t.shape
  4.     for a in range(A):
  5.         for b in range(a, B):
  6.             for c in range(b, C):
  7.                 r = (1.0 / 6) * (
  8.                     t[a, b, c]
  9.                     + t[a, c, b]
  10.                     + t[b, a, c]
  11.                     + t[b, c, a]
  12.                     + t[c, a, b]
  13.                     + t[c, b, a]
  14.                 )
  15.                 res[a, b, c] = r
  16.                 res[a, c, b] = r
  17.                 res[b, a, c] = r
  18.                 res[b, c, a] = r
  19.                 res[c, a, b] = r
  20.                 res[c, b, a] = r
  21.     return res

答案2

得分: 1

A one-liner, using np.transpose to vectorize:

  1. np.stack([t.transpose(i) for i in itertools.permutations(np.arange(3), 3)]).mean(0)

Those 3s can be changed or functionalized if you want higher dimensionality than 3

  1. import numpy as np
  2. from itertools import permutations as perm
  4. import numpy as np
  5. from itertools import permutations as perm
  7. def foo(T):
  8.     dims = len(T.shape)
  9.     assert np.all(np.array(T.shape) == T.shape[0])
  10.     all_perms = np.stack([T.transpose(i) for i in perm(np.arange(dims), dims)])
  11.     return all_perms.mean(0)
英文:

A one-liner, using np.transpose to vectorize:

  1. np.stack([t.transpose(i) for i in itertools.permutations(np.arange(3), 3)]).mean(0)

Those 3s can be changed or functionalized if you want higher dimensionality than 3

import numpy as np
from itertools import permutations as perm

  1. import numpy as np
  2. from itertools import permutations as perm
  4. def foo(T):
  5.     dims = len(T.shape)
  6.     assert np.all(np.array(T.shape) == T.shape[0])
  7.     all_perms = np.stack([T.transpose(i) for i in perm(np.arange(dims), dims)])
  8.     return all_perms.mean(0)

答案3

得分: 0

我不知道这是否比三重嵌套循环更快。无论如何，你可以尝试这样做：

  1. import itertools
  2. import numpy as np
  4. t = np.random.rand(3,3,3)
  6. def foo(T):    
  7.     t = np.random.rand(3,3,3) 
  8.     res = np.zeros((3,3,3))
  9.     
  10.     for a, b, c in set(itertools.permutations((0, 0, 0, 1, 1, 1, 2, 2, 2), 3)):
  11.         idx = [a,b,c]
  12.         combinations = list(itertools.permutations(idx, len(idx)))
  13.         idx_arrays = tuple(np.array(idx) for idx in zip(*combinations))     
  14.         res[a, b, c] = (1.0/len(combinations))*np.sum(t[idx_arrays])
  15.         
  16.         
  17.     return res
  19. sol = foo(t)
英文:

I don't know whether this is faster than 3 nested for loops. Anyways you can try this:

  1. import itertools
  2. import numpy as np
  4. t = np.random.rand(3,3,3)
  6. def foo(T):    
  7.     t = np.random.rand(3,3,3) 
  8.     res = np.zeros((3,3,3))
  9.     
  10.     for a, b, c in set(itertools.permutations((0, 0, 0, 1, 1, 1, 2, 2, 2), 3)):
  11.         idx = [a,b,c]
  12.         combinations = list(itertools.permutations(idx, len(idx)))
  13.         idx_arrays = tuple(np.array(idx) for idx in zip(*combinations))     
  14.         res[a, b, c] = (1.0/len(combinations))*np.sum(t[idx_arrays])
  15.         
  16.         
  17.     return res
  19. sol = foo(t)

Here, instead of using 3 nested for loops I am creating unique permutations of 3 elements from the tuple (0, 0, 0, 1, 1, 1, 2, 2, 2).

答案4

得分: 0

这段代码的大致含义是：通过改变np.indices的输出形状来生成idx_arrays，然后使用高级索引从原始数组中汇总值，并除以排列组合的数量。最后，将结果重新形状为一个3x3x3的数组。

这段代码以更高效的方式产生与原始代码相同的结果，利用了NumPy的向量化操作。

在我的机器上，foo(t) 耗时约261微秒，foo2(t) 耗时约36.7微秒。

英文:

What about this solution?

  1. def foo2(T):
  2.     indices = np.indices((3, 3, 3)).reshape(3, -1)
  3.     all_permutations = np.array(list(itertools.permutations(indices, len(indices))))
  4.     idx_arrays = tuple(np.array(idx) for idx in zip(*all_permutations))
  5.     summed = np.sum(T[tuple(idx_arrays)], axis=0)
  6.     res = (1.0 / len(all_permutations)) * summed
  7.     return res.reshape(3, 3, 3)

In short, you reshape the output of np.indices to generate the idx_arrays. Then you sum the values from the original array using advanced indexing and divide by the number of permutations. Finally, reshape the result back to a 3x3x3 array.

This code produces the same results as the original one but in a more efficient manner by leveraging the power of NumPy's vectorized operations.

On my machine:

  1. %%timeit foo(t)
  2. %%timeit foo2(t)
  3. 261 µs ± 4.36 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
  4. 36.7 µs ± 296 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

huangapple
  • 本文由 发表于 2023年5月11日 17:31:47
  • 转载请务必保留本文链接：https://go.coder-hub.com/76226149.html
  • numpy
  • python
匿名

发表评论

匿名网友

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen:

确定