File size: 3,736 Bytes
07fcdfe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | import random
import unittest
from torch import load as torch_load
from pdgrapher._models import GCNArgs, GCNBase, ResponsePredictionModel, PerturbationDiscoveryModel
import torch
import os
torch.set_num_threads(5)
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
class TestGCNArgs(unittest.TestCase):
DEFAULT_GCNARGS_DICT = {
"positional_features_dims": 16,
"embedding_layer_dim": 16,
"dim_gnn": 16,
"num_vars": 1,
"n_layers_gnn": 1,
"n_layers_nn": 2,
}
def test_empty_args(self):
gcnargs_obj = GCNArgs()
def test_from_to_dict(self):
gcnargs_dict = {}
gcnargs_obj = GCNArgs()
self.assertEqual(GCNArgs.from_dict(gcnargs_dict), gcnargs_obj)
gcnargs_dict = {"n_layers_gnn": 3, "embedding_layer_dim": 4, "n_layers_nn": 1}
gcnargs_obj = GCNArgs(n_layers_gnn=3, embedding_layer_dim=4, n_layers_nn=1)
obj = GCNArgs.from_dict(gcnargs_dict)
self.assertEqual(obj, gcnargs_obj)
gcnargs_dict = {**self.DEFAULT_GCNARGS_DICT, **gcnargs_dict}
self.assertDictEqual(obj.to_dict(), gcnargs_dict)
class TestGCNBase(unittest.TestCase):
edge_index = torch_load("data/processed/torch_data/real_lognorm/edge_index_A549.pt")
def test_empty_args(self):
gcnargs = GCNArgs()
gcnbase = GCNBase(gcnargs, "response", self.edge_index)
with self.assertRaises(NotImplementedError):
gcnbase.forward(None)
def test_some_args(self):
gcnargs = GCNArgs(n_layers_gnn=4, n_layers_nn=4)
gcnbase = GCNBase(gcnargs, "response", self.edge_index)
# test if out_fun is identity
for i in range(100):
x = random.randint(0, 100)
self.assertEqual(gcnbase.out_fun(x), x)
# test number of GNN layers
self.assertEqual(len(gcnbase.convs), 4)
# test number of NN layers
self.assertEqual(len(gcnbase.mlp), 1+4+1) # input + specified + output layers
class TestResponsePredictionModel(unittest.TestCase):
edge_index = torch_load("data/processed/torch_data/real_lognorm/edge_index_A549.pt")
def test_empty_args(self):
gcnargs = GCNArgs()
model = ResponsePredictionModel(gcnargs, self.edge_index)
def test_some_args(self):
gcnargs = GCNArgs(n_layers_gnn=2, n_layers_nn=2)
model = ResponsePredictionModel(gcnargs, self.edge_index)
# test if out_fun is identity
for i in range(100):
x = random.randint(0, 100)
self.assertEqual(model.out_fun(x), x)
# test number of GNN layers
self.assertEqual(len(model.convs), 2)
# test number of NN layers
self.assertEqual(len(model.mlp), 1+2+1) # input + specified + output layers
class TestPerturbationDiscoveryModel(unittest.TestCase):
edge_index = torch_load("data/processed/torch_data/real_lognorm/edge_index_A549.pt")
def test_empty_args(self):
gcnargs = GCNArgs()
model = PerturbationDiscoveryModel(gcnargs, self.edge_index)
def test_some_args(self):
gcnargs = GCNArgs(n_layers_gnn=2, n_layers_nn=2)
model = PerturbationDiscoveryModel(gcnargs, self.edge_index)
# test if out_fun is identity
for i in range(100):
x = random.randint(0, 100)
self.assertEqual(model.out_fun(x), x)
# test number of GNN layers
self.assertEqual(len(model.convs), 2)
# test number of NN layers
self.assertEqual(len(model.mlp), 1+2+1) # input + specified + output layers
if __name__ == "__main__":
unittest.TestLoader.sortTestMethodsUsing = None
unittest.main()
|