在一个函数（能量轴和时间轴）中使用两个不同长度的数组进行广播以构建二维数组

Question

我有一个将两个轴作为输入的函数 - 一个时间轴和一个能量轴 - 我想弄清楚是否有办法通过广播而不是在for循环其中一个轴。

函数看起来是这样的，也包含在我的代码示例中：diffusion equation

这是我认为是我尝试过的天真的方法，两个轴具有不同的长度并在 for 循环中处理时间数组：

import numpy as np

def function(zeta, tau, alpha):
    return 1/np.sqrt(2*np.pi*alpha**2*tau)*np.exp(-zeta**2/(2*alpha**2*tau))

zeta = np.linspace(-2, 2, 100)
tau  = np.logspace(1e1, 1e9, 200)
alpha = 1e-6

res = np.zeros((len(tau), len(zeta)))
for i, t in enumerate(tau):
    res[i, :] = function(zeta, t, alpha)

但我想做的是：

res = function(zeta, tau, alpha)

这给出了（我猜是预期的）错误：

return 1/np.sqrt(2*np.pi*alpha**2*tau)*np.exp(-zeta**2/(2*alpha**2*tau))
ValueError: operands could not be broadcast together with shapes (100,) (200,)

那么有没有办法在 zeta 和 tau 之间同时广播该函数，并加快构建 100x200 数组 res？

Answer 1

首先，对 np.logspace 的评论：我相信您可能想使用 tau = np.logspace(1, 9, 200) 而不是 tau = np.logspace(1e1, 1e9, 200)。文档说：

In linear space, the sequence starts at base ** start (base to the power of start) and ends with base ** stop

base的默认值为10.0，因此1和9的start和stop将从[=24=得到一个序列] 到 10.0 ** 9 我相信这可能是你的意图。

关于2D广播的问题，我想这可能是你要找的：

X, Y = np.meshgrid(zeta, tau)
result = function(X, Y, alpha)

这个想法来自对类似问题的回答。

UPDATE: This should give the same result and save some space:
X, Y = np.meshgrid(zeta, tau, sparse=True)
From the docs:

sparse coordinate grids are intended to be use with Broadcasting. When all coordinates are used in an expression, broadcasting still leads to a fully-dimensonal result array.

import numpy as np
def function(zeta, tau, alpha):
    return 1/np.sqrt(2*np.pi*alpha**2*tau)*np.exp(-zeta**2/(2*alpha**2*tau))

zeta = np.linspace(-2, 2, 100)
tau  = np.logspace(1, 9, 200)
alpha = 1e-6

X, Y = np.meshgrid(zeta, tau)
result = function(X, Y, alpha)
[print(v, eval(v+'.shape')) for v in ['zeta', 'tau', 'result']]
print(f"\nresult:\n{result}")

res = np.zeros((len(tau), len(zeta)))
for i, t in enumerate(tau):
    res[i, :] = function(zeta, t, alpha)
print(f"\nres.shape {res.shape}")
print(f"\nres:\n{res}")

print(f"\nallclose()? {np.allclose(result, res)}")

输出：

zeta (100,)
tau (200,)
result (200, 100)

result:
[[0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 ...
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]]

res.shape (200, 100)

res:
[[0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 ...
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]]

allclose()? True

性能： 为您的原始方法计时，上面的方法显示：

Timeit results:
foo_1 ran in 0.001397762499982491 seconds using 1000 iterations
foo_2 ran in 0.004021247100026813 seconds using 1000 iterations

所以 meshgrid 看起来您的输入速度大约快 3 倍。

完整的测试代码在这里：

import numpy as np
def function(zeta, tau, alpha):
    return 1/np.sqrt(2*np.pi*alpha**2*tau)*np.exp(-zeta**2/(2*alpha**2*tau))

zeta = np.linspace(-2, 2, 100)
tau  = np.logspace(1, 9, 200)
alpha = 1e-6

from timeit import timeit

def foo_1(zeta, tau, alpha):
    X, Y = np.meshgrid(zeta, tau)
    result = function(X, Y, alpha)
    return result
def foo_2(zeta, tau, alpha):
    res = np.zeros((len(tau), len(zeta)))
    for i, t in enumerate(tau):
        res[i, :] = function(zeta, t, alpha)
    return res

n = 1000
print(f'Timeit results:')
for foo in ['foo_' + str(i + 1) for i in range(2)]:
    t = timeit(f"{foo}(zeta, tau, alpha)", setup=f"from __main__ import zeta, tau, alpha, {foo}", number=n) / n
    print(f'{foo} ran in {t} seconds using {n} iterations')

在一个函数（能量轴和时间轴）中使用两个不同长度的数组进行广播以构建二维数组

Broadcasting with two arrays of different lengths in a function (energy and time axis) to build a 2D array

python

axes

multidimensional-array

array-broadcasting