feat: adds optional categorical paramater as option to onehot trans… (#161)

* feat: adds optional `categorical` paramater as option to onehot transform

* fix: fixes variable name spelling on line 53

* adds tests for categorical=false category

* refactor: modifications for the new OneHot function

* refactor: incorporates reviews on previous push

* doc: summarises description of categorical parameter

* refactor: categorical as keyword parameter

* modifies existing tests to accomodate change in the OneHot function (only for categorical=false)

* doc: summarises description of categorical parameter

* Update src/transforms/onehot.jl

* Update src/transforms/onehot.jl

* Update src/transforms/onehot.jl

* Update src/transforms/onehot.jl

* Update src/transforms/onehot.jl

* Update test/transforms/onehot.jl

* Update test/transforms/onehot.jl

* Add kwarg constructor

* Fix show test

* Add 'categorical=true' tests

* Add levels and fix revert bug when 'categorical=true'

* Change order of levels

* Rename the 'categorical' keyword to 'categ'; Check if onehot columns are of type CategoricalArray{Bool} when 'categ=true'

* Fix show test

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

Author: 3 people

src/transforms/onehot.jl

@@ -3,34 +3,40 @@
 # ------------------------------------------------------------------
 
 """
-    OneHot(col)
-    
+    OneHot(col; categ=true)
+
 Transforms categorical column `col` into one-hot columns of levels
 returned by the `levels` function of CategoricalArrays.jl.
+The `categ` option can be used to convert resulting
+columns to categorical arrays as opposed to boolean vectors.
 
 # Examples
 
 ```julia
 OneHot(1)
 OneHot(:a)
 OneHot("a")
+OneHot("a", categ=false)
 ```
 """
 struct OneHot{S<:ColSpec} <: StatelessFeatureTransform
   colspec::S
-  function OneHot(col::Col)
+  categ::Bool
+  function OneHot(col, categ)
     cs = colspec([col])
-    new{typeof(cs)}(cs)
+    new{typeof(cs)}(cs, categ)
   end
 end
 
+OneHot(col; categ=true) = OneHot(col, categ)
+
 isrevertible(::Type{<:OneHot}) = true
 
 function applyfeat(transform::OneHot, feat, prep)
   cols = Tables.columns(feat)
   names = Tables.columnnames(cols) |> collect
   columns = Any[Tables.getcolumn(cols, nm) for nm in names]
-  
+
   name = choose(transform.colspec, names)[1]
   ind = findfirst(==(name), names)
   x = columns[ind]
@@ -48,6 +54,9 @@ function applyfeat(transform::OneHot, feat, prep)
 
   newnms, newcols = first.(onehot), last.(onehot)
 
+  # convert to categorical arrays if necessary
+  newcols = transform.categ ? categorical.(newcols, levels=[false, true]) : newcols
+  
   splice!(names, ind, newnms)
   splice!(columns, ind, newcols)
 
@@ -65,7 +74,7 @@ function revertfeat(::OneHot, newfeat, fcache)
 
   oname, inds, levels, ordered = fcache
   x = map(zip(columns[inds]...)) do row
-    levels[findfirst(row)]
+    levels[findfirst(==(true), row)]
   end
 
   ocolumn = categorical(x; levels, ordered)

test/shows.jl

@@ -200,13 +200,14 @@
 
     # compact mode
     iostr = sprint(show, T)
-    @test iostr == "OneHot([:a])"
+    @test iostr == "OneHot([:a], true)"
 
     # full mode
     iostr = sprint(show, MIME("text/plain"), T)
     @test iostr == """
     OneHot transform
-    └─ colspec = [:a]"""
+    ├─ colspec = [:a]
+    └─ categ = true"""
   end
 
   @testset "Identity" begin

test/transforms/onehot.jl

@@ -4,23 +4,55 @@
   c = categorical([3, 2, 2, 1, 1, 3])
   t = Table(; a, b, c)
 
-  T = OneHot(1)
+  T = OneHot(1; categ=true)
+  n, c = apply(T, t)
+  @test Tables.columnnames(n) == (:a_false, :a_true, :b, :c)
+  @test n.a_false == categorical(Bool[1, 0, 0, 1, 0, 0])
+  @test n.a_true  == categorical(Bool[0, 1, 1, 0, 1, 1])
+  @test n.a_false isa CategoricalVector{Bool}
+  @test n.a_true  isa CategoricalVector{Bool}
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  T = OneHot(:b; categ=true)
+  n, c = apply(T, t)
+  @test Tables.columnnames(n) == (:a, :b_f, :b_m, :c)
+  @test n.b_f == categorical(Bool[0, 1, 0, 0, 0, 1])
+  @test n.b_m == categorical(Bool[1, 0, 1, 1, 1, 0])
+  @test n.b_f isa CategoricalVector{Bool}
+  @test n.b_m isa CategoricalVector{Bool}
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  T = OneHot("c"; categ=true)
+  n, c = apply(T, t)
+  @test Tables.columnnames(n) == (:a, :b, :c_1, :c_2, :c_3)
+  @test n.c_1 == categorical(Bool[0, 0, 0, 1, 1, 0])
+  @test n.c_2 == categorical(Bool[0, 1, 1, 0, 0, 0])
+  @test n.c_3 == categorical(Bool[1, 0, 0, 0, 0, 1])
+  @test n.c_1 isa CategoricalVector{Bool}
+  @test n.c_2 isa CategoricalVector{Bool}
+  @test n.c_3 isa CategoricalVector{Bool}
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  T = OneHot(1; categ=false)
   n, c = apply(T, t)
   @test Tables.columnnames(n) == (:a_false, :a_true, :b, :c)
   @test n.a_false == Bool[1, 0, 0, 1, 0, 0]
   @test n.a_true  == Bool[0, 1, 1, 0, 1, 1]
   tₒ = revert(T, n, c)
   @test t == tₒ
 
-  T = OneHot(:b)
+  T = OneHot(:b; categ=false)
   n, c = apply(T, t)
   @test Tables.columnnames(n) == (:a, :b_f, :b_m, :c)
   @test n.b_f == Bool[0, 1, 0, 0, 0, 1]
   @test n.b_m == Bool[1, 0, 1, 1, 1, 0]
   tₒ = revert(T, n, c)
   @test t == tₒ
 
-  T = OneHot("c")
+  T = OneHot("c"; categ=false)
   n, c = apply(T, t)
   @test Tables.columnnames(n) == (:a, :b, :c_1, :c_2, :c_3)
   @test n.c_1 == Bool[0, 0, 0, 1, 1, 0]
@@ -35,7 +67,7 @@
   b_m = rand(10)
   t   = Table(; b, b_f, b_m)
 
-  T = OneHot(:b)
+  T = OneHot(:b; categ=false)
   n, c = apply(T, t)
   @test Tables.columnnames(n) == (:b_f_, :b_m_, :b_f, :b_m)
   @test n.b_f_ == Bool[0, 1, 0, 0, 0, 1]
@@ -50,7 +82,7 @@
   b_m_ = rand(10)
   t    = Table(; b, b_f, b_m, b_f_, b_m_)
 
-  T = OneHot(:b)
+  T = OneHot(:b; categ=false)
   n, c = apply(T, t)
   @test Tables.columnnames(n) == (:b_f__, :b_m__, :b_f, :b_m, :b_f_, :b_m_)
   @test n.b_f__ == Bool[0, 1, 0, 0, 0, 1]
@@ -70,4 +102,4 @@
   # invalid column selection
   @test_throws AssertionError apply(OneHot(:c), t)
   @test_throws AssertionError apply(OneHot("c"), t)
-end
+end