Generalize KNNRegressor to multitarget case

src/NearestNeighbors.jl

@@ -128,10 +128,11 @@ end
 function MMI.predict(m::KNNRegressor, (tree, y, w), X)
     Xmatrix     = MMI.matrix(X, transpose=true) # NOTE: copies the data
     idxs, dists = NN.knn(tree, Xmatrix, m.K)
+    return _predict(m, y, idxs, dists, w)
+end
+function _predict(m::KNNRegressor, y::AbstractVector, idxs, dists, w)
     preds       = zeros(length(idxs))
-
     w_ = ones(m.K)
-
     for i in eachindex(idxs)
         idxs_  = idxs[i]
         dists_ = dists[i]
@@ -148,6 +149,26 @@ function MMI.predict(m::KNNRegressor, (tree, y, w), X)
     return preds
 end
 
+function _predict(m::KNNRegressor, y, idxs, dists, w)
+    ymat    = MMI.matrix(y)
+    preds   = similar(ymat, length(idxs), size(ymat, 2))
+    w_ = ones(m.K)
+    for i in eachindex(idxs)
+        idxs_  = idxs[i]
+        dists_ = dists[i]
+        values = [view(ymat, j, :) for j in idxs_]
+        if w !== nothing
+            w_ = w[idxs_]
+        end
+        if m.weights == :uniform
+            preds[i,:] .= sum(values .* w_) / sum(w_)
+        else
+            preds[i,:] .= sum(values .* w_ .* (1.0 .- dists_ ./ sum(dists_))) / (sum(w_) - 1)
+        end
+    end
+    return MMI.table(preds, names=MMI.schema(y).names, prototype=y)
+end
+
 # ====
 
 metadata_pkg.((KNNRegressor, KNNClassifier),
@@ -161,7 +182,7 @@ metadata_pkg.((KNNRegressor, KNNClassifier),
 
 metadata_model(KNNRegressor,
     input   = Table(Continuous),
-    target  = AbstractVector{Continuous},
+    target  = Union{AbstractVector{Continuous}, Table(Continuous)},
     weights = true,
     descr   = KNNRegressorDescription
     )

src/registry/Metadata.toml

@@ -3319,4 +3319,4 @@
 ":implemented_methods" = []
 ":hyperparameters" = "`(:builder, :optimiser, :loss, :epochs, :batch_size, :lambda, :alpha, :optimiser_changes_trigger_retraining)`"
 ":hyperparameter_types" = "`(\"Any\", \"Any\", \"Any\", \"Int64\", \"Int64\", \"Float64\", \"Float64\", \"Bool\")`"
-":hyperparameter_ranges" = "`(nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing)`"
+":hyperparameter_ranges" = "`(nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing)`"

src/registry/src/check_registry.jl

@@ -1,5 +1,6 @@
 function check_registry() 
-    basedir = joinpath(dirname(pathof(MLJModels)), "registry")
+#    basedir = joinpath(dirname(pathof(MLJModels)), "registry")
+    basedir = Registry.environment_path
     Pkg.activate(basedir)
 
     # Read Metadata.toml

test/NearestNeighbors.jl

@@ -7,6 +7,7 @@ using MLJModels.NearestNeighbors_
 using CategoricalArrays
 using MLJBase
 using Random
+using Tables
 
 Random.seed!(5151)
 
@@ -109,7 +110,19 @@ p2 = predict(knnr, f2, xtest)
 @test all(p[ntest+1:2*ntest] .≈ 2.0)
 @test all(p[2*ntest+1:end] .≈ -2.0)
 
+ymat = vcat(fill( 0.0, n, 2), fill(2.0, n, 2), fill(-2.0, n, 2))
+yv = Tables.table(ymat; header = [:a, :b])
 
+fv,_,_ = fit(knnr, 1, x, yv)
+f2v,_,_ = fit(knnr, 1, x, yv, w)
+
+pv = predict(knnr, fv, xtest)
+
+for col in [:a, :b]
+    @test all(pv[col][1:ntest] .≈ [0.0])
+    @test all(pv[col][ntest+1:2*ntest] .≈ [2.0])
+    @test all(pv[col][2*ntest+1:end] .≈ [-2.0])
+end
 
 
 
@@ -128,8 +141,7 @@ infos[:docstring]
 infos = info_dict(knnr)
 
 @test infos[:input_scitype] == Table(Continuous)
-@test infos[:target_scitype] == AbstractVector{Continuous}
-
+@test infos[:target_scitype] == Union{AbstractVector{Continuous}, Table(Continuous)}
 infos[:docstring]
 
 end