Add GLM.jl models (#5)

* Add GLM.jl models

* Update src/StatsLearnModels.jl

---------

Co-authored-by: Júlio Hoffimann <[email protected]>

Author: eliascarv

Project.toml

@@ -6,6 +6,8 @@ version = "0.2.0"
 [deps]
 ColumnSelectors = "9cc86067-7e36-4c61-b350-1ac9833d277f"
 DecisionTree = "7806a523-6efd-50cb-b5f6-3fa6f1930dbb"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+GLM = "38e38edf-8417-5370-95a0-9cbb8c7f171a"
 TableTransforms = "0d432bfd-3ee1-4ac1-886a-39f05cc69a3e"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
@@ -18,6 +20,8 @@ StatsLearnModelsMLJModelInterfaceExt = "MLJModelInterface"
 [compat]
 ColumnSelectors = "0.1"
 DecisionTree = "0.12"
+Distributions = "0.25"
+GLM = "1.9"
 MLJModelInterface = "1.9"
 TableTransforms = "1.15"
 Tables = "1.11"

src/StatsLearnModels.jl

@@ -9,21 +9,28 @@ using ColumnSelectors: selector
 using TableTransforms: StatelessFeatureTransform
 import TableTransforms: applyfeat, isrevertible
 
+import GLM
 import DecisionTree as DT
 using DecisionTree: AdaBoostStumpClassifier, DecisionTreeClassifier, RandomForestClassifier
 using DecisionTree: DecisionTreeRegressor, RandomForestRegressor
+using Distributions: UnivariateDistribution
 
 include("interface.jl")
 include("models/decisiontree.jl")
+include("models/glm.jl")
 include("learn.jl")
 
 export
-  # models
+  # DecisionTree.jl
   AdaBoostStumpClassifier,
   DecisionTreeClassifier,
   RandomForestClassifier,
   DecisionTreeRegressor,
   RandomForestRegressor,
+  
+  # GLM.jl
+  LinearRegressor,
+  GeneralizedLinearRegressor,
 
   # transform
   Learn

src/models/glm.jl

@@ -0,0 +1,43 @@
+abstract type GLMModel end
+
+struct LinearRegressor{K} <: GLMModel
+  kwargs::K
+end
+
+LinearRegressor(; kwargs...) = LinearRegressor(values(kwargs))
+
+struct GeneralizedLinearRegressor{D<:UnivariateDistribution,L<:Union{GLM.Link,Nothing},K} <: GLMModel
+  dist::D
+  link::L
+  kwargs::K
+end
+
+GeneralizedLinearRegressor(dist::UnivariateDistribution, link=nothing; kwargs...) =
+  GeneralizedLinearRegressor(dist, link, values(kwargs))
+
+function fit(model::GLMModel, input, output)
+  cols = Tables.columns(output)
+  names = Tables.columnnames(cols)
+  outcol = first(names)
+  X = Tables.matrix(input)
+  y = Tables.getcolumn(cols, outcol)
+  fitted = _fit(model, X, y)
+  FittedModel(model, (fitted, outcol))
+end
+
+function predict(fmodel::FittedModel{<:GLMModel}, table)
+  model, outcol = fmodel.cache
+  X = Tables.matrix(table)
+  ŷ = GLM.predict(model, X)
+  (; outcol => ŷ) |> Tables.materializer(table)
+end
+
+_fit(model::LinearRegressor, X, y) = GLM.lm(X, y; model.kwargs...)
+
+function _fit(model::GeneralizedLinearRegressor, X, y)
+  if isnothing(model.link) 
+    GLM.glm(X, y, model.dist; model.kwargs...)
+  else
+    GLM.glm(X, y, model.dist, model.link; model.kwargs...)
+  end
+end

test/Project.toml

@@ -1,5 +1,7 @@
 [deps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+GLM = "38e38edf-8417-5370-95a0-9cbb8c7f171a"
 MLJ = "add582a8-e3ab-11e8-2d5e-e98b27df1bc7"
 MLJDecisionTreeInterface = "c6f25543-311c-4c74-83dc-3ea6d1015661"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

test/runtests.jl

@@ -4,6 +4,9 @@ using DataFrames
 using Random
 using Test
 
+using GLM: ProbitLink
+using Distributions: Binomial
+
 import MLJ, MLJDecisionTreeInterface
 
 const SLM = StatsLearnModels
@@ -17,7 +20,7 @@ const SLM = StatsLearnModels
   @testset "interface" begin
     @testset "MLJ" begin
       Random.seed!(123)
-      Tree = MLJ.@load(DecisionTreeClassifier, pkg=DecisionTree, verbosity=0)
+      Tree = MLJ.@load(DecisionTreeClassifier, pkg = DecisionTree, verbosity = 0)
       fmodel = SLM.fit(Tree(), input[train, :], output[train, :])
       pred = SLM.predict(fmodel, input[test, :])
       accuracy = count(pred.target .== output.target[test]) / length(test)
@@ -32,6 +35,25 @@ const SLM = StatsLearnModels
       accuracy = count(pred.target .== output.target[test]) / length(test)
       @test accuracy > 0.9
     end
+
+    @testset "GLM" begin
+      x = [1, 2, 3]
+      y = [2, 4, 7]
+      input = DataFrame(; ones=ones(length(x)), x)
+      output = DataFrame(; y)
+      model = LinearRegressor()
+      fmodel = SLM.fit(model, input, output)
+      pred = SLM.predict(fmodel, input)
+      @test all(isapprox.(pred.y, output.y, atol=0.5))
+      x = [1, 2, 2]
+      y = [1, 0, 1]
+      input = DataFrame(; ones=ones(length(x)), x)
+      output = DataFrame(; y)
+      model = GeneralizedLinearRegressor(Binomial(), ProbitLink())
+      fmodel = SLM.fit(model, input, output)
+      pred = SLM.predict(fmodel, input)
+      @test all(isapprox.(pred.y, output.y, atol=0.5))
+    end
   end
 
   @testset "Learn" begin