|
| 1 | +############################### matrix valued version ########################## |
| 2 | +# function BarnesHutFactorization(k::GradientKernel, x, y = x, D = nothing; |
| 3 | +# θ::Real = 1/4, leafsize::Int = BARNES_HUT_DEFAULT_LEAFSIZE) |
| 4 | +# xs = vector_of_static_vectors(x) |
| 5 | +# ys = x === y ? xs : vector_of_static_vectors(y) |
| 6 | +# Tree = BallTree(ys, leafsize = leafsize) |
| 7 | +# m = length(y) |
| 8 | +# XT, YT, KT, TT, DT, RT = typeof.((xs, ys, k, Tree, D, θ)) |
| 9 | +# # w = zeros(length(m)) |
| 10 | +# # i = zeros(Bool, m) |
| 11 | +# # WT, BT = typeof(w), typeof(i) |
| 12 | +# T = gramian_eltype(k, xs, ys) |
| 13 | +# F = BarnesHutFactorization{T, KT, XT, YT, TT, DT, RT}(k, xs, ys, Tree, D, θ) #, w, i) |
| 14 | +# end |
| 15 | + |
| 16 | +function LinearAlgebra.mul!(y::AbstractVector{<:Real}, F::BarnesHutFactorization{<:Any, <:GradientKernel}, |
| 17 | + x::AbstractVector{<:Real}, α::Real = 1, β::Real = 0) |
| 18 | + d = length(F.x[1]) |
| 19 | + X, Y = reshape(x, d, :), reshape(y, d, :) # converting vector of reals to vector of vectors |
| 20 | + xx, yy = [c for c in eachcol(X)], [c for c in eachcol(Y)] |
| 21 | + mul!(yy, F, xx, α, β) |
| 22 | + return y |
| 23 | +end |
| 24 | + |
| 25 | +function LinearAlgebra.mul!(y::AbstractVector{<:AbstractVector}, F::BarnesHutFactorization{<:Any, <:GradientKernel}, |
| 26 | + x::AbstractVector{<:AbstractVector}, α::Real = 1, β::Real = 0) |
| 27 | + taylor!(y, F, x, α, β) |
| 28 | +end |
| 29 | + |
| 30 | +# function Base.:*(F::BarnesHutFactorization{<:Number}, x::AbstractVector{<:Number}) |
| 31 | +# T = promote_type(eltype(F), eltype(x)) |
| 32 | +# y = zeros(T, size(F, 1)) |
| 33 | +# mul!(y, F, x) |
| 34 | +# end |
| 35 | + |
| 36 | +function taylor!(b::AbstractVector{<:AbstractVector}, F::BarnesHutFactorization{<:Any, <:GradientKernel}, |
| 37 | + w::AbstractVector{<:AbstractVector}, α::Number = 1, β::Number = 0, θ::Real = F.θ) |
| 38 | + size(F, 2) == length(w) || throw(DimensionMismatch("length of w does not match second dimension of F: $(length(w)) ≠ $(size(F, 2))")) |
| 39 | + # eltype(b) == promote_type(eltype(F), eltype(w)) || |
| 40 | + # throw(TypeError("eltype of target vector b not equal to eltype of matrix-vector product: $(eltype(b)) and $(promote_type(eltype(F), eltype(w)))")) |
| 41 | + f_orders(r²) = derivatives(F.k.k, r², 3) |
| 42 | + sums_w = node_sums(w, F.Tree) # IDEA: could pre-allocate, worth it? is several orders of magnitude less expensive than multiply |
| 43 | + sums_w_r = weighted_node_sums(adjoint.(w), adjoint.(F.y), F.Tree) # need sum of outer products of F.y and w |
| 44 | + centers = get_hyper_centers(F) |
| 45 | + @. sums_w_r -= sums_w * adjoint(centers) # need to center the moments |
| 46 | + Gijs = [F[1, 1] for _ in 1:Base.Threads.nthreads()] |
| 47 | + for i in eachindex(F.x) # exactly 4 * length(y) allocations? |
| 48 | + if β == 0 |
| 49 | + @. b[i] = 0 # this avoids trouble if b is initialized with NaN's, e.g. thorugh "similar" |
| 50 | + else |
| 51 | + @. b[i] *= β |
| 52 | + end |
| 53 | + Gij = Gijs[Base.Threads.threadid()] |
| 54 | + taylor_recursion!(b[i], Gij, 1, F.k, f_orders, F.x[i], F.y, |
| 55 | + w, sums_w, sums_w_r, θ, F.Tree, centers, α) # x[i] creates an allocation |
| 56 | + end |
| 57 | + if !isnothing(F.D) # if there is a diagonal correction, need to add it |
| 58 | + mul!(b, F.D, w, α, 1) |
| 59 | + end |
| 60 | + return b |
| 61 | +end |
| 62 | + |
| 63 | +# barnes hut recursion for matrix-valued kernels, could merge with scalar version |
| 64 | +# bi is target vector corresponding to input point xi |
| 65 | +# Gij is temporary storage for evaluation of k(xi, y[j]), important if it is matrix valued |
| 66 | +# to avoid allocations |
| 67 | +function taylor_recursion!(bi::AbstractVector, Gij, |
| 68 | + index, k::GradientKernel, f_orders, |
| 69 | + xi, y::AbstractVector, |
| 70 | + w::AbstractVector{<:AbstractVector}, |
| 71 | + sums_w::AbstractVector{<:AbstractVector}, |
| 72 | + sums_w_r::AbstractVector{<:AbstractMatrix}, |
| 73 | + θ::Real, T::BallTree, centers, α::Number) |
| 74 | + h = T.hyper_spheres[index] |
| 75 | + if isleaf(T.tree_data.n_internal_nodes, index) # do direct computation |
| 76 | + for i in get_leaf_range(T.tree_data, index) |
| 77 | + j = T.indices[i] |
| 78 | + # @time Gij = evaluate_block!(Gij, k, xi, y[j]) # k(xi, y[j]) |
| 79 | + Gij = IsotropicGradientKernelElement{eltype(bi)}(k.k, xi, y[j]) |
| 80 | + wj = w[j] |
| 81 | + mul!(bi, Gij, wj, α, 1) |
| 82 | + end |
| 83 | + return bi |
| 84 | + |
| 85 | + elseif h.r < θ * euclidean(xi, centers[index]) # compress |
| 86 | + S_index = sums_w_r[index] |
| 87 | + ri = difference(xi, centers[index]) |
| 88 | + sum_abs2_ri = sum(abs2, ri) |
| 89 | + # NOTE: this line is the only one that still allocates ~688 bytes |
| 90 | + f0, f1, f2, f3 = f_orders(sum_abs2_ri) # contains first and second order derivative evaluation |
| 91 | + |
| 92 | + # Gij = evaluate_block!(Gij, k, xi, centers[index]) # k(xi, centers_of_mass[index]) |
| 93 | + Gij = IsotropicGradientKernelElement{eltype(bi)}(k.k, xi, centers[index]) |
| 94 | + |
| 95 | + # zeroth order |
| 96 | + mul!(bi, Gij, sums_w[index], α, 1) # bi .+= α * k(xi, centers_of_mass[index]) * sums_w[index] |
| 97 | + # first order |
| 98 | + # this block has zero allocations |
| 99 | + mul!(bi, 2*f3*dot(ri, S_index, ri) + f2 * tr(S_index), ri, 4α, 1) |
| 100 | + mul!(bi, S_index, ri, 4α*f2, 1) |
| 101 | + mul!(bi, S_index', ri, 4α*f2, 1) |
| 102 | + return bi |
| 103 | + else # recurse NOTE: parallelizing here is not as efficient as parallelizing over target points |
| 104 | + taylor_recursion!(bi, Gij, getleft(index), k, f_orders, xi, y, w, sums_w, sums_w_r, θ, T, centers, α) |
| 105 | + taylor_recursion!(bi, Gij, getright(index), k, f_orders, xi, y, w, sums_w, sums_w_r, θ, T, centers, α) |
| 106 | + end |
| 107 | +end |
| 108 | + |
| 109 | +# function node_mapreduce(x::AbstractVector, w::AbstractVector, T::BallTree, index::Int = 1) |
| 110 | +# length(x) == 0 && return zero(eltype(x)) |
| 111 | +# sums = fill(w[1]*x[1]', length(T.hyper_spheres)) |
| 112 | +# node_outer_products!(sums, x, w, T) |
| 113 | +# end |
| 114 | +# |
| 115 | +# # NOTE: x should either be vector of numbers or vector of static arrays |
| 116 | +# function node_mapreduce!(sums::AbstractVector{<:AbstractMatrix}, x::AbstractVector, |
| 117 | +# w::AbstractVector{<:Number}, T::BallTree, index::Int = 1) |
| 118 | +# if isleaf(T.tree_data.n_internal_nodes, index) |
| 119 | +# i = get_leaf_range(T.tree_data, index) |
| 120 | +# wi, xi = @views w[T.indices[i]], x[T.indices[i]] |
| 121 | +# sums[index] = wi * xi' |
| 122 | +# # adjoint.(w)' * adjoint.(x) |
| 123 | +# else |
| 124 | +# task = @spawn weighted_node_sums!(sums, x, w, T, getleft(index)) |
| 125 | +# weighted_node_sums!(sums, x, w, T, getright(index)) |
| 126 | +# wait(task) |
| 127 | +# sums[index] = sums[getleft(index)] + sums[getright(index)] |
| 128 | +# end |
| 129 | +# return sums |
| 130 | +# end |
0 commit comments