mlj interface draft done. untested wip

Author: PyDataBlog

Project.toml

@@ -4,18 +4,19 @@ authors = ["Bernard Brenyah", "Andrey Oskin"]
 version = "0.1.0"
 
 [deps]
+Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
 StatsBase = "0.32, 0.33"
-julia = "1.3"
+julia = "1.3, 1.4"
 
 [extras]
+MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 
 [targets]
 test = ["Test", "Random", "Suppressor", "MLJBase"]

src/mlj_interface.jl

@@ -1,6 +1,7 @@
 # TODO 1: a using MLJModelInterface or import MLJModelInterface statement
 using MLJModelInterface
 using ParallelKMeans
+using Distances
 
 
 ####
@@ -16,7 +17,7 @@ using ParallelKMeans
     max_iters::Int                          = 300::(_ > 0)
     #transpose_type::String                  = "permute"::(_ in ("permute", "transpose"))
     threads::Int                            = Threads.nthreads()::(_ > 0)
-    verbosity::Int                          = 0::(_ in (0, 1))
+    verbosity::Int                          = 0::(_ in (0, 1))  # Temp fix. Do we need to follow mlj verbosity style?
     init = nothing
 end
 
@@ -32,11 +33,13 @@ const ParallelKMeans_Desc = "Parallel & lightning fast implementation of all var
 ####
 """
     TODO 3.1: Docs
+
+    See also the [package documentation](https://pydatablog.github.io/ParallelKMeans.jl/stable).
 """
 function MLJModelInterface.fit(m::ParaKMeans, verbosity::Int, X)
     # fit the specified struct as a ParaKMeans model
 
-    # assumes user supplied table with columns as features
+    # convert tabular input data into the matrix model expects. Column assumed as features so input data is permuted
     DMatrix = MLJModelInterface.matrix(X; transpose=true)
 
     # fit model and get results
@@ -53,7 +56,9 @@ function MLJModelInterface.fit(m::ParaKMeans, verbosity::Int, X)
     end
 
     cache = nothing
-    report = NamedTuple{}()
+    report = (cluster_centers=fitresult.centers, iterations=fitresult.iterations, 
+              converged=fitresult.converged, totalcost=fitresult.totalcost,
+              labels=fitresult.assignments)
 
     return (fitresult, cache, report)
 end
@@ -64,12 +69,15 @@ end
 """
 function MLJModelInterface.fitted_params(model::KMeansModel, fitresult)
     # extract what's relevant from `fitresult`
-    centres = fitresult.centres
-    converged = fitresult.converged
-    iters = fitresult.iterations
-    totalcost = fitresult.totalcost
+    results, _, _ = fitresult  # unpack fitresult
+    centers = results.centers
+    converged = results.converged
+    iters = results.iterations
+    totalcost = results.totalcost
+
     # then return as a NamedTuple
-    return (centres = centres, totalcost = totalcost, iterations = iters, converged = converged)
+    return (cluster_centers = centers, totalcost = totalcost,
+            iterations = iters, converged = converged)
 end
 
 
@@ -79,8 +87,17 @@ end
 """
     TODO 3.3: Docs
 """
-function MLJModelInterface.predict(m::ParaKMeans, fitresult, Xnew)
-    # ...
+function MLJModelInterface.transform(m::ParaKMeans, fitresult, Xnew)
+    # make predictions/assignments using the learned centroids
+    results = fitresult[1]
+    DMatrix = MLJModelInterface.matrix(Xnew, transpose=true)
+
+    # TODO 3.3.1: Warn users if fitresult is from a `non-converged` fit.
+    # use centroid matrix to assign clusters for new data
+    centroids = results.centers
+    distances = pairwise(SqEuclidean(), DMatrix, centroids; dims=2)
+    preds = argmin.(eachrow(distances))
+    return MLJModelInterface.table(reshape(preds, :, 1), prototype=Xnew)
 end