def __init__(self, in_size, hidden_size=800, angles=512, fragment_size=5): super(StructuredGenerator, self).__init__() self.in_size = in_size self.angle_lookup = AngleProject(2 * hidden_size, 3) self.preproc = AutoregressiveAngle(in_size, 3)#nn.GRUCell(in_size, in_size) self.rnn = nn.LSTM(in_size, hidden_size, 2, bidirectional=True, batch_first=True) self.position_lookup = PositionLookup(fragment_size=fragment_size)
def each_generate(self, data, generated, sample):
with torch.no_grad():
angles, distances, protein = generated
angs = angles.tensor.cpu()
dists = distances.tensor.cpu()
first_length = (protein.indices.cpu() == 0).sum()
expanded = first_length * (first_length - 1) // 2
first_dist = dists[:expanded]
dist = torch.zeros(first_length, first_length)
count = 0
for idx in range(first_length):
for idy in range(idx + 1, first_length):
dist[idx, idy] = dists[count]
dist[idy, idx] = dists[count]
count += 1
lookup = PositionLookup()
c_alpha, _ = lookup(
angs[protein.indices.cpu() == 0],
torch.zeros_like(protein.indices[protein.indices == 0].cpu()))
c_alpha = c_alpha[:, 1]
c_alpha = c_alpha.numpy()
dist = dist.numpy()
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot(c_alpha[:, 0], c_alpha[:, 1], c_alpha[:, 2])
self.writer.add_figure("output", fig, self.step_id)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.scatter(angs[:, 1], angs[:, 2])
self.writer.add_figure("rama", fig, self.step_id)
fig, ax = plt.subplots()
ax.matshow(dist)
self.writer.add_figure("dist", fig, self.step_id)
def __init__(self, in_size, size, distance_size, sequence_size=20,
attention_size=128, heads=128, hidden_size=128, mlp_depth=3,
depth=3, max_distance=20, distance_kernels=16, neighbours=15,
activation=func.relu_, batch_norm=True, conditional=False,
angles=False, dropout=0.1, connected=attention_connected,
normalization=spectral_norm):
super().__init__()
distance_size = distance_size + distance_kernels - 1
self.encoder = StructuredTransformerEncoder(
size, size, distance_size,
attention_size=attention_size, heads=heads, hidden_size=hidden_size,
depth=depth, mlp_depth=mlp_depth, activation=activation,
batch_norm=batch_norm, dropout=dropout, normalization=normalization,
connected=connected
)
self.angles = angles
self.lookup = PositionLookup(fragment_size=10)
self.conditional = conditional
self.neighbours = neighbours
self.activation = activation
self.rbf = (0, max_distance, distance_kernels)
self.preprocess = LocalFeatures(6, size)
self.postprocess = LocalFeatures(size, size)
self.result = nn.Linear(2 * size, 1)
self.angle_result = MLP(6, 1, hidden_size=32, depth=3, batch_norm=False, normalization=spectral_norm)
def __init__(self, in_size, size, distance_size, mix=10, schedule=2, radius=8,
attention_size=128, heads=128, hidden_size=128, mlp_depth=3,
depth=3, max_distance=20, distance_kernels=16, neighbours=15,
activation=func.relu_, batch_norm=True, conditional=False,
angles=False, dropout=0.1, connected=attention_connected,
normalization=lambda x: x):
super().__init__()
distance_size = distance_size + distance_kernels - 1
self.encoder = StructuredTransformerEncoder(
size, size, distance_size,
attention_size=attention_size, heads=heads, hidden_size=hidden_size,
depth=depth, mlp_depth=mlp_depth, activation=activation,
batch_norm=batch_norm, dropout=0.1, normalization=lambda x: x,
connected=connected
)
self.schedule = schedule
self.mix = mix
self.radius = radius
self.angles = angles
self.lookup = PositionLookup(fragment_size=10)
self.conditional = conditional
self.neighbours = neighbours
self.activation = activation
self.rbf = (0, max_distance, distance_kernels)
self.preprocess = nn.Linear(6, size)
self.mean = AngleProject(size, 3 * self.mix)#nn.Linear(size, 3 * self.mix)
self.log_concentration = nn.Linear(size, 3 * self.mix)
self.weights = nn.Linear(size, self.mix)
self.factor = nn.Linear(size, 3 * self.mix)
def prepare(self):
position_lookup = PositionLookup()
index = random.randrange(0, len(self.data))
(positions, ground_truth, protein) = self.data[index]
angles = torch.randn(positions.tensor.size(0), 3)
positions, _ = position_lookup(angles, protein.indices)
return (PackedTensor(positions), ground_truth, protein)
def __init__(self, in_size, hidden_size=128, depth=3, fragment_size=5):
super(TransformerGN, self).__init__()
self.preprocess = nn.Linear(in_size, hidden_size)
self.angle_lookup = AngleSample(hidden_size, 60)#nn.Linear(hidden_size, 3)
self.position_lookup = PositionLookup(fragment_size=fragment_size)
self.blocks = nn.ModuleList([
Transformer(hidden_size, hidden_size, hidden_size // 2, attention_size=8)
for idx in range(depth)
])
def __init__(self, depth=4, shape=64):
super().__init__()
self.preprocess = nn.Conv2d(1, 128, 5, padding=2)
self.blocks = nn.ModuleList([
nn.Sequential(nn.Conv2d(128, 128, 3, padding=1), nn.LeakyReLU(),
nn.Conv2d(128, 128, 3, dilation=1, padding=1),
nn.LeakyReLU()) for idx in range(depth)
])
self.predict = nn.Linear(128, 1)
self.shape = shape
self.lookup = PositionLookup()
def __init__(self, in_size, hidden_size=128, depth=3,
fragment_size=5, angles=60, integrate=5):
super(ResidualGN, self).__init__()
self.state = nn.Linear(in_size, hidden_size)
self.angles = nn.Linear(in_size, hidden_size)
self.angle_unlookup = nn.Linear(6, hidden_size)
self.angle_lookup = nn.Linear(hidden_size, 3)#AngleLookup(hidden_size, angles)
self.position_lookup = PositionLookup(fragment_size=fragment_size)
self.blocks = nn.ModuleList([
InteractionConv(hidden_size, hidden_size, hidden_size // 2, depth=1)
for idx in range(depth)
])
self.integrate = integrate
def each_generate(self, data, gt, protein):
with torch.no_grad():
lookup = PositionLookup()
angs = data.tensor
c_alpha, _ = lookup(data.tensor[protein.indices == 0], torch.zeros_like(protein.indices[protein.indices == 0]))
c_alpha = c_alpha[:, 1].numpy()
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot(c_alpha[:, 0], c_alpha[:, 1], c_alpha[:, 2])
self.writer.add_figure("output", fig, self.step_id)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.scatter(angs[:, 1], angs[:, 2])
self.writer.add_figure("rama", fig, self.step_id)
def __init__(self, depth=4, down=2, shape=64):
super().__init__()
self.preprocess = nn.Conv2d(1 + 42, 128, 5, padding=2)
self.blocks = nn.ModuleList([
nn.Sequential(nn.Conv2d(128, 128, 3, padding=1), nn.LeakyReLU(),
nn.Conv2d(128, 128, 3, dilation=1, padding=1),
nn.LeakyReLU()) for idx in range(depth)
])
self.predict = nn.Linear(128, 1)
self.shape = shape
self.down = down
self.lookup = PositionLookup()
self.gaussian = nn.Parameter(
torch.rand(1, 1, shape, shape, requires_grad=True))
def __init__(self, in_size, hidden_size=128, depth=3,
fragment_size=5, angles=60, integrate=5):
super(ConvolutionalGN, self).__init__()
self.state = nn.Linear(in_size, hidden_size)
self.angles = nn.Linear(in_size, hidden_size)
self.angle_unlookup = nn.Linear(6, hidden_size)
self.angle_lookup = AngleSample(hidden_size, angles)
self.position_lookup = PositionLookup(fragment_size=fragment_size)
self.blocks = nn.ModuleList([
InteractionConv(hidden_size, hidden_size, hidden_size // 2, depth=1)
for idx in range(depth)
])
self.transformer = Transformer(hidden_size, hidden_size, hidden_size // 2, attention_size=8)
self.integrate = integrate
def __init__(self,
in_size,
size,
distance_size,
sequence_size=20,
attention_size=128,
heads=128,
hidden_size=128,
mlp_depth=3,
depth=3,
max_distance=20,
distance_kernels=16,
neighbours=15,
activation=func.relu_,
batch_norm=False,
conditional=False,
angles=False,
dropout=0.1,
connected=linear_connected,
normalization=lambda x: x):
super().__init__()
distance_size = distance_size + distance_kernels - 1
self.encoder = StructuredTransformerEncoder(
size,
size,
distance_size,
attention_size=attention_size,
heads=heads,
hidden_size=hidden_size,
depth=depth,
mlp_depth=mlp_depth,
activation=activation,
batch_norm=batch_norm,
normalization=spectral_norm,
dropout=0.1,
connected=connected)
self.angles = angles
self.lookup = PositionLookup()
self.conditional = conditional
self.neighbours = neighbours
self.activation = activation
self.rbf = (0, max_distance, distance_kernels)
self.local_features = LocalFeatures(6, size)
self.reweighting = spectral_norm(nn.Linear(size, 1))
self.energy = spectral_norm(nn.Linear(size, 1))
self.pool = Pool(size, 4)
def __init__(self, depth=4, down=2, shape=64):
super().__init__()
self.preprocess = nn.Conv2d(1 + 1 + 42, 128, 5, padding=2)
self.blocks = nn.ModuleList([
nn.Sequential(
spectral_norm(nn.Conv2d(128, 128, 3, padding=1)),
nn.LeakyReLU(),
spectral_norm(
nn.Conv2d(128,
128,
3,
dilation=2**(idx % 5),
padding=2**(idx % 5))), nn.LeakyReLU())
for idx in range(depth)
])
self.predict = spectral_norm(nn.Linear(128, 1))
self.shape = shape
self.down = down
self.lookup = PositionLookup()
def __init__(self, in_size, distance_size=128, hidden_size=128, angles=512,
fragment_size=5, attention_size=128, heads=8, depth=3,
mlp_depth=3, dropout=0.1, activation=func.relu_, batch_norm=False,
pre_norm=True, normalization=lambda x: x, connected=attention_connected):
super(DistanceGenerator, self).__init__()
self.in_size = in_size
self.distance_size = 128
self.angle_lookup = AngleProject(hidden_size, 3)
self.distance_lookup = nn.Linear(distance_size, 1)
self.preproc = LocalFeatures(in_size + 1, hidden_size)
self.transformer = DistanceTransformer(
in_size, hidden_size, distance_size,
hidden_size=hidden_size, attention_size=attention_size,
heads=8, depth=depth, mlp_depth=mlp_depth,
dropout=dropout, activation=activation,
batch_norm=batch_norm, pre_norm=pre_norm,
normalization=normalization, connected=connected
)
self.position_lookup = PositionLookup(fragment_size=fragment_size)
def each_generate(self, data, generated, sample):
with torch.no_grad():
data, protein = generated
lookup = PositionLookup()
angs = data.tensor.cpu()
c_alpha, _ = lookup(angs[protein.indices.cpu() == 0], torch.zeros_like(protein.indices[protein.indices == 0].cpu()))
c_alpha = c_alpha[:, 1]
dist = (c_alpha[:, None] - c_alpha[None, :]).norm(dim=-1)
c_alpha = c_alpha.numpy()
dist = dist.numpy()
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot(c_alpha[:, 0], c_alpha[:, 1], c_alpha[:, 2])
self.writer.add_figure("output", fig, self.step_id)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.scatter(angs[:, 1], angs[:, 2])
self.writer.add_figure("rama", fig, self.step_id)
fig, ax = plt.subplots()
ax.matshow(dist)
self.writer.add_figure("dist", fig, self.step_id)
def each_generate(self, data, gt, protein):
with torch.no_grad():
lookup = PositionLookup()
angs = data.tensor
c_alpha, _ = lookup(
data.tensor[protein.indices == 0],
torch.zeros_like(protein.indices[protein.indices == 0]))
dist = (c_alpha[None, :, 1] - c_alpha[:, None, 1]).norm(dim=-1)
dist = dist.cpu().numpy()
fig, ax = plt.subplots(figsize=(10, 10))
ax.matshow(dist)
self.writer.add_figure("dist", fig, self.step_id)
c_alpha = c_alpha[:, 1].numpy()
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot(c_alpha[:, 0], c_alpha[:, 1], c_alpha[:, 2])
self.writer.add_figure("output", fig, self.step_id)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.scatter(angs[:, 1], angs[:, 2])
self.writer.add_figure("rama", fig, self.step_id)
def valid_callback(trn, inputs, outputs):
with torch.no_grad():
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
positions, pos_target = outputs[1]
ter, mask, struc = pos_target
angles, angle_target = outputs[0]
ang = angles
angles = angle_target
positions = positions.view(-1, 3)
ter = ter.view(-1, 3)
positions = positions[(struc.indices == 0).nonzero().view(-1)].numpy()
ter = ter[(struc.indices == 0).nonzero().view(-1)].numpy()
dst = np.linalg.norm((positions[None, :] - positions[:, None]),
axis=-1)
dstt = np.linalg.norm((ter[None, :] - ter[:, None]), axis=-1)
ax.plot(positions[:, 0], positions[:, 1], positions[:, 2])
trn.writer.add_figure("output", fig, trn.step_id)
plt.close("all")
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot(ter[:, 0], ter[:, 1], ter[:, 2])
trn.writer.add_figure("input", fig, trn.step_id)
plt.close("all")
fig = plt.figure()
ax = fig.add_subplot(111)
ax.imshow(dst)
trn.writer.add_figure("heat out", fig, trn.step_id)
plt.close("all")
fig = plt.figure()
ax = fig.add_subplot(111)
ax.imshow(dstt)
trn.writer.add_figure("heat in", fig, trn.step_id)
plt.close("all")
fig = plt.figure()
ax = fig.add_subplot(111)
ax.imshow(abs(dstt - dst))
trn.writer.add_figure("heat del", fig, trn.step_id)
plt.close("all")
fig = plt.figure()
ax = fig.add_subplot(111)
position_lookup = PositionLookup()
re_ter = position_lookup(
angles[(struc.indices == 0).nonzero().view(-1)],
struc.indices[(struc.indices == 0).nonzero().view(-1)])[0].view(
-1, 3).numpy()
dsttt = np.linalg.norm((re_ter[None, :] - re_ter[:, None]), axis=-1)
ax.imshow(dsttt)
trn.writer.add_figure("heat expected", fig, trn.step_id)
plt.close("all")
trn.writer.add_scalar("size", float(ter.shape[0]), trn.step_id)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.scatter(angles[:, 1].numpy() % 6.3, angles[:, 2].numpy() % 6.3)
ax.scatter(ang[:, 1].numpy() % 6.3, ang[:, 2].numpy() % 6.3)
trn.writer.add_figure("rama", fig, trn.step_id)
plt.close("all")
def __init__(self, in_size, hidden_size=800, angles=512, fragment_size=5): super(RGN, self).__init__() self.angle_lookup = AngleProject(2 * hidden_size, 3) self.rnn = nn.LSTM(in_size, hidden_size, 2, bidirectional=True, batch_first=True) self.position_lookup = PositionLookup(fragment_size=fragment_size)
