Julia中的数组子集选择

Question

使用Julia语言，我定义了一个函数，使用拒绝抽样方法，在半径为3.14的球内均匀采样点，代码如下：

function spherical_sample(N::Int64)

    # generate N points uniformly distributed inside sphere 
    # using rejection sampling:

    points = pi*(2*rand(5*N,3).-1.0)

    ind = sum(points.^2,dims=2) .<= pi^2

    ## ideally I wouldn't have to do this: 
    ind_ = dropdims(ind,dims=2)

    return points[ind_,:][1:N,:]

end

我找到了一个有关数组子集的技巧：

ind = sum(points.^2,dims=2) .<= pi^2

## ideally I wouldn't have to do this: 
ind_ = dropdims(ind,dims=2)

但原则上数组索引应该只需要一行代码。我如何在Julia中更好地实现这个功能呢？

- Aidan Rocke

Answer 1

问题在于您正在创建一个二维索引向量。您可以通过使用eachrow来避免这个问题:

ind = sum.(eachrow(points.^2)) .<= pi^2

因此，您的完整回答将是：

function spherical_sample(N::Int64)
    points = pi*(2*rand(5*N,3).-1.0)
    ind = sum.(eachrow(points.^2)) .<= pi^2
    return points[ind,:][1:N,:]
end

- varelabuilds

Answer 2

下面是一个一行代码：

points[(sum(points.^2,dims=2) .<= pi^2)[:],:][1:N, :]

注意[:]会降低一个维度，这样BitArray就可以用于索引。

- Przemyslaw Szufel

---

在这个答案中，和David的类似，最好从ind中取前N个元素，如sum(x->x^2, points, dims=2) .<= pi^2)[1:N]（它还会丢弃一个维度并避免过多分配临时矩阵）。但仍然不能保证ind有N个元素，因此这通常会产生错误 - 请参见我的下面评论。 - Bogumił Kamiński

(sum(x->x^2, points, dims=2) .<= pi^2)[1:N] 并不能提供帮助，因为这个选择向量需要包含超过 N 个元素（而且有时仍可能失败）。 - Przemyslaw Szufel

同意 - 这就是我在评论中写的。 - Bogumił Kamiński

Answer 3

这并不是直接回答你的问题（因为你已经得到了两个建议），但我想给你一些提示，告诉你如何以更高效的方式实现整个过程。

第一点是避免生成5*N行的数据 - 问题在于很可能不足以生成N个有效样本。关键在于，你的模型中有效样本的概率约为50%，因此有可能没有足够的数据来选择，[1:N,:]的选择会导致错误。

以下是我会使用的代码，可以避免这个问题：

function spherical_sample(N::Integer) # no need to require Int64 only here
    points = 2 .* pi .* rand(N, 3) .- 1.0 # note that all operations are vectorized to avoid excessive allocations
    while N > 0 # we will run the code until we have N valid rows
        v = @view points[N, :] # use view to avoid allocating
        if sum(x -> x^2, v) <= pi^2 # sum accepts a transformation function as a first argument
            N -= 1 # row is valid - move to the previous one
        else
            rand!(v) # row is invalid - resample it in place
            @. v = 2 * pi * v - 1.0 # again - do the computation in place via broadcasting
        end
    end
    return points
end

- Bogumił Kamiński

---

你如何得到19.7%的概率？实际上，如果我只使用一行数据，则概率接近50％，因为我正在从[-pi，pi] ^ 3中抽样点，球体的体积为4/3 * pi * r ^ 3 ≈ 4 * r ^ 3 ，而立方体的体积为8 * r ^ 3，因此vol（sphere）/ vol（cube）≈4/8 = 50％。 - Aidan Rocke

没错 - 我一定是打错了什么。大约是50%左右。即使使用N=5，当我运行您的函数的@benchmark时，由于找不到足够的有效点的概率很高，我仍然经常会遇到错误。 - Bogumił Kamiński

在我的使用情况中，N始终大于或等于10，并且由于球体的体积/立方体的体积≈50％，我们可以将样本建模为二项式随机变量，并且在50次投掷中观察到少于10个正面的概率小于三百万分之一。 - Aidan Rocke

Answer 4

这个实现非常快，并且使用了StaticArrays。你可能也可以使用普通的元组来实现类似的东西：

using StaticArrays

function sphsample(N)
    T = SVector{3, Float64}
    v = Vector{T}(undef, N)
    n = 1
    while n <= N
        p = rand(T) .- 0.5
        @inbounds v[n] = p .* 2π
        n += (sum(abs2, p) <= 0.25)
    end
    return v
end

在我的笔记本电脑上，这种解决方案比带有“视图”功能的解决方案快约9倍。

- DNF