def test_center():
assert Center().__repr__() == 'Center()'
pos = torch.Tensor([[0, 0], [2, 0], [4, 0]])
data = Data(pos=pos)
out = Center()(data).pos
assert out.tolist() == [[-2, 0], [0, 0], [2, 0]]
def test_center():
assert Center().__repr__() == 'Center()'
pos = torch.Tensor([[0, 0], [2, 0], [4, 0]])
data = Data(pos=pos)
data = Center()(data)
assert len(data) == 1
assert data.pos.tolist() == [[-2, 0], [0, 0], [2, 0]]
def _obs(self) -> Tuple[Batch, List[List[int]]]:
"""
returns
-------
Tuple[Batch, List[List[int]]
The Batch object contains the Pytorch Geometric graph representing the molecule. The list of lists of integers
is a list of all the torsions of the molecule, where each torsion is represented by a list of four integers, where the integers
are the indices of the four atoms making up the torsion.
"""
mol = Chem.rdmolops.RemoveHs(self.mol)
conf = mol.GetConformer()
atoms = mol.GetAtoms()
bonds = mol.GetBonds()
node_features = [molecule_features.atom_type_CO(atom) + molecule_features.atom_coords(atom, conf) for atom in atoms]
edge_indices = molecule_features.get_bond_pairs(mol)
edge_attributes = [molecule_features.bond_type(bond) for bond in bonds] * 2
data = Data(
x=torch.tensor(node_features, dtype=torch.float),
edge_index=torch.tensor(edge_indices, dtype=torch.long),
edge_attr=torch.tensor(edge_attributes,dtype=torch.float),
pos=torch.Tensor(conf.GetPositions())
)
data = Center()(data)
data = NormalizeRotation()(data)
data.x[:,-3:] = data.pos
data = Batch.from_data_list([data])
return data, self.nonring
def __init__(self):
self.center = Center()
model.to(device)
optimizer = torch.optim.Adam(model.parameters(),
lr=learn_rate,
weight_decay=p.weight_decay)
writer = SummaryWriter(
comment='model:exp2_3hops_lr:{}_shuffle:{}_seed:{}'.format(
learn_rate, p.shuffle_dataset, p.random_seed))
max_roc_auc = 0
# ---- Training ----
print('Training...')
for epoch in range(1, epochs + 1):
trainset.transform = Compose(
(Center(), RandomRotate(90, epoch % 3), AddPositionalData(), TwoHop()))
validset.transform = Compose(
(Center(), RandomRotate(90, epoch % 3), AddPositionalData(), TwoHop()))
train_loader = DataLoader(trainset,
shuffle=p.shuffle_dataset,
batch_size=p.batch_size)
val_loader = DataLoader(validset,
shuffle=False,
batch_size=p.test_batch_size)
learn_rate = optimizer.param_groups[0][
'lr'] # for when it may be modified during run
model.train()
pred = torch.Tensor()
tr_weights = torch.Tensor()
loss = []
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())
model.to(device)
optimizer = torch.optim.Adam(model.parameters(),
lr=learn_rate,
weight_decay=p.weight_decay)
writer = SummaryWriter(
comment='model:exp2_4hops_lr:{}_shuffle:{}_seed:{}'.format(
learn_rate, p.shuffle_dataset, p.random_seed))
max_roc_auc = 0
# ---- Training ----
print('Training...')
for epoch in range(1, epochs + 1):
trainset.transform = Compose((Center(), RandomRotate(90, epoch % 3),
AddPositionalData(), TwoHop(), TwoHop()))
validset.transform = Compose((Center(), RandomRotate(90, epoch % 3),
AddPositionalData(), TwoHop(), TwoHop()))
train_loader = DataLoader(trainset,
shuffle=p.shuffle_dataset,
batch_size=p.batch_size)
val_loader = DataLoader(validset,
shuffle=False,
batch_size=p.test_batch_size)
learn_rate = optimizer.param_groups[0][
'lr'] # for when it may be modified during run
model.train()
pred = torch.Tensor()
tr_weights = torch.Tensor()
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]
print('Setting up model...{}'.format(n_features))
k_ij
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
if 'pos' in p.dataset: # Is there positional data in the features?
degrees = 0
if epoch > 200:
if epoch < 700:
degrees = 90*((epoch-200)/500)
else:
degrees = 90
rotation_axis = axes[epoch % 3] # only for structural data.
trainset.transform = Compose((RemovePositionalData(),
RandomRotate(degrees, axis=rotation_axis),
AddPositionalData(),
RemoveXYZ()))
validset.transform = RemoveXYZ()
# Using shape index data:
trainset.transform = Compose((Center(), FaceAttributes(),
NodeCurvature(), FaceToEdge(),
TwoHop()), AddShapeIndex())
validset.transform = Compose((Center(), FaceAttributes(),
NodeCurvature(), FaceToEdge(),
TwoHop()), AddShapeIndex())
train_loader = DataListLoader(trainset, shuffle=p.shuffle_dataset, batch_size=p.batch_size)
val_loader = DataListLoader(validset, shuffle=False, batch_size=p.test_batch_size)
learn_rate = optimizer.param_groups[0]['lr'] # for when it may be modified during run
# rotator()
model.train()
first_batch_labels = torch.Tensor()
pred = torch.Tensor()
tr_weights = torch.Tensor()
loss = []