|  |  | @@ -0,0 +1,50 @@ | 
		
	
		
			
			|  |  |  | from icosagon.fastmodel import FastModel | 
		
	
		
			
			|  |  |  | from icosagon.data import Data | 
		
	
		
			
			|  |  |  | from icosagon.trainprep import prepare_training, \ | 
		
	
		
			
			|  |  |  | TrainValTest | 
		
	
		
			
			|  |  |  | import torch | 
		
	
		
			
			|  |  |  | import time | 
		
	
		
			
			|  |  |  |  | 
		
	
		
			
			|  |  |  |  | 
		
	
		
			
			|  |  |  | def _make_symmetric(x: torch.Tensor): | 
		
	
		
			
			|  |  |  | x = (x + x.transpose(0, 1)) / 2 | 
		
	
		
			
			|  |  |  | return x | 
		
	
		
			
			|  |  |  |  | 
		
	
		
			
			|  |  |  |  | 
		
	
		
			
			|  |  |  | def _symmetric_random(n_rows, n_columns): | 
		
	
		
			
			|  |  |  | return _make_symmetric(torch.rand((n_rows, n_columns), | 
		
	
		
			
			|  |  |  | dtype=torch.float32).round().to_sparse()) | 
		
	
		
			
			|  |  |  |  | 
		
	
		
			
			|  |  |  |  | 
		
	
		
			
			|  |  |  | def _some_data_with_interactions(): | 
		
	
		
			
			|  |  |  | d = Data() | 
		
	
		
			
			|  |  |  | d.add_node_type('Gene', 1000) | 
		
	
		
			
			|  |  |  | d.add_node_type('Drug', 100) | 
		
	
		
			
			|  |  |  |  | 
		
	
		
			
			|  |  |  | fam = d.add_relation_family('Drug-Gene', 1, 0, True) | 
		
	
		
			
			|  |  |  | fam.add_relation_type('Target', | 
		
	
		
			
			|  |  |  | torch.rand((100, 1000), dtype=torch.float32).round().to_sparse()) | 
		
	
		
			
			|  |  |  |  | 
		
	
		
			
			|  |  |  | fam = d.add_relation_family('Gene-Gene', 0, 0, True) | 
		
	
		
			
			|  |  |  | fam.add_relation_type('Interaction', | 
		
	
		
			
			|  |  |  | _symmetric_random(1000, 1000)) | 
		
	
		
			
			|  |  |  |  | 
		
	
		
			
			|  |  |  | fam = d.add_relation_family('Drug-Drug', 1, 1, True) | 
		
	
		
			
			|  |  |  | for i in range(500): | 
		
	
		
			
			|  |  |  | fam.add_relation_type('Side Effect: Nausea %d' % i, | 
		
	
		
			
			|  |  |  | _symmetric_random(100, 100)) | 
		
	
		
			
			|  |  |  | fam.add_relation_type('Side Effect: Infertility %d' % i, | 
		
	
		
			
			|  |  |  | _symmetric_random(100, 100)) | 
		
	
		
			
			|  |  |  | fam.add_relation_type('Side Effect: Death %d' % i, | 
		
	
		
			
			|  |  |  | _symmetric_random(100, 100)) | 
		
	
		
			
			|  |  |  | return d | 
		
	
		
			
			|  |  |  |  | 
		
	
		
			
			|  |  |  |  | 
		
	
		
			
			|  |  |  | def test_fast_model_01(): | 
		
	
		
			
			|  |  |  | d = _some_data_with_interactions() | 
		
	
		
			
			|  |  |  | prep_d = prepare_training(d, TrainValTest(.8, .1, .1)) | 
		
	
		
			
			|  |  |  | model = FastModel(prep_d) | 
		
	
		
			
			|  |  |  | for i in range(10): | 
		
	
		
			
			|  |  |  | t = time.time() | 
		
	
		
			
			|  |  |  | _ = model(None) | 
		
	
		
			
			|  |  |  | print('Model forward took:', time.time() - t) |