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)