Add StatsLearnModel (#9)

* Add 'StatsLearnModel'

* Update 'Learn'

Author: eliascarv

src/StatsLearnModels.jl

@@ -8,7 +8,7 @@ using Tables
 using Distances
 using DataScienceTraits
 using StatsBase: mode, mean
-using ColumnSelectors: selector
+using ColumnSelectors: ColumnSelector, selector
 using TableTransforms: StatelessFeatureTransform
 
 import DataScienceTraits as DST

src/interface.jl

@@ -31,3 +31,22 @@ struct FittedModel{M,C}
 end
 
 Base.show(io::IO, ::FittedModel{M}) where {M} = print(io, "FittedModel{$(nameof(M))}")
+
+"""
+    StatsLearnModels.StatsLearnModel(model, incols, outcols)
+
+Wrapper type for learning models used for dispatch purposes.
+"""
+struct StatsLearnModel{M,I<:ColumnSelector,O<:ColumnSelector}
+  model::M
+  input::I
+  output::O
+end
+
+StatsLearnModel(model, incols, outcols) = StatsLearnModel(model, selector(incols), selector(outcols))
+
+function Base.show(io::IO, model::StatsLearnModel{M}) where {M}
+  println(io, "StatsLearnModel{$(nameof(M))}")
+  println(io, "├─ input: $(model.input)")
+  print(io, "└─ output: $(model.output)")
+end

src/learn.jl

@@ -26,20 +26,22 @@ struct Learn{M<:FittedModel} <: StatelessFeatureTransform
   input::Vector{Symbol}
 end
 
-function Learn(train, model, (incols, outcols)::Pair)
+Learn(train, model, (incols, outcols)::Pair) = Learn(train, StatsLearnModel(model, incols, outcols))
+
+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)
-  innms = selector(incols)(names)
-  outnms = selector(outcols)(names)
+  innms = lmodel.input(names)
+  outnms = lmodel.output(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)
+  fmodel = fit(lmodel.model, input, output)
   Learn(fmodel, innms)
 end
 

test/runtests.jl

@@ -23,6 +23,14 @@ const SLM = StatsLearnModels
     @test sprint(show, fmodel) == "FittedModel{DecisionTreeClassifier}"
   end
 
+  @testset "StatsLearnModel" begin
+    model = SLM.StatsLearnModel(DecisionTreeClassifier(), [:a, :b], :c)
+    @test sprint(show, model) == """
+    StatsLearnModel{DecisionTreeClassifier}
+    ├─ input: [:a, :b]
+    └─ output: :c"""
+  end
+
   @testset "models" begin
     @testset "MLJ" begin
       Random.seed!(123)