Python优化嵌套for循环中的重塑操作

Question

我正在寻求帮助以找到一种更 Pythonic/广播方式来优化以下两个数组重塑函数：

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

例如，在 k=11、n=64 和 m=3 的情况下，A 的形状可能为 (33,10)，B 的形状可能为 (192,10)。

任何提高我对 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