Handle MatrixTerm when calculating column-to-term indexes in ModelMatrix (#134)

* handle MatrixTerm in asgn model matrix calculation

* add test for #133

* bump version (bugfix)

Author: kleinschmidt

Project.toml

@@ -1,6 +1,6 @@
 name = "StatsModels"
 uuid = "3eaba693-59b7-5ba5-a881-562e759f1c8d"
-version = "0.6.1"
+version = "0.6.2"
 
 [deps]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"

src/modelframe.jl

@@ -172,11 +172,16 @@ end
 
 Base.size(mm::ModelMatrix, dim...) = size(mm.m, dim...)
 
+asgn(f::FormulaTerm) = asgn(f.rhs)
+asgn(mt::MatrixTerm) = asgn(mt.terms)
+asgn(t) = mapreduce(((i,t), ) -> i*ones(width(t)),
+                    append!,
+                    enumerate(vectorize(t)),
+                    init=Int[])
+
 function ModelMatrix{T}(mf::ModelFrame) where T<:AbstractMatrix{<:AbstractFloat}
     mat = modelmatrix(mf)
-    asgn = mapreduce((it)->first(it)*ones(width(last(it))), append!,
-                     enumerate(vectorize(mf.f.rhs)), init=Int[])
-    ModelMatrix(convert(T, mat), asgn)
+    ModelMatrix(convert(T, mat), asgn(mf.f))
 end
 
 ModelMatrix(mf::ModelFrame) = ModelMatrix{Matrix{Float64}}(mf)

test/modelframe.jl

@@ -0,0 +1,13 @@
+@testset "ModelFrame (legacy API)" begin
+
+    @testset "#133" begin
+        df = (x = rand(12),
+              y = categorical(repeat(1:3, inner=4)),
+              z = categorical(repeat(1:2, outer=6)));
+        f = @formula(x ~ y * z);
+        mf = ModelFrame(f, df)
+        mm = ModelMatrix(mf)
+        @test mm.assign == [1, 2, 2, 3, 4, 4]
+    end
+    
+end

test/runtests.jl

@@ -13,6 +13,7 @@ my_tests = ["formula.jl",
             "temporal_terms.jl",
             "schema.jl",
             "modelmatrix.jl",
+            "modelframe.jl",
             "statsmodel.jl",
             "contrasts.jl",
             "extension.jl"]