NumPy：平衡不同深度的垃圾箱之间的水位

Question

我正在 NumPy 中实现模拟，作为模拟更新步骤的一部分，我需要解决以下问题。有一些水将被分配到一组不同深度的相连容器中。我需要找出要放入每个垃圾箱的水量，以使整体水位相同。

因此，例如，如果有 4 个深度分别为 1、2、3 和 5 的垃圾箱，我必须分配 5 个单位的水，我将分配 3 1/3 单位到深度为 5 的垃圾箱，1深度 3 bin 的 1/3 单位，深度 2 bin 的 1/3 单位。

因此，给定一组水槽深度和要分配的水量，我需要计算一个数组来描述要分配到每个水槽中的水量。

我需要一个非常高性能的解决方案，因为这个特定的例程在模拟过程中被调用了很多次。计算此函数可能是程序中最慢的部分。所以，我有点需要一个 NumPy 解决方案，出于性能原因尽可能避免 Python。

编辑：在实际问题中，总有少于 9 个箱子，但确切数量可能会有所不同。如果有帮助，我可以限制垃圾箱的深度。

有人有什么想法吗？

Answer 1

我的 numpy 实现工作正常。不过可以使用 numba 对其进行优化。

对于其他人的回答：有一个不便之处 - 当一些箱子具有相同的深度时，代码需要解决这个问题。

import numpy as np


def find_levels_numpy(bins, n_units):
    depth_diff = np.diff(bins)
    bin_bottom_level = np.concatenate([[0], np.cumsum(-depth_diff)])
    volume_for_bottom_bin = np.cumsum(np.concatenate([[0], -depth_diff]) * np.arange(len(bins)))
    # NOTE: results of the code above could be cached for given bins and reused for different n_units!

    idx_of_last_filled_bin = np.searchsorted(volume_for_bottom_bin, n_units) - 1
    levels = bin_bottom_level[idx_of_last_filled_bin] - bin_bottom_level
    remainder = n_units - volume_for_bottom_bin[idx_of_last_filled_bin]
    levels += remainder / (idx_of_last_filled_bin + 1)
    levels[idx_of_last_filled_bin+1:] = 0
    return levels


if __name__ == '__main__':
    # sorting bins for convenience
    bins = np.array([0.5, 5, 3, 2, 1, 1])
    n_units = 11
    sort_indices = np.argsort(bins)[::-1]
    bins = bins[sort_indices]
    level_by_bin = find_levels_numpy(bins, n_units)

    # undoing sorting
    level_by_bin[sort_indices] = level_by_bin[np.arange(len(sort_indices))]
    print(level_by_bin)
    print(np.sum(level_by_bin))

[0.25 4.75 2.75 1.75 0.75 0.75]
11.0

NumPy：平衡不同深度的垃圾箱之间的水位

NumPy: balancing water levels between bins of different depths

python

performance

numpy