Merge pull request #46 from r-causal/numerics

add extra tests

Author: malcolmbarrett

tests/testthat/test-extract_weight_data.R

@@ -0,0 +1,91 @@
+# Tests for extract_weight_data() utility function
+
+test_that("extract_weight_data handles NULL input", {
+  result <- extract_weight_data(NULL)
+  expect_null(result)
+})
+
+test_that("extract_weight_data passes through numeric vectors unchanged", {
+  # Test simple numeric vector
+  numeric_weights <- c(1, 2, 3, 4, 5)
+  result <- extract_weight_data(numeric_weights)
+  expect_identical(result, numeric_weights)
+  expect_type(result, "double")
+  
+  # Test with different numeric values
+  uniform_weights <- runif(10, 0.5, 1.5)
+  result2 <- extract_weight_data(uniform_weights)
+  expect_identical(result2, uniform_weights)
+  
+  # Test with extreme values
+  extreme_weights <- c(0.001, 100, 0.5, 2.5)
+  result3 <- extract_weight_data(extreme_weights)
+  expect_identical(result3, extreme_weights)
+})
+
+test_that("extract_weight_data extracts data from psw objects", {
+  skip_if_not_installed("propensity")
+  
+  # Create psw object from numeric data
+  numeric_data <- c(1.2, 0.8, 1.5, 0.9, 1.1)
+  psw_obj <- propensity::psw(numeric_data, estimand = "ate")
+  
+  # Extract should return the underlying numeric data
+  result <- extract_weight_data(psw_obj)
+  expect_equal(result, numeric_data)
+  expect_type(result, "double")
+  
+  # Should not have psw class anymore
+  expect_false(inherits(result, "psw"))
+})
+
+test_that("extract_weight_data preserves values from different psw estimands", {
+  skip_if_not_installed("propensity")
+  
+  numeric_data <- c(2.1, 1.3, 0.7, 1.8, 1.0)
+  
+  # Test different estimands produce same numeric extraction
+  psw_ate <- propensity::psw(numeric_data, estimand = "ate")
+  psw_att <- propensity::psw(numeric_data, estimand = "att")
+  psw_ato <- propensity::psw(numeric_data, estimand = "ato")
+  
+  result_ate <- extract_weight_data(psw_ate)
+  result_att <- extract_weight_data(psw_att)
+  result_ato <- extract_weight_data(psw_ato)
+  
+  expect_equal(result_ate, numeric_data)
+  expect_equal(result_att, numeric_data)
+  expect_equal(result_ato, numeric_data)
+  
+  # All should be identical
+  expect_identical(result_ate, result_att)
+  expect_identical(result_ate, result_ato)
+})
+
+test_that("extract_weight_data works with integer input", {
+  # Test with integer vector (passes through unchanged)
+  int_weights <- c(1L, 2L, 3L, 4L)
+  result <- extract_weight_data(int_weights)
+  expect_equal(result, c(1L, 2L, 3L, 4L))
+  expect_type(result, "integer")
+})
+
+test_that("extract_weight_data preserves NA values", {
+  # Test with NA values in numeric vector
+  numeric_with_na <- c(1.2, NA, 0.8, NA, 1.5)
+  result <- extract_weight_data(numeric_with_na)
+  expect_identical(result, numeric_with_na)
+  expect_equal(sum(is.na(result)), 2)
+})
+
+test_that("extract_weight_data works with edge case values", {
+  # Test with zeros, very small, and very large values
+  edge_weights <- c(0, 1e-10, 1e10, Inf, -Inf)
+  result <- extract_weight_data(edge_weights)
+  expect_identical(result, edge_weights)
+  
+  # Test with single value
+  single_weight <- 1.5
+  result_single <- extract_weight_data(single_weight)
+  expect_identical(result_single, single_weight)
+})

tests/testthat/test-numeric-psw-equivalence.R

@@ -0,0 +1,293 @@
+# Tests for numeric vs psw weight equivalence and backward compatibility
+
+test_that("backward compatibility: all functions work with numeric weights", {
+  # Create numeric weight vectors
+  set.seed(123)
+  n <- 100
+  x <- rnorm(n)
+  g <- sample(c(0, 1), n, replace = TRUE)
+  numeric_weights <- runif(n, 0.5, 2.0)
+  
+  # Test all balance functions with numeric weights
+  expect_no_error(smd_result <- bal_smd(x, g, weights = numeric_weights))
+  expect_true(is.finite(smd_result))
+  
+  expect_no_error(vr_result <- bal_vr(x, g, weights = numeric_weights))
+  expect_true(is.finite(vr_result) && vr_result > 0)
+  
+  expect_no_error(ks_result <- bal_ks(x, g, weights = numeric_weights))
+  expect_true(is.finite(ks_result) && ks_result >= 0 && ks_result <= 1)
+  
+  # Test correlation with numeric weights
+  y <- 2 * x + rnorm(n, sd = 0.5)
+  expect_no_error(corr_result <- bal_corr(x, y, weights = numeric_weights))
+  expect_true(is.finite(corr_result) && corr_result >= -1 && corr_result <= 1)
+  
+  # Test energy balance with numeric weights
+  covariates <- data.frame(x1 = x, x2 = rnorm(n))
+  expect_no_error(energy_result <- bal_energy(covariates, g, weights = numeric_weights))
+  expect_true(is.finite(energy_result) && energy_result >= 0)
+  
+  # Test check_balance with numeric weights
+  test_data <- data.frame(
+    x = x,
+    y = y,
+    g = g,
+    w = numeric_weights
+  )
+  expect_no_error(
+    balance_results <- check_balance(test_data, c(x, y), g, .wts = w, .metrics = "smd")
+  )
+  expect_s3_class(balance_results, "data.frame")
+  expect_true(all(is.finite(balance_results$estimate)))
+})
+
+test_that("numeric and psw weights produce identical results", {
+  skip_if_not_installed("propensity")
+  
+  # Create test data
+  set.seed(42)
+  n <- 100
+  x <- rnorm(n)
+  g <- sample(c(0, 1), n, replace = TRUE)
+  
+  # Create numeric weights
+  numeric_weights <- runif(n, 0.5, 2.0)
+  
+  # Create equivalent psw objects
+  psw_ate <- propensity::psw(numeric_weights, estimand = "ate")
+  psw_att <- propensity::psw(numeric_weights, estimand = "att")
+  psw_ato <- propensity::psw(numeric_weights, estimand = "ato")
+  
+  # Test bal_smd equivalence
+  smd_numeric <- bal_smd(x, g, weights = numeric_weights)
+  smd_psw_ate <- bal_smd(x, g, weights = psw_ate)
+  smd_psw_att <- bal_smd(x, g, weights = psw_att) 
+  smd_psw_ato <- bal_smd(x, g, weights = psw_ato)
+  
+  expect_identical(smd_numeric, smd_psw_ate)
+  expect_identical(smd_numeric, smd_psw_att)
+  expect_identical(smd_numeric, smd_psw_ato)
+  
+  # Test bal_vr equivalence
+  vr_numeric <- bal_vr(x, g, weights = numeric_weights)
+  vr_psw_ate <- bal_vr(x, g, weights = psw_ate)
+  vr_psw_att <- bal_vr(x, g, weights = psw_att)
+  vr_psw_ato <- bal_vr(x, g, weights = psw_ato)
+  
+  expect_identical(vr_numeric, vr_psw_ate)
+  expect_identical(vr_numeric, vr_psw_att)
+  expect_identical(vr_numeric, vr_psw_ato)
+  
+  # Test bal_ks equivalence
+  ks_numeric <- bal_ks(x, g, weights = numeric_weights)
+  ks_psw_ate <- bal_ks(x, g, weights = psw_ate)
+  ks_psw_att <- bal_ks(x, g, weights = psw_att)
+  ks_psw_ato <- bal_ks(x, g, weights = psw_ato)
+  
+  expect_identical(ks_numeric, ks_psw_ate)
+  expect_identical(ks_numeric, ks_psw_att)
+  expect_identical(ks_numeric, ks_psw_ato)
+})
+
+test_that("numeric and psw weights produce identical results for correlation", {
+  skip_if_not_installed("propensity")
+  
+  # Create test data
+  set.seed(42)
+  n <- 100
+  x <- rnorm(n)
+  y <- 2 * x + rnorm(n, sd = 0.5)
+  
+  # Create numeric weights
+  numeric_weights <- runif(n, 0.5, 2.0)
+  
+  # Create equivalent psw object
+  psw_weights <- propensity::psw(numeric_weights, estimand = "ate")
+  
+  # Test bal_corr equivalence
+  corr_numeric <- bal_corr(x, y, weights = numeric_weights)
+  corr_psw <- bal_corr(x, y, weights = psw_weights)
+  
+  expect_identical(corr_numeric, corr_psw)
+})
+
+test_that("numeric and psw weights produce identical results for energy balance", {
+  skip_if_not_installed("propensity")
+  
+  # Create test data
+  set.seed(42)
+  n <- 100
+  covariates <- data.frame(
+    x1 = rnorm(n),
+    x2 = rnorm(n)
+  )
+  g <- sample(c(0, 1), n, replace = TRUE)
+  
+  # Create numeric weights
+  numeric_weights <- runif(n, 0.5, 2.0)
+  
+  # Create equivalent psw object
+  psw_weights <- propensity::psw(numeric_weights, estimand = "ate")
+  
+  # Test bal_energy equivalence
+  energy_numeric <- bal_energy(covariates, g, weights = numeric_weights)
+  energy_psw <- bal_energy(covariates, g, weights = psw_weights)
+  
+  expect_identical(energy_numeric, energy_psw)
+})
+
+test_that("check_balance works equivalently with numeric and psw weights", {
+  skip_if_not_installed("propensity")
+  
+  # Use nhefs_weights dataset
+  data(nhefs_weights)
+  
+  # Extract numeric data from existing psw weights for comparison
+  w_ate_numeric <- vctrs::vec_data(nhefs_weights$w_ate)
+  w_att_numeric <- vctrs::vec_data(nhefs_weights$w_att)
+  
+  # Create new psw objects from the same numeric data
+  w_ate_psw <- propensity::psw(w_ate_numeric, estimand = "ate")
+  w_att_psw <- propensity::psw(w_att_numeric, estimand = "att")
+  
+  # Add numeric weights to test data
+  test_data <- nhefs_weights[1:100, ]
+  test_data$w_ate_numeric <- w_ate_numeric[1:100]
+  test_data$w_att_numeric <- w_att_numeric[1:100]
+  test_data$w_ate_psw <- w_ate_psw[1:100]
+  test_data$w_att_psw <- w_att_psw[1:100]
+  
+  # Test single weight comparison
+  result_numeric <- check_balance(
+    test_data, 
+    c(age, wt71), 
+    qsmk, 
+    .wts = w_ate_numeric,
+    .metrics = "smd",
+    include_observed = FALSE
+  )
+  
+  result_psw <- check_balance(
+    test_data,
+    c(age, wt71),
+    qsmk,
+    .wts = w_ate_psw,
+    .metrics = "smd", 
+    include_observed = FALSE
+  )
+  
+  # Results should be identical except for method names
+  expect_equal(result_numeric$estimate, result_psw$estimate)
+  expect_equal(result_numeric$variable, result_psw$variable)
+  # Method names will differ (w_ate_numeric vs w_ate_psw) but that's expected
+})
+
+test_that("mixed numeric and psw weights work in same function call", {
+  skip_if_not_installed("propensity")
+  
+  # Use nhefs_weights dataset with mixed weight types
+  data(nhefs_weights)
+  test_data <- nhefs_weights[1:100, ]
+  
+  # Create a numeric version of one weight
+  test_data$w_ate_numeric <- vctrs::vec_data(test_data$w_ate)
+  
+  # This should work without error - one numeric, one psw
+  result <- check_balance(
+    test_data,
+    c(age, wt71),
+    qsmk,
+    .wts = c(w_ate_numeric, w_att), # Mix numeric and psw
+    .metrics = "smd",
+    include_observed = FALSE
+  )
+  
+  expect_s3_class(result, "data.frame")
+  expect_true(nrow(result) > 0)
+  expect_true(all(is.finite(result$estimate)))
+  expect_true("w_ate_numeric" %in% result$method)
+  expect_true("w_att" %in% result$method)
+})
+
+test_that("edge cases work identically for numeric and psw weights", {
+  skip_if_not_installed("propensity")
+  
+  # Test with extreme weights
+  set.seed(42)
+  n <- 50
+  x <- rnorm(n)
+  g <- sample(c(0, 1), n, replace = TRUE)
+  
+  # Extreme numeric weights
+  extreme_numeric <- c(rep(0.001, n/2), rep(100, n/2))
+  extreme_psw <- propensity::psw(extreme_numeric, estimand = "ate")
+  
+  # Both should handle extreme weights the same way
+  smd_numeric <- bal_smd(x, g, weights = extreme_numeric)
+  smd_psw <- bal_smd(x, g, weights = extreme_psw)
+  expect_identical(smd_numeric, smd_psw)
+  
+  # Test with uniform weights (should be close to unweighted)
+  uniform_numeric <- rep(1, n)
+  uniform_psw <- propensity::psw(uniform_numeric, estimand = "ate")
+  
+  smd_unweighted <- bal_smd(x, g)
+  smd_uniform_numeric <- bal_smd(x, g, weights = uniform_numeric)
+  smd_uniform_psw <- bal_smd(x, g, weights = uniform_psw)
+  
+  expect_identical(smd_uniform_numeric, smd_uniform_psw)
+  expect_equal(smd_unweighted, smd_uniform_numeric, tolerance = 1e-10)
+})
+
+test_that("NA handling is identical for numeric and psw weights", {
+  skip_if_not_installed("propensity")
+  
+  # Create test data with NAs
+  set.seed(42)
+  n <- 100
+  x <- rnorm(n)
+  x[1:5] <- NA  # Add some NAs
+  g <- sample(c(0, 1), n, replace = TRUE) 
+  
+  # Create weights
+  numeric_weights <- runif(n, 0.5, 2.0)
+  psw_weights <- propensity::psw(numeric_weights, estimand = "ate")
+  
+  # Test na.rm = FALSE (should return NA)
+  smd_numeric_na <- bal_smd(x, g, weights = numeric_weights, na.rm = FALSE)
+  smd_psw_na <- bal_smd(x, g, weights = psw_weights, na.rm = FALSE)
+  
+  expect_identical(smd_numeric_na, smd_psw_na)
+  expect_true(is.na(smd_numeric_na))
+  
+  # Test na.rm = TRUE
+  smd_numeric_narm <- bal_smd(x, g, weights = numeric_weights, na.rm = TRUE)
+  smd_psw_narm <- bal_smd(x, g, weights = psw_weights, na.rm = TRUE)
+  
+  expect_identical(smd_numeric_narm, smd_psw_narm)
+  expect_true(is.finite(smd_numeric_narm))
+})
+
+test_that("all helper functions work with both numeric and psw weights", {
+  skip_if_not_installed("propensity")
+  
+  # Test weighted_quantile
+  set.seed(42)
+  n <- 100
+  x <- rnorm(n)
+  
+  numeric_weights <- runif(n, 0.5, 2.0)
+  psw_weights <- propensity::psw(numeric_weights, estimand = "ate")
+  
+  q_numeric <- weighted_quantile(x, c(0.25, 0.5, 0.75), numeric_weights)
+  q_psw <- weighted_quantile(x, c(0.25, 0.5, 0.75), psw_weights)
+  
+  expect_identical(q_numeric, q_psw)
+  
+  # Test ess function
+  ess_numeric <- ess(numeric_weights)
+  ess_psw <- ess(psw_weights)
+  
+  expect_identical(ess_numeric, ess_psw)
+})