[Feature Request] Return both Energy Distance and p value when return_all is set to True #6
Description
Currently, when return_all is set to true, pted returns test and permute.
Can it return test, permute, and p-value instead? I'm interested in both the energy distance and p-value.
The proposed code change in pted.py would be:
permute = np.array(permute)
# Compute p-value
if two_tailed:
q = 2 * min(np.sum(permute >= test), np.sum(permute <= test))
else:
q = np.sum(permute >= test)
p = (1.0 + q) / (1.0 + permutations)
if return_all:
return test, permute, p
else:
return p
And the new return signature would become Union[float, tuple[float, np.ndarray, float]]
I can create a pull request if it's more preferable.
Many thanks for putting this on GitHub Connor, very clean API, and it's lightning fast with PyTorch tensors, will definitely cite.