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
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-21 13:49 -0400
1"""
2Unit tests for analysis utilities.
3"""
5import unittest
6import copy
7import torch
8import numpy as np
9import pandas as pd
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)
20from contextualized.easy import ContextualizedRegressor
23class TestTestEachContext(unittest.TestCase):
25 def __init__(self, *args, **kwargs):
26 super().__init__(*args, **kwargs)
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
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"])
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 )
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"]))
68 pval_c0_x1 = pvals.loc[1, "Pvals"]
69 self.assertTrue(pval_c0_x1 < 0.2, "C0 X1 p-value is not significant.")
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 )
78class TestSelectGoodBootstraps(unittest.TestCase):
80 def __init__(self, *args, **kwargs):
81 super().__init__(*args, **kwargs)
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))
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))
106class TestPvals(unittest.TestCase):
107 def __init__(self, *args, **kwargs):
108 super().__init__(*args, **kwargs)
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])
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
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
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()
148if __name__ == "__main__":
149 unittest.main()