Class GVPGNN

s3f/gvp.py:162–242

Modified based on https://github.com/drorlab/gvp-pytorch/blob/main/gvp/models.py GVP-GNN for Model Quality Assessment as described in manuscript. Takes in protein structure graphs of type `torchdrug.data.Graph` or `torchdrug.data.PackedGraph` and returns a scalar representatio

@R.register("models.GVPGNN")
class GVPGNN(nn.Module, core.Configurable):
    '''
    Modified based on https://github.com/drorlab/gvp-pytorch/blob/main/gvp/models.py
    GVP-GNN for Model Quality Assessment as described in manuscript.
    
    Takes in protein structure graphs of type `torchdrug.data.Graph` 
    or `torchdrug.data.PackedGraph` and returns a scalar representation for
    each graph and node in the batch in a `torch.Tensor` of shapes [n_nodes, d]
    and [batch_size, d]
    
    :param node_in_dim: node dimensions in input graph, should be
                        (6, 3) if using original features
    :param node_h_dim: node dimensions to use in GVP-GNN layers
    :param node_in_dim: edge dimensions in input graph, should be
                        (32, 1) if using original features
    :param edge_h_dim: edge dimensions to embed to before use
                       in GVP-GNN layers
    :seq_in: if `True`, sequences will also be passed in with
             the forward pass; otherwise, sequence information
             is assumed to be part of input node embeddings
    :param num_layers: number of GVP-GNN layers
    :param drop_rate: rate to use in all dropout layers
    '''
    def __init__(self, node_in_dim, node_h_dim, 
                 edge_in_dim, edge_h_dim, readout="sum",
                 num_layers=3, drop_rate=0.1,
                 activations=(F.relu, None), vector_gate=True):

        super().__init__()
        self.output_dim = node_h_dim[0]
        self.rbf_dim = edge_in_dim[0]

        self.residue_embdding = nn.Linear(node_in_dim[0], node_in_dim[0], bias=False)
        self.W_v = nn.Sequential(
            layer.GVPLayerNorm(node_in_dim),
            layer.GVP(node_in_dim, node_h_dim, activations=(None, None), vector_gate=vector_gate)
        )
        self.W_e = nn.Sequential(
            layer.GVPLayerNorm(edge_in_dim),
            layer.GVP(edge_in_dim, edge_h_dim, activations=(None, None), vector_gate=vector_gate)
        )

        self.layers = nn.ModuleList(
                layer.GVPConvLayer(node_h_dim, edge_h_dim, drop_rate=drop_rate,
                                 activations=activations, vector_gate=vector_gate)
            for _ in range(num_layers))
        
        ns, _ = node_h_dim
        self.W_out = nn.Sequential(
            layer.GVPLayerNorm(node_h_dim),
            layer.GVP(node_h_dim, (ns, 0), activations=activations, vector_gate=vector_gate)
        )

        if readout == "sum":
            self.readout = layers.SumReadout()
        elif readout == "mean":
            self.readout = layers.MeanReadout()
        else: