diff --git a/R/RcppExports.R b/R/RcppExports.R index c6d54ba4..1dea53c7 100644 --- 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 index 0772e4d6..abfdefa2 100644 --- 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 index 7fcba591..ef14c0ff 100644 --- 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 index 89ff4fc9..98b48abe 100644 --- 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"