Coverage for contextualized/analysis/tests.py: 99%

86 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-21 13:49 -0400

1""" 

2Unit tests for analysis utilities. 

3""" 

4 

5import unittest 

6import copy 

7import torch 

8import numpy as np 

9import pandas as pd 

10 

11 

12from contextualized.analysis import ( 

13 test_each_context, 

14 select_good_bootstraps, 

15 calc_heterogeneous_predictor_effects_pvals, 

16 calc_homogeneous_context_effects_pvals, 

17 calc_homogeneous_predictor_effects_pvals, 

18) 

19 

20from contextualized.easy import ContextualizedRegressor 

21 

22 

23class TestTestEachContext(unittest.TestCase): 

24 

25 def __init__(self, *args, **kwargs): 

26 super().__init__(*args, **kwargs) 

27 

28 def setUp(self): 

29 """ 

30 Shared data setup. 

31 """ 

32 torch.manual_seed(0) 

33 np.random.seed(0) 

34 n_samples = 1000 

35 C = np.random.uniform(0, 1, size=(n_samples, 2)) 

36 X = np.random.uniform(0, 1, size=(n_samples, 2)) 

37 beta = np.concatenate([np.ones((n_samples, 1)), C], axis=1) 

38 Y = np.sum( 

39 beta[:, :2] * X, axis=1 

40 ) # X1 changes effect under C0. C1 has no effect, X0 is constant 

41 

42 self.C_train_df = pd.DataFrame(C, columns=["C0", "C1"]) 

43 self.X_train_df = pd.DataFrame(X, columns=["X0", "X1"]) 

44 self.Y_train_df = pd.DataFrame(Y, columns=["Y"]) 

45 

46 def test_test_each_context(self): 

47 """ 

48 Test that the output shape of the test_each_context function is as expected. 

49 """ 

50 pvals = test_each_context( 

51 ContextualizedRegressor, 

52 self.C_train_df, 

53 self.X_train_df, 

54 self.Y_train_df, 

55 model_kwargs={"encoder_type": "mlp", "layers": 1}, 

56 fit_kwargs={"max_epochs": 1, "learning_rate": 1e-2, "n_bootstraps": 10}, 

57 ) 

58 

59 expected_shape = ( 

60 self.C_train_df.shape[1] 

61 * self.X_train_df.shape[1] 

62 * self.Y_train_df.shape[1], 

63 4, 

64 ) 

65 self.assertEqual(pvals.shape, expected_shape) 

66 self.assertTrue(all(0 <= pval <= 1 for pval in pvals["Pvals"])) 

67 

68 pval_c0_x1 = pvals.loc[1, "Pvals"] 

69 self.assertTrue(pval_c0_x1 < 0.2, "C0 X1 p-value is not significant.") 

70 

71 other_pvals = pvals.drop(1) 

72 self.assertTrue( 

73 all(pval >= 0.2 for pval in other_pvals["Pvals"]), 

74 "Other p-values are significant.", 

75 ) 

76 

77 

78class TestSelectGoodBootstraps(unittest.TestCase): 

79 

80 def __init__(self, *args, **kwargs): 

81 super().__init__(*args, **kwargs) 

82 

83 def setUp(self): 

84 self.C = np.random.uniform(0, 1, size=(100, 2)) 

85 self.X = np.random.uniform(0, 1, size=(100, 2)) 

86 self.Y = np.random.uniform(0, 1, size=(100, 2)) 

87 

88 def test_model_has_fewer_bootstraps(self): 

89 """ 

90 Test that the model has fewer bootstraps after calling select_good_bootstraps. 

91 """ 

92 model = ContextualizedRegressor(n_bootstraps=3) 

93 model.fit(self.C, self.X, self.Y) 

94 Y_pred = model.predict(self.C, self.X, individual_preds=True) 

95 train_errs = np.zeros_like((self.Y - Y_pred) ** 2) 

96 train_errs[0] = 0.1 

97 train_errs[1] = 0.2 

98 train_errs[2] = 0.3 

99 model_copy = copy.deepcopy(model) 

100 select_good_bootstraps(model, train_errs) 

101 self.assertEqual(len(model.models), 1) 

102 self.assertEqual(len(model_copy.models), 3) 

103 self.assertLess(len(model.models), len(model_copy.models)) 

104 

105 

106class TestPvals(unittest.TestCase): 

107 def __init__(self, *args, **kwargs): 

108 super().__init__(*args, **kwargs) 

109 

110 def setUp(self): 

111 np.random.seed(0) 

112 torch.manual_seed(0) 

113 # X1 is a heterogeneous predictor under C0, X0 is a homogeneous predictor 

114 # C0 is a homogeneous context predictor on Y0, C1 is a heterogeneous context predictor on Y1 

115 self.C = np.random.uniform(-1, 1, size=(100, 2)) 

116 self.X = np.random.uniform(-1, 1, size=(100, 2)) 

117 betas = np.concatenate([np.ones((100, 1)), self.C[:, 0, None]], axis=1) 

118 Y0 = np.sum(betas * self.X, axis=1) + self.C[:, 0] 

119 Y1 = np.sum(betas * self.X, axis=1) + self.C[:, 1] 

120 self.Y = np.column_stack([Y0, Y1]) 

121 

122 def test_homogeneous_context_effect_pvals(self): 

123 model = ContextualizedRegressor(n_bootstraps=10) 

124 model.fit(self.C, self.X, self.Y) 

125 pvals = calc_homogeneous_context_effects_pvals(model, self.C) 

126 assert pvals.shape == (self.C.shape[1], self.Y.shape[1]) 

127 assert pvals[0, 0] < 0.2 and pvals[1, 1] < 0.2 

128 assert pvals[0, 1] > 0.2 and pvals[1, 0] > 0.2 

129 

130 def test_homogeneous_predictor_effect_pvals(self): 

131 model = ContextualizedRegressor(n_bootstraps=10) 

132 model.fit(self.C, self.X, self.Y) 

133 pvals = calc_homogeneous_predictor_effects_pvals(model, self.X) 

134 assert pvals.shape == (self.X.shape[1], self.Y.shape[1]) 

135 assert pvals[0, 0] < 0.2 and pvals[0, 1] < 0.2 

136 assert pvals[1, 0] > 0.2 and pvals[1, 1] > 0.2 

137 

138 def test_heterogeneous_predictor_effect_pvals(self): 

139 model = ContextualizedRegressor(n_bootstraps=10) 

140 model.fit(self.C, self.X, self.Y) 

141 pvals = calc_heterogeneous_predictor_effects_pvals(model, self.C) 

142 assert pvals.shape == (self.C.shape[1], self.X.shape[1], self.Y.shape[1]) 

143 assert pvals[0, 1, 0] < 0.2 and pvals[0, 1, 1] < 0.2 

144 pvals[0, 1, 0] = pvals[0, 1, 1] = 1 

145 assert (pvals > 0.2).all() 

146 

147 

148if __name__ == "__main__": 

149 unittest.main()