时间戳

首页 > 解决方案 > 使用 numpy 广播计算系列的部分总和

python - 使用 numpy 广播计算系列的部分总和

问题描述

我必须计算以下总和，S定义为：

5 项之和

我尝试了以下功能：

import numpy as np
a = np.array([0.4,2.6,3, 1.2, 3.4])
b = np.array([-5,7,2,1.1,1.8])
c = np.array([3.3,30,15,0.4,28])
t = np.linspace(0, np.pi, 600)
def S(t):
  return np.sum([a[i]*np.cos(b[i]*t + c[i]) for i in range(5)], axis=0)

这工作得很好。但我想知道是否有纯 numpy 版本使用广播而不依赖 Python 的列表理解？

我努力了：

def S(t):
  return np.sum(a*np.cos(b*t + c), axis=0)

当我计算时，S(t)我收到以下错误：

...
ValueError: operands could not be broadcast together with shapes (5,) (600,)

我怎样才能让它正常工作？

标签： pythonnumpyarray-broadcasting

解决方案

a我想我通过向、b和数组添加一个 newaxis 找到了答案，c如下所示：

def S(t):
   return np.sum(a[:,None]*cos(b[:,None]*t + c[:,None]), axis=0)

或者

def S(t):
   return np.sum(a[:,np.newaxis]*cos(b[:,np.newaxis]*t + c[:,np.newaxis]), axis=0)

编辑：

正如迈克尔在评论中所建议的那样，更好的解决方案是添加一个 newaxis 来t代替（请注意，在这种情况下，总和是在 上执行的axis=1）：

def S(t):
   return np.sum(a*cos(b*t[:,None] + c), axis=1)

推荐阅读