Python 优化嵌套 for 循环中的整形操作

Question

我正在寻求帮助以找到更多pythonic/broadcasting 方法来优化以下两个数组整形函数：

import numpy

def A_reshape(k,m,A):
    """
    Reshaping input float ndarray A of shape (x,y) 
    to output array A_r of shape (k,y,m)
    where k,m are user known dimensions
    """
    if A.ndim == 1: # in case A is flat make it (len(A),1)
        A = A.reshape((len(A),1))
    y = A.shape[1]
    A_r = np.zeros((k,y,m))
    for i in range(0,y,1):
        u = A[:,i].reshape((k,m))
        for j in range(0,m,1):
            A_r[:,i,j] = u[:,j]
    return A_r

def B_reshape(n,m,B):
    """
    Reshaping input float ndarray B of shape (z,y) 
    to output array B_r of shape (n,y,m)
    where n,m are user known dimensions
    """
    if B.ndim == 1: # in case B is flat make it (len(A),1)
        B = B.reshape((len(B),1))
    y = B.shape[1]
    B_r = np.zeros((n,y,m))
    for i in range(0,y,1):
        v = B[:,i]
        for j in range(0,m,1):
            B_r[:,i,j] = v[j*n:(j+1)*n]
    return B_r

A 的形状可能是 (33,10)，B 的形状可能是 (192,10)，例如给定 k=11、n=64 和 m=3。

如果有任何建议可以提高我对 numpy 重塑技术的理解并避免使用 for 循环，我们将不胜感激。谢谢

Answer 1

尝试：

def A_reshape(k,m,A):
    A2 = A.reshape(k,m,-1)
    A2 = np.moveaxis(A2, 2, 1)
    return A2

假设A的形状是(x,y)。最初，展开第一个维度：

(x,y) -> (k,m,y)

接下来，大小为y的轴从位置2移动到位置1。

(k,m,y) -> (k,y,m)

B_reshape的情况比较复杂，因为维度变换是：

(x,y) -> (m,n,y) # not (n,m,y)

(m,n,y) -> (n,y,m) # m is moved to the end

密码是：

def B_reshape(n,m,B):
    B2 = B.reshape(m,n,-1)
    B2 = np.moveaxis(B2, 0, 2)
    return B2

Python 优化嵌套 for 循环中的整形操作

Python optimizing reshape operations in nested for loops

python

optimization

numpy

reshape

python-3.x