#!/usr/bin/env Rscript

# Visualize the power of a correlation test as a function of sample
# size and effect size (correlation coefficient)
#
# Author: Denes Turei (turei.denes@gmail.com)

packages <- c('pwr', 'tidyr', 'dplyr', 'magrittr', 'withr', 'ggplot2')
local_lib <- '~/r_packages'
.libPaths(new = local_lib)

for(pkg in packages){

    if(!require(pkg, character.only = TRUE)){

        dir.create(local_lib, showWarnings = FALSE)
        install.packages(pkg, lib = local_lib)
        library(pkg, character.only = TRUE)

    }

}


sample_sizes <- seq(10L, 100L, 10L)
corr_coeffs <- seq(.1, 1.0, .1)

powers <-
    tidyr::crossing(
        sample_size = sample_sizes,
        corr_coeff = corr_coeffs
    ) %>%
    dplyr::mutate(
        power = pwr::pwr.r.test(
            n = sample_size,
            r = corr_coeff,
            sig.level = .05
        )$power
    )

power_heatmap <-
    ggplot2::ggplot(
        powers,
        ggplot2::aes(
            x = sample_size,
            y = corr_coeff,
            fill = power
        )
    ) +
    ggplot2::geom_tile() +
    ggplot2::xlab('Sample size (n)') +
    ggplot2::ylab('Correlation coefficient (r)') +
    ggplot2::scale_fill_viridis_c(
        guide = ggplot2::guide_colorbar(
            # if cairo_pdf fails to render the greek letters,
            # remove them from the text below:
            title = 'Power of\ncorrelation test\n(1 - β; α = 0.05)')
    ) +
    ggplot2::theme_bw()

# if cairo_pdf is not available on your system, change it to with_pdf
# and remove the greek letters from the legend title above
withr::with_cairo_pdf(
    'correlation_power.pdf',
    width = 4,
    height = 2.7,
    print(power_heatmap)
)