Add keyword arg to modelmatrix; define momentmatrix

src/regressionmodel.jl

@@ -34,19 +34,25 @@ Return the mean of the response.
 function meanresponse end
 
 """
-    modelmatrix(model::RegressionModel)
+    modelmatrix(model::RegressionModel; weighted::Bool=false)
 
-Return the model matrix (a.k.a. the design matrix).
+Return the model matrix (a.k.a. the design matrix) or, if `weighted=true` the weighted 
+model matrix, i.e. `X' * sqrt.(W)`, where `X` is the model matrix and
+`W` is the diagonal matrix whose elements are the model weights. 
 """
-function modelmatrix end
+function modelmatrix(model::RegressionModel; weighted::Bool=false) end
 
 """
-    crossmodelmatrix(model::RegressionModel)
+    crossmodelmatrix(model::RegressionModel; weighted::Bool=false)
 
-Return `X'X` where `X` is the model matrix of `model`.
+Return `X'X` where `X` is the model matrix of `model` or, if `weighted=true`, `X'WX`, 
+where `W` is the diagonal matrix whose elements are the model weights. 
 This function will return a pre-computed matrix stored in `model` if possible.
 """
-crossmodelmatrix(model::RegressionModel) = (x = modelmatrix(model); Symmetric(x' * x))
+function crossmodelmatrix(model::RegressionModel; weighted::Bool=false) 
+    x = weighted ? modelmatrix(model; weighted=weighted) : modelmatrix(model)
+    return Symmetric(x' * x)
+end
 
 """
     leverage(model::RegressionModel)

src/statisticalmodel.jl

@@ -314,3 +314,13 @@ function adjr2(model::StatisticalModel, variant::Symbol)
 end
 
 const adjr² = adjr2
+
+"""
+    momentmatrix(model::StatisticalModel)
+
+Return the matrix containing the estimating equation.
+
+For linear regression models, the moment matrix is given by `u*X`,
+where `u` is the vector of residuals and `X` is the model matrix.
+"""
+function momentmatrix end

test/regressionmodel.jl

@@ -6,13 +6,31 @@ using StatsAPI: RegressionModel, crossmodelmatrix
 struct MyRegressionModel <: RegressionModel
 end
 
+struct MyWeightedRegressionModel <: RegressionModel
+    wts::AbstractVector
+end
+
 StatsAPI.modelmatrix(::MyRegressionModel) = [1 2; 3 4]
 
+function StatsAPI.modelmatrix(r::MyWeightedRegressionModel; weighted::Bool=false)
+    X = [1 2; 3 4]
+    weighted ? sqrt.(r.wts).*X : X
+end
+
+w = [0.3, 0.2]
+
 @testset "TestRegressionModel" begin
     m = MyRegressionModel()
+    r = MyWeightedRegressionModel(w)
 
     @test crossmodelmatrix(m) == [10 14; 14 20]
+    @test crossmodelmatrix(m; weighted=false) == [10 14; 14 20]
     @test crossmodelmatrix(m) isa Symmetric
+
+    @test crossmodelmatrix(r) == [10 14; 14 20]
+    @test crossmodelmatrix(r; weighted=false) == [10 14; 14 20]
+    @test crossmodelmatrix(r; weighted=true) ≈ [2.1 3.0; 3.0 4.4]
+    @test crossmodelmatrix(r; weighted=true) isa Symmetric
 end
 
 end # module TestRegressionModel

test/statisticalmodel.jl

@@ -36,4 +36,4 @@ StatsAPI.nobs(::MyStatisticalModel) = 100
     @test adjr2 === adjr²
 end
 
-end # module TestStatisticalModel
+end # module TestStatisticalModel