Major cleanup

Author: juliohm

ext/StatsLearnModelsMLJModelInterfaceExt.jl

@@ -18,10 +18,10 @@ function SLM.fit(model::MI.Model, input, output)
   y = Tables.getcolumn(cols, target)
   data = MI.reformat(model, input, y)
   fitresult, _... = MI.fit(model, 0, data...)
-  SLM.FittedModel(model, (fitresult, target))
+  SLM.FittedStatsLearnModel(model, (fitresult, target))
 end
 
-function SLM.predict(fmodel::SLM.FittedModel{<:MI.Model}, table)
+function SLM.predict(fmodel::SLM.FittedStatsLearnModel{<:MI.Model}, table)
   (; model, cache) = fmodel
   fitresult, target = cache
   data = MI.reformat(model, table)

src/interface.jl

@@ -3,71 +3,58 @@
 # ------------------------------------------------------------------
 
 """
-    StatsLearnModels.fit(model, input, output) -> FittedModel
+    StatsLearnModel(model, invars, outvars)
 
-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
+Wrap a (possibly external) `model` with selectors of
+input variables `invars` and output variables `outvars`.
 
-"""
-    StatsLearnModels.predict(model::FittedModel, table)
+## Examples
 
-Predict the target values using the fitted statistical learning `model`
-and a new `table` of features.  
+```julia
+StatsLearnModel(DecisionTreeClassifier(), ["x1","x2"], "y")
+StatsLearnModel(DecisionTreeClassifier(), 1:3, "target")
+```
 """
-function predict end
-
-"""
-    StatsLearnModels.FittedModel(model, cache)
-
-Wrapper type used to save learning model and auxiliary
-variables needed for prediction.
-"""
-struct FittedModel{M,C}
+struct StatsLearnModel{M,I<:ColumnSelector,O<:ColumnSelector}
   model::M
-  cache::C
+  invars::I
+  outvars::O
 end
 
-Base.show(io::IO, ::FittedModel{M}) where {M} = print(io, "FittedModel{$(nameof(M))}")
+StatsLearnModel(model, invars, outvars) = StatsLearnModel(model, selector(invars), selector(outvars))
 
 """
-    StatsLearnModels.StatsLearnModel(model, incols, outcols)
+    fit(model, input, output)
 
-Wrapper type for learning models used for dispatch purposes.
+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.
 """
-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 fit end
 
 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)")
+  println(io, "├─ features: $(model.invars)")
+  print(io, "└─ targets: $(model.outvars)")
 end
 
 """
-    StatsLearnModels.model(lmodel::StatsLearnModel)
-  
-Returns the model of the `lmodel`.
-"""
-model(lmodel::StatsLearnModel) = lmodel.model
+    FittedStatsLearnModel(model, cache)
 
+Wrap the statistical learning `model` with the `cache`
+produced during the [`fit`](@ref) stage.
 """
-    StatsLearnModels.input(lmodel::StatsLearnModel)
-  
-Returns the input column selection of the `lmodel`.
-"""
-input(lmodel::StatsLearnModel) = lmodel.input
+struct FittedStatsLearnModel{M,C}
+  model::M
+  cache::C
+end
 
 """
-    StatsLearnModels.output(lmodel::StatsLearnModel)
-  
-Returns the output column selection of the `lmodel`.
+    predict(model::FittedStatsLearnModel, table)
+
+Predict targets using the fitted statistical
+learning `model` and a new `table` of features.
 """
-output(lmodel::StatsLearnModel) = lmodel.output
+function predict end
+
+Base.show(io::IO, ::FittedStatsLearnModel{M}) where {M} = print(io, "FittedStatsLearnModel{$(nameof(M))}")

src/learn.jl

@@ -3,13 +3,11 @@
 # ------------------------------------------------------------------
 
 """
-    Learn(train, model, incols => outcols)
+    Learn(train, model, invars => outvars)
 
-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).
+Fits the statistical learning `model` to `train` table,
+using the selectors of input variables `invars` and
+output variables `outvars`.
 
 # Examples
 
@@ -21,12 +19,12 @@ Learn(train, model, [1, 2, 3] => [:d, :e])
 Learn(train, model, r"[abc]" => ["d", "e"])
 ```
 """
-struct Learn{M<:FittedModel} <: StatelessFeatureTransform
+struct Learn{M<:FittedStatsLearnModel} <: StatelessFeatureTransform
   model::M
-  input::Vector{Symbol}
+  invars::Vector{Symbol}
 end
 
-Learn(train, model, (incols, outcols)::Pair) = Learn(train, StatsLearnModel(model, incols, outcols))
+Learn(train, model, (invars, outvars)::Pair) = Learn(train, StatsLearnModel(model, invars, outvars))
 
 function Learn(train, lmodel::StatsLearnModel)
   if !Tables.istable(train)
@@ -35,21 +33,26 @@ function Learn(train, lmodel::StatsLearnModel)
 
   cols = Tables.columns(train)
   names = Tables.columnnames(cols)
-  innms = lmodel.input(names)
-  outnms = lmodel.output(names)
+  invars = lmodel.invars(names)
+  outvars = lmodel.outvars(names)
 
-  input = (; (nm => Tables.getcolumn(cols, nm) for nm in innms)...)
-  output = (; (nm => Tables.getcolumn(cols, nm) for nm in outnms)...)
+  input = (; (var => Tables.getcolumn(cols, var) for var in invars)...)
+  output = (; (var => Tables.getcolumn(cols, var) for var in outvars)...)
 
   fmodel = fit(lmodel.model, input, output)
-  Learn(fmodel, innms)
+
+  Learn(fmodel, invars)
 end
 
 isrevertible(::Type{<:Learn}) = false
 
 function applyfeat(transform::Learn, feat, prep)
+  model = transform.model
+  vars = transform.invars
+
   cols = Tables.columns(feat)
-  pairs = (nm => Tables.getcolumn(cols, nm) for nm in transform.input)
+  pairs = (var => Tables.getcolumn(cols, var) for var in vars)
   test = (; pairs...) |> Tables.materializer(feat)
-  predict(transform.model, test), nothing
+
+  predict(model, test), nothing
 end

src/models/decisiontree.jl

@@ -17,10 +17,10 @@ function fit(model::DecisionTreeModel, input, output)
   y = Tables.getcolumn(cols, outnm)
   X = Tables.matrix(input)
   DT.fit!(model, X, y)
-  FittedModel(model, outnm)
+  FittedStatsLearnModel(model, outnm)
 end
 
-function predict(fmodel::FittedModel{<:DecisionTreeModel}, table)
+function predict(fmodel::FittedStatsLearnModel{<:DecisionTreeModel}, table)
   outnm = fmodel.cache
   X = Tables.matrix(table)
   ŷ = DT.predict(fmodel.model, X)

src/models/glm.jl

@@ -26,10 +26,10 @@ function fit(model::GLMModel, input, output)
   X = Tables.matrix(input)
   y = Tables.getcolumn(cols, outnm)
   fitted = _fit(model, X, y)
-  FittedModel(model, (fitted, outnm))
+  FittedStatsLearnModel(model, (fitted, outnm))
 end
 
-function predict(fmodel::FittedModel{<:GLMModel}, table)
+function predict(fmodel::FittedStatsLearnModel{<:GLMModel}, table)
   model, outnm = fmodel.cache
   X = Tables.matrix(table)
   ŷ = GLM.predict(model, X)

src/models/nearestneighbors.jl

@@ -34,10 +34,10 @@ function fit(model::NearestNeighborsModel, input, output)
   else
     NN.BallTree(data, metric; leafsize, reorder)
   end
-  FittedModel(model, (tree, outnm, outcol))
+  FittedStatsLearnModel(model, (tree, outnm, outcol))
 end
 
-function predict(fmodel::FittedModel{<:NearestNeighborsModel}, table)
+function predict(fmodel::FittedStatsLearnModel{<:NearestNeighborsModel}, table)
   (; model, cache) = fmodel
   tree, outnm, outcol = cache
   data = Tables.matrix(table, transpose=true)

test/runtests.jl

@@ -21,24 +21,24 @@ const SLM = StatsLearnModels
   @testset "show" begin
     model = DecisionTreeClassifier()
     fmodel = SLM.fit(model, input[train, :], output[train, :])
-    @test sprint(show, fmodel) == "FittedModel{DecisionTreeClassifier}"
+    @test sprint(show, fmodel) == "FittedStatsLearnModel{DecisionTreeClassifier}"
   end
 
   @testset "StatsLearnModel" begin
     # accessor functions
     model = DecisionTreeClassifier()
-    incols = selector([:a, :b])
-    outcols = selector(:c)
-    lmodel = SLM.StatsLearnModel(model, incols, outcols)
-    @test SLM.model(lmodel) === model
-    @test SLM.input(lmodel) === incols
-    @test SLM.output(lmodel) === outcols
+    invars = selector([:a, :b])
+    outvars = selector(:c)
+    lmodel = SLM.StatsLearnModel(model, invars, outvars)
+    @test lmodel.model === model
+    @test lmodel.invars === invars
+    @test lmodel.outvars === outvars
     # show
     lmodel = SLM.StatsLearnModel(DecisionTreeClassifier(), [:a, :b], :c)
     @test sprint(show, lmodel) == """
     StatsLearnModel{DecisionTreeClassifier}
-    ├─ input: [:a, :b]
-    └─ output: :c"""
+    ├─ features: [:a, :b]
+    └─ targets: :c"""
   end
 
   @testset "models" begin
@@ -105,17 +105,18 @@ const SLM = StatsLearnModels
 
   @testset "Learn" begin
     Random.seed!(123)
-    outcol = :target
-    incols = setdiff(propertynames(iris), [outcol])
+    outvar = :target
+    invars = setdiff(propertynames(iris), [outvar])
+    outvars = outvar
     model = DecisionTreeClassifier()
-    transform = Learn(iris[train, :], model, incols => outcol)
+    transform = Learn(iris[train, :], model, invars => outvars)
     @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, incols => outcol)
+    @test_throws ArgumentError Learn(nothing, model, invars => outvars)
   end
 end