def test_two_hop(): assert TwoHop().__repr__() == 'TwoHop()' edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]]) edge_attr = torch.tensor([1, 2, 3, 1, 2, 3], dtype=torch.float) data = Data(edge_index=edge_index, edge_attr=edge_attr) data = TwoHop()(data) edge_index, edge_attr = data.edge_index, data.edge_attr assert edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], [1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2]] assert edge_attr.tolist() == [1, 2, 3, 1, 0, 0, 2, 0, 0, 3, 0, 0]
def test_two_hop(): assert TwoHop().__repr__() == 'TwoHop()' edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]]) edge_attr = torch.tensor([1, 2, 3, 1, 2, 3], dtype=torch.float) data = Data(edge_index=edge_index, edge_attr=edge_attr, num_nodes=4) data = TwoHop()(data) assert len(data) == 2 assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], [1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2]] assert data.edge_attr.tolist() == [1, 2, 3, 1, 0, 0, 2, 0, 0, 3, 0, 0] data = Data(edge_index=edge_index, num_nodes=4) data = TwoHop()(data) assert len(data) == 1 assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], [1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2]]
def apply_pretransforms(pre_transforms=None): from dataset import Structures # Structures should check already whether these pre_transforms have been computed if pre_transforms is None: trainset = Structures(root='./datasets/{}_train/'.format(p.dataset), pre_transform=Compose( (FaceAttributes(), NodeCurvature(), FaceToEdge(), TwoHop()))) testset = Structures(root='./datasets/{}_test/'.format(p.dataset), pre_transform=Compose( (FaceAttributes(), NodeCurvature(), FaceToEdge(), TwoHop()))) else: trainset = Structures(root='./datasets/{}_train/'.format(p.dataset), pre_transform=pre_transforms) testset = Structures(root='./datasets/{}_test/'.format(p.dataset), pre_transform=pre_transforms) return trainset, testset
np.random.seed(p.random_seed) learn_rate = p.learn_rate modelpath = make_model_directory() if str(device) == 'cuda:0': epochs = p.epochs else: epochs = 20 # ---- Importing and structuring Datasets and Model ---- print('Importing structures.') # Remember!!! Shape Index can only be computed on local. Add other transforms after # Pre_tranform step to not contaminate the data. trainset = Structures(root='./datasets/masif_site_train/', pre_transform=Compose((FaceAttributes(), NodeCurvature(), FaceToEdge(), TwoHop()))) # Define transform in epoch, so that rotation occurs around Δ axis every time. validset = Structures(root='./datasets/masif_site_test/', pre_transform=Compose((FaceAttributes(), NodeCurvature(), FaceToEdge(), TwoHop()))) if p.shuffle_dataset: trainset = trainset.shuffle() n_features = trainset.get(0).x.shape[1] # ---- Import previous model to allow deep network to train ------------- model = p.model_type(9, heads=p.heads).to(cpu) model.to(device) optimizer = torch.optim.Adam(model.parameters(),
np.random.seed(p.random_seed) learn_rate = p.learn_rate modelpath = make_model_directory() if str(device) == 'cuda:0': epochs = p.epochs else: epochs = 20 # ---- Importing and structuring Datasets and Model ---- print('Importing structures.') # Remember!!! Shape Index can only be computed on local. Add other transforms after # Pre_tranform step to not contaminate the data. dataset = Structures(root='./datasets/{}_train/'.format(p.dataset), pre_transform=Compose((FaceAttributes(), NodeCurvature(), FaceToEdge(), TwoHop())), transform=AddShapeIndex()) samples = len(dataset) assert (p.validation_split < 0.3) cutoff = int(np.floor(samples * (1 - p.validation_split))) trainset = dataset[:cutoff] validset = dataset[cutoff:] maskedset = validset[:int(len(validset) / 2)] validset = validset[int(len(validset) / 2):] if p.shuffle_dataset: trainset = trainset.shuffle() n_features = trainset.get(0).x.shape[1] print('Setting up model...') models = [
def generate_surface(model_type, model_path, pdb_code, use_structural_data=False): ''' Save the surface prediction for a particular structure. ''' from dataset import read_ply ''' converter = Compose((Center(), FaceAttributes(), NodeCurvature(), FaceToEdge(), TwoHop(), AddShapeIndex())) ''' converter = Compose((Center(), FaceAttributes(), NodeCurvature(), FaceToEdge(), TwoHop(), AddShapeIndex())) path = glob('./structures/test/{}.ply'.format(pdb_code))[0] name = path.split('/')[-1] structure = read_ply(path) face = structure.face structure = converter(structure) device = torch.device('cpu') structure.x.shape[1] if p.heads is not None: model = model_type(structure.x.shape[1], heads=p.heads) else: model = model_type(structure.x.shape[1]) model.load_state_dict(torch.load(model_path, map_location=device)) model.eval() prediction = model(structure) rounded = prediction.round() # ---- Make directory --- dir = model_path.split('models/', 1)[1].split('.')[0] full_path = os.path.expanduser( '~/Desktop/Drawer/LPDI/masif-tools/surfaces/' + dir) folder = full_path.rsplit('/', 1)[0] if not os.path.exists(folder): os.mkdir(folder) if not os.path.exists(full_path): os.mkdir(full_path) save_ply(filename='./surfaces/{}/{}'.format(dir, name), vertices=structure.pos.detach().numpy(), normals=structure.norm.detach().numpy(), faces=face.t().detach().numpy(), charges=structure.x[:, 0].reshape(-1, 1).detach().numpy(), hbond=structure.x[:, 1].reshape(-1, 1).detach().numpy(), hphob=structure.x[:, 2].reshape(-1, 1).detach().numpy(), iface=prediction.detach().numpy()) save_ply(filename='./surfaces/{}/r_{}'.format(dir, name), vertices=structure.pos.detach().numpy(), normals=structure.norm.detach().numpy(), faces=face.t().detach().numpy(), charges=structure.x[:, 0].reshape(-1, 1).detach().numpy(), hbond=structure.x[:, 1].reshape(-1, 1).detach().numpy(), hphob=structure.x[:, 2].reshape(-1, 1).detach().numpy(), iface=rounded.detach().numpy())
np.random.seed(p.random_seed) learn_rate = p.learn_rate modelpath = make_model_directory() if str(device) == 'cuda:0': epochs = p.epochs else: epochs = 20 # ---- Importing and structuring Datasets and Model ---- print('Importing structures.') # Remember!!! Shape Index can only be computed on local. Add other transforms after # Pre_tranform step to not contaminate the data. trainset = Structures(root='./datasets/masif_site_train/', pre_transform=Compose((FaceAttributes(), NodeCurvature(), FaceToEdge(), TwoHop()))) # Define transform in epoch, so that rotation occurs around Δ axis every time. validset = Structures(root='./datasets/masif_site_test/', pre_transform=Compose((FaceAttributes(), NodeCurvature(), FaceToEdge(), TwoHop()))) if p.shuffle_dataset: trainset = trainset.shuffle() n_features = trainset.get(0).x.shape[1] # ---- Import previous model to allow deep network to train ------------- model = p.model_type(9, heads=p.heads).to(cpu) model.to(device) optimizer = torch.optim.Adam(model.parameters(),
learn_rate = p.learn_rate modelpath = make_model_directory() if str(device) == 'cuda:0': epochs = p.epochs else: epochs = 20 # ---- Importing and structuring Datasets and Model ---- print('Importing structures.') # Remember!!! Shape Index can only be computed on local. Add other transforms after # Pre_tranform step to not contaminate the data. dataset = Structures(root='./datasets/{}_train/'.format(p.dataset), pre_transform=Compose((FaceAttributes(), NodeCurvature(), FaceToEdge(), TwoHop())), transform=Compose((AddShapeIndex(), Center(), AddPositionalData()))) samples = len(dataset) assert(p.validation_split < 0.3) cutoff = int(np.floor(samples*(1-p.validation_split))) trainset = dataset[:cutoff] validset = dataset[cutoff:] maskedset = validset[:int(len(validset)/2)] validset = validset[int(len(validset)/2):] if p.shuffle_dataset: trainset = trainset.shuffle() n_features = trainset.get(0).x.shape[1]
def __init__(self, hops): super(MultiHop, self).__init__() self.nhops = hops self.converter = TwoHop() assert hops > 1
dataset # Applying pretransformation once and saving datasets! # New strategy... load the data first, without pretransforms, # then load into a dataloader and batch apply the transforms. # Why is it not using gpu? from utils import apply_pretransforms from dataset import Structures_SI_mem, Structures from torch_geometric.transforms import * from transforms import * dataset = Structures_SI(root='./datasets/thous_train/', pre_transform=Compose( (Center(), FaceAttributes(), NodeCurvature(), FaceToEdge(), TwoHop(), AddShapeIndex()))) dataset = Structures_SI_mem(root='./datasets/thous_train/', pre_transform=Compose( (Center(), FaceAttributes(), NodeCurvature(), FaceToEdge(), TwoHop(), AddShapeIndex()))) dataset[0] dataset[0] apply_pretransforms() # Checking why my transformations keep crashing... import torch from torch_geometric.transforms import Compose, FaceToEdge, TwoHop, Center from transforms import *
device = torch.device('cuda:0') # reproducibility torch.manual_seed(p.random_seed) np.random.seed(p.random_seed) learn_rate = p.learn_rate modelpath = make_model_directory() if str(device) == 'cuda:0': epochs = p.epochs else: epochs = 20 # ---- Importing and structuring Datasets and Model ---- if p.twohop is True: print("Adding two-hop edges to data graphs") converter = TwoHop() else: converter = None print('Importing structures.') trainset = Structures(root='./datasets/{}_train/'.format(p.dataset), prefix=p.dataset) samples = len(trainset) cutoff = int(np.floor(samples*(1-p.validation_split))) validset = trainset[cutoff:] trainset = trainset[:cutoff] if p.shuffle_dataset: trainset = trainset.shuffle() n_features = trainset.get(0).x.shape[1]