use DataAPI.levels instead of levels field on contrasts, terms as hints (#181)

* use AbstractTerm hints directly into schema

* introduce levels(::AbstractContrasts) instead of property access

* fix ambiguity error and also ambiguity in levels/levels()

* add placeholder Grouping "contrasts"

* Apply suggestions from code review

Co-authored-by: Milan Bouchet-Valat <[email protected]>

* Revert "add placeholder Grouping "contrasts""

This reverts commit d01aff9.

* tests for hints (and fix hint bug I missed...)

* tests for base= and levels= constructors and levels()/baselevel()

Co-authored-by: Milan Bouchet-Valat <[email protected]>

Author: kleinschmidt

src/StatsModels.jl

@@ -5,6 +5,7 @@ using StatsBase
 using ShiftedArrays
 using ShiftedArrays: lag, lead
 using DataStructures
+using DataAPI
 using DataAPI: levels
 using Printf: @sprintf
 using Distributions: Chisq, ccdf

src/contrasts.jl

@@ -159,12 +159,12 @@ function ContrastsMatrix(contrasts::C, levels::AbstractVector{T}) where {C<:Abst
     # if levels are defined on contrasts, use those, validating that they line up.
     # what does that mean? either:
     #
-    # 1. contrasts.levels == levels (best case)
+    # 1. DataAPI.levels(contrasts) == levels (best case)
     # 2. data levels missing from contrast: would generate empty/undefined rows.
     #    better to filter data frame first
     # 3. contrast levels missing from data: would have empty columns, generate a
     #    rank-deficient model matrix.
-    c_levels = something(contrasts.levels, levels)
+    c_levels = something(DataAPI.levels(contrasts), levels)
     if eltype(c_levels) != eltype(levels)
         throw(ArgumentError("mismatching levels types: got $(eltype(levels)), expected " *
                             "$(eltype(c_levels)) based on contrasts levels."))
@@ -242,11 +242,13 @@ for contrastType in [:DummyCoding, :EffectsCoding, :HelmertCoding, :SeqDiffCodin
         ## constructor with optional keyword arguments, defaulting to nothing
         $contrastType(; base=nothing, levels::Union{AbstractVector,Nothing}=nothing) = $contrastType(base, levels)
         baselevel(c::$contrastType) = c.base
+        DataAPI.levels(c::$contrastType) = c.levels
     end
 end
 
-# fallback method for other types that might not have base field
+# fallback method for other types that might not have base or level fields
 baselevel(c::AbstractContrasts) = nothing
+DataAPI.levels(c::AbstractContrasts) = nothing
 
 """
     FullDummyCoding()
@@ -586,6 +588,8 @@ end
 termnames(C::HypothesisCoding, levels::AbstractVector, baseind::Int) =
     something(C.labels, levels[1:length(levels) .!= baseind])
 
+DataAPI.levels(c::HypothesisCoding) = c.levels
+
 """
     StatsModels.ContrastsCoding(mat::AbstractMatrix[, levels]])
     StatsModels.ContrastsCoding(mat::AbstractMatrix[; levels=nothing])
@@ -628,6 +632,8 @@ function contrasts_matrix(C::ContrastsCoding, baseind, n)
     C.mat
 end
 
+DataAPI.levels(c::ContrastsCoding) = c.levels
+
 ## hypothesis matrix
 """
     needs_intercept(mat::AbstractMatrix)
@@ -721,4 +727,3 @@ function pretty_mat(mat::AbstractMatrix; tol::Real=10*eps(eltype(mat)))
         return fracs
     end
 end
-    

src/schema.jl

@@ -172,6 +172,11 @@ concrete_term(t::Term, dt::ColumnTable, hints::Dict{Symbol}) =
     concrete_term(t, getproperty(dt, t.sym), get(hints, t.sym, nothing))
 concrete_term(t::Term, d) = concrete_term(t, d, nothing)
 
+# if the "hint" is already an AbstractTerm, use that
+# need this specified to avoid ambiguity
+concrete_term(t::Term, d::ColumnTable, hint::AbstractTerm) = hint
+concrete_term(t::Term, x, hint::AbstractTerm) = hint
+
 # second possible fix for #97
 concrete_term(t, d, hint) = t
 

test/contrasts.jl

@@ -283,4 +283,53 @@
         cmat3p[1] += 1e-3
         @test needs_intercept(cmat3p) == true
     end
+
+    @testset "levels and baselevel" begin
+        using DataAPI: levels
+        using StatsModels: baselevel, FullDummyCoding, ContrastsCoding
+
+        levs = [:a, :b, :c, :d]
+        base = [:a]
+        for C in [DummyCoding, EffectsCoding, SeqDiffCoding, HelmertCoding]
+            c = C()
+            @test levels(c) == nothing
+            @test baselevel(c) == nothing
+
+            c = C(levels=levs)
+            @test levels(c) == levs
+            @test baselevel(c) == nothing
+
+            c = C(base=base)
+            @test levels(c) == nothing
+            @test baselevel(c) == base
+
+            c = C(levels=levs, base=base)
+            @test levels(c) == levs
+            @test baselevel(c) == base
+        end
+
+        c = FullDummyCoding()
+        @test baselevel(c) == nothing
+        @test levels(c) == nothing
+
+        @test_throws MethodError FullDummyCoding(levels=levs)
+        @test_throws MethodError FullDummyCoding(base=base)
+
+        c = HypothesisCoding(rand(3,4))
+        @test baselevel(c) == levels(c) == nothing
+        c = HypothesisCoding(rand(3,4), levels=levs)
+        @test baselevel(c) == nothing
+        @test levels(c) == levs
+        # no notion of base level for HypothesisCoding
+        @test_throws MethodError HypothesisCoding(rand(3,4), base=base)
+
+        c = ContrastsCoding(rand(4,3))
+        @test baselevel(c) == levels(c) == nothing
+        c = ContrastsCoding(rand(4,3), levels=levs)
+        @test baselevel(c) == nothing
+        @test levels(c) == levs
+        # no notion of base level for ContrastsCoding
+        @test_throws MethodError ContrastsCoding(rand(4,3), base=base)
+        
+    end
 end

test/schema.jl

@@ -2,9 +2,37 @@
 
     using StatsModels: schema, apply_schema, FullRank
 
-    f = @formula(y ~ 1 + a + b + c + b&c)
-    df = (y = rand(9), a = 1:9, b = rand(9), c = repeat(["d","e","f"], 3))
-    f = apply_schema(f, schema(f, df))
-    @test f == apply_schema(f, schema(f, df))
+    @testset "no-op apply_schema" begin
+        f = @formula(y ~ 1 + a + b + c + b&c)
+        df = (y = rand(9), a = 1:9, b = rand(9), c = repeat(["d","e","f"], 3))
+        f = apply_schema(f, schema(f, df))
+        @test f == apply_schema(f, schema(f, df))
+    end
 
+    @testset "hints" begin
+        f = @formula(y ~ 1 + a)
+        d = (y = rand(10), a = repeat([1,2], outer=2))
+
+        sch = schema(f, d)
+        @test sch[term(:a)] isa ContinuousTerm
+
+        sch1 = schema(f, d, Dict(:a => CategoricalTerm))
+        @test sch1[term(:a)] isa CategoricalTerm{DummyCoding}
+        f1 = apply_schema(f, sch1)
+        @test f1.rhs.terms[end] == sch1[term(:a)]
+        
+        sch2 = schema(f, d, Dict(:a => DummyCoding()))
+        @test sch2[term(:a)] isa CategoricalTerm{DummyCoding}
+        f2 = apply_schema(f, sch2)
+        @test f2.rhs.terms[end] == sch2[term(:a)]
+
+        hint = deepcopy(sch2[term(:a)])
+        sch3 = schema(f, d, Dict(:a => hint))
+        # if an <:AbstractTerm is supplied as hint, it's included as is
+        @test sch3[term(:a)] === hint !== sch2[term(:a)]
+        f3 = apply_schema(f, sch3)
+        @test f3.rhs.terms[end] === hint
+
+    end
+    
 end