以不同方式切片 numpy 数组的不同行

Question

我正在研究蒙特卡洛辐射传输代码，该代码模拟通过介质发射光子并对其随机游走进行统计建模。它运行缓慢，一次发射一个光子，所以我想对其进行矢量化，一次运行可能 1000 个光子。

我已将光子通过的平板分成nlayers 光学深度0 和depth 之间的切片。实际上，这意味着我有nlayers + 2 区域（nlayers 加上平板上方的区域和平板下方的区域）。在每一步，我都必须跟踪每个光子通过的层。

假设我已经知道两个光子从第 0 层开始。一个迈出一步并在第 2 层结束，另一个迈出一步并在第 6 层结束。这由数组 pastpresent 表示看起来像这样：

[[ 0 2]
 [ 0 6]]

我想生成一个数组traveled_through，其中包含(nlayers + 2) 列和2 行，描述光子i 是否通过层j（包括端点）。它看起来像这样（nlayers = 10）：

[[ 1 1 1 0 0 0 0 0 0 0 0 0]
 [ 1 1 1 1 1 1 1 0 0 0 0 0]]

我可以通过迭代光子并单独生成traveled_through 的每一行来做到这一点，但这相当慢，并且有点破坏了一次运行许多光子的意义，所以我宁愿不这样做。

我尝试将数组定义如下：

traveled_through = np.zeros((2, nlayers)).astype(int)
traveled_through[ : , np.min(pastpresent, axis = 1) : np.max(pastpresent, axis = 1) + ] = 1

这个想法是，在给定光子的行中，从起始层到结束层的索引将设置为 1，所有其他保持为 0。但是，我收到以下错误：

traveled_through[ : , np.min(pastpresent, axis = 1) : np.max(pastpresent, axis = 1) + 1 ] = 1
IndexError: invalid slice

我最好的猜测是 numpy 不允许使用这种方法对数组的不同行进行不同的索引。有没有人建议如何为任意数量的光子和任意数量的层生成traveled_through？

Answer 1

如果两个光子总是从 0 开始，你也许可以按如下方式构建你的数组。

首先设置变量...

>>> pastpresent = np.array([[0, 2], [0, 6]])
>>> nlayers = 10

...然后构造数组：

>>> (pastpresent[:,1][:,np.newaxis] + 1 > np.arange(nlayers+2)).astype(int)
array([[1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0]])

或者如果光子具有任意起始层：

>>> pastpresent2 = np.array([[1, 7], [3, 9]])
>>> (pastpresent2[:,0][:,np.newaxis] < np.arange(nlayers+2)) & 
    (pastpresent2[:,1][:,np.newaxis] + 1 > np.arange(nlayers+2)).astype(int)       
array([[0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
       [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0]])

Answer 2

我有点喜欢这种事情的一个小技巧涉及logical_xor ufunc的accumulate方法：

>>> a = np.zeros(10, dtype=int)
>>> b = [3, 7]
>>> a[b] = 1
>>> a
array([0, 0, 0, 1, 0, 0, 0, 1, 0, 0])
>>> np.logical_xor.accumulate(a, out=a)
array([0, 0, 0, 1, 1, 1, 1, 0, 0, 0])

请注意，这会将1 中的位置之间的条目设置为b，第一个索引包括在内，最后一个索引不包括在内，因此您必须根据您的具体目标处理 1 个错误。

有几行，你可以让它工作：

>>> a = np.zeros((3, 10), dtype=int)
>>> b = np.array([[1, 7], [0, 4], [3, 8]])
>>> b[:, 1] += 1  # handle the off by 1 error
>>> a[np.arange(len(b))[:, None], b] = 1
>>> a
array([[0, 1, 0, 0, 0, 0, 0, 0, 1, 0],
       [1, 0, 0, 0, 0, 1, 0, 0, 0, 0],
       [0, 0, 0, 1, 0, 0, 0, 0, 0, 1]])
>>> np.logical_xor.accumulate(a, axis=1, out=a)
array([[0, 1, 1, 1, 1, 1, 1, 1, 0, 0],
       [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
       [0, 0, 0, 1, 1, 1, 1, 1, 1, 0]])