Allow multivariate prediction with univariate Kriging

src/krig.jl

@@ -147,14 +147,24 @@ predictmean(fitted::FittedKriging, weights::KrigingWeights, var::Symbol) = first
 function krigmean(fitted::FittedKriging, weights::KrigingWeights, vars)
   d = fitted.state.data
   λ = weights.λ
-  k = length(vars)
+  k = size(λ, 2)
+  n = length(vars)
 
-  @assert size(λ, 2) == k "invalid number of variables for Kriging model"
+  @assert (k == n || k == 1) "invalid number of variables for Kriging model"
 
   cols = Tables.columns(values(d))
-  @inbounds map(1:k) do j
-    sum(1:k) do p
-      λₚ = @view λ[p:k:end, j]
+
+  if k == n
+    @inbounds map(1:k) do j
+      sum(1:n) do p
+        λₚ = @view λ[p:k:end, j]
+        zₚ = Tables.getcolumn(cols, vars[p])
+        sum(i -> λₚ[i] * zₚ[i], eachindex(λₚ, zₚ))
+      end
+    end
+  elseif k == 1
+    @inbounds map(1:n) do p
+      λₚ = @view λ[:, 1]
       zₚ = Tables.getcolumn(cols, vars[p])
       sum(i -> λₚ[i] * zₚ[i], eachindex(λₚ, zₚ))
     end

src/krig/simple.jl

@@ -36,14 +36,24 @@ function krigmean(fitted::FittedKriging{<:SimpleKriging}, weights::KrigingWeight
   d = fitted.state.data
   μ = fitted.model.mean
   λ = weights.λ
-  k = length(vars)
+  k = size(λ, 2)
+  n = length(vars)
 
-  @assert size(λ, 2) == k "invalid number of variables for Kriging model"
+  @assert (k == n || k == 1) "invalid number of variables for Kriging model"
 
   cols = Tables.columns(values(d))
-  @inbounds map(1:k) do j
-    sum(1:k) do p
-      λₚ = @view λ[p:k:end, j]
+
+  if k == n
+    @inbounds map(1:k) do j
+      sum(1:n) do p
+        λₚ = @view λ[p:k:end, j]
+        zₚ = Tables.getcolumn(cols, vars[p])
+        μ[p] + sum(i -> λₚ[i] * (zₚ[i] - μ[p]), eachindex(λₚ, zₚ))
+      end
+    end
+  elseif k == 1
+    @inbounds map(1:n) do p
+      λₚ = @view λ[:, 1]
       zₚ = Tables.getcolumn(cols, vars[p])
       μ[p] + sum(i -> λₚ[i] * (zₚ[i] - μ[p]), eachindex(λₚ, zₚ))
     end

test/misc.jl

@@ -14,16 +14,22 @@
   vpred = GeoStatsModels.fitpredict(IDW(), vgtb, pset, neighbors=false)
   @test vpred == pred
 
-  # fitpredict with Kriging
-  gtb = georef((; z=[1.0, 0.0, 1.0]), [(25.0, 25.0), (50.0, 75.0), (75.0, 50.0)])
+  # fitpredict with multiple variables and Kriging
+  gtb = georef((; a=[1.0, 0.0, 0.0], b=[0.0, 1.0, 0.0], c=[0.0, 0.0, 1.0]), [(25.0, 25.0), (50.0, 75.0), (75.0, 50.0)])
   grid = CartesianGrid((100, 100), (0.5, 0.5), (1.0, 1.0))
-  pred = GeoStatsModels.fitpredict(Kriging(GaussianVariogram(range=35.0)), gtb, grid, maxneighbors=3)
+  pred = GeoStatsModels.fitpredict(Kriging(SphericalVariogram(range=35.0)), gtb, grid, maxneighbors=3)
   inds = LinearIndices(size(grid))
-  @test isapprox(pred.z[inds[25, 25]], 1.0, atol=1e-3)
-  @test isapprox(pred.z[inds[50, 75]], 0.0, atol=1e-3)
-  @test isapprox(pred.z[inds[75, 50]], 1.0, atol=1e-3)
+  @test isapprox(pred.a[inds[25, 25]], 1.0, atol=1e-3)
+  @test isapprox(pred.a[inds[50, 75]], 0.0, atol=1e-3)
+  @test isapprox(pred.a[inds[75, 50]], 0.0, atol=1e-3)
+  @test isapprox(pred.b[inds[25, 25]], 0.0, atol=1e-3)
+  @test isapprox(pred.b[inds[50, 75]], 1.0, atol=1e-3)
+  @test isapprox(pred.b[inds[75, 50]], 0.0, atol=1e-3)
+  @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)
 
-  # fitpredict with CoKriging
+  # fitpredict with multiple variables and CoKriging
   gtb = georef((; a=[1.0, 0.0, 0.0], b=[0.0, 1.0, 0.0], c=[0.0, 0.0, 1.0]), [(25.0, 25.0), (50.0, 75.0), (75.0, 50.0)])
   grid = CartesianGrid((100, 100), (0.5, 0.5), (1.0, 1.0))
   model = Kriging([1.0 0.3 0.1; 0.3 1.0 0.2; 0.1 0.2 1.0] * SphericalVariogram(range=35.0))