Haj

DeclareDesign · Mar 16, 2018 · 76ade51 · 76ade51
1 parent 3f22860
commit 76ade51
Show file tree

Hide file tree

Showing 4 changed files with 82 additions and 0 deletions.
diff --git a/R/RcppExports.R b/R/RcppExports.R
@@ -5,6 +5,14 @@ ht_covar_partial <- function(y1, y0, p10, p1, p0) {
     .Call(`_estimatr_ht_covar_partial`, y1, y0, p10, p1, p0)
 }
 
+cov_a <- function(y1, y0, p10, p1, p0) {
+    .Call(`_estimatr_cov_a`, y1, y0, p10, p1, p0)
+}
+
+var_a <- function(y, p) {
+    .Call(`_estimatr_var_a`, y, p)
+}
+
 ht_var_partial <- function(y, p) {
     .Call(`_estimatr_ht_var_partial`, y, p)
 }

diff --git a/src/RcppExports.cpp b/src/RcppExports.cpp
@@ -21,6 +21,33 @@ BEGIN_RCPP
     return rcpp_result_gen;
 END_RCPP
 }
+// cov_a
+double cov_a(const Eigen::VectorXd& y1, const Eigen::VectorXd& y0, const Eigen::MatrixXd& p10, const Eigen::VectorXd& p1, const Eigen::VectorXd& p0);
+RcppExport SEXP _estimatr_cov_a(SEXP y1SEXP, SEXP y0SEXP, SEXP p10SEXP, SEXP p1SEXP, SEXP p0SEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< const Eigen::VectorXd& >::type y1(y1SEXP);
+    Rcpp::traits::input_parameter< const Eigen::VectorXd& >::type y0(y0SEXP);
+    Rcpp::traits::input_parameter< const Eigen::MatrixXd& >::type p10(p10SEXP);
+    Rcpp::traits::input_parameter< const Eigen::VectorXd& >::type p1(p1SEXP);
+    Rcpp::traits::input_parameter< const Eigen::VectorXd& >::type p0(p0SEXP);
+    rcpp_result_gen = Rcpp::wrap(cov_a(y1, y0, p10, p1, p0));
+    return rcpp_result_gen;
+END_RCPP
+}
+// var_a
+double var_a(const Eigen::VectorXd& y, const Eigen::MatrixXd& p);
+RcppExport SEXP _estimatr_var_a(SEXP ySEXP, SEXP pSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< const Eigen::VectorXd& >::type y(ySEXP);
+    Rcpp::traits::input_parameter< const Eigen::MatrixXd& >::type p(pSEXP);
+    rcpp_result_gen = Rcpp::wrap(var_a(y, p));
+    return rcpp_result_gen;
+END_RCPP
+}
 // ht_var_partial
 double ht_var_partial(const Eigen::VectorXd& y, const Eigen::MatrixXd& p);
 RcppExport SEXP _estimatr_ht_var_partial(SEXP ySEXP, SEXP pSEXP) {
@@ -109,6 +136,8 @@ END_RCPP
 
 static const R_CallMethodDef CallEntries[] = {
     {"_estimatr_ht_covar_partial", (DL_FUNC) &_estimatr_ht_covar_partial, 5},
+    {"_estimatr_cov_a", (DL_FUNC) &_estimatr_cov_a, 5},
+    {"_estimatr_var_a", (DL_FUNC) &_estimatr_var_a, 2},
     {"_estimatr_ht_var_partial", (DL_FUNC) &_estimatr_ht_var_partial, 2},
     {"_estimatr_AtA", (DL_FUNC) &_estimatr_AtA, 1},
     {"_estimatr_Kr", (DL_FUNC) &_estimatr_Kr, 2},

diff --git a/src/horvitz_thompson_variance.cpp b/src/horvitz_thompson_variance.cpp
@@ -28,6 +28,50 @@ double ht_covar_partial(const Eigen::VectorXd & y1,
   return cov_total;
 }
 
+// [[Rcpp::export]]
+double cov_a(const Eigen::VectorXd & y1,
+             const Eigen::VectorXd & y0,
+             const Eigen::MatrixXd & p10,
+             const Eigen::VectorXd & p1,
+             const Eigen::VectorXd & p0) {
+  double cov_total = 0.0;
+
+  for (int i = 0; i < y1.size(); ++i) {
+    for(int j = 0; j < y0.size(); ++j) {
+      if(p10(i, j) == 0) {
+        cov_total += (std::pow(y1(i), 2) / (2.0 * p1(i))) + (std::pow(y0(j), 2) / (2.0 * p0(j)));
+      } else {
+        cov_total += (y1(i) / p1(i)) * (y0(j) / p0(j)) * (p10(i, j) - p1(i) * p0(j)) / p10(i, j);
+      }
+    }
+  }
+
+  return cov_total;
+}
+
+// [[Rcpp::export]]
+double var_a(const Eigen::VectorXd & y,
+             const Eigen::MatrixXd & p) {
+  double var_total = 0.0;
+
+  for (int i = 0; i < y.size(); ++i) {
+    for(int j = 0; j < y.size(); ++j) {
+      if (i != j) {
+        var_total += ((p(i, j) - p(i, i) * p(j, j)) / p(i, j)) *
+          (y(i) / p(i,i)) * (y(j) / p(j,j));
+      } else {
+        if (p(i, j) == 0) {
+          var_total += (std::pow(y(i), 2) / (2.0 * p(i, i))) + (std::pow(y(j), 2) / (2.0 * p(j, j)));
+        } else {
+          var_total += (1 - p(i, i)) * std::pow((y(i) / p(i, i)), 2);
+        }
+      }
+    }
+  }
+
+  return var_total;
+}
+
 // [[Rcpp::export]]
 double ht_var_partial(const Eigen::VectorXd & y,
                       const Eigen::MatrixXd & p) {

diff --git a/tests/testthat/test-difference-in-means.R b/tests/testthat/test-difference-in-means.R
@@ -624,6 +624,7 @@ test_that("DIM matches lm_robust under certain conditions", {
   bldat <- data.frame(bl = rep(1:3, each = 4),
                       z = c(0, 0, 0, 0, rep(0:1, times = 4)),
                       y = rnorm(12))
+
   expect_warning(
     dim_miss <- difference_in_means(y ~ z, blocks = bl, data = bldat),
     "1 block\\(s\\) dropped for containing no variation in treatment"