def forward(self, sample):
angles, subgraph = sample
for idx in range(self.repeats):
tertiary, _ = self.position_lookup(angles, torch.zeros_like(subgraph.indices))
asin = angles.sin()
acos = angles.cos()
afeat = torch.cat((asin, acos), dim=1)
features = ts.scatter.batched(self.preprocess_correct, afeat, subgraph.indices)
ors = self.orientations(tertiary)
pos = tertiary[:, 1]
inds = torch.arange(0, pos.size(0), dtype=torch.float, device=pos.device).view(-1, 1)
distances = torch.cat((pos, ors, inds), dim=1)
dist, structure = self.knn_structure(tertiary, subgraph)
neighbour_pos = (pos[:, None] - pos[structure.connections] + 1e-6)
dist = (neighbour_pos).contiguous()
dist = dist.norm(dim=2, keepdim=True)
dist = gaussian_rbf(dist, *self.rbf)
distance_data = RelativeStructure(structure, self.rbf)
relative_data = distance_data.message(
distances, distances
)
relative_structure = OrientationStructure(structure, relative_data)
correction = self.correction(features, relative_structure)
angles = angles + self.corrected_angles(correction)
result = ts.PackedTensor(angles, lengths=list(subgraph.counts))
return result, subgraph
def forward(self, inputs):
tertiary, subgraph = inputs
angles = tertiary
asin = angles.sin()
acos = angles.cos()
afeat = torch.cat((asin, acos), dim=1)
angle_result = self.angle_result(afeat)
features = ts.scatter.batched(self.preprocess, afeat, subgraph.indices)
tertiary, _ = self.lookup(tertiary, torch.zeros_like(subgraph.indices))
ors = self.orientations(tertiary)
pos = tertiary[:, 1]
inds = torch.arange(0, pos.size(0), dtype=torch.float, device=pos.device).view(-1, 1)
distances = torch.cat((pos, ors, inds), dim=1)
dist, structure = self.knn_structure(tertiary, subgraph)
neighbour_pos = (pos[:, None] - pos[structure.connections] + 1e-6)
dist = (neighbour_pos).contiguous()
dist = dist.norm(dim=2, keepdim=True)
dist = gaussian_rbf(dist, *self.rbf)
distance_data = RelativeStructure(structure, self.rbf)
relative_data = distance_data.message(
distances, distances
)
relative_structure = OrientationStructure(structure, relative_data)
encoding = self.encoder(features, relative_structure)
encoding = ts.scatter.batched(self.postprocess, encoding, subgraph.indices)
encoding = torch.cat((features, encoding), dim=1)
result = self.result(encoding)#self.result(ts.scatter.mean(encoding, subgraph.indices))
result = torch.cat((result, angle_result), dim=0)
return result
def forward(self, features, sequence, pair_features, structure, protein=None):
# featurize distances
distance = pair_features[:, :, 0]
distance = gaussian_rbf(distance.view(-1, 1), *self.rbf).reshape(distance.size(0), distance.size(1), -1)
pair_features = torch.cat((distance, pair_features[:, :, 1:]), dim=2)
relative_structure = OrientationStructure(structure, pair_features)
encoding = self.encoder(features, relative_structure)
# initial evaluation
sequence = self.prepare_sequence(sequence)
masked_structure = MaskedStructure(
structure, pair_features, sequence, encoding
)
result = self.decoder(encoding, masked_structure)
# differentiable scheduled sampling
for idx in range(self.schedule):
samples = self.sequence_embedding(hard_one_hot(result))
mask = torch.rand(samples.size(0), device=samples.device) < 0.25
mask = mask.unsqueeze(1).float()
sequence = mask * samples + (1 - mask) * sequence
masked_structure = MaskedStructure(
structure, relative_data, sequence, encoding
)
result = self.decoder(encoding, masked_structure)
return result
def forward(self,
angle_features,
sequence,
pair_features,
structure,
protein=None):
features = torch.cat((angle_features, sequence), dim=1)
# featurize distances
distance = pair_features[:, :, 0]
distance = gaussian_rbf(distance.view(-1, 1),
*self.rbf).reshape(distance.size(0),
distance.size(1), -1)
print(distance.shape, pair_features.shape)
pair_features = torch.cat((distance, pair_features[:, :, 1:]), dim=2)
relative_structure = ...
if self.sequence:
relative_structure = SequenceOrientationStructure(
structure, pair_features, sequence)
else:
relative_structure = OrientationStructure(structure, pair_features)
encoding = self.encoder(features, relative_structure)
if self.local is not None and protein is not None:
combined = torch.cat((encoding, features), dim=1)
encoding = ts.scatter.batched(self.local, combined,
protein.indices)
result = self.decoder(encoding)
return result
def compare(self, source, target):
x, y = source[:, :3], target[:, :3]
x_i, y_i = target[:, -1], source[:, -1]
distance = (x - y).norm(dim=1, keepdim=True)
distance_sin = torch.sin((x_i - y_i) / 10)[:, None]
distance_cos = torch.cos((x_i - y_i) / 10)[:, None]
return torch.cat(
(gaussian_rbf(distance, *self.rbf), distance_sin, distance_cos),
dim=1)
def compare(self, source, target):
x, y = target[:, :3], source[:, :3]
x_o, y_o = target[:, 3:-1].reshape(-1, 3,
3), source[:,
3:-1].reshape(-1, 3, 3)
x_i, y_i = target[:, -1], source[:, -1]
distance, direction, rotation = relative_orientation(x, y, x_o, y_o)
distance_sin = torch.sin((x_i - y_i) / 10)[:, None]
distance_cos = torch.cos((x_i - y_i) / 10)[:, None]
return torch.cat((gaussian_rbf(distance, *self.rbf), direction,
rotation, distance_sin, distance_cos),
dim=1)
def forward(self, tertiary, noise, sequence, subgraph):
noise = noise / 3.15
offset = (torch.log(noise) / torch.log(torch.tensor(0.60))).long()
condition = torch.zeros(tertiary.size(0),
10,
device=tertiary.device,
dtype=torch.float)
print(condition.shape, tertiary.shape, offset.shape, noise.shape)
condition[torch.arange(offset.size(0)), offset.view(-1)] = 1
angles = tertiary
asin = angles.sin()
acos = angles.cos()
tertiary, _ = self.lookup(tertiary, torch.zeros_like(subgraph.indices))
if self.angles:
afeat = torch.cat((asin, acos, condition), dim=1)
else:
afeat = torch.cat(
(tertiary.reshape(tertiary.size(0), -1), condition), dim=1)
features = ts.scatter.batched(self.preprocess, afeat, subgraph.indices)
if self.conditional:
features = sequence
ors = self.orientations(tertiary)
pos = tertiary[:, 1]
inds = torch.arange(0,
pos.size(0),
dtype=torch.float,
device=pos.device).view(-1, 1)
distances = torch.cat((pos, ors, inds), dim=1)
dist, structure = self.knn_structure(tertiary, subgraph)
neighbour_pos = (pos[:, None] - pos[structure.connections] + 1e-6)
dist = (neighbour_pos).contiguous()
dist = dist.norm(dim=2, keepdim=True)
dist = gaussian_rbf(dist, *self.rbf)
distance_data = RelativeStructure(structure, self.rbf)
relative_data = distance_data.message(distances, distances)
relative_structure = OrientationStructure(structure, relative_data)
encoding = self.encoder(features, relative_structure)
encoding = ts.scatter.batched(self.postprocess, encoding,
subgraph.indices)
encoding = torch.cat((features, encoding), dim=1)
result = self.result(encoding)
return result
def forward(self, angle_features, sequence, pair_features, structure):
features = torch.cat((angle_features, sequence), dim=1)
# featurize distances
distance = pair_features[:, :, 0]
distance = gaussian_rbf(distance.view(-1, 1),
*self.rbf).reshape(distance.size(0),
distance.size(1), -1)
print(distance.shape, pair_features.shape)
pair_features = torch.cat((distance, pair_features[:, :, 1:]), dim=2)
relative_structure = OrientationStructure(structure, pair_features)
encoding = self.encoder(features, relative_structure)
result = self.decoder(encoding)
return result
def forward(self, tertiary, subgraph):
features = torch.zeros(tertiary.size(0),
self.size,
dtype=tertiary.dtype,
device=tertiary.device)
# mask = torch.rand(tertiary.size(0), device=tertiary.device) < 0.2
# features = tertiary.clone()
# features[mask] = 0
# features = torch.cat((features, mask[:, None].float()), dim=1)
#features = self.preprocess(features)
tertiary, _ = self.lookup(tertiary, torch.zeros_like(subgraph.indices))
ors = self.orientations(tertiary)
pos = tertiary[:, 1]
inds = torch.arange(0,
pos.size(0),
dtype=torch.float,
device=pos.device).view(-1, 1)
distances = torch.cat((pos, ors, inds), dim=1)
dist, structure = self.knn_structure(tertiary, subgraph)
neighbour_pos = (pos[:, None] - pos[structure.connections] + 1e-6)
dist = (neighbour_pos).contiguous()
dist = dist.norm(dim=2, keepdim=True)
dist = gaussian_rbf(dist, *self.rbf)
distance_data = RelativeStructure(structure, self.rbf)
relative_data = distance_data.message(distances, distances)
relative_structure = OrientationStructure(structure, relative_data)
encoding = self.encoder(features, relative_structure)
encoding = ts.scatter.batched(self.postprocess, encoding,
subgraph.indices)
#encoding, _ = ts.scatter.sequential(self.rnn, encoding, subgraph.indices)
result = self.result(encoding)
result_tertiary, _ = self.lookup(result,
torch.zeros_like(subgraph.indices))
return result, result_tertiary
def forward(self, inputs):
angles, distances, structure = inputs
print(distances.mean(), distances.min(), distances.max())
print(distances.shape, angles.shape)
node_struc, dist_struc = get_distance_structure(structure)
node_out = self.angle_lookup(angles)
dist_out = self.distance_lookup(gaussian_rbf(distances, *self.rbf))
node_out = ts.scatter.batched(self.preproc, node_out, structure.indices)
node_out, dist_out = self.transformer(node_out, dist_out, node_struc, dist_struc, structure)
indices = torch.repeat_interleave(structure.unique, structure.counts * (structure.counts - 1) // 2, dim=0)
angles = self.angle_result(ts.scatter.mean(node_out, structure.indices))
distances = self.distance_result(dist_out)
result = torch.cat((angles, distances), dim=0)
return result