不同形状的 numpy 数组的矢量化函数

Question

我想将矢量化操作用于高效的代码，在这里我对不同大小的数组进行操作，因此需要对数组进行切片。但是，我想调整函数，以便我可以使用平面 numpy 数组以及具有 n 行的数组进行计算。我在这个例子中使用了简单的数字，但将在我的项目中使用随机数生成。目前我已经使用 if 语句解决了它，以根据形状区分切片。没有 if 语句有没有更有效的方法？

import numpy as np

def func(a, b, c, d):
    if len(a.shape) == 1:
        cf = -c
        cf[:len(b)] += a - b
        cf[-len(d)] -= d
    else:
        cf = -c
        cf[:, :b.shape[1]] += a - b
        cf[:, -d.shape[1]:] -= d
    return cf


a = np.array([1, 1, 1, 1, 1, 1, 1])
b = np.array([2, 2, 2, 2, 2, 2, 2])
c = np.array([3, 3, 3, 3, 3, 3, 3, 3])
d = np.array([4])
func(a, b, c, d)

n = 10
a = np.zeros((n, 7))+1
b = np.zeros((n, 7))+2
c = np.zeros((n, 8))+3
d = np.zeros((n, 1))+4
func(a, b, c, d)

可能是我这里没有找到通用的切片方法？！

Answer 1

此代码有效。

def func(a, b, c, d):
    cf = -c
    cf[..., :b.shape[-1]] += a - b
    cf[..., -d.shape[-1]:] -= d
    return cf

测试：

## Test 1
a = np.array([1, 1, 1, 1, 1, 1, 1])
b = np.array([2, 2, 2, 2, 2, 2, 2])
c = np.array([3, 3, 3, 3, 3, 3, 3, 3])
d = np.array([4])
func(a, b, c, d)
"""
array([-4, -4, -4, -4, -4, -4, -4, -7])
"""

## Test 2
n = 2
a = np.zeros((n, 7))+1
b = np.zeros((n, 7))+2
c = np.zeros((n, 8))+3
d = np.zeros((n, 1))+4
func(a, b, c, d)
"""
array([[-4., -4., -4., -4., -4., -4., -4., -7.],
       [-4., -4., -4., -4., -4., -4., -4., -7.]])
"""