Return MvNormal in multivariate Kriging

Author: juliohm

src/krig.jl

@@ -178,14 +178,20 @@ predict(fitted::FittedKriging, vars, gₒ) = predictmean(fitted, weights(fitted,
 
 predictprob(fitted::FittedKriging, var::AbstractString, gₒ) = predictprob(fitted, Symbol(var), gₒ)
 
-predictprob(fitted::FittedKriging, var::Symbol, gₒ) = predictprob(fitted, (var,), gₒ) |> first
+function predictprob(fitted::FittedKriging, var::Symbol, gₒ)
+  w = weights(fitted, gₒ)
+  μ = predictmean(fitted, w, (var,)) |> first
+  σ² = predictvar(fitted, w, gₒ) |> first
+  # https://github.com/JuliaStats/Distributions.jl/issues/1413
+  Normal(ustrip.(μ), √σ²)
+end
 
 function predictprob(fitted::FittedKriging, vars, gₒ)
   w = weights(fitted, gₒ)
   μ = predictmean(fitted, w, vars)
-  σ² = predictvar(fitted, w, gₒ)
+  Σ = predictvar(fitted, w, gₒ)
   # https://github.com/JuliaStats/Distributions.jl/issues/1413
-  @. Normal(ustrip(μ), √σ²)
+  MvNormal(ustrip.(μ), Σ)
 end
 
 predictmean(fitted::FittedKriging, weights::KrigingWeights, vars) = krigmean(fitted, weights, vars)
@@ -213,11 +219,12 @@ function predictvar(fitted::FittedKriging, weights::KrigingWeights, gₒ)
   RHS = fitted.state.RHS
   fun = fitted.model.fun
 
-  # variance formula for given function
-  σ² = krigvar(fun, weights, RHS, gₒ)
+  # covariance formula for given function
+  Σ = krigvar(fun, weights, RHS, gₒ)
 
   # treat numerical issues
-  max.(zero(σ²), σ²)
+  ϵ = eltype(Σ)(1e-10)
+  Symmetric(Σ + ϵ * I)
 end
 
 function krigvar(fun::GeoStatsFunction, weights::KrigingWeights, RHS, gₒ)
@@ -230,7 +237,7 @@ function krigvar(fun::GeoStatsFunction, weights::KrigingWeights, RHS, gₒ)
   # compute cov(0) considering change of support
   Cₒ = ustrip.(covzero(fun, gₒ)) * I(k)
 
-  diag(Cₒ) - diag(Cλ) - diag(Cν)
+  Cₒ - Cλ - Cν
 end
 
 function krigvar(t::Transiogram, weights::KrigingWeights, RHS, gₒ)
@@ -252,7 +259,7 @@ function krigvar(t::Transiogram, weights::KrigingWeights, RHS, gₒ)
   Tₒ = t(gₒ, gₒ)
   Cₒ = @inbounds [p[i] * (Tₒ[i, j] - p[j]) for i in 1:k, j in 1:k]
 
-  diag(Cₒ) - diag(Cλ) - diag(Cν)
+  Cₒ - Cλ - Cν
 end
 
 function covzero(fun::GeoStatsFunction, gₒ)

src/models.jl

@@ -133,7 +133,7 @@ function fitpredictneigh(model, gtb, dom, path, point, prob, minneighbors, maxne
   inds = traverse(dom, path)
 
   # prediction function
-  predfun = prob ? predictprob : predict
+  predfun = prob ? _marginals ∘ predictprob : predict
 
   # predict variables
   cols = Tables.columns(values(gtb))
@@ -177,7 +177,7 @@ function fitpredictfull(model, gtb, dom, path, point, prob)
   inds = traverse(dom, path)
 
   # prediction function
-  predfun = prob ? predictprob : predict
+  predfun = prob ? _marginals ∘ predictprob : predict
 
   # fit model to data
   fmodel = fit(model, gtb)
@@ -197,6 +197,9 @@ function fitpredictfull(model, gtb, dom, path, point, prob)
   georef(predtab, dom)
 end
 
+_marginals(dist::UnivariateDistribution) = (dist,)
+_marginals(dist::MvNormal) = Normal.(mean(dist), var(dist))
+
 # ----------------
 # IMPLEMENTATIONS
 # ----------------

test/krig.jl

@@ -331,14 +331,14 @@
       okdist = GeoStatsModels.predictprob(ok, (:a, :b, :c), pset[i])
       ukdist = GeoStatsModels.predictprob(uk, (:a, :b, :c), pset[i])
       dkdist = GeoStatsModels.predictprob(dk, (:a, :b, :c), pset[i])
-      @test mean.(skdist) ≈ [j == i for j in 1:3]
-      @test mean.(okdist) ≈ [j == i for j in 1:3]
-      @test mean.(ukdist) ≈ [j == i for j in 1:3]
-      @test mean.(dkdist) ≈ [j == i for j in 1:3]
-      @test isapprox(var.(skdist), [0.0, 0.0, 0.0], atol=1e-10)
-      @test isapprox(var.(okdist), [0.0, 0.0, 0.0], atol=1e-10)
-      @test isapprox(var.(ukdist), [0.0, 0.0, 0.0], atol=1e-10)
-      @test isapprox(var.(dkdist), [0.0, 0.0, 0.0], atol=1e-10)
+      @test mean(skdist) ≈ [j == i for j in 1:3]
+      @test mean(okdist) ≈ [j == i for j in 1:3]
+      @test mean(ukdist) ≈ [j == i for j in 1:3]
+      @test mean(dkdist) ≈ [j == i for j in 1:3]
+      @test isapprox(var(skdist), [0.0, 0.0, 0.0], atol=1e-8)
+      @test isapprox(var(okdist), [0.0, 0.0, 0.0], atol=1e-8)
+      @test isapprox(var(ukdist), [0.0, 0.0, 0.0], atol=1e-8)
+      @test isapprox(var(dkdist), [0.0, 0.0, 0.0], atol=1e-8)
     end
 
     # predict on a specific point
@@ -347,14 +347,14 @@
     okdist = GeoStatsModels.predictprob(ok, (:a, :b, :c), pₒ)
     ukdist = GeoStatsModels.predictprob(uk, (:a, :b, :c), pₒ)
     dkdist = GeoStatsModels.predictprob(dk, (:a, :b, :c), pₒ)
-    @test all(μ -> 0 ≤ μ ≤ 1, mean.(skdist))
-    @test all(μ -> 0 ≤ μ ≤ 1, mean.(okdist))
-    @test all(μ -> 0 ≤ μ ≤ 1, mean.(ukdist))
-    @test all(μ -> 0 ≤ μ ≤ 1, mean.(dkdist))
-    @test all(≥(0), var.(skdist))
-    @test all(≥(0), var.(okdist))
-    @test all(≥(0), var.(ukdist))
-    @test all(≥(0), var.(dkdist))
+    @test all(μ -> 0 ≤ μ ≤ 1, mean(skdist))
+    @test all(μ -> 0 ≤ μ ≤ 1, mean(okdist))
+    @test all(μ -> 0 ≤ μ ≤ 1, mean(ukdist))
+    @test all(μ -> 0 ≤ μ ≤ 1, mean(dkdist))
+    @test all(≥(0), var(skdist))
+    @test all(≥(0), var(okdist))
+    @test all(≥(0), var(ukdist))
+    @test all(≥(0), var(dkdist))
   end
 
   @testset "Transiogram" begin
@@ -369,30 +369,26 @@
 
     # interpolation property
     for i in 1:3
-      skdist = GeoStatsModels.predictprob(sk, (:a, :b, :c), pset[i])
-      okdist = GeoStatsModels.predictprob(ok, (:a, :b, :c), pset[i])
-      ukdist = GeoStatsModels.predictprob(uk, (:a, :b, :c), pset[i])
-      dkdist = GeoStatsModels.predictprob(dk, (:a, :b, :c), pset[i])
-      @test mean.(skdist) ≈ [j == i for j in 1:3]
-      @test mean.(okdist) ≈ [j == i for j in 1:3]
-      @test mean.(ukdist) ≈ [j == i for j in 1:3]
-      @test mean.(dkdist) ≈ [j == i for j in 1:3]
+      skmean = GeoStatsModels.predict(sk, (:a, :b, :c), pset[i])
+      okmean = GeoStatsModels.predict(ok, (:a, :b, :c), pset[i])
+      ukmean = GeoStatsModels.predict(uk, (:a, :b, :c), pset[i])
+      dkmean = GeoStatsModels.predict(dk, (:a, :b, :c), pset[i])
+      @test skmean ≈ [j == i for j in 1:3]
+      @test okmean ≈ [j == i for j in 1:3]
+      @test ukmean ≈ [j == i for j in 1:3]
+      @test dkmean ≈ [j == i for j in 1:3]
     end
 
     # predict on a specific point
     pₒ = Point(50.0, 50.0)
-    skdist = GeoStatsModels.predictprob(sk, (:a, :b, :c), pₒ)
-    okdist = GeoStatsModels.predictprob(ok, (:a, :b, :c), pₒ)
-    ukdist = GeoStatsModels.predictprob(uk, (:a, :b, :c), pₒ)
-    dkdist = GeoStatsModels.predictprob(dk, (:a, :b, :c), pₒ)
-    @test all(μ -> 0 ≤ μ ≤ 1, mean.(skdist))
-    @test all(μ -> 0 ≤ μ ≤ 1, mean.(okdist))
-    @test all(μ -> 0 ≤ μ ≤ 1, mean.(ukdist))
-    @test all(μ -> 0 ≤ μ ≤ 1, mean.(dkdist))
-    @test all(≥(0), var.(skdist))
-    @test all(≥(0), var.(okdist))
-    @test all(≥(0), var.(ukdist))
-    @test all(≥(0), var.(dkdist))
+    skmean = GeoStatsModels.predict(sk, (:a, :b, :c), pₒ)
+    okmean = GeoStatsModels.predict(ok, (:a, :b, :c), pₒ)
+    ukmean = GeoStatsModels.predict(uk, (:a, :b, :c), pₒ)
+    dkmean = GeoStatsModels.predict(dk, (:a, :b, :c), pₒ)
+    @test all(μ -> 0 ≤ μ ≤ 1, skmean)
+    @test all(μ -> 0 ≤ μ ≤ 1, okmean)
+    @test all(μ -> 0 ≤ μ ≤ 1, ukmean)
+    @test all(μ -> 0 ≤ μ ≤ 1, dkmean)
   end
 
   @testset "Fallbacks" begin

test/misc.jl

@@ -28,4 +28,14 @@
   @test isapprox(pred.c[inds[25, 25]], 0.0, atol=1e-3)
   @test isapprox(pred.c[inds[50, 75]], 0.0, atol=1e-3)
   @test isapprox(pred.c[inds[75, 50]], 1.0, atol=1e-3)
+  pred = GeoStatsModels.fitpredict(model, gtb, grid, maxneighbors=3, prob=true)
+  @test isapprox(mean(pred.a[inds[25, 25]]), 1.0, atol=1e-3)
+  @test isapprox(mean(pred.a[inds[50, 75]]), 0.0, atol=1e-3)
+  @test isapprox(mean(pred.a[inds[75, 50]]), 0.0, atol=1e-3)
+  @test isapprox(mean(pred.b[inds[25, 25]]), 0.0, atol=1e-3)
+  @test isapprox(mean(pred.b[inds[50, 75]]), 1.0, atol=1e-3)
+  @test isapprox(mean(pred.b[inds[75, 50]]), 0.0, atol=1e-3)
+  @test isapprox(mean(pred.c[inds[25, 25]]), 0.0, atol=1e-3)
+  @test isapprox(mean(pred.c[inds[50, 75]]), 0.0, atol=1e-3)
+  @test isapprox(mean(pred.c[inds[75, 50]]), 1.0, atol=1e-3)
 end