Refactor Learn transform

Author: juliohm

src/interface.jl

@@ -2,47 +2,21 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-"""
-    StatsLearnModel(model, features, targets)
-
-Wrap a statistical learning `model` with selectors
-of `features` and `targets`.
-
-## Examples
-
-```julia
-StatsLearnModel(DecisionTreeClassifier(), ["x1","x2"], "y")
-StatsLearnModel(DecisionTreeClassifier(), 1:3, "target")
-```
-"""
-struct StatsLearnModel{M,F<:ColumnSelector,T<:ColumnSelector}
-  model::M
-  feats::F
-  targs::T
-end
-
-StatsLearnModel(model, feats, targs) = StatsLearnModel(model, selector(feats), selector(targs))
-
 """
     fit(model, input, output)
 
-Fit statistical learning `model` using features in `input` table
-and targets in `output` table. Returns a fitted model with all
-the necessary information for prediction with the `predict` function.
+Fit statistical learning `model` using predictors
+in `input` table and targets in `output` table.
+Returns a fitted model with all the necessary
+information for prediction with [`predict`](@ref).
 """
 function fit end
 
-function Base.show(io::IO, model::StatsLearnModel{M}) where {M}
-  println(io, "StatsLearnModel{$(nameof(M))}")
-  println(io, "├─ features: $(model.feats)")
-  print(io, "└─ targets: $(model.targs)")
-end
-
 """
     FittedStatsLearnModel(model, cache)
 
-Wrap the statistical learning `model` with the `cache`
-produced during the [`fit`](@ref) stage.
+Wrap the statistical learning `model` with the
+`cache` produced during the [`fit`](@ref) stage.
 """
 struct FittedStatsLearnModel{M,C}
   model::M
@@ -53,7 +27,9 @@ end
     predict(model::FittedStatsLearnModel, table)
 
 Predict targets using the fitted statistical
-learning `model` and a new `table` of features.
+learning `model` and a new `table` containing
+the same predictors used during the [`fit`](@ref)
+stage.
 """
 function predict end
 

src/learn.jl

@@ -3,55 +3,69 @@
 # ------------------------------------------------------------------
 
 """
-    Learn(train, model, features => targets)
+    Learn(table; [model])
 
-Fits the statistical learning `model` to `train` table,
-using the selectors of `features` and `targets`.
+Perform supervised learning with labeled `table` and
+statistical learning `model`.
+
+Uses `KNNClassifier(1)` or `KNNRegressor(1)` model by
+default depending on the scientific type of the labels
+stored in the table.
 
 ## Examples
 
 ```julia
-Learn(train, model, [1, 2, 3] => "d")
-Learn(train, model, [:a, :b, :c] => :d)
-Learn(train, model, ["a", "b", "c"] => 4)
-Learn(train, model, [1, 2, 3] => [:d, :e])
-Learn(train, model, r"[abc]" => ["d", "e"])
+Learn(label(table, "y"))
+Learn(label(table, ["y1", "y2"]))
+Learn(label(table, 3), model=KNNClassifier(5))
 ```
+
+See also [`label`](@ref).
 """
-struct Learn{M<:FittedStatsLearnModel} <: StatelessFeatureTransform
+struct Learn{T<:LabeledTable,M} <: StatelessFeatureTransform
+  table::T
   model::M
-  feats::Vector{Symbol}
 end
 
-Learn(train, model, (feats, targs)::Pair) = Learn(train, StatsLearnModel(model, feats, targs))
-
-function Learn(train, lmodel::StatsLearnModel)
-  if !Tables.istable(train)
-    throw(ArgumentError("training data must be a table"))
-  end
-
-  cols = Tables.columns(train)
-  names = Tables.columnnames(cols)
-  feats = lmodel.feats(names)
-  targs = lmodel.targs(names)
+Learn(table::LabeledTable; model=_defaultmodel(table)) = Learn(table, model)
 
-  input = (; (var => Tables.getcolumn(cols, var) for var in feats)...)
-  output = (; (var => Tables.getcolumn(cols, var) for var in targs)...)
+function applyfeat(transform::Learn, feat, prep)
+  # labeled table and model
+  table = transform.table
+  model = transform.model
 
-  fmodel = fit(lmodel.model, input, output)
+  # variables in labeled table
+  cols = Tables.columns(table)
+  vars = Tables.columnnames(cols)
 
-  Learn(fmodel, feats)
-end
+  # split targets and predictors
+  targs = table.labels
+  preds = setdiff(vars, targs)
 
-isrevertible(::Type{<:Learn}) = false
+  # learn function with statistical model
+  input = (; (pred => Tables.getcolumn(cols, pred) for pred in preds)...)
+  output = (; (targ => Tables.getcolumn(cols, targ) for targ in targs)...)
+  fmodel = fit(model, input, output)
 
-function applyfeat(transform::Learn, feat, prep)
-  model = transform.model
-  vars = transform.feats
+  # predict labels with new predictors
+  fcols = Tables.columns(feat)
+  fvars = Tables.columnnames(fcols)
+  preds ⊆ fvars || throw(ArgumentError("predictors $preds not found in input table"))
+  finput = (; (pred => Tables.getcolumn(fcols, pred) for pred in preds)...)
+  foutput = predict(fmodel, finput) |> Tables.materializer(feat)
 
-  cols = Tables.columns(feat)
-  pairs = (var => Tables.getcolumn(cols, var) for var in vars)
-  test = (; pairs...) |> Tables.materializer(feat)
+  foutput, nothing
+end
 
-  predict(model, test), nothing
+function _defaultmodel(table::LabeledTable)
+  cols = Tables.columns(table)
+  vals = Tables.getcolumn(cols, only(table.labels))
+  type = elscitype(vals)
+  if type <: Categorical
+    KNNClassifier(1)
+  elseif type <: Continuous
+    KNNRegressor(1)
+  else
+    throw(ErrorException("no default learning model for $type labels"))
+  end
 end

test/Project.toml

@@ -1,5 +1,4 @@
 [deps]
-ColumnSelectors = "9cc86067-7e36-4c61-b350-1ac9833d277f"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 GLM = "38e38edf-8417-5370-95a0-9cbb8c7f171a"

test/runtests.jl

@@ -6,7 +6,6 @@ using Test
 
 using GLM: ProbitLink
 using Distributions: Binomial
-using ColumnSelectors: selector
 
 import MLJ, MLJDecisionTreeInterface
 
@@ -24,22 +23,6 @@ const SLM = StatsLearnModels
     model = DecisionTreeClassifier()
     fmodel = SLM.fit(model, input[train, :], output[train, :])
     @test sprint(show, fmodel) == "FittedStatsLearnModel{DecisionTreeClassifier}"
-
-    # show method
-    lmodel = SLM.StatsLearnModel(DecisionTreeClassifier(), [:a, :b], :c)
-    @test sprint(show, lmodel) == """
-    StatsLearnModel{DecisionTreeClassifier}
-    ├─ features: [:a, :b]
-    └─ targets: :c"""
-
-    # accessor functions
-    model = DecisionTreeClassifier()
-    feats = selector([:a, :b])
-    targs = selector(:c)
-    lmodel = SLM.StatsLearnModel(model, feats, targs)
-    @test lmodel.model === model
-    @test lmodel.feats === feats
-    @test lmodel.targs === targs
   end
 
   @testset "Models" begin
@@ -109,19 +92,12 @@ const SLM = StatsLearnModels
     input = iris[:, Not(:target)]
     output = iris[:, [:target]]
     train, test = MLJ.partition(1:nrow(input), 0.7, rng=123)
-    outvar = :target
-    feats = setdiff(propertynames(iris), [outvar])
-    targs = outvar
     model = DecisionTreeClassifier()
-    transform = Learn(iris[train, :], model, feats => targs)
+    transform = Learn(label(iris[train, :], :target); model)
     @test !isrevertible(transform)
     pred = transform(iris[test, :])
     accuracy = count(pred.target .== iris.target[test]) / length(test)
     @test accuracy > 0.9
-
-    # throws
-    # training data is not a table
-    @test_throws ArgumentError Learn(nothing, model, feats => targs)
   end
 
   @testset "MLJ" begin