moved specializations for BlockFactorizations into gramian file

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.7.2"
 manifest_format = "2.0"
-project_hash = "f2af4d41979a22255d208735c0b40625066093e6"
+project_hash = "f4736cace86674bc7adec0defb4812b15ace0d38"
 
 [[deps.AbstractFFTs]]
 deps = ["ChainRulesCore", "LinearAlgebra"]
@@ -27,9 +27,9 @@ uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
 [[deps.BlockFactorizations]]
 deps = ["LinearAlgebra", "Test"]
-git-tree-sha1 = "0abb1500c57376d219d58492c833de3c7a949b80"
+git-tree-sha1 = "78517ad2331c4c2f27ee4a4c363deb38626cd9eb"
 uuid = "5c499583-5bfe-4591-9b59-c1e192d48697"
-version = "1.2.0"
+version = "1.2.1"
 
 [[deps.ChainRulesCore]]
 deps = ["Compat", "LinearAlgebra", "SparseArrays"]
@@ -51,9 +51,9 @@ version = "0.3.0"
 
 [[deps.Compat]]
 deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
-git-tree-sha1 = "96b0bc6c52df76506efc8a441c6cf1adcb1babc4"
+git-tree-sha1 = "b153278a25dd42c65abbf4e62344f9d22e59191b"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
-version = "3.42.0"
+version = "3.43.0"
 
 [[deps.CompilerSupportLibraries_jll]]
 deps = ["Artifacts", "Libdl"]
@@ -192,9 +192,9 @@ uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [[deps.LogExpFunctions]]
 deps = ["ChainRulesCore", "ChangesOfVariables", "DocStringExtensions", "InverseFunctions", "IrrationalConstants", "LinearAlgebra"]
-git-tree-sha1 = "a7e100b068a6cbead98b9f4e5c8b488934b7aea0"
+git-tree-sha1 = "a970d55c2ad8084ca317a4658ba6ce99b7523571"
 uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
-version = "0.3.11"
+version = "0.3.12"
 
 [[deps.Logging]]
 uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
@@ -338,9 +338,9 @@ uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [[deps.StatsAPI]]
 deps = ["LinearAlgebra"]
-git-tree-sha1 = "c3d8ba7f3fa0625b062b82853a7d5229cb728b6b"
+git-tree-sha1 = "8d7530a38dbd2c397be7ddd01a424e4f411dcc41"
 uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0"
-version = "1.2.1"
+version = "1.2.2"
 
 [[deps.StatsBase]]
 deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"]

Project.toml

@@ -1,7 +1,7 @@
 name = "CovarianceFunctions"
 uuid = "b3329135-7958-41d4-bb02-e084c5a526bf"
 authors = ["sebastianament"]
-version = "0.2.3"
+version = "0.2.4"
 
 [deps]
 BlockFactorizations = "5c499583-5bfe-4591-9b59-c1e192d48697"
@@ -22,7 +22,7 @@ ToeplitzMatrices = "c751599d-da0a-543b-9d20-d0a503d91d24"
 WoodburyFactorizations = "9f1bac23-581c-4ebc-bd36-df60d764636d"
 
 [compat]
-BlockFactorizations = "1.2"
+BlockFactorizations = "1.2.1"
 DiffResults = "1.0"
 FillArrays = "0.12, 0.13"
 ForwardDiff = "0.10"

src/gradient.jl

@@ -21,9 +21,9 @@ function (G::GradientKernel)(x::AbstractVector, y::AbstractVector)
     gradient_kernel(G.k, x, y, input_trait(G.k))
 end
 
-# necessary for blockmul! of BlockFactorization
-function BlockFactorizations.evaluate_block!(K::AbstractMatOrFac, G::GradientKernel,
-            x::AbstractVector, y::AbstractVector, T::InputTrait = input_trait(G))
+# need this to work with BlockFactorizations, see blockmul! in gramian
+function evaluate_block!(K::AbstractMatOrFac, G::GradientKernel,
+            x::AbstractVector, y::AbstractVector, T::InputTrait = input_trait(G.k))
     gradient_kernel!(K, G.k, x, y, T)
 end
 
@@ -70,7 +70,9 @@ function allocate_gradient_kernel(k, x, y, T::IsotropicInput)
     r = reshape(x - y, :, 1) # do these work without reshape?
     kxy = k(x, y)
     d = length(x)
-    D = Diagonal(MVector{d, typeof(kxy)}(undef)) # not using Fill here, because value can't be changed
+    # IDEA: have mutable Fill, since this is all we need, will save further allocations, check mutable_fill.jl in doodles
+    # D = Diagonal(MVector{d, typeof(kxy)}(undef)) # not using Fill here, because value can't be changed
+    D = Diagonal(zeros(typeof(kxy), d)) # not using Fill here, because value can't be changed
     C = MMatrix{1, 1}(kxy)
     K = Woodbury(D, r, C, r')
 end
@@ -389,7 +391,8 @@ function value_gradient_kernel!(K::DerivativeKernelElement, k, x::AbstractVector
     return K
 end
 
-function BlockFactorizations.evaluate_block!(K::DerivativeKernelElement, G::ValueGradientKernel,
+# need this to work with BlockFactorizations, see blockmul! in gramian
+function evaluate_block!(K::DerivativeKernelElement, G::ValueGradientKernel,
             x::AbstractVector, y::AbstractVector, T::InputTrait = input_trait(G.k))
     value_gradient_kernel!(K, G.k, x, y, T)
 end

src/gramian.jl

@@ -85,21 +85,14 @@ Base.AbstractMatrix(G::Gramian) = Matrix(G)
 Base.adjoint(G::Gramian) = Gramian(G.k, G.y, G.x)
 Base.transpose(G::Gramian) = Gramian(G.k, G.y, G.x)
 
-# need this for blockmul! to work in BlockFactorization
-# specialization for Gramians of matrix-valued kernels
-# IDEA: precompute input_trait?
-function BlockFactorizations.evaluate_block!(Gij, G::Gramian{<:Any, <:MultiKernel}, i::Int, j::Int)
-    BlockFactorizations.evaluate_block!(Gij, G.k, G.x[i], G.y[j], input_trait(G.k))
-end
-
 # by default, Gramians of matrix-valued kernels are BlockFactorizations, O(1) memory complexity
 function gramian(k::MultiKernel, x::AbstractVector, y::AbstractVector, lazy::Val{true} = Val(true))
     G = Gramian(k, x, y)
     BlockFactorization(G, isstrided = true) # strided because every block has same size
 end
 
 # instantiates the blocks but respects structure O(n^2d) memory complexity for gradient kernel
-function gramian(k::MultiKernel, x::AbstractVector, y::AbstractVector, lazy::Val{false})
+function gramian(k::MultiKernel, x::AbstractVector, y::AbstractVector, lazy::Val{false}) # = Val(false))
     G = Gramian(k, x, y)
     G = Matrix(G)
     BlockFactorization(G, isstrided = true)
@@ -185,11 +178,37 @@ end
 #     return -1
 # end
 
+###################### Specializations for BlockGramians #######################
 function LinearAlgebra.:\(B::BlockGramian, b::AbstractVector)
     T = promote_type(eltype(B), eltype(b))
     x = zeros(T, size(B, 1))
     ldiv!(x, B, b)
 end
+# solve general BlockGramian via conjugate gradient solver
 function LinearAlgebra.ldiv!(x::AbstractVector, B::BlockGramian, b::AbstractVector)
     cg!(x, B, b)
 end
+
+# carries out multiplication for general BlockFactorization
+function BlockFactorizations.blockmul!(y::AbstractVecOfVecOrMat, G::Gramian, x::AbstractVecOfVecOrMat, α::Real = 1, β::Real = 0)
+    Gijs = [G[1, 1] for _ in 1:Base.Threads.nthreads()] # pre-allocate storage for elements
+    IT = input_trait(G.k)
+    @threads for i in eachindex(y)
+        @. y[i] = β * y[i]
+        Gij = Gijs[Base.Threads.threadid()]
+        for j in eachindex(x)
+            Gij = evaluate_block!(Gij, G, i, j, IT) # this is change to original blockmul!, allowing evaluation of Gramian's elements without additional allocations
+            mul!(y[i], Gij, x[j], α, 1)
+        end
+    end
+    return y
+end
+
+function evaluate_block!(Gij, G::Gramian, i::Int, j::Int, T = input_trait(G.k))
+    evaluate_block!(Gij, G.k, G.x[i], G.y[j], T)
+end
+
+# fallback
+function evaluate_block!(Gij, k, x, y, T = input_trait(G.k))
+    k(x, y)
+end

src/hessian.jl

@@ -64,7 +64,7 @@ end
 #     return K
 # end
 
-# function BlockFactorizations.evaluate_block!(K::HessianKernelElement, G::HessianKernel, x::AbstractVector, y::AbstractVector, T::InputTrait = input_trait(G))
+# function evaluate_block!(K::HessianKernelElement, G::HessianKernel, x::AbstractVector, y::AbstractVector, T::InputTrait = input_trait(G))
 #     hessian_kernel!(K, G.k, x, y, T)
 # end
 
@@ -118,8 +118,8 @@ function allocate_hessian_kernel(k, x, y, ::IsotropicInput)
     IsotropicHessianKernelElement(k, copy(x), copy(y))
 end
 
-function BlockFactorizations.evaluate_block!(K::IsotropicHessianKernelElement,
-    G::HessianKernel, x::AbstractVector, y::AbstractVector, ::IsotropicInput = IsotropicInput())
+function evaluate_block!(K::IsotropicHessianKernelElement, G::HessianKernel,
+    x::AbstractVector, y::AbstractVector, ::IsotropicInput = IsotropicInput())
     hessian_kernel!(K, G.k, x, y)
 end
 
@@ -223,8 +223,8 @@ function allocate_hessian_kernel(k, x, y, ::DotProductInput)
     DotProductHessianKernelElement(k, copy(x), copy(y))
 end
 # necessary for BlockFactorization
-function BlockFactorizations.evaluate_block!(K::DotProductHessianKernelElement,
-    G::HessianKernel, x::AbstractVector, y::AbstractVector, ::DotProductInput = DotProductInput())
+function evaluate_block!(K::DotProductHessianKernelElement, G::HessianKernel,
+    x::AbstractVector, y::AbstractVector, ::DotProductInput = DotProductInput())
     hessian_kernel!(K, G.k, x, y)
 end
 function hessian_kernel!(K::DotProductHessianKernelElement, k, x::AbstractVector, y::AbstractVector, ::DotProductInput = DotProductInput())
@@ -350,8 +350,8 @@ function value_gradient_hessian_kernel!(K::ValueGradientHessianKernelElement, k,
     return K
 end
 
-function BlockFactorizations.evaluate_block!(K::ValueGradientHessianKernelElement,
-    G::ValueGradientHessianKernel, x::AbstractVector, y::AbstractVector, T::InputTrait = input_trait(G.k))
+function evaluate_block!(K::ValueGradientHessianKernelElement, G::ValueGradientHessianKernel,
+        x::AbstractVector, y::AbstractVector, T::InputTrait = input_trait(G.k))
     value_gradient_hessian_kernel!(K, G.k, x, y, T)
 end