Created
November 7, 2019 21:05
-
-
Save dyerrington/24295c5570ebd0107ac13bb1da0e4885 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from statsmodels.stats.power import tt_ind_solve_power | |
from scipy.interpolate import interp1d | |
import matplotlib.pyplot as plt | |
def test_ttest_power_diff(mean, std, sample1_size=None, alpha=0.05, desired_power=0.8, mean_diff_percentages=[0.1, 0.05]): | |
''' | |
calculates the power function for a given mean and std. the function plots a graph showing the comparison between desired mean differences | |
:param mean: the desired mean | |
:param std: the std value | |
:param sample1_size: if None, it is assumed that both samples (first and second) will have same size. The function then will | |
walk through possible sample sizes (up to 100, hardcoded). | |
If this value is not None, the function will check different alternatives for sample 2 sizes up to sample 1 size. | |
:param alpha: alpha default value is 0.05 | |
:param desired_power: will use this value in order to mark on the graph | |
:param mean_diff_percentages: iterable list of percentages. A line per value will be calculated and plotted. | |
:return: None | |
''' | |
fig, ax = plt.subplots() | |
for mean_diff_percent in mean_diff_percentages: | |
mean_diff = mean_diff_percent * mean | |
effect_size = mean_diff / std | |
print('Mean diff: ', mean_diff) | |
print('Effect size: ', effect_size) | |
powers = [] | |
max_size = sample1_size | |
if sample1_size is None: | |
max_size = 100 | |
sizes = np.arange(1, max_size, 2) | |
for sample2_size in sizes: | |
if(sample1_size is None): | |
n = tt_ind_solve_power(effect_size=effect_size, nobs1=sample2_size, alpha=alpha, ratio=1.0, alternative='two-sided') | |
#print('tt_ind_solve_power(alpha=', alpha, 'sample2_size=', sample2_size, '): sample size in *second* group: {:.5f}'.format(n)) | |
else: | |
n = tt_ind_solve_power(effect_size=effect_size, nobs1=sample1_size, alpha=alpha, ratio=(1.0*sample2_size/sample1_size), alternative='two-sided') | |
#print('tt_ind_solve_power(alpha=', alpha, 'sample2_size=', sample2_size, '): sample size *each* group: {:.5f}'.format(n)) | |
powers.append(n) | |
try: # mark the desired power on the graph | |
z1 = interp1d(powers, sizes) | |
results = z1(desired_power) | |
plt.plot([results], [desired_power], 'gD') | |
except Exception as e: | |
print("Error: ", e) | |
#ignore | |
plt.title('Power vs. Sample Size') | |
plt.xlabel('Sample Size') | |
plt.ylabel('Power') | |
plt.plot(sizes, powers, label='diff={:2.0f}%'.format(100*mean_diff_percent)) #, '-gD') | |
plt.legend() | |
plt.show() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment