Project: s41467-024-51563-8 | Function: forward (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements the core GearBind multi-level message passing framework that combines atom-level geometric relational graph convolution, edge-level spatial line graph processing, and residue-level geometric attention to extract comprehensive protein-protein interaction features for binding affinity prediction

**KEY TERMS:** {
  "layer_input": "atom representations a_i^(l-1)",
  "hidden": "updated atom features a_i^(l)",
  "edge_hidden": "edge representations e_(j,i,r)^(l)",
  "line_graph": "spatial line graph for edge message passing",
  "residue_hidden": "residue-level attention output ResAttn",
  "node_mask": "interface atom mask for multi-level processing"
}

**CONTEXT CODE:** from torchdrug.layers import functional


class DDGAttention(nn.Module):

    def __init__(self, input_dim, output_dim, value_dim=16, query_key_dim=16, num_heads=12):
        super(DDGAttention, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.value_dim = value_dim
        self.query_key_dim = query_key_dim
        self.num_heads = num_heads

        self.query = nn.Linear(input_dim, query_key_dim*num_heads, bias=False)
        self.key   = nn.Linear(input_dim, query_key_dim*num_heads, bias=False)
        self.value = nn.Linear(input_dim, value_dim*num_heads, bias=False)

        self.out_transform = nn.Linear(
            in_features = (num_heads*value_dim) + (num_heads*(3+3+1)),
            out_features = output_dim,
        )
        self.layer_norm = nn.LayerNorm(output_dim)

    def _alpha_from_logits(self, logits, mask, inf=1e5):
        """
        Args:
            logits: Logit matrices, (N, L_i, L_j, num_heads).
            mask:   Masks, (N, L).
        Returns:
            alpha:  Attention weights.
        """
        N, L, _, _ = logits.size()
        mask_row = mask.view(N, L, 1, 1).expand_as(logits)      # (N, L, *, *)
        mask_pair = mask_row * mask_row.permute(0, 2, 1, 3)     # (N, L, L, *)
        
        logits = torch.where(mask_pair, logits, logits-inf)
        alpha = torch.softmax(logits, dim=2)  # (N, L, L, num_heads)
        alpha = torch.where(mask_row, alpha, torch.zeros_like(alpha))
        return alpha

    def _heads(self, x, n_heads, n_ch):
        """
        Args:
            x:  (..., num_heads * num_channels)
        Returns:
            (..., num_heads, num_channels)
        """
        s = list(x.size())[:-1] + [n_heads, n_ch]
        return x.view(*s)

    def forward(self, x, pos_CA, pos_CB, frame, mask):
        raise NotImplementedError("Function implementation has been masked")


class GeometricAttention(nn.Module):

    def __init__(self, node_feat_dim, pair_feat_dim, value_dim=16, query_key_dim=16, 
                num_query_points=8, num_value_points=8, num_heads=12, activation=None):
        super().__init__()
        self.node_feat_dim = node_feat_dim
        self.pair_feat_dim = pair_feat_dim
        self.value_dim = value_dim
        self.query_key_dim = query_key_dim
        self.num_query_points = num_query_points
        self.num_value_points = num_value_points
        self.num_heads = num_heads

        # Node
        self.proj_query = nn.Linear(node_feat_dim, query_key_dim*num_heads, bias=False)
        self.proj_key   = nn.Linear(node_feat_dim, query_key_dim*num_heads, bias=False)
        self.proj_value = nn.Linear(node_feat_dim, value_dim*num_heads, bias=False)

        # Pair
        self.proj_pair_bias = nn.Linear(pair_feat_dim, num_heads, bias=False)

        self.spatial_coef = nn.Parameter(torch.full([1, 1, 1, self.num_heads], fill_value=np.log(np.exp(1.) - 1.)), requires_grad=True)

        # Output
        self.out_transform = nn.Linear(
            in_features = (num_heads*pair_feat_dim) + (num_heads*value_dim) + (num_heads*(3+3+1)),
            out_features = node_feat_dim,
        )
        self.layer_norm = nn.LayerNorm(node_feat_dim)

        if isinstance(activation, str):
            self.activation = getattr(F, activation)
        else:
            self.activation = activation

    def _alpha_from_logits(self, logits, mask, inf=1e5):
        """
        Args:
            logits: Logit matrices, (N, L_i, L_j, num_heads).
            mask:   Masks, (N, L).
        Returns:
            alpha:  Attention weights.
        """
        N, L, _, _ = logits.size()
        mask_row = mask.view(N, L, 1, 1).expand_as(logits)      # (N, L, *, *)
        mask_pair = mask_row * mask_row.permute(0, 2, 1, 3)     # (N, L, L, *)
        
        logits = torch.where(mask_pair, logits, logits-inf)
        alpha = torch.softmax(logits, dim=2)  # (N, L, L, num_heads)
        alpha = torch.where(mask_row, alpha, torch.zeros_like(alpha))
        return alpha

    def _heads(self, x, n_heads, n_ch):
        """
        Args:
            x:  (..., num_heads * num_channels)
        Returns:
            (..., num_heads, num_channels)
        """
        s = list(x.size())[:-1] + [n_heads, n_ch]
        return x.view(*s)

    def _normalize_vector(self, v, dim, eps=1e-6):
        return v / (torch.linalg.norm(v, ord=2, dim=dim, keepdim=True) + eps)

    def _node_logits(self, x):
        query_l = self._heads(self.proj_query(x), self.num_heads, self.query_key_dim)    # (N, L, n_heads, qk_ch)
        key_l = self._heads(self.proj_key(x), self.num_heads, self.query_key_dim)      # (N, L, n_heads, qk_ch)

        query_l = query_l.permute(0, 2, 1, 3)   # (N,L1,H,C) -> (N,H,L1,C)
        key_l = key_l.permute(0, 2, 3, 1)       # (N,L2,H,C) -> (N,H,C,L2)

        logits = torch.matmul(query_l, key_l)   # (N,H,L1,L2)
        logits = logits.permute(0, 2, 3, 1)     # (N,L1,L2,H)

        return logits

    def _pair_logits(self, z):
        logits_pair = self.proj_pair_bias(z)
        return logits_pair

    def _beta_logits(self, p_CB):
        qk = p_CB[:, :, None, :].expand(-1, -1, self.num_heads, -1)
        sum_sq_dist = ((qk.unsqueeze(2) - qk.unsqueeze(1)) ** 2).sum(-1)    # (N, L, L, n_heads)
        gamma = F.softplus(self.spatial_coef)
        logtis_beta = sum_sq_dist * ((-1 * gamma * np.sqrt(2 / 9)) / 2)
        return logtis_beta

    def _pair_aggregation(self, alpha, z):
        N, L = z.shape[:2]
        feat_p2n = torch.einsum("blkh, blkc->blhc", alpha, z)
        return feat_p2n.reshape(N, L, -1)

    def _node_aggregation(self, alpha, x):
        N, L = x.shape[:2]
        value_l = self._heads(self.proj_value(x), self.num_heads, self.query_key_dim)  # (N, L, n_heads, v_ch)


Complete the function implementation replacing "raise NotImplementedError".