Title:	Hardware-Accelerated Rerandomization for Improved Balance
Version:	0.3
Description:	Provides hardware-accelerated tools for performing rerandomization and randomization testing in experimental research. Using a 'JAX' backend, the package enables exact rerandomization inference even for large experiments with hundreds of billions of possible randomizations. Key functionalities include generating pools of acceptable rerandomizations based on covariate balance, conducting exact randomization tests, and performing pre-analysis evaluations to determine optimal rerandomization acceptance thresholds. The package supports various hardware acceleration frameworks including 'CPU', 'CUDA', and 'METAL', making it versatile across accelerated computing environments. This allows researchers to efficiently implement stringent rerandomization designs and conduct valid inference even with large sample sizes. The package is partly based on Jerzak and Goldstein (2023) <doi:10.48550/arXiv.2310.00861>.
URL:	https://github.com/cjerzak/fastrerandomize-software
BugReports:	https://github.com/cjerzak/fastrerandomize-software/issues
Depends:	R (≥ 3.5.0)
License:	GPL-3
Encoding:	UTF-8
LazyData:	false
Imports:	reticulate, assertthat, utils, stats, graphics
Suggests:	knitr, rmarkdown
VignetteBuilder:	knitr
RoxygenNote:	7.3.3
NeedsCompilation:	no
Packaged:	2025-12-22 22:12:44 UTC; cjerzak
Author:	Fucheng Warren Zhu [aut], Aniket Sachin Kamat [aut], Connor Jerzak [aut, cre], Rebecca Goldstein [aut]
Maintainer:	Connor Jerzak <connor.jerzak@gmail.com>
Repository:	CRAN
Date/Publication:	2025-12-22 22:30:02 UTC

QJEData: Agricultural Treatment Experiment Data

Description

Data from a field experiment studying moral hazard in tenancy contracts in agriculture.

After subsetting, this dataset includes observations on 968 experimental units with the following variables of interest: household composition, treatment assignment, and agricultural outcomes.

Usage

data(QJEData)

Format

A data frame with 968 rows and 7 columns:

children

Numeric (integer). Number of children in the household. Larger numbers may reflect increased household labor needs and different investment or effort incentives.

married

Numeric/binary. Whether the household head is currently married (1) or not (0). Marital status may influence decision-making and risk preferences in farming.

hh_size

Numeric (integer). Household size. Differences in family labor availability or consumption needs can influence effort levels and thus relate to moral hazard in production decisions.

hh_sexrat

Numeric. The ratio of adult men to adult women in the household. Imbalances in the male–female ratio can affect labor division and investment decisions.

treat1

Numeric/binary. Primary treatment indicator (e.g., whether a farmer is offered a specific tenancy contract or cost-sharing arrangement).

R_yield_ELA_sqm

Numeric. Crop yield per square meter (e.g., kilograms of output per square meter). This is a principal outcome measure for evaluating productivity and treatment impact on farm performance.

ELA_Fertil_D

Numeric/binary. Indicator for whether fertilizer was used (1) or not (0). This measures input investment—a key mechanism in moral hazard models (farmers may alter input use under different contracts).

Source

Burchardi, K.B., et al. (2019). Moral hazard: Experimental evidence from tenancy contracts. The Quarterly Journal of Economics, 134(1), 281-347.

YOPData

Description

Data from a re-analysis of the Youth Opportunities Program anti-poverty RCT in Uganda, with satellite imagery neural representations linked to RCT units.

Usage

data(YOPData)

Format

A list containing two data frames:

RCTData

Treatment, outcome, and geolocation information

ImageEmbeddings

CLIP-RSICD neural embeddings of satellite imagery

Source

• Blattman, C., Fiala, N. and Martinez, S. (2020). The Long-term Impacts of Grants on Poverty: Nine-year Evidence from Uganda's Youth Opportunities Program. American Economic Review: Insights, 2(3), 287-304.

• Jerzak, C.T., Johansson, F.D. and Daoud, A. (2023). Image-based Treatment Effect Heterogeneity. Conference on Causal Learning and Reasoning, 531-552. PMLR.

A function to build the environment for fastrerandomize. Builds a conda environment in which 'JAX' and 'np' are installed. Users can also create a conda environment where 'JAX' and 'np' are installed themselves.

Description

A function to build the environment for fastrerandomize. Builds a conda environment in which 'JAX' and 'np' are installed. Users can also create a conda environment where 'JAX' and 'np' are installed themselves.

Usage

build_backend(conda_env = "fastrerandomize_env", conda = "auto")

Arguments

Value

Invisibly returns NULL; this function is used for its side effects of creating and configuring a conda environment for fastrerandomize. This function requires an Internet connection. You can find out a list of conda Python paths via: Sys.which("python")

Examples

## Not run: 
# Create a conda environment named "fastrerandomize_env"
# and install the required Python packages (jax, numpy, etc.)
build_backend(conda_env = "fastrerandomize_env", conda = "auto")

# If you want to specify a particular conda path:
# build_backend(conda_env = "fastrerandomize_env", conda = "/usr/local/bin/conda")

## End(Not run)

Check if 'Python' and 'JAX' are available

Description

This function checks if 'Python' and 'JAX' can be accessed via 'reticulate'. If not, it returns 'NULL' and prints a message suggesting to run 'build_backend()'.

Usage

check_jax_availability(conda_env = "fastrerandomize_env", conda = "auto")

Arguments

Value

Returns ‘TRUE' (invisibly) if both ’Python' and 'JAX' are available; otherwise returns 'NULL'.

Examples

## Not run: 
  check_jax_availability()

## End(Not run)

Compute potential outcome difference in means for a single assignment under a hypothesized tau in base R

Description

Compute potential outcome difference in means for a single assignment under a hypothesized tau in base R

Usage

compute_diff_at_tau_for_oneW_R(Wprime, obsY, obsW, tau)

Arguments

Value

Scalar difference in means for the assignment Wprime.

Diagnostic map from observed (or targeted) balance to precision and stringency

Description

Implements the calculations in Theorem 1 and Appendix D of the paper involving: (1) Realized RMSE from an observed Mahalanobis distance M (or SMDs); (2) Ex-ante RMSE when accepting assignments with M < a (equivalently, with acceptance probability q under complete randomization); (3) largest acceptance probability q that attains a user-specified precision goal, provided via an RMSE target or via a power target (alpha, 1-beta, |tau|).

Usage

diagnose_rerandomization(
  smd = NULL,
  M = NULL,
  d = NULL,
  n_T,
  n_C,
  sigma = NULL,
  R2 = NULL,
  rmse_goal = NULL,
  tau = NULL,
  alpha = 0.05,
  power = 0.8,
  two_sided = TRUE,
  q_min = 1e-09,
  q_tol = 1e-10
)

Arguments

Details

Realized (conditional) RMSE: with standardized/whitened X and typical orientation,

\mathrm{RMSE}_{\text{realized}} \approx \sqrt{\;\sigma^2\!\left(\frac{1}{n_T}+\frac{1}{n_C}\right) + \frac{\sigma_{\text{Prog}}^2}{d}\, M\;} \;=\; \sigma\,\sqrt{\left(\frac{1}{n_T}+\frac{1}{n_C}\right) + \frac{R^2}{1-R^2}\,\frac{M}{d}}\,,

and the conservative upper bound replaces \sigma_{\text{Prog}}^2/d by \sigma_{\text{Prog}}^2.

Ex-ante (design-stage) RMSE under thresholding: with acceptance rule M <= a (acceptance probability q),

\mathbb{E}[\mathrm{MSE}\mid M\le a] = \left(\frac{1}{n_T}+\frac{1}{n_C}\right)\!\left(\sigma^2 + v_a(d)\,\sigma_{\text{Prog}}^2\right),

where v_a(d) = \Pr(\chi^2_{d+2}\le c)/\Pr(\chi^2_d\le c) and c = a / (\tfrac{1}{n_T}+\tfrac{1}{n_C}). Since q = \Pr(\chi^2_d \le c), we can parameterize by q: v(q;d) = \Pr(\chi^2_{d+2}\le \chi^2_{d;q})/q, with \chi^2_{d;q} the q-th quantile.

Power inversion (Appendix D): for two-sided size \alpha and power 1-\beta, a normal approximation suggests the RMSE goal |\tau| / (z_{1-\alpha/2}+z_{1-\beta}).

Value

A list of class '"fastrerandomize_diagnostic"' with elements:

• inputs: Echo of parsed inputs and derived quantities (M, d, S=1/n_T+1/n_C).

• realized: rmse_factor (dimensionless, per sigma), rmse, and conservative rmse_upper_factor, rmse_upper.

• power_check: If tau, alpha, power, sigma, R2 are given, includes z_needed, z_realized, and already_sufficient.

• recommendation: If a target is supplied (via rmse_goal or tau), returns q_star, a_star, v_star, rmse_exante, expected_M_accepted, and expected_draws_per_accept = 1/q_star.

Examples

# Example 1: observed SMDs, realized precision only (dimensionless factors)
smd <- c(0.10, -0.05, 0.08, 0.02)  # standardized mean differences
out1 <- diagnose_rerandomization(smd = smd, n_T = 100, n_C = 100)
print(out1)

# Example 2: same, but supply sigma and R^2 for absolute RMSE
out2 <- diagnose_rerandomization(smd = smd, n_T = 100, n_C = 100, sigma = 1.2, R2 = 0.4)

# Example 3: choose q to hit a power target (two-sided alpha=.05, 80% power, |tau|=0.2)
out3 <- diagnose_rerandomization(smd = smd, n_T = 100, n_C = 100, sigma = 1.2, R2 = 0.4,
                tau = 0.2, alpha = 0.05, power = 0.80)
                
# Analyze rerandomization recommendation given contextual factors
out3$recommendation

# Example 4: choose q to hit an absolute RMSE goal directly
out4 <- diagnose_rerandomization(M = sum(smd^2), d = length(smd), n_T = 100, n_C = 100,
                sigma = 1.2, R2 = 0.4, rmse_goal = 0.25)

Simple difference in means in base R

Description

Simple difference in means in base R

Usage

diff_in_means_R(Y, W)

Arguments

Value

Scalar difference in means: mean(Y|W=1) - mean(Y|W=0).

JAX-accelerated distance calculations

Description

Compute pairwise distances between the rows of one matrix A or two matrices A and B, using JAX-backed, JIT-compiled kernels. Supports common metrics: Euclidean, squared Euclidean, Manhattan, Chebyshev, Cosine, Minkowski (with optional feature weights), and Mahalanobis (with full or diagonal inverse covariance).

The function automatically batches computations to avoid excessive device memory use.

Usage

fast_distance(
  A,
  B = NULL,
  metric = c("euclidean", "sqeuclidean", "manhattan", "chebyshev", "cosine", "minkowski",
    "mahalanobis"),
  p = 2,
  weights = NULL,
  cov_inv = NULL,
  approximate_inv = TRUE,
  squared = FALSE,
  row_batch_size = NULL,
  as_dist = FALSE,
  return_type = "R",
  verbose = FALSE,
  conda_env = "fastrerandomize_env",
  conda_env_required = TRUE
)

Arguments

Details

- **Mahalanobis**: with approximate_inv = TRUE, the diagonal of the pooled covariance is used (variance stabilizer); otherwise a full inverse covariance is used. - **Weighted distances**: supply weights (length p) for "minkowski" and "manhattan" (the latter uses p = 1). - Computations run in float32 and are JIT-compiled with JAX; where applicable, GPU/Metal/CPU device selection follows your existing backend.

Value

An n \times m distance matrix in the format specified by return_type. If as_dist = TRUE and B = NULL (symmetric case), returns a dist object.

See Also

dist

Examples

## Not run: 
# Simple Euclidean within-matrix distances (returns an n x n matrix)
X <- matrix(rnorm(50 * 8), 50, 8)
D <- fast_distance(X, metric = "euclidean")

# Cosine distance between two sets
A <- matrix(rnorm(100 * 16), 100, 16)
B <- matrix(rnorm(120 * 16), 120, 16)
Dcos <- fast_distance(A, B, metric = "cosine")

# Minkowski with p = 3 and feature weights
w <- runif(ncol(A))
Dm3 <- fast_distance(A, B, metric = "minkowski", p = 3, weights = w)

# Mahalanobis (diagonal approx, fast & robust)
Dmah_diag <- fast_distance(X, metric = "mahalanobis", approximate_inv = TRUE)

# Mahalanobis with full inverse (computed internally)
Dmah_full <- fast_distance(X, metric = "mahalanobis", approximate_inv = FALSE)

# Return a base R 'dist' object
D_dist <- fast_distance(X, metric = "euclidean", as_dist = TRUE)

## End(Not run)

Constructor for fastrerandomize randomizations

Description

Create an S3 object of class fastrerandomize_randomizations that stores the randomizations (and optionally balance statistics) generated by functions such as generate_randomizations.

Usage

fastrerandomize_class(
  randomizations,
  balance = NULL,
  fastrr_env = NULL,
  call = NULL
)

Arguments

Value

An object of class fastrerandomize_randomizations.

Constructor for fastrerandomize randomization test objects

Description

Constructor for fastrerandomize randomization test objects

Usage

fastrerandomize_test(p_value, FI, tau_obs, fastrr_env = NULL, call = NULL, ...)

Arguments

Value

An object of class fastrerandomize_test.

Fiducial interval logic in base R, for randomization test

Description

Fiducial interval logic in base R, for randomization test

Usage

find_fiducial_interval_R(
  obsW,
  obsY,
  allW,
  tau_obs,
  alpha = 0.05,
  c_initial = 2,
  n_search_attempts = 500
)

Arguments

Value

2-element numeric vector [lower, upper] or [NA, NA] if none accepted.

Generate randomizations for a rerandomization-based experimental design

Description

This function generates randomizations for experimental design using either exact enumeration or Monte Carlo sampling methods. It provides a unified interface to both approaches while handling memory and computational constraints appropriately.

Usage

generate_randomizations(
  n_units,
  n_treated,
  X = NULL,
  randomization_accept_prob,
  threshold_func = NULL,
  max_draws = 10^6,
  batch_size = 1000,
  randomization_type = "monte_carlo",
  approximate_inv = TRUE,
  file = NULL,
  return_type = "R",
  verbose = TRUE,
  conda_env = "fastrerandomize_env",
  conda_env_required = TRUE
)

Arguments

Details

The function supports two methods of generating randomizations:

1. Exact enumeration: Generates all possible randomizations (memory intensive but exact).

2. Monte Carlo sampling: Generates randomizations through sampling (more memory efficient).

For large problems (e.g., X with >20 rows), Monte Carlo sampling is recommended.

Value

Returns an S3 object with slots:

• assignments An array where each row represents one possible treatment assignment vector containing the accepted randomizations.

• balance_measures A numeric vector containing the balance measure for each corresponding randomization.

• fastrr_env The fastrerandomize environment.

• file_output If file is specified, results are saved to the given file path instead of being returned.

See Also

generate_randomizations_exact for the exact enumeration method. generate_randomizations_mc for the Monte Carlo sampling method.

Examples

## Not run: 
# Generate synthetic data 
X <- matrix(rnorm(20*5), 20, 5)

# Generate randomizations using exact enumeration
RandomizationSet_Exact <- generate_randomizations(
               n_units = nrow(X), 
               n_treated = round(nrow(X)/2), 
               X = X, 
               randomization_accept_prob=0.1,
               randomization_type="exact")

# Generate randomizations using Monte Carlo sampling
RandomizationSet_MC <- generate_randomizations(
               n_units = nrow(X), 
               n_treated = round(nrow(X)/2), 
               X = X,
               randomization_accept_prob = 0.1,
               randomization_type = "monte_carlo",
               max_draws = 100000,
               batch_size = 1000)
 
## End(Not run)

Generate randomizations in base R, filtering by Hotelling's T^2 acceptance

Description

Base R function to either do exact enumeration or Monte Carlo random permutations, then keep the fraction whose T^2 is below the acceptance cutoff.

Usage

generate_randomizations_R(
  n_units,
  n_treated,
  X,
  accept_prob,
  random_type,
  max_draws,
  batch_size
)

Arguments

Value

A list with:

• randomizations: a matrix (rows = accepted assignments).

• balance: numeric vector of T^2 values for each accepted assignment.

Generate Complete Randomizations with Optional Balance Constraints

Description

Generates all possible treatment assignments for a completely randomized experiment, optionally filtering them based on covariate balance criteria. The function can generate either all possible randomizations or a subset that meets specified balance thresholds using Hotelling's T-squared statistic.

Usage

generate_randomizations_exact(
  n_units,
  n_treated,
  X = NULL,
  randomization_accept_prob = 1,
  approximate_inv = TRUE,
  threshold_func = NULL,
  verbose = TRUE,
  conda_env = "fastrerandomize_env",
  conda_env_required = TRUE
)

Arguments

Details

The function works in two main steps: 1. Generates all possible combinations of treatment assignments given n_units and n_treated 2. If covariates (X) are provided, filters these combinations based on balance criteria using the specified threshold function

The balance filtering process uses Hotelling's T-squared statistic by default to measure multivariate balance between treatment and control groups. Randomizations are accepted if their balance measure is below the specified quantile threshold.

Value

The function returns a list with two elements: candidate_randomizations: an array of randomization vectors M_candidate_randomizations: an array of their balance measures.

Note

This function requires 'JAX' and 'NumPy' to be installed and accessible through the reticulate package.

References

Hotelling, H. (1931). The generalization of Student's ratio. The Annals of Mathematical Statistics, 2(3), 360-378.

See Also

generate_randomizations for full randomization generation function. generate_randomizations_mc for the Monte Carlo version.

Examples

## Not run: 
# Generate synthetic data 
X <- matrix(rnorm(60), nrow = 10)  # 10 units, 6 covariates

# Generate balanced randomizations with covariates
BalancedRandomizations <- generate_randomizations_exact(
  n_units = 10,
  n_treated = 5,
  X = X,
  randomization_accept_prob = 0.25  # Keep top 25% most balanced
)

## End(Not run)

Draws a random sample of acceptable randomizations from all possible complete randomizations using Monte Carlo sampling

Description

This function performs sampling with replacement to generate randomizations in a memory-efficient way. It processes randomizations in batches to avoid memory issues and filters them based on covariate balance. The function uses JAX for fast computation and memory management.

Usage

generate_randomizations_mc(
  n_units,
  n_treated,
  X,
  randomization_accept_prob = 1,
  threshold_func = NULL,
  max_draws = 1e+05,
  batch_size = 1000,
  approximate_inv = TRUE,
  verbose = TRUE,
  conda_env = "fastrerandomize_env",
  conda_env_required = TRUE
)

Arguments

Details

The function works by:

1. Generating batches of random permutations.

2. Computing balance measures for each permutation using the provided threshold function.

3. Keeping only the top permutations that meet the acceptance probability threshold.

4. Managing memory by clearing unused objects and caches between batches.

The function uses smaller data types (int8, float16) where possible to reduce memory usage. It also includes assertions to verify array shapes and dimensions throughout.

Value

The function returns a list with two elements: candidate_randomizations: an array of randomization vectors M_candidate_randomizations: an array of their balance measures.

See Also

generate_randomizations for full randomization generation function. generate_randomizations_exact for the exact version.

Examples

## Not run: 
# Generate synthetic data 
X <- matrix(rnorm(100*5), 100, 5) # 5 covariates

# Generate 1000 randomizations for 100 units with 50 treated
rand_less_strict <- generate_randomizations_mc(
               n_units = 100, 
               n_treated = 50, 
               X = X, 
               randomization_accept_prob=0.01, 
               max_draws = 100000,
               batch_size = 1000)

# Use a stricter balance criterion
rand_more_strict <- generate_randomizations_mc(
               n_units = 100, 
               n_treated = 50, 
               X = X, 
               randomization_accept_prob=0.001, 
               max_draws = 1000000,
               batch_size = 1000)

## End(Not run)

Compute Hotelling's T-squared statistic in base R

Description

This function provides a base R implementation of Hotelling's T-squared balance measure, renamed with '_R' for clarity that it is the R-based analog to the JAX version in fastrerandomize.

Usage

hotellingT2_R(X, W)

Arguments

Value

A numeric scalar: Hotelling's T-squared statistic for that assignment.

Plot method for fastrerandomize_randomizations objects

Description

Generates a histogram of the balance measures for accepted randomizations.

Usage

## S3 method for class 'fastrerandomize_randomizations'
plot(x, ...)

Arguments

Value

No return value. This function is called for the side effect of generating a histogram of the accepted balance measures of object with class fastrerandomize_randomizations.

Plot method for fastrerandomize_test objects

Description

Plots a simple visualization of the observed effect and the fiducial interval (if present) on a horizontal axis.

Usage

## S3 method for class 'fastrerandomize_test'
plot(x, ...)

Arguments

Value

No output returned. Performs side effect of plotting fastrerandomize_test class objects.

Print method for fastrerandomize_randomizations objects

Description

Print method for fastrerandomize_randomizations objects

Usage

## S3 method for class 'fastrerandomize_randomizations'
print(x, ...)

Arguments

Value

Prints an object of class fastrerandomize_randomizations.

Print method for fastrerandomize_test objects

Description

Print method for fastrerandomize_test objects

Usage

## S3 method for class 'fastrerandomize_test'
print(x, ...)

Arguments

Value

No return value, prints object of class fastrerandomize_test.

Print timestamped messages with optional quieting

Description

This function prints messages prefixed with the current timestamp in a standardized format. Messages can be suppressed using the quiet parameter.

Usage

print2(text, quiet = FALSE)

Arguments

Details

The function prepends the current timestamp in "YYYY-MM-DD HH:MM:SS" format to the provided message.

Value

No return value, called for side effect of printing with timestamp.

See Also

Sys.time for the underlying timestamp functionality.

Examples

# Print a basic message with timestamp
print2("Processing started")

# Suppress output
print2("This won't show", quiet = TRUE)

# Use in a loop
for(i in 1:3) {
  print2(sprintf("Processing item %d", i))
}

Fast randomization test

Description

Fast randomization test

Usage

randomization_test(
  obsW = NULL,
  obsY = NULL,
  alpha = 0.05,
  candidate_randomizations = NULL,
  candidate_randomizations_array = NULL,
  n0_array = NULL,
  n1_array = NULL,
  findFI = FALSE,
  c_initial = 2,
  conda_env = "fastrerandomize_env",
  conda_env_required = TRUE
)

Arguments

Value

Returns an S3 object with slots:

• p_value A numeric value or vector representing the p-value of the test (or the expected p-value under the prior structure specified in the function inputs).

• FI A numeric vector representing the fiducial interval if findFI=TRUE.

• tau_obs A numeric value or vector representing the estimated treatment effect(s).

• fastrr_env The fastrerandomize environment.

References

• Zhang, Y. and Zhao, Q., 2023. What is a randomization test?. Journal of the American Statistical Association, 118(544), pp.2928-2942.

See Also

generate_randomizations for randomization generation function.

Examples

## Not run: 
# A small synthetic demonstration with 6 units, 3 treated and 3 controls:

# Generate pre-treatment covariates
X <- matrix(rnorm(24*2), ncol = 2)

# Generate candidate randomizations
RandomizationSet_MC <- generate_randomizations(
  n_units = nrow(X),
  n_treated = round(nrow(X)/2),
  X = X,
  randomization_accept_prob = 0.1,
  randomization_type = "monte_carlo",
  max_draws = 100000,
  batch_size = 1000
)

# Generate outcome
W <- RandomizationSet_MC$randomizations[1,]
obsY <- rnorm(nrow(X), mean = 2 * W)

# Perform randomization test
results_base <- randomization_test(
  obsW = W,
  obsY = obsY,
  candidate_randomizations = RandomizationSet_MC$randomizations
)
print(results_base)

# Perform randomization test with fiducial interval
result_fi <- randomization_test(
  obsW = W,
  obsY = obsY,
  candidate_randomizations = RandomizationSet_MC$randomizations,
  findFI = TRUE
)
print(result_fi)

## End(Not run)

Base R randomization test: difference in means + optional fiducial interval

Description

Base R randomization test: difference in means + optional fiducial interval

Usage

randomization_test_R(obsW, obsY, allW, findFI = FALSE, alpha = 0.05)

Arguments

Value

A list with p_value, tau_obs, and (optionally) FI if 'findFI=TRUE'.

Summary method for fastrerandomize_randomizations objects

Description

Summary method for fastrerandomize_randomizations objects

Usage

## S3 method for class 'fastrerandomize_randomizations'
summary(object, ...)

Arguments

Value

A list with summary statistics, printed by default.

Summary method for fastrerandomize_test objects

Description

Summary method for fastrerandomize_test objects

Usage

## S3 method for class 'fastrerandomize_test'
summary(object, ...)

Arguments

Value

Returns an (invisible) list with a summary of fastrerandomize_test class objects.