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 collections.abc import Sequence

import torch
from torch import nn
from torch import autograd

from torch_scatter import scatter_add

from torchdrug import core, layers, utils
from torchdrug.layers import functional
from torchdrug.core import Registry as R

from . import layer, geometry


@R.register("models.BindModel")
class BindModel(nn.Module, core.Configurable):

    def __init__(self, model, num_mlp_layer=2):
        super(BindModel, self).__init__()
        self.model = model
        self.num_mlp_layer = num_mlp_layer

        hidden_dims = [model.output_dim] * (num_mlp_layer - 1)
        self.mlp = layers.MLP(model.output_dim * 2, hidden_dims + [1])

    def forward(self, batch, all_loss=None, metric=None):
        mutant = batch["mutant"]
        mutant_output = self.model(mutant, mutant.node_feature.float(), all_loss=all_loss, metric=metric)

        wild_type = batch["wild_type"]
        wild_type_output = self.model(wild_type, wild_type.node_feature.float(), all_loss=all_loss, metric=metric)
        
        wild_type_output = wild_type_output["graph_feature"]
        mutant_output = mutant_output["graph_feature"]

        outputs = torch.cat([mutant_output, wild_type_output], dim=-1)
        pred = self.mlp(outputs)
        outputs = torch.cat([wild_type_output, mutant_output], dim=-1)
        pred = pred - self.mlp(outputs)
        
        return {
            'ddG': pred, 
            'wild_type_feature': wild_type_output, 
            'mutant_feature': mutant_output
        }


@R.register("models.GearBind")
class GearBind(nn.Module, core.Configurable):

    def __init__(self, input_dim, hidden_dims, num_relation, edge_input_dim=None, num_angle_bin=None,
                 short_cut=False, batch_norm=False, activation="relu", concat_hidden=False, readout="sum",
                 use_attn=True):
        super(GearBind, self).__init__()
        if not isinstance(hidden_dims, Sequence):
            hidden_dims = [hidden_dims]
        self.input_dim = input_dim
        self.output_dim = sum(hidden_dims) if concat_hidden else hidden_dims[-1]
        self.dims = [input_dim] + list(hidden_dims)
        self.edge_dims = [edge_input_dim] + self.dims[:-1]
        self.num_relation = num_relation
        self.num_angle_bin = num_angle_bin
        self.short_cut = short_cut
        self.concat_hidden = concat_hidden
        self.batch_norm = batch_norm
        self.use_attn = use_attn

        self.layers = nn.ModuleList()
        for i in range(len(self.dims) - 1):
            self.layers.append(layers.GeometricRelationalGraphConv(self.dims[i], self.dims[i + 1], num_relation,
                                                                   None, batch_norm, activation))
        self.atom_position_gather = geometry.AtomPositionGather()
        if use_attn:
            self.attn_layers = nn.ModuleList()
            for i in range(len(self.dims) - 1):
                self.attn_layers.append(layer.DDGAttention(self.dims[i+1], self.dims[i+1]))
        if num_angle_bin:
            self.spatial_line_graph = layers.SpatialLineGraph(num_angle_bin)
            self.edge_layers = nn.ModuleList()
            for i in range(len(self.edge_dims) - 1):
                self.edge_layers.append(layers.GeometricRelationalGraphConv(
                    self.edge_dims[i], self.edge_dims[i + 1], num_angle_bin, None, batch_norm, activation))

        if batch_norm:
            self.batch_norm_layers = nn.ModuleList()
            for i in range(len(self.dims) - 1):
                self.batch_norm_layers.append(nn.BatchNorm1d(self.dims[i + 1]))

        if readout == "sum":
            self.readout = layers.SumReadout()
        elif readout == "mean":
            self.readout = layers.MeanReadout()
        else:
            raise ValueError("Unknown readout `%s`" % readout)

    def forward(self, graph, input, all_loss=None, metric=None):
        raise NotImplementedError("Function implementation has been masked")


@R.register("models.DDGPredictor")
class DDGPredictor(nn.Module, core.Configurable):

    def __init__(self, hidden_dim=512, num_layers=4, pair_dim=128, max_relpos=100, 
                num_neighbors=128, num_heads=12, activation=None):
        super(DDGPredictor, self).__init__()
        self.hidden_dim = hidden_dim
        self.pair_dim = pair_dim
        self.max_relpos = max_relpos
        self.num_layers = num_layers
        self.output_dim = 2
        self.atom_position_gather = geometry.AtomPositionGather()
        self.mutation_site_graph = geometry.KNNMutationSite(k=num_neighbors)

        self.relpos_embedding = nn.Embedding(max_relpos*2+2, pair_dim)
        self.residue_encoder = layer.PerResidueEncoder(hidden_dim)

        self.blocks = nn.ModuleList([
            layer.GeometricAttention(hidden_dim, pair_dim, num_heads=num_heads, activation=activation) for _ in range(num_layers)
        ])

        # Readout
        self.mlp = nn.Sequential(
            nn.Linear(hidden_dim*2, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),


Complete the function implementation replacing "raise NotImplementedError".