bottom up accumulation

[weymouth]

I've hacked together the first part of the Barnes-Hut approach, which is a method to accumulate leaf properties down the tree onto the nodes. I think that I've got it working, but my indexing is ugly.

1. Is there a way to link the leaves to their parent node? My method below seems safe...
2. I'm surely wasting time with all the calls to memory_index for the nodes. Should I be computing an offset instead?

using ImplicitBVH
using ImplicitBVH: BBox, BSphere

# Generate some simple bounding spheres
bounding_spheres = [
    BSphere{Float32}([1., 0., 0.], 1.),
    BSphere{Float32}([2., 0., 0.], 1.),
    BSphere{Float32}([3., 0., 0.], 1.),
    BSphere{Float32}([1., 1., 1.], 1.),
    BSphere{Float32}([2., 1., 1.], 1.),
    BSphere{Float32}([3., 1., 1.], 1.),
    BSphere{Float32}([1., 1., 2.], 1.),
    BSphere{Float32}([2., 1., 2.], 1.),
    BSphere{Float32}([3., 1., 2.], 1.),
]

# Build BVH
bvh = BVH(bounding_spheres)

# Accumulate on the nodes
using ImplicitBVH: level_indices,memory_index,pow2
function accumulate_up!(node_values, leaf_values, bvh)
    tree = bvh.tree; levels = tree.levels
    # leaf level 
    leaf = 0
    for i in range(level_indices(tree,levels-1)...)
        @inbounds node_values[i] = leaf_values[bvh.leaves[leaf+=1].index]
        leaf==length(leaf_values) && break
        @inbounds node_values[i] += leaf_values[bvh.leaves[leaf+=1].index]
    end

    # nodes levels
    for level in levels-2:-1:1
        for i in pow2(level-1):pow2(level)-1-bvh.skips[level]
            @inbounds node_values[memory_index(tree,i)] = node_values[memory_index(tree,2i)]
            ImplicitBVH.unsafe_isvirtual(tree,2i+1) && break
            @inbounds node_values[memory_index(tree,i)] += node_values[memory_index(tree,2i+1)]
        end
    end
    node_values
end

# Test it
leafmass = Float32.(1:length(bounding_spheres))
nodemass = zeros(Float32,length(bvh.nodes))
accumulate!(nodemass,leafmass,bvh)

[weymouth]

    # nodes levels
    for level in levels-2:-1:1
        skip = @inbounds bvh.skips[level+1]; offset = sum(@view bvh.skips[1:level])
        for i in pow2(level-1):pow2(level)-1-bvh.skips[level]
            @inbounds node_values[i-offset] = node_values[2i-offset-skip]
            unsafe_isvirtual(tree,2i+1) && continue
            @inbounds node_values[i-offset] += node_values[2i+1-offset-skip]
        end
    end

That's a bit faster I imagine.

[anicusan]

I see what you're doing now, I have some thoughts about using the BVH.skips and CPU / GPU parallelisation - I'll get back to you as soon as I can :)

[weymouth]

You're right that it would be a good idea to do these summations in parallel if possible! Happy to get your tips.

I've also coded up the function for the Barnes-Hut kernel evaluation. That one should be in serial since I'll be calling it in a parallel loop, but the stack I'm using for the depth-first-search might not be the best way to do it.

Merry Christmas!

using ImplicitBVH
using ImplicitBVH: BBox, BSphere

# Generate some simple bounding spheres
bounding_spheres = [
    BSphere{Float32}([1., 0., 0.], 1.),
    BSphere{Float32}([2., 0., 0.], 1.),
    BSphere{Float32}([3., 0., 0.], 1.),
    BSphere{Float32}([1., 1., 1.], 1.),
    BSphere{Float32}([2., 1., 1.], 1.),
    BSphere{Float32}([3., 1., 1.], 1.),
    BSphere{Float32}([1., 1., 2.], 1.),
    BSphere{Float32}([2., 1., 2.], 1.),
    BSphere{Float32}([3., 1., 2.], 1.),
]

# Build BVH
bvh = BVH(bounding_spheres)

# Accumulate leaf values onto nodes
using ImplicitBVH: level_indices,pow2,unsafe_isvirtual
function accumulate!(node_values, leaf_values, bvh)
    tree = bvh.tree; levels = tree.levels
    # leaf level
    leaf = 0
    for i in range(level_indices(tree,levels-1)...)
        @inbounds node_values[i] = leaf_values[bvh.leaves[leaf+=1].index]
        leaf==length(leaf_values) && break
        @inbounds node_values[i] += leaf_values[bvh.leaves[leaf+=1].index]
    end

    # node levels
    for level in levels-2:-1:1
        parent,child = @inbounds bvh.skips[level:level+1]
        for i in pow2(level-1):pow2(level) - 1 - (child-parent)
            @inbounds node_values[i-parent] = node_values[2i-child]
            unsafe_isvirtual(tree,2i+1) && continue
            @inbounds node_values[i-parent] += node_values[2i+1-child]
        end
    end
    node_values
end
accumulate(leaf_values::T,bvh) where T = accumulate!(T(undef,length(bvh.nodes)), leaf_values, bvh)

# Test it
leafm = Float32.(1:length(bounding_spheres))
nodem = accumulate(leafm,bvh)
using StaticArrays
leafx = [SA[bb.x...] for bb in bounding_spheres]
nodex = accumulate(leafm .* leafx,bvh) ./ nodem

# Relative squared-distance from bounding volumes
reldist(x,bb::BSphere) = max(sum(abs2,x .- bb.x)/bb.r^2-1,0)
function reldist(x,bb::BBox)
    c = (bb.up .+ bb.lo) ./ 2
    r = (bb.up .- bb.lo) ./ 2
    q = abs.(x .- c) .- r
    sum(abs2,max.(q,0))/sum(abs2,r)
end
reldist(x,bb::ImplicitBVH.BoundingVolume) = reldist(x,bb.volume)

# Barnes-Hut kernel evaluation!
using ImplicitBVH: memory_index
function evaluate(fnc,x,bvh,node_values,leaf_values;d²=1,stack=Vector{Int}(undef,bvh.tree.levels))
    tree = bvh.tree; length_nodes = length(bvh.nodes)
    top = 1; stack[top] = 1
    val = zero(fnc(x,node_values[1]))
    while top>0
        i = stack[top]; top-=1
        j = memory_index(tree,i)
        if j ≤ length_nodes
            if reldist(x,bvh.nodes[j])>d²
                val += fnc(x,node_values[j])
            else
                stack[top+=1] = 2i;
                !unsafe_isvirtual(tree,2i+1) && (stack[top+=1] = 2i+1)
            end
        else
            val += fnc(x,leaf_values[bvh.leaves[j-length_nodes].index])
        end
    end
    val
end

using TypedTables
nodeT = Table(m=nodem,x=nodex)
leafT = Table(m=leafm,x=leafx)
fnc(x,data) = data.m/sum(abs2,x-data.x)

x = zero(SVector{3,Float32})
evaluate(fnc,x,bvh,nodeT,leafT)

using ForwardDiff
dx = ForwardDiff.Dual.(x,ones(SVector{3,Float32}))
evaluate(fnc,dx,bvh,nodeT,leafT)