Add Learn transform (#1)

* Add 'Learn' transform

* Add compat

* Fix typo

* Update tests

Author: eliascarv

Project.toml

@@ -4,6 +4,8 @@ authors = ["Elias Carvalho <[email protected]> and contributors"]
 version = "0.1.0"
 
 [deps]
+ColumnSelectors = "9cc86067-7e36-4c61-b350-1ac9833d277f"
+TableTransforms = "0d432bfd-3ee1-4ac1-886a-39f05cc69a3e"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [weakdeps]
@@ -13,6 +15,8 @@ MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
 StatsLearnModelsMLJModelInterfaceExt = "MLJModelInterface"
 
 [compat]
+ColumnSelectors = "0.1"
 MLJModelInterface = "1.9"
+TableTransforms = "1.15"
 Tables = "1.11"
 julia = "1.9"

src/StatsLearnModels.jl

@@ -1,31 +1,13 @@
 module StatsLearnModels
 
-"""
-    StatsLearnModels.fit(model, input, output) -> FittedModel
+using Tables
+using ColumnSelectors: selector
+using TableTransforms: StatelessFeatureTransform
+import TableTransforms: applyfeat, isrevertible
 
-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.
-"""
-function fit end
+export Learn
 
-"""
-    StatsLearnModels.predict(model::FittedModel, table)
-
-Predict the target values using the fitted statistical learning `model`
-and a new `table` of features.  
-"""
-function predict end
-
-"""
-    StatsLearnModels.FittedModel(model, cache)
-
-Wrapper type used to save learning model and auxiliary
-variables needed for prediction.
-"""
-struct FittedModel{M,C}
-  model::M
-  cache::C
-end
+include("interface.jl")
+include("learn.jl")
 
 end

src/interface.jl

@@ -0,0 +1,27 @@
+"""
+    StatsLearnModels.fit(model, input, output) -> FittedModel
+
+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.
+"""
+function fit end
+
+"""
+    StatsLearnModels.predict(model::FittedModel, table)
+
+Predict the target values using the fitted statistical learning `model`
+and a new `table` of features.  
+"""
+function predict end
+
+"""
+    StatsLearnModels.FittedModel(model, cache)
+
+Wrapper type used to save learning model and auxiliary
+variables needed for prediction.
+"""
+struct FittedModel{M,C}
+  model::M
+  cache::C
+end

src/learn.jl

@@ -0,0 +1,49 @@
+"""
+    Learn(model, train, incols, outcols)
+
+Fits the statistical learning `model` using the input columns, selected by `incols`,
+and the output columns, selected by `outcols`, from the `train` table.
+
+The column selection can be a single column identifier (index or name),
+a collection of identifiers or a regular expression (regex).
+
+# Examples
+
+```julia
+Learn(model, train, [1, 2, 3], "d")
+Learn(model, train, [:a, :b, :c], :d)
+Learn(model, train, ["a", "b", "c"], 4)
+Learn(model, train, [1, 2, 3], [:d, :e])
+Learn(model, train, r"[abc]", ["d", "e"])
+```
+"""
+struct Learn{M<:FittedModel} <: StatelessFeatureTransform
+  model::M
+  input::Vector{Symbol}
+end
+
+function Learn(model, train, incols, outcols)
+  if !Tables.istable(train)
+    throw(ArgumentError("training data must be a table"))
+  end
+
+  cols = Tables.columns(train)
+  names = Tables.columnnames(cols)
+  innms = selector(incols)(names)
+  outnms = selector(outcols)(names)
+
+  input = (; (nm => Tables.getcolumn(cols, nm) for nm in innms)...)
+  output = (; (nm => Tables.getcolumn(cols, nm) for nm in outnms)...)
+
+  fmodel = fit(model, input, output)
+  Learn(fmodel, innms)
+end
+
+isrevertible(::Type{<:Learn}) = false
+
+function applyfeat(transform::Learn, feat, prep)
+  cols = Tables.columns(feat)
+  pairs = (nm => Tables.getcolumn(cols, nm) for nm in transform.input)
+  test = (; pairs...) |> Tables.materializer(feat)
+  predict(transform.model, test), nothing
+end

test/Project.toml

@@ -3,4 +3,5 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 MLJ = "add582a8-e3ab-11e8-2d5e-e98b27df1bc7"
 MLJDecisionTreeInterface = "c6f25543-311c-4c74-83dc-3ea6d1015661"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+TableTransforms = "0d432bfd-3ee1-4ac1-886a-39f05cc69a3e"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/runtests.jl

@@ -1,18 +1,41 @@
-import StatsLearnModels as SLM
+using StatsLearnModels
 using MLJ, MLJDecisionTreeInterface
+using TableTransforms
 using DataFrames
 using Random
 using Test
 
+const SLM = StatsLearnModels
+
 @testset "StatsLearnModels.jl" begin
-  Random.seed!(123)
-  iris = DataFrame(load_iris())
-  input = iris[:, Not(:target)]
-  output = iris[:, [:target]]
-  Tree = @load(DecisionTreeClassifier, pkg=DecisionTree, verbosity=0)
-  train, test = partition(1:nrow(input), 0.7, rng=123)
-  fmodel = SLM.fit(Tree(), input[train, :], output[train, :])
-  pred = SLM.predict(fmodel, input[test, :])
-  accuracy = count(pred.target .== output.target[test]) / length(test)
-  @test accuracy > 0.9
+  @testset "interface" begin
+    Random.seed!(123)
+    iris = DataFrame(load_iris())
+    input = iris[:, Not(:target)]
+    output = iris[:, [:target]]
+    Tree = @load(DecisionTreeClassifier, pkg=DecisionTree, verbosity=0)
+    train, test = partition(1:nrow(input), 0.7, rng=123)
+    fmodel = SLM.fit(Tree(), input[train, :], output[train, :])
+    pred = SLM.predict(fmodel, input[test, :])
+    accuracy = count(pred.target .== output.target[test]) / length(test)
+    @test accuracy > 0.9
+  end
+
+  @testset "Learn" begin
+    Random.seed!(123)
+    iris = DataFrame(load_iris())
+    outcol = :target
+    incols = setdiff(propertynames(iris), [outcol])
+    Tree = @load(DecisionTreeClassifier, pkg=DecisionTree, verbosity=0)
+    train, test = partition(1:nrow(iris), 0.7, rng=123)
+    transform = Learn(Tree(), iris[train, :], incols, outcol)
+    @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(Tree(), nothing, incols, outcol)
+  end
 end