Make sure data is compatible with Kriging model (#26)

* Make sure data is compatible with Kriging model

* Refactor predict e predictprob for Kriging

Author: juliohm

src/krig.jl

@@ -74,6 +74,12 @@ function initkrig(model::KrigingModel, data)
   ncon = nconstraints(model, nvar)
   nrow = nobs * nvar + ncon
 
+  # make sure data is compatible with model
+  nfeat = ncol(data) - 1
+  if nfeat != nvar
+    throw(ArgumentError("$nfeat data column(s) provided to $nvar-variate Kriging model"))
+  end
+
   # pre-allocate memory for LHS
   LHS = Matrix{V}(undef, nrow, nrow)
 
@@ -168,10 +174,14 @@ end
 
 predict(fitted::FittedKriging, var::AbstractString, gₒ) = predict(fitted, Symbol(var), gₒ)
 
+predict(fitted::FittedKriging, var::Symbol, gₒ) = predict(fitted, (var,), gₒ) |> first
+
 predict(fitted::FittedKriging, vars, gₒ) = predictmean(fitted, weights(fitted, gₒ), vars)
 
 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, vars, gₒ)
   w = weights(fitted, gₒ)
   μ = predictmean(fitted, w, vars)
@@ -182,29 +192,15 @@ end
 
 predictmean(fitted::FittedKriging, weights::KrigingWeights, vars) = krigmean(fitted, weights, vars)
 
-predictmean(fitted::FittedKriging, weights::KrigingWeights, var::Symbol) = first(predictmean(fitted, weights, (var,)))
-
 function krigmean(fitted::FittedKriging, weights::KrigingWeights, vars)
   d = fitted.state.data
   λ = weights.λ
   k = size(λ, 2)
-  n = length(vars)
-
-  @assert (k == n || k == 1) "invalid number of variables for Kriging model"
 
   cols = Tables.columns(values(d))
-
-  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
-  else # k == 1
-    @inbounds map(1:n) do p
-      λₚ = @view λ[:, 1]
+  @inbounds map(1:k) do j
+    sum(1:k) do p
+      λₚ = @view λ[p:k:end, j]
       zₚ = Tables.getcolumn(cols, vars[p])
       sum(i -> λₚ[i] ⦿ zₚ[i], eachindex(λₚ, zₚ))
     end
@@ -223,10 +219,7 @@ function predictvar(fitted::FittedKriging, weights::KrigingWeights, gₒ)
   σ² = krigvar(fun, weights, RHS, gₒ)
 
   # treat numerical issues
-  σ²₊ = max.(zero(σ²), σ²)
-
-  # treat scalar case
-  length(σ²₊) == 1 ? first(σ²₊) : σ²₊
+  max.(zero(σ²), σ²)
 end
 
 krigvar(fun::Variogram, weights::KrigingWeights, RHS, gₒ) = covvar(fun, weights, RHS, gₒ)

src/krig/simple.jl

@@ -37,23 +37,11 @@ function krigmean(fitted::FittedKriging{<:SimpleKriging}, weights::KrigingWeight
   μ = fitted.model.mean
   λ = weights.λ
   k = size(λ, 2)
-  n = length(vars)
-
-  @assert (k == n || k == 1) "invalid number of variables for Kriging model"
 
   cols = Tables.columns(values(d))
-
-  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
-  else # k == 1
-    @inbounds map(1:n) do p
-      λₚ = @view λ[:, 1]
+  @inbounds map(1:k) do j
+    sum(1:k) do p
+      λₚ = @view λ[p:k:end, j]
       zₚ = Tables.getcolumn(cols, vars[p])
       μ[p] + sum(i -> λₚ[i] ⦿ (zₚ[i] - μ[p]), eachindex(λₚ, zₚ))
     end

test/misc.jl

@@ -14,21 +14,6 @@
   vpred = GeoStatsModels.fitpredict(IDW(), vgtb, pset, neighbors=false)
   @test vpred == pred
 
-  # 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(SphericalVariogram(range=35.0)), gtb, grid, maxneighbors=3)
-  inds = LinearIndices(size(grid))
-  @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 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))