Coverage for /opt/conda/lib/python3.14/site-packages/medil/_independence_testing.py: 83%

1"""Independence testing on samples of random variables.""" 

2 

3from multiprocessing import Pool, cpu_count 

4from typing import NamedTuple, Optional 

5 

6import numpy as np 

7import numpy.typing as npt 

8from scipy.spatial.distance import pdist, squareform 

9from scipy.stats import chatterjeexi, chi2, chi2_contingency 

10 

11 

12def _dcov(samples): 

13 r"""Compute sample distance covariance matrix. 

14 

15 Parameters 

16 ---------- 

17 samples : 2d numpy array of floats 

18 A :math:`N \times M` matrix with :math:`N` samples of 

19 :math:`M` random variables. 

20 

21 Returns 

22 ------- 

23 2d numpy array 

24 A square matrix :math:`C`, where :math:`C_{i,j}` is the sample 

25 distance covariance between random variables :math:`R_i` and 

26 :math:`R_j`. 

27 

28 Notes 

29 ----- 

30 Trades time complexity for space complexity, so it can be too 

31 memory-intensive for larger datasets, in which case recommend to 

32 use xicor or distance_correlation_t_test from the dcor package. 

33 """ 

34 num_samps, num_feats = samples.shape 

35 num_pairs = num_samps * (num_samps - 1) // 2 

36 dists = np.zeros((num_feats, num_pairs)) 

37 d_bars = np.zeros(num_feats) 

38 # compute doubly centered distance matrix for every feature: 

39 for feat_idx in range(num_feats): 

40 n = num_samps 

41 t = np.tile 

42 # raw distance matrix: 

43 d = squareform(pdist(samples[:, feat_idx].reshape(-1, 1), "cityblock")) 

44 # doubly centered: 

45 d_bar = d.mean() 

46 d -= t(d.mean(0), (n, 1)) + t(d.mean(1), (n, 1)).T - t(d_bar, (n, n)) 

47 d = squareform(d, checks=False) # ignore assymmetry due to numerical error 

48 dists[feat_idx] = d 

49 d_bars[feat_idx] = d_bar 

50 return dists @ dists.T / num_samps**2, d_bars 

51 

52 

53def _estimate_UDG(sample, method="dcov_fast", significance_level=0.05): 

54 samp_size, num_feats = sample.shape 

55 

56 if isinstance(method, np.ndarray): 

57 p_vals = method 

58 udg = p_vals < significance_level 

59 elif method == "dcov_fast": 

60 cov, d_bars = _dcov(sample) 

61 crit_val = chi2(1).ppf(1 - significance_level) 

62 test_val = samp_size * cov / np.outer(d_bars, d_bars) 

63 udg = test_val >= crit_val 

64 p_vals = None 

65 elif method == "g-test": 

66 if not np.issubdtype(sample.dtype, np.integer): 

67 raise ValueError( 

68 f"g-test requires integer-valued data; got dtype {sample.dtype!r}" 

69 ) 

70 p_vals = np.zeros((num_feats, num_feats), float) 

71 idxs, jdxs = np.triu_indices(num_feats, 1) 

72 sample_iter = (sample[:, i_j].T for i_j in zip(idxs, jdxs)) 

73 with Pool(max(1, int(0.75 * cpu_count()))) as p: 

74 p_vals[idxs, jdxs] = p_vals[jdxs, idxs] = np.fromiter( 

75 p.imap(_g_test, sample_iter, 100), float 

76 ) 

77 udg = p_vals < significance_level 

78 else: 

79 p_vals = np.zeros((num_feats, num_feats), float) 

80 idxs, jdxs = np.triu_indices(num_feats, 1) 

81 zipped = zip(idxs, jdxs) 

82 sample_iter = (sample[:, i_j].T for i_j in zipped) 

83 # if method == "dcov_big": 

84 # can use distance_correlation_t_test from dcor package 

85 if method == "xicor": 

86 test = _xicor_test 

87 with Pool(max(1, int(0.75 * cpu_count()))) as p: 

88 p_vals[idxs, jdxs] = p_vals[jdxs, idxs] = np.fromiter( 

89 p.imap(test, sample_iter, 100), float 

90 ) 

91 udg = p_vals < significance_level 

92 np.fill_diagonal(udg, False) 

93 return udg, p_vals 

94 

95 

96def _g_test(x_y): 

97 x, y = x_y 

98 x_vals, xi = np.unique(x, return_inverse=True) 

99 y_vals, yi = np.unique(y, return_inverse=True) 

100 table = np.zeros((len(x_vals), len(y_vals)), dtype=int) 

101 np.add.at(table, (xi, yi), 1) 

102 _, p, _, _ = chi2_contingency(table, lambda_="log-likelihood") 

103 return p 

104 

105 

106def _xicor_test(x_y): 

107 x, y = x_y 

108 xi, pvalue = _xicorr(x, y) 

109 return pvalue 

110 

111 

112class _XiCorrResult(NamedTuple): 

113 correlation: float 

114 pvalue: Optional[float] 

115 

116 

117def _xicorr(x: npt.ArrayLike, y: npt.ArrayLike, ties: bool = True) -> _XiCorrResult: 

118 """Compute Chatterjee's xi correlation coefficient.""" 

119 x = np.asarray(x).ravel() 

120 y = np.asarray(y).ravel() 

121 if x.size != y.size: 

122 raise ValueError( 

123 "All inputs to `_xicorr` must be of the same " 

124 f"size, found x-size {x.size} and y-size {y.size}" 

125 ) 

126 result = chatterjeexi(x, y, y_continuous=not ties) 

127 return _XiCorrResult(result.statistic, result.pvalue)