working version of fast taylor decomposition for gradient kernels

Author: SebastianAment

src/CovarianceFunctions.jl

@@ -60,7 +60,7 @@ include("taylor.jl")
 include("gradient.jl")
 include("gradient_algebra.jl")
 include("hessian.jl")
-# include("taylor_gradient.jl") # fast MVM algorithm for isotropic GradientKernel Gramians
+include("taylor_gradient.jl") # fast MVM algorithm for isotropic GradientKernel Gramians
 include("separable.jl")
 
 end # CovarianceFunctions

src/barneshut.jl

@@ -30,7 +30,7 @@ function BarnesHutFactorization(k, x, y = x, D = nothing; θ::Real = 1/4, leafsi
     # w = zeros(length(m))
     # i = zeros(Bool, m)
     # WT, BT = typeof(w), typeof(i)
-    T = gramian_eltype(k, xs, ys)
+    T = gramian_eltype(k, xs[1], ys[1])
     BarnesHutFactorization{T, KT, XT, YT, TT, DT, RT}(k, xs, ys, Tree, D, θ) #, w, i)
 end
 function BarnesHutFactorization(G::Gramian, θ::Real = 1/2; leafsize::Int = BARNES_HUT_DEFAULT_LEAFSIZE)
@@ -49,7 +49,7 @@ function LinearAlgebra.mul!(y::AbstractVector, F::BarnesHutFactorization, x::Abs
         taylor!(y, F, x, α, β)
     end
 end
-function Base.:*(F::BarnesHutFactorization, x::AbstractVector)
+function Base.:*(F::BarnesHutFactorization{<:Number}, x::AbstractVector{<:Number})
     T = promote_type(eltype(F), eltype(x))
     y = zeros(T, size(F, 1))
     mul!(y, F, x)
@@ -148,45 +148,45 @@ end
 
 ############################# centers of mass ##################################
 # this is a weighted sum, could be generalized to incorporate node_sums
-function compute_centers_of_mass(x::AbstractVector, w::AbstractVector, T::BallTree)
+function compute_centers_of_mass(w::AbstractVector, x::AbstractVector, T::BallTree)
     D = eltype(x) <: StaticVector ? length(eltype(x)) : length(x[1]) # if x is static vector
     com = [zero(MVector{D, Float64}) for _ in 1:length(T.hyper_spheres)]
-    compute_centers_of_mass!(com, x, w, T)
+    compute_centers_of_mass!(com, w, x, T)
 end
 
 function compute_centers_of_mass(F::BarnesHutFactorization, w::AbstractVector)
-    compute_centers_of_mass(F.y, w, F.Tree)
+    compute_centers_of_mass(w, F.y, F.Tree)
 end
 
-function compute_centers_of_mass!(com::AbstractVector, x::AbstractVector, w::AbstractVector, T::BallTree)
+function compute_centers_of_mass!(com::AbstractVector, w::AbstractVector, x::AbstractVector, T::BallTree)
     abs_w = abs.(w)
-    weighted_node_sums!(com, x, abs_w, T)
+    weighted_node_sums!(com, abs_w, x, T)
     sum_w = node_sums(abs_w, T)
     ε = eps(eltype(w)) # ensuring division by zero it not a problem
     @. com ./= sum_w + ε
 end
 
-node_sums(x::AbstractVector, T::BallTree) = weighted_node_sums(x, Ones(length(x)), T)
+node_sums(x::AbstractVector, T::BallTree) = weighted_node_sums(Ones(length(x)), x, T)
 function node_sums!(sums, x::AbstractVector, T::BallTree)
-    weighted_node_sums!(sums, x, Ones(length(x)), T)
+    weighted_node_sums!(sums, Ones(length(x)), x, T)
 end
 
-function weighted_node_sums(x::AbstractVector, w::AbstractVector, T::BallTree, index::Int = 1)
+function weighted_node_sums(w::AbstractVector, x::AbstractVector, T::BallTree, index::Int = 1)
     length(x) == 0 && return zero(eltype(x))
-    sums = zeros(typeof(w[1]'x[1]), length(T.hyper_spheres))
-    weighted_node_sums!(sums, x, w, T)
+    sums = fill(zero(w[1]'x[1]), length(T.hyper_spheres))
+    weighted_node_sums!(sums, w, x, T)
 end
 
 # NOTE: x should either be vector of numbers or vector of static arrays
-function weighted_node_sums!(sums::AbstractVector, x::AbstractVector,
-                            w::AbstractVector{<:Number}, T::BallTree, index::Int = 1)
+function weighted_node_sums!(sums::AbstractVector, w::AbstractVector,
+                            x::AbstractVector, T::BallTree, index::Int = 1)
     if isleaf(T.tree_data.n_internal_nodes, index)
         i = get_leaf_range(T.tree_data, index)
         wi, xi = @views w[T.indices[i]], x[T.indices[i]]
         sums[index] = wi'xi
     else
-        task = @spawn weighted_node_sums!(sums, x, w, T, getleft(index))
-        weighted_node_sums!(sums, x, w, T, getright(index))
+        task = @spawn weighted_node_sums!(sums, w, x, T, getleft(index))
+        weighted_node_sums!(sums, w, x, T, getright(index))
         wait(task)
         sums[index] = sums[getleft(index)] + sums[getright(index)]
     end

test/barneshut.jl

@@ -1,12 +1,14 @@
 module TestBarnesHut
 using LinearAlgebra
+using WoodburyFactorizations
 using CovarianceFunctions
 using CovarianceFunctions: BarnesHutFactorization, barneshut!, vector_of_static_vectors,
-        node_sums, euclidean, GradientKernel, taylor!
+        node_sums, euclidean, GradientKernel, taylor!, IsotropicGradientKernelElement
 using NearestNeighbors
 using NearestNeighbors: isleaf, getleft, getright, get_leaf_range
 using Test
 
+# like barnes hut but puts 0 as far field contribution
 function barneshut_no_far_field!(b::AbstractVector, F::BarnesHutFactorization, w::AbstractVector,
                     α::Number = 1, β::Number = 0, θ::Real = F.θ)
     D = length(eltype(F.x))
@@ -95,10 +97,9 @@ verbose = false
         nexp = 16
         err = zeros(nexp)
         err_no_split = zeros(nexp)
-        err_hyp = zeros(nexp)
-        err_hyp_no_split = zeros(nexp)
         err_nff = zeros(nexp)
         err_taylor = zeros(nexp)
+        err_taylor_hyp = zeros(nexp)
         theta_array = range(1e-1, 1, length = nexp)
         for (i, θ) in enumerate(theta_array)
             barneshut!(b_bh, F, w, 1, 0, θ, split = true)
@@ -110,8 +111,11 @@ verbose = false
             barneshut_no_far_field!(b_bh, F, w, 1, 0, θ) # compare against pseudo barnes hut where far field = 0
             err_nff[i] = norm(b - b_bh)
 
-            taylor!(b_bh, F, w, 1, 0, θ) # compare against pseudo barnes hut where far field = 0
+            taylor!(b_bh, F, w, 1, 0, θ, use_com = true)
             err_taylor[i] = norm(b - b_bh)
+
+            taylor!(b_bh, F, w, 1, 0, θ, use_com = false)
+            err_taylor_hyp[i] = norm(b - b_bh)
         end
 
         rel_err = err / norm(b)
@@ -122,28 +126,58 @@ verbose = false
         # using Plots
         # plotly()
         #
-        # rel_err_nff = err_nff / norm(b)
-        # rel_err_no_split = err_no_split / norm(b)
-        # rel_err_taylor = err_taylor / norm(b)
-        #
+        # norm_b = norm(b)
+        # rel_err_nff = err_nff / norm_b
+        # rel_err_no_split = err_no_split / norm_b
+        # rel_err_taylor = err_taylor / norm_b
+        # rel_err_taylor_hyp = err_taylor_hyp / norm_b
+        # #
         # plot(theta_array, rel_err, yscale = :log10, label = "barneshut", ylabel = "relative error", xlabel = "θ")
         # plot!(theta_array, rel_err_no_split, yscale = :log10, label = "no split")
         # plot!(theta_array, rel_err_nff, yscale = :log10, label = "sparse")
         # plot!(theta_array, rel_err_taylor, yscale = :log10, label = "taylor")
+        # plot!(theta_array, rel_err_taylor_hyp, yscale = :log10, label = "taylor hyper-sphere centers")
         # gui()
 
     end # testset weight vectors
 
-
     # @testset "gradient kernels" begin
-    #     k = CovarianceFunctions.Cauchy()
+    #
+    #     n = 1024
+    #     d = 2
+    #     x = randn(d, n)
+    #     # k = CovarianceFunctions.Cauchy()
+    #     k = CovarianceFunctions.EQ()
     #     g = CovarianceFunctions.GradientKernel(k)
     #
     #     F = BarnesHutFactorization(g, x)
     #     @test F isa BarnesHutFactorization
-    #     @test eltype(F) <: Diagonal
+    #     @test eltype(F) <: IsotropicGradientKernelElement
     #     @test size(F) == (n, n)
-    #     @test size(F[1, 1]) == (d_out, d_out)
+    #     @test size(F[1, 1]) == (d, d)
+    #
+    #     a = [randn(d) for _ in 1:n]
+    #     b = [zeros(d) for _ in 1:n]
+    #     # F*a
+    #     G = gramian(g, x)
+    #     b_truth = deepcopy(b)
+    #     # @time b_truth = G.A * a
+    #     # @time b_truth = G.A * a
+    #     mul!(b_truth, G.A, a)
+    #     norm_b = sqrt(sum(sum.(abs2, b_truth)))
+    #
+    #     α, β = 1, 0
+    #     θ = 0
+    #     taylor!(b, F, a, α, β, θ)
+    #     err = sqrt(sum(sum.(abs2, b - b_truth)))
+    #     rel_err = err / norm_b
+    #     @test rel_err < 1e-10
+    #
+    #     θ = 1/10
+    #     taylor!(b, F, a, α, β, θ)
+    #     err = sqrt(sum(sum.(abs2, b - b_truth)))
+    #     rel_err = err / norm_b
+    #     @test rel_err < 1e-3
     # end # testset matrix valued bh
 
 end # testset