Merge pull request #68 from mousum-github/main

update NegativeBinomial and corresponding test cases
xKDR · Dec 20, 2022 · fcb1f2b · fcb1f2b
2 parents 2e008d6 + 9d0d10b
commit fcb1f2b
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 14 deletions.
diff --git a/src/CRRao.jl b/src/CRRao.jl
@@ -385,7 +385,7 @@ Cauchit() = Cauchit(Cauchit_Link)
 export LinearRegression, LogisticRegression, PoissonRegression, NegBinomRegression
 export Prior_Ridge, Prior_Laplace, Prior_Cauchy, Prior_TDist, Prior_HorseShoe, Prior_Gauss
 export CRRaoLink, Logit, Probit, Cloglog, Cauchit, fit
-export coeftable, r2, adjr2, loglikelihood, aic, bic, sigma, predict, residuals, cooksdistance
+export coef, coeftable, r2, adjr2, loglikelihood, aic, bic, sigma, predict, residuals, cooksdistance
 export FrequentistRegression, BayesianRegression
 
 include("random_number_generator.jl")

diff --git a/src/frequentist/getter.jl b/src/frequentist/getter.jl
@@ -1,3 +1,29 @@
+"""
+```julia
+coef(container::FrequentistRegression)
+```
+
+Estimated coefficients of the model. Extends the `coef` method from [StatsAPI.jl](https://github.com/JuliaStats/StatsAPI.jl).
+
+# Example
+
+```julia
+using CRRao, RDatasets, StatsModels
+
+# Get the dataset
+mtcars = dataset("datasets", "mtcars")
+
+# Train the model
+container = fit(@formula(MPG ~ HP + WT + Gear), mtcars, LinearRegression())
+
+# Get table of coefficients
+coef(container)
+```
+"""
+function coef(container::FrequentistRegression)
+    return StatsBase.coef(container.model)
+end
+
 """
 ```julia
 coeftable(container::FrequentistRegression)

diff --git a/src/frequentist/negativebinomial_regression.jl b/src/frequentist/negativebinomial_regression.jl
@@ -1,5 +1,5 @@
 function negativebinomial_reg(formula::FormulaTerm, data::DataFrame, Link::GLM.Link)
-    model = glm(formula, data, NegativeBinomial(), Link)
+    model = negbin(formula, data, Link)
     return FrequentistRegression(:NegativeBinomialRegression, model, formula, typeof(Link))
 end
 
@@ -37,16 +37,16 @@ Model Class: Count Regression
 Likelihood Mode: Negative Binomial
 Link Function: Log
 Computing Method: Optimization
-─────────────────────────────────────────────────────────────────────────────────
-                         Coef.  Std. Error      z  Pr(>|z|)  Lower 95%  Upper 95%
-─────────────────────────────────────────────────────────────────────────────────
-(Intercept)         -1.10939      0.459677  -2.41    0.0158  -2.01034   -0.208444
-Target               0.0117398    0.142779   0.08    0.9345  -0.268101   0.291581
-Coop                 1.0506       0.111556   9.42    <1e-20   0.831949   1.26924
-NCost: major loss   -0.204244     0.508156  -0.40    0.6877  -1.20021    0.791723
-NCost: modest loss   1.27142      0.290427   4.38    <1e-04   0.702197   1.84065
-NCost: net gain      0.176797     0.254291   0.70    0.4869  -0.321604   0.675197
-─────────────────────────────────────────────────────────────────────────────────
+──────────────────────────────────────────────────────────────────────────────────
+                          Coef.  Std. Error      z  Pr(>|z|)  Lower 95%  Upper 95%
+──────────────────────────────────────────────────────────────────────────────────
+(Intercept)         -1.14517       0.480887  -2.38    0.0172  -2.0877    -0.202652
+Target               0.00862527    0.145257   0.06    0.9527  -0.276074   0.293324
+Coop                 1.06397       0.115995   9.17    <1e-19   0.836621   1.29131
+NCost: major loss   -0.23511       0.511443  -0.46    0.6457  -1.23752    0.7673
+NCost: modest loss   1.30767       0.276012   4.74    <1e-05   0.766698   1.84865
+NCost: net gain      0.183453      0.275387   0.67    0.5053  -0.356296   0.723202
+──────────────────────────────────────────────────────────────────────────────────
 ```
 """
 function fit(formula::FormulaTerm, data::DataFrame, modelClass::NegBinomRegression)

diff --git a/test/basic/NegBinomialRegression.jl b/test/basic/NegBinomialRegression.jl
@@ -8,8 +8,16 @@ priors = [
 ]
 
 CRRao.set_rng(StableRNG(123))
-model = fit(@formula(Num ~ Target + Coop + NCost), sanction, NegBinomRegression())
-@test sizeof(model) > 0
+
+@testset "Frequentist NegativeBinomial Regression" begin
+    model = fit(@formula(Num ~ Target + Coop + NCost), sanction, NegBinomRegression())
+    @test sizeof(model) > 0
+    @test coef(model) ≈ [-1.1450754826987306, 0.008593046159347997, 
+                         1.0639242654967949, -0.23501846904848583,
+                         1.307960123972135, 0.18348159198531477]
+    @test aic(model) ≈ 344.8795490118083
+    @test bic(model) ≈ 361.3765107986354
+end
 
 for prior in priors
     CRRao.set_rng(StableRNG(123))