mlj interface unittests ready for refactoring

Author: PyDataBlog

src/mlj_interface.jl

@@ -15,15 +15,15 @@ import Distances
     k::Int                                  = 3::(_ > 0)
     tol::Float64                            = 1e-6::(_ < 1)
     max_iters::Int                          = 300::(_ > 0)
-    copy::Bool                              = true::(_ in (true, false))
+    copy::Bool                              = true
     threads::Int                            = Threads.nthreads()::(_ > 0)
     verbosity::Int                          = 0::(_ in (0, 1))  # Temp fix. Do we need to follow mlj verbosity style?
     init = nothing
 end
 
 
 # Expose all instances of user specified structs and package artifcats.
-const KMeansModel = Union{KMeans}
+#const KMeansModel = Union{KMeans}
 const ParallelKMeans_Desc = "Parallel & lightning fast implementation of all variants of the KMeans clustering algorithm in native Julia."
 
 
@@ -72,7 +72,7 @@ end
 """
     TODO 3.2: Docs
 """
-function MLJModelInterface.fitted_params(model::KMeansModel, fitresult)
+function MLJModelInterface.fitted_params(model::KMeans, fitresult)
     # extract what's relevant from `fitresult`
     results, _, _ = fitresult  # unpack fitresult
     centers = results.centers
@@ -111,7 +111,7 @@ end
 ####
 
 # TODO 4: metadata for the package and for each of the model interfaces
-metadata_pkg.(KMeansModel,
+metadata_pkg.(KMeans,
     name = "ParallelKMeans",
     uuid = "42b8e9d4-006b-409a-8472-7f34b3fb58af", # see your Project.toml
     url  = "https://github.com/PyDataBlog/ParallelKMeans.jl",  # URL to your package repo

test/test07_mlj_interface.jl

@@ -11,7 +11,7 @@ using MLJBase
 @testset "Test struct construction" begin
     model = ParallelKMeans.KMeans()
 
-    @test model.algo            == Lloyd()
+    @test model.algo            == :Lloyd
     @test model.init            == nothing
     @test model.k               == 3
     @test model.k_init          == "k-means++"
@@ -24,17 +24,40 @@ end
 
 
 @testset "Test model fitting" begin
+    X = table([1 2; 1 4; 1 0; 10 2; 10 4; 10 0])
+    model = ParallelKMeans.KMeans(k=2)
+    results = fit(model, 0, X)
 
+    @test results[2]             == nothing
+    @test results[end].converged == true
+    @test results[end].totalcost == 16
 end
 
 
 @testset "Test fitted params" begin
+    X = table([1 2; 1 4; 1 0; 10 2; 10 4; 10 0])
+    model = ParallelKMeans.KMeans(k=2)
+    results = fit(model, 0, X)
+
+    params = fitted_params(model, results)
+    @test params.converged == true
+    @test params.totalcost == 16
 
 end
 
 
 @testset "Test transform" begin
+    Random.seed!(2020)
+    X = table([1 2; 1 4; 1 0; 10 2; 10 4; 10 0])
+    X_test = table([10 1])
 
+    # Train model using training data X
+    model = ParallelKMeans.KMeans(k=2)
+    results = fit(model, 0, X)
+
+    # Use trained model to cluster new data X_test
+    preds = transform(model, results, X_test)
+    @test preds[:x1][1] == 2
 end