models/PD.py

import copy
import sys

sys.path.append("..")
import numpy as np
from rdkit import RDConfig
import random
import os
import torch
import torch.nn as nn
from torch.nn import Module, Linear, Embedding
from torch.nn import functional as F
from torch_scatter import scatter_add, scatter_mean
from torch_geometric.data import Data, Batch
from rdkit.Chem import ChemicalFeatures
from rdkit import Chem
from rdkit.Chem import rdchem

from .encoders import get_encoder, MLP
from .encoders.cftfm import residue_atom_mask
from .common import *
from .protein_features import *
from .esmadapter import *
from .esm2adapter import *
from utils.pdb_utils import VOCAB
from utils.rmsd import kabsch_torch
from utils.protein_ligand import PDBProtein
from utils.relax import openmm_relax

ATOM_FAMILIES = ['Acceptor', 'Donor', 'Aromatic', 'Hydrophobe', 'LumpedHydrophobe', 'NegIonizable', 'PosIonizable',
                 'ZnBinder']
ATOM_FAMILIES_ID = {s: i for i, s in enumerate(ATOM_FAMILIES)}
NUM_ATOMS = [4, 5, 11, 8, 8, 6, 9, 9, 4, 10, 8, 8, 9, 8, 11, 7, 6, 7, 14, 12, 7]

ATOM_TYPES = [
    '', 'N', 'CA', 'C', 'O', 'CB', 'CG', 'CG1', 'CG2', 'OG', 'OG1', 'SG', 'CD',
    'CD1', 'CD2', 'ND1', 'ND2', 'OD1', 'OD2', 'SD', 'CE', 'CE1', 'CE2', 'CE3',
    'NE', 'NE1', 'NE2', 'OE1', 'OE2', 'CH2', 'NH1', 'NH2', 'OH', 'CZ', 'CZ2',
    'CZ3', 'NZ', 'OXT'
]
RES_ATOM14 = [
    [''] * 14,
    ['N', 'CA', 'C', 'O', 'CB', '', '', '', '', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD', 'NE', 'CZ', 'NH1', 'NH2', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'OD1', 'ND2', '', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'OD1', 'OD2', '', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'SG', '', '', '', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD', 'OE1', 'NE2', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD', 'OE1', 'OE2', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', '', '', '', '', '', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'ND1', 'CD2', 'CE1', 'NE2', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', 'CD1', '', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD1', 'CD2', '', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD', 'CE', 'NZ', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'SD', 'CE', '', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD1', 'CD2', 'CE1', 'CE2', 'CZ', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD', '', '', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'OG', '', '', '', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'OG1', 'CG2', '', '', '', '', '', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD1', 'CD2', 'NE1', 'CE2', 'CE3', 'CZ2', 'CZ3', 'CH2'],
    ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD1', 'CD2', 'CE1', 'CE2', 'CZ', 'OH', '', ''],
    ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', '', '', '', '', '', '', ''],
]

RES_ATOM_TYPE = [[ATOM_TYPES.index(a) for a in res]for res in RES_ATOM14]

AA_NUMBER_NAME = {1: 'ALA', 2: 'ARG', 3: 'ASN', 4: 'ASP', 5: 'CYS', 6: 'GLN', 7: 'GLU', 8: 'GLY', 9: 'HIS', 10: 'ILE',
                  11: 'LEU', 12: 'LYS', 13: 'MET', 14: 'PHE', 15: 'PRO', 16: 'SER', 17: 'THR', 18: 'TRP', 19: 'TYR',
                  20: 'VAL'}

RES_ATOMS = [[ATOM_TYPES.index(i) for i in res if i != ''] for res in RES_ATOM14]

BOND_TYPE = {1: rdchem.BondType.SINGLE, 2: rdchem.BondType.DOUBLE, 3: rdchem.BondType.TRIPLE,
             12: rdchem.BondType.AROMATIC}

def quaternion_to_matrix(q):
    """Convert a quaternion to its corresponding rotation matrix."""
    q = q / q.norm()
    w, x, y, z = q
    R = torch.zeros((3, 3), device=q.device)
    R[0, 0] = 1 - 2 * (y ** 2 + z ** 2)
    R[0, 1] = 2 * (x * y - z * w)
    R[0, 2] = 2 * (x * z + y * w)
    R[1, 0] = 2 * (x * y + z * w)
    R[1, 1] = 1 - 2 * (x ** 2 + z ** 2)
    R[1, 2] = 2 * (y * z - x * w)
    R[2, 0] = 2 * (x * z - y * w)
    R[2, 1] = 2 * (y * z + x * w)
    R[2, 2] = 1 - 2 * (x ** 2 + y ** 2)
    return R


def nearest(residue_mask):
    index = [[0, 0] for _ in range(len(residue_mask))]
    p, q = 0, len(residue_mask)
    for i in range(len(residue_mask)):
        if residue_mask[i] == 0:
            p = i
        else:
            index[i][0] = p
    for i in range(len(residue_mask) - 1, -1, -1):
        if residue_mask[i] == 0:
            q = i
        else:
            index[i][1] = q
    return index


def interpolation_init(pred_X, residue_mask, backbone_pos, atom2residue, protein_atom_batch, residue_batch):
    num_protein = protein_atom_batch.max().item() + 1
    offset = 0
    for i in range(num_protein):
        residue_mask_i = residue_mask[residue_batch == i]
        backbone_pos_i = backbone_pos[residue_batch == i]
        if (~residue_mask_i).sum() <= 2:
            offset += len(residue_mask_i)
            continue
        else:
            residue_index = torch.arange(len(residue_mask_i)).to(protein_atom_batch.device)
            front = residue_index[~residue_mask_i][:2]
            end = residue_index[~residue_mask_i][-2:]
            near = nearest(residue_mask_i)
            for k in range(len(residue_mask_i)):
                if residue_mask_i[k]:
                    mask = atom2residue == (k + offset)
                    if k < front[0]:
                        pred_X[mask] = backbone_pos_i[front[0]] + (k - front[0]) / (front[0] - front[1]) * (
                                    backbone_pos_i[front[0]] - backbone_pos_i[front[1]])
                    elif k > end[1]:
                        pred_X[mask] = backbone_pos_i[end[1]] + (k - end[1]) / (end[1] - end[0]) * (
                                    backbone_pos_i[end[1]] - backbone_pos_i[end[0]])
                    else:
                        pred_X[mask] = ((k - near[k][0]) * backbone_pos_i[near[k][1]] + (near[k][1] - k) *
                                        backbone_pos_i[near[k][0]]) * 1 / (near[k][1] - near[k][0])
            offset += len(residue_mask_i)

    return pred_X


def interpolation_init_new(res_X, residue_mask, backbone_pos, residue_batch):
    num_protein = residue_batch.max().item() + 1
    offset = 0
    backbone = torch.tensor([[-0.525, 1.363, 0.0], [0.0, 0.0, 0.0], [1.526, 0.0, 0.0], [0.627, 1.062, 0.0]],
                            device=res_X.device)
    for i in range(num_protein):
        residue_mask_i = residue_mask[residue_batch == i]
        backbone_pos_i = backbone_pos[residue_batch == i]
        if (~residue_mask_i).sum() <= 2:
            offset += len(residue_mask_i)
            continue
        else:
            residue_index = torch.arange(len(residue_mask_i)).to(res_X.device)
            front = residue_index[~residue_mask_i][:2]
            end = residue_index[~residue_mask_i][-2:]
            near = nearest(residue_mask_i)
            for k in range(len(residue_mask_i)):
                if residue_mask_i[k]:
                    ind = k + offset
                    if k < front[0]:
                        alpha = (backbone_pos_i[front[0]] + (k - front[0]) / (front[0] - front[1]) * (backbone_pos_i[front[0]] - backbone_pos_i[front[1]]))[1: 2]
                    elif k > end[1]:
                        alpha = (backbone_pos_i[end[1]] + (k - end[1]) / (end[1] - end[0]) * (backbone_pos_i[end[1]] - backbone_pos_i[end[0]]))[1: 2]
                    else:
                        alpha = (((k - near[k][0]) * backbone_pos_i[near[k][1]] + (near[k][1] - k) * backbone_pos_i[near[k][0]]) * 1 / (near[k][1] - near[k][0]))[1: 2]
                    res_X[ind][:4] = alpha + backbone @ quaternion_to_matrix(q=torch.randn(4, device=res_X.device)).t()
            offset += len(residue_mask_i)

    return res_X


class Pocket_Design_new(Module):
    def __init__(self, config, protein_atom_feature_dim, ligand_atom_feature_dim, device):
        super().__init__()
        self.config = config
        self.device = device
        self.hidden_channels = config.hidden_channels
        self.protein_atom_emb = nn.Embedding(protein_atom_feature_dim, int(config.hidden_channels/2-8))
        self.ligand_atom_emb = Linear(ligand_atom_feature_dim, config.hidden_channels)
        self.encoder = get_encoder(config.encoder, device)
        self.residue_mlp = Linear(config.hidden_channels, 20)
        self.Softmax = nn.Softmax(dim=1)
        self.huber_loss = nn.SmoothL1Loss(reduction='mean')
        self.dist_loss = torch.nn.MSELoss(reduction='mean')
        self.pred_loss = nn.CrossEntropyLoss(reduction='mean')
        self.interpolate_steps = 3
        self.atom_pos_embedding = nn.Embedding(14, 8)
        self.residue_embedding = nn.Embedding(21, int(config.hidden_channels/2 - 16))  # one embedding for mask
        self.standard2alphabet = torch.tensor([1, 6, 13, 9, 19, 12, 5, 2, 17, 8, 0, 11, 16, 14, 10, 4, 7, 18, 15, 3]).to(device)
        self.alphabet2standard = torch.tensor([10, 0, 7, 19, 15, 6, 1, 16, 9, 3, 14, 11, 5, 2, 13, 18, 12, 8, 17, 4]).to(device)
        self.residue_atom_mask = residue_atom_mask.to(device)
        self.write_pdb = True
        self.write_whole_pdb = True
        self.generate_id = 0
        self.generate_id1 = 0
        self.proteinloss = ProteinFeature()
        self.pe = PositionalEncodings(16)
        self.res_atom_type = torch.tensor(RES_ATOM_TYPE).to(device)
        self.orig_data_path = config.orig_data_path
        self.pocket10_path = config.pocket10_path
        if config.encoder.esm[:4] == 'esm2':
            encoder_args = {'_target_': 'esm2_adapter',
                            'encoder': {'d_model': 128,
                                        'use_esm_alphabet': True},
                            'dropout': 0.1,
                            'adapter_layer_indices': [6, 20, 32]}
            self.esmadapter = ESM2WithStructuralAdatper.from_pretrained(args=encoder_args, name=config.encoder.esm).to(device)
        else:
            encoder_args = {'_target_': 'esm_adapter',
                            'encoder': {'d_model': 128,
                                        'n_enc_layers': 3,
                                        'n_dec_layers': 3,
                                        'use_esm_alphabet': True},
                            'adapter_layer_indices': [6, 20, 32]}
            self.esmadapter = ProteinBertModelWithStructuralAdatper.from_pretrained(args=encoder_args).to(device)

    def forward_(self, batch):
        loss_list = [0., 0., 0.]
        residue_mask = batch['protein_edit_residue']
        full_seq = batch['seq']
        ligand_mask = batch['ligand_mask'].bool()
        label_ligand, pred_ligand = copy.deepcopy(batch['ligand_pos']), copy.deepcopy(batch['ligand_pos'])

        # init res_X
        label_X, res_X = copy.deepcopy(batch['residue_pos']), copy.deepcopy(batch['residue_pos'])
        res_X = interpolation_init_new(res_X, residue_mask, copy.deepcopy(batch['backbone_pos']),batch['amino_acid_batch'])
        for k in range(len(batch['amino_acid'])):
            if residue_mask[k]:
                pos = res_X[k]
                pos[4:] = (pos[1].repeat(10, 1) + 0.1 * torch.randn(10, 3, device=self.device))
                res_X[k] = pos
        pred_ligand = label_ligand + torch.randn_like(label_ligand).to(self.device) * 0.5

        ligand_feat = self.ligand_atom_emb(batch['ligand_feat'])

        for t in range(self.interpolate_steps):
            print(t)
            res_S = copy.deepcopy(batch['amino_acid_processed'])
            if t > 1:
                '''
                res_H[residue_mask] = res_H[residue_mask] + torch.matmul(pred_res_type[:, self.alphabet2standard].detach().float(), self.residue_embedding(torch.arange(1, 21).to(self.device))).unsqueeze(1)
                res_H[~residue_mask] = res_H[~residue_mask] + self.residue_embedding(res_S[~residue_mask]).unsqueeze(-2)
                '''
                res_S[residue_mask] = self.alphabet2standard[sampled_type.detach().clone()] + 1
                atom_emb = self.protein_atom_emb(self.res_atom_type[res_S])  # atom embedding
                atom_pos_emb = self.atom_pos_embedding(torch.arange(14).to(self.device)).unsqueeze(0).repeat(res_S.shape[0], 1, 1)  # pos embedding
                res_emb = self.residue_embedding(res_S).unsqueeze(-2).repeat(1, 14, 1)  # res embedding
                res_pos_emb = self.pe(batch['res_idx']).unsqueeze(-2).repeat(1, 14, 1)  # res pos embedding
                res_H = torch.cat([atom_emb, atom_pos_emb, res_emb, res_pos_emb], dim=-1)
            elif t <= 1:
                atom_emb = self.protein_atom_emb(self.res_atom_type[res_S])  # atom embedding
                atom_pos_emb = self.atom_pos_embedding(torch.arange(14).to(self.device)).unsqueeze(0).repeat(res_S.shape[0], 1, 1)  # pos embedding
                res_emb = self.residue_embedding(res_S).unsqueeze(-2).repeat(1, 14, 1)  # res embedding
                res_pos_emb = self.pe(batch['res_idx']).unsqueeze(-2).repeat(1, 14, 1)  # res pos embedding
                res_H = torch.cat([atom_emb, atom_pos_emb, res_emb, res_pos_emb], dim=-1)

            _, res_X, pred_res_type, pred_ligand = self.encoder(res_H, res_X.detach().clone(), res_S, batch['amino_acid_batch'], full_seq, pred_ligand.detach().clone(),
                             ligand_feat, batch['ligand_mask'], batch['edit_residue_num'], residue_mask, self.esmadapter, batch['full_seq_mask'], batch['r10_mask'])

            atom_mask = self.residue_atom_mask[batch['amino_acid'][residue_mask]].bool()

            loss_list[0] += 2*self.huber_loss(res_X[residue_mask][atom_mask],label_X[residue_mask][atom_mask]) + self.huber_loss(pred_ligand[ligand_mask], label_ligand[ligand_mask])
            loss_list[1] += self.pred_loss(pred_res_type, self.standard2alphabet[batch['amino_acid'][residue_mask] - 1])
            # bond and angle loss
            # loss_list[2] += 3*self.proteinloss.structure_loss(res_X[residue_mask], label_X[residue_mask], batch['amino_acid'][residue_mask] - 1, batch['res_idx'][residue_mask], batch['amino_acid_batch'][residue_mask])
            loss_list[2] += 0.

            sampled_type, _ = sample_from_categorical(pred_res_type.detach())
        aar = (self.standard2alphabet[batch['amino_acid'][residue_mask] - 1] == sampled_type).sum() / len(res_S[residue_mask])
        rmsd = torch.sqrt((res_X[residue_mask][:, :4].reshape(-1, 3) - label_X[residue_mask][:, :4].reshape(-1, 3)).norm(dim=1).sum() / len(res_S[residue_mask]) / 4)

        return loss_list[1] + loss_list[0] + loss_list[2], loss_list, aar, rmsd

    def init(self, batch):
        residue_mask = batch['protein_edit_residue']
        label_ligand, pred_ligand = copy.deepcopy(batch['ligand_pos']), copy.deepcopy(batch['ligand_pos'])
        pred_ligand = label_ligand + torch.randn_like(label_ligand).to(self.device) * 0.5
        res_X = copy.deepcopy(batch['residue_pos'])  # init res_X
        res_X = interpolation_init_new(res_X, residue_mask, copy.deepcopy(batch['backbone_pos']),
                                       batch['amino_acid_batch'])
        res_S = copy.deepcopy(batch['amino_acid_processed'])
        for k in range(len(batch['amino_acid'])):  # init side chain atoms of masked residues
            if residue_mask[k]:
                pos = res_X[k]
                pos[4:] = (pos[1].repeat(10, 1) + 0.1 * torch.randn(10, 3, device=self.device))
                res_X[k] = pos

        ligand_feat = self.ligand_atom_emb(batch['ligand_feat'])

        atom_emb = self.protein_atom_emb(self.res_atom_type[res_S])  # atom embedding
        atom_pos_emb = self.atom_pos_embedding(torch.arange(14).to(self.device)).unsqueeze(0).repeat(res_S.shape[0], 1,
                                                                                                     1)  # pos embedding
        res_emb = self.residue_embedding(res_S).unsqueeze(-2).repeat(1, 14, 1)  # res embedding
        res_pos_emb = self.pe(batch['res_idx']).unsqueeze(-2).repeat(1, 14, 1)  # res pos embedding
        res_H = torch.cat([atom_emb, atom_pos_emb, res_emb, res_pos_emb], dim=-1)
        self.seq = batch['seq']
        self.full_seq_mask = batch['full_seq_mask']
        self.r10_mask = batch['r10_mask']
        return res_H, res_X, res_S, batch['amino_acid_batch'], pred_ligand, ligand_feat, batch['ligand_mask'], batch['edit_residue_num'], residue_mask

    def forward(self, res_H, res_X, res_S, res_batch, pred_ligand, ligand_feat, ligand_mask, edit_residue_num, residue_mask, use_esm=True):
        '''
        res_H[residue_mask] = res_H[residue_mask] + torch.matmul(pred_res_type[:, self.alphabet2standard].detach().float(), self.residue_embedding(torch.arange(1, 21).to(self.device))).unsqueeze(1)
        res_H[~residue_mask] = res_H[~residue_mask] + self.residue_embedding(res_S[~residue_mask]).unsqueeze(-2)
        '''

        res_H, res_X, ligand_pos, ligand_feat, pred_res_type = self.encoder(res_H, res_X, res_S, res_batch, pred_ligand, ligand_feat, ligand_mask, edit_residue_num, residue_mask)

        if use_esm and self.seq.shape[1] <= 1000:
            h_residue = res_H.sum(-2)
            batch_size = res_batch.max().item() + 1
            encoder_out = {'feats': torch.zeros(batch_size, self.seq.shape[1], self.hidden_channels).to(self.device)}
            encoder_out['feats'][self.r10_mask] = h_residue.view(-1, self.hidden_channels)
            init_pred = self.seq
            decode_logits = self.esmadapter(init_pred, encoder_out)['logits']
            pred_res_type = decode_logits[self.full_seq_mask][:, 4:24]

        return res_H, res_X, ligand_pos, ligand_feat, pred_res_type

    def generate(self, batch, target_path='./generate'):
        print('Start Generating')
        residue_mask = batch['protein_edit_residue']
        res_S = batch['amino_acid_processed']
        full_seq = batch['seq']
        label_S = copy.deepcopy(batch['amino_acid'])
        label_X, res_X = copy.deepcopy(batch['residue_pos']), copy.deepcopy(batch['residue_pos'])
        label_ligand, pred_ligand = copy.deepcopy(batch['ligand_pos']), copy.deepcopy(batch['ligand_pos'])
        res_X = interpolation_init_new(res_X, residue_mask, copy.deepcopy(batch['backbone_pos']), batch['amino_acid_batch'])
        res_batch = batch['amino_acid_batch']
        for k in range(len(batch['amino_acid'])):
            if residue_mask[k]:
                pos = res_X[k]
                pos[4:] = (pos[1].repeat(10, 1) + 0.1 * torch.randn(10, 3, device=self.device))
                res_X[k] = pos

        ligand_feat = self.ligand_atom_emb(batch['ligand_feat'])

        for t in range(self.interpolate_steps):
            if t < -1:
                res_S[residue_mask] = self.alphabet2standard[sampled_type.detach().clone()] + 1
                atom_emb = self.protein_atom_emb(self.res_atom_type[res_S])  # atom embedding
                atom_pos_emb = self.atom_pos_embedding(torch.arange(14).to(self.device)).unsqueeze(0).repeat(res_S.shape[0], 1, 1)  # pos embedding
                res_emb = self.residue_embedding(res_S).unsqueeze(-2).repeat(1, 14, 1)  # res embedding
                res_pos_emb = self.pe(batch['res_idx']).unsqueeze(-2).repeat(1, 14, 1)  # res pos embedding
                res_H = torch.cat([atom_emb, atom_pos_emb, res_emb, res_pos_emb], dim=-1)
            elif t == 0:
                atom_emb = self.protein_atom_emb(self.res_atom_type[res_S])  # atom embedding
                atom_pos_emb = self.atom_pos_embedding(torch.arange(14).to(self.device)).unsqueeze(0).repeat(res_S.shape[0], 1, 1)  # pos embedding
                res_emb = self.residue_embedding(res_S).unsqueeze(-2).repeat(1, 14, 1)  # res embedding
                res_pos_emb = self.pe(batch['res_idx']).unsqueeze(-2).repeat(1, 14, 1)  # res pos embedding
                res_H = torch.cat([atom_emb, atom_pos_emb, res_emb, res_pos_emb], dim=-1)

            res_H, res_X, pred_ligand, ligand_feat, pred_res_type, attend_logits = self.encoder(res_H, res_X, res_S, res_batch, pred_ligand, ligand_feat, batch['ligand_mask'], batch['edit_residue_num'], residue_mask)
            if full_seq.shape[1] <= 1000:
                h_residue = res_H.sum(-2)
                batch_size = res_batch.max().item() + 1
                encoder_out = {
                    'feats': torch.zeros(batch_size, full_seq.shape[1], self.hidden_channels).to(self.device)}
                encoder_out['feats'][batch['r10_mask']] = h_residue.view(-1, self.hidden_channels)
                init_pred = full_seq
                decode_logits = self.esmadapter(init_pred, encoder_out)['logits']
                pred_res_type = decode_logits[batch['full_seq_mask']][:, 4:24]

            sampled_type, _ = sample_from_categorical(pred_res_type)

        aar = (self.standard2alphabet[batch['amino_acid'][residue_mask] - 1] == sampled_type).sum() / len(label_S[residue_mask])
        rmsd = torch.sqrt((res_X[residue_mask][:, :4].reshape(-1, 3) - label_X[residue_mask][:, :4].reshape(-1, 3)).norm(dim=1).sum() / len(label_S[residue_mask]) / 4)
        
        if self.write_pdb:
            res_S[residue_mask] = self.alphabet2standard[sampled_type.detach().clone()] + 1
            to_sdf(pred_ligand, batch['ligand_element'].long(), batch['ligand_mask'].bool(), batch['ligand_batch'],batch['ligand_bond_type'].long(), batch['ligand_bond_index'].long(), batch['edge_batch'], self.generate_id, target_path)
            to_pdb(label_X, batch['amino_acid'], batch['res_idx'], batch['amino_acid_batch'], self.generate_id, batch['protein_filename'], target_path, original=True)
            self.generate_id = to_pdb(res_X, res_S, batch['res_idx'], batch['amino_acid_batch'], self.generate_id, batch['protein_filename'], target_path, original = False)
            
        if self.write_whole_pdb:
            self.generate_id1 = to_whole_pdb(res_X, res_S, batch['res_idx'], batch['amino_acid_batch'], self.generate_id1, batch['protein_filename'], batch['r10_mask'], self.orig_data_path, target_path)
        return aar, rmsd, attend_logits


def sample_from_categorical(logits=None, temperature=3.0):
    if temperature:
        dist = torch.distributions.Categorical(logits=logits.div(temperature))
        tokens = dist.sample()
        scores = dist.log_prob(tokens)
    else:
        scores, tokens = logits.log_softmax(dim=-1).max(dim=-1)
    return tokens, scores


def sample_from_topk(tensor, k=3):
    """
    Apply softmax to the tensor, then randomly sample an index from the top k values.

    :param tensor: Input tensor.
    :param k: Number of top values to consider.
    :return: Index of the sampled value.
    """

    # Apply softmax
    probs = torch.nn.functional.softmax(tensor, dim=0)

    # Get top k values and their indices
    _, top_indices = torch.topk(probs, k)

    sampled_indices = torch.randint(0, k, (top_indices.shape[0],))

    # Use advanced indexing to gather the sampled elements from each row
    sampled_elements = top_indices[torch.arange(top_indices.shape[0]), sampled_indices]

    return sampled_elements, None


def random_mask(batch, device, mask=True):
    if mask:
        tmp = []
        num_protein = batch['protein_atom_batch'].max() + 1
        for i in range(num_protein):
            mask = batch['amino_acid_batch'] == i
            ind = torch.multinomial(batch['protein_edit_residue'][mask].float(), 1)
            selected = torch.zeros_like(batch['protein_edit_residue'][mask], dtype=bool)
            selected[ind] = 1
            tmp.append(selected)
        batch['random_mask_residue'] = torch.cat(tmp, dim=0)

        # remove side chains for the masked atoms
        index = torch.arange(len(batch['amino_acid']))[batch['random_mask_residue']]
        for key in ['protein_pos', 'protein_atom_feature']:
            tmp = []
            for k in range(batch['atom2residue'].max() + 1):
                mask = batch['atom2residue'] == k
                if k in index:
                    if key == 'protein_atom_feature':
                        feature_mask = batch['protein_atom_feature'][mask]
                        feature_mask[:, -20:] = torch.zeros(20, device=device)
                        feature_mask[:, -21] = 1
                        batch['protein_atom_feature'][mask] = feature_mask
                    tmp.append(batch[key][mask][:4])
                else:
                    tmp.append(batch[key][mask])
            batch[key] = torch.cat(tmp, dim=0)
        batch['residue_natoms'][batch['random_mask_residue']] = 4
        batch['atom2residue'] = torch.repeat_interleave(torch.arange(len(batch['residue_natoms']), device=device),
                                                        batch['residue_natoms'])
        batch['protein_edit_atom'] = torch.repeat_interleave(batch['protein_edit_residue'], batch['residue_natoms'],
                                                             dim=0)
        batch['random_mask_atom'] = torch.repeat_interleave(batch['random_mask_residue'], batch['residue_natoms'],
                                                            dim=0)
    else:
        # reset protein pos and feature
        index = torch.arange(len(batch['amino_acid']))[batch['random_mask_residue']]
        num_residues = batch['atom2residue'].max() + 1
        pos_tmp, feature_tmp, natoms_tmp = [], [], []
        for k in range(num_residues):
            mask = batch['atom2residue'] == k
            res_type = batch['amino_acid'][k]
            sidechain_size = NUM_ATOMS[res_type] - 4
            if k in index:
                pos_tmp.append(batch['protein_pos'][mask][:4])
                if sidechain_size > 0:
                    pos_tmp.append(
                        batch['protein_pos'][mask][1:2].repeat(sidechain_size, 1) + 0.1 * torch.randn(sidechain_size, 3,
                                                                                                      device=device))
                feature_tmp.append(atom_feature(res_type, device))
                natoms_tmp.append(NUM_ATOMS[res_type])
            else:
                pos_tmp.append(batch['protein_pos'][mask])
                feature_tmp.append(batch['protein_atom_feature'][mask])
                natoms_tmp.append(batch['protein_pos'][mask].shape[0])
        batch['protein_pos'], batch['protein_atom_feature'] = torch.cat(pos_tmp, dim=0), torch.cat(feature_tmp, dim=0)
        batch['protein_atom_feature'][:, -21] = 0

        batch['residue_natoms'] = torch.tensor(natoms_tmp, device=device)
        batch['atom2residue'] = torch.repeat_interleave(torch.arange(len(batch['residue_natoms']), device=device),
                                                        batch['residue_natoms'])
        batch['protein_edit_atom'] = torch.repeat_interleave(batch['protein_edit_residue'], batch['residue_natoms'],
                                                             dim=0)

    # follow batch
    num_protein = batch['protein_atom_batch'].max() + 1
    repeats = torch.tensor([batch['residue_natoms'][batch['amino_acid_batch'] == i].sum() for i in range(num_protein)])
    batch['protein_atom_batch'] = torch.repeat_interleave(torch.arange(num_protein), repeats).to(device)
    batch['edit_backbone'] = copy.deepcopy(batch['protein_edit_atom'])
    index = torch.arange(len(batch['amino_acid']))[batch['protein_edit_residue']]
    for k in range(len(batch['amino_acid'])):
        mask = batch['atom2residue'] == k
        if k in index:
            data_mask = batch['edit_backbone'][mask]
            data_mask[4:] = 0
            batch['edit_backbone'][mask] = data_mask

    return batch


def atom_feature(res_type, device):
    atom_types = torch.arange(38)
    max_num_aa = 21
    atom_type = torch.tensor(RES_ATOMS[res_type]).view(-1, 1) == atom_types.view(1, -1)
    amino_acid = F.one_hot(res_type, num_classes=max_num_aa).repeat(NUM_ATOMS[res_type], 1)
    x = torch.cat([atom_type.to(device), amino_acid], dim=-1)
    return x


def to_pdb(res_X, amino_acid, res_idx, res_batch, index, pocket_filename, target_path, original):
    lines = ['HEADER    POCKET', 'COMPND    POCKET\n']
    num_protein = res_batch.max().item() + 1
    for n in range(num_protein):
        #pdb_path = os.path.join(orig_data_path, pocket_filename[n])
        pdb_path = pocket_filename[n]
        with open(pdb_path, 'r') as f:
            pdb_block = f.read()
        protein = PDBProtein(pdb_block)
        residues, atoms = protein.return_residues()
        mask = (res_batch == n)
        res_X_protein = res_X[mask]
        amino_acid_protein = amino_acid[mask]
        res_idx_protein = res_idx[mask]
        atom_count = 0
        if original:
            path = os.path.join(target_path, str(index + n) + '_orig.pdb')
        else:
            path = os.path.join(target_path, str(index + n) + '.pdb')
        with open(path, 'w') as f:
            f.writelines(lines)
            for k in range(len(res_X_protein)):
                atom_type = RES_ATOM14[amino_acid_protein[k]]
                chain = residues[k]['chain']
                for i in range(NUM_ATOMS[amino_acid_protein[k]]):
                    j0 = str('ATOM').ljust(6)  # atom#6s
                    j1 = str(atom_count).rjust(5)  # aomnum#5d
                    j2 = str(atom_type[i]).center(4)  # atomname$#4s
                    j3 = AA_NUMBER_NAME[amino_acid_protein[k].item()].ljust(3)  # resname#1s
                    j4 = str(chain).rjust(1)  # Astring
                    j5 = str(res_idx_protein[k].item()).rjust(4)  # resnum
                    j6 = str('%8.3f' % (float(res_X_protein[k, i, 0]))).rjust(8)  # x
                    j7 = str('%8.3f' % (float(res_X_protein[k, i, 1]))).rjust(8)  # y
                    j8 = str('%8.3f' % (float(res_X_protein[k, i, 2]))).rjust(8)  # z\
                    j9 = str('%6.2f' % (1.00)).rjust(6)  # occ
                    j10 = str('%6.2f' % (25.02)).ljust(6)  # temp
                    j11 = str(atom_type[i][0]).rjust(12)  # elname
                    f.write("%s%s %s %s %s%s    %s%s%s%s%s%s\n" % (j0, j1, j2, j3, j4, j5, j6, j7, j8, j9, j10, j11))
                    atom_count += 1
            f.write('END')
            f.write('\n')
        openmm_relax(path)
    return index + num_protein


def to_whole_pdb(res_X, amino_acid, res_idx, res_batch, index, protein_filename, r10_mask, orig_data_path, target_path):
    lines = ['HEADER    POCKET', 'COMPND    POCKET\n']
    num_protein = res_batch.max().item() + 1
    for n in range(num_protein):
        pdb_path = protein_filename[n]
        with open(pdb_path, 'r') as f:
            pdb_block = f.read()
        protein = PDBProtein(pdb_block)
        residues, atoms = protein.return_residues()

        mask = (res_batch == n)
        res_X_protein = res_X[mask]
        amino_acid_protein = amino_acid[mask]
        res_idx_protein = res_idx[mask]
        assert r10_mask[n].sum() == len(amino_acid_protein)

        path = os.path.join(target_path, str(index + n) + '_whole.pdb')
        atom_count = 0
        stored_res_count = 0
        with open(path, 'w') as f:
            f.writelines(lines)
            for k in range(len(residues)):
                if r10_mask[n, k+1]:
                    chain = atoms[residues[k]['atoms'][0]]['line'][21:22].strip()
                    atom_type = RES_ATOM14[amino_acid_protein[stored_res_count]]
                    for i in range(NUM_ATOMS[amino_acid_protein[stored_res_count]]):
                        j0 = str('ATOM').ljust(6)  # atom#6s
                        j1 = str(atom_count).rjust(5)  # aomnum#5d
                        j2 = str(atom_type[i]).center(4)  # atomname$#4s
                        j3 = AA_NUMBER_NAME[amino_acid_protein[stored_res_count].item()].ljust(3)  # resname#1s
                        j4 = str(chain).rjust(1)  # Astring
                        j5 = str(res_idx_protein[stored_res_count].item()).rjust(4)  # resnum
                        j6 = str('%8.3f' % (float(res_X_protein[stored_res_count, i, 0]))).rjust(8)  # x
                        j7 = str('%8.3f' % (float(res_X_protein[stored_res_count, i, 1]))).rjust(8)  # y
                        j8 = str('%8.3f' % (float(res_X_protein[stored_res_count, i, 2]))).rjust(8)  # z\
                        j9 = str('%6.2f' % (1.00)).rjust(6)  # occ
                        j10 = str('%6.2f' % (25.02)).ljust(6)  # temp
                        j11 = str(atom_type[i][0]).rjust(12)  # elname
                        f.write("%s%s %s %s %s%s    %s%s%s%s%s%s\n" % (j0, j1, j2, j3, j4, j5, j6, j7, j8, j9, j10, j11))
                        atom_count += 1
                    stored_res_count += 1
                else:
                    for atom_idx in residues[k]['atoms']:
                        line = atoms[atom_idx]['line']
                        line = line[:6] + str(atom_count).rjust(5) + line[11:] + "\n"
                        atom_count += 1
                        f.write(line)
            f.write('END')
            f.write('\n')
        openmm_relax(path)
    return index + num_protein


def to_sdf(pred_pos, elements, mask, ligand_batch, bond_types, bond_index, edge_batch, id, target_path):
    num_ligand = edge_batch.max().item() + 1
    for l in range(num_ligand):
        filename = os.path.join(target_path, str(id + l) + '.sdf')
        positions = pred_pos[l][mask[l]]
        elements_protein = elements[ligand_batch == l]
        bond_types_protein = bond_types[edge_batch == l]
        bond_index_protein = bond_index[:, edge_batch == l].transpose(0, 1)

        mol = rdchem.EditableMol(Chem.Mol())

        # Add atoms to molecule
        for element in elements_protein:
            atom = Chem.Atom(element.item())
            mol.AddAtom(atom)

        # Add bonds to molecule
        edge_set = set()
        for k, (bond_type, (start_idx, end_idx)) in enumerate(zip(bond_types_protein, bond_index_protein)):
            if (start_idx.item(), end_idx.item()) not in edge_set:
                edge_set.add((start_idx.item(), end_idx.item()))
                edge_set.add((end_idx.item(), start_idx.item()))
                mol.AddBond(start_idx.item(), end_idx.item(), BOND_TYPE[bond_type.item()])

        # Set 3D coordinates (assuming positions are in 3D)
        mol = mol.GetMol()
        conf = Chem.Conformer(mol.GetNumAtoms())
        for i, position in enumerate(positions):
            conf.SetAtomPosition(i, position.tolist())
        mol.AddConformer(conf)

        writer = Chem.SDWriter(filename)
        writer.write(mol)
        writer.close()

    return mol


def init_weight(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_normal_(m.weight)
    elif isinstance(m, nn.Conv1d):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
    elif isinstance(m, nn.LayerNorm):
        nn.init.constant_(m.weight, 1)


class AminoAcidFeature(nn.Module):
    def __init__(self, backbone_only=False) -> None:
        super().__init__()

        self.backbone_only = backbone_only

        # number of classes
        self.num_aa_type = len(VOCAB)
        self.num_atom_type = VOCAB.get_num_atom_type()
        self.num_atom_pos = VOCAB.get_num_atom_pos()

        # atom-level special tokens
        self.atom_mask_idx = VOCAB.get_atom_mask_idx()
        self.atom_pad_idx = VOCAB.get_atom_pad_idx()
        self.atom_pos_mask_idx = VOCAB.get_atom_pos_mask_idx()
        self.atom_pos_pad_idx = VOCAB.get_atom_pos_pad_idx()

        # global nodes and mask nodes
        self.boa_idx = VOCAB.symbol_to_idx(VOCAB.BOA)
        self.boh_idx = VOCAB.symbol_to_idx(VOCAB.BOH)
        self.bol_idx = VOCAB.symbol_to_idx(VOCAB.BOL)
        self.mask_idx = VOCAB.get_mask_idx()

        # atoms encoding
        residue_atom_type, residue_atom_pos = [], []
        backbone = [VOCAB.atom_to_idx(atom[0]) for atom in VOCAB.backbone_atoms]
        n_channel = VOCAB.MAX_ATOM_NUMBER if not backbone_only else 4
        special_mask = VOCAB.get_special_mask()
        for i in range(len(VOCAB)):
            if i == self.boa_idx or i == self.boh_idx or i == self.bol_idx or i == self.mask_idx:
                # global nodes
                residue_atom_type.append([self.atom_mask_idx for _ in range(n_channel)])
                residue_atom_pos.append([self.atom_pos_mask_idx for _ in range(n_channel)])
            elif special_mask[i] == 1:
                # other special token (pad)
                residue_atom_type.append([self.atom_pad_idx for _ in range(n_channel)])
                residue_atom_pos.append([self.atom_pos_pad_idx for _ in range(n_channel)])
            else:
                # normal amino acids
                sidechain_atoms = VOCAB.get_sidechain_info(VOCAB.idx_to_symbol(i))
                atom_type = backbone
                atom_pos = [VOCAB.atom_pos_to_idx(VOCAB.atom_pos_bb) for _ in backbone]
                if not backbone_only:
                    sidechain_atoms = VOCAB.get_sidechain_info(VOCAB.idx_to_symbol(i))
                    atom_type = atom_type + [VOCAB.atom_to_idx(atom[0]) for atom in sidechain_atoms]
                    atom_pos = atom_pos + [VOCAB.atom_pos_to_idx(atom[1]) for atom in sidechain_atoms]
                num_pad = n_channel - len(atom_type)
                residue_atom_type.append(atom_type + [self.atom_pad_idx for _ in range(num_pad)])
                residue_atom_pos.append(atom_pos + [self.atom_pos_pad_idx for _ in range(num_pad)])

        # mapping from residue to atom types and positions
        self.residue_atom_type = nn.parameter.Parameter(
            torch.tensor(residue_atom_type, dtype=torch.long),
            requires_grad=False)
        self.residue_atom_pos = nn.parameter.Parameter(
            torch.tensor(residue_atom_pos, dtype=torch.long),
            requires_grad=False)

        # sidechain geometry
        if not backbone_only:
            sc_bonds, sc_bonds_mask = [], []
            sc_chi_atoms, sc_chi_atoms_mask = [], []
            for i in range(len(VOCAB)):
                if special_mask[i] == 1:
                    sc_bonds.append([])
                    sc_chi_atoms.append([])
                else:
                    symbol = VOCAB.idx_to_symbol(i)
                    atom_type = VOCAB.backbone_atoms + VOCAB.get_sidechain_info(symbol)
                    atom2channel = {atom: i for i, atom in enumerate(atom_type)}
                    chi_atoms, bond_atoms = VOCAB.get_sidechain_geometry(symbol)
                    sc_chi_atoms.append(
                        [[atom2channel[atom] for atom in atoms] for atoms in chi_atoms]
                    )
                    bonds = []
                    for src_atom in bond_atoms:
                        for dst_atom in bond_atoms[src_atom]:
                            bonds.append((atom2channel[src_atom], atom2channel[dst_atom]))
                    sc_bonds.append(bonds)
            max_num_chis = max([len(chis) for chis in sc_chi_atoms])
            max_num_bonds = max([len(bonds) for bonds in sc_bonds])
            for i in range(len(VOCAB)):
                num_chis, num_bonds = len(sc_chi_atoms[i]), len(sc_bonds[i])
                num_pad_chis, num_pad_bonds = max_num_chis - num_chis, max_num_bonds - num_bonds
                sc_chi_atoms_mask.append(
                    [1 for _ in range(num_chis)] + [0 for _ in range(num_pad_chis)]
                )
                sc_bonds_mask.append(
                    [1 for _ in range(num_bonds)] + [0 for _ in range(num_pad_bonds)]
                )
                sc_chi_atoms[i].extend([[-1, -1, -1, -1] for _ in range(num_pad_chis)])
                sc_bonds[i].extend([(-1, -1) for _ in range(num_pad_bonds)])

            # mapping residues to their sidechain chi angle atoms and bonds
            self.sidechain_chi_angle_atoms = nn.parameter.Parameter(
                torch.tensor(sc_chi_atoms, dtype=torch.long),
                requires_grad=False)
            self.sidechain_chi_mask = nn.parameter.Parameter(
                torch.tensor(sc_chi_atoms_mask, dtype=torch.bool),
                requires_grad=False
            )
            self.sidechain_bonds = nn.parameter.Parameter(
                torch.tensor(sc_bonds, dtype=torch.long),
                requires_grad=False
            )
            self.sidechain_bonds_mask = nn.parameter.Parameter(
                torch.tensor(sc_bonds_mask, dtype=torch.bool),
                requires_grad=False
            )

    def _construct_residue_pos(self, S):
        # construct residue position. global node is 1, the first residue is 2, ... (0 for padding)
        glbl_node_mask = self._is_global(S)
        glbl_node_idx = torch.nonzero(glbl_node_mask).flatten()  # [batch_size * 3] (boa, boh, bol)
        shift = F.pad(glbl_node_idx[:-1] - glbl_node_idx[1:] + 1, (1, 0), value=1)  # [batch_size * 3]
        residue_pos = torch.ones_like(S)
        residue_pos[glbl_node_mask] = shift
        residue_pos = torch.cumsum(residue_pos, dim=0)
        return residue_pos

    def _construct_segment_ids(self, res_idx, batch):
        consecutive = (res_idx[1:] == res_idx[:-1]) & (batch[1:] == batch[:-1])
        segment_ids = torch.zeros_like(res_idx).long()
        id = 0
        for i in range(1, len(segment_ids)):
            if consecutive[i - 1]:
                segment_ids[i] = id
            else:
                id += 1
                segment_ids[i] = id
        return segment_ids

    def _construct_atom_type(self, S):
        # construct atom types
        return self.residue_atom_type[S]

    def _construct_atom_pos(self, S):
        # construct atom positions
        return self.residue_atom_pos[S]

    @torch.no_grad()
    def get_sidechain_chi_angles_atoms(self, S):
        chi_angles_atoms = self.sidechain_chi_angle_atoms[S]  # [N, max_num_chis, 4]
        chi_mask = self.sidechain_chi_mask[S]  # [N, max_num_chis]
        return chi_angles_atoms, chi_mask

    @torch.no_grad()
    def get_sidechain_bonds(self, S):
        bonds = self.sidechain_bonds[S]  # [N, max_num_bond, 2]
        bond_mask = self.sidechain_bonds_mask[S]
        return bonds, bond_mask

    def forward(self, X, S, batch_id, k_neighbors):
        H, (_, _, atom_pos) = self.embedding(S)
        ctx_edges, inter_edges = self.construct_edges(
            X, S, batch_id, k_neighbors, atom_pos=atom_pos)
        return H, (ctx_edges, inter_edges)


class ProteinFeature(nn.Module):
    def __init__(self, backbone_only=False):
        super().__init__()
        self.backbone_only = backbone_only
        self.aa_feature = AminoAcidFeature()

    def _cal_sidechain_bond_lengths(self, S, X):
        bonds, bonds_mask = self.aa_feature.get_sidechain_bonds(S)
        n = torch.nonzero(bonds_mask)[:, 0]  # [Nbonds]
        src, dst = bonds[bonds_mask].T
        src_X, dst_X = X[(n, src)], X[(n, dst)]  # [Nbonds, 3]
        bond_lengths = torch.norm(dst_X - src_X, dim=-1)
        return bond_lengths

    def _cal_sidechain_chis(self, S, X):
        chi_atoms, chi_mask = self.aa_feature.get_sidechain_chi_angles_atoms(S)
        n = torch.nonzero(chi_mask)[:, 0]  # [Nchis]
        a0, a1, a2, a3 = chi_atoms[chi_mask].T  # [Nchis]
        x0, x1, x2, x3 = X[(n, a0)], X[(n, a1)], X[(n, a2)], X[(n, a3)]  # [Nchis, 3]
        u_0, u_1, u_2 = (x1 - x0), (x2 - x1), (x3 - x2)  # [Nchis, 3]
        # normals of the two planes
        n_1 = F.normalize(torch.cross(u_0, u_1), dim=-1)  # [Nchis, 3]
        n_2 = F.normalize(torch.cross(u_1, u_2), dim=-1)  # [Nchis, 3]
        cosChi = (n_1 * n_2).sum(-1)  # [Nchis]
        eps = 1e-7
        cosChi = torch.clamp(cosChi, -1 + eps, 1 - eps)
        return cosChi

    def _cal_backbone_bond_lengths(self, X, seg_id):
        # loss of backbone (...N-CA-C(O)-N...) bond length
        # N-CA, CA-C, C=O
        bl1 = torch.norm(X[:, 1:4] - X[:, :3], dim=-1)  # [N, 3], (N-CA), (CA-C), (C=O)
        # C-N
        bl2 = torch.norm(X[1:, 0] - X[:-1, 2], dim=-1)  # [N-1]
        same_chain_mask = seg_id[1:] == seg_id[:-1]
        bl2 = bl2[same_chain_mask]
        bl = torch.cat([bl1.flatten(), bl2], dim=0)
        return bl

    def _cal_angles(self, X, seg_id):
        ori_X = X
        X = X[:, :3].reshape(-1, 3)  # [N * 3, 3], N, CA, C
        U = F.normalize(X[1:] - X[:-1], dim=-1)  # [N * 3 - 1, 3]

        # 1. dihedral angles
        u_2, u_1, u_0 = U[:-2], U[1:-1], U[2:]  # [N * 3 - 3, 3]
        # backbone normals
        n_2 = F.normalize(torch.cross(u_2, u_1), dim=-1)
        n_1 = F.normalize(torch.cross(u_1, u_0), dim=-1)
        # angle between normals
        eps = 1e-7
        cosD = (n_2 * n_1).sum(-1)  # [(N-1) * 3]
        cosD = torch.clamp(cosD, -1 + eps, 1 - eps)
        # D = torch.sign((u_2 * n_1).sum(-1)) * torch.acos(cosD)
        seg_id_atom = seg_id.repeat(1, 3).flatten()  # [N * 3]
        same_chain_mask = sequential_and(
            seg_id_atom[:-3] == seg_id_atom[1:-2],
            seg_id_atom[1:-2] == seg_id_atom[2:-1],
            seg_id_atom[2:-1] == seg_id_atom[3:]
        )  # [N * 3 - 3]
        # D = D[same_chain_mask]
        cosD = cosD[same_chain_mask]

        # 2. bond angles (C_{n-1}-N, N-CA), (N-CA, CA-C), (CA-C, C=O), (CA-C, C-N_{n+1}), (O=C, C-Nn)
        u_0, u_1 = U[:-1], U[1:]  # [N*3 - 2, 3]
        cosA1 = ((-u_0) * u_1).sum(-1)  # [N*3 - 2], (C_{n-1}-N, N-CA), (N-CA, CA-C), (CA-C, C-N_{n+1})
        same_chain_mask = sequential_and(
            seg_id_atom[:-2] == seg_id_atom[1:-1],
            seg_id_atom[1:-1] == seg_id_atom[2:]
        )
        cosA1 = cosA1[same_chain_mask]  # [N*3 - 2 * num_chain]
        u_co = F.normalize(ori_X[:, 3] - ori_X[:, 2], dim=-1)  # [N, 3], C=O
        u_cca = -U[1::3]  # [N, 3], C-CA
        u_cn = U[2::3]  # [N-1, 3], C-N_{n+1}
        cosA2 = (u_co * u_cca).sum(-1)  # [N], (C=O, C-CA)
        cosA3 = (u_co[:-1] * u_cn).sum(-1)  # [N-1], (C=O, C-N_{n+1})
        same_chain_mask = (seg_id[:-1] == seg_id[1:])  # [N-1]
        cosA3 = cosA3[same_chain_mask]
        cosA = torch.cat([cosA1, cosA2, cosA3], dim=-1)
        cosA = torch.clamp(cosA, -1 + eps, 1 - eps)

        return cosD, cosA

    def coord_loss(self, pred_X, true_X, batch_id, atom_mask, reference=None):
        pred_bb, true_bb = pred_X[:, :4], true_X[:, :4]
        bb_mask = atom_mask[:, :4]
        true_X = true_X.clone()
        ops = []

        align_obj = pred_bb if reference is None else reference[:, :4]

        for i in range(torch.max(batch_id) + 1):
            is_cur_graph = batch_id == i
            cur_bb_mask = bb_mask[is_cur_graph]
            _, R, t = kabsch_torch(
                true_bb[is_cur_graph][cur_bb_mask],
                align_obj[is_cur_graph][cur_bb_mask],
                requires_grad=True)
            true_X[is_cur_graph] = torch.matmul(true_X[is_cur_graph], R.T) + t
            ops.append((R.detach(), t.detach()))

        xloss = F.smooth_l1_loss(
            pred_X[atom_mask], true_X[atom_mask],
            reduction='sum') / atom_mask.sum()  # atom-level loss
        bb_rmsd = torch.sqrt(((pred_X[:, :4] - true_X[:, :4]) ** 2).sum(-1).mean(-1))  # [N]
        return xloss, bb_rmsd, ops

    def structure_loss(self, pred_X, true_X, S, res_idx, batch, full_profile=True):
        seg_id = self.aa_feature._construct_segment_ids(res_idx, batch)

        # loss of backbone (...N-CA-C(O)-N...) bond length
        true_bl = self._cal_backbone_bond_lengths(true_X, seg_id)
        pred_bl = self._cal_backbone_bond_lengths(pred_X, seg_id)
        bond_loss = F.smooth_l1_loss(pred_bl, true_bl)

        # loss of backbone dihedral angles
        if full_profile:
            true_cosD, true_cosA = self._cal_angles(true_X, seg_id)
            pred_cosD, pred_cosA = self._cal_angles(pred_X, seg_id)
            angle_loss = F.smooth_l1_loss(pred_cosD, true_cosD)
            bond_angle_loss = F.smooth_l1_loss(pred_cosA, true_cosA)

        # loss of sidechain bonds
        true_sc_bl = self._cal_sidechain_bond_lengths(S, true_X)
        pred_sc_bl = self._cal_sidechain_bond_lengths(S, pred_X)
        sc_bond_loss = F.smooth_l1_loss(pred_sc_bl, true_sc_bl)

        # loss of sidechain chis
        if full_profile:
            true_sc_chi = self._cal_sidechain_chis(S, true_X)
            pred_sc_chi = self._cal_sidechain_chis(S, pred_X)
            sc_chi_loss = F.smooth_l1_loss(pred_sc_chi, true_sc_chi)

        # exerting constraints on bond lengths only is sufficient
        loss = bond_loss + sc_bond_loss

        if full_profile:
            details = (loss, bond_loss, bond_angle_loss, angle_loss, sc_bond_loss, sc_chi_loss)
        else:
            details = (loss, bond_loss, sc_bond_loss)

        return loss


def sequential_and(*tensors):
    res = tensors[0]
    for mat in tensors[1:]:
        res = torch.logical_and(res, mat)
    return res