python - 根据最后一列累积 NumPy 数组的行

Question

我有以下问题。我有一个包含坐标数组的数组，前三个条目是 x、y、z 坐标，第四个条目是轨道的 id。我想在轨道上添加漂移，该漂移在第一个时间点之后开始。是否有一种简单的方法可以将漂移动态添加到带有 ID 的轨道(可以具有不同的长度)，立即添加到整个数组？ (所以你可以看到，id为3的轨迹只有3个坐标条目，而id为3的轨迹有6个)

import numpy as np
drift=np.array([1,1,0])
a = np.array([[1,1,1,0],[1,1,1,0],[1,1,1,0],
              [1,1,1,2],[1,1,1,2],[1,1,1,3],
              [1,1,1,3],[1,1,1,3],[1,1,1,3],
              [1,1,1,3],[1,1,1,3]])

输出:

output = np.array([[1,1,1,0],[2,2,1,0],[3,3,1,0],
                   [1,1,1,2],[2,2,1,2],[1,1,1,3],
                   [2,2,1,3],[3,3,1,3],[4,4,1,3],
                   [5,5,1,3],[6,6,1,3]])

Answer 1

以下是如何以矢量化方式完成此操作的示例:

import numpy as np


drift = np.array([1, 1, 0])
a = np.array([[1, 1, 1, 0], [1, 1, 1, 0], [1, 1, 1, 0], [1, 1, 1, 2], 
              [1, 1, 1, 2], [1, 1, 1, 3], [1, 1, 1, 3], [1, 1, 1, 3], 
              [1, 1, 1, 3], [1, 1, 1, 3], [1, 1, 1, 3]])


def multirange(counts: np.ndarray) -> np.ndarray:
    """
    Calculates concatenated ranges. Code was taken at:
    https://stackoverflow.com/questions/20027936/how-to-efficiently-concatenate-many-arange-calls-in-numpy
    """
    counts = counts[counts != 0]
    counts1 = counts[:-1]
    reset_index = np.cumsum(counts1)
    incr = np.ones(counts.sum(), dtype=int)
    incr[0] = 0
    incr[reset_index] = 1 - counts1
    incr.cumsum(out=incr)
    return incr


def drifts(ids: np.ndarray,
           drift: np.ndarray) -> np.ndarray:
    diffs = np.diff(ids)
    max_drifts_per_id = np.concatenate((np.where(diffs)[0], [len(ids) - 1])) + 1
    max_drifts_per_id[1:] = max_drifts_per_id[1:] - max_drifts_per_id[:-1]
    multipliers = multirange(max_drifts_per_id)
    drifts = np.tile(drift, (len(ids), 1))
    return drifts * multipliers[:, np.newaxis]


a[:, :-1] += drifts(a[:, -1], drift)
print(a)

输出:

array([[0, 0, 0, 0],
       [1, 1, 0, 0],
       [2, 2, 0, 0],
       [0, 0, 0, 2],
       [1, 1, 0, 2],
       [0, 0, 0, 3],
       [1, 1, 0, 3],
       [2, 2, 0, 3],
       [3, 3, 0, 3],
       [4, 4, 0, 3],
       [5, 5, 0, 3]])

<小时/>

说明:

drifts 函数的思想是获取一个 ids 数组(在我们的例子中，我们可以通过 a[:, -1] 的形式获取: array([0, 0, 0, 2, 2, 3, 3, 3, 3, 3, 3])) 和 drift (np.array([1 , 1, 0])) 获取以下数组，然后可以将其附加到原始数组:

array([[0, 0, 0],
       [1, 1, 0],
       [2, 2, 0],
       [0, 0, 0],
       [1, 1, 0],
       [0, 0, 0],
       [1, 1, 0],
       [2, 2, 0],
       [3, 3, 0],
       [4, 4, 0],
       [5, 5, 0]])

<小时/>

逐行:

diffs = np.diff(ids)

这里我们得到一个数组，其中所有非零元素都将具有第一个数组中最后一个 id 的索引:

array([0, 0, 2, 0, 1, 0, 0, 0, 0, 0]) 

参见np.diff了解详情。

<小时/>

max_drifts_per_id = np.concatenate((np.where(diffs)[0], [len(ids) - 1])) + 1

np.where(diffs)[0] 将给出前一个数组中那些非零元素的索引。我们附加最后一个元素的索引并将结果索引增加 1，以便稍后获取范围。请参阅np.where了解详情。连接后 max_drifts_per_id 将是:

array([ 3,  5, 11])

<小时/>

max_drifts_per_id[1:] = max_drifts_per_id[1:] - max_drifts_per_id[:-1]

从前面的结果中，我们得到了范围最终值的数组:

array([3, 2, 6])

<小时/>

multipliers = multirange(max_drifts_per_id)

我们使用 multirange 作为串联 np.arange 调用的有效替代方案。 。请参阅How to efficiently concatenate many arange calls in numpy?了解详情。结果乘数将是:

array([0, 1, 2, 0, 1, 0, 1, 2, 3, 4, 5])

<小时/>

drifts = np.tile(drift, (len(ids), 1))

作者:np.tile我们扩展 drift 使其具有与 ids 相同的行数:

array([[1, 1, 0],
       [1, 1, 0],
       [1, 1, 0],
       [1, 1, 0],
       [1, 1, 0],
       [1, 1, 0],
       [1, 1, 0],
       [1, 1, 0],
       [1, 1, 0],
       [1, 1, 0],
       [1, 1, 0]])

<小时/>

return drifts * multipliers[:, np.newaxis]

我们将其乘以乘数并得到:

array([[0, 0, 0],
       [1, 1, 0],
       [2, 2, 0],
       [0, 0, 0],
       [1, 1, 0],
       [0, 0, 0],
       [1, 1, 0],
       [2, 2, 0],
       [3, 3, 0],
       [4, 4, 0],
       [5, 5, 0]])

<小时/>

最后可以将返回值添加到原始数组中:

a[:, :-1] += drifts(a[:, -1], drift)