| 25 | class SurfGVP(nn.Module, core.Configurable): |
| 26 | |
| 27 | def __init__(self, node_in_dim, node_h_dim, |
| 28 | edge_in_dim, edge_h_dim, surf_in_dim, surf_edge_in_dim, |
| 29 | num_surf_graph_neighbor=16, |
| 30 | num_surf_res_neighbor=3, readout="sum", |
| 31 | num_layers=3, drop_rate=0.1, |
| 32 | activations=(F.relu, None), vector_gate=True): |
| 33 | |
| 34 | super().__init__() |
| 35 | self.output_dim = node_h_dim[0] |
| 36 | self.rbf_dim = edge_in_dim[0] |
| 37 | self.surf_in_dim = surf_in_dim[0] |
| 38 | self.surf_rbf_dim = surf_edge_in_dim[0] |
| 39 | self.k = num_surf_res_neighbor |
| 40 | self.num_surf_graph_neighbor = num_surf_graph_neighbor |
| 41 | |
| 42 | self.residue_embdding = nn.Linear(node_in_dim[0], node_in_dim[0], bias=False) |
| 43 | self.W_v = nn.Sequential( |
| 44 | layer.GVPLayerNorm(node_in_dim), |
| 45 | layer.GVP(node_in_dim, node_h_dim, activations=(None, None), vector_gate=vector_gate) |
| 46 | ) |
| 47 | self.W_e = nn.Sequential( |
| 48 | layer.GVPLayerNorm(edge_in_dim), |
| 49 | layer.GVP(edge_in_dim, edge_h_dim, activations=(None, None), vector_gate=vector_gate) |
| 50 | ) |
| 51 | |
| 52 | self.layers = nn.ModuleList( |
| 53 | layer.GVPConvLayer(node_h_dim, edge_h_dim, drop_rate=drop_rate, |
| 54 | activations=activations, vector_gate=vector_gate) |
| 55 | for _ in range(num_layers)) |
| 56 | |
| 57 | ns, _ = node_h_dim |
| 58 | self.W_out = nn.Sequential( |
| 59 | layer.GVPLayerNorm(node_h_dim), |
| 60 | layer.GVP(node_h_dim, (ns, 0), activations=activations, vector_gate=vector_gate) |
| 61 | ) |
| 62 | |
| 63 | # Similar layer for surface graphs |
| 64 | self.surf_in_linear = nn.Linear(node_in_dim[0]+1, node_in_dim[0], bias=False) |
| 65 | self.surf_in_mlp = nn.Sequential( |
| 66 | nn.Linear(node_in_dim[0] + surf_in_dim[0], node_in_dim[0] * 2), |
| 67 | nn.Dropout(drop_rate), |
| 68 | nn.LayerNorm(node_in_dim[0] * 2), |
| 69 | nn.ReLU(), |
| 70 | nn.Linear(node_in_dim[0] * 2, node_in_dim[0]), |
| 71 | ) |
| 72 | self.surf_W_v = nn.Sequential( |
| 73 | layer.GVPLayerNorm(node_in_dim), |
| 74 | layer.GVP(node_in_dim, node_h_dim, activations=(None, None), vector_gate=vector_gate) |
| 75 | ) |
| 76 | self.surf_W_e = nn.Sequential( |
| 77 | layer.GVPLayerNorm(edge_in_dim), |
| 78 | layer.GVP(edge_in_dim, edge_h_dim, activations=(None, None), vector_gate=vector_gate) |
| 79 | ) |
| 80 | self.surf_layers = nn.ModuleList( |
| 81 | layer.GVPConvLayer(node_h_dim, edge_h_dim, drop_rate=drop_rate, |
| 82 | activations=activations, vector_gate=vector_gate) |
| 83 | for _ in range(num_layers)) |
| 84 | ns, _ = node_h_dim |