Add median_heuristic_transform (#245)

* Add `median_heuristic_transform` * Move `median_heuristic_transform` to convenience fcns in the docs * Only compute pairwise distances between different elements * Update src/transform/scaletransform.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Fix tests * Update test/transform/scaletransform.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Update Project.toml Co-authored-by: Théo Galy-Fajou <[email protected]> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
JuliaGaussianProcesses · Apr 12, 2022 · 7143f87 · 7143f87 · devmotion · Apr 12, 2022
1 parent 99b53c6
commit 7143f87
Show file tree

Hide file tree

Showing 8 changed files with 61 additions and 1 deletion.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.10.34"
+version = "0.10.35"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -15,6 +15,7 @@ LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 TensorCore = "62fd8b95-f654-4bbd-a8a5-9c27f68ccd50"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

diff --git a/docs/Project.toml b/docs/Project.toml
@@ -1,10 +1,13 @@
 [deps]
+Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
 Kronecker = "2c470bb0-bcc8-11e8-3dad-c9649493f05e"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
+Distances = "0.10"
 Documenter = "0.27"
 KernelFunctions = "0.10"
 Kronecker = "0.4, 0.5"

diff --git a/docs/src/transform.md b/docs/src/transform.md
@@ -41,4 +41,5 @@ PeriodicTransform
 
 ```@docs
 with_lengthscale
+median_heuristic_transform
 ```
diff --git a/src/KernelFunctions.jl b/src/KernelFunctions.jl
@@ -30,6 +30,8 @@ export Transform,
     PeriodicTransform
 export with_lengthscale
 
+export median_heuristic_transform
+
 export NystromFact, nystrom
 
 export gaborkernel
@@ -63,6 +65,8 @@ using ZygoteRules: ZygoteRules, AContext, literal_getproperty, literal_getfield
 # Hack to work around Zygote type inference problems.
 const Distances_pairwise = Distances.pairwise
 
+using Statistics: median!
+
 abstract type Kernel end
 abstract type SimpleKernel <: Kernel end
 

diff --git a/src/transform/scaletransform.jl b/src/transform/scaletransform.jl
@@ -33,3 +33,39 @@ _map(t::ScaleTransform, x::RowVecs) = RowVecs(only(t.s) .* x.X)
 Base.isequal(t::ScaleTransform, t2::ScaleTransform) = isequal(only(t.s), only(t2.s))
 
 Base.show(io::IO, t::ScaleTransform) = print(io, "Scale Transform (s = ", only(t.s), ")")
+
+# Helpers
+
+"""
+    median_heuristic_transform(distance, x::AbstractVector)
+
+Create a [`ScaleTransform`](@ref) that divides the input elementwise by the median
+`distance` of the data points in `x`.
+
+The `distance` has to support pairwise evaluation with `KernelFunctions.pairwise`. All
+`PreMetric`s of the package [Distances.jl](https://github.com/JuliaStats/Distances.jl) such
+as `Euclidean` satisfy this requirement automatically.
+
+# Examples
+
+```jldoctest
+julia> using Distances, Statistics
+
+julia> x = ColVecs(rand(100, 10));
+
+julia> t = median_heuristic_transform(Euclidean(), x);
+
+julia> y = map(t, x);
+
+julia> median(euclidean(y[i], y[j]) for i in 1:10, j in 1:10 if i != j) ≈ 1
+true
+```
+"""
+function median_heuristic_transform(f, x::AbstractVector)
+    # Compute pairwise distances between **different** elements
+    n = length(x)
+    distances = vec(pairwise(f, x))
+    deleteat!(distances, 1:(n + 1):(n^2))
+
+    return ScaleTransform(inv(median!(distances)))
+end
diff --git a/test/Project.toml b/test/Project.toml
@@ -13,6 +13,7 @@ PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -9,6 +9,7 @@ using LogExpFunctions
 using PDMats
 using Random
 using SpecialFunctions
+using Statistics
 using Test
 using Zygote: Zygote
 using ForwardDiff: ForwardDiff

diff --git a/test/transform/scaletransform.jl b/test/transform/scaletransform.jl
@@ -20,4 +20,17 @@
     @test isequal(ScaleTransform(s), ScaleTransform(s))
     @test repr(t) == "Scale Transform (s = $(s2))"
     test_ADs(x -> SEKernel() ∘ ScaleTransform(exp(x[1])), randn(rng, 1))
+
+    @testset "median heuristic" begin
+        for x in (x, XV, XC, XR), dist in (Euclidean(), Cityblock())
+            n = length(x)
+            t = median_heuristic_transform(dist, x)
+            @test t isa ScaleTransform
+            @test first(t.s) ≈
+                inv(median(dist(x[i], x[j]) for i in 1:n, j in 1:n if i != j))
+
+            y = map(t, x)
+            @test median(dist(y[i], y[j]) for i in 1:n, j in 1:n if i != j) ≈ 1
+        end
+    end
 end