Refactor pairwise for multivariate case (#51)

* Refactor pairwise for multivariate case

* Add pairwise tests with transiograms

* Reduce allocations in symmetric entries

* Leverage symmetry of variogram in pairwise only

* Add more tests for pairwise! with Variogram

src/theoretical.jl

@@ -83,84 +83,81 @@ function (f::GeoStatsFunction)(g₁::Geometry, g₂::Geometry)
   mean(f(p₁, p₂) for p₁ in s₁, p₂ in s₂)
 end
 
-"""
-    returntype(f, g₁, g₂)
-
-Return type of f(g₁, g₂).
-"""
-returntype(f::GeoStatsFunction, g₁, g₂) = typeof(f(g₁, g₂))
-
 """
     pairwise(f, domain)
 
-Evaluate geostatistical function `f` between all elements in the `domain`.
+Evaluate geostatistical function `f` between all elements of the `domain`.
 """
 function pairwise(f::GeoStatsFunction, domain)
-  g = first(domain)
-  n = length(domain)
-  T = returntype(f, g, g)
-  F = Matrix{T}(undef, n, n)
+  T, (m, n, k) = matrixparams(f, domain)
+  F = Matrix{T}(undef, (m * k, n * k))
   pairwise!(F, f, domain)
 end
 
-"""
-    pairwise!(F, f, domain)
-
-Evaluates geostatistical function `f` between all elements in the `domain` in-place, filling the matrix `F`.
-"""
-function pairwise!(F, f::GeoStatsFunction, domain)
-  n = length(domain)
-  @inbounds for j in 1:n
-    gⱼ = domain[j]
-    sⱼ = _sample(f, gⱼ)
-    for i in (j + 1):n
-      gᵢ = domain[i]
-      sᵢ = _sample(f, gᵢ)
-      F[i, j] = mean(f(pᵢ, pⱼ) for pᵢ in sᵢ, pⱼ in sⱼ)
-    end
-    F[j, j] = mean(f(pⱼ, pⱼ) for pⱼ in sⱼ, pⱼ in sⱼ)
-    for i in 1:(j - 1)
-      F[i, j] = F[j, i] # leverage the symmetry
-    end
-  end
-  F
-end
-
 """
     pairwise(f, domain₁, domain₂)
 
 Evaluate geostatistical function `f` between all elements of `domain₁` and `domain₂`.
 """
 function pairwise(f::GeoStatsFunction, domain₁, domain₂)
-  g₁ = first(domain₁)
-  g₂ = first(domain₂)
-  m = length(domain₁)
-  n = length(domain₂)
-  T = returntype(f, g₁, g₂)
-  F = Matrix{T}(undef, m, n)
+  T, (m, n, k) = matrixparams(f, domain₁, domain₂)
+  F = Matrix{T}(undef, (m * k, n * k))
   pairwise!(F, f, domain₁, domain₂)
 end
 
+"""
+    pairwise!(F, f, domain)
+
+Evaluates geostatistical function `f` between all elements in the `domain` in-place, filling the matrix `F`.
+"""
+pairwise!(F, f::GeoStatsFunction, domain) = pairwise!(F, f, domain, domain)
+
 """
     pairwise!(F, f, domain₁, domain₂)
 
 Evaluates geostatistical function `f` between all elements of `domain₁` and `domain₂` in-place, filling the matrix `F`.
 """
 function pairwise!(F, f::GeoStatsFunction, domain₁, domain₂)
-  m = length(domain₁)
-  n = length(domain₂)
+  _, (m, n, k) = matrixparams(f, domain₁, domain₂)
   @inbounds for j in 1:n
     gⱼ = domain₂[j]
     sⱼ = _sample(f, gⱼ)
     for i in 1:m
       gᵢ = domain₁[i]
       sᵢ = _sample(f, gᵢ)
-      F[i, j] = mean(f(pᵢ, pⱼ) for pᵢ in sᵢ, pⱼ in sⱼ)
+      M = mean(f(pᵢ, pⱼ) for pᵢ in sᵢ, pⱼ in sⱼ)
+      F[((i - 1) * k + 1):(i * k), ((j - 1) * k + 1):(j * k)] .= M
     end
   end
   F
 end
 
+"""
+    matrixparams(f, domain)
+
+Return the parameters used to assemble the matrix for the
+geostatistical function `f` between all elements in the `domain`.
+"""
+matrixparams(f::GeoStatsFunction, domain) = matrixparams(f, domain, domain)
+
+"""
+    matrixparams(f, domain₁, domain₂)
+
+Return the parameters used to assemble the matrix for the
+geostatistical function `f` between all elements in `domain₁`
+and `domain₂`.
+"""
+function matrixparams(f::GeoStatsFunction, domain₁, domain₂)
+  m = length(domain₁)
+  n = length(domain₂)
+  g₁ = first(domain₁)
+  g₂ = first(domain₂)
+  R = f(g₁, g₂)
+  k = size(R, 1)
+  T = eltype(R)
+  T, (m, n, k)
+end
+
 # -----------
 # IO METHODS
 # -----------

src/theoretical/variogram.jl

@@ -63,6 +63,34 @@ function structures(γ::Variogram)
   cₒ, (c,), (γ,)
 end
 
+# leverage symmetry of variograms
+function pairwise!(Γ, γ::Variogram, domain)
+  _, (_, n, k) = matrixparams(γ, domain, domain)
+  @inbounds for j in 1:n
+    gⱼ = domain[j]
+    sⱼ = _sample(γ, gⱼ)
+    # lower triangular entries
+    for i in (j + 1):n
+      gᵢ = domain[i]
+      sᵢ = _sample(γ, gᵢ)
+      M = mean(γ(pᵢ, pⱼ) for pᵢ in sᵢ, pⱼ in sⱼ)
+      Γ[((i - 1) * k + 1):(i * k), ((j - 1) * k + 1):(j * k)] .= M
+    end
+    # diagonal entries
+    M = mean(γ(pⱼ, pⱼ) for pⱼ in sⱼ, pⱼ in sⱼ)
+    Γ[((j - 1) * k + 1):(j * k), ((j - 1) * k + 1):(j * k)] .= M
+  end
+
+  # upper triangular entries
+  @inbounds for j in 1:n*k
+    for i in 1:(j - 1)
+      Γ[i, j] = Γ[j, i]
+    end
+  end
+
+  Γ
+end
+
 # ----------------
 # IMPLEMENTATIONS
 # ----------------

test/theoretical/transiogram.jl

@@ -21,6 +21,13 @@
   for t in ts
     @test t(1000.0) ≈ [0.5 0.5; 0.5 0.5]
   end
+
+  # pairwise evaluation
+  ps = [Point(0, 0), Point(1, 1), Point(2, 2)]
+  for t in ts
+    T = GeoStatsFunctions.pairwise(t, ps)
+    @test size(T) == (6, 6)
+  end
 end
 
 @testset "MatrixExponentialTransiogram" begin
@@ -49,14 +56,28 @@ end
   A = rand(3, 2)
   R = A ./ sum(A, dims=2)
   t = @test_throws ArgumentError MatrixExponentialTransiogram(R)
+
+  # pairwise evaluation
+  ps = [Point(0, 0), Point(1, 1), Point(2, 2)]
+  A = rand(3, 3)
+  R = A ./ sum(A, dims=2)
+  t = MatrixExponentialTransiogram(R)
+  T = GeoStatsFunctions.pairwise(t, ps)
+  @test size(T) == (9, 9)
 end
 
 @testset "PiecewiseLinearTransiogram" begin
+  # basic tests
   csv = CSV.File(joinpath(datadir, "facies5.csv"))
   gtb = georef(csv, ("X", "Y", "Z"))
   t = EmpiricalTransiogram(gtb, "FACIES", maxlag=20, nlags=20)
   τ = PiecewiseLinearTransiogram(t.abscissas, t.ordinates)
   @test τ(0.0u"m") == I(5)
   @test all(x -> 0 < x < 1, τ(5.0u"m"))
   @test all(allequal, eachcol(τ(100.0u"m")))
+
+  # pairwise evaluation
+  ps = [Point(0, 0, 0), Point(1, 1, 1), Point(2, 2, 2)]
+  T = GeoStatsFunctions.pairwise(τ, ps)
+  @test size(T) == (15, 15)
 end

test/theoretical/variogram.jl

@@ -169,6 +169,14 @@
   @test size(Γ) == (3, 4)
   @test all(Γ .> 0)
 
+  # non-allocating pairwise!
+  Γ = rand(100, 100)
+  γ = GaussianVariogram()
+  𝒫 = rand(Point, 100)
+  GeoStatsFunctions.pairwise!(Γ, γ, 𝒫)
+  @test (@allocated GeoStatsFunctions.pairwise!(Γ, γ, 𝒫)) == 0
+  @test issymmetric(Γ)
+
   # constructor
   for γ in [
     CircularVariogram(),
@@ -368,9 +376,9 @@ end
   # result type is defined for nested models
   # see https://github.com/JuliaEarth/GeoStats.jl/issues/121 
   γ = GaussianVariogram() + ExponentialVariogram()
-  @test GeoStatsFunctions.returntype(γ, Point(0.0, 0.0, 0.0), Point(0.0, 0.0, 0.0)) == Float64
+  @test typeof(γ(Point(0.0, 0.0, 0.0), Point(0.0, 0.0, 0.0))) == Float64
   γ = GaussianVariogram(sill=1.0f0, range=1.0f0, nugget=0.1f0)
-  @test GeoStatsFunctions.returntype(γ, Point(0.0f0, 0.0f0, 0.0f0), Point(0.0f0, 0.0f0, 0.0f0)) == Float32
+  @test typeof(γ(Point(0.0f0, 0.0f0, 0.0f0), Point(0.0f0, 0.0f0, 0.0f0))) == Float32
 
   # nested model with matrix coefficients
   C₁ = [1.0 0.5; 0.5 2.0]