numba_array_impl.py 文件源码

def potential_numba_array(cluster):
    d = distances_numba_array(cluster)
    # Original: dtri = np.triu(d)
    # np.triu is not supported; so write my own loop to clear the
    # lower triangle
    for i in range(d.shape[0]):
        for j in range(d.shape[1]):
            if i > j:
                d[i, j] = 0
    # Original: lj_numba_array(d[d > 1e-6]).sum()
    # d[d > 1e-6] is not supported due to the indexing with boolean
    # array.  Replace with custom loop.
    energy = 0.0
    for v in d.flat:
        if v > 1e-6:
            energy += lj_numba_array(v)
    return energy