//  [python]
def forward(self, node_attr, pos, batch_ctx, batch, pred_res_type, mask_protein, external_index, backbone=True,
                mask=True):
        S_id, R = batch['res_idx'], batch['amino_acid']
        residue_batch, atom2residue = batch['amino_acid_batch'], batch['atom2residue']
        edit_residue, edit_atom = batch['protein_edit_residue'], batch['protein_edit_atom']
        if mask:
            R[batch['random_mask_residue']] = 0
        R = F.one_hot(R, num_classes=21).float()
        if backbone:
            atom2residue, edit_atom = batch['atom2residue_backbone'], batch['protein_edit_atom_backbone']
            R_edit = R[edit_residue]
            R_edit[:, 1:] = pred_res_type
            R[edit_residue] = R_edit
            R[:, 0] = 0

        edge_index = knn_graph(pos, k=self.k, batch=batch_ctx, flow='target_to_source')
        edge_length = torch.linalg.norm(pos[edge_index[0]] - pos[edge_index[1]], dim=1)
        edge_attr = self.distance_expansion(edge_length)

        h = node_attr
        for interaction in self.interactions:
            h = interaction(h, edge_index, edge_attr)

        h_ligand_coarse = scatter_sum(h[~mask_protein], batch['ligand_atom_batch'], dim=0)
        pos_ligand_coarse = scatter_sum(batch['ligand_pos'], batch['ligand_atom_batch'], dim=0)
        E, residue_edge_index, residue_edge_length, edge_index_new, E_new = self.features(pos_ligand_coarse,
                                                                                          batch['protein_edit_residue'],
                                                                                          batch['residue_pos'], S_id,
                                                                                          residue_batch)
        h_protein = h[mask_protein]
        V = torch.cat([self.residue_feat.soft_forward(R), scatter_sum(h_protein, atom2residue, dim=0)], dim=-1)
        h_res = self.W_v(V)

        h_res = torch.cat([h_res, h_ligand_coarse])
        edge_index_combined = torch.cat([residue_edge_index, edge_index_new], 1)
        E = torch.cat([E, E_new], 0)
        h_E = self.W_e(E)

        for layer in self.residue_encoder_layers:
            h_res = layer(h_res, h_E, edge_index_combined)

        # update X:
        h_res = h_res[:len(residue_batch)]
        h_E = h_E[:residue_edge_index.size(1)]
        # protein internal update
        mij = torch.cat([h_res[residue_edge_index[0]], h_res[residue_edge_index[1]], h_E], dim=-1)
        if backbone:
            protein_pos = pos[mask_protein]
            ligand_pos = pos[~mask_protein]
            N = atom2residue.max() + 1
            X_bb = torch.zeros(N, 4, 3).to(pos.device)
            for j in range(N):
                X_bb[j] = protein_pos[atom2residue == j][:4]  # 4 backbone atoms [N,4,3]
            xij = X_bb[residue_edge_index[0]] - X_bb[residue_edge_index[1]]  # [N,4,3]
            dij = xij.norm(dim=-1) + 1e-6  # [N,4]
            fij = torch.maximum(self.T_x(mij)[:, :4], 3.8 - dij)  # break term [N,4]
            xij = xij / dij.unsqueeze(-1) * fij.unsqueeze(-1)
            f_res = scatter_mean(xij, residue_edge_index[0], dim=0)  # [N,4,3]
            X_bb[edit_residue] += f_res.clamp(min=-5.0, max=5.0)[edit_residue]

            # Clash correction
            for _ in range(2):
                xij = X_bb[residue_edge_index[0]] - X_bb[residue_edge_index[1]]  # [N,4,3]
                dij = xij.norm(dim=-1) + 1e-6  # [N,4]
                fij = F.relu(3.8 - dij)  # repulsion term [N,4]
                xij = xij / dij.unsqueeze(-1) * fij.unsqueeze(-1)
                f_res = scatter_mean(xij, residue_edge_index[0], dim=0)  # [N,4,3]
                X_bb[edit_residue] += f_res.clamp(min=-5.0, max=5.0)[edit_residue]

            # protein-ligand external update
            protein_pos[edit_atom] = X_bb[edit_residue].view(-1, 3)
            pos[mask_protein] = protein_pos
            dij = torch.linalg.norm(pos[mask_protein][external_index[0]] - pos[~mask_protein][external_index[1]], dim=1) + 1e-6
            mij = torch.cat(
                [h[mask_protein][external_index[0]], h[~mask_protein][external_index[1]], self.distance_expansion(dij)],
                dim=-1)
            xij = pos[mask_protein][external_index[0]] - pos[~mask_protein][external_index[1]]
            fij = torch.maximum(self.T_a(mij).squeeze(-1), 1.5 - dij)
            xij = xij / dij.unsqueeze(-1) * fij.unsqueeze(-1)
            f_atom = scatter_mean(xij, external_index[0], dim=0, dim_size=protein_pos.size(0))
            protein_pos += f_atom
            f_ligand_atom = scatter_mean(xij, external_index[1], dim=0, dim_size=ligand_pos.size(0))
            ligand_pos -= f_ligand_atom * 0.05

        else:
            protein_pos = pos[mask_protein]
            ligand_pos = pos[~mask_protein]
            X_avg = scatter_mean(protein_pos, atom2residue, dim=0)
            X = X_avg.unsqueeze(1).repeat(1, 14, 1)
            N = atom2residue.max() + 1
            mask = torch.zeros(N, 14, dtype=bool).to(protein_pos.device)
            residue_natoms = atom2residue.bincount()
            for j in range(N):
                mask[j][:residue_natoms[j]] = 1  # all atoms mask [N,14]
                X[j][:residue_natoms[j]] = protein_pos[atom2residue == j]

            xij = X[residue_edge_index[0]] - X_avg[residue_edge_index[1]].unsqueeze(1)  # [N,14,3]
            dij = xij.norm(dim=-1) + 1e-6  # [N,14]
            fij = torch.maximum(self.T_x(mij), 3.8 - dij)  # break term [N,14]
            xij = xij / dij.unsqueeze(-1) * fij.unsqueeze(-1)
            f_res = scatter_mean(xij, residue_edge_index[0], dim=0)  # [N,14,3]
            f_res[:, :4] *= 0.1
            X[edit_residue] += f_res.clamp(min=-5.0, max=5.0)[edit_residue]

            for _ in range(2):
                protein_pos = X[mask]
                X_avg = scatter_mean(protein_pos, atom2residue, dim=0)
                xij = X[residue_edge_index[0]] - X_avg[residue_edge_index[1]].unsqueeze(1)  # [N,14,3]
                dij = xij.norm(dim=-1) + 1e-6  # [N,14]
                fij = F.relu(3.8 - dij)  # repulsion term [N,14]
                xij = xij / dij.unsqueeze(-1) * fij.unsqueeze(-1)
                f_res = scatter_mean(xij, residue_edge_index[0], dim=0)  # [N,14,3]
                X[edit_residue] += f_res.clamp(min=-5.0, max=5.0)[edit_residue]

            # protein-ligand external update
            protein_pos = X[mask]
            pos[mask_protein] = protein_pos
            dij = torch.linalg.norm(pos[mask_protein][external_index[0]] - pos[~mask_protein][external_index[1]], dim=1) + 1e-6
            mij = torch.cat(
                [h[mask_protein][external_index[0]], h[~mask_protein][external_index[1]], self.distance_expansion(dij)],
                dim=-1)
            xij = pos[mask_protein][external_index[0]] - pos[~mask_protein][external_index[1]]
            fij = torch.maximum(self.T_a(mij).squeeze(-1), 1.5 - dij)
            xij = xij / dij.unsqueeze(-1) * fij.unsqueeze(-1)
            f_atom = scatter_mean(xij, external_index[0], dim=0, dim_size=protein_pos.size(0))
            f_atom[batch['edit_backbone']] *= 0.1
            protein_pos += f_atom
            f_ligand_atom = scatter_mean(xij, external_index[1], dim=0, dim_size=ligand_pos.size(0))
            ligand_pos -= f_ligand_atom * 0.05

        return h, h_res, protein_pos, ligand_pos